diff --git a/.github/scripts/generate_test_matrix.py b/.github/scripts/generate_test_matrix.py index 78d062a2..fa129547 100644 --- a/.github/scripts/generate_test_matrix.py +++ b/.github/scripts/generate_test_matrix.py @@ -172,5 +172,3 @@ def main(): if __name__ == "__main__": main() - -# Made with Bob diff --git a/.github/scripts/generate_weekly_shards.py b/.github/scripts/generate_weekly_shards.py index 8d8f61ec..1d5986e5 100644 --- a/.github/scripts/generate_weekly_shards.py +++ b/.github/scripts/generate_weekly_shards.py @@ -20,13 +20,27 @@ of) much smaller ones. See push-to-clickhouse.yaml's weekly-model-scan job for how `matrix.runner` selects the actual runs-on label. +Each model type's list is PRE-FILTERED before any chunking: models with no +adapter for their config class, models too large for Spyre, and MoE models are +all removed here, and a terminal verdict row for each is written straight to +ClickHouse (so this script needs the CLICKHOUSE_* env vars, or --write-to-csv to +record them in a file instead). + +That ordering is the point. The dropped models cluster heavily — config-class +families cluster by download count — so filtering inside each worker, as +weekly_test.py used to, left surviving counts wildly uneven: some CI jobs +finished in minutes while others ran for hours. Filtering first means a shard's +size is a count of real evaluations, and shard durations become comparable. + Per-tier shard sizes do NOT need to be small: weekly_test.py already re-chunks whatever list it's given into fresh-OS-process batches of GENERATIVE_NUMBER_OF_MODEL_PER_PROCESS/EMBEDDING_NUMBER_OF_MODEL_PER_PROCESS regardless of shard size, which is what actually bounds how many models' memory can accumulate in one process before a clean restart. A tiny shard size buys no extra safety over a large one — it only multiplies GitHub -Actions job count, and matrices are hard-capped at 256 jobs total. +Actions job count, and matrices are hard-capped at 256 jobs total. Since the +sizes now apply to filtered lists, the same values yield fewer, fuller jobs +than they used to. Usage (called by the GHA workflow): python .github/scripts/generate_weekly_shards.py \ @@ -44,16 +58,27 @@ import json import os import sys +from datetime import date from pathlib import Path -# Add the project root to the Python path so we can import from utils/ +# Add the project root to sys.path BEFORE importing from tests/ or utils/ — this +# script lives in .github/scripts/, so neither is importable from its own +# directory, and it is run as a plain script (not `python -m`), which puts that +# directory on sys.path rather than the repo root. Must stay above the imports +# below; the workflow happens to invoke it from the repo root, which would mask +# a wrong order here until someone ran it from anywhere else. project_root = Path(__file__).parent.parent.parent sys.path.insert(0, str(project_root)) -from utils.fetch_top_embedding_models import fetch_top_embedding_models # noqa: E402 -from utils.fetch_top_generative_models import fetch_top_generative_models # noqa: E402 - -MODEL_TYPES = ("generative", "embedding") +from tests.spyre.weekly_generation.failure_categories import ( # noqa: E402 + MAX_NUMBER_PARAMS, +) +from tests.spyre.weekly_generation.model_prefilter import ( # noqa: E402 + fetch_and_filter, +) +from tests.spyre.weekly_generation.model_type import ModelType # noqa: E402 +from tests.spyre.weekly_generation.sink.result_sink import ResultSink # noqa: E402 +from tests.spyre.weekly_generation.sink.sink_factory import create_sink # noqa: E402 def _chunk(rows: list[dict], shard_size: int) -> list[list[dict]]: @@ -81,6 +106,7 @@ def _tier_for(row: dict, x1_max_params: int, x2_max_params: int) -> str: def generate_shards( top_k: int, + max_params: int, shard_size_generative: int, shard_size_embedding: int, x1_max_params: int, @@ -88,63 +114,97 @@ def generate_shards( x2_shard_size: int, x4_shard_size: int, output_dir: Path, - model_types: tuple[str, ...] = MODEL_TYPES, + model_types: list[ModelType], + snapshot_date: date | None = None, + write_to_csv: Path | None = None, ) -> list[dict]: - """Fetch each requested mode's top-K list once, write shard JSON files, + """Fetch each requested model type's top-K list once, write shard JSON files, and return the combined matrix (list of {mode, shard_index, shard_file, runner} dicts). - Within each mode, models are split into three parameter-count tiers (see - module docstring), each chunked at its own shard size and tagged with - the runner ("x1"/"x2"/"x4") that ends up handling it. + Each model type's list is pre-filtered (see ``fetch_and_filter``) before any + chunking, so a shard's size is a count of real evaluations rather than of + fetched candidates. Within each type, the survivors are then split into three + parameter-count tiers (see module docstring), each chunked at its own shard + size and tagged with the runner ("x1"/"x2"/"x4") that handles it. + + *model_types* restricts which ``ModelType`` members to fetch/shard — used by + workflow_dispatch's model_type input so a manual run can scan just embedding + models (much quicker, less resource-hungry) without the schedule-triggered + full scan having to change. + + *max_params* is the ceiling above which a model is rejected outright as too + large for Spyre — distinct from *x1_max_params*/*x2_max_params*, which only + route surviving models between runner tiers. + + *write_to_csv* records the terminal verdicts in a new CSV (one per model type) + instead of ClickHouse, so the whole fetch → filter → route → chunk path can be + exercised without credentials. That sink is write-only, so the emitted shards + then include models a real run would have dropped as recently-scanned. - *model_types* restricts which of MODEL_TYPES to fetch/shard — used by - workflow_dispatch's model_type input so a manual run can scan just - embedding models (much quicker, less resource-hungry) without the - schedule-triggered full scan having to change. + The matrix's ``mode`` key keeps its name because push-to-clickhouse.yaml + reads ``matrix.mode`` and passes it to ``weekly_test.py --mode``. """ + snapshot_date = snapshot_date or date.today() output_dir.mkdir(parents=True, exist_ok=True) - all_fetchers = { - "generative": (fetch_top_generative_models, shard_size_generative), - "embedding": (fetch_top_embedding_models, shard_size_embedding), + # Only the x1 tier's shard size is per-model-type; x2/x4 hold far fewer, + # larger models, so one size each is enough. + x1_shard_sizes = { + ModelType.GENERATIVE: shard_size_generative, + ModelType.EMBEDDING: shard_size_embedding, } - fetchers = {model_type: all_fetchers[model_type] for model_type in model_types} matrix: list[dict] = [] - for mode, (fetch_fn, shard_size) in fetchers.items(): - rows: list[dict] = fetch_fn(limit=top_k) - # model_info is a live huggingface_hub.ModelInfo object attached by - # build_catalog — not JSON-serializable, and no longer needed since - # is_moe is precomputed onto each row (see utils/hf_model_catalog.py). - for row in rows: - row.pop("model_info", None) + for model_type in model_types: + # One sink per model type — each binds a single table (or file) — and + # closed here because this function is what constructed it. Closing is + # what flushes the ClickHouse sink's buffered verdict rows. + sink: ResultSink = create_sink( + model_type=model_type, + write_to_csv=write_to_csv, + ) + with sink: + rows: list[dict] = fetch_and_filter( + model_type=model_type, + snapshot_date=snapshot_date, + top_k=top_k, + sink=sink, + max_params=max_params, + ) by_tier: dict[str, list[dict]] = {"x1": [], "x2": [], "x4": []} for row in rows: by_tier[_tier_for(row, x1_max_params, x2_max_params)].append(row) - tier_shard_sizes = {"x1": shard_size, "x2": x2_shard_size, "x4": x4_shard_size} - mode_shard_count = 0 + tier_shard_sizes = { + "x1": x1_shard_sizes[model_type], + "x2": x2_shard_size, + "x4": x4_shard_size, + } + model_type_shard_count = 0 for runner, group_rows in by_tier.items(): group_shard_size = tier_shard_sizes[runner] shards = _chunk(group_rows, group_shard_size) - mode_shard_count += len(shards) + model_type_shard_count += len(shards) print( - f"{mode} ({runner}): {len(group_rows)} model(s), split into " + f"{model_type} ({runner}): {len(group_rows)} model(s), split into " f"{len(shards)} shard(s) of up to {group_shard_size} each" ) for shard_index, shard_rows in enumerate(shards): - shard_file = f"{mode}-{runner}-shard-{shard_index:03d}.json" + shard_file = f"{model_type}-{runner}-shard-{shard_index:03d}.json" (output_dir / shard_file).write_text(json.dumps(shard_rows)) matrix.append( { - "mode": mode, + "mode": model_type, "shard_index": shard_index, "shard_file": shard_file, "runner": runner, } ) - print(f"{mode}: {len(rows)} model(s) total, {mode_shard_count} shard(s)") + print( + f"{model_type}: {len(rows)} model(s) total, " + f"{model_type_shard_count} shard(s)" + ) return matrix @@ -168,7 +228,16 @@ def main() -> None: "--top-k", type=int, default=10000, - help="Number of top models to fetch per mode (by downloads).", + help="Number of top models to fetch per model type (by downloads).", + ) + parser.add_argument( + "--max-params", + type=int, + default=MAX_NUMBER_PARAMS, + help=( + "Reject models above this parameter count " + f"(default: {MAX_NUMBER_PARAMS:,})." + ), ) parser.add_argument( "--shard-size-generative", @@ -220,20 +289,34 @@ def main() -> None: default=Path("shards"), help="Directory to write shard JSON files into.", ) + parser.add_argument( + "--write-to-csv", + type=Path, + default=None, + metavar="VERDICTS_CSV", + help=( + "Record the terminal verdicts in a new CSV per mode (suffixed " + "-generative / -embedding) instead of ClickHouse, so this script can " + "run without credentials." + ), + ) parser.add_argument( "--model-type", - choices=("all", *MODEL_TYPES), + choices=("all", *(model_type.value for model_type in ModelType)), default="all", - help="Restrict the scan to one mode (e.g. 'embedding' for a quick, " + help="Restrict the scan to one model type (e.g. 'embedding' for a quick, " "low-resource manual run). 'all' (the default, and what the " - "scheduled run always uses) fetches/shards both model-types.", + "scheduled run always uses) fetches/shards every model type.", ) args = parser.parse_args() - model_types = MODEL_TYPES if args.model_type == "all" else (args.model_type,) + model_types = ( + list(ModelType) if args.model_type == "all" else [ModelType(args.model_type)] + ) matrix = generate_shards( top_k=args.top_k, + max_params=args.max_params, shard_size_generative=args.shard_size_generative, shard_size_embedding=args.shard_size_embedding, x1_max_params=args.x1_max_params, @@ -242,9 +325,10 @@ def main() -> None: x4_shard_size=args.x4_shard_size, output_dir=args.output_dir, model_types=model_types, + write_to_csv=args.write_to_csv, ) - print(f"\nTotal shards across both model-types: {len(matrix)}") + print(f"\nTotal shards across {len(model_types)} model type(s): {len(matrix)}") # Split by runner tier so push-to-clickhouse.yaml's three per-tier jobs # can each cap strategy.max-parallel in cards (x1=1, x2=2, x4=4/shard). diff --git a/.github/workflows/push-to-clickhouse.yaml b/.github/workflows/push-to-clickhouse.yaml index 1b38aefe..502d29fe 100644 --- a/.github/workflows/push-to-clickhouse.yaml +++ b/.github/workflows/push-to-clickhouse.yaml @@ -7,17 +7,29 @@ # the manual dev-instance upload + periodic dev->prod migration. # # Unlike torch-spyre's push-to-clickhouse.yaml, this is not a post-hoc ingest -# of a separate test run's artifact: weekly_test.py's skip-window dedup guard -# (don't re-test a model verified in the last 10 days) needs to read -# ClickHouse live during the scan, so fetch + test + push both happen in the -# weekly-model-scan job below rather than being split across a test workflow -# and a separate ingest one. +# of a separate test run's artifact: the scan reads and writes ClickHouse +# directly. # # At top_k=10000 per mode, a single job processing the whole list serially # would take far too long, so the fetch itself IS split out: generate-matrix # fetches the top-K list once (no Spyre hardware needed), shards it, and # weekly-model-scan fans out over those shards in parallel (mirrors # torch-spyre's generate_matrix -> matrix job pattern). +# +# generate-matrix also does all the FILTERING, which is why it needs ClickHouse +# credentials of its own: it applies the no-adapter / too-large / MoE checks and +# writes those terminal verdict rows itself. weekly_test.py used to do this +# inside each shard, but the dropped models cluster — config-class families +# cluster by download count — so surviving counts per shard varied wildly and +# some jobs finished in minutes while others ran for hours. Filtering before +# chunking makes shard size a count of real evaluations. +# +# Every model that survives the filter is evaluated: there is no "recently +# scanned" check, so a re-run repeats the full scan rather than resuming a +# partial one. If a run is cancelled after generate-matrix succeeds +# (cancel-in-progress below makes that routine), the terminal verdict rows it +# already wrote stay — they are terminal properties of the checkpoint, so a +# re-run simply overwrites them with the same values. # ============================================================================= name: push-to-clickhouse @@ -69,12 +81,16 @@ env: jobs: # --------------------------------------------------------------------------- - # Fetch the top-K list once per mode and shard it for the parallel matrix - # job below. No Spyre hardware needed — utils/hf_model_catalog.py's import - # chain (via hf_adapters/auto_spyre_model.py) has no module-level - # torch_spyre dependency; it's lazily imported inside a device-guarded - # function body in hf_common.py. The now-deprecated test_weekly_DEPRECATED - # .yaml is existing proof this exact fetch path runs fine on ubuntu-latest. + # Fetch the top-K list once per mode, filter it, and shard the survivors for + # the parallel matrix job below. No Spyre hardware needed — + # utils/hf_model_catalog.py's import chain (via hf_adapters/auto_spyre_model + # .py) has no module-level torch_spyre dependency; it's lazily imported inside + # a device-guarded function body in hf_common.py. The now-deprecated + # test_weekly_DEPRECATED.yaml is existing proof this exact fetch path runs fine + # on ubuntu-latest. + # + # This job DOES need ClickHouse: it writes the terminal verdict rows for + # everything it filters out. See the header comment. # --------------------------------------------------------------------------- generate-matrix: name: Generate weekly shards @@ -95,6 +111,13 @@ jobs: matrix_x4: ${{ steps.generate.outputs.matrix_x4 }} env: HF_TOKEN: ${{ secrets.HF_TOKEN }} + # Needed because this job writes the terminal verdict rows for the models + # it filters out. + CLICKHOUSE_HOST: ${{ secrets.CLICKHOUSE_HOST }} + CLICKHOUSE_PORT: ${{ secrets.CLICKHOUSE_PORT }} + CLICKHOUSE_USER: ${{ secrets.CLICKHOUSE_USER }} + CLICKHOUSE_PASS: ${{ secrets.CLICKHOUSE_PASS }} + CLICKHOUSE_DB: ${{ secrets.CLICKHOUSE_DB }} # Overrides the workflow-level HF_HOME (/storage-1/... — a Spyre-pod-only # mount). Every AutoConfig.from_pretrained() call was failing with a # PermissionError trying to write there, which contains_remote_code() @@ -115,7 +138,12 @@ jobs: set -euo pipefail python -m pip install --upgrade pip pip install -e . - pip install huggingface_hub tqdm transformers + # clickhouse-connect and python-dotenv are named explicitly because + # `pip install -e .` does NOT pull them: they live in the `dev` and + # `models-ops` dependency groups respectively (pyproject.toml), and + # dependency groups are not installed by a plain editable install. + # Omitting them fails this job at import time on the sink. + pip install huggingface_hub tqdm transformers clickhouse-connect python-dotenv # TEMPORARY debug step: confirms HF_TOKEN actually reaches this job as a # working credential, and re-runs the real keep()-filter checks against @@ -157,6 +185,38 @@ jobs: print(f"AutoConfig.from_pretrained() FAILED: {type(e).__name__}: {e}") PY + # Fail fast if ClickHouse secrets are missing/wrong: this job now writes the + # terminal verdict rows, and the fetch below can take 90+ minutes before it + # would otherwise reach the sink. Never prints + # secret values, only presence/length and a live SELECT 1. Same check the + # scan jobs run, minus their `cd`/`uv run` wrapper — this job uses plain + # python with working-directory: '.'. + - name: Verify ClickHouse connection + run: | + for VAR in CLICKHOUSE_HOST CLICKHOUSE_PORT CLICKHOUSE_USER CLICKHOUSE_PASS CLICKHOUSE_DB; do + VAL="${!VAR}" + if [[ -n "$VAL" ]]; then + echo "$VAR: set (${#VAL} chars)" + else + echo "$VAR: NOT SET" + fi + done + python - <<'PY' + import os + import clickhouse_connect + + client = clickhouse_connect.get_client( + host=os.environ["CLICKHOUSE_HOST"], + port=int(os.environ.get("CLICKHOUSE_PORT", 443)), + user=os.environ.get("CLICKHOUSE_USER", "default"), + password=os.environ["CLICKHOUSE_PASS"], + database=os.environ.get("CLICKHOUSE_DB", "spyre"), + secure=True, + ) + client.command("SELECT 1") + print("ClickHouse connectivity check: OK") + PY + - name: Generate shards id: generate run: | @@ -182,9 +242,15 @@ jobs: # Split by runner tier so max-parallel caps total cards held at once: # 10*1 + 4*2 + 1*4 = 22 cards (~2 nodes), so PR/daily CI isn't starved. + # + # The `if` guards an empty matrix: now that generate-matrix filters before + # chunking, a tier can legitimately produce zero shards — no fetched model + # falls in its parameter range, or a --model-type run skipped it entirely. + # This showed up in testing as matrix_x4=[]. weekly-model-scan-x1: name: Weekly model scan x1 (${{ matrix.mode }} shard ${{ matrix.shard_index }}) needs: generate-matrix + if: needs.generate-matrix.outputs.matrix_x1 != '[]' runs-on: [x86_64, spyre_pf_x1, linux, image_torch_spyre] timeout-minutes: 4320 strategy: @@ -262,6 +328,7 @@ jobs: weekly-model-scan-x2: name: Weekly model scan x2 (${{ matrix.mode }} shard ${{ matrix.shard_index }}) needs: generate-matrix + if: needs.generate-matrix.outputs.matrix_x2 != '[]' runs-on: [x86_64, spyre_pf_x2, linux, image_torch_spyre] timeout-minutes: 4320 strategy: @@ -339,6 +406,7 @@ jobs: weekly-model-scan-x4: name: Weekly model scan x4 (${{ matrix.mode }} shard ${{ matrix.shard_index }}) needs: generate-matrix + if: needs.generate-matrix.outputs.matrix_x4 != '[]' runs-on: [x86_64, spyre_pf_x4, linux, image_torch_spyre] timeout-minutes: 4320 strategy: diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/spyre/__init__.py b/tests/spyre/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/spyre/weekly_generation/add_past_rows.py b/tests/spyre/weekly_generation/add_past_rows.py index b548249f..ff807068 100644 --- a/tests/spyre/weekly_generation/add_past_rows.py +++ b/tests/spyre/weekly_generation/add_past_rows.py @@ -12,7 +12,7 @@ Arguments: * ``csv_file`` Path to a CSV file whose columns match ``TABLE_COLUMNS`` - (see ``clickhouse_db.py``). + (see ``table_schema.py``). * ``mode`` Either ``embedding`` or ``generative``. """ diff --git a/tests/spyre/weekly_generation/clickhouse_db.py b/tests/spyre/weekly_generation/clickhouse_db.py index 05ba41c1..94ec540a 100644 --- a/tests/spyre/weekly_generation/clickhouse_db.py +++ b/tests/spyre/weekly_generation/clickhouse_db.py @@ -1,40 +1,47 @@ -#!/usr/bin/env python3 -""" -Creates or drops the model_spyre_support table in ClickHouse. +"""ClickHouse client factory for the model_spyre_support tables. -Columns: - - model_name (String) – unique model identifier - - config_class (String) – model config_class (e.g. BertConfig, Qwen3Config) - - adapter_name (String) – adapter name (e.g. hf_bert, hf_gemma4) - - added_date (Date?) – date the adapter was added to the git repo (optional - None if not existing) - - snapshot_date (Date) – date this weekly snapshot was taken - - verified_on_cpu (Bool) – passes on CPU - - verified_on_gpu (Bool) – passes on GPU - - verified_on_spyre (Bool) – passes on Spyre - - num_downloads (UInt64) – number of downloads - - family (String) – model family reported by the catalog - - architecture (String) – model architecture reported by the catalog - - parameters_number (UInt64) – number of model parameters - - failure_category (String?) – classification for failures (optional - None if model passed) - - error (String?) – error message if the model failed (optional - None if model passed) +The table shape itself lives in ``table_schema`` — a leaf module with no +dependencies — and is re-exported here so existing callers keep working. Import +from ``table_schema`` directly when you only need the column list or DDL: +importing *this* module pulls in ``clickhouse_connect`` and reads ``.env`` off +disk, neither of which a consumer of a column tuple should have to pay for. -Credentials are loaded from a .env file at the repo root (two levels above this -script), then fall back to environment variables already set in the shell. -Copy .env.example → .env and fill in the values before running. +Credentials are loaded from a .env file at the repo root, then fall back to +environment variables already set in the shell. Copy .env.example → .env and fill +in the values before running. CLICKHOUSE_HOST, CLICKHOUSE_PORT, CLICKHOUSE_USER, CLICKHOUSE_PASS, CLICKHOUSE_DB - """ import os -from datetime import date from pathlib import Path import clickhouse_connect from dotenv import load_dotenv -# Locate the repo root (.env lives two directories above this script: -# repo_root/.github/scripts/clickhouse_db.py) +# Re-exported for backwards compatibility; table_schema is the source of truth. +from tests.spyre.weekly_generation.table_schema import ( + DATABASE, + EMBEDDING_CREATE_TABLE_SQL, + EMBEDDING_TABLE_NAME, + GENERATIVE_CREATE_TABLE_SQL, + GENERATIVE_TABLE_NAME, + TABLE_COLUMNS, +) + +__all__ = [ + "DATABASE", + "EMBEDDING_CREATE_TABLE_SQL", + "EMBEDDING_TABLE_NAME", + "GENERATIVE_CREATE_TABLE_SQL", + "GENERATIVE_TABLE_NAME", + "TABLE_COLUMNS", + "get_client", + "table_exists", +] + +# Locate the repo root (.env lives three directories above this module: +# repo_root/tests/spyre/weekly_generation/clickhouse_db.py) _REPO_ROOT = Path(__file__).resolve().parent.parent.parent.parent load_dotenv(_REPO_ROOT / ".env") @@ -50,61 +57,6 @@ def get_client(): ) -EMBEDDING_TABLE_NAME = "embedding_model_spyre_support" -GENERATIVE_TABLE_NAME = "generative_model_spyre_support" -DATABASE = "spyre" - -# Single source of truth for the Python-facing column list. Order matches the -# CREATE TABLE DDL below and the positional values used by insert_model_row — -# keep the three in sync when adding columns. -TABLE_COLUMNS: tuple[str, ...] = ( - "model_name", - "config_class", - "adapter_name", - "added_date", - "snapshot_date", - "verified_on_cpu", - "verified_on_gpu", - "verified_on_spyre", - "num_downloads", - "family", - "architecture", - "parameters_number", - "failure_category", - "error", -) - - -def _make_create_table_sql(table_name: str) -> str: - return f""" -CREATE TABLE IF NOT EXISTS {DATABASE}.{table_name} -( - model_name String, - config_class String, - adapter_name String, - added_date Nullable(Date), - snapshot_date Date, - verified_on_cpu Bool, - verified_on_gpu Bool, - verified_on_spyre Bool, - num_downloads UInt64, - family String, - architecture String, - parameters_number UInt64, - failure_category Nullable(String), - error Nullable(String) -) -ENGINE = ReplacingMergeTree(snapshot_date) -ORDER BY (model_name, snapshot_date) -""" - - -# The column names and order below MUST match ``TABLE_COLUMNS`` above. -# When adding/removing/renaming a column, update both in the same change. -EMBEDDING_CREATE_TABLE_SQL = _make_create_table_sql(EMBEDDING_TABLE_NAME) -GENERATIVE_CREATE_TABLE_SQL = _make_create_table_sql(GENERATIVE_TABLE_NAME) - - def table_exists(client, table_name: str) -> bool: result = client.query( "SELECT count() FROM system.tables " @@ -112,203 +64,3 @@ def table_exists(client, table_name: str) -> bool: parameters={"db": DATABASE, "tbl": table_name}, ) return result.result_rows[0][0] > 0 - - -def print_table(client, table_name: str) -> None: - result = client.query( - "SELECT name, type FROM system.columns " - "WHERE database = {db:String} AND table = {tbl:String} " - "ORDER BY position", - parameters={"db": DATABASE, "tbl": table_name}, - ) - print(f"Table '{DATABASE}.{table_name}' already exists with columns:") - for col_name, col_type in result.result_rows: - print(f" {col_name:<25} {col_type}") - - -def insert_model_row( - client, - *, - table_name: str, - model_name: str, - config_class: str, - adapter_name: str, - added_date: date | None, - snapshot_date: date, - verified_on_cpu: bool, - verified_on_gpu: bool, - verified_on_spyre: bool, - num_downloads: int, - family: str, - architecture: str, - parameters_number: int, - failure_category: str | None, - error: str | None, -) -> bool: - """Insert a single row into the given table. - - The caller is responsible for any duplicate-suppression guard (e.g. - ``ResultSink.should_insert_row``). This function always writes. - """ - client.insert( - table_name, - [ - [ - model_name, - config_class, - adapter_name, - added_date, - snapshot_date, - verified_on_cpu, - verified_on_gpu, - verified_on_spyre, - num_downloads, - family, - architecture, - parameters_number, - failure_category, - error, - ] - ], - column_names=list(TABLE_COLUMNS), - ) - return True - - -def _parse_bool(value: str) -> bool: - return value.strip().lower() in ("1", "true", "yes") - - -def _parse_nullable_date(value: str | None) -> date | None: - v = (value or "").strip() - return date.fromisoformat(v) if v else None - - -def _parse_nullable_str(value: str | None) -> str | None: - v = (value or "").strip() - return v if v else None - - -def import_csv(sink, csv_path: str) -> tuple[int, int]: - """Read *csv_path* and insert rows into *sink*, respecting its dedup guard. - - Uses ``sink.add_entry()`` so ``should_insert_row`` is applied for every row. - Returns a ``(inserted, skipped)`` tuple. - """ - import csv - - inserted = skipped = malformed = 0 - with open(csv_path, newline="", encoding="utf-8") as fh: - reader = csv.DictReader(fh) - for row in reader: - snapshot_raw: str = (row.get("snapshot_date") or "").strip() - model_name: str = (row.get("model_name") or "").strip() - if not snapshot_raw or not model_name: - malformed += 1 - print( - f" warn: skipping CSV row {reader.line_num} " - f"(model_name={model_name!r}, snapshot_date={snapshot_raw!r}): " - "missing required field" - ) - continue - try: - snapshot_date_val: date = date.fromisoformat(snapshot_raw) - except ValueError: - malformed += 1 - print( - f" warn: skipping CSV row {reader.line_num} " - f"(model_name={model_name!r}): " - f"invalid snapshot_date={snapshot_raw!r}" - ) - continue - written = sink.add_entry( - model_name=model_name, - config_class=(row.get("config_class") or "").strip(), - adapter_name=(row.get("adapter_name") or "").strip(), - added_date=_parse_nullable_date(row.get("added_date")), - snapshot_date=snapshot_date_val, - verified_on_cpu=_parse_bool(row.get("verified_on_cpu") or ""), - verified_on_gpu=_parse_bool(row.get("verified_on_gpu") or ""), - verified_on_spyre=_parse_bool(row.get("verified_on_spyre") or ""), - num_downloads=int((row.get("num_downloads") or "0").strip() or "0"), - family=(row.get("family") or "").strip(), - architecture=(row.get("architecture") or "").strip(), - parameters_number=int( - (row.get("parameters_number") or "0").strip() or "0" - ), - failure_category=_parse_nullable_str(row.get("failure_category")), - error=_parse_nullable_str(row.get("error")), - ) - if written: - inserted += 1 - else: - skipped += 1 - - if malformed: - print(f" import_csv: {malformed} malformed row(s) skipped.") - return inserted, skipped - - -if __name__ == "__main__": - import argparse - - parser = argparse.ArgumentParser(description="ClickHouse table management utility.") - parser.add_argument( - "--drop", - metavar="TABLE_NAME", - help="Drop the specified table after confirmation.", - ) - - add_csv_group = parser.add_argument_group("import CSV") - add_csv_group.add_argument( - "--add_csv", metavar="CSV_FILE", help="CSV file to import into the table." - ) - add_csv_group.add_argument( - "--table_name", metavar="TABLE_NAME", help="Target table for --add_csv." - ) - - args = parser.parse_args() - - if args.add_csv or args.table_name: - if not args.add_csv or not args.table_name: - parser.error("--add_csv and --table_name must be used together.") - csv_file = args.add_csv - table = args.table_name - # Lazy import to avoid a circular dependency (result_sink imports from this module). - from tests.spyre.weekly_generation.result_sink import ( - ClickHouseResultSink, - EmbeddingGenerativeMode, - ) - - if table == EMBEDDING_TABLE_NAME: - mode = EmbeddingGenerativeMode.EMBEDDING - elif table == GENERATIVE_TABLE_NAME: - mode = EmbeddingGenerativeMode.GENERATIVE - else: - parser.error( - f"Unknown table '{table}'. Expected one of: " - f"{EMBEDDING_TABLE_NAME}, {GENERATIVE_TABLE_NAME}." - ) - with ClickHouseResultSink(mode) as sink: - inserted, skipped = import_csv(sink, csv_file) - print( - f"Inserted {inserted} row(s) into '{DATABASE}.{table}' ({skipped} skipped by dedup guard)." - ) - elif args.drop: - table = args.drop - answer = ( - input(f"Are you sure you want to drop table '{DATABASE}.{table}'? [y/N] ") - .strip() - .lower() - ) - if answer == "y": - client = get_client() - if not table_exists(client, table): - print(f"Table '{DATABASE}.{table}' does not exist.") - else: - client.command(f"DROP TABLE {DATABASE}.{table}") - print(f"Table '{DATABASE}.{table}' dropped.") - else: - print("Aborted.") - else: - parser.print_help() diff --git a/tests/spyre/weekly_generation/failure_categories.py b/tests/spyre/weekly_generation/failure_categories.py new file mode 100644 index 00000000..474a75c7 --- /dev/null +++ b/tests/spyre/weekly_generation/failure_categories.py @@ -0,0 +1,31 @@ +"""Failure categories and size limits shared across the weekly-scan pipeline. + +Deliberately a leaf module: it imports nothing from this repo and nothing that +pulls in a database driver. That is what lets ``model_prefilter`` — and its unit +tests — be imported on a machine with no ``clickhouse_connect`` installed, while +still sharing one definition of each category string with ``weekly_test``, +``result_sink`` and the CI shard producer (``generate_weekly_shards``). + +The category strings are persisted verbatim in the ClickHouse +``failure_category`` column, so changing a value silently splits historical rows +from new ones. Add categories rather than rename them. +""" + +# Models above this parameter count are rejected before a worker is spawned: +# they cannot be brought up on Spyre. Producers apply this while building the +# model list; weekly_test keeps an in-worker backstop for rows whose parameter +# count was unknown at fetch time. +MAX_NUMBER_PARAMS = 60_000_000_000 + +FAILURE_CATEGORY_NOT_IMPLEMENTED_ADAPTER = "not-implemented-adapter" +FAILURE_CATEGORY_MODEL_TOO_LARGE = "model_too_large" +FAILURE_CATEGORY_CPU_LOAD_FAILED = "cpu_load_failed" +FAILURE_CATEGORY_CPU_GENERATE_FAILED = "cpu_generate_failed" +FAILURE_CATEGORY_QUANTIZED_MODEL = "quantized_model" +FAILURE_CATEGORY_HARDWARE_EXCEPTION = "hardware_exception" +FAILURE_CATEGORY_MISFORMED_HF_FAILED = "misformed_hf_failed" +FAILURE_CATEGORY_TEST_EXECUTION_EXCEPTION = "test_execution_exception" +FAILURE_CATEGORY_VERIFICATION_FAILED = "verification_failed" +FAILURE_CATEGORY_WORKER_CRASHED = "worker_crashed" +FAILURE_CATEGORY_WORKER_TIMEOUT = "worker_timeout" +FAILURE_CATEGORY_MOE = "moe" diff --git a/tests/spyre/weekly_generation/model_fetcher.py b/tests/spyre/weekly_generation/model_fetcher.py new file mode 100644 index 00000000..ff2e9e17 --- /dev/null +++ b/tests/spyre/weekly_generation/model_fetcher.py @@ -0,0 +1,49 @@ +"""Fetch a model type's top-K catalog from the HuggingFace Hub. + +One indirection over the two ``utils.fetch_top_*_models`` functions, so callers +select a catalog with a ``ModelType`` instead of branching on a string. Both +producers (``generate_weekly_shards`` and ``weekly_test --fetch``) reach the Hub +through here, which is also the single place the rows get normalized for the +JSON round-trip they are about to take. +""" + +from collections.abc import Callable + +from tests.spyre.weekly_generation.model_type import ModelType +from utils.fetch_top_embedding_models import fetch_top_embedding_models +from utils.fetch_top_generative_models import fetch_top_generative_models +from utils.utilities import ts + +all_fetchers: dict[ModelType, Callable[..., list[dict]]] = { + ModelType.GENERATIVE: fetch_top_generative_models, + ModelType.EMBEDDING: fetch_top_embedding_models, +} +"""Catalog fetcher per model type. Module-level and mutable so tests can +monkeypatch it and exercise the pipeline without hitting the network.""" + + +def fetch(model_type: ModelType, top_k: int) -> list[dict]: + """Return the top *top_k* models of *model_type*, ordered by downloads. + + Descending download order is a contract, not an incidental: the tier router + and shard chunker downstream both assume it, and the pre-filter preserves it. + + Each row is a plain dict keyed as the catalog CSV header is (``model_id``, + ``downloads``, ``parameters``, ``is_supported``, ``is_moe``, + ``config_class``, ``model_type``, ``architectures``) — JSON-serializable, so + it survives being written to a shard file and read back by another process. + """ + models: list[dict] = all_fetchers[model_type](limit=top_k) + + # model_info is a live huggingface_hub.ModelInfo object attached by + # build_catalog — not JSON-serializable, and no longer needed since + # is_moe is precomputed onto each model (see utils/hf_model_catalog.py). + # Dropped here rather than at each call site because every consumer either + # JSON-dumps these rows into a shard file or hands them to a spawned child + # (which pickles them); both break on a ModelInfo. + for model in models: + model.pop("model_info", None) + + print(f"{ts()} Fetched {len(models)} {model_type} model(s).") + + return models diff --git a/tests/spyre/weekly_generation/model_prefilter.py b/tests/spyre/weekly_generation/model_prefilter.py new file mode 100644 index 00000000..664a0783 --- /dev/null +++ b/tests/spyre/weekly_generation/model_prefilter.py @@ -0,0 +1,202 @@ +"""Decide which fetched models are worth handing to a Spyre worker. + +Three checks, applied in the order ``weekly_test.main`` used to apply them +in-process. All three are terminal properties of the checkpoint itself +(no adapter, too large, mixture-of-experts) and produce a row recording that +verdict, so a dropped model is never silently absent from a run's output. + +Running this **upstream of sharding** is the point. ``generate_weekly_shards`` +chunks a downloads-ordered list into fixed-size shards, and filtered-out models +cluster heavily — config-class families cluster by download count, so a single +unsupported family can hollow out one shard while leaving the next untouched. +Filtering after chunking left shards with wildly different amounts of real work: +some CI jobs finished in minutes, others ran for hours. Filtering first means +shard size maps to evaluations. + +``weekly_test --fetch`` calls this too, for manual runs with no shard file, so +both entry points share one definition of "worth handing to a Spyre worker". +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from datetime import date + +from tests.spyre.weekly_generation import model_fetcher +from tests.spyre.weekly_generation.failure_categories import ( + FAILURE_CATEGORY_MODEL_TOO_LARGE, + FAILURE_CATEGORY_MOE, + FAILURE_CATEGORY_NOT_IMPLEMENTED_ADAPTER, +) +from tests.spyre.weekly_generation.model_type import ModelType +from tests.spyre.weekly_generation.sink.result_sink import ResultSink + + +@dataclass(frozen=True) +class SkippedModel: + """A model rejected for a terminal reason, with the row to record for it.""" + + row: dict + failure_category: str + reason: str + """Human-readable detail for the log line; never persisted.""" + + +@dataclass +class PrefilterResult: + """Two-way partition of the fetched models: work to do, and verdicts to record. + + Every fetched model lands in exactly one of the two lists, which is what lets + ``counts`` reconcile against the input length. + """ + + keep: list[dict] = field(default_factory=list) + skipped: list[SkippedModel] = field(default_factory=list) + + @property + def counts(self) -> dict[str, int]: + """Per-category tallies, for one-line run summaries.""" + tally: dict[str, int] = { + "keep": len(self.keep), + } + for item in self.skipped: + tally[item.failure_category] = tally.get(item.failure_category, 0) + 1 + return tally + + +def _parameter_count(row: dict) -> int | None: + """Parameter count for *row*, or None when the fetcher could not size it. + + Mirrors ``weekly_test``'s original guard, which tested + ``params not in (None, "")`` before coercing, so an unsizable row is neither + treated as zero-parameter nor assumed oversized — it goes to a worker and is + judged by whether it actually loads. + + Note there is no worker-side size check to fall back on. A comment in + ``weekly_test`` used to claim ``_process_batch`` kept one "as a defensive + backstop for rows where parameters were unknown at fetch time"; no such check + existed. An oversized model the fetcher could not size therefore surfaces as + cpu_load_failed or a worker timeout rather than model_too_large — which is + what happened before this filter moved upstream, too. + """ + params = row.get("parameters") + if params in (None, ""): + return None + try: + return int(params) + except (TypeError, ValueError): + return None + + +def prefilter_models( + models: list[dict], + max_params: int, +) -> PrefilterResult: + """Partition *models* into work to do and verdicts to record. + + Pure: decides, and writes nothing. Recording the terminal verdicts is + ``write_skipped_rows``' job, so this needs no sink and its tests need no + storage backend. + + Args: + models: one dict per model as fetched from the HuggingFace Hub by + ``build_catalog``, keyed as the catalog CSV header is (``model_id``, + ``downloads``, ``parameters``, ``is_supported``, ``is_moe``, + ``config_class``, ``model_type``, ``architectures``). Callers should + ``pop("model_info")`` first — the returned lists hold the same dict + objects, and that field is not JSON-serializable. + max_params: parameter ceiling above which a model cannot be brought up + on Spyre. + + Returns: + A ``PrefilterResult``. ``keep`` preserves the input order, which the + tier router and shard chunker both depend on being downloads-descending. + """ + result = PrefilterResult() + + for row in models: + # No adapter registered for this config class — the same terminal + # decision resolve_adapter_module_for_test would reach in the worker, + # reached without spawning one. `is False` rather than falsy: a missing + # key or None means the fetcher could not determine the config class, + # which is not the same as knowing it is unsupported. + if row.get("is_supported") is False: + result.skipped.append( + SkippedModel( + row=row, + failure_category=FAILURE_CATEGORY_NOT_IMPLEMENTED_ADAPTER, + reason=f"no adapter for config_class={row.get('config_class')!r}", + ) + ) + continue + + params = _parameter_count(row) + if params is not None and params > max_params: + result.skipped.append( + SkippedModel( + row=row, + failure_category=FAILURE_CATEGORY_MODEL_TOO_LARGE, + reason=(f"{params:,} parameters exceeds the {max_params:,} limit"), + ) + ) + continue + + # MoE models aren't supported on Spyre yet. is_moe is precomputed at + # fetch time (utils/hf_model_catalog.py) so it survives the JSON + # round-trip through the model-list file. + if row.get("is_moe"): + result.skipped.append( + SkippedModel( + row=row, + failure_category=FAILURE_CATEGORY_MOE, + reason="MoE model", + ) + ) + continue + + result.keep.append(row) + + return result + + +def fetch_and_filter( + model_type: ModelType, + snapshot_date: date, + top_k: int, + sink: ResultSink, + max_params: int, +) -> list[dict]: + """Fetch *model_type*'s top-*top_k* catalog, filter it, record the verdicts. + + The one entry point both producers share, which is what keeps "worth handing + to a Spyre worker" a single definition: ``generate_weekly_shards`` calls it + before chunking (see the module docstring for why order matters there), and + ``weekly_test --fetch`` calls it instead of reading a shard file. + + Terminal verdicts are written to *sink* as they are decided, so a model + dropped here still gets its row and is not silently absent from the run's + output. + + Does NOT close *sink*: the caller constructed it and keeps writing + evaluation results to it afterwards. ``weekly_test.main`` in particular + hands in the same sink it uses for the rest of the run, and closing it here + left that path writing to a closed file. + + Returns: + The models to evaluate, in the fetched (downloads-descending) order that + the tier router and shard chunker both rely on. + """ + # Deferred so that importing this module — and running prefilter_models, + # which is pure — needs neither skip_writer nor anything it pulls in. + from tests.spyre.weekly_generation.skip_writer import write_skipped_rows + + models: list[dict] = model_fetcher.fetch(model_type=model_type, top_k=top_k) + + result: PrefilterResult = prefilter_models(models, max_params=max_params) + written: int = write_skipped_rows(sink, result.skipped, snapshot_date=snapshot_date) + + print( + f"{model_type}: {len(models)} fetched -> {len(result.keep)} to evaluate " + f"({written} terminal row(s) written for skipped models) {result.counts}" + ) + return result.keep diff --git a/tests/spyre/weekly_generation/model_type.py b/tests/spyre/weekly_generation/model_type.py new file mode 100644 index 00000000..9e3b6fd9 --- /dev/null +++ b/tests/spyre/weekly_generation/model_type.py @@ -0,0 +1,23 @@ +"""The two kinds of model the weekly scan evaluates. + +``generative`` and ``embedding`` each name a whole bundle of per-kind choices +that used to travel as a bare string: which catalog is fetched, which ClickHouse +table the sink binds, which model class ``_load_on_cpu`` instantiates, which +verification pipeline ``eval_model`` runs, and how many models one worker +process handles. + +``StrEnum``, not ``(str, Enum)``. These members reach shard filenames, the GHA +matrix, and the log lines an operator reads, so a member must *format* as its +value. ``class ModelType(str, Enum)`` inherits ``Enum.__str__``/``__format__``, +making ``f"{ModelType.GENERATIVE}"`` render ``"ModelType.GENERATIVE"`` — which +is how shard files briefly came out named +``ModelType.GENERATIVE-x1-shard-000.json``. ``StrEnum`` (3.11+, and this project +requires >=3.11) is the variant whose ``__str__`` is ``str.__str__``. +""" + +from enum import StrEnum + + +class ModelType(StrEnum): + GENERATIVE = "generative" + EMBEDDING = "embedding" diff --git a/tests/spyre/weekly_generation/result_sink.py b/tests/spyre/weekly_generation/result_sink.py deleted file mode 100644 index fd82f658..00000000 --- a/tests/spyre/weekly_generation/result_sink.py +++ /dev/null @@ -1,499 +0,0 @@ -"""Abstract result sink for the weekly Spyre test suite. - -Two implementations: -- ``CsvResultSink`` — appends rows to a CSV file; loads existing rows once so - the skip guard is O(1) per call. -- ``ClickHouseResultSink`` — inserts rows into ClickHouse; the skip guard runs - a single bulk SELECT on construction so all per-row checks are O(1), and rows - are buffered in memory then flushed in a single bulk INSERT on ``close()``. - -Skip rule (both sinks): insert a row for *model_name* when either - - * no prior row for *model_name* exists, OR - * the most-recent prior row has ``failure_category == 'hardware_exception'`` - (accelerator problem, worth retrying now) OR the snapshot is older than - ``_SKIP_WINDOW_DAYS`` days (long enough since last run). - -Rows with `hardware_exception` are always re-run because the accelerator's -availability is a transient property — a failure yesterday says nothing about -today. Everything else (verified success, non-implemented adapter, quantized -model, moe, cpu load/generate failure, …) is treated as terminal for the -window. -""" - -from __future__ import annotations - -import csv -from abc import ABC, abstractmethod -from datetime import date, datetime, timedelta -from enum import Enum -from pathlib import Path -from typing import Any - -from tests.spyre.weekly_generation.clickhouse_db import ( - DATABASE, - EMBEDDING_CREATE_TABLE_SQL, - EMBEDDING_TABLE_NAME, - GENERATIVE_CREATE_TABLE_SQL, - GENERATIVE_TABLE_NAME, - TABLE_COLUMNS, - get_client, - table_exists, -) - -# Constant value -_SKIP_WINDOW_DAYS: int = 10 - -# Duplicated (as a string literal) from weekly_test.FAILURE_CATEGORY_HARDWARE_EXCEPTION -# to avoid a circular import: weekly_test already imports EmbeddingGenerativeMode -# from this module. Keep the string values in sync. -_HARDWARE_EXCEPTION_CATEGORY: str = "hardware_exception" - - -def _within_skip_window(existing_snapshot: date, today: date) -> bool: - return (today - existing_snapshot).days < _SKIP_WINDOW_DAYS - - -_SNAPSHOT_DATE_FORMATS = ( - "%Y-%m-%d", # ISO 8601 — primary format written by this module - "%d/%m/%Y", # DD/MM/YYYY - "%m/%d/%Y", # MM/DD/YYYY - "%Y/%m/%d", # YYYY/MM/DD -) - - -def _coerce_snapshot(value: object) -> date | None: - if isinstance(value, date): - return value - if isinstance(value, str) and value: - for fmt in _SNAPSHOT_DATE_FORMATS: - try: - return datetime.strptime(value.strip(), fmt).date() - except ValueError: - continue - return None - - -def _require_non_empty(value: str, field_name: str) -> str: - stripped: str = value.strip() - if not stripped: - raise ValueError(f"{field_name} must be a non-empty string") - return stripped - - -class ResultSink(ABC): - """Abstract destination for weekly-test result rows. - - Implementations must be usable as a context manager; ``__exit__`` should - release any external resources (CSV file handle, DB client). - """ - - _today: date - - def __init__(self, today: date | None = None) -> None: - """Store the reference *today* used by the skip-window guard. - - Subclasses must call ``super().__init__(today=today)`` before touching - anything that depends on ``self._today``. - """ - self._today = today or date.today() - - @abstractmethod - def get_recent_blocking_entries(self, model_name: str) -> list[dict[str, Any]]: - """Return prior rows for *model_name* that block a new insert. - - A row blocks re-insertion when BOTH: - - * its ``snapshot_date`` is within the skip window - (``today - snapshot_date < _SKIP_WINDOW_DAYS``), AND - * its ``failure_category`` is NOT ``hardware_exception`` — hardware - failures are treated as transient and always re-run. - - Sorted by ``snapshot_date`` descending. Each row is a dict keyed by - column name. Empty list when no blocking prior entry exists — the - caller can treat empty as "insert away" without inspecting the rows. - """ - - @abstractmethod - def get_all_models(self) -> list[dict[str, Any]]: - """Return one row per known ``model_name``, reflecting its most recent snapshot. - - When a model appears in multiple rows (one per weekly run), only the row - with the greatest ``snapshot_date`` is returned. The result is a flat - list of dicts keyed by column name (same keys as ``TABLE_COLUMNS``), one - dict per distinct model, in no guaranteed order. - """ - - @abstractmethod - def _insert_entry( - self, - *, - model_name: str, - config_class: str, - adapter_name: str, - added_date: date | None, - snapshot_date: date, - verified_on_cpu: bool, - verified_on_gpu: bool, - verified_on_spyre: bool, - num_downloads: int, - family: str, - architecture: str, - parameters_number: int, - failure_category: str | None, - error: str | None, - ) -> None: - """Storage-specific write of one row's normalized fields. - - Called by ``add_entry`` after the skip guard has passed. Subclasses must - not perform any deduplication here — that is the responsibility of - ``should_insert_row``. - """ - - def add_entry( - self, - *, - model_name: str, - config_class: str, - adapter_name: str, - added_date: date | None, - snapshot_date: date, - verified_on_cpu: bool, - verified_on_gpu: bool, - verified_on_spyre: bool, - num_downloads: int, - family: str, - architecture: str, - parameters_number: int, - failure_category: str | None, - error: str | None, - ) -> bool: - """Persist one row when the skip guard allows it. - - Returns True if the row was written, False if ``should_insert_row`` - rejected it. Idempotent to call for every row in the driver loop. - """ - model_name = _require_non_empty(model_name, "model_name") - if not self.should_insert_row(model_name): - return False - self._insert_entry( - model_name=model_name, - config_class=config_class, - adapter_name=adapter_name, - added_date=added_date, - snapshot_date=snapshot_date, - verified_on_cpu=verified_on_cpu, - verified_on_gpu=verified_on_gpu, - verified_on_spyre=verified_on_spyre, - num_downloads=num_downloads, - family=family, - architecture=architecture, - parameters_number=parameters_number, - failure_category=failure_category, - error=error, - ) - return True - - def should_insert_row(self, model_name: str) -> bool: - """Return True when *model_name* should be re-run. - - See module docstring for the full rule. In short: absent OR the most - recent prior row is a ``hardware_exception`` (retry) OR the prior - snapshot has aged past the skip window. - """ - model_name: str = _require_non_empty(model_name, "model_name") - return not self.get_recent_blocking_entries(model_name) - - def __enter__(self) -> ResultSink: - return self - - def __exit__(self, exc_type: object, exc: object, tb: object) -> None: - self.close() - - def flush(self) -> None: - """Persist any rows buffered in memory. Default is a no-op. - - Subclasses that buffer rows (e.g. ``ClickHouseResultSink``) override - this so callers can force a durable write at safe checkpoints — for - example, between batches in the weekly-test driver, so a hard crash - of the parent loses at most one batch instead of the whole run. - """ - - def close(self) -> None: - """Release resources. Default is a no-op; subclasses override.""" - - -class CsvResultSink(ResultSink): - """Append rows to a CSV file. - - On construction, any existing rows are read once into an in-memory index of - ``{model_name: list[dict]}`` so lookups are O(1). - """ - - def __init__(self, path: Path, today: date | None = None) -> None: - super().__init__(today=today) - self._path: Path = path - self._rows_by_model: dict[str, list[dict[str, Any]]] = {} - path.parent.mkdir(parents=True, exist_ok=True) - file_exists: bool = path.exists() and path.stat().st_size > 0 - if file_exists: - self._load_index() - self._fh = open(path, "a", newline="") - self._writer = csv.DictWriter(self._fh, fieldnames=list(TABLE_COLUMNS)) - if not file_exists: - self._writer.writeheader() - self._fh.flush() - - def _load_index(self) -> None: - with open(self._path, newline="") as fh: - reader = csv.DictReader(fh) - for raw_row in reader: - model_name: str = (raw_row.get("model_name") or "").strip() - if not model_name: - print( - f" warn: skipping CSV row with empty model_name in " - f"'{self._path}' (line {reader.line_num})" - ) - continue - self._rows_by_model.setdefault(model_name, []).append(dict(raw_row)) - loaded = sum(len(v) for v in self._rows_by_model.values()) - print( - f" index: loaded {loaded} row(s) for " - f"{len(self._rows_by_model)} model(s) from '{self._path}'" - ) - - def get_recent_blocking_entries(self, model_name: str) -> list[dict[str, Any]]: - key: str = _require_non_empty(model_name, "model_name") - rows: list[dict[str, Any]] = list(self._rows_by_model.get(key, ())) - filtered: list[tuple[date, dict[str, Any]]] = [] - for row in rows: - # hardware_exception is transient — do NOT let it block a re-run. - if ( - row.get("failure_category") or "" - ).strip() == _HARDWARE_EXCEPTION_CATEGORY: - continue - snap: date | None = _coerce_snapshot(row.get("snapshot_date")) - if snap is None: - continue - if not _within_skip_window(snap, self._today): - continue - filtered.append((snap, row)) - filtered.sort(key=lambda item: item[0], reverse=True) - return [row for _, row in filtered] - - def get_all_models(self) -> list[dict[str, Any]]: - result: list[dict[str, Any]] = [] - for rows in self._rows_by_model.values(): - best: dict[str, Any] = max( - rows, - key=lambda r: _coerce_snapshot(r.get("snapshot_date")) or date.min, - ) - result.append(best) - return result - - def _insert_entry( - self, - *, - model_name: str, - config_class: str, - adapter_name: str, - added_date: date | None, - snapshot_date: date, - verified_on_cpu: bool, - verified_on_gpu: bool, - verified_on_spyre: bool, - num_downloads: int, - family: str, - architecture: str, - parameters_number: int, - failure_category: str | None, - error: str | None, - ) -> None: - rec: dict[str, Any] = { - "model_name": model_name, - "config_class": config_class, - "adapter_name": adapter_name, - "added_date": added_date, - "snapshot_date": snapshot_date, - "verified_on_cpu": verified_on_cpu, - "verified_on_gpu": verified_on_gpu, - "verified_on_spyre": verified_on_spyre, - "num_downloads": num_downloads, - "family": family, - "architecture": architecture, - "parameters_number": parameters_number, - "failure_category": failure_category, - "error": error, - } - self._writer.writerow(rec) - self._fh.flush() - self._rows_by_model.setdefault(model_name, []).append(dict(rec)) - - def close(self) -> None: - if self._fh is not None: - self._fh.close() - - -class ClickHouseResultSink(ResultSink): - """Insert rows into ClickHouse. - - On construction the table is created if missing, and a single bulk SELECT - pre-fetches every model name that currently blocks a re-run (see the - module docstring for the rule) so that all subsequent - ``should_insert_row`` calls are O(1) with no network I/O. - - Rows accepted by the skip guard are accumulated in ``_pending`` and flushed - to ClickHouse in one bulk INSERT on ``close()`` (or ``__exit__``), so the - total network round-trips for N rows is 2 (one SELECT, one INSERT) instead - of 2 × N. - """ - - def __init__( - self, embedding_generative: EmbeddingGenerativeMode, today: date | None = None - ) -> None: - super().__init__(today=today) - self._embedding_generative = embedding_generative - if embedding_generative is EmbeddingGenerativeMode.EMBEDDING: - self._table_name = EMBEDDING_TABLE_NAME - create_sql = EMBEDDING_CREATE_TABLE_SQL - else: - self._table_name = GENERATIVE_TABLE_NAME - create_sql = GENERATIVE_CREATE_TABLE_SQL - self._client = get_client() - if not table_exists(self._client, self._table_name): - self._client.command(create_sql) - print(f"ClickHouse: table '{self._table_name}' created.\n") - else: - print(f"ClickHouse: table '{self._table_name}' already exists.\n") - - # Bulk pre-fetch: model names whose most-recent-in-window row blocks a - # re-run. Populated once here; used by get_recent_blocking_entries() - # for O(1) per-row checks. - self._skip_model_names: set[str] = self._fetch_blocking_names() - print( - f"ClickHouse: {len(self._skip_model_names)} model(s) already have a " - f"non-hardware-exception snapshot within the last " - f"{_SKIP_WINDOW_DAYS} days — will be skipped.\n" - ) - - # Rows waiting to be flushed; each entry is a list matching TABLE_COLUMNS order. - self._pending: list[list[Any]] = [] - - def _fetch_blocking_names(self) -> set[str]: - """One SELECT to get all model names that block a re-run. - - A model blocks iff it has any row in the skip window whose - ``failure_category`` is not ``hardware_exception`` — matching the - semantics of ``get_recent_blocking_entries``. Rows with - ``failure_category = 'hardware_exception'`` (including a lone row - older than the window) do NOT block. - """ - cutoff: date = self._today - timedelta(days=_SKIP_WINDOW_DAYS - 1) - result = self._client.query( - "SELECT DISTINCT model_name " - "FROM {db:Identifier}.{tbl:Identifier} " - "WHERE snapshot_date >= {cutoff:Date} " - "AND (failure_category IS NULL OR failure_category != {hw:String})", - parameters={ - "db": DATABASE, - "tbl": self._table_name, - "cutoff": cutoff, - "hw": _HARDWARE_EXCEPTION_CATEGORY, - }, - ) - return {row[0] for row in result.result_rows} - - def get_recent_blocking_entries(self, model_name: str) -> list[dict[str, Any]]: - """Return a non-empty sentinel list when *model_name* is in the skip set. - - The actual row data is not needed by the caller — it only tests - ``bool(result)`` — so we return a lightweight placeholder instead of - re-querying ClickHouse. - """ - key: str = _require_non_empty(model_name, "model_name") - if key in self._skip_model_names: - # Return a truthy non-empty list so should_insert_row returns False. - return [{"model_name": key}] - return [] - - def get_all_models(self) -> list[dict[str, Any]]: - columns_sql: str = ", ".join( - f"argMax({col}, snapshot_date) AS {col}" if col != "model_name" else col - for col in TABLE_COLUMNS - ) - result = self._client.query( - f"SELECT {columns_sql} " - "FROM {db:Identifier}.{tbl:Identifier} " - "GROUP BY model_name", - parameters={"db": DATABASE, "tbl": self._table_name}, - ) - return [dict(zip(TABLE_COLUMNS, row)) for row in result.result_rows] - - def _insert_entry( - self, - *, - model_name: str, - config_class: str, - adapter_name: str, - added_date: date | None, - snapshot_date: date, - verified_on_cpu: bool, - verified_on_gpu: bool, - verified_on_spyre: bool, - num_downloads: int, - family: str, - architecture: str, - parameters_number: int, - failure_category: str | None, - error: str | None, - ) -> None: - """Buffer the row; the actual INSERT happens in ``close()``.""" - # failure_category/error are non-nullable in the live table (String/ - # LowCardinality(String) DEFAULT '') — None must become '' here, or - # clickhouse_connect raises DataError on the bulk insert for any - # fully-passing row (which always has both fields as None). - self._pending.append( - [ - model_name, - config_class, - adapter_name, - added_date, - snapshot_date, - verified_on_cpu, - verified_on_gpu, - verified_on_spyre, - num_downloads, - family, - architecture, - parameters_number, - failure_category or "", - error or "", - ] - ) - - def flush(self) -> None: - """Flush all buffered rows to ClickHouse in a single bulk INSERT. - - Safe to call repeatedly; a no-op when the buffer is empty. The driver - loop calls this after every batch so a hard parent crash loses at most - one batch of rows rather than the entire run. - """ - if not self._pending: - return - print(f"ClickHouse: flushing {len(self._pending)} buffered row(s)…", flush=True) - self._client.insert( - self._table_name, - self._pending, - column_names=list(TABLE_COLUMNS), - ) - print("ClickHouse: bulk insert complete.") - self._pending.clear() - - def close(self) -> None: - """Flush any remaining buffered rows on shutdown.""" - self.flush() - - -class EmbeddingGenerativeMode(str, Enum): - EMBEDDING = "embedding" - GENERATIVE = "generative" diff --git a/tests/spyre/weekly_generation/sink/__init__.py b/tests/spyre/weekly_generation/sink/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/spyre/weekly_generation/sink/clickhouse_sink.py b/tests/spyre/weekly_generation/sink/clickhouse_sink.py new file mode 100644 index 00000000..2477b238 --- /dev/null +++ b/tests/spyre/weekly_generation/sink/clickhouse_sink.py @@ -0,0 +1,148 @@ +"""ClickHouse result sink: the real destination for weekly-scan rows. + +The only module in the weekly pipeline that reaches ``clickhouse_connect`` (via +``clickhouse_db``), which is why it lives apart from the ABC in ``result_sink``: +``--write-to-csv`` runs must reach that ABC, and the CSV sink, without the driver. +Note the two imports below are deliberately split — the client and credentials +come from ``clickhouse_db``, the table shape from the dependency-free +``table_schema``, so a schema consumer never drags in the driver. +""" + +from __future__ import annotations + +from datetime import date +from typing import Any + +from tests.spyre.weekly_generation.clickhouse_db import get_client, table_exists +from tests.spyre.weekly_generation.failure_categories import ( + FAILURE_CATEGORY_HARDWARE_EXCEPTION as _HARDWARE_EXCEPTION_CATEGORY, +) +from tests.spyre.weekly_generation.model_type import ModelType +from tests.spyre.weekly_generation.sink.result_sink import ResultSink +from tests.spyre.weekly_generation.table_schema import ( + DATABASE, + EMBEDDING_CREATE_TABLE_SQL, + EMBEDDING_TABLE_NAME, + GENERATIVE_CREATE_TABLE_SQL, + GENERATIVE_TABLE_NAME, + TABLE_COLUMNS, +) + + +class ClickHouseResultSink(ResultSink): + """Insert rows into ClickHouse. + + On construction the table is created if it does not already exist. + + Rows are accumulated in ``_pending`` and flushed to ClickHouse in one bulk + INSERT on ``close()`` (or ``__exit__``), so N rows cost one round-trip rather + than N. ``flush()`` forces that write early, which the weekly driver uses at + batch boundaries so a crash loses at most one batch. + """ + + def __init__(self, model_type: ModelType) -> None: + self._model_type = model_type + if model_type is ModelType.EMBEDDING: + self._table_name = EMBEDDING_TABLE_NAME + create_sql = EMBEDDING_CREATE_TABLE_SQL + else: + self._table_name = GENERATIVE_TABLE_NAME + create_sql = GENERATIVE_CREATE_TABLE_SQL + self._client = get_client() + if not table_exists(self._client, self._table_name): + self._client.command(create_sql) + print(f"ClickHouse: table '{self._table_name}' created.\n") + else: + print(f"ClickHouse: table '{self._table_name}' already exists.\n") + + # Rows waiting to be flushed; each entry is a list matching TABLE_COLUMNS order. + self._pending: list[list[Any]] = [] + + def fetch_hw_failure_models(self, snapshot_date: date) -> set[str]: + """Model names whose row on *snapshot_date* is a ``hardware_exception``. + + Groundwork for recovery runs: when a scan aborts because the accelerator + went away (see ``weekly_test.HardwareExceptionAbortError``), the affected + models are the only ones worth re-testing — the failure says nothing about + the checkpoint, so every other verdict from that run still stands. + + Not called yet. Nothing upstream exposes a recovery run, so this is + reachable only by hand until that is wired up. + """ + result = self._client.query( + "SELECT DISTINCT model_name " + "FROM {db:Identifier}.{tbl:Identifier} " + "WHERE snapshot_date = {snapshot_date:Date} " + "AND (failure_category = {hw:String})", + parameters={ + "db": DATABASE, + "tbl": self._table_name, + "snapshot_date": snapshot_date, + "hw": _HARDWARE_EXCEPTION_CATEGORY, + }, + ) + return {row[0] for row in result.result_rows} + + def _insert_entry( + self, + *, + model_name: str, + config_class: str, + adapter_name: str, + added_date: date | None, + snapshot_date: date, + verified_on_cpu: bool, + verified_on_gpu: bool, + verified_on_spyre: bool, + num_downloads: int, + family: str, + architecture: str, + parameters_number: int, + failure_category: str | None, + error: str | None, + ) -> None: + """Buffer the row; the actual INSERT happens in ``close()``.""" + # failure_category/error are non-nullable in the live table (String/ + # LowCardinality(String) DEFAULT '') — None must become '' here, or + # clickhouse_connect raises DataError on the bulk insert for any + # fully-passing row (which always has both fields as None). + self._pending.append( + [ + model_name, + config_class, + adapter_name, + added_date, + snapshot_date, + verified_on_cpu, + verified_on_gpu, + verified_on_spyre, + num_downloads, + family, + architecture, + parameters_number, + failure_category or "", + error or "", + ] + ) + + def flush(self) -> None: + """Flush all buffered rows to ClickHouse in a single bulk INSERT. + + Safe to call repeatedly; a no-op when the buffer is empty. The driver + loop calls this after every batch so a hard parent crash loses at most + one batch of rows rather than the entire run. + """ + if not self._pending: + return + print(f"ClickHouse: flushing {len(self._pending)} buffered row(s)…", flush=True) + self._client.insert( + self._table_name, + self._pending, + column_names=list(TABLE_COLUMNS), + ) + print("ClickHouse: bulk insert complete.") + self._pending.clear() + + def close(self) -> None: + """Flush any remaining buffered rows on shutdown.""" + self.flush() diff --git a/tests/spyre/weekly_generation/sink/csv_sink.py b/tests/spyre/weekly_generation/sink/csv_sink.py new file mode 100644 index 00000000..7ec23869 --- /dev/null +++ b/tests/spyre/weekly_generation/sink/csv_sink.py @@ -0,0 +1,83 @@ +"""CSV result sink: one run's rows to one new file, no database. + +What ``--write-to-csv`` selects, so the weekly pipeline can be exercised on a +host with no ClickHouse credentials — and, because ``TABLE_COLUMNS`` comes from +the dependency-free ``table_schema`` rather than from ``clickhouse_db``, with no +ClickHouse driver installed either. The header stays column-for-column identical +to the live table, since the point of this sink is to show what *would* be +inserted. +""" + +from __future__ import annotations + +import csv +from datetime import date +from pathlib import Path +from typing import Any + +from tests.spyre.weekly_generation.sink.result_sink import ResultSink +from tests.spyre.weekly_generation.table_schema import TABLE_COLUMNS + + +class CsvResultSink(ResultSink): + """Write rows to a fresh CSV file. Write-only, for no-database test runs. + + Nothing is read back, and the file must not already exist — so one file holds + exactly one run's rows. + """ + + def __init__(self, path: Path) -> None: + """Open *path* for writing. Raises if it already exists and is non-empty.""" + self._path: Path = path + if path.exists() and path.stat().st_size > 0: + raise FileExistsError( + f"'{path}' already exists and is not empty. This sink writes a " + f"single run's results to a new file and never reads one back; " + f"choose a different path or remove the existing file." + ) + path.parent.mkdir(parents=True, exist_ok=True) + self._fh = open(path, "w", newline="") + self._writer = csv.DictWriter(self._fh, fieldnames=list(TABLE_COLUMNS)) + self._writer.writeheader() + self._fh.flush() + + def _insert_entry( + self, + *, + model_name: str, + config_class: str, + adapter_name: str, + added_date: date | None, + snapshot_date: date, + verified_on_cpu: bool, + verified_on_gpu: bool, + verified_on_spyre: bool, + num_downloads: int, + family: str, + architecture: str, + parameters_number: int, + failure_category: str | None, + error: str | None, + ) -> None: + rec: dict[str, Any] = { + "model_name": model_name, + "config_class": config_class, + "adapter_name": adapter_name, + "added_date": added_date, + "snapshot_date": snapshot_date, + "verified_on_cpu": verified_on_cpu, + "verified_on_gpu": verified_on_gpu, + "verified_on_spyre": verified_on_spyre, + "num_downloads": num_downloads, + "family": family, + "architecture": architecture, + "parameters_number": parameters_number, + "failure_category": failure_category, + "error": error, + } + self._writer.writerow(rec) + self._fh.flush() + + def close(self) -> None: + if self._fh is not None: + self._fh.close() diff --git a/tests/spyre/weekly_generation/sink/result_sink.py b/tests/spyre/weekly_generation/sink/result_sink.py new file mode 100644 index 00000000..2842f8a4 --- /dev/null +++ b/tests/spyre/weekly_generation/sink/result_sink.py @@ -0,0 +1,133 @@ +"""Abstract result sink for the weekly Spyre test suite. + +This module holds only the ABC. The two concrete implementations live in the +``sink`` package, and callers build one through ``sink.sink_factory.create_sink`` +rather than importing them directly: + +- ``sink.csv_sink.CsvResultSink`` — write-only, for runs with no database access. + Writes one run's rows to a new file and never reads one back. +- ``sink.clickhouse_sink.ClickHouseResultSink`` — inserts rows into ClickHouse; + rows are buffered in memory then flushed in a single bulk INSERT on ``close()``. + +Keeping the ABC here, apart from its subclasses, is what lets callers type against +``ResultSink`` without dragging in a storage backend. Together with +``table_schema`` holding the column list, it is also what makes ``--write-to-csv`` +genuinely runnable on a host with no ``clickhouse_connect`` installed: only +``clickhouse_sink`` reaches the driver. + +A sink is write-only: every row handed to ``add_entry`` is recorded. Deciding +*which* models to evaluate happens upstream, in ``model_prefilter``, so a row +reaching here is one the caller already decided to record. Filtering at write +time would discard rows the caller chose to write, leaving a run with fewer rows +than the models it handled and no accounting for the difference. +""" + +from __future__ import annotations + +from abc import ABC, abstractmethod +from datetime import date + + +def _require_non_empty(value: str, field_name: str) -> str: + stripped: str = value.strip() + if not stripped: + raise ValueError(f"{field_name} must be a non-empty string") + return stripped + + +class ResultSink(ABC): + """Abstract destination for weekly-test result rows. + + Implementations must be usable as a context manager; ``__exit__`` should + release any external resources (CSV file handle, DB client). + + Closing is the owner's job, and only the owner's. A sink is closed exactly + once, by whoever constructed it — a helper that is *handed* a sink must not + wrap it in ``with``, because ``close()`` is not idempotent in the way that + would need: the CSV sink closes its file handle, so a later ``add_entry`` + raises ``ValueError: I/O operation on closed file``. ``weekly_test.main`` + owns the sink for a whole run and passes it down to the pre-filter, so this + is a live constraint rather than a hypothetical one. Use ``flush()`` for + intermediate durability instead. + """ + + @abstractmethod + def _insert_entry( + self, + *, + model_name: str, + config_class: str, + adapter_name: str, + added_date: date | None, + snapshot_date: date, + verified_on_cpu: bool, + verified_on_gpu: bool, + verified_on_spyre: bool, + num_downloads: int, + family: str, + architecture: str, + parameters_number: int, + failure_category: str | None, + error: str | None, + ) -> None: + """Storage-specific write of one row's normalized fields. + + Called by ``add_entry`` once *model_name* has been validated. Subclasses + must not deduplicate or drop rows here: a run's row count is meant to + match the models it handled. + """ + + def add_entry( + self, + *, + model_name: str, + config_class: str, + adapter_name: str, + added_date: date | None, + snapshot_date: date, + verified_on_cpu: bool, + verified_on_gpu: bool, + verified_on_spyre: bool, + num_downloads: int, + family: str, + architecture: str, + parameters_number: int, + failure_category: str | None, + error: str | None, + ) -> None: + """Persist one row. Always writes; rejects an empty *model_name*.""" + model_name = _require_non_empty(model_name, "model_name") + self._insert_entry( + model_name=model_name, + config_class=config_class, + adapter_name=adapter_name, + added_date=added_date, + snapshot_date=snapshot_date, + verified_on_cpu=verified_on_cpu, + verified_on_gpu=verified_on_gpu, + verified_on_spyre=verified_on_spyre, + num_downloads=num_downloads, + family=family, + architecture=architecture, + parameters_number=parameters_number, + failure_category=failure_category, + error=error, + ) + + def __enter__(self) -> ResultSink: + return self + + def __exit__(self, exc_type: object, exc: object, tb: object) -> None: + self.close() + + def flush(self) -> None: + """Persist any rows buffered in memory. Default is a no-op. + + Subclasses that buffer rows (e.g. ``ClickHouseResultSink``) override + this so callers can force a durable write at safe checkpoints — for + example, between batches in the weekly-test driver, so a hard crash + of the parent loses at most one batch instead of the whole run. + """ + + def close(self) -> None: + """Release resources. Default is a no-op; subclasses override.""" diff --git a/tests/spyre/weekly_generation/sink/sink_factory.py b/tests/spyre/weekly_generation/sink/sink_factory.py new file mode 100644 index 00000000..39dace04 --- /dev/null +++ b/tests/spyre/weekly_generation/sink/sink_factory.py @@ -0,0 +1,55 @@ +"""Build the result sink a run should write to. + +The single place that decides between the two backends, so no caller has to +import a concrete sink — or know that importing one has a cost. +""" + +from __future__ import annotations + +from pathlib import Path + +from tests.spyre.weekly_generation.model_type import ModelType +from tests.spyre.weekly_generation.sink.csv_sink import CsvResultSink +from tests.spyre.weekly_generation.sink.result_sink import ResultSink + + +def csv_path_for(base: Path, model_type: ModelType) -> Path: + """Per-model-type CSV path derived from a ``--write-to-csv`` argument. + + A single invocation can cover both model types and each needs its own file, + so the stem gets a ``-`` suffix. Exposed rather than inlined + because callers also report the path they wrote to, and computing it twice is + how the log line and the actual file get to disagree. + """ + return base.with_name(f"{base.stem}-{model_type}{base.suffix}") + + +def create_sink( + model_type: ModelType, + write_to_csv: Path | None, +) -> ResultSink: + """Return the sink for a *model_type* run, keyed on whether a CSV was asked for. + + One sink instance binds one destination — a single ClickHouse table, or a + single file — so a run covering both model types needs one per type rather + than one overall. + + With *write_to_csv* the verdicts go to a new CSV instead of ClickHouse, which + needs no credentials. Its path gets a ``-`` suffix, since a single + invocation can cover both types and each needs its own file. + + Otherwise a ``ClickHouseResultSink`` is built, which connects during + construction and creates its table if missing. + """ + if write_to_csv is not None: + path = csv_path_for(write_to_csv, model_type) + print(f"{model_type}: recording verdicts in '{path}' (no DB access)") + return CsvResultSink(path=path) + + # Deferred, so reaching the CSV branch neither imports clickhouse_connect nor + # reads .env. Together with csv_sink taking its column list from the + # dependency-free table_schema, this is what lets --write-to-csv run on a host + # with no driver installed at all. + from tests.spyre.weekly_generation.sink.clickhouse_sink import ClickHouseResultSink + + return ClickHouseResultSink(model_type=model_type) diff --git a/tests/spyre/weekly_generation/skip_writer.py b/tests/spyre/weekly_generation/skip_writer.py new file mode 100644 index 00000000..1538d098 --- /dev/null +++ b/tests/spyre/weekly_generation/skip_writer.py @@ -0,0 +1,59 @@ +"""Record the verdict rows for models the pre-filter rejected. + +Split from ``model_prefilter`` so that module stays free of database imports: +this one touches a ``ResultSink``, that one is pure. + +Every ``SkippedModel`` in ``PrefilterResult.skipped`` represents a terminal +verdict — no adapter, too large, MoE — and gets exactly one row written here. +The field mapping mirrors the ``add_entry`` signature on ``ResultSink``. +""" + +from __future__ import annotations + +from datetime import date +from typing import TYPE_CHECKING + +if TYPE_CHECKING: # pragma: no cover - import-cycle-free typing only + from tests.spyre.weekly_generation.model_prefilter import SkippedModel + from tests.spyre.weekly_generation.sink.result_sink import ResultSink + + +def write_skipped_rows( + sink: ResultSink, + skipped: list[SkippedModel], + *, + snapshot_date: date, + verbose: bool = True, +) -> int: + """Write one terminal row per entry in *skipped*. Returns rows written. + + Takes ``PrefilterResult.skipped``; every entry there is a terminal verdict and + gets a row. + """ + written = 0 + for item in skipped: + row = item.row + model_id = str(row["model_id"]) + sink.add_entry( + model_name=model_id, + config_class=str(row.get("config_class") or ""), + adapter_name="", + added_date=None, + snapshot_date=snapshot_date, + verified_on_cpu=False, + verified_on_gpu=False, + verified_on_spyre=False, + num_downloads=int(row.get("downloads") or 0), + family=str(row.get("model_type") or ""), + architecture=str(row.get("architectures") or ""), + parameters_number=int(row.get("parameters") or 0), + failure_category=item.failure_category, + error=None, + ) + written += 1 + if verbose: + print( + f" skip-row: '{model_id}' → {item.failure_category} " + f"({item.reason})" + ) + return written diff --git a/tests/spyre/weekly_generation/table_schema.py b/tests/spyre/weekly_generation/table_schema.py new file mode 100644 index 00000000..ee890eb1 --- /dev/null +++ b/tests/spyre/weekly_generation/table_schema.py @@ -0,0 +1,94 @@ +"""Table shape for the model_spyre_support tables. Pure data, no dependencies. + +Deliberately a leaf module: nothing here imports ``clickhouse_connect``, +``dotenv``, or anything else outside the standard library, and importing it has no +side effects. That is the point — ``csv_sink`` needs ``TABLE_COLUMNS`` so its +output is column-for-column identical to the live table, but a ``--write-to-csv`` +run should not need the ClickHouse driver installed or read a ``.env`` file off +disk to get a tuple of column names. ``clickhouse_db`` is where the client and +credentials live, and it re-exports everything below for its existing callers. + +The DDL and ``TABLE_COLUMNS`` are kept in this one file, adjacent, because they +must agree: the column list is what ``ClickHouseResultSink`` passes as +``column_names`` to a positional bulk insert, so a mismatch against the CREATE +TABLE body silently writes values into the wrong columns. ``test_table_schema.py`` +asserts the two match by parsing the DDL, so the requirement is enforced rather +than merely commented. + +Columns: + - model_name (String) – unique model identifier + - config_class (String) – model config_class (e.g. BertConfig, Qwen3Config) + - adapter_name (String) – adapter name (e.g. hf_bert, hf_gemma4) + - added_date (Date?) – date the adapter was added to the git repo (optional - None if not existing) + - snapshot_date (Date) – date this weekly snapshot was taken + - verified_on_cpu (Bool) – passes on CPU + - verified_on_gpu (Bool) – passes on GPU + - verified_on_spyre (Bool) – passes on Spyre + - num_downloads (UInt64) – number of downloads + - family (String) – model family reported by the catalog + - architecture (String) – model architecture reported by the catalog + - parameters_number (UInt64) – number of model parameters + - failure_category (String?) – classification for failures (optional - None if model passed) + - error (String?) – error message if the model failed (optional - None if model passed) +""" + +from __future__ import annotations + +EMBEDDING_TABLE_NAME = "embedding_model_spyre_support" +GENERATIVE_TABLE_NAME = "generative_model_spyre_support" +DATABASE = "spyre" + +# Single source of truth for the Python-facing column list. Order matches the +# CREATE TABLE DDL below and the positional values the sinks write — enforced by +# test_table_schema.py rather than left to reviewers. +TABLE_COLUMNS: tuple[str, ...] = ( + "model_name", + "config_class", + "adapter_name", + "added_date", + "snapshot_date", + "verified_on_cpu", + "verified_on_gpu", + "verified_on_spyre", + "num_downloads", + "family", + "architecture", + "parameters_number", + "failure_category", + "error", +) + + +def _make_create_table_sql(table_name: str) -> str: + """DDL for one snapshot table. Column order must match ``TABLE_COLUMNS``. + + ``ReplacingMergeTree(snapshot_date)`` is what makes a duplicate row a + tolerable failure mode: same-day rows for one model collapse on merge, which + is why the weekly scan prefers writing a possible duplicate over dropping a + result it was asked to produce. + """ + return f""" +CREATE TABLE IF NOT EXISTS {DATABASE}.{table_name} +( + model_name String, + config_class String, + adapter_name String, + added_date Nullable(Date), + snapshot_date Date, + verified_on_cpu Bool, + verified_on_gpu Bool, + verified_on_spyre Bool, + num_downloads UInt64, + family String, + architecture String, + parameters_number UInt64, + failure_category Nullable(String), + error Nullable(String) +) +ENGINE = ReplacingMergeTree(snapshot_date) +ORDER BY (model_name, snapshot_date) +""" + + +EMBEDDING_CREATE_TABLE_SQL = _make_create_table_sql(EMBEDDING_TABLE_NAME) +GENERATIVE_CREATE_TABLE_SQL = _make_create_table_sql(GENERATIVE_TABLE_NAME) diff --git a/tests/spyre/weekly_generation/weekly_sub_process.py b/tests/spyre/weekly_generation/weekly_sub_process.py new file mode 100644 index 00000000..eb338d6d --- /dev/null +++ b/tests/spyre/weekly_generation/weekly_sub_process.py @@ -0,0 +1,382 @@ +"""The child-process half of the weekly scan: evaluate models, report rows. + +Everything here runs in a ``multiprocessing`` "spawn" child started by +``weekly_test.main``, one child per batch. Two consequences shape the module: + +* **Torch stays out of module scope.** The heavy imports (``hf_adapters``, the + Spyre test entry points, ``tests.conftest``) happen inside the functions that + need them, so the parent — which imports this module only to name + ``_process_batch`` as the process target — never pays for them. +* **A failure is a row, not an exception.** One bad model must not cost the + other N-1 in its batch, so every model is evaluated inside its own + ``try``/``except`` and errors are recorded in the returned dict's + ``failure_category``/``error`` fields. The one exception is + ``hardware_exception``, which ends the batch early because the accelerator + itself is gone. + +The child never touches the sink. It returns plain dicts over the queue and the +parent does all the writing, which keeps database credentials and connection +state in one process. +""" + +import os +import sys +import traceback as _traceback +from datetime import date +from multiprocessing.queues import SimpleQueue + +from huggingface_hub.errors import HfHubHTTPError + +from tests.spyre.weekly_generation.failure_categories import ( + FAILURE_CATEGORY_CPU_GENERATE_FAILED, + FAILURE_CATEGORY_CPU_LOAD_FAILED, + FAILURE_CATEGORY_HARDWARE_EXCEPTION, + FAILURE_CATEGORY_MISFORMED_HF_FAILED, + FAILURE_CATEGORY_MODEL_TOO_LARGE, + FAILURE_CATEGORY_NOT_IMPLEMENTED_ADAPTER, + FAILURE_CATEGORY_QUANTIZED_MODEL, + FAILURE_CATEGORY_TEST_EXECUTION_EXCEPTION, + FAILURE_CATEGORY_VERIFICATION_FAILED, +) +from tests.spyre.weekly_generation.model_type import ModelType +from utils.utilities import ts + + +def _process_batch( + batch: list[dict], + adapter_dates: dict[str, str | None], + result_queue: SimpleQueue, + model_type: ModelType, + snapshot_date: date, +) -> None: + """Worker target: evaluate up to one batch of models in a single spawned child. + + Amortizes the per-child fixed cost (spawn + module imports + kernel + teardown on exit) across the batch; ``weekly_test``'s + ``{GENERATIVE,EMBEDDING}_NUMBER_OF_MODEL_PER_PROCESS`` set how many. Puts a + ``list[dict]`` on the queue — one full result dict per row, in the same order + as *batch*. If a single model errors, its ``error`` field is populated and + the loop continues to the next model; the child does NOT abort. + + The queue is the ``multiprocessing.SimpleQueue`` the parent created via its + spawn context, not an ``asyncio`` one — nothing here is coroutine-based. + + Exits via ``os._exit(0)`` rather than returning; see the comment at the end + for why skipping interpreter shutdown is what actually frees the card. + + Each returned dict has the same shape ``main`` expects for a rec plus an + ``error`` field (str or None): + + { + "model_name": ..., + "config_class": ..., + "adapter_name": ..., + "added_date": ..., # ISO 8601 str or None + "snapshot_date": ..., # date object + "verified_on_cpu": bool, + "verified_on_gpu": False, + "verified_on_spyre": bool, + "num_downloads": int, + "family": str, + "architecture": str, + "parameters_number": int, + "error": None or str, + "failure_category": None or str, + } + """ + import time as _t + + _child_entered: float = _t.monotonic() + print( + f"{ts()} child[{os.getpid()}] entered _process_batch with {len(batch)} model(s)", + flush=True, + ) + + from tests.conftest import resolve_adapter_module_for_test + + results: list[dict] = [] + for row in batch: + model_path: str = str(row["model_id"]) + rec: dict = { + "model_name": model_path, + "config_class": row.get("config_class"), + "adapter_name": "", + "added_date": None, + "snapshot_date": snapshot_date, + "verified_on_cpu": False, + "verified_on_gpu": False, + "verified_on_spyre": False, + "num_downloads": int(row.get("downloads") or 0), + "family": str(row.get("model_type") or ""), + "architecture": str(row.get("architectures") or ""), + "parameters_number": int(row.get("parameters") or 0), + "error": None, + "failure_category": None, + } + try: + try: + adapter_module = resolve_adapter_module_for_test(model_path) + except Exception: + rec["failure_category"] = FAILURE_CATEGORY_NOT_IMPLEMENTED_ADAPTER + raise + adapter_name: str = os.path.splitext( + os.path.basename(adapter_module.__file__) + )[0] + rec["adapter_name"] = adapter_name + rec["added_date"] = adapter_dates.get(adapter_name) + + metrics = eval_model(model_path, adapter_module, model_type) + rec["verified_on_cpu"] = bool(metrics.get("load", False)) + rec["verified_on_spyre"] = bool(metrics.get("correct", False)) + rec["error"] = metrics.get("error") or None + rec["failure_category"] = metrics.get("failure_category") or None + if not rec["verified_on_cpu"] and rec["failure_category"] is None: + rec["failure_category"] = _classify_failure( + rec["error"] or "", FAILURE_CATEGORY_CPU_LOAD_FAILED + ) + except Exception as e: + # Skip the error/traceback for shallow failure categories where the + # failure_category itself is fully self-describing. + if rec["failure_category"] not in ( + FAILURE_CATEGORY_NOT_IMPLEMENTED_ADAPTER, + FAILURE_CATEGORY_MODEL_TOO_LARGE, + ): + rec["error"] = ( + f"{type(e).__name__}: {e}\n" + f"{''.join(_traceback.format_exc().splitlines(keepends=True)[-6:])}" + ) + if rec["failure_category"] is None: + rec["failure_category"] = FAILURE_CATEGORY_TEST_EXECUTION_EXCEPTION + results.append(rec) + print( + f"{ts()} child[{os.getpid()}] finished model " + f"{len(results)}/{len(batch)}: {model_path!r} " + f"(verified_on_cpu={rec['verified_on_cpu']}, " + f"verified_on_spyre={rec['verified_on_spyre']}, " + f"failure_category={rec['failure_category']}, " + f"error={rec['error']})", + flush=True, + ) + # Bail out of the batch immediately on a hardware exception — the + # Spyre device is unreachable, so every remaining model in this + # batch would hit the same wall. The parent picks up the signal + # from the returned results and aborts the outer loop. + if rec["failure_category"] == FAILURE_CATEGORY_HARDWARE_EXCEPTION: + print( + f"{ts()} child[{os.getpid()}] aborting batch — " + f"hardware_exception detected; " + f"{len(batch) - len(results)} model(s) not attempted", + flush=True, + ) + break + + result_queue.put(results) + print( + f"{ts()} child[{os.getpid()}] done in " + f"{_t.monotonic() - _child_entered:.2f}s ({len(results)} results)", + flush=True, + ) + + # Skip Python's graceful shutdown: no atexit handlers, no thread + # finalization, no torch/torch_spyre destructors walking the tensor graph + # that the kernel is about to reclaim in bulk anyway. Closing the Spyre + # device FD on _exit(2) triggers the driver's own release path (VFIO + # unmap-all + IOMMU teardown), which is what actually returns the + # accelerator memory. Prior measurements: leaving Python's graceful + # shutdown in place cost ~30 s per child; running gc.collect() here on + # top of that added another ~20 s. + sys.stdout.flush() + sys.stderr.flush() + os._exit(0) + + +def _classify_failure(err: str, default: str) -> str: + """Bucket a raw error/traceback string into a failure_category. + + Signals in order of specificity: + + * ``"Failed to open the IBM Spyre VFIO device"`` — the accelerator itself + is unreachable (driver, permissions, another process holding it, …); + the model under test is not to blame, so tag as hardware_exception. + * ``"quantiz"`` / ``"optimum"`` — bitsandbytes / AWQ / GPTQ error text + almost always contains ``quantiz``, and ``optimum`` catches the + optimum-quanto / optimum-neuron loaders. + + Anything unrecognised falls through to *default* (usually the surrounding + context's fallback: cpu_load_failed at load time, test_execution_exception + at eval time). + """ + if not err: + return default + if "Failed to open the IBM Spyre VFIO device" in err or "Replace card" in err: + return FAILURE_CATEGORY_HARDWARE_EXCEPTION + if "does not appear to have files named ('model" in err: + return FAILURE_CATEGORY_MISFORMED_HF_FAILED + lowered: str = err.lower() + if "quantiz" in lowered or "optimum" in lowered: + return FAILURE_CATEGORY_QUANTIZED_MODEL + return default + + +def eval_model(model_id: str, adapter, model_type: ModelType) -> dict: + """Load *model_id* on CPU then run the mode's verification pipeline. + + Generative mode: CPU-load → CPU-generate (single-prompt HF forward pass) + → Spyre smoke + token-compare. The intermediate CPU-generate step catches + lazy shape errors, tokenizer/config mismatches, and custom-code bugs that + don't surface at ``from_pretrained`` time; on failure the row is tagged + ``cpu_generate_failed`` and the Spyre steps are skipped. + + Embedding mode: CPU-load → Spyre cosine-compare (no generate step — + embedders don't have a ``.generate()`` method). + + Returns a metrics dict with keys ``load``, ``correct``, ``error``, + ``failure_category``. ``correct`` is ``smoke_passed and not mismatches`` + — in embedding mode there is no smoke step, so ``smoke_passed`` is + treated as True and the outcome reduces to ``not mismatches``. + """ + load_on_cpu = False + smoke_passed = model_type == ModelType.EMBEDDING + mismatches = True + result: dict = {"error": "", "failure_category": None} + + try: + if adapter is not None: + load_on_cpu, load_error = _load_on_cpu( + model_path=model_id, model_type=model_type + ) + if load_error and not result["error"]: + result["error"] = load_error + if load_on_cpu: + if model_type == ModelType.GENERATIVE: + # Extra CPU-generate step — a load that succeeds but crashes + # here means the checkpoint is malformed in a way that only + # surfaces during forward. Stop before we waste Spyre time. + generate_ok, generate_error = _cpu_generate(model_path=model_id) + if not generate_ok: + if generate_error and not result["error"]: + result["error"] = generate_error + result["failure_category"] = _classify_failure( + generate_error or "", + FAILURE_CATEGORY_CPU_GENERATE_FAILED, + ) + else: + from tests.spyre.test_e2e_smoke_spyre import run_smoke_test + from tests.spyre.test_e2e_token_compare_spyre import ( + token_compare_spyre, + ) + + smoke_passed = ( + run_smoke_test(model_path=model_id)["status"] == "PASS" + ) + mismatches, _ = token_compare_spyre(model_id) + else: + from tests.spyre.test_e2e_embed_compare_spyre import ( + embed_compare_spyre, + ) + + mismatches, _ = embed_compare_spyre(model_id) + except Exception as e: + err: str = ( + f"{type(e).__name__}: {e}\n" + f"{''.join(_traceback.format_exc().splitlines(keepends=True)[-6:])}" + ) + result["error"] = err + result["failure_category"] = _classify_failure( + err, FAILURE_CATEGORY_TEST_EXECUTION_EXCEPTION + ) + finally: + result["correct"] = smoke_passed and not mismatches + result["load"] = load_on_cpu + if result["failure_category"] is None and load_on_cpu and not result["correct"]: + result["failure_category"] = FAILURE_CATEGORY_VERIFICATION_FAILED + return result + + +def _load_on_cpu( + model_path: str, + model_type: ModelType, +) -> tuple[bool, str | None]: + """Try to load *model_path* on CPU. Returns ``(loaded, error_message)``. + + ``error_message`` is ``None`` on success. On failure, it carries a + ``"ExcType: message\\n"`` string that the caller can + stash into the row's ``error`` field. Transient HF 5xx propagate — the + driver retries at a higher level. + """ + import hf_adapters.hf_common as _hf_common + from hf_adapters import AutoSpyreModelForCausalLM + from hf_adapters.auto_spyre_model import AutoSpyreModel + from tests.conftest import get_dtype_for_cpu + + _orig_device = _hf_common.DEVICE # save + _hf_common.DEVICE = "cpu" # patch + try: + dtype = get_dtype_for_cpu(model_path) + model = None + match model_type: + case ModelType.EMBEDDING: + model = AutoSpyreModel.from_pretrained(model_path, dtype=dtype) + case ModelType.GENERATIVE: + model = AutoSpyreModelForCausalLM.from_pretrained( + model_path, dtype=dtype + ) + + return model is not None, None + except HfHubHTTPError as e: + if e.response is not None and e.response.status_code >= 500: + raise + err: str = ( + f"{type(e).__name__}: {e}\n" + f"{''.join(_traceback.format_exc().splitlines(keepends=True)[-6:])}" + ) + print(f"_load_on_cpu exception - {e}") + return False, err + except Exception as e: + err = ( + f"{type(e).__name__}: {e}\n" + f"{''.join(_traceback.format_exc().splitlines(keepends=True)[-6:])}" + ) + print(f"_load_on_cpu exception - {e}") + return False, err + finally: + _hf_common.DEVICE = _orig_device # restore + + +def _cpu_generate(model_path: str) -> tuple[bool, str | None]: + """Run a single-prompt HF ``generate()`` on CPU for *model_path*. + + Separate from ``_load_on_cpu``: load succeeding tells us the checkpoint + is well-formed, generate succeeding tells us the forward pass runs + end-to-end (catches lazy shape errors, tokenizer/config mismatches, and + custom-code bugs that don't surface at ``from_pretrained`` time). + + Returns ``(ok, error_message)`` with the same convention as + ``_load_on_cpu``. + """ + import hf_adapters.hf_common as _hf_common + from tests.cpu._generate_helpers import simple_generate + + _orig_device = _hf_common.DEVICE + _hf_common.DEVICE = "cpu" + try: + simple_generate(model_path=model_path) + return True, None + except HfHubHTTPError as e: + if e.response is not None and e.response.status_code >= 500: + raise + err: str = ( + f"{type(e).__name__}: {e}\n" + f"{''.join(_traceback.format_exc().splitlines(keepends=True)[-6:])}" + ) + print(f"_cpu_generate exception - {e}") + return False, err + except Exception as e: + err = ( + f"{type(e).__name__}: {e}\n" + f"{''.join(_traceback.format_exc().splitlines(keepends=True)[-6:])}" + ) + print(f"_cpu_generate exception - {e}") + return False, err + finally: + _hf_common.DEVICE = _orig_device diff --git a/tests/spyre/weekly_generation/weekly_test.py b/tests/spyre/weekly_generation/weekly_test.py index 417462f5..e514caaa 100644 --- a/tests/spyre/weekly_generation/weekly_test.py +++ b/tests/spyre/weekly_generation/weekly_test.py @@ -1,65 +1,91 @@ """Weekly Spyre evaluation suite. -Evaluates the top-k generative or embedding models against Spyre hardware. -Run directly with a required ``--mode`` argument:: - - python tests/spyre/weekly_generation/weekly_test.py --mode generative [--top-k 200] - python tests/spyre/weekly_generation/weekly_test.py --mode embedding [--top-k 200] - -``--mode generative`` fetches top causal-LM models and runs the CPU-load + -token-compare pipeline. ``--mode embedding`` fetches top embedding models -and runs the CPU-load + cosine-compare pipeline. - -Additional flags: - -* ``--top-k N`` Number of top models to fetch by download count (default: 200). -* ``--write-to-csv F`` Write results to a CSV file instead of ClickHouse. -* ``--model-list-file F`` Load the model list from this JSON file (as produced by - ``.github/scripts/generate_weekly_shards.py``) instead of fetching it. Used by the - sharded CI workflow, where the top-K list is fetched once and split across parallel jobs. +Two ways to say which models to evaluate — exactly one is required:: + + # CI: evaluate a shard prepared upstream + python tests/spyre/weekly_generation/weekly_test.py \\ + --mode generative --model-list-file shards/generative-x1-shard-000.json + + # Manual: fetch, filter and evaluate in one pass + python tests/spyre/weekly_generation/weekly_test.py \\ + --mode embedding --fetch --top-k 200 + +Either way the same three pre-filters apply — no adapter for the config class, +too large for Spyre, MoE — and each dropped model gets a terminal row recording +why. With ``--model-list-file`` that happened upstream in +``.github/scripts/generate_weekly_shards.py``; with ``--fetch`` it happens here, +through the same ``fetch_and_filter``. So everything that reaches the evaluation +loop needs a Spyre card. + +Filtering before the list is sharded is what keeps shard durations comparable: +the dropped models cluster by download count, so filtering per-shard used to +leave some CI jobs finishing in minutes and others running for hours. + +``--mode`` is parsed into a ``ModelType``, which selects the per-process batch +size, the model class loaded in ``_load_on_cpu``, the verification pipeline in +``eval_model`` (token-compare for generative, cosine-compare for embedding), and +which ClickHouse table the sink reads and writes. With ``--fetch`` it also picks +the catalog to fetch. + +Process model +------------- +Each batch is evaluated in a freshly spawned child (``weekly_sub_process``) that +exits when the batch ends, which is what actually returns the accelerator's +memory. The parent owns the sink for the whole run and does all the writing; the +child only returns plain dicts over a queue. + +Flags: + +* ``--mode {generative,embedding}`` Required. See above. +* ``--model-list-file F`` Evaluate this JSON list. Mutually exclusive + with ``--fetch``; one of the two is required. +* ``--fetch`` Fetch and filter here instead. +* ``--top-k N`` With ``--fetch``: how many to fetch + (default: 10000). +* ``--max-params N`` With ``--fetch``: parameter ceiling. +* ``--write-to-csv F`` Record results in a new CSV instead of + ClickHouse, for runs with no database access. Write-only: the file must not + already exist, and nothing is read back. + +Result rows +----------- +One row is recorded per model handled, so a run's row count never silently +disagrees with its input. Nothing is filtered at write time — a row reaching the +sink is one the caller decided to record, and dropping it there would leave a run +with fewer rows than the models it handled and no accounting for the difference. +Same-day duplicates are collapsed on merge by +``ReplacingMergeTree(snapshot_date)``. """ import argparse import json import logging import multiprocessing -import os import subprocess import sys import time -import traceback as _traceback -from asyncio import Queue from datetime import date from pathlib import Path -from huggingface_hub.errors import HfHubHTTPError - -from tests.spyre.weekly_generation.result_sink import ( - EmbeddingGenerativeMode, +from tests.spyre.weekly_generation.failure_categories import ( + FAILURE_CATEGORY_HARDWARE_EXCEPTION, + FAILURE_CATEGORY_WORKER_CRASHED, + FAILURE_CATEGORY_WORKER_TIMEOUT, + MAX_NUMBER_PARAMS, ) -from utils.utilities import ts +from tests.spyre.weekly_generation.model_prefilter import fetch_and_filter +from tests.spyre.weekly_generation.model_type import ModelType +from tests.spyre.weekly_generation.sink.sink_factory import create_sink, csv_path_for +from tests.spyre.weekly_generation.weekly_sub_process import _process_batch +from utils.utilities import human_bytes, ts logging.getLogger("transformers").setLevel(logging.ERROR) -MAX_NUMBER_PARAMS = 60_000_000_000 - -FAILURE_CATEGORY_NOT_IMPLEMENTED_ADAPTER = "not-implemented-adapter" -FAILURE_CATEGORY_MODEL_TOO_LARGE = "model_too_large" -FAILURE_CATEGORY_CPU_LOAD_FAILED = "cpu_load_failed" -FAILURE_CATEGORY_CPU_GENERATE_FAILED = "cpu_generate_failed" -FAILURE_CATEGORY_QUANTIZED_MODEL = "quantized_model" -FAILURE_CATEGORY_HARDWARE_EXCEPTION = "hardware_exception" -FAILURE_CATEGORY_MISFORMED_HF_FAILED = "misformed_hf_failed" -FAILURE_CATEGORY_TEST_EXECUTION_EXCEPTION = "test_execution_exception" -FAILURE_CATEGORY_VERIFICATION_FAILED = "verification_failed" -FAILURE_CATEGORY_WORKER_CRASHED = "worker_crashed" -FAILURE_CATEGORY_WORKER_TIMEOUT = "worker_timeout" -FAILURE_CATEGORY_MOE = "moe" - -# Hard wall-clock cap for a single worker process (in seconds). If a batch -# takes longer than this, the parent kills the child, marks the entire batch -# as failed with FAILURE_CATEGORY_WORKER_TIMEOUT, and moves on. Prevents a -# single hung model from stalling the whole run indefinitely. +# Per-model wall-clock allowance for a worker process, in seconds. A batch's cap +# is this times its model count; see the timeout guard in main(), which kills the +# child and marks the batch FAILURE_CATEGORY_WORKER_TIMEOUT once it is exceeded, +# so one hung model cannot stall the run indefinitely. +_WORKER_TIMEOUT_SECONDS_PER_MODEL: int = 10 * 60 class HardwareExceptionAbortError(RuntimeError): @@ -67,41 +93,17 @@ class HardwareExceptionAbortError(RuntimeError): The Spyre accelerator is unreachable and no subsequent work in this process can succeed, so the run aborts. Bubbling this up to the - ``__main__`` block means the script exits with a non-zero code — - CI / GHA can alert on it, and a subsequent scheduled run picks the - aborted rows up automatically via the sink's retry-on-hardware_exception - skip rule. + ``__main__`` block means the script exits with a non-zero code, so CI / GHA + can alert on it. + + The aborted models are not picked up automatically: every model handed to a + run is evaluated, so the next scheduled scan re-evaluates the whole list + rather than just these. Re-running only the hardware-exception rows is what + ``ClickHouseResultSink.fetch_hw_failure_models`` is for, but nothing calls it + yet — a recovery run has to be driven by hand until that is wired up. """ -def _classify_failure(err: str, default: str) -> str: - """Bucket a raw error/traceback string into a failure_category. - - Signals in order of specificity: - - * ``"Failed to open the IBM Spyre VFIO device"`` — the accelerator itself - is unreachable (driver, permissions, another process holding it, …); - the model under test is not to blame, so tag as hardware_exception. - * ``"quantiz"`` / ``"optimum"`` — bitsandbytes / AWQ / GPTQ error text - almost always contains ``quantiz``, and ``optimum`` catches the - optimum-quanto / optimum-neuron loaders. - - Anything unrecognised falls through to *default* (usually the surrounding - context's fallback: cpu_load_failed at load time, test_execution_exception - at eval time). - """ - if not err: - return default - if "Failed to open the IBM Spyre VFIO device" in err or "Replace card" in err: - return FAILURE_CATEGORY_HARDWARE_EXCEPTION - if "does not appear to have files named ('model" in err: - return FAILURE_CATEGORY_MISFORMED_HF_FAILED - lowered: str = err.lower() - if "quantiz" in lowered or "optimum" in lowered: - return FAILURE_CATEGORY_QUANTIZED_MODEL - return default - - _REPO_ROOT = Path(__file__).resolve().parents[3] _SPYRE_TESTS_DIR = _REPO_ROOT / "tests" / "spyre" _TESTS_DIR = _REPO_ROOT / "tests" @@ -210,12 +212,6 @@ def _parse_args(argv: list[str] | None = None) -> argparse.Namespace: "causal-LM load + token-compare pipeline." ), ) - parser.add_argument( - "--top-k", - type=int, - default=200, - help="Number of top models to fetch (by downloads).", - ) parser.add_argument( "--write-to-csv", type=Path, @@ -226,321 +222,57 @@ def _parse_args(argv: list[str] | None = None) -> argparse.Namespace: "into ClickHouse. No DB connection is made when this flag is set." ), ) - parser.add_argument( + source = parser.add_mutually_exclusive_group(required=True) + source.add_argument( "--model-list-file", type=Path, default=None, metavar="MODEL_LIST_JSON", help=( - "Load the model list from this JSON file instead of calling " - "fetch_top_generative_models/fetch_top_embedding_models. --top-k is " - "ignored when this is set. Used by the sharded CI workflow." + "Evaluate the models in this JSON file, as produced by " + ".github/scripts/generate_weekly_shards.py. It applies the " + "pre-filters and records a terminal row for each model it drops, so " + "everything in the file is expected to need a Spyre card. This is " + "how CI runs: the list is fetched and filtered once, then sharded " + "across parallel jobs." ), ) - return parser.parse_args(argv) - - -def _load_on_cpu( - model_path: str, mode: EmbeddingGenerativeMode -) -> tuple[bool, str | None]: - """Try to load *model_path* on CPU. Returns ``(loaded, error_message)``. - - ``error_message`` is ``None`` on success. On failure it carries a - ``"ExcType: message\\n"`` string that the caller can - stash into the row's ``error`` field. Transient HF 5xx propagate — the - driver retries at a higher level. - """ - import hf_adapters.hf_common as _hf_common - from hf_adapters import AutoSpyreModelForCausalLM - from hf_adapters.auto_spyre_model import AutoSpyreModel - from tests.conftest import get_dtype_for_cpu - - _orig_device = _hf_common.DEVICE # save - _hf_common.DEVICE = "cpu" # patch - try: - dtype = get_dtype_for_cpu(model_path) - model = None - match mode: - case "embedding": - model = AutoSpyreModel.from_pretrained(model_path, dtype=dtype) - case "generative": - model = AutoSpyreModelForCausalLM.from_pretrained( - model_path, dtype=dtype - ) - - return model is not None, None - except HfHubHTTPError as e: - if e.response is not None and e.response.status_code >= 500: - raise - err: str = ( - f"{type(e).__name__}: {e}\n" - f"{''.join(_traceback.format_exc().splitlines(keepends=True)[-6:])}" - ) - print(f"_load_on_cpu exception - {e}") - return False, err - except Exception as e: - err = ( - f"{type(e).__name__}: {e}\n" - f"{''.join(_traceback.format_exc().splitlines(keepends=True)[-6:])}" - ) - print(f"_load_on_cpu exception - {e}") - return False, err - finally: - _hf_common.DEVICE = _orig_device # restore - - -def _cpu_generate(model_path: str) -> tuple[bool, str | None]: - """Run a single-prompt HF ``generate()`` on CPU for *model_path*. - - Separate from ``_load_on_cpu``: load succeeding tells us the checkpoint - is well-formed, generate succeeding tells us the forward pass runs - end-to-end (catches lazy shape errors, tokenizer/config mismatches, and - custom-code bugs that don't surface at ``from_pretrained`` time). - - Returns ``(ok, error_message)`` with the same convention as - ``_load_on_cpu``. - """ - import hf_adapters.hf_common as _hf_common - from tests.cpu._generate_helpers import simple_generate - - _orig_device = _hf_common.DEVICE - _hf_common.DEVICE = "cpu" - try: - simple_generate(model_path=model_path) - return True, None - except HfHubHTTPError as e: - if e.response is not None and e.response.status_code >= 500: - raise - err: str = ( - f"{type(e).__name__}: {e}\n" - f"{''.join(_traceback.format_exc().splitlines(keepends=True)[-6:])}" - ) - print(f"_cpu_generate exception - {e}") - return False, err - except Exception as e: - err = ( - f"{type(e).__name__}: {e}\n" - f"{''.join(_traceback.format_exc().splitlines(keepends=True)[-6:])}" - ) - print(f"_cpu_generate exception - {e}") - return False, err - finally: - _hf_common.DEVICE = _orig_device - - -def eval_model(model_id: str, adapter, mode: EmbeddingGenerativeMode) -> dict: - """Load *model_id* on CPU then run the mode's verification pipeline. - - Generative mode: CPU-load → CPU-generate (single-prompt HF forward pass) - → Spyre smoke + token-compare. The intermediate CPU-generate step catches - lazy shape errors, tokenizer/config mismatches, and custom-code bugs that - don't surface at ``from_pretrained`` time; on failure the row is tagged - ``cpu_generate_failed`` and the Spyre steps are skipped. - - Embedding mode: CPU-load → Spyre cosine-compare (no generate step — - embedders don't have a ``.generate()`` method). - - Returns a metrics dict with keys ``load``, ``correct``, ``error``, - ``failure_category``. ``correct`` is ``smoke_passed and not mismatches`` - — in embedding mode there is no smoke step, so ``smoke_passed`` is - treated as True and the outcome reduces to ``not mismatches``. - """ - load_on_cpu = False - smoke_passed = mode == EmbeddingGenerativeMode.EMBEDDING - mismatches = True - result: dict = {"error": "", "failure_category": None} - - try: - if adapter is not None: - load_on_cpu, load_error = _load_on_cpu(model_path=model_id, mode=mode) - if load_error and not result["error"]: - result["error"] = load_error - if load_on_cpu: - if mode == EmbeddingGenerativeMode.GENERATIVE: - # Extra CPU-generate step — a load that succeeds but crashes - # here means the checkpoint is malformed in a way that only - # surfaces during forward. Stop before we waste Spyre time. - generate_ok, generate_error = _cpu_generate(model_path=model_id) - if not generate_ok: - if generate_error and not result["error"]: - result["error"] = generate_error - result["failure_category"] = _classify_failure( - generate_error or "", - FAILURE_CATEGORY_CPU_GENERATE_FAILED, - ) - else: - from tests.spyre.test_e2e_smoke_spyre import run_smoke_test - from tests.spyre.test_e2e_token_compare_spyre import ( - token_compare_spyre, - ) - - smoke_passed = ( - run_smoke_test(model_path=model_id)["status"] == "PASS" - ) - mismatches, _ = token_compare_spyre(model_id) - else: - from tests.spyre.test_e2e_embed_compare_spyre import ( - embed_compare_spyre, - ) - - mismatches, _ = embed_compare_spyre(model_id) - except Exception as e: - err: str = ( - f"{type(e).__name__}: {e}\n" - f"{''.join(_traceback.format_exc().splitlines(keepends=True)[-6:])}" - ) - result["error"] = err - result["failure_category"] = _classify_failure( - err, FAILURE_CATEGORY_TEST_EXECUTION_EXCEPTION - ) - finally: - result["correct"] = smoke_passed and not mismatches - result["load"] = load_on_cpu - if result["failure_category"] is None and load_on_cpu and not result["correct"]: - result["failure_category"] = FAILURE_CATEGORY_VERIFICATION_FAILED - return result - - -def _process_batch( - batch: list[dict], - adapter_dates: dict[str, str | None], - result_queue: Queue, - mode: EmbeddingGenerativeMode, - snapshot_date: date, -) -> None: - """Worker target: evaluate up to ``NUMBER_OF_MODEL_PER_PROCESS`` models - in a single spawned child. - - Amortizes the per-child fixed cost (spawn + module imports + kernel - teardown on exit) across N models. Puts a ``list[dict]`` on the queue — - one full result dict per row, in the same order as *batch*. If a single - model errors, its ``error`` field is populated and the loop continues to - the next model; the child does NOT abort. - - Each returned dict has the same shape ``main`` expects for a rec plus an - ``error`` field (str or None): - - { - "model_name": ..., - "config_class": ..., - "adapter_name": ..., - "added_date": ..., # ISO 8601 str or None - "snapshot_date": ..., # date object - "verified_on_cpu": bool, - "verified_on_gpu": False, - "verified_on_spyre": bool, - "num_downloads": int, - "family": str, - "architecture": str, - "parameters_number": int, - "error": None or str, - "failure_category": None or str, - } - """ - import time as _t - - _child_entered: float = _t.monotonic() - print( - f"{ts()} child[{os.getpid()}] entered _process_batch with {len(batch)} model(s)", - flush=True, + source.add_argument( + "--fetch", + action="store_true", + help=( + "Fetch the top --top-k models for --mode, apply the same " + "pre-filters, record the terminal rows, and evaluate the survivors — " + "all in one pass. For manual runs where there is no shard file." + ), ) - - from tests.conftest import resolve_adapter_module_for_test - - results: list[dict] = [] - for row in batch: - model_path: str = str(row["model_id"]) - rec: dict = { - "model_name": model_path, - "config_class": row.get("config_class"), - "adapter_name": "", - "added_date": None, - "snapshot_date": snapshot_date, - "verified_on_cpu": False, - "verified_on_gpu": False, - "verified_on_spyre": False, - "num_downloads": int(row.get("downloads") or 0), - "family": str(row.get("model_type") or ""), - "architecture": str(row.get("architectures") or ""), - "parameters_number": int(row.get("parameters") or 0), - "error": None, - "failure_category": None, - } - try: - try: - adapter_module = resolve_adapter_module_for_test(model_path) - except Exception: - rec["failure_category"] = FAILURE_CATEGORY_NOT_IMPLEMENTED_ADAPTER - raise - adapter_name: str = os.path.splitext( - os.path.basename(adapter_module.__file__) - )[0] - rec["adapter_name"] = adapter_name - rec["added_date"] = adapter_dates.get(adapter_name) - - metrics = eval_model(model_path, adapter_module, mode) - rec["verified_on_cpu"] = bool(metrics.get("load", False)) - rec["verified_on_spyre"] = bool(metrics.get("correct", False)) - rec["error"] = metrics.get("error") or None - rec["failure_category"] = metrics.get("failure_category") or None - if not rec["verified_on_cpu"] and rec["failure_category"] is None: - rec["failure_category"] = _classify_failure( - rec["error"] or "", FAILURE_CATEGORY_CPU_LOAD_FAILED - ) - except Exception as e: - # Skip the error/traceback for shallow failure categories where the - # failure_category itself is fully self-describing. - if rec["failure_category"] not in ( - FAILURE_CATEGORY_NOT_IMPLEMENTED_ADAPTER, - FAILURE_CATEGORY_MODEL_TOO_LARGE, - ): - rec["error"] = ( - f"{type(e).__name__}: {e}\n" - f"{''.join(_traceback.format_exc().splitlines(keepends=True)[-6:])}" - ) - if rec["failure_category"] is None: - rec["failure_category"] = FAILURE_CATEGORY_TEST_EXECUTION_EXCEPTION - results.append(rec) - print( - f"{ts()} child[{os.getpid()}] finished model " - f"{len(results)}/{len(batch)}: {model_path!r} " - f"(verified_on_cpu={rec['verified_on_cpu']}, " - f"verified_on_spyre={rec['verified_on_spyre']}, " - f"failure_category={rec['failure_category']}, " - f"error={rec['error']})", - flush=True, - ) - # Bail out of the batch immediately on a hardware exception — the - # Spyre device is unreachable, so every remaining model in this - # batch would hit the same wall. The parent picks up the signal - # from the returned results and aborts the outer loop. - if rec["failure_category"] == FAILURE_CATEGORY_HARDWARE_EXCEPTION: - print( - f"{ts()} child[{os.getpid()}] aborting batch — " - f"hardware_exception detected; " - f"{len(batch) - len(results)} model(s) not attempted", - flush=True, - ) - break - - result_queue.put(results) - print( - f"{ts()} child[{os.getpid()}] done in " - f"{_t.monotonic() - _child_entered:.2f}s ({len(results)} results)", - flush=True, + fetch_opts = parser.add_argument_group("--fetch options") + fetch_opts.add_argument( + "--top-k", + type=int, + default=10_000, + help=( + "With --fetch: how many top models to fetch by downloads " + "(default: 10000). Ignored with --model-list-file." + ), ) - - # Skip Python's graceful shutdown: no atexit handlers, no thread - # finalization, no torch/torch_spyre destructors walking the tensor graph - # that the kernel is about to reclaim in bulk anyway. Closing the Spyre - # device FD on _exit(2) triggers the driver's own release path (VFIO - # unmap-all + IOMMU teardown), which is what actually returns the - # accelerator memory. Prior measurements: leaving Python's graceful - # shutdown in place cost ~30 s per child; running gc.collect() here on - # top of that added another ~20 s. - sys.stdout.flush() - sys.stderr.flush() - os._exit(0) + fetch_opts.add_argument( + "--max-params", + type=int, + default=MAX_NUMBER_PARAMS, + help=( + "With --fetch: reject models above this parameter count " + f"(default: {MAX_NUMBER_PARAMS:,}). Ignored with --model-list-file, " + "whose producer applied its own limit." + ), + ) + args = parser.parse_args(argv) + # Reject the silent no-op of passing a fetch tuning flag without --fetch. + if not args.fetch: + for flag, dest in (("--top-k", "top_k"), ("--max-params", "max_params")): + if getattr(args, dest) != parser.get_default(dest): + parser.error(f"{flag} only applies with --fetch") + return args def _repo_cache_dir(repo_id: str) -> Path: @@ -603,17 +335,10 @@ def _cleanup_batch_weights( freed = _delete_repo_weights(to_delete) total_freed += freed if freed: - print(f" freed {_human_bytes(freed)} (total {_human_bytes(total_freed)})") + print(f" freed {human_bytes(freed)} (total {human_bytes(total_freed)})") return total_freed -def _human_bytes(n): - for unit in ("B", "KB", "MB", "GB", "TB"): - if n < 1024 or unit == "TB": - return f"{n:.1f}{unit}" - n /= 1024 - - def _chunk_into_batches(rows: list[dict], batch_size: int) -> list[list[dict]]: """Split *rows* into consecutive sub-lists of length *batch_size* (the last batch may be shorter). @@ -622,195 +347,90 @@ def _chunk_into_batches(rows: list[dict], batch_size: int) -> list[list[dict]]: def main( - mode: EmbeddingGenerativeMode, - write_to_csv: Path | str | None, + model_type: ModelType, + model_list_file: Path | None, + write_to_csv: Path | None, + fetch: bool, top_k: int, - model_list_file: Path | None = None, + max_params: int, ) -> None: - from tests.spyre.weekly_generation.result_sink import ( - ClickHouseResultSink, - CsvResultSink, - ResultSink, - ) - from utils.fetch_top_embedding_models import fetch_top_embedding_models - from utils.fetch_top_generative_models import fetch_top_generative_models + """Evaluate a model list on Spyre, one spawned worker per batch. + + Exactly one of *model_list_file* (a shard prepared by + ``generate_weekly_shards``) and *fetch* supplies the list; argparse enforces + that. Both arrive pre-filtered, so every model reaching the batch loop is + expected to need a card. + + Owns the sink for the whole run: it is created here, written to by both the + ``--fetch`` pre-filter and the evaluation loop, flushed at each batch + boundary, and closed once in the ``finally`` below. ``top_k``/``max_params`` + apply only under *fetch*. + + Raises: + HardwareExceptionAbortError: a batch reported ``hardware_exception``, so + the accelerator is unreachable and the remaining batches are skipped. + """ + from tests.spyre.weekly_generation.sink.result_sink import ResultSink print(f"{ts()} Starting main.") total_freed: int = 0 snapshot_date = date.today() - if mode == EmbeddingGenerativeMode.GENERATIVE: - number_of_model_per_process = GENERATIVE_NUMBER_OF_MODEL_PER_PROCESS - elif mode == EmbeddingGenerativeMode.EMBEDDING: - number_of_model_per_process = EMBEDDING_NUMBER_OF_MODEL_PER_PROCESS - else: - raise Exception(f"Unknown mode: {mode}") - - if model_list_file is not None: - print(f"{ts()} Loading model list from '{model_list_file}' (--top-k ignored).") - to_process_list = json.loads(model_list_file.read_text()) - elif mode == EmbeddingGenerativeMode.GENERATIVE: - to_process_list = fetch_top_generative_models(limit=top_k) - else: - to_process_list = fetch_top_embedding_models(limit=top_k) - # Must run after to_process_list is resolved — the cache check is scoped to - # this run's models rather than walking the whole (network-mounted) cache. - preexisting: set[str] = _repos_with_weights( - [str(row["model_id"]) for row in to_process_list] - ) + + models_per_process = { + ModelType.GENERATIVE: GENERATIVE_NUMBER_OF_MODEL_PER_PROCESS, + ModelType.EMBEDDING: EMBEDDING_NUMBER_OF_MODEL_PER_PROCESS, + } + adapter_dates: dict[str, str | None] = _get_adapter_dates() - sink: ResultSink - if write_to_csv: - sink = CsvResultSink(path=write_to_csv, today=snapshot_date) - print( - f"CSV mode: results will be appended to '{write_to_csv}' (no DB access).\n" - ) - else: - sink = ClickHouseResultSink(today=snapshot_date, embedding_generative=mode) - print("DB mode: results will be appended to the DB.\n") - total = len(to_process_list) + sink: ResultSink = create_sink( + model_type=model_type, + write_to_csv=write_to_csv, + ) + + # All six are read by the finally block, so they are bound before the try + # opens — otherwise a failure while building the model list would raise + # UnboundLocalError from the cleanup path and mask the real error. processed = 0 + total = 0 overall_start = time.monotonic() + batch_paths: list[str] = [] + had_weights_map: dict[str, bool] = {} + preexisting: set[str] = set() - # Early-stop guard runs in the parent (fast: dict lookup for CSV, - # single SELECT for CH), so we drop already-recent models BEFORE - # batching. That keeps batch sizes uniform relative to real work. - prefiltered: list[dict] = [] - early_skipped: int = 0 - moe_skipped: int = 0 - too_large_skipped: int = 0 - unsupported_skipped: int = 0 - print(f"{ts()} Will process {len(to_process_list)} models in total.") - for row in to_process_list: - model_path = str(row["model_id"]) - if not sink.should_insert_row(model_path): - early_skipped += 1 - print( - f"{ts()} sink: '{model_path}' skipped early — " - f"recent snapshot exists within the " - f"{sink.__class__.__name__} skip window" - ) - continue - # No adapter registered for this model's config class — same terminal - # decision resolve_adapter_module_for_test would reach in the worker, - # but reached here without spawning one. Uses the fetcher-computed - # is_supported flag (True iff config_class is in the adapter mapping). - - if row.get("is_supported") is False: - unsupported_skipped += 1 - sink.add_entry( - model_name=model_path, - config_class=str(row.get("config_class") or ""), - adapter_name="", - added_date=None, - snapshot_date=snapshot_date, - verified_on_cpu=False, - verified_on_gpu=False, - verified_on_spyre=False, - num_downloads=int(row.get("downloads") or 0), - family=str(row.get("model_type") or ""), - architecture=str(row.get("architectures") or ""), - parameters_number=int(row.get("parameters") or 0), - failure_category=FAILURE_CATEGORY_NOT_IMPLEMENTED_ADAPTER, - error=None, - ) - print( - f"{ts()} sink: '{model_path}' skipped early — " - f"no adapter for config_class={row.get('config_class')!r}" - ) - continue - # Reject models too large to bring up on Spyre BEFORE spawning a - # worker. Same intent as the in-worker guard in _process_batch, but - # catches everything the fetcher already sized so no worker time is - # wasted. _process_batch keeps its own check as a defensive backstop - # for rows where parameters were unknown at fetch time. - params = row.get("parameters") - if params not in (None, "") and int(params) > MAX_NUMBER_PARAMS: - too_large_skipped += 1 - sink.add_entry( - model_name=model_path, - config_class=str(row.get("config_class") or ""), - adapter_name="", - added_date=None, - snapshot_date=snapshot_date, - verified_on_cpu=False, - verified_on_gpu=False, - verified_on_spyre=False, - num_downloads=int(row.get("downloads") or 0), - family=str(row.get("model_type") or ""), - architecture=str(row.get("architectures") or ""), - parameters_number=int(params), - failure_category=FAILURE_CATEGORY_MODEL_TOO_LARGE, - error=None, - ) - print( - f"{ts()} sink: '{model_path}' skipped early — " - f"{int(params):,} parameters exceeds the " - f"{MAX_NUMBER_PARAMS:,} limit" - ) - continue - # MoE models aren't supported on Spyre yet — write the row up-front - # with failure_category=moe and don't send it to the workers. - # is_moe is precomputed at fetch time (utils/hf_model_catalog.py) so - # it survives a JSON round-trip through --model-list-file, unlike the - # raw (non-serializable) model_info object. - if row.get("is_moe"): - moe_skipped += 1 - sink.add_entry( - model_name=model_path, - config_class=str(row.get("config_class") or ""), - adapter_name="", - added_date=None, + # The try opens before the model list is built so that a failure while + # fetching (a Hub outage, say) still closes the sink — under --fetch the + # pre-filter has by then already written terminal rows worth keeping. + try: + if fetch: + rows: list[dict] = fetch_and_filter( + model_type=model_type, snapshot_date=snapshot_date, - verified_on_cpu=False, - verified_on_gpu=False, - verified_on_spyre=False, - num_downloads=int(row.get("downloads") or 0), - family=str(row.get("model_type") or ""), - architecture=str(row.get("architectures") or ""), - parameters_number=int(row.get("parameters") or 0), - failure_category=FAILURE_CATEGORY_MOE, - error=None, + top_k=top_k, + sink=sink, + max_params=max_params, ) - print(f"{ts()} sink: '{model_path}' skipped early — MoE model") - continue - prefiltered.append(row) - if early_skipped: - print( - f"\n{ts()} Early-skip: {early_skipped}/{total} models already have a " - f"recent snapshot; {len(prefiltered)} left to evaluate.\n" - ) - if unsupported_skipped: - print( - f"{ts()} Unsupported-skip: {unsupported_skipped}/{total} models have " - f"no adapter for their config_class and were written directly to the " - f"sink.\n" - ) - if too_large_skipped: - print( - f"{ts()} Too-large-skip: {too_large_skipped}/{total} models exceed " - f"the {MAX_NUMBER_PARAMS:,} parameter limit and were written directly " - f"to the sink.\n" - ) - if moe_skipped: - print( - f"{ts()} MoE-skip: {moe_skipped}/{total} models tagged as moe and " - f"written directly to the sink.\n" + else: + assert model_list_file is not None # argparse guarantees one of the two + print(f"{ts()} Loading model list from '{model_list_file}'.") + rows = json.loads(model_list_file.read_text()) + + total = len(rows) + print(f"{ts()} Will evaluate {total} model(s).") + + # Must run after *rows* is resolved — the cache check is scoped to this + # run's models rather than walking the whole (network-mounted) cache. + preexisting = _repos_with_weights([str(row["model_id"]) for row in rows]) + + batch_size = models_per_process[model_type] + batches: list[list[dict]] = _chunk_into_batches( + rows=rows, + batch_size=batch_size, ) + total_batches: int = len(batches) - batches: list[list[dict]] = _chunk_into_batches( - prefiltered, number_of_model_per_process - ) - total_batches: int = len(batches) - - ctx = multiprocessing.get_context("spawn") - - # Initialised here so the finally block can clean up the in-flight batch - # even when KeyboardInterrupt fires mid-batch. - batch_paths: list[str] = [] - had_weights_map: dict[str, bool] = {} + ctx = multiprocessing.get_context("spawn") - try: for batch_idx, batch in enumerate(batches, start=1): batch_start = time.monotonic() batch_paths = [str(r["model_id"]) for r in batch] @@ -832,12 +452,12 @@ def main( batch, adapter_dates, result_queue, - mode, + model_type, snapshot_date, ), ) proc.start() - timeout = 10 * 60 * number_of_model_per_process + timeout = _WORKER_TIMEOUT_SECONDS_PER_MODEL * batch_size proc.join(timeout=timeout) @@ -880,7 +500,7 @@ def main( "family": str(row.get("model_type") or ""), "architecture": str(row.get("architectures") or ""), "parameters_number": int(row.get("parameters") or 0), - "error": (f"worker exceeded " f"{timeout}s timeout"), + "error": f"worker exceeded {timeout}s timeout", "failure_category": FAILURE_CATEGORY_WORKER_TIMEOUT, } for row, path in zip(batch, batch_paths) @@ -930,7 +550,7 @@ def main( except ValueError: rec["added_date"] = None - if sink.add_entry( + sink.add_entry( model_name=str(rec["model_name"]), config_class=str(rec["config_class"]), adapter_name=str(rec["adapter_name"]), @@ -949,17 +569,13 @@ def main( else str(rec["failure_category"]) ), error=(None if rec.get("error") is None else str(rec["error"])), - ): - print( - f"{ts()} sink: row written for '{model_path}' " - f"(verified_on_cpu={rec.get('verified_on_cpu')}, " - f"verified_on_spyre={rec.get('verified_on_spyre')}, " - f"failure_category={rec.get('failure_category')}, )" - ) - else: - print( - f"{ts()} sink: row skipped for '{model_path}' (guard rejected)" - ) + ) + print( + f"{ts()} sink: row written for '{model_path}' " + f"(verified_on_cpu={rec.get('verified_on_cpu')}, " + f"verified_on_spyre={rec.get('verified_on_spyre')}, " + f"failure_category={rec.get('failure_category')}, )" + ) # Cache cleanup: delete weights downloaded during this batch, # regardless of whether the worker processed each model. @@ -1008,9 +624,10 @@ def main( sink.close() if write_to_csv: - print( - f"\n{ts()} CSV: '{write_to_csv}' closed ({processed} rows processed)." - ) + # Report the file actually written, not the bare --write-to-csv + # argument — the factory suffixes it with the model type. + written_to = csv_path_for(write_to_csv, model_type) + print(f"\n{ts()} CSV: '{written_to}' closed ({processed} rows processed).") overall_elapsed = time.monotonic() - overall_start mins, secs = divmod(int(overall_elapsed), 60) @@ -1027,10 +644,12 @@ def main( args = _parse_args() try: main( - mode=EmbeddingGenerativeMode(args.mode), + model_type=ModelType(args.mode), + model_list_file=args.model_list_file, write_to_csv=args.write_to_csv, + fetch=args.fetch, top_k=args.top_k, - model_list_file=args.model_list_file, + max_params=args.max_params, ) except HardwareExceptionAbortError as e: # Non-zero exit so CI / GHA scheduled runs can alert. main()'s diff --git a/tests/test_table_schema.py b/tests/test_table_schema.py new file mode 100644 index 00000000..19bd9499 --- /dev/null +++ b/tests/test_table_schema.py @@ -0,0 +1,197 @@ +"""Tests for the weekly-scan table schema. + +Two things are pinned here. + +**The DDL and ``TABLE_COLUMNS`` agree, in order.** ``ClickHouseResultSink`` +buffers each row as a positional list and passes ``column_names=TABLE_COLUMNS`` +to a bulk insert, so a mismatch against the CREATE TABLE body does not raise — +it writes values into the wrong columns. Three comments in the source used to +warn about keeping them in sync; this asserts it instead. + +**``table_schema`` stays a leaf.** Importing it must not pull in +``clickhouse_connect`` or ``dotenv``, because ``csv_sink`` imports it and +``--write-to-csv`` is meant to work on a host with neither installed. + +Run with ``pytest --noconftest`` (no torch needed) — see +``test_weekly_prefilter.py`` for the same pattern. +""" + +from __future__ import annotations + +import re +from pathlib import Path + +from tests.spyre.weekly_generation.table_schema import ( + DATABASE, + EMBEDDING_CREATE_TABLE_SQL, + EMBEDDING_TABLE_NAME, + GENERATIVE_CREATE_TABLE_SQL, + GENERATIVE_TABLE_NAME, + TABLE_COLUMNS, +) + + +def _columns_in_ddl(ddl: str) -> list[str]: + """The column names from a CREATE TABLE body, in declaration order. + + Takes the parenthesised block that opens after the table name and closes + before ENGINE — matched by paren depth, since the ``ORDER BY (...)`` tuple + later in the statement also contains parentheses. + """ + start = ddl.index("(") + depth = 0 + for index in range(start, len(ddl)): + if ddl[index] == "(": + depth += 1 + elif ddl[index] == ")": + depth -= 1 + if depth == 0: + body = ddl[start + 1 : index] + break + else: # pragma: no cover - unbalanced DDL is a syntax error, not a test case + raise AssertionError("unbalanced parentheses in the DDL") + + names: list[str] = [] + for line in body.splitlines(): + stripped = line.strip().rstrip(",") + if not stripped: + continue + match = re.match(r"^(\w+)\s+\S", stripped) + assert match, f"could not parse a column name from {stripped!r}" + names.append(match.group(1)) + return names + + +class TestDdlMatchesTableColumns: + """A positional bulk insert makes any drift here a silent data-corruption bug.""" + + def test_embedding_ddl_columns_match_in_order(self) -> None: + assert _columns_in_ddl(EMBEDDING_CREATE_TABLE_SQL) == list(TABLE_COLUMNS) + + def test_generative_ddl_columns_match_in_order(self) -> None: + assert _columns_in_ddl(GENERATIVE_CREATE_TABLE_SQL) == list(TABLE_COLUMNS) + + def test_the_two_tables_share_one_shape(self) -> None: + """Both model types write the same 14 columns; only the table name differs.""" + assert _columns_in_ddl(EMBEDDING_CREATE_TABLE_SQL) == _columns_in_ddl( + GENERATIVE_CREATE_TABLE_SQL + ) + + def test_column_names_are_unique(self) -> None: + assert len(set(TABLE_COLUMNS)) == len(TABLE_COLUMNS) + + def test_the_parser_would_notice_a_reordering(self) -> None: + """Guard the guard: a swapped pair in the DDL must fail the comparison. + + Without this, a parser that silently returned [] or a sorted list would + make the tests above pass against any DDL at all. + """ + swapped = EMBEDDING_CREATE_TABLE_SQL.replace( + " model_name String,\n config_class String,", + " config_class String,\n model_name String,", + ) + assert swapped != EMBEDDING_CREATE_TABLE_SQL, "the replace found no match" + assert _columns_in_ddl(swapped) != list(TABLE_COLUMNS) + + +class TestDdlShape: + def test_replacing_merge_tree_on_snapshot_date(self) -> None: + """Same-day duplicates must collapse on merge. + + The weekly scan deliberately prefers writing a possible duplicate over + dropping a result it was asked to produce; that trade-off is only safe + because of this engine choice. + """ + for ddl in (EMBEDDING_CREATE_TABLE_SQL, GENERATIVE_CREATE_TABLE_SQL): + assert "ReplacingMergeTree(snapshot_date)" in ddl + assert "ORDER BY (model_name, snapshot_date)" in ddl + + def test_is_idempotent(self) -> None: + """The sink runs this whenever the table is missing.""" + for ddl in (EMBEDDING_CREATE_TABLE_SQL, GENERATIVE_CREATE_TABLE_SQL): + assert "CREATE TABLE IF NOT EXISTS" in ddl + + def test_each_ddl_targets_its_own_qualified_table(self) -> None: + assert f"{DATABASE}.{EMBEDDING_TABLE_NAME}" in EMBEDDING_CREATE_TABLE_SQL + assert f"{DATABASE}.{GENERATIVE_TABLE_NAME}" in GENERATIVE_CREATE_TABLE_SQL + assert EMBEDDING_TABLE_NAME != GENERATIVE_TABLE_NAME + + def test_nullable_columns_are_the_optional_ones(self) -> None: + """added_date/failure_category/error are None for some rows; the rest never are. + + ``ClickHouseResultSink`` coerces failure_category/error to '' before + inserting, so this documents the DDL's intent rather than the sink's + behaviour — but added_date genuinely arrives as None. + """ + nullable = set( + re.findall(r"^\s*(\w+)\s+Nullable\(", EMBEDDING_CREATE_TABLE_SQL, re.M) + ) + assert nullable == {"added_date", "failure_category", "error"} + + +class TestTableSchemaIsALeafModule: + """``--write-to-csv`` must not need the ClickHouse driver or a .env file.""" + + def test_imports_nothing_outside_the_standard_library(self) -> None: + import ast + + import tests.spyre.weekly_generation.table_schema as ts + + # Located via the module's own __file__ so the test does not depend on + # pytest's cwd. + src = Path(ts.__file__).read_text() + roots: set[str] = set() + for node in ast.walk(ast.parse(src)): + if isinstance(node, ast.Import): + roots.update(alias.name.split(".")[0] for alias in node.names) + elif isinstance(node, ast.ImportFrom) and node.module: + roots.add(node.module.split(".")[0]) + # __future__ only; anything else means the leaf has grown a dependency. + assert roots <= {"__future__"}, f"table_schema gained imports: {roots}" + + def test_importing_it_does_not_load_the_driver(self) -> None: + """Fresh subprocess: neither module may appear in sys.modules afterwards.""" + import subprocess + import sys + + program = ( + "import sys; " + "import tests.spyre.weekly_generation.table_schema as ts; " + "assert ts.TABLE_COLUMNS, 'schema did not load'; " + "leaked = [m for m in ('clickhouse_connect', 'dotenv') " + " if m in sys.modules]; " + "assert not leaked, f'table_schema pulled in {leaked}'; " + "print('clean')" + ) + proc = subprocess.run( + [sys.executable, "-c", program], + capture_output=True, + text=True, + cwd=Path(__file__).resolve().parent.parent, + ) + assert proc.returncode == 0, proc.stderr + assert "clean" in proc.stdout + + +class TestClickhouseDbReExports: + """Existing importers of clickhouse_db keep working after the split.""" + + def test_re_exports_are_the_same_objects(self) -> None: + import tests.spyre.weekly_generation.clickhouse_db as db + import tests.spyre.weekly_generation.table_schema as ts + + for name in ( + "DATABASE", + "EMBEDDING_TABLE_NAME", + "GENERATIVE_TABLE_NAME", + "TABLE_COLUMNS", + "EMBEDDING_CREATE_TABLE_SQL", + "GENERATIVE_CREATE_TABLE_SQL", + ): + assert getattr(db, name) is getattr(ts, name), name + + def test_still_exposes_the_client_helpers(self) -> None: + import tests.spyre.weekly_generation.clickhouse_db as db + + assert callable(db.get_client) + assert callable(db.table_exists) diff --git a/tests/test_weekly_prefilter.py b/tests/test_weekly_prefilter.py new file mode 100644 index 00000000..51120196 --- /dev/null +++ b/tests/test_weekly_prefilter.py @@ -0,0 +1,813 @@ +"""Unit tests for the weekly-scan pre-filter. + +Lives at the ``tests/`` root rather than under ``tests/spyre/`` (Spyre hardware) +or ``tests/cpu/`` (torch fixtures), following ``test_adapter_coverage.py``: run +with ``pytest --noconftest`` so the torch-importing root conftest is bypassed and +only ``pytest`` itself is needed. See test_pull_request.yaml's adapter-coverage +job for the same pattern. + +``prefilter_models`` takes no sink — it decides and writes nothing — so the +pre-filter tests need no storage backend at all, which is what keeps them +runnable on an interpreter with no database driver (the concrete sinks import +``clickhouse_connect`` transitively). Tests that do need a real sink are gated on +``requires_sink`` and import one inside the test body. +""" + +from __future__ import annotations + +import importlib.util +import inspect +from datetime import date +from pathlib import Path + +import pytest + +from tests.spyre.weekly_generation.failure_categories import ( + FAILURE_CATEGORY_MODEL_TOO_LARGE, + FAILURE_CATEGORY_MOE, + FAILURE_CATEGORY_NOT_IMPLEMENTED_ADAPTER, + MAX_NUMBER_PARAMS, +) +from tests.spyre.weekly_generation.model_prefilter import prefilter_models +from tests.spyre.weekly_generation.model_type import ModelType + +# The concrete sinks pull in clickhouse_connect transitively; the pre-filter and +# its stub sink must not. Gating at class level (rather than importing +# CsvResultSink at module scope) keeps the pure tests runnable with no DB driver. +requires_sink = pytest.mark.skipif( + importlib.util.find_spec("clickhouse_connect") is None, + reason="clickhouse_connect not installed (the ClickHouse sink imports it)", +) + +# Subprocess tests below need the repo root as cwd to import from tests/. +_REPO_ROOT = Path(__file__).resolve().parent.parent + + +def _filter( + rows: list[dict], + *, + max_params: int = MAX_NUMBER_PARAMS, +): + """Run the pre-filter over *rows*.""" + return prefilter_models(rows, max_params=max_params) + + +def _row(model_id: str, **overrides: object) -> dict: + """A catalog row that passes every filter unless *overrides* say otherwise.""" + row: dict = { + "model_id": model_id, + "downloads": 100, + "parameters": 1_000_000_000, + "is_supported": True, + "is_moe": False, + "config_class": "LlamaConfig", + "model_type": "llama", + "architectures": "LlamaForCausalLM", + } + row.update(overrides) + return row + + +class TestFilterBranches: + def test_clean_row_is_kept(self) -> None: + result = _filter([_row("org/ok")]) + assert [r["model_id"] for r in result.keep] == ["org/ok"] + assert result.skipped == [] + + def test_unsupported_config_class(self) -> None: + result = _filter([_row("org/unsup", is_supported=False)]) + assert result.keep == [] + assert len(result.skipped) == 1 + assert result.skipped[0].failure_category == ( + FAILURE_CATEGORY_NOT_IMPLEMENTED_ADAPTER + ) + + def test_missing_is_supported_is_not_treated_as_unsupported(self) -> None: + """The check is ``is False``: unknown support status still gets evaluated. + + A missing key or None means the fetcher could not determine the config + class, which is not the same as knowing there is no adapter for it. + """ + assert [ + r["model_id"] for r in _filter([_row("org/a", is_supported=None)]).keep + ] == ["org/a"] + no_key = _row("org/b") + del no_key["is_supported"] + assert [r["model_id"] for r in _filter([no_key]).keep] == ["org/b"] + + def test_too_large(self) -> None: + result = _filter([_row("org/huge", parameters=999)], max_params=100) + assert result.keep == [] + assert result.skipped[0].failure_category == FAILURE_CATEGORY_MODEL_TOO_LARGE + + def test_exactly_at_the_limit_is_kept(self) -> None: + """The guard is ``>`` not ``>=`` — a model exactly at the cap is fine.""" + result = _filter([_row("org/edge", parameters=100)], max_params=100) + assert [r["model_id"] for r in result.keep] == ["org/edge"] + + def test_moe(self) -> None: + result = _filter([_row("org/moe", is_moe=True)]) + assert result.keep == [] + assert result.skipped[0].failure_category == FAILURE_CATEGORY_MOE + + +class TestPrefilterIsPure: + """``prefilter_models`` decides and writes nothing. + + Recording the terminal rows is ``write_skipped_rows``' job, called separately + by ``fetch_and_filter``. That split is what lets this function be called with + no sink at all — and these tests run with no database driver installed. + """ + + def test_takes_no_sink(self) -> None: + """A sink parameter would reintroduce the dependency this split removed.""" + params = inspect.signature(prefilter_models).parameters + assert "sink" not in params + assert set(params) == {"models", "max_params"} + + def test_every_model_lands_in_exactly_one_list(self) -> None: + rows = [ + _row("org/a"), + _row("org/b", is_supported=False), + _row("org/c", is_moe=True), + ] + result = _filter(rows) + assert len(result.keep) + len(result.skipped) == len(rows) + assert [r["model_id"] for r in result.keep] == ["org/a"] + assert [s.row["model_id"] for s in result.skipped] == ["org/b", "org/c"] + + +class TestPrecedence: + def test_unsupported_wins_over_moe(self) -> None: + """Both apply; the reported category is the first check that fires.""" + result = _filter([_row("org/both", is_supported=False, is_moe=True)]) + assert result.skipped[0].failure_category == ( + FAILURE_CATEGORY_NOT_IMPLEMENTED_ADAPTER + ) + + +class TestParameterCoercion: + @pytest.mark.parametrize("value", [None, "", "not-a-number", object()]) + def test_unknown_size_is_kept_for_the_in_worker_backstop( + self, value: object + ) -> None: + """Unsizable rows must NOT be treated as zero-parameter or as too large. + + There is no worker-side size check to fall back on, so such a model is + judged by whether it actually loads. That matches the behaviour before + this filter moved upstream: the parent's guard skipped unsizable rows + too, and no in-worker check ever existed despite a comment claiming one. + """ + result = _filter([_row("org/unknown", parameters=value)], max_params=100) + assert [r["model_id"] for r in result.keep] == ["org/unknown"] + + def test_numeric_string_is_compared_as_a_number(self) -> None: + """A stringified count must not be compared lexicographically.""" + result = _filter([_row("org/str", parameters="999")], max_params=100) + assert result.keep == [] + assert result.skipped[0].failure_category == FAILURE_CATEGORY_MODEL_TOO_LARGE + + def test_zero_parameters_is_kept(self) -> None: + result = _filter([_row("org/zero", parameters=0)]) + assert [r["model_id"] for r in result.keep] == ["org/zero"] + + +class TestOrderAndTallies: + def test_keep_preserves_input_order(self) -> None: + """The tier router and shard chunker rely on downloads-descending order.""" + rows = [_row(f"org/m{i}", downloads=1000 - i) for i in range(20)] + rows[3]["is_moe"] = True + rows[11]["is_supported"] = False + result = _filter(rows) + + kept = [r["model_id"] for r in result.keep] + assert kept == sorted(kept, key=lambda m: -(1000 - int(m.split("m")[1]))) + assert "org/m3" not in kept and "org/m11" not in kept + assert len(kept) == 18 + + def test_counts_reconcile_with_the_input(self) -> None: + rows = [ + _row("org/a"), + _row("org/b", is_supported=False), + _row("org/c", is_moe=True), + _row("org/d", parameters=10**15), + _row("org/e"), + ] + result = _filter(rows) + counts = result.counts + assert counts["keep"] == 2 + assert counts[FAILURE_CATEGORY_NOT_IMPLEMENTED_ADAPTER] == 1 + assert counts[FAILURE_CATEGORY_MOE] == 1 + assert counts[FAILURE_CATEGORY_MODEL_TOO_LARGE] == 1 + assert sum(counts.values()) == len(rows) + + def test_empty_input(self) -> None: + result = _filter([]) + assert result.keep == [] + assert result.counts == {"keep": 0} + + def test_rows_are_returned_by_identity(self) -> None: + """The same dict objects come back, so a caller's pop() still applies.""" + row = _row("org/same") + result = _filter([row]) + assert result.keep[0] is row + + +class TestModelType: + """ModelType must format as its bare value, not as ``ModelType.GENERATIVE``. + + It reaches shard filenames, the GHA matrix, and operator-facing log lines. A + plain ``(str, Enum)`` inherits ``Enum.__str__`` and renders the qualified + name, which is how shard files were once written as + ``ModelType.GENERATIVE-x1-shard-000.json``. + """ + + @pytest.mark.parametrize("member", list(ModelType)) + def test_str_and_format_are_the_value(self, member: ModelType) -> None: + assert str(member) == member.value + assert f"{member}" == member.value + assert f"{member}-x1-shard-000.json" == f"{member.value}-x1-shard-000.json" + + @pytest.mark.parametrize("member", list(ModelType)) + def test_json_serializable_as_the_value(self, member: ModelType) -> None: + """The GHA matrix is emitted as JSON and read back by the workflow.""" + import json + + assert json.dumps({"mode": member}) == f'{{"mode": "{member.value}"}}' + + def test_round_trips_from_the_cli_string(self) -> None: + """--mode / --model-type pass the raw string through ModelType(...).""" + assert ModelType("generative") is ModelType.GENERATIVE + assert ModelType("embedding") is ModelType.EMBEDDING + + def test_values_match_the_cli_choices(self) -> None: + assert {m.value for m in ModelType} == {"generative", "embedding"} + + +@requires_sink +class TestWriteSkippedRows: + def test_writes_one_row_per_skipped_model(self, tmp_path) -> None: + from tests.spyre.weekly_generation.sink.csv_sink import CsvResultSink + from tests.spyre.weekly_generation.skip_writer import write_skipped_rows + + csv_path = tmp_path / "out.csv" + rows = [ + _row("org/unsup", is_supported=False), + _row("org/moe", is_moe=True), + ] + result = _filter(rows) + today = date.today() + + with CsvResultSink(path=csv_path) as sink: + written = write_skipped_rows( + sink, result.skipped, snapshot_date=today, verbose=False + ) + + assert written == 2 + text = csv_path.read_text() + assert "org/unsup" in text and "org/moe" in text + assert FAILURE_CATEGORY_NOT_IMPLEMENTED_ADAPTER in text + assert FAILURE_CATEGORY_MOE in text + + def test_field_mapping_matches_the_replaced_add_entry_calls(self, tmp_path) -> None: + """Pin the 14 column values the three deleted branches used to write.""" + import csv as _csv + + from tests.spyre.weekly_generation.sink.csv_sink import CsvResultSink + from tests.spyre.weekly_generation.skip_writer import write_skipped_rows + + csv_path = tmp_path / "out.csv" + today = date.today() + result = _filter( + [ + _row( + "org/unsup", + is_supported=False, + downloads=42, + parameters=7, + model_type="mistral", + architectures="MistralForCausalLM", + config_class="MistralConfig", + ) + ] + ) + with CsvResultSink(path=csv_path) as sink: + write_skipped_rows(sink, result.skipped, snapshot_date=today, verbose=False) + + written_row = next(iter(_csv.DictReader(csv_path.open()))) + assert written_row["model_name"] == "org/unsup" + assert written_row["config_class"] == "MistralConfig" + assert written_row["adapter_name"] == "" + assert written_row["added_date"] == "" + assert written_row["snapshot_date"] == str(today) + assert written_row["verified_on_cpu"] == "False" + assert written_row["verified_on_gpu"] == "False" + assert written_row["verified_on_spyre"] == "False" + assert written_row["num_downloads"] == "42" + assert written_row["family"] == "mistral" + assert written_row["architecture"] == "MistralForCausalLM" + assert written_row["parameters_number"] == "7" + assert written_row["failure_category"] == ( + FAILURE_CATEGORY_NOT_IMPLEMENTED_ADAPTER + ) + assert written_row["error"] == "" + + def test_empty_skipped_list_writes_nothing(self, tmp_path) -> None: + from tests.spyre.weekly_generation.sink.csv_sink import CsvResultSink + from tests.spyre.weekly_generation.skip_writer import write_skipped_rows + + csv_path = tmp_path / "out.csv" + with CsvResultSink(path=csv_path) as sink: + assert write_skipped_rows(sink, [], snapshot_date=date.today()) == 0 + + +@requires_sink +class TestFetchAndFilter: + """The shared entry point: fetch, filter, record verdicts, return the rest. + + Both producers go through this, so the ownership contract it keeps matters: + it must write the terminal rows and must NOT close the sink, because + ``weekly_test.main`` hands in the sink it uses for the whole run and keeps + writing evaluation results to it afterwards. + """ + + @pytest.fixture + def _stub_fetcher(self, monkeypatch): + """Replace the Hub fetchers so no network call happens.""" + import tests.spyre.weekly_generation.model_fetcher as model_fetcher + + rows = [ + _row("org/keep"), + _row("org/moe", is_moe=True), + _row("org/keep-too"), + ] + + def _fake(limit: int, **_kw) -> list[dict]: + # Fresh dicts per call, and one carrying the non-serializable + # model_info the real fetcher attaches, so the pop is exercised. + out = [dict(r) for r in rows] + out[0]["model_info"] = object() + return out[:limit] + + monkeypatch.setattr( + model_fetcher, + "all_fetchers", + {ModelType.GENERATIVE: _fake, ModelType.EMBEDDING: _fake}, + ) + return rows + + def test_keeps_survivors_records_verdicts_and_leaves_sink_open( + self, tmp_path, _stub_fetcher + ) -> None: + import csv as _csv + + from tests.spyre.weekly_generation.model_prefilter import fetch_and_filter + from tests.spyre.weekly_generation.sink.csv_sink import CsvResultSink + + path = tmp_path / "out.csv" + today = date.today() + sink = CsvResultSink(path=path) + + kept = fetch_and_filter( + model_type=ModelType.GENERATIVE, + snapshot_date=today, + top_k=10, + sink=sink, + max_params=MAX_NUMBER_PARAMS, + ) + + # The MoE model is dropped and recorded; the other two survive. + assert [r["model_id"] for r in kept] == ["org/keep", "org/keep-too"] + + # The sink must still be writable — main() writes every evaluation + # result through this same object after fetch_and_filter returns. + with sink: + sink.add_entry( + model_name="org/keep", + config_class="LlamaConfig", + adapter_name="hf_llama", + added_date=None, + snapshot_date=today, + verified_on_cpu=True, + verified_on_gpu=False, + verified_on_spyre=True, + num_downloads=100, + family="llama", + architecture="LlamaForCausalLM", + parameters_number=1, + failure_category=None, + error=None, + ) + + written = list(_csv.DictReader(path.open())) + assert [r["model_name"] for r in written] == ["org/moe", "org/keep"] + assert written[0]["failure_category"] == FAILURE_CATEGORY_MOE + + def test_drops_the_unserializable_model_info_field( + self, tmp_path, _stub_fetcher + ) -> None: + """Shards are JSON-dumped, and ModelInfo would break that dump.""" + import json + + from tests.spyre.weekly_generation.model_prefilter import fetch_and_filter + from tests.spyre.weekly_generation.sink.csv_sink import CsvResultSink + + with CsvResultSink(path=tmp_path / "o.csv") as sink: + kept = fetch_and_filter( + model_type=ModelType.GENERATIVE, + snapshot_date=date.today(), + top_k=10, + sink=sink, + max_params=MAX_NUMBER_PARAMS, + ) + assert all("model_info" not in r for r in kept) + json.dumps(kept) # must not raise + + def test_only_terminal_verdicts_get_a_row(self, tmp_path, _stub_fetcher) -> None: + """Models handed on for evaluation are not recorded here. + + Their row comes later, from the evaluation loop, and writing one now would + double-count them. + """ + import csv as _csv + + from tests.spyre.weekly_generation.model_prefilter import fetch_and_filter + from tests.spyre.weekly_generation.sink.csv_sink import CsvResultSink + + path = tmp_path / "o.csv" + with CsvResultSink(path=path) as sink: + kept = fetch_and_filter( + model_type=ModelType.GENERATIVE, + snapshot_date=date.today(), + top_k=10, + sink=sink, + max_params=MAX_NUMBER_PARAMS, + ) + + assert [r["model_id"] for r in kept] == ["org/keep", "org/keep-too"] + assert [r["model_name"] for r in _csv.DictReader(path.open())] == ["org/moe"] + + def test_max_params_is_honoured(self, tmp_path, _stub_fetcher) -> None: + from tests.spyre.weekly_generation.model_prefilter import fetch_and_filter + from tests.spyre.weekly_generation.sink.csv_sink import CsvResultSink + + with CsvResultSink(path=tmp_path / "o.csv") as sink: + kept = fetch_and_filter( + model_type=ModelType.GENERATIVE, + snapshot_date=date.today(), + top_k=10, + sink=sink, + max_params=1, # every stub row is 1e9 params + ) + assert kept == [] + + +def _fake_sink_cls(): + """Build a minimal concrete ResultSink subclass, importing the base lazily. + + ``add_entry``'s behaviour lives entirely in the base class, so testing it + needs a sink, not a *storage backend*. This fake records the ``_insert_entry`` + calls it receives, which is exactly what the base class's contract is about, + and defining it inside a function keeps ``result_sink`` out of this module's + import-time dependencies. + """ + from tests.spyre.weekly_generation.sink.result_sink import ResultSink + + class _FakeSink(ResultSink): + def __init__(self) -> None: + self.written: list[str] = [] + + def _insert_entry(self, *, model_name, **_rest) -> None: + self.written.append(model_name) + + return _FakeSink + + +def _add(sink, name: str) -> bool: + return sink.add_entry( + model_name=name, + config_class="LlamaConfig", + adapter_name="hf_llama", + added_date=None, + snapshot_date=date.today(), + verified_on_cpu=True, + verified_on_gpu=False, + verified_on_spyre=True, + num_downloads=1, + family="llama", + architecture="LlamaForCausalLM", + parameters_number=1, + failure_category=None, + error=None, + ) + + +@requires_sink +class TestAddEntryAlwaysWrites: + """add_entry records every row it is handed. + + Which models to evaluate is decided upstream, in ``model_prefilter``, so a row + reaching add_entry is one the caller already decided to record. Filtering here + would silently drop results a run was asked to produce, leaving its row count + lower than its input with no accounting. + """ + + @pytest.fixture(autouse=True) + def _bind(self): + self.FakeSink = _fake_sink_cls() + + def test_writes_the_row_it_is_handed(self) -> None: + sink = self.FakeSink() + _add(sink, "org/fresh") + assert sink.written == ["org/fresh"] + + def test_repeated_writes_of_the_same_model_all_land(self) -> None: + """One row per call, so a run's count matches the models it handled.""" + sink = self.FakeSink() + _add(sink, "org/dup") + _add(sink, "org/dup") + assert sink.written == ["org/dup", "org/dup"] + + def test_the_abc_exposes_no_filtering_hook(self) -> None: + """A sink decides nothing about which models to run.""" + sink = self.FakeSink() + assert not hasattr(sink, "should_insert_row") + + def test_empty_model_name_is_still_rejected(self) -> None: + sink = self.FakeSink() + with pytest.raises(ValueError, match="model_name"): + _add(sink, " ") + + +@requires_sink +class TestSinkConstructors: + """Pin each sink's constructor signature and where it is imported from. + + The sinks moved out of ``result_sink`` into the ``sink`` package, and the + ClickHouse sink's first parameter was renamed from ``embedding_generative`` + to ``model_type`` (now a ``ModelType``, not a parallel enum). Both are + load-bearing for ``sink_factory``, which passes them by keyword. + """ + + @pytest.mark.parametrize( + "module_path, cls_name, expected_positional", + [ + ( + "tests.spyre.weekly_generation.sink.csv_sink", + "CsvResultSink", + ["path"], + ), + ( + "tests.spyre.weekly_generation.sink.clickhouse_sink", + "ClickHouseResultSink", + ["model_type"], + ), + ], + ) + def test_constructor_signatures( + self, module_path: str, cls_name: str, expected_positional: list[str] + ) -> None: + """No stray keyword-only params, and the positional order is unchanged.""" + import importlib + + module = importlib.import_module(module_path) + params = inspect.signature(getattr(module, cls_name).__init__).parameters + positional = [ + name + for name, p in params.items() + if name != "self" and p.kind in (p.POSITIONAL_ONLY, p.POSITIONAL_OR_KEYWORD) + ] + assert positional == expected_positional + assert [p for p in params.values() if p.kind is p.KEYWORD_ONLY] == [] + + def test_sinks_are_no_longer_exported_from_result_sink(self) -> None: + """result_sink holds only the ABC now; the factory is the way in.""" + import tests.spyre.weekly_generation.sink.result_sink as rs + + assert not hasattr(rs, "CsvResultSink") + assert not hasattr(rs, "ClickHouseResultSink") + + def test_both_sinks_implement_the_abc(self) -> None: + from tests.spyre.weekly_generation.sink.clickhouse_sink import ( + ClickHouseResultSink, + ) + from tests.spyre.weekly_generation.sink.csv_sink import CsvResultSink + from tests.spyre.weekly_generation.sink.result_sink import ResultSink + + assert issubclass(CsvResultSink, ResultSink) + assert issubclass(ClickHouseResultSink, ResultSink) + + +@requires_sink +class TestSinkFactory: + """``create_sink`` picks the backend and, for CSV, the per-model-type path.""" + + def test_write_to_csv_yields_a_csv_sink_with_a_suffixed_path( + self, tmp_path + ) -> None: + from tests.spyre.weekly_generation.sink.csv_sink import CsvResultSink + from tests.spyre.weekly_generation.sink.sink_factory import create_sink + + base = tmp_path / "verdicts.csv" + with create_sink( + model_type=ModelType.EMBEDDING, + write_to_csv=base, + ) as sink: + assert isinstance(sink, CsvResultSink) + + # One file per model type, since a single run can cover both — and named + # from the enum's *value*, not "verdicts-ModelType.EMBEDDING.csv". + assert (tmp_path / "verdicts-embedding.csv").exists() + assert not base.exists() + + def test_csv_path_for_matches_the_file_the_sink_creates(self, tmp_path) -> None: + """The helper weekly_test logs from must agree with what got written. + + ``main`` reports its output path via ``csv_path_for`` while the sink was + built by ``create_sink``; if the two computed the name separately they + could drift, and the run would name a file that does not exist. + """ + from tests.spyre.weekly_generation.sink.sink_factory import ( + create_sink, + csv_path_for, + ) + + base = tmp_path / "results.csv" + expected = csv_path_for(base, ModelType.GENERATIVE) + with create_sink( + model_type=ModelType.GENERATIVE, + write_to_csv=base, + ): + pass + assert expected.exists() + assert expected.name == "results-generative.csv" + + def test_each_model_type_gets_its_own_file(self, tmp_path) -> None: + from tests.spyre.weekly_generation.sink.sink_factory import create_sink + + base = tmp_path / "v.csv" + for model_type in ModelType: + with create_sink( + model_type=model_type, + write_to_csv=base, + ): + pass + assert {p.name for p in tmp_path.iterdir()} == { + "v-generative.csv", + "v-embedding.csv", + } + + def test_no_csv_path_builds_the_clickhouse_sink(self, monkeypatch) -> None: + """Without --write-to-csv the factory must reach for ClickHouse. + + The constructor is stubbed out: instantiating the real one would connect + and create its table, and what is under test is the branch, not the + driver. Patched on ``clickhouse_sink`` rather than on the factory, + because the factory imports it lazily inside the function body — see + ``test_csv_branch_runs_with_no_clickhouse_driver_installed`` for why. + """ + import tests.spyre.weekly_generation.sink.clickhouse_sink as ch_module + from tests.spyre.weekly_generation.sink.sink_factory import create_sink + + seen: dict = {} + + class _FakeClickHouseSink: + def __init__(self, *, model_type) -> None: + seen["model_type"] = model_type + + monkeypatch.setattr(ch_module, "ClickHouseResultSink", _FakeClickHouseSink) + sink = create_sink(model_type=ModelType.GENERATIVE, write_to_csv=None) + assert isinstance(sink, _FakeClickHouseSink) + assert seen == {"model_type": ModelType.GENERATIVE} + + def test_csv_branch_does_not_connect_to_clickhouse( + self, tmp_path, monkeypatch + ) -> None: + """--write-to-csv must not open a DB connection. + + ``get_client`` is the only thing that reaches the network, so binding it + to a raiser is the check that matters: the CSV branch must complete + without it. + """ + import tests.spyre.weekly_generation.clickhouse_db as db + from tests.spyre.weekly_generation.sink.sink_factory import create_sink + + def _boom(*_a, **_k): + raise AssertionError("--write-to-csv must not connect to ClickHouse") + + monkeypatch.setattr(db, "get_client", _boom) + + with create_sink( + model_type=ModelType.GENERATIVE, + write_to_csv=tmp_path / "v.csv", + ) as sink: + _add(sink, "org/x") + assert (tmp_path / "v-generative.csv").exists() + + def test_csv_branch_runs_with_no_clickhouse_driver_installed( + self, tmp_path + ) -> None: + """--write-to-csv must work on a host with neither the driver nor dotenv. + + That is the whole point of the flag, and it is easy to lose: taking + ``TABLE_COLUMNS`` from ``clickhouse_db`` instead of ``table_schema``, or + hoisting the ClickHouseResultSink import to ``sink_factory``'s module + scope, would each reintroduce the dependency. + + Runs in a subprocess with both modules blocked at the finder level. + In-process patching cannot test this — ``sink_factory`` and its imports + are already in ``sys.modules`` by then, so blocking them later proves + nothing. + """ + import subprocess + import sys + import textwrap + + program = textwrap.dedent(""" + import sys + + BLOCKED = ("clickhouse_connect", "dotenv") + + class _Blocker: + def find_spec(self, name, path=None, target=None): + if name.split(".")[0] in BLOCKED: + raise ModuleNotFoundError(f"No module named {name!r}") + return None + + sys.meta_path.insert(0, _Blocker()) + + from datetime import date + from pathlib import Path + + from tests.spyre.weekly_generation.model_type import ModelType + from tests.spyre.weekly_generation.sink.sink_factory import create_sink + + with create_sink( + model_type=ModelType.GENERATIVE, + write_to_csv=Path(sys.argv[1]) / "v.csv", + ) as sink: + sink.add_entry( + model_name="org/x", config_class="LlamaConfig", + adapter_name="hf_llama", added_date=None, + snapshot_date=date.today(), verified_on_cpu=True, + verified_on_gpu=False, verified_on_spyre=True, + num_downloads=1, family="llama", + architecture="LlamaForCausalLM", parameters_number=1, + failure_category=None, error=None, + ) + print("OK") + """) + proc = subprocess.run( + [sys.executable, "-c", program, str(tmp_path)], + capture_output=True, + text=True, + cwd=_REPO_ROOT, + ) + assert ( + proc.returncode == 0 + ), f"CSV branch failed without the driver installed:\n{proc.stderr}" + assert "OK" in proc.stdout + written = (tmp_path / "v-generative.csv").read_text() + assert "model_name" in written and "org/x" in written + + +@requires_sink +class TestCsvSinkIsWriteOnly: + """The CSV sink writes one run to a new file and never reads one back.""" + + def test_refuses_a_non_empty_existing_file(self, tmp_path) -> None: + from tests.spyre.weekly_generation.sink.csv_sink import CsvResultSink + + existing = tmp_path / "already.csv" + existing.write_text("model_name\norg/x\n") + with pytest.raises(FileExistsError, match="already exists"): + CsvResultSink(path=existing) + + def test_accepts_an_empty_existing_file(self, tmp_path) -> None: + """A zero-byte file is not a previous run's results.""" + from tests.spyre.weekly_generation.sink.csv_sink import CsvResultSink + + empty = tmp_path / "empty.csv" + empty.touch() + sink = CsvResultSink(path=empty) + sink.close() + assert "model_name" in empty.read_text() + + def test_creates_parent_directories(self, tmp_path) -> None: + from tests.spyre.weekly_generation.sink.csv_sink import CsvResultSink + + sink = CsvResultSink(path=tmp_path / "a" / "b" / "out.csv") + sink.close() + assert (tmp_path / "a" / "b" / "out.csv").exists() + + def test_writes_every_row_including_repeats(self, tmp_path) -> None: + import csv as _csv + + from tests.spyre.weekly_generation.sink.csv_sink import CsvResultSink + + path = tmp_path / "o.csv" + sink = CsvResultSink(path=path) + _add(sink, "org/x") + _add(sink, "org/x") + sink.close() + assert len(list(_csv.DictReader(path.open()))) == 2 diff --git a/utils/hf_model_catalog.py b/utils/hf_model_catalog.py index 9b8a511c..7a14a64a 100644 --- a/utils/hf_model_catalog.py +++ b/utils/hf_model_catalog.py @@ -373,9 +373,15 @@ def _safe_filter(model: ModelInfo) -> bool: logging.warning("filter_fn failed for %s: %s", model.id, e) return False + timings: dict[str, float] = {} + t_total = time.perf_counter() + + t0 = time.perf_counter() candidates: list[ModelInfo] = list(fetch_fn(limit)) + timings["fetch (HF list_models)"] = time.perf_counter() - t0 print(f"Retrieved {len(candidates)} raw {label} candidates.") + t0 = time.perf_counter() with ThreadPoolExecutor(max_workers=16) as ex: keep_flags: list[bool] = list( tqdm( @@ -385,6 +391,7 @@ def _safe_filter(model: ModelInfo) -> bool: ) ) models: list[ModelInfo] = [m for m, keep in zip(candidates, keep_flags) if keep] + timings["filter (filter_fn)"] = time.perf_counter() - t0 print(f"Kept {len(models)} {label} models after filtering.") models = models[:limit] @@ -406,6 +413,7 @@ def _safe_filter(model: ModelInfo) -> bool: tail_head: list[str] = ["is_custom_code", "config_class", "is_supported", "Year"] header: list[str] = base_head + extra_head + tail_head + t0 = time.perf_counter() with ThreadPoolExecutor(max_workers=16) as ex: config_classes: list[str | None] = list( tqdm( @@ -414,7 +422,9 @@ def _safe_filter(model: ModelInfo) -> bool: desc="Fetching config classes", ) ) + timings["config classes (AutoConfig)"] = time.perf_counter() - t0 + t0 = time.perf_counter() rows: list[dict[str, object]] = [] for rank, (m, config_class) in enumerate(zip(models, config_classes), start=1): architectures: list[str] | None = (m.config or {}).get("architectures") @@ -447,12 +457,16 @@ def _safe_filter(model: ModelInfo) -> bool: ) ) + timings["build rows"] = time.perf_counter() - t0 + + t0 = time.perf_counter() if output_csv is not None: print(f"Writing top {len(rows)} to {output_csv}") with open(output_csv, "w", newline="", encoding="utf-8") as f: writer = csv.DictWriter(f, fieldnames=header) writer.writeheader() writer.writerows(rows) + timings["write CSV"] = time.perf_counter() - t0 # Attach the source ModelInfo to each row AFTER the CSV write. It is a # runtime-only field (not serializable, and never part of the schema), @@ -462,8 +476,19 @@ def _safe_filter(model: ModelInfo) -> bool: # is_moe is precomputed here (a pure function of data already fetched — # tags, config.model_type, config.architectures) so callers that need it # don't have to carry the non-serializable ModelInfo object forward. + t0 = time.perf_counter() for row, m in zip(rows, models): row["model_info"] = m row["is_moe"] = is_moe(m) + timings["attach model_info / is_moe"] = time.perf_counter() - t0 + + timings["other"] = (time.perf_counter() - t_total) - sum(timings.values()) + + total = sum(timings.values()) + print(f"\nTiming breakdown for {label} catalog ({total:.1f}s total):") + width = max(len(name) for name in timings) + for name, secs in timings.items(): + share = secs / total if total else 0.0 + print(f" {name:<{width}} {secs:7.2f}s {share:5.1%}") return rows diff --git a/utils/utilities.py b/utils/utilities.py index c2219ed0..c5f657ec 100644 --- a/utils/utilities.py +++ b/utils/utilities.py @@ -4,3 +4,17 @@ def ts() -> str: """Local-time timestamp prefix, e.g. '[2026-07-17 14:32:05]'.""" return datetime.now().strftime("[%Y-%m-%d %H:%M:%S]") + + +def human_bytes(n: float) -> str: + """Format a byte count for logs, e.g. 1536 -> '1.5KB'. + + Binary units (1 KB = 1024 B), one decimal place, no space before the unit. + Saturates at TB rather than continuing to PB — the caller is reporting freed + HuggingFace cache space, which never reaches that scale. + """ + for unit in ("B", "KB", "MB", "GB", "TB"): + if n < 1024 or unit == "TB": + return f"{n:.1f}{unit}" + n /= 1024 + raise AssertionError("unreachable: the TB branch always returns")