Snowflake Notebook Setup

Run Kumo in your Snowflake notebook’s Python process, using Snowflake tables as the source for a local graph. The notebook runtime supplies the Snowflake identity; no Kumo account or API key is needed. This guide prepares the environment and connector for the Snowflake TPC-H notebook example.

Prepare the Notebook

Use Snowflake GPU Container Runtime 2.9 with Python 3.12. Choose an x86 compute pool and review GPU sizing and working storage. Your notebook role needs access to:

Package access through Snowflake Artifact Repository or a PyPI external access integration.
The query warehouse, database, schema, and source tables you will read.
Output tables or a model registry if you plan to publish results there.

The compute pool runs the notebook, graph ingestion, and model computation. The query warehouse supplies Snowflake SQL compute for source access. They are separate resources. See Snowflake’s compute pool configuration for available instance families and account configuration.

Install the Package

Install Kumo Engine with the Snowflake extra in a notebook cell before importing it:

!pip install "kumo-engine[snowflake]"

Follow Snowflake’s container runtime package guidance to configure package access. Restart the kernel if prompted after installation. After a restart, rerun the import, scratch configuration, and connector cells below.

Verify the installed packages and GPU access before configuring storage:

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

Save this output when reporting an installation problem. This checks package imports and a built-in PyTorch CUDA operation. It does not verify source access or a full training job. If it fails, confirm that your package build is compatible with the runtime’s PyTorch and CUDA versions.

Configure Working Storage

Choose a writable directory on local storage with room for staged source data, the graph index, training artifacts, and predictions. The example uses /tmp/kumoai-local; confirm that this path is on the intended working disk in your compute pool, or replace it with the appropriate local path. A directory name alone does not identify its storage backing.

from pathlib import Path
import shutil
import tempfile
SCRATCH_DIR = "/tmp/kumoai-local"
scratch = Path(SCRATCH_DIR)
scratch.mkdir(parents=True, exist_ok=True)
with tempfile.TemporaryFile(dir=scratch) as probe:
    probe.write(b"kumo scratch check")
    probe.flush()
kumo.configure(scratch_dir=SCRATCH_DIR)
print("Kumo version:", kumo.__version__)
print("Free working space (GiB):", round(shutil.disk_usage(scratch).free / 2**30, 1))

This checks write access and available space, not whether an entire workload will fit. Keep source data in Snowflake and export results you need beyond the runtime’s storage lifetime. Publish and Retrieve Models explains what a saved model contains, and Making Predictions explains how to retain prediction outputs.

Connect Your Data

Prepare the views described in the Snowflake TPC-H notebook example, then replace the example warehouse, database, and schema with the target containing those views:

WAREHOUSE = "COMPUTE_WH"
DATABASE = "ANALYTICS"
SCHEMA = "TPCH"
connector = kumo.SnowflakeConnector.from_spcs_env(
    warehouse=WAREHOUSE,
    database=DATABASE,
    schema_name=SCHEMA,
)

from_spcs_env() creates both the connector and its identity from the runtime’s account, host, and mounted OAuth token. Do not copy the token into the notebook. The environment must provide the Snowflake runtime credentials; an active Snowpark session by itself is not sufficient. If the helper reports missing runtime credentials, check that the notebook is running in the intended container environment. See Snowflake’s service-user credential documentation.

Inspect the catalog and one table’s schema to verify source access:

print(connector.table_names())
print([
    (column.name, column.dtype, column.stype)
    for column in connector["ORDERS"].columns
])

For the example, the schema must contain CUSTOMER, PART, ORDERS, and LINEITEM with the columns specified in the notebook, including the synthetic LINEITEM_ID column in the LINEITEM view. Use the source table names and casing reported by your connector. A successful schema read establishes access to that table’s metadata; the later graph snapshot reads and stages its data.

Continue to the Snowflake TPC-H notebook example to define tables, build the graph, train a model, and predict customer orders.