Typically, when we think about machine learning, we think along the lines of ‘grab a dataset on kaggle, choose and train an appropriate model for the use case, tune the hyperparameters, evaluate on the hold-out, report the 98% accuracy, the job is done’. This is typically true for the research part of ML, however, once the model actually should be used in some real application, the story becomes a bit more complicated.
MLFlow provides an ecosystem to manage the surrounding artefacts of the models: tracking the experiments, metrics, hyperparameters, and other aspects. In this topic, we will look at the general architecture and the capabilities of MLFlow.
The components
The design of MLflow is mostly centered around three key components: MLflow Tracking, MLflow Projects, and MLflow Models.
MLflow Tracking is responsible for logging and tracking run metrics, parameters, code versions, and output files during the experimentation and development phase of a machine learning project. It allows data scientists to organize and compare multiple runs, visualize performance metrics, and reproduce successful runs with ease.
MLflow Projects provides a convention for packaging data science code in a reusable and reproducible way. It allows developers to specify their code, data, and environment dependencies in a simple format, making it easy to run the same code on different platforms or collaborate with others.
MLflow Models module streamlines the process of packaging and deploying machine learning models to various production environments. It provides a uniform model format and tools for logging models, loading them for serving or batch inference, and deploying them to diverse platforms, including cloud and edge environments.
A closer look at the tracking module
MLflow Tracking is a core component of MLflow that enables tracking and managing the experiments. It provides a centralized location for logging and querying various artifacts generated during the experimentation phase, such as metrics, parameters, code versions, and output files. Here's a more detailed look at the design and features of MLflow Tracking:
The module allows you to log and track individual runs of your machine learning code. Each run is identified by a unique run ID and contains metadata such as the start and end time, source code version, and various parameters. This metadata helps in reproducing and comparing different runs. During a run, you can log various metrics (e.g., accuracy, loss, F1-score) at different points in time. MLflow Tracking stores these metrics and their values, allowing you to visualize and compare metric trends across multiple runs. In addition to metrics, MLflow Tracking enables logging of arbitrary output files (artifacts) generated during a run. These artifacts can include model checkpoints, data samples, images, or any other relevant files. This feature facilitates reproducibility and sharing of results. The runs are organized into logical groups called experiments. Each experiment has a unique identifier and can contain multiple runs. This hierarchical structure makes it easier to manage and compare related experiments.
Additionally, the Tracking module integrates with MLflow Models. This integration allows you to log and track trained models alongside their corresponding run metadata.
Projects and models
MLflow Projects is a component of MLflow that simplifies the packaging and reproducible execution of code. It provides a convention for organizing and describing code, data, and dependencies, making it easier to reproduce and share complete machine learning projects.
The Projects module uses a simple MLproject file to define the entry point of a project, as well as any necessary dependencies (e.g., Python packages, conda environments, or system requirements). This file serves as a single source for reproducing the project's execution environment. It allows you to encapsulate your code, data, and environment dependencies into a self-contained package. This package can be shared and executed consistently across different environments, ensuring reproducibility. MLflow Projects also supports running projects locally or remotely, enabling collaboration and distributed execution. Remote execution can be performed on platforms like Kubernetes, Databricks, or other cloud-based environments.
MLflow Models is a component that simplifies the process of transitioning models from experimentation to production environments by providing a standard format for packaging machine learning models, regardless of the framework used for training. This format includes the model artifacts (e.g., pickle files, TensorFlow SavedModel, or ONNX models) and any necessary dependencies or configurations.
It also supports multiple representations of the same model, catering to different deployment environments or inference scenarios (this is referred to as MLFlow flavors). For example, a model can have a Python function for batch inference, as well as a Docker image for containerized deployment. There is also a Model Registry for managing and tracking model versions, transitions, and deployments. The registry allows you to store, version, and annotate models, facilitating model governance and lineage tracking.
The Models module enforces model input and output signatures during serving, ensuring that the deployed model adheres to the expected input and output data formats. This helps catch issues early and maintains consistency between training and serving environments.
Conclusion
As a result, you are now familiar with the capabilities of MLFlow and it’s main components.