Python SDK

kumoai.trainer

View as Markdown

The kumoai.trainer module provides the Trainer class for training custom GNN models on a Graph and TrainingTable, and generating predictions with a PredictionTable. Models can be customized with ModelPlan, though the plan suggested by PredictiveQuery.suggest_model_plan() is typically sufficient for strong out-of-the-box performance.

Model Plan

A ModelPlan defines the full parameter specification for training a Kumo model. It is composed of five sub-plans:

  • ColumnProcessingPlan — encoder overrides for individual columns
  • ModelArchitecturePlan — GNN or Graph Transformer parameters
  • NeighborSamplingPlan — subgraph sampling parameters
  • OptimizationPlan — learning rate, batch size, epochs, and related settings
  • TrainingJobPlan — AutoML-level settings

After generating a default model plan with PredictiveQuery.suggest_model_plan(), no further changes are required to train your first model. These options are available for fine-tuning.

ModelPlan

The top-level model configuration object. Each sub-plan is accessible as an attribute.

1model_plan = pq.suggest_model_plan()
2
3# Customize a sub-plan:
4model_plan.optimization.max_epochs = 50
5model_plan.optimization.base_lr = [1e-3, 3e-3]
training_job
TrainingJobPlanDefaults to TrainingJobPlan()

AutoML job-level settings.

column_processing
ColumnProcessingPlanDefaults to ColumnProcessingPlan()

Encoder overrides for individual columns.

neighbor_sampling
NeighborSamplingPlanDefaults to NeighborSamplingPlan()

Subgraph sampling configuration.

optimization
OptimizationPlanDefaults to OptimizationPlan()

Optimization hyperparameters.

model_architecture
ModelArchitecturePlanDefaults to ModelArchitecturePlan()

GNN or Graph Transformer architecture parameters.

ColumnProcessingPlan

Specifies encoder overrides and missing value strategy overrides for individual table columns.

1from kumoai.encoder import GloVe
2from kumoai.trainer import ColumnProcessingPlan
3
4plan = ColumnProcessingPlan(
5    encoder_overrides={"products.description": GloVe(model_name="glove-wiki-gigaword-50")},
6)
encoder_overrides
Optional[Dict[str, Encoder]]Defaults to None

A mapping from "table.column" to an encoder instance. Overrides Kumo’s auto-inferred encoder for that column.

na_strategy
Optional[Dict[Stype, NAStrategy]]Defaults to None

A mapping from semantic type to NAStrategy. Overrides the default imputation strategy for all columns of that semantic type.

ModelArchitecturePlan

Base class for architecture plans. Use GNNModelPlan or GraphTransformerModelPlan to configure a specific architecture.

GNNModelPlan

Configures a Graph Neural Network architecture.

1from kumoai.trainer import GNNModelPlan
2
3arch = GNNModelPlan(channels=[64, 128], aggregation=[["sum", "mean"]])
channels
List[int]Defaults to inferred

Candidate hidden channel sizes for AutoML search.

aggregation
List[List[AggregationType]]Defaults to inferred

Candidate aggregation function combinations for AutoML search.

dropout
List[float]Defaults to inferred

Candidate dropout rates for AutoML search.

GraphTransformerModelPlan

Configures a Graph Transformer architecture.

1from kumoai.trainer import GraphTransformerModelPlan
2
3arch = GraphTransformerModelPlan(num_layers=[2, 4], num_heads=[4, 8])
channels
List[int]Defaults to inferred

Candidate hidden channel sizes.

num_layers
List[int]Defaults to inferred

Candidate number of transformer layers.

num_heads
List[int]Defaults to inferred

Candidate number of attention heads.

dropout
List[float]Defaults to inferred

Candidate dropout rates.

positional_encodings
List[List[PositionalEncodingType]]Defaults to inferred

Candidate positional encoding combinations.

NeighborSamplingPlan

Controls how Kumo samples subgraphs during training.

num_neighbors
List[List[int]]Defaults to inferred

Candidate per-hop neighbor counts for AutoML search. Each inner list specifies the number of neighbors to sample at each hop.

sample_from_entity_table
boolDefaults to inferred

Whether to sample neighbors from the entity table.

OptimizationPlan

Controls learning rate, batch size, epochs, and other training optimization parameters.

max_epochs
intDefaults to inferred

Maximum number of training epochs.

max_steps_per_epoch
intDefaults to inferred

Maximum number of training steps per epoch.

max_val_steps
intDefaults to inferred

Maximum number of validation steps.

max_test_steps
intDefaults to inferred

Maximum number of test steps.

loss
List[Union[str, LossConfig]]Defaults to inferred

Candidate loss functions for AutoML search.

base_lr
List[float]Defaults to inferred

Candidate base learning rates.

weight_decay
List[float]Defaults to inferred

Candidate weight decay values.

batch_size
List[int]Defaults to inferred

Candidate batch sizes.

early_stopping
List[Optional[EarlyStoppingConfig]]Defaults to inferred

Candidate early stopping configurations.

lr_scheduler
List[Optional[LRSchedulerConfig]]Defaults to inferred

Candidate learning rate scheduler configurations.

majority_sampling_ratio
List[Optional[float]]Defaults to inferred

Candidate majority class sampling ratios for imbalanced classification.

weight_mode
List[Optional[WeightMode]]Defaults to inferred

Candidate sample weighting modes.

TrainingJobPlan

AutoML job-level settings controlling the number of experiments and evaluation metrics.

num_experiments
intDefaults to inferred

Number of hyperparameter experiments to run during AutoML.

metrics
List[str]Defaults to inferred

Evaluation metrics to compute.

tune_metric
strDefaults to inferred

The primary metric used to select the best model.

refit_trainval
boolDefaults to True

Whether to refit the best model on the combined train+validation set.

refit_full
boolDefaults to False

Whether to additionally refit on the full dataset (train+validation+test).

Training

Trainer

Trains a Kumo GNN model on a PredictiveQuery. The two primary methods are fit() (training) and predict() (batch inference).

1from kumoai.trainer import Trainer
2
3trainer = Trainer(model_plan=model_plan)
4result = trainer.fit(graph=graph, train_table=train_table)
model_plan
ModelPlanRequired

The model plan specifying architecture, optimization, and sampling parameters.

model_plan property

Returns Optional[ModelPlan]

encoders property

Returns Optional[Dict[str, str]] — The encoder configuration used during training.

is_trained property

Returns bool — True if this trainer has been successfully fit and is ready for prediction.

fit()

Trains a model on the provided graph and training table.

graph
GraphRequired

The relational graph.

train_table
Union[TrainingTable, TrainingTableJob]Defaults to None

An optional pre-generated training table. Provide exactly one of train_table or pquery.

pquery
Optional[PredictiveQuery]Defaults to None

An optional PredictiveQuery. When provided without train_table, Kumo generates the training table as an inline child job sharing the same graph snapshot. Recommended for live-streaming data to ensure consistency.

non_blocking
boolDefaults to False

If True, returns a TrainingJob immediately rather than blocking.

custom_tags
Mapping[str, str]Defaults to {}

Optional key-value tags attached to the training job.

warm_start_job_id
strDefaults to None

Training job ID to warm-start from (initializes from an existing model’s weights).

Returns Union[TrainingJob, TrainingJobResult]

predict()

Generates batch predictions using the trained model.

graph
GraphRequired

The relational graph.

prediction_table
Union[PredictionTable, PredictionTableJob]Defaults to None

The prediction table generated from a PredictiveQuery.

output_config
Union[OutputConfig, Dict[str, Any]]Required

Output configuration for the batch prediction job (OutputConfig instance or dict). Key fields: output_connector (connector to write to), output_table_name (table name, or schema/table tuple for Databricks), output_types (predictions or embeddings), output_metadata_fields (JOB_TIMESTAMP or ANCHOR_TIMESTAMP).

output_connector
ConnectorDefaults to None

Deprecated. Raises ValueError when passed. Use output_config instead.

output_table_name
strDefaults to None

Deprecated. Raises ValueError when passed. Use output_config instead.

non_blocking
boolDefaults to False

If True, returns a BatchPredictionJob immediately.

custom_tags
Mapping[str, str]Defaults to {}

Optional key-value tags attached to the prediction job.

training_job_id
strDefaults to None

The job ID of the training job whose model will be used for prediction. If None, uses the model from the most recent fit() call on this Trainer instance.

binary_classification_threshold
floatDefaults to None

For binary classification models, the score threshold above which predictions are classified as 1. If None, the raw probability score is returned.

num_classes_to_return
Optional[int]Defaults to None

For ranking task models, the number of classes to return in the prediction output.

num_workers
intDefaults to 1

Number of parallel workers for batch prediction. Values greater than 1 partition the prediction table and process in parallel.

prediction_time
datetimeDefaults to None

The point in time at which to generate predictions. Only valid when prediction_table is None. If None, the anchor time is inferred from the latest available data. Cannot be specified together with prediction_table.

explanations
boolDefaults to False

If True, generates per-prediction feature attributions in addition to prediction outputs.

Returns Union[BatchPredictionJob, BatchPredictionJobResult]

load() classmethod

Loads a trained Trainer from a completed training job.

training_job_id
strRequired

The training job ID.

Returns Trainer

TrainingJob

Represents an ongoing training job.

result()

Blocks until complete and returns the TrainingJobResult.

Returns TrainingJobResult

status()

Returns JobStatusReport

cancel()

Cancels the running training job.

metrics_so_far()

Returns Optional[ModelEvaluationMetrics] — Metrics computed so far during training.

progress()

Returns AutoTrainerProgress — Detailed progress information.

TrainingJobResult

Represents a completed training job.

1result = trainer.fit(graph=graph, train_table=train_table)
2metrics = result.metrics()
job_id
TrainingJobIDRequired

The training job ID.

id property

Returns TrainingJobID

model_plan property

Returns ModelPlan — The model plan used in this training job.

training_table property

Returns Union[TrainingTableJob, TrainingTable]

predictive_query property

Returns PredictiveQuery — The predictive query that defined this training job.

tracking_url property

Returns str — URL to the training job in the Kumo UI.

metrics()

Returns ModelEvaluationMetrics — Evaluation metrics for the completed job.

holdout_df()

Returns pd.DataFrame — The holdout dataset as a DataFrame.

explain()

Returns per-entity feature importances for a predictive query.

query
strRequired

The PQL query string.

indices
Sequence[Union[str, float, int]]Defaults to None

Entity indices to explain. Explains all entities if None.

run_mode
RunModeDefaults to RunMode.FAST

The run mode for explanation computation.

num_neighbors
List[int]Defaults to None

Per-hop neighbor counts for subgraph sampling.

anchor_time
Union[pd.Timestamp, Literal['entity']]Defaults to None

The anchor time for temporal explanation.

Returns pd.DataFrame

Batch Prediction

BatchPredictionJob

Represents an ongoing batch prediction job.

result()

Returns BatchPredictionJobResult

status()

Returns JobStatusReport

cancel()

Cancels the running batch prediction job.

BatchPredictionJobResult

Represents a completed batch prediction job.

data_df()

Returns pd.DataFrame — Prediction results.

data_urls()

Returns List[str] — Download URLs for prediction results.

summary()

Returns BatchPredictionJobSummary

Online Serving and Distillation

Distillation training and export_model() produce a serving bundle (online model directory and embeddings.parquet from batch prediction) in storage you control. Inference uses NVIDIA Triton Inference Server to load that bundle. See the Online Serving guide for the end-to-end flow.

DistillationTrainer

Trains a shallow model for online serving by reusing representations (embeddings) from a base GNN training job.

1from kumoai.trainer import DistillationTrainer
2
3distil = DistillationTrainer(model_plan=distil_plan, base_training_job_id="<job_id>")
4result = distil.fit(graph=graph, train_table=train_table)
model_plan
DistilledModelPlanRequired

The distilled model plan.

base_training_job_id
strRequired

The training job ID of the base GNN model to distill from.

is_trained property

Returns bool

fit()

graph
GraphRequired

The relational graph.

train_table
Union[TrainingTable, TrainingTableJob]Required

The training table.

non_blocking
boolDefaults to False

If True, returns a TrainingJob immediately.

custom_tags
Mapping[str, str]Defaults to {}

Optional job tags.

Returns Union[TrainingJob, TrainingJobResult]

load() classmethod

job_id
strRequired

The training job ID.

Returns DistillationTrainer

DistilledModelPlan

Model plan for distillation. Composed of TrainingJobPlan, ColumnProcessingPlan, OptimizationPlan, DistillationPlan, and a distillation-specific architecture plan.

DistillationPlan

Configuration for the distillation process, specifying embedding keys, time offsets, and real-time interaction settings.

embedding_keys
List[str]Defaults to inferred

Column keys used as embedding inputs.

max_embedding_offset
TimeOffsetDefaults to inferred

Maximum time offset for embedding lookups.

min_embedding_offset
TimeOffsetDefaults to inferred

Minimum time offset for embedding lookups.

real_time_interactions
Dict[str, int]Defaults to {}

Real-time interaction table configuration.

export_model()

Exports online serving model files and batch prediction embeddings to external storage for use with Triton Inference Server.

1from kumoai.trainer import export_model, ModelOutputConfig
2
3result = export_model(config=output_config, non_blocking=False)
config
ModelOutputConfigRequired

Specifies the training job, output path, and batch prediction job to bundle.

non_blocking
boolDefaults to True

If True, returns an ArtifactExportJob immediately.

online_embedding_format
Optional[Literal['pandas', 'mphf', 'lmdb']]Defaults to None

Override for the storage backend used for online-embedding artifacts. One of pandas, mphf, or lmdb. If omitted, uses the value set on config. mphf is the server-side default and recommended for large catalogs. pandas is simple but limited to catalogs below 100M IDs. lmdb has the lowest serving RAM footprint but higher cold-cache latency.

payload_version
Optional[int]Defaults to None

Override for the wire format version of the exported online-serving model. This keyword is accepted and forwarded to the export config; it only takes effect where the deployment backend supports payload-version selection. When omitted, uses the value set on config.

Returns Union[ArtifactExportJob, ArtifactExportResult]

ModelOutputConfig

Output configuration for export_model(). Specifies output types, destination connector, and table name.

output_types
Set[str]Required

The output types to produce. Valid values: "predictions", "embeddings", or both.

output_connector
ConnectorDefaults to None

The connector to write outputs to. Local download only if None.

output_table_name
Union[str, Tuple[str, str]]Defaults to None

Table name in the output connector. For Databricks, provide a (schema, table) tuple.

output_metadata_fields
List[MetadataField]Defaults to None

Additional metadata columns to include in prediction output. Options: JOB_TIMESTAMP, ANCHOR_TIMESTAMP.

ArtifactExportJob

Represents an ongoing model artifact export job.

id property

Returns str

result()

Returns ArtifactExportResult

status()

Returns JobStatus

cancel()

Returns bool — True if the job was successfully cancelled.

ArtifactExportResult

Represents a completed model artifact export.

tracking_url()

Returns str — URL to the export job in the Kumo UI.