Hyperparameter Optimization (HPO)

hpo

In-node hyperparameter optimization — an Optuna study over the aligned backtest.

Gated on cfg.hpo.enabled; a strict no-op otherwise (a model with no search space, or HPO off, tunes nothing and yields {} — exactly today's behavior). Two granularities, the DS-facing knob in HpoConfig:

• fleetwide (default): tune each model once on a representative sample of series and apply the winning params across all series — the only granularity affordable at 100k. Resolved on the driver (resolve_fleetwide) before the engine fans out, then threaded to run_cell as pre-resolved params (never via cfg — the config is the run_id identity key, so putting tuned params in it would shift the run_id and break reproducibility/idempotency; see make_run_id).
• per_series: tune on each series inside run_cell (heavier; a DS opt-in for the tail of hard series).

The objective reuses the aligned backtest: for one trial's params it runs backtest_cell on each sampled series, averages the per-fold panels, and scores on cfg.backtest.decision_metric. HPO therefore requires backtesting (require_backtest) — there are no folds to tune on otherwise.

Pure + offline: no GCP, no Spark, deterministic (fixed-seed TPE sampler). The engines only add a tiny driver-side sample-and-resolve call in front of their existing fan-out.

require_backtest

require_backtest(cfg)

Raise if HPO is enabled without backtesting — the study has no folds to score on.

Source code in src/scale_forecasting/hpo.py

def require_backtest(cfg: RunConfig) -> None:
    """Raise if HPO is enabled without backtesting — the study has no folds to score on."""
    if cfg.hpo.enabled and not cfg.backtest.enabled:
        raise ConfigError(
            "hpo.enabled requires backtest.enabled: HPO tunes on the backtest folds "
            "(decision_metric), so there is nothing to optimize with backtesting off."
        )

resolve_fleetwide

resolve_fleetwide(sample, cfg, ctx=None)

Tune every model in cfg.models once on the shared sample → {model: params}.

The driver-side fleetwide pre-pass (the opinionated default): the winning params for each model are applied across all series in the run. Models with no search space are simply absent from the mapping (they keep their {} defaults in run_cell). Pure — the caller supplies the pandas sample (a small, driver-collected set of series); this function does no I/O.

Source code in src/scale_forecasting/hpo.py

def resolve_fleetwide(
    sample: list[pd.DataFrame], cfg: RunConfig, ctx: ModelContext | None = None
) -> dict[str, dict[str, Any]]:
    """Tune every model in ``cfg.models`` once on the shared sample → ``{model: params}``.

    The driver-side fleetwide pre-pass (the opinionated default): the winning params for each model
    are applied across *all* series in the run. Models with no search space are simply absent from
    the mapping (they keep their ``{}`` defaults in ``run_cell``). Pure — the caller supplies the
    pandas ``sample`` (a small, driver-collected set of series); this function does no I/O.
    """
    require_backtest(cfg)
    ctx = ctx if ctx is not None else _context(cfg)
    resolved: dict[str, dict[str, Any]] = {}
    for name in cfg.models:
        params = tune_model(name, sample, cfg, ctx)
        if params:
            resolved[name] = params
    return resolved

tune_model

tune_model(model_name, sample, cfg, ctx=None, tally=None)

Tune one model on sample and return its winning params ({} if nothing to tune).

Builds a deterministic Optuna study (fixed-seed TPE) of cfg.hpo.n_trials trials whose objective is _score_params. Returns {} immediately — creating no study — when the model has no search space (_has_search_space), so an all-defaults model costs nothing.

cfg.model_params[model_name] sits underneath every trial. A model tuned without the hyperparameters its author pinned is a different model from the one that will be fitted, so the study would optimize the wrong thing — an authored n_lags changes what learning_rate is best for. Where the two name the same key the trial wins, and the winner is returned with the authored layer still beneath it, so the returned dict is exactly the params the cell will build with and exactly what lands in forecast_metadata.best_params.

A model with no search space returns {} rather than the authored params: nothing was tuned, and worker._resolve_params applies the authored layer at the cell either way.

tally counts the fits the search paid for, and only the per_series granularity passes one. A search costs n_trials × sample × folds fits and none of them produce a shipped forecast, so they are counted separately from the cell's own — worker.run_cell keeps two tallies and writes them to n_hpo_fits and n_fits. The fleetwide pre-pass passes nothing: it runs on the driver before any cell exists, so its fits belong to the run rather than to any one row, and there is nowhere honest to put them.

Source code in src/scale_forecasting/hpo.py

def tune_model(
    model_name: str,
    sample: list[pd.DataFrame],
    cfg: RunConfig,
    ctx: ModelContext | None = None,
    tally: FitTally | None = None,
) -> dict[str, Any]:
    """Tune one model on ``sample`` and return its winning params (``{}`` if nothing to tune).

    Builds a deterministic Optuna study (fixed-seed TPE) of ``cfg.hpo.n_trials`` trials whose
    objective is `_score_params`. Returns ``{}`` immediately — creating no study — when the
    model has no search space (`_has_search_space`), so an all-defaults model costs nothing.

    **``cfg.model_params[model_name]`` sits underneath every trial.** A model tuned without the
    hyperparameters its author pinned is a different model from the one that will be fitted, so the
    study would optimize the wrong thing — an authored ``n_lags`` changes what ``learning_rate`` is
    best for. Where the two name the same key the trial wins, and the winner is returned with the
    authored layer still beneath it, so the returned dict is exactly the params the cell will build
    with and exactly what lands in ``forecast_metadata.best_params``.

    A model with no search space returns ``{}`` rather than the authored params: nothing was tuned,
    and `worker._resolve_params` applies the authored layer at the cell either way.

    ``tally`` counts the fits the *search* paid for, and only the ``per_series`` granularity passes
    one. A search costs ``n_trials × sample × folds`` fits and none of them produce a shipped
    forecast, so they are counted separately from the cell's own — `worker.run_cell` keeps two
    tallies and writes them to ``n_hpo_fits`` and ``n_fits``. The fleetwide pre-pass passes nothing:
    it runs on the driver before any cell exists, so its fits belong to the run rather than to any
    one row, and there is nowhere honest to put them.
    """
    from .features import clean_model_params, resolve_model_variant

    model_cls = get_model(model_name)
    spec = resolve_model_variant(cfg, model_name, model_cls=model_cls)
    authored: dict[str, Any] = dict(spec.clean_params)
    variant_pins: dict[str, Any] = (
        dict(clean_model_params(cfg.model_params[model_name]))
        if "@" in model_name and model_name in cfg.model_params
        else {}
    )
    if not _has_search_space(model_cls):
        return {}

    import optuna

    optuna.logging.set_verbosity(optuna.logging.WARNING)
    ctx = ctx if ctx is not None else _context(cfg)
    metric = cfg.backtest.decision_metric

    sampler = optuna.samplers.TPESampler(seed=ctx.seed)
    study = optuna.create_study(direction="minimize", sampler=sampler)

    def objective(trial: optuna.Trial) -> float:
        return _score_params(
            model_name,
            {
                "training_mode": spec.training_mode,
                **authored,
                **model_cls.search_space(trial),
                **variant_pins,
            },
            sample,
            cfg,
            ctx,
            tally,
        )

    study.optimize(objective, n_trials=cfg.hpo.n_trials)
    _log.info(
        "hpo %s: best=%s (%s→%.4g over %d trials)",
        model_name,
        study.best_params,
        metric,
        study.best_value,
        cfg.hpo.n_trials,
    )
    return {**authored, **study.best_params, **variant_pins}