Configuration

config

Load, validate, and freeze the run config — the single source of run behavior.

A run is one JSON file. It is validated here before anything executes, and the frozen, normalized object is what gets logged verbatim to run_registry.raw_config — so the config is the experiment record.

Public surface: - RunConfig — the frozen pydantic model. - load_config(path) -> RunConfig — read + validate a JSON file. - estimate_workload(cfg, obs_counts=...) -> Workload — the dry-run work estimate: cells always, fits and fold cohorts when the caller supplies per-series observation counts. - estimate_fanout(cfg) -> Fanout — the same estimate narrowed to its four count fields.

BacktestConfig

Bases: BaseModel

Time-series cross-validation. Off by default (cheapest first run).

Every field here is now honoured. The five that arrived ahead of the code that reads them — short_series, min_folds, min_train_floor, gap and window — were declared in one commit because run_id is a digest of the whole config: adding a field moves every identity ever recorded, so landing them together cost one identity break instead of five. test_declared_ahead_fields.py holds the other half of that bargain: each field reaches the digest.

short_series decides what happens to a series too short for the requested fold grid, and every branch of it is a different trade rather than a different amount of the same thing. adapt (the default) holds the geometry exactly and drops the oldest folds, so a short series is scored on the most recent window it can reach. overlap holds the fold count and shrinks the step to buy it, which means validation windows share observations and the per-fold scores stop being independent evidence. shrink_train holds the count and the step and lowers the training requirement instead, never past min_train_floor — which is why that field is required with this mode rather than optional. skip leaves any series short of the full grid unscored, so every series on the leaderboard was measured identically. error refuses the run outright; it is checked before any cell runs and never raises inside one, because a scoring shortfall must never cost a forecast. min_folds is the give-up floor the first three respect: a series that cannot reach it is left unscored rather than weakly scored.

gap and window are also honoured. gap is an embargo: training stops gap observations before the validation window starts, so train_end + gap == val_start and the fold measures a forecast issued with a reporting lag. It moves the training end, never the validation window — the folds of a gap=14 run cover the same dates as the folds of a gap=0 run, so the two are comparable — and it costs history, so a series may achieve fewer folds under it. window is the sliding scheme's training width, defaulting to min_train, which frees min_train to mean only the feasibility floor. Both are mirrored into the BigQuery-native fold SQL, so the two engines score the same windows.

scheme is the one that is fully honoured, and it decides what a fold's score is a score of. expanding and sliding refit at every origin, so a score is about a freshly-trained model. expanding_frozen fits once and then feeds the model the observations that arrive between origins with its parameters held fixed, which measures what refitting less often costs. expanding_stale fits once and never tells the model what happened next, which measures how fast it decays untouched — the only one of the three every model can answer identically, and so the one where a cross-model leaderboard is comparing like with like. Widening this Literal moves no existing run_id — the digest hashes dumped values, not the schema.

control_arm is what lets a refit scheme answer the frozen schemes' question without becoming one. Both frozen schemes score a blind arm — one fit, walked forward untouched — beside their primary arm, and expanding and sliding did not, so the run most people actually make could not say what refitting was buying it. The arm costs one extra fit per cell, not per fold, plus a forecast, which is why it is affordable on the default path. On expanding_stale it is refused rather than ignored: that scheme's primary arm already is the blind arm. See backtest.BacktestOutcome for what the two arms produce.

One asymmetry worth knowing: short_series is a Python-path policy. The BigQuery-native models count their folds back from one global MAX(ds) rather than from each series' own last observation, so they have no per-series grid to adapt; a short series there simply contributes fewer scored rows. error is the exception, because it is checked against the panel before either engine starts.

CapacityConfig

Bases: BaseModel

How hard to look for room when a service says it has none — per service (G2).

"Resources are not available" is a state, not an exception: a run walks its candidate places, and if none has room it waits and walks them again, until an attempt budget or a clock runs out. scale_forecasting.capacity implements that; this is where a run tunes it.

Not part of the run_id digest — excluded in registry.ids._NOT_IDENTITY, and there is a test that says so. Same rule as compute.profile.source, for the same reason: a run's identity is what was asked for, and patience is an operational knob. If it moved the digest, waiting longer for a GPU would fork your run id and break dedupe-on-read.

enabled: false restores the pre-retry behaviour exactly — one pass over the candidates, no back-off — which is the escape hatch if a retry loop ever misbehaves in production. It does not disable classification: a config fault still stops immediately and a quota ceiling is still named as one in the ledger, because those were improvements to the diagnosis, not to the patience.

BigQuery has no entry and that omission is deliberate (capacity.UNMANAGED_SERVICES): slot contention is resolved BigQuery-side and surfaces as latency, not as a create that failed somewhere and could be retried elsewhere. There is no candidate list to walk.

policy_for

policy_for(service)

Resolve this config into the runtime policy for service (pure).

The shipped per-service default with any authored field laid over it, plus the enabled: false collapse to a single pass. Raises KeyError for a service with no candidate walk (BigQuery), because asking for its retry policy is a programming error rather than a configuration one.

Source code in src/scale_forecasting/config.py

def policy_for(self, service: str) -> CapacityPolicy:
    """Resolve this config into the runtime policy for ``service`` (pure).

    The shipped per-service default with any authored field laid over it, plus the
    ``enabled: false`` collapse to a single pass. Raises `KeyError` for a service with no
    candidate walk (BigQuery), because asking for its retry policy is a programming error rather
    than a configuration one.
    """
    base = DEFAULT_POLICIES[service]
    override: CapacityServicePolicy = getattr(self, service)
    fields = {
        key: value
        for key, value in override.model_dump().items()
        if value is not None and key != "max_passes"
    }
    max_passes = 1 if not self.enabled else (override.max_passes or base.max_passes)
    return replace(base, max_passes=max_passes, **fields)

CapacityServicePolicy

Bases: BaseModel

A partial override of one service's shipped retry policy — every field optional.

Partial on purpose. A user who wants to wait longer for a GPU should be able to say only {"max_wall_seconds": 7200} and inherit the rest; requiring the whole policy would mean copying four numbers they have no opinion about and silently freezing them against future default changes. Unset fields fall through to capacity.DEFAULT_POLICIES[service].

Bounds are validated here as well as in capacity.CapacityPolicy, so a bad config fails at load with a pydantic error naming the field rather than at provisioning time with a ValueError.

ComputeConfig

Bases: BaseModel

Runtime scale, dependency delivery, and cost guardrails.

DataConfig

Bases: BaseModel

Where the series come from and their shape.

source_table and the three column names are written into BigQuery SQL inside backticks (SELECT `{ts_id_col}` ... FROM `{source_table}```) by the sizing probes, the SDK, and the native engines. They are validated here, once, at config-load time rather than at each of those sites: a column must be a BigQuery column name (letter or underscore, then letters, digits, or underscores, at most 300 characters) and a table path may betable,dataset.table, orproject.dataset.table(project ids may carry hyphens) — but neither may contain a backtick, whitespace, or a control character, which are the characters that could end the quoted identifier early. Validation adds no field, so it moves norun_id``.

EnsembleCompute

Bases: BaseModel

When the ensemble DAG node runs. Not where — there is only one where.

Distinct from EnsembleConfig, which selects the ensemble strategies. mode="barrier" ensembles once after every base model finishes; mode="microbatch" ensembles each series as soon as its upstream base models complete, so the ensemble overlaps the families instead of queueing behind the slowest one.

Both modes are live. The microbatch shape was measured on Airflow in smoke 15 (2026-09-20): the ensemble task started in the same second as the four family tasks, ran 2,879 s alongside them, and finished 20 s after the last member — which is not a shape barrier mode can produce.

This model used to carry runtime, spark_mode and spark_cluster_name too. They were accepted by the loader, documented as inert, and read by nothing: the ensemble node is hard-wired to the driver (dag.build_dag_nodes stamps it runtime="bigquery", and job_launch.run_ensemble blends in driver pandas, taking no cluster of its own). Because they were still in the config digest, setting one started a new run and changed nothing about it. They are gone rather than documented-as-inert so the surface stops offering a choice that does not exist. Setting one now fails at load, which beats quietly starting a differently-keyed run that behaves identically.

No run_id moved. Identity hashes model_dump(), defaults included, so dropping three defaulted fields would ordinarily re-key every config in existence and stale the whole validation ledger to buy a tidier surface. registry.ids._REMOVED_DEFAULTS pins the three keys back into the digest at the values they used to carry, which is legitimate here precisely because nothing ever read them: no id on record ever meant anything different from what it means now.

EnsembleConfig

Bases: BaseModel

Consensus across base models.

strategies is a list so several ensembles can run at once. The singular strategy string is accepted as shorthand for a one-element list.

FamilyCompute

Bases: BaseModel

A sparse per-family compute override, layered over the flat ComputeConfig defaults.

Every field is optional: an unset field inherits the run-level default (python_runtime, Spark serverless, CPU, the flat gpu_type), so a config sets only what a family needs to differ on. There is no native family here — native models always run in BigQuery. See RunConfig.resolve_family_compute for how these layer onto the defaults; the block is inert until the DAG orchestrator consumes it.

Hardware constraints (validated): only the deep_learning family may request a GPU (enforced where the family key is known, in ComputeConfig); Dataproc Serverless offers L4 only (no T4/A100 — use spark_mode="cluster" or runtime="ray"/"vertex"/"gce"); Spark-only fields (spark_mode/spark_cluster_name) are rejected on runtime in ("ray", "vertex", "gce"); VM fields (machine_type/workers) are rejected on runtime in ("spark", "ray"); and runtime="gce" is strictly single-VM (workers cannot exceed 1).

Fanout dataclass

Fanout(n_series, n_models, n_folds, n_cells)

The count half of a Workload, kept for callers that predate it.

estimate_fanout is still in the public __init__ and still returns this, so nothing downstream had to move. One number did change meaning: n_cells is now n_series × n_models, matching every other n_cells in the system (the leaderboard's, the Ray planner's, the quota estimator's). It used to multiply by folds, which is why a two-fold dry run reported three times the cells it would ever write a row for.

FeaturesConfig

Bases: BaseModel

Optional feature engineering for the Python models.

Defaults are conservative/generic (no transform, no holidays); the shipped example_config.json turns on holidays + log1p.

There is no lags field, and the reason is a design rule rather than an omission. "Lag" means two different things, and only one of them can be built out here:

• A lag of the target is only honest at predict time if the model rolls its own forecasts forward one step at a time. That is recursion, it belongs to the model, and the models that need it already own it — _lag_forecaster does it for lightgbm / xgboost / regression_lags, and NeuralProphet does it behind n_lags. A config cannot supply one, because nobody knows the future target; the removed field filled the horizon with a flat line at the last observation and handed the result to models that had fitted coefficients against real history.
• A lag of a covariate is honest, because a covariate's future is as knowable as the covariate itself. No model in the suite lags one internally, so this is the gap that exog_lags fills.

exog_lags maps an exog column to the lags to build from it: {"promo": [1, 7]} adds promo_lag_1 and promo_lag_7. They are ordinary columns by the time any model sees them, so every supports_exog model can use them.

Precedence, when the two rules ever meet: the model wins. A model that lags covariates internally sets lags_covariates_internally and is handed the unlagged columns only, so a covariate is never lagged twice. Every shipped model leaves that flag False today.

all_covariates property

all_covariates

Every declared covariate column (dynamic_covariates + static_covariates), order-preserving and deduplicated.

dynamic_covariates property

dynamic_covariates

Time-varying numeric covariates (exog + future_covariates + past_covariates), order-preserving and deduplicated.

known_future_covariates property

known_future_covariates

Time-varying covariates known across the forecast horizon (exog + future_covariates), order-preserving and deduplicated.

HierarchyConfig

Bases: BaseModel

Hierarchical and grouped time-series aggregation and coherent forecast reconciliation.

Off by default (enabled: false). When enabled, levels defines the cross-sectional aggregation hierarchy from top to bottom (for example [["region"], ["region", "category"]]), with the total aggregate ("__total__") prepended automatically and the bottom-level series (data.ts_id_col) at the leaves.

Implements the Hyndman & Athanasopoulos (FPP3 Ch. 11) reconciliation projection y_tilde = S @ G @ y_hat across eight methods:

• bottom_up — sum bottom-level forecasts upward via the summing matrix S.
• top_down — disaggregate the top-level forecast downward by average historical proportions (FPP3 §11.2).
• middle_out — anchor on middle_level (sum upward above it, disaggregate downward below it by historical proportions within each middle-level node).
• ols — ordinary least squares MinT reconciliation (W_h = I).
• wls_struct — structural scaling WLS (W_h = diag(S @ 1)), requiring only the hierarchy structure S.
• wls_var — variance scaling WLS (W_h = diag(W_hat_1)) weighted by inverse residual variances.
• mint_shrink — Wickramasuriya et al. (2019) Minimum Trace with Schäfer-Strimmer shrinkage covariance of residuals, positive-definite even when n_series >> T_obs.
• mint_nonneg — active-set non-negative Minimum Trace reconciliation guaranteeing bottom-level forecasts >= 0 while preserving exact additive coherence across all levels.

HpoConfig

Bases: BaseModel

Hyperparameter optimization on the aligned backtest (optional).

Off by default. When enabled, an Optuna study tunes each model's search_space on the backtest folds and the winning params are stamped to forecast_metadata.best_params (see scale_forecasting.hpo). HPO therefore requires backtest.enabled.

granularity is the DS-facing cost knob:

• fleetwide (default) — tune each model once on a sample_size sample of series and apply the winner across all series. The only granularity affordable at the 100k hero scale: the study runs on the driver before fan-out (a handful of series × n_trials fits), not per cell.
• per_series — tune inside every cell (n_trials fits per series). Accurate for the tail of hard series but multiplies fit cost by n_trials fleet-wide; an explicit opt-in.

OutputConfig

Bases: BaseModel

What the numbers we ship mean — the functional behind yhat and physical bounds.

Its own section rather than a field on BacktestConfig, even though both of its rules involve the backtest, because a reader asking "how do I control the point forecast?" will not look under backtest, and the docs are a product surface here.

Every model emits one number per future date, so something decides what that number is. For a long time this project decided by accident: ten of sixteen models built their band from residual quantiles, the frame assembler took the 0.5 quantile as yhat, and the shipped forecast was silently prediction + median(in-sample residual) — un-named, un-configurable, and applied to some models and not others. Measurement said the correction was worth keeping (it moved fleet WAPE by 5.7%), so it is the default. This field is what turns it from an opinion into a default: the alternative is now sayable.

yhat_raw is written alongside yhat whatever this is set to, so the choice is never destructive and the two arms can always be compared after the fact. Physical constraints (floor, cap, floor_col, cap_col, and integer_rounding) are applied post-calibration so calibrate_from_oof estimates empirical residuals on unclamped raw model outputs while yhat, yhat_lower, yhat_upper, quantiles, and scored OOF predictions respect physical reality. See calibration.py.

effective_quantiles property

effective_quantiles

Sorted quantile levels to calibrate and emit for each horizon step.

effective_target_quantile property

effective_target_quantile

Resolved operational quantile $q^* = c_u / (c_u + c_o)$ or target_quantile.

ProfileConfig

Bases: BaseModel

Whether to measure what a run costs before sizing it, and with how much headroom.

Sizing today is a pure cell count: n_series x n_models x n_folds, turned into nodes by a flat cells-per-slot constant. Nothing in that arithmetic knows that a deep-learning fit and a naive mean differ by orders of magnitude, so a fleet is provisioned for the count and not for the work. Turning this on replaces the guess with a short instrumented pre-pass (scale_forecasting.profiling): fit a stratified sample of series, measure what they actually consumed, and size each family's slot from the measurement.

The two margins are deliberately different, and the asymmetry is the point: over-estimating time buys extra slots, which costs money, while under-estimating memory OOM-kills the task, which costs the run. So memory carries the larger headroom, and the two are applied to different tails — memory_margin to the sample max (a slot must hold the worst case that lands in it), time_margin to the median (a fleet is sized for typical work, and sizing it for the worst case over-provisions every run).

Part of the run_id digest — unlike the two fields that are not (its own source, and compute.capacity, both of which are resolved or operational rather than authored). It changes the resource shape rather than the numbers a run produces, so it is arguable — but the config is the experiment record, and a run whose fleet was sized differently is not the same run for performance purposes. Silently varying the shape under a stable id would be the worse trade.

A profile is produced by one run and consumed by later ones, which is why mode is split into a source (what evidence to consume) and a measure (what evidence to produce) below. Every runtime consumes: submit.sizing_properties, dataproc_cluster.cluster_sizing and the Ray sizing paths all resolve one through profiling.source.profile_for_run.

What differs per runtime is when the measurement can be taken, and only Ray gets a choice. spark.executor.cores and spark.task.cpus are fixed at submit (Serverless) or at create (cluster), before any of our code runs out there, so on Spark the only evidence that can reach a sizing decision is prior evidence — a harvest, or the shipped baseline. ray_engine can additionally call profiling.source.resolve_profile and measure in-run, because a Ray task's num_cpus/num_gpus is a request made against a pool that already exists. So "auto" and "always" still differ only on the Ray path; "off" is the setting that is load-bearing everywhere, where it also suppresses the static-arithmetic overlay.

consumes_evidence property

consumes_evidence

Should this run try to size itself from measurements? (mode="off" vetoes it.)

needs_source_resolution property

needs_source_resolution

Is the source still a question rather than an answer? (i.e. must plan time lock it.)

records_measurements property

records_measurements

Should this run write per-cell measurements? (mode="off" vetoes measure.)

unpins_threads property

unpins_threads

Should the fleet leave native thread pools uncapped so effective_cores is real?

ResolvedFamilyCompute dataclass

ResolvedFamilyCompute(
    family,
    runtime,
    spark_mode,
    spark_cluster_name,
    hardware,
    gpu_type,
    machine_type=None,
    workers=None,
    min_workers=None,
    max_workers=None,
    accelerator_count=0,
    gke_mode=None,
    gke_cluster_name=None,
    ray_mode=None,
    automl_mode=None,
)

One family's effective compute after layering its override on the flat defaults (pure).

The fully-resolved plan the DAG orchestrator acts on: runtime is where the family's job runs; spark_mode/spark_cluster_name are None unless runtime == "spark"; gke_mode/gke_cluster_name apply when runtime == "gke" (or runtime == "ray" with ray_mode == "gke"); ray_mode applies when runtime == "ray"; automl_mode applies when runtime == "vertex_automl"; gpu_type is None unless hardware == "gpu". Produced by RunConfig.resolve_family_compute.

RetryResources

Bases: BaseModel

What a repair may claim differently from the attempt it repairs.

A repair re-submits a handful of cells out of a run that was sized for all of them. The fleet arithmetic does not know that: it sizes from the config's fan-out, so a forty-cell repair of a hundred-thousand-cell family asks for the hundred-thousand-cell fleet and pays for it. This is where a run says how wide its repairs should be instead.

It lives here, under compute.capacity, for the reason the parent's docstring already gives: everything under capacity is excluded from the run_id digest, and a repair that resized itself must stay the same run. A sibling of compute.max_executors would fork the id, which would mean the repair wrote its rows under a run nobody was looking at.

It is a config block rather than only a CLI flag for the same reason compute.max_executors is: the unattended repair paths (airflow_tasks.retry_families, main's --retry) launch from a config and nothing else, so a ceiling reachable only through an argument is a ceiling an orchestrated repair can never set. An explicit max_executors= argument still wins over it (job_launch.submit_retry) — the operator at the terminal overrides the file.

Why only this one knob, and what may join it. A repair reuses the config its run staged, so the only settings it can vary are the ones the driver reads at launch: anything the worker reads out of the staged config (bucket_target_cells, max_parallelism, the model list, the data block) is the same bytes for the repair as for the attempt, by construction. Of the driver-side knobs, max_executors is the one the launchers already take as an argument (job_launch.launch_family_job). Note its reach: it caps a Spark batch's spark.dynamicAllocation.maxExecutors and is ignored by the Ray and in-process paths, so a Ray family's repair still provisions from the fan-out.

RunConfig

Bases: BaseModel

A complete, validated, frozen run specification.

max_horizon property

max_horizon

The largest horizon any predict call in this run will be asked for.

Two different horizons exist in a config and it is easy to reach for the wrong one. The forward forecast uses data.horizon; every backtest fold predicts backtest.horizon, which may be larger. Anything sizing itself against "the horizon" — a params validator refusing a model that cannot emit enough steps, a context handed to a model — has to mean the larger of the two, because the run will ask for both. A property rather than a field: it is derived, so it stays out of model_dump and no run_id moves.

The embargo counts, and leaving it out was a bug. A fold's model stops at train_end while its validation window starts gap observations later, so the model has to forecast across the embargo before it reaches anything that gets scored: the Python fold asks for gap + val_size steps and throws the first gap away, and the BigQuery fold asks its ML.FORECAST for gap + backtest.horizon. Both engines were already right; this number, which is supposed to describe them, was short by exactly gap. What that cost: NeuralProphet emits exactly n_forecasts direct steps and does not recurse, so a gap=3 run passed its params validator and then came up three steps short at the tail of every fold.

resolve_family_compute

resolve_family_compute(family)

Resolve one family's effective compute by layering its override on the flat defaults.

Pure and deterministic — the single resolver the DAG orchestrator also uses, so a config validated at load needs no re-check at submit. family is a compute family (statistical/ml/deep_learning/automl); native has no compute choice (it always runs in BigQuery) and raises. Resolution rules for unset override fields:

• runtime → python_runtime (or "vertex_automl" for automl).
• hardware → gpu only for deep_learning and automl (when compute.use_gpu or an explicit override); every other family is cpu.
• Spark: spark_mode → serverless; spark_cluster_name applies only under cluster. On ray, vertex, gce, gke, and vertex_automl both are None.
• GKE & Ray modes: gke_mode → compute.gke_mode ("job" | "ray") on gke; ray_mode → compute.ray_mode ("vertex" | "gke") on ray (or "gke" when runtime == "gke" and gke_mode == "ray").
• AutoML mode: automl_mode → compute.automl_mode ("tabular_workflow" | "training_job") on vertex_automl.
• gpu_type (when hardware == "gpu") → the flat compute.gpu_type, but forced to L4 on Dataproc Serverless (no T4/A100 there). Inheriting T4/A100 on Serverless raises.
• machine_type → compute.machine_type (auto-resolved via resolve_vm_machine_type from (hardware, gpu_type, accelerator_count) on vertex/gce/gke job, or when explicitly set on ray/spark cluster/gke ray/vertex_automl); workers → compute.workers on vertex and gke job, 1 on gce, or the family override when set on ray/spark/vertex_automl; min_workers / max_workers apply on autoscaling runtimes (ray/spark/gke ray/ vertex_automl).

Source code in src/scale_forecasting/config.py

def resolve_family_compute(self, family: str) -> ResolvedFamilyCompute:
    """Resolve one family's effective compute by layering its override on the flat defaults.

    Pure and deterministic — the single resolver the DAG orchestrator also uses, so a config
    validated at load needs no re-check at submit. ``family`` is a compute family
    (``statistical``/``ml``/``deep_learning``/``automl``); ``native`` has no compute choice (it
    always runs in BigQuery) and raises. Resolution rules for unset override fields:

    * ``runtime`` → ``python_runtime`` (or ``"vertex_automl"`` for ``automl``).
    * ``hardware`` → ``gpu`` only for ``deep_learning`` and ``automl`` (when
      ``compute.use_gpu`` or an explicit override); every other family is ``cpu``.
    * Spark: ``spark_mode`` → ``serverless``; ``spark_cluster_name`` applies only under
      ``cluster``. On ``ray``, ``vertex``, ``gce``, ``gke``, and ``vertex_automl`` both are
      ``None``.
    * GKE & Ray modes: ``gke_mode`` → ``compute.gke_mode`` (``"job"`` | ``"ray"``) on ``gke``;
      ``ray_mode`` → ``compute.ray_mode`` (``"vertex"`` | ``"gke"``) on ``ray`` (or ``"gke"``
      when ``runtime == "gke"`` and ``gke_mode == "ray"``).
    * AutoML mode: ``automl_mode`` → ``compute.automl_mode`` (``"tabular_workflow"`` |
      ``"training_job"``) on ``vertex_automl``.
    * ``gpu_type`` (when ``hardware == "gpu"``) → the flat ``compute.gpu_type``, but **forced to
      L4** on Dataproc Serverless (no T4/A100 there). Inheriting T4/A100 on Serverless raises.
    * ``machine_type`` → ``compute.machine_type`` (auto-resolved via `resolve_vm_machine_type`
      from ``(hardware, gpu_type, accelerator_count)`` on ``vertex``/``gce``/``gke`` ``job``, or
      when explicitly set on ``ray``/``spark`` ``cluster``/``gke`` ``ray``/``vertex_automl``);
      ``workers`` → ``compute.workers`` on ``vertex`` and ``gke`` ``job``, ``1`` on ``gce``, or
      the family override when set on ``ray``/``spark``/``vertex_automl``; ``min_workers`` /
      ``max_workers`` apply on autoscaling runtimes (``ray``/``spark``/``gke`` ``ray``/
      ``vertex_automl``).
    """
    if family == "native":
        raise ValueError(
            "native models always run in BigQuery; they have no per-family compute"
        )
    ov = self.compute.families.get(_as_compute_family(family)) or FamilyCompute()
    if family == "automl":
        runtime = ov.runtime or "vertex_automl"
        if runtime != "vertex_automl":
            raise ValueError(
                f"family 'automl' models require runtime='vertex_automl' (got {runtime!r})"
            )
    else:
        runtime = ov.runtime or self.python_runtime
        if runtime == "vertex_automl":
            raise ValueError(
                f"family '{family}' cannot run on runtime='vertex_automl'; "
                f"set compute.families.{family}.runtime or use "
                "python_runtime='spark'/'ray'/'vertex'/'gce'/'gke'"
            )

    if family in ("deep_learning", "automl"):
        hardware = ov.hardware or ("gpu" if self.compute.use_gpu else "cpu")
    else:
        hardware = "cpu"

    if runtime == "spark":
        spark_mode = ov.spark_mode or "serverless"
        spark_cluster_name = ov.spark_cluster_name if spark_mode == "cluster" else None
    else:
        if ov.spark_mode is not None or ov.spark_cluster_name is not None:
            raise ValueError(
                f"family '{family}': spark_mode/spark_cluster_name require runtime='spark' "
                f"(got {runtime!r})"
            )
        spark_mode = None
        spark_cluster_name = None

    gke_mode: str | None = None
    gke_cluster_name: str | None = None
    ray_mode: str | None = None
    automl_mode: str | None = None
    if runtime == "gke":
        if ov.ray_mode is not None:
            raise ValueError(
                f"family '{family}': ray_mode requires runtime='ray' (got {runtime!r}); "
                "use gke_mode='ray' instead"
            )
        if ov.automl_mode is not None:
            raise ValueError(
                f"family '{family}': automl_mode requires runtime='vertex_automl' "
                f"(got {runtime!r})"
            )
        gke_mode = ov.gke_mode or self.compute.gke_mode
        gke_cluster_name = ov.gke_cluster_name or self.compute.gke_cluster_name
        ray_mode = "gke" if gke_mode == "ray" else None
    elif runtime == "ray":
        if ov.gke_mode is not None:
            raise ValueError(
                f"family '{family}': gke_mode requires runtime='gke' (got {runtime!r}); "
                "use ray_mode='gke' instead"
            )
        if ov.automl_mode is not None:
            raise ValueError(
                f"family '{family}': automl_mode requires runtime='vertex_automl' "
                f"(got {runtime!r})"
            )
        ray_mode = ov.ray_mode or self.compute.ray_mode
        if ray_mode == "gke":
            gke_mode = "ray"
            gke_cluster_name = (
                ov.gke_cluster_name
                or self.compute.gke_cluster_name
                or self.compute.ray_cluster_name
            )
        elif ov.gke_cluster_name is not None:
            raise ValueError(
                f"family '{family}': gke_cluster_name requires runtime='gke' or ray_mode='gke'"
            )
    elif runtime == "vertex_automl":
        if ov.gke_mode is not None or ov.gke_cluster_name is not None:
            raise ValueError(
                f"family '{family}': gke_mode/gke_cluster_name require runtime='gke' "
                f"(got {runtime!r})"
            )
        if ov.ray_mode is not None:
            raise ValueError(
                f"family '{family}': ray_mode requires runtime='ray' (got {runtime!r})"
            )
        automl_mode = ov.automl_mode or self.compute.automl_mode
    else:
        if ov.gke_mode is not None or ov.gke_cluster_name is not None:
            raise ValueError(
                f"family '{family}': gke_mode/gke_cluster_name require runtime='gke' "
                f"(got {runtime!r})"
            )
        if ov.ray_mode is not None:
            raise ValueError(
                f"family '{family}': ray_mode requires runtime='ray' (got {runtime!r})"
            )
        if ov.automl_mode is not None:
            raise ValueError(
                f"family '{family}': automl_mode requires runtime='vertex_automl' "
                f"(got {runtime!r})"
            )

    gpu_type: str | None
    accelerator_count: int
    if hardware == "gpu":
        if runtime == "spark" and spark_mode == "serverless":
            if ov.gpu_type in ("T4", "A100", "A100_80GB"):
                raise ValueError(
                    f"family '{family}': Dataproc Serverless supports L4 only, not "
                    f"{ov.gpu_type}; use spark_mode='cluster' or "
                    "runtime='ray'/'vertex'/'gce'/'gke'"
                )
            gpu_type = "L4"
        else:
            gpu_type = ov.gpu_type or self.compute.gpu_type
        accelerator_count = ov.accelerator_count or self.compute.accelerator_count
        _default_gpu_machine_type(gpu_type, accelerator_count)
    else:
        if ov.accelerator_count is not None:
            raise ValueError(
                f"family '{family}': accelerator_count requires hardware='gpu' "
                f"(got hardware={hardware!r})"
            )
        gpu_type = None
        accelerator_count = 0

    # ``None`` on the autoscaling runtimes means "no fixed pool" — the ceiling is max_workers.
    workers: int | None
    min_workers: int | None = None
    max_workers: int | None = None
    if runtime in ("vertex", "gce") or (runtime == "gke" and gke_mode == "job"):
        if (
            ov.min_workers is not None
            or ov.max_workers is not None
            or (
                ov.runtime is None
                and (
                    self.compute.min_workers is not None or self.compute.max_workers is not None
                )
            )
        ):
            raise ValueError(
                f"family '{family}': min_workers/max_workers are only supported on "
                f"autoscaling runtimes ('ray' or 'spark'), not {runtime!r}; use 'workers'"
            )
        machine_type = resolve_vm_machine_type(
            hardware,
            gpu_type,
            self.compute.machine_type,
            ov.machine_type,
            accelerator_count=accelerator_count if hardware == "gpu" else 1,
        )
        if runtime == "gce":
            if (ov.workers is not None and ov.workers > 1) or (
                ov.runtime is None and self.compute.workers > 1
            ):
                raise ValueError(
                    f"family '{family}': runtime='gce' supports single-VM execution only "
                    "(workers=1); use runtime='vertex' or runtime='gke' for multi-worker pools"
                )
            workers = 1
        else:
            workers = ov.workers or self.compute.workers
    else:
        if runtime == "spark" and spark_mode == "serverless" and ov.machine_type is not None:
            raise ValueError(
                f"family '{family}': Dataproc Serverless does not use VM machine_type; "
                "use spark_mode='cluster' or runtime='vertex'/'gce'/'gke'/'ray'"
            )
        if ov.machine_type is not None or (
            self.compute.machine_type != "auto"
            and not (runtime == "spark" and spark_mode == "serverless")
        ):
            machine_type = resolve_vm_machine_type(
                hardware,
                gpu_type,
                self.compute.machine_type,
                ov.machine_type,
                accelerator_count=accelerator_count if hardware == "gpu" else 1,
            )
        else:
            machine_type = None
        workers = ov.workers
        min_workers = ov.min_workers if ov.min_workers is not None else self.compute.min_workers
        max_workers = ov.max_workers if ov.max_workers is not None else self.compute.max_workers
        if min_workers is not None and max_workers is not None and min_workers > max_workers:
            raise ValueError(
                f"family '{family}': min_workers ({min_workers}) cannot exceed "
                f"max_workers ({max_workers})"
            )

    return ResolvedFamilyCompute(
        family=family,
        runtime=runtime,
        spark_mode=spark_mode,
        spark_cluster_name=spark_cluster_name,
        hardware=hardware,
        gpu_type=gpu_type,
        machine_type=machine_type,
        workers=workers,
        min_workers=min_workers,
        max_workers=max_workers,
        accelerator_count=accelerator_count,
        gke_mode=gke_mode,
        gke_cluster_name=gke_cluster_name,
        ray_mode=ray_mode,
        automl_mode=automl_mode,
    )

with_available_models

with_available_models()

Return a copy with models filtered to those whose upstream packages are installed.

Raises ConfigError if none of the requested models have their required packages installed in the current environment.

Source code in src/scale_forecasting/config.py

def with_available_models(self) -> RunConfig:
    """Return a copy with ``models`` filtered to those whose upstream packages are installed.

    Raises `ConfigError` if none of the requested models have their required packages
    installed in the current environment.
    """
    from .models import filter_available_models

    kept = filter_available_models(self.models)
    if not kept:
        raise ConfigError(
            f"none of the configured models ({list(self.models)}) have their required "
            f"packages installed in this environment"
        )
    if list(kept) == list(self.models):
        return self
    return self.model_copy(update={"models": kept})

with_series_limit

with_series_limit(n_series)

Return a copy with data.series_limit overridden (self if n_series is None).

The scale knob every submit path shares (the 10 → 100 → 1k → 100k story). Because it changes the config, the copy yields a distinct run_id, so each scale is its own queryable run.

Source code in src/scale_forecasting/config.py

def with_series_limit(self, n_series: int | None) -> RunConfig:
    """Return a copy with ``data.series_limit`` overridden (``self`` if ``n_series`` is None).

    The scale knob every submit path shares (the 10 → 100 → 1k → 100k story). Because it
    changes the config, the copy yields a distinct ``run_id``, so each scale is its own
    queryable run.
    """
    if n_series is None:
        return self
    return self.model_copy(
        update={"data": self.data.model_copy(update={"series_limit": n_series})}
    )

Workload dataclass

Workload(
    n_series,
    n_models,
    n_folds,
    n_cells,
    n_fits,
    full_fit_equivalents,
    train_rows_total,
    fold_histogram,
    n_unscored,
)

Dry-run estimate of the work a run will schedule.

Two halves, and the split is the point. The count half (n_series … n_cells) is a function of the config alone, so a plain --dry-run fills it with no environment and no data read. The geometry half (n_fits onward) depends on how long each series actually is, which only the panel knows — so it is None/empty unless the caller supplies obs_counts (what plan --feasibility reads). Reporting a guess there would be worse than reporting nothing: on a ragged panel the guess is wrong in the direction that flatters the plan.

n_cells is n_series × n_models — one cell is one (series, model) pair, which is what run_cell runs, what engines.ray_io.plan_cluster sizes against, and what the leaderboard's n_cells counts. It deliberately does not multiply by folds: folds happen inside a cell, and the old fold-multiplied number made a backtested run look like it scheduled three times the work when it schedules the same work three times as deep.

This is the estimate; the run records the measurement. forecast_metadata.n_fits and train_rows_total are the same two quantities counted at the fit site by worker.run_cell. The two disagree on four of the six refit paths, because the estimate below assumes a fresh fit per fold and only expanding/sliding do that — see backtest.FitTally. That gap is not an error in either number. It is what the scheme cost or saved, and a completed run can be asked for it.

corrected_arm_for

corrected_arm_for(decision_metric)

Which corrected point-forecast arm a decision metric implies — "mean" or "median".

A metric whose loss is quadratic in the error is minimised by the mean; everything that is absolute-error-shaped, and every proper interval score, is minimised by the median. The pairing is a theorem, not a preference, so it is what a run falls back to whenever no measurement is available to do better.

One rule, two callers. calibration.select_arm uses it under point_forecast="auto" to know which arm it is weighing raw against; RunConfig._normalize's warning about an explicit median under a squared-error metric is the same rule read the other way. Splitting it across those two would let a fleetwide judgement and a per-series selection disagree about what "corrected" means, which is the kind of drift nobody notices until a leaderboard reads oddly.

The answer is declared on the metric (BaseMetric.mean_optimal), not held here as a list of names. It used to be a frozenset of four, which was correct for the fifteen shipped metrics and unreachable for a sixteenth: a deployment that adds a squared-error metric would have been given the median, silently and wrongly, with nowhere to say otherwise. An unregistered name still answers "median" — the same answer the frozenset gave it — so this reads identically for every caller that was ever correct.

No longer what an unset output.point_forecast resolves to: with a backtest present the default is auto, which measures the choice per series rather than deducing it fleetwide. This function is what auto falls back to when a cell has nothing to measure.

Says nothing about whether a backtest exists — that is the caller's guard.

Source code in src/scale_forecasting/config.py

def corrected_arm_for(decision_metric: str) -> str:
    """Which *corrected* point-forecast arm a decision metric implies — `"mean"` or `"median"`.

    A metric whose loss is quadratic in the error is minimised by the **mean**; everything that is
    absolute-error-shaped, and every proper interval score, is minimised by the **median**. The
    pairing is a theorem, not a preference, so it is what a run falls back to whenever no
    measurement is available to do better.

    One rule, two callers. `calibration.select_arm` uses it under `point_forecast="auto"` to know
    which arm it is weighing `raw` against; `RunConfig._normalize`'s warning about an explicit
    `median` under a squared-error metric is the same rule read the other way. Splitting it across
    those two would let a fleetwide judgement and a per-series selection disagree about what
    "corrected" means, which is the kind of drift nobody notices until a leaderboard reads oddly.

    **The answer is declared on the metric** (`BaseMetric.mean_optimal`), not held here as a list of
    names. It used to be a frozenset of four, which was correct for the fifteen shipped metrics and
    unreachable for a sixteenth: a deployment that adds a squared-error metric would have been given
    the median, silently and wrongly, with nowhere to say otherwise. An unregistered name still
    answers `"median"` — the same answer the frozenset gave it — so this reads identically for every
    caller that was ever correct.

    No longer what an unset `output.point_forecast` resolves to: with a backtest present the
    default is `auto`, which measures the choice per series rather than deducing it fleetwide. This
    function is what `auto` falls back to when a cell has nothing to measure.

    Says nothing about whether a backtest exists — that is the caller's guard.
    """
    from .metrics.base_metric import _REGISTRY

    metric_cls = _REGISTRY.get(decision_metric)
    return "mean" if metric_cls is not None and metric_cls.mean_optimal else "median"

estimate_fanout

estimate_fanout(cfg)

The config-only cell-count estimate — estimate_workload narrowed to its four count fields.

When data.series_limit is unset, series count isn't known offline, so n_series and n_cells are None (the CLI reports "all series").

Source code in src/scale_forecasting/config.py

def estimate_fanout(cfg: RunConfig) -> Fanout:
    """The config-only cell-count estimate — `estimate_workload` narrowed to its four count fields.

    When ``data.series_limit`` is unset, series count isn't known offline, so
    ``n_series`` and ``n_cells`` are ``None`` (the CLI reports "all series").
    """
    w = estimate_workload(cfg)
    return Fanout(n_series=w.n_series, n_models=w.n_models, n_folds=w.n_folds, n_cells=w.n_cells)

estimate_workload

estimate_workload(cfg, *, obs_counts=None)

Estimate the work a run will schedule, in cells and in fits (pure).

obs_counts is the observation count of each series that will be forecast — one entry per series, in any order. Supply it (from SELECT ts_id, COUNT(*) … GROUP BY ts_id) and the fold geometry is resolved exactly, per series, through backtest.fit_rows. Omit it and only the cell counts come back; n_series then falls back to data.series_limit, which is None for an unbounded run, and n_cells follows it.

Why fits and not just cells. A cell with two backtest folds does three fits, and the two fold fits train on shorter windows than the final one — so the honest cost multiplier for backtesting is full_fit_equivalents (2.94 for four years of daily history, two folds, a 28-day horizon), not n_folds + 1 (3). full_fit_equivalents is dimensionless: it is what one cell costs relative to the same cell with backtesting off, so multiply it by n_cells to compare two plans. n_fits and train_rows_total are absolute and both scale with models.

With backtesting off, every cell does exactly one full-history fit, so n_fits == n_cells and the multiplier is 1.0 without needing to know a single series length.

Source code in src/scale_forecasting/config.py

def estimate_workload(cfg: RunConfig, *, obs_counts: Sequence[int] | None = None) -> Workload:
    """Estimate the work a run will schedule, in cells and in fits (pure).

    ``obs_counts`` is the observation count of each series that will be forecast — one entry per
    series, in any order. Supply it (from ``SELECT ts_id, COUNT(*) … GROUP BY ts_id``) and the fold
    geometry is resolved exactly, per series, through `backtest.fit_rows`. Omit it and only the
    cell counts come back; ``n_series`` then falls back to ``data.series_limit``, which is ``None``
    for an unbounded run, and ``n_cells`` follows it.

    **Why fits and not just cells.** A cell with two backtest folds does three fits, and the two
    fold fits train on shorter windows than the final one — so the honest cost multiplier for
    backtesting is ``full_fit_equivalents`` (2.94 for four years of daily history, two folds, a
    28-day horizon), not ``n_folds + 1`` (3). ``full_fit_equivalents`` is dimensionless: it is what
    one cell costs relative to the same cell with backtesting off, so multiply it by ``n_cells`` to
    compare two plans. ``n_fits`` and ``train_rows_total`` are absolute and both scale with models.

    With backtesting off, every cell does exactly one full-history fit, so ``n_fits == n_cells`` and
    the multiplier is ``1.0`` without needing to know a single series length.
    """
    from .backtest import fit_rows

    n_models = len(cfg.models)
    n_series = len(obs_counts) if obs_counts is not None else cfg.data.series_limit
    n_folds = cfg.backtest.n_folds if cfg.backtest.enabled else 1
    n_cells = None if n_series is None else n_series * n_models

    n_fits: int | None = None
    equivalents: float | None = None
    rows_total: int | None = None
    histogram: dict[int, int] = {}
    n_unscored: int | None = None

    if obs_counts is None:
        # The one geometry fact that needs no data: no backtest means one fit per cell.
        if not cfg.backtest.enabled:
            n_fits, equivalents = n_cells, 1.0
    else:
        per_series = [fit_rows(int(n), cfg) for n in obs_counts]
        achieved = [len(rows) - 1 for rows in per_series]
        rows_total = n_models * sum(sum(rows) for rows in per_series)
        n_fits = n_models * sum(len(rows) for rows in per_series)
        baseline = n_models * sum(int(n) for n in obs_counts)
        equivalents = rows_total / baseline if baseline else None
        histogram = {k: achieved.count(k) for k in sorted(set(achieved))}
        n_unscored = histogram.get(0, 0) if cfg.backtest.enabled else None

    return Workload(
        n_series=n_series,
        n_models=n_models,
        n_folds=n_folds,
        n_cells=n_cells,
        n_fits=n_fits,
        full_fit_equivalents=equivalents,
        train_rows_total=rows_total,
        fold_histogram=histogram,
        n_unscored=n_unscored,
    )

load_config

load_config(path)

Read a JSON config file and return a validated, frozen RunConfig.

All failure modes (missing file, bad JSON, invalid schema) surface as a single ConfigError with a clear message, so callers fail fast and never log an invalid run.

Source code in src/scale_forecasting/config.py

def load_config(path: str | Path) -> RunConfig:
    """Read a JSON config file and return a validated, frozen ``RunConfig``.

    All failure modes (missing file, bad JSON, invalid schema) surface as a single
    ``ConfigError`` with a clear message, so callers fail fast and never log an
    invalid run.
    """
    p = Path(path)
    try:
        raw = p.read_text()
    except OSError as e:
        raise ConfigError(f"cannot read config file '{p}': {e}") from e
    try:
        data = json.loads(raw)
    except json.JSONDecodeError as e:
        raise ConfigError(f"config file '{p}' is not valid JSON: {e}") from e
    try:
        return RunConfig(**data)
    except ValidationError as e:
        raise ConfigError(f"invalid config '{p}':\n{e}") from e

load_config_uri

load_config_uri(uri)

Load a validated config from a local path or a gs:// URI (the portable source).

A gs:// URI is the staged config an emitted launch command references, so any ADC-authenticated machine can re-run from it without a local file. Anything else is treated as a filesystem path (delegates to load_config). Failure modes surface as a single ConfigError.

Source code in src/scale_forecasting/config.py

def load_config_uri(uri: str) -> RunConfig:
    """Load a validated config from a local path **or** a ``gs://`` URI (the portable source).

    A ``gs://`` URI is the staged config an emitted launch command references, so any
    ADC-authenticated machine can re-run from it without a local file. Anything else is treated as a
    filesystem path (delegates to `load_config`). Failure modes surface as a single ``ConfigError``.
    """
    if not uri.startswith("gs://"):
        return load_config(uri)
    bucket, _, blob = uri[len("gs://") :].partition("/")
    if not bucket or not blob:
        raise ConfigError(f"malformed config URI '{uri}' (expected gs://bucket/path.json)")
    from google.cloud import storage

    try:
        raw = storage.Client().bucket(bucket).blob(blob).download_as_text()
    except Exception as e:  # surface any fetch failure as one ConfigError
        raise ConfigError(f"cannot read config URI '{uri}': {e}") from e
    try:
        data = json.loads(raw)
    except json.JSONDecodeError as e:
        raise ConfigError(f"config URI '{uri}' is not valid JSON: {e}") from e
    try:
        return RunConfig(**data)
    except ValidationError as e:
        raise ConfigError(f"invalid config '{uri}':\n{e}") from e

resolve_vertex_machine_type

resolve_vertex_machine_type(
    hardware,
    gpu_type,
    cpu_machine_type="auto",
    gpu_machine_type="auto",
    override_machine_type=None,
    *,
    accelerator_count=1,
)

Backward-compatible wrapper around resolve_vm_machine_type.

Source code in src/scale_forecasting/config.py

def resolve_vertex_machine_type(
    hardware: str,
    gpu_type: str | None,
    cpu_machine_type: str = "auto",
    gpu_machine_type: str = "auto",
    override_machine_type: str | None = None,
    *,
    accelerator_count: int = 1,
) -> str:
    """Backward-compatible wrapper around `resolve_vm_machine_type`."""
    default_mt = (
        gpu_machine_type if hardware == "gpu" and gpu_machine_type != "auto" else cpu_machine_type
    )
    return resolve_vm_machine_type(
        hardware,
        gpu_type,
        default_mt,
        override_machine_type,
        accelerator_count=accelerator_count,
    )

resolve_vm_machine_type

resolve_vm_machine_type(
    hardware,
    gpu_type=None,
    default_machine_type="auto",
    override_machine_type=None,
    *,
    accelerator_count=1,
)

Resolve and validate the GCE/Vertex VM machine type for a family's hardware (pure).

Enforces GCP's three structural GPU-to-VM compatibility patterns at config load time: * T4 (NVIDIA_TESLA_T4): attaches to n1-* in counts of {1, 2, 4} (1-2 GPUs cap at 48 vCPUs; default n1-standard-8 for 1-2 GPUs, n1-standard-16 for 4 GPUs). * L4 (NVIDIA_L4): bundled into g2-standard-* in counts of {1, 2, 4, 8} (1 GPU allows g2-standard-{4,8,12,16,32}, default g2-standard-8; 2/4/8 GPUs require fixed shapes g2-standard-{24,48,96}). * A100 (NVIDIA_TESLA_A100, 40 GiB) & A100_80GB (NVIDIA_A100_80GB, 80 GiB): strict 1:1 mapping from accelerator_count to a2-highgpu-{1,2,4,8}g / a2-megagpu-16g or a2-ultragpu-{1,2,4,8}g.

Source code in src/scale_forecasting/config.py

def resolve_vm_machine_type(
    hardware: str,
    gpu_type: str | None = None,
    default_machine_type: str = "auto",
    override_machine_type: str | None = None,
    *,
    accelerator_count: int = 1,
) -> str:
    """Resolve and validate the GCE/Vertex VM machine type for a family's hardware (pure).

    Enforces GCP's three structural GPU-to-VM compatibility patterns at config load time:
    * ``T4`` (``NVIDIA_TESLA_T4``): attaches to ``n1-*`` in counts of ``{1, 2, 4}`` (1-2 GPUs cap
      at 48 vCPUs; default ``n1-standard-8`` for 1-2 GPUs, ``n1-standard-16`` for 4 GPUs).
    * ``L4`` (``NVIDIA_L4``): bundled into ``g2-standard-*`` in counts of ``{1, 2, 4, 8}`` (1 GPU
      allows ``g2-standard-{4,8,12,16,32}``, default ``g2-standard-8``; 2/4/8 GPUs require fixed
      shapes ``g2-standard-{24,48,96}``).
    * ``A100`` (``NVIDIA_TESLA_A100``, 40 GiB) & ``A100_80GB`` (``NVIDIA_A100_80GB``, 80 GiB):
      strict 1:1 mapping from ``accelerator_count`` to ``a2-highgpu-{1,2,4,8}g`` /
      ``a2-megagpu-16g`` or ``a2-ultragpu-{1,2,4,8}g``.
    """
    from .resources.catalog import machine_cores

    if hardware == "gpu":
        auto_gpu = _default_gpu_machine_type(gpu_type, accelerator_count)
        if override_machine_type is not None and override_machine_type != "auto":
            chosen = override_machine_type
        elif default_machine_type in ("auto", "n2-standard-8"):
            chosen = auto_gpu
        else:
            chosen = default_machine_type

        eff_gpu = gpu_type or "T4"
        if eff_gpu == "T4":
            if not chosen.startswith("n1-"):
                raise ValueError(
                    f"Vertex AI / GCE T4 GPUs require an n1-* machine type (got {chosen!r})"
                )
            cores = machine_cores(chosen) or 0
            if accelerator_count in (1, 2) and cores > 48:
                raise ValueError(
                    f"T4 with accelerator_count={accelerator_count} allows at most 48 vCPUs "
                    f"on n1-* (got {chosen!r} with {cores} vCPUs)"
                )
        elif eff_gpu == "L4":
            if not chosen.startswith("g2-"):
                raise ValueError(
                    f"Vertex AI / GCE L4 GPUs require a g2-* machine type (got {chosen!r})"
                )
            if accelerator_count == 1 and chosen not in _L4_SINGLE_GPU_MACHINES:
                raise ValueError(
                    f"1x L4 GPU requires one of {sorted(_L4_SINGLE_GPU_MACHINES)} (got {chosen!r})"
                )
            if accelerator_count > 1 and chosen != _L4_MULTI_GPU_MACHINES.get(accelerator_count):
                expected = _L4_MULTI_GPU_MACHINES[accelerator_count]
                raise ValueError(
                    f"{accelerator_count}x L4 GPUs require machine_type={expected!r} "
                    f"(got {chosen!r})"
                )
        elif eff_gpu == "A100":
            expected = _A100_MACHINES[accelerator_count]
            if chosen != expected:
                raise ValueError(
                    f"{accelerator_count}x A100 GPUs require machine_type={expected!r} "
                    f"(got {chosen!r})"
                )
        elif eff_gpu == "A100_80GB":
            expected = _A100_80GB_MACHINES[accelerator_count]
            if chosen != expected:
                raise ValueError(
                    f"{accelerator_count}x A100_80GB GPUs require machine_type={expected!r} "
                    f"(got {chosen!r})"
                )
        return chosen

    if override_machine_type is not None and override_machine_type != "auto":
        chosen = override_machine_type
    elif default_machine_type == "auto" or default_machine_type.startswith(("g2-", "a2-")):
        chosen = "n2-standard-8"
    else:
        chosen = default_machine_type

    if chosen.startswith(("g2-", "a2-")):
        raise ValueError(
            f"CPU workloads cannot use GPU-attached machine family {chosen!r}; "
            "set hardware='gpu' or choose a CPU machine family (e.g. 'n2-standard-8')"
        )
    return chosen