|
| 1 | +"""Read per-task operator declarations from ``formulacode_task_overrides``. |
| 2 | +
|
| 3 | +The overrides table holds facts about a task that cannot be derived from the |
| 4 | +repo or the PR. Two of its columns are invariant inputs, and both needed a |
| 5 | +host-side reader because the table is RLS-locked with no ``anon`` grant -- |
| 6 | +neither the trial container nor the image build can query it themselves: |
| 7 | +
|
| 8 | +``benchmark_dest`` |
| 9 | + Passed into the image build as the ``BENCHMARK_DEST`` build arg, which |
| 10 | + ``docker_build_run.sh`` turns into the ``benchmark_dest_present_post_clean`` |
| 11 | + breadcrumb. Input to the FATAL ``benchmark_dest_missing`` invariant. |
| 12 | +
|
| 13 | +``expected_n`` |
| 14 | + Injected into the trial container as ``FORMULACODE_EXPECTED_N`` via the |
| 15 | + task's ``[verifier.env]`` section. Input to the ``dilution_ratio`` |
| 16 | + invariant. |
| 17 | +
|
| 18 | +Lookup failure is never fatal. Stage 6 and stage 7 must run on a machine where |
| 19 | +the migration has not been applied, and "no override" is a legitimate state for |
| 20 | +all but a handful of tasks -- the corresponding invariants then skip, which is |
| 21 | +the correct answer rather than a degraded one. |
| 22 | +""" |
| 23 | + |
| 24 | +from __future__ import annotations |
| 25 | + |
| 26 | +from typing import Any |
| 27 | + |
| 28 | +from datasmith.utils.core import get_logger |
| 29 | +from datasmith.utils.db import fetch_all |
| 30 | + |
| 31 | +logger = get_logger("utils.overrides") |
| 32 | + |
| 33 | +TABLE = "formulacode_task_overrides" |
| 34 | + |
| 35 | +# The canonical task identity, per CLAUDE.md: (owner, repo, issue_number). |
| 36 | +TaskKey = tuple[str, str, int] |
| 37 | + |
| 38 | + |
| 39 | +_ALL_CACHE: dict[TaskKey, dict[str, Any]] | None = None |
| 40 | + |
| 41 | + |
| 42 | +def all_overrides(*, refresh: bool = False) -> dict[TaskKey, dict[str, Any]]: |
| 43 | + """Every override row, read once per process. |
| 44 | +
|
| 45 | + The table is tiny (5 rows today, and it is hand-maintained), while its |
| 46 | + consumers sit in per-item loops -- stage 6 synthesises one PR at a time and |
| 47 | + enqueues neighbours mid-flight, so there is no single point at which the |
| 48 | + full task set is known. Caching turns what would be one round-trip per |
| 49 | + task into one per run. |
| 50 | +
|
| 51 | + A failed read caches ``{}`` too: if the table is absent, it will still be |
| 52 | + absent on the next call, and retrying per item would be a slow way to |
| 53 | + reach the same answer. |
| 54 | + """ |
| 55 | + global _ALL_CACHE |
| 56 | + if _ALL_CACHE is not None and not refresh: |
| 57 | + return _ALL_CACHE |
| 58 | + |
| 59 | + try: |
| 60 | + rows = fetch_all(TABLE, select="owner, repo, issue_number, benchmark_dest, expected_n") |
| 61 | + except Exception: |
| 62 | + logger.warning( |
| 63 | + "Could not read %s; proceeding without overrides. The " |
| 64 | + "benchmark_dest_missing and dilution_ratio invariants will skip.", |
| 65 | + TABLE, |
| 66 | + exc_info=True, |
| 67 | + ) |
| 68 | + _ALL_CACHE = {} |
| 69 | + return _ALL_CACHE |
| 70 | + |
| 71 | + out: dict[TaskKey, dict[str, Any]] = {} |
| 72 | + for row in rows or []: |
| 73 | + try: |
| 74 | + out[(row["owner"], row["repo"], int(row["issue_number"]))] = row |
| 75 | + except (KeyError, TypeError, ValueError): |
| 76 | + continue |
| 77 | + _ALL_CACHE = out |
| 78 | + return out |
| 79 | + |
| 80 | + |
| 81 | +def fetch_overrides(tasks: list[TaskKey]) -> dict[TaskKey, dict[str, Any]]: |
| 82 | + """Return ``{(owner, repo, issue_number): override_row}`` for *tasks*. |
| 83 | +
|
| 84 | + Tasks with no override row are simply absent from the mapping. A missing |
| 85 | + table, an unreachable database, or any other read failure yields ``{}`` |
| 86 | + with a warning -- callers treat that identically to "no overrides exist", |
| 87 | + so a fresh checkout behaves the same as a populated one minus the extra |
| 88 | + checks. |
| 89 | + """ |
| 90 | + if not tasks: |
| 91 | + return {} |
| 92 | + |
| 93 | + wanted = set(tasks) |
| 94 | + return {key: row for key, row in all_overrides().items() if key in wanted} |
| 95 | + |
| 96 | + |
| 97 | +def benchmark_dest_for(overrides: dict[TaskKey, dict[str, Any]], key: TaskKey) -> str: |
| 98 | + """The declared benchmark path for *key*, or "" when undeclared. |
| 99 | +
|
| 100 | + Empty string is deliberate: ``docker_build_run.sh`` gates its breadcrumb on |
| 101 | + a non-empty ``$BENCHMARK_DEST``, so "" keeps the FATAL invariant skipping |
| 102 | + rather than firing against a value nobody declared. |
| 103 | + """ |
| 104 | + row = overrides.get(key) or {} |
| 105 | + return str(row.get("benchmark_dest") or "") |
| 106 | + |
| 107 | + |
| 108 | +def expected_n_for(overrides: dict[TaskKey, dict[str, Any]], key: TaskKey) -> int | None: |
| 109 | + """The declared benchmark count for *key*, or None when undeclared. |
| 110 | +
|
| 111 | + None (not 0) is the "not judged yet" signal: the dilution invariant skips |
| 112 | + on None, and 0 would be a live comparison against a number nobody chose. |
| 113 | + """ |
| 114 | + row = overrides.get(key) or {} |
| 115 | + raw = row.get("expected_n") |
| 116 | + if raw is None: |
| 117 | + return None |
| 118 | + try: |
| 119 | + return int(raw) |
| 120 | + except (TypeError, ValueError): |
| 121 | + return None |
0 commit comments