Publish and Retrieve Models

Publish a trained model when you need to reuse it after the current Python process or notebook session ends. Publication copies the complete model bundle to durable storage or a supported model registry. Retrieval validates that bundle and restores a fitted Trainer. Connector retrieval first fetches a local copy.

This page continues from Training. Choose the retention method by destination. Linux deployments commonly use a retained filesystem, S3, or GCS. Snowflake notebooks can use Snowflake Model Registry. The demonstrated Databricks notebook copies the complete bundle to a Unity Catalog volume and loads it from that path in a later process.

Understand What a Model Bundle Contains

Every successful fit creates a reloadable bundle at fit_result.bundle_dir. The bundle includes the trained checkpoint, model plan, query, effective training-table generation settings, metrics, and manifests required by Trainer.load().

The bundle does not include:

• Source tables or credentials
• The materialized graph snapshot or index
• Generated training or prediction inputs
• Prediction or embedding outputs, including the local holdout prediction file
• Complete epoch-by-epoch training history
• Post-training explanation files

Retain these artifacts separately before deleting scratch if you need to reuse them. Keep the entire bundle unchanged. A checkpoint by itself is not sufficient, and a model bundle cannot recreate missing source data or graph storage. Use Training Diagnostics to retain holdout predictions and detailed training history separately, and Model Explainability for explanation files produced for supported tasks.

Publish During Training

Supply ModelOutputConfig to the fit call when you want a durable model copy. This replaces the ordinary fit call; it does not require an additional training run.

ModelOutputConfig(connector, path=None, location=None) requires exactly one destination:

• Use path for filesystem or object-storage publication.
• Use location with ModelRegistryLocation for a supported model registry.

from kumo.engine import FileConnector, ModelOutputConfig
destination = FileConnector(root_dir="/path/to/retained-models")
fit_result = trainer.fit(
    training_table,
    snapshot=snapshot,
    export_config=ModelOutputConfig(destination, path="model-v1"),
)
print(fit_result.published_model)

Databricks Notebook

DatabricksConnector does not support model publication through ModelOutputConfig in this local Python path. After a successful fit, copy the complete bundle to a Unity Catalog volume. Use ARTIFACT_ROOT from Databricks Notebook Setup and give each run a separate directory:

from pathlib import Path
import shutil
from uuid import uuid4
run_root = ARTIFACT_ROOT / uuid4().hex
run_root.mkdir(parents=True, exist_ok=False)
model_dir = run_root / "model"
shutil.copytree(Path(fit_result.bundle_dir), model_dir)
print("Record this model path:", model_dir)

Keep the complete copied directory, including its manifests. Record the path and verify that you can load it before deleting local scratch. The demonstrated Databricks run used this copy procedure; it did not publish a model registry entry.

The volume copy survives a cluster restart while the volume and its grants are retained. The graph snapshot, generated training table, working model files, and detailed training history remain in driver local scratch unless retained separately. A model bundle does not contain source data, credentials, or the graph index. See Training Diagnostics to retain history files.

Retrieve a Published Model

Load only model bundles from publishers you trust. Retrieval validates the bundle and restores a fitted trainer for local execution; it does not perform registry-hosted inference.

For models with text embeddings, keep the matching word vectors available in the prediction environment. Retain the text package for its installed models, or recreate the registered embedding store and pass embedding_models=store when loading. Keep any custom embedding module importable after restart. Follow Reload a Trained Text Model and Custom Word Vectors before using the retrieval examples below.

Trainer.load() and Trainer.load_from_connector() restore a completed model for inference. To initialize a new training run from the saved weights, pass the loaded trainer explicitly as warm_start. Loading alone does not resume interrupted training. Calling fit() without warm_start starts fresh.

from kumo.engine import Trainer, configure
configure(scratch_dir="/path/to/local-storage/kumo-scratch")
trainer = Trainer.load("/path/to/retained-models/model-v1")

from kumo.engine import FileConnector
file_connector = FileConnector(root_dir="/path/to/retained-models")
trainer = Trainer.load_from_connector(file_connector, path="model-v1")

trainer = Trainer.load_from_connector(
    storage_connector,
    path="models/model-v1",
)

FileConnector, S3, GCS, and Snowflake support connector-based model retrieval. The local Databricks connector does not publish or retrieve model bundles; use a retained Unity Catalog volume and Trainer.load() as shown above. BigQuery does not support model publication or retrieval through this API. It does support prediction export.

Find a Published Snowflake Model

Use a Snowflake connector whose identity can read the registry in the intended database and schema. List the visible Kumo model versions before choosing one to retrieve:

models = connector.sink_backend.list_models(connector)
for model in models[:10]:
    print(model.name, model.version, model.uri)
    print("Metadata:", dict(list(model.metadata.items())[:5]))
    print("Metrics:", dict(list(model.metrics.items())[:5]))

The listing returns published-model records ordered by name and version, with recorded training metadata and metrics for each version. The slice limits displayed records, not the registry listing itself. Discovery does not download model bundles. Only versions carrying Kumo’s metadata marker are included; an empty list can mean no visible Kumo versions even when other models exist in the registry. Check the connector’s database, schema, and registry permissions if an expected version is missing.

Copy the chosen record’s name and version into the explicit ModelRegistryLocation retrieval example above. After retrieval, inspect the retained configuration and lineage to confirm the query and training snapshot associated with the model.

Recreate the Prediction Context

Loading a model restores its fitted state, but not its data. Before making predictions in a new process:

1. Reconnect to the source data.
2. Recreate compatible table, graph, and predictive-query definitions.
3. Materialize or restore a compatible graph snapshot.
4. Generate the population to score using the restored trainer’s generation settings.

The snapshot can contain updated data, provided the trained target, entity, relationships, and feature requirements remain compatible. Continue to Making Predictions to score entities, generate embeddings, and save the outputs. To repeat this workflow in a scheduled process, follow Production Workflow.