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
¶
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
¶
Exactly which stuck headers close_runs would finalize, and to what.
DoctorReport
dataclass
¶
A read-only picture of one registry: what is in it, and what looks wrong.
DropPlan
dataclass
¶
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.
SweepPlan
dataclass
¶
Artifact prefixes under this registry's root with no run_registry row.
TableStat
dataclass
¶
One registry table's row count, or None rows if the table is absent.
all_header_statuses
¶
{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
blocking_runs
¶
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
borrow_model_results
¶
Copy forecast_predictions, backtest_oof, and forecast_metadata for reuse.
Source code in src/scale_forecasting/registry/ops.py
close_runs
¶
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
doctor
¶
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
drop_run
¶
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
export
¶
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
find_reusable_models
¶
Discover completed prior runs in run_registry whose cells can satisfy cfg.models.
Source code in src/scale_forecasting/registry/ops.py
1356 1357 1358 1359 1360 1361 1362 1363 1364 1365 1366 1367 1368 1369 1370 1371 1372 1373 1374 1375 1376 1377 1378 1379 1380 1381 1382 1383 1384 1385 1386 1387 1388 1389 1390 1391 1392 1393 1394 1395 1396 1397 1398 1399 1400 1401 1402 1403 1404 1405 1406 1407 1408 1409 1410 1411 1412 1413 1414 1415 1416 1417 1418 1419 1420 1421 1422 1423 1424 1425 1426 1427 1428 1429 1430 1431 1432 1433 1434 1435 1436 1437 1438 1439 1440 1441 1442 1443 1444 1445 1446 1447 1448 1449 1450 1451 1452 1453 1454 1455 | |
format_close_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
format_doctor
¶
Render a DoctorReport as readable text.
Source code in src/scale_forecasting/registry/ops.py
format_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
human_bytes
¶
1536 → 1.5 KiB. Preview text only — never parsed.
Source code in src/scale_forecasting/registry/ops.py
init
¶
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
main
¶
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
1572 1573 1574 1575 1576 1577 1578 1579 1580 1581 1582 1583 1584 1585 1586 1587 1588 1589 1590 1591 1592 1593 1594 1595 1596 1597 1598 1599 1600 1601 1602 1603 1604 1605 1606 1607 1608 1609 1610 1611 1612 1613 1614 1615 1616 1617 1618 1619 1620 1621 1622 1623 1624 1625 1626 1627 1628 1629 1630 1631 1632 1633 1634 1635 1636 1637 1638 1639 1640 1641 1642 1643 1644 1645 1646 1647 1648 1649 1650 1651 1652 1653 1654 1655 1656 1657 1658 1659 1660 1661 1662 | |
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
1183 1184 1185 1186 1187 1188 1189 1190 1191 1192 1193 1194 1195 1196 1197 1198 1199 1200 1201 1202 1203 1204 1205 1206 1207 1208 1209 1210 1211 1212 1213 1214 1215 1216 1217 1218 1219 1220 1221 1222 1223 1224 1225 1226 1227 1228 1229 1230 1231 1232 1233 1234 1235 1236 1237 1238 1239 1240 1241 1242 1243 1244 1245 1246 1247 1248 1249 1250 1251 1252 1253 1254 1255 1256 1257 1258 1259 1260 1261 1262 1263 1264 1265 1266 1267 1268 1269 1270 1271 1272 1273 1274 1275 1276 1277 | |
plan_close_runs
¶
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
plan_drop_run
¶
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
plan_sweep_orphans
¶
Artifact prefixes under this registry's root with no run_registry row.
Source code in src/scale_forecasting/registry/ops.py
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
render_borrow_model_sql
¶
Render parameterized INSERT INTO ... SELECT ... statements for borrowing one model.
Source code in src/scale_forecasting/registry/ops.py
1280 1281 1282 1283 1284 1285 1286 1287 1288 1289 1290 1291 1292 1293 1294 1295 1296 1297 1298 1299 1300 1301 1302 1303 1304 1305 1306 1307 1308 1309 1310 1311 1312 1313 1314 1315 1316 1317 1318 1319 1320 1321 1322 1323 1324 1325 1326 1327 1328 1329 1330 1331 1332 1333 1334 1335 1336 1337 1338 1339 1340 1341 1342 1343 1344 1345 1346 1347 1348 1349 1350 1351 1352 1353 | |
render_delete_rows
¶
{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
render_export_sql
¶
{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
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
roll_up_against_plan
¶
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, andcombined_run_statuswould call itFAILED— 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
roll_up_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 jobCANCELLED⇒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 isprobes.reconcile's call, made against the runtime, not this one's.AWAITING_CAPACITYlands 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.FAILEDrather thanCOMPLETEDbecause nothing completed, and rather thanCANCELLEDbecause 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
snapshot
¶
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
sweep_orphans
¶
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.