[Vedaant-J]

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

Summary

Prediction exports and model bundles can only go to S3, GCS, ADLS or local paths: the engine checks the connector type. Databricks volumes and Snowflake stages hold files too, but one upload is capped (5 GiB on volumes) or has no multipart. This PR lets any sink say which of its paths are files, and gives such stores one shared filesystem:

• kumoapi.plugin.FileStoreSink: an optional protocol, holds_files(connector, path) (from the path alone) and filesystem(connector, path) (a pyarrow.fs.FileSystem).
• kumo_core.chunked_fs.ChunkedFileHandler: a pyarrow filesystem over five store calls (size, list, put, ranged read, remove). A file up to part_bytes is one plain object; a larger one is stored as parts under <name>.chunks/<generation>/ and committed by an empty _SUCCESS written last.
• kumo_core.object_store_sink.FilesystemSink: the existing object-store writers over any sink's filesystem(). Bundles upload _provenance.json last.
• kumo_core.object_store_sink.FileStoreRouting: a mixin that sends a sink's file paths to FilesystemSink and its other paths to its own table methods.
• ObjectStoreSinkBackend.preflight_write_files: an export into an occupied output directory is refused before scoring, on every object store.

No volume or stage sink ships here; #114 and #115 add them.

Usage

class StageSink:                                    # an existing table sink
    def holds_files(self, connector, path):
        return path.startswith('@')
    def filesystem(self, connector, path):
        return PyFileSystem(StageHandler(connector)), path.rstrip('/')

PredictionOutputConfig(connector, output_paths={"predictions": "@stage/kumo/scores/"})
trainer.fit(..., export_config=ModelOutputConfig(connector, path="@stage/kumo/models/m1"))
Trainer.load_from_connector(connector, path="@stage/kumo/models/m1")

Back-compat

Situation	Behaviour
S3, GCS, ADLS or local publish and fetch	Unchanged (ObjectStoreSinkBackend holds every path).
output_paths on a sink that does not hold the path	ValueError, as before; the message now mentions volume and stage paths.
Sink whose package extra is not installed	The install hint is raised, not the file-destination error.
Export into an output directory that already holds files (S3, GCS, ADLS or a sink's file store)	Refused before scoring, whatever the write mode, with the same Output directory ... already contains data from another export error that was raised after scoring.
A sink without FileStoreSink	A table destination, as before.
A chunked file read while it is rewritten	The previous file until the new _SUCCESS lands.
A write that fails part way	Nothing committed; a stream that supports abort is aborted; the previous file stays readable; the error is raised.
A .chunks/ file read by Spark or an older engine	A directory of parts; files up to part_bytes are plain objects.

Testing

• kumo-core/tests/test_chunked_fs.py: empty, sub-part, one-part and multi-part files; reads across part boundaries; newest committed generation wins; failed writes keep the old file and abort the stream; listings, deletes, streamed writes, spool location; concurrent rewrites.
• kumo-engine/test/test_filestore_sink.py: output configs, a bundle round trip, trainer publish and load, and prediction export, the occupied-directory refusal before scoring, and a file-only sink refusing a table through a filesystem-only sink over an in-memory store.

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

🤖 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]

Superseded by the change map below.

[Vedaant-J]

Change map

File	Function / scope	What changed	Why
packages/kumo-api/kumoapi/plugin.py	FileStoreSink	New optional protocol: holds_files, filesystem	Lets any sink hold files for some paths
packages/kumo-core/src/kumo_core/artifacts.py	PROVENANCE_FILE	New constant _provenance.json	One name for every artifact's manifest
packages/kumo-core/src/kumo_core/chunked_fs.py	ChunkedFileHandler, set_spool_dir, spool_root	New pyarrow filesystem over put/list/read/remove; large files as committed parts	Stores that cap object size or lack multipart
packages/kumo-core/src/kumo_core/object_store_sink.py	holds_files, filesystem, preflight_write_files, FilesystemSink, FileStoreRouting, open_sized, upload_file, upload_bundle	Object-store writers over any sink's filesystem(); occupied output directory refused before scoring on every store; manifest uploaded last	Reuse export and bundle code for volumes and stages
packages/kumo-engine/kumo/engine/connector/__init__.py	is_file_destination, file_store	Pick file store or table per path	Replaces the connector-type check
packages/kumo-engine/kumo/engine/trainer/output_config.py	PredictionOutputConfig, TrainingTableOutputConfig	Accept any path a sink holds; new writer()	Exports to volume and stage paths
packages/kumo-engine/kumo/engine/trainer/trainer.py	_require_publisher, _preflight, fit, predict, PredictionResult._export	Publisher check, preflight and export writer come from the sink protocols	Publish and export by file path on any sink
packages/kumo-engine/kumo/engine/_scratch.py	_spool_parts	Chunked writers spool under <scratch>/spool	Follows kumo.configure(scratch_dir=)
packages/kumo-core/tests/test_chunked_fs.py	new tests	Layouts, boundary reads, failed writes, generations, concurrency	Covers the chunked filesystem contract
packages/kumo-engine/test/test_filestore_sink.py	new tests	Configs, bundles, trainer, export over a filesystem-only sink	End to end without a real store
packages/kumo-engine/test/test_prediction_completion.py	fake sink	Fake file sink implements FileStoreSink	Export path now asks the sink
tests/public/Dockerfile.dockerignore	public image allowlist	Ships chunked_fs.py	object_store_sink imports it
packages/kumo-adls/tests/test_native_output.py	ADLS sink tests	Copy through the shared upload_file	Follows the upload helper's move

[coderabbitai]

Actionable comments posted: 2

Note

Quiet mode is enabled, so only the most important comments were posted inline. Other review comments are grouped below.

Other comments: 2

packages/kumo-core/src/kumo_core/chunked_fs.py: 2

104-114: Minor · Stability & Availability · Quick win

Fix the stale-layout race between _layout and _forget.

_layout computes a layout without a lock and then stores it in the cache. A writer can finish and call _forget(path) between the compute step and the store step. The reader then stores the old layout after the invalidation.

Nothing removes that entry afterward. The _Reader uses the cached sizes. If the new file is larger, a read returns truncated data without an error. If the sizes differ in another way, the read raises OSError.

The race needs a reader and a writer on the same path in one process. test_layout_cache_survives_concurrent_opens_and_writes does not cover it. In that test, each thread reads and writes only its own file.

Use an invalidation counter. Store the layout only if no _forget call ran during the compute step.

Proposed fix

     def _layout(self, path: str) -> Optional[List[Tuple[str, int]]]:
         cache = self.__dict__.setdefault('_layouts', {})
         layout = cache.get(path)
         if layout is not None:
             return layout
+        epoch = self.__dict__.get('_epoch', 0)
         size = self._size(path)
         layout = ([(path, size)] if size is not None else _logical_parts(
             self._list(path + CHUNKS), path + CHUNKS))
-        if layout is not None:
+        if layout is not None and self.__dict__.get('_epoch', 0) == epoch:
             cache[path] = layout
         return layout

In _forget, increment the counter before it removes entries:

        self.__dict__['_epoch'] = self.__dict__.get('_epoch', 0) + 1

If you need exact atomicity across threads, guard the counter update and the cache store with one threading.Lock.

Add a test in which one thread writes a path while other threads read the same path.

---

321-341: Minor · Data Integrity & Integration · Quick win

Guard failed _Writer construction from finalization.

When _open_part() fails after _clear(), _Writer remains partially initialized. io.RawIOBase finalization calls _Writer.close(), whose final branch uploads the empty spool as the plain object. _clear() has already removed the old chunked representation.

super().close() directly calls the inherited close method and does not re-enter _Writer.close(). Keep the resource setup inside the guard:

Suggested fix

-        self._dir = tempfile.mkdtemp(prefix='kumo-chunked-',
-                                     dir=spool_root(handler))
-        self._spool = open(os.path.join(self._dir, 'spool'), 'wb')
-        self._filled = 0
-        self._parts = 0
-        self._written = 0
-        self._part: Optional[BinaryIO] = None
-        self._streams = type(handler)._stream is not ChunkedFileHandler._stream
-        self._limit = (min(handler.plain_bytes, handler.part_bytes)
-                       if self._streams else handler.part_bytes)
-        self._plain: Optional[BinaryIO] = None
-        if self._streams and size is not None:
-            if size <= handler.part_bytes:
-                self._plain = handler._stream(path)
-            else:
-                self._clear()
-                self._part = self._open_part()
+        directory: Optional[str] = None
+        spool: Optional[BinaryIO] = None
+        try:
+            directory = tempfile.mkdtemp(prefix='kumo-chunked-',
+                                         dir=spool_root(handler))
+            self._dir = directory
+            spool = open(os.path.join(directory, 'spool'), 'wb')
+            self._spool = spool
+            self._filled = 0
+            self._parts = 0
+            self._written = 0
+            self._part: Optional[BinaryIO] = None
+            self._streams = type(handler)._stream is not ChunkedFileHandler._stream
+            self._limit = (min(handler.plain_bytes, handler.part_bytes)
+                           if self._streams else handler.part_bytes)
+            self._plain: Optional[BinaryIO] = None
+            if self._streams and size is not None:
+                if size <= handler.part_bytes:
+                    self._plain = handler._stream(path)
+                else:
+                    self._clear()
+                    self._part = self._open_part()
+        except BaseException:
+            if spool is not None:
+                spool.close()
+            if directory is not None:
+                shutil.rmtree(directory, ignore_errors=True)
+            super().close()
+            raise

This prevents the empty-object replacement. It does not restore data already removed by _clear().

---

Review info

Run configuration

• Source mode: Auto
• Repository: NVIDIA/kumo-engine
• Branch: vedaant/mlops-holdout-fit-mode
• File: .coderabbit.yaml
• Commit: 893be81
• Layers: Repository YAML
• Review profile: QUIET
• Plan: Enterprise
• Run ID: aca98211-94a5-45c1-8517-a0cd53672083

Commits

Reviewing files that changed from the base of the PR and between 893be81 and d2153fe.

Files selected for processing: 10

• packages/kumo-api/kumoapi/plugin.py
• packages/kumo-core/src/kumo_core/artifacts.py
• packages/kumo-core/src/kumo_core/chunked_fs.py
• packages/kumo-core/src/kumo_core/object_store_sink.py
• packages/kumo-core/tests/test_chunked_fs.py
• packages/kumo-engine/kumo/engine/_scratch.py
• packages/kumo-engine/kumo/engine/connector/__init__.py
• packages/kumo-engine/kumo/engine/trainer/output_config.py
• packages/kumo-engine/test/test_filestore_sink.py
• packages/kumo-engine/test/test_prediction_completion.py

Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 6 remain after this review.

[coderabbitai]

Actionable comments posted: 1

Note

Quiet mode is enabled, so only the most important comments were posted inline. Other review comments are grouped below.

Other comments: 1

packages/kumo-core/src/kumo_core/object_store_sink.py: 1

414-415: Minor · Functional Correctness · Quick win

Route file-path training-table exports through FilesystemSink when the raw sink has no write_files.

TrainingTableOutputConfig accepts a file path for a FileStoreSink. TrainingTable.export then reads write_files from the raw sink and raises NotImplementedError if it is missing. The adapter that provides write_files for file paths is not selected on this path.

Use the adapter only when the raw sink lacks write_files and the destination is a file location. This keeps the existing raw write_files path for sinks that already provide it.

🐛 Suggested fix

--- a/packages/kumo-engine/kumo/engine/pquery/training_table.py
+++ b/packages/kumo-engine/kumo/engine/pquery/training_table.py
@@ -4,6 +4,8 @@
 from pathlib import Path
 from typing import TYPE_CHECKING, Optional
 
+from kumo_core.object_store_sink import FilesystemSink
 from kumo_core.model_plan import TrainingTableGenerationPlan
 from kumoapi.model_plan import \
     TrainingTableGenerationPlan as PublicGenerationPlan
 from kumoapi.train import TrainingTableSpec
 
+from kumo.engine.connector import is_file_destination
 from kumo.engine.connector.source_table import SourceTable
@@ -54,7 +56,11 @@
         connector = output_config.connector
-        write_files = getattr(connector.sink_backend, "write_files", None)
+        sink = connector.sink_backend
+        write_files = getattr(sink, "write_files", None)
+        if not callable(write_files) and is_file_destination(
+                connector, output_config.table_name):
+            write_files = FilesystemSink(sink).write_files
         if not callable(write_files):

---

Review info

Run configuration

• Source mode: Auto
• Repository: NVIDIA/kumo-engine
• Branch: vedaant/mlops-filestore-sink
• File: .coderabbit.yaml
• Commit: fc4d709
• Layers: Repository YAML
• Review profile: QUIET
• Plan: Enterprise
• Run ID: dc4767da-d88b-4cf6-9095-139246272709

Commits

Reviewing files that changed from the base of the PR and between db74c6c and fc4d709.

Files selected for processing: 6

• packages/kumo-api/kumoapi/json_serde.py
• packages/kumo-api/kumoapi/plugin.py
• packages/kumo-core/src/kumo_core/artifacts.py
• packages/kumo-core/src/kumo_core/object_store_sink.py
• packages/kumo-engine/kumo/engine/connector/__init__.py
• packages/kumo-engine/kumo/engine/trainer/output_config.py

Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 11 remain after this review.