def run(
cfg: RunConfig,
*,
dry_run: bool = False,
spark: object | None = None,
settings: Settings | None = None,
force: bool = False,
n_series: int | None = None,
max_executors: int | None = None,
ignore_unavailable_models: bool = False,
) -> str:
"""Execute one run as a DAG: every model family in parallel under one run_id; return that id.
Plans the run's DAG (`dag.plan_dag`) — one job per model family present in the config
(statistical / ml / deep_learning each on its resolved runtime, native in BigQuery) plus the
downstream ensemble node — then:
* ``dry_run=True`` → delegate to `launch_plan.plan_run` (the offline "what would this
schedule" path): resolve the run_id + `estimate_fanout` + emit the launch-command templates,
touching no GCP, and return the run_id.
* otherwise → resolve `Settings`, ``ensure_tables`` + ``write_header`` (RUNNING, the one shared
header), launch each Python family's job on its own worker thread
(`job_launch.launch_family_job`) and run the BigQuery-native family inline
(`job_launch.launch_native_job`) so they all overlap — each in contributor mode with its own
``run_jobs`` row. Once every family job joins and *all* of them succeeded, and
``cfg.ensemble.enabled``, run the ensemble DAG node (`job_launch.launch_ensemble_job`, which
reads the just-written base predictions/OOF and scores each consensus onto the leaderboard
under its own ``run_jobs`` row), then finalize the header with the combined status +
wall-clock
``runtime_seconds`` + ``bq_models`` (+ the native engine's observed ``n_series`` when it ran).
On any family *or* ensemble failure the header is finalized FAILED/PARTIAL before the first
error re-raises, so the run stays queryable and the CLI exits non-zero.
``spark`` is an optional injected `SparkSession` (incl. a Spark Connect
``DataprocSparkSession``). When supplied, a Spark-runtime family runs **in-process against that
session** instead of a remote Dataproc batch — the notebook / Connect demo path — using the
identical engine code (the injectable-session seam). The default (``None``) keeps the
remote-batch behavior every CLI/Composer caller relies on.
``settings`` optionally injects a resolved `Settings` (the GCP infra identity); the default
(``None``) resolves it from the ``SF_*`` environment. The SDK ``Forecaster`` uses this to thread
an explicit identity instead of relying on process env. ``force`` bumps each family's
``run_jobs`` attempt (a fresh, distinctly-keyed job under the same run_id) and shapes the
``dry_run`` plan's re-run guidance.
``n_series`` overrides ``data.series_limit`` (the 10→100→1k→100k scale knob) before anything
else, so it changes the ``run_id`` and every family sees the same limit. ``max_executors`` caps
a remote Spark batch's dynamic-allocation ceiling (ignored by the in-process/Ray paths).
Idempotent: the config-pinned run_id + append-only/dedupe-on-read cell writes mean re-running
the same config lands byte-identical rows. ``compute.profile.source: "auto"`` is locked to what
it resolves to (`launch_plan.lock_profile_source`), exactly as the plan/stage verbs do it, but
that resolved pointer is excluded from the id (`registry.ids`) — otherwise a re-run sized from
newer evidence would land under a new id and the dedupe above would never fire. A source an
operator pinned by hand is checked for drift first (`profiling.source.check_pinned_source`) and
fails the run.
"""
import threading
from concurrent.futures import ThreadPoolExecutor
from .backtest import hpo_scoring_claim
from .dag import plan_dag, preflight
from .errors import EngineError, require_extra
from .job_outcome import combined_run_status
from .profiling.source import check_pinned_source
from .registry.header import header_status, merge_header_telemetry
from .registry.lifecycle import run_header
from .settings import Settings
# The series-limit override is applied first so it flows into the run_id and every family — a
# different scale is a distinct, independently-queryable run.
cfg = cfg.with_series_limit(n_series)
if ignore_unavailable_models:
cfg = cfg.with_available_models()
# Ahead of the dry-run branch on purpose: a config a model cannot honour, or a plan that would
# buy a device nothing routes to, is exactly what a dry run exists to catch. The checks need the
# registry loaded, which `plan_dag` deliberately does not do on the caller's behalf.
preflight(cfg)
if dry_run:
# Single-source the offline plan: plan_run resolves the id + fanout + runtime split, reports
# the exists-vs-new verdict, and emits the launch-command templates. It returns the run_id,
# so this contract is unchanged.
return launch_plan.plan_run(cfg, settings=settings, force=force).run_id
# Everything below reaches BigQuery and a runtime. A bare ``pip install scale-forecasting`` is
# the pure layer (the playground runs; this does not), so name the extra before the environment
# check — a core-only install has no use for ``SF_PROJECT_ID`` yet.
require_extra("gcp", purpose="Launching a run")
settings = settings or Settings.resolve()
# A pinned ``compute.profile.source`` is a human assertion that one specific run's measurements
# apply here; if the data has moved under it, say so once, now, rather than six times from
# inside each family's launch. Before the lock below, which would otherwise turn the system's
# own ``auto`` choice into something that looks pinned. ``force`` overrides, as it does the
# idempotency guard.
check_pinned_source(cfg, settings=settings, force=force)
# `backtest.short_series: "error"` is the one policy that refuses a run rather than deciding
# what a short series gets, and this is where it is worth refusing: the message lands in the
# operator's terminal, nothing has been provisioned, and no registry rows exist yet. A no-op
# under every other policy, and under this one on a panel that fits. Each Python engine keeps
# its own backstop for the launch paths that never come through here.
launch_plan.preflight_short_series(cfg, settings=settings)
# Lock `source: "auto"` before the digest — the same step `launch_plan.plan_run` / ``stage_run``
# take, and it has to happen here too or the three verbs would disagree about this run's
# identity: ``--dry-run`` would report an id the real run then never uses.
cfg = launch_plan.lock_profile_source(cfg, settings=settings)
run_dag = plan_dag(cfg)
run_id = run_dag.run_id
# Idempotency guard: the run_id is config-pinned, so re-running the same config is a no-op once
# it has COMPLETED — return without relaunching. Relaunching would resubmit each family's
# deterministic platform job id (a Dataproc batch_id / Ray submission_id) and collide
# (AlreadyExists), and the reused-attempt terminal write would clobber the completed run's job
# rows. ``force`` re-executes as a fresh attempt (distinct job ids); a run that never completed
# (never ran, or a prior FAILED/PARTIAL) falls through and runs.
if not force and header_status(run_id, settings=settings) == "COMPLETED":
_log.info("run %s already COMPLETED; skipping relaunch (pass force=True to re-run)", run_id)
return run_id
_log.info(
"run %s start: families=%s ensemble=%s",
run_id,
run_dag.families,
run_dag.ensemble_enabled,
)
bq_outcome: BqOutcome | None = None
# One error slot per family job (keyed by family name), plus the ensemble node's.
job_errors: dict[str, BaseException] = {}
ensemble_error: BaseException | None = None
# And one terminal status per job, keyed the same way — what its ``run_jobs`` row now says. The
# two are not the same fact and neither implies the other: a job can fail without raising (its
# cells all errored on a remote driver that reported the *submission* as fine), and the run
# header has to hear about it. See `job_outcome`.
job_statuses: dict[str, str | None] = {}
orig_native = run_dag.native_job
native = orig_native
python_jobs = run_dag.python_jobs
# Microbatch ensemble overlaps the base jobs: it drains series as each one's full base set lands
# (rather than waiting for the join like the barrier), so it runs as a concurrent pool task with
# an ``upstream_done`` predicate the join flips. ``base_done`` is set once every family (Python
# + native) has joined, telling the drain loop no further base predictions will land. A family
# that fails leaves its models absent, so no series ever reaches the full-base-set readiness bar
# and the concurrent node drains nothing — same "no ensembles when a base family failed" outcome
# as the barrier's post-join skip, just reached by readiness rather than an up-front gate.
ensemble_concurrent = run_dag.ensemble_enabled and cfg.compute.ensemble.mode == "microbatch"
base_done = threading.Event()
# One header owner: run_header writes RUNNING on entry and finalizes once, after every family
# job joins, with the combined status computed below. Every job runs with manage_header=False so
# nothing else touches this row. Per-family errors are captured (not raised through the block)
# so the finalize records the right status; the first is re-raised after, for a non-zero exit.
with run_header(cfg, run_id, settings=settings, manage=True) as hdr:
if not force and (
cfg.compute.reuse_prior_runs or cfg.compute.reuse_from_run_id is not None
):
from dataclasses import replace
from .registry.ops import (
borrow_model_results,
find_reusable_models,
record_borrowed_family_job,
)
borrowed = find_reusable_models(cfg, run_id, hdr.snapshot_millis, settings=settings)
if borrowed:
borrow_model_results(cfg, run_id, borrowed, settings=settings)
try:
merge_header_telemetry(
run_id,
{
"borrowed_models": {
k: {
"source_run_id": v.source_run_id,
"source_model_type": v.source_model_type,
"model_spec_hash": v.model_spec_hash,
"snapshot_millis": v.snapshot_millis,
}
for k, v in borrowed.items()
}
},
settings=settings,
)
except Exception as exc: # noqa: BLE001 - telemetry is best-effort
_log.warning("borrowed_models telemetry not recorded for %s: %s", run_id, exc)
if native is not None:
rem_native = tuple(m for m in native.models if m not in borrowed)
if not rem_native:
record_borrowed_family_job(
cfg,
run_id,
"native",
"bigquery",
None,
None,
None,
[borrowed[m] for m in native.models],
force=force,
settings=settings,
)
job_statuses["native"] = "COMPLETED"
native = None
elif len(rem_native) < len(native.models):
native = replace(native, models=rem_native)
remaining_py = []
for job in python_jobs:
rem_py = tuple(m for m in job.models if m not in borrowed)
if not rem_py:
record_borrowed_family_job(
cfg,
run_id,
job.family,
job.runtime,
None if job.compute is None else job.compute.spark_mode,
None if job.compute is None else job.compute.hardware,
None if job.compute is None else job.compute.gpu_type,
[borrowed[m] for m in job.models],
force=force,
settings=settings,
)
job_statuses[job.family] = "COMPLETED"
elif len(rem_py) < len(job.models):
remaining_py.append(replace(job, models=rem_py))
else:
remaining_py.append(job)
python_jobs = remaining_py
new_jobs = [*remaining_py, *([native] if native is not None else [])]
run_dag = replace(run_dag, jobs=tuple(new_jobs))
# Launch each Python family on its own worker thread, and run the BigQuery-native family
# inline on the main thread, so all families overlap. Each family carries the same
# contributor-mode contract (its model subset + shared header owned here) and its own
# run_jobs row, so N heterogeneous families run under one run_id.
# When the run has several ephemeral Ray (or Dataproc-cluster) families, provision one
# shared cluster per runtime for the duration of the launch block (each family submits its
# own job to it, torn down once on exit); otherwise these yield None and each family
# self-provisions as before.
# +1 pool worker for the concurrent microbatch ensemble so it never contends with a family
# for a thread; barrier mode keeps the exact family-only pool it always had.
max_workers = max(1, len(python_jobs)) + (1 if ensemble_concurrent else 0)
with (
shared_clusters.shared_ray_cluster(cfg, run_dag, run_id, settings) as ray_cluster,
shared_clusters.shared_spark_cluster(cfg, run_dag, run_id, settings) as spark_cluster,
ThreadPoolExecutor(max_workers=max_workers) as pool,
):
futures = {
pool.submit(
job_launch.launch_family_job,
cfg,
job,
run_id,
settings,
spark,
force=force,
max_executors=max_executors,
ray_cluster=ray_cluster,
spark_cluster=spark_cluster,
): job
for job in python_jobs
}
# Microbatch: fire the ensemble now, concurrently with the base jobs, draining ready
# series until the join flips base_done. Barrier: it stays a post-join step (below).
ensemble_future = (
pool.submit(
job_launch.launch_ensemble_job,
cfg,
run_id,
settings,
force=force,
upstream_done=base_done.is_set,
)
if ensemble_concurrent
else None
)
if native is not None:
try:
bq_outcome = job_launch.launch_native_job(
cfg, native, run_id, settings, force=force
)
job_statuses["native"] = "COMPLETED"
except Exception as exc: # noqa: BLE001 - captured, finalized below, re-raised
job_errors["native"] = exc
job_statuses["native"] = "FAILED"
for future, job in futures.items():
try:
# `launch_family_job` returns the status it wrote on the row, or ``None`` when
# the cell audit could not be read — in which case the row kept the finalizer's
# COMPLETED, so that is what the header must roll up. Reading it back out of
# BigQuery instead would be a second chance to be wrong about a row we wrote.
job_statuses[job.family] = future.result() or "COMPLETED"
except Exception as exc: # noqa: BLE001 - captured, finalized below, re-raised
job_errors[job.family] = exc
job_statuses[job.family] = "FAILED"
# Every base family has joined: no more base predictions will land, so the concurrent
# drain loop can stop after its final ready-series pass. Set before the ensemble join.
base_done.set()
if ensemble_future is not None:
try:
ensemble_future.result()
job_statuses["ensemble"] = "COMPLETED"
except Exception as exc: # noqa: BLE001 - captured, finalized below, re-raised
ensemble_error = exc
job_statuses["ensemble"] = "FAILED"
# Barrier ensemble: it reads every family's base predictions / backtest_oof, so it runs
# strictly after the join and only when every family succeeded. (Microbatch already ran
# concurrently above.) A failure here is captured like a family error — the ensembles are
# part of the run's success contract.
#
# "Succeeded" is the row's status, not merely "did not raise". A family whose every cell
# errored raises nothing and leaves no base predictions to blend, so an ensemble launched
# over it can only fail or produce a consensus of nothing — and until the cell audit existed
# there was no way to tell that family apart from a healthy one.
families_green = all(status == "COMPLETED" for status in job_statuses.values())
if not ensemble_concurrent and families_green and run_dag.ensemble_enabled:
try:
job_launch.launch_ensemble_job(cfg, run_id, settings, force=force)
job_statuses["ensemble"] = "COMPLETED"
except Exception as exc: # noqa: BLE001 - captured, finalized below, re-raised
ensemble_error = exc
job_statuses["ensemble"] = "FAILED"
# Combined status across the family jobs that ran: all green → COMPLETED, all failed →
# FAILED, some but not all → PARTIAL (surviving families' forecasts stay usable). An
# ensemble failure on top of all-green families fails the run: full output undelivered.
# The same roll-up the Airflow `finalize_run` task applies, over the same statuses.
status = combined_run_status(job_statuses, ensemble_enabled=run_dag.ensemble_enabled)
fields: dict[str, object] = {"bq_models": list(orig_native.models) if orig_native else []}
if bq_outcome is not None:
fields["n_series"] = bq_outcome.n_series
# The run-level honesty claim, alongside the per-cell `hpo_scoring` / `ensemble_scoring`
# columns: one place a reader can look to see whether anything this run *learned* was
# scored on a fold it had already optimised against. Best-effort — a telemetry merge that
# fails must not turn a green run red, and the per-cell columns still carry the answer.
try:
merge_header_telemetry(
run_id, {"scoring.hpo": hpo_scoring_claim(cfg)}, settings=settings
)
except Exception as exc: # noqa: BLE001 - telemetry is never worth failing a run over
_log.warning("scoring telemetry not recorded for run %s: %s", run_id, exc)
hdr.finalize(status=status, **fields)
first_error = next(iter(job_errors.values()), None) or ensemble_error
if first_error is not None:
# Re-raise the first failure so the CLI exits non-zero; the header already records the
# combined status (FAILED or PARTIAL).
raise first_error
if status != "COMPLETED":
# A run can now finish badly without anything having raised: every job's launch call
# returned, and the cell audit is what found that a family produced nothing. There is no
# exception to re-raise, so one is made here — otherwise the header would say FAILED while
# the CLI exited zero, and every unattended caller (Composer, a shell script, the SDK)
# would read the run as fine. The registry rows carry the detail; this only sets the exit.
# Say "did not complete", not "produced nothing". A family here may be FAILED (no usable
# cell) or PARTIAL (most cells fine, some errored), and `job_statuses` cannot tell them
# apart — it holds one status per family, not the tallies behind it. Claiming the stronger
# of the two sends an operator away from a run that mostly worked. The tally is one query
# from here, so the message points at it rather than guessing.
failed = sorted(f for f, s in job_statuses.items() if s != "COMPLETED")
raise EngineError(
f"run {run_id} finished {status}: {', '.join(failed)} did not complete "
f"(per-family cell tallies in run_jobs.job_telemetry.$.cells)"
)
_log.info("run %s done: status=%s", run_id, status)
return run_id