Databricks Notebook Setup

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.

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. See the Serverless GPU environment version 6 release notes.

The required compute size depends on your data and workload. Review GPU sizing and working storage 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.

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

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:

%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:

%pip install "kumo-engine[databricks]"

For either Databricks S3 connection path, include the S3 integration in the same installation:

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

After installation, restart Python in a separate cell:

dbutils.library.restartPython()

The restart clears Python variables. Rerun the configuration cells below and recreate your connector using Databricks connectors after it completes. Verify that the runtime imports and a CUDA operation succeed:

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):

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:

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.

Continue to Databricks connectors to choose a source, configure its connector, and verify data access.