> ## Documentation Index
> Fetch the complete documentation index at: https://kumo.ai/docs/llms.txt
> Use this file to discover all available pages before exploring further.

# kumoai.trainer

> Trainer, ModelPlan, training jobs, batch prediction, and online serving

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

<Note>
  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.
</Note>

### `ModelPlan`

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

```python theme={null}
model_plan = pq.suggest_model_plan()

# Customize a sub-plan:
model_plan.optimization.max_epochs = 50
model_plan.optimization.base_lr = [1e-3, 3e-3]
```

<ParamField body="training_job" type="TrainingJobPlan" default="TrainingJobPlan()">
  AutoML job-level settings.
</ParamField>

<ParamField body="column_processing" type="ColumnProcessingPlan" default="ColumnProcessingPlan()">
  Encoder overrides for individual columns.
</ParamField>

<ParamField body="neighbor_sampling" type="NeighborSamplingPlan" default="NeighborSamplingPlan()">
  Subgraph sampling configuration.
</ParamField>

<ParamField body="optimization" type="OptimizationPlan" default="OptimizationPlan()">
  Optimization hyperparameters.
</ParamField>

<ParamField body="model_architecture" type="ModelArchitecturePlan" default="ModelArchitecturePlan()">
  GNN or Graph Transformer architecture parameters.
</ParamField>

***

### `ColumnProcessingPlan`

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

```python theme={null}
from kumoai.encoder import GloVe
from kumoai.trainer import ColumnProcessingPlan

plan = ColumnProcessingPlan(
    encoder_overrides={"products.description": GloVe(model_name="glove-wiki-gigaword-50")},
)
```

<ParamField body="encoder_overrides" type="Optional[Dict[str, Encoder]]" default="None">
  A mapping from `"table.column"` to an encoder instance. Overrides Kumo's auto-inferred encoder for that column.
</ParamField>

<ParamField body="na_strategy" type="Optional[Dict[Stype, NAStrategy]]" default="None">
  A mapping from semantic type to `NAStrategy`. Overrides the default imputation strategy for all columns of that semantic type.
</ParamField>

***

### `ModelArchitecturePlan`

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

***

### `GNNModelPlan`

Configures a Graph Neural Network architecture.

```python theme={null}
from kumoai.trainer import GNNModelPlan

arch = GNNModelPlan(channels=[64, 128], aggregation=[["sum", "mean"]])
```

<ParamField body="channels" type="List[int]" default="inferred">
  Candidate hidden channel sizes for AutoML search.
</ParamField>

<ParamField body="aggregation" type="List[List[AggregationType]]" default="inferred">
  Candidate aggregation function combinations for AutoML search.
</ParamField>

<ParamField body="dropout" type="List[float]" default="inferred">
  Candidate dropout rates for AutoML search.
</ParamField>

***

### `GraphTransformerModelPlan`

Configures a Graph Transformer architecture.

```python theme={null}
from kumoai.trainer import GraphTransformerModelPlan

arch = GraphTransformerModelPlan(num_layers=[2, 4], num_heads=[4, 8])
```

<ParamField body="channels" type="List[int]" default="inferred">
  Candidate hidden channel sizes.
</ParamField>

<ParamField body="num_layers" type="List[int]" default="inferred">
  Candidate number of transformer layers.
</ParamField>

<ParamField body="num_heads" type="List[int]" default="inferred">
  Candidate number of attention heads.
</ParamField>

<ParamField body="dropout" type="List[float]" default="inferred">
  Candidate dropout rates.
</ParamField>

<ParamField body="positional_encodings" type="List[List[PositionalEncodingType]]" default="inferred">
  Candidate positional encoding combinations.
</ParamField>

***

### `NeighborSamplingPlan`

Controls how Kumo samples subgraphs during training.

<ParamField body="num_neighbors" type="List[List[int]]" default="inferred">
  Candidate per-hop neighbor counts for AutoML search. Each inner list specifies the number of neighbors to sample at each hop.
</ParamField>

<ParamField body="sample_from_entity_table" type="bool" default="inferred">
  Whether to sample neighbors from the entity table.
</ParamField>

***

### `OptimizationPlan`

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

<ParamField body="max_epochs" type="int" default="inferred">
  Maximum number of training epochs.
</ParamField>

<ParamField body="max_steps_per_epoch" type="int" default="inferred">
  Maximum number of training steps per epoch.
</ParamField>

<ParamField body="max_val_steps" type="int" default="inferred">
  Maximum number of validation steps.
</ParamField>

<ParamField body="max_test_steps" type="int" default="inferred">
  Maximum number of test steps.
</ParamField>

<ParamField body="loss" type="List[Union[str, LossConfig]]" default="inferred">
  Candidate loss functions for AutoML search.
</ParamField>

<ParamField body="base_lr" type="List[float]" default="inferred">
  Candidate base learning rates.
</ParamField>

<ParamField body="weight_decay" type="List[float]" default="inferred">
  Candidate weight decay values.
</ParamField>

<ParamField body="batch_size" type="List[int]" default="inferred">
  Candidate batch sizes.
</ParamField>

<ParamField body="early_stopping" type="List[Optional[EarlyStoppingConfig]]" default="inferred">
  Candidate early stopping configurations.
</ParamField>

<ParamField body="lr_scheduler" type="List[Optional[LRSchedulerConfig]]" default="inferred">
  Candidate learning rate scheduler configurations.
</ParamField>

<ParamField body="majority_sampling_ratio" type="List[Optional[float]]" default="inferred">
  Candidate majority class sampling ratios for imbalanced classification.
</ParamField>

<ParamField body="weight_mode" type="List[Optional[WeightMode]]" default="inferred">
  Candidate sample weighting modes.
</ParamField>

***

### `TrainingJobPlan`

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

<ParamField body="num_experiments" type="int" default="inferred">
  Number of hyperparameter experiments to run during AutoML.
</ParamField>

<ParamField body="metrics" type="List[str]" default="inferred">
  Evaluation metrics to compute.
</ParamField>

<ParamField body="tune_metric" type="str" default="inferred">
  The primary metric used to select the best model.
</ParamField>

<ParamField body="refit_trainval" type="bool" default="True">
  Whether to refit the best model on the combined train+validation set.
</ParamField>

<ParamField body="refit_full" type="bool" default="False">
  Whether to additionally refit on the full dataset (train+validation+test).
</ParamField>

***

## Training

### `Trainer`

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

```python theme={null}
from kumoai.trainer import Trainer

trainer = Trainer(model_plan=model_plan)
result = trainer.fit(graph=graph, train_table=train_table)
```

<ParamField body="model_plan" type="ModelPlan" required>
  The model plan specifying architecture, optimization, and sampling parameters.
</ParamField>

#### `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.

<ParamField body="graph" type="Graph" required>
  The relational graph.
</ParamField>

<ParamField body="train_table" type="Union[TrainingTable, TrainingTableJob]" default="None">
  An optional pre-generated training table. Provide exactly one of `train_table` or `pquery`.
</ParamField>

<ParamField body="pquery" type="Optional[PredictiveQuery]" default="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.
</ParamField>

<ParamField body="non_blocking" type="bool" default="False">
  If `True`, returns a `TrainingJob` immediately rather than blocking.
</ParamField>

<ParamField body="custom_tags" type="Mapping[str, str]" default="{}">
  Optional key-value tags attached to the training job.
</ParamField>

<ParamField body="warm_start_job_id" type="str" default="None">
  Training job ID to warm-start from (initializes from an existing model's weights).
</ParamField>

**Returns** `Union[TrainingJob, TrainingJobResult]`

#### `predict()`

Generates batch predictions using the trained model.

<ParamField body="graph" type="Graph" required>
  The relational graph.
</ParamField>

<ParamField body="prediction_table" type="Union[PredictionTable, PredictionTableJob]" default="None">
  The prediction table generated from a `PredictiveQuery`.
</ParamField>

<ParamField body="output_config" type="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).
</ParamField>

<ParamField body="output_connector" type="Connector" default="None">
  **Deprecated.** Raises ValueError when passed. Use output\_config instead.
</ParamField>

<ParamField body="output_table_name" type="str" default="None">
  **Deprecated.** Raises ValueError when passed. Use output\_config instead.
</ParamField>

<ParamField body="non_blocking" type="bool" default="False">
  If `True`, returns a `BatchPredictionJob` immediately.
</ParamField>

<ParamField body="custom_tags" type="Mapping[str, str]" default="{}">
  Optional key-value tags attached to the prediction job.
</ParamField>

<ParamField body="training_job_id" type="str" default="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.
</ParamField>

<ParamField body="binary_classification_threshold" type="float" default="None">
  For binary classification models, the score threshold above which predictions are classified as 1. If None, the raw probability score is returned.
</ParamField>

<ParamField body="num_classes_to_return" type="Optional[int]" default="None">
  For ranking task models, the number of classes to return in the prediction output.
</ParamField>

<ParamField body="num_workers" type="int" default="1">
  Number of parallel workers for batch prediction. Values greater than 1 partition the prediction table and process in parallel.
</ParamField>

<ParamField body="prediction_time" type="datetime" default="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.
</ParamField>

<ParamField body="explanations" type="bool" default="False">
  If True, generates per-prediction feature attributions in addition to prediction outputs.
</ParamField>

**Returns** `Union[BatchPredictionJob, BatchPredictionJobResult]`

#### `load()` `classmethod`

Loads a trained `Trainer` from a completed training job.

<ParamField body="training_job_id" type="str" required>
  The training job ID.
</ParamField>

**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.

```python theme={null}
result = trainer.fit(graph=graph, train_table=train_table)
metrics = result.metrics()
```

<ParamField body="job_id" type="TrainingJobID" required>
  The training job ID.
</ParamField>

#### `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.

<ParamField body="query" type="str" required>
  The PQL query string.
</ParamField>

<ParamField body="indices" type="Sequence[Union[str, float, int]]" default="None">
  Entity indices to explain. Explains all entities if `None`.
</ParamField>

<ParamField body="run_mode" type="RunMode" default="RunMode.FAST">
  The run mode for explanation computation.
</ParamField>

<ParamField body="num_neighbors" type="List[int]" default="None">
  Per-hop neighbor counts for subgraph sampling.
</ParamField>

<ParamField body="anchor_time" type="Union[pd.Timestamp, Literal['entity']]" default="None">
  The anchor time for temporal explanation.
</ParamField>

**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](/fine-tuning/online-serving) for the end-to-end flow.

### `DistillationTrainer`

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

```python theme={null}
from kumoai.trainer import DistillationTrainer

distil = DistillationTrainer(model_plan=distil_plan, base_training_job_id="<job_id>")
result = distil.fit(graph=graph, train_table=train_table)
```

<ParamField body="model_plan" type="DistilledModelPlan" required>
  The distilled model plan.
</ParamField>

<ParamField body="base_training_job_id" type="str" required>
  The training job ID of the base GNN model to distill from.
</ParamField>

#### `is_trained` `property`

**Returns** `bool`

#### `fit()`

<ParamField body="graph" type="Graph" required>
  The relational graph.
</ParamField>

<ParamField body="train_table" type="Union[TrainingTable, TrainingTableJob]" required>
  The training table.
</ParamField>

<ParamField body="non_blocking" type="bool" default="False">
  If `True`, returns a `TrainingJob` immediately.
</ParamField>

<ParamField body="custom_tags" type="Mapping[str, str]" default="{}">
  Optional job tags.
</ParamField>

**Returns** `Union[TrainingJob, TrainingJobResult]`

#### `load()` `classmethod`

<ParamField body="job_id" type="str" required>
  The training job ID.
</ParamField>

**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.

<ParamField body="embedding_keys" type="List[str]" default="inferred">
  Column keys used as embedding inputs.
</ParamField>

<ParamField body="max_embedding_offset" type="TimeOffset" default="inferred">
  Maximum time offset for embedding lookups.
</ParamField>

<ParamField body="min_embedding_offset" type="TimeOffset" default="inferred">
  Minimum time offset for embedding lookups.
</ParamField>

<ParamField body="real_time_interactions" type="Dict[str, int]" default="{}">
  Real-time interaction table configuration.
</ParamField>

***

### `export_model()`

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

```python theme={null}
from kumoai.trainer import export_model, ModelOutputConfig

result = export_model(config=output_config, non_blocking=False)
```

<ParamField body="config" type="ModelOutputConfig" required>
  Specifies the training job, output path, and batch prediction job to bundle.
</ParamField>

<ParamField body="non_blocking" type="bool" default="True">
  If `True`, returns an `ArtifactExportJob` immediately.
</ParamField>

<ParamField body="online_embedding_format" type="Optional[Literal['pandas', 'mphf', 'lmdb']]" default="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.
</ParamField>

<ParamField body="payload_version" type="Optional[int]" default="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`.
</ParamField>

**Returns** `Union[ArtifactExportJob, ArtifactExportResult]`

***

### `ModelOutputConfig`

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

<ParamField body="output_types" type="Set[str]" required>
  The output types to produce. Valid values: `"predictions"`, `"embeddings"`, or both.
</ParamField>

<ParamField body="output_connector" type="Connector" default="None">
  The connector to write outputs to. Local download only if `None`.
</ParamField>

<ParamField body="output_table_name" type="Union[str, Tuple[str, str]]" default="None">
  Table name in the output connector. For Databricks, provide a `(schema, table)` tuple.
</ParamField>

<ParamField body="output_metadata_fields" type="List[MetadataField]" default="None">
  Additional metadata columns to include in prediction output. Options: `JOB_TIMESTAMP`, `ANCHOR_TIMESTAMP`.
</ParamField>

***

### `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.