3 min read

    Lab 10: MLflow & MLOps Expansion

    #databricks#mlflow#mlops#machine-learning

    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.

    python
    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!

    python
    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.

    python
    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 →**