> For clean Markdown of any page, append .md to the page URL. > For a complete documentation index, see https://docs.nvidia.com/sdgm/llms.txt. > For AI client integration (Claude Code, Cursor, etc.), connect to the MCP server at https://docs.nvidia.com/sdgm/_mcp/server. # 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 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. ```python 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] ``` AutoML job-level settings. Encoder overrides for individual columns. Subgraph sampling configuration. Optimization hyperparameters. GNN or Graph Transformer architecture parameters. *** ### `ColumnProcessingPlan` Specifies encoder overrides and missing value strategy overrides for individual table columns. ```python from kumoai.encoder import GloVe from kumoai.trainer import ColumnProcessingPlan plan = ColumnProcessingPlan( encoder_overrides={"products.description": GloVe(model_name="glove-wiki-gigaword-50")}, ) ``` A mapping from `"table.column"` to an encoder instance. Overrides Kumo's auto-inferred encoder for that column. 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. ```python from kumoai.trainer import GNNModelPlan arch = GNNModelPlan(channels=[64, 128], aggregation=[["sum", "mean"]]) ``` Candidate hidden channel sizes for AutoML search. Candidate aggregation function combinations for AutoML search. Candidate dropout rates for AutoML search. *** ### `GraphTransformerModelPlan` Configures a Graph Transformer architecture. ```python from kumoai.trainer import GraphTransformerModelPlan arch = GraphTransformerModelPlan(num_layers=[2, 4], num_heads=[4, 8]) ``` Candidate hidden channel sizes. Candidate number of transformer layers. Candidate number of attention heads. Candidate dropout rates. Candidate positional encoding combinations. *** ### `NeighborSamplingPlan` Controls how Kumo samples subgraphs during training. Candidate per-hop neighbor counts for AutoML search. Each inner list specifies the number of neighbors to sample at each hop. Whether to sample neighbors from the entity table. *** ### `OptimizationPlan` Controls learning rate, batch size, epochs, and other training optimization parameters. Maximum number of training epochs. Maximum number of training steps per epoch. Maximum number of validation steps. Maximum number of test steps. Candidate loss functions for AutoML search. Candidate base learning rates. Candidate weight decay values. Candidate batch sizes. Candidate early stopping configurations. Candidate learning rate scheduler configurations. Candidate majority class sampling ratios for imbalanced classification. Candidate sample weighting modes. *** ### `TrainingJobPlan` AutoML job-level settings controlling the number of experiments and evaluation metrics. Number of hyperparameter experiments to run during AutoML. Evaluation metrics to compute. The primary metric used to select the best model. Whether to refit the best model on the combined train+validation set. 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). ```python from kumoai.trainer import Trainer trainer = Trainer(model_plan=model_plan) result = trainer.fit(graph=graph, train_table=train_table) ``` 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. The relational graph. An optional pre-generated training table. Provide exactly one of `train_table` or `pquery`. 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. If `True`, returns a `TrainingJob` immediately rather than blocking. Optional key-value tags attached to the training job. 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. The relational graph. The prediction table generated from a `PredictiveQuery`. 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). **Deprecated.** Raises ValueError when passed. Use output\_config instead. **Deprecated.** Raises ValueError when passed. Use output\_config instead. If `True`, returns a `BatchPredictionJob` immediately. Optional key-value tags attached to the prediction job. 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. For binary classification models, the score threshold above which predictions are classified as 1. If None, the raw probability score is returned. For ranking task models, the number of classes to return in the prediction output. Number of parallel workers for batch prediction. Values greater than 1 partition the prediction table and process in parallel. 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. 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. 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. ```python result = trainer.fit(graph=graph, train_table=train_table) metrics = result.metrics() ``` 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. The PQL query string. Entity indices to explain. Explains all entities if `None`. The run mode for explanation computation. Per-hop neighbor counts for subgraph sampling. 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](/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 from kumoai.trainer import DistillationTrainer distil = DistillationTrainer(model_plan=distil_plan, base_training_job_id="") result = distil.fit(graph=graph, train_table=train_table) ``` The distilled model plan. The training job ID of the base GNN model to distill from. #### `is_trained` `property` **Returns** `bool` #### `fit()` The relational graph. The training table. If `True`, returns a `TrainingJob` immediately. Optional job tags. **Returns** `Union[TrainingJob, TrainingJobResult]` #### `load()` `classmethod` 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. Column keys used as embedding inputs. Maximum time offset for embedding lookups. Minimum time offset for embedding lookups. 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. ```python from kumoai.trainer import export_model, ModelOutputConfig result = export_model(config=output_config, non_blocking=False) ``` Specifies the training job, output path, and batch prediction job to bundle. If `True`, returns an `ArtifactExportJob` immediately. 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. 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. The output types to produce. Valid values: `"predictions"`, `"embeddings"`, or both. The connector to write outputs to. Local download only if `None`. Table name in the output connector. For Databricks, provide a `(schema, table)` tuple. 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.