Skip to navigation
User Guide: Kumo StudioRun Models

Predict with a Different Graph

View as Markdown

Predicting against a graph that differs from the one used during training is only supported through the Kumo Predict 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 with the new graph

pquery.graph = new_graph
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

Define an output_connector for the destination that receives the prediction output, then run the prediction. Replace output_connector with your configured connector object before calling predict:

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_connector=output_connector,
        output_table_name="<your-output-table-name>",
    ),
    non_blocking=False,
)
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 fails. Make sure to include rows with a missing target alongside any reference data when constructing your inference graph.

Refer to the kumoai.trainer reference for the current prediction interface.