Training & Predictions | Structured Data and Graph Models

Starting with raw data (as SourceTable objects), you created a Graph of Kumo Table objects and defined a PredictiveQuery. Now you can train a model and generate predictions.

Creating a Trainer

Training requires a ModelPlan, which defines the search space of model configurations. Generate one from your predictive query:

1 model_plan = pquery.suggest_model_plan()
2 print(model_plan)

The model plan can be edited with full granularity — see the Model Plan reference for customizable attributes. Once ready, create a Trainer:

1 trainer = kumo.Trainer(model_plan)

Training a Model

Call fit() with your graph and training table:

1 training_job = trainer.fit(
2     graph=graph,
3     train_table=pquery.generate_training_table(non_blocking=True),
4     non_blocking=False,
5 )
6 print(f"Training metrics: {training_job.metrics()}")

The Kumo SDK makes extensive use of non_blocking as an optional parameter for long-running operations. Setting this flag to True lets a long-running operation return immediately, returning a Future object that tracks the operation as it runs in the background. Setting this flag to False lets it wait until completion before returning.

Viewing Metrics and Artifacts

A completed training job exposes metrics and artifacts:

1 # View evaluation metrics:
2 print(training_job.metrics())
3 
4 # Download holdout dataset as a DataFrame:
5 holdout_df = training_job.holdout_df()
6 
7 # Open the job in the Kumo UI:
8 print(training_job.tracking_url)

Generating Batch Predictions

Once a model has been trained, generate batch predictions and write to an external data source:

1 # For v1.4 and above:
2 from kumoai.artifact_export.config import OutputConfig
3 # For v1.3 and below (backward compatible):
4 # from kumoai.trainer.config import OutputConfig
5 
6 prediction_job = trainer.predict(
7     graph=graph,
8     prediction_table=pquery.generate_prediction_table(non_blocking=True),
9     output_config=OutputConfig(
10         output_types={'predictions', 'embeddings'},
11         output_connector=connector,
12         output_table_name='kumo_predictions',
13     ),
14     training_job_id=training_job.job_id,
15     non_blocking=False,
16 )
17 print(f'Batch prediction job summary: {prediction_job.summary()}')

Full Code Example

A complete end-to-end example on the CustomerLTV dataset:

1 import kumoai as kumo
2 
3 # Initialize the SDK:
4 kumo.init(url="https://<customer_id>.kumoai.cloud/api", api_key=API_KEY)
5 
6 # Create a Connector:
7 connector = kumo.S3Connector("s3://kumo-public-datasets/customerltv_mini_integ_test/")
8 
9 # Create Tables from SourceTables:
10 customer = kumo.Table.from_source_table(
11     source_table=connector.table('customer'),
12     primary_key='CustomerID',
13 ).infer_metadata()
14 
15 stock = kumo.Table.from_source_table(
16     source_table=connector.table('stock'),
17     primary_key='StockCode',
18 ).infer_metadata()
19 
20 transaction = kumo.Table.from_source_table(
21     source_table=connector.table('transaction'),
22     time_column='InvoiceDate',
23 ).infer_metadata()
24 
25 # Create a Graph:
26 graph = kumo.Graph(
27     tables={
28         'customer': customer,
29         'stock': stock,
30         'transaction': transaction,
31     },
32     edges=[
33         dict(src_table='transaction', fkey='StockCode', dst_table='stock'),
34         dict(src_table='transaction', fkey='CustomerID', dst_table='customer'),
35     ],
36 )
37 graph.validate(verbose=True)
38 
39 # Create a Predictive Query:
40 pquery = kumo.PredictiveQuery(
41     graph=graph,
42     query=(
43         "PREDICT MAX(transaction.Quantity, 0, 30, days)\n"
44         "FOR EACH customer.CustomerID\n"
45         "ASSUMING SUM(transaction.UnitPrice, 0, 7, days) > 15"
46     ),
47 )
48 pquery.validate(verbose=True)
49 
50 # Train a model:
51 model_plan = pquery.suggest_model_plan()
52 trainer = kumo.Trainer(model_plan)
53 training_job = trainer.fit(
54     graph=graph,
55     train_table=pquery.generate_training_table(non_blocking=True),
56     non_blocking=False,
57 )
58 print(f"Training metrics: {training_job.metrics()}")
59 
60 # Generate predictions:
61 from kumoai.artifact_export.config import OutputConfig
62 
63 prediction_job = trainer.predict(
64     graph=graph,
65     prediction_table=pquery.generate_prediction_table(non_blocking=True),
66     output_config=OutputConfig(
67         output_types={'predictions', 'embeddings'},
68         output_connector=connector,
69         output_table_name='kumo_predictions',
70     ),
71     training_job_id=training_job.job_id,
72     non_blocking=False,
73 )
74 print(f'Batch prediction job summary: {prediction_job.summary()}')

Polling Job Status

Any job scheduled with non_blocking=True returns a Future object with methods to track progress:

1 # Schedule a training job and get a future:
2 training_future = trainer.fit(graph=graph, train_table=training_table, non_blocking=True)
3 
4 # Print the job ID:
5 print(f"Training job ID: {training_future.id}")
6 
7 # Attach to watch logs live (detach anytime without cancelling the job):
8 training_future.attach()
9 
10 # Or poll status in a loop:
11 print(training_future.status())
12 
13 # Or block until complete:
14 training_job = training_future.result()

The same pattern applies to prediction jobs, training table jobs, and prediction table jobs.

Next Steps

While this example covered the core concepts, the SDK provides much more advanced functionality to help improve model iteration speed, evaluate champion/challenger models in production use-cases, integrate cleanly with upstream and downstream data pipelines, and more.