How do I generate predictions with a different graph than my training graph?

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:

Initialize the SDK and load your predictive query from an existing training job.
Load the new tables you want to run inference on.
Define a new graph using those tables and the appropriate edges.
Generate a prediction table from the predictive query.
Run the prediction job against the new graph using the original trainer.

Step 1: Initialize and load the predictive query

1 import kumoai
2 
3 kumoai.init(url="https://<your-instance>.kumoai.cloud/api", api_key="<your-api-key>")
4 
5 TRAINING_JOB_ID = "<your-training-job-id>"
6 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.

1 from kumoai.graph import Table
2 
3 main_table = Table.load("<your-inference-main-table>")
4 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.

1 from kumoai.graph import Graph, Edge
2 
3 new_graph = Graph(
4     tables={
5         "<original-training-table-name>": main_table,
6         "<original-training-related-table-name>": related_table,
7     },
8     edges=[
9         Edge(
10             src_table="<original-training-table-name>",
11             fkey="<foreign-key-column>",
12             dst_table="<original-training-related-table-name>",
13         ),
14     ],
15 )

Step 4: Generate the prediction table with the new graph

1 pquery.graph = new_graph
2 prediction_table_plan = pquery.suggest_prediction_table_plan()
3 prediction_table = pquery.generate_prediction_table(prediction_table_plan, non_blocking=True)

Step 5: Run the prediction job

1 from kumoai.artifact_export.config import OutputConfig
2 
3 trainer = kumoai.Trainer.load(TRAINING_JOB_ID)
4 
5 prediction_job = trainer.predict(
6     graph=new_graph,
7     prediction_table=prediction_table,
8     training_job_id=TRAINING_JOB_ID,
9     output_config=OutputConfig(
10         output_types={"predictions"},
11         output_connector=<your_connector>,
12         output_table_name="<your-output-table-name>",
13     ),
14     non_blocking=False,
15 )
16 
17 print(prediction_job.summary())

For Binary Classification, Multiclass Classification, Regression, and Forecast tasks, 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.