diff --git a/CLAUDE.md b/CLAUDE.md index 046ba550..08b8bbb5 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -62,6 +62,7 @@ fc-data --stage 7 --harbor-limit 10 # Smoke t 6. **synthesize_images** — Agent-based Docker build context synthesis (uses env_payload/python_version from stage 4) 7. **harbor_healthcheck** — Run every synthesized container through Harbor's oracle agent, record per-benchmark speedups to `harbor_runs`. Supports local Docker and Daytona via `--harbor-environment`; the row records which one in `harbor_runs.environment`. Local runs are useful for iteration; only Daytona runs gate stage 8. 8. **publish** — Build, verify, and publish Docker images to DockerHub. Only publishes PRs with at least one successful **Daytona** `harbor_runs` row whose `max_speedup >= 1.05`. +9. **scrape_benchmark_source** — For each `(owner, repo)` in `candidate_containers`, check out the repo at its container SHA, AST-parse every ASV-style benchmark function under the repo's `benchmark_dir`, and upsert one row per `(owner, repo, benchmark_without_params)` into `benchmark_codes` for the FormulaCode website's data sync. ### Dataset verification (`dataset/`) @@ -112,7 +113,8 @@ import time (`src/datasmith/__init__.py` → `dotenv.load_dotenv`), so reading - **Existing uses** (non-exhaustive, grep `DATASMITH_` for the full list): `DATASMITH_RL_DEFAULT_PAUSE_S`, `DATASMITH_RL_PAUSE_JITTER_S`, `DATASMITH_RL_MAX_RETRIES`, `DATASMITH_NEIGHBOR_WINDOW_DAYS`, - `DATASMITH_NEIGHBOR_CAP`. + `DATASMITH_NEIGHBOR_CAP`, `DATASMITH_BENCH_SCRAPE_MAX_FILE_BYTES`, + `DATASMITH_BENCH_SCRAPE_DIRS`. ## Supabase @@ -138,7 +140,7 @@ Local vLLM servers (8123, 8124) are exposed via a LiteLLM proxy on `https://mode ### Public read-only access (RLS) -Four tables are readable by the `anon` role: `repositories`, `pull_requests`, `candidate_containers`, `harbor_runs`. Migration `00012_public_read_rls.sql` enables RLS with a `public_read` SELECT policy on those tables; migration `00015_revoke_anon_select.sql` revokes Supabase's default broad `anon` `SELECT` grant and re-grants it only on the four, so every other table returns `permission denied`. The service-role key bypasses both layers, so pipeline processes are unaffected. Public anon access is served on `https://api.formulacode.org` (no Cloudflare Access gate); pipeline operators continue to use `https://db.formulacode.org` with CF Access + service-role key. +Six tables are readable by the `anon` role: `repositories`, `pull_requests`, `candidate_containers`, `harbor_runs`, `benchmark_information`, `benchmark_codes`. Migration `00012_public_read_rls.sql` enables RLS with a `public_read` SELECT policy on the original four; migration `00016_benchmark_information.sql` adds the same policy to `benchmark_information`; migration `00017_benchmark_codes.sql` adds it to `benchmark_codes`. Migration `00015_revoke_anon_select.sql` revokes Supabase's default broad `anon` `SELECT` grant and re-grants it only on the public set, so every other table returns `permission denied`. The service-role key bypasses both layers, so pipeline processes are unaffected. Public anon access is served on `https://api.formulacode.org` (no Cloudflare Access gate); pipeline operators continue to use `https://db.formulacode.org` with CF Access + service-role key. ### Key tables @@ -148,6 +150,8 @@ Four tables are readable by the `anon` role: `repositories`, `pull_requests`, `c | `packages` | Resolved `env_payload` (pinned deps) and `python_version` per commit | Stage 4 | | `candidate_containers` | Successful agent-generated `build_pkg_sh` / `build_run_sh` per SHA | Stage 6 (on success) | | `harbor_runs` | One row per Harbor oracle trial for a synthesized container: `max_speedup`, `geomean_speedup`, `n_benchmarks`, `wallclock_sec`, `reward_payload`, `status`. One-to-many FK on `candidate_containers(owner, repo, sha)`. | Stage 7 | +| `benchmark_information` | Per-benchmark speedup measurements from terminal-bench eval runs: one row per (run, owner/repo/issue, benchmark, agent, model). `speedup` is `(agent/nop)/(oracle/nop)` so 1.0 = parity with the human expert. `benchmark_type` (`time`/`mem`/`peakmem`/`track`) is a generated column derived from the ASV naming convention. Loaded out-of-band via `scripts/load_benchmark_information.py`. | (manual) | +| `benchmark_codes` | One row per `(owner, repo, benchmark_without_params)` carrying the Python source of each ASV benchmark function plus its setup. Joined to `benchmark_information` on `(owner, repo, benchmark_name)` by the FormulaCode website. | Stage 9 | | `error_logs` | Per-attempt synthesis results: agent output, failure stage/return code, error messages | Stage 6 (`Synthesizer._log_attempt`) | | `runner_progress` | Live progress counters (total/completed/failed) per pipeline run | `BaseRunner` (all stages) | | `runner_failures` | One row per item failure with error message + traceback | `BaseRunner._log_failure` | @@ -170,6 +174,21 @@ conn.cursor().execute(open('supabase/migrations/00007_error_logs.sql').read()) Use `datasmith.utils.db.fetch_all(table, select=..., filters=..., gte_filters=..., ...)` for paginated reads, or `get_client()` for direct Supabase client access. +### Level aggregation + +Per-benchmark speedups are rolled up into four levels via **geometric mean** in `src/datasmith/harbor_adapter/template/parser.py` (`geometric_mean()` + `aggregate_by_hierarchy()`): + +- **level1** — identity, one entry per benchmark (`module.Class.method`) +- **level2** — grouped by `module.Class` (drop the last dotted segment), geomean within each group +- **level3** — grouped by `module` (top dotted segment), geomean within each group +- **level4** — a single overall geomean across every benchmark + +`benchmark_information.benchmark_name` is already param-stripped (e.g. `benchmarks.ConstructorsSuite.time_point`), so consumers can replicate the rollup locally without re-parsing. The FormulaCode website CSV's `1-Params` row is equivalent to datasmith's `level1`; the four upper website rows map onto `level2`/`level3`/`level4` modulo column naming. + +### Task identity + +The canonical identifier for one PR / one task is the tuple `(owner, repo, issue_number)`. Tables that need a single-column join key expose a `task_id` field whose value equals `issue_number`; never construct it as a derived string. `candidate_containers.task_id` is a STORED generated column (see migration `00019_candidate_containers_task_id.sql`). + ## Environment setup Requires a `tokens.env` file in the repo root with `GH_TOKEN`, `CACHE_LOCATION`, `SUPABASE_URL`, `SUPABASE_KEY`, and optionally `DSPY_*` vars for LLM backends and `DOCKERHUB_*` vars for publishing. See README.md for the full template. diff --git a/scripts/harbor_tasks_migration.sql b/scripts/harbor_tasks_migration.sql new file mode 100644 index 00000000..a670c917 --- /dev/null +++ b/scripts/harbor_tasks_migration.sql @@ -0,0 +1,110 @@ +-- ============================================================================ +-- Harbor `tasks` table migration: composite (owner, repo, issue_number) PK +-- ============================================================================ +-- +-- Target: Harbor's Supabase project (see HARBOR_SUPABASE_URL in tokens.env), +-- NOT datasmith's local Supabase. Apply via: +-- +-- psql "$HARBOR_DATABASE_URL" -f scripts/harbor_tasks_migration.sql +-- +-- This script realigns Harbor's tasks table with the canonical FormulaCode +-- identity tuple `(owner, repo, issue_number)`. The legacy `task_id` column +-- was constructed inconsistently across the codebase (`owner__repo-N` in +-- the publish path, `owner_repo_N` in harbor_adapter), so we extract its +-- components and key the table on the tuple instead. The `task_id` column +-- is retained for one release as an integer mirror of `issue_number`. +-- +-- Companion change: scripts/migrate_snapshot_keys.py walks Supabase Storage +-- and renames `snapshots/{old_task_id}/oracle.tar.gz` → +-- `snapshots/{owner}/{repo}/{issue_number}/oracle.tar.gz` once this script +-- has run successfully (so the new tasks rows already carry the triple). +-- +-- Re-runnable: every statement is idempotent. + +BEGIN; + +-- 1. Add the three target columns (nullable initially so the backfill can +-- run before we add NOT NULL constraints). +ALTER TABLE tasks ADD COLUMN IF NOT EXISTS owner TEXT; +ALTER TABLE tasks ADD COLUMN IF NOT EXISTS repo TEXT; +ALTER TABLE tasks ADD COLUMN IF NOT EXISTS issue_number INT; + +-- 2. Backfill from the legacy task_id string. +-- Two patterns observed in the wild: +-- Format A: owner__repo-N (publish/HuggingFace path) +-- Format B: owner_repo_N (harbor_adapter path) +-- Format A uses a double-underscore between owner and repo and a dash +-- before the PR number; Format B uses single underscores throughout. +-- +-- We try Format A first (more specific), then fall back to Format B. +-- Rows that match neither stay NULL and will be flagged below. +UPDATE tasks + SET owner = split_part(task_id, '__', 1), + repo = split_part(split_part(task_id, '__', 2), '-', 1), + issue_number = NULLIF( + regexp_replace( + split_part(task_id, '__', 2), + '^[^-]*-(\d+)$', + '\1' + ), + split_part(task_id, '__', 2) + )::INT + WHERE owner IS NULL + AND task_id ~ '^[^_]+__[^_-]+-\d+$'; + +-- Fallback: Format B (owner_repo_N) — but only for rows still unbackfilled. +-- This is brittle because owner and repo can both contain underscores, so +-- we assume the last underscore-separated segment is the integer issue. +UPDATE tasks + SET owner = regexp_replace(task_id, '^(.+)_([^_]+)_(\d+)$', '\1'), + repo = regexp_replace(task_id, '^(.+)_([^_]+)_(\d+)$', '\2'), + issue_number = regexp_replace(task_id, '^(.+)_([^_]+)_(\d+)$', '\3')::INT + WHERE owner IS NULL + AND task_id ~ '^.+_[^_]+_\d+$'; + +-- Surface anything that didn't match either pattern. +DO $$ +DECLARE n_orphaned INT; +BEGIN + SELECT count(*) INTO n_orphaned FROM tasks WHERE owner IS NULL; + IF n_orphaned > 0 THEN + RAISE WARNING 'tasks: % rows have task_id values that do not parse — leaving them with NULL owner/repo/issue_number; review manually before enforcing NOT NULL', n_orphaned; + END IF; +END $$; + +-- 3. Once the backfill is verified clean, enforce NOT NULL + swap the PK. +-- Skip this block if any orphans remain — operator handles them first. +DO $$ +DECLARE n_orphaned INT; +BEGIN + SELECT count(*) INTO n_orphaned FROM tasks WHERE owner IS NULL OR repo IS NULL OR issue_number IS NULL; + IF n_orphaned = 0 THEN + ALTER TABLE tasks ALTER COLUMN owner SET NOT NULL; + ALTER TABLE tasks ALTER COLUMN repo SET NOT NULL; + ALTER TABLE tasks ALTER COLUMN issue_number SET NOT NULL; + + -- Replace the PK only if it isn't already on the triple. + IF NOT EXISTS ( + SELECT 1 + FROM pg_constraint + WHERE conrelid = 'tasks'::regclass + AND contype = 'p' + AND pg_get_constraintdef(oid) ILIKE '%(owner, repo, issue_number)%' + ) THEN + ALTER TABLE tasks DROP CONSTRAINT IF EXISTS tasks_pkey; + ALTER TABLE tasks ADD CONSTRAINT tasks_pkey PRIMARY KEY (owner, repo, issue_number); + END IF; + ELSE + RAISE WARNING 'tasks: skipping PK swap — % rows still have NULL owner/repo/issue_number', n_orphaned; + END IF; +END $$; + +-- 4. Reduce the legacy task_id column to an integer mirror of issue_number. +-- Keep the column for one release so any code we missed still reads +-- something sensible. The column is no longer a unique key. +ALTER TABLE tasks ALTER COLUMN task_id DROP NOT NULL; +UPDATE tasks SET task_id = issue_number::TEXT WHERE issue_number IS NOT NULL; +COMMENT ON COLUMN tasks.task_id IS + 'deprecated: equal to issue_number::text — use (owner, repo, issue_number) instead'; + +COMMIT; diff --git a/scripts/load_benchmark_information.py b/scripts/load_benchmark_information.py new file mode 100644 index 00000000..6170768e --- /dev/null +++ b/scripts/load_benchmark_information.py @@ -0,0 +1,249 @@ +"""Load per-benchmark speedups from terminal-bench runs into ``benchmark_information``. + +Scope: only the tasks listed in ``analysis/tasks.txt`` (one ``runs//`` +path per line). For each line we read the run-level ``run_metadata.json`` and the +per-task ``tests/config.json`` from the dataset the run consumed, look up the +canonical owner/repo/issue_number in ``pull_requests`` via ``merge_commit_sha = gt_hash``, +then explode ``parser_extra_metrics.per_benchmark_speedups_by_agent`` into one row +per (agent, model, benchmark). + +The per-benchmark dict is identical across trials within a task, so we only +process the first trial entry that carries a non-empty dict. + +speedup = (agent/nop) / (oracle/nop) — i.e. oracle_time / agent_time, so 1.0 = parity +with the human expert and >1.0 means the agent beat the human. +""" + +from __future__ import annotations + +import argparse +import json +import logging +import sys +from functools import cache +from pathlib import Path +from typing import Any + +from datasmith.utils.db import get_client + +logger = logging.getLogger("load_benchmark_information") + +DEFAULT_TB_ROOT = Path("/mnt/sdd1/atharvas/formulacode/eval_frameworks/terminal-bench") +DEFAULT_TASKS_FILE = DEFAULT_TB_ROOT / "analysis" / "tasks.txt" + + +def parse_agent_key(key: str) -> tuple[str, str | None]: + """``"terminus-2:anthropic-claude-sonnet-4-20250514"`` -> ("terminus-2", "...").""" + if ":" not in key: + return key, None + agent, model = key.split(":", 1) + if agent == "oracle" and model == "oracle": + return "oracle", None + return agent, model + + +@cache +def load_run_metadata(tb_root: Path, run_id: str) -> dict[str, Any]: + p = tb_root / "runs" / run_id / "run_metadata.json" + return json.loads(p.read_text()) + + +@cache +def load_results(tb_root: Path, run_id: str) -> dict[str, Any]: + p = tb_root / "runs" / run_id / "results.json" + return json.loads(p.read_text()) + + +def load_task_config(tb_root: Path, dataset_path: Path, task_id: str) -> dict[str, Any] | None: + """Read ``//tests/config.json``. + + ``run_metadata.dataset_path`` may point to a host path that doesn't exist on + this machine; if so we fall back to scanning ``/dataset/*`` for a + matching task directory. + """ + candidates: list[Path] = [] + if dataset_path.exists(): + candidates.append(dataset_path / task_id / "tests" / "config.json") + for ds in (tb_root / "dataset").glob("*"): + candidates.append(ds / task_id / "tests" / "config.json") + for cand in candidates: + if cand.exists(): + return json.loads(cand.read_text()) + return None + + +@cache +def lookup_pr_by_sha(gt_hash: str) -> tuple[str, str, int] | None: + client = get_client() + resp = ( + client.table("pull_requests") + .select("owner,repo,issue_number") + .eq("merge_commit_sha", gt_hash) + .limit(1) + .execute() + ) + if not resp.data: + return None + row = resp.data[0] + return row["owner"], row["repo"], int(row["issue_number"]) + + +def find_per_benchmark_dict(results: dict[str, Any], task_id: str) -> tuple[dict[str, Any], dict[str, Any]] | None: + """Return (per_benchmark_speedups_by_agent, sample_trial) for the task. + + Picks the first trial whose ``per_benchmark_speedups_by_agent`` is non-empty, + since the dict is identical across trials within a task. + """ + sample: dict[str, Any] | None = None + for trial in results.get("results", []): + if trial.get("task_id") != task_id: + continue + sample = sample or trial + pem = trial.get("parser_extra_metrics") or {} + pbs = pem.get("per_benchmark_speedups_by_agent") or {} + if pbs: + return pbs, trial + return ({}, sample) if sample is not None else None + + +def build_rows( + *, + run_id: str, + measured_at: str, + commit_hash: str | None, + owner: str, + repo: str, + issue_number: int, + pbs: dict[str, dict[str, dict[str, Any]]], + sample_trial: dict[str, Any], +) -> list[dict[str, Any]]: + rows: list[dict[str, Any]] = [] + trial_started = sample_trial.get("trial_started_at") + trial_ended = sample_trial.get("trial_ended_at") + for agent_key, benchmarks in pbs.items(): + agent_name, model_name = parse_agent_key(agent_key) + for bm_name, bm in benchmarks.items(): + agent_vs_nop = bm.get("agent/nop") + oracle_vs_nop = bm.get("oracle/nop") + if agent_vs_nop is None or oracle_vs_nop in (None, 0): + continue + speedup = agent_vs_nop / oracle_vs_nop + rows.append({ + "measured_at": measured_at, + "run_id": run_id, + "owner": owner, + "repo": repo, + "issue_number": issue_number, + "benchmark_name": bm_name, + "agent_name": agent_name, + "model_name": model_name, + "speedup": speedup, + "agent_speedup_vs_nop": agent_vs_nop, + "oracle_speedup_vs_nop": oracle_vs_nop, + "advantage": bm.get("advantage"), + "significant": bm.get("significant"), + "commit_hash": commit_hash, + "trial_started_at": trial_started, + "trial_ended_at": trial_ended, + "raw_payload": bm, + }) + return rows + + +def upsert_rows(rows: list[dict[str, Any]], batch: int = 500) -> None: + if not rows: + return + client = get_client() + for i in range(0, len(rows), batch): + chunk = rows[i : i + batch] + client.table("benchmark_information").upsert( + chunk, + on_conflict="run_id,owner,repo,issue_number,benchmark_name,agent_name,model_name", + ).execute() + logger.info("upserted %d rows (%d/%d)", len(chunk), i + len(chunk), len(rows)) + + +def process_task_line(tb_root: Path, line: str) -> list[dict[str, Any]]: + parts = Path(line).parts + if len(parts) < 3 or parts[0] != "runs": + logger.warning("skip malformed line: %s", line) + return [] + run_id, task_id = parts[1], parts[2] + + meta = load_run_metadata(tb_root, run_id) + measured_at = meta.get("start_time") or meta.get("end_time") + commit_hash = meta.get("commit_hash") + dataset_path = Path(meta.get("dataset_path", "")) + + config = load_task_config(tb_root, dataset_path, task_id) + if config is None: + logger.warning("no tests/config.json found for %s/%s", run_id, task_id) + return [] + gt_hash = config.get("gt_hash") + if not gt_hash: + logger.warning("no gt_hash for %s/%s", run_id, task_id) + return [] + + pr = lookup_pr_by_sha(gt_hash) + if pr is None: + logger.warning("no pull_requests row for gt_hash=%s (%s/%s)", gt_hash, run_id, task_id) + return [] + owner, repo, issue_number = pr + + results = load_results(tb_root, run_id) + found = find_per_benchmark_dict(results, task_id) + if found is None: + logger.warning("no trial found for %s/%s", run_id, task_id) + return [] + pbs, sample_trial = found + if not pbs: + logger.info("no per_benchmark_speedups_by_agent for %s/%s", run_id, task_id) + return [] + + return build_rows( + run_id=run_id, + measured_at=measured_at, + commit_hash=commit_hash, + owner=owner, + repo=repo, + issue_number=issue_number, + pbs=pbs, + sample_trial=sample_trial, + ) + + +def main() -> int: + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument("--tb-root", type=Path, default=DEFAULT_TB_ROOT) + ap.add_argument("--tasks-file", type=Path, default=DEFAULT_TASKS_FILE) + ap.add_argument("--dry-run", action="store_true", help="Build rows but don't upsert.") + ap.add_argument("--limit", type=int, default=0, help="Process only first N task lines.") + args = ap.parse_args() + + logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(name)s %(message)s") + + lines = [ln.strip() for ln in args.tasks_file.read_text().splitlines() if ln.strip()] + if args.limit: + lines = lines[: args.limit] + logger.info("processing %d task lines from %s", len(lines), args.tasks_file) + + all_rows: list[dict[str, Any]] = [] + for ln in lines: + try: + all_rows.extend(process_task_line(args.tb_root, ln)) + except Exception: + logger.exception("failed on %s", ln) + + logger.info("built %d rows", len(all_rows)) + if args.dry_run: + if all_rows: + logger.info("sample row: %s", json.dumps(all_rows[0], default=str, indent=2)[:500]) + return 0 + + upsert_rows(all_rows) + logger.info("done") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/migrate_snapshot_keys.py b/scripts/migrate_snapshot_keys.py new file mode 100644 index 00000000..023387b6 --- /dev/null +++ b/scripts/migrate_snapshot_keys.py @@ -0,0 +1,183 @@ +"""One-shot rename of Harbor snapshot tarballs to the (owner, repo, issue_number) layout. + +Companion to ``scripts/harbor_tasks_migration.sql``. After the SQL has +backfilled ``tasks.owner / repo / issue_number`` and the new harbor_adapter +code is writing to the triple-keyed path, this script walks the existing +``snapshots/`` bucket in Harbor's Supabase Storage and moves every +``snapshots/{old_task_id}/oracle.tar.gz`` to +``snapshots/{owner}/{repo}/{issue_number}/oracle.tar.gz``. + +Safe to run repeatedly: objects already at the new key are skipped. +Defaults to ``--dry-run``; pass ``--apply`` to actually move objects. + +Required env (read from tokens.env via ``datasmith.utils``): +- HARBOR_SUPABASE_URL +- HARBOR_SUPABASE_SERVICE_KEY (move + delete need the service role) + +Usage: + python scripts/migrate_snapshot_keys.py [--apply] [--limit N] +""" + +from __future__ import annotations + +import argparse +import json +import os +import sys +import urllib.error +import urllib.request +from typing import Any + +# Make sure tokens.env is loaded. +import datasmith # noqa: F401 -- side-effect import; dotenv.load_dotenv + +BUCKET = "snapshots" +PAGE_SIZE = 1000 +OBJECT_NAME = "oracle.tar.gz" + + +def _service_headers() -> dict[str, str]: + key = os.environ.get("HARBOR_SUPABASE_SERVICE_KEY", "") + if not key: + raise SystemExit("HARBOR_SUPABASE_SERVICE_KEY not set") + return {"Authorization": f"Bearer {key}", "apikey": key} + + +def _request(url: str, method: str, headers: dict[str, str], data: bytes | None = None) -> bytes: + # URLs are always constructed from HARBOR_SUPABASE_URL + a literal path, never user input. + req = urllib.request.Request(url, data=data, headers=headers, method=method) # noqa: S310 + with urllib.request.urlopen(req, timeout=60) as resp: # noqa: S310 + return resp.read() + + +def list_top_level_dirs(base_url: str) -> list[str]: + """Return the top-level 'directory' names under the snapshots bucket. + + Supabase Storage doesn't have real directories — ``list`` with a prefix + returns objects whose name begins with that prefix. We list with an empty + prefix and a ``/`` delimiter to get the unique first-segment names. + """ + names: list[str] = [] + offset = 0 + while True: + body = { + "prefix": "", + "limit": PAGE_SIZE, + "offset": offset, + "sortBy": {"column": "name", "order": "asc"}, + } + try: + raw = _request( + f"{base_url}/storage/v1/object/list/{BUCKET}", + "POST", + {**_service_headers(), "Content-Type": "application/json"}, + data=json.dumps(body).encode(), + ) + except urllib.error.HTTPError as exc: + sys.stderr.write(f"list failed: {exc.code} {exc.read().decode(errors='replace')}\n") + raise SystemExit(2) from exc + page: list[dict[str, Any]] = json.loads(raw) + if not page: + break + names.extend(entry["name"] for entry in page if entry.get("name")) + if len(page) < PAGE_SIZE: + break + offset += PAGE_SIZE + return names + + +def fetch_task_map(base_url: str) -> dict[str, tuple[str, str, int]]: + """Build legacy_task_id → (owner, repo, issue_number) from the tasks table.""" + raw = _request( + f"{base_url}/rest/v1/tasks?select=task_id,owner,repo,issue_number", + "GET", + {**_service_headers(), "Accept": "application/json"}, + ) + out: dict[str, tuple[str, str, int]] = {} + for row in json.loads(raw): + legacy = row.get("task_id") + owner = row.get("owner") + repo = row.get("repo") + issue_number = row.get("issue_number") + if not legacy or not owner or not repo or issue_number is None: + continue + out[str(legacy)] = (owner, repo, int(issue_number)) + return out + + +def move_object(base_url: str, source_key: str, dest_key: str) -> None: + body = {"bucketId": BUCKET, "sourceKey": source_key, "destinationKey": dest_key} + _request( + f"{base_url}/storage/v1/object/move", + "POST", + {**_service_headers(), "Content-Type": "application/json"}, + data=json.dumps(body).encode(), + ) + + +def object_exists(base_url: str, key: str) -> bool: + """HEAD-style probe via list — cheaper than fetching the blob.""" + parent, _, name = key.rpartition("/") + body = {"prefix": parent + "/", "limit": PAGE_SIZE, "search": name} + try: + raw = _request( + f"{base_url}/storage/v1/object/list/{BUCKET}", + "POST", + {**_service_headers(), "Content-Type": "application/json"}, + data=json.dumps(body).encode(), + ) + except urllib.error.HTTPError: + return False + return any(entry.get("name") == name for entry in json.loads(raw)) + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--apply", action="store_true", help="Actually move objects (default: dry run).") + ap.add_argument("--limit", type=int, default=None, help="Stop after this many candidates.") + args = ap.parse_args() + + base_url = os.environ.get("HARBOR_SUPABASE_URL", "") + if not base_url: + raise SystemExit("HARBOR_SUPABASE_URL not set") + + task_map = fetch_task_map(base_url) + if not task_map: + raise SystemExit("tasks table returned no rows with the triple — run harbor_tasks_migration.sql first") + print(f"loaded {len(task_map)} task rows", flush=True) + + legacy_dirs = list_top_level_dirs(base_url) + candidates = [d for d in legacy_dirs if "/" not in d and d in task_map] + print(f"found {len(candidates)} legacy snapshot dirs to migrate", flush=True) + + n_renamed = 0 + n_skipped = 0 + n_missing = 0 + for legacy in candidates: + if args.limit is not None and n_renamed >= args.limit: + break + owner, repo, issue = task_map[legacy] + source = f"{legacy}/{OBJECT_NAME}" + dest = f"{owner}/{repo}/{issue}/{OBJECT_NAME}" + + if not object_exists(base_url, source): + n_missing += 1 + continue + if object_exists(base_url, dest): + print(f"skip (dest exists): {source} → {dest}") + n_skipped += 1 + continue + + print(f"{'MOVE' if args.apply else 'DRY '} {source} → {dest}") + if args.apply: + move_object(base_url, source, dest) + n_renamed += 1 + + print( + f"done — moved={n_renamed} skipped_dest_exists={n_skipped} no_oracle_tarball={n_missing}", + flush=True, + ) + + +if __name__ == "__main__": + main() diff --git a/src/datasmith/github/models.py b/src/datasmith/github/models.py index 7fc65fba..a22c8e50 100644 --- a/src/datasmith/github/models.py +++ b/src/datasmith/github/models.py @@ -165,7 +165,7 @@ def to_record(self) -> FormulaCodeRecord | None: owner=self.owner, repo=self.repo, issue_number=self.issue_number, - task_id=f"{self.owner}__{self.repo}-{self.issue_number}", + task_id=self.issue_number, gt_hash=self.merge_commit_sha, base_commit=self.base_sha, date=self.merged_at, @@ -194,7 +194,7 @@ class FormulaCodeRecord(BaseModel): owner: str repo: str issue_number: int - task_id: str + task_id: int # mirror of issue_number; the canonical id is (owner, repo, issue_number) gt_hash: str = "" base_commit: str = "" date: datetime | None = None diff --git a/src/datasmith/harbor_adapter/adapter.py b/src/datasmith/harbor_adapter/adapter.py index 820b1831..7d5a48fd 100644 --- a/src/datasmith/harbor_adapter/adapter.py +++ b/src/datasmith/harbor_adapter/adapter.py @@ -19,7 +19,9 @@ class FormulaCodeRecord: container_name: str # repo_name-{base_sha}:final patch: str # diff(base_sha, merge_sha) - task_id: str # repo_name_{ID} + owner: str # repository owner (canonical id is (owner, repo, issue_number)) + repo: str # repository name + issue_number: int # PR number; the canonical row identifier gt_hash: str # merge_sha base_commit: str # base_sha instructions: str @@ -28,6 +30,24 @@ class FormulaCodeRecord: difficulty: str = "hard" # difficulty level repo_name: str | None = None + @property + def task_id(self) -> int: + """Single-column join key (= issue_number). Repo qualifier is implicit.""" + return self.issue_number + + @property + def task_dir_name(self) -> str: + """Filesystem-safe name used as Harbor's task / trial directory. + + Harbor discovers tasks by scanning ``LocalDatasetConfig.path`` for + ``task.toml`` files (flat or one-level deep) and uses the parent + directory name as the trial name. We can't nest by ``owner/repo`` or + Harbor's task IDs become unique-per-issue but ambiguous-across-repos. + The triple-segment string keeps each trial dir globally unique while + staying within Harbor's flat-discovery contract. + """ + return f"{self.owner}__{self.repo}__{self.issue_number}" + class HarborTaskPaths: """Convenience paths for writing a Harbor task.""" @@ -130,7 +150,9 @@ def _write_test_files( # test.sh test_sh_content = render_test_sh( base_commit=rec.base_commit, - task_id=rec.task_id, + owner=rec.owner, + repo=rec.repo, + issue_number=rec.issue_number, run_pytest=run_pytest, rounds=rounds, ) @@ -146,6 +168,7 @@ def _write_test_files( # up with identical content; Harbor's later upload_dir() either overwrites # with the same bytes or merges harmlessly. cfg = rec.__dict__.copy() + cfg["task_id"] = rec.task_id # surface the property so legacy readers see it cfg_json = json.dumps(cfg, indent=2) (paths.tests_dir / "config.json").write_text(cfg_json) (paths.environment_dir / "config.json").write_text(cfg_json) @@ -166,7 +189,9 @@ def _write_solution_files( # setup.sh extra_setup_commands = "" setup_sh_content = render_run_setup_sh( - task_id=rec.task_id, + owner=rec.owner, + repo=rec.repo, + issue_number=rec.issue_number, rounds=rounds, extra_setup_commands=extra_setup_commands, ) @@ -186,7 +211,7 @@ def generate_task( verifier_env: dict[str, str] | None = None, ) -> Path: """Generate a complete Harbor task directory for the given FormulaCodeRecord.""" - out_dir = self.out_root / rec.task_id + out_dir = self.out_root / rec.task_dir_name out_dir.mkdir(parents=True, exist_ok=True) # Create harbor task paths diff --git a/src/datasmith/harbor_adapter/records.py b/src/datasmith/harbor_adapter/records.py index c375dad0..b2db4c98 100644 --- a/src/datasmith/harbor_adapter/records.py +++ b/src/datasmith/harbor_adapter/records.py @@ -27,11 +27,13 @@ def to_record(pr: dict[str, Any]) -> FormulaCodeRecord: """ owner = pr["owner"] repo = pr["repo"] - issue_number = pr["issue_number"] + issue_number = int(pr["issue_number"]) return FormulaCodeRecord( container_name=str(pr.get("container_name") or ""), patch=str(pr.get("patch") or ""), - task_id=f"{owner}_{repo}_{issue_number}", + owner=owner, + repo=repo, + issue_number=issue_number, gt_hash=str(pr.get("merge_commit_sha") or ""), base_commit=str(pr.get("base_sha") or ""), instructions=str(pr.get("rendered_problem") or ""), diff --git a/src/datasmith/harbor_adapter/template/parser.py b/src/datasmith/harbor_adapter/template/parser.py index 95db3397..95a38840 100644 --- a/src/datasmith/harbor_adapter/template/parser.py +++ b/src/datasmith/harbor_adapter/template/parser.py @@ -6,7 +6,7 @@ and writes reward.json + reward.txt. Usage: - python /tests/parser.py --task-id ID [--agent-key KEY] + python /tests/parser.py --owner OWNER --repo REPO --issue-number N [--agent-key KEY] """ from __future__ import annotations @@ -70,10 +70,12 @@ def _supabase_get(url: str, headers: dict) -> dict | list | None: return None -def fetch_oracle_benchmarks(base_url: str, task_id: str) -> dict[str, float] | None: +def fetch_oracle_benchmarks( + base_url: str, owner: str, repo: str, issue_number: int +) -> dict[str, float] | None: """Fetch the oracle's per-benchmark timings from Supabase. - 1. GET tasks?task_id=eq.{id} to find baseline_run_id + 1. GET tasks?owner=eq.{}&repo=eq.{}&issue_number=eq.{} to find baseline_run_id 2. GET runs?uuid=eq.{baseline_run_id} to get payload 3. Extract payload.lsv_results.measure.benchmarks → {name: current_seconds} @@ -89,16 +91,19 @@ def fetch_oracle_benchmarks(base_url: str, task_id: str) -> dict[str, float] | N "Accept": "application/json", } + task_filter = f"owner=eq.{owner}&repo=eq.{repo}&issue_number=eq.{issue_number}" + task_label = f"{owner}/{repo}#{issue_number}" + # Step 1: Get baseline_run_id from tasks table - tasks_url = f"{base_url}/rest/v1/tasks?task_id=eq.{task_id}&select=baseline_run_id" + tasks_url = f"{base_url}/rest/v1/tasks?{task_filter}&select=baseline_run_id" tasks = _supabase_get(tasks_url, headers) if not tasks or not isinstance(tasks, list) or len(tasks) == 0: - print(f"[parser] No task row found for {task_id}") + print(f"[parser] No task row found for {task_label}") return None baseline_run_id = tasks[0].get("baseline_run_id") if not baseline_run_id: - print(f"[parser] No baseline_run_id set for task {task_id}") + print(f"[parser] No baseline_run_id set for task {task_label}") return None # Step 2: Get the baseline run's payload @@ -434,7 +439,9 @@ def write_reward( def main() -> None: parser = argparse.ArgumentParser(description="Compute reward from LSV results") - parser.add_argument("--task-id", required=True, help="Task identifier") + parser.add_argument("--owner", required=True, help="Repository owner") + parser.add_argument("--repo", required=True, help="Repository name") + parser.add_argument("--issue-number", required=True, type=int, help="PR number") parser.add_argument( "--agent-key", default="agent", help="Agent key (e.g., oracle, terminus-2)" ) @@ -490,7 +497,9 @@ def main() -> None: # Fetch oracle data from Supabase base_url = os.environ.get("SUPABASE_URL", "") if base_url: - oracle_benchmarks = fetch_oracle_benchmarks(base_url, args.task_id) + oracle_benchmarks = fetch_oracle_benchmarks( + base_url, args.owner, args.repo, args.issue_number + ) if oracle_benchmarks: advantages = compute_per_benchmark_advantages( benchmarks, oracle_benchmarks diff --git a/src/datasmith/harbor_adapter/template/setup.sh b/src/datasmith/harbor_adapter/template/setup.sh index 006a5da8..b9a38d41 100644 --- a/src/datasmith/harbor_adapter/template/setup.sh +++ b/src/datasmith/harbor_adapter/template/setup.sh @@ -4,8 +4,13 @@ cd /workspace/repo || exit 1 ts() { date -u "+%Y-%m-%dT%H:%M:%SZ"; } -TASK_ID="{{ task_id }}" -export TASK_ID +OWNER="{{ owner }}" +REPO="{{ repo }}" +ISSUE_NUMBER="{{ issue_number }}" +# TASK_ID retained as a single-column alias (= issue_number) for back-compat +# with anything that still reads $TASK_ID. The canonical id is the triple. +TASK_ID="${ISSUE_NUMBER}" +export OWNER REPO ISSUE_NUMBER TASK_ID # Harbor's [verifier.env] section only reaches test.sh, not setup.sh, so we # bake the agent name directly into setup.sh at render time. lsv_init.py diff --git a/src/datasmith/harbor_adapter/template/test.sh b/src/datasmith/harbor_adapter/template/test.sh index 1bda79d4..f8034f94 100644 --- a/src/datasmith/harbor_adapter/template/test.sh +++ b/src/datasmith/harbor_adapter/template/test.sh @@ -17,8 +17,11 @@ eval "$(micromamba shell hook --shell=bash)" micromamba activate "$ENV_NAME" || true set -u -TASK_ID="{{ task_id }}" -export TASK_ID +OWNER="{{ owner }}" +REPO="{{ repo }}" +ISSUE_NUMBER="{{ issue_number }}" +TASK_ID="${ISSUE_NUMBER}" +export OWNER REPO ISSUE_NUMBER TASK_ID AGENT_MODEL_NAME_INPUT="${1:-agent}" AGENT_KEY="$(printf '%s' "${AGENT_MODEL_NAME_INPUT}" | tr '[:upper:]' '[:lower:]' | sed -E 's/[^a-z0-9_-]+/-/g; s/^-+//; s/-+$//')" @@ -85,7 +88,7 @@ if [ "${AGENT_KEY}" != "oracle" ] && [ -n "${SUPABASE_URL:-}" ] && [ -n "${SUPAB import sys; sys.path.insert(0, '/opt/lsv') from upload import download_snapshots import os -download_snapshots(os.environ['SUPABASE_URL'], '${TASK_ID}', '${SNAPSHOT_DIR}') +download_snapshots(os.environ['SUPABASE_URL'], os.environ['OWNER'], os.environ['REPO'], int(os.environ['ISSUE_NUMBER']), '${SNAPSHOT_DIR}') " || echo "WARNING: snapshot download failed" fi @@ -134,7 +137,7 @@ PYEOF # ── Parser: compute reward ─────────────────────────────────────────────── echo "[$(ts)] [test] Computing reward..." -python /opt/lsv/parser.py --task-id "${TASK_ID}" --agent-key "${AGENT_KEY}" +python /opt/lsv/parser.py --owner "${OWNER}" --repo "${REPO}" --issue-number "${ISSUE_NUMBER}" --agent-key "${AGENT_KEY}" # ── Upload to Supabase (if configured) ─────────────────────────────────── if [ -n "${SUPABASE_URL:-}" ] && [ -n "${SUPABASE_ANON_KEY:-}" ] && [ -z "${FORMULACODE_NO_UPLOAD:-}" ]; then @@ -143,7 +146,7 @@ if [ -n "${SUPABASE_URL:-}" ] && [ -n "${SUPABASE_ANON_KEY:-}" ] && [ -z "${FORM if [ "${AGENT_KEY}" = "oracle" ]; then oracle_flag="--oracle" fi - python /opt/lsv/upload.py --task-id "${TASK_ID}" --agent-key "${AGENT_KEY}" ${oracle_flag} || \ + python /opt/lsv/upload.py --owner "${OWNER}" --repo "${REPO}" --issue-number "${ISSUE_NUMBER}" --agent-key "${AGENT_KEY}" ${oracle_flag} || \ echo "WARNING: Supabase upload failed" fi diff --git a/src/datasmith/harbor_adapter/template/upload.py b/src/datasmith/harbor_adapter/template/upload.py index 9558e0d4..6b54de4b 100644 --- a/src/datasmith/harbor_adapter/template/upload.py +++ b/src/datasmith/harbor_adapter/template/upload.py @@ -7,7 +7,7 @@ Uses urllib.request (stdlib) — no extra dependencies required. Usage: - python /tests/upload.py --task-id ID --agent-key KEY [--oracle] + python /tests/upload.py --owner OWNER --repo REPO --issue-number N --agent-key KEY [--oracle] """ from __future__ import annotations @@ -76,13 +76,25 @@ def _request( # ── Supabase operations ───────────────────────────────────────────────────── +def _task_filter(owner: str, repo: str, issue_number: int) -> str: + """PostgREST filter snippet matching the canonical (owner, repo, issue_number) triple.""" + return f"owner=eq.{owner}&repo=eq.{repo}&issue_number=eq.{issue_number}" + + +def _task_storage_prefix(owner: str, repo: str, issue_number: int) -> str: + """Storage key prefix for snapshot tarballs.""" + return f"{owner}/{repo}/{issue_number}" + + def ensure_task_exists( base_url: str, - task_id: str, + owner: str, + repo: str, + issue_number: int, ) -> bool: """Ensure a task row exists. Returns True if inserted, False if exists.""" resp = _request( - f"{base_url}/rest/v1/tasks?task_id=eq.{task_id}&select=task_id", + f"{base_url}/rest/v1/tasks?{_task_filter(owner, repo, issue_number)}&select=owner", "GET", {**_anon_headers(), "Accept": "application/json"}, ) @@ -94,7 +106,9 @@ def ensure_task_exists( f"{base_url}/rest/v1/tasks", "POST", {**_service_headers(), "Prefer": "return=minimal"}, - data=json.dumps({"task_id": task_id}).encode(), + data=json.dumps( + {"owner": owner, "repo": repo, "issue_number": issue_number, "task_id": issue_number} + ).encode(), content_type="application/json", ) return True @@ -116,7 +130,9 @@ def upload_tarball(base_url: str, tarball_path: Path, object_key: str) -> str: def insert_run( base_url: str, run_uuid: str, - task_id: str, + owner: str, + repo: str, + issue_number: int, agent_key: str, payload: dict, pytest_success_ratio: float, @@ -125,7 +141,10 @@ def insert_run( """Insert a run row into Supabase.""" row = { "uuid": run_uuid, - "task_id": task_id, + "owner": owner, + "repo": repo, + "issue_number": issue_number, + "task_id": issue_number, # deprecated single-column alias "agent_name": agent_key, "model_name": agent_key, "model_agent_signature": f"{agent_key}:{agent_key}", @@ -147,7 +166,9 @@ def insert_run( def update_task_oracle_metadata( base_url: str, - task_id: str, + owner: str, + repo: str, + issue_number: int, run_uuid: str, snapshot_url: str | None = None, ) -> None: @@ -157,7 +178,7 @@ def update_task_oracle_metadata( patch_data["snapshot_storage_url"] = snapshot_url _request( - f"{base_url}/rest/v1/tasks?task_id=eq.{task_id}", + f"{base_url}/rest/v1/tasks?{_task_filter(owner, repo, issue_number)}", "PATCH", {**_service_headers(), "Prefer": "return=minimal"}, data=json.dumps(patch_data).encode(), @@ -166,11 +187,11 @@ def update_task_oracle_metadata( def upload_snapshots( - base_url: str, task_id: str, snapshot_dir: str | Path + base_url: str, owner: str, repo: str, issue_number: int, snapshot_dir: str | Path ) -> str | None: """Upload .snapshots/ directory to Supabase Storage (oracle only). - Creates a tarball and uploads to snapshots/{task_id}/oracle.tar.gz. + Creates a tarball and uploads to snapshots/{owner}/{repo}/{issue_number}/oracle.tar.gz. Uses POST; falls back to PUT on 409 (object already exists). Returns a public URL on success, else None. """ @@ -190,7 +211,7 @@ def _snap_filter(info: tarfile.TarInfo) -> tarfile.TarInfo | None: tar.add(str(snapshot_dir), arcname=".snapshots", filter=_snap_filter) - object_key = f"{task_id}/oracle.tar.gz" + object_key = f"{_task_storage_prefix(owner, repo, issue_number)}/oracle.tar.gz" url = f"{base_url}/storage/v1/object/snapshots/{object_key}" data = tarball_path.read_bytes() headers = _service_headers() @@ -215,11 +236,13 @@ def _snap_filter(info: tarfile.TarInfo) -> tarfile.TarInfo | None: tarball_path.unlink(missing_ok=True) -def download_snapshots(base_url: str, task_id: str, snapshot_dir: str | Path) -> bool: +def download_snapshots( + base_url: str, owner: str, repo: str, issue_number: int, snapshot_dir: str | Path +) -> bool: """Download oracle snapshots from URL stored on tasks.snapshot_storage_url. Reads snapshot_storage_url from tasks row using anon key, then downloads/extracts - into snapshot_dir. Falls back to legacy path snapshots/{task_id}/oracle.tar.gz + into snapshot_dir. Falls back to snapshots/{owner}/{repo}/{issue_number}/oracle.tar.gz when snapshot_storage_url is missing. """ snapshot_dir = Path(snapshot_dir) @@ -227,7 +250,7 @@ def download_snapshots(base_url: str, task_id: str, snapshot_dir: str | Path) -> try: task_resp = _request( - f"{base_url}/rest/v1/tasks?task_id=eq.{task_id}&select=snapshot_storage_url", + f"{base_url}/rest/v1/tasks?{_task_filter(owner, repo, issue_number)}&select=snapshot_storage_url", "GET", {**_anon_headers(), "Accept": "application/json"}, ) @@ -239,10 +262,10 @@ def download_snapshots(base_url: str, task_id: str, snapshot_dir: str | Path) -> if snapshot_url: download_url = snapshot_url else: - object_key = f"{task_id}/oracle.tar.gz" + object_key = f"{_task_storage_prefix(owner, repo, issue_number)}/oracle.tar.gz" download_url = f"{base_url}/storage/v1/object/snapshots/{object_key}" print( - "[upload] WARNING: tasks.snapshot_storage_url missing; using legacy snapshot path" + "[upload] WARNING: tasks.snapshot_storage_url missing; using default snapshot path" ) data = _request(download_url, "GET", _anon_headers()) @@ -252,7 +275,7 @@ def download_snapshots(base_url: str, task_id: str, snapshot_dir: str | Path) -> with tarfile.open(tarball_path, "r:gz") as tar: tar.extractall(path=snapshot_dir.parent, filter="data") - print(f"[upload] Downloaded oracle snapshots for {task_id}") + print(f"[upload] Downloaded oracle snapshots for {owner}/{repo}#{issue_number}") return True except Exception as e: print(f"[upload] WARNING: snapshot download failed ({e})") @@ -319,7 +342,9 @@ def build_payload( def main() -> None: parser = argparse.ArgumentParser(description="Upload results to Supabase") - parser.add_argument("--task-id", required=True, help="Task identifier") + parser.add_argument("--owner", required=True, help="Repository owner") + parser.add_argument("--repo", required=True, help="Repository name") + parser.add_argument("--issue-number", required=True, type=int, help="PR number") parser.add_argument("--agent-key", required=True, help="Agent key") parser.add_argument( "--oracle", @@ -340,7 +365,7 @@ def main() -> None: return run_uuid = str(uuid4()) - object_key = f"{args.agent_key}-{args.task_id}-{run_uuid}.tar.gz" + object_key = f"{args.agent_key}-{args.owner}-{args.repo}-{args.issue_number}-{run_uuid}.tar.gz" # Load data files lsv_results = {} @@ -369,7 +394,10 @@ def main() -> None: lsv_mean_speedup_ratio = reward_data.get("lsv_mean_speedup", 0.0) - print(f"[upload] task_id={args.task_id} agent={args.agent_key} uuid={run_uuid}") + print( + f"[upload] {args.owner}/{args.repo}#{args.issue_number} " + f"agent={args.agent_key} uuid={run_uuid}" + ) # 1. Create and upload tarball (non-fatal if storage bucket not configured) storage_url = "" @@ -386,7 +414,7 @@ def main() -> None: if not os.environ.get("SUPABASE_SERVICE_KEY"): raise RuntimeError("SUPABASE_SERVICE_KEY required for oracle operations") - inserted = ensure_task_exists(base_url, args.task_id) + inserted = ensure_task_exists(base_url, args.owner, args.repo, args.issue_number) print(f"[upload] Task: {'inserted' if inserted else 'already exists'}") # 4. Build payload and insert run @@ -395,7 +423,9 @@ def main() -> None: insert_run( base_url, run_uuid, - args.task_id, + args.owner, + args.repo, + args.issue_number, args.agent_key, payload, pytest_success_ratio, @@ -410,9 +440,13 @@ def main() -> None: snapshot_url = None snapshot_dir = ARTIFACTS_DIR / ".snapshots" if snapshot_dir.exists(): - snapshot_url = upload_snapshots(base_url, args.task_id, snapshot_dir) + snapshot_url = upload_snapshots( + base_url, args.owner, args.repo, args.issue_number, snapshot_dir + ) - update_task_oracle_metadata(base_url, args.task_id, run_uuid, snapshot_url) + update_task_oracle_metadata( + base_url, args.owner, args.repo, args.issue_number, run_uuid, snapshot_url + ) print(f"[upload] Set baseline_run_id={run_uuid}") if snapshot_url: print(f"[upload] Set snapshot_url={snapshot_url}") diff --git a/src/datasmith/harbor_adapter/utils.py b/src/datasmith/harbor_adapter/utils.py index 5f10f3da..c00673a7 100644 --- a/src/datasmith/harbor_adapter/utils.py +++ b/src/datasmith/harbor_adapter/utils.py @@ -143,7 +143,9 @@ def render_dockerfile(base_image: str) -> str: def render_test_sh( base_commit: str, - task_id: str, + owner: str, + repo: str, + issue_number: int, run_pytest: bool = True, rounds: int = 1, ) -> str: @@ -151,7 +153,9 @@ def render_test_sh( return _render_template( "test.sh", base_commit=base_commit, - task_id=task_id, + owner=owner, + repo=repo, + issue_number=issue_number, run_pytest=run_pytest, rounds=rounds, ) @@ -164,14 +168,18 @@ def render_solution_sh(solution_patch: str) -> str: def render_run_setup_sh( *, - task_id: str, + owner: str, + repo: str, + issue_number: int, rounds: int = 1, extra_setup_commands: str = "", ) -> str: """Render setup.sh with task metadata and extra setup commands.""" return _render_template( "setup.sh", - task_id=task_id, + owner=owner, + repo=repo, + issue_number=issue_number, rounds=rounds, extra_setup_commands=extra_setup_commands or "", ) diff --git a/src/datasmith/publish/huggingface.py b/src/datasmith/publish/huggingface.py index 5b5abfb5..a02205a5 100644 --- a/src/datasmith/publish/huggingface.py +++ b/src/datasmith/publish/huggingface.py @@ -87,10 +87,10 @@ def create_dataset_card(self, version: str) -> str: | Field | Type | Description | |-------|------|-------------| -| task_id | string | Unique task identifier (owner__repo-issue_number) | | owner | string | Repository owner | | repo | string | Repository name | -| issue_number | int | PR number | +| issue_number | int | PR number; the canonical row identity is (owner, repo, issue_number) | +| task_id | int | Mirror of `issue_number` for single-column joins | | gt_hash | string | Ground truth merge commit SHA | | base_commit | string | Base commit SHA | | date | string | Merge date | diff --git a/src/datasmith/publish/pipeline.py b/src/datasmith/publish/pipeline.py index 951c0ef1..b76a5091 100644 --- a/src/datasmith/publish/pipeline.py +++ b/src/datasmith/publish/pipeline.py @@ -54,7 +54,7 @@ async def publish_pipeline( "repo", record.repo ).eq("issue_number", record.issue_number).execute() except Exception: - logger.warning("Failed to mark %s as published", record.task_id) + logger.warning("Failed to mark %s/%s#%d as published", record.owner, record.repo, record.issue_number) logger.info("Published %d records as %s", len(records), version) return len(records) diff --git a/src/datasmith/publish/records.py b/src/datasmith/publish/records.py index e12ba811..de6ff0ca 100644 --- a/src/datasmith/publish/records.py +++ b/src/datasmith/publish/records.py @@ -110,7 +110,7 @@ def records_from_supabase( # noqa: C901 owner=row["owner"], repo=row["repo"], issue_number=row["issue_number"], - task_id=f"{row['owner']}__{row['repo']}-{row['issue_number']}", + task_id=int(row["issue_number"]), gt_hash=sha, base_commit=row.get("base_sha", ""), date=row.get("merged_at"), diff --git a/src/datasmith/runners/harbor_healthcheck.py b/src/datasmith/runners/harbor_healthcheck.py index 19bb86eb..15b54131 100644 --- a/src/datasmith/runners/harbor_healthcheck.py +++ b/src/datasmith/runners/harbor_healthcheck.py @@ -34,9 +34,9 @@ def _patch_harbor_trial_name() -> None: ``f"{task_name[:32]}__{ShortUUID().random(length=7)}"`` (see ``harbor/models/trial/config.py``), which makes trial directories non-deterministic across runs and defeats simple re-triage. We want the - directory under ``jobs//`` to be exactly our ``task_id`` - (``owner_repo_prnumber``) so a second run of the same task lands in the - same path. Override at import time — only affects this process. + directory under ``jobs//`` to be exactly the task name + (``owner__repo__issue_number``) so a second run of the same task lands + in the same path. Override at import time — only affects this process. """ from harbor.models.trial.config import TrialConfig @@ -90,9 +90,9 @@ def _materialize_tasks( *, rounds: int, ) -> dict[str, dict[str, Any]]: - """Write one Harbor task directory per PR. Returns a mapping from - ``task_id`` (the directory name Harbor sees) back to the datasmith row - metadata we need when inserting harbor_runs.""" + """Write one Harbor task directory per PR. Returns a mapping from the + Harbor task directory name back to the datasmith row metadata we need + when inserting harbor_runs.""" adapter = FormulaCodeAdapter(harbor_tasks_root=task_dir, force=True) verifier_env = _build_verifier_env() or None @@ -116,9 +116,9 @@ def _materialize_tasks( verifier_env=verifier_env, ) except Exception: - logger.exception("generate_task failed for %s", rec.task_id) + logger.exception("generate_task failed for %s/%s#%d", rec.owner, rec.repo, rec.issue_number) continue - task_id_map[rec.task_id] = { + task_id_map[rec.task_dir_name] = { "owner": pr["owner"], "repo": pr["repo"], "sha": pr["merge_commit_sha"], @@ -201,10 +201,10 @@ def _row_from_trial( # noqa: C901 trial can't be mapped back to a datasmith PR.""" from harbor.models.trial.paths import TrialPaths - task_id = trial.task_id.get_name() - meta = task_id_map.get(task_id) + task_name = trial.task_id.get_name() + meta = task_id_map.get(task_name) if meta is None: - logger.warning("Trial %s has no matching datasmith PR — skipping row", task_id) + logger.warning("Trial %s has no matching datasmith PR — skipping row", task_name) return None trial_dir = _trial_dir_from_uri(trial.trial_uri) diff --git a/src/datasmith/runners/scrape_benchmark_source.py b/src/datasmith/runners/scrape_benchmark_source.py new file mode 100644 index 00000000..439b9aaa --- /dev/null +++ b/src/datasmith/runners/scrape_benchmark_source.py @@ -0,0 +1,76 @@ +"""Runner for stage 9: scrape ASV benchmark source code per (owner, repo, sha). + +Reads candidate_containers for the date window, checks out each repo at its +SHA via :func:`prepare_repo_checkout`, parses ``benchmarks/*.py`` with the +extractor in :mod:`datasmith.scrape.benchmark_source`, and upserts rows into +``benchmark_codes``. +""" + +from __future__ import annotations + +import asyncio +import functools +import tempfile +from pathlib import Path +from typing import Any + +from datasmith.runners.base import BaseRunner +from datasmith.scrape.benchmark_source import BenchmarkSource, extract_benchmarks +from datasmith.utils import get_client, get_logger + +logger = get_logger("runners.scrape_benchmark_source") + + +def _scrape_one(owner: str, repo: str, sha: str) -> list[BenchmarkSource]: + """Synchronous worker — runs in a thread pool because GitPython is blocking.""" + from datasmith.resolution.git_utils import prepare_repo_checkout + + repo_name = f"{owner}/{repo}" + with tempfile.TemporaryDirectory(prefix="fc-bench-scrape-") as tmp: + _, repo_dir, cleanup = prepare_repo_checkout(repo_name, sha, Path(tmp)) + try: + return extract_benchmarks(Path(repo_dir)) + finally: + try: + cleanup() + except Exception: + logger.debug("worktree cleanup failed for %s@%s", repo_name, sha[:8]) + + +class ScrapeBenchmarkSourceRunner(BaseRunner): + """Stage 9 — populate ``benchmark_codes`` from candidate containers.""" + + def __init__(self, n_concurrent: int = 8) -> None: + super().__init__(name="scrape_benchmark_source", n_concurrent=n_concurrent) + + async def _process_item(self, item: Any) -> None: + owner = item["owner"] + repo = item["repo"] + sha = item["sha"] + + loop = asyncio.get_running_loop() + benches = await loop.run_in_executor(None, functools.partial(_scrape_one, owner, repo, sha)) + + if not benches: + logger.info("no ASV benchmarks for %s/%s@%s", owner, repo, sha[:8]) + return + + rows = [ + { + "owner": owner, + "repo": repo, + "benchmark_without_params": b.benchmark_without_params, + "source": b.source, + "setup_source": b.setup_source, + "last_scraped_sha": sha, + } + for b in benches + ] + + client = get_client() + # Upsert in chunks; PostgREST request size is bounded. + for start in range(0, len(rows), 200): + chunk = rows[start : start + 200] + client.table("benchmark_codes").upsert(chunk, on_conflict="owner,repo,benchmark_without_params").execute() + + logger.info("scraped %d benchmarks for %s/%s@%s", len(rows), owner, repo, sha[:8]) diff --git a/src/datasmith/scrape/__init__.py b/src/datasmith/scrape/__init__.py new file mode 100644 index 00000000..c5d85a30 --- /dev/null +++ b/src/datasmith/scrape/__init__.py @@ -0,0 +1 @@ +"""Scraping modules for the FormulaCode dataset.""" diff --git a/src/datasmith/scrape/benchmark_source.py b/src/datasmith/scrape/benchmark_source.py new file mode 100644 index 00000000..680cdba8 --- /dev/null +++ b/src/datasmith/scrape/benchmark_source.py @@ -0,0 +1,216 @@ +"""Extract ASV benchmark source code from a checked-out repository. + +Used by pipeline stage 9 (``scrape_benchmark_source``) to populate the +``benchmark_codes`` table, which the FormulaCode website joins against +``benchmark_information`` on (owner, repo, benchmark_without_params). + +The naming convention mirrors ASV: classes named ``Time*`` / ``Mem*`` / +``Peakmem*`` whose methods are ``time_*`` / ``mem_*`` / ``peakmem_*`` / +``track_*``, and module-level functions matching the same method prefixes. +The fully-qualified benchmark name is ``..`` or +``.`` where ```` is the dotted file path +relative to the ASV ``benchmark_dir``. +""" + +from __future__ import annotations + +import ast +import json +import os +import re +from dataclasses import dataclass +from pathlib import Path + +from datasmith.utils import get_logger + +logger = get_logger("scrape.benchmark_source") + +# Tunable knobs (see CLAUDE.md tunable-constants rule). +DATASMITH_BENCH_SCRAPE_MAX_FILE_BYTES: int = int(os.environ.get("DATASMITH_BENCH_SCRAPE_MAX_FILE_BYTES", "1000000")) +DATASMITH_BENCH_SCRAPE_DIRS: tuple[str, ...] = tuple( + d.strip() + for d in os.environ.get("DATASMITH_BENCH_SCRAPE_DIRS", "benchmarks,asv_bench/benchmarks").split(",") + if d.strip() +) + +_ASV_METHOD_PREFIXES = ("time_", "mem_", "peakmem_", "track_") +_ASV_CLASS_PREFIXES = ("Time", "Mem", "Peakmem") +_SETUP_NAMES = frozenset({"setup", "setup_cache", "teardown"}) + + +@dataclass(frozen=True) +class BenchmarkSource: + """One benchmark function's source (and its setup, if any).""" + + benchmark_without_params: str + source: str + setup_source: str | None + + +def _resolve_benchmark_dir(repo_root: Path) -> Path | None: # noqa: C901 + """Resolve the ASV ``benchmark_dir`` from ``asv.conf.json`` if present. + + Falls back to ``DATASMITH_BENCH_SCRAPE_DIRS`` candidates if no config + exists or the configured dir is missing. + """ + conf_path = repo_root / "asv.conf.json" + if conf_path.is_file(): + try: + raw = conf_path.read_text(encoding="utf-8", errors="replace") + # asv.conf.json sometimes carries // comments; strip them defensively. + raw_no_comments = re.sub(r"^\s*//.*$", "", raw, flags=re.MULTILINE) + conf = json.loads(raw_no_comments) + except (OSError, json.JSONDecodeError) as exc: + logger.debug("asv.conf.json unreadable at %s: %s", conf_path, exc) + else: + bench_dir = conf.get("benchmark_dir") + if isinstance(bench_dir, str) and bench_dir: + candidate = (repo_root / bench_dir).resolve() + try: + candidate.relative_to(repo_root.resolve()) + except ValueError: + logger.debug("benchmark_dir %s escapes repo root", bench_dir) + else: + if candidate.is_dir(): + return candidate + + for fallback in DATASMITH_BENCH_SCRAPE_DIRS: + candidate = (repo_root / fallback).resolve() + try: + candidate.relative_to(repo_root.resolve()) + except ValueError: + continue + if candidate.is_dir(): + return candidate + return None + + +def _module_path(bench_dir: Path, py_file: Path) -> str: + """Dotted module path of ``py_file`` relative to ``bench_dir``. + + ``benchmarks/arithmetic.py`` → ``benchmarks.arithmetic`` + ``benchmarks/sub/foo.py`` → ``benchmarks.sub.foo`` + """ + bench_root_name = bench_dir.name # the leading segment users see in benchmark_name + rel = py_file.resolve().relative_to(bench_dir.resolve()) + parts = list(rel.with_suffix("").parts) + if rel.name == "__init__.py": + parts = parts[:-1] + return ".".join([bench_root_name, *parts]) + + +def _is_asv_method(name: str) -> bool: + return any(name.startswith(p) for p in _ASV_METHOD_PREFIXES) + + +def _is_asv_class_name(name: str) -> bool: + return any(name.startswith(p) for p in _ASV_CLASS_PREFIXES) + + +def _source_of(node: ast.AST, text: str) -> str | None: + src = ast.get_source_segment(text, node) + return src.rstrip() + "\n" if src else None + + +def _extract_from_file(py_file: Path, module_dotted: str) -> list[BenchmarkSource]: # noqa: C901 + try: + size = py_file.stat().st_size + except OSError: + return [] + if size > DATASMITH_BENCH_SCRAPE_MAX_FILE_BYTES: + logger.debug("skip oversize bench file %s (%d bytes)", py_file, size) + return [] + + try: + text = py_file.read_text(encoding="utf-8", errors="replace") + except OSError as exc: + logger.debug("unreadable bench file %s: %s", py_file, exc) + return [] + + try: + tree = ast.parse(text, filename=str(py_file)) + except SyntaxError as exc: + logger.debug("syntax error parsing %s: %s", py_file, exc) + return [] + + module_setup: list[str] = [] + for node in tree.body: + if isinstance(node, ast.FunctionDef | ast.AsyncFunctionDef) and node.name in _SETUP_NAMES: + src = _source_of(node, text) + if src: + module_setup.append(src) + + results: list[BenchmarkSource] = [] + + for node in tree.body: + if isinstance(node, ast.FunctionDef | ast.AsyncFunctionDef): + if not _is_asv_method(node.name): + continue + body = _source_of(node, text) + if not body: + continue + setup_src = "\n".join(module_setup) if module_setup else None + results.append( + BenchmarkSource( + benchmark_without_params=f"{module_dotted}.{node.name}", + source=body, + setup_source=setup_src, + ) + ) + + elif isinstance(node, ast.ClassDef): + if not _is_asv_class_name(node.name): + continue + class_setup: list[str] = [] + for child in node.body: + if isinstance(child, ast.FunctionDef | ast.AsyncFunctionDef) and child.name in _SETUP_NAMES: + src = _source_of(child, text) + if src: + class_setup.append(src) + + setup_parts = [*module_setup, *class_setup] + setup_src = "\n".join(setup_parts) if setup_parts else None + + class_src = _source_of(node, text) + for child in node.body: + if not isinstance(child, ast.FunctionDef | ast.AsyncFunctionDef): + continue + if not _is_asv_method(child.name): + continue + method_src = _source_of(child, text) + if not method_src: + continue + # The website shows the method body, but the class header (and + # any class-level attrs like params/param_names) is needed for + # context. Concatenate them when both are available. + if class_src: + combined = class_src + else: + combined = method_src + results.append( + BenchmarkSource( + benchmark_without_params=f"{module_dotted}.{node.name}.{child.name}", + source=combined, + setup_source=setup_src, + ) + ) + + return results + + +def extract_benchmarks(repo_root: Path) -> list[BenchmarkSource]: + """Walk a checked-out repo and return every ASV benchmark's source. + + Returns an empty list if no benchmark directory can be located. Safe to + call on any repo; idempotent. + """ + bench_dir = _resolve_benchmark_dir(repo_root) + if bench_dir is None: + logger.debug("no benchmark dir found under %s", repo_root) + return [] + + out: list[BenchmarkSource] = [] + for py_file in sorted(bench_dir.rglob("*.py")): + module_dotted = _module_path(bench_dir, py_file) + out.extend(_extract_from_file(py_file, module_dotted)) + return out diff --git a/src/datasmith/update/cli.py b/src/datasmith/update/cli.py index aba69cbd..31b2a901 100644 --- a/src/datasmith/update/cli.py +++ b/src/datasmith/update/cli.py @@ -23,6 +23,7 @@ 6: "synthesize_images — Generate Docker build contexts for confirmed performance commits", 7: "harbor_healthcheck — Run synthesized containers through Harbor oracle; record speedups to harbor_runs", 8: "publish — Build, verify, and publish Docker images to DockerHub", + 9: "scrape_benchmark_source — Extract ASV benchmark source code into benchmark_codes (website data sync)", } @@ -55,7 +56,7 @@ def parse_args(argv: list[str] | None = None) -> argparse.Namespace: default=None, action="append", metavar="N", - help="Run only stage N (1-8); repeat to run multiple stages (e.g. --stage 1 --stage 2)", + help="Run only stage N (1-9); repeat to run multiple stages (e.g. --stage 1 --stage 2)", ) parser.add_argument("--dry-run", action="store_true", help="Log what each stage would do without executing") parser.add_argument( diff --git a/src/datasmith/update/pipeline.py b/src/datasmith/update/pipeline.py index 4a41992f..4ce2238f 100644 --- a/src/datasmith/update/pipeline.py +++ b/src/datasmith/update/pipeline.py @@ -76,6 +76,7 @@ def _fetch_repo_descriptions(rows: list[dict[str, Any]]) -> dict[tuple[str, str] "synthesize_images", "harbor_healthcheck", "publish", + "scrape_benchmark_source", ] @@ -244,6 +245,8 @@ async def _run_stage(self, stage_name: str, start_date: str, end_date: str) -> N await self._harbor_healthcheck(start_date, end_date) elif stage_name == "publish": await self._publish(start_date, end_date) + elif stage_name == "scrape_benchmark_source": + await self._scrape_benchmark_source(start_date, end_date) async def _scrape_repos(self) -> None: from datasmith.github.client import GitHubClient @@ -760,6 +763,43 @@ async def _publish(self, start_date: str, end_date: str) -> None: await publish_pipeline(start_date, end_date) + async def _scrape_benchmark_source(self, start_date: str, end_date: str) -> None: + # Pull every (owner, repo, sha) we've successfully synthesized a container + # for. We don't filter by date — the website wants the full corpus of + # benchmark sources, and bench source rarely changes per commit, so + # re-scraping is mostly a no-op on the upsert path. + rows = fetch_all("candidate_containers", select="owner, repo, sha") + + # Dedup by (owner, repo) — one SHA per repo is enough to populate the + # benchmark_codes rows; ASV bench files rarely diverge across commits + # within a single repo and we always keep the newest scrape via the + # last_scraped column. + seen: set[tuple[str, str]] = set() + items: list[dict[str, Any]] = [] + for r in rows: + key = (r["owner"], r["repo"]) + if key in seen: + continue + seen.add(key) + items.append({"owner": r["owner"], "repo": r["repo"], "sha": r["sha"]}) + + logger.info("Scraping benchmark source for %d repos", len(items)) + + if self._dry_run: + self._log_dry_run_summary( + "scrape_benchmark_source", + items, + extra={"Date range": f"{start_date} to {end_date}"}, + ) + return + + from datasmith.runners.scrape_benchmark_source import ScrapeBenchmarkSourceRunner + + runner = ScrapeBenchmarkSourceRunner( + **({"n_concurrent": self._n_concurrent} if self._n_concurrent else {}), + ) + await runner.run(items) + def _get_completed_stages(self) -> list[str]: try: rows = fetch_all("runner_progress", select="runner_name, completed, total") diff --git a/supabase/migrations/00016_benchmark_information.sql b/supabase/migrations/00016_benchmark_information.sql new file mode 100644 index 00000000..07d0e0a8 --- /dev/null +++ b/supabase/migrations/00016_benchmark_information.sql @@ -0,0 +1,63 @@ +-- Per-benchmark speedup measurements from terminal-bench eval runs. +-- Each row is one (benchmark, agent, run) triple: how an LLM agent performed +-- on a single ASV benchmark within a given task (owner/repo/issue), relative +-- to the human-expert oracle. +-- +-- Source: terminal-bench `runs//results.json`, field +-- `parser_extra_metrics.per_benchmark_speedups_by_agent[agent:model][benchmark]`. +-- The human expert is not stored as its own agent row; its baseline lives in +-- `oracle_speedup_vs_nop` on every agent row. + +CREATE TABLE IF NOT EXISTS benchmark_information ( + id BIGSERIAL PRIMARY KEY, + measured_at TIMESTAMPTZ NOT NULL, -- run_metadata.start_time + run_id TEXT NOT NULL, -- e.g. "2026-01-05__09-47-55" + owner TEXT NOT NULL, + repo TEXT NOT NULL, + issue_number INT NOT NULL, + benchmark_name TEXT NOT NULL, -- "benchmarks.ConstructorsSuite.time_point" + agent_name TEXT NOT NULL, -- "openhands", "terminus-2", "oracle", ... + model_name TEXT, -- full model id; NULL for oracle/human + speedup DOUBLE PRECISION NOT NULL, -- human_time / agent_time (1.0 = parity) + agent_speedup_vs_nop DOUBLE PRECISION, -- raw agent/nop from parser + oracle_speedup_vs_nop DOUBLE PRECISION, -- raw oracle/nop from parser + advantage DOUBLE PRECISION, -- parser `advantage` field + significant BOOLEAN, -- parser `significant` field + commit_hash TEXT, + trial_started_at TIMESTAMPTZ, + trial_ended_at TIMESTAMPTZ, + raw_payload JSONB, -- full per-benchmark dict verbatim + created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + CONSTRAINT benchmark_information_unique + UNIQUE (run_id, owner, repo, issue_number, benchmark_name, agent_name, model_name) +); + +CREATE INDEX IF NOT EXISTS idx_benchmark_information_repo + ON benchmark_information (owner, repo, issue_number); + +CREATE INDEX IF NOT EXISTS idx_benchmark_information_agent + ON benchmark_information (agent_name, model_name); + +CREATE INDEX IF NOT EXISTS idx_benchmark_information_speedup + ON benchmark_information (speedup); + +CREATE INDEX IF NOT EXISTS idx_benchmark_information_run + ON benchmark_information (run_id); + +CREATE INDEX IF NOT EXISTS idx_benchmark_information_measured_at + ON benchmark_information (measured_at); + +GRANT ALL ON benchmark_information TO anon, authenticated, service_role; +GRANT USAGE, SELECT ON SEQUENCE benchmark_information_id_seq TO anon, authenticated, service_role; + +-- Public read access (mirrors 00012_public_read_rls.sql / 00015_revoke_anon_select.sql). +-- Anon role can SELECT but not write; service-role bypasses RLS. +ALTER TABLE benchmark_information ENABLE ROW LEVEL SECURITY; + +DROP POLICY IF EXISTS public_read ON benchmark_information; +CREATE POLICY public_read ON benchmark_information + FOR SELECT + TO anon + USING (true); + +REVOKE INSERT, UPDATE, DELETE ON benchmark_information FROM anon; diff --git a/supabase/migrations/00017_benchmark_codes.sql b/supabase/migrations/00017_benchmark_codes.sql new file mode 100644 index 00000000..2d77a2cb --- /dev/null +++ b/supabase/migrations/00017_benchmark_codes.sql @@ -0,0 +1,44 @@ +-- Per-benchmark source code, keyed by (owner, repo, benchmark_without_params). +-- Populated by pipeline stage 9 (scrape_benchmark_source), which checks out +-- each repo at its candidate-container SHA and AST-parses every ASV-style +-- function under `benchmarks/`. The FormulaCode website consumes this via +-- api.formulacode.org to render the player page's "Benchmark code" tabs +-- (~1.3 MB of the website CSV historically). +-- +-- The `benchmark_without_params` value matches the param-stripped form +-- stored in `benchmark_information.benchmark_name` +-- (e.g. "benchmarks.ConstructorsSuite.time_point"), so the website can join +-- benchmark_codes ⋈ benchmark_information on +-- (owner, repo, benchmark_without_params = benchmark_name). +-- +-- If a benchmark moves or is deleted upstream the row is kept; the website +-- falls back gracefully and `last_scraped_sha` is the commit the source +-- was read from. + +CREATE TABLE IF NOT EXISTS benchmark_codes ( + owner TEXT NOT NULL, + repo TEXT NOT NULL, + benchmark_without_params TEXT NOT NULL, + source TEXT NOT NULL, + setup_source TEXT, + last_scraped TIMESTAMPTZ NOT NULL DEFAULT NOW(), + last_scraped_sha TEXT, + PRIMARY KEY (owner, repo, benchmark_without_params) +); + +CREATE INDEX IF NOT EXISTS idx_benchmark_codes_repo + ON benchmark_codes (owner, repo); + +GRANT ALL ON benchmark_codes TO anon, authenticated, service_role; + +-- Public read access (mirrors 00016_benchmark_information.sql). +-- Anon role can SELECT but not write; service-role bypasses RLS. +ALTER TABLE benchmark_codes ENABLE ROW LEVEL SECURITY; + +DROP POLICY IF EXISTS public_read ON benchmark_codes; +CREATE POLICY public_read ON benchmark_codes + FOR SELECT + TO anon + USING (true); + +REVOKE INSERT, UPDATE, DELETE ON benchmark_codes FROM anon; diff --git a/supabase/migrations/00019_candidate_containers_task_id.sql b/supabase/migrations/00019_candidate_containers_task_id.sql new file mode 100644 index 00000000..9df9781f --- /dev/null +++ b/supabase/migrations/00019_candidate_containers_task_id.sql @@ -0,0 +1,16 @@ +-- Add `task_id` to candidate_containers so the FormulaCode website can join +-- its rows to the pipeline by a single column instead of the (owner, repo, +-- issue_number) tuple. Per the project-wide convention we adopted alongside +-- this change, `task_id = issue_number` — the (owner, repo) qualifier is +-- always carried in adjacent columns. A generated column avoids any backfill +-- or write-path changes. +-- +-- candidate_containers is already public-read (00012, 00015) so no grant +-- changes are needed. + +ALTER TABLE candidate_containers + ADD COLUMN IF NOT EXISTS task_id INT + GENERATED ALWAYS AS (issue_number) STORED; + +CREATE INDEX IF NOT EXISTS idx_candidate_containers_task_id + ON candidate_containers (task_id); diff --git a/supabase/migrations/00020_benchmark_information_type.sql b/supabase/migrations/00020_benchmark_information_type.sql new file mode 100644 index 00000000..9e078861 --- /dev/null +++ b/supabase/migrations/00020_benchmark_information_type.sql @@ -0,0 +1,24 @@ +-- Expose `benchmark_type` (time / mem / peakmem / track) on +-- benchmark_information. Derived from the ASV naming convention in +-- `benchmark_name`: classes are `Time*` / `Mem*` / `Peakmem*` and methods +-- are `time_*` / `mem_*` / `peakmem_*` / `track_*`. A STORED generated +-- column backfills automatically; the website can filter without parsing. +-- +-- See https://asv.readthedocs.io/en/stable/writing_benchmarks.html for the +-- convention. The patterns are case-sensitive to match ASV's discovery. + +ALTER TABLE benchmark_information + ADD COLUMN IF NOT EXISTS benchmark_type TEXT + GENERATED ALWAYS AS ( + CASE + WHEN benchmark_name ~ '(^|\.)Peakmem[A-Z]' OR benchmark_name ~ '\.peakmem_' THEN 'peakmem' + WHEN benchmark_name ~ '(^|\.)Mem[A-Z]' OR benchmark_name ~ '\.mem_' THEN 'mem' + -- ASV defines `time_*` (in-process) and `timeraw_*` (fresh subprocess); both are time benchmarks. + WHEN benchmark_name ~ '(^|\.)Time[A-Z]' OR benchmark_name ~ '\.time(raw)?_' THEN 'time' + WHEN benchmark_name ~ '\.track_' THEN 'track' + ELSE NULL + END + ) STORED; + +CREATE INDEX IF NOT EXISTS idx_benchmark_information_type + ON benchmark_information (benchmark_type); diff --git a/tests/github/test_client.py b/tests/github/test_client.py index 4226f483..41c4f239 100644 --- a/tests/github/test_client.py +++ b/tests/github/test_client.py @@ -175,7 +175,7 @@ async def test_get_pr_astropy_16222(self, client: GitHubClient) -> None: record = pr.to_record() assert record is not None assert record.gt_hash == "1ff8068f4378c64c15dc7a37cfd05e6ad1d69f93" - assert record.task_id == "astropy__astropy-16222" + assert record.task_id == 16222 await client.close() diff --git a/tests/github/test_models.py b/tests/github/test_models.py index ae2c26c2..0fef7286 100644 --- a/tests/github/test_models.py +++ b/tests/github/test_models.py @@ -100,7 +100,7 @@ def test_to_record_success(self) -> None: assert record.owner == "myorg" assert record.repo == "mylib" assert record.issue_number == 42 - assert record.task_id == "myorg__mylib-42" + assert record.task_id == 42 assert record.gt_hash == "sha123" assert record.base_commit == "base456" assert record.date == now @@ -268,7 +268,7 @@ def test_record_fields(self) -> None: owner="org", repo="lib", issue_number=1, - task_id="org__lib-1", + task_id=1, gt_hash="abc", patch="diff content", ) diff --git a/tests/publish/test_huggingface.py b/tests/publish/test_huggingface.py index 0134adcb..4e79dc72 100644 --- a/tests/publish/test_huggingface.py +++ b/tests/publish/test_huggingface.py @@ -14,7 +14,7 @@ def _make_record(**kwargs): "owner": "test-org", "repo": "test-repo", "issue_number": 42, - "task_id": "test-org__test-repo-42", + "task_id": 42, "gt_hash": "abc123", } defaults.update(kwargs) diff --git a/tests/publish/test_pipeline.py b/tests/publish/test_pipeline.py index 74fa30f1..19a4e6f8 100644 --- a/tests/publish/test_pipeline.py +++ b/tests/publish/test_pipeline.py @@ -21,7 +21,7 @@ async def test_marks_published(self): mock_record.owner = "org" mock_record.repo = "repo" mock_record.issue_number = 1 - mock_record.task_id = "org__repo-1" + mock_record.task_id = 1 mock_record.container_name = "" with ( @@ -40,9 +40,7 @@ async def test_skips_already_published(self): assert count == 0 async def test_returns_record_count(self): - records = [ - MagicMock(owner="o", repo="r", issue_number=i, task_id=f"o__r-{i}", container_name="") for i in range(5) - ] + records = [MagicMock(owner="o", repo="r", issue_number=i, task_id=i, container_name="") for i in range(5)] mock_client = MagicMock() mock_table = MagicMock() diff --git a/tests/publish/test_records.py b/tests/publish/test_records.py index 87207536..141efd3b 100644 --- a/tests/publish/test_records.py +++ b/tests/publish/test_records.py @@ -13,7 +13,7 @@ def _make_record(**kwargs): "owner": "test-org", "repo": "test-repo", "issue_number": 42, - "task_id": "test-org__test-repo-42", + "task_id": 42, "gt_hash": "abc123", "base_commit": "def456", } @@ -23,7 +23,7 @@ def _make_record(**kwargs): class TestRecordsToParquet: def test_roundtrip(self): - records = [_make_record(), _make_record(issue_number=43, task_id="test-org__test-repo-43")] + records = [_make_record(), _make_record(issue_number=43, task_id=43)] data = records_to_parquet(records) assert len(data) > 0 restored = records_from_parquet(data) @@ -45,9 +45,9 @@ def test_parquet_bytes_valid(self): table = pq.read_table(io.BytesIO(data)) assert table.num_rows == 1 - def test_task_id_format(self): + def test_task_id_is_issue_number(self): r = _make_record() - assert r.task_id == "test-org__test-repo-42" + assert r.task_id == r.issue_number == 42 def test_required_fields_validation(self): with pytest.raises((TypeError, Exception)): @@ -82,7 +82,7 @@ def test_queries_supabase(self): assert len(records) == 1 assert records[0].owner == "org" - assert records[0].task_id == "org__repo-1" + assert records[0].task_id == 1 # Verify fetch_all was called with correct filters (first call = pull_requests) first_call_kwargs = mock_fetch.call_args_list[0] diff --git a/tests/scrape/__init__.py b/tests/scrape/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/scrape/test_benchmark_source.py b/tests/scrape/test_benchmark_source.py new file mode 100644 index 00000000..b9ed0240 --- /dev/null +++ b/tests/scrape/test_benchmark_source.py @@ -0,0 +1,106 @@ +"""Unit tests for the ASV benchmark source extractor.""" + +from __future__ import annotations + +import json +from pathlib import Path + +from datasmith.scrape.benchmark_source import extract_benchmarks + +ARITH_PY = ''' +def setup(*args, **kwargs): + """Module-level setup; should attach to every benchmark in this file.""" + return None + + +class TimeArithmetic: + params = [(100, 1000), (0, 1)] + param_names = ["shape", "axis"] + + def setup(self, shape, axis): + self.df = make_df(shape) + + def time_abs(self, shape, axis): + execute(self.df.abs()) + + def time_neg(self, shape, axis): + execute(-self.df) + + +class MemFootprint: + def mem_total(self): + return self.df.memory_usage().sum() + + +def track_loose_metric(): + return 1.0 + + +def some_helper_that_is_not_a_benchmark(): + return None +''' + + +def _write_repo(tmp_path: Path, files: dict[str, str], conf: dict | None = None) -> Path: + if conf is not None: + (tmp_path / "asv.conf.json").write_text(json.dumps(conf)) + for rel, content in files.items(): + p = tmp_path / rel + p.parent.mkdir(parents=True, exist_ok=True) + p.write_text(content) + return tmp_path + + +def test_extract_class_and_module_benchmarks(tmp_path: Path) -> None: + repo = _write_repo(tmp_path, {"benchmarks/arithmetic.py": ARITH_PY}) + out = extract_benchmarks(repo) + + names = {b.benchmark_without_params for b in out} + assert "benchmarks.arithmetic.TimeArithmetic.time_abs" in names + assert "benchmarks.arithmetic.TimeArithmetic.time_neg" in names + assert "benchmarks.arithmetic.MemFootprint.mem_total" in names + assert "benchmarks.arithmetic.track_loose_metric" in names + assert "benchmarks.arithmetic.some_helper_that_is_not_a_benchmark" not in names + + +def test_setup_source_includes_module_and_class_setup(tmp_path: Path) -> None: + repo = _write_repo(tmp_path, {"benchmarks/arithmetic.py": ARITH_PY}) + out = {b.benchmark_without_params: b for b in extract_benchmarks(repo)} + + time_abs = out["benchmarks.arithmetic.TimeArithmetic.time_abs"] + assert time_abs.setup_source is not None + assert "def setup(*args, **kwargs)" in time_abs.setup_source # module-level + assert "self.df = make_df(shape)" in time_abs.setup_source # class-level + + track = out["benchmarks.arithmetic.track_loose_metric"] + assert track.setup_source is not None + assert "def setup(*args, **kwargs)" in track.setup_source + # No class-level setup for module-level function: + assert "self.df = make_df" not in track.setup_source + + +def test_honours_asv_conf_benchmark_dir(tmp_path: Path) -> None: + repo = _write_repo( + tmp_path, + {"asv_bench/benchmarks/arithmetic.py": ARITH_PY}, + conf={"benchmark_dir": "asv_bench/benchmarks"}, + ) + out = extract_benchmarks(repo) + names = {b.benchmark_without_params for b in out} + # Module prefix should be the leaf dir name from benchmark_dir, not the full path. + assert "benchmarks.arithmetic.TimeArithmetic.time_abs" in names + + +def test_no_benchmark_dir_returns_empty(tmp_path: Path) -> None: + repo = _write_repo(tmp_path, {"src/lib.py": "x = 1"}) + assert extract_benchmarks(repo) == [] + + +def test_class_methods_include_class_header_in_source(tmp_path: Path) -> None: + """Website needs class-level attrs like `params` to render meaningful code blocks.""" + repo = _write_repo(tmp_path, {"benchmarks/arithmetic.py": ARITH_PY}) + out = {b.benchmark_without_params: b for b in extract_benchmarks(repo)} + body = out["benchmarks.arithmetic.TimeArithmetic.time_abs"].source + assert "class TimeArithmetic" in body + assert "params = [(100, 1000), (0, 1)]" in body + assert "def time_abs" in body diff --git a/tests/test_website_snippets.py b/tests/test_website_snippets.py index b626ee8f..55ca229b 100644 --- a/tests/test_website_snippets.py +++ b/tests/test_website_snippets.py @@ -558,11 +558,11 @@ def test_ds_update_multiple_stages(self) -> None: args = _parse(["--start-date", "2026-02-01", "--end-date", "2026-03-01", "--stage", "5", "--stage", "6"]) assert args.stage == [5, 6] - def test_pipeline_has_8_stages(self) -> None: - """Website documents 8 pipeline stages.""" + def test_pipeline_has_9_stages(self) -> None: + """Website documents 9 pipeline stages.""" from datasmith.update.pipeline import STAGES - assert len(STAGES) == 8 + assert len(STAGES) == 9 # ---------------------------------------------------------------------------