> This page is for Kumo Engine. > For clean Markdown of any page, append .md to the page URL. > For a complete documentation index, see https://docs.nvidia.com/sdgm/llms.txt. > For AI client integration (Claude Code, Cursor, etc.), connect to the MCP server at https://docs.nvidia.com/sdgm/_mcp/server. # Databricks Jobs Recipe > Deploy a Databricks Asset Bundle that trains a Kumo Engine model on demand and scores daily with Unity Catalog Deploy two Databricks jobs that train and score a Kumo Engine model with MLOps turned on. When you finish this page, a train job fits the model on your Unity Catalog tables, stores it on a Unity Catalog volume, records the fit in your workspace's MLflow, and registers it in Unity Catalog under the `candidate` alias. A score job loads the version that holds `champion`, scores your users with `trainer.predict()` as in [Score in Batch](/sdgm/kumo-engine/user-code-guide/operate-models/score), and appends the scores to a Delta table. It is set to run daily at 03:00 UTC once you turn its schedule on. The jobs come from the Databricks Asset Bundle in [`examples/mlops/databricks`](https://github.com/NVIDIA/kumo-engine/tree/main/examples/mlops/databricks) of the Kumo Engine repository. They run the same `kumo.MLOps` workflow as [Turn On MLOps](/sdgm/kumo-engine/user-code-guide/operate-models/turn-on) and [Register and Promote Models](/sdgm/kumo-engine/user-code-guide/operate-models/registry), as two Python scripts with command line flags. ## Before You Start You need: * The [Databricks CLI](https://docs.databricks.com/aws/en/dev-tools/cli/install) with bundle support, authenticated to your workspace as the identity the jobs run as. * Three tables in one Unity Catalog schema, as in the example graph: `users` with primary key `USER_ID`, `items` with primary key `ITEM_ID`, and `orders` with primary key `ORDER_ID`, time column `ORDER_TIME`, and the foreign keys `USER_ID` and `ITEM_ID`. To use other table names with a common prefix, such as `demo_users`, set the `table_prefix` variable. To use your own graph, see [Adapt the Template](#adapt-the-template). * A Unity Catalog volume in that schema, which holds the artifact root and the installation files. * A SQL warehouse, which reads the source tables and writes the predictions table. * A workspace folder for the MLflow experiment. The experiment is created on the first run, but its parent folder must already exist. * The `kumo-engine` wheel, from PyPI or the [Kumo Engine releases](https://github.com/NVIDIA/kumo-engine/releases) page. [Turn On MLOps](/sdgm/kumo-engine/user-code-guide/operate-models/turn-on#install-the-mlops-extras) shows how to check that a build includes MLOps. The identity that runs the jobs needs these privileges: | For | Privileges | | ------------------------ | --------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | | Source tables | `USE CATALOG`, `USE SCHEMA`, and `SELECT` on the three tables | | Predictions table | `CREATE TABLE` on the schema. The score job creates and owns the predictions table. If the table already exists under another owner, grant `SELECT` and `MODIFY` on that table only | | Volume | `READ VOLUME` and `WRITE VOLUME` on the volume, and `CREATE VOLUME` on the schema or a volume named `_kumo_staging` in the schema, which the score job uses to stage table writes | | SQL warehouse and MLflow | `CAN USE` on the warehouse, and permission to edit the experiment | | Model | `CREATE MODEL` on the schema. The first fit creates the model, and its owner can register each later version and move its aliases. A model that already exists must be owned by this identity | ## Get the Template Clone the Kumo Engine repository and change to the bundle directory: ```bash git clone https://github.com/NVIDIA/kumo-engine.git cd kumo-engine/examples/mlops/databricks ``` The directory holds `databricks.yml`, the bundle definition. The bundle deploys the whole `examples/mlops` directory: `fit_job.py`, `score_job.py`, and `jobs.py`, the module both scripts share. ## Put Kumo Engine on the Volume Each job task installs Kumo Engine on its job cluster from two files on a volume: the `kumo-engine` wheel, and a `requirements.txt` that installs the wheel with its dependencies. Create a `requirements.txt` that names the wheel's volume path with the `databricks` and `mlflow` extras. Replace ``, ``, ``, and `` with your volume and the wheel's file name: ```text --prefer-binary /Volumes////wheels/.whl[databricks,mlflow] ``` Leave `torch` out: the Databricks Runtime ML image already provides it. Copy both files to the volume: ```bash databricks fs cp .whl dbfs:/Volumes////wheels/ databricks fs cp requirements.txt dbfs:/Volumes////wheels/ ``` `databricks fs ls dbfs:/Volumes////wheels/` lists both files. ## Set the Bundle Variables Pass the bundle variables with `--var`, or replace the `<...>` defaults in `databricks.yml`. These have no usable default: | Variable | Value | | ----------------------------------- | --------------------------------------------------------------------------------------------------------------------------------------------- | | `catalog`, `schema` | The catalog and schema that hold the source tables. The model and the predictions table are created there too | | `volume` | The volume in that schema | | `warehouse_id` | The SQL warehouse ID | | `kumo_engine_wheel`, `requirements` | The volume paths of the two files from the previous step | | `spark_version` | Databricks Runtime 17.3 LTS ML or later, such as `17.3.x-gpu-ml-scala2.13` on a GPU node type or `17.3.x-cpu-ml-scala2.13` on a CPU node type | | `node_type_id` | A GPU node type, such as `g5.xlarge` on AWS, or a CPU node type with local NVMe disks, such as `m5d.xlarge` on AWS | These have defaults that you can override: | Variable | Default | | ------------------- | ------------------------------------------------------------------------ | | `artifact_root` | `/Volumes////kumo` | | `experiment` | `/Shared/kumo-engine/users-churn` | | `model_name` | `users_churn`, registered as `..users_churn` | | `predictions_table` | `users_churn_predictions`, created in `.` | | `table_prefix` | Empty (no prefix) | | `query` | `PREDICT COUNT(orders.*, 0, 7, days) > 0 FOR EACH users.USER_ID` | | `alias` | `candidate`, the alias each fit moves to its new version | | `score_alias` | `champion`, the alias the score job loads | | `on_error` | `raise`, so a failed store, tracking call, or registration fails the job | Keep the variables in a shell variable so that every bundle command uses the same values. This example uses the catalog `main`, the schema `ml`, and the volume `kumo`, and the commands below reuse those names; replace them with yours throughout. Replace `` and ``, and set `experiment` to a path whose parent folder exists, such as a folder under your workspace user folder, where `` is your workspace user name: ```bash VARS="--var catalog=main --var schema=ml --var volume=kumo \ --var warehouse_id= \ --var kumo_engine_wheel=/Volumes/main/ml/kumo/wheels/.whl \ --var requirements=/Volumes/main/ml/kumo/wheels/requirements.txt \ --var spark_version=17.3.x-gpu-ml-scala2.13 --var node_type_id=g5.xlarge \ --var experiment=/Users//users-churn" ``` ## Deploy the Jobs Validate the bundle, then deploy it to the `dev` target, which places the bundle under your user folder: ```bash databricks bundle validate $VARS databricks bundle deploy $VARS ``` `validate` prints `Validation OK!`, and `deploy` creates two jobs, `kumo-engine-train` and `kumo-engine-score`. `databricks bundle summary $VARS` lists both with their URLs. Deploying the same bundle again as the same identity updates these two jobs instead of creating new ones. To keep two deployments apart, change `bundle.name` or the target's `root_path` in `databricks.yml`. ## Train and Register a Model Run the train job once: ```bash databricks bundle run kumo_train $VARS ``` The command waits for the run and prints its URL and final state, `TERMINATED SUCCESS`. Open the URL: the task's output ends with `model: models/`. The fit is stored under the artifact root, recorded as a run in the experiment, and registered as version `1` of `main.ml.users_churn` with the `candidate` alias. The run is named after the version, for example `fit-users_churn-v1`. Run the train job again whenever you want a new model. Each run registers a new version and moves `candidate` to it. ## Promote a Version The score job loads the version that holds `champion`, so a new version reaches the score job only when you move `champion` to it. For the first version, set the alias with the Databricks CLI: ```bash databricks registered-models set-alias main.ml.users_churn champion 1 ``` To check the aliases, run `databricks registered-models get main.ml.users_churn --include-aliases`: its `aliases` list `candidate` and `champion`, each with `version_num` 1. For later versions, compare each candidate with the current champion before you move the alias, as described in [Promote a Version](/sdgm/kumo-engine/user-code-guide/operate-models/registry#promote-a-version). To skip the review and score every new fit, set `--var alias=champion`. Each fit then moves `champion` itself. ## Score Run the score job once to check it: ```bash databricks bundle run kumo_score $VARS ``` The run ends with `TERMINATED SUCCESS`, and the task's output ends with `predictions: rows`. The job builds a fresh snapshot of the source tables, scores every user with the champion, stores the prediction run under the artifact root, and records it in the experiment, named for example `predict-users_churn-v1-`. It also appends the scores to the predictions table, which it creates on the first run. Query the table in the SQL editor: ```sql SELECT PREDICTION_ID, BUNDLE_ID, ANCHOR_TIMESTAMP, COUNT(*) AS users FROM main.ml.users_churn_predictions GROUP BY ALL; ``` Each run adds one group of rows. `PREDICTION_ID` is the prediction run's id, `BUNDLE_ID` is the model id, and `ANCHOR_TIMESTAMP` is the time the scores predict from. The other columns are `ENTITY`, `TARGET_PRED`, `False_PROB`, and `True_PROB`. ## Turn On the Daily Schedule The score job is deployed with its schedule paused, so it runs only when you start it. To score daily at 03:00 UTC, change `pause_status: PAUSED` to `pause_status: UNPAUSED` under `kumo_score` in `databricks.yml` and deploy again with `databricks bundle deploy $VARS`. To run at another time, change `quartz_cron_expression` in the same block. The train job has no schedule. Run it when you want a new model, or add a `schedule` block to it in the same way. ## Adapt the Template The two scripts share one module, `jobs.py`. Edit it to use your own data: * `build_graph` defines the tables, primary keys, time columns, and edges. Keep the alias `users` for the entity table if you keep the default query. * `QUERY` is the default predictive query; the `query` variable overrides it per deployment. * `TEXT_COLUMNS` lists columns to encode as text, such as `('users.BIO',)`. Text columns also need the `text` extra: install the wheel with `[databricks,mlflow,text]` in `requirements.txt`. To score a stored snapshot instead of a fresh one, add `--snapshot ` to the score task's parameters. ## Remove the Jobs To delete both jobs and the deployed files, run: ```bash databricks bundle destroy $VARS --auto-approve ``` The command lists what it deletes and ends with a line that confirms the deletion. It leaves the artifact root, the registered model, the experiment, and the predictions table in place. ## Next Steps #### [Register and Promote Models](/sdgm/kumo-engine/user-code-guide/operate-models/registry) Compare a new version with the champion before you move the alias. #### [MLOps Reference](/sdgm/kumo-engine/user-code-guide/operate-models/reference) Look up what each stored artifact records and where it is stored.