> 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 Notebook Setup

> Install Kumo in a Databricks notebook and configure local working storage

Run Kumo in a Python notebook attached to Databricks GPU compute. Kumo prepares graph data and trains on the notebook compute. After setup, choose your data source in [Connect Your Data](/kumo-engine/1.0.0/user-code-guide/connect-your-data-databricks#choose-a-source).

## Prepare the Notebook

Choose either **Databricks Serverless GPU environment version 6** or a single node GPU cluster with **Databricks Runtime 17.3 LTS ML GPU** and dedicated access. Both use Python 3.12. Runtime 17.3 LTS ML is not required when you use Serverless GPU.

For Serverless GPU, select a GPU accelerator and the version 6 AI base environment in the notebook's [Environment panel](https://docs.databricks.com/aws/en/compute/serverless/dependencies). See the [Serverless GPU environment version 6 release notes](https://docs.databricks.com/aws/en/release-notes/serverless/environment-version/six-gpu).

The required compute size depends on your data and workload. Review [GPU sizing and working storage](/kumo-engine/1.0.0/get-started-with-kumo-engine/gpu-sizing) before a larger run. Databricks Runtime 16.4 LTS ML GPU with Python 3.12 is also supported after the manual PyTorch upgrade below.

Databricks permits writes to `/local_disk0` on dedicated or no isolation compute, but not on standard or serverless compute. Use `/tmp/kumoai-local` for scratch on Serverless GPU and `/local_disk0/kumo-scratch` on a classic dedicated GPU cluster. See [where Databricks writes temporary data](https://docs.databricks.com/aws/en/files/write-data).

> **Info**
>
> Confirm that scratch on your selected GPU compute is backed by local storage.
> For classic compute on AWS, disable elastic disk and EBS volumes as described in
> [prerequisites](/kumo-engine/1.0.0/get-started-with-kumo-engine/prerequisites#local-storage).
> Those AWS settings do not apply to Azure compute.

Provide:

* A Unity Catalog volume where you can retain completed artifacts.
* A writable directory on the driver's local storage for Kumo scratch files.
* Credentials and permissions for your chosen data source.

The notebook user needs `USE CATALOG` and `USE SCHEMA` on the volume's catalog and schema, plus `READ VOLUME` and `WRITE VOLUME` on the volume. Source credentials and permissions depend on the connection you choose. See [Databricks connectors](/kumo-engine/1.0.0/user-code-guide/connect-your-data-databricks#choose-a-source) for source-specific configuration.

## Install Kumo

On **Runtime 17.3 LTS ML GPU**, retain the preinstalled PyTorch `2.7.0`, which Kumo supports. PyTorch supplies CUDA support through its CUDA PyPI dependencies, so use `torch.version.cuda` below to inspect the effective target. See the [17.3 runtime library notes](https://learn.microsoft.com/en-us/azure/databricks/release-notes/runtime/17.3lts-ml#python-libraries-on-gpu-clusters).

On **Runtime 16.4 LTS ML GPU only**, manually upgrade the preinstalled PyTorch `2.6` stack before installing Kumo. Run this matched PyTorch `2.9.0`, torchvision, and torchaudio command for CUDA 12.6 from the [official PyTorch installation instructions](https://pytorch.org/get-started/previous-versions/#v290):

```python
%pip install torch==2.9.0 torchvision==0.24.0 torchaudio==2.9.0 \
    --index-url https://download.pytorch.org/whl/cu126
```

After any runtime-specific preparation, install Kumo Engine with the Databricks integration. Run this command in the notebook that will execute Kumo:

```python
%pip install "kumo-engine[databricks]"
```

For either [Databricks S3 connection path](/kumo-engine/1.0.0/user-code-guide/connect-your-data-databricks#databricks-s3), include the S3 integration in the same installation:

```python
%pip install "kumo-engine[databricks,s3]"
```

After installation, restart Python in a separate cell:

```python
dbutils.library.restartPython()
```

The restart clears Python variables. Rerun the configuration cells below and recreate your connector using [Databricks connectors](/kumo-engine/1.0.0/user-code-guide/connect-your-data-databricks#choose-a-source) after it completes. Verify that the runtime imports and a CUDA operation succeed:

```python
import kumo.engine as kumo
import torch
import torch_geometric

if not torch.cuda.is_available():
    raise RuntimeError("No GPU is available to this notebook.")

values = torch.tensor([1.0, 2.0], device="cuda")
groups = torch.tensor([0, 0], device="cuda")
total = torch.zeros(1, device="cuda").scatter_add_(0, groups, values)
assert total.item() == 3.0

print("Kumo:", kumo.__version__)
print("PyTorch:", torch.__version__, "CUDA:", torch.version.cuda)
print("PyG:", torch_geometric.__version__)
print("GPU:", torch.cuda.get_device_name(0))
```

Kumo no longer requires the separate `pyg-lib` or `torch-scatter` packages. Do not install them to satisfy an older setup check. This cell checks imports and GPU execution; it does not verify data access or a full training job.

Save this output when reporting an installation problem. If installation or import fails, check access to the package index and confirm that it supplies a native build compatible with the effective PyTorch and CUDA versions. Attaching a notebook to the GPU cluster does not install packages into its Python kernel.

## Configure Working Storage

Kumo needs local scratch space for staged source data, graph indexes, training files, and local prediction results. Use `/tmp/kumoai-local` on Serverless GPU or `/local_disk0/kumo-scratch` on a classic dedicated GPU cluster. Keep retained model and prediction files in a Unity Catalog volume. Confirm that the chosen scratch path is backed by the intended local storage. Its name alone does not prove the storage type.

If you do not have an artifact volume, have the catalog or schema owner create one in **SQL Editor**. Replace the catalog and schema with existing names, choose a volume name, and use the notebook execution principal for `<notebook-principal>` (the user's email address or the service principal's application ID):

```sql
CREATE VOLUME IF NOT EXISTS `<catalog>`.`<schema>`.`<volume>`;
GRANT USE CATALOG ON CATALOG `<catalog>` TO `<notebook-principal>`;
GRANT USE SCHEMA ON SCHEMA `<catalog>`.`<schema>` TO `<notebook-principal>`;
GRANT READ VOLUME, WRITE VOLUME
  ON VOLUME `<catalog>`.`<schema>`.`<volume>` TO `<notebook-principal>`;
```

Replace the volume path with one the notebook execution principal can read and write. `ARTIFACT_ROOT` is a destination for later explicit model publication; assigning this variable does not copy or retain artifacts:

```python
from pathlib import Path
import shutil
import tempfile

SCRATCH_DIR = Path("/tmp/kumoai-local")  # Serverless GPU.
# Classic dedicated GPU cluster: Path("/local_disk0/kumo-scratch").
ARTIFACT_ROOT = Path("/Volumes/<catalog>/<schema>/<volume>/runs")

SCRATCH_DIR.mkdir(parents=True, exist_ok=True)
ARTIFACT_ROOT.mkdir(parents=True, exist_ok=True)
with tempfile.NamedTemporaryFile(dir=SCRATCH_DIR) as probe:
    probe.write(b"scratch check")
    probe.flush()
with tempfile.NamedTemporaryFile(dir=ARTIFACT_ROOT) as probe:
    probe.write(b"volume check")
    probe.flush()

kumo.configure(scratch_dir=str(SCRATCH_DIR))
print("Free scratch space (GiB):", round(shutil.disk_usage(SCRATCH_DIR).free / 2**30, 1))
print("Retained artifact root:", ARTIFACT_ROOT)
```

These checks establish write access and reported free space. They do not prove that a full workload will fit or identify the underlying disk. Local scratch is temporary; on classic compute, Databricks removes driver-local files when the cluster restarts or terminates. Keep the files you need later in a Unity Catalog volume, and retain the complete model bundle rather than a checkpoint alone. See [Publish and Retrieve Models](/kumo-engine/1.0.0/user-code-guide/train-a-model/publish-and-retrieve-models#azure-databricks-notebook).

Continue to [Databricks connectors](/kumo-engine/1.0.0/user-code-guide/connect-your-data-databricks#choose-a-source) to choose a source, configure its connector, and verify data access.