Databricks Machine Learning Associate

Databricks Runtime ML

The standard runtime plus the ML stack, preinstalled and version-tested: scikit-learn, XGBoost, LightGBM, PyTorch, TensorFlow, MLflow and the Feature Engineering client. GPU variants add CUDA.

• No dependency resolving of your own; distributed training and autologging ready to use. AutoML is included up to Runtime 17.x ML and was removed in 18.0 ML.
• %pip install extras in a notebook, scoped to that session.
• AutoML runs on Runtime ML with dedicated access mode for Unity Catalog data; standard access mode isn't supported.

AutoML

AutoML trains many models on a table, logs every trial to MLflow and hands you code, not a black box: automl.classify(dataset, target_col=…, timeout_minutes=…) (or regress, forecast).

• Problems: classification, regression and forecasting.
• Output: a data exploration notebook (distributions, missing values, correlations) and an editable notebook per trial, with the preprocessing it did (imputation, type detection) written out.
• Default split: random 60% train / 20% validation / 20% test; a time_col splits chronologically.

import mlflow
from databricks.feature_engineering import (
    FeatureEngineeringClient, FeatureLookup,
)

fe = FeatureEngineeringClient()
name, key = "main.ml.customer_features", "customer_id"

fe.create_table(name=name, primary_keys=[key],
                schema=features_df.schema)
# "merge", the only mode, upserts on the key
fe.write_table(name=name, df=features_df,
               mode="merge")

# labels_df: customer_id, churned
training_set = fe.create_training_set(
    df=labels_df,
    feature_lookups=[FeatureLookup(
        table_name=name, lookup_key=key,
        feature_names=["tenure_months", "avg_spend"])],
    label="churned", exclude_columns=[key],
)
model = train(training_set.load_df())  # as usual

# the lookups are stored with the model
fe.log_model(model=model, artifact_path="model",
             flavor=mlflow.sklearn,
             training_set=training_set,
             registered_model_name="main.ml.churn")

# pass keys only: the features are looked up again
preds = fe.score_batch(
    model_uri="models:/main.ml.churn@champion",
    df=ids_df)

Offline vs online feature tables

Publish to an online store when an endpoint must look up features that aren't sent in the request.

MLflow tracking

An experiment holds runs; a run holds parameters, metrics (optionally per step), tags, artifacts and models.

import mlflow

mlflow.set_experiment("/Users/me/churn")
with mlflow.start_run(run_name="gbt-depth5") as run:
    mlflow.log_param("max_depth", 5)
    mlflow.log_metric("val_f1", 0.81)
    mlflow.log_figure(fig, "roc.png")
    mlflow.sklearn.log_model(model, name="model")

Autologging is on for notebook training runs on Runtime ML, capturing parameters, metrics and the model with no logging code. Elsewhere (serverless included) call mlflow.autolog(); mlflow.autolog(disable=True) stops it.

The MLflow UI shows, per run: parameters, metrics, tags, artifacts, the logged model and its signature (its own object in MLflow 3), the source notebook and revision, the dataset, and a sortable table to compare runs.

Find the best run with the client API

best = mlflow.search_runs(
    experiment_names=["/Users/me/churn"],
    order_by=["metrics.val_f1 DESC"],
    max_results=1,
)
run_id = best.loc[0, "run_id"]

mlflow.search_runs returns a pandas DataFrame; MlflowClient().search_runs(...) returns Run objects. Both take a filter_string, such as "metrics.val_f1 > 0.8".

Register a model in Unity Catalog

mlflow.set_registry_uri("databricks-uc")

# name is catalog.schema.model
mlflow.register_model(
    model_uri=f"runs:/{run_id}/model",
    name="main.ml.churn",
)

Or register as you log, with registered_model_name=. Each registration adds the next version. A model in Unity Catalog needs a signature; autologging and log_model(input_example=…) add one.

Privileges: USE CATALOG + USE SCHEMA + CREATE MODEL to create a model; ownership or CREATE MODEL VERSION to add versions; EXECUTE on the model to load or serve it.

Aliases and tags

from mlflow import MlflowClient
c, m = MlflowClient(), "main.ml.churn"

c.set_registered_model_alias(m, "champion", 3)
c.set_registered_model_alias(m, "challenger", 4)
# promote: repoint the alias, don't retrain
c.set_registered_model_alias(m, "champion", 4)
c.delete_registered_model_alias(m, "challenger")

c.set_registered_model_tag(m, "team", "growth")
c.delete_registered_model_tag(m, "team")
c.set_model_version_tag(m, "4", "validated", "true")

Load models:/main.ml.churn@champion (alias) or models:/main.ml.churn/4 (version). Batch jobs that load the alias pick up the new version on their next run; an endpoint serves a fixed version, so update it to the version the alias now points to.

Unity Catalog vs workspace model registry