Skip to content

XGBoost Forecaster API

XGBoost Substation Forecaster

This package implements an XGBoost-based model to forecast power flows at NGED primary substations using numerical weather prediction (NWP) forecasts. It implements the BaseForecaster interface defined in ml_core.

How it works

One xgb.Booster is trained per time_series_id, so each substation's model can learn its own relationship between weather and power. Features are passed via the AllFeatures schema (see contracts), which joins NWP variables, power lag/rolling features, and static metadata. Categorical and string columns are encoded as integer codes before being handed to XGBoost; all features are cast to Float32, and missing values are left as NaN so XGBoost handles them natively. The model is deterministic; the ensemble spread comes from each NWP ensemble member's weather flowing through the relevant booster — predict() scores every ensemble member present in its input in one call.

train() collects its input once and groups in memory by time_series_id, feeding each series to an xgb.QuantileDMatrix (8-bit quantile bins, not an uncompressed Float32 copy). Keeping that collect bounded is the caller's job — the dominant cost is the multi-tens-of-GB NWP scan, which must be pruned at the inputs (control member, the relevant H3 cells, the window's init_time partitions) and streamed; filtering the engineered output cannot prune the upstream join/upsample. predict() likewise collects once and groups by time_series_id; at validation the full ~51-member ensemble is too large to collect whole, so the caller (cv_power_forecasts) predicts one init_time chunk at a time, writing to Delta incrementally. See Bounding feature-engineering memory for the dataset sizes and the table of which predicates actually prune the NWP scan.

Save format

XGBoostForecaster.save(path) writes:

  • {time_series_id}.ubj — one XGBoost native binary model per trained substation
  • meta.json — the full XGBoostConfig serialised via Pydantic, so load() is completely self-contained

Configuration

XGBoostConfig extends BaseForecasterConfig with XGBoost hyperparameters (n_estimators, learning_rate, max_depth, etc.) and the model identity fields (power_fcst_model_name, power_fcst_model_version, ml_flow_experiment_id). Identity fields are stamped onto every row of the PowerForecast output so the Delta Lake table is self-describing.

xgboost_forecaster.forecaster

XGBoost-based power forecasting model.

Classes

XGBoostConfig

Bases: BaseForecasterConfig

Configuration for XGBoostForecaster.

Inherits the universal fields (selected_features, MLflow experiment id) from BaseForecasterConfig and adds XGBoost-specific hyperparameters.

Source code in packages/xgboost_forecaster/src/xgboost_forecaster/forecaster.py 
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
class XGBoostConfig(BaseForecasterConfig):
    """Configuration for XGBoostForecaster.

    Inherits the universal fields (selected_features, MLflow experiment id) from
    BaseForecasterConfig and adds XGBoost-specific hyperparameters.
    """

    n_estimators: int = 1000
    learning_rate: float = 0.05
    max_depth: int = 6
    min_child_weight: int = 1
    subsample: float = 0.8
    colsample_bytree: float = 0.8
    device: str = "cpu"
    objective: str = "reg:squarederror"

    def to_xgb_params(self) -> dict[str, Any]:
        """Return the params dict accepted by xgb.train()."""
        return {
            "eta": self.learning_rate,
            "max_depth": self.max_depth,
            "min_child_weight": self.min_child_weight,
            "subsample": self.subsample,
            "colsample_bytree": self.colsample_bytree,
            "device": self.device,
            "objective": self.objective,
            "seed": self.random_seed,
        }
Attributes
n_estimators = 1000 class-attribute instance-attribute
learning_rate = 0.05 class-attribute instance-attribute
max_depth = 6 class-attribute instance-attribute
min_child_weight = 1 class-attribute instance-attribute
subsample = 0.8 class-attribute instance-attribute
colsample_bytree = 0.8 class-attribute instance-attribute
device = 'cpu' class-attribute instance-attribute
objective = 'reg:squarederror' class-attribute instance-attribute
selected_features instance-attribute
ml_flow_experiment_id = None class-attribute instance-attribute
experiment_name = '' class-attribute instance-attribute
weather_source = '' class-attribute instance-attribute
training_strategy = '' class-attribute instance-attribute
random_seed = 0 class-attribute instance-attribute
Methods:
to_xgb_params()

Return the params dict accepted by xgb.train().

Source code in packages/xgboost_forecaster/src/xgboost_forecaster/forecaster.py 
50
51
52
53
54
55
56
57
58
59
60
61
def to_xgb_params(self) -> dict[str, Any]:
    """Return the params dict accepted by xgb.train()."""
    return {
        "eta": self.learning_rate,
        "max_depth": self.max_depth,
        "min_child_weight": self.min_child_weight,
        "subsample": self.subsample,
        "colsample_bytree": self.colsample_bytree,
        "device": self.device,
        "objective": self.objective,
        "seed": self.random_seed,
    }

XGBoostForecaster

Bases: BaseForecaster

Trains and serves one XGBoost Booster per time_series_id.

All lead times for a given time_series_id are handled by a single Booster. The model is deterministic; ensemble forecasts arise because each NWP ensemble member's weather is a separate row through the relevant Booster. predict scores every ensemble member present in its input in one call (it groups by time_series_id and dispatches each group to its Booster).

Save layout: a directory containing one {time_series_id}.ubj file per trained Booster plus a meta.json that stores the full XGBoostConfig so that load() is self-contained.

Source code in packages/xgboost_forecaster/src/xgboost_forecaster/forecaster.py 
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
class XGBoostForecaster(BaseForecaster):
    """Trains and serves one XGBoost Booster per time_series_id.

    All lead times for a given time_series_id are handled by a single Booster. The model is
    deterministic; ensemble forecasts arise because each NWP ensemble member's weather is a separate
    row through the relevant Booster. ``predict`` scores every ensemble member present in its input
    in one call (it groups by ``time_series_id`` and dispatches each group to its Booster).

    Save layout: a directory containing one ``{time_series_id}.ubj`` file per trained
    Booster plus a ``meta.json`` that stores the full XGBoostConfig so that load() is
    self-contained.
    """

    MODEL_NAME = "xgboost"
    MODEL_VERSION = 1

    model_params: XGBoostConfig  # narrows the base class annotation for type checkers

    def __init__(self, model_params: XGBoostConfig) -> None:
        super().__init__(model_params)
        self._models: dict[int, xgb.Booster] = {}

    @property
    def _feature_cols(self) -> list[str]:
        return sorted(self.model_params.selected_features)

    @property
    def trained_time_series_ids(self) -> list[int]:
        """The sorted ``time_series_id``s this forecaster will serve a ``predict`` for.

        For ``XGBoostForecaster`` this is exactly the set of series it holds a trained Booster for
        (one Booster per ``time_series_id``), so ``predict`` raises ``KeyError`` if asked for any
        other series. ``save()`` records this set in ``meta.json`` and ``load()`` reconstructs it
        from the ``.ubj`` files on disk. See ``BaseForecaster.trained_time_series_ids`` for the
        model-agnostic contract this implements (the train==predict population invariant).
        """
        return sorted(self._models.keys())

    def train(self, data: pt.LazyFrame[AllFeatures], time_series_ids: list[int]) -> None:
        """Fit one Booster per ``time_series_id`` in ``time_series_ids``.

        ``data`` is collected once and grouped in memory by ``time_series_id``; each group's rows
        feed an ``xgb.QuantileDMatrix`` (compressed to 8-bit quantile bins, not an uncompressed
        Float32 copy). Keeping this bounded is the *caller's* job — the NWP scan must be pruned at
        the inputs (control member, the relevant H3 cells, the window's ``init_time`` partitions),
        because filtering the engineered output cannot prune the upstream join/upsample. See
        ``_load_engineering_inputs`` and the "NWP scan pruning" notes in
        ``docs/architecture/overview.md``.

        Only the requested ``time_series_ids`` are trained; a requested series with no non-null
        ``power`` rows (e.g. none in the training window) simply does not appear and gets no Booster.
        """
        feature_cols = self._feature_cols
        requested = set(time_series_ids)
        # Stream the collect. init_time prunes whole NWP partitions, and the member-early sort
        # (delta_store.nwp.NWP_SORT_COLS) lets row-group stats skip most of each partition for the
        # control-member read — but h3_index is not a sort-early column, so cell filtering is still
        # decode-then-filter within the surviving row groups. The streaming engine applies those
        # predicates per morsel, holding peak memory to a few GB where the in-memory engine would
        # materialise every surviving row first. See docs/architecture/overview.md.
        df = data.drop_nulls(subset=["power"]).collect(engine="streaming")
        for group_key, group in df.group_by(["time_series_id"]):
            ts_id = int(group_key[0])
            if ts_id not in requested:
                continue
            group = cast(pt.DataFrame[AllFeatures], group)
            features = _prepare_features(group, feature_cols)
            label = group["power"].cast(pl.Float32)
            dtrain = xgb.QuantileDMatrix(features, label=label)
            booster = xgb.train(
                self.model_params.to_xgb_params(),
                dtrain,
                num_boost_round=self.model_params.n_estimators,
            )
            self._models[ts_id] = booster

    def predict(
        self, data: pt.LazyFrame[AllFeatures], *, fold_id: str = "live"
    ) -> pt.DataFrame[PowerForecast]:
        """Generate one power_fcst per row, dispatching by time_series_id to the right Booster.

        ``data`` is collected once and grouped in memory by ``time_series_id``. Rows for a
        ``time_series_id`` this model was not trained on are ignored (the model only scores its own
        trained population — see ``trained_time_series_ids``). Keeping the collect bounded is the
        caller's job: at validation every NWP ensemble member is present, so the caller engineers one
        H3 cell at a time (see ``cv_power_forecasts`` and ``docs/architecture/overview.md``).

        ``fold_id`` is stamped onto every output row (the model has no inherent fold; the caller
        supplies it). Defaults to the ``"live"`` production sentinel.
        """
        feature_cols = self._feature_cols
        cfg = self.model_params

        def _build_part(
            group_df: pt.DataFrame[AllFeatures], predictions: np.ndarray
        ) -> pl.DataFrame:
            return group_df.select(
                [
                    "valid_time",
                    "time_series_id",
                    "ensemble_member",
                    "nwp_init_time",
                    "power_fcst_init_time",
                ]
            ).with_columns(
                pl.col("ensemble_member").cast(pl.Int8),
                pl.Series("power_fcst", predictions, dtype=pl.Float32),
                power_fcst_model_name=pl.lit(self.MODEL_NAME),
                power_fcst_model_version=pl.lit(self.MODEL_VERSION, dtype=pl.Int16),
                ml_flow_experiment_id=pl.lit(cfg.ml_flow_experiment_id, dtype=pl.Int32),
                experiment_name=pl.lit(cfg.experiment_name),
                fold_id=pl.lit(fold_id),
            )

        df = data.collect(engine="streaming")  # stream the NWP scan — see train() / overview.md
        parts: list[pl.DataFrame] = []
        for group_key, group in df.group_by(["time_series_id"]):
            booster = self._models.get(int(group_key[0]))
            if booster is None:
                continue  # ignore series this model was not trained on
            group = cast(pt.DataFrame[AllFeatures], group)
            X = _prepare_features(group, feature_cols)
            predictions: np.ndarray = booster.predict(xgb.DMatrix(X))
            parts.append(_build_part(group, predictions))

        if not parts:
            empty = cast(pt.DataFrame[AllFeatures], df.head(0))
            parts.append(_build_part(empty, np.empty(0, dtype=np.float32)))

        # The identity columns (power_fcst_model_name / experiment_name / fold_id) are built as
        # ``pl.lit(str)`` above, so they are already ``String`` — the dtype PowerForecast declares —
        # and need no cast before validation.
        return PowerForecast.validate(pl.concat(parts))

    def save(self, path: Path) -> None:
        """Save all Boosters as .ubj files plus a meta.json with the full config."""
        path.mkdir(parents=True, exist_ok=True)
        for ts_id, booster in self._models.items():
            booster.save_model(str(path / f"{ts_id}.ubj"))
        (path / "meta.json").write_text(
            json.dumps(
                {
                    "model_params": self.model_params.model_dump(mode="json"),
                    "trained_time_series_ids": self.trained_time_series_ids,
                    "model_class": f"{type(self).__module__}.{type(self).__qualname__}",
                }
            )
        )

    @classmethod
    def load(cls, path: Path) -> Self:
        """Reconstruct an XGBoostForecaster from a saved directory."""
        meta = json.loads((path / "meta.json").read_text())
        config = XGBoostConfig.model_validate(meta["model_params"])
        instance = cls(config)
        for ubj_file in sorted(path.glob("*.ubj")):
            booster = xgb.Booster()
            booster.load_model(str(ubj_file))
            instance._models[int(ubj_file.stem)] = booster
        return instance
Attributes
MODEL_NAME = 'xgboost' class-attribute instance-attribute
MODEL_VERSION = 1 class-attribute instance-attribute
model_params instance-attribute
trained_time_series_ids property

The sorted time_series_ids this forecaster will serve a predict for.

For XGBoostForecaster this is exactly the set of series it holds a trained Booster for (one Booster per time_series_id), so predict raises KeyError if asked for any other series. save() records this set in meta.json and load() reconstructs it from the .ubj files on disk. See BaseForecaster.trained_time_series_ids for the model-agnostic contract this implements (the train==predict population invariant).

feature_engineer = TabularFeatureEngineer() class-attribute

The feature pipeline this forecaster's data is engineered through.

Associated by composition (the forecaster references a feature engineer rather than implementing feature engineering), so a forecaster can swap the whole pipeline by overriding this with a different FeatureEngineer. The default produces the tabular AllFeatures frame that train/predict consume.

Methods:
__init__(model_params)
Source code in packages/xgboost_forecaster/src/xgboost_forecaster/forecaster.py 
82
83
84
def __init__(self, model_params: XGBoostConfig) -> None:
    super().__init__(model_params)
    self._models: dict[int, xgb.Booster] = {}
train(data, time_series_ids)

Fit one Booster per time_series_id in time_series_ids.

data is collected once and grouped in memory by time_series_id; each group's rows feed an xgb.QuantileDMatrix (compressed to 8-bit quantile bins, not an uncompressed Float32 copy). Keeping this bounded is the caller's job — the NWP scan must be pruned at the inputs (control member, the relevant H3 cells, the window's init_time partitions), because filtering the engineered output cannot prune the upstream join/upsample. See _load_engineering_inputs and the "NWP scan pruning" notes in docs/architecture/overview.md.

Only the requested time_series_ids are trained; a requested series with no non-null power rows (e.g. none in the training window) simply does not appear and gets no Booster.

Source code in packages/xgboost_forecaster/src/xgboost_forecaster/forecaster.py 
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
def train(self, data: pt.LazyFrame[AllFeatures], time_series_ids: list[int]) -> None:
    """Fit one Booster per ``time_series_id`` in ``time_series_ids``.

    ``data`` is collected once and grouped in memory by ``time_series_id``; each group's rows
    feed an ``xgb.QuantileDMatrix`` (compressed to 8-bit quantile bins, not an uncompressed
    Float32 copy). Keeping this bounded is the *caller's* job — the NWP scan must be pruned at
    the inputs (control member, the relevant H3 cells, the window's ``init_time`` partitions),
    because filtering the engineered output cannot prune the upstream join/upsample. See
    ``_load_engineering_inputs`` and the "NWP scan pruning" notes in
    ``docs/architecture/overview.md``.

    Only the requested ``time_series_ids`` are trained; a requested series with no non-null
    ``power`` rows (e.g. none in the training window) simply does not appear and gets no Booster.
    """
    feature_cols = self._feature_cols
    requested = set(time_series_ids)
    # Stream the collect. init_time prunes whole NWP partitions, and the member-early sort
    # (delta_store.nwp.NWP_SORT_COLS) lets row-group stats skip most of each partition for the
    # control-member read — but h3_index is not a sort-early column, so cell filtering is still
    # decode-then-filter within the surviving row groups. The streaming engine applies those
    # predicates per morsel, holding peak memory to a few GB where the in-memory engine would
    # materialise every surviving row first. See docs/architecture/overview.md.
    df = data.drop_nulls(subset=["power"]).collect(engine="streaming")
    for group_key, group in df.group_by(["time_series_id"]):
        ts_id = int(group_key[0])
        if ts_id not in requested:
            continue
        group = cast(pt.DataFrame[AllFeatures], group)
        features = _prepare_features(group, feature_cols)
        label = group["power"].cast(pl.Float32)
        dtrain = xgb.QuantileDMatrix(features, label=label)
        booster = xgb.train(
            self.model_params.to_xgb_params(),
            dtrain,
            num_boost_round=self.model_params.n_estimators,
        )
        self._models[ts_id] = booster
predict(data, *, fold_id='live')

Generate one power_fcst per row, dispatching by time_series_id to the right Booster.

data is collected once and grouped in memory by time_series_id. Rows for a time_series_id this model was not trained on are ignored (the model only scores its own trained population — see trained_time_series_ids). Keeping the collect bounded is the caller's job: at validation every NWP ensemble member is present, so the caller engineers one H3 cell at a time (see cv_power_forecasts and docs/architecture/overview.md).

fold_id is stamped onto every output row (the model has no inherent fold; the caller supplies it). Defaults to the "live" production sentinel.

Source code in packages/xgboost_forecaster/src/xgboost_forecaster/forecaster.py 
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
def predict(
    self, data: pt.LazyFrame[AllFeatures], *, fold_id: str = "live"
) -> pt.DataFrame[PowerForecast]:
    """Generate one power_fcst per row, dispatching by time_series_id to the right Booster.

    ``data`` is collected once and grouped in memory by ``time_series_id``. Rows for a
    ``time_series_id`` this model was not trained on are ignored (the model only scores its own
    trained population — see ``trained_time_series_ids``). Keeping the collect bounded is the
    caller's job: at validation every NWP ensemble member is present, so the caller engineers one
    H3 cell at a time (see ``cv_power_forecasts`` and ``docs/architecture/overview.md``).

    ``fold_id`` is stamped onto every output row (the model has no inherent fold; the caller
    supplies it). Defaults to the ``"live"`` production sentinel.
    """
    feature_cols = self._feature_cols
    cfg = self.model_params

    def _build_part(
        group_df: pt.DataFrame[AllFeatures], predictions: np.ndarray
    ) -> pl.DataFrame:
        return group_df.select(
            [
                "valid_time",
                "time_series_id",
                "ensemble_member",
                "nwp_init_time",
                "power_fcst_init_time",
            ]
        ).with_columns(
            pl.col("ensemble_member").cast(pl.Int8),
            pl.Series("power_fcst", predictions, dtype=pl.Float32),
            power_fcst_model_name=pl.lit(self.MODEL_NAME),
            power_fcst_model_version=pl.lit(self.MODEL_VERSION, dtype=pl.Int16),
            ml_flow_experiment_id=pl.lit(cfg.ml_flow_experiment_id, dtype=pl.Int32),
            experiment_name=pl.lit(cfg.experiment_name),
            fold_id=pl.lit(fold_id),
        )

    df = data.collect(engine="streaming")  # stream the NWP scan — see train() / overview.md
    parts: list[pl.DataFrame] = []
    for group_key, group in df.group_by(["time_series_id"]):
        booster = self._models.get(int(group_key[0]))
        if booster is None:
            continue  # ignore series this model was not trained on
        group = cast(pt.DataFrame[AllFeatures], group)
        X = _prepare_features(group, feature_cols)
        predictions: np.ndarray = booster.predict(xgb.DMatrix(X))
        parts.append(_build_part(group, predictions))

    if not parts:
        empty = cast(pt.DataFrame[AllFeatures], df.head(0))
        parts.append(_build_part(empty, np.empty(0, dtype=np.float32)))

    # The identity columns (power_fcst_model_name / experiment_name / fold_id) are built as
    # ``pl.lit(str)`` above, so they are already ``String`` — the dtype PowerForecast declares —
    # and need no cast before validation.
    return PowerForecast.validate(pl.concat(parts))
save(path)

Save all Boosters as .ubj files plus a meta.json with the full config.

Source code in packages/xgboost_forecaster/src/xgboost_forecaster/forecaster.py 
198
199
200
201
202
203
204
205
206
207
208
209
210
211
def save(self, path: Path) -> None:
    """Save all Boosters as .ubj files plus a meta.json with the full config."""
    path.mkdir(parents=True, exist_ok=True)
    for ts_id, booster in self._models.items():
        booster.save_model(str(path / f"{ts_id}.ubj"))
    (path / "meta.json").write_text(
        json.dumps(
            {
                "model_params": self.model_params.model_dump(mode="json"),
                "trained_time_series_ids": self.trained_time_series_ids,
                "model_class": f"{type(self).__module__}.{type(self).__qualname__}",
            }
        )
    )
load(path) classmethod

Reconstruct an XGBoostForecaster from a saved directory.

Source code in packages/xgboost_forecaster/src/xgboost_forecaster/forecaster.py 
213
214
215
216
217
218
219
220
221
222
223
@classmethod
def load(cls, path: Path) -> Self:
    """Reconstruct an XGBoostForecaster from a saved directory."""
    meta = json.loads((path / "meta.json").read_text())
    config = XGBoostConfig.model_validate(meta["model_params"])
    instance = cls(config)
    for ubj_file in sorted(path.glob("*.ubj")):
        booster = xgb.Booster()
        booster.load_model(str(ubj_file))
        instance._models[int(ubj_file.stem)] = booster
    return instance
save_to_mlflow(run_id)

Upload this trained model's artifacts to the given MLflow run.

Writes the model to a temporary directory via save (the subclass's own format), then uploads that directory to the run's artifact store under model/. The caller is responsible for setting the tracking URI (mlflow.set_tracking_uri) beforehand.

Parameters:

Name Type Description Default
run_id str

The MLflow run to attach the artifacts to.

 required
Source code in packages/ml_core/src/ml_core/base_forecaster.py 
131
132
133
134
135
136
137
138
139
140
141
142
143
144
def save_to_mlflow(self, run_id: str) -> None:
    """Upload this trained model's artifacts to the given MLflow run.

    Writes the model to a temporary directory via ``save`` (the subclass's own format), then
    uploads that directory to the run's artifact store under ``model/``. The caller is
    responsible for setting the tracking URI (``mlflow.set_tracking_uri``) beforehand.

    Args:
        run_id: The MLflow run to attach the artifacts to.
    """
    with tempfile.TemporaryDirectory() as tmp_dir:
        self.save(Path(tmp_dir))
        with mlflow.start_run(run_id=run_id):
            mlflow.log_artifacts(tmp_dir, artifact_path=_MLFLOW_ARTIFACT_PATH)
load_from_mlflow(run_id, cache_base_path) classmethod

Load a trained model for run_id, serving from the local cache when possible.

On a cache hit ({cache_base_path}/{run_id}/model already exists) the model is loaded straight from disk and MLflow is never contacted — this is what lets the live service keep serving during an MLflow outage. On a cache miss the artifacts are downloaded from the run into the cache, then loaded. The cache key is the immutable run ID, so a cached model never goes stale. The caller sets the tracking URI (mlflow.set_tracking_uri) beforehand.

Parameters:

Name Type Description Default
run_id str

The MLflow run the model was saved under.

 required
cache_base_path Path

Root of the local cache; the model lives at {cache_base_path}/{run_id}.

 required

Returns:

Type Description
Self

The reconstructed, trained forecaster.
Source code in packages/ml_core/src/ml_core/base_forecaster.py 
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
@classmethod
def load_from_mlflow(cls, run_id: str, cache_base_path: Path) -> Self:
    """Load a trained model for ``run_id``, serving from the local cache when possible.

    On a cache hit (``{cache_base_path}/{run_id}/model`` already exists) the model is loaded
    straight from disk and MLflow is never contacted — this is what lets the live service keep
    serving during an MLflow outage. On a cache miss the artifacts are downloaded from the run
    into the cache, then loaded. The cache key is the immutable run ID, so a cached model never
    goes stale. The caller sets the tracking URI (``mlflow.set_tracking_uri``) beforehand.

    Args:
        run_id: The MLflow run the model was saved under.
        cache_base_path: Root of the local cache; the model lives at
            ``{cache_base_path}/{run_id}``.

    Returns:
        The reconstructed, trained forecaster.
    """
    run_cache_dir = cache_base_path / run_id
    model_dir = run_cache_dir / _MLFLOW_ARTIFACT_PATH
    if not model_dir.exists():
        run_cache_dir.mkdir(parents=True, exist_ok=True)
        mlflow.artifacts.download_artifacts(
            run_id=run_id,
            artifact_path=_MLFLOW_ARTIFACT_PATH,
            dst_path=str(run_cache_dir),
        )
    return cls.load(model_dir)