Documentation · Quickstart · Performance · Examples
Kumo Engine is an accelerator for context generation over structured data. It stores your tables and their relationships as a graph on disk and, for each entity, samples the related records up to its prediction time. Describe what to predict in one query, and Kumo Engine uses that context to train and run predictive models directly on your relational data.
Example: predicting whether a customer will order in the next 30 days. Keys link each customer to their orders, and the model sees only orders placed before the prediction time.
- Scales beyond memory: Kumo Engine samples context from graphs with billions of records on one machine, even when the graph is larger than memory.
- Training data from one query: A predictive query generates labeled training examples at points in each entity's history. Each example sees only records timestamped at or before its point in time, with no hand-written label logic or point-in-time joins.
- Trains directly on relational data: A built-in graph neural network (GNN) learns from the related records without requiring hand-built features. It handles classification, regression, forecasting, and recommendation tasks.
- End-to-end workflow: Hyperparameter search, explanations, embedding export, and batch prediction are built in. Kumo Engine also works with MLflow, Databricks, SageMaker, and Snowflake to log each run, version models so scheduled prediction jobs use the one you promote, and monitor prediction drift.
- Runs locally or in your data platform: Run Kumo Engine in Python on Linux, macOS, or Windows (WSL2), or in a Snowflake or Databricks notebook. Read data from, and write predictions to, local files, cloud storage, and data warehouses.
- Works with coding agents: The bundled agent skill teaches your coding agent the Kumo Engine API, so it can build, train, and run models on your data.
Kumo Engine is part of NVIDIA Kumo, a family of open models and accelerated libraries for predictive AI on tabular, relational, and time-series data, and a core part of the NVIDIA Agent Knowledge Foundation. To make predictions without training a model, use the pretrained Kumo foundation models (Kumo Tabular, Kumo Relational, Kumo Forecast, and Kumo Anomaly) with Structured Data Models or Kumo-TS.
Install Kumo Engine from PyPI:
pip install kumo-engine| Requirement | Details |
|---|---|
| Python | CPython 3.11 or later |
| Operating system | Linux x86-64 or ARM64 with glibc 2.34 or later, macOS 14 or later on Apple silicon, or Windows through WSL2 (x86-64) |
| CPU | AES and SSE2 instructions on x86-64, or AES and NEON on ARM64 |
| GPU | An NVIDIA GPU for production workloads. The quickstart and other small datasets also run on a CPU or an Apple silicon GPU. |
| Storage | A local NVMe SSD for production workloads. Network-attached volumes, such as Amazon EBS, make sampling much slower. |
The base package reads and writes local CSV and Parquet files. Optional extras add connectors and additional features; combine them as needed, for example pip install "kumo-engine[snowflake,text]".
| Extra | Adds |
|---|---|
s3 |
Amazon S3 |
gcs |
Google Cloud Storage |
adls |
Azure Data Lake Storage Gen2 |
bigquery |
BigQuery |
snowflake |
Snowflake, including Snowflake notebooks |
databricks |
Databricks, including Databricks notebooks |
text |
Encoding for free-text columns, such as descriptions or reviews, with GloVe 6B 50d word vectors |
mlflow |
MLflow run tracking and model registry |
sagemaker |
SageMaker managed MLflow |
Choose your hardware and runtime environment:
- Hardware: Prerequisites · GPU sizing
- Setup guides: Linux · macOS · Windows (WSL2) · Snowflake notebook · Databricks notebook
- Platform details: the package README covers ARM64 Linux setup, including Faiss for recommendation tasks.
This example trains a model that predicts which users will place an order in the next 30 days, using a sample dataset of users and orders that ships with the package.
from kumo.engine import (Edge, FileConnector, Graph, PredictiveQuery, Table,
Trainer, datasets)
files = FileConnector(root_dir=datasets.shop())
users = Table.from_source_table(files["users"], primary_key="user_id")
orders = Table.from_source_table(files["orders"], primary_key="order_id",
time_column="order_time")
graph = Graph(
tables={"users": users, "orders": orders},
edges=[Edge(src_table="orders", fkey="user_id", dst_table="users")],
)
snapshot = graph.snapshot()
pquery = PredictiveQuery(
graph, "PREDICT COUNT(orders.*, 0, 30, days) > 0 FOR EACH users.user_id")
trainer = Trainer(pquery.suggest_model_plan(snapshot))
training_result = trainer.fit(pquery.generate_training_table(snapshot),
snapshot=snapshot)
result = trainer.predict(pquery.generate_prediction_table(snapshot),
snapshot=snapshot, binary_classification_threshold=0.5)
print(result.frame.head())When training completes, the example prints the first five predictions:
ENTITY TARGET_PRED False_PROB True_PROB
0 625 True 0.118336 0.881664
1 631 True 0.110448 0.889552
2 505 True 0.172911 0.827089
3 519 False 0.939077 0.060923
4 752 False 0.939077 0.060923
result.frame has one row per user, with the predicted label TARGET_PRED and the class probabilities False_PROB and True_PROB. The prediction time defaults to the latest order_time in the snapshot, so predictions cover the 30 days after it. Exact values can differ between runs. To use your own data, point FileConnector at a directory of CSV or Parquet files.
Next steps:
- Follow a complete walkthrough in a Google Colab, Snowflake, or Databricks notebook.
- Explore more tasks in Predictive Query Examples, and browse the example collection for industry notebooks, standalone scripts, and synthetic-data generators.
- To have a coding agent build the workflow for you, install the agent skill by running
python -m kumo.engine.skills. See Get Started with Agents.
Kumo Engine converts your tables into a graph once, at ingest. Its on-disk layout and its sampler are designed together for high-performance sampling:
- Links resolved at ingest: Minimal perfect hashing maps every primary key to a compact integer ID, and each foreign key is resolved once. Following a link while sampling is a direct lookup rather than a join.
- Bounded context: By default, each hop samples at most a set number of related records, so each entity's context has a fixed maximum size, however large the graph grows.
- Contiguous recent history: Each record's links are grouped by relationship and stored newest first. By default, sampling takes the most recent links: a binary search finds the prediction time, and the sampler reads the links that follow in order rather than gathering scattered rows.
- Ordered reads for large batches: Sampling expands a batch one hop at a time. For large batches, the sampler sorts each hop's records by their position on disk so reads move through each file in order, and spreads the work across CPU threads.
- Larger than memory: The graph files are read-only and memory-mapped. The operating system keeps frequently read pages in RAM and reads the rest from disk on demand, so the graph does not have to fit in memory, and many threads read it in parallel without locks.
See the benchmark report for measured sampling and end-to-end results, with hardware, cache conditions, and methodology.
| Goal | Start Here |
|---|---|
| Connect data, define graphs, train, and predict | User Code Guide |
| Configure sampling and training | Model Plan Guide |
| Run batch predictions in production | Production Workflow |
| Use the in-process Python API | Local Python interface, imported as kumo.engine |
| Version, promote, and monitor models | MLOps guide |
| Check compatibility before upgrading | Release notes |
| Explore the source code | Architecture and engineering docs |
Note
This Kumo Studio information is prerelease and is not yet supported. Stay tuned for the Docker and Kubernetes releases.
Kumo Studio allows enterprise teams to build and run predictive models together. Users can connect data sources, train and evaluate models, run batch predictions, and review job history, all through a web UI or Python SDK. Admins control access to data and models with role-based access control.
Deploy Studio on Kubernetes via Helm or on any compatible Linux host with Docker Compose. See the preliminary documentation for prerequisites and setup.
Report bugs and request features in GitHub Issues. To report a security vulnerability, follow SECURITY.md instead of opening an issue.
This project is currently not accepting contributions.
Development guides and source-checkout validation
Kumo-owned source is licensed under the Apache License 2.0. Vendored and incorporated third-party code retains its applicable licenses and attribution.
Component licenses and third-party attribution
Apache-licensed Python subprojects retain a copy of the root license at each package build root so it accompanies independently built wheels and source distributions. The nested component license files are:
Third-party attribution is recorded in the repository notice, Python distribution notice, UI catalog notice, and Helm stack notice.