{"id":49012,"date":"2025-04-06T15:30:59","date_gmt":"2025-04-06T15:30:59","guid":{"rendered":"https:\/\/www.devopsschool.com\/blog\/?p=49012"},"modified":"2025-04-06T16:09:29","modified_gmt":"2025-04-06T16:09:29","slug":"mlflow-lab-end-to-end-workflow-on-databricks","status":"publish","type":"post","link":"https:\/\/www.devopsschool.com\/blog\/mlflow-lab-end-to-end-workflow-on-databricks\/","title":{"rendered":"MLflow Lab: End-to-End Workflow on Databricks"},"content":{"rendered":"\n

Absolutely! Let\u2019s walk through a complete, step-by-step tutorial<\/strong> to help you understand MLflow on Databricks<\/strong> from start to finish<\/strong>, using revised and working code<\/strong>.<\/p>\n\n\n\n

We’ll cover all key components of MLflow: \u2705 Tracking
\u2705 Models
\u2705 Model Registry
\u2705 Signature + Input Example<\/p>\n\n\n\n

\n\n\n\n

\ud83d\ude80 Objective:<\/strong><\/h2>\n\n\n\n

Train a classification model on the Iris dataset<\/strong>, log everything with MLflow, and register the model.<\/p>\n\n\n\n

\n\n\n\n

\u2705 Step-by-Step MLflow Lab on Databricks<\/strong><\/h2>\n\n\n\n
\n\n\n\n

\ud83c\udfaf Step 1: Setup (Install Required Libraries)<\/strong><\/h3>\n\n\n\n

Run this in a cell:<\/p>\n\n\n

%pip install scikit-learn pandas mlflow\n<\/code><\/span><\/pre>\n\n\n
\n\n\n\n
\ud83c\udfaf Step 2: Import Libraries and Load Data<\/strong><\/h3>\n\n\n
import<\/span> pandas as<\/span> pd\nfrom<\/span> sklearn import<\/span> datasets\nfrom<\/span> sklearn.ensemble import<\/span> RandomForestClassifier\nfrom<\/span> sklearn.model_selection import<\/span> train_test_split\nfrom<\/span> sklearn.metrics import<\/span> accuracy_score\nimport<\/span> mlflow\nimport<\/span> mlflow.sklearn\nfrom<\/span> mlflow.models.signature import<\/span> infer_signature\n<\/code><\/span>Code language:<\/span> JavaScript<\/span> (<\/span>javascript<\/span>)<\/span><\/small><\/pre>\n\n\n
\n\n\n\n
\ud83c\udfaf Step 3: Prepare the Data<\/strong><\/h3>\n\n\n
# Load Iris dataset<\/span>\niris = datasets.load_iris()\nX = iris.data\ny = iris.target\n\n# Split into training and testing<\/span>\nX_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3<\/span>, random_state=42<\/span>)\n<\/code><\/span>Code language:<\/span> PHP<\/span> (<\/span>php<\/span>)<\/span><\/small><\/pre>\n\n\n
\n\n\n\n
\ud83c\udfaf Step 4: Start MLflow Run and Train Model<\/strong><\/h3>\n\n\n
# Start MLflow experiment run<\/span>\nwith mlflow.start_run() as<\/span> run:\n\n    # Train the model<\/span>\n    model = RandomForestClassifier(n_estimators=100<\/span>, max_depth=5<\/span>, random_state=42<\/span>)\n    model.fit(X_train, y_train)\n\n    # Make predictions and calculate accuracy<\/span>\n    predictions = model.predict(X_test)\n    acc = accuracy_score(y_test, predictions)\n\n    # Log parameters and metric<\/span>\n    mlflow.log_param(\"n_estimators\"<\/span>, 100<\/span>)\n    mlflow.log_param(\"max_depth\"<\/span>, 5<\/span>)\n    mlflow.log_metric(\"accuracy\"<\/span>, acc)\n\n    # Create sample input and signature<\/span>\n    input_example = X_test[:5<\/span>]\n    signature = infer_signature(X_train, model.predict(X_train))\n\n    # Log model with signature and input example<\/span>\n    mlflow.sklearn.log_model(\n        sk_model=model,\n        artifact_path=\"iris_rf_model\"<\/span>,\n        input_example=input_example,\n        signature=signature\n    )\n\n    # Save run_id for model registration<\/span>\n    run_id = run.info.run_id\n    print<\/span>(f\"Run ID: {run_id}\"<\/span>)\n    print<\/span>(f\"Accuracy: {acc}\"<\/span>)\n<\/code><\/span>Code language:<\/span> PHP<\/span> (<\/span>php<\/span>)<\/span><\/small><\/pre>\n\n\n
\n\n\n\n
\ud83c\udfaf Step 5: Register the Model<\/strong><\/h3>\n\n\n\n
Paste this in a new cell:<\/p>\n\n\n
model_uri = f\"runs:\/{run_id}\/iris_rf_model\"<\/span>\n\n# Register the model under a name<\/span>\nmodel_details = mlflow.register_model(\n    model_uri=model_uri,\n    name=\"IrisClassifierModel\"<\/span>\n)\n<\/code><\/span>Code language:<\/span> PHP<\/span> (<\/span>php<\/span>)<\/span><\/small><\/pre>\n\n\n
\n
\u2705 Now go to “Models” tab in Databricks<\/strong>, and you’ll see IrisClassifierModel<\/code><\/strong> with versioning.<\/p>\n<\/blockquote>\n\n\n\n
\n\n\n\n
\ud83c\udfaf Step 6: Promote the Model (via UI)<\/strong><\/h3>\n\n\n\n
Go to:<\/p>\n\n\n\n
    \n
  • Models > IrisClassifierModel<\/strong><\/li>\n\n\n\n
  • Click on the version (e.g., Version 1)<\/li>\n\n\n\n
  • Click Stage<\/code> \u2192 Choose Staging<\/code> or Production<\/code><\/li>\n<\/ul>\n\n\n\n
    \n\n\n
    # To Promore Version1 to Production<\/span>\n\nfrom mlflow.tracking import MlflowClient\n\nclient = MlflowClient()\nclient.transition_model_version_stage(\n    name=\"IrisClassifierModel\"<\/span>,\n    version=1<\/span>,  # or the actual version you created<\/span>\n    stage=\"Production\"<\/span>  # or \"Staging\", \"Archived\"<\/span>\n)<\/code><\/span>Code language:<\/span> PHP<\/span> (<\/span>php<\/span>)<\/span><\/small><\/pre>\n\n\n
    \n\n\n\n
    OR<\/p>\n\n\n\n
    \n\n\n
    # Assign an alias (like @production) to the version<\/span>\nfrom mlflow.tracking import MlflowClient\n\nclient = MlflowClient()\n\nclient.set_registered_model_alias(\n    name=\"IrisClassifierModel\"<\/span>,\n    alias=\"production\"<\/span>,\n    version=1<\/span>\n)\n\n<\/code><\/span>Code language:<\/span> PHP<\/span> (<\/span>php<\/span>)<\/span><\/small><\/pre>\n\n\n
    \n\n\n
    # Run this to list your registered models:<\/span>\nclient = MlflowClient()\nmodels = client.list_registered_models()\nfor<\/span> m in models:\n    print<\/span>(m.name)<\/code><\/span>Code language:<\/span> PHP<\/span> (<\/span>php<\/span>)<\/span><\/small><\/pre>\n\n\n
    \n\n\n\n
    \ud83c\udfaf Step 7: Load Model from Registry and Predict<\/strong><\/h3>\n\n\n\n
    Option A: If your model is in “Production” stage<\/strong> (via UI):<\/p>\n\n\n\n
    Now that the model is registered and staged, let\u2019s load it and use it:<\/p>\n\n\n
    from mlflow.pyfunc import load_model\n\n# Load model from registry using stage name<\/span>\nmodel = load_model(\"models:\/IrisClassifierModel\/Production\"<\/span>)\n\n# Predict<\/span>\npreds = model.predict(X_test)\nprint<\/span>(preds)<\/code><\/span>Code language:<\/span> PHP<\/span> (<\/span>php<\/span>)<\/span><\/small><\/pre>\n\n\n
    \n\n\n\n
    Option B: If you used an alias<\/strong> (e.g., production<\/code>) instead of a stage:<\/h2>\n\n\n\n
    \n\n\n
    from mlflow.pyfunc import load_model\n\n# Load model using alias<\/span>\nmodel = load_model(\"models:\/IrisClassifierModel@production\"<\/span>)\n\n# Predict<\/span>\npreds = model.predict(X_test)\nprint<\/span>(preds)<\/code><\/span>Code language:<\/span> PHP<\/span> (<\/span>php<\/span>)<\/span><\/small><\/pre>\n\n\n
    \n\n\n\n
    \ud83c\udf89 You Did It!<\/h2>\n\n\n\n
    MLflow Component<\/th>Covered \u2705<\/th><\/tr><\/thead>
    Tracking<\/strong><\/td>Yes (log params, metrics, artifacts)<\/td><\/tr>
    Models<\/strong><\/td>Yes (logged with signature + input example)<\/td><\/tr>
    Model Registry<\/strong><\/td>Yes (registered and promoted via UI)<\/td><\/tr>
    Serving\/Loading<\/strong><\/td>Yes (loaded from registry and used for predictions)<\/td><\/tr><\/tbody><\/table><\/figure>\n\n\n\n
    \n\n\n\n
    \ud83e\udde0 Bonus Tips:<\/h2>\n\n\n\n
      \n
    • Want to serve the model as REST API<\/strong>? MLflow on Databricks supports model serving.<\/li>\n\n\n\n
    • Want to run this as a Project<\/strong>? We can convert this into an MLflow Project next.<\/li>\n<\/ul>\n\n\n\n
      Let me know if you want this exported as a .dbc<\/code> Databricks Notebook or want to integrate with CI\/CD!<\/p>\n\n\n\n
      Objective: Serve Your MLflow Model via REST API on Databricks<\/h2>\n\n\n\n
      Absolutely! Since you’ve already registered the model and assigned an alias like production<\/code>, you\u2019re ready to serve the model<\/strong> using Databricks Model Serving<\/strong>.<\/p>\n\n\n\n
      Let\u2019s go through this step-by-step<\/strong>, assuming you’re using a Databricks Trial or Paid workspace (not Community Edition)<\/strong>.<\/p>\n\n\n\n
      \n\n\n\n
      \ud83d\ude80 Objective: Serve Your MLflow Model via REST API on Databricks<\/strong><\/h2>\n\n\n\n
      We\u2019ll go from registered model with alias<\/strong> \u279c to real-time REST endpoint<\/strong> \u279c to making predictions from code<\/strong>.<\/p>\n\n\n\n
      \n\n\n\n
      \u2705 Step 1: Confirm Model Registration & Alias<\/strong><\/h2>\n\n\n\n
      You’ve already done this, but here\u2019s a quick reference for future:<\/p>\n\n\n
      from mlflow.tracking import MlflowClient\n\nclient = MlflowClient()\n\n# Register model (already done)<\/span>\nmodel_uri = f\"runs:\/{run_id}\/iris_rf_model\"<\/span>\nclient.register_model(model_uri=model_uri, name=\"IrisClassifierModel\"<\/span>)\n\n# Set alias<\/span>\nclient.set_registered_model_alias(\n    name=\"IrisClassifierModel\"<\/span>,\n    alias=\"production\"<\/span>,\n    version=1<\/span>\n)\n<\/code><\/span>Code language:<\/span> PHP<\/span> (<\/span>php<\/span>)<\/span><\/small><\/pre>\n\n\n
      \n\n\n\n
      \u2705 Step 2: Enable Model Serving from Databricks UI<\/strong><\/h2>\n\n\n\n
        \n
      1. Go to Databricks Workspace<\/strong>.<\/li>\n\n\n\n
      2. In the left sidebar, click “Models”<\/strong>.<\/li>\n\n\n\n
      3. Click on IrisClassifierModel<\/code><\/strong>.<\/li>\n\n\n\n
      4. Click on Version 1<\/strong> (or the version you aliased).<\/li>\n\n\n\n
      5. You should see a \u201cServing\u201d<\/strong> or \u201cEnable Serving\u201d<\/strong> button.<\/li>\n\n\n\n
      6. Click it, then:\n
          \n
        • Choose Real-time serving<\/strong><\/li>\n\n\n\n
        • Click Start serving<\/strong><\/li>\n<\/ul>\n<\/li>\n<\/ol>\n\n\n\n
          \n
          \u2705 Once serving is enabled, you\u2019ll see the endpoint URL<\/strong> (copy it!).<\/p>\n<\/blockquote>\n\n\n\n
          \n\n\n\n
          \u2705 Step 3: Use the REST Endpoint for Predictions<\/strong><\/h2>\n\n\n\n
          Here\u2019s a full Python example to send test data and get predictions:<\/p>\n\n\n
          import requests\nimport json\n\n# Replace with your actual endpoint from Databricks<\/span>\nurl = \"https:\/\/<your-databricks-instance>\/serving-endpoints\/IrisClassifierModel\/invocations\"<\/span>\n\n# If needed, generate a Personal Access Token from Databricks User Settings<\/span>\ntoken = \"dapiXXXXXXXXXXXXXXXXXXXX\"<\/span>\n\n# Headers<\/span>\nheaders = {\n    \"Authorization\"<\/span>: f\"Bearer {token}\"<\/span>,\n    \"Content-Type\"<\/span>: \"application\/json\"<\/span>\n}\n\n# Input payload (match your model\u2019s input structure)<\/span>\ndata = {\n    \"dataframe_split\"<\/span>: {\n        \"columns\"<\/span>: [\"sepal length (cm)\"<\/span>, \"sepal width (cm)\"<\/span>, \"petal length (cm)\"<\/span>, \"petal width (cm)\"<\/span>],\n        \"data\"<\/span>: [[5.1<\/span>, 3.5<\/span>, 1.4<\/span>, 0.2<\/span>]]\n    }\n}\n\n# Send request<\/span>\nresponse = requests.post(url, headers=headers, json=data)\n\n# Print response<\/span>\nprint<\/span>(\"Prediction:\"<\/span>, response.json())\n<\/code><\/span>Code language:<\/span> PHP<\/span> (<\/span>php<\/span>)<\/span><\/small><\/pre>\n\n\n
          \n\n\n\n
          \u2705 Step 4: Test It!<\/strong><\/h2>\n\n\n\n
          Run the above Python code in:<\/p>\n\n\n\n
            \n
          • Databricks notebook<\/li>\n\n\n\n
          • Jupyter notebook<\/li>\n\n\n\n
          • Any Python script<\/li>\n<\/ul>\n\n\n\n
            \n\n\n\n
            \ud83d\udd10 Generate a Personal Access Token (If Needed)<\/strong><\/h2>\n\n\n\n
              \n
            1. Click on your profile icon in the top-right corner of Databricks.<\/li>\n\n\n\n
            2. Go to “User Settings” > “Access Tokens”<\/strong><\/li>\n\n\n\n
            3. Click Generate New Token<\/strong><\/li>\n\n\n\n
            4. Copy it and use it in your token<\/code> variable<\/li>\n<\/ol>\n\n\n\n
              \n\n\n\n
              \ud83e\uddea Example Output:<\/strong><\/h2>\n\n\n
              Prediction<\/span>: [0]<\/span>\n<\/code><\/span>Code language:<\/span> CSS<\/span> (<\/span>css<\/span>)<\/span><\/small><\/pre>\n\n\n
              This means it predicted class 0<\/code> (e.g., Setosa for Iris dataset).<\/p>\n\n\n\n
              \n\n\n\n
              \ud83e\udde0 Summary of Steps:<\/h2>\n\n\n\n
              Step<\/th>Action<\/th><\/tr><\/thead>
              \u2705 1<\/td>Register model and set alias (production<\/code>)<\/td><\/tr>
              \u2705 2<\/td>Enable model serving in Databricks UI<\/td><\/tr>
              \u2705 3<\/td>Copy REST endpoint URL<\/td><\/tr>
              \u2705 4<\/td>Send test prediction via Python using requests<\/code><\/td><\/tr><\/tbody><\/table><\/figure>\n\n\n\n
              \n\n\n\n
              Would you like me to generate a ready-to-run notebook (.dbc) with this entire process? Or help you test it directly with your live Databricks instance?<\/p>\n","protected":false},"excerpt":{"rendered":"
              Absolutely! Let\u2019s walk through a complete, step-by-step tutorial to help you understand MLflow on Databricks from start to finish, using revised and working code. We’ll cover all key components of… <\/p>\n","protected":false},"author":1,"featured_media":0,"comment_status":"open","ping_status":"","sticky":false,"template":"","format":"standard","meta":{"_joinchat":[],"footnotes":""},"categories":[2],"tags":[],"class_list":["post-49012","post","type-post","status-publish","format-standard","hentry","category-uncategorised"],"_links":{"self":[{"href":"https:\/\/www.devopsschool.com\/blog\/wp-json\/wp\/v2\/posts\/49012","targetHints":{"allow":["GET"]}}],"collection":[{"href":"https:\/\/www.devopsschool.com\/blog\/wp-json\/wp\/v2\/posts"}],"about":[{"href":"https:\/\/www.devopsschool.com\/blog\/wp-json\/wp\/v2\/types\/post"}],"author":[{"embeddable":true,"href":"https:\/\/www.devopsschool.com\/blog\/wp-json\/wp\/v2\/users\/1"}],"replies":[{"embeddable":true,"href":"https:\/\/www.devopsschool.com\/blog\/wp-json\/wp\/v2\/comments?post=49012"}],"version-history":[{"count":9,"href":"https:\/\/www.devopsschool.com\/blog\/wp-json\/wp\/v2\/posts\/49012\/revisions"}],"predecessor-version":[{"id":49021,"href":"https:\/\/www.devopsschool.com\/blog\/wp-json\/wp\/v2\/posts\/49012\/revisions\/49021"}],"wp:attachment":[{"href":"https:\/\/www.devopsschool.com\/blog\/wp-json\/wp\/v2\/media?parent=49012"}],"wp:term":[{"taxonomy":"category","embeddable":true,"href":"https:\/\/www.devopsschool.com\/blog\/wp-json\/wp\/v2\/categories?post=49012"},{"taxonomy":"post_tag","embeddable":true,"href":"https:\/\/www.devopsschool.com\/blog\/wp-json\/wp\/v2\/tags?post=49012"}],"curies":[{"name":"wp","href":"https:\/\/api.w.org\/{rel}","templated":true}]}}