Integrate darts models with MLflow (logging, autologging, loading)

[pelekhs]

Is your feature request related to a current problem? Please describe.
Missing API to log / load / autolog ML models to MLflow.

Describe proposed solution
Provide an API that integrates darts models with MLflow models and provides model logging and loading capabilities. (check this as an example for pytorch)

Describe potential alternatives
Pyfunc models and model flavors can be used right now but this quite a time consuming process to handle darts models with MLflow.

Additional context
Would it be maybe good to handle darts models as pytorch lightning models during this process? And maybe build up on the prexisting pytorch integration of MLflow?

[dennisbader]

@pelekhs, acknowledged. Give us some time to go over this and come back to you.

[turbotimon]

I would really appreciate that too. I am currently using MLFlow together with Darts for a project, but have to do it manually

[dennisbader]

Hi again. We discussed and think this would be a great idea :)
Can you help us understand what such an integration would take?

In general we are open to have something like this in place as long as it doesn't disrupt our current API. Some compromises might be fine though if they make sense.

From [FR] U8darts Integration I see that you're already in contact with the MLFlow team. We would be happy to offer guidance/supervision as well!

As an extension to the feature request description from [FR] U8darts Integration:

• All of Darts TorchForecastingModels (TFM, deep learning) are built on top of PyTorch Lightning (and underlying PyTorch). So callbacks such as MLFlow loggers are available to the user.
• In Darts, the TFM models perform the PyTorch-Lightning-related setup (trainer, datasets, etc) behind the hood. This means that our API is different to the standard PyTorch-Lightning API.

[turbotimon]

In my opinion, the main advantage is to get the mlflow.autolog() functionality working

• As it already works with Pytorch Lightning, it maybe needs just a thin wrapper
• When I run it with e.g. TCNModel, I get the following error (I didn't investigate further):
  WARNING mlflow.utils.autologging_utils: Encountered unexpected error during pytorch autologging: 'LightningDoublePrecisionModule' object has no attribute 'example_input_array'.
• Customizations I can think of:
  • Log (save/export) the full darts model, not only the underlying Pytorch model.
  • Logging some darts-specific or useful metrics and parameters.

[Sonta]

Same here, would appreciate this.
@turbotimon if it helps you, I had the exact same error until i converted my time series to 32-bit.

Getting another error about a NoneType instead, but the artifact logging seems to work nevertheless.

My main issue is the loading of the model which I would need to implement manually through a custom mlflow pyfunc (haven't done it because I have a workaround through another model registry).

Can offer some assistance if someone wants to take the lead on this, don't think I'll have the capacity to do it by myself though. Looking at the responses from the mlflow guys it would need to be some kind of 3rd-party package I guess.

[turbotimon]

@Sonta Thanks for the hint with 32-bit! I'll try it out

[solalatus]

@Sonta when you talk about the artefact logger working, you mean, you "just" instantiated a Lightning based MLFlowLogger, and that was it? Or did you have to do anything special?

Any help is much appreciated!

[chartsengrafs]

We need mlflow for Darts and time series in general so badly.

[GitHunter0]

We need mlflow for Darts and time series in general so badly.

Yes, this is essential to enable darts to be used in a production setting.

PS: For reference, sktime has this function https://www.sktime.net/en/latest/api_reference/deployment.html#mlflow

[GitHunter0]

In the meanwhile, could some please provide an example converting a darts model into an MLflow model using mlflow.pyfunc.PythonModel?

[kusumy]

I would really appreciate that too. I am currently using MLFlow together with Darts for a project, but have to do it manually

Hi, can you give me an example of this? I need it so badly, primarily when registering the model.

Many thanks for considering my request.