Training

Train a model using the query, training_table, and snapshot created in Predictive Queries. These calls use Kumo Engine and run synchronously in your Python process on Linux, in a Snowflake notebook, or in a Databricks notebook. The training API is shared. Scratch storage and model publication differ by deployment.

Model Plan

Request a plan suited to the query, then inspect it before starting training:

from kumoapi.model_plan import RunMode
plan = query.suggest_model_plan(snapshot, run_mode=RunMode.FAST)
print(plan.as_json())

Suggestion uses graph and query metadata; it does not start training. For weighted supplied training rows, pass the same TrainingTableSpec as train_table_spec= so the suggestion accounts for sample weights.

The plan is a kumoapi.model_plan.ModelPlan. Its inferred values are resolved for your task during training. Useful sections include:

For a short initial run, restrict the number of experiments and epochs and disable refits:

plan.training_job.num_experiments = 1
plan.training_job.refit_trainval = False
plan.training_job.refit_full = False
plan.optimization.max_epochs = 5

Choose your training budget based on the task and dataset; these settings are a starting point, not a model-quality target. training_job.metrics selects built-in metrics, and training_job.tune_metric chooses the metric used to select the model. For example, a regression task can use metrics=["mae", "mse", "rmse"] and tune_metric="mae". To calculate a custom metric function, score held-out predictions after training and compare them with separately retained ground truth.

Training-table generation has its own plan passed to query generation. Changing the model plan does not regenerate the training table’s labels or splits.

For parameter meanings, supported values, and encoder configuration, see the Model Plan Guide.

Train a Model

Construct a trainer and select the graph snapshot explicitly:

from kumo.engine import Trainer
trainer = Trainer(plan)
fit_result = trainer.fit(training_table, snapshot=snapshot)
print("Metrics:", fit_result.metrics)
print("Model bundle:", fit_result.bundle_dir)

trainer.fit(
    train_table,
    *,
    snapshot,
    export_config=None,
    warm_start=None,
) -> TrainingResult

The training table supplies rows, labels, splits, and generation settings. The snapshot supplies graph data and features. You can reuse the same training table with a compatible updated snapshot without regenerating its labels.

trainer.model_plan is the public configuration for the next fit. You can modify it or assign another ModelPlan before fitting again. Each fit uses its own run directory, and the trainer’s fitted state then represents the most recent successful fit. Keep earlier results if you need their model bundles.

Without warm_start, each fit() starts fresh, including when called on a loaded trainer. To reuse a trained model for inference, load its saved bundle.

Start a New Run from a Trained Model

Pass a fitted or loaded Trainer as warm_start to initialize a new run from its selected model. Continue from the fitted trainer, training_table, and snapshot above:

previous = trainer
next_plan = query.suggest_model_plan(snapshot, run_mode=RunMode.FAST)
next_plan.training_job.num_experiments = 1
next_plan.training_job.refit_trainval = False
next_plan.training_job.refit_full = False
next_plan.optimization.max_epochs = 5
next_trainer = Trainer(next_plan)
next_result = next_trainer.fit(
    training_table,
    snapshot=snapshot,
    warm_start=previous,
)
print("New run metrics:", next_result.metrics)
print("New model bundle:", next_result.bundle_dir)

For a retained model, replace previous = trainer with previous = Trainer.load("/path/to/retained-models/model-v1"), or retrieve it through its publication connector. Keep the source bundle available throughout the fit. The new run reuses the source’s selected weights, optimizer state, fitted encoders, model architecture, and neighbor sampling. The destination plan’s optimization and training controls, including the epoch budget above, still apply. Changing the destination’s encoder or architecture settings does not replace the inherited fitted model structure.

The source and destination must have the same task type; forecasting does not support warm starts. Supply a compatible training table and graph snapshot, including compatible targets and features. Use a complete compatible bundle containing the selected checkpoint and resolved model plan; an older bundle containing only unresolved search settings is not a reliable warm-start source.

This creates a separate training run and leaves the source trainer available for inference. It does not resume an interrupted job at its last step. Use next_trainer and next_result for the new model’s predictions and artifacts.

Metrics and Training Results

A completed fit returns attributes, rather than a hosted job handle:

Inspect the available metric names instead of assuming every task returns the same keys:

for split in ("train", "val", "test"):
    metrics = {
        name: value
        for name, value in fit_result.metrics.items()
        if name.startswith(f"{split}_")
    }
    print(split, metrics)

Metrics depend on the task, configured evaluation metrics, and splits evaluated by the run. For example, classification can report AUROC, while regression can report MAE. Unavailable or nonfinite metrics can be omitted, so an empty split dictionary is not a score of zero. See Training Diagnostics to read metrics from a retained bundle and read holdout predictions to inspect and retain individual evaluation rows.

Every successful fit saves the metadata needed to reload the bundle in another process, including its query, model plan, and effective generation settings. An export is not required for reload. The bundle lives under execution scratch unless you retain it elsewhere. On Linux, copy or publish it to retained filesystem or object storage. In a Snowflake notebook, publish it to Snowflake Model Registry or preserve the complete local bundle before notebook storage is removed. In a Databricks notebook, copy the complete bundle from scratch to a Unity Catalog volume before the cluster stops.

Publish or Retrieve a Model

For a supported connector destination, supply ModelOutputConfig to the fit call and publish the complete bundle to durable storage or a model registry. In a Databricks notebook, copy the bundle to a Unity Catalog volume after fitting; DatabricksConnector does not publish models to a registry through this local API. Record the returned publication name, version, or path so another process can retrieve the same bundle.

See Publish and Retrieve Models for Linux, object storage, Snowflake Registry, and Databricks volume procedures.

Continue to Training Diagnostics to inspect epoch and experiment history, Model Explainability to inspect any explanation files produced by this fit, or Making Predictions to score entities and generate embeddings.