MLflow Integration#

Custom MLflow model flavor for Darts forecasting models. Supports saving, loading, and logging any Darts ForecastingModel (statistical, ML-based, and PyTorch-based) to MLflow, as well as automatic logging via autolog().

See the MLflow quickstart for an end-to-end walkthrough.

Here’s a quick start example
import os
import tempfile

import mlflow

import darts.metrics as metrics
from darts.datasets import AirPassengersDataset
from darts.models import LinearRegressionModel
from darts.utils.mlflow import autolog

# dummy temporary directory for local MLflow tracking;
# use a permanent location for real use cases
tmpdir = tempfile.mkdtemp()
mlflow_db = os.path.join(tmpdir, "mlflow.db")
artifact_root = os.path.join(tmpdir, "mlruns")

mlflow.set_tracking_uri(f"sqlite:///{mlflow_db}")
# SQLite stores run metadata; artifacts default to ./mlruns unless we set a location
mlflow.set_experiment(
    experiment_id=mlflow.create_experiment(
        "darts-quickstart",
        artifact_location=artifact_root,
    )
)

# load series and create train and val splits
series = AirPassengersDataset().load()
train, val = series[:-36], series[-36:]

# two models to compare with different lookback windows
model_lags, names = [12, 24], ["one-year", "two-years"]
horizon = 12

# activate autologging and try out the models
autolog()
for lags, name in zip(model_lags, names):
    with mlflow.start_run(run_name=name) as run:
        model = LinearRegressionModel(lags=lags)
        model.fit(train)  # autolog logs params and covariate metadata

        # log standalone metrics from manual predictions
        pred = model.predict(n=horizon)
        metrics.mae(val, pred)  # time-aggregated logged as scaler metric
        metrics.ae(val, pred)  # time-dependent logged as stepped metric

        # log backtest metrics from historical forecasts over val set
        bt = model.backtest(
            series=series,
            start=val.start_time(),
            retrain=False,
            forecast_horizon=horizon,
            reduction=None,  # `None` logs as stepped metric, scalar otherwise
            metric=[metrics.mae, metrics.rmse],
        )
autolog(disable=True)

# you can launch the MLflow UI with the command below from your terminal;
# then open the returned address for example in a browser (similar to http://localhost:5000)
print(f"mlflow ui --backend-store-uri {mlflow.get_tracking_uri()}")

When autolog() is enabled, the following functionalities emit detailed logs when inside an active MLflow run (e.g. within with mlflow.start_run():):

  • Calling ForecastingModel.fit():

    • Logs model creation parameters (model.model_params), both as MLflow params and as a model_params.json artifact.

    • Logs target series info and covariate usage information (past, future, and static covariates) as a series_info.json artifact.

    • Stores the trained model artifact when log_models=True (default: False), and sets darts.model_is_pretrained to True.

    • Logs per-epoch training and validation metrics for PyTorch-based models.

  • Calling ForecastingModel.historical_forecasts(retrain=True):

    • Logs the same model creation parameters and series_info.json as fit() (overwriting any prior fit() artifacts in the same run).

    • Stores an untrained model template (via untrained_model()) when log_models=True and no model was logged yet in the run. Sets the run tag and model metadata darts.model_is_pretrained to False so downstream code knows the artifact must be refit before inference. If a model was already logged (e.g. via fit() or manual log_model()), the existing artifact is kept unchanged.

  • Calling any Darts metric:

    • Logs the result of that metric call as an MLflow metric. More information in the notes below.

  • Calling ForecastingModel.backtest():

    • Logs all evaluation metrics under backtest_* keys. More information in the notes below.

Important

Cross-Model Run Comparability: Metric values are only comparable across runs when the evaluation settings match. Use the same evaluation time frame, forecast horizon, and evaluation start date for every backtest() / metric call you intend to compare.

Note

Metric Naming Convention: Logged metric keys follow the pattern {metric_name}{component}{quantile_or_label} (e.g., "mae_target0_q0_500"), where each part is included only when the corresponding axis is present:

  • metric_name - the metric function name, or the name metric keyword argument when provided (e.g., "mae").

  • component - the component name: e.g., "_target0" when component_reduction=None.

  • quantile_or_label:

    • the quantile label: e.g., "_q0.500" for quantile metrics with keyword argument q=[0.5]

    • the quantile interval label: e.g., "_qi0.800" for quantile interval metrics with keyword argument q_interval=[(0.1, 0.9)] (80% interval between quantiles 0.1 and 0.9).

    • the class label: e.g., "_label1" for classification metrics with keyword argument labels when label_reduction=None.

Backtest metrics are prefixed by "backtest_".

Note

Metric Logging and Display: Darts offers many ways of evaluating forecasting models. To offer the highest value, we adapt what is logged to MLflow based on the use case scenario.

  • Metric type:

    • Time-aggregated metrics (e.g. mae()): Logged as scaler values.

    • Per-time step metrics (e.g. ae() where time_reduction=None): Logged as stepped metrics (one value per step in the forecast horizon).

  • Single or Multi-series:

    • Single series: Logged as explained above.

    • Multiple (a list of) series: When the metric’s series_reduction=None, the logged metric is aggregated over all series using autolog’s agg_func. The detailed per-series metrics / backtest metrics are logged under a single metrics_per_series.json table run artifact.

  • Standalone or backtest metric:

    • Standalone metric: Logged as explained above.

    • Backtest metric: If backtest’s reduction is other than None, the windowed forecast metrics are aggregated and logged as explained above. If None, metrics are logged as stepped MLflow metrics.

      • Time-aggregated metrics: each step represents a specifc forecast (window) metric. Steps represent the historical forecast windows (0, 1, …, n_windows - 1).

      • Per-time step metrics: each step represents a step in the forecast horizon aggregated over all windows. Steps represent the steps in the forecast horizon (0, 1, …, horizon - 1).

  • Multi-series alignment (series_align autolog option):

    • "end" (default): shorter series skip early steps/windows so their last points overlap the tail of longer series.

    • "start": shorter series contribute at early steps/windows; longer series have fewer contributing series at tail steps.

  • Backtest aggregate (log_backtest_aggregate autolog option):

    • When enabled, logs an additional scalar per backtest metric key under backtest_agg_{metric_key} (e.g. backtest_agg_mae), computed as agg_func over all logged steps for that key.

When components are preserved (component_reduction=None), all series scored together must have the same number of components; names are taken from the first series.

darts.utils.mlflow.autolog(log_models=False, log_params=True, log_metrics=True, log_torch_metrics=True, agg_func=<function nanmean>, series_align='end', log_backtest_aggregate=False, disable=False, silent=False)[source]#

Enable (or disable) automatic MLflow logging for Darts.

For a detailed overview of logged params, metrics, and artifacts, see the detailed documentation.

Parameters:

  • log_models (bool) – If True, log a model artifact when calling fit() (trained model) or historical_forecasts(retrain=True) (untrained template via untrained_model()). Defaults to False. If a model was already logged in the run (including via manual log_model()), autolog does not add another artifact.

  • log_params (bool) – If True (default), log model creation parameters when calling fit() or historical_forecasts().

  • log_metrics (bool) – If True (default), log the result of any Darts metric call made inside an active MLflow run, including standalone metric calls or via backtest.

  • log_torch_metrics (bool) – If True (default), enable mlflow.pytorch.autolog(log_models=False) around PyTorch-based model training to automatically log per-epoch training and validation metrics. Only effective for PyTorch-based models.

  • agg_func (Callable) – Function used to aggregate a metric’s per-series values into the single value logged for a list of series (e.g. np.nanmean, the default, or np.median). Called as agg_func(values) on a list of floats. Also used when log_backtest_aggregate=True to collapse all logged steps of a backtest metric into a single scalar.

  • series_align (Literal['start', 'end']) – How to align shorter series when multiple series differ in window or time length. "end" (default) aligns from the end so shorter series skip early steps/windows; "start" aligns from the start so shorter series contribute at early steps/windows.

  • log_backtest_aggregate (bool) – If True, log an additional scalar per backtest metric key under backtest_agg_{metric_key} (e.g. backtest_agg_mae), computed as agg_func over all logged steps for that key. Defaults to False.

  • disable (bool) – If True, restore the original fit() methods and stop autologging.

  • silent (bool) – If True (default False), suppress all event logging and warnings from MLflow during autologging.

Return type:

None

darts.utils.mlflow.get_default_conda_env()[source]#

Return a default conda environment dict for a Darts model.

Returns:

A conda environment specification dictionary.

Return type:

dict

darts.utils.mlflow.get_default_pip_requirements()[source]#

Return the default pip requirements for logging a Darts model.

Returns:

A list of pip requirement strings.

Return type:

list[str]

darts.utils.mlflow.load_model(model_uri, dst_path=None, **kwargs)[source]#

Load a Darts model from an MLflow model URI.

Parameters:

  • model_uri (str) – An MLflow model URI, e.g. "runs:/<run_id>/model", "models:/<name>/<version>", or a local file:///... path.

  • dst_path (str | None) – Optional local path for downloading remote artifacts.

  • **kwargs – Additional keyword arguments forwarded to the model’s load() method (e.g. map_location for a TorchForecastingModel).

Returns:

The loaded Darts forecasting model.

Return type:

ForecastingModel

darts.utils.mlflow.log_model(model, **kwargs)[source]#

Log a Darts model to the current MLflow run, using the Darts MLflow flavor.

This is a thin wrapper around mlflow.models.Model.log() that supplies the Darts flavor for saving/loading; every other argument is forwarded as-is. See the MLflow documentation for the full list of accepted parameters (e.g. name, registered_model_name, conda_env, pip_requirements, metadata, tags, …).

Parameters:

  • model (ForecastingModel) – A fitted Darts ForecastingModel instance.

  • **kwargs – Forwarded to mlflow.models.Model.log(). Use name to set the run-relative artifact path. artifact_path parameter is deprecated by MLflow and not exposed here.

Returns:

MLflow ModelInfo object containing model_uri, run_id, artifact_path, model_id, timestamps, and other metadata about the logged model.

Return type:

ModelInfo

Notes

signature and input_example are currently not supported, as they are used to support serving and input validation in the MLflow pyfunc flavor, which is not implemented for Darts models.

darts.utils.mlflow.save_model(model, path, conda_env=None, code_paths=None, mlflow_model=None, signature=None, input_example=None, pip_requirements=None, extra_pip_requirements=None, metadata=None)[source]#

Save a Darts forecasting model in MLflow format.

Produces an MLflow model directory at path containing:

  • The serialized Darts model (delegated to the model’s own save() method).

  • An MLmodel YAML file with flavor metadata.

  • conda.yaml and requirements.txt environment files.

Parameters:

  • model (ForecastingModel) – A fitted Darts ForecastingModel instance.

  • path (str) – Local filesystem path where the model directory will be created.

  • conda_env (dict | str | None) – A conda environment specification (dict or path to a conda.yaml). If None, a default environment is generated.

  • code_paths (list[str] | None) – A list of local filesystem paths to Python file dependencies (or directories containing file dependencies). These files are prepended to the system path when the model is loaded.

  • mlflow_model (Model | None) – Optional MLflow Model object to use for saving. When provided (typically by Model.log()), this model instance is used instead of creating a new one.

  • signature (ModelSignature | None) – Unsupported, see notes. An mlflow.models.ModelSignature instance describing model input/output. Use mlflow.models.infer_signature() to automatically generate from example inputs.

  • input_example (Any | None) – Unsupported, see notes. An example input for the model (used by MLflow UI).

  • pip_requirements (list[str] | None) – A list of pip requirement strings. Overrides conda_env pip section when provided.

  • extra_pip_requirements (list[str] | None) – A list of additional pip requirement strings to add to the model’s environment, in addition to the default requirements.

  • metadata (dict[str, Any] | None) – Optional dictionary of custom metadata to store in the MLmodel file.

Return type:

None

Notes

signature and input_example params are currently not supported, as they are used to support serving and input validation in the MLflow pyfunc flavor, which is not implemented for Darts models. They are accepted as params for simplifying potential future extensibility, and to keep in line with MLflow API conventions.