Model Training & Tuning

Master MLflow experiment tracking, Databricks AutoML, Hyperopt for hyperparameter tuning, and distributed training patterns — covering approximately 30% of the exam (Experimentation domain).

MLflow Experiment Tracking

MLflow is the backbone of ML experimentation on Databricks. The exam heavily tests your knowledge of the MLflow Tracking API, including runs, parameters, metrics, artifacts, and experiment organization.

Core MLflow Concepts

• Experiment — A named container for related runs (e.g., one experiment per ML project or model)
• Run — A single execution of ML code that logs parameters, metrics, and artifacts
• Parameters — Input settings (e.g., learning rate, number of trees)
• Metrics — Output measurements (e.g., accuracy, RMSE, F1 score)
• Artifacts — Output files (e.g., model files, plots, data samples)
• Tags — Key-value metadata for organizing and searching runs

Logging an Experiment Run

import mlflow
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import accuracy_score, f1_score

# Set the experiment (creates if not exists)
mlflow.set_experiment("/Users/team/churn-prediction")

with mlflow.start_run(run_name="rf-baseline") as run:
    # Log parameters
    mlflow.log_param("n_estimators", 100)
    mlflow.log_param("max_depth", 10)
    mlflow.log_param("min_samples_split", 5)

    # Train model
    rf = RandomForestClassifier(n_estimators=100, max_depth=10)
    rf.fit(X_train, y_train)

    # Log metrics
    y_pred = rf.predict(X_test)
    mlflow.log_metric("accuracy", accuracy_score(y_test, y_pred))
    mlflow.log_metric("f1_score", f1_score(y_test, y_pred))

    # Log the model
    mlflow.sklearn.log_model(rf, "model")

    # Log artifacts
    mlflow.log_artifact("confusion_matrix.png")

    print(f"Run ID: {run.info.run_id}")

MLflow Autologging

Databricks supports automatic logging for popular frameworks:

# Enable autologging for scikit-learn
mlflow.sklearn.autolog()

# Enable autologging for all supported frameworks
mlflow.autolog()

# Now any model training automatically logs params, metrics, and artifacts
rf = RandomForestClassifier(n_estimators=200)
rf.fit(X_train, y_train)  # automatically logged!

Databricks AutoML

Databricks AutoML automatically prepares data, trains multiple models, and generates notebooks with the best-performing approaches. It is a glass-box solution — all code is visible and editable.

Using AutoML

from databricks import automl

# Classification
summary = automl.classify(
    dataset=train_df,
    target_col="churned",
    primary_metric="f1",
    timeout_minutes=30,
    max_trials=50
)

# Access the best model
best_model = summary.best_trial
print(f"Best F1: {best_model.metrics['test_f1_score']}")
print(f"Best run: {best_model.mlflow_run_id}")

AutoML Key Features

• Glass-box approach — Generates editable notebooks for the best trials, not black-box models
• Automatic feature engineering — Handles missing values, one-hot encoding, and feature selection
• Multiple algorithms — Tests XGBoost, LightGBM, sklearn, and other frameworks automatically
• MLflow integration — All trials are logged to MLflow for comparison
• Data exploration notebook — Generates a summary statistics and visualization notebook

Hyperopt for Hyperparameter Tuning

Hyperopt is the recommended library for distributed hyperparameter optimization on Databricks. It uses Bayesian optimization (Tree of Parzen Estimators) to efficiently search the parameter space.

Core Hyperopt Components

• fmin() — The main function that minimizes an objective function
• hp.choice() — Choose from a list of discrete options
• hp.uniform() — Uniform distribution over a continuous range
• hp.loguniform() — Log-uniform distribution (for learning rates)
• hp.quniform() — Quantized uniform (for integer parameters)
• SparkTrials — Distributes trials across Spark workers
• Trials — Tracks results on a single machine

Distributed Hyperopt Example

from hyperopt import fmin, tpe, hp, SparkTrials, STATUS_OK
import numpy as np

def objective(params):
    with mlflow.start_run(nested=True):
        rf = RandomForestClassifier(
            n_estimators=int(params["n_estimators"]),
            max_depth=int(params["max_depth"]),
            min_samples_split=int(params["min_samples_split"])
        )
        rf.fit(X_train, y_train)
        accuracy = accuracy_score(y_test, rf.predict(X_test))

        mlflow.log_params(params)
        mlflow.log_metric("accuracy", accuracy)

        # Hyperopt MINIMIZES, so return negative accuracy
        return {"loss": -accuracy, "status": STATUS_OK}

search_space = {
    "n_estimators": hp.quniform("n_estimators", 50, 500, 50),
    "max_depth": hp.quniform("max_depth", 3, 20, 1),
    "min_samples_split": hp.quniform("min_samples_split", 2, 10, 1)
}

# Distribute across Spark workers
spark_trials = SparkTrials(parallelism=4)

with mlflow.start_run(run_name="hyperopt-tuning"):
    best_params = fmin(
        fn=objective,
        space=search_space,
        algo=tpe.suggest,
        max_evals=50,
        trials=spark_trials
    )

Distributed Training

Databricks supports several approaches to distributed model training:

Spark ML (Native Distributed)

Spark ML algorithms are natively distributed. They use the Spark cluster for both data processing and training:

from pyspark.ml.classification import RandomForestClassifier as SparkRF

spark_rf = SparkRF(
    featuresCol="features",
    labelCol="label",
    numTrees=100,
    maxDepth=10
)
model = spark_rf.fit(train_spark_df)  # distributed across workers

Pandas API on Spark (pandas UDFs)

Train single-node models on partitions of data:

import pandas as pd
from pyspark.sql.functions import pandas_udf

@pandas_udf("double")
def predict_udf(features: pd.Series) -> pd.Series:
    # Model is broadcast to each worker
    return pd.Series(model.predict(features.tolist()))

# Apply predictions in parallel across partitions
predictions_df = spark_df.withColumn("prediction", predict_udf("features"))

Horovod / TorchDistributor

For deep learning, Databricks integrates with distributed training frameworks:

from databricks.connect import DatabricksSession
from pyspark.ml.torch.distributor import TorchDistributor

def train_fn():
    import torch
    # Training logic here
    return model

distributor = TorchDistributor(
    num_processes=4,
    local_mode=False,
    use_gpu=True
)
trained_model = distributor.run(train_fn)

Practice Questions

Question 1 — MLflow Tracking

Question 2 — Hyperopt

Question 3 — AutoML

Question 4 — Distributed Training

Question 5 — MLflow Logging