Kumo Engine

An open accelerator for context generation over structured data.

NVIDIA Kumo Engine prepares connected records and their history for predictive modeling at scale. Its Rust-based DiskGraph engine indexes data on local SSDs and retrieves time-correct neighborhoods for training and batch prediction, so teams can run structured data models over enterprise-scale datasets at lower cost by feeding the GPU from NVMe instead of RAM.

Define your tables, relationships, and prediction task in Python. Kumo Engine handles ingestion from the data lake, graph construction, context sampling, training, evaluation, and prediction for a wide range of predictive use cases. Examples include recommendations, fraud detection, churn prediction, propensity modeling, and risk scoring, among many others.

Benchmark highlight: One million three-hop neighborhoods in 19 seconds.

Documentation · Installation · Quick Tour · Benchmarks

NVIDIA Kumo is a family of open models, libraries, and tools for predictive AI across tabular, relational, and time-series data, and a core part of the NVIDIA Agent Knowledge Foundation.

• NVIDIA Kumo Structured: pretrained foundation models (Kumo Tabular, Kumo Relational, Kumo Forecast, and Kumo Anomaly) that make predictions without training a model per task. Refer to NVIDIA/structured-data-models.
• NVIDIA Kumo Engine (this repository): an open accelerator for preparing connected structured data and training predictive models at production scale.

The published 1.0.0 wheel targets Linux x86-64 with glibc 2.34 or later and CPython 3.11 or later. The documented Linux setup uses Ubuntu 24.04 and Python 3.12. Prepare an NVIDIA GPU and local SSD working storage for the documented training workflows.

pip install kumo-engine

The base package reads local CSV and Parquet files. To read from a remote data source, install the matching optional dependency:

Refer to prerequisites and GPU sizing for a compatible runtime and storage capacity. Follow the setup guide for Linux, Snowflake notebooks, or Databricks notebooks. For an ARM64 source build, refer to the DGX Spark guide.

• Context at scale: Retrieve connected records from an NVMe-backed graph index without keeping the full index in RAM.
• End-to-end workflows: Train and evaluate models, run batch predictions, generate embeddings, and save reloadable model bundles.
• Production ready: Ingest in parallel from S3, Snowflake, Databricks, BigQuery, and GCS.
• Existing data and infrastructure: Read files or connect to cloud storage and warehouses using existing compute and credentials.
• Compatible with agentic workflows: Agent skills let AI agents set up and run Kumo Engine workflows.

Predict which customers will order in the next 30 days. This example creates two small synthetic CSV tables, builds their graph, trains a classifier, and prints predictions. It uses the installed kumo.engine API without a Kumo account or warehouse credentials.

First, create the sample data in a new working directory:

import csv
from datetime import datetime, timedelta
from pathlib import Path
from random import Random

data = Path("kumo-readme-data")
data.mkdir(exist_ok=False)
rng = Random(42)

with (data / "customers.csv").open("w", newline="") as f:
    writer = csv.writer(f)
    writer.writerow(["customer_id", "segment"])
    writer.writerows(
        (i, rng.choice(["consumer", "business"])) for i in range(1000)
    )

with (data / "orders.csv").open("w", newline="") as f:
    writer = csv.writer(f)
    writer.writerow(["order_id", "customer_id", "order_time"])
    for i in range(12000):
        customer_id = rng.randrange(1000)
        time = datetime(2025, 1, 1) + timedelta(
            seconds=rng.randrange(365 * 24 * 60 * 60)
        )
        writer.writerow([i, customer_id, time.strftime("%Y-%m-%d %H:%M:%S")])

Run the following in the same working directory after completing setup. Replace the scratch path with a writable directory on your local SSD, separate from the sample-data directory.

from pathlib import Path

from kumo.engine import (
    Edge, FileConnector, Graph, PredictiveQuery, Table, Trainer, configure,
)
from kumoapi.model_plan import RunMode
from kumoapi.typing import Dtype

configure(scratch_dir="/path/to/local-ssd/kumo-scratch")
files = FileConnector(root_dir=str(Path("kumo-readme-data").resolve()))
customers = Table.from_source_table(
    files["customers"], primary_key="customer_id"
)
orders = Table.from_source_table(
    files["orders"], primary_key="order_id", time_column="order_time"
).with_column(
    "order_time", dtype=Dtype.time, timestamp_format="%Y-%m-%d %H:%M:%S"
)

graph = Graph(
    tables={"customers": customers, "orders": orders},
    edges=[Edge(
        src_table="orders", fkey="customer_id", dst_table="customers"
    )],
)
snapshot = graph.snapshot()

query = PredictiveQuery(
    graph,
    "PREDICT COUNT(orders.*, 0, 30, days) > 0 "
    "FOR EACH customers.customer_id",
)
training_table = query.generate_training_table(snapshot)
plan = query.suggest_model_plan(snapshot, run_mode=RunMode.FAST)
plan.training_job.num_experiments = 1
plan.training_job.refit_trainval = False
plan.training_job.refit_full = False
plan.optimization.max_epochs = 3

trainer = Trainer(plan)
trained = trainer.fit(training_table, snapshot=snapshot)
predictions = trainer.predict(
    query.generate_prediction_table(snapshot), snapshot=snapshot,
    binary_classification_threshold=0.5,
)
print(predictions.frame.head())
print("Saved model bundle:", trained.bundle_dir)

The prediction frame contains entity identifiers and model outputs. These scores cover a future 30-day window from the snapshot's latest event time. Keep the saved bundle and any results you need on durable storage.

For a complete walkthrough with data preparation, diagnostics, and retained outputs, follow the Google Colab, Snowflake, or Databricks notebook. Explore more tasks in Predictive Query Examples.

One million starting entities; synthetic graph: 50 billion nodes, 175.5 billion edges:

Both ran with warm caches on one AWS i4i.32xlarge: 128 vCPUs, 1 TiB RAM, and 30 TB local NVMe. Kumo Engine's 2.6 TiB index exceeded RAM. Timings exclude preparation and model training. Initial preparation took 49.8 minutes for DiskGraph and 40.1 minutes for DuckDB.

On a separate 111-million-node citation graph, Kumo Engine completed 15 training/validation epochs in 41 minutes 14 seconds, compared with 64 minutes 13 seconds for matched-model PyG. Test accuracy was 63.26% versus 63.93%; both used one g5.16xlarge with one A10G GPU.

Refer to the full benchmark report for experimental scope, preparation costs, memory measurements, and additional comparisons. Results depend on the workload and configuration.

For repository internals, refer to the architecture and engineering documentation index.

Report bugs or documentation problems in GitHub Issues, including the affected path and observed behavior. Keep credentials and private deployment details out of reports.

This project is currently not accepting contributions.

Kumo-owned source is licensed under the Apache License 2.0. Vendored and incorporated third-party code retains its applicable licenses and attribution.