Lab 10: MLflow & MLOps Expansion
Before we write code, let's understand our Goal: Data Scientists train hundreds of models with different algorithms and parameters. We need a way to track which model performed best, lock its schema so it doesn't crash in production, and easily switch between models without rewriting our application code.
The Tool: MLflow is an open-source platform, natively built into Databricks, designed to manage the entire Machine Learning lifecycle (MLOps).
1. Advanced MLflow Logging
Real-World Analogy Mapping: Imagine a scientist in a laboratory mixing chemicals.
- The Long Way: The scientist mixes chemicals randomly, creates a cure, but forgets to write down the recipe. They can never recreate the cure.
- The Smart Way (MLflow): The scientist uses a strict Laboratory Notebook (MLflow). Every time they run an experiment, they log the exact parameters, the accuracy score, and save the final model artifact.
However, in production, we need to go beyond basic logging. We must define a Model Signature.
Anatomy Breakdown: A Signature is a strict contract. It tells the model exactly what the incoming data columns and data types should look like. If an app sends a String instead of an Integer, the Signature blocks it and prevents a catastrophic crash.
import mlflow
from mlflow.models.signature import infer_signature
with mlflow.start_run():
# 1. Train the model
model.fit(X_train, y_train)
predictions = model.predict(X_test)
# 2. Create the Model Signature (defines expected input/output data types)
signature = infer_signature(X_train, predictions)
# 3. Provide an Input Example (helps downstream engineers know what the data looks like)
input_example = X_train.head(3)
# 4. Log the model securely to Unity Catalog Registry
mlflow.sklearn.log_model(
sk_model=model,
artifact_path="model",
signature=signature,
input_example=input_example,
registered_model_name="enterprise_data.ml_schema.churn_predictor"
)
2. Model Aliases (The Smart Deployment)
Problems Faced: If a web app is hardcoded to query "Model Version 5", what happens when the Data Scientist trains a better "Version 6"? The software engineers have to rewrite and redeploy the web app code to point to the new version.
How Present Technology Solves It:
Unity Catalog uses Aliases (e.g., @Champion, @Challenger). The web app is programmed to simply query the @Champion alias. When a new model is ready, you just move the alias tag to the new version, instantly upgrading the app with zero code changes!
from mlflow import MlflowClient
client = MlflowClient()
# Assign the 'Champion' alias to Version 6 of the model
client.set_registered_model_alias(
name="enterprise_data.ml_schema.churn_predictor",
alias="Champion",
version="6"
)
3. Invoking a Serverless Model Endpoint
Once the model is registered and assigned an alias, you can click "Serve" in the Databricks UI. This wraps your model in a highly scalable REST API.
Follow the Data: Any external application (a website, a mobile app, or a simple python script) can now send data to this URL and instantly receive a prediction.
import requests
# 1. Setup the connection
token = "dapi_your_secure_token"
url = "https://your-workspace.cloud.databricks.com/serving-endpoints/churn_endpoint/invocations"
headers = {"Authorization": f"Bearer {token}", "Content-Type": "application/json"}
# 2. Send data (e.g., a customer's age and tenure)
data = {"dataframe_records": [{"age": 45, "tenure": 12}]}
# 3. Receive the prediction (Will they churn?)
response = requests.post(url, headers=headers, json=data)
print(response.json())
← Previous: Lab 09: Testing & Monitoring | Next: Lab 11: Deep Spark UI & Memory Tuning →**