[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 Unity Catalog volumes as a file destination of DatabricksConnector.

Summary

This PR adds Unity Catalog volumes as a file destination of DatabricksConnector. Model bundles and prediction and training-table exports can live on a volume when the job runs off Databricks: a laptop, an EC2 box, SageMaker, CI. Once #50 is in, so can snapshot exports and the MLOps artifact root, with no change here. It needs only the FileStoreSink protocol and chunked filesystem from #112, and sits directly on it.

Problem. Teams whose data and models live in Unity Catalog ask:

• Can a training job on a GPU box outside Databricks publish its model (and, with Record provenance for every artifact under one MLOps artifact root #50, its snapshot) to the same volume the Databricks scoring job reads from?
• Can an MLOps artifact root (Record provenance for every artifact under one MLOps artifact root #50) be a volume, next to the Delta tables the run read?
• Can Spark read the exported prediction Parquet straight off the volume?

Today none of this works off-cluster. A /Volumes/... path is only reachable through the FUSE mount on Databricks compute, and the Databricks sink holds no files, so every file export to it raises.

Cause. DatabricksSinkBackend writes tables only (COPY INTO through a staging volume). It does not implement FileStoreSink, so is_file_destination(databricks, "/Volumes/...") is False, and the engine treats the path as a table name.

before                                              after
DatabricksSinkBackend                               DatabricksSinkBackend
  write_table / write_files   -> Delta table          write_table / write_files   -> Delta table, or files for /Volumes/... # new
  publish_model / fetch_model -> UC registry          publish_model / fetch_model -> UC registry, or path=/Volumes/...      # new
                                                      holds_files / filesystem    -> FileStoreSink for /Volumes/...         # new

Change.

1. kumo_databricks.volume.VolumeHandler is a chunked_fs.ChunkedFileHandler (Let any sink hold files: FileStoreSink protocol and a chunked filesystem #112) over WorkspaceClient.files, built from the connector's own identity. Calls used:

   Operation	Files API call
   write	files.upload(path, stream, overwrite=True) fed from a background thread (no local disk for sized writes)
   size	files.get_metadata (HEAD; a directory reads as not found)
   list	files.list_directory_contents, recursed
   ranged read	GET /api/2.0/fs/files<path> with Range: bytes=a- (one open-ended download serves consecutive reads); the Content-Range start is checked, so a server that ignores the range raises instead of returning wrong bytes
   delete	files.delete; a tree reuses staging.clear_volume

   The Files API caps a single upload at 5 GiB, so part_bytes is 4 GiB. A file whose size is known when it is written (bundle files, exported Parquet parts and, with Record provenance for every artifact under one MLOps artifact root #50, snapshot index files, via open_sized) is one plain object up to 4 GiB, so Spark and other tools read it directly. A larger file streams into parts. A write without a known size (write_table of a frame) stays plain up to 256 MiB, spooled under the engine scratch directory, and streams into parts above that. Errors are mapped to FileNotFoundError / PermissionError / OSError with no chained cause. Any non-SDK error, such as a transport error from a presigned upload, is named by type only, so a signed URL never appears in a message.

2. Layout on the volume, captured with Record provenance for every artifact under one MLOps artifact root #50 on top (snapshot exports and _provenance.json come from it). 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. The part size was lowered to 1 MiB for the test, so one index file splits; hashes are shortened:

   /Volumes/<catalog>/<schema>/<volume>/<root>/
     snapshots/0e3bd315…/2584a783…/
       kumo_snapshot.json
       _provenance.json
       index/_MANIFEST.json
       index/col_stats.json.chunks/00000      # 1048576 bytes   (new: only above the part size)
       index/col_stats.json.chunks/00001      #  656664 bytes
       index/col_stats.json.chunks/_SUCCESS   # 0 bytes, written last
       index/nodes/orders/shard-00000/edges.bin
       ...
     models/users-churn/                       # publish_model(path=...): bundle, _provenance.json last
     predictions/                              # PredictionOutputConfig export
       part-00000-00000.parquet
       _provenance.json
       _SUCCESS
   

   With the default 4 GiB part size, every file in that run is a plain object. Under an MLOps root (Record provenance for every artifact under one MLOps artifact root #50) the series layout applies (models/<bundle id>/, predictions/<id>/predictions/).

3. Sink routing. holds_files(connector, path) is True only for /Volumes/<catalog>/<schema>/<volume>/...; everything else is still a table name and keeps today's code. For volume 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. 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 to it unchanged. No second copy of the bundle or export logic is added.

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

   root = "/Volumes/main/kumo/artifacts"
   trainer.fit(training_table, snapshot=snapshot, export_config=ModelOutputConfig(
       databricks, path=f"{root}/models/churn"))
   trainer = Trainer.load_from_connector(databricks, path=f"{root}/models/churn")
   trainer.predict(prediction_table, snapshot=snapshot, output_config=PredictionOutputConfig(
       databricks, output_paths={"predictions": f"{root}/scores/"}))

   On Databricks compute, a plain /Volumes/... string still works as before, through FUSE.

5. Client. A WorkspaceClient is built per filesystem() call, as the existing session code does for credential rotation. The engine holds one filesystem per transfer phase (export, load, publish), so it is one client per phase, not per file.

Grants: USE CATALOG, USE SCHEMA, READ VOLUME (load, fetch), WRITE VOLUME (export, publish). No SQL warehouse is used.

Destination check before scoring. The sink's preflight_write_files, which Trainer.predict calls before any scorer starts, refuses an occupied volume 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.

Back-compat

Record and expose: nothing that works today changes behaviour.

Situation	Behaviour
Table destination on a DatabricksConnector	Unchanged: COPY INTO through the staging volume.
UC registry publish/fetch (location=)	Unchanged.
/Volumes/... path on a DatabricksConnector	Was refused (not a file destination); now written through the Files API.
/Volumes/... string root on Databricks compute	Unchanged: a FileConnector over FUSE.
Volume export read by Spark	Files up to 4 GiB written with a known size are plain objects; larger ones are .chunks/ parts.
Volume export read by an older engine	Plain files read as before; a .chunks file is unknown to it.
Upload or download failure	OSError subclasses with no chained vendor text; snapshot transfers retry as for object stores.
Path that isn't a volume (dbfs:/..., ../, fewer than 3 components)	Treated as a table name (no file I/O).

Known limitations

• A write without a known size (write_table of an in-memory frame) above 256 MiB becomes .chunks/ parts, which Spark can't read directly.
• A write that raises inside with still commits what was written. This matches pyarrow's S3 behaviour; the series guards readers with _SUCCESS, provenance-last and per-range CRCs.

Testing Done

• New packages/kumo-databricks/tests/test_volume.py (fake Files API in tests/fakes.py):

  • reads across part boundaries;
  • sized writes at the part size stay plain and larger ones become parts;
  • a failed streamed upload raises OSError with no signed URL and no cause, and does not hang;
  • a download that ignores Range is refused;
  • directories, missing paths and tree deletes;
  • a refused upload raises PermissionError;
  • holds_files / filesystem accept only volume paths;
  • bundle publish and fetch, and Parquet and frame writes, through a volume path.

• tests/public/Dockerfile.dockerignore admits chunked_fs.py, object_store_sink.py and artifact_export.py, so the public connector image can import the volume sink.

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

• On Record prediction runs and monitor score drift #64 with this PR and Hold MLOps artifacts in Snowflake stages and score registry versions with run_batch #115 applied: the merge is clean, and the volume tests also pass with the _provenance.json sidecar passed through write_files.

• Live, on an earlier head of this PR that sat on the full series (same volume code), on a test workspace with synthetic users/items/orders data, from an EC2 box (off-cluster) and in a one-time cluster job (single-node m5d.xlarge, DBR 16.4), with wheels built from this head:

  Check	Off-cluster	In-cluster
  3.0 GiB file written with its size → one plain object, full CRC match, mid-file range exact	write 12 s, read 50 s	write 20 s, read 39 s
  4.29 GiB file written with its size → 2 parts (4 GiB + 300 MiB) + _SUCCESS, full CRC match	write 17 s, read 57 s	write 30 s, read 56 s
  6000-byte read across the 4 GiB part boundary	exact	exact
  snapshot export / load via SnapshotOutputConfig (part size lowered to 1 MiB for the test so col_stats.json splits; read across its boundary exact)	11.0 s / 2.9 s	9.0 s / 2.7 s
  bundle publish + fetch by path	ok	ok
  prediction export with _provenance.json recording the volume path	ok	ok
  kumo.mlops(kumo.MLOps(artifact_root=(databricks, "/Volumes/..."))) fit + predict: snapshots/, training_tables/, models/, predictions/ under the root, 4 _provenance.json	ok	ok

  Read throughput on 2 GiB of one part (MB/s), handler vs raw Files API:

  	Off-cluster	In-cluster
  raw, one streamed GET	72	82
  raw, 16 parallel ranged GETs	663	410
  handler, sequential 64 MiB reads	88	89
  handler, snapshot-load pattern (16 threads × 8 MiB reads)	541	355

Checklist

• No breaking changes.
• Unit tests added.

🤖 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

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-dbx-volume
• File: .coderabbit.yaml
• Commit: e2d71ad
• Layers: Repository YAML
• Review profile: QUIET
• Plan: Enterprise
• Run ID: 3465e1ea-7c1f-426e-9411-2685bc41db47

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

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

[Vedaant-J]

Change map

File	Function / scope	What changed	Why
packages/kumo-databricks/src/kumo_databricks/volume.py	VolumeHandler, is_volume_path, volume_filesystem, _Part, _Upload, _Feed, _translated	New chunked_fs handler over WorkspaceClient.files: size (HEAD), recursive list, delete (tree via staging.clear_volume), 4 GiB parts, 256 MiB plain limit for unsized writes; one open-ended Range download per part serves consecutive reads; checks the Content-Range start; streams writes into files.upload on a helper thread; failures surface at write/close; SDK errors map to FileNotFoundError/PermissionError/OSError; other errors named by type only; never chained	Volumes as a file store off Databricks compute; ranged reads at raw-API speed; refuse a server that ignores the range; sized and large writes need no local disk; signed URLs never appear in errors
packages/kumo-databricks/src/kumo_databricks/sink.py	holds_files, filesystem, write_table, write_files, publish_model, fetch_model	FileStoreSink for /Volumes/<c>/<s>/<v>/...; other paths are tables; volume paths go through FilesystemSink, forwarding extra write options (**options) to it; tables and registry locations unchanged; preflight_write_files refuses an occupied volume directory before scoring unless OVERWRITE	Engine file-vs-table gate; reuse #112's writers instead of a second copy; later write options (manifest sidecars) reach volumes without a change here
packages/kumo-databricks/src/kumo_databricks/registry.py	_address	Error for a non-volume path names the /Volumes/... form	Points users to the new option
packages/kumo-databricks/README.md	Artifacts on a volume	Usage through output configs and Trainer.load_from_connector, a pointer to docs/mlops.md for MLOps roots, layout (plain up to 4 GiB, .chunks/ above), grants	Docs
packages/kumo-databricks/tests/fakes.py	FakeVolume, FakeMetadata, FakeStreamingResponse, FakeWorkspaceClient.api_client	In-memory volume with directories, uploads and Range downloads	Unit tests without a workspace
packages/kumo-databricks/tests/test_volume.py	9 tests	Cross-part reads; sized plain vs parts; failed upload hides the URL; range refusal; directories/missing/delete; refused upload; path gate; bundle, Parquet and frame writes through a volume path; needs only #112; an occupied volume directory is refused before scoring, OVERWRITE and tables pass	Behaviour tests
tests/public/Dockerfile.dockerignore	connector source closure	Admits kumo_core/chunked_fs.py, object_store_sink.py, artifact_export.py	The public connector image runs the Databricks tests, which now import them

Requirement: a DatabricksConnector holds files on a Unity Catalog volume from any machine. Existing behaviour: the sink writes tables only, and /Volumes/... works only through FUSE on Databricks compute. Smallest change: one chunked_fs handler over the Files API, plus holds_files/filesystem and routing of volume paths to #112's FilesystemSink. Verification: the tests above on #112 alone, the same tests with manifest sidecars on the full series, and the live runs in the body. Stopping condition: volume paths work for every file writer and reader; tables and registry locations are unchanged.