ML Core API

ml_core

Unified ML model interface and shared utilities for substation forecasting.

Why this package exists

To enable rapid experimentation with different model architectures (XGBoost, GNNs, etc.), we need a common interface that shields the model developer from the underlying data engineering.

ml_core.features

Feature engineering package.

Public interface:

• FeatureEngineer — abstract base; implement to swap the full feature pipeline.
• TabularFeatureEngineer — default implementation: nearest-cell NWP spatial join followed by the declarative tabular pipeline.

Private sub-modules (importable by tests, not part of the public API):

• _parsed_features — typed feature descriptors and ParsedFeatures parser.
• _nwp — NWP upsampling, processing, and power/NWP join helpers.
• _lags — power-lag, weather-lag (dual-strategy), and leaky-lag nullification.
• tabular_feature_engineer — _engineer_features orchestrator and tabular pipeline helpers.

Attributes

__all__ = ['FeatureEngineer', 'TabularFeatureEngineer'] module-attribute

Classes

FeatureEngineer

Bases: ABC

Turns raw power + NWP + metadata into a model-ready AllFeatures frame.

Source code in packages/ml_core/src/ml_core/features/feature_engineer.py

class FeatureEngineer(ABC):
    """Turns raw power + NWP + metadata into a model-ready ``AllFeatures`` frame."""

    @abstractmethod
    def engineer(
        self,
        *,
        selected_features: set[str],
        power_time_series: pt.LazyFrame[PowerTimeSeries],
        time_series_metadata: pt.DataFrame[TimeSeriesMetadata],
        nwp: pt.LazyFrame[Nwp],
        power_fcst_init_time: datetime | None = None,
        nwp_init_time: datetime | None = None,
        nwp_publication_delay_hours: int = NWP_PUBLICATION_DELAY_HOURS,
        local_timezone: str = DEFAULT_LOCAL_TIMEZONE,
    ) -> pt.LazyFrame[AllFeatures]:
        """Engineer features for training/inference.

        Two operating modes, matching ``tabular_feature_engineer._engineer_features`` (see
        that function for the full detail this docstring summarises):

        - **Bulk mode** (``power_fcst_init_time=None``, the default): NWP-centric, vectorised
          over every NWP run in the input, one forecast per ``(nwp_init_time, valid_time)``
          pair. Used for training and multi-run backtesting.
        - **Single-run mode** (``power_fcst_init_time`` given): power-centric, joins exactly
          one NWP run (``nwp_init_time``, or derived from ``power_fcst_init_time`` minus
          ``nwp_publication_delay_hours`` if omitted) and stamps a constant
          ``power_fcst_init_time`` on every row. Used for production inference and replay
          backfilling.

        Args:
            selected_features: The feature names to produce.
            power_time_series: Observed power, one row per ``(time_series_id, time)``.
            time_series_metadata: Per-time-series metadata (carries ``h3_res_5`` for the
                spatial mapping).
            nwp: Gridded NWP in physical units, keyed by ``h3_index``.
            power_fcst_init_time: ``None`` for bulk mode; a single datetime for single-run
                mode. See ``_engineer_features`` for the full contract.
            nwp_init_time: The NWP run to join in single-run mode. Must be ``None`` in bulk
                mode. See ``_engineer_features``.
            nwp_publication_delay_hours: Delay used to derive whichever of
                ``power_fcst_init_time``/``nwp_init_time`` is not supplied.
            local_timezone: IANA zone the local-time features are computed in.

        Returns:
            A lazy ``AllFeatures`` frame, ready to hand to ``BaseForecaster.train``/``predict``.
        """

Methods:

engineer(*, selected_features, power_time_series, time_series_metadata, nwp, power_fcst_init_time=None, nwp_init_time=None, nwp_publication_delay_hours=NWP_PUBLICATION_DELAY_HOURS, local_timezone=DEFAULT_LOCAL_TIMEZONE) abstractmethod

Engineer features for training/inference.

Two operating modes, matching tabular_feature_engineer._engineer_features (see that function for the full detail this docstring summarises):

• Bulk mode (power_fcst_init_time=None, the default): NWP-centric, vectorised over every NWP run in the input, one forecast per (nwp_init_time, valid_time) pair. Used for training and multi-run backtesting.
• Single-run mode (power_fcst_init_time given): power-centric, joins exactly one NWP run (nwp_init_time, or derived from power_fcst_init_time minus nwp_publication_delay_hours if omitted) and stamps a constant power_fcst_init_time on every row. Used for production inference and replay backfilling.

Parameters:

Name	Type	Description	Default
selected_features	set[str]	The feature names to produce.	required
power_time_series	LazyFrame[PowerTimeSeries]	Observed power, one row per (time_series_id, time).	required
time_series_metadata	DataFrame[TimeSeriesMetadata]	Per-time-series metadata (carries h3_res_5 for the spatial mapping).	required
nwp	LazyFrame[Nwp]	Gridded NWP in physical units, keyed by h3_index.	required
power_fcst_init_time	datetime | None	None for bulk mode; a single datetime for single-run mode. See _engineer_features for the full contract.	None
nwp_init_time	datetime | None	The NWP run to join in single-run mode. Must be None in bulk mode. See _engineer_features.	None
nwp_publication_delay_hours	int	Delay used to derive whichever of power_fcst_init_time/nwp_init_time is not supplied.	NWP_PUBLICATION_DELAY_HOURS
local_timezone	str	IANA zone the local-time features are computed in.	DEFAULT_LOCAL_TIMEZONE

Returns:

Type	Description
LazyFrame[AllFeatures]	A lazy AllFeatures frame, ready to hand to BaseForecaster.train/predict.

Source code in packages/ml_core/src/ml_core/features/feature_engineer.py

@abstractmethod
def engineer(
    self,
    *,
    selected_features: set[str],
    power_time_series: pt.LazyFrame[PowerTimeSeries],
    time_series_metadata: pt.DataFrame[TimeSeriesMetadata],
    nwp: pt.LazyFrame[Nwp],
    power_fcst_init_time: datetime | None = None,
    nwp_init_time: datetime | None = None,
    nwp_publication_delay_hours: int = NWP_PUBLICATION_DELAY_HOURS,
    local_timezone: str = DEFAULT_LOCAL_TIMEZONE,
) -> pt.LazyFrame[AllFeatures]:
    """Engineer features for training/inference.

    Two operating modes, matching ``tabular_feature_engineer._engineer_features`` (see
    that function for the full detail this docstring summarises):

    - **Bulk mode** (``power_fcst_init_time=None``, the default): NWP-centric, vectorised
      over every NWP run in the input, one forecast per ``(nwp_init_time, valid_time)``
      pair. Used for training and multi-run backtesting.
    - **Single-run mode** (``power_fcst_init_time`` given): power-centric, joins exactly
      one NWP run (``nwp_init_time``, or derived from ``power_fcst_init_time`` minus
      ``nwp_publication_delay_hours`` if omitted) and stamps a constant
      ``power_fcst_init_time`` on every row. Used for production inference and replay
      backfilling.

    Args:
        selected_features: The feature names to produce.
        power_time_series: Observed power, one row per ``(time_series_id, time)``.
        time_series_metadata: Per-time-series metadata (carries ``h3_res_5`` for the
            spatial mapping).
        nwp: Gridded NWP in physical units, keyed by ``h3_index``.
        power_fcst_init_time: ``None`` for bulk mode; a single datetime for single-run
            mode. See ``_engineer_features`` for the full contract.
        nwp_init_time: The NWP run to join in single-run mode. Must be ``None`` in bulk
            mode. See ``_engineer_features``.
        nwp_publication_delay_hours: Delay used to derive whichever of
            ``power_fcst_init_time``/``nwp_init_time`` is not supplied.
        local_timezone: IANA zone the local-time features are computed in.

    Returns:
        A lazy ``AllFeatures`` frame, ready to hand to ``BaseForecaster.train``/``predict``.
    """

TabularFeatureEngineer

Bases: FeatureEngineer

Nearest res-5 NWP cell per time series, then the declarative tabular pipeline.

Source code in packages/ml_core/src/ml_core/features/tabular_feature_engineer.py

class TabularFeatureEngineer(FeatureEngineer):
    """Nearest res-5 NWP cell per time series, then the declarative tabular pipeline."""

    def engineer(
        self,
        *,
        selected_features: set[str],
        power_time_series: pt.LazyFrame[PowerTimeSeries],
        time_series_metadata: pt.DataFrame[TimeSeriesMetadata],
        nwp: pt.LazyFrame[Nwp],
        power_fcst_init_time: datetime | None = None,
        nwp_init_time: datetime | None = None,
        nwp_publication_delay_hours: int = NWP_PUBLICATION_DELAY_HOURS,
        local_timezone: str = DEFAULT_LOCAL_TIMEZONE,
    ) -> pt.LazyFrame[AllFeatures]:
        """Map each NWP cell to its nearest time series, then run the tabular feature pipeline.

        See :meth:`FeatureEngineer.engineer` for the argument and operating-mode contract, and
        ``_engineer_features`` in this module for the pipeline itself — including
        ``local_timezone``, the IANA zone the local-time features are computed in.
        """
        nwp_per_time_series = _attach_nearest_nwp_cell(nwp, time_series_metadata)
        return _engineer_features(
            selected_features,
            power_time_series,
            time_series_metadata,
            nwp=nwp_per_time_series,
            power_fcst_init_time=power_fcst_init_time,
            nwp_init_time=nwp_init_time,
            nwp_publication_delay_hours=nwp_publication_delay_hours,
            local_timezone=local_timezone,
        )

Methods:

engineer(*, selected_features, power_time_series, time_series_metadata, nwp, power_fcst_init_time=None, nwp_init_time=None, nwp_publication_delay_hours=NWP_PUBLICATION_DELAY_HOURS, local_timezone=DEFAULT_LOCAL_TIMEZONE)

Map each NWP cell to its nearest time series, then run the tabular feature pipeline.

See :meth:FeatureEngineer.engineer for the argument and operating-mode contract, and _engineer_features in this module for the pipeline itself — including local_timezone, the IANA zone the local-time features are computed in.

Source code in packages/ml_core/src/ml_core/features/tabular_feature_engineer.py

def engineer(
    self,
    *,
    selected_features: set[str],
    power_time_series: pt.LazyFrame[PowerTimeSeries],
    time_series_metadata: pt.DataFrame[TimeSeriesMetadata],
    nwp: pt.LazyFrame[Nwp],
    power_fcst_init_time: datetime | None = None,
    nwp_init_time: datetime | None = None,
    nwp_publication_delay_hours: int = NWP_PUBLICATION_DELAY_HOURS,
    local_timezone: str = DEFAULT_LOCAL_TIMEZONE,
) -> pt.LazyFrame[AllFeatures]:
    """Map each NWP cell to its nearest time series, then run the tabular feature pipeline.

    See :meth:`FeatureEngineer.engineer` for the argument and operating-mode contract, and
    ``_engineer_features`` in this module for the pipeline itself — including
    ``local_timezone``, the IANA zone the local-time features are computed in.
    """
    nwp_per_time_series = _attach_nearest_nwp_cell(nwp, time_series_metadata)
    return _engineer_features(
        selected_features,
        power_time_series,
        time_series_metadata,
        nwp=nwp_per_time_series,
        power_fcst_init_time=power_fcst_init_time,
        nwp_init_time=nwp_init_time,
        nwp_publication_delay_hours=nwp_publication_delay_hours,
        local_timezone=local_timezone,
    )

ml_core.base_forecaster

BaseForecaster: the interface every forecasting model implements.

Also holds BaseForecasterConfig, which carries a trained model's experiment identity.

Attributes

TRAINED_METADATA_FILENAME = 'time_series_metadata.parquet' module-attribute

Name of the TimeSeriesMetadata rows a saved model carries for its trained population.

Production inference reads each series' location from here, never from the TimeSeriesMetadata roster, so an unreadable or thinned roster cannot fail a live slot or silently drop a series from it. Being the model's own frozen copy of what it trained against also keeps a series' H3 cell and static feature values identical between training and serving. See https://openclimatefix.github.io/nged-substation-forecast/design-philosophy/inherent-stability/#the-rules.

Logging it as a second MLflow artifact instead of putting it in the model directory would reopen the merge problem _MLFLOW_MODEL_ARTIFACT documents.

Classes

BaseForecasterConfig

Bases: BaseModel

Universal configuration for all forecasting models.

Subclasses add model-specific hyperparameters. Having a shared base ensures that every forecaster carries its own feature list and optional MLflow experiment id in one serialisable object, simplifying save/load and the conf/model/*.yaml config wiring.

The tag fields (weather_source, training_strategy) are stamped onto MLflow runs for leaderboard grouping. They live here so that constructing the config from a model YAML's model_params validates them at load time — they are present in every conf/model/*.yaml file under model_params.

experiment_name is the per-experiment key (set to the MLflow experiment name at registration and stored in the saved config); it is stamped onto every PowerForecast row, distinct from the model-family MODEL_NAME. random_seed is threaded into each model's training so that re-training a fold reproduces the same model — keeping retries and the leaderboard stable.

Model identity (name and version) lives on the BaseForecaster class itself as MODEL_NAME and MODEL_VERSION — those are properties of the implementation, not the experiment config.

Serialisation must be canonical. A config is compared and stored as its serialised form: register_experiment stamps model_dump_json() onto the MLflow experiment as the config tag and compares a re-registration against it, and logs flatten_config(...) as write-once MLflow params. Python's per-process string hash randomisation means a set iterates in a different order in every process, so a set dumped straight to a list would make two dumps of the same config differ — a re-registration would then look like a config change and its param write would be rejected. selected_features is therefore serialised sorted. A subclass that adds a set-valued (or otherwise unordered) field must do the same.

Unknown keys are rejected, not ignored (extra="forbid"). A key no field declares raises ValidationError, so a misspelled hyperparameter in a run's config_overrides fails at registration, and a stored config carrying a key the current code no longer declares is refused rather than silently losing it — the recovery is to re-train, never to hand-edit. Why this is worth failing over: https://openclimatefix.github.io/nged-substation-forecast/ml_experimentation/model-configuration/#tweaking-a-config-for-an-experiment.

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

class BaseForecasterConfig(BaseModel):
    """Universal configuration for all forecasting models.

    Subclasses add model-specific hyperparameters. Having a shared base ensures that every
    forecaster carries its own feature list and optional MLflow experiment id in one
    serialisable object, simplifying save/load and the ``conf/model/*.yaml`` config wiring.

    The tag fields (weather_source, training_strategy) are stamped onto MLflow runs for
    leaderboard grouping. They live here so that constructing the config from a model YAML's
    ``model_params`` validates them at load time — they are present in every
    ``conf/model/*.yaml`` file under ``model_params``.

    ``experiment_name`` is the per-experiment key (set to the MLflow experiment name at
    registration and stored in the saved config); it is stamped onto every ``PowerForecast``
    row, distinct from the model-family ``MODEL_NAME``. ``random_seed`` is threaded into each
    model's training so that re-training a fold reproduces the same model — keeping retries
    and the leaderboard stable.

    Model identity (name and version) lives on the ``BaseForecaster`` class itself as
    ``MODEL_NAME`` and ``MODEL_VERSION`` — those are properties of the implementation, not
    the experiment config.

    **Serialisation must be canonical.** A config is compared and stored as its serialised form:
    ``register_experiment`` stamps ``model_dump_json()`` onto the MLflow experiment as the
    ``config`` tag and compares a re-registration against it, and logs ``flatten_config(...)`` as
    write-once MLflow params. Python's per-process string hash randomisation means a ``set``
    iterates in a different order in every process, so a set dumped straight to a list would make
    two dumps of the *same* config differ — a re-registration would then look like a config change
    and its param write would be rejected. ``selected_features`` is therefore serialised sorted. A
    subclass that adds a set-valued (or otherwise unordered) field must do the same.

    **Unknown keys are rejected, not ignored** (``extra="forbid"``). A key no field declares raises
    ``ValidationError``, so a misspelled hyperparameter in a run's ``config_overrides`` fails at
    registration, and a *stored* config carrying a key the current code no longer declares is
    refused rather than silently losing it — the recovery is to re-train, never to hand-edit.
    Why this is worth failing over:
    <https://openclimatefix.github.io/nged-substation-forecast/ml_experimentation/model-configuration/#tweaking-a-config-for-an-experiment>.
    """

    model_config = ConfigDict(extra="forbid")

    selected_features: set[str]
    ml_flow_experiment_id: int | None = None
    experiment_name: str = ""
    weather_source: str = ""
    training_strategy: str = ""
    random_seed: int = 0

    @field_serializer("selected_features")
    def _serialise_selected_features(self, selected_features: set[str]) -> list[str]:
        """Serialise the feature set sorted, so two dumps of the same config always match."""
        return sorted(selected_features)

Attributes

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

BaseForecaster

Bases: ABC

Defines the universal interface for all energy forecasting ML models.

Every forecasting model subclasses this abstract base to allow shared Dagster assets and evaluation code to remain completely agnostic to the underlying model implementation.

Subclasses must define MODEL_NAME, MODEL_VERSION and CONFIG_CLASS as class-level constants. MODEL_NAME and MODEL_VERSION are stamped onto every PowerForecast row at predict time and used as the MLflow experiment name. Bumping MODEL_VERSION requires a code change (intentional), not a config edit.

Lazy evaluation contract: train and predict both accept a pt.LazyFrame[AllFeatures]. Callers must not collect before passing data in; doing so wastes memory and prevents Polars from optimising the full query plan. A subclass materialises the data at the model boundary (typically a single .collect(), streamed). Keeping that bounded is the caller's responsibility: it prunes the inputs (NWP control member, the relevant H3 cells, the window's init_time partitions) and, where the full ensemble is needed, processes one init_time chunk at a time — filtering the engineered output cannot prune the upstream join/upsample. See the NWP scan-pruning notes in https://openclimatefix.github.io/nged-substation-forecast/architecture/overview/.

Persistence has two layers. Subclasses implement save/load for their own on-disk format and need know nothing about MLflow. The concrete save_to_mlflow/load_from_mlflow methods, shared by all subclasses, wrap that disk format with MLflow's artifact store, so the same trained model can be shared across machines by round-tripping through a run.

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

class BaseForecaster(ABC):
    """Defines the universal interface for all energy forecasting ML models.

    Every forecasting model subclasses this abstract base to allow shared Dagster assets and
    evaluation code to remain completely agnostic to the underlying model implementation.

    Subclasses must define ``MODEL_NAME``, ``MODEL_VERSION`` and ``CONFIG_CLASS`` as class-level
    constants. ``MODEL_NAME`` and ``MODEL_VERSION`` are stamped onto every ``PowerForecast`` row at
    predict time and used as the MLflow experiment name. Bumping ``MODEL_VERSION`` requires a code
    change (intentional), not a config edit.

    Lazy evaluation contract: `train` and `predict` both accept a `pt.LazyFrame[AllFeatures]`.
    Callers must not collect before passing data in; doing so wastes memory and prevents Polars
    from optimising the full query plan. A subclass materialises the data at the model boundary
    (typically a single `.collect()`, streamed). Keeping that bounded is the *caller's*
    responsibility: it prunes the inputs (NWP control member, the relevant H3 cells, the window's
    `init_time` partitions) and, where the full ensemble is needed, processes one `init_time` chunk
    at a time — filtering the engineered output cannot prune the upstream join/upsample. See the NWP
    scan-pruning notes in
    <https://openclimatefix.github.io/nged-substation-forecast/architecture/overview/>.

    Persistence has two layers. Subclasses implement ``save``/``load`` for their own on-disk
    format and need know nothing about MLflow. The concrete ``save_to_mlflow``/``load_from_mlflow``
    methods, shared by all subclasses, wrap that disk format with MLflow's artifact store, so the
    same trained model can be shared across machines by round-tripping through a run.
    """

    MODEL_NAME: ClassVar[str]
    MODEL_VERSION: ClassVar[int]

    CONFIG_CLASS: ClassVar[type[BaseForecasterConfig]]
    """The config class ``load`` rebuilds from a saved ``model_params`` mapping.

    A subclass's ``load`` must build its config through this attribute rather than naming its
    config class again, so the class ``ml_core.production_helpers._check_meta_is_servable``
    validates a saved config against is always the class ``load`` uses.
    """

    feature_engineer: ClassVar[FeatureEngineer] = TabularFeatureEngineer()
    """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.
    """

    def __init__(self, model_params: BaseForecasterConfig) -> None:
        """Store the config that carries this model's hyper-parameters and experiment identity."""
        self.model_params = model_params

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

        This is the model's own frozen record of who it trained on — *not* a statement about how
        the model is internally structured. A model may hold one sub-model per series, a single
        model spanning many series, or anything in between; the contract is only about which
        ``time_series_id``s it will score.

        **Why this is load-bearing.** The **train==predict population invariant**: the
        population a model is scored on must equal the population it was trained on, *even if* the
        live eligibility set has drifted since training (power coverage changes, so a series may
        newly qualify or drop out). Consumers (``cv_power_forecasts``, ``live_forecasts``) filter
        their inputs to this set, so a model scores exactly the population it learned — never
        whatever eligibility says *today*. That is what keeps the leaderboard apples-to-apples and
        stops production forecasting a series the model never saw.

        Subclasses persist and reconstruct this set through their own ``save``/``load`` (e.g. in
        ``meta.json``), so it survives a round-trip through MLflow.
        """
        ...

    @abstractmethod
    def save(self, path: Path) -> None:
        """Save the trained model state to a directory, **replacing** anything already there.

        Two requirements on every implementation:

        - The saved directory must contain a ``meta.json`` with a ``model_class`` field — the
          fully-qualified ``{module}.{qualname}`` of the concrete subclass (e.g.
          ``"xgboost_forecaster.forecaster.XGBoostForecaster"``) — so that production inference
          (``ml_core.production_helpers.load_forecaster_from_dir``) can reconstruct the correct
          class from a plain model directory with no other context (issue #221).
        - ``path`` must be cleared first, so that saving over a directory holding a *larger*
          model's files leaves none of them behind. Merging instead of replacing is how a dropped
          time series' weights survive a re-train (issue #197); ``XGBoostForecaster.save`` clears
          with ``shutil.rmtree(path, ignore_errors=True)``.

        The clearing requirement makes ``path`` the model's to own while it saves, so anything a
        caller left there is gone afterwards. (Depositing a file *after* a save is fine, and is how
        ``production_helpers.fetch_model_artifacts`` puts ``promotion.json`` beside the model.)
        """

    @classmethod
    @abstractmethod
    def load(cls, path: Path) -> Self:
        """Reconstruct a trained instance from a previously saved directory."""

    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))

    @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.",
                )
            )

    @abstractmethod
    def train(self, data: pt.LazyFrame[AllFeatures], time_series_ids: list[int]) -> None:
        """Fit the model on ``data`` for the given ``time_series_id`` population.

        Args:
            data: The engineered features (lazy; the caller must not pre-collect).
            time_series_ids: The population the model must train on — the caller's eligible set.
                The model decides how to map it onto boosters: one booster per id
                (``XGBoostForecaster``), or one booster per group of ids (e.g. all solar sites) for
                a future model. It is the model's frozen record of who it trained on (the
                train==predict invariant), and bounds ``predict`` to that same population.
        """

    @abstractmethod
    def predict(
        self, data: pt.LazyFrame[AllFeatures], *, fold_id: str = "live"
    ) -> pt.DataFrame[PowerForecast]:
        """Return power forecasts for all rows in the given AllFeatures data.

        Args:
            data: The engineered features to forecast from.
            fold_id: The value stamped onto every row's ``fold_id`` column. The model has no
                inherent notion of which CV fold it is serving (``fold_id`` is orchestration
                context), so the caller supplies it: ``cv_power_forecasts`` passes the fold's
                label, while production inference keeps the ``"live"`` default.
        """

Attributes

MODEL_NAME class-attribute

MODEL_VERSION class-attribute

CONFIG_CLASS class-attribute

The config class load rebuilds from a saved model_params mapping.

A subclass's load must build its config through this attribute rather than naming its config class again, so the class ml_core.production_helpers._check_meta_is_servable validates a saved config against is always the class load uses.

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.

model_params = model_params instance-attribute

trained_time_series_ids abstractmethod property

The sorted time_series_id population this model will serve a predict for.

This is the model's own frozen record of who it trained on — not a statement about how the model is internally structured. A model may hold one sub-model per series, a single model spanning many series, or anything in between; the contract is only about which time_series_ids it will score.

Why this is load-bearing. The train==predict population invariant: the population a model is scored on must equal the population it was trained on, even if the live eligibility set has drifted since training (power coverage changes, so a series may newly qualify or drop out). Consumers (cv_power_forecasts, live_forecasts) filter their inputs to this set, so a model scores exactly the population it learned — never whatever eligibility says today. That is what keeps the leaderboard apples-to-apples and stops production forecasting a series the model never saw.

Subclasses persist and reconstruct this set through their own save/load (e.g. in meta.json), so it survives a round-trip through MLflow.

Methods:

__init__(model_params)

Store the config that carries this model's hyper-parameters and experiment identity.

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

def __init__(self, model_params: BaseForecasterConfig) -> None:
    """Store the config that carries this model's hyper-parameters and experiment identity."""
    self.model_params = model_params

save(path) abstractmethod

Save the trained model state to a directory, replacing anything already there.

Two requirements on every implementation:

• The saved directory must contain a meta.json with a model_class field — the fully-qualified {module}.{qualname} of the concrete subclass (e.g. "xgboost_forecaster.forecaster.XGBoostForecaster") — so that production inference (ml_core.production_helpers.load_forecaster_from_dir) can reconstruct the correct class from a plain model directory with no other context (issue #221).
• path must be cleared first, so that saving over a directory holding a larger model's files leaves none of them behind. Merging instead of replacing is how a dropped time series' weights survive a re-train (issue #197); XGBoostForecaster.save clears with shutil.rmtree(path, ignore_errors=True).

The clearing requirement makes path the model's to own while it saves, so anything a caller left there is gone afterwards. (Depositing a file after a save is fine, and is how production_helpers.fetch_model_artifacts puts promotion.json beside the model.)

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

@abstractmethod
def save(self, path: Path) -> None:
    """Save the trained model state to a directory, **replacing** anything already there.

    Two requirements on every implementation:

    - The saved directory must contain a ``meta.json`` with a ``model_class`` field — the
      fully-qualified ``{module}.{qualname}`` of the concrete subclass (e.g.
      ``"xgboost_forecaster.forecaster.XGBoostForecaster"``) — so that production inference
      (``ml_core.production_helpers.load_forecaster_from_dir``) can reconstruct the correct
      class from a plain model directory with no other context (issue #221).
    - ``path`` must be cleared first, so that saving over a directory holding a *larger*
      model's files leaves none of them behind. Merging instead of replacing is how a dropped
      time series' weights survive a re-train (issue #197); ``XGBoostForecaster.save`` clears
      with ``shutil.rmtree(path, ignore_errors=True)``.

    The clearing requirement makes ``path`` the model's to own while it saves, so anything a
    caller left there is gone afterwards. (Depositing a file *after* a save is fine, and is how
    ``production_helpers.fetch_model_artifacts`` puts ``promotion.json`` beside the model.)
    """

load(path) abstractmethod classmethod

Reconstruct a trained instance from a previously saved directory.

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

@classmethod
@abstractmethod
def load(cls, path: Path) -> Self:
    """Reconstruct a trained instance from a previously saved directory."""

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.",
            )
        )

train(data, time_series_ids) abstractmethod

Fit the model on data for the given time_series_id population.

Parameters:

Name	Type	Description	Default
data	LazyFrame[AllFeatures]	The engineered features (lazy; the caller must not pre-collect).	required
time_series_ids	list[int]	The population the model must train on — the caller's eligible set. The model decides how to map it onto boosters: one booster per id (XGBoostForecaster), or one booster per group of ids (e.g. all solar sites) for a future model. It is the model's frozen record of who it trained on (the train==predict invariant), and bounds predict to that same population.	required

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

@abstractmethod
def train(self, data: pt.LazyFrame[AllFeatures], time_series_ids: list[int]) -> None:
    """Fit the model on ``data`` for the given ``time_series_id`` population.

    Args:
        data: The engineered features (lazy; the caller must not pre-collect).
        time_series_ids: The population the model must train on — the caller's eligible set.
            The model decides how to map it onto boosters: one booster per id
            (``XGBoostForecaster``), or one booster per group of ids (e.g. all solar sites) for
            a future model. It is the model's frozen record of who it trained on (the
            train==predict invariant), and bounds ``predict`` to that same population.
    """

predict(data, *, fold_id='live') abstractmethod

Return power forecasts for all rows in the given AllFeatures data.

Parameters:

Name	Type	Description	Default
data	LazyFrame[AllFeatures]	The engineered features to forecast from.	required
fold_id	str	The value stamped onto every row's fold_id column. The model has no inherent notion of which CV fold it is serving (fold_id is orchestration context), so the caller supplies it: cv_power_forecasts passes the fold's label, while production inference keeps the "live" default.	'live'

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

@abstractmethod
def predict(
    self, data: pt.LazyFrame[AllFeatures], *, fold_id: str = "live"
) -> pt.DataFrame[PowerForecast]:
    """Return power forecasts for all rows in the given AllFeatures data.

    Args:
        data: The engineered features to forecast from.
        fold_id: The value stamped onto every row's ``fold_id`` column. The model has no
            inherent notion of which CV fold it is serving (``fold_id`` is orchestration
            context), so the caller supplies it: ``cv_power_forecasts`` passes the fold's
            label, while production inference keeps the ``"live"`` default.
    """

Functions:

write_trained_metadata(model_dir, time_series_metadata)

Write a model's frozen TimeSeriesMetadata copy into its saved directory.

Call this after a subclass's save, which clears the directory first.

Parameters:

Name	Type	Description	Default
model_dir	Path	The directory a subclass's save just wrote.	required
time_series_metadata	DataFrame[TimeSeriesMetadata]	The roster rows the model was engineered against. _UNPERSISTED_METADATA_COLUMN is dropped; everything else is kept.	required

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

def write_trained_metadata(
    model_dir: Path, time_series_metadata: pt.DataFrame[TimeSeriesMetadata]
) -> None:
    """Write a model's frozen ``TimeSeriesMetadata`` copy into its saved directory.

    Call this *after* a subclass's ``save``, which clears the directory first.

    Args:
        model_dir: The directory a subclass's ``save`` just wrote.
        time_series_metadata: The roster rows the model was engineered against.
            ``_UNPERSISTED_METADATA_COLUMN`` is dropped; everything else is kept.
    """
    # `pl.exclude` rather than `drop`: Patito overrides `DataFrame.drop` with a signature that
    # takes no `strict=False`, and the column is `allow_missing`, so it may not be there to drop.
    time_series_metadata.select(pl.exclude(_UNPERSISTED_METADATA_COLUMN)).write_parquet(
        model_dir / TRAINED_METADATA_FILENAME
    )

load_trained_metadata(model_dir)

Read back the TimeSeriesMetadata rows a saved model carries.

set_model rather than validate, matching how the R&D assets read the roster itself: this reads back what was written, so re-checking it would only reject rosters the rest of the system already accepts.

Parameters:

Name	Type	Description	Default
model_dir	Path	A directory written by save_to_mlflow (via write_trained_metadata) and typically unpacked by fetch_model_artifacts.	required

Returns:

Type	Description
DataFrame[TimeSeriesMetadata]	One row per series in trained_time_series_ids — the population the model will serve a
DataFrame[TimeSeriesMetadata]	predict for, not the wider one it was engineered over — without
DataFrame[TimeSeriesMetadata]	_UNPERSISTED_METADATA_COLUMN.

Raises:

Type	Description
FileNotFoundError	The directory holds no such file, so it was saved by code predating this file or assembled by hand. This is a promotion fault rather than a data outage — the model on disk is not one this code can serve — so it raises rather than degrading, as a missing meta.json does.

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

def load_trained_metadata(model_dir: Path) -> pt.DataFrame[TimeSeriesMetadata]:
    """Read back the ``TimeSeriesMetadata`` rows a saved model carries.

    ``set_model`` rather than ``validate``, matching how the R&D assets read the roster itself:
    this reads back what was written, so re-checking it would only reject rosters the rest of the
    system already accepts.

    Args:
        model_dir: A directory written by ``save_to_mlflow`` (via ``write_trained_metadata``) and
            typically unpacked by ``fetch_model_artifacts``.

    Returns:
        One row per series in ``trained_time_series_ids`` — the population the model will serve a
        ``predict`` for, not the wider one it was engineered over — without
        ``_UNPERSISTED_METADATA_COLUMN``.

    Raises:
        FileNotFoundError: The directory holds no such file, so it was saved by code predating
            this file or assembled by hand. This is a promotion fault rather than a data outage —
            the model on disk is not one this code can serve — so it raises rather than degrading,
            as a missing ``meta.json`` does.
    """
    path = model_dir / TRAINED_METADATA_FILENAME
    if not path.exists():
        raise FileNotFoundError(
            f"The model directory {model_dir} has no {TRAINED_METADATA_FILENAME}, so there is no "
            "record of where its time series are and it cannot be served. Re-train against the "
            "current code and promote that run."
        )
    return pt.DataFrame(pl.read_parquet(path)).set_model(TimeSeriesMetadata)

ml_core.metrics

Metric computation for cross-validation results.

Attributes

Classes

NoOverlappingActualsError

Bases: ValueError

Raised by compute_metrics when no forecast row joins to any observed actual.

A distinct subclass so callers that score a fold in per-series batches can treat "this batch's series have no overlapping actuals" as skippable — mirroring how such series silently vanish from the inner join when the whole fold is scored in one call — while every other ValueError (negative lead times, missing capacity) still propagates.

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

class NoOverlappingActualsError(ValueError):
    """Raised by ``compute_metrics`` when no forecast row joins to any observed actual.

    A distinct subclass so callers that score a fold in per-series batches can treat "this
    batch's series have no overlapping actuals" as skippable — mirroring how such series
    silently vanish from the inner join when the whole fold is scored in one call — while
    every other ``ValueError`` (negative lead times, missing capacity) still propagates.
    """

Functions:

compute_effective_capacity(power_lf)

Compute the v0.1 effective capacity (full-history P99 of |power|) per time series.

One row per time_series_id: effective_capacity_mw is the 99th percentile of abs(power) over all non-null observations. Series whose P99 is null or non-positive (e.g. all-null or all-zero power) are dropped, since EffectiveCapacity requires effective_capacity_mw > 0.

time is set to that series' latest observed timestep (time.max()). The v0.1 capacity is a single scalar per series, so time is really an "as of" marker — it stamps the estimate as current to the end of the observed history — rather than a timestep the value varies over. The v0.7 upgrade makes capacity genuinely time-varying (one row per (time_series_id, time)), and only then does time carry per-row meaning. v0.1 stays one scalar row per series rather than the value repeated at every half-hour: densifying a constant adds rows without information, and the metrics join is by time_series_id alone until capacity varies.

Kept as a pure helper (no Dagster, no IO) so the P99 logic is unit-testable in isolation.

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

def compute_effective_capacity(
    power_lf: pt.LazyFrame[PowerTimeSeries],
) -> pt.DataFrame[EffectiveCapacity]:
    """Compute the v0.1 effective capacity (full-history P99 of ``|power|``) per time series.

    One row per ``time_series_id``: ``effective_capacity_mw`` is the 99th percentile of
    ``abs(power)`` over all non-null observations. Series whose P99 is null or non-positive (e.g.
    all-null or all-zero power) are dropped, since ``EffectiveCapacity`` requires
    ``effective_capacity_mw > 0``.

    ``time`` is set to that series' **latest** observed timestep (``time.max()``). The v0.1 capacity
    is a single scalar per series, so ``time`` is really an "as of" marker — it stamps the estimate
    as current to the end of the observed history — rather than a timestep the value varies over.
    The v0.7 upgrade makes capacity genuinely time-varying (one row per ``(time_series_id, time)``),
    and only then does ``time`` carry per-row meaning. v0.1 stays one scalar row per series rather
    than the value repeated at every half-hour: densifying a constant adds rows without information,
    and the metrics join is by ``time_series_id`` alone until capacity varies.

    Kept as a pure helper (no Dagster, no IO) so the P99 logic is unit-testable in isolation.
    """
    capacity = (
        power_lf.filter(pl.col("power").is_not_null())
        .group_by("time_series_id")
        .agg(
            effective_capacity_mw=pl.col("power").abs().quantile(0.99),
            time=pl.col("time").max(),
        )
        .filter(pl.col("effective_capacity_mw") > 0)
        .sort("time_series_id")
        .collect()
        .cast({"effective_capacity_mw": pl.Float32})
    )
    return EffectiveCapacity.validate(capacity)

compute_metrics(cv_forecasts, actuals, metadata, capacity)

Compute evaluation metrics from CV predictions and observed power.

Full metric definitions, equations, and design rationale: https://openclimatefix.github.io/nged-substation-forecast/techniques/evaluation-metrics/

For each (time_series_id, fold_id, power_fcst_model_name) group:

1. Joins predictions to observed power on (time_series_id, valid_time).
2. Assigns each row a horizon_slice from its lead time (valid_time − power_fcst_init_time) — see _horizon_slice_expr for the bands.
3. Collapses the ensemble members within each forecast run (per power_fcst_init_time) into per-timestamp quantities: the deterministic ensemble mean, the fair CRPS, the Fortin-corrected ensemble variance, and the empirical DELIVERY_QUANTILES. Each run covering a valid_time is scored independently, exactly as a production consumer would experience it — runs at different lead times are never pooled.
4. Aggregates per horizon_slice, plus the "all" aggregate over every lead time: MAE, NMAE, RMSE, and MBE on the ensemble mean; CRPS and the spread-skill ratio from the member-aware quantities; pinball loss at each delivery quantile (plus their mean); and PICP and interval width for each symmetric quantile band.
5. Joins time_series_type from metadata onto each row.
6. Returns one row per (time_series_id, fold_id, power_fcst_model_name, horizon_slice, metric_name, metric_param) in the tall Metrics format.

NMAE is normalised by the pre-computed full-history effective_capacity_mw (joined per time_series_id from capacity) — a capacity-like denominator, computed over the full history so it stays stable across folds. This time_series_id-only join is correct while capacity is one scalar row per series (the v0.1 shape). When the v0.7 upgrade makes capacity time-varying (one row per (time_series_id, time)), this join must become a temporal as-of join on (time_series_id, valid_time).

Single-member "ensembles" (e.g. a deterministic baseline forecaster) are scored unconditionally: their fair CRPS equals their MAE, their spread-skill ratio is 0, and their quantile bands are degenerate (all quantiles coincide, so PICP ≈ 0 and interval width = 0). Those are honest descriptions of a deterministic forecast, not errors.

Parameters:

Name	Type	Description	Default
cv_forecasts	DataFrame[PowerForecast]	CV predictions to evaluate. Currently only CV fold rows are handled; fold_id="live" support will be added with the production_monitoring scope in Phase 8.	required
actuals	LazyFrame[PowerTimeSeries]	Observed half-hourly power (lazy — only the joined subset is collected). Deduplicated on (time_series_id, time) before the join — a duplicated actual would otherwise double the ensemble members and corrupt the member-aware metrics.	required
metadata	DataFrame[TimeSeriesMetadata]	Substation metadata used to join time_series_type onto each metric row. Must cover every scored time_series_id.	required
capacity	DataFrame[EffectiveCapacity]	Pre-computed per-series effective capacity; effective_capacity_mw is the NMAE denominator. Must cover every scored time_series_id.	required

Returns:

Type	Description
DataFrame[Metrics]	A validated tall Metrics DataFrame with time_series_type populated.

Raises:

Type	Description
NoOverlappingActualsError	If no rows survive the inner join (forecasts cover a period with no observed data).
ValueError	If any forecast row has a negative lead time (valid_time before power_fcst_init_time — an undeliverable hindcast row; see issue #346), if any scored series has no row in capacity or in metadata, or if any computed metric value is non-finite (NaN/inf — which Metrics.validate would otherwise accept, since NaN is not null, and which would poison the MLflow aggregate means).

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

def compute_metrics(
    cv_forecasts: pt.DataFrame[PowerForecast],
    actuals: pt.LazyFrame[PowerTimeSeries],
    metadata: pt.DataFrame[TimeSeriesMetadata],
    capacity: pt.DataFrame[EffectiveCapacity],
) -> pt.DataFrame[Metrics]:
    """Compute evaluation metrics from CV predictions and observed power.

    Full metric definitions, equations, and design rationale:
    <https://openclimatefix.github.io/nged-substation-forecast/techniques/evaluation-metrics/>

    For each ``(time_series_id, fold_id, power_fcst_model_name)`` group:

    1. Joins predictions to observed ``power`` on ``(time_series_id, valid_time)``.
    2. Assigns each row a ``horizon_slice`` from its lead time
       (``valid_time − power_fcst_init_time``) — see ``_horizon_slice_expr`` for the bands.
    3. Collapses the ensemble members *within each forecast run* (per
       ``power_fcst_init_time``) into per-timestamp quantities: the deterministic ensemble
       mean, the fair CRPS, the Fortin-corrected ensemble variance, and the empirical
       ``DELIVERY_QUANTILES``. Each run covering a ``valid_time`` is scored independently,
       exactly as a production consumer would experience it — runs at different lead times
       are never pooled.
    4. Aggregates per ``horizon_slice``, plus the ``"all"`` aggregate over every lead time:
       MAE, NMAE, RMSE, and MBE on the ensemble mean; CRPS and the spread-skill ratio from
       the member-aware quantities; pinball loss at each delivery quantile (plus their
       mean); and PICP and interval width for each symmetric quantile band.
    5. Joins ``time_series_type`` from ``metadata`` onto each row.
    6. Returns one row per ``(time_series_id, fold_id, power_fcst_model_name,
       horizon_slice, metric_name, metric_param)`` in the tall ``Metrics`` format.

    NMAE is normalised by the pre-computed full-history ``effective_capacity_mw`` (joined per
    ``time_series_id`` from ``capacity``) — a capacity-like denominator, computed over the full
    history so it stays stable across folds. This ``time_series_id``-only join is correct while
    ``capacity`` is one scalar row per series (the v0.1 shape). When the v0.7
    upgrade makes capacity time-varying (one row per ``(time_series_id,
    time)``), this join must become a temporal as-of join on ``(time_series_id, valid_time)``.

    Single-member "ensembles" (e.g. a deterministic baseline forecaster) are scored
    unconditionally: their fair CRPS equals their MAE, their spread-skill ratio is 0, and
    their quantile bands are degenerate (all quantiles coincide, so PICP ≈ 0 and interval
    width = 0). Those are honest descriptions of a deterministic forecast, not errors.

    Args:
        cv_forecasts: CV predictions to evaluate.
            Currently only CV fold rows are handled; ``fold_id="live"`` support will
            be added with the ``production_monitoring`` scope in Phase 8.
        actuals: Observed half-hourly power (lazy — only the joined subset is collected).
            Deduplicated on ``(time_series_id, time)`` before the join — a duplicated actual
            would otherwise double the ensemble members and corrupt the member-aware metrics.
        metadata: Substation metadata used to join ``time_series_type`` onto each metric row.
            Must cover every scored ``time_series_id``.
        capacity: Pre-computed per-series effective capacity; ``effective_capacity_mw`` is the
            NMAE denominator. Must cover every scored ``time_series_id``.

    Returns:
        A validated tall ``Metrics`` DataFrame with ``time_series_type`` populated.

    Raises:
        NoOverlappingActualsError: If no rows survive the inner join (forecasts cover a
            period with no observed data).
        ValueError: If any forecast row has a negative lead time (``valid_time`` before
            ``power_fcst_init_time`` — an undeliverable hindcast row; see issue #346), if
            any scored series has no row in ``capacity`` or in ``metadata``, or if any
            computed metric value is
            non-finite (NaN/inf — which ``Metrics.validate`` would otherwise accept, since
            NaN is not null, and which would poison the MLflow aggregate means).
    """
    # A negative lead time means hindcast rows — valid times already in the past at
    # power_fcst_init_time, which a live forecast could never deliver. Scoring them would
    # silently flatter the model (they'd land in "intraday" via the left-closed bands), so
    # fail loudly. The CV inference pass currently emits such rows for valid times inside
    # the NWP publication-delay window; issue #346 tracks removing them at the source.
    n_negative = cv_forecasts.filter(pl.col("valid_time") < pl.col("power_fcst_init_time")).height
    if n_negative > 0:
        raise ValueError(
            f"{n_negative} forecast row(s) have valid_time before power_fcst_init_time "
            "(negative lead time). These are undeliverable hindcast rows and must not be "
            "scored — regenerate the forecasts without them (see issue #346)."
        )

    # Join forecasts to actuals; rename power → power_actual to avoid shadowing.
    # Strip the Patito model subclass from actuals so that Polars' cross-subclass
    # type check (assert_same_type) doesn't reject a join between two differently-typed
    # pt.LazyFrame objects.
    # Dedupe actuals on the join key: a duplicated (time_series_id, time) row would double
    # every ensemble member through the join, silently corrupting the member-aware metrics
    # (CRPS, spread, quantiles) while leaving the deterministic ones — which only see the
    # mean — untouched. Nothing enforces this uniqueness upstream, so guard it here.
    actuals_plain = pl.LazyFrame._from_pyldf(actuals._ldf)
    joined = cv_forecasts.lazy().join(
        actuals_plain.select(["time_series_id", "time", "power"])
        .unique(subset=["time_series_id", "time"], keep="any")
        .rename({"power": "power_actual"}),
        left_on=["time_series_id", "valid_time"],
        right_on=["time_series_id", "time"],
        how="inner",
    )

    # Collapse the ensemble members of each forecast run into per-timestamp quantities: the
    # deterministic ensemble mean plus the member-aware values (fair CRPS, Fortin-corrected
    # variance, empirical delivery quantiles). power_fcst_init_time is a group key so that
    # runs covering the same valid_time at different lead times are scored independently,
    # never pooled into a lagged-ensemble blend. horizon_slice is constant within a
    # (power_fcst_init_time, valid_time) group.
    quantile_aggs = {
        _quantile_column(q): pl.col("power_fcst").cast(pl.Float64).quantile(q, "linear")
        for q in DELIVERY_QUANTILES
    }
    per_run = (
        joined.with_columns(horizon_slice=_horizon_slice_expr())
        .group_by(
            [
                "time_series_id",
                "fold_id",
                "power_fcst_model_name",
                "power_fcst_init_time",
                "valid_time",
            ]
        )
        .agg(
            power_fcst=pl.col("power_fcst").mean(),
            power_actual=pl.col("power_actual").first(),
            horizon_slice=pl.col("horizon_slice").first(),
            crps=_fair_crps_expr(),
            corrected_var=_corrected_variance_expr(),
            **quantile_aggs,
        )
    )

    # Compute the ensemble-mean error column once, reuse in all aggregations.
    with_error = per_run.with_columns(error=pl.col("power_fcst") - pl.col("power_actual"))

    # Wide metrics: one row per (time_series_id, fold_id, power_fcst_model_name, horizon_slice),
    # where horizon_slice covers the four lead-time bands plus the "all" aggregate over every
    # lead time.
    base_keys = ["time_series_id", "fold_id", "power_fcst_model_name"]
    wide_columns = [*base_keys, "horizon_slice", *_wide_metric_columns()]
    per_slice = _wide_metrics(with_error, [*base_keys, "horizon_slice"])
    all_slice = _wide_metrics(with_error, base_keys).with_columns(
        horizon_slice=pl.lit("all").cast(pl.Enum(HORIZON_SLICES))
    )
    metrics_wide = pl.concat([per_slice.select(wide_columns), all_slice.select(wide_columns)])

    # Join the pre-computed full-history effective capacity — the NMAE denominator. Strip the
    # Patito model so Polars' cross-subclass join type check doesn't reject the join.
    capacity_denom = pl.LazyFrame._from_pyldf(capacity.lazy()._ldf).select(
        ["time_series_id", "effective_capacity_mw"]
    )
    wide = metrics_wide.join(capacity_denom, on="time_series_id", how="left").collect()

    # Every scored series must have a capacity row; fail loudly rather than silently emitting a
    # null NMAE (which Metrics.metric_value, a non-nullable Float32, would reject anyway).
    missing = wide.filter(pl.col("effective_capacity_mw").is_null())["time_series_id"]
    if missing.len() > 0:
        raise ValueError(
            f"No effective_capacity row for time_series_id(s) {sorted(set(missing.to_list()))}; "
            "materialise the effective_capacity asset for these series before scoring."
        )
    wide = wide.with_columns(nmae=pl.col("mae") / pl.col("effective_capacity_mw"))

    # Pivot to tall format, then split the "name:param" encoding of the parametric wide
    # columns into (metric_name, metric_param); scalar metrics have no separator, so their
    # split field_1 is null and fills to "all". The leftover effective_capacity_mw column is
    # dropped by the unpivot.
    name_parts = pl.col("metric_name").str.split_exact(":", 1)
    metrics_tall = (
        wide.unpivot(
            on=[*_wide_metric_columns(), "nmae"],
            index=[*base_keys, "horizon_slice"],
            variable_name="metric_name",
            value_name="metric_value",
        )
        .with_columns(
            metric_name=name_parts.struct.field("field_0"),
            metric_param=name_parts.struct.field("field_1").fill_null("all"),
        )
        .cast(
            {
                "metric_name": pl.Enum(METRIC_NAMES),
                "metric_param": pl.Enum(METRIC_PARAMS),
                "metric_value": pl.Float32,
            }
        )
    )

    if metrics_tall.is_empty():
        raise NoOverlappingActualsError(
            "No rows in the joined forecast/actuals data. "
            "Check that cv_forecasts and actuals overlap in time."
        )

    # NaN/inf are not null, so Metrics.validate would accept them — and a single non-finite
    # value poisons every downstream MLflow mean. Fail loudly instead, naming the offenders.
    non_finite = metrics_tall.filter(~pl.col("metric_value").is_finite())
    if not non_finite.is_empty():
        offenders = non_finite.select(
            ["time_series_id", "horizon_slice", "metric_name", "metric_param", "metric_value"]
        )
        raise ValueError(
            f"{non_finite.height} metric row(s) have a non-finite metric_value, which would "
            f"silently poison the MLflow aggregate means. First offenders:\n{offenders.head(10)}"
        )

    # Join time_series_type from metadata. Left, so a series with no metadata row surfaces as a
    # null to be named below rather than vanishing from the leaderboard.
    type_map = metadata.select(["time_series_id", "time_series_type"])
    metrics_tall = metrics_tall.join(type_map, on="time_series_id", how="left").with_columns(
        time_series_type=pl.col("time_series_type").cast(pl.Enum(TIME_SERIES_TYPE_SLICES))
    )

    # Every scored series must have a metadata row. A null here would drop the series out of every
    # per-type MLflow aggregate while still counting towards the overall mean, so two experiments
    # scored over the same population would not be comparable.
    missing_type = metrics_tall.filter(pl.col("time_series_type").is_null())["time_series_id"]
    if missing_type.len() > 0:
        raise ValueError(
            f"No metadata row for time_series_id(s) {sorted(set(missing_type.to_list()))}; "
            "materialise the time-series metadata for these series before scoring."
        )

    return Metrics.validate(metrics_tall, allow_superfluous_columns=True)

build_mlflow_aggregate_metrics(metrics_df)

Return a flat {metric_key: value} dict for mlflow.log_metrics.

Computes mean metric_value across all series ("all" aggregate), per time_series_type, and per horizon_slice. The key token is {metric_name} for metric_param="all" metrics and {metric_name}_{metric_param} for parametric ones (e.g. pinball_loss_p10, picp_p10_p90); parametric metrics are restricted to the _MLFLOW_LOGGED_PARAMETRIC headline subset. Key formats:

• "{token}__all" — overall aggregate (horizon_slice="all").
• "{token}__{type_slug}" — per-type aggregates (horizon_slice="all").
• "{token}__all__{horizon_slice}" — overall aggregate per lead-time band (e.g. "nmae__all__day_ahead"). Per-type sliced aggregates are deliberately not logged — that detail stays queryable in the forecast_metrics Delta table, as does the full 13-quantile / 6-band parametric detail.

Parameters:

Name	Type	Description	Default
metrics_df	DataFrame	Per-series Metrics rows with time_series_type populated.	required

Returns:

Type	Description
dict[str, float]	Flat dict of MLflow metric key → mean float value.

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

def build_mlflow_aggregate_metrics(
    metrics_df: pl.DataFrame,
) -> dict[str, float]:
    """Return a flat ``{metric_key: value}`` dict for ``mlflow.log_metrics``.

    Computes mean ``metric_value`` across all series (``"all"`` aggregate), per
    ``time_series_type``, and per ``horizon_slice``. The key token is ``{metric_name}`` for
    ``metric_param="all"`` metrics and ``{metric_name}_{metric_param}`` for parametric ones
    (e.g. ``pinball_loss_p10``, ``picp_p10_p90``); parametric metrics are restricted to the
    ``_MLFLOW_LOGGED_PARAMETRIC`` headline subset. Key formats:

    - ``"{token}__all"`` — overall aggregate (``horizon_slice="all"``).
    - ``"{token}__{type_slug}"`` — per-type aggregates (``horizon_slice="all"``).
    - ``"{token}__all__{horizon_slice}"`` — overall aggregate per lead-time band
      (e.g. ``"nmae__all__day_ahead"``). Per-type sliced aggregates are deliberately not
      logged — that detail stays queryable in the ``forecast_metrics`` Delta table, as does
      the full 13-quantile / 6-band parametric detail.

    Args:
        metrics_df: Per-series ``Metrics`` rows with ``time_series_type`` populated.

    Returns:
        Flat dict of MLflow metric key → mean float value.
    """
    logged = metrics_df.filter(_mlflow_logged_expr()).with_columns(
        metric_key=_metric_key_token_expr()
    )
    base = logged.filter(pl.col("horizon_slice") == "all")

    result: dict[str, float] = {}

    # Per-type aggregates. The column is `allow_missing` on `Metrics`, so guard its presence —
    # but never its nullability: `compute_metrics` raises rather than emit a null type.
    if "time_series_type" in base.columns:
        per_type = base.group_by(["metric_key", "time_series_type"]).agg(
            mean_value=pl.col("metric_value").mean()
        )
        for row in per_type.iter_rows(named=True):
            key = f"{row['metric_key']}__{_type_slug(str(row['time_series_type']))}"
            result[key] = float(row["mean_value"])

    # Overall "all" aggregate (includes all series regardless of type).
    overall = base.group_by("metric_key").agg(mean_value=pl.col("metric_value").mean())
    for row in overall.iter_rows(named=True):
        result[f"{row['metric_key']}__all"] = float(row["mean_value"])

    # Overall aggregate per lead-time band (includes all series regardless of type).
    per_slice = (
        logged.filter(pl.col("horizon_slice") != "all")
        .group_by(["metric_key", "horizon_slice"])
        .agg(mean_value=pl.col("metric_value").mean())
    )
    for row in per_slice.iter_rows(named=True):
        result[f"{row['metric_key']}__all__{row['horizon_slice']}"] = float(row["mean_value"])

    return result

enrich_metrics_rows(per_series_metrics, experiment_name, evaluation_scope, window_start, window_end, window_label, computed_at, mlflow_run_id)

Add scope and evaluation-window provenance columns to a per-series Metrics frame.

Called by the metrics Dagster asset after compute_metrics() returns, once the window bounds and MLflow run ID are known. Kept here so the enrichment logic is unit-testable without Dagster.

Parameters:

Name	Type	Description	Default
per_series_metrics	DataFrame[Metrics]	Frame produced by compute_metrics().	required
experiment_name	str	Experiment that produced these forecasts.	required
evaluation_scope	EvalScopeType	"leaderboard" or "ad_hoc".	required
window_start	datetime	Inclusive start of the evaluated valid_time window.	required
window_end	datetime	Inclusive end of the evaluated valid_time window.	required
window_label	str	Human-readable label (fold_id for leaderboard; "ad_hoc").	required
computed_at	datetime	UTC timestamp when this metric batch was computed.	required
mlflow_run_id	str | None	MLflow fold run ID; None for ad_hoc.	required

Returns:

Type	Description
DataFrame[Metrics]	A validated Metrics DataFrame with all columns fully populated.

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

def enrich_metrics_rows(
    per_series_metrics: pt.DataFrame[Metrics],
    experiment_name: str,
    evaluation_scope: EvalScopeType,
    window_start: datetime,
    window_end: datetime,
    window_label: str,
    computed_at: datetime,
    mlflow_run_id: str | None,
) -> pt.DataFrame[Metrics]:
    """Add scope and evaluation-window provenance columns to a per-series Metrics frame.

    Called by the ``metrics`` Dagster asset after ``compute_metrics()`` returns, once the
    window bounds and MLflow run ID are known. Kept here so the enrichment logic is
    unit-testable without Dagster.

    Args:
        per_series_metrics: Frame produced by ``compute_metrics()``.
        experiment_name: Experiment that produced these forecasts.
        evaluation_scope: ``"leaderboard"`` or ``"ad_hoc"``.
        window_start: Inclusive start of the evaluated ``valid_time`` window.
        window_end: Inclusive end of the evaluated ``valid_time`` window.
        window_label: Human-readable label (``fold_id`` for leaderboard; ``"ad_hoc"``).
        computed_at: UTC timestamp when this metric batch was computed.
        mlflow_run_id: MLflow fold run ID; ``None`` for ``ad_hoc``.

    Returns:
        A validated ``Metrics`` DataFrame with all columns fully populated.
    """
    return Metrics.validate(
        per_series_metrics.with_columns(
            experiment_name=pl.lit(experiment_name, dtype=pl.String),
            evaluation_scope=pl.lit(evaluation_scope).cast(pl.Enum(EVALUATION_SCOPES)),
            window_start=pl.lit(window_start).cast(UTC_DATETIME_DTYPE),
            window_end=pl.lit(window_end).cast(UTC_DATETIME_DTYPE),
            window_label=pl.lit(window_label, dtype=pl.String),
            computed_at=pl.lit(computed_at).cast(UTC_DATETIME_DTYPE),
            mlflow_run_id=pl.lit(mlflow_run_id, dtype=pl.String),
        )
    )

ml_core.cv_helpers

Pure, data-only helpers for the cross-validation Dagster assets.

Every function here is deliberately free of I/O (no Delta, MLflow, or Dagster imports) so it can be unit-tested in isolation. The CV asset bodies stay thin by delegating their logic here.

Attributes

CV_PARTITION_KEY_SEPARATOR = '__' module-attribute

Separator between experiment name and fold id in a CV partition key.

A double-underscore reduces collision risk with experiment names that contain single underscores.

Classes

Functions:

date_to_utc_datetime(d, *, end_of_day=False)

Return a tz-aware UTC datetime at the start (or inclusive end) of the given date.

Used to turn a fold's [start, end] calendar dates into the half-open-free, inclusive [start 00:00:00, end 23:59:59] UTC window that the training and validation data loads filter on (and that eligibility uses for val_end).

Parameters:

Name	Type	Description	Default
d	date	The calendar date.	required
end_of_day	bool	If True, return d at 23:59:59 (the inclusive end-of-day used by both the training/validation windows and val_end); otherwise 00:00:00.	False

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

def date_to_utc_datetime(d: date, *, end_of_day: bool = False) -> datetime:
    """Return a tz-aware UTC datetime at the start (or inclusive end) of the given date.

    Used to turn a fold's ``[start, end]`` calendar dates into the half-open-free, inclusive
    ``[start 00:00:00, end 23:59:59]`` UTC window that the training and validation data loads filter
    on (and that eligibility uses for ``val_end``).

    Args:
        d: The calendar date.
        end_of_day: If True, return ``d`` at ``23:59:59`` (the inclusive end-of-day used by both
            the training/validation windows and ``val_end``); otherwise ``00:00:00``.
    """
    clock = time(23, 59, 59) if end_of_day else time(0, 0, 0)
    return datetime.combine(d, clock, tzinfo=UTC)

eligible_time_series_ids(coverage, fold, min_training_months)

Return the sorted time_series_ids eligible for a fold, from data coverage alone.

A time series is eligible when it has at least min_training_months of observations before the fold's val_start and observations through the fold's val_end.

Eligibility is a function of the data only — it does not depend on any model or experiment config — so every experiment evaluates a fold on the identical population, which is what makes leaderboard comparisons fair.

Parameters:

Name	Type	Description	Default
coverage	DataFrame	One row per time series with the columns time_series_id, first_time, and last_time (the min and max observation timestamps, tz-aware UTC).	required
fold	CvFoldConfig	The CV fold whose eligibility is being computed.	required
min_training_months	int	Minimum months of pre-val_start history required.	required

Returns:

Type	Description
list[int]	Sorted list of eligible time_series_id values.

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

def eligible_time_series_ids(
    coverage: pl.DataFrame,
    fold: CvFoldConfig,
    min_training_months: int,
) -> list[int]:
    """Return the sorted ``time_series_id``s eligible for a fold, from data coverage alone.

    A time series is eligible when it has at least ``min_training_months`` of observations
    before the fold's ``val_start`` **and** observations through the fold's ``val_end``.

    Eligibility is a function of the **data only** — it does not depend on any model or
    experiment config — so every experiment evaluates a fold on the identical population,
    which is what makes leaderboard comparisons fair.

    Args:
        coverage: One row per time series with the columns ``time_series_id``, ``first_time``,
            and ``last_time`` (the min and max observation timestamps, tz-aware UTC).
        fold: The CV fold whose eligibility is being computed.
        min_training_months: Minimum months of pre-``val_start`` history required.

    Returns:
        Sorted list of eligible ``time_series_id`` values.
    """
    earliest_required_first_time = date_to_utc_datetime(
        _subtract_months(fold.val_start, min_training_months)
    )
    val_end_dt = date_to_utc_datetime(fold.val_end, end_of_day=True)
    eligible = coverage.filter(
        (pl.col("first_time") <= earliest_required_first_time) & (pl.col("last_time") >= val_end_dt)
    )
    return sorted(eligible["time_series_id"].to_list())

parse_cv_partition_key(partition_key)

Return (experiment_name, fold_id) from a CV partition key.

Partition key format: "{experiment_name}__{fold_id}". The separator is a double-underscore to reduce collision risk with experiment names that contain single underscores. Splitting from the right keeps __ inside the experiment name intact.

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

def parse_cv_partition_key(partition_key: str) -> tuple[str, str]:
    """Return ``(experiment_name, fold_id)`` from a CV partition key.

    Partition key format: ``"{experiment_name}__{fold_id}"``. The separator is a
    double-underscore to reduce collision risk with experiment names that contain single
    underscores. Splitting from the right keeps ``__`` inside the experiment name intact.
    """
    experiment_name, fold_id = partition_key.rsplit(CV_PARTITION_KEY_SEPARATOR, maxsplit=1)
    return experiment_name, fold_id

flatten_config(config, prefix='')

Flatten a (possibly nested) config into dotted-key string params for MLflow.

MLflow params are scalar strings, so nested dicts are flattened with dotted keys and every leaf value is stringified. A pydantic BaseModel is first dumped to a plain dict.

Parameters:

Name	Type	Description	Default
config	BaseModel | dict[str, Any]	A pydantic model or (possibly nested) dict to flatten.	required
prefix	str	Internal recursion prefix; callers should not set it.	''

Returns:

Type	Description
dict[str, str]	A flat {dotted_key: str_value} mapping suitable for mlflow.log_params.

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

def flatten_config(config: BaseModel | dict[str, Any], prefix: str = "") -> dict[str, str]:
    """Flatten a (possibly nested) config into dotted-key string params for MLflow.

    MLflow params are scalar strings, so nested dicts are flattened with dotted keys and every
    leaf value is stringified. A pydantic ``BaseModel`` is first dumped to a plain dict.

    Args:
        config: A pydantic model or (possibly nested) dict to flatten.
        prefix: Internal recursion prefix; callers should not set it.

    Returns:
        A flat ``{dotted_key: str_value}`` mapping suitable for ``mlflow.log_params``.
    """
    if isinstance(config, BaseModel):
        config = config.model_dump(mode="json")
    flat: dict[str, str] = {}
    for key, value in config.items():
        full_key = f"{prefix}{key}"
        if isinstance(value, dict):
            flat.update(flatten_config(value, prefix=f"{full_key}."))
        else:
            flat[full_key] = str(value)
    return flat

ml_core.mlflow_runs

Idempotent MLflow run-resolution helpers, used by every CV/registration code path.

Single concern: resolving MLflow runs by tag. The CV assets (trained_cv_model, cv_power_forecasts, metrics) and register_experiment_job run in separate processes (and, on retries, at separate times), so a live Run handle cannot be passed between them. Instead, every code path discovers or creates the run it needs by tag and resumes it by ID.

Each helper returns an ID string, never an open run handle. The caller wraps the returned ID in with mlflow.start_run(run_id=...) to log and close the run within its own process. Because lookup is by tag, every helper is safe to call from any process and idempotent under Dagster retries (a re-run resumes the same run rather than duplicating it).

The tracking URI is set by the caller (mlflow.set_tracking_uri(...)); these helpers simply use mlflow / MlflowClient, which honour whatever URI is in effect. This is what lets the helpers run against a file-based MLflow in tests with no server.

Attributes

CV_PARENT_RUN_NAME = 'cv_summary' module-attribute

Run name of an experiment's parent run, which carries the aggregate leaderboard metrics.

Classes

PromotableRun dataclass

One MLflow fold run (cv_role=fold) — a valid promoted_model promotion candidate.

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

@dataclass(frozen=True)
class PromotableRun:
    """One MLflow fold run (``cv_role=fold``) — a valid ``promoted_model`` promotion candidate."""

    run_id: str
    experiment_name: str
    fold_id: str
    start_time: datetime

Attributes

run_id instance-attribute

experiment_name instance-attribute

fold_id instance-attribute

start_time instance-attribute

Methods:

__init__(run_id, experiment_name, fold_id, start_time)

Functions:

load_experiment_forecaster(experiment_name)

Reconstruct the forecaster class + resolved config from an experiment's MLflow tags.

register_experiment stamps the config JSON and the forecaster class's fully-qualified import path (forecaster_target) on the experiment. The config JSON carries no class identity of its own, so the BaseForecasterConfig subclass to deserialise it into is reached through the forecaster's CONFIG_CLASS — the same class its load uses. The caller is responsible for setting the tracking URI (mlflow.set_tracking_uri) beforehand.

Parameters:

Name	Type	Description	Default
experiment_name	str	The MLflow experiment name (also the partition-key prefix).	required

Returns:

Type	Description
type[BaseForecaster]	A (forecaster_cls, forecaster_config) tuple — the same pair
BaseForecasterConfig	register_experiment resolved, reconstructed from the stored tags.

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

def load_experiment_forecaster(
    experiment_name: str,
) -> tuple[type[BaseForecaster], BaseForecasterConfig]:
    """Reconstruct the forecaster class + resolved config from an experiment's MLflow tags.

    ``register_experiment`` stamps the config JSON and the forecaster class's fully-qualified
    import path (``forecaster_target``) on the experiment. The config JSON carries no class
    identity of its own, so the ``BaseForecasterConfig`` subclass to deserialise it into is reached
    through the forecaster's ``CONFIG_CLASS`` — the same class its ``load`` uses. The caller is
    responsible for setting the tracking URI (``mlflow.set_tracking_uri``) beforehand.

    Args:
        experiment_name: The MLflow experiment name (also the partition-key prefix).

    Returns:
        A ``(forecaster_cls, forecaster_config)`` tuple — the same pair
        ``register_experiment`` resolved, reconstructed from the stored tags.
    """
    experiment = mlflow.get_experiment_by_name(experiment_name)
    if experiment is None:
        raise ValueError(f"No MLflow experiment named {experiment_name!r}.")
    tags = experiment.tags
    forecaster_cls = cast(type[BaseForecaster], import_class(tags["forecaster_target"]))
    forecaster_config = forecaster_cls.CONFIG_CLASS.model_validate_json(tags["config"])
    return forecaster_cls, forecaster_config

get_or_create_experiment(experiment_name)

Return the MLflow experiment id for experiment_name, creating it if absent.

This is the self-healing fallback path; the canonical creator is register_experiment_job, which also stamps the experiment's config/description tags. Created here untagged so a stray asset call cannot race the job into a tagless, half-registered experiment.

Parameters:

Name	Type	Description	Default
experiment_name	str	Human-readable, unique experiment name.	required

Returns:

Type	Description
str	The MLflow experiment id.

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

def get_or_create_experiment(experiment_name: str) -> str:
    """Return the MLflow experiment id for ``experiment_name``, creating it if absent.

    This is the self-healing fallback path; the canonical creator is
    ``register_experiment_job``, which also stamps the experiment's ``config``/``description``
    tags. Created here untagged so a stray asset call cannot race the job into a tagless,
    half-registered experiment.

    Args:
        experiment_name: Human-readable, unique experiment name.

    Returns:
        The MLflow experiment id.
    """
    experiment = mlflow.get_experiment_by_name(experiment_name)
    if experiment is not None:
        return experiment.experiment_id
    return mlflow.create_experiment(experiment_name)

get_or_create_parent_run(experiment_id)

Return the experiment's parent run id (tag cv_role=parent), creating it if absent.

The parent run holds the experiment's aggregate (leaderboard) metrics and the flattened config params. Resolved by tag so any process finds the same run.

Parameters:

Name	Type	Description	Default
experiment_id	str	The MLflow experiment id.	required

Returns:

Type	Description
str	The parent run id.

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

def get_or_create_parent_run(experiment_id: str) -> str:
    """Return the experiment's parent run id (tag ``cv_role=parent``), creating it if absent.

    The parent run holds the experiment's aggregate (leaderboard) metrics and the flattened
    config params. Resolved by tag so any process finds the same run.

    Args:
        experiment_id: The MLflow experiment id.

    Returns:
        The parent run id.
    """
    client = MlflowClient()
    runs = client.search_runs(
        experiment_ids=[experiment_id],
        filter_string="tags.cv_role = 'parent'",
        max_results=1,
    )
    if runs:
        return runs[0].info.run_id
    with mlflow.start_run(
        experiment_id=experiment_id,
        run_name=CV_PARENT_RUN_NAME,
        tags={"cv_role": "parent"},
    ) as run:
        return run.info.run_id

get_or_create_fold_run(experiment_id, parent_run_id, fold_id)

Return the fold's child run id (tags cv_role=fold, fold_id=...), creating it if absent.

The fold run holds that fold's per-fold tags and metrics (never MLflow params — a fold run is reused across every re-materialisation of its partition, and params are write-once, so nothing that can legitimately change between materialisations may be logged as one; see trained_cv_model in defs/cv_assets.py). It is created nested under the experiment's parent run, so the MLflow UI groups folds beneath cv_summary. Resolved by (cv_role, fold_id) so any process — and any Dagster retry of the fold — finds the same run.

Parameters:

Name	Type	Description	Default
experiment_id	str	The MLflow experiment id.	required
parent_run_id	str	The experiment's parent run id (from get_or_create_parent_run).	required
fold_id	str	The fold identifier (e.g. "2022").	required

Returns:

Type	Description
str	The fold run id.

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

def get_or_create_fold_run(experiment_id: str, parent_run_id: str, fold_id: str) -> str:
    """Return the fold's child run id (tags ``cv_role=fold, fold_id=...``), creating it if absent.

    The fold run holds that fold's per-fold tags and metrics (never MLflow params — a fold run is
    reused across every re-materialisation of its partition, and params are write-once, so
    nothing that can legitimately change between materialisations may be logged as one; see
    ``trained_cv_model`` in ``defs/cv_assets.py``). It is created **nested** under the
    experiment's parent run, so the MLflow UI groups folds beneath ``cv_summary``. Resolved by
    ``(cv_role, fold_id)`` so any process — and any Dagster retry of the fold — finds the same
    run.

    Args:
        experiment_id: The MLflow experiment id.
        parent_run_id: The experiment's parent run id (from ``get_or_create_parent_run``).
        fold_id: The fold identifier (e.g. ``"2022"``).

    Returns:
        The fold run id.
    """
    client = MlflowClient()
    runs = client.search_runs(
        experiment_ids=[experiment_id],
        filter_string=f"tags.cv_role = 'fold' and tags.fold_id = '{fold_id}'",
        max_results=1,
    )
    if runs:
        return runs[0].info.run_id
    # Resume the parent so the new run nests beneath it (MLflow nests under the active run).
    # experiment_id must be passed explicitly: resuming a run does not switch the active
    # experiment, so without it the child would land in the default experiment ("0").
    with (
        mlflow.start_run(run_id=parent_run_id),
        mlflow.start_run(
            experiment_id=experiment_id,
            run_name=fold_id,
            nested=True,
            tags={"cv_role": "fold", "fold_id": fold_id},
        ) as fold_run,
    ):
        return fold_run.info.run_id

list_promotable_runs()

List every fold run (cv_role=fold) across all MLflow experiments, newest first.

A read-only convenience for the promotable_model_runs asset (defs/production_assets.py), which logs this as a metadata table in the Dagster UI so a promoted_model promotion candidate's run id can be copy-pasted into that asset's launchpad rather than retyped from memory. The champion is still picked by eye off the MLflow leaderboard; this only lists candidates. The caller is responsible for setting the tracking URI (mlflow.set_tracking_uri) beforehand.

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

def list_promotable_runs() -> list[PromotableRun]:
    """List every fold run (``cv_role=fold``) across all MLflow experiments, newest first.

    A read-only convenience for the ``promotable_model_runs`` asset
    (``defs/production_assets.py``), which logs this as a metadata table in the Dagster UI so a
    ``promoted_model`` promotion candidate's run id can be copy-pasted into that asset's
    launchpad rather than retyped from memory. The champion is still picked by eye off the MLflow
    leaderboard; this only lists candidates. The caller is responsible for setting the tracking
    URI (``mlflow.set_tracking_uri``) beforehand.
    """
    client = MlflowClient()
    runs = [
        PromotableRun(
            run_id=run.info.run_id,
            experiment_name=experiment.name,
            fold_id=run.data.tags.get("fold_id", "unknown"),
            start_time=datetime.fromtimestamp(run.info.start_time / 1000, tz=UTC),
        )
        for experiment in client.search_experiments()
        for run in client.search_runs(
            experiment_ids=[experiment.experiment_id],
            filter_string="tags.cv_role = 'fold'",
            max_results=1000,
        )
    ]
    return sorted(runs, key=lambda run: run.start_time, reverse=True)

ml_core.production_helpers

Pure, IO-light helpers for production (live) inference.

Every function here is unit-testable in isolation: the two data-shaping helpers (select_nwp_init_time, build_live_power_frame) take power_fcst_init_time as an explicit parameter rather than calling datetime.now() internally, so a test can pass any fixed time and get a deterministic result; the two disk/MLflow helpers (load_forecaster_from_dir, fetch_model_artifacts) do the IO and check that the saved model is one this code can still build a config for and parse the features of. The live_forecasts and promoted_model Dagster assets (src/nged_substation_forecast/defs/production_assets.py) stay thin shells over these.

Attributes

AvailabilityModeType = Literal['live', 'replay'] module-attribute

Which NWP-availability rule select_nwp_init_time applies.

• "live": the scheduled path. No modelled publication delay — the Delta table only contains runs that have genuinely been published, so the cutoff is power_fcst_init_time itself.
• "replay": re-running a past slot. The cutoff is power_fcst_init_time - nwp_publication_delay_hours, reconstructing what was actually available at that historical power_fcst_init_time (without the delay we would leak runs that only landed afterwards).

Classes

Functions:

select_nwp_init_time(available_init_times, *, power_fcst_init_time, availability_mode, nwp_publication_delay_hours=NWP_PUBLICATION_DELAY_HOURS)

Return the freshest NWP init_time available at power_fcst_init_time.

Which runs count as available depends on availability_mode.

Parameters:

Name	Type	Description	Default
available_init_times	Sequence[datetime]	The init_times genuinely present in the NWP Delta table (e.g. from DeltaTable(...).partitions()).	required
power_fcst_init_time	datetime	The scheduled forecast time (the partition's window end).	required
availability_mode	AvailabilityModeType	"live" uses cutoff power_fcst_init_time; "replay" uses cutoff power_fcst_init_time - nwp_publication_delay_hours.	required
nwp_publication_delay_hours	int	Only used in "replay" mode.	NWP_PUBLICATION_DELAY_HOURS

Returns:

Type	Description
datetime	The freshest init_time that is <= the cutoff.

Raises:

Type	Description
ValueError	If no available init_time qualifies.

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

def select_nwp_init_time(
    available_init_times: Sequence[datetime],
    *,
    power_fcst_init_time: datetime,
    availability_mode: AvailabilityModeType,
    nwp_publication_delay_hours: int = NWP_PUBLICATION_DELAY_HOURS,
) -> datetime:
    """Return the freshest NWP ``init_time`` available at ``power_fcst_init_time``.

    Which runs count as available depends on ``availability_mode``.

    Args:
        available_init_times: The ``init_time``s genuinely present in the NWP Delta table
            (e.g. from ``DeltaTable(...).partitions()``).
        power_fcst_init_time: The scheduled forecast time (the partition's window end).
        availability_mode: ``"live"`` uses cutoff ``power_fcst_init_time``; ``"replay"`` uses
            cutoff ``power_fcst_init_time - nwp_publication_delay_hours``.
        nwp_publication_delay_hours: Only used in ``"replay"`` mode.

    Returns:
        The freshest ``init_time`` that is ``<=`` the cutoff.

    Raises:
        ValueError: If no available ``init_time`` qualifies.
    """
    cutoff = (
        power_fcst_init_time
        if availability_mode == "live"
        else power_fcst_init_time - timedelta(hours=nwp_publication_delay_hours)
    )
    qualifying = [init_time for init_time in available_init_times if init_time <= cutoff]
    if not qualifying:
        raise ValueError(
            f"No NWP run available at or before cutoff {cutoff.isoformat()} "
            f"(power_fcst_init_time={power_fcst_init_time.isoformat()}, "
            f"availability_mode={availability_mode!r}). Available init times: "
            f"{sorted(available_init_times)}"
        )
    return max(qualifying)

build_live_power_frame(observed_power, time_series_ids, *, power_fcst_init_time, history, horizon)

Build a dense half-hourly (time_series_id, time) spine for live inference.

Needed because ml_core.features._nwp._join_nwp_single_run is power-centric — with no future power rows a live run would emit zero forecast rows. Left-joins observed power onto a spine covering (power_fcst_init_time - history, power_fcst_init_time + horizon] for every requested time_series_id, so rows beyond the last observation are present with power = null. Also harmless for replay (future observations already exist there; _nullify_leaky_lags prevents lag leakage regardless).

Parameters:

Name	Type	Description	Default
observed_power	LazyFrame[PowerTimeSeries]	Lazy observed power, one row per (time_series_id, time).	required
time_series_ids	list[int]	The series to build a spine for (typically forecaster.trained_time_series_ids).	required
power_fcst_init_time	datetime	The forecast init time. The spine's window is anchored on this.	required
history	timedelta	How far before power_fcst_init_time the spine extends (exclusive) — must cover the longest power lag feature the model uses.	required
horizon	timedelta	How far after power_fcst_init_time the spine extends (inclusive) — the forecast horizon.	required

Returns:

Type	Description
LazyFrame[PowerTimeSeries]	A lazy PowerTimeSeries frame with one row per (time_series_id, time) on the
LazyFrame[PowerTimeSeries]	half-hourly grid, observed values joined in, future/missing rows null.

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

def build_live_power_frame(
    observed_power: pt.LazyFrame[PowerTimeSeries],
    time_series_ids: list[int],
    *,
    power_fcst_init_time: datetime,
    history: timedelta,
    horizon: timedelta,
) -> pt.LazyFrame[PowerTimeSeries]:
    """Build a dense half-hourly ``(time_series_id, time)`` spine for live inference.

    Needed because ``ml_core.features._nwp._join_nwp_single_run`` is power-centric — with no
    future power rows a live run would emit zero forecast rows. Left-joins observed power onto
    a spine covering ``(power_fcst_init_time - history, power_fcst_init_time + horizon]`` for
    every requested ``time_series_id``, so rows beyond the last observation are present with
    ``power = null``. Also harmless for replay (future observations already exist there;
    ``_nullify_leaky_lags`` prevents lag leakage regardless).

    Args:
        observed_power: Lazy observed power, one row per ``(time_series_id, time)``.
        time_series_ids: The series to build a spine for (typically
            ``forecaster.trained_time_series_ids``).
        power_fcst_init_time: The forecast init time. The spine's window is anchored on this.
        history: How far before ``power_fcst_init_time`` the spine extends (exclusive) — must
            cover the longest power lag feature the model uses.
        horizon: How far after ``power_fcst_init_time`` the spine extends (inclusive) — the
            forecast horizon.

    Returns:
        A lazy ``PowerTimeSeries`` frame with one row per ``(time_series_id, time)`` on the
        half-hourly grid, observed values joined in, future/missing rows null.
    """
    grid_start = power_fcst_init_time - history + timedelta(minutes=30)
    grid_end = power_fcst_init_time + horizon
    grid_times = pl.datetime_range(
        grid_start, grid_end, interval="30m", time_zone="UTC", eager=True
    )

    ids_lf = pl.LazyFrame({"time_series_id": time_series_ids}, schema={"time_series_id": pl.Int32})
    times_lf = pl.LazyFrame({"time": grid_times}, schema={"time": UTC_DATETIME_DTYPE})
    spine = ids_lf.join(times_lf, how="cross")

    # Strip the Patito subclass before joining (see the `polars-patito-gotchas` skill).
    power_plain = pl.LazyFrame._from_pyldf(observed_power._ldf)
    dense = spine.join(power_plain, on=["time_series_id", "time"], how="left").sort(
        ["time_series_id", "time"]
    )
    return pt.LazyFrame.from_existing(dense).set_model(PowerTimeSeries)

load_forecaster_from_dir(path)

Load the production model from a plain disk directory (no MLflow at inference time).

Reads meta.json and resolves model_class via contracts.config_schemas.import_class (the same mechanism ml_core.mlflow_runs.load_experiment_forecaster uses), then calls the concrete subclass's load(path).

The forecaster returned is one this code can actually serve, not merely one it could deserialise: a config it cannot rebuild, or a feature vocabulary it cannot parse, is rejected here rather than partway through a live tick's feature engineering.

Parameters:

Name	Type	Description	Default
path	Path	Directory previously populated by fetch_model_artifacts (the promoted_model asset's output).	required

Returns:

Type	Description
BaseForecaster	The reconstructed, trained forecaster.

Raises:

Type	Description
FileNotFoundError	path or its meta.json does not exist — materialise the promoted_model asset first.
ValueError	This code cannot serve the saved model — see _check_meta_is_servable. Promotion applies the same check, so this fires only when the code changed after the champion was promoted.

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

def load_forecaster_from_dir(path: Path) -> BaseForecaster:
    """Load the production model from a plain disk directory (no MLflow at inference time).

    Reads ``meta.json`` and resolves ``model_class`` via ``contracts.config_schemas.import_class``
    (the same mechanism ``ml_core.mlflow_runs.load_experiment_forecaster`` uses), then calls the
    concrete subclass's ``load(path)``.

    The forecaster returned is one this code can actually serve, not merely one it could
    deserialise: a config it cannot rebuild, or a feature vocabulary it cannot parse, is rejected
    here rather than partway through a live tick's feature engineering.

    Args:
        path: Directory previously populated by ``fetch_model_artifacts`` (the
            ``promoted_model`` asset's output).

    Returns:
        The reconstructed, trained forecaster.

    Raises:
        FileNotFoundError: ``path`` or its ``meta.json`` does not exist — materialise the
            ``promoted_model`` asset first.
        ValueError: This code cannot serve the saved model — see ``_check_meta_is_servable``.
            Promotion applies the same check, so this fires only when the code changed after the
            champion was promoted.
    """
    meta_path = path / "meta.json"
    if not meta_path.exists():
        raise FileNotFoundError(
            f"No production model found at {path} (missing meta.json). Materialise the "
            "`promoted_model` asset first."
        )
    meta = json.loads(meta_path.read_text())
    # Before `load`, not after: `load` reads every serialised sub-model off disk, which is a lot of
    # IO to do on the way to rejecting the directory over fields already parsed here.
    forecaster_cls = _check_meta_is_servable(meta=meta, source=str(path))
    return forecaster_cls.load(path)

fetch_model_artifacts(run_id, dest)

Download and unpack an MLflow run's saved model into dest, replacing it atomically.

Downloads and unpacks into a temporary directory first, so a failed or interrupted download never touches dest — only a fully-downloaded model is moved into place (via rmtree + move). dest is local disk by convention — Settings.production_model_path derives from local_artifacts_path, though nothing enforces that — so unlike the Delta tables this is a directory of many files with no commit protocol over it, and a part-written one would be served. The run holds the model as a single archive artifact (ml_core.base_forecaster._MLFLOW_MODEL_ARTIFACT), so dest gets exactly the files the last save_to_mlflow wrote and can never inherit a stale file from an earlier, larger model.

The downloaded model's saved config is checked against the running code before the swap, reading the staged meta.json rather than loading the model, so a model this code cannot serve is refused while the previous champion stays in dest and keeps serving. Reading the JSON is deliberate: it applies the same validation the subclass's load would, without pulling every booster into memory to do it.

Also writes a promotion.json ({"mlflow_run_id", "promoted_at"}) into dest for provenance; a BaseForecaster.load implementation reads its own population from its saved record (e.g. XGBoostForecaster from meta.json's trained_time_series_ids), never from a directory listing, so this extra file is harmless.

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 (via BaseForecaster.save_to_mlflow).	required
dest	Path	Directory to populate — typically Settings.production_model_path.	required

Raises:

Type	Description
MlflowException	run_id names a run holding no model archive — most often a mistyped or stale run id, since a run that trained a model has one. Raised by ml_core.base_forecaster._download_and_unpack_model, before dest is touched.
ValueError	The run holds no meta.json, or this code cannot serve the model it describes — see _check_meta_is_servable — or it carries no readable frozen metadata — see _check_trained_metadata_is_readable. Re-train against the current code and promote that run instead.

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

def fetch_model_artifacts(run_id: str, dest: Path) -> None:
    """Download and unpack an MLflow run's saved model into ``dest``, replacing it atomically.

    Downloads and unpacks into a temporary directory first, so a failed or interrupted download
    never touches ``dest`` — only a fully-downloaded model is moved into place (via ``rmtree`` +
    ``move``). ``dest`` is local disk by convention — ``Settings.production_model_path`` derives
    from ``local_artifacts_path``, though nothing enforces that — so unlike the Delta tables this is
    a directory of many files with no commit protocol over it, and a part-written one would be
    served. The run holds the model as a single archive artifact
    (``ml_core.base_forecaster._MLFLOW_MODEL_ARTIFACT``), so ``dest`` gets exactly the files the
    last ``save_to_mlflow`` wrote and can never inherit a stale file from an earlier, larger model.

    The downloaded model's saved config is checked against the running code *before* the swap,
    reading the staged ``meta.json`` rather than loading the model, so a model this code cannot
    serve is refused while the previous champion stays in ``dest`` and keeps serving. Reading the
    JSON is deliberate: it applies the same validation the subclass's ``load`` would, without
    pulling every booster into memory to do it.

    Also writes a ``promotion.json`` (``{"mlflow_run_id", "promoted_at"}``) into ``dest`` for
    provenance; a ``BaseForecaster.load`` implementation reads its own population from its saved
    record (e.g. ``XGBoostForecaster`` from ``meta.json``'s ``trained_time_series_ids``), never
    from a directory listing, so this extra file is harmless.

    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 (via ``BaseForecaster.save_to_mlflow``).
        dest: Directory to populate — typically ``Settings.production_model_path``.

    Raises:
        MlflowException: ``run_id`` names a run holding no model archive — most often a mistyped
            or stale run id, since a run that trained a model has one. Raised by
            ``ml_core.base_forecaster._download_and_unpack_model``, before ``dest`` is touched.
        ValueError: The run holds no ``meta.json``, or this code cannot serve the model it
            describes — see ``_check_meta_is_servable`` — or it carries no readable frozen
            metadata — see ``_check_trained_metadata_is_readable``. Re-train against the current
            code and promote that run instead.
    """
    with tempfile.TemporaryDirectory() as tmp_dir:
        downloaded_dir = _download_and_unpack_model(
            run_id=run_id,
            work_dir=Path(tmp_dir),
            remedy="check the run id, and pick one whose training completed.",
        )
        meta_path = downloaded_dir / "meta.json"
        if not meta_path.exists():
            raise ValueError(
                f"The model saved under run {run_id} has no meta.json, so no forecaster here can "
                "load it. Re-train against the current code and promote that run (see "
                "BaseForecaster.save)."
            )
        meta = json.loads(meta_path.read_text())
        _check_meta_is_servable(meta=meta, source=f"run {run_id}")
        _check_trained_metadata_is_readable(model_dir=downloaded_dir, run_id=run_id)

        promotion = {
            "mlflow_run_id": run_id,
            "promoted_at": datetime.now(UTC).isoformat(),
        }
        (downloaded_dir / "promotion.json").write_text(json.dumps(promotion))

        if dest.exists():
            shutil.rmtree(dest)
        dest.parent.mkdir(parents=True, exist_ok=True)
        shutil.move(str(downloaded_dir), str(dest))

ml_core.repro

Reproducibility provenance: the git SHA and Delta-table versions behind an MLflow run.

Answers "exactly which code and which data produced this?" for any MLflow run. The git SHA pins the code — stamped explicitly because MLflow's mlflow.source.git.commit auto-detection needs gitpython installed and the working directory inside the repo, neither of which holds in a production container. Each Delta table's version() pins the data: Delta Lake time travel makes data versioning one integer per table, so a run can later be replayed with pl.scan_delta(path, version=N) after git checkout {sha}.

Every function here is deliberately non-raising: provenance is metadata, and a missing .git directory (containers) or an absent Delta table must never fail the surrounding training or forecasting run. Such failures degrade to the sentinels "unknown" / "absent".

provenance_tags stage-prefixes its keys (register_, train_, predict_, metrics_) because a single MLflow fold run is written by three separate assets — trained_cv_model, cv_power_forecasts and metrics — each potentially on a different code revision and data state. Un-prefixed keys would clobber one another; the prefix preserves all three provenance snapshots side by side.

Attributes

logger = logging.getLogger(__name__) module-attribute

MlflowTags = dict[str, str] module-attribute

A {tag_key: tag_value} mapping ready to hand to mlflow.set_tags.

StageType = Literal['register', 'train', 'predict', 'metrics'] module-attribute

The assets that stamp provenance; each value becomes a tag-key prefix (see provenance_tags). Add a new member when a new asset starts stamping.

TableNameType = Literal['power_time_series', 'nwp_data', 'eligible_time_series', 'power_forecasts', 'effective_capacity'] module-attribute

Logical names of the Delta tables whose versions get stamped — the keys of a delta_paths mapping. Add a new member when a stage starts reading (and stamping) another table.

UNKNOWN = 'unknown' module-attribute

Sentinel git SHA / dirty flag returned when no git repository is reachable (e.g. a container).

ABSENT = 'absent' module-attribute

Sentinel Delta version returned when a table does not exist or cannot be read.

Classes

Functions:

get_git_info(cwd=None)

Return {"git_sha", "git_dirty"} for the current checkout, never raising.

git_dirty is "true" when the working tree has uncommitted changes, "false" when clean. When the SHA cannot be read (no .git, no git binary, a timeout) both values are UNKNOWN; when the SHA is read but the dirty check fails, the SHA is kept and only git_dirty degrades to UNKNOWN — a good SHA is never discarded.

Parameters:

Name	Type	Description	Default
cwd	Path | None	Directory the git commands run from. Defaults to this module's directory (_GIT_CWD) — inside the repo for an editable/workspace install, so the SHA is captured regardless of the process's working directory. Overridable for testing.	None

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

def get_git_info(cwd: Path | None = None) -> MlflowTags:
    """Return ``{"git_sha", "git_dirty"}`` for the current checkout, never raising.

    ``git_dirty`` is ``"true"`` when the working tree has uncommitted changes, ``"false"`` when
    clean. When the SHA cannot be read (no ``.git``, no ``git`` binary, a timeout) both values are
    ``UNKNOWN``; when the SHA is read but the dirty check fails, the SHA is kept and only
    ``git_dirty`` degrades to ``UNKNOWN`` — a good SHA is never discarded.

    Args:
        cwd: Directory the ``git`` commands run from. Defaults to this module's directory
            (``_GIT_CWD``) — inside the repo for an editable/workspace install, so the SHA is
            captured regardless of the process's working directory. Overridable for testing.
    """
    run_from = cwd if cwd is not None else _GIT_CWD

    def _git(*args: str) -> str:
        return subprocess.run(
            ["git", *args],
            cwd=run_from,
            capture_output=True,
            text=True,
            errors="replace",  # non-UTF-8 bytes (e.g. an odd filename) must not raise on decode.
            check=True,
            timeout=_GIT_TIMEOUT_S,
        ).stdout

    try:
        sha = _git("rev-parse", "HEAD").strip()
    except Exception:  # noqa: BLE001 — provenance must never fail the surrounding run.
        return {"git_sha": UNKNOWN, "git_dirty": UNKNOWN}
    try:
        porcelain = _git("status", "--porcelain")
    except Exception:  # noqa: BLE001 — keep the SHA we already have; only dirtiness is unknown.
        return {"git_sha": sha, "git_dirty": UNKNOWN}
    return {"git_sha": sha, "git_dirty": "true" if porcelain.strip() else "false"}

get_delta_versions(paths, storage_options=None)

Return {f"delta_version__{name}": str(version)} for each named Delta table.

Never raises.

Parameters:

Name	Type	Description	Default
paths	dict[TableNameType, str]	{logical_name: table_uri}. The URI may be local or a remote object-store URI.	required
storage_options	ObjectStoreOptions | None	object-store options for a remote URI; None/empty for local. Widened to the plain dict delta-rs expects at the call boundary.	None

Returns:

Type	Description
MlflowTags	One entry per input path. A table that does not exist (or cannot be read) maps to
MlflowTags	ABSENT rather than raising, so provenance capture never fails the calling run.

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

def get_delta_versions(
    paths: dict[TableNameType, str], storage_options: ObjectStoreOptions | None = None
) -> MlflowTags:
    """Return ``{f"delta_version__{name}": str(version)}`` for each named Delta table.

    Never raises.

    Args:
        paths: ``{logical_name: table_uri}``. The URI may be local or a remote object-store URI.
        storage_options: object-store options for a remote URI; ``None``/empty for local. Widened
            to the plain dict delta-rs expects at the call boundary.

    Returns:
        One entry per input path. A table that does not exist (or cannot be read) maps to
        ``ABSENT`` rather than raising, so provenance capture never fails the calling run.
    """
    options = typeddict_to_dict(storage_options) or {}
    versions: MlflowTags = {}
    for name, path in paths.items():
        key = f"delta_version__{name}"
        try:
            if not DeltaTable.is_deltatable(path, storage_options=options):
                versions[key] = ABSENT
                continue
            versions[key] = str(DeltaTable(path, storage_options=options).version())
        except Exception:  # Provenance must never fail the surrounding run.
            logger.warning("Could not read Delta version for %r at %s", name, path, exc_info=True)
            versions[key] = ABSENT
    return versions

provenance_tags(stage, delta_paths=None, storage_options=None)

Build stage-prefixed MLflow tags stamping code (git) and data (Delta version) provenance.

Parameters:

Name	Type	Description	Default
stage	StageType	Prefix identifying the writing asset. Keeps the tags of assets that share one MLflow run from clobbering.	required
delta_paths	dict[TableNameType, str] | None	{logical_name: table_uri} for the Delta tables this stage reads; omit for a stage that reads no data (registration).	None
storage_options	ObjectStoreOptions | None	object-store options for remote table URIs; None/empty for local.	None

Returns:

Type	Description
MlflowTags	e.g. for stage="train": ``{"train_git_sha", "train_git_dirty",
MlflowTags	"train_delta_version__power_time_series", ...}``.

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

def provenance_tags(
    stage: StageType,
    delta_paths: dict[TableNameType, str] | None = None,
    storage_options: ObjectStoreOptions | None = None,
) -> MlflowTags:
    """Build stage-prefixed MLflow tags stamping code (git) and data (Delta version) provenance.

    Args:
        stage: Prefix identifying the writing asset. Keeps the tags of assets that share one MLflow
            run from clobbering.
        delta_paths: ``{logical_name: table_uri}`` for the Delta tables this stage reads; omit for
            a stage that reads no data (registration).
        storage_options: object-store options for remote table URIs; ``None``/empty for local.

    Returns:
        e.g. for ``stage="train"``: ``{"train_git_sha", "train_git_dirty",
        "train_delta_version__power_time_series", ...}``.
    """
    git = get_git_info()
    tags: MlflowTags = {f"{stage}_git_sha": git["git_sha"], f"{stage}_git_dirty": git["git_dirty"]}
    if delta_paths:
        for key, version in get_delta_versions(delta_paths, storage_options).items():
            tags[f"{stage}_{key}"] = version
    return tags