Hold MLOps artifacts in Snowflake stages and score registry versions with run_batch by Vedaant-J · Pull Request #115 · NVIDIA/kumo-engine

Vedaant-J · 2026-10-24T19:29:28Z

Stack: #53 → #112 → #50 → #60 → #64; #114 on #112; #115 on #64. Merge bottom to top.

Summary

Teams whose data lives only in Snowflake have nowhere to keep Kumo artifacts and cannot score a registered Kumo model inside Snowflake. This PR adds both:

Stages as a file store. A @DB.SCHEMA.STAGE/... path is a file destination of SnowflakeConnector (Let any sink hold files: FileStoreSink protocol and a chunked filesystem #112's FileStoreSink, as Hold MLOps artifacts on Unity Catalog volumes through the Files API #114 does for volumes): bundles, snapshot and prediction exports, and the MLOps artifact root. Writes use PUT; reads use presigned ranged URLs, falling back to one GET per file. Files up to 4 GiB are plain objects; larger ones are parts only kumo reads.
Runnable registry versions. With kumo-engine installed, a version registered in the Snowflake Model Registry is a KumoModel that scores entity ids against a stage snapshot in Snowpark Container Services (ModelVersion.run_batch on a GPU pool). kumo-snowflake finds the builder through the kumo_snowflake.registry_models entry point. Registering a bundle again reuses its version when the serving environment matches.

In the container, the snapshot and the stage assets are read with the container's own Snowflake identity (_snowflake_serving.py).

Usage

sf = kumo.SnowflakeConnector(warehouse="WH", database="ML", schema_name="MODELS")
cfg = kumo.MLOps((sf, "@ML.MODELS.KUMO"), tracking=SnowflakeTracking(sf, experiment_name="CHURN"),
                 registry=kumo.ModelRegistryLocation("CHURN_GNN", alias="champion"))
with kumo.mlops(cfg):                    # the version serves the fit's snapshot on the stage
    result = trainer.fit(train_table, snapshot=snapshot)

mv = Registry(session, database_name="ML", schema_name="MODELS").get_model("CHURN_GNN").version("champion")
mv.run_batch(session.sql("SELECT ID::STRING AS ENTITY FROM ML.MODELS.USERS"), compute_pool="KUMO_GPU_POOL",
             output_spec=OutputSpec(stage_location="@ML.MODELS.KUMO_OUT/churn/"))

Back-compat

Situation	Behaviour
Table name on a `SnowflakeConnector`	Unchanged.
`@DB.SCHEMA.STAGE/...` path	Was refused; now written to the stage.
Export into a stage directory that already holds files	Refused before scoring by default or with `APPEND`; `OVERWRITE` empties it first, as on every file store.
`SNOWFLAKE_FULL` stage	Writes work; reads raise `OSError` naming `SNOWFLAKE_SSE`.
Registry publish without kumo-engine, or if the builder fails	The non-scoring carrier, as before (with a warning).
Same bundle registered again	Existing version if the serving environment matches; else a new `V<n>`.
`ModelOutputConfig`, `SnowflakeTracking`	New optional `serving_snapshot` / `serving_options`; existing calls unchanged.
`run_batch` output	Scores only; not a Kumo prediction run for `evaluate` or `score_drift`.

Testing

kumo-snowflake/tests/test_stage_fs.py (in-memory stage): routing, bundles and Parquet through a stage path, listings, exact removes, ranged and GET reads, the occupied-directory rule, refusals, no signed URL in errors.
kumo-snowflake/tests/test_registry.py: version reuse, builder arguments, environment refusal, carrier fallback, bundle-only fetch.
kumo-engine/test/test_snowflake_model.py and test_model_agnostic_mlops.py: a version packaged by snowflake-ml scores like Trainer.predict; serving options; staged assets load through the loading connector; a stage root registers versions serving its snapshot.
Live on a Snowflake test account with synthetic data: stage file store and an MLOps root on a stage; run_batch binary (300/300 rows, labels equal to local predict) and regression on a GPU pool; all steps as a role with only the documented grants.

Change map: see the kumo-pr-change-map comment.

🤖 Generated with Claude Code

copy-pr-bot · 2026-10-24T19:29:31Z

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

coderabbitai · 2026-10-24T19:29:34Z

Important

Review skipped

Auto reviews are disabled on base/target branches other than the default branch.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

Run configuration

Source mode: Auto
Repository: NVIDIA/kumo-engine
Branch: vedaant/mlops-sf-stage
File: .coderabbit.yaml
Commit: c9f10fc
Layers: Repository YAML
Review profile: QUIET
Plan: Enterprise
Run ID: 39684a7e-39ca-4deb-955d-f7ca8fab62ef

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

Trigger review

Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

Vedaant-J · 2026-10-24T19:30:01Z