Predict with a Different Graph

Prerelease documentation

This Kumo Studio information is prerelease and is not yet supported. Stay tuned for the Docker and Kubernetes releases.

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

import kumoai_studio
kumoai_studio.init(url="https://<your-instance>.kumoai.cloud/api", api_key="<your-api-key>")
TRAINING_JOB_ID = "<your-training-job-id>"
pquery = kumoai_studio.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_studio.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, rather than the names of the new tables you loaded above.

from kumoai_studio.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_studio.artifact_export.config import OutputConfig
trainer = kumoai_studio.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.