MLflow for Darts implementation

.gitignore

@@ -25,3 +25,4 @@ docs_env
 .venv
 .env
 uv.lock
+*mlruns/

CHANGELOG.md

@@ -11,6 +11,12 @@ but cannot always guarantee backwards compatibility. Changes that may **break co
 
 **Improved**
 
+- 🚀🚀 Added a custom MLflow model flavor for Darts, available under `darts.utils.mlflow`. It provides an MLflow integration for any Darts `ForecastingModel` (statistical, scikit-learn-like, and PyTorch-based). [#3022](https://github.com/unit8co/darts/pull/3022) by [Jakub Chłapek](https://github.com/jakubchlapek), [Zhihao Dai](https://github.com/daidahao) and [Michel Zeller](https://github.com/mizeller).
+  - Added `save_model()`, `load_model()`, and `log_model()` to persist and reload Darts models as MLflow models, including model and covariate metadata.
+  - Added `autolog()` to automatically log model creation parameters, covariate usage, metrics, and the trained model artifact when `log_models=True` (default `False`) on every `fit()`, `backtest()` and `historical_forecasts()` call within an active MLflow run.
+  - Metric functions from `darts.metrics` are automatically logged when called inside an active run, with keys reflecting the metric output shape (per-component, per-quantile/interval, per-label, and per-timestep results charted across MLflow steps). Multi-series results log the aggregate over series and write the full per-series breakdown to a JSON artifact.
+  - Added a new notebook for [MLflow quickstart](https://unit8co.github.io/darts/examples/29-MLflow-quickstart.html) with detailed usage examples.
+  - Note: model serving and deployment (MLflow's `pyfunc` flavor, model signatures, and input examples) are not yet supported.
 - Added support for per-timestep (non-aggregated) encoder and decoder variable importances in `TFTExplainer`, exposed as `TimeSeries` via `TFTExplainabilityResult.get_encoder_importance_over_time()` and `get_decoder_importance_over_time()`. [#3170](https://github.com/unit8co/darts/pull/3170) by [exactml](https://github.com/exactml).
 - Calling `TFTModel.fit_from_dataset()` on a dataset that does not have future covariates now raises an informative exception. [#3149](https://github.com/unit8co/darts/pull/3149) by [YOON KIWOONG](https://github.com/kiwoongyoon).
 - 🔴 Percentage and range-based metrics (`ape`, `mape`, `sape`, `smape`, `wmape`, `ope`, `arre`, `marre`, `coefficient_of_variation`) no longer raise a hard `ValueError` when the denominator is exactly zero. A new `zero_division` parameter controls the behavior: [#3122](https://github.com/unit8co/darts/pull/3122) by [Mahimn](https://github.com/mahimn01).
@@ -26,6 +32,8 @@ but cannot always guarantee backwards compatibility. Changes that may **break co
 
 **Dependencies**
 
+- Added `mlflow>=3.0` to the optional dependency group, enabling the MLflow integration. [#3022](https://github.com/unit8co/darts/pull/3022) by [Jakub Chłapek](https://github.com/jakubchlapek), [Zhihao Dai](https://github.com/daidahao) and [Michel Zeller](https://github.com/mizeller).
+
 ### For developers of the library:
 
 ## [0.46.1](https://github.com/unit8co/darts/tree/0.46.1) (2026-07-20)

darts/metrics/utils.py

@@ -155,6 +155,22 @@ def wrapper_classification_support(*args, **kwargs):
     return wrapper_classification_support
 
 
+_metric_callbacks: list[Callable] = []
+
+
+def register_metric_callback(callback: Callable) -> None:
+    """Register a callback to be invoked after every top-level metric call."""
+    _metric_callbacks.append(callback)
+
+
+def unregister_metric_callback(callback: Callable) -> None:
+    """Remove a previously registered metric callback, if present."""
+    try:
+        _metric_callbacks.remove(callback)
+    except ValueError:
+        pass
+
+
 def multi_ts_support(func) -> Callable[..., METRIC_OUTPUT_TYPE]:
     """
     This decorator further adapts the metrics that took as input two (or three for scaled metrics with `insample`)
@@ -169,6 +185,9 @@ def multi_ts_support(func) -> Callable[..., METRIC_OUTPUT_TYPE]:
 
     @wraps(func)
     def wrapper_multi_ts_support(*args, **kwargs):
+        original_args = args
+        original_kwargs = dict(kwargs)
+
         actual_series = (
             kwargs["actual_series"] if "actual_series" in kwargs else args[0]
         )
@@ -312,6 +331,10 @@ def wrapper_multi_ts_support(*args, **kwargs):
         elif series_seq_type == SeriesType.SINGLE:
             vals = vals[0]
 
+        # invoke registered callbacks (e.g. MLflow autologging)
+        for cb in _metric_callbacks:
+            cb(func=func, result=vals, args=original_args, kwargs=original_kwargs)
+
         # flatten along series axis if n series == 1
         return vals
 

darts/models/forecasting/pl_forecasting_module.py

@@ -450,8 +450,11 @@ def _update_metrics(self, output, target, metrics):
             pred = output.squeeze(dim=-1)
 
         # torch metrics require 2D targets of shape (batch size * ocl, num targets)
-        target = target.reshape(-1, self.n_targets)
-        pred = pred.reshape(-1, self.n_targets)
+        # contiguous() is needed because model outputs can be non-contiguous views
+        # (e.g. NBEATS slices the last dimension), and some torchmetrics implementations
+        # call .view() internally which requires a contiguous tensor.
+        target = target.reshape(-1, self.n_targets).contiguous()
+        pred = pred.reshape(-1, self.n_targets).contiguous()
 
         metrics.update(pred, target)
 

darts/tests/conftest.py

@@ -32,6 +32,7 @@ def _package_available(*names: str) -> bool:
 PLOTLY_AVAILABLE = _package_available("plotly")
 IPYTHON_AVAILABLE = _package_available("IPython")
 TIREX_AVAILABLE = _package_available("tirex")
+MLFLOW_AVAILABLE = _package_available("mlflow")
 
 tfm_kwargs: dict[str, Any] = {
     "pl_trainer_kwargs": {