[Vedaant-J]

Stack: main (#53 merged) → #112 → #114 → #50 → #60 → #64 → #115. 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; a commit drops only older committed generations and reads skip incomplete ones	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 and opens, generations, overlapping writers, 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
packages/kumo-engine/test/test_adls_connector.py	ADLS connector test	An ADLS Gen2 root is a file destination and a bundle round-trips through it	ADLS artifact roots use the shared file-store path

[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.

[Victsm]

Two remaining data-integrity blockers. The chunked-filesystem implementation is unchanged from the independently reproduced version; the latest ADLS test does not cover these cases.

1. Failed constructor: start with a.bin = b"preserve-me"; make _stream() raise while constructing _Writer(..., size=1). Finalization leaves a.bin = b"", because _failed is false and close() publishes the empty spool. Keep construction failed until fallible setup succeeds, clean up on failure, and test that finalization preserves the old bytes.
2. Overlapping writers: with 16-byte parts, B writes 19 bytes; A commits the same destination and deletes B's unfinished first part; B writes 18 more bytes and commits. B's intended 37-byte file reads as 21 bytes despite _SUCCESS. Cleanup must preserve other unfinished/newer generations; readers should reject a committed part layout with gaps. Add that exact same-path regression.

These corroborate the two open bot findings: #112 (comment) and #112 (comment) . No new transactional filesystem subsystem is requested.

Approval prerequisites: fix both cases, pass the focused regressions plus the existing chunked-filesystem suite, propagate the correction through the stack, and complete required CI. The filesystem-only model/prediction adapter issues are already addressed; those are not additional open blockers.

Reviewed head: 3d3c9d896b8d3a0490f19864419457828c15dd91.