As an MLOps engineer, I’ve had the opportunity to work with numerous tools and frameworks designed to simplify and enhance the machine learning (ML) lifecycle. One tool that stands out in the MLOps landscape is MLflow. In this blog, I’ll dive deep into MLflow, its features, benefits, and how it can revolutionize your ML workflows.
What is MLflow?
MLflow is an open-source platform designed to manage the end-to-end machine learning lifecycle. Developed by Databricks, MLflow provides a suite of tools to streamline the process of developing, deploying, and managing ML models. It is framework-agnostic, meaning it can be used with any ML library or language.
Key Features of MLflow
MLflow offers four main components that cater to different aspects of the ML lifecycle:
MLflow Tracking: This component allows you to log and query experiments. It tracks parameters, metrics, and artifacts for your experiments, making it easy to compare different runs and select the best-performing model.
MLflow Projects: With MLflow Projects, you can package your code into a reusable and reproducible format. Each project can specify its dependencies and environment settings, ensuring that your experiments are easily shareable and reproducible.
MLflow Models: This component provides a standard format for packaging and deploying ML models. It supports multiple flavors (e.g., Python, R, Java), allowing you to deploy models to different environments seamlessly.
MLflow Registry: The MLflow Registry is a centralized model repository that helps manage the full lifecycle of an ML model, including versioning, stage transitions (e.g., staging to production), and annotations.
Why Use MLflow?
MLflow addresses several pain points commonly faced in the ML lifecycle, including:
Reproducibility: By tracking experiments and packaging code, MLflow ensures that your results are reproducible and easily shareable.
Scalability: MLflow’s framework-agnostic design allows you to scale your ML workflows across different environments and teams.
Collaboration: The centralized model registry and tracking capabilities enhance collaboration between data scientists, ML engineers, and operations teams.
Efficiency: Automation of logging, packaging, and deployment tasks streamlines the ML lifecycle, freeing up time for more strategic work.
Getting Started with MLflow
Here’s a step-by-step guide to getting started with MLflow:
Step 1: Install MLflow
You can install MLflow using pip:
bash
pip install mlflow
Step 2: Set Up Your Environment
Create a new project directory and set up your environment. For this example, we’ll use a simple ML project with a Python environment:
bash
mkdir mlflow_project
cd mlflow_project
python -m venv venv
source venv/bin/activate
pip install mlflow pandas scikit-learn
Step 3: Log an Experiment
Create a Python script to log an experiment. For this example, we’ll use a simple linear regression model:
python
import os
import warnings
import sys
import pandas as pd
import numpy as np
from sklearn.metrics import mean_squared_error, mean_absolute_error, r2_score
from sklearn.model_selection import train_test_split
from sklearn.linear_model import ElasticNet
from urllib.parse import urlparse
import mlflow
from mlflow.models.signature import infer_signature
import mlflow.sklearn
import dagshub
import logging
import dagshub
dagshub.init(repo_owner='bappymalik4161', repo_name='mlflow-test', mlflow=True)
logging.basicConfig(level=logging.WARN)
logger = logging.getLogger(__name__)
def eval_metrics(actual, pred):
rmse = np.sqrt(mean_squared_error(actual, pred))
mae = mean_absolute_error(actual, pred)
r2 = r2_score(actual, pred)
return rmse, mae, r2
if __name__ == "__main__":
warnings.filterwarnings("ignore")
np.random.seed(40)
# Read the wine-quality csv file from the URL
csv_url = (
"https://raw.githubusercontent.com/mlflow/mlflow/master/tests/datasets/winequality-red.csv"
)
try:
data = pd.read_csv(csv_url, sep=";")
except Exception as e:
logger.exception(
"Unable to download training & test CSV, check your internet connection. Error: %s", e
)
# Split the data into training and test sets. (0.75, 0.25) split.
train, test = train_test_split(data)
# The predicted column is "quality" which is a scalar from [3, 9]
train_x = train.drop(["quality"], axis=1)
test_x = test.drop(["quality"], axis=1)
train_y = train[["quality"]]
test_y = test[["quality"]]
alpha = float(sys.argv[1]) if len(sys.argv) > 1 else 0.5
l1_ratio = float(sys.argv[2]) if len(sys.argv) > 2 else 0.5
with mlflow.start_run():
lr = ElasticNet(alpha=alpha, l1_ratio=l1_ratio, random_state=42)
lr.fit(train_x, train_y)
predicted_qualities = lr.predict(test_x)
(rmse, mae, r2) = eval_metrics(test_y, predicted_qualities)
print("Elasticnet model (alpha={:f}, l1_ratio={:f}):".format(alpha, l1_ratio))
print(" RMSE: %s" % rmse)
print(" MAE: %s" % mae)
print(" R2: %s" % r2)
mlflow.log_param("alpha", alpha)
mlflow.log_param("l1_ratio", l1_ratio)
mlflow.log_metric("rmse", rmse)
mlflow.log_metric("r2", r2)
mlflow.log_metric("mae", mae)
# For remote server only (Dagshub)
remote_server_uri = "https://dagshub.com/bappymalik4161/mlflow-test.mlflow"
mlflow.set_tracking_uri(remote_server_uri)
tracking_url_type_store = urlparse(mlflow.get_tracking_uri()).scheme
# Model registry does not work with file store
if tracking_url_type_store != "file":
# Register the model
# There are other ways to use the Model Registry, which depends on the use case,
# please refer to the doc for more information:
# https://mlflow.org/docs/latest/model-registry.html#api-workflow
mlflow.sklearn.log_model(
lr, "model", registered_model_name="ElasticnetWineModel")
else:
mlflow.sklearn.log_model(lr, "model")
Follow me on github: https://github.com/bittush8789/MLOps-Foundation-/tree/main/04.MLflow
Follow me LinkedIn: https://www.linkedin.com/in/bittu-kumar-54ab13254/