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.). BaseForecasterConfig contributes selected_features, random_seed (threaded into XGBoost's own seed parameter for deterministic training), the experiment-identity fields experiment_name and ml_flow_experiment_id, and the leaderboard tag fields weather_source and training_strategy. Model-family identity — MODEL_NAME ("xgboost") and MODEL_VERSION — lives on the XGBoostForecaster class itself, not in the config; both a config's experiment identity and the class's model-family identity are stamped onto every row of the PowerForecast output, so the Delta Lake table is self-describing. XGBoostConfig inherits extra="forbid" from BaseForecasterConfig, so a key that names neither an XGBoostConfig field nor an inherited BaseForecasterConfig field raises ValidationError rather than being silently ignored.

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

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

model_config = ConfigDict(extra='forbid') 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

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

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
    CONFIG_CLASS: ClassVar[type[XGBoostConfig]] = XGBoostConfig

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

    def __init__(self, model_params: XGBoostConfig) -> None:
        """Start with no trained boosters; ``train`` populates one per ``time_series_id``."""
        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``, which is what ``load()`` reads
        it back from. 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
        <https://openclimatefix.github.io/nged-substation-forecast/architecture/overview/>.

        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
        # <https://openclimatefix.github.io/nged-substation-forecast/architecture/overview/>.
        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
        <https://openclimatefix.github.io/nged-substation-forecast/architecture/overview/>).

        ``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.

        Clears ``path`` first, so re-saving a smaller model over a directory that already holds a
        bigger one's ``.ubj`` files can never leave the dropped series' boosters behind on disk
        (issue #197). The same replace-don't-merge property holds through MLflow, because
        ``BaseForecaster.save_to_mlflow`` uploads this directory as a single archive artifact.
        """
        shutil.rmtree(path, ignore_errors=True)
        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": class_target(self),
                }
            )
        )

    @classmethod
    def load(cls, path: Path) -> Self:
        """Reconstruct an XGBoostForecaster from a saved directory.

        ``meta.json``'s ``trained_time_series_ids`` — not whatever ``.ubj`` files happen to be in
        the directory — decides the population, because a model's population is the model's own
        frozen record and not a directory listing (issue #197). A directory can hold files this
        model did not write: ``ml_core.base_forecaster.save_to_mlflow`` adds a
        ``time_series_metadata.parquet`` and ``ml_core.production_helpers.fetch_model_artifacts``
        a ``promotion.json``, and a hand-assembled directory can hold anything. Globbing would let
        such a file enlarge the population, silently scoring a series with a model that was never
        trained for it and breaking the train==predict invariant (see
        ``BaseForecaster.trained_time_series_ids``).
        """
        meta = json.loads((path / "meta.json").read_text())
        config = cls.CONFIG_CLASS.model_validate(meta["model_params"])
        instance = cls(config)
        for ts_id in meta["trained_time_series_ids"]:
            booster = xgb.Booster()
            booster.load_model(str(path / f"{ts_id}.ubj"))
            instance._models[int(ts_id)] = booster
        return instance

Attributes

MODEL_NAME = 'xgboost' class-attribute instance-attribute

MODEL_VERSION = 1 class-attribute instance-attribute

CONFIG_CLASS = XGBoostConfig class-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, which is what load() reads it back from. 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)

Start with no trained boosters; train populates one per time_series_id.

Source code in packages/xgboost_forecaster/src/xgboost_forecaster/forecaster.py

def __init__(self, model_params: XGBoostConfig) -> None:
    """Start with no trained boosters; ``train`` populates one per ``time_series_id``."""
    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 https://openclimatefix.github.io/nged-substation-forecast/architecture/overview/.

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

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
    <https://openclimatefix.github.io/nged-substation-forecast/architecture/overview/>.

    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
    # <https://openclimatefix.github.io/nged-substation-forecast/architecture/overview/>.
    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 https://openclimatefix.github.io/nged-substation-forecast/architecture/overview/).

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

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
    <https://openclimatefix.github.io/nged-substation-forecast/architecture/overview/>).

    ``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.

Clears path first, so re-saving a smaller model over a directory that already holds a bigger one's .ubj files can never leave the dropped series' boosters behind on disk (issue #197). The same replace-don't-merge property holds through MLflow, because BaseForecaster.save_to_mlflow uploads this directory as a single archive artifact.

Source code in packages/xgboost_forecaster/src/xgboost_forecaster/forecaster.py

def save(self, path: Path) -> None:
    """Save all Boosters as .ubj files plus a meta.json with the full config.

    Clears ``path`` first, so re-saving a smaller model over a directory that already holds a
    bigger one's ``.ubj`` files can never leave the dropped series' boosters behind on disk
    (issue #197). The same replace-don't-merge property holds through MLflow, because
    ``BaseForecaster.save_to_mlflow`` uploads this directory as a single archive artifact.
    """
    shutil.rmtree(path, ignore_errors=True)
    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": class_target(self),
            }
        )
    )

load(path) classmethod

Reconstruct an XGBoostForecaster from a saved directory.

meta.json's trained_time_series_ids — not whatever .ubj files happen to be in the directory — decides the population, because a model's population is the model's own frozen record and not a directory listing (issue #197). A directory can hold files this model did not write: ml_core.base_forecaster.save_to_mlflow adds a time_series_metadata.parquet and ml_core.production_helpers.fetch_model_artifacts a promotion.json, and a hand-assembled directory can hold anything. Globbing would let such a file enlarge the population, silently scoring a series with a model that was never trained for it and breaking the train==predict invariant (see BaseForecaster.trained_time_series_ids).

Source code in packages/xgboost_forecaster/src/xgboost_forecaster/forecaster.py

@classmethod
def load(cls, path: Path) -> Self:
    """Reconstruct an XGBoostForecaster from a saved directory.

    ``meta.json``'s ``trained_time_series_ids`` — not whatever ``.ubj`` files happen to be in
    the directory — decides the population, because a model's population is the model's own
    frozen record and not a directory listing (issue #197). A directory can hold files this
    model did not write: ``ml_core.base_forecaster.save_to_mlflow`` adds a
    ``time_series_metadata.parquet`` and ``ml_core.production_helpers.fetch_model_artifacts``
    a ``promotion.json``, and a hand-assembled directory can hold anything. Globbing would let
    such a file enlarge the population, silently scoring a series with a model that was never
    trained for it and breaking the train==predict invariant (see
    ``BaseForecaster.trained_time_series_ids``).
    """
    meta = json.loads((path / "meta.json").read_text())
    config = cls.CONFIG_CLASS.model_validate(meta["model_params"])
    instance = cls(config)
    for ts_id in meta["trained_time_series_ids"]:
        booster = xgb.Booster()
        booster.load_model(str(path / f"{ts_id}.ubj"))
        instance._models[int(ts_id)] = booster
    return instance

save_to_mlflow(run_id, *, time_series_metadata)

Upload this trained model to the given MLflow run, as one replaceable archive.

Writes the model to a temporary directory via save (the subclass's own format), adds the frozen metadata copy (write_trained_metadata), packs that directory into a single model.tar.gz and logs that one file to the run's artifact root. Logging one archive rather than a directory of files is what makes a re-upload replace the previous model instead of merging with it — see _MLFLOW_MODEL_ARTIFACT. 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 artifact to.	required
time_series_metadata	DataFrame[TimeSeriesMetadata]	The roster rows this model was engineered against. Required, because a model uploaded without them cannot be promoted — see TRAINED_METADATA_FILENAME. Narrowed to trained_time_series_ids before it is written: callers engineer over a wider population than they end up training (an eligible series with no usable power gets no model), and carrying the extras would widen the NWP scan every consumer prunes with these rows.	required

Source code in packages/ml_core/src/ml_core/base_forecaster.py

def save_to_mlflow(
    self, run_id: str, *, time_series_metadata: pt.DataFrame[TimeSeriesMetadata]
) -> None:
    """Upload this trained model to the given MLflow run, as one replaceable archive.

    Writes the model to a temporary directory via ``save`` (the subclass's own format), adds
    the frozen metadata copy (``write_trained_metadata``), packs that directory into a single
    ``model.tar.gz`` and logs *that one file* to the run's artifact root. Logging one archive
    rather than a directory of files is what makes a re-upload **replace** the previous model
    instead of merging with it — see ``_MLFLOW_MODEL_ARTIFACT``. The caller is responsible for
    setting the tracking URI (``mlflow.set_tracking_uri``) beforehand.

    Args:
        run_id: The MLflow run to attach the artifact to.
        time_series_metadata: The roster rows this model was engineered against. Required,
            because a model uploaded without them cannot be promoted — see
            ``TRAINED_METADATA_FILENAME``. Narrowed to ``trained_time_series_ids`` before it is
            written: callers engineer over a wider population than they end up training (an
            eligible series with no usable power gets no model), and carrying the extras would
            widen the NWP scan every consumer prunes with these rows.
    """
    with tempfile.TemporaryDirectory() as tmp_dir:
        model_dir = Path(tmp_dir) / "model"
        model_dir.mkdir()
        self.save(model_dir)
        write_trained_metadata(
            model_dir=model_dir,
            time_series_metadata=time_series_metadata.filter(
                pl.col("time_series_id").is_in(self.trained_time_series_ids)
            ),
        )
        archive_path = Path(tmp_dir) / _MLFLOW_MODEL_ARTIFACT
        _archive_model_dir(model_dir, archive_path)
        with mlflow.start_run(run_id=run_id):
            mlflow.log_artifact(str(archive_path))

load_from_mlflow(run_id) classmethod

Download a trained model's archive from an MLflow run, unpack it and load it.

Downloads into a temporary directory and loads from there. There is deliberately no local-disk cache: a CV fold run is reused across re-materialisations (get_or_create_fold_run resolves the same run for every re-run of a fold's partition), so a cache keyed by run_id would not be unique for its contents — and production inference makes no MLflow call at all. Full rationale: https://openclimatefix.github.io/nged-substation-forecast/architecture/ml-orchestration/#why-there-is-no-local-cache.

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 the model was saved under.	required

Returns:

Type	Description
Self	The reconstructed, trained forecaster.

Source code in packages/ml_core/src/ml_core/base_forecaster.py

@classmethod
def load_from_mlflow(cls, run_id: str) -> Self:
    """Download a trained model's archive from an MLflow run, unpack it and load it.

    Downloads into a temporary directory and loads from there. There is deliberately no
    local-disk cache: a CV fold run is **reused** across re-materialisations
    (``get_or_create_fold_run`` resolves the same run for every re-run of a fold's partition),
    so a cache keyed by ``run_id`` would not be unique for its contents — and production
    inference makes no MLflow call at all. Full rationale:
    <https://openclimatefix.github.io/nged-substation-forecast/architecture/ml-orchestration/#why-there-is-no-local-cache>.

    The caller is responsible for setting the tracking URI (``mlflow.set_tracking_uri``)
    beforehand.

    Args:
        run_id: The MLflow run the model was saved under.

    Returns:
        The reconstructed, trained forecaster.
    """
    with tempfile.TemporaryDirectory() as tmp_dir:
        return cls.load(
            _download_and_unpack_model(
                run_id=run_id,
                work_dir=Path(tmp_dir),
                remedy="re-materialise `trained_cv_model` for this fold.",
            )
        )

Functions: