[haoweiyu-kumo]

Stacked on #144 (refactor/native-batch-runtime).

Serve complete DiskGraph batches through the Rust Flight gateway, embedded in the Python/Ray actor process. diskgraph.SamplingServer exposes endpoint discovery, native health, close, and context management through PyO3. Rust owns task uploads, sampling, Arrow encoding, credits and cancellation. Python starts/stops the handle and consumes batches through PyArrow Flight. Start and close release the GIL; the standalone CLI uses the same Rust implementation.

Include trainer endpoint/thread settings, replica tail accounting, feature projections and ordered RHS scans. Native lifecycle/progress logs reach Python logging. The unused gateway graph/feature ticket API is replaced; graph publication and cluster provisioning remain separate PRs.

Maintained change map

Validation: 45 Rust tests, 491 focused Python tests and 14 notice tests passed. Clippy, rustfmt, scoped mypy, Ruff, CI flake8/isort, generated-notice/license checks and documentation checks passed. Rebuilt the extension and checked source archive completeness/compilation. A full release wheel and cloud/Ray deployment were not run locally.

[haoweiyu-kumo]

Change map

Requirement: implement #146 serving in Rust and embed the native server inside a Python/Ray actor through PyO3. Previously Python dispatched Flight requests and encoded sampled batches. The shared Rust gateway now owns task registration, complete epochs, Arrow encoding, credits and cancellation. Python owns a lifecycle handle and a Flight client. The unused per-request gateway API is replaced. Verify batch parity, bounded credits, task integrity, GIL-safe cleanup, same-process ownership, logs and source packaging. Stop at prepared graph replicas; graph publication and cluster provisioning remain separate changes.

File	Function / scope	What changed	Why
3rdparty/notices/inventory.json	Extension artifact membership	Refresh the Rust components shipped by the extension, including native Flight serving.	Cover the embedded gateway dependency closure.
THIRD-PARTY.txt	Generated notices	Regenerate third-party notices for the native extension.	Distribute licenses for embedded serving dependencies.
packages/diskgraph/Cargo.lock	Native dependencies	Resolve the extension's gateway dependency and remove obsolete gateway serialization libraries.	Match the shared embedded and standalone implementation.
packages/diskgraph/README.md	Native runtime guide	Describe embedded Flight serving and its standalone CLI.	Expose both entry points accurately.
packages/diskgraph/diskgraph-gateway/Cargo.toml	Gateway dependencies	Use the shared runtime and native Arrow/Flight/task storage dependencies.	Serve complete task epochs in Rust.
packages/diskgraph/diskgraph-gateway/README.md	Gateway contract	Document complete task epochs, embedded lifecycle, native logging, and standalone use.	Make process ownership and the replacement protocol explicit.
packages/diskgraph/diskgraph-gateway/src/encoding.rs	BatchEncoder	Encode native task, graph, and feature payloads into bounded Arrow IPC components.	Avoid Python objects and Python serialization on sampling workers.
packages/diskgraph/diskgraph-gateway/src/handler.rs	Deleted legacy request handler	Remove per-request graph/feature dispatch.	Replace the unused gateway API with complete task epochs.
packages/diskgraph/diskgraph-gateway/src/lib.rs	SamplerGateway	Serve DoPut, DoExchange, task status and health; export the owned server lifecycle.	Keep task transfer, sampling, Arrow encoding and flow control native.
packages/diskgraph/diskgraph-gateway/src/main.rs	Standalone entry point	Start the same owned server used by PyO3 and close it on SIGTERM/Ctrl-C.	Share serving behavior without Python or launcher-specific flags.
packages/diskgraph/diskgraph-gateway/src/protocol.rs	Versioned envelopes	Validate typed frame shape, metadata, binary components, and byte budgets.	Share the existing Flight client wire contract safely.
packages/diskgraph/diskgraph-gateway/src/registry.rs	Deleted request registry	Remove the old per-request ID registry.	The shared runtime owns complete task iteration.
packages/diskgraph/diskgraph-gateway/src/runtime.rs	Native task options	Deserialize normalized options, rebase task sidecars, and open shared runtime iterators.	Preserve local sampling semantics and error propagation.
packages/diskgraph/diskgraph-gateway/src/sender.rs	Deleted legacy sender	Remove graph-only serialization and transmission.	Encode complete RuntimeBatch values through the task protocol.
packages/diskgraph/diskgraph-gateway/src/server.rs	SamplingServer, ServerConfig	Bind and own the native runtime/listener, report health, cancel streams and close safely.	Embed serving in an actor process; avoid joining threads during Python garbage collection.
packages/diskgraph/diskgraph-gateway/src/storage.rs	Immutable task upload	Validate paths, file order, lengths and checksums before atomic publication.	Reject partial or corrupted training-table extensions.
packages/diskgraph/diskgraph-gateway/tests/integration_gateway.rs	Native listener lifecycle	Exercise action discovery, process health, port ownership, repeated close and invalid budgets.	Validate shared lifecycle independently of Python.
packages/diskgraph/diskgraph-pybind/Cargo.toml	Extension dependencies	Link the gateway library into the Python extension.	Expose native serving without a worker subprocess.
packages/diskgraph/diskgraph-pybind/README.md	Embedded API guide	Document SamplingServer, context management, endpoint discovery and logs.	Show the lifecycle used by Python/Ray workers.
packages/diskgraph/diskgraph-pybind/python/diskgraph/__init__.py	Public exports	Export SamplingServer from the native extension.	Provide a direct Python lifecycle handle.
packages/diskgraph/diskgraph-pybind/python/tests/test_sampling_server.py	Public native handle tests	Check same-process serving, native health, startup failures, endpoint formatting, context exit and garbage collection.	Prevent leaked listeners and GIL-sensitive cleanup regressions.
packages/diskgraph/diskgraph-pybind/src/lib.rs	Extension registration and logger bridge	Register SamplingServer and retain the existing native logger reset handle.	Expose the class and refresh log levels when a server starts.
packages/diskgraph/diskgraph-pybind/src/sampling_server.rs	PyO3 lifecycle bindings	Expose start, endpoint, health, port, close and context management; release the GIL during start/close.	Keep sampled requests and responses entirely inside Rust on the server.
packages/diskgraph/third_party/artifact-manifest.json	Generated artifact graph	Resolve the embedded gateway as part of diskgraph._c.	Track the actual shipped extension dependency graph.
packages/kumo-core/README.md	Transport setup	Document the embedded server, actor lifecycle, native logging and client stream semantics.	Make the process and data flow explicit.
packages/kumo-core/pyproject.toml	Runtime dependencies	Declare orjson.	Encode bounded protocol metadata and task identities.
packages/kumo-core/src/kumo_core/diskgraph/lightning_data.py	DiskgraphTaskBatchIterable, DiskgraphLightningDataModule	Select remote indexes, close epoch streams and shared clients, count replica batch tails.	Integrate push sampling without dropping partial shards or leaking producers.
packages/kumo-core/src/kumo_core/diskgraph/remote_sampling.py	RemoteTaskIterator, RemotePushIterator, RemoteRuntimeIndex	Register tasks and consume versioned credit-controlled replica streams with bounded prefetch.	Offload complete batches while preserving native manifest/option resolution and propagating incomplete epochs.
packages/kumo-core/src/kumo_core/diskgraph/sampling_protocol.py	Client envelope codec	Encode control/upload frames and decode native Arrow components.	Preserve typed batches while removing Python batch encoding.
packages/kumo-core/src/kumo_core/diskgraph/sampling_service.py	SamplingService	Own the embedded PyO3 server handle and expose its endpoint, health and close.	Retain the worker integration surface without subprocess supervision.
packages/kumo-core/src/kumo_core/diskgraph/sampling_task.py	task_files, upload_task	Hash and upload immutable task extensions in bounded chunks with acknowledgement.	Reuse task registration across epochs and reject incomplete/corrupt transfers.
packages/kumo-core/src/kumo_core/train/diskgraph_nested_trainer.py	Sampler option forwarding and init	Forward remote settings and select inline consumption instead of local index auto-configuration.	Keep sampling work on the configured workers.
packages/kumo-core/src/kumo_core/training/config.py	SamplerKwargs	Add optional endpoints and a positive sampling thread budget.	Validate and carry per-run remote sampling settings.
packages/kumo-core/src/kumo_core/training/runtime.py	LocalDiskgraphTrainingRuntime, local bindings	Carry optional sampling configuration through the local training runtime.	Make fresh and reloaded internal trainers use the selected backend.
packages/kumo-engine/README.md	Trainer configuration	Document remote endpoints and link worker setup.	Provide the current public entry point.
packages/kumo-engine/kumo/engine/trainer/_runtime.py	local_training_bindings	Forward optional sampling configuration alongside accelerator settings.	Connect engine selection to core execution.
packages/kumo-engine/kumo/engine/trainer/trainer.py	Trainer constructor and fit	Add optional endpoints/thread settings and forward them during fit.	Enable remote sampling while training remains on the calling node.
packages/kumo-engine/test/test_training_explanations.py	Fit/explanation orchestration	Exercise accelerator forwarding with both local and remote sampling settings.	Guard the engine-to-core runtime boundary.
test/train/datamodule/test_persistent_sampling.py	Module and epoch lifetime tests	Verify stage sharing, teardown, early exit, and frame bounds.	Prevent leaked streams and cross-stage cancellation.
test/train/datamodule/test_push_sampling.py	Epoch and transfer integration	Verify parity, credits, projections, task integrity, replica tails, ordered RHS scans, cancellation and embedded shutdown.	Preserve complete batch semantics and bounded production inside an actor.
test/train/datamodule/test_remote_sampling.py	Client framing and option tests	Cover Arrow types, incomplete epochs, RPC failures, option normalization, and cleanup.	Reject corrupted/truncated delivery and preserve caller configuration.
test/train/datamodule/test_remote_sampling_integration.py	Synthetic graph/task fixtures and parity	Serve an independent graph replica and compare complete batches with the local iterator.	Exercise the real Flight/native boundary without cloud services.
test/train/datamodule/test_sampling_service.py	Embedded service integration	Verify native batch parity, same-process health, replicas, validation and Python logging.	Exercise the native listener through the worker-facing wrapper.
test/train/test_diskgraph_nested_trainer.py	Remote init regressions	Verify remote settings skip local auto-configuration and use an inline loader.	Keep trainer initialization compatible with offloaded sampling.
uv.lock	Generated dependency metadata	Record the core orjson dependency.	Keep resolution consistent with package declarations.

Validation

45 Rust tests (gateway and shared runtime), 491 focused Python tests (embedded server, transport, loaders, nested trainer, engine bindings), and 14 notice tests passed. Clippy with warnings denied, rustfmt, scoped mypy, Ruff, CI flake8/isort, generated artifact/notice checks, license checks, and documentation checks passed. The native extension was rebuilt; a source archive includes the embedded server and compiles. A full release wheel and cloud/Ray deployment were not run locally.

[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: feat/flight-push-sampling
• File: .coderabbit.yaml
• Commit: 13fbc77
• Layers: Repository YAML
• Review profile: QUIET
• Plan: Enterprise
• Run ID: 59dfd94b-6eb1-4c1f-baf6-e83bbbd58278

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