Score in Batch

Score a registered Kumo model in batch with trainer.predict() inside a kumo.mlops block. When you finish this page, you can score the champion version over its whole population, and you know when a platform’s own model API is the right choice instead.

This page continues from Register and Promote Models: you need a model that kumo.MLOps registered, with a champion alias.

Score With trainer.predict

Use trainer.predict() for batch prediction. The entities to score come from the predictive query, so you do not pass a list of ids. The scores are written to Parquet files on local disk, not held in memory, and stored under the artifact root when the run finishes. The prediction run is recorded in your tracker, so Monitor Predictions can compare and evaluate it.

The following example loads the champion version and scores today’s snapshot. cfg is the kumo.MLOps configuration and graph the graph from Register and Promote Models:

champion = kumo.Trainer.load_from_connector(
    cfg, location=kumo.ModelRegistryLocation("users-churn", alias="champion")
)
with kumo.mlops(cfg):
    snapshot = graph.snapshot()
    prediction = champion.predict(snapshot)
print(prediction.id, prediction.num_rows)

The scores are stored under predictions/<prediction id>/ in the artifact root. To choose the population, the prediction time, and an output table or path, see Making Predictions. To run this on a schedule, see Schedule Prediction Jobs.

Score Through a Platform’s Model API

Every model that Kumo registers through MlflowTracking is also an MLflow pyfunc model, and a version that Kumo registers in the Snowflake Model Registry can run in a Snowflake batch inference job (preview). Both paths differ from trainer.predict():

• They score the entity ids you pass, not the query’s population.
• Each serving process first copies the whole graph snapshot to its local disk, once per worker, before it scores anything. Size each worker’s local disk for the full snapshot. trainer.predict() needs the snapshot once, on the job’s machine.
• Each call returns its results in memory, and each call has a fixed setup cost. Pass ids in large chunks rather than a few at a time.
• The scores reflect the snapshot you name, not the latest data, unless you store and pass a newer snapshot.
• Nothing is recorded. Called outside a kumo.mlops block, as in the examples below, these paths store no prediction run under the artifact root and log no run in your tracker, so Monitor Predictions cannot compare or evaluate their scores.

These paths therefore suit small graphs only.

MLflow Pyfunc

Pass the model a DataFrame of entity ids and the graph snapshot to score against. The process that calls it needs Kumo Engine installed with the extra for your tracker, as listed in Turn On MLOps.

The following example loads the champion version of users-churn and scores three users. snapshot is a graph snapshot stored under the artifact root: one that graph.snapshot() created inside a kumo.mlops block, or one loaded with kumo.GraphSnapshot.load_from_connector(cfg, id="<snapshot id>"). Replace the prediction_time value with the time to score as of:

import mlflow
import pandas as pd
with cfg.tracking.client():
    model = mlflow.pyfunc.load_model("models:/users-churn@champion")
scores = model.predict(
    pd.DataFrame({"ENTITY": [1, 2, 3]}),
    params={
        "snapshot": cfg.snapshot_uri(snapshot),
        "prediction_time": "2026-11-01",
    },
)
print(scores)

cfg.tracking.client() points MLflow at your tracker’s server and credentials while the model loads. cfg.snapshot_uri(snapshot) returns where the snapshot is stored under the artifact root; it also takes a snapshot id.

The result has one row per scored input row, in input order, with the same columns as a prediction run, such as ENTITY, TARGET_PRED, False_PROB, and True_PROB for binary classification. For the same entities, snapshot, and prediction time, the scores match those of trainer.predict().

The input DataFrame needs an ENTITY column and can have a TIMESTAMP column with a prediction time per row. params takes:

A snapshot on Amazon S3, Google Cloud Storage, or local disk is read with your environment’s default credentials. To read a snapshot on a Unity Catalog volume from outside a Databricks cluster, or on a Snowflake stage, load the model with a connector that can read it, such as the connector of your artifact root:

with cfg.tracking.client():
    model = mlflow.pyfunc.load_model(
        "models:/users-churn@champion", model_config={"assets_from": dbx}
    )

dbx is the kumo.DatabricksConnector of a (dbx, "/Volumes/...") artifact root, as in Turn On MLOps; for a stage, pass the kumo.SnowflakeConnector instead.

Load only models from registries and publishers you trust.

Snowflake Batch Inference (Preview)

A version that kumo.MLOps registers in the Snowflake Model Registry, with a Snowflake stage as the artifact root, can score in a Snowflake ML batch inference job on a GPU compute pool. The following example scores the users in ML.MODELS.USERS with the champion version of ML.MODELS.CHURN_GNN. Replace <snapshot stage path> with the value cfg.snapshot_uri(snapshot) returns, and KUMO_GPU_POOL with your GPU compute pool. Run it in a Snowpark session, such as get_active_session() in a Snowflake notebook, whose role has USAGE on the compute pool and a warehouse, CREATE SERVICE on the schema, ownership or USAGE on the model, and READ and WRITE on the artifact root’s stage and the output stage:

from snowflake.ml.model.batch_inference import InferenceSpec, InputSpec, OutputSpec, ResourcesSpec
from snowflake.ml.registry import Registry
model_version = (
    Registry(session, database_name="ML", schema_name="MODELS")
    .get_model("CHURN_GNN")
    .version("champion")
)
job = model_version.run_batch(
    session.sql("SELECT USER_ID::STRING AS ENTITY FROM ML.MODELS.USERS"),
    compute_pool="KUMO_GPU_POOL",
    output_spec=OutputSpec(stage_location="@ML.MODELS.KUMO_OUT/churn/2026-11-01/"),
    input_spec=InputSpec(params={"snapshot": "<snapshot stage path>", "prediction_time": "2026-11-01"}),
    resources_spec=ResourcesSpec(gpu_requests="1"),
    inference_spec=InferenceSpec(num_workers=1),
)
job.wait()

params takes the same values as the pyfunc. Without snapshot, the version scores against the snapshot the fit stored, so its scores reflect the data as of training. When job.wait() returns, the scores are Parquet files under the stage_location you gave, with the same columns as a prediction run. The first job also builds the inference image and can take 12 to 17 minutes. These batch inference results come from runs on an earlier Kumo Engine build; this section has not been rerun on the current release.

Next Steps