Skip to content

Registry operations

ops

Registry lifecycle operations — the manage-only operator surface.

A deployment's registry is not write-once. Runs accumulate, experiments end, a bad run wants deleting, and the operator wants to know what is actually in there before touching any of it. This module is that surface: eight bounded verbs over exactly one registry — the one Settings.registry_dataset_ref resolves to.

Manage only, by design. There is no wipe here. A full teardown is bq rm -r -f <dataset> or the BigQuery console — nobody needs us to write that, and shipping it invites the accident. Every verb below touches only what the caller named:

verb what it does
init create this registry's tables + views (idempotent; does not touch the source panel)
doctor read-only health report — row counts, runs stuck RUNNING, orphaned artifacts
close_runs finalize abandoned RUNNING headers to what their job rows already imply
drop_run delete named run(s) from every tier: rows, GCS artifacts, and BQML model objects
sweep_orphans delete artifact prefixes under this registry with no run_registry row
reap_clusters delete Ray clusters belonging to this registry whose run has finished
snapshot BigQuery table snapshots of the five registry tables (cheap, expirable)
export dump the registry to GCS (Parquet or newline-delimited JSON) for offline analysis

reap_clusters is the odd one out — it deletes compute, not registry contents — and it is here because the registry is what decides. A Vertex Ray cluster has no idle TTL to set (the Dataproc clusters do, which is why only the Ray side leaks), so the only way to know a cluster is abandoned is to ask whether its run is over, and that answer lives in run_registry. Scoped exactly like sweep_orphans: clusters are labelled with the registry that owns them, so another deployment in the same project is outside its world. The policy and the Vertex calls are ray_reaper's; this module contributes the verb and all_header_statuses.

The ordering rule. A registry row is the only index of which GCS objects exist. Delete the rows first and the artifacts become unidentifiable garbage forever — which is exactly how the old reset path accumulated an unbounded orphan pile. So every destructive verb goes:

1. enumerate artifact prefixes   <- FIRST, while the index still exists
2. delete the GCS objects
3. delete the rows / model objects

sweep_orphans is the cleanup for every prefix that was stranded before that rule existed. It is correctly scopeable only because the artifact root now carries the registry key (Settings.artifact_root): a prefix under this root with no run_registry row is garbage by definition, and no other registry sharing the bucket can be caught by it.

Safety model (the shape reset and the probe's cancel path already established): preview is the default, yes=True executes, and the preview prints the exact objects — run ids, row counts, object counts, byte totals — that will be touched. A mutating verb also refuses to run against a registry with an in-flight run unless forced; use monitor(probe=True) to tell a live run from a dead row still marked RUNNING.

close_runs is the deliberate exception to that refusal, because in-flight-looking runs are exactly what it exists to repair. It earns the exception by being the one mutating verb that deletes nothing: it writes a header status and nothing else, and it refuses any run whose job rows are not already all terminal — so it can only ever record a conclusion the rows had already reached.

This module is the policy, not the plumbing. Step 1 of the ordering rule — reading the GCS layout back to find out which objects a run owns — is artifacts, which owns that layout in both directions (it is the module that composed the paths on the way in). What lives here is the decision of when to enumerate, what the answer blocks, and in which order the deletes go.

Split along the pure/I-O seam, like the rest of the registry package: the planners and SQL renderers below are pure strings and dataclasses (tested offline, no client), and the six verbs are thin I/O wrappers that execute what a planner produced.

Public surface: the verbs init, doctor, drop_run, sweep_orphans, snapshot, export; the plan/report types DropPlan, SweepPlan, DoctorReport, TableStat; and the pure helpers blocking_runs, render_delete_rows, render_snapshot_sql, render_export_sql, format_plan, format_doctor. The artifact-prefix helpers the destructive verbs run on (ArtifactPrefix, split_gcs_uri, run_id_from_blob, orphan_run_ids, list_prefixes, delete_prefixes) are artifacts'.

BorrowedModelSource dataclass

BorrowedModelSource(
    target_model_type,
    source_run_id,
    source_model_type,
    model_spec_hash,
    snapshot_millis,
)

Provenance record for a model specification borrowed from a completed prior run.

CloseCandidate dataclass

CloseCandidate(
    run_id, header_status, job_statuses, new_status, reason
)

One run close_runs looked at: its stuck header, its jobs, and the verdict.

new_status is None when the run must be left alone, and reason always says why — a skipped run is as much of an answer as a closed one, so both are reported.

ClosePlan dataclass

ClosePlan(registry, candidates=(), unknown=())

Exactly which stuck headers close_runs would finalize, and to what.

DoctorReport dataclass

DoctorReport(
    registry,
    artifact_root,
    tables=(),
    views=(),
    live_runs=(),
    orphans=(),
    notes=(),
)

A read-only picture of one registry: what is in it, and what looks wrong.

healthy property

healthy

No missing tables, no in-flight runs, no orphaned artifacts.

DropPlan dataclass

DropPlan(
    registry,
    run_ids=(),
    prefixes=(),
    models=(),
    blocked=(),
    unknown=(),
)

Exactly what drop_run would delete — the preview, and the thing drop_run executes.

blocked is the subset of run_ids whose header is still in a LIVE_STATUSES state. A non-empty blocked stops the whole plan (not just those runs): deleting half a batch because the other half was busy is a worse outcome than doing nothing and saying why.

is_empty property

is_empty

Nothing to do — every named run is unknown to this registry.

SweepPlan dataclass

SweepPlan(registry, root, prefixes=(), known_runs=0)

Artifact prefixes under this registry's root with no run_registry row.

TableStat dataclass

TableStat(table, rows)

One registry table's row count, or None rows if the table is absent.

all_header_statuses

all_header_statuses(settings)

{run_id: latest header status} for every run in this registry (one query).

The whole-registry form of _statuses, for a caller that has to answer "is this run finished?" without knowing the run ids up front. ray_reaper is that caller: it reads run ids off cluster names, which are clamped to 63 characters and so may hold only a prefix of the real id — there is nothing to bind into a WHERE run_id IN UNNEST(...), and prefix matching happens in pure code against this map instead. One header row per run makes it a small read even on a busy registry.

Source code in src/scale_forecasting/registry/ops.py
def all_header_statuses(
    settings: Settings,
) -> dict[str, str | None]:  # pragma: no cover - GCP I/O, @gcp smoke
    """``{run_id: latest header status}`` for **every** run in this registry (one query).

    The whole-registry form of `_statuses`, for a caller that has to answer "is this run finished?"
    without knowing the run ids up front. `ray_reaper` is that caller: it reads run ids off cluster
    *names*, which are clamped to 63 characters and so may hold only a prefix of the real id — there
    is nothing to bind into a ``WHERE run_id IN UNNEST(...)``, and prefix matching happens in pure
    code against this map instead. One header row per run makes it a small read even on a busy
    registry.
    """

    from ..errors import RegistryError

    sql = (
        "SELECT run_id, ARRAY_AGG(status ORDER BY created_at DESC LIMIT 1)[OFFSET(0)] AS status\n"
        f"FROM `{settings.registry_table_ref('run_registry')}`\nGROUP BY run_id"
    )
    client = bigquery_client(settings.project_id)
    try:
        return {str(r["run_id"]): r["status"] for r in client.query(sql).result()}
    except Exception as exc:  # re-raised with registry context
        raise RegistryError(
            f"could not read run statuses from {settings.registry_dataset_ref}: {exc}"
        ) from exc

blocking_runs

blocking_runs(status_by_run)

The subset of runs whose header is in a LIVE_STATUSES state, sorted.

A None status (no header row) is not blocking: an unknown run cannot be in flight, and DropPlan.unknown reports it separately.

Source code in src/scale_forecasting/registry/ops.py
def blocking_runs(status_by_run: dict[str, str | None]) -> tuple[str, ...]:
    """The subset of runs whose header is in a `LIVE_STATUSES` state, sorted.

    A ``None`` status (no header row) is *not* blocking: an unknown run cannot be in flight, and
    `DropPlan.unknown` reports it separately.
    """
    return tuple(
        sorted(r for r, s in status_by_run.items() if s is not None and s.upper() in LIVE_STATUSES)
    )

borrow_model_results

borrow_model_results(
    cfg, run_id, borrowed, *, settings=None
)

Copy forecast_predictions, backtest_oof, and forecast_metadata for reuse.

Source code in src/scale_forecasting/registry/ops.py
def borrow_model_results(
    cfg: object,
    run_id: str,
    borrowed: dict[str, BorrowedModelSource],
    *,
    settings: Settings | None = None,
) -> None:  # pragma: no cover - GCP I/O, @gcp smoke
    """Copy ``forecast_predictions``, ``backtest_oof``, and ``forecast_metadata`` for reuse."""
    if not borrowed:
        return
    from google.cloud import bigquery

    from ..config import RunConfig
    from ..features import variant_diagnostics_for

    assert isinstance(cfg, RunConfig)
    resolved = _resolved(settings)
    client = bigquery_client(resolved.project_id)
    stmts = render_borrow_model_sql(resolved.registry_dataset_ref)

    for target_m, src in borrowed.items():
        vdiag = variant_diagnostics_for(target_m, cfg)
        qparams = [
            bigquery.ScalarQueryParameter("target_run_id", "STRING", run_id),
            bigquery.ScalarQueryParameter("target_model_type", "STRING", target_m),
            bigquery.ScalarQueryParameter("source_run_id", "STRING", src.source_run_id),
            bigquery.ScalarQueryParameter("source_model_type", "STRING", src.source_model_type),
            bigquery.ScalarQueryParameter("base_model", "STRING", str(vdiag["base_model"])),
            bigquery.ScalarQueryParameter(
                "scope_capability", "STRING", str(vdiag["scope_capability"])
            ),
            bigquery.ScalarQueryParameter("training_mode", "STRING", str(vdiag["training_mode"])),
            bigquery.ScalarQueryParameter("cov_profile", "STRING", str(vdiag["cov_profile"])),
            bigquery.ScalarQueryParameter("is_univariate", "BOOL", bool(vdiag["is_univariate"])),
            bigquery.ScalarQueryParameter(
                "covariates_used_json", "STRING", json.dumps(vdiag["covariates_used"])
            ),
            bigquery.ScalarQueryParameter("model_spec_hash", "STRING", src.model_spec_hash),
        ]
        job_cfg = bigquery.QueryJobConfig(query_parameters=qparams)
        for table_name in ("forecast_predictions", "backtest_oof", "forecast_metadata"):
            client.query(stmts[table_name], job_config=job_cfg).result()
        _log.info(
            "run %s: borrowed model '%s' from prior run %s (model '%s', spec_hash=%s)",
            run_id,
            target_m,
            src.source_run_id,
            src.source_model_type,
            src.model_spec_hash,
        )

close_runs

close_runs(run_ids=None, *, settings=None, yes=False)

Finalize abandoned headers to the status their own job rows already imply.

The verb that was missing. A run whose driver died after writing its header leaves a RUNNING row forever, and nothing else could clear it: --cancel stamps CANCELLED over work that actually completed, drop_run destroys real predictions to fix a status field, and --probe only reads. This closes the header and touches nothing else — no artifact, no job row, no prediction is written or deleted, so the worst case is a header status that is wrong in a different way and can be re-closed.

It never decides that a job is dead. A run with a non-terminal job row is skipped with a reason, because only a runtime probe can tell a live job from a stale row — that is probes.reconcile's job, and conflating the two is how an operator "cleans up" a running job. So the safe order is sdk.Forecaster.monitor (probe=True) first, then this.

What it closes to is decided against the run's own plan, not just its rows — see roll_up_against_plan. Reading the rows alone cannot see a family that was planned and never submitted, which is precisely the hole an abandoned run leaves, so a run whose last family happened to finish used to close COMPLETED while claiming an ensemble that never ran.

Preview is the default; yes=True executes. Returns the plan either way.

Source code in src/scale_forecasting/registry/ops.py
def close_runs(
    run_ids: Sequence[str] | None = None, *, settings: Settings | None = None, yes: bool = False
) -> ClosePlan:  # pragma: no cover - GCP I/O, @gcp smoke
    """Finalize abandoned headers to the status their own job rows already imply.

    **The verb that was missing.** A run whose driver died after writing its header leaves a
    ``RUNNING`` row forever, and nothing else could clear it: ``--cancel`` stamps ``CANCELLED`` over
    work that actually completed, `drop_run` destroys real predictions to fix a status field, and
    ``--probe`` only reads. This closes the header and touches nothing else — no artifact, no job
    row, no prediction is written or deleted, so the worst case is a header status that is wrong in
    a different way and can be re-closed.

    **It never decides that a job is dead.** A run with a non-terminal job row is skipped with a
    reason, because only a runtime probe can tell a live job from a stale row — that is
    `probes.reconcile`'s job, and conflating the two is how an operator "cleans up" a running job.
    So the safe order is `sdk.Forecaster.monitor` (``probe=True``) first, then this.

    **What it closes to is decided against the run's own plan,** not just its rows — see
    `roll_up_against_plan`. Reading the rows alone cannot see a family that was planned and never
    submitted, which is precisely the hole an abandoned run leaves, so a run whose last family
    happened to finish used to close ``COMPLETED`` while claiming an ensemble that never ran.

    Preview is the default; ``yes=True`` executes. Returns the plan either way.
    """
    from .header import update_header

    resolved = _resolved(settings)
    plan = plan_close_runs(run_ids, settings=resolved)
    _log.warning("%s", format_close_plan(plan))

    if plan.is_empty:
        _log.warning("nothing to close in %s", plan.registry)
        return plan
    if not yes:
        _log.warning("DRY RUN — no header changed. Re-run with yes=True to execute.")
        return plan

    for candidate in plan.closable:
        update_header(candidate.run_id, settings=resolved, status=candidate.new_status)
        _log.warning(
            "closed %s: %s → %s (%s)",
            candidate.run_id,
            candidate.header_status,
            candidate.new_status,
            candidate.reason,
        )
    _log.warning("closed %d header(s) in %s", len(plan.closable), plan.registry)
    return plan

doctor

doctor(*, settings=None)

A read-only health report for this registry. Touches nothing.

Answers the three questions an operator actually has: is it there (tables present, row counts per tier), is anything stuck (headers still in a LIVE_STATUSES state), and is anything leaking (artifact prefixes under this registry's root with no run_registry row). The last one is the state every pre-isolation reset left behind, and it is only answerable now that the artifact path carries the registry key.

Source code in src/scale_forecasting/registry/ops.py
def doctor(*, settings: Settings | None = None) -> DoctorReport:  # pragma: no cover - GCP I/O
    """A read-only health report for this registry. Touches nothing.

    Answers the three questions an operator actually has: is it *there* (tables present, row counts
    per tier), is anything *stuck* (headers still in a `LIVE_STATUSES` state), and is anything
    *leaking* (artifact prefixes under this registry's root with no ``run_registry`` row). The last
    one is the state every pre-isolation reset left behind, and it is only answerable now that the
    artifact path carries the registry key.
    """

    from .views import render_create_views

    resolved = _resolved(settings)
    client = bigquery_client(resolved.project_id)
    notes: list[str] = []

    stats: list[TableStat] = []
    for name in REGISTRY_TABLE_NAMES:
        ref = resolved.registry_table_ref(name)
        try:
            rows = list(client.query(f"SELECT COUNT(*) AS n FROM `{ref}`").result())
            stats.append(TableStat(name, int(rows[0]["n"])))
        except Exception:  # noqa: BLE001 - an absent table is a finding, not a crash
            stats.append(TableStat(name, None))

    live: tuple[tuple[str, str], ...] = ()
    orphans: tuple[ArtifactPrefix, ...] = ()
    if all(s.exists for s in stats):
        statuses = ", ".join(f"'{s}'" for s in sorted(LIVE_STATUSES))
        sql = (
            "SELECT run_id, status FROM (\n"
            "  SELECT run_id, status, ROW_NUMBER() OVER "
            "(PARTITION BY run_id ORDER BY created_at DESC) AS rn\n"
            f"  FROM `{resolved.registry_table_ref('run_registry')}`\n"
            f") WHERE rn = 1 AND status IN ({statuses}) ORDER BY run_id"
        )
        live = tuple((str(r["run_id"]), str(r["status"])) for r in client.query(sql).result())
        known = _known_run_ids(resolved)
        prefixes = artifacts.list_prefixes(resolved)
        orphans = tuple(prefixes[r] for r in artifacts.orphan_run_ids(prefixes, known))
    else:
        notes.append("skipped the run/orphan checks — run init first, some tables are missing")

    return DoctorReport(
        registry=resolved.registry_dataset_ref,
        artifact_root=resolved.artifact_root,
        tables=tuple(stats),
        views=tuple(render_create_views(resolved.registry_dataset_ref)),
        live_runs=live,
        orphans=orphans,
        notes=tuple(notes),
    )

drop_run

drop_run(run_ids, *, settings=None, yes=False, force=False)

Delete named run(s) from every tier: GCS artifacts, BQML models, then registry rows.

The daily verb, and the one that absorbs "prune" — pass as many ids as you like. Preview is the default: without yes=True this plans, logs the preview, and returns without touching anything. force=True overrides the in-flight refusal; check with sdk.Forecaster.monitor (probe=True) first, because a registry row marked RUNNING can mean either a live job or a dead one that never got to write its status.

Order is artifacts → models → rows, never the reverse: the rows are the only index of which objects belong to the run, so deleting them first strands the artifacts permanently.

Returns the executed plan. Raises RegistryError if the delete fails — most usefully when BigQuery rejects a DELETE against rows still in the streaming buffer, which happens for roughly 90 minutes after a run finishes writing.

Source code in src/scale_forecasting/registry/ops.py
def drop_run(
    run_ids: Sequence[str],
    *,
    settings: Settings | None = None,
    yes: bool = False,
    force: bool = False,
) -> DropPlan:  # pragma: no cover - GCP I/O, @gcp smoke
    """Delete named run(s) from every tier: GCS artifacts, BQML models, then registry rows.

    The daily verb, and the one that absorbs "prune" — pass as many ids as you like. Preview is the
    default: without ``yes=True`` this plans, logs the preview, and returns without touching
    anything. ``force=True`` overrides the in-flight refusal; check with `sdk.Forecaster.monitor`
    (``probe=True``) first, because a registry row marked ``RUNNING`` can mean either a live job or
    a dead one that never got to write its status.

    Order is artifacts → models → rows, never the reverse: the rows are the only index of which
    objects belong to the run, so deleting them first strands the artifacts permanently.

    Returns the executed plan. Raises `RegistryError` if the delete fails — most usefully when
    BigQuery rejects a ``DELETE`` against rows still in the streaming buffer, which happens for
    roughly 90 minutes after a run finishes writing.
    """
    from google.cloud import bigquery

    from ..errors import RegistryError

    resolved = _resolved(settings)
    plan = plan_drop_run(run_ids, settings=resolved)
    _log.warning("%s", format_plan(plan))

    if plan.blocked and not force:
        raise RegistryError(
            f"refusing to drop: {', '.join(plan.blocked)} still in flight. Check with "
            "monitor(probe=True) — a RUNNING row can also be a dead job — then pass force=True."
        )
    if plan.is_empty:
        _log.warning("nothing to drop — none of the named runs are in %s", plan.registry)
        return plan
    if not yes:
        _log.warning("DRY RUN — nothing deleted. Re-run with yes=True to execute.")
        return plan

    # 1+2. Artifacts first, while the rows still say which objects exist.
    deleted = artifacts.delete_prefixes(resolved, plan.run_ids)
    _log.warning("deleted %d artifact objects", deleted)

    # 3a. BQML model objects (invisible to the registry — matched by name).
    client = bigquery_client(resolved.project_id)
    for model_id in plan.models:
        ref = f"{resolved.registry_dataset_ref}.{model_id}"
        try:
            client.query(f"DROP MODEL IF EXISTS `{ref}`;").result()
        except Exception as exc:  # noqa: BLE001 - best-effort, reported and continued
            _log.warning("could not drop model %s: %s", ref, exc)

    # 3b. Rows.
    job_config = bigquery.QueryJobConfig(
        query_parameters=[bigquery.ArrayQueryParameter("run_ids", "STRING", list(plan.run_ids))]
    )
    for name, statement in render_delete_rows(resolved.registry_dataset_ref).items():
        try:
            client.query(statement, job_config=job_config).result()
        except Exception as exc:  # re-raised with table context
            raise RegistryError(
                f"drop_run failed deleting from {name}: {exc}. If the run just finished, its rows "
                "may still be in the Storage Write API streaming buffer (~90 min), which BigQuery "
                "will not let a DELETE touch — wait and retry."
            ) from exc
    _log.warning("dropped %d run(s) from %s", len(plan.run_ids), plan.registry)
    return plan

export

export(destination_root, *, settings=None, fmt='PARQUET')

Export each registry table to GCS; return {table: destination prefix}.

Non-destructive but overwriting — each table's destination prefix is replaced, which is what makes re-exporting to a stable path idempotent. Use a run-stamped destination_root if you want a history of exports rather than a current one.

Source code in src/scale_forecasting/registry/ops.py
def export(
    destination_root: str, *, settings: Settings | None = None, fmt: str = "PARQUET"
) -> dict[str, str]:  # pragma: no cover - GCP I/O, @gcp smoke
    """Export each registry table to GCS; return ``{table: destination prefix}``.

    Non-destructive but **overwriting** — each table's destination prefix is replaced, which is what
    makes re-exporting to a stable path idempotent. Use a run-stamped ``destination_root`` if you
    want a history of exports rather than a current one.
    """

    from ..errors import RegistryError

    resolved = _resolved(settings)
    statements = render_export_sql(resolved.registry_dataset_ref, destination_root, fmt=fmt)
    client = bigquery_client(resolved.project_id)
    root = destination_root.rstrip("/")
    out: dict[str, str] = {}
    for name, statement in statements.items():
        try:
            client.query(statement).result()
        except Exception as exc:  # re-raised with table context
            raise RegistryError(f"export failed for {name}: {exc}") from exc
        out[name] = f"{root}/{name}/"
    _log.info("exported %d tables to %s", len(out), root)
    return out

find_reusable_models

find_reusable_models(
    cfg, run_id, snapshot_millis, *, settings=None
)

Discover completed prior runs in run_registry whose cells can satisfy cfg.models.

Source code in src/scale_forecasting/registry/ops.py
def find_reusable_models(
    cfg: object,
    run_id: str,
    snapshot_millis: int | None,
    *,
    settings: Settings | None = None,
) -> dict[str, BorrowedModelSource]:  # pragma: no cover - GCP I/O, @gcp smoke
    """Discover completed prior runs in ``run_registry`` whose cells can satisfy ``cfg.models``."""
    from google.cloud import bigquery

    from ..config import RunConfig
    from ..engines.bigquery_names import source_ref

    assert isinstance(cfg, RunConfig)
    explicit_source = cfg.compute.reuse_from_run_id
    if not cfg.compute.reuse_prior_runs and explicit_source is None:
        return {}

    resolved = _resolved(settings)
    client = bigquery_client(resolved.project_id)
    ds = resolved.registry_dataset_ref

    source_last_mod_ms: int | None = None
    if explicit_source is None and snapshot_millis is None:
        try:
            tbl = client.get_table(source_ref(cfg, resolved.dataset_ref))
            if tbl.modified is not None:
                source_last_mod_ms = int(tbl.modified.timestamp() * 1000)
        except Exception as exc:  # noqa: BLE001 - best-effort table metadata check
            _log.debug("reuse_prior_runs: could not read source table modified time: %s", exc)
        if source_last_mod_ms is None:
            return {}

    if explicit_source is not None:
        where_clause = "WHERE run_id = @explicit_run_id AND status = 'COMPLETED'"
    elif snapshot_millis is not None:
        where_clause = (
            "WHERE run_id != @run_id AND status = 'COMPLETED' "
            "AND snapshot_millis = @snapshot_millis"
        )
    else:
        where_clause = (
            "WHERE run_id != @run_id AND status = 'COMPLETED' "
            "AND snapshot_millis >= @source_last_mod_ms"
        )
    sql_headers = (
        f"SELECT run_id, status, snapshot_millis, TO_JSON_STRING(raw_config) AS raw_config, "
        f"created_at FROM `{ds}.run_registry` {where_clause} "
        "QUALIFY ROW_NUMBER() OVER (PARTITION BY run_id ORDER BY created_at DESC NULLS LAST) = 1 "
        "ORDER BY created_at DESC NULLS LAST LIMIT 25"
    )
    params = [
        bigquery.ScalarQueryParameter("run_id", "STRING", run_id),
        bigquery.ScalarQueryParameter("explicit_run_id", "STRING", explicit_source),
        bigquery.ScalarQueryParameter("snapshot_millis", "INT64", snapshot_millis),
        bigquery.ScalarQueryParameter("source_last_mod_ms", "INT64", source_last_mod_ms),
    ]
    try:
        hdr_rows = [
            dict(r)
            for r in client.query(
                sql_headers, job_config=bigquery.QueryJobConfig(query_parameters=params)
            ).result()
        ]
    except Exception as exc:  # noqa: BLE001 - best-effort reuse discovery; fall back to normal fit
        _log.warning("reuse_prior_runs: header lookup failed (%s); running all models fresh", exc)
        return {}
    if not hdr_rows:
        return {}

    cand_ids = [str(r["run_id"]) for r in hdr_rows if r.get("run_id")]
    sql_avail = (
        "SELECT m.run_id, m.model_type "
        f"FROM `{ds}.forecast_metadata` AS m "
        f"JOIN (SELECT DISTINCT run_id, model_type FROM `{ds}.forecast_predictions` "
        "WHERE run_id IN UNNEST(@cand_ids) AND ensemble_id IS NULL) AS p "
        "USING (run_id, model_type) "
        "WHERE m.run_id IN UNNEST(@cand_ids) AND m.fold_id IS NULL AND m.ensemble_id IS NULL "
        "GROUP BY m.run_id, m.model_type "
        "HAVING COUNTIF(m.cell_status = 'ok') > 0 AND COUNTIF(m.cell_status = 'error') = 0"
    )
    avail_params = [bigquery.ArrayQueryParameter("cand_ids", "STRING", cand_ids)]
    available_by_run: dict[str, set[str]] = {}
    try:
        for r in client.query(
            sql_avail, job_config=bigquery.QueryJobConfig(query_parameters=avail_params)
        ).result():
            available_by_run.setdefault(str(r["run_id"]), set()).add(str(r["model_type"]))
    except Exception as exc:  # noqa: BLE001 - best-effort availability check
        _log.warning("reuse_prior_runs: cell availability check failed (%s)", exc)
        return {}

    return match_reusable_models(
        cfg,
        run_id,
        snapshot_millis,
        hdr_rows,
        available_by_run,
        source_last_modified_millis=source_last_mod_ms,
    )

format_close_plan

format_close_plan(plan)

Render a ClosePlan as the operator-facing preview — the exact text a dry run prints.

Skipped runs are printed as prominently as closable ones. The reason a header is stuck is the thing the operator actually needs to see: "still has a RUNNING job" means go probe it, and burying that under a count of what was closed is how the one run that mattered gets missed.

Source code in src/scale_forecasting/registry/ops.py
def format_close_plan(plan: ClosePlan) -> str:
    """Render a `ClosePlan` as the operator-facing preview — the exact text a dry run prints.

    Skipped runs are printed as prominently as closable ones. The reason a header is stuck is the
    thing the operator actually needs to see: "still has a RUNNING job" means go probe it, and
    burying that under a count of what *was* closed is how the one run that mattered gets missed.
    """
    lines = [f"close-runs against registry {plan.registry}"]
    if plan.unknown:
        lines.append(f"  not stuck (skipped): {', '.join(plan.unknown)}")
    if plan.closable:
        lines.append(f"  will close ({len(plan.closable)}):")
        for c in plan.closable:
            lines.append(f"    {c.run_id}  {c.header_status} -> {c.new_status}  [{c.reason}]")
    if plan.skipped:
        lines.append(f"  left alone ({len(plan.skipped)}):")
        for c in plan.skipped:
            lines.append(f"    {c.run_id}  {c.header_status}  [{c.reason}]")
    if not plan.closable and not plan.skipped:
        lines.append("  no stuck headers")
    return "\n".join(lines)

format_doctor

format_doctor(report)

Render a DoctorReport as readable text.

Source code in src/scale_forecasting/registry/ops.py
def format_doctor(report: DoctorReport) -> str:
    """Render a `DoctorReport` as readable text."""
    lines = [f"registry {report.registry}", f"  artifacts: {report.artifact_root}"]
    for stat in report.tables:
        rows = "MISSING" if stat.rows is None else f"{stat.rows:,} rows"
        lines.append(f"  {stat.table:<22} {rows}")
    lines.append(f"  views: {', '.join(report.views) or '(none)'}")
    if report.live_runs:
        pairs = ", ".join(f"{r} ({s})" for r, s in report.live_runs)
        lines.append(f"  IN FLIGHT ({len(report.live_runs)}): {pairs}")
    if report.orphans:
        objects = sum(p.object_count for p in report.orphans)
        total = human_bytes(sum(p.byte_total for p in report.orphans))
        lines.append(
            f"  ORPHANED artifacts: {len(report.orphans)} prefixes, {objects} objects, {total}"
            " — run sweep_orphans"
        )
    lines.extend(f"  note: {n}" for n in report.notes)
    if report.healthy:
        lines.append("  healthy")
    return "\n".join(lines)

format_plan

format_plan(plan)

Render a plan as the operator-facing preview — the exact text a dry run prints.

Source code in src/scale_forecasting/registry/ops.py
def format_plan(plan: DropPlan | SweepPlan) -> str:
    """Render a plan as the operator-facing preview — the exact text a dry run prints."""
    lines: list[str] = []
    if isinstance(plan, DropPlan):
        lines.append(f"drop-run against registry {plan.registry}")
        if plan.unknown:
            lines.append(f"  not in this registry (skipped): {', '.join(plan.unknown)}")
        if plan.blocked:
            lines.append(f"  IN FLIGHT — refusing: {', '.join(plan.blocked)}")
        if plan.run_ids:
            lines.append(f"  runs ({len(plan.run_ids)}): {', '.join(plan.run_ids)}")
            lines.append(f"  rows: DELETE from {len(REGISTRY_TABLE_NAMES)} tables")
        if plan.models:
            lines.append(f"  BQML model objects ({len(plan.models)}): {', '.join(plan.models)}")
    else:
        lines.append(f"sweep-orphans against registry {plan.registry}")
        lines.append(f"  artifact root: {plan.root}")
        lines.append(f"  runs known to the registry: {plan.known_runs}")
        lines.append(
            f"  orphan prefixes ({len(plan.prefixes)}): "
            f"{', '.join(p.run_id for p in plan.prefixes) or '(none)'}"
        )
    lines.append(f"  GCS: {plan.object_count} objects, {human_bytes(plan.byte_total)}")
    return "\n".join(lines)

human_bytes

human_bytes(n)

1536 → 1.5 KiB. Preview text only — never parsed.

Source code in src/scale_forecasting/registry/ops.py
def human_bytes(n: int) -> str:
    """``1536`` → ``1.5 KiB``. Preview text only — never parsed."""
    size = float(n)
    for unit in ("B", "KiB", "MiB", "GiB", "TiB"):
        if size < 1024 or unit == "TiB":
            return f"{size:.0f} {unit}" if unit == "B" else f"{size:.1f} {unit}"
        size /= 1024
    return f"{size:.1f} TiB"  # pragma: no cover - unreachable, the loop returns at TiB

init

init(*, settings=None, create_dataset=False)

Create this registry's five tables + five views in Settings.registry_dataset_ref.

Idempotent (CREATE TABLE IF NOT EXISTS + CREATE OR REPLACE VIEW) and registry-only: the source panel is a separate concern with a separate lifetime and is never created here. Point SF_REGISTRY_DATASET_ID at a new dataset and this is how you stand up a second registry.

Terraform owns dataset creation, so by default this expects the dataset to exist and fails with BigQuery's own 404 if it does not. create_dataset=True is the opt-in escape hatch for a registry outside the Terraform state — convenient, and it means the dataset's location and labels are not what Terraform would have given it. Returns the registry ref.

Source code in src/scale_forecasting/registry/ops.py
def init(
    *, settings: Settings | None = None, create_dataset: bool = False
) -> str:  # pragma: no cover - GCP I/O, @gcp smoke
    """Create this registry's five tables + five views in `Settings.registry_dataset_ref`.

    Idempotent (``CREATE TABLE IF NOT EXISTS`` + ``CREATE OR REPLACE VIEW``) and **registry-only**:
    the source panel is a separate concern with a separate lifetime and is never created here. Point
    ``SF_REGISTRY_DATASET_ID`` at a new dataset and this is how you stand up a second registry.

    Terraform owns dataset creation, so by default this expects the dataset to exist and fails with
    BigQuery's own 404 if it does not. ``create_dataset=True`` is the opt-in escape hatch for a
    registry outside the Terraform state — convenient, and it means the dataset's location and
    labels are not what Terraform would have given it. Returns the registry ref.
    """

    from ..errors import RegistryError
    from .ddl import render_create_tables, render_migrations
    from .tables import ensure_views

    resolved = _resolved(settings)
    client = bigquery_client(resolved.project_id)
    if create_dataset:
        _log.warning(
            "creating dataset %s outside Terraform — its location/labels will not match the "
            "Terraform-managed datasets in this deployment",
            resolved.registry_dataset_ref,
        )
        client.create_dataset(resolved.registry_dataset_ref, exists_ok=True)

    # The registry half of what `tables.ensure_tables` does — CREATE, then the additive ALTERs, so
    # this both stands up a new registry and brings an older one up to the current column set. The
    # five run-collection tables are always native BigQuery, hence `iceberg=False` and no
    # connection. Two passes, not one merged dict: both renderers key on the table name, so merging
    # would drop a CREATE for every table that also has a migration.
    creates = render_create_tables(
        resolved.registry_dataset_ref, iceberg=False, tables=REGISTRY_TABLE_NAMES
    )
    migrations = render_migrations(resolved.registry_dataset_ref, tables=REGISTRY_TABLE_NAMES)
    for stage, statements in (("creating", creates), ("migrating", migrations)):
        for name, statement in statements.items():
            try:
                client.query(statement).result()
            except Exception as exc:  # re-raised with table context
                raise RegistryError(f"init failed {stage} {name}: {exc}") from exc
    ensure_views(settings=resolved)
    _log.info("registry ready: %s", resolved.registry_dataset_ref)
    return resolved.registry_dataset_ref

main

main(argv=None)

CLI: python -m scale_forecasting.registry.ops <verb> [...].

One subcommand per verb, the same eight the SDK exposes (G1 — one implementation, three entry points). Mutating verbs preview by default and need --yes.

Source code in src/scale_forecasting/registry/ops.py
def main(argv: list[str] | None = None) -> None:
    """CLI: ``python -m scale_forecasting.registry.ops <verb> [...]``.

    One subcommand per verb, the same eight the SDK exposes (G1 — one implementation, three
    entry points). Mutating verbs preview by default and need ``--yes``.
    """
    import argparse

    configure_cli_logging()
    require_extra("gcp", purpose="Registry operations")
    p = argparse.ArgumentParser(
        prog="registry-ops",
        description="Manage one scale-forecasting registry (the dataset SF_REGISTRY_DATASET_ID, "
        "or SF_DATASET_ID, resolves to). No full-wipe verb — delete the dataset in BigQuery.",
    )
    sub = p.add_subparsers(dest="verb", required=True)

    p_init = sub.add_parser("init", help="create this registry's tables + views (idempotent)")
    p_init.add_argument(
        "--create-dataset",
        action="store_true",
        help="also create the dataset if absent (Terraform normally owns this)",
    )
    sub.add_parser("doctor", help="read-only health report")

    p_close = sub.add_parser(
        "close-runs", help="finalize abandoned RUNNING headers from their own job rows"
    )
    p_close.add_argument(
        "run_ids", nargs="*", help="run ids to close (default: every stuck header)"
    )
    p_close.add_argument("--yes", action="store_true", help="execute (default is a preview)")

    p_drop = sub.add_parser("drop-run", help="delete named run(s): artifacts, models, rows")
    p_drop.add_argument("run_ids", nargs="+", help="one or more run ids")
    p_drop.add_argument("--yes", action="store_true", help="execute (default is a preview)")
    p_drop.add_argument("--force", action="store_true", help="drop even if a run looks in-flight")

    p_sweep = sub.add_parser("sweep-orphans", help="delete artifacts with no run_registry row")
    p_sweep.add_argument("--yes", action="store_true", help="execute (default is a preview)")

    p_reap = sub.add_parser("reap-clusters", help="delete Ray clusters whose run has finished")
    p_reap.add_argument(
        "--region",
        action="append",
        dest="regions",
        default=None,
        help="a region to sweep; repeatable (default: the data-plane region alone)",
    )
    p_reap.add_argument(
        "--min-age-seconds",
        type=float,
        default=None,
        help="how old a cluster with no run header must be before it counts as garbage",
    )
    p_reap.add_argument("--yes", action="store_true", help="execute (default is a preview)")

    p_snap = sub.add_parser("snapshot", help="BigQuery table snapshots of the registry")
    p_snap.add_argument("suffix", help="names the snapshot set, e.g. 20260831")
    p_snap.add_argument("--into", default=None, help="target dataset (default: the registry's own)")
    p_snap.add_argument("--expiration-days", type=int, default=None, help="snapshot TTL in days")

    p_exp = sub.add_parser("export", help="dump the registry to GCS")
    p_exp.add_argument("destination", help="gs:// prefix to export under")
    p_exp.add_argument("--format", default="PARQUET", choices=EXPORT_FORMATS, dest="fmt")

    ns = p.parse_args(argv)
    if ns.verb == "init":
        init(create_dataset=ns.create_dataset)
    elif ns.verb == "doctor":
        _log.warning("%s", format_doctor(doctor()))
    elif ns.verb == "close-runs":
        close_runs(ns.run_ids or None, yes=ns.yes)
    elif ns.verb == "drop-run":
        drop_run(ns.run_ids, yes=ns.yes, force=ns.force)
    elif ns.verb == "sweep-orphans":
        sweep_orphans(yes=ns.yes)
    elif ns.verb == "reap-clusters":
        from ..ray_reaper import DEFAULT_MIN_AGE_SECONDS, reap_clusters

        reap_clusters(
            regions=ns.regions,
            min_age_seconds=(
                DEFAULT_MIN_AGE_SECONDS if ns.min_age_seconds is None else ns.min_age_seconds
            ),
            yes=ns.yes,
        )
    elif ns.verb == "snapshot":
        snapshot(ns.suffix, into=ns.into, expiration_days=ns.expiration_days)
    elif ns.verb == "export":
        export(ns.destination, fmt=ns.fmt)

match_reusable_models

match_reusable_models(
    cfg,
    run_id,
    snapshot_millis,
    prior_headers,
    available_models_by_run,
    *,
    source_last_modified_millis=None,
)

Match requested models against candidate prior runs by make_model_spec_hash (pure).

Returns {target_model_type: BorrowedModelSource} for every model in cfg.models whose exact model_spec_hash (and matching source_table + series_limit + verified snapshot_millis) is already present with completed cells in a candidate prior run.

Source code in src/scale_forecasting/registry/ops.py
def match_reusable_models(
    cfg: object,
    run_id: str,
    snapshot_millis: int | None,
    prior_headers: Sequence[dict[str, object]],
    available_models_by_run: dict[str, set[str]],
    *,
    source_last_modified_millis: int | None = None,
) -> dict[str, BorrowedModelSource]:
    """Match requested models against candidate prior runs by ``make_model_spec_hash`` (pure).

    Returns ``{target_model_type: BorrowedModelSource}`` for every model in ``cfg.models`` whose
    exact ``model_spec_hash`` (and matching ``source_table`` + ``series_limit`` + verified
    ``snapshot_millis``) is already present with completed cells in a candidate prior run.
    """
    from ..config import RunConfig
    from .ids import make_model_spec_hash

    assert isinstance(cfg, RunConfig)
    explicit_source = cfg.compute.reuse_from_run_id
    if not cfg.compute.reuse_prior_runs and explicit_source is None:
        return {}
    if explicit_source is None and snapshot_millis is None and source_last_modified_millis is None:
        return {}

    matched: dict[str, BorrowedModelSource] = {}
    want_hashes = {m: make_model_spec_hash(m, cfg) for m in cfg.models}

    for hdr in prior_headers:
        cand_run_id = str(hdr.get("run_id") or "")
        if not cand_run_id or cand_run_id == run_id:
            continue
        if explicit_source is not None and cand_run_id != explicit_source:
            continue
        if str(hdr.get("status") or "") != "COMPLETED":
            continue
        cand_snap_raw = hdr.get("snapshot_millis")
        cand_snap = int(str(cand_snap_raw)) if cand_snap_raw is not None else None
        if explicit_source is None:
            if snapshot_millis is not None:
                if cand_snap != snapshot_millis:
                    continue
            else:
                if (
                    cand_snap is None
                    or source_last_modified_millis is None
                    or cand_snap < source_last_modified_millis
                ):
                    continue

        raw_cfg = hdr.get("raw_config")
        if isinstance(raw_cfg, str):
            try:
                raw_cfg = json.loads(raw_cfg)
            except Exception:  # noqa: BLE001 - skip malformed prior config
                continue
        if not isinstance(raw_cfg, dict):
            continue
        try:
            prior_cfg = RunConfig.model_validate(raw_cfg)
        except Exception:  # noqa: BLE001 - skip prior config that no longer validates
            continue

        if (
            prior_cfg.data.source_table != cfg.data.source_table
            or prior_cfg.data.series_limit != cfg.data.series_limit
        ):
            continue

        avail = available_models_by_run.get(cand_run_id, set())
        if not avail:
            continue

        prior_by_hash: dict[str, str] = {}
        for pm in prior_cfg.models:
            if pm in avail:
                h = make_model_spec_hash(pm, prior_cfg)
                prior_by_hash.setdefault(h, pm)

        for target_m, want_h in want_hashes.items():
            if target_m in matched:
                continue
            source_m = prior_by_hash.get(want_h)
            if source_m is not None:
                matched[target_m] = BorrowedModelSource(
                    target_model_type=target_m,
                    source_run_id=cand_run_id,
                    source_model_type=source_m,
                    model_spec_hash=want_h,
                    snapshot_millis=cand_snap if cand_snap is not None else snapshot_millis,
                )
        if len(matched) == len(cfg.models):
            break

    return matched

plan_close_runs

plan_close_runs(run_ids=None, *, settings=None)

Which stuck headers close_runs would finalize, and to what — without writing anything.

run_ids=None means "every run whose header is stuck", which is the usual call: the operator wants the registry consistent, not a particular run. Naming runs explicitly narrows it, and a named run whose header is not stuck is reported in unknown rather than silently rewritten — this verb repairs abandoned headers and must never touch a settled one.

Three bulk reads, whatever the number of runs: the headers, the job rows, and the configs. The configs are what let roll_up_against_plan notice a family that was planned and never submitted — the exact hole an abandoned run leaves, and invisible to the job rows alone. A run whose config cannot be read falls back to the row-only roll_up_job_statuses.

Source code in src/scale_forecasting/registry/ops.py
def plan_close_runs(
    run_ids: Sequence[str] | None = None, *, settings: Settings | None = None
) -> ClosePlan:  # pragma: no cover - GCP I/O, @gcp smoke
    """Which stuck headers `close_runs` would finalize, and to what — without writing anything.

    ``run_ids=None`` means "every run whose header is stuck", which is the usual call: the operator
    wants the registry consistent, not a particular run. Naming runs explicitly narrows it, and a
    named run whose header is *not* stuck is reported in ``unknown`` rather than silently rewritten
    — this verb repairs abandoned headers and must never touch a settled one.

    Three bulk reads, whatever the number of runs: the headers, the job rows, and the configs. The
    configs are what let `roll_up_against_plan` notice a family that was planned and never submitted
    — the exact hole an abandoned run leaves, and invisible to the job rows alone. A run whose
    config cannot be read falls back to the row-only `roll_up_job_statuses`.
    """
    resolved = _resolved(settings)
    targets = list(_stuck_run_ids(resolved) if run_ids is None else run_ids)
    if not targets:
        return ClosePlan(registry=resolved.registry_dataset_ref)

    header = _statuses(resolved, targets)
    stuck = [r for r in targets if (header.get(r) or "").upper() in LIVE_STATUSES]
    unknown = tuple(sorted(set(targets) - set(stuck)))
    if not stuck:
        return ClosePlan(registry=resolved.registry_dataset_ref, unknown=unknown)

    jobs = _job_rows(resolved, stuck)
    plans = _run_plans(resolved, stuck)
    candidates = []
    for run_id in stuck:
        job_rows = jobs.get(run_id, ())
        plan = plans.get(run_id)
        if plan is None:
            new_status, reason = roll_up_job_statuses([s for _f, s in job_rows])
        else:
            families, ensemble_enabled = plan
            new_status, reason = roll_up_against_plan(
                job_rows, planned_families=families, ensemble_enabled=ensemble_enabled
            )
        candidates.append(
            CloseCandidate(
                run_id=run_id,
                header_status=header.get(run_id),
                job_statuses=tuple(s for _f, s in job_rows),
                new_status=new_status,
                reason=reason,
            )
        )
    return ClosePlan(
        registry=resolved.registry_dataset_ref,
        candidates=tuple(candidates),
        unknown=unknown,
    )

plan_drop_run

plan_drop_run(run_ids, *, settings=None)

Everything drop_run would delete, without deleting any of it.

Step 1 of the ordering rule: enumerate while the index still exists. A run absent from run_registry lands in unknown and is dropped from the plan — including its artifacts, because an unattributable prefix is sweep_orphans' job, not this verb's.

Source code in src/scale_forecasting/registry/ops.py
def plan_drop_run(
    run_ids: Sequence[str], *, settings: Settings | None = None
) -> DropPlan:  # pragma: no cover - GCP I/O, @gcp smoke
    """Everything `drop_run` would delete, without deleting any of it.

    Step 1 of the ordering rule: enumerate while the index still exists. A run absent from
    ``run_registry`` lands in ``unknown`` and is dropped from the plan — including its artifacts,
    because an unattributable prefix is `sweep_orphans`' job, not this verb's.
    """
    resolved = _resolved(settings)
    statuses = _statuses(resolved, list(run_ids))
    unknown = tuple(sorted(r for r, s in statuses.items() if s is None))
    live = blocking_runs(statuses)
    targets = tuple(sorted(r for r in run_ids if statuses.get(r) is not None))

    prefixes = artifacts.list_prefixes(resolved, targets) if targets else {}
    return DropPlan(
        registry=resolved.registry_dataset_ref,
        run_ids=targets,
        prefixes=tuple(prefixes[r] for r in targets if r in prefixes),
        models=_run_models(resolved, targets) if targets else (),
        blocked=live,
        unknown=unknown,
    )

plan_sweep_orphans

plan_sweep_orphans(*, settings=None)

Artifact prefixes under this registry's root with no run_registry row.

Source code in src/scale_forecasting/registry/ops.py
def plan_sweep_orphans(
    *, settings: Settings | None = None
) -> SweepPlan:  # pragma: no cover - GCP I/O, @gcp smoke
    """Artifact prefixes under this registry's root with no ``run_registry`` row."""
    resolved = _resolved(settings)
    known = _known_run_ids(resolved)
    prefixes = artifacts.list_prefixes(resolved)
    return SweepPlan(
        registry=resolved.registry_dataset_ref,
        root=resolved.artifact_root,
        prefixes=tuple(prefixes[r] for r in artifacts.orphan_run_ids(prefixes, known)),
        known_runs=len(known),
    )

record_borrowed_family_job

record_borrowed_family_job(
    cfg,
    run_id,
    family,
    runtime,
    spark_mode,
    hardware,
    gpu_type,
    borrowed_sources,
    *,
    force=False,
    settings=None,
)

Write a COMPLETED run_jobs row for a family satisfied 100% by borrowed prior cells.

Source code in src/scale_forecasting/registry/ops.py
def record_borrowed_family_job(
    cfg: object,
    run_id: str,
    family: str,
    runtime: str,
    spark_mode: str | None,
    hardware: str | None,
    gpu_type: str | None,
    borrowed_sources: Sequence[BorrowedModelSource],
    *,
    force: bool = False,
    settings: Settings | None = None,
) -> None:  # pragma: no cover - GCP I/O, @gcp smoke
    """Write a COMPLETED ``run_jobs`` row for a family satisfied 100% by borrowed prior cells."""
    from datetime import UTC, datetime

    from .ids import make_job_key
    from .jobs import next_job_attempt, update_job, write_job
    from .rows import assemble_job_row

    attempt, _ = next_job_attempt(run_id, family, force=force, settings=settings)
    now = datetime.now(UTC)
    source_ids = ",".join(sorted({s.source_run_id for s in borrowed_sources}))
    telemetry = {
        "borrowed": {
            s.target_model_type: {
                "source_run_id": s.source_run_id,
                "source_model_type": s.source_model_type,
                "model_spec_hash": s.model_spec_hash,
                "snapshot_millis": s.snapshot_millis,
            }
            for s in borrowed_sources
        }
    }
    job_id = make_job_key(run_id, family, attempt)
    row = assemble_job_row(
        run_id,
        family,
        attempt,
        now,
        runtime=runtime,
        spark_mode=spark_mode,
        hardware=hardware,
        gpu_type=gpu_type,
        system_job_id=f"borrowed:{source_ids}",
        status="COMPLETED",
        started_at=now,
    )
    write_job(row, settings=settings)
    update_job(
        job_id,
        settings=settings,
        status="COMPLETED",
        runtime_seconds=0.0,
        ended_at=now,
        merge_telemetry=telemetry,
    )

render_borrow_model_sql

render_borrow_model_sql(dataset)

Render parameterized INSERT INTO ... SELECT ... statements for borrowing one model.

Source code in src/scale_forecasting/registry/ops.py
def render_borrow_model_sql(dataset: str) -> dict[str, str]:
    """Render parameterized ``INSERT INTO ... SELECT ...`` statements for borrowing one model."""
    from ..metrics import METRIC_NAMES

    metric_cols = ", ".join(METRIC_NAMES)
    return {
        "forecast_predictions": (
            f"INSERT INTO `{dataset}.forecast_predictions` ("
            "run_id, ts_id, model_type, compute_engine, ensemble_id, forecast_date, "
            "yhat, yhat_raw, yhat_adjusted, yhat_lower, yhat_upper, quantiles, created_at, "
            "explanations) "
            "SELECT @target_run_id AS run_id, ts_id, @target_model_type AS model_type, "
            "compute_engine, ensemble_id, forecast_date, yhat, yhat_raw, yhat_adjusted, "
            "yhat_lower, yhat_upper, quantiles, CURRENT_TIMESTAMP() AS created_at, explanations "
            f"FROM (SELECT * FROM `{dataset}.forecast_predictions` "
            "WHERE run_id = @source_run_id AND model_type = @source_model_type "
            "AND ensemble_id IS NULL "
            "QUALIFY ROW_NUMBER() OVER ("
            "PARTITION BY ts_id, forecast_date ORDER BY created_at DESC NULLS LAST) = 1)"
        ),
        "backtest_oof": (
            f"INSERT INTO `{dataset}.backtest_oof` ("
            "run_id, ts_id, model_type, fold_id, forecast_date, y_true, yhat, yhat_raw, "
            "yhat_adjusted, cutoff_date, horizon_step, yhat_lower, yhat_upper, created_at, "
            "ensemble_id, yhat_stale) "
            "SELECT @target_run_id AS run_id, ts_id, @target_model_type AS model_type, "
            "fold_id, forecast_date, y_true, yhat, yhat_raw, yhat_adjusted, cutoff_date, "
            "horizon_step, yhat_lower, yhat_upper, CURRENT_TIMESTAMP() AS created_at, "
            "ensemble_id, yhat_stale "
            f"FROM (SELECT * FROM `{dataset}.backtest_oof` "
            "WHERE run_id = @source_run_id AND model_type = @source_model_type "
            "AND ensemble_id IS NULL "
            "QUALIFY ROW_NUMBER() OVER ("
            "PARTITION BY ts_id, fold_id, forecast_date ORDER BY created_at DESC NULLS LAST) = 1)"
        ),
        "forecast_metadata": (
            f"INSERT INTO `{dataset}.forecast_metadata` ("
            f"run_id, ts_id, model_type, compute_engine, model_hash, ensemble_id, fold_id, "
            f"{metric_cols}, fit_seconds, best_params, model_artifact, created_at, worker_id, "
            "cell_started_at, cell_ended_at, cpu_seconds, process_rss_bytes, peak_gpu_bytes, "
            "intraop_threads, n_obs, cell_status, error_class, error_detail, backtest_status, "
            "backtest_note, n_folds_achieved, achieved_step, achieved_min_train, first_val_date, "
            "last_val_date, interval_source, point_forecast_source, point_forecast_decision, "
            "interval_calibration, point_forecast_margin, ensemble_scoring, hpo_scoring, n_fits, "
            "train_rows_total, n_hpo_fits, fit_diagnostics, device_requested, device_available, "
            "device_used, device_name, backtest_refit, staleness_gap) "
            "SELECT @target_run_id AS run_id, ts_id, @target_model_type AS model_type, "
            f"compute_engine, @model_spec_hash AS model_hash, ensemble_id, fold_id, {metric_cols}, "
            "fit_seconds, best_params, model_artifact, CURRENT_TIMESTAMP() AS created_at, "
            "worker_id, cell_started_at, cell_ended_at, cpu_seconds, process_rss_bytes, "
            "peak_gpu_bytes, intraop_threads, n_obs, cell_status, error_class, error_detail, "
            "backtest_status, backtest_note, n_folds_achieved, achieved_step, achieved_min_train, "
            "first_val_date, last_val_date, interval_source, point_forecast_source, "
            "point_forecast_decision, interval_calibration, point_forecast_margin, "
            "ensemble_scoring, hpo_scoring, n_fits, train_rows_total, n_hpo_fits, "
            "JSON_SET(COALESCE(fit_diagnostics, JSON '{}'), "
            "'$.borrowed_from_run_id', @source_run_id, "
            "'$.borrowed_from_model_type', @source_model_type, "
            "'$.base_model', @base_model, "
            "'$.scope_capability', @scope_capability, "
            "'$.training_mode', @training_mode, "
            "'$.cov_profile', @cov_profile, "
            "'$.is_univariate', @is_univariate, "
            "'$.covariates_used', PARSE_JSON(@covariates_used_json), "
            "'$.model_spec_hash', @model_spec_hash) AS fit_diagnostics, "
            "device_requested, device_available, device_used, device_name, backtest_refit, "
            "staleness_gap "
            f"FROM (SELECT * FROM `{dataset}.forecast_metadata` "
            "WHERE run_id = @source_run_id AND model_type = @source_model_type "
            "AND ensemble_id IS NULL "
            "QUALIFY ROW_NUMBER() OVER ("
            "PARTITION BY ts_id, COALESCE(fold_id, -1) ORDER BY created_at DESC NULLS LAST) = 1)"
        ),
    }

render_delete_rows

render_delete_rows(dataset, *, tables=REGISTRY_TABLE_NAMES)

{table: DELETE ... WHERE run_id IN UNNEST(@run_ids)} for a per-run delete.

Every registry table carries run_id, so one parameterized statement per table covers the whole tier. The ids are bound as an array query parameter, never interpolated.

A caveat worth surfacing rather than hiding: BigQuery rejects a DELETE that matches rows still in the Storage Write API streaming buffer (~90 minutes after the append). Dropping a run that just finished can therefore fail; drop_run re-raises that with a message saying to wait.

Source code in src/scale_forecasting/registry/ops.py
def render_delete_rows(
    dataset: str, *, tables: Sequence[str] = REGISTRY_TABLE_NAMES
) -> dict[str, str]:
    """``{table: DELETE ... WHERE run_id IN UNNEST(@run_ids)}`` for a per-run delete.

    Every registry table carries ``run_id``, so one parameterized statement per table covers the
    whole tier. The ids are bound as an **array query parameter**, never interpolated.

    A caveat worth surfacing rather than hiding: BigQuery rejects a ``DELETE`` that matches rows
    still in the Storage Write API streaming buffer (~90 minutes after the append). Dropping a run
    that just finished can therefore fail; `drop_run` re-raises that with a message saying to wait.
    """
    return {
        name: f"DELETE FROM `{dataset}.{name}` WHERE run_id IN UNNEST(@run_ids);" for name in tables
    }

render_export_sql

render_export_sql(
    dataset,
    destination_root,
    *,
    fmt="PARQUET",
    tables=REGISTRY_TABLE_NAMES,
)

{table: EXPORT DATA ...} — dump the registry to GCS for offline analysis.

A different job from render_snapshot_sql, and both earn their place: a snapshot stays inside BigQuery (fast, queryable, cheap) while an export leaves it (readable by anyone with the bucket and no BigQuery access at all). Each table lands under <destination_root>/<table>/ with a sharded wildcard, which is how BigQuery writes an export.

fmt is one of EXPORT_FORMATS. Prefer JSON when the registry's JSON columns matter downstream — a Parquet export flattens them to strings.

Source code in src/scale_forecasting/registry/ops.py
def render_export_sql(
    dataset: str,
    destination_root: str,
    *,
    fmt: str = "PARQUET",
    tables: Sequence[str] = REGISTRY_TABLE_NAMES,
) -> dict[str, str]:
    """``{table: EXPORT DATA ...}`` — dump the registry to GCS for offline analysis.

    A different job from `render_snapshot_sql`, and both earn their place: a snapshot stays inside
    BigQuery (fast, queryable, cheap) while an export leaves it (readable by anyone with the bucket
    and no BigQuery access at all). Each table lands under
    ``<destination_root>/<table>/`` with a sharded wildcard, which is how BigQuery writes an export.

    ``fmt`` is one of `EXPORT_FORMATS`. Prefer ``JSON`` when the registry's ``JSON`` columns matter
    downstream — a Parquet export flattens them to strings.
    """
    fmt_up = fmt.upper()
    if fmt_up not in EXPORT_FORMATS:
        raise ValueError(f"unsupported export format {fmt!r}; expected one of {EXPORT_FORMATS}")
    bq_format = "NEWLINE_DELIMITED_JSON" if fmt_up == "JSON" else "PARQUET"
    ext = "json" if fmt_up == "JSON" else "parquet"
    root = destination_root.rstrip("/")
    return {
        name: (
            "EXPORT DATA OPTIONS (\n"
            f"  uri = '{root}/{name}/{name}-*.{ext}',\n"
            f"  format = '{bq_format}',\n"
            "  overwrite = true\n"
            ") AS\n"
            f"SELECT * FROM `{dataset}.{name}`;"
        )
        for name in tables
    }

render_snapshot_sql

render_snapshot_sql(
    source_dataset,
    target_dataset,
    suffix,
    *,
    tables=REGISTRY_TABLE_NAMES,
    expiration_days=None,
)

{table: CREATE SNAPSHOT TABLE ...} — a point-in-time copy of the registry.

A BigQuery table snapshot, not a CREATE TABLE AS SELECT: a snapshot is metadata plus the delta from its base, so it costs almost nothing at creation and only accrues storage as the base diverges. It is read-only and can carry an expiration, which is what makes "snapshot before the risky thing" a habit rather than a budget decision.

suffix is appended to each table name (run_registry → run_registry_20260831), so a snapshot can live in the registry's own dataset without colliding; pass a different target_dataset to keep the registry uncluttered.

Source code in src/scale_forecasting/registry/ops.py
def render_snapshot_sql(
    source_dataset: str,
    target_dataset: str,
    suffix: str,
    *,
    tables: Sequence[str] = REGISTRY_TABLE_NAMES,
    expiration_days: int | None = None,
) -> dict[str, str]:
    """``{table: CREATE SNAPSHOT TABLE ...}`` — a point-in-time copy of the registry.

    A BigQuery **table snapshot**, not a ``CREATE TABLE AS SELECT``: a snapshot is metadata plus the
    delta from its base, so it costs almost nothing at creation and only accrues storage as the base
    diverges. It is read-only and can carry an expiration, which is what makes "snapshot before the
    risky thing" a habit rather than a budget decision.

    ``suffix`` is appended to each table name (``run_registry`` → ``run_registry_20260831``), so a
    snapshot can live in the registry's own dataset without colliding; pass a different
    ``target_dataset`` to keep the registry uncluttered.
    """
    options = ""
    if expiration_days is not None:
        if expiration_days <= 0:
            raise ValueError(f"expiration_days must be positive, got {expiration_days}")
        options = (
            "\nOPTIONS (expiration_timestamp = "
            f"TIMESTAMP_ADD(CURRENT_TIMESTAMP(), INTERVAL {expiration_days} DAY))"
        )
    return {
        name: (
            f"CREATE SNAPSHOT TABLE IF NOT EXISTS `{target_dataset}.{name}_{suffix}`\n"
            f"CLONE `{source_dataset}.{name}`{options};"
        )
        for name in tables
    }

roll_up_against_plan

roll_up_against_plan(
    job_rows, *, planned_families, ensemble_enabled
)

Roll a run's (family, status) job rows up against the families its config planned (pure).

roll_up_job_statuses reads the rows and nothing else, so a family that was planned and never submitted is invisible to it — and that is exactly the shape the dead-driver case takes. A run whose kernel died after its last family finished, before the ensemble node could start, reads as "every job COMPLETED" and closes to COMPLETED, claiming an ensemble that never ran. Two live runs were mis-reported that way on 2026-09-20.

So the plan comes in too: every planned family with no row enters the picture as None, and the whole map goes to job_outcome.combined_run_status — the same function main.run and the Airflow finalize_run task use to write the header in the first place. One policy, in one place, so a repaired header says what the run would have said for itself. A missing ensemble on an otherwise-complete run is therefore FAILED (the requested output does not exist), and a missing base family alongside completed ones is PARTIAL (what did land is still usable).

Once a plan is in hand the delegation is unconditional, not just for the missing-family case, because the two readings had drifted in a second place too: a present FAILED ensemble row alongside completed base families rolls up to PARTIAL on the rows alone, where the run itself would have written FAILED for the same reason a missing one does.

Three answers are decided before the plan is consulted, and all three are the row-only ones:

  • no rows at all ⇒ FAILED, unchanged — there is no partial anything to describe.
  • any present non-terminal row ⇒ None, unchanged, for the same reason as ever: only a runtime probe can tell a live job from a stale row. A missing row must not refuse, or the abandoned runs this verb exists to repair could never be closed at all.
  • every present row CANCELLED ⇒ CANCELLED. Cancelling is a normal way for a planned family to end up with no row, and combined_run_status would call it FAILED — it never sees a cancelled run, because cancelling skips the finalizer that calls it.

A family with more than one row (sf-<run_id>-<family>-a<attempt>: retries share the family, not the job_id) folds to COMPLETED only if every one of its rows is, which never upgrades a family and matches how the row-only roll-up already reads a failed-then-retried pair.

Source code in src/scale_forecasting/registry/ops.py
def roll_up_against_plan(
    job_rows: Sequence[tuple[str, str | None]],
    *,
    planned_families: Sequence[str],
    ensemble_enabled: bool,
) -> tuple[str | None, str]:
    """Roll a run's ``(family, status)`` job rows up against the families its config planned (pure).

    `roll_up_job_statuses` reads the rows and nothing else, so a family that was planned and never
    submitted is invisible to it — and that is exactly the shape the dead-driver case takes. A run
    whose kernel died after its last family finished, before the ensemble node could start, reads as
    "every job COMPLETED" and closes to ``COMPLETED``, claiming an ensemble that never ran. Two live
    runs were mis-reported that way on 2026-09-20.

    So the plan comes in too: every planned family with no row enters the picture as ``None``, and
    the whole map goes to `job_outcome.combined_run_status` — the same function `main.run` and the
    Airflow ``finalize_run`` task use to write the header in the first place. One policy, in one
    place, so a repaired header says what the run would have said for itself. A missing ensemble on
    an otherwise-complete run is therefore ``FAILED`` (the requested output does not exist), and a
    missing *base* family alongside completed ones is ``PARTIAL`` (what did land is still usable).

    Once a plan is in hand the delegation is unconditional, not just for the missing-family case,
    because the two readings had drifted in a second place too: a *present* ``FAILED`` ensemble row
    alongside completed base families rolls up to ``PARTIAL`` on the rows alone, where the run
    itself would have written ``FAILED`` for the same reason a missing one does.

    Three answers are decided before the plan is consulted, and all three are the row-only ones:

    * **no rows at all ⇒ ``FAILED``**, unchanged — there is no partial anything to describe.
    * **any *present* non-terminal row ⇒ ``None``**, unchanged, for the same reason as ever: only a
      runtime probe can tell a live job from a stale row. A *missing* row must not refuse, or the
      abandoned runs this verb exists to repair could never be closed at all.
    * **every present row ``CANCELLED`` ⇒ ``CANCELLED``.** Cancelling is a normal way for a planned
      family to end up with no row, and `combined_run_status` would call it ``FAILED`` — it never
      sees a cancelled run, because cancelling skips the finalizer that calls it.

    A family with more than one row (``sf-<run_id>-<family>-a<attempt>``: retries share the family,
    not the ``job_id``) folds to ``COMPLETED`` only if *every* one of its rows is, which never
    upgrades a family and matches how the row-only roll-up already reads a failed-then-retried pair.
    """
    from ..job_outcome import combined_run_status

    statuses = [status for _family, status in job_rows]
    rolled, reason = roll_up_job_statuses(statuses)
    if not statuses or rolled is None or rolled == "CANCELLED" or not planned_families:
        return rolled, reason

    status_by_family: dict[str, str | None] = {}
    for family, status in job_rows:
        if status_by_family.get(family, "COMPLETED") == "COMPLETED":
            status_by_family[family] = (status or "").upper()  # a second row never upgrades

    missing = tuple(f for f in planned_families if f not in status_by_family)
    merged = {**status_by_family, **dict.fromkeys(missing, None)}
    status = combined_run_status(merged, ensemble_enabled=ensemble_enabled)
    if not missing:
        return status, reason
    plural = "y" if len(missing) == 1 else "ies"
    note = f"{len(missing)} planned famil{plural} never recorded a row: {', '.join(missing)}"
    return status, f"{reason}; {note}"

roll_up_job_statuses

roll_up_job_statuses(job_statuses)

Roll a run's job-row statuses into the header status it should have had (pure).

Returns (status, reason); status is None when the run must not be closed and reason says why. This verb's whole job is to write the status the run itself failed to, never a different one, so the policy tracks job_outcome.combined_run_status — what a run that finished normally would have written:

  • every job COMPLETED ⇒ COMPLETED
  • every job FAILED ⇒ FAILED; every job CANCELLED ⇒ CANCELLED
  • a mix of terminals ⇒ PARTIAL
  • any non-terminal job ⇒ None. The job may still be running. Nothing here probes a runtime, so a header is only closed when the rows themselves already prove the work settled; deciding a job is dead is probes.reconcile's call, made against the runtime, not this one's. AWAITING_CAPACITY lands here by construction (it is simply not in the terminal set), and that is the right answer for the right reason: a job between capacity attempts has neither finished nor failed, so closing its run would write a verdict on work still queued to happen.
  • no job rows at all ⇒ FAILED. The run wrote a header and then never recorded a single family — it died in the submit path. FAILED rather than COMPLETED because nothing completed, and rather than CANCELLED because nobody stopped it.

It can only see rows that exist, which is the one thing a status list cannot tell you: a family that was planned and never submitted leaves nothing behind to be counted. That gap is roll_up_against_plan's to close, and it is the entry point plan_close_runs uses whenever the run's config is readable. Everything here stays the row-only answer, for when it isn't.

Source code in src/scale_forecasting/registry/ops.py
def roll_up_job_statuses(job_statuses: Sequence[str | None]) -> tuple[str | None, str]:
    """Roll a run's job-row statuses into the header status it should have had (pure).

    Returns ``(status, reason)``; ``status`` is ``None`` when the run must **not** be closed and
    ``reason`` says why. This verb's whole job is to write the status the run itself failed to,
    never a different one, so the policy tracks `job_outcome.combined_run_status` — what a run that
    finished normally would have written:

    * every job ``COMPLETED`` ⇒ ``COMPLETED``
    * every job ``FAILED`` ⇒ ``FAILED``; every job ``CANCELLED`` ⇒ ``CANCELLED``
    * a mix of terminals ⇒ ``PARTIAL``
    * **any non-terminal job ⇒ ``None``.** The job may still be running. Nothing here probes a
      runtime, so a header is only closed when the *rows themselves* already prove the work settled;
      deciding a job is dead is `probes.reconcile`'s call, made against the runtime, not this one's.
      ``AWAITING_CAPACITY`` lands here by construction (it is simply not in the terminal set), and
      that is the right answer for the right reason: a job between capacity attempts has neither
      finished nor failed, so closing its run would write a verdict on work still queued to happen.
    * **no job rows at all ⇒ ``FAILED``.** The run wrote a header and then never recorded a single
      family — it died in the submit path. ``FAILED`` rather than ``COMPLETED`` because nothing
      completed, and rather than ``CANCELLED`` because nobody stopped it.

    **It can only see rows that exist**, which is the one thing a status list cannot tell you: a
    family that was planned and never submitted leaves nothing behind to be counted. That gap is
    `roll_up_against_plan`'s to close, and it is the entry point `plan_close_runs` uses whenever the
    run's config is readable. Everything here stays the row-only answer, for when it isn't.
    """
    if not job_statuses:
        return "FAILED", "no job rows — the run never recorded a family"
    normalized = [(s or "").upper() for s in job_statuses]
    unsettled = sorted({s or "(none)" for s in normalized if s not in _TERMINAL_JOB_STATUSES})
    if unsettled:
        return None, f"{len(unsettled)} job status(es) not terminal: {', '.join(unsettled)}"
    distinct = set(normalized)
    if len(distinct) == 1:
        only = distinct.pop()
        return only, f"every job {only}"
    return "PARTIAL", f"mixed terminal statuses: {', '.join(sorted(distinct))}"

snapshot

snapshot(
    suffix,
    *,
    settings=None,
    into=None,
    expiration_days=None,
)

Take a BigQuery table snapshot of each registry table; return {table: snapshot ref}.

Non-destructive, so there is no dry run — the whole point is to take one before the risky thing. suffix names the snapshot set (a date, a ticket, before_migration); into redirects to another dataset so the registry itself stays uncluttered; expiration_days sets a TTL so a habit of snapshotting does not become a storage bill.

Source code in src/scale_forecasting/registry/ops.py
def snapshot(
    suffix: str,
    *,
    settings: Settings | None = None,
    into: str | None = None,
    expiration_days: int | None = None,
) -> dict[str, str]:  # pragma: no cover - GCP I/O, @gcp smoke
    """Take a BigQuery table snapshot of each registry table; return ``{table: snapshot ref}``.

    Non-destructive, so there is no dry run — the whole point is to take one *before* the risky
    thing. ``suffix`` names the snapshot set (a date, a ticket, ``before_migration``);
    ``into`` redirects to another dataset so the registry itself stays uncluttered;
    ``expiration_days`` sets a TTL so a habit of snapshotting does not become a storage bill.
    """

    from ..errors import RegistryError

    resolved = _resolved(settings)
    target = into or resolved.registry_dataset_ref
    statements = render_snapshot_sql(
        resolved.registry_dataset_ref, target, suffix, expiration_days=expiration_days
    )
    client = bigquery_client(resolved.project_id)
    out: dict[str, str] = {}
    for name, statement in statements.items():
        try:
            client.query(statement).result()
        except Exception as exc:  # re-raised with table context
            raise RegistryError(f"snapshot failed for {name}: {exc}") from exc
        out[name] = f"{target}.{name}_{suffix}"
    _log.info("snapshotted %d tables into %s", len(out), target)
    return out

sweep_orphans

sweep_orphans(*, settings=None, yes=False)

Delete artifact prefixes under this registry with no run_registry row.

The cleanup for everything stranded before the artifact path carried a registry key. Scoping is what makes it safe: the sweep only ever looks under Settings.artifact_root, so another registry's objects in the same bucket are outside its world, and an object that cannot be attributed to a run at all is skipped rather than guessed at.

Preview by default; yes=True executes. Returns the plan.

Source code in src/scale_forecasting/registry/ops.py
def sweep_orphans(
    *, settings: Settings | None = None, yes: bool = False
) -> SweepPlan:  # pragma: no cover - GCP I/O, @gcp smoke
    """Delete artifact prefixes under *this* registry with no ``run_registry`` row.

    The cleanup for everything stranded before the artifact path carried a registry key. Scoping is
    what makes it safe: the sweep only ever looks under `Settings.artifact_root`, so another
    registry's objects in the same bucket are outside its world, and an object that cannot be
    attributed to a run at all is skipped rather than guessed at.

    Preview by default; ``yes=True`` executes. Returns the plan.
    """
    resolved = _resolved(settings)
    plan = plan_sweep_orphans(settings=resolved)
    _log.warning("%s", format_plan(plan))
    if not plan.prefixes:
        return plan
    if not yes:
        _log.warning("DRY RUN — nothing deleted. Re-run with yes=True to execute.")
        return plan
    deleted = artifacts.delete_prefixes(resolved, [p.run_id for p in plan.prefixes])
    _log.warning("swept %d orphan prefixes (%d objects)", len(plan.prefixes), deleted)
    return plan