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

# Training

> Configure a model plan, fit it against a graph snapshot, and inspect its metrics and saved artifacts

Train a model using the `query`, `training_table`, and `snapshot` created in [Predictive Queries](/kumo-engine/1.0.0/user-code-guide/train-a-model/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:

```python
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:

| Section              | Controls                                                                               |
| -------------------- | -------------------------------------------------------------------------------------- |
| `training_job`       | Number of experiments, built-in evaluation metrics, tuning metric, and refit behavior. |
| `optimization`       | Epoch and step limits, candidate batch sizes, and optimization settings.               |
| `column_processing`  | Column encoder overrides and missing-value strategies.                                 |
| `neighbor_sampling`  | How the model samples graph neighborhoods.                                             |
| `model_architecture` | Model architecture settings.                                                           |

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

```python
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](/kumo-engine/1.0.0/user-code-guide/train-a-model/predictive-queries#generate-a-training-table).
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](/kumo-engine/1.0.0/user-code-guide/train-a-model/model-plan-guide).

## Train a Model

Construct a trainer and select the graph snapshot explicitly:

```python
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)
```

```text
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](/kumo-engine/1.0.0/user-code-guide/train-a-model/publish-and-retrieve-models#retrieve-a-published-model).

## 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:

```python
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](/kumo-engine/1.0.0/user-code-guide/train-a-model/publish-and-retrieve-models#retrieve-a-published-model).
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:

| Attribute           | Contents                                                                                |
| ------------------- | --------------------------------------------------------------------------------------- |
| `metrics`           | Dictionary of metric names and values. Available keys depend on the task and run.       |
| `artifact_paths`    | Dictionary mapping artifact roles to local working paths.                               |
| `bundle_dir`        | Path to the complete model bundle, ready for `Trainer.load()`.                          |
| `published_model`   | Publication name, version, and URI when an export was requested; otherwise `None`.      |
| `holdout_path`      | Local holdout prediction file, or `None` when no file was produced.                     |
| `holdout_truncated` | Whether saved holdout rows were truncated; `None` when size metadata could not be read. |
| `holdout_df()`      | Reads the whole saved holdout file into a pandas DataFrame.                             |

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

```python
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](/kumo-engine/1.0.0/user-code-guide/train-a-model/training-diagnostics#read-retained-model-metrics) to read metrics from a retained bundle and [read holdout predictions](/kumo-engine/1.0.0/user-code-guide/train-a-model/training-diagnostics#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](/kumo-engine/1.0.0/user-code-guide/train-a-model/publish-and-retrieve-models) for Linux, object storage, Snowflake Registry, and Databricks volume procedures.

Continue to [Training Diagnostics](/kumo-engine/1.0.0/user-code-guide/train-a-model/training-diagnostics) to inspect epoch and experiment history, [Model Explainability](/kumo-engine/1.0.0/user-code-guide/train-a-model/model-explainability) to inspect any explanation files produced by this fit, or [Making Predictions](/kumo-engine/1.0.0/user-code-guide/making-predictions) to score entities and generate embeddings.