> 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.

# Production on Databricks

> Integrate Engine model artifacts into production jobs, provisioning and observability

Run Engine in a single-node GPU Lakeflow Job. Use MLflow for tracking and Unity Catalog for model registration and retained artifacts.

## Prepare Databricks

Adapt the following steps to your workspace configuration. This guide uses dedicated, single-node GPU compute and the MLflow Unity Catalog registration path.

### Provision resources

1. Choose a compatible GPU runtime from [Databricks setup](/sdgm/kumo-engine/get-started-with-kumo-engine/installation/databricks-notebook).
2. Size the node using [GPU sizing](/sdgm/kumo-engine/get-started-with-kumo-engine/gpu-sizing). Include CPU, RAM, GPU memory and local SSD capacity. For example, the guidance for 100 GiB of source Parquet recommends 32 vCPUs, at least 150 GiB RAM and 400 GiB free local SSD scratch, increasing to 800 GiB for graphs with many edges. GPU memory depends on the model.
3. Place scratch on the node's local SSD and retained artifacts in a Unity Catalog volume. Network storage and cloud block volumes are not supported as Engine working storage.
4. Provision a run-as service principal, an MLflow experiment, a model namespace, an artifact volume and an output schema. Identify the source tables the job will read.

### Configure access

1. Grant the run-as identity the permissions needed for its tasks:

   | Resource                 | Permissions                                                                                            |
   | ------------------------ | ------------------------------------------------------------------------------------------------------ |
   | Catalog and schema       | `USE CATALOG` and `USE SCHEMA`                                                                         |
   | Source tables            | `SELECT`                                                                                               |
   | Artifact volume          | `READ VOLUME` and `WRITE VOLUME` where the job writes artifacts                                        |
   | Output schema and tables | Table creation and write privileges for the selected delivery method                                   |
   | Registered model         | `EXECUTE` for scoring; creation/version privileges for registration and ownership for alias management |
   | MLflow experiment        | Permission to write tracking runs                                                                      |

2. Configure the connector's SQL warehouse or cluster query target using [Databricks connectors](/sdgm/kumo-engine/user-code-guide/connect-your-data-databricks#choose-a-source). This query target is separate from the GPU job compute.

3. Configure authentication through the job identity or managed secrets. Use the service principal for scheduled execution rather than notebook-user credentials.

### Install dependencies

1. Follow the runtime-specific preparation in [Install Kumo](/sdgm/kumo-engine/get-started-with-kumo-engine/installation/databricks-notebook#install-kumo). Select an Engine build compatible with the runtime and containing the MLOps APIs used by the job.
2. Add `kumo-engine[databricks,mlflow]` to the job's Python library requirements. The `databricks` extra supplies the connector and Databricks SDK; the `mlflow` extra supplies MLflow tracking dependencies and GPU metrics support. Use the same approved package version for training and scoring.
3. Add the `s3` extra if the job reads S3 through the Engine connector: `kumo-engine[databricks,mlflow,s3]`. Include any other source-specific extras identified in the connector setup.
4. Install these requirements on job compute. For notebook tasks, follow the installation guide's Python restart instructions after changing packages.
5. Run the setup guide's import and GPU check in the job environment. Record the installed versions with the deployment configuration.

## Prepare the model and prediction script

1. Select an evaluated model and retain its complete bundle, exact version and required text embedding assets. Store these outside temporary scratch. See [Publish and Retrieve Models](/sdgm/kumo-engine/user-code-guide/train-a-model/publish-and-retrieve-models).

2. Move the working prediction code into `predict_batch.py`, including imports, connector setup, graph definitions and predictive-query definitions. The [production script](/sdgm/kumo-engine/user-code-guide/deploy-a-model/overview#prepare-the-prediction-script) illustrates the structure; replace its Snowflake connectors with your Databricks connectors.

3. Supply the following inputs through your script arguments or job configuration:

   | Input               | Value                                                             |
   | ------------------- | ----------------------------------------------------------------- |
   | Model               | Exact accepted registry version or bundle path                    |
   | Data                | Source tables and recorded source version or retained snapshot    |
   | Prediction time     | Fixed UTC anchor for temporal queries; omitted for static queries |
   | Batch identifier    | One identifier for the scoring request, retained across retries   |
   | Attempt destination | A separate output table or path for each attempt                  |
   | Scratch             | Writable node-local SSD directory                                 |

   These are application inputs, not built-in Engine CLI flags. A model bundle does not contain the source data or graph snapshot.

4. Load the model and build or reload the intended snapshot. Regenerate the prediction population for the batch.

5. Call `Trainer.predict` inside the Kumo MLOps context to retain prediction records for monitoring. Configure a durable artifact root and MLflow tracking destination.

6. Export predictions to the attempt destination. Check the saved keys, population and task-specific row counts before making the output available to the application.

## Deploy the job

1. Create a Python task in Lakeflow Jobs that runs `predict_batch.py`.
2. Configure its run-as identity, GPU compute, Python libraries and application inputs. Start with one concurrent run.
3. Set a task timeout and failure notifications. Add bounded retries after configuring duplicate-safe output delivery.
4. Run the job manually from fresh compute. Review the terminal job status, saved predictions, MLflow run and Kumo prediction records.
5. Measure peak RAM, GPU memory, scratch use and job duration with representative data. Adjust node capacity before increasing concurrency. Additional workers do not automatically distribute an Engine graph.
6. Test an interrupted run and retry before publishing its output. Keep the same model, data version, population and prediction anchor across attempts. Only one attempt should be accepted for each logical batch.

The job runs under the service principal and requires access to the saved model, prediction output location, and MLflow and Kumo tracking stores.

## Schedule predictions and monitor results

1. Schedule the job after the upstream data refresh, or configure a dependency on the refresh task.
2. Record the model version, data or snapshot version, prediction anchor, batch identifier, accepted output location and job status together.
3. Configure operational alerts for failed jobs and missed freshness deadlines. Track job duration, resource use and requested, scored and skipped rows.
4. Compute score drift from retained Kumo prediction records. Evaluate model performance when labels become available and the prediction outcome window has completed.
5. Export monitoring results to governed tables and configure their refresh schedule, alert thresholds and responsible operator.

Job logs describe execution health; drift and labeled outcomes describe model behavior. Investigate data freshness and population changes before deciding to retrain. Spark UDF or pyfunc scoring does not automatically create Kumo prediction records and can require a full snapshot per executor.

For output delivery, keep each attempt separate until validation succeeds. Publish the accepted location once per batch and retain the previous accepted output on failure. See [Schedule Prediction Jobs](/sdgm/kumo-engine/user-code-guide/deploy-a-model/schedule-prediction-jobs).

## Promote and refresh models

1. Train and register a candidate without changing the production alias.
2. Compare the candidate with the current model using the same evaluation population and outcome definition.
3. Run the candidate in the production job environment, then promote the accepted version. Record the exact version used for each prediction batch.
4. Retain the previous model, embedding assets and runtime configuration. To roll back, select that version for subsequent jobs and restore the previous accepted output where needed.

Follow [MLOps workflows on Databricks](https://docs.databricks.com/aws/en/machine-learning/mlops/mlops-workflow) to separate training, validation, promotion and inference. See [Refresh a Model](/sdgm/kumo-engine/user-code-guide/deploy-a-model/refresh-a-model) for refresh cadence and evaluation guidance.

## Assign ISV operating responsibilities

1. Assign the following responsibilities before deployment:

   | Owner                  | Responsibility                                                |
   | ---------------------- | ------------------------------------------------------------- |
   | Customer administrator | Compute, storage, networking, identities and resource grants  |
   | ISV deployer           | Compatible runtime, job code and deployment configuration     |
   | Production operator    | Schedules, alerts, promotion, rollback, retention and support |

2. For ISV-operated environments, separate customer credentials, artifact roots, outputs and execution resources. A shared Python process does not provide tenant isolation.

3. Store credentials in managed secrets or use workload identity. Define credential rotation and revocation procedures, and exclude tokens and customer rows from diagnostic logs.

4. Agree on compute ownership, artifact retention and deletion permissions.

## Troubleshoot jobs

| Symptom                              | Action                                                                                                  |
| ------------------------------------ | ------------------------------------------------------------------------------------------------------- |
| Import or CUDA failure               | Compare the installed Engine, Python, PyTorch and CUDA versions with the selected runtime requirements. |
| Access denied                        | Review the run-as identity's grants on the affected source, model, volume or output.                    |
| Scratch exhaustion                   | Increase local SSD capacity or reduce overlapping jobs.                                                 |
| Job submitted but no output          | Check terminal job status, execution logs and the export destination.                                   |
| Scores exist but monitoring is empty | Check the Kumo MLOps context, artifact root and monitoring access.                                      |