knowledge_base: add ml_experiments_and_monitoring example

knowledge_base/ml_experiments_and_monitoring/.gitignore

@@ -0,0 +1 @@
+.databricks

knowledge_base/ml_experiments_and_monitoring/README.md

@@ -0,0 +1,60 @@
+# ML experiments and monitoring
+
+This Declarative Automation Bundles example defines an MLflow experiment, a Unity Catalog registered model with grants, a serverless training job, and a quality monitor.
+The training job sets the `champion` alias and writes an inference table that Lakehouse Monitoring watches.
+
+For more information, please refer to the [documentation](https://docs.databricks.com/aws/en/dev-tools/bundles/resources).
+
+## Prerequisites
+
+* Databricks CLI v1.20.0 or above
+* A Unity Catalog schema that you can use to create models and tables
+* An account group for model consumers
+
+## Usage
+
+| Resource | Purpose |
+| --- | --- |
+| `experiments.training` | Stores the training runs. |
+| `registered_models.iris` | Stores model versions and grants `EXECUTE` to the consumers group. |
+| `jobs.train` | Trains the model, sets the `champion` alias, and writes the inference table. |
+| `job_runs.train_on_deploy` | Runs the training job during deploy. |
+| `quality_monitors.predictions` | Monitors the inference table and sends email on failure. |
+
+1. Set the `catalog`, `schema`, and `consumers_group` variables in `databricks.yml`.
+
+2. Deploy the bundle. The deploy runs the training job once, then creates the monitor:
+
+   ```
+   databricks bundle deploy
+   ```
+
+3. Run the training job again to register a new model version:
+
+   ```
+   databricks bundle run train
+   ```
+
+4. Refresh the monitor. Use the inference table name from the job output:
+
+   ```
+   databricks quality-monitors run-refresh <catalog>.<schema>.<model_name>_predictions
+   ```
+
+5. Deploy to production. The `prod` target keeps the daily refresh schedule active:
+
+   ```
+   databricks bundle deploy --target prod
+   ```
+
+6. Delete each model version that the training job registered:
+
+   ```
+   databricks model-versions delete <catalog>.<schema>.<model_name> <version>
+   ```
+
+7. Remove the resources:
+
+   ```
+   databricks bundle destroy
+   ```

knowledge_base/ml_experiments_and_monitoring/databricks.yml

@@ -0,0 +1,32 @@
+bundle:
+  name: ml_experiments_and_monitoring
+  databricks_cli_version: ">= 1.20.0"
+
+include:
+  - resources/*.yml
+
+variables:
+  catalog:
+    description: Catalog for the model and the inference table.
+    default: main
+  schema:
+    description: Existing schema for the model, the inference table, and the monitor output.
+    default: ml_monitoring
+  consumers_group:
+    description: Account group that can use the registered model.
+    default: ml-consumers
+  notification_email:
+    description: Email address that gets monitor failure notifications.
+    default: ${workspace.current_user.userName}
+
+targets:
+  dev:
+    default: true
+    mode: development
+    workspace:
+      root_path: /Workspace/bundle-examples/${bundle.name}/${bundle.target}/${workspace.current_user.short_name}
+
+  prod:
+    mode: production
+    workspace:
+      root_path: /Workspace/bundle-examples/${bundle.name}/${bundle.target}

knowledge_base/ml_experiments_and_monitoring/resources/iris.registered_model.yml

@@ -0,0 +1,12 @@
+resources:
+  registered_models:
+    iris:
+      # Unity Catalog model. The training job registers versions and sets the "champion" alias.
+      catalog_name: ${var.catalog}
+      schema_name: ${var.schema}
+      name: iris_classifier
+      comment: Iris classifier trained by the ml_experiments_and_monitoring example.
+      grants:
+        - principal: ${var.consumers_group}
+          privileges:
+            - EXECUTE

knowledge_base/ml_experiments_and_monitoring/resources/predictions.quality_monitor.yml

@@ -0,0 +1,23 @@
+resources:
+  quality_monitors:
+    predictions:
+      # Lakehouse Monitoring on the inference table that the training job writes.
+      table_name: ${resources.job_runs.train_on_deploy.job_parameters.inference_table}
+      output_schema_name: ${var.catalog}.${var.schema}
+      assets_dir: ${workspace.root_path}/monitoring
+      inference_log:
+        problem_type: PROBLEM_TYPE_CLASSIFICATION
+        timestamp_col: timestamp
+        model_id_col: model_version
+        prediction_col: prediction
+        label_col: label
+        granularities:
+          - 1 day
+      # Refresh daily at 06:00 UTC. Development mode turns off this schedule.
+      schedule:
+        quartz_cron_expression: 0 0 6 * * ?
+        timezone_id: UTC
+      notifications:
+        on_failure:
+          email_addresses:
+            - ${var.notification_email}

knowledge_base/ml_experiments_and_monitoring/resources/train.job.yml

@@ -0,0 +1,30 @@
+resources:
+  jobs:
+    train:
+      name: Train iris classifier
+
+      parameters:
+        - name: experiment_name
+          default: ${resources.experiments.training.name}
+        - name: model_name
+          default: ${var.catalog}.${var.schema}.${resources.registered_models.iris.name}
+        - name: inference_table
+          default: ${var.catalog}.${var.schema}.${resources.registered_models.iris.name}_predictions
+
+      tasks:
+        - task_key: train
+          spark_python_task:
+            python_file: ../src/train.py
+            parameters:
+              - "{{job.parameters.experiment_name}}"
+              - "{{job.parameters.model_name}}"
+              - "{{job.parameters.inference_table}}"
+          environment_key: default
+
+      environments:
+        - environment_key: default
+          spec:
+            environment_version: "6"
+            dependencies:
+              - mlflow>=3
+              - scikit-learn

knowledge_base/ml_experiments_and_monitoring/resources/train.job_run.yml

@@ -0,0 +1,7 @@
+resources:
+  job_runs:
+    # Run training during deploy, so the inference table exists before the monitor is created.
+    train_on_deploy:
+      job_id: ${resources.jobs.train.id}
+      job_parameters:
+        inference_table: ${var.catalog}.${var.schema}.${resources.registered_models.iris.name}_predictions

knowledge_base/ml_experiments_and_monitoring/resources/training.experiment.yml

@@ -0,0 +1,5 @@
+resources:
+  experiments:
+    training:
+      # MLflow experiment that stores the training runs.
+      name: ${workspace.root_path}/training

knowledge_base/ml_experiments_and_monitoring/src/train.py

@@ -0,0 +1,46 @@
+import sys
+from datetime import datetime, timezone
+
+import mlflow
+import pandas as pd
+from databricks.connect import DatabricksSession
+from mlflow import MlflowClient
+from sklearn.datasets import load_iris
+from sklearn.linear_model import LogisticRegression
+from sklearn.model_selection import train_test_split
+
+experiment_name, model_name, inference_table = sys.argv[1:4]
+
+mlflow.set_registry_uri("databricks-uc")
+mlflow.set_experiment(experiment_name)
+
+X, y = load_iris(return_X_y=True, as_frame=True)
+X_train, X_test, y_train, y_test = train_test_split(X, y, random_state=42)
+
+with mlflow.start_run():
+    model = LogisticRegression(max_iter=200).fit(X_train, y_train)
+    mlflow.log_metric("accuracy", model.score(X_test, y_test))
+    info = mlflow.sklearn.log_model(
+        model,
+        name="model",
+        input_example=X_train.head(5),
+        registered_model_name=model_name,
+    )
+
+# Point the "champion" alias at the new version.
+version = info.registered_model_version
+MlflowClient().set_registered_model_alias(model_name, "champion", version)
+print(f"Registered {model_name} version {version} as champion")
+
+# Write scored test rows to the inference table that the monitor watches.
+predictions = pd.DataFrame(
+    {
+        "timestamp": datetime.now(timezone.utc),
+        "model_version": str(version),
+        "prediction": model.predict(X_test).astype(int),
+        "label": y_test.astype(int).to_numpy(),
+    }
+)
+spark = DatabricksSession.builder.getOrCreate()
+spark.createDataFrame(predictions).write.mode("append").saveAsTable(inference_table)
+print(f"Wrote predictions to {inference_table}")