You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
{{ message }}
Repository navigation
Serve remote sampling with an embedded Rust Flight gateway - #146
Builds on the native batch runtime from merged #144. Targets main.
Serve complete DiskGraph batches through an embedded Rust Flight server and native client. The shared public diskgraph.RuntimeIndex resolves manifests and normalizes options for both modes: omit sampling_endpoints for local sampling or supply replica URIs for the Rust remote backend. Workers serve prepared graph/training-table replicas; the trainer only needs a portable manifest. Artifact distribution, cluster provisioning and public Trainer integration remain separate changes. Trainer has no remote sampling constructor arguments and passes sampling=None to local_training_bindings.
Rust owns lazy replica connections across stages/extensions, epoch sharding, weighted-sidecar path translation, one-batch prefetch per replica, deterministic round-robin merging and cancellation. Ordered RHS catalog scans use the first replica. Remove the Python remote iterator module and thread pool. Epoch close joins native readers and discards queued batches while preserving connections; index close cancels all epochs and is terminal. GC signals cancellation without joining under the GIL. A process-local RuntimeOwner lazily creates one local or remote index and is passed explicitly to all module datasets. Epoch cleanup closes only epoch iterators; module teardown closes the shared owner. Standalone datasets close their own owner explicitly. Pickle excludes native state; a PID change opens a process-local handle. This removes remote-specific binding, managed flags and private cache injection. Existing local constructor/iterator calls retain their behavior. Both local and remote iterators always expose idempotent close(); local close is a synchronous no-op and loader cleanup registers it directly without getattr/callable checks.
Local sampling and the server consume RuntimeBatch into a shared diskgraph_runtime::ArrowBatch. The Flight client returns that same representation, and one Python exporter restores the complete minibatch. Remove the separate local dictionary builder: numeric vectors move into Arrow arrays and existing target/feature buffers stay shared. Local conversion performs no IPC/JSON serialization and has no transport byte limit.
Python builds immutable SamplingRequest objects; Rust owns ticket serialization, request kind and protocol version. Version 9 groups aligned seed/target, node/assignment, edge-endpoint and positive-label columns. Column references return arrays; half-open column ranges return record batches without target struct wrappers or array/as markers. Each DoGet returns an epoch shard: initial schema application metadata describes counts; a metadata-only message starts each minibatch with its layout/count; component schemas, dictionaries and Arrow data follow. Validate final RPC status, EOF/counts, references and the cumulative 64 MiB minibatch budget. Export through the Arrow C Data Interface; no nested IPC streams or Python wire decoding.
Server production reserves a one-slot output queue before sampling; gRPC propagates backpressure. Transport buffers/native working memory are additional. Listener shutdown retains its five-second drain grace. Native lifecycle/progress events reach Python logging.
Requirement: serve complete native sampling epochs through DoGet over caller-prepared graph/training-table replicas, preserving local/remote batch parity. The earlier Python path duplicated remote orchestration; subsequent review found completed replay iterators retained until epoch end, a single-stream default incompatible with interleaved refit, and positive training lengths with zero batches after worker/replica sharding and drop_last.
The implementation keeps request handling, Arrow component construction, decoding and bounded replica prefetch in Rust, with one shared Python exporter and explicit process-local runtime ownership. Follow-up fixes reuse per-stage cleanup stacks, allow three concurrent streams by default, force remote DataLoaders to num_workers=0 without an override, and reject flat remote training epochs with no native batches before applying the minimum epoch target. Local DataLoader worker settings remain supported.
Verification covers parity, grouped and multi-column targets, shared buffers, malformed streams, byte limits, final RPC errors, backpressure, replay retention, paused refit streams, loader worker policy and shutdown. Stopping condition: keep this PR within native transport and its internal training integration, with focused regressions for those behaviors. Artifact distribution, cluster provisioning and public Trainer integration remain separate changes. Trainer has no remote sampling constructor arguments and passes sampling=None to local_training_bindings.
Open review findings: the bounded task-index cache still needs inactive-entry eviction, and closing a client must cancel retained paused stream readers. These are not claimed fixed by the loader or stream-budget 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
Resolved native dependencies
Lock native Flight/runtime dependencies and the shared runtime Arrow/anyhow dependency edges.
Keep reproducible builds consistent with the crate manifests.
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 dependencies.
Serve prepared artifacts without an upload or hashing implementation.
packages/diskgraph/diskgraph-gateway/README.md
Native component protocol and admission
Document prepared artifacts, version 9 grouped components, shared ArrowBatch, three-stream default and sizing for concurrent clients.
Explain transport boundaries, backpressure and capacity needed for overlapping refit stages.
Replace the unused gateway API with complete task epochs.
packages/diskgraph/diskgraph-gateway/src/lib.rs
SamplerGateway and bounded output
Convert RuntimeBatch into the runtime ArrowBatch, then send minibatch metadata and native schemas/dictionaries/component batches through DoGet with bounded production.
Share component grouping with local consumers while keeping transport encoding and byte limits in the gateway.
packages/diskgraph/diskgraph-gateway/src/main.rs
Standalone entry point and stream budget
Start the shared owned server, default to three streams and close on SIGTERM/Ctrl-C.
Support train/validation/test overlap for one client without a separate serving implementation.
packages/diskgraph/diskgraph-gateway/src/paths.rs
Prepared artifact path resolution
Validate worker-local extensions and reject traversal/symlinks; keep moved path-containment tests beside their implementation.
Protect prepared artifact boundaries without coupling the shared runtime batch module to gateway paths.
Define version 9 stream/minibatch/component metadata and document message order, grouped examples, completion, budgets and flow control with references to the shared runtime modules.
Keep the wire contract explicit while locating grouping and selection in their transport-independent owner.
Convert native output into ArrowBatch with the GIL released and call the shared exporter; remove the separate local dictionary builder and vector-cloning helpers; expose uniform close.
Unify minibatch construction and export without adding local IPC serialization or transport byte limits.
Expose start, endpoint, health, port, close and context management; release the GIL during start/close and default to three streams.
Keep serving native and admit overlapping train/validation/test streams for one client.
packages/diskgraph/diskgraph-runtime/Cargo.toml
Runtime dependencies
Declare the existing Arrow array/schema and anyhow dependencies used by the shared minibatch representation.
Keep component construction independent of transport and Python.
packages/diskgraph/diskgraph-runtime/README.md
Native batch runtime contract
Document consuming RuntimeBatch into ArrowBatch and sharing component construction, reference resolution and Python export across backends.
Explain ownership, buffer sharing, GIL release and the absence of local IPC serialization or wire limits.
packages/diskgraph/diskgraph-runtime/src/batch.rs
ArrowBatch and component construction
Move grouping from the gateway into a transport-independent consuming conversion; move numeric vectors and feature chunks into Arrow components; test grouped targets, ragged domains, optional fields, row alignment and shared buffers.
Define minibatch structure once and remove vector cloning from local Python export.
Runtime ownership, replay cleanup and remote loaders
Pass a shared RuntimeOwner to normal/RHS/replay datasets; reuse per-stage cleanup stacks and release completed passes; force all remote loaders to zero subprocess workers; reject zero-natural-batch flat remote training before applying minimum steps.
Bound replay retention, preserve active-stage cleanup and remove the worker-induced remote sharding mismatch without rejecting valid grouped or weighted epochs.
Runtime doubles, replay retention and loader policy
Give doubles uniform close methods; check bounded iterator retention and cleanup on completion/error/early close; verify every remote loader mode ignores requested subprocess workers and rejects empty flat training epochs.
Prevent replay leaks and positive-length zero-batch remote training while preserving local worker behavior.
test/train/datamodule/test_persistent_sampling.py
Shared runtime lifecycle
Verify both backends share one owner across stages, retain indexes across epoch/phase cleanup and close/reopen at job boundaries, including failed worker shutdown.
Prevent premature index shutdown and leaks with explicit constructor ownership.
test/train/datamodule/test_push_sampling.py
Native shard integration and remote loader regressions
Verify parity, selections, backpressure and epoch cleanup; test paused stage streams with one/two logical clients, default/native service budgets and real remote loaders with requested subprocess workers.
Exercise native admission and prove forced inline loaders yield their advertised batches, preserve seed coverage and reject empty replica shards before requests.
test/train/datamodule/test_remote_sampling.py
Wrapper and standalone cleanup
Verify shared wrapper routing, native requests and endpoint validation; early iterator close cancels epochs while explicit dataset close releases its index.
Keep standalone lifetime consistent with borrowed module ownership.
Mock native serving and filesystem checks to verify the three-stream default and explicit budgets of one, four and six; keep imports compatible with CI isort.
Cover default refit capacity, preserve explicit limits and repair the import-order lint failure.
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
Latest follow-up validation at 6c82945036: 249 Python tests passed across the loader, native streaming, remote wrapper and persistent-sampling suites, including the new inline-loader and empty-epoch regressions. The preceding stream-budget checks passed four configuration cases, 18 public native-handle tests and 32 gateway tests (29 unit, three integration) against a rebuilt native extension. The nine focused native overlap/shared-backend cases are included in the subsequent full streaming suite; counts are not additive.
The CI failure was import ordering in test_push_sampling.py and test_sampling_service_config.py. The exact repository-wide check, python -m isort --extend-skip .venv --check-only ., now passes locally with isort 5.12.0; Flake8 passes on touched test files. Scoped Ruff and rustfmt checks and repository documentation validation passed. Full-file mypy and Ruff checks on the loader and its tests retain 92 and 87 baseline findings respectively, with no new diagnostics.
Earlier PR validation recorded 180 Rust tests and 539 Python tests, plus Clippy with warnings denied, generated artifact/notice, license and documentation checks. Those earlier totals were not rerun in full for the loader follow-ups. A production release wheel and cloud deployment were not validated locally. The cache-eviction and paused-client-close review findings above remain open.
It looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the reviews.auto_review.auto_pause_after_reviewed_commits setting.
The change adds a native Arrow Flight gateway for DiskGraph epoch sampling. It defines a versioned protocol, validates requests and artifact paths, streams bounded Arrow components, and adds native client and server lifecycle APIs. Python bindings expose sampling clients, servers, streams, and remote runtime indexes. Training configuration supports remote endpoints, replica-aware batching, shared runtime indexes, cancellation, and cleanup. Tests cover protocol validation, parity, backpressure, lifecycle, and failure handling. Dependency manifests and third-party notices are updated.
Remote training can reject multi-stage streams, fail during long epochs, or complete an epoch without training updates. Resolve these failures and the remaining cache, replay-memory, and setup-documentation risks before merging.
Pre-merge checks: 1 · 2
Failed checks: 2 warnings
Check name
Status
Explanation
Resolution
Change Map Design Review
The required maintained comment beginning with <!-- kumo-pr-change-map --> is unavailable. Repository inspection finds only template and policy references, and the supplied conversation does not inc…
Provide the single maintained conversation comment with the observed failures, current behavior, smallest sufficient changes, focused verification, and stopping condition. Remove or replace the whole-epoch timeout with a suitable idle polic…
Pr Size
The pull request changes 53 files. The authoritative review-scoped diff lists 53 changed paths, which exceeds the limit of 40.
Split the pull request into smaller, reviewable units.
Passed checks: 1 passed
Check name
Status
Explanation
Public Resource Access
All added external references are publicly accessible. The new dependency origins in 3rdparty/notices/inventory.json resolve on GitHub, the added `registry+https://github.com/rust-lang/crates.io-ind…
Full details: Public Resource Access
Explanation
All added external references are publicly accessible. The new dependency origins in 3rdparty/notices/inventory.json resolve on GitHub, the added registry+https://github.com/rust-lang/crates.io-index sources in the artifact manifest and lock data are public, and the added Apache license URLs return HTTP 200. Added relative repository links and path dependencies also resolve within the pull request.
Full details: Change Map Design Review
Explanation
The required maintained comment beginning with <!-- kumo-pr-change-map --> is unavailable. Repository inspection finds only template and policy references, and the supplied conversation does not include the comment body. Therefore the required design evidence cannot be reviewed. The changed code also retains the cited concerns: fixed one-hour whole-stream deadlines in client.rs; permanent four-entry cache rejection while holding the cache lock during RuntimeIndex::new in lib.rs; a default max_streams=1 with one-second admission timeout; misleading artifact-transfer wording in packages/kumo-engine/README.md; and cleanup ownership that retains replay iterators while local close() is a no-op.
Resolution
Provide the single maintained conversation comment with the observed failures, current behavior, smallest sufficient changes, focused verification, and stopping condition. Remove or replace the whole-epoch timeout with a suitable idle policy. Evict idle cache entries and avoid holding the cache lock during index construction, with per-key coordination. Set or document a safe max_streams value and test multi-stage admission. State that matching artifacts must already exist on each worker. Release completed local iterators from cleanup ownership while retaining cleanup for active remote iterators.
Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts
haoweiyu-kumo
changed the title
Push complete DiskGraph task batches over Arrow Flight
Serve remote DiskGraph sampling through the native Flight gateway
Oct 31, 2026
haoweiyu-kumo
changed the title
Serve remote DiskGraph sampling through the native Flight gateway
Serve remote sampling with an embedded Rust Flight gateway
Oct 31, 2026
The reason will be displayed to describe this comment to others. Learn more.
Follow-up Comments
Record the upstream copyright for tonic-prost.
11589-11589: 🗄️ Data Integrity & Integration | 🟡 Minor | ⚡ Quick win
Record the upstream copyright for tonic-prost.
The upstream release has a copyright notice in its MIT license. This record instead points to an SPDX template with placeholders. (docs.rs)
The generated notice attributes that copyright to tonic, but not to tonic-prost. (raw.githubusercontent.com) Record the upstream license and copyright for tonic-prost, then regenerate the notice.
Clarify who provides the training-table extension to the workers.
Line 133 says "Workers receive the prepared training-table extension". The core README says sampling does not upload or register artifacts. It also says a missing extension fails with a preparation error. Trainer.fit creates the extension with a run-specific tt_id and does not distribute it.
State that the extension must be visible to every worker at the same relative train_tables/<id> path. Name the mechanism that provides it, for example shared storage. Remove the wording that implies the engine transfers it.
The Quickstart imports Trainer with from kumo.engine import ... Trainer. This snippet calls kumo.Trainer, but kumo is never imported. The snippet raises NameError when copied.
The reason will be displayed to describe this comment to others. Learn more.
Major · Stability & Availability · Quick win
Remove the fixed one-hour limit on whole epoch streams.
Line 171 sets a gRPC timeout of 3600 s on the DoGet call. Line 188 sets a reader deadline of the same length. Both values cover the whole epoch stream, not one read.
The server applies backpressure. A producer waits while the trainer does not read. The stream lifetime therefore equals the training time for the shard. This includes pauses for validation, checkpointing and slow GPU steps.
If a shard takes longer than one hour, next_message returns DeadlineExceeded. Training then fails with FlightTimedOutError in the middle of the epoch, and no retry or resume exists. Large graphs and long epochs reach this limit in normal training.
Remove the whole-stream timeout. If a hang guard is necessary, apply it per message: reset reader.deadline before each next_message read, with a configurable idle limit. Remove ticket.set_timeout, or make it configurable.
🐛 Proposed per-message idle deadline
- let mut ticket = Request::new(Ticket { ticket: request.into() });- ticket.set_timeout(Duration::from_secs(3600));+ let ticket = Request::new(Ticket { ticket: request.into() });
async fn next_message(
reader: &mut Reader,
stop: &mut watch::Receiver<bool>,
client_stop: &mut watch::Receiver<bool>,
) -> anyhow::Result<Option<DecodedFlightData>> {
if *stop.borrow() || *client_stop.borrow() {
return Err(Status::cancelled("Sampling stream closed").into());
}
+ // Bound the wait for one message, not the whole epoch.+ reader.deadline = tokio::time::Instant::now() + IDLE_TIMEOUT;
Note: a per-message idle limit still fails if the trainer itself pauses longer than the limit. Choose the limit with this in mind, or apply it only while a read is waiting for data.
# Find any other mutation of the task index cache in the gateway.
rg -n -e 'indexes\.' -e 'indexes =' -- packages/diskgraph/diskgraph-gateway/src/lib.rs ||test"$?" -eq 1
# Find the client error mapping for ResourceExhausted.
rg -n -F -e 'ResourceExhausted' -e 'client_error' -e 'ArrowInvalid' -- packages/diskgraph packages/kumo-core ||test"$?" -eq 1
Repository: NVIDIA/kumo-engine
Length of output: 1494
Evict idle task indexes instead of rejecting new keys permanently.
The gateway caches at most four RuntimeIndex entries keyed by (extension_ref, engine_seed, threads). Once the cache holds four entries, a request with a new key fails with ResourceExhausted("Task index cache is full"). No code in lib.rs removes entries, so the failure persists until the server state is dropped.
Requests can reach the limit when one gateway serves more than four distinct extension_ref values, or when callers request different num_threads values. client_error in packages/diskgraph/diskgraph-pybind/src/sampling_client.rs maps ResourceExhausted to ArrowInvalid. The trainer then receives an argument error.
The cache lock also stays held while RuntimeIndex::new(...) runs. Opening one large index can block lookups for other streams.
Evict entries that no stream references. The cache holds one Arc per entry, so a strong count of 1 means that only the cache uses the entry. Return the error only when all four entries are still in use.
♻️ Proposed eviction of idle indexes
} else {
if indexes.len() >= 4 {
- return Err(Status::resource_exhausted("Task index cache is full"));+ // Entries referenced only by the cache have no active stream.+ indexes.retain(|_, index| Arc::strong_count(index) > 1);+ if indexes.len() >= 4 {+ return Err(Status::resource_exhausted("Task index cache is full"));+ }
}
Committable suggestion
IMPORTANT
Review the replacement and indentation, then test before committing.
Suggested change
let key = (request.extension_ref.clone(), request.engine_seed, threads);
let index = {
letmut indexes = state.indexes.lock();
ifletSome(index) = indexes.get(&key){
index.clone()
}else{
if indexes.len() >= 4{
returnErr(Status::resource_exhausted("Task index cache is full"));
}
let index = Arc::new(
RuntimeIndex::new(
state.root.to_string_lossy().into_owned(),
request.extension_ref,
request.engine_seed,
Some(threads),
)
.map_err(runtime_error)?,
);
indexes.insert(key, index.clone());
index
}
};
let key = (request.extension_ref.clone(), request.engine_seed, threads);
let index = {
letmut indexes = state.indexes.lock();
ifletSome(index) = indexes.get(&key){
index.clone()
}else{
if indexes.len() >= 4{
// Entries referenced only by the cache have no active stream.
Name SamplingRequest explicitly in this description.
The sentence follows the SamplingClient(endpoint) example, but SamplingClient accepts only endpoint. SamplingRequest accepts the keyword iterator settings and validates batch, thread, and shard values. The current wording can lead readers to pass these keywords to SamplingClient and receive a TypeError. (raw.githubusercontent.com)
Proposed wording
-The native constructor accepts keyword iterator settings, validates batch/thread-budgets and shard bounds, and owns the request kind and protocol version.+`SamplingRequest` accepts keyword iterator settings and validates batch/thread+budgets and shard bounds. The native layer owns the request kind and protocol version.
The reason will be displayed to describe this comment to others. Learn more.
Caution
Some comments are outside the diff and can’t be posted inline due to GitHub limitations.
⚠️ Outside diff range comments (1)
Major · Reject capped epochs that have no native batches. · lightning_data.py:2570-2576
packages/kumo-core/src/kumo_core/diskgraph/lightning_data.py:2570-2576 Major · Functional Correctness · Quick win
Reject capped epochs that have no native batches.
_max_train_batches reports two for the eight-seed, two-endpoint, two-DataLoader-worker configuration, while _iter_batches yields zero. The iterator exits on had_any == False, so the workflow does not hang; Lightning instead receives a positive len(train_dataloader) and performs no training updates.
Pass the effective loader_num_workers from _make_dataloader into the count calculation. Count batches after splitting across runtime_num_workers * loader_num_workers * replicas. If that count is zero, raise before min_steps_per_epoch changes it to one. Using self.num_workers would be incorrect because _make_dataloader can override it to zero.
Make _natural_batches shard-aware with the same effective worker count, and reject a zero natural count before applying min_steps_per_epoch. Apply that argument to the weighted and multi-stage train-count call sites as well.
Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 8 remain after this review.
The reason will be displayed to describe this comment to others. Learn more.
Source review at 6c82945036f1; inline comments cover the epoch deadline, historical index-cache exhaustion and whole-batch byte cap.
Two bounded validation requests:
Concurrent-stream memory: please add a slow-consumer/multiworker test at representative fanout/features, measuring peak server/client RSS, admission-limit behavior and cleanup after cancellation. A one-message server queue and one decoded prefetch per replica do not cap aggregate memory across loader workers, native sampling/IPC working buffers, gRPC and PyTorch prefetch. Please document stream sizing; this is not a request to introduce a new memory-management subsystem without evidence.
Matched performance evidence: please compare local-before/local-after, current local versus loopback Flight, and one representative constrained-local versus remote case with the same graph, workload and sampling settings. Report batches/s, wire volume, decode cost and trainer wait or epoch time. The historical dgfeed results motivate offloading but do not qualify this Flight implementation or establish that sharing ArrowBatch has no local-path overhead.
Public Trainer integration and direct-local-plus-remote hybrid execution are follow-ups, not requested changes for this PR.
The reason will be displayed to describe this comment to others. Learn more.
The fixed 3,600-second RPC timeout and the separate absolute Reader.deadline expire an entire epoch shard, including time the consumer spends training between reads. A healthy long epoch can therefore fail on its next read. Please remove this hard-coded lifetime limit or make an intentional caller-owned overall budget explicit, changing both deadline mechanisms. Keep connection establishment/cancellation bounded separately. A short injectable-deadline regression can demonstrate that intentional consumer pauses do not expire a healthy epoch; no hour-long test is needed. This is application policy, not an Arrow Flight requirement.
The reason will be displayed to describe this comment to others. Learn more.
Please also report concrete feeder-failure/cancellation test results and distinguish those failures from a healthy trainer spending a long time between reads. Any replacement timeout must respect that distinction; an idle timeout that counts intentional consumer pauses would recreate the problem.
Restarting the training job from a saved epoch/checkpoint is an acceptable initial recovery approach. Please clarify whether the installed Kumo Engine provides automatic retry/restart above the iterator, what existing job/loader timeout controls apply, and what recovery the caller must perform. This is not a request to implement transparent batch replay or failover here. Preserve explicit error propagation, cleanup and checkpoint/restart documentation.
The reason will be displayed to describe this comment to others. Learn more.
This cache retains historical (extension_ref, engine_seed, threads) keys after streams finish. Serve four distinct configurations sequentially, let all streams close, then request a fifth: it returns Task index cache is full even with zero active streams. New-key admission stays blocked until server restart. Please reclaim unused entries before rejecting a new key while preserving active users, and cover five sequential configurations in a focused regression.
The reason will be displayed to describe this comment to others. Learn more.
Clarifying the intended deployment: remote feeders start/stop with their training job and serve that job's graph; they are not a general multi-job service. Please state that lifecycle and explain which runtime configurations need caching within it (for example, train/validation/test or TT/PT extensions).
This qualifies my earlier eviction request: I am not asking for a general-purpose multi-job cache. Please establish whether the four-key limit can interrupt the intended single-job workflow, then use the smallest appropriate change. Cache reuse can be useful without promising a long-lived shared service.
The reason will be displayed to describe this comment to others. Learn more.
This limit applies cumulatively to the complete logical minibatch, not each frame/component. A batch that works through local PyBind can consequently fail remotely with instructions to reduce batch size/fanout, changing the model workload to satisfy transport. Please establish representative local/remote batch-size compatibility and distinguish bounded in-flight transport bytes from maximum logical-batch size. If larger batches are intentionally outside the initial supported workload, state and test that boundary explicitly. Splitting into more components does not bypass the current check, and simply raising the constant is not proof of memory safety.
The reason will be displayed to describe this comment to others. Learn more.
Concrete workload: embedding features can make one logical minibatch several GiB, so 64 MiB is not a sufficient supported-batch boundary. Local PyBind has no equivalent transport cap: Python demands a batch and receives shared native buffers.
Please separate logical batch size, outstanding/prefetched batches and transport-message size. Here, the one-slot server queue counts FlightData messages, and transport draining is not acknowledgement that Python consumed the complete batch; serialization/client reconstruction add memory too. The configured 128 MiB per-message limit is another distinct constraint.
I am not requiring a chunking implementation. Whole-batch transfer is acceptable if it supports the intended batch sizes with justified demand/backpressure and resource use. Choose the smallest solution that preserves local batch-size capability; chunking is only an option if actual transport/memory constraints require it. If chosen, splitting only at component boundaries would not handle an oversized embedding component. Please demonstrate a representative large-batch case rather than forcing smaller model batches solely for this transport.
Replica consistency — reviewed at 6c82945036f1. Please check that the trainer and every feeder refer to the same graph index and prepared TT/PT extension before yielding batches. Can we reuse existing graph/snapshot and task-artifact identifiers that distinguish the intended versions, rather than relying only on matching paths and wire version?
A lightweight mismatch check and focused wrong-graph/wrong-extension test would be valuable here. This is data consistency, not a request for a new authentication token, content-hashing subsystem or service registry; the current deployment assumes trusted network peers.
Default stream admission effectively limits a feeder to three loader workers — reviewed at 6c82945036f1.
For one active epoch, each DataLoader worker opens its own DoGet stream to the feeder. The server's default max_streams=3 admits only three of those workers. Each permit is held for the entire epoch request, not released after each minibatch. A fourth worker waits at most one second on the server, then receives ResourceExhausted: Sampling service is busy unless an admitted stream finishes or is cancelled in that interval.
Thus an ordinary configuration with four or more concurrent loader workers can fail admission; the three slots do not rotate among workers as minibatches are consumed. Trainer pauses can prolong occupancy, but are not required to trigger this problem. Local PyBind has no equivalent three-worker admission restriction.
Please ensure admission capacity matches the intended job's worker/stage/replica concurrency, and add a focused four-worker test that keeps the first three epoch streams open while the fourth requests admission. Verify the supported configuration can run all intended workers, plus cancellation/permit release. A configurable resource limit can be reasonable; this is a request to correct or justify the default behavior, not to introduce unlimited concurrency or a new admission subsystem.
Refining the performance validation request at 6c82945036f1: please measure (1) per-component schema/IPC/message overhead through the one-message queue, and (2) deterministic round-robin head-of-line blocking when one feeder is temporarily slow while another has a ready batch.
Retain local-before/local-after measurements because this PR also changes local batch conversion; verify that the existing tensor-mapper input contract is preserved. Compare current local and Flight on matched batches, and include the earlier custom TCP prototype where reproducible rather than assuming it avoids these costs.
One representative workload can establish throughput, decode cost and trainer wait. Packing relation edges with offsets/counts or changing merge order are possible follow-ups if measurements justify them, not mandatory rewrites in this comment.
The reason will be displayed to describe this comment to others. Learn more.
There is a second large-batch limit here, independent of the cumulative 64 MiB check: the server caps each Flight/gRPC message at 128 MiB, and client.rs:140–141 sets the same limit.
The mini-batch is already divided into components, but lib.rs encodes each component's entire Arrow record batch into one record-batch message (with separate schema/dictionary messages); it does not split that component further. A large sampled subgraph with embedding features can therefore exceed 128 MiB in a single component. For example, 100,000 sampled nodes × 768 float32 values is about 293 MiB for the embedding values alone.
Consequently, removing/raising the 64 MiB mini-batch cap alone does not preserve local-mode batch capability. This strengthens the case for splitting oversized components into bounded record-batch messages and reconstructing the original component client-side, rather than merely making the total-batch cap configurable. Please include a >128 MiB single-component regression when addressing large-batch support. The present 64 MiB cap masks this additional failure.
The reason will be displayed to describe this comment to others. Learn more.
Optional naming clarification: this new SamplingRequest describes an epoch-shard batch iterator, including stage/epoch/sharding, targets, topology and feature projections; it is not a request for one graph-sampling operation. Consider naming it BatchIterationRequest to make that contract clearer and distinguish it from the existing low-level sampling requests. This is a clarity suggestion, not a correctness blocker.
Following up on component/message overhead and the 128 MiB per-message limit: consider packing compatible topology components across relationships, while allowing oversized feature components to be split into bounded transport chunks.
For example, packed src/dst arrays plus relationship offsets/counts could reduce repeated schemas and small messages. Feature components can retain their own schemas and be chunked when large. Both unpacking topology and reassembling feature chunks should stay in the native Rust client/shared exporter, preserving the existing complete mini-batch contract for Python's tensor mapper, including per-hop counts and relationship identity.
Please assess topology packing against its concatenation/copy cost; it should not impose extra packing copies on the local PyBind path without measured benefit. This is an implementation option for the existing overhead/large-batch concerns, not a requirement to force the entire mini-batch into one schema or message.
why are we restricting num_workers for dataloaders to 0 for remote sampling? that will kill performance as the remote machines then cannot prepare enough batches and the gpu trainer is bound by GIL too
this is because then the trainer does the runtime_batch_to_hetero conversion in python with GIL, so one at a time, on the local mode the num_workers on dataloaders, has each worker with its own GIL which amortizes this python conversion time and lets training not be gil, cpu bound.
@haoweiyu-kumo some designs from my earlier prototype (kumo#31656) that I think make sense and should be kept, along with many of your changes.
The current dataloading works like this: the trainer makes a DataLoader with W workers (W processes), each with T threads of internal parallelism. For the trainer the API is still for batch in dataloader. The important speedup is that each worker process has its own GIL, because Python holds the GIL for the conversion to HeteroData, feature encoding, queueing and other PyTorch overhead (runtime_batch_to_hetero alone was ~43
ms/batch in my measurements). With num_workers=0 all of that sits on the training thread.
I think the workers are the right boundary for distributed feeders: the trainer gets batches from the DataLoader irrespective of where each worker's batches come from. Worker 1 can be local, worker 2 can read from remote 1, worker 3 from remote 2. So let's not restrict remote mode to num_workers=0; each worker can open its own remote client after the fork, and the remote feeders only need to return what a worker builds today, the RuntimeBatch, which your implementation already does. So all the hetero conversions happen on the trainer but for them to be fast and not a bottleneck we need num_workers to be large. This does introduce a dependency of having quite a bit of vCPUs with the gpu instance but even 16-32 vcpus would be enough. Also, some more useful designs I think we should keep and let's discuss it if needed:
Stateless feeders, state on the trainer. The worker's iterator decides which seeds go in each batch and asks a feeder to sample that batch; the feeder samples it from its replica and returns the RuntimeBatch. Epoch, shuffle, replay, stages and quotas stay on the trainer, in the same code as the local path. This makes failures easy (a failed feeder's outstanding batches go to another feeder holding the same replica,
instead of losing a whole shard) and balances load (any feeder serves any batch, so a slow one doesn't set the pace). Latency is hidden by keeping a few requests in flight per feeder. We only need batches to be valid samples of the graph, not to match a local run, so ordering across feeders doesn't have to be fixed.
Compression on the wire. Arrow IPC supports per-buffer zstd/LZ4; this gets important with large batches.
The scheduling logic (shards, replay, quotas) would stay in the Python iterator on the trainer, rather than being reimplemented server-side. That also shrinks the Rust server.
Fail loudly. If no feeder can serve a batch, raise instead of ending the epoch short.
Server limits should scale with workers. max_streams=3 assumes one process per trainer. With W workers it has to be W × stages, or connections should be pooled per worker.
let me know if this makes sense, other @Victsm talk about the other problems which may occur
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
GitHub Actions
Error
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Victsm
left a comment
Builds on the native batch runtime from merged #144. Targets
main.Serve complete DiskGraph batches through an embedded Rust Flight server and native client. The shared public
diskgraph.RuntimeIndexresolves manifests and normalizes options for both modes: omitsampling_endpointsfor local sampling or supply replica URIs for the Rust remote backend. Workers serve prepared graph/training-table replicas; the trainer only needs a portable manifest. Artifact distribution, cluster provisioning and public Trainer integration remain separate changes. Trainer has no remote sampling constructor arguments and passes sampling=None to local_training_bindings.Rust owns lazy replica connections across stages/extensions, epoch sharding, weighted-sidecar path translation, one-batch prefetch per replica, deterministic round-robin merging and cancellation. Ordered RHS catalog scans use the first replica. Remove the Python remote iterator module and thread pool. Epoch close joins native readers and discards queued batches while preserving connections; index close cancels all epochs and is terminal. GC signals cancellation without joining under the GIL. A process-local RuntimeOwner lazily creates one local or remote index and is passed explicitly to all module datasets. Epoch cleanup closes only epoch iterators; module teardown closes the shared owner. Standalone datasets close their own owner explicitly. Pickle excludes native state; a PID change opens a process-local handle. This removes remote-specific binding, managed flags and private cache injection. Existing local constructor/iterator calls retain their behavior. Both local and remote iterators always expose idempotent close(); local close is a synchronous no-op and loader cleanup registers it directly without getattr/callable checks.
Local sampling and the server consume RuntimeBatch into a shared diskgraph_runtime::ArrowBatch. The Flight client returns that same representation, and one Python exporter restores the complete minibatch. Remove the separate local dictionary builder: numeric vectors move into Arrow arrays and existing target/feature buffers stay shared. Local conversion performs no IPC/JSON serialization and has no transport byte limit.
Python builds immutable SamplingRequest objects; Rust owns ticket serialization, request kind and protocol version. Version 9 groups aligned seed/target, node/assignment, edge-endpoint and positive-label columns. Column references return arrays; half-open column ranges return record batches without target struct wrappers or array/as markers. Each DoGet returns an epoch shard: initial schema application metadata describes counts; a metadata-only message starts each minibatch with its layout/count; component schemas, dictionaries and Arrow data follow. Validate final RPC status, EOF/counts, references and the cumulative 64 MiB minibatch budget. Export through the Arrow C Data Interface; no nested IPC streams or Python wire decoding.
Server production reserves a one-slot output queue before sampling; gRPC propagates backpressure. Transport buffers/native working memory are additional. Listener shutdown retains its five-second drain grace. Native lifecycle/progress events reach Python logging.
Maintained change map
Validation: 180 Rust tests passed: 43 runtime, 29 gateway, 105 PyO3 unit and 3 gateway integration. Rebuilt the native extension; 539 Python tests passed: 106 transport/ownership/public handles, 16 shared-wrapper, 45 local e2e, 317 trainer/loader, 46 engine and 9 local-runtime. Clippy with warnings denied, scoped rustfmt, generated artifact/notice, license and documentation checks passed. The shared constructor tests verify moved seed buffers, shared feature chunks, grouped target metadata and ragged row domains. Previously checked scoped mypy/Ruff and CI flake8/isort passed; full-file lightning_data.py mypy has eight unchanged baseline errors. A full release wheel and cloud deployment were not run locally.