Skip to main content
Currently, running predictions on a graph that differs from the one used during training is only supported via the Kumo Python SDK — it is not available through the REST API. The basic steps are:
  1. Initialize the SDK and load your predictive query from an existing training job.
  2. Load the new tables you want to run inference on.
  3. Define a new graph using those tables and the appropriate edges.
  4. Generate a prediction table from the predictive query.
  5. Run the prediction job against the new graph using the original trainer.

​
Step 1: Initialize and load the predictive query

import kumoai

kumoai.init(url="https://<your-instance>.kumoai.cloud/api", api_key="<your-api-key>")

TRAINING_JOB_ID = "<your-training-job-id>"
pquery = kumoai.PredictiveQuery.load_from_training_job(TRAINING_JOB_ID)

​
Step 2: Load the new tables

Load the tables you want to use for inference. These are typically different from the tables used during training — for example, a held-out dataset or a new batch of data. 
from kumoai.graph import Table

main_table = Table.load("<your-inference-main-table>")
related_table = Table.load("<your-inference-related-table>")

​
Step 3: Define a new graph

Construct a Graph using the new tables and the same edge structure as the training graph.
The keys in the tables dictionary must exactly match the names of the tables used in the original training graph. These are the table names Kumo learned from during training — not the names of the new tables you loaded above.
from kumoai.graph import Graph, Edge

new_graph = Graph(
    tables={
        "<original-training-table-name>": main_table,
        "<original-training-related-table-name>": related_table,
    },
    edges=[
        Edge(
            src_table="<original-training-table-name>",
            fkey="<foreign-key-column>",
            dst_table="<original-training-related-table-name>",
        ),
    ],
)

​
Step 4: Generate the prediction table

prediction_table_plan = pquery.suggest_prediction_table_plan()
prediction_table = pquery.generate_prediction_table(prediction_table_plan, non_blocking=True)

​
Step 5: Run the prediction job

from kumoai.artifact_export.config import OutputConfig

trainer = kumoai.Trainer.load(TRAINING_JOB_ID)

prediction_job = trainer.predict(
    graph=new_graph,
    prediction_table=prediction_table,
    training_job_id=TRAINING_JOB_ID,
    output_config=OutputConfig(
        output_types={"predictions"},
        output_table_name="<your-output-table-name>",
    ),
    non_blocking=False,
)

print(prediction_job.summary())
The rows you want predictions for must have an empty (null) target column. If all rows already have a value in the target column, the prediction job will fail. Make sure to include rows with a missing target alongside any reference data when constructing your inference graph.