Adding a model

A model is one file in src/scale_forecasting/models/ that ends with a register(...) call, plus one import line in models/__init__.py. Nothing else changes — not the worker, not the engines, not the registry. That is the whole point of the factory: the system discovers models by name at import time.

flowchart LR
    T["1. Copy docs/model_template.py<br/>→ models/my_model.py"] --> I["2. Implement BaseModel<br/>fit(y, X) & predict(horizon, X, quantiles)"]
    I --> R["3. Call register(MyModel)<br/>+ import in models/__init__.py"]
    R --> W["worker.run_cell<br/>(Spark / Ray / Vertex / GCE / GKE / Playground)"]
    W --> C["calibration.apply_calibration<br/>(Point & Interval Calibration)"]
    C --> BQ["Registry Tables & Leaderboard"]

The 4-step checklist

1. Copy the template. docs/model_template.py → src/scale_forecasting/models/my_model.py.
2. Name it. Rename the class and set name = "my_model" (unique, lowercase, snake_case). This string is how users select the model in a run config and in the playground.
3. Fill in fit and predict. See the contract below.
4. Register it. Add my_model, to the import block in src/scale_forecasting/models/__init__.py. (The file already ends with register(...); this import is what triggers it.)

Then try it immediately, offline:

python -m scale_forecasting.playground --list           # my_model now appears
python -m scale_forecasting.playground --model my_model --backtest

or open notebooks/00_model_playground.ipynb (or notebooks/09_custom_models_and_metrics.ipynb for a full custom plugin walkthrough), set MODEL = "my_model", and run the cells.

The contract every model owes

A model is a BaseModel subclass (see models/base_model.py). The seams:

Piece	What it is
name	unique selector string
runtime	"python" (runs in a Spark/Ray/Vertex/GCE/GKE cell), "vertex_automl", or "bigquery" (SQL)
family	"statistical" | "ml" | "deep_learning" | "automl" | "native"
supports_exog	True if fit/predict use the X frame
supports_explainability	True if the model populates feature_attributions and explanations
supports_native_intervals	True if you produce your own prediction bounds
fit(y, X)	fit on one series; y is indexed by ds, already transformed per config
predict(horizon, X, quantiles)	return the canonical frame in original units

fit(y, X) — y is one series' target (a pd.Series indexed by datetime ds), already transformed per the run config (e.g. log1p). X is the aligned feature/exog frame, or None. Stash what predict needs on self. Per-run context is on self.ctx (freq, horizon, seed, transform, transform_lambda); HPO hyperparameters are on self.params. Raise ModelError on a genuinely unfittable series — the worker captures it into an error cell and the batch survives; it never propagates.

predict(horizon, X, quantiles) — return the canonical prediction frame: columns ds, yhat, yhat_lower, yhat_upper, quantiles, in original units, with ordered bounds (lower ≤ yhat ≤ upper). The base class does the assembly for you:

• build a point forecast array,
• turn it into a {quantile: array} map — either from your own native bounds, or via self.residual_intervals(mean, quantiles) if you recorded residuals in fit with self._set_residuals(...),
• invert the transform with invert_transform(values, self.ctx.transform, self.ctx.transform_lambda) so values are in original units (the λ is fit once per cell by the worker and passed on ctx; it's None for none/log1p),
• call self._assemble_frame(ds, qmap) (use self._future_index(last_date, horizon) for the dates).

What the system does with your frame afterwards

Your yhat is the raw arm. Downstream, calibration.apply_calibration may ship a corrected number in its place — yhat plus a residual shift, per output.point_forecast — writing both arms to yhat_raw and yhat_adjusted so the choice is recorded and reversible. Two consequences for a model author:

• yhat should be your model's honest central estimate, not a pre-corrected one. If you shift it yourself, the system's correction stacks on top of yours and neither is visible in the row. Where the model library you wrap offers a choice, return the untouched forecast.
• Your band may be replaced. When a backtest runs, yhat_lower / yhat_upper / quantiles are re-estimated per horizon step from out-of-fold residuals, whatever your model produced. forecast_metadata.interval_source still records that your model computed its own bounds (supports_native_intervals), and interval_calibration records that they were re-estimated; they are separate columns because they are separate facts. With no backtest, your band ships untouched.

Nothing here changes what predict must return. It is worth knowing because it is why the honest version of both numbers is the useful one: the system can always correct a clean estimate, and can never recover one from an estimate that was already adjusted.

Rules that keep the product coherent

• One model, one file. No model imports another model. Shared machinery goes in a helper module (see _lag_forecaster.py, used by both tree models).
• Never read global config. Everything a model needs arrives via self.ctx and self.params. This is what lets the identical code run local, on Spark, and on Ray.
• Never hard-code a seasonal period. Use seasonality.seasonal_period(self.ctx.freq) — the one shared freq→period source. (pandas 3 spellings: D, W, MS, ME, h.)
• Original units out. The frame predict returns must be in original units; invert any transform yourself.
• Optional heavy deps import lazily inside fit. If your model needs an extra (torch, xgboost, …), import it inside fit and raise ModelError("… install the 'models' extra") on ImportError. That way the model still registers without the dep, and the contract test skips it cleanly instead of breaking the suite. See xgboost_model.py.
• Avoid shadowing the library you wrap. When your model's name collides with the package it imports, keep them distinct: either suffix the file _model.py (prophet_model.py imports prophet) or alias the import (theta.py does from statsmodels...theta import ThetaModel as _StatsmodelsTheta). Models with no such collision use the bare name (sarimax.py). The registered name string is unaffected either way.

What you get for free

• python -m scale_forecasting.playground and the notebook list and run your model with no changes — they read the factory registry.
• The contract test (tests/unit/test_models_contract.py) is parametrized over every registered model, so your model is automatically checked for: canonical frame shape, ordered bounds, original-units output, determinism under a fixed seed, and supports_exog honored. Run pytest tests/unit/test_models_contract.py -k my_model.
• HPO (optional): implement search_space(cls, trial) to expose an Optuna search space; it's used only when hpo.enabled in the config.
• Two-tier explainability (optional): set supports_explainability = True, implement feature_attributions (returning a normalized dict[str, float] of Tier 1 global/series-level driver importances that the worker merges into forecast_metadata.fit_diagnostics["feature_attributions"]), and attach an "explanations" column on the frame returned by predict() containing per-horizon-step dicts {"baseline_score": float, "attributions": {feature: float}} (Tier 2 local step-by-step attributions streamed to forecast_predictions.explanations). Both tiers are immediately queryable and visualizable via Forecaster.attributions_df() and Forecaster.plot_attributions() (or review.build_attributions_frame() offline).
• Fit diagnostics (optional): implement diagnostics() to return whatever your library says about the fit — an AIC, the order an auto_arima picked, the epoch an early stop landed on. It is called once per cell after the final fit and lands in forecast_metadata.fit_diagnostics, a JSON bag beside best_params.

Do not put a scored metric here. A metric is a pure function of (y_true, yhat, y_train, bounds) that the framework computes, so it means the same thing across every model and a leaderboard can rank on it; a diagnostic is whatever your library happened to expose, exists for some models and not others, and two models' "AIC" are not on the same scale. A key named after a metric is dropped with a warning rather than written. So is a value json.dumps cannot carry — cast numpy scalars to float/int before returning them. Anything the method raises is caught and the cell simply reports no diagnostics: describing a fit must never cost the forecast.