[Vedaant-J]

Series: kumo-engine MLOps series: every snapshot, training table, model and prediction run carries a self-contained _provenance.json; with kumo.mlops(...) stores each artifact under one artifact root, tracks fits and prediction runs in MLflow or Snowflake and registers models; volumes and stages can hold files; prediction runs are monitored for drift and realized performance.
Stack: #53 → #112 → #50 → #60 → #64; #114 on #112; #115 on #64
Merge order: one linear stack; merge bottom to top (#53 first); #114 sits on #112 and merges any time after it; #115 (draft) sits on #64 and merges after it
This PR: Adds Snowflake named internal stages as a file destination of SnowflakeConnector, and a Snowflake Model Registry version that scores in Snowflake.

Draft: every check passed live on a Snowflake test account (synthetic data), including a final round as a role holding only the documented grants, pip_repository through Snowflake's shared PyPI repository, and a regression model. Still unverified: multiclass, ranking and text-embedding models. Details under Testing Done.

Update (10-28): registry through the tracker, snapshots from the stage root, live-run fix

• Registration: kumo.MLOps((sf, "@DB.SCHEMA.STAGE/kumo"), tracking=SnowflakeTracking(sf, experiment_name=...), registry=ModelRegistryLocation(name, alias=...)) registers the runnable version through the tracker's connector (no registry=ModelOutputConfig(...)).
• Serving snapshot by default: with a stage root, the version scores against the fit's snapshot already stored under <root>/snapshots/; no separate snapshot.export is needed. run_batch's snapshot param overrides it per job. SnowflakeTracking(serving_options={...}) sets cuda_version, python_version, function_type, pip_repository.
• Snapshots by object: ModelOutputConfig(serving_snapshot=...) also takes a GraphSnapshot or snapshot id stored under the block's stage root.
• Fix from the live validation: Trainer.load_from_connector(sf, location=...) outside a kumo.mlops block raised MissingDependencyError for stage-stored assets; it now reads them through the loading connector, and a stage location is never read as a local path.
• docs/mlops.md: the Snowflake quickstart uses a stage root; run_batch details move to Reference > Snowflake serving.

Examples further down that show registry=ModelOutputConfig(sf, ..., serving_snapshot=location) predate this update.

Summary

This PR adds Snowflake named internal stages as a file destination of SnowflakeConnector, and a Snowflake Model Registry version that scores in Snowflake. Part A: model bundles, snapshot exports, prediction and training-table exports and the MLOps artifact root (#50) can live on a stage, next to the tables the job reads; it uses the FileStoreSink protocol and chunked filesystem from #112, as #114 does for Databricks volumes. Part B: a model registered through a SnowflakeConnector (#50/#60) becomes a runnable custom model that Snowflake ML batch inference jobs run on a Snowpark Container Services GPU pool, loading its snapshot from a stage through Part A. The PR sits on #64 because Part B builds on the series' registry and scoring code.

Problem. Teams whose data lives in Snowflake and who have no S3/GCS bucket of their own ask:

• Where does a training job running on a GPU box keep its snapshot and model so that a scoring job elsewhere can load them, without a separate object store?
• Can an MLOps artifact root (Record provenance for every artifact under one MLOps artifact root #50) live inside the Snowflake account, under the same grants as the data?
• Can exported prediction Parquet be loaded into a table with one COPY INTO?
• Can a model registered in the Snowflake Model Registry be scored by Snowflake itself (mv.run_batch(...)), on Snowflake compute, without a separate scoring job?

Today none of this works. The Snowflake sink writes tables and registry models only. Any file export to a Snowflake connector raises, and so does an MLOps root on a stage. A registered version is a carrier whose predict raises NotImplementedError, and registering the same bundle twice logs a second version.

Cause. SnowflakeSinkBackend does not implement FileStoreSink, so is_file_destination(snowflake, "@DB.SCHEMA.STAGE/...") is False, and the engine treats the path as a table name. The package already had the stage transport (PUT/GET/LIST/REMOVE in artifacts._StageFiles, GET_PRESIGNED_URL minting in staging.py), but nothing exposed it as a filesystem.

before                                              after
SnowflakeSinkBackend                                SnowflakeSinkBackend
  write_table / write_files   -> table                write_table / write_files   -> table, or files for @DB.SCHEMA.STAGE/... # new
  publish_model / fetch_model -> Model Registry       publish_model / fetch_model -> Model Registry, or path=@...            # new
                                                      holds_files / filesystem    -> FileStoreSink for @DB.SCHEMA.STAGE/...  # new

Change.

1. kumo_snowflake.stage_fs.StageFileSystemHandler is a chunked_fs.ChunkedFileHandler (Let any sink hold files: FileStoreSink protocol and a chunked filesystem #112) over the connector's Snowpark session. It reuses the existing transport. The Snowflake calls it uses:

   Operation	Call
   write	session.file.put(local, '@stage/dir', auto_compress=False, source_compression='NONE', overwrite=True): no .gz suffix, bytes unchanged
   size, list	LIST @stage/path with sizes. LIST matches by prefix, so results are filtered to exact names
   delete a tree	REMOVE @stage/dir/
   delete one file	REMOVE @stage/parent/ PATTERN = '.*/<path>' (Snowflake matches PATTERN against an internal storage path longer than LIST shows, so only the stage-relative suffix is anchored), then a LIST check that raises OSError if the file is still there. A plain REMOVE @stage/a/part-1 would also delete part-10
   ranged read	GET_PRESIGNED_URL (existing presigned_urls, 24 h) plus HTTP Range: bytes=a-b. The response must be a 206 whose Content-Range starts at a
   fallback read	session.file.get('@stage/parent/', dir, pattern=<exact name>) once per file per filesystem, then local reads

   Ranged reads need plaintext over presigned URLs. A stage created with ENCRYPTION = (TYPE = 'SNOWFLAKE_SSE') serves plaintext; the default SNOWFLAKE_FULL serves client-side ciphertext. Once per (account, stage), a 23-byte probe is PUT under @stage/.kumo-preflight/<uuid>/ and bytes 1–4 are fetched by Range, then the probe is removed. A 206 with the right bytes enables ranged reads. If the probe's listed size differs from its 23 bytes, the stage is SNOWFLAKE_FULL: LIST reports padded ciphertext sizes there, so files cannot be read back exactly, and reads raise OSError naming SNOWFLAKE_SSE. Any other failure falls back to GET, with one warning: "create the stage with SNOWFLAKE_SSE", or "could not verify (the role may lack WRITE)" when the probe could not be written. Errors carry a type and status code only, never a signed URL.

   PUT has no small size cap, so a file up to 4 GiB (the same bound as Hold MLOps artifacts on Unity Catalog volumes through the Files API #114) is one staged object, and COPY INTO reads exported Parquet directly. A larger file becomes <name>.chunks/00000… parts plus _SUCCESS (Let any sink hold files: FileStoreSink protocol and a chunked filesystem #112's layout). Writes spool through local disk under the engine scratch directory, at most one part per open writer, because PUT takes a local file.

2. Layout on the stage, captured with the full series on top (snapshot exports, training tables and _provenance.json come from Record provenance for every artifact under one MLOps artifact root #50). In Record provenance for every artifact under one MLOps artifact root #50's current layout each snapshot export sits one level deeper, in a generation directory next to kumo_snapshot.json; this PR's code is the same either way. This is an MLOps fit on synthetic users/events data against the in-memory stage, with the part size lowered to 1 MiB for the run so the model checkpoint splits. Hashes are shortened, and index node files are omitted:

   @DB.SCHEMA.KUMO_ARTIFACTS/kumo/
     snapshots/3c87be5e…/d170475c5bb30521/
       kumo_snapshot.json                 #  2660 bytes
       _provenance.json                   #  1014
       index/_MANIFEST.json               #  2114
       index/col_stats.json               # 29502
       index/nodes/...
     training_tables/tt-2e8c1fb2…/
       part-00000.parquet … part-00002.parquet
       _provenance.json  _KUMO_EXPORT  _SUCCESS
     models/5e03240c…/
       model/checkpoint.chunks/00000 … 00013   # 13 × 1048576 + 423781 bytes (only above the part size)
       model/checkpoint.chunks/_SUCCESS        # 0 bytes, written last
       metrics.json  training-spec.json  publication-manifest.json  prediction-schema.json
       _provenance.json                        # written last
   

   With the default 4 GiB part size, every file in that run is one plain object.

3. Sink routing. holds_files(connector, path) is True only for a fully qualified named stage. It is decided from the path alone, with no I/O. Table names, user stages (@~) and table stages (@%t) keep today's code. For stage paths, write_table, write_files, publish_model(path=) and fetch_model(path=) go through Let any sink hold files: FileStoreSink protocol and a chunked filesystem #112's FilesystemSink, so no second copy of the export or bundle logic is added. Write options the sink does not use itself (such as the manifest sidecars that Record provenance for every artifact under one MLOps artifact root #50's training-table export and Record prediction runs and monitor score drift #64's prediction export pass; the table path does not use them) are forwarded unchanged. A non-stage path= still reaches the registry's own refusal.

4. Usage (documented in the kumo-snowflake README, which points to docs/mlops.md for MLOps roots):

   trainer.fit(training_table, snapshot=snapshot, export_config=ModelOutputConfig(
       snowflake, path="@ML.ARTIFACTS.KUMO/models/m1"))
   output = PredictionOutputConfig(
       snowflake, output_paths={"predictions": "@ML.ARTIFACTS.KUMO/scores/"})

   CREATE STAGE ML.ARTIFACTS.KUMO ENCRYPTION = (TYPE = 'SNOWFLAKE_SSE');
   COPY INTO SCORES FROM @ML.ARTIFACTS.KUMO/scores/
     PATTERN = '.*[.]parquet' FILE_FORMAT = (TYPE = PARQUET)
     MATCH_BY_COLUMN_NAME = CASE_INSENSITIVE;

Grants: USAGE on the database and schema, READ and WRITE on the stage. No warehouse runs: on the live account PUT, LIST, GET, REMOVE and GET_PRESIGNED_URL completed while the session's warehouse stayed suspended.

Destination check before scoring. The sink's preflight_write_files, which Trainer.predict calls before any scorer starts, refuses an occupied stage directory unless write_mode=WriteMode.OVERWRITE; table names are not checked. With #64's default of one directory per run (<path>/<prediction id>/), a daily job's second run writes to a new directory and passes.

Part B: a runnable Snowflake Model Registry version

Cause. kumo_snowflake.registry logged a _BundleModel carrier with a one-column signature (entity → score) whose predict raised. Scoring needs the bundle, a graph snapshot and the engine, and the registry version had no way to reach a snapshot or install the engine.

before                                            after
registry.log_model(                               registry.log_model(
  _BundleModel(bundle),     # predict raises        KumoModel(bundle, serving.json),  # new, scores ENTITY rows
  signatures={'predict': entity -> score},          signatures={'predict': ENTITY -> ENTITY + prediction columns, params},
  metrics=..., task=...)                            pip_requirements=['kumo-engine[snowflake]==<version>'],
                                                    target_platforms=['SNOWPARK_CONTAINER_SERVICES'],
                                                    python_version='<registering Python>',
                                                    code_paths=[kumo, kumo_core, kumoapi, kumopql, kumo_snowflake, diskgraph, torchmetrics],
                                                    options={'cuda_version': '12.6', 'function_type': 'TABLE_FUNCTION'},
                                                    metrics=..., task=...)

Change.

1. kumo.engine.integrations._snowflake_model.KumoModel is a snowflake.ml.model.custom_model.CustomModel. predict(self, df, *, snapshot='', prediction_time='', binary_classification_threshold=-1.0, num_classes_to_return=0) takes an ENTITY column; the keyword parameters are what InputSpec(params=...) passes. snowflake-ml pickles the class by value and rebuilds it in the container with ModelClass(ModelContext(artifacts=...)), so the module imports no kumo code at module level, and the first predict loads the bundle (Trainer.load) and the snapshot and caches both per process (_serving._trainers, _serving._snapshots). An @ snapshot is read with SnowflakeConnector.from_spcs_env(...), the container's own identity, through Part A's stage filesystem. Ids are cast to the snapshot's primary-key type; ENTITY and CLASS come back as strings.

2. Registration (kumo_snowflake.registry.publish_model). kumo-snowflake may not import kumo-engine, so kumo-engine registers its builder under a new entry-point group, kumo_snowflake.registry_models (kumo.engine.integrations._snowflake_model:log_model_args). With no builder (a hosted export), or if the builder fails, the carrier is registered as before, with a warning. The second artifact, serving.json, records the snapshot and the registering connector's warehouse, database and schema:

   {"snapshot": "@ML.MODELS.KUMO/snapshots/3c87be5e…", "warehouse": "WH", "database": "ML", "schema": "MODELS"}

3. ModelOutputConfig(..., serving_snapshot=..., serving_options={...}). _publish passes them to the sink as metadata (kumo.serving_snapshot, kumo.serving.<option>). Options and defaults:

   Option	Default	Logged as
   cuda_version	12.6	options={"cuda_version": ...}
   python_version	the registering Python	python_version=...; another value raises (see below)
   function_type	TABLE_FUNCTION	options={"function_type": ...}; a table function may return a different number of rows than it gets (multiclass, ranking, entities missing from the snapshot)
   pip_repository	none (PyPI)	artifact_repository_map={"pip": ...}: the index for the third-party requirements

   An unknown option raises when the config is built.

4. Dedup. Registering a bundle whose kumo_bundle_id is already a Kumo version of the model (kumo_engine=1 and the same id in the version metrics) returns that version and logs <model> version <v> already holds bundle <id>; reused it, the same shape as the MLflow (Record fits in experiment tracking and publish models to registries #60) and Unity Catalog (Record prediction runs and monitor score drift #64) registries.

5. Shared scoring. predict, _require_inputs, _in_input_order and snapshot_location move from _mlflow_model.py to integrations/_serving.py, unchanged, so the MLflow pyfunc and KumoModel score the same way and the Snowflake model does not import mlflow.

6. Usage (docs/mlops.md, "Batch inference in Snowflake"):

   job = 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/", mode=SaveMode.OVERWRITE),
       input_spec=InputSpec(params={"prediction_time": "2026-10-01", "binary_classification_threshold": 0.5}),
       resources_spec=ResourcesSpec(gpu_requests="1"),
       inference_spec=InferenceSpec(num_workers=1, max_batch_rows=1_000_000))

   The doc lists the GPU pool (GPU_NV_S/GPU_NV_M, GPU_NV_SM on Azure), the grants (USAGE on pool, database, schema, warehouse and model; CREATE SERVICE; stage READ/WRITE; an image repository) and the output layout (<stage_location>/<job name>/ Parquet). The Model Monitor recipe now notes that its SOURCE can be built from a batch job's output; the monitor still never calls the model.

Wheel delivery. Snowflake builds the inference image from the version (pip requirements plus code_paths), and run_batch takes neither a custom image nor an external access integration for the build. kumo's wheels are not on PyPI, so by default the version embeds kumo-engine, kumo-core, kumo-api, kumo-pql, kumo-snowflake and DiskGraph as code (about 100 MB), plus torchmetrics: the build runs uv pip install --extra-index-url https://download.pytorch.org/whl/cu<ver> and uv takes a package only from the first index that has it, where torchmetrics is 1.0.3 (and requests 2.28.1). The pip requirements are the third-party requirements of whichever installed distribution provides kumo_core and kumopql (kumo-core and kumo-pql in a workspace install, kumo-engine in the bundled wheel), merged to one line per package, with ==X.* pins rewritten as ranges (Snowflake rejects them), torch at this installation's minor version, and any faiss pin replaced by faiss-cpu at the same version (kumo's faiss-gpu wheel is not on PyPI; the container only imports faiss). Models with text embeddings also require kumo-glove-6b-50d. serving_options={"pip_repository": ...} makes these third-party requirements install from a Snowflake pip artifact repository instead of PyPI (artifact_repository_map={"pip": ...}); kumo itself stays embedded, because installing kumo-engine from a repository would need an index that holds kumo's wheels.

Python version. snowflake-ml pickles the model class by value (custom.py: cloudpickle.register_pickle_by_value), so the container must run the registering Python: on the live account a 3.14-registered model segfaulted in predict on the 3.12 image. python_version defaults to the registering Python (the image has 3.12 and 3.14), and another value raises; the registry then registers the carrier, with a warning. This replaces the planned 3.12 default.

Any bundle's predictor

score_registered scores with the predictor the bundle's manifest names (#60), and takes its output columns from the serving record, which registration now writes. It casts ENTITY only for the built-in Trainer predictor. When the model was fitted with a stage as its artifact root, its assets (such as word vectors) are read from <root>/assets/ on that stage, with the container's identity, through _deps.asset_locations. test_model_agnostic_mlops.py scores a toy bundle through the Snowflake model and checks that staged asset roots reach the resolver.

Back-compat

Record and expose: nothing that works today changes behaviour.

Situation	Behaviour
Table destination on a SnowflakeConnector	Unchanged: PUT to the user stage and COPY INTO.
Model Registry publish/fetch (location=)	Unchanged.
@DB.SCHEMA.STAGE/... path on a SnowflakeConnector	Was refused (not a file destination); now written to the stage.
StageArtifactLocation (existing snapshot publication)	Unchanged; this PR reuses its transport helpers. The only edit in artifacts.py makes one import lazy.
Stage export read by COPY INTO	Files up to 4 GiB are plain objects; larger ones are .chunks/ parts, which only kumo reads.
Stage export read by an older engine	An older engine has no stage filesystem and cannot read it.
Stage on SNOWFLAKE_FULL	Writes work; reads raise OSError naming SNOWFLAKE_SSE (encrypted sizes cannot be read back).
Read-only role on an SSE stage	Reads work, downloading each file whole once per filesystem, with a warning.
Single-file delete that leaves the file behind	Raises OSError instead of letting a stale file shadow the new one.
File name outside A-Z a-z 0-9 _ $ = + . , / -	ValueError at write time.
Registry publish with kumo-engine installed	New: a runnable KumoModel version with pip requirements, SPCS target and the options above; same name, version, alias, metrics and task as before.
Registry publish without kumo-engine, or if the builder fails	Unchanged: the carrier, with a warning when the builder failed.
Registry publish of a bundle that is already a Kumo version	Returns that version instead of logging a duplicate.
Trainer.load_from_connector(sf, location=...) on either kind of version	Unchanged: the bundle is found by its marker files; serving.json sits beside it.
Existing ModelOutputConfig(...) calls	Unchanged; serving_snapshot and serving_options default to None.
MLflow pyfunc scoring	Unchanged; its helpers moved to _serving.py.

Known limitations

• Fallback disk use. The GET fallback keeps each downloaded file under the scratch directory until the filesystem is released. A snapshot load by a read-only role needs the snapshot's size again in scratch.
• Write spooling. Writes spool up to one part (4 GiB) per concurrent writer, because PUT needs a local file.
• Presigned URLs are minted one per file (one small query each). Batching them would need a hook in Let any sink hold files: FileStoreSink protocol and a chunked filesystem #112.
• Not verified live: multiclass, ranking and text-embedding models; examples/mlops/snowflake/run_batch_smoke.py itself (the same steps ran from scripts). The Model Monitor checks used synthetic actuals for the scored window.
• First job builds the image (about 5 minutes) and each version carries about 100 MB of embedded code, which Trainer.load_from_connector also downloads (about 110 s on the live account).
• One prediction time per job (prediction_time param); per-row timestamps are not part of the declared input.
• A write that raises inside with still commits what was written (Let any sink hold files: FileStoreSink protocol and a chunked filesystem #112's behaviour, as for S3). Readers are guarded by _SUCCESS, provenance-last and per-range CRCs.

Testing Done

• New packages/kumo-snowflake/tests/test_stage_fs.py runs against an in-memory stage (FakeStage in tests/fakes.py). LIST and REMOVE match by prefix and report names under the stage name, as Snowflake does. It covers:

  • upload flags;
  • holds_files (named stages only);
  • bundle publish/fetch, write_files and write_table;
  • listings with sizes;
  • exact removes that spare part-10, part-1.x and nested names, and a REMOVE that matches nothing raising;
  • chunked ranged reads, with the probe under the stage root and removed;
  • a client-side encrypted (FULL) stage refused on read, before any download; the GET fallback for a server that ignores Range: each file is downloaded once, and scratch is cleaned when the filesystem is released;
  • a read-only role gets the right warning;
  • a bad file name fails at write;
  • no signed URL in errors;
  • a registry path= is not routed to the stage.

• On Let any sink hold files: FileStoreSink protocol and a chunked filesystem #112 alone: the kumo-snowflake suite passes (191 passed, 1 skipped). flake8 and isort are clean.

• On Record prediction runs and monitor score drift #64 with this PR and Hold MLOps artifacts on Unity Catalog volumes through the Files API #114 applied: the merge is clean, and the stage tests also pass with the _provenance.json sidecar passed through write_files.

• End-to-end on the in-memory stage, on an earlier head that sat on the full series (same stage code), under kumo.mlops with a (connector, "@…") root, in both read modes:

  • a fit stored its snapshot, training table and model, each with _provenance.json;
  • the model was fetched back by path and the snapshot was loaded back;
  • a predict with a stage output_paths was read back with PredictionResult.load_from_connector.

• Part B, on Record prediction runs and monitor score drift #64 at 779f9339 with Part A:

  • New packages/kumo-engine/test/test_snowflake_model.py (5 tests): a fitted model is logged with log_model_args and round-tripped through snowflake-ml 2.4.0's own ModelPackager.save/load (no session), which exercises by-value pickling, artifact remapping and the container constructor; its predict on string ids equals Trainer.predict, the snapshot is cached, an empty batch returns the declared columns; defaults (12.6, 3.12, TABLE_FUNCTION, no repository) and every override reach log_model; a version without a snapshot names serving_snapshot=; an @ snapshot is read through from_spcs_env with the recorded target; the entry point resolves to the builder; _publish passes the serving metadata and an unknown option raises.
  • packages/kumo-snowflake/tests/test_registry.py: dedup (no model, other bundle, a non-Kumo version, a match), the builder's arguments reaching log_model with the connector's target, and the carrier fallback when no builder is installed or it raises.
  • kumo-snowflake non-live suite: 201 passed, 1 skipped. kumo-engine test_snowflake_model, test_trainer_registry, test_mlflow_model, test_mlops: 66 passed. flake8 and isort are clean.

• Live, round 1 (administrator role):

  Check	Result
  Stage file store	Write, LIST sizes, exact delete (part-1 removed; part-10, part-1.x, sub/part-1 kept), ranged reads on SSE, GET fallback on SSE, FULL refused on read; MLOps fit/predict/model load with the artifact root and export paths on a stage
  Warehouse	PUT, LIST, GET, REMOVE, GET_PRESIGNED_URL with the warehouse suspended
  Registry	Register with alias, load_from_connector(alias=...) returns the same bundle, re-registering the bundle returns the existing version (3 versions for 3 bundles)
  run_batch	Version registered through kumo.MLOps on an existing GPU_NV_S pool, GPU used; 300/300 rows, labels identical to local predict, max probability difference 4.5e-5; 690 s including the image build
  Model Monitor	Created on that version from the batch output; refreshed; ROC_AUC, COUNT and drift metrics returned

• Final wheels (bundled kumo-engine, Python 3.12, cu126 torch) on a local A10G: fit tracked, registered and loaded back with holdout_df(); the first run fell back to the carrier because requirements were read from kumo-core's own metadata, which the bundled wheel does not have — fixed in this head and re-run with the fixed module (registered runnable).

• Live, final round, as KUMO_MLOPS_VALIDATION_ROLE holding only the grants in the docs, on a dedicated GPU_NV_S pool, at this head:

  Check	Result
  Grants	Missing from the earlier docs, found from errors (must have CREATE MODEL granted on SCHEMA, ... CREATE MODEL MONITOR ..., ... CREATE TABLE ...): added. CREATE IMAGE REPOSITORY revoked: a forced image rebuild still ran (images go to the account default repository), so removed from the docs
  Stages, tracking, registry	Stage ops, fit/prediction/drift/evaluate runs, register with alias, load back, re-register returns the existing version, all as the role
  run_batch, binary	300/300 rows, labels identical to local, max probability difference 3e-6
  run_batch, regression, pip_repository=SNOWFLAKE.SNOWPARK.PYPI_SHARED_REPOSITORY	The version records the map; a forced rebuild used the repository's uv config; 300/300 rows, max difference 1.6e-3 on values up to 16
  Model Monitor, both tasks	Baseline holdout_df() of the registry-loaded model (Record prediction runs and monitor score drift #64); ROC_AUC / RMSE, COUNT and drift returned; PREDICTION_SCORE_COLUMNS=(TARGET_PRED) ACTUAL_SCORE_COLUMNS=(TARGET) for regression
  Laptop route	Fit on a local A10G (GPU available: True (cuda), used: True) with a key-pair SnowflakeIdentity in memory, no SNOWFLAKE_* variables or connections.toml; tracked, registered, loaded back

  The live runs found the PATTERN, FULL-size, delivery, Python-version and grant issues fixed above; the in-memory stage now reproduces the first two. Unit tests at this head (rebased on Record prediction runs and monitor score drift #64 with the bundled holdout): kumo-snowflake non-live suite 201 passed, 1 skipped; kumo-engine test_snowflake_model, test_trainer_registry, test_mlflow_model, test_mlops, test_snowflake_tracking, test_tracking_contract, test_holdout_df 113 passed; after the final-round changes test_snowflake_model, test_trainer_registry, test_mlflow_model 62 passed and kumo-snowflake 201 passed.

Checklist

• No breaking changes.
• Unit tests added.
• Live test on a Snowflake account.
• Least-privilege grants and pip_repository live.

🤖 Generated with Claude Code

[copy-pr-bot]

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]

Important

Draft PR not reviewed

Draft PRs are not automatically reviewed by default.

• Trigger a manual review

To automatically review draft PRs, update your CodeRabbit configuration:

reviews:
  auto_review:
    drafts: true

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

_{Comment @coderabbitai help to get the list of available commands.}

[Vedaant-J]

Change map

File	Function / scope	What changed	Why
packages/kumo-snowflake/src/kumo_snowflake/stage_fs.py	new: StageFileSystemHandler, stage_filesystem, _pattern, _RangedPart, _DownloadedPart	Chunked pyarrow filesystem over a named stage: PUT writes (one object up to 4 GiB), LIST sizes, tree REMOVE, exact-file REMOVE … PATTERN '.*/<path>' with a LIST check that raises if the file remains; presigned-URL Range reads (206 + Content-Range checked) after a one-time SSE probe under the stage root, else a once-per-file GET … pattern= fallback into the spool dir with a cause-specific warning; the probe detects a SNOWFLAKE_FULL stage from its listed (padded) size and reads there raise naming SNOWFLAKE_SSE (_read_mode, modes ranged/get/encrypted)	Stages become a file store through #112's ChunkedFileHandler, reusing _StageFiles and presigned_urls
packages/kumo-snowflake/src/kumo_snowflake/sink.py	SnowflakeSinkBackend.holds_files, filesystem, write_table, write_files, publish_model, fetch_model	FileStoreSink for @DB.SCHEMA.STAGE/... (path-only check); stage paths in the four writers/readers go through FilesystemSink, forwarding extra write options (**options); table names and registry locations keep today's code; preflight_write_files refuses an occupied stage directory before scoring unless OVERWRITE	One routing rule, shared export and bundle code with object stores and #114
packages/kumo-snowflake/README.md	stage section; registry paragraph	Stage paths in output configs and a pointer to docs/mlops.md for MLOps roots, file-name characters, the 4 GiB plain bound, a COPY INTO … PATTERN recipe, grants, SNOWFLAKE_SSE recommendation and the GET fallback; runnable registry versions via the kumo_snowflake.registry_models entry point and dedup; no warehouse needed (verified live); FULL stages refused on read	Docs
packages/kumo-snowflake/tests/fakes.py	new: FakeStage, _stage_session	In-memory named stage behind a Snowpark-shaped session: prefix-matching LIST/REMOVE with PATTERN, PUT/GET with pattern, presigned URLs and Range responses; switches for ciphertext, ignored Range and a REMOVE that matches nothing; REMOVE/GET PATTERN matched against an internal path behind the listed name (_matched), and ciphertext lists padded sizes, as found on a live account	Unit tests need no Snowflake account
packages/kumo-snowflake/tests/test_stage_fs.py	new, 15 tests	Upload flags; holds_files; bundle, file and frame exports; listings; exact and failed removes; ranged chunked reads and probe placement; GET fallback (ciphertext, ignored Range, read-only role) downloading once and cleaning up; bad names; no signed URL in errors; registry paths not routed to the stage; an occupied stage directory is refused before scoring, OVERWRITE and tables pass; a FULL stage is refused on read	Behaviour of the new file store at the sink boundary
tests/public/Dockerfile.dockerignore	kumo-core allowlist	Admits chunked_fs.py, object_store_sink.py, artifact_export.py (same lines as #114)	The public connector image must import the stage filesystem the Snowflake tests exercise
packages/kumo-snowflake/src/kumo_snowflake/artifacts.py	StageArtifactLocation.restore	Imports kumo_core.diskgraph.snapshot.rebased inside restore instead of at module level	The public connector image has no kumo_core.diskgraph; the sink and stage filesystem now import this module
packages/kumo-engine/kumo/engine/integrations/_snowflake_model.py	new: KumoModel.predict, log_model_args, embedded_code, pip_requirements, text_requirements, _pinned, signature, DEFAULTS, EMBEDDED, REQUIRING, PROVIDED, TEXT_REQUIREMENT	CustomModel whose keyword-only params mirror InputSpec(params=...) and whose first predict loads bundle and snapshot through _serving; log_model_args writes serving.json, declares ENTITY → ENTITY + prediction_output_columns(task) with params, pip requirement kumo-engine[snowflake]==<version>, SPCS target, and cuda_version/python_version/function_type/pip_repository from kumo.serving.<option> metadata over the defaults (12.6, 3.12, TABLE_FUNCTION); by default embeds kumo's packages and torchmetrics as code_paths with third-party pip requirements from kumo-core/kumo-pql metadata (==X.* as ranges, torch at the local minor, faiss-cpu, snowpark/connector); pip_repository maps those requirements to a Snowflake pip artifact repository (kumo stays embedded); kumo-glove-6b-50d added for text models; python_version defaults to the registering Python and another value raises; requirements from the distribution providing kumo_core/kumopql (workspace or bundled wheel), merged per package, faiss→faiss-cpu; snowflake-ml imported through import_provider; serving.json records the output columns	A registry version Snowflake ML batch inference can run; module-level imports kept kumo-free because snowflake-ml pickles it by value
packages/kumo-engine/kumo/engine/integrations/_serving.py	new: predict, _require_inputs, _in_input_order, snapshot_location (moved), score_registered, load_snapshot, output_columns, task_type, requirement, _cached, _entities	Shared row scoring moved from _mlflow_model.py unchanged; per-process caches; @ snapshots read with SnowflakeConnector.from_spcs_env and the recorded warehouse/database/schema; ids cast to the snapshot key type; empty batch returns the declared columns; the pip requirement helper is gone (kumo is always embedded); score_registered scores with the predictor the bundle names (load_predictor), takes output columns from the serving record, and casts ENTITY only for the Trainer predictor; new _staged reads @-stage assets with the container identity; cache _predictors	One scoring path for the MLflow pyfunc and the Snowflake model, without importing mlflow in Snowflake
packages/kumo-engine/kumo/engine/integrations/_mlflow_model.py	module imports	Imports the moved helpers from _serving	Helpers moved
packages/kumo-engine/kumo/engine/trainer/output_config.py	ModelOutputConfig.serving_snapshot, serving_options, SERVING_SNAPSHOT_KEY, SERVING_OPTION_PREFIX, SERVING_OPTIONS	Two optional fields (default None); unknown option keys raise in __post_init__; docstring: python_version must equal the registering Python, pip_repository is the index for third-party requirements	Name the serving snapshot and override CUDA, Python, function type and pip repository at registration
packages/kumo-engine/kumo/engine/trainer/trainer.py	_publish	Adds kumo.serving_snapshot and kumo.serving.<option> to the publication metadata when set	Carry the serving settings to the sink through the existing publish_model contract
packages/kumo-engine/pyproject.toml	entry points	[project.entry-points."kumo_snowflake.registry_models"] kumo-engine = "kumo.engine.integrations._snowflake_model:log_model_args"	kumo-snowflake may not import kumo-engine (package boundary test), so it discovers the builder
packages/kumo-snowflake/src/kumo_snowflake/registry.py	publish_model, new _model_args, _model_builder, _carrier_args, _serving_target, _version_holding, MODEL_BUILDER_GROUP; module docstring	Uses the entry-point builder's log_model arguments, else the carrier (with a warning if the builder raised); returns an existing Kumo version holding the same kumo_bundle_id instead of logging again, logging ... already holds bundle <id>; reused it	Runnable versions with engine installed, unchanged hosted exports, no duplicate versions (same shape as #60/#64 dedup)
packages/kumo-snowflake/tests/test_registry.py	existing publish tests; new dedup, builder-arguments and carrier-fallback tests	Existing tests patch _model_args; new tests cover four dedup cases, builder arguments plus connector target reaching log_model, and fallback with no builder or a raising one	Registry boundary
packages/kumo-engine/test/test_snowflake_model.py	new, 5 tests	Round-trip through snowflake-ml's ModelPackager and score equality with Trainer.predict; defaults and overrides; missing snapshot error; stage snapshot via from_spcs_env; entry point resolution; _publish metadata and unknown option; embedded code and requirements by default; repository path without code; another Python raises; pip_repository keeps the embedded code and requirements; text models require the embeddings package; requirements from a bundled wheel's metadata; resets _serving._predictors; calls _publish(bundle_dir, bundle_id, metrics, config)	The model and its plumbing without a Snowflake account
packages/kumo-engine/test/test_mlflow_model.py	_in_input_order test import	Imports from _serving; imports predict from _serving	Helper moved
packages/kumo-engine/docs/mlops.md	new "Batch inference in Snowflake"; Snowflake Model Monitor bullet	Registration with serving_snapshot, a run_batch example, input/output contract, environment and wheel delivery, defaults/overrides table, GPU pool, grants, snapshot reads, dedup, carrier for hosted exports, marked untested live; monitor SOURCE from batch output; verified-live banner, code embedding and Python-version notes, defaults table updated, corrected MLOps((sf, root)) and serving_snapshot=location example, SSE required; grants table verified with a least-privilege role (CREATE MODEL, CREATE MODEL MONITOR, CREATE TABLE added; CREATE IMAGE REPOSITORY removed); pip_repository row; banner and monitor note cover regression; "Model and assets" bullet: the version scores with the bundle's predictor and reads its assets from their stage	Docs
examples/mlops/snowflake/run_batch_smoke.py	new	Live smoke: fit on synthetic users/orders, export the snapshot to a stage, register with serving_snapshot, run_batch on a GPU pool, compare with local scores; serving_snapshot is the path export returned; --pip-repository passed as serving_options	Live verification once Snowflake auth works (untested)
packages/kumo-engine/kumo/engine/_deps.py	asset_locations	Context of extra asset roots that Dependencies look in first	The Snowflake model reads staged assets with an explicit connector
packages/kumo-engine/kumo/engine/_predictor.py	TrainerPredictor.predict	Imports predict from _serving	Helper moved
packages/kumo-engine/test/test_model_agnostic_mlops.py	new test_the_snowflake_model_scores_with_the_bundles_predictor	A toy bundle scores through score_registered; staged asset roots reach the resolver	Proves the Snowflake model takes any predictor
packages/kumo-engine/kumo/engine/trainer/trainer.py	_connector_registry, _publish, Trainer.load_from_connector	10-28: a stage root registers versions serving the stored training snapshot; snapshot objects resolve to their stored location; stage assets load through the loading connector	No extra snapshot export; live run: load outside a block failed
packages/kumo-engine/kumo/engine/_deps.py	connector_roots, Dependencies.candidates	10-28: stage and volume asset roots read through a connector; @ locations never read as local paths	Live run: MissingDependencyError outside a block
packages/kumo-engine/kumo/engine/integrations/snowflake.py	SnowflakeTracking	10-28: serving_options=	Registry follows the tracker
packages/kumo-engine/kumo/engine/trainer/output_config.py	ModelOutputConfig	10-28: serving_snapshot takes a GraphSnapshot or snapshot id	Snapshots by object
packages/kumo-engine/kumo/engine/integrations/_serving.py	score_registered	10-28: no-snapshot message points at the stage root	Message matches the default
packages/kumo-engine/test/test_model_agnostic_mlops.py	stage tests	10-28: staged assets through the loading connector; serving snapshot from the stage root	Reproduce the live-run failure
packages/kumo-engine/docs/mlops.md	Snowflake quickstart, Snowflake serving	10-28: stage-root quickstart; run_batch details in Reference	Short main path

Requirement: a SnowflakeConnector holds files in a named internal stage from any machine. Existing behaviour: the sink writes tables and registry models only, and the stage transport exists only inside StageArtifactLocation. Smallest change: one chunked_fs handler reusing that transport, plus holds_files/filesystem and routing of stage paths to #112's FilesystemSink. Verification: the tests above on #112 alone, and the same tests with manifest sidecars on the full series. The live account run is still pending. Stopping condition: stage paths work for every file writer and reader; tables and registry locations are unchanged. Part B requirement: a registered version runs under ModelVersion.run_batch. Smallest change: one CustomModel reusing the MLflow pyfunc's scoring, found through an entry point, with serving settings carried as publication metadata; carrier kept for hosted exports. Verification: snowflake-ml packager round-trip and registry fakes; live run_batch pending.