Repository navigation
Conversation
GitHub Actions
Inactive
Walkthrough
Priority: Merge Risk:
 |
| Check name | Status | Explanation | Resolution |
|---|---|---|---|
| Change Map Design Review |  |
The required conversation comment beginning <!-- kumo-pr-change-map --> is not present in the supplied context or repository checkout. The available summary describes the change map but does not pro… |
Provide the complete PR change-map comment, including its file/function scope, linked problem evidence, existing behavior, smallest sufficient change, focused verification, and stopping condition. Then reassess the implementation against th… |
-
Check name Status Explanation Pr Size The pull request changes 11 files. The limit is 40 files, so the PR passes this check. Public Resource Access The pull request adds no external URL, repository location, or package source location. The only URI-like added text is code that strips a sink scheme, not a resource reference. The @stage/...value…Full details: Public Resource Access
-
Explanation
The pull request adds no external URL, repository location, or package source location. The only URI-like added text is code that strips a sink scheme, not a resource reference. The
@stage/...values occur in an in-memory test fixture and illustrative test paths, which the check excludes. No dependency or package-source manifest changed.
-
Full details: Change Map Design Review
-
Explanation
The required conversation comment beginning
<!-- kumo-pr-change-map -->is not present in the supplied context or repository checkout. The available summary describes the change map but does not provide its required rationale, observed failure, verification, or stopping condition, so the design cannot be assessed against that comment.Resolution
Provide the complete PR change-map comment, including its file/function scope, linked problem evidence, existing behavior, smallest sufficient change, focused verification, and stopping condition. Then reassess the implementation against that source.
- Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts
Comment @coderabbitai help to get the list of available commands.
Change map
|
…cope Snapshot export took (connector, path), gated on object-store sinks, and was the only place a snapshot reached durable storage. - GraphSnapshot.export(SnapshotOutputConfig(connector, path)), in the style of the other output configs; any sink that holds files at the path (FileStoreSink, #112) can take it, and every upload tells a chunked store its size so files up to one part stay plain. - Graph.snapshot(mlops=): with MLOps on, the snapshot is exported under <root>/snapshots/ unless a complete export of the same format and id is there; failures warn. - ConcurrentExportError is a RuntimeError (no retry role). GraphSnapshot.graph_key/data_version are private. The README section moves to docs/mlops.md. Verified: test_snapshot_export.py (30: config form, a filesystem-only chunked sink round trip, stored once in a block, existing export, retry and repair cases), snapshot manifest and training-table suites. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
Note
Quiet mode is enabled, so only the most important comments were posted inline. Other review comments are grouped below.
packages/kumo-core/src/kumo_core/chunked_fs.py: 2
104-114:Minor · Stability & Availability · Quick win Fix the stale-layout race between
_layoutand_forget.
_layoutcomputes 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
_Readeruses 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 raisesOSError.The race needs a reader and a writer on the same path in one process.
test_layout_cache_survives_concurrent_opens_and_writesdoes 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
_forgetcall 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 layoutIn
_forget, increment the counter before it removes entries:self.__dict__['_epoch'] = self.__dict__.get('_epoch', 0) + 1If 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
_Writerconstruction from finalization.When
_open_part()fails after_clear(),_Writerremains partially initialized.io.RawIOBasefinalization 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
This prevents the empty-object replacement. It does not restore data already removed by
_clear().
-
-
- 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
-
packages/kumo-api/kumoapi/plugin.pypackages/kumo-core/src/kumo_core/artifacts.pypackages/kumo-core/src/kumo_core/chunked_fs.pypackages/kumo-core/src/kumo_core/object_store_sink.pypackages/kumo-core/tests/test_chunked_fs.pypackages/kumo-engine/kumo/engine/_scratch.pypackages/kumo-engine/kumo/engine/connector/__init__.pypackages/kumo-engine/kumo/engine/trainer/output_config.pypackages/kumo-engine/test/test_filestore_sink.pypackages/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.
-
| def _clear(self) -> None: | ||
| self._handler._remove(self._path + CHUNKS, tree=True) | ||
| if self._handler._size(self._path) is not None: | ||
| self._handler._remove(self._path, tree=False) |
There was a problem hiding this comment.
🔎 Supported by static analysis
-
Script executed:
set -eu printf '%s\n' '--- writer implementation ---' nl -ba packages/kumo-core/src/kumo_core/chunked_fs.py | sed -n '180,455p' printf '%s\n' '--- commit/overwrite/failure tests ---' rg -n -F --glob 'test_chunked_fs.py' -e '_SUCCESS' -e 'overwrite' -e 'crash' -e 'fail' -e 'stream' -e 'spool' packages/kumo-core/tests || test "$?" -eq 1 printf '%s\n' '--- relevant test blocks ---' nl -ba packages/kumo-core/tests/test_chunked_fs.py | sed -n '1,520p' printf '%s\n' '--- PR diff for chunked_fs.py ---' git diff 893be81a8ee4aecad9f2dd4f183bb20bb82589af d2153fe2d59e6e47a1ac81a214f3cf634ec72ded -- packages/kumo-core/src/kumo_core/chunked_fs.py | sed -n '1,520p'
Repository: NVIDIA/kumo-engine
Length of output: 43353
Do not clear the committed representation before chunk commit.
_start_stream() and _flush_part() remove the old representation before writing the replacement, while _SUCCESS is uploaded only during close(). If the stream, part upload, or later write fails, readers find neither the old file nor a committed new file. Stage each replacement separately and remove the old representation only after the replacement is committed. Apply this shared correction to sized, unsized-streaming, and spooled writes.
| if not all( | ||
| is_file_destination(self.connector, path) | ||
| for path in paths.values()): |
There was a problem hiding this comment.
🔎 Supported by static analysis
-
Script executed:
sed -n '35,65p' packages/kumo-engine/kumo/engine/trainer/output_config.py sed -n '145,175p' packages/kumo-engine/kumo/engine/trainer/trainer.py sed -n '650,675p' packages/kumo-engine/kumo/engine/trainer/trainer.py sed -n '390,430p' packages/kumo-core/src/kumo_core/object_store_sink.py
Repository: NVIDIA/kumo-engine
Length of output: 5767
Script executed:
rg -n -F --glob '*.py' -e 'FilesystemSink' -e 'sink_backend' -e 'class FileStoreSink' -e 'class FileArtifactSink' packages/kumo-api packages/kumo-core packages/kumo-engine/kumo/engine | head -220 printf '\n--- connector capability boundary ---\n' sed -n '1,120p' packages/kumo-engine/kumo/engine/connector/__init__.py printf '\n--- protocol declarations ---\n' sed -n '570,635p' packages/kumo-api/kumoapi/plugin.py printf '\n--- FilesystemSink surrounding implementation ---\n' sed -n '360,425p' packages/kumo-core/src/kumo_core/object_store_sink.py
Repository: NVIDIA/kumo-engine
Length of output: 10786
Adapt FileStoreSink before prediction export.
PredictionOutputConfig only validates the destination. It does not invoke write_files. In both prediction setup and PredictionResult._export, the code reads write_files directly from the connector's raw sink_backend. Wrap a FileStoreSink that lacks FileArtifactSink with FilesystemSink at this prediction export boundary, and use the adapter for the file writes.
d2153fe to
db52381
Compare
db52381 to
db74c6c
Compare
Snapshot exports, model bundles and export manifests could only go to object stores: the engine checked for ObjectStoreSinkBackend, and output configs accepted output_paths only for S3, GCS, ADLS and local connectors. A store that is not an object store (a Databricks volume, a Snowflake stage) had no way in, even with a pyarrow filesystem to offer. kumoapi.plugin.FileStoreSink is an optional sink capability: holds_files(connector, path) says from the path alone whether it is a file location rather than a table, and filesystem(connector, path) returns the pyarrow filesystem. ObjectStoreSinkBackend implements it. kumo.engine.connector.is_file_destination/file_store are the engine's one file-vs-table gate; PredictionOutputConfig(output_paths=) and the TrainingTableOutputConfig directory slash use it, so object stores behave as before. kumo_core gains FilesystemSink (the object-store file writers over any filesystem()), upload_bundle/download_bundle (the manifest file last), open_sized (tells a chunked filesystem the size) and PROVENANCE_FILE. The chunked filesystem streams a sized file of up to one part as one plain object, writes frames through a local spool so they stay plain too, spools under the engine scratch directory read at each write, and reads parts into the caller's buffer with one copy. Verified: test_chunked_fs.py (26), test_filestore_sink.py (protocol, output configs, bundles and Parquet over a filesystem-only sink), file connector, completion, preflight, registry and kumo-core export suites. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
db74c6c to
fc4d709
Compare
There was a problem hiding this comment.
Actionable comments posted: 1
Note
Quiet mode is enabled, so only the most important comments were posted inline. Other review comments are grouped below.
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
FilesystemSinkwhen the raw sink has nowrite_files.
TrainingTableOutputConfigaccepts a file path for aFileStoreSink.TrainingTable.exportthen readswrite_filesfrom the raw sink and raisesNotImplementedErrorif it is missing. The adapter that provideswrite_filesfor file paths is not selected on this path.Use the adapter only when the raw sink lacks
write_filesand the destination is a file location. This keeps the existing rawwrite_filespath for sinks that already provide it.🐛 Suggested fix
-
-
- 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
-
packages/kumo-api/kumoapi/json_serde.pypackages/kumo-api/kumoapi/plugin.pypackages/kumo-core/src/kumo_core/artifacts.pypackages/kumo-core/src/kumo_core/object_store_sink.pypackages/kumo-engine/kumo/engine/connector/__init__.pypackages/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.
-
| def publish_model(self, *args, **kwargs): | ||
| published = super().publish_model(*args, **kwargs) | ||
| return replace(published, uri=self._bare(published.uri)) |
There was a problem hiding this comment.
Use FilesystemSink for file-path model publication and loading.
If a sink provides holds_files and filesystem but no model methods, these adapter methods work only when callers invoke FilesystemSink manually. Trainer.fit checks the raw sink for publish_model, and Trainer.load_from_connector checks it for ModelFetchingSink. Both reject that sink before using the adapter. Select FilesystemSink for accepted file paths at both entrypoints. Keep registry-location operations on the raw sink. (raw.githubusercontent.com)
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 a sink capability for holding files (
FileStoreSink) and a chunked pyarrow filesystem for stores that can only put, list, read and remove whole objects.Summary
This PR adds a sink capability for holding files (
FileStoreSink) and a chunked pyarrow filesystem for stores that can only put, list, read and remove whole objects. With these, Databricks volumes and Snowflake stages can hold the same exports, bundles and manifests that object stores hold today. The volume and stage sinks themselves are #114 and #115, which build on this one.Problem. Teams that keep everything inside Databricks or Snowflake ask:
PredictionOutputConfig(output_paths=...)write Parquet to a/Volumes/...or@stage/...path?PUT)?Today the answer to all three is no.
Cause.
connector.source_type in {'S3', 'GCS', 'ADLS', 'LOCAL_FILE'}in the output configs. A Databricks or Snowflake connector is always a table destination.Change.
kumoapi.plugin.FileStoreSinkis an optional protocol any sink can implement:ObjectStoreSinkBackendimplements it, withholds_filesalways True.kumo.engine.connector.is_file_destination(connector, path)andfile_store(connector, path)are the engine's single file-vs-table check.PredictionOutputConfig(output_paths=)accepts paths the sink holds as files.TrainingTableOutputConfigadds the directory slash to such paths.kumo_core.object_store_sink:FilesystemSink(sink)runs the existing object-storewrite_files,write_table,publish_model(path=)andfetch_model(path=)over any sink'sfilesystem(). A volume or stage sink delegates its file paths to it and keeps its table code.upload_bundle/download_bundleare the publish/fetch bodies, now shared; the manifest file goes last.open_sized(fs, path, size)tells a chunked filesystem a file's size.kumo_core.artifacts.PROVENANCE_FILE = '_provenance.json'is the one manifest file name. The leading underscore keeps Spark/Hive Parquet readers off it.kumo_core.chunked_fs.ChunkedFileHandleris apyarrow.fs.FileSystemHandlerover five store primitives (size, list, put, ranged part read, remove), plus an optional streaming upload. A file up topart_bytesis one plain object; a larger one is stored as numbered parts committed by a trailing_SUCCESS:part_bytesas one plain object, readable by Spark, and a larger one as streamed parts, with no local disk. Frames written throughwrite_tablego via a local spool, so they stay plain too.Type check.
kumoapi.json_serde.from_jsonpassescast(Any, cls)toTypeAdapter(no behaviour change): this PR's kumo-api change runs the Kumo API ty job, and ty 0.0.77 rejects the existingUnion[Type[T], object]value as a type expression.Back-compat
Record and expose, never enforce: object-store and local connectors take the same paths as before, and nothing new raises on them.
upload_bundle/download_bundle)PredictionOutputConfig(output_paths=)on an object storeoutput_pathson a warehouse connector whose sink does not hold the pathValueError, as today (message unchanged at the start)FileStoreSink.chunksdirectory is invisible. A write that stays plain leaves the previous plain object readable until it lands; a write that becomes chunked deletes the previous object first, so a concurrent reader can getFileNotFoundErroruntil_SUCCESSis writtenKnown limitations
holds_files/filesystemon top.withbody raises still commits what was written, as with pyarrow's S3 filesystem. Readers are guarded by_SUCCESS, provenance-last and per-range CRCs.Testing Done
kumo-core/tests/test_chunked_fs.py, 26 tests:kumo-engine/test/test_filestore_sink.py:FilesystemSinkover a chunked in-memory store.Checklist
🤖 Generated with Claude Code