MLflow for DistributedMLForecast

[expertise-team]

I can see in the documentation how we can use MLFlow with MLForecast.
Is there any way to use it for the distributed (like Spark) version? The logging and loading of the model would be very important.

[jmoralez]

Hey. Can you elaborate on your use case a bit? Would this be for like batch inference on a cluster?

[expertise-team]

This is our current scenario:

• We use your library for forecasting demand with LightGBM for 500 retail stores with 6000 products.
• Our training database is over 4 billion records with sales and stock data, promotion
• We expect the training dataset will grow, since there are 1.500 stores in total and the products are also limited now.
• We use Databricks and the distributed version of your library for the training and the prediction.
• We use it in a batch process and we use the forecast for calculating orders for each store and product.

Goal:

• We want to finetune our featuring and training process because there are timeserieses where we overestimate the demand.
• In order to do that, we want to use MLflow for tracking our experiments and their metrics and so that we can decide which experiment to use for the prediction.
• We want to save and track the trained models in MLflow, if a problem occurs, we can rollback to another version of the model.

Currently we just save the model in Azure blob storage with the Distributed version save method and we can add this folder as an artifact.
We are not able to track it as a model and register it with versions into MLFlow. Like that it would be much easier to use the model for different batch predictions in Databricks.

As far as I see, the machine learning libraries are more advanced for single machine scenarios, not for a distributed environment. I am also thinking to separate the processes:

• we could use the distributed environment for calculating the features and save them
• we would try a single machine with enough memory for learning, like that MLflow wouldn't be a problem, since many examples are about a single node version. I don't know if this database would fit into a single machines memory, I have never tried it.

I hope, I could add more detail to my problem.

[jmoralez]

How are you planning on using the model? If you just need distributed for training but once it's trained you plan to use it to forecast a smaller subset you could convert the model to local and save that instead. You can reduce the size of the local model by setting the keep_last_n argument of the distributed fit to the number of samples that the features require (we're working on automatic that so if you use the built-in lag transformations you could skip that step).

[expertise-team]

Thank you for your quick answer, we will test to convert it to a local model and use it for forecasting.