Production Workflow

Put an evaluated model into production by loading a fixed model version, scoring fresh data, and delivering results to your application. Start with a model from Training and a successful batch from Making Predictions.

Prepare the Production Job

Make the prediction script or notebook independent of the training session. In every new process:

1. Prepare the runtime. Install the required packages, configure local SSD scratch, and establish source and output access. For models with text columns, make the trained text embedding models available again, including any custom vectors.
2. Load the model. Retain the complete bundle outside scratch and load its exact path or registry version.
3. Read fresh data. Reconnect to production tables, recreate compatible graph and predictive-query definitions, and materialize a snapshot after the upstream refresh completes. A model bundle does not contain source data or the graph snapshot.
4. Score the population. Generate new prediction rows using the loaded trainer’s training_table_generation_plan. Set the temporal anchor where applicable and the task’s threshold or candidate count. A saved prediction table keeps its original population and anchors; a new snapshot alone does not regenerate them.
5. Deliver the results. Export predictions to the application’s table or files. Check saved counts, keyed scores, and application read access before accepting the batch.

Retain the model version, source or snapshot version, prediction anchor, output location, and batch status together. Keep the previous accepted output available until the new batch passes its checks. Run the complete script once in its intended execution environment before enabling its schedule.

Example: Productionize Your Model in Snowflake

This example turns a temporal binary classification model for customer purchases into a recurring batch job. Keep the graph and predictive query from your successful training workflow. The example uses these illustrative destinations:

Replace these locations with schemas and a stage your execution role can access. The model registry is durable model storage. The payload stage holds executable job files. Scratch is recreated for each job and must not be the only copy of a model you need to retain.

Publish the Model to Its Registry Schema

In the training notebook, create a registry connector for ML.MODELS. Set WAREHOUSE to the query warehouse your role can use. This is separate from the compute pool that runs Python and GPU work.

from kumo.engine import (
    ModelOutputConfig,
    ModelRegistryLocation,
    SnowflakeConnector,
)
WAREHOUSE = "COMPUTE_WH"
registry_connector = SnowflakeConnector.from_spcs_env(
    warehouse=WAREHOUSE,
    database="ML",
    schema_name="MODELS",
)

Use trainer, training_table, and snapshot from Training. Add publication to that fit call instead of running a second fit solely to save the model:

fit_result = trainer.fit(
    training_table,
    snapshot=snapshot,
    export_config=ModelOutputConfig(
        registry_connector,
        location=ModelRegistryLocation("CUSTOMER_ORDER_MODEL"),
    ),
)
published = fit_result.published_model
print("Model:", published.name)
print("Version:", published.version)

This saves the complete bundle as a version of ML.MODELS.CUSTOMER_ORDER_MODEL. Record the returned name and version with the evaluation results. Select an accepted version for the production job; do not assume the registry’s default version is the one you evaluated. See Publish and Retrieve Models for publication permissions and retained artifacts.

Prepare the Prediction Script

Create predict_batch.py from your working prediction script. Include its imports and graph and query definitions so it can run without notebook state. Use the same compatible Engine build and required text embedding assets as the training environment.

Provide stable values through the job configuration and the batch anchor and output name as arguments to each invocation:

At the start of the script, configure scratch and load the selected model. The registry connector must use the same database and schema as publication. For custom text embeddings, first recreate the matching store as described in Encoder Selection, then add embedding_models=store to the load call below.

import argparse
import os
from datetime import datetime
from pathlib import Path
from kumo.engine import (
    ModelRegistryLocation,
    PredictionOutputConfig,
    SnowflakeConnector,
    Trainer,
    configure,
)
from kumoapi.jobs import MetadataField, PredictionArtifactType, WriteMode
parser = argparse.ArgumentParser()
parser.add_argument("--prediction-time", required=True)
parser.add_argument("--output-table", required=True)
batch = parser.parse_args()
scratch = Path(os.environ["KUMO_SCRATCH"])
scratch.mkdir(parents=True, exist_ok=True)
configure(scratch_dir=str(scratch))
warehouse = os.environ["KUMO_WAREHOUSE"]
connector = SnowflakeConnector.from_spcs_env(
    warehouse=warehouse, database="PROD", schema_name="SALES",
)
registry_connector = SnowflakeConnector.from_spcs_env(
    warehouse=warehouse, database="ML", schema_name="MODELS",
)
output_connector = SnowflakeConnector.from_spcs_env(
    warehouse=warehouse, database="ML", schema_name="PREDICTIONS",
)
trainer = Trainer.load_from_connector(
    registry_connector,
    location=ModelRegistryLocation(
        os.environ["KUMO_MODEL_NAME"],
        version=os.environ["KUMO_MODEL_VERSION"],
    ),
)

from_spcs_env() uses the runtime’s account, host, and mounted OAuth token. Supply the warehouse explicitly. No token needs to be copied into the script. The connection uses the job service’s execution identity; grant access to that role, not just to the person who tested the notebook. See Snowflake service-user credentials.

Retrieval downloads the bundle into job-local scratch. Grant the job role USAGE on the registry database and schema and READ on the model to retrieve its files. Model USAGE alone does not grant access to those files. See Snowflake model privileges.

In this workflow, Kumo Engine loads the bundle and runs prediction in the ML Job on a GPU-enabled compute pool. The query warehouse handles data access through the connectors.

Next, include the table, graph, and predictive-query construction from your training workflow, using connector for the production tables. This must define graph and query before the following code runs. Preserve the trained target, entity types, relationships, and feature requirements. Then create fresh prediction inputs and write a batch table:

prediction_time = datetime.fromisoformat(batch.prediction_time)
if prediction_time.tzinfo is not None:
    raise ValueError("--prediction-time must be UTC without a timezone suffix")
snapshot = graph.snapshot()
prediction_table = query.generate_prediction_table(
    snapshot,
    prediction_time=prediction_time,
    plan=trainer.training_table_generation_plan,
)
result = trainer.predict(
    prediction_table,
    snapshot=snapshot,
    binary_classification_threshold=0.5,
    output_config=PredictionOutputConfig(
        connector=output_connector,
        artifact_type=PredictionArtifactType.PREDICTIONS,
        table_name=batch.output_table,
        write_mode=WriteMode.OVERWRITE,
        extra_fields=[MetadataField.ANCHOR_TIMESTAMP],
    ),
)
print("Predicted rows:", result.num_rows)

Choose the classification threshold using your evaluation results; 0.5 is illustrative. For other tasks, use the task-specific prediction options. For a static query, remove the prediction-time argument and parsing, omit prediction_time from table generation, and omit MetadataField.ANCHOR_TIMESTAMP from output configuration.

The output role needs the Snowflake output permissions. Read the saved table through the application’s access path and compare its row count and sample keyed scores with the local result before accepting the batch. This example overwrites its named table on rerun. Use a new table name when you need to preserve an earlier output.

Run the Script as an ML Job

Register the completed predict_batch.py from a Snowflake notebook or a client with an active Snowpark session. The submitting environment needs snowflake-ml-python version 1.26.0 or later for ML Job Definitions. Set COMPUTE_POOL to a pool meeting your compute requirements, RUNTIME_VERSION to its compatible Snowflake Container Runtime version, ENGINE_REQUIREMENT to the pinned kumo-engine[snowflake] version used for training, and JOB_SCRATCH to its local SSD directory. Replace the illustrative names below, including the external access integration used for package installation.

from snowflake.ml.jobs import MLJobDefinition
from snowflake.snowpark.context import get_active_session
definition = MLJobDefinition.register(
    source="predict_batch.py",
    compute_pool=COMPUTE_POOL,
    stage_name="ML.JOBS.PREDICTION_PAYLOADS",
    database="ML",
    schema="JOBS",
    session=get_active_session(),
    runtime_environment=RUNTIME_VERSION,
    pip_requirements=[ENGINE_REQUIREMENT],
    external_access_integrations=["PACKAGE_FEED_EAI"],
    env_vars={
        "KUMO_SCRATCH": JOB_SCRATCH,
        "KUMO_WAREHOUSE": WAREHOUSE,
        "KUMO_MODEL_NAME": published.name,
        "KUMO_MODEL_VERSION": published.version,
    },
)
job = definition(
    "--prediction-time", "2026-10-13T00:00:00",
    "--output-table", "CUSTOMER_PURCHASE_SCORES_20261013",
)
job.result()

The published record above comes from the accepted training run. If submitting from another session, replace its fields with the recorded model name and version. Add the text extra or custom embedding dependencies to the job environment when needed. Include any imported graph-building modules with the payload; directory submission and imports describe how to package them.

The definition retains stable configuration, including the selected model version. Pass the anchor and output table on each invocation so a daily schedule does not repeatedly score the first day’s anchor or overwrite its table. For a retry, use the intended batch’s original anchor and deliberate output policy.

Configure package-feed access using Snowflake’s dependency instructions. The job role also needs the ML Job privileges, source-read access, and output-write access. The compute pool runs prediction; the query warehouse serves connector operations.

job.result() waits for completion and raises if the job fails. Check the saved table before delivering the batch. After the complete batch succeeds, connect the definition to Snowflake Tasks and arrange for each scheduled invocation to receive its batch arguments.

Schedule Predictions and Model Refresh

Choose a prediction cadence that meets the application’s freshness deadline, and define how failed batches affect delivery. Keep prediction and retraining policies separate: an existing model can score newer, compatible data without a new fit.