Databricks Jobs Recipe

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, 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 of the Kumo Engine repository. They run the same kumo.MLOps workflow as Turn On MLOps and Register and Promote Models, as two Python scripts with command line flags.

Before You Start

You need:

• The Databricks CLI 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.
• 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 page. Turn On MLOps shows how to check that a build includes MLOps.

The identity that runs the jobs needs these privileges:

Get the Template

Clone the Kumo Engine repository and change to the bundle directory:

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 <catalog>, <schema>, <volume>, and <wheel> with your volume and the wheel’s file name:

--prefer-binary
/Volumes/<catalog>/<schema>/<volume>/wheels/<wheel>.whl[databricks,mlflow]

Leave torch out: the Databricks Runtime ML image already provides it.

Copy both files to the volume:

databricks fs cp <wheel>.whl dbfs:/Volumes/<catalog>/<schema>/<volume>/wheels/
databricks fs cp requirements.txt dbfs:/Volumes/<catalog>/<schema>/<volume>/wheels/

databricks fs ls dbfs:/Volumes/<catalog>/<schema>/<volume>/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:

These have defaults that you can override:

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 <warehouse-id> and <wheel>, and set experiment to a path whose parent folder exists, such as a folder under your workspace user folder, where <you> is your workspace user name:

VARS="--var catalog=main --var schema=ml --var volume=kumo \
  --var warehouse_id=<warehouse-id> \
  --var kumo_engine_wheel=/Volumes/main/ml/kumo/wheels/<wheel>.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/<you>/users-churn"

Deploy the Jobs

Validate the bundle, then deploy it to the dev target, which places the bundle under your user folder:

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:

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/<model id>. 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:

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.

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:

databricks bundle run kumo_score $VARS

The run ends with TERMINATED SUCCESS, and the task’s output ends with predictions: <n> 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-<date>.

It also appends the scores to the predictions table, which it creates on the first run. Query the table in the SQL editor:

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 <snapshot id> to the score task’s parameters.

Remove the Jobs

To delete both jobs and the deployed files, run:

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