diff --git a/knowledge_base/model_serving_endpoint/.gitignore b/knowledge_base/model_serving_endpoint/.gitignore new file mode 100644 index 0000000..15bcc6d --- /dev/null +++ b/knowledge_base/model_serving_endpoint/.gitignore @@ -0,0 +1 @@ +.databricks diff --git a/knowledge_base/model_serving_endpoint/README.md b/knowledge_base/model_serving_endpoint/README.md new file mode 100644 index 0000000..e641ec4 --- /dev/null +++ b/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__`: + + ```sh + databricks serving-endpoints query dev__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 .. + databricks registered-models delete .. + ``` + + In `dev`, `` has the prefix `dev__`. Also drop the `_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. diff --git a/knowledge_base/model_serving_endpoint/databricks.yml b/knowledge_base/model_serving_endpoint/databricks.yml new file mode 100644 index 0000000..fc85f80 --- /dev/null +++ b/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 diff --git a/knowledge_base/model_serving_endpoint/resources/model.model_serving_endpoint.yml b/knowledge_base/model_serving_endpoint/resources/model.model_serving_endpoint.yml new file mode 100644 index 0000000..c864e2f --- /dev/null +++ b/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 .._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 diff --git a/knowledge_base/model_serving_endpoint/resources/model.registered_model.yml b/knowledge_base/model_serving_endpoint/resources/model.registered_model.yml new file mode 100644 index 0000000..aa2cbe9 --- /dev/null +++ b/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) diff --git a/knowledge_base/model_serving_endpoint/resources/train_model.experiment.yml b/knowledge_base/model_serving_endpoint/resources/train_model.experiment.yml new file mode 100644 index 0000000..76546e0 --- /dev/null +++ b/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 diff --git a/knowledge_base/model_serving_endpoint/resources/train_model.job.yml b/knowledge_base/model_serving_endpoint/resources/train_model.job.yml new file mode 100644 index 0000000..6ba2e11 --- /dev/null +++ b/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 diff --git a/knowledge_base/model_serving_endpoint/src/train_model.py b/knowledge_base/model_serving_endpoint/src/train_model.py new file mode 100644 index 0000000..a81e383 --- /dev/null +++ b/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}")