As a Data scientist you are constantly working on Machine Learning problems, mostly collaborating with others. To do so successfully, vendors are developing ML workflow tools which can assist teams in managing the lifecycle of Machine Learning. In this insight I focus on one of the leading tools: MLflow and how MLflow can facilitate and improve your ML lifecycle.
Why even bother about ML workflow tools?
We see that a lot of companies are working on ML projects in teams without having a standardized or centralized solution which can handle or manage ML models. This means that lots of teams are working on different model versions that are stored on different locations.
Recognize this?
This also means that teams find it hard to compare performance of models and that they lack continuous monitoring. Missing these standards leads to inefficiencies and frustrations.
To prevent the aboven ML workflow tools are being developed. In this insight we focus on one of the leading tools: MLflow.
What can MLflow bring to the table?
MLflow is an open-source solution which can be leveraged from your local notebooks, local scripts and it is even natively integrated in Databricks (out-of-box available). The main advantages of including MLflow in your ML lifecycle is the transparency and standardization it brings to the table when it comes to training, tuning and deploying ML models.
- MLflow Tracking: Enables teams to record, have & query a clear overview of the performance of the training runs as all models (& performance metrics) are stored centrally in an “experiment”
- MLflow Projects: Enables teams to easily reproduce model versions with its corresponding code on any platform.
- MLflow (model) Registry: Enables teams to manage ML model versions, modify lifecycle stages and deploy models to production.
- MLflow Models: Packages models in a standard format so that they can e.g. be served as an endpoint through a REST API. Also, MLflow models are compatible with the most important ML libraries.
What components can pick the low hanging fruit?
Firstly, start to leverage MLflow tracking to centralize all the information on training runs (metrics, parameters, model artifacts) in an experiment. Those experiments are accessible through different machines, so that teams can easily train, tune and compare models during the Model Training phase of the machine learning project. For us, the greatest advantages of having the centralized experiments are:
- A user-friendly UI
- Centralized environment to log model training runs
- The models are compatible with ML’s most important libraries and are logged in MLflow in a standard format, which means that they can be used e.g. in real-time serving through a REST API or in batch in Databricks
- Easy to compare performance metrics or model parameters in the UI. Not only available in tables, but also possibility to visualize the differences in graphs.
- The experiments are highly scalable, thanks to its smart search option. Even when having lots and lots of models, thanks to the search option, co-workers can look for models through tags, developer name or accuracy levels.
- No entry barrier. You can integrate the experiments immediately in existing e.g. python or spark code and start tracking your training runs.
Overview of MLFlow process flow: the model registry serves as thé central repository for all data scientists & ML engineers
Example of MLflow experiment tracking: every run of the model is versioned & stored with all run parameters & metrics
Secondly, enable your team to manage ML model lifecycles easily by using the MLflow model registry. In short, the registry is a central hub, where all the models are stored under a unique model name, model version and model stage. Models can then be easily deployed to production: from the registry UI or through its API. The greatest advantages of leveraging the model registry as a team are:
- The UI provides all team members with clear information on what model versions (all versions remain) are used for experiments, for testing or for production.
- You could allocate roles in the registry where some could request models to transition from “Staging” to “Production” and others can validate that request. Having the ability to set those stages can be very powerful, as you could just refer to the production model from your code and be certain that you will always refer to the correct model.
- A nice recent add-on is that you can now choose to serve your model through a REST API endpoint, which can be done straight from the UI and is directly deployed and served through Azure Databricks
MLflow model can be called directly within a Spark UDF function for use
In conclusion, both components can improve efficiency and reduce frustration when collaborating in team on various ML projects.
How long before your team can benefit from MLflow?
Quick. For me, one of the main advantages of working with MLflow is its easiness: just import the package, configure your client & off you go!
import mlflow from mlflow.tracking.client import MlflowClient
You could get started within a couple of hours, if you have some coding experience and know how to install libraries to e.g. python or spark. During your discovery of MLflow, you will observe that tracking training runs is as simple as running certain predefined lines of code, which can be found in the clear documentation of the MLflow project.
Can we leverage MLflow in cloud solutions?
Yes! Databricks offers a seamless integration with MLflow, which means that MLflow is a fully managed solution. No need to worry about the right configuration of the platform, you as a Data Science team can immediately start with tracking your experiments from the Databricks notebooks or manage and deploy models from the model registry. As an intensive Databricks user, I really applaud that integration!
Another integration that we, at element61, like is the integration with Azure ML. Azure ML is a fully managed cloud solution, which supports the entire machine learning lifecycle thanks to its MLOps. We mostly advice Azure ML to be used at companies where the maturity of machine learning is still growing, because it has automation and visual features which can help citizen data scientists to more easily work on ML projects. To integrate the best of both worlds we propose that Citizen Data Scientists work in Azure ML, while expert Data Scientists can train, tune and compare models at scale using MLflow experiments, while models are managed through their lifecycle towards deployment leveraging the MLflow registry. Models from MLflow are compatible in Azure ML which means that models from the model registry can be stored and deployed as an endpoint in Azure ML.
How to get started with MLflow?
You can leverage it locally or in the cloud. You will notice that it takes minimal effort before you can start!
- For those who want to use MLflow locally
- Install the MLflow package by running: pip install mlflow
- Add MLflow commands to your existing code & execute your code to start tracking experiments!
- To check your experiments, run mlflow ui in command line
- To use the model registry, you will need to set up a server using a
database-backed store - After successfully setting up the server, you can start leveraging the model registry by accessing it through the mlflow ui
- For those who want to use MLflow in the cloud (So simple!)
- Create a new Azure Databricks resource (or use an existing one)
- Add MLflow commands to your existing code & execute your code to start tracking experiments!
- Check your experiments immediately via the top right experiment icon
- Check out your model registry directly via the Models icon on the left
We love how MLflow allows users to get started with minimal effort. Try it out now and enjoy it yourself!
For more information on how to improve your Machine Learning lifecycle, do not hesitate to contact me or one of my colleagues.
Continue reading
Building on our partnership with both Databricks; element61 has some other great articles & pages which might enrich above insight.
Continue reading or contact us to get started:
