knowledge_base: add model_serving_endpoint example

knowledge_base/model_serving_endpoint/.gitignore

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

knowledge_base/model_serving_endpoint/README.md

@@ -0,0 +1,66 @@
+# Model serving endpoint
+
+This Declarative Automation Bundles example shows how to serve a Unity Catalog model on a Model Serving endpoint.
+A serverless job trains a small scikit-learn model, registers a version, and sets the alias `champion`.
+The endpoint uses AI Gateway rate limits, an inference table, usage tracking, and `CAN_QUERY` permissions.
+
+For more information, refer to the [Model Serving documentation](https://docs.databricks.com/aws/en/machine-learning/model-serving/).
+
+## Prerequisites
+
+* Databricks CLI v1.20.0 or above
+* A workspace with Unity Catalog, Model Serving, and serverless jobs enabled
+* An existing catalog and schema (default: `main.default`)
+* `USE CATALOG`, plus `USE SCHEMA`, `CREATE MODEL`, and `CREATE TABLE` on the schema
+
+## Usage
+
+The endpoint needs a model version, but a new model has no versions.
+Thus, the first deploy has two deploy phases: steps 1 and 4.
+
+1. Deploy only the model, the experiment, and the training job:
+
+   ```sh
+   databricks bundle deploy --select registered_models.model,experiments.train_model_experiment,jobs.train_model
+   ```
+
+2. Run the training job:
+
+   ```sh
+   databricks bundle run train_model
+   ```
+
+3. Read the version number in the job output, for example `--var model_version=1`.
+   The job sets the `champion` alias on this version.
+
+4. Deploy the full bundle. Set `model_version` to the `champion` version:
+
+   ```sh
+   databricks bundle deploy --var model_version=1
+   ```
+
+5. Query the endpoint. In `dev`, the endpoint name has the prefix `dev_<your_user_name>_`:
+
+   ```sh
+   databricks serving-endpoints query dev_<your_user_name>_iris-classifier --json '{
+     "dataframe_records": [
+       {"sepal length (cm)": 5.1, "sepal width (cm)": 3.5, "petal length (cm)": 1.4, "petal width (cm)": 0.2}
+     ]
+   }'
+   ```
+
+   The response has one class number for each record, for example `{"predictions": [0]}`.
+
+6. Remove the resources. `bundle destroy` keeps the registered model while it has versions, so delete the versions and the model:
+
+   ```sh
+   databricks bundle destroy
+   databricks model-versions delete <catalog>.<schema>.<model_name> <version>
+   databricks registered-models delete <catalog>.<schema>.<model_name>
+   ```
+
+   In `dev`, `<model_name>` has the prefix `dev_<your_user_name>_`. Also drop the `<model_name>_payload` inference table.
+
+To serve a new model version, do steps 2 to 4 again with the new version number.
+To roll back, deploy again with the prior version number.
+Use `-t prod` for the `prod` target.

knowledge_base/model_serving_endpoint/databricks.yml

@@ -0,0 +1,62 @@
+# This is a Declarative Automation Bundle definition for model_serving_endpoint.
+# See https://docs.databricks.com/dev-tools/bundles/index.html for documentation.
+bundle:
+  name: model_serving_endpoint
+  engine: direct
+
+variables:
+  catalog:
+    description: Existing Unity Catalog catalog for the model and the inference table
+    default: main
+  schema:
+    description: Existing Unity Catalog schema for the model and the inference table
+    default: default
+  model_name:
+    description: Name of the Unity Catalog registered model (without catalog and schema)
+    default: iris_classifier
+  model_version:
+    description: >-
+      Model version that the endpoint serves. Serving endpoints take a version number,
+      not an alias. Set it to the version that holds the `champion` alias.
+    default: "1"
+  endpoint_name:
+    description: Name of the serving endpoint. Must be unique in the workspace.
+    default: iris-classifier
+  consumers_group:
+    description: Workspace group that gets CAN_QUERY on the endpoint
+    default: users
+  workload_size:
+    description: Provisioned concurrency range of the served model (Small, Medium, Large)
+    default: Small
+  scale_to_zero:
+    description: Scale the served model down to zero when it gets no traffic
+    default: true
+  rate_limit_calls:
+    description: Maximum calls per minute for the full endpoint
+    default: 100
+  rate_limit_user_calls:
+    description: Maximum calls per minute for each user
+    default: 20
+
+include:
+  - resources/*.yml
+
+targets:
+  dev:
+    mode: development
+    default: true
+    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}
+    permissions:
+      - user_name: ${workspace.current_user.userName}
+        level: CAN_MANAGE
+    variables:
+      workload_size: Medium
+      scale_to_zero: false
+      rate_limit_calls: 1000
+      rate_limit_user_calls: 100

knowledge_base/model_serving_endpoint/resources/model.model_serving_endpoint.yml

@@ -0,0 +1,40 @@
+resources:
+  model_serving_endpoints:
+    model_endpoint:
+      name: ${var.endpoint_name}
+
+      config:
+        served_entities:
+          - name: iris-classifier
+            # A Unity Catalog model needs the full three-level name.
+            entity_name: ${var.catalog}.${var.schema}.${resources.registered_models.model.name}
+            # The endpoint takes a version number, not an alias such as `champion`.
+            entity_version: ${var.model_version}
+            workload_size: ${var.workload_size}
+            scale_to_zero_enabled: ${var.scale_to_zero}
+
+      ai_gateway:
+        rate_limits:
+          - key: endpoint
+            calls: ${var.rate_limit_calls}
+            renewal_period: minute
+          - key: user
+            calls: ${var.rate_limit_user_calls}
+            renewal_period: minute
+        # Log requests and responses to <catalog>.<schema>.<prefix>_payload.
+        inference_table_config:
+          enabled: true
+          catalog_name: ${var.catalog}
+          schema_name: ${var.schema}
+          table_name_prefix: ${resources.registered_models.model.name}
+        # Write usage to the system.serving system tables.
+        usage_tracking_config:
+          enabled: true
+
+      tags:
+        - key: example
+          value: model_serving_endpoint
+
+      permissions:
+        - group_name: ${var.consumers_group}
+          level: CAN_QUERY

knowledge_base/model_serving_endpoint/resources/model.registered_model.yml

@@ -0,0 +1,7 @@
+resources:
+  registered_models:
+    model:
+      catalog_name: ${var.catalog}
+      schema_name: ${var.schema}
+      name: ${var.model_name}
+      comment: Iris classifier served by the model_serving_endpoint example (${bundle.target} target)

knowledge_base/model_serving_endpoint/resources/train_model.experiment.yml

@@ -0,0 +1,4 @@
+resources:
+  experiments:
+    train_model_experiment:
+      name: ${workspace.root_path}/train_model

knowledge_base/model_serving_endpoint/resources/train_model.job.yml

@@ -0,0 +1,36 @@
+resources:
+  jobs:
+    train_model:
+      name: train_model
+
+      parameters:
+        - name: model_full_name
+          default: ${var.catalog}.${var.schema}.${resources.registered_models.model.name}
+        - name: alias
+          default: champion
+        - name: experiment_name
+          default: ${resources.experiments.train_model_experiment.name}
+
+      tasks:
+        - task_key: train
+          description: Train a small scikit-learn model, register a new version in Unity Catalog, and set the alias
+          environment_key: serverless_env
+          spark_python_task:
+            python_file: ../src/train_model.py
+            parameters:
+              - --model-full-name
+              - "{{job.parameters.model_full_name}}"
+              - --alias
+              - "{{job.parameters.alias}}"
+              - --experiment-name
+              - "{{job.parameters.experiment_name}}"
+
+      environments:
+        - environment_key: serverless_env
+          spec:
+            environment_version: "6"
+            dependencies:
+              - mlflow>=3.1
+              - scikit-learn
+
+      max_concurrent_runs: 1

knowledge_base/model_serving_endpoint/src/train_model.py

@@ -0,0 +1,42 @@
+"""Train a small scikit-learn model and register it in Unity Catalog.
+
+Each run creates a new model version and moves the alias to that version.
+"""
+
+import argparse
+
+import mlflow
+from mlflow import MlflowClient
+from mlflow.models import infer_signature
+from sklearn.datasets import load_iris
+from sklearn.linear_model import LogisticRegression
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--model-full-name", required=True)
+parser.add_argument("--alias", default="champion")
+parser.add_argument("--experiment-name", required=True)
+args = parser.parse_args()
+
+mlflow.set_registry_uri("databricks-uc")
+mlflow.set_experiment(args.experiment_name)
+
+X, y = load_iris(return_X_y=True, as_frame=True)
+model = LogisticRegression(max_iter=200).fit(X, y)
+
+# Unity Catalog requires a model signature.
+signature = infer_signature(X, model.predict(X))
+
+with mlflow.start_run():
+    model_info = mlflow.sklearn.log_model(
+        sk_model=model,
+        name="model",
+        signature=signature,
+        input_example=X.head(2),
+        registered_model_name=args.model_full_name,
+    )
+
+version = model_info.registered_model_version
+MlflowClient().set_registered_model_alias(args.model_full_name, args.alias, version)
+
+print(f"Registered {args.model_full_name} version {version} with alias '{args.alias}'.")
+print(f"Deploy this version with: --var model_version={version}")