Serve remote sampling with an embedded Rust Flight gateway

packages/diskgraph/README.md

@@ -126,7 +126,7 @@ Use `--customer-prefix` if the storage directory differs from the customer name.
 The repository's CI builds the native Python wheel for every relevant push and
 pull request (`.github/workflows/_diskgraph-wheel.yml`), and the kumo-engine
 wheel bundles it. The wheel contains the extension and the standalone `ingest`
-binary.
+binary. The extension embeds the native Flight server.
 
 Kumo builds and distributes that artifact from its in-tree
 `packages/diskgraph` source. This repository is the sole source of truth for
@@ -143,10 +143,16 @@ maturin:
 uv pip install .
 ```
 
-The base package installs the native extension and bundled standalone ingest
-binary. Use the `kumo` extra only when calling Kumo compiler APIs outside an
+The base package installs the native extension and standalone `ingest` binary.
+Use the `kumo` extra only when calling Kumo compiler APIs outside an
 environment that already installs `kumo-api`.
 
+## Native Flight gateway
+
+[`diskgraph-gateway`](diskgraph-gateway/README.md) serves worker-owned epoch shards
+from the shared native runtime. Python exposes the same Rust server through
+`diskgraph.SamplingServer`; the standalone CLI can be built with Cargo.
+
 ## Native batch runtime
 
 [`diskgraph-runtime`](diskgraph-runtime/README.md) provides reusable index ownership

packages/diskgraph/diskgraph-gateway/Cargo.toml

@@ -9,29 +9,21 @@ edition = "2021"
 license.workspace = true
 
 [dependencies]
-diskgraph-core = { path = "../diskgraph-core" }
+diskgraph-runtime = { path = "../diskgraph-runtime" }
 diskgraph-engine = { path = "../diskgraph-engine" }
-diskgraph-index = { path = "../diskgraph-index" }
-diskgraph-rowpack = { path = "../diskgraph-rowpack" }
 arrow-flight = "58.4.0"
 arrow-array = { workspace = true }
 arrow-schema = { workspace = true }
 arrow-ipc = "58.4.0"
 tonic = "0.14"
 tokio = { workspace = true }
-tokio-stream = "0.1"
+tokio-stream = { version = "0.1", features = ["net"] }
+futures = { workspace = true }
 anyhow = { workspace = true }
 tracing = { workspace = true }
 tracing-subscriber = { workspace = true }
 serde = { workspace = true }
 serde_json = { workspace = true }
-futures = { workspace = true }
-bytes = { workspace = true }
-prost = "0.13"
-prost-derive = "0.13"
-gxhash = "3.5"
-bincode = "1.3"
-uuid = { version = "1.4", features = ["v4"] }
 parking_lot = "0.12"
 clap = { version = "4.4", features = ["derive"] }
 

packages/diskgraph/diskgraph-gateway/README.md

@@ -0,0 +1,149 @@
+# Native Flight gateway
+
+`diskgraph-gateway` serves complete task batches from `diskgraph-runtime` without
+loading Python. It owns epoch/shard iteration over prepared local artifacts,
+Arrow serialization, bounded production, cancellation, health counters, and
+structured lifecycle/progress logs. Minibatch component construction and layout
+reference resolution live in `diskgraph-runtime`: the server sends its
+`ArrowBatch`, and the client reassembles that same representation. Local and
+remote Python iterators use one shared Arrow C Data exporter.
+
+```bash
+cargo run --release -p diskgraph-gateway -- \
+  --data /local/graph --host 0.0.0.0 --port 8815 \
+  --max-streams 6 --max-threads 2
+```
+
+`--max-streams` defaults to three for overlapping train/validation/test streams
+from one loader worker. Budget overlapping stages times concurrent loader workers
+using this server, across all trainers; the example covers three stages for two
+workers. Paused streams retain their permits under backpressure.
+
+The graph must already exist. Before requesting an epoch, the caller must
+materialize and verify its training-table extension beneath the graph at
+`train_tables/<id>` or `pred_tables/<id>`. Artifact distribution and manifest
+verification belong to the preparation layer. This gateway reads those files;
+it does not upload or register them. Missing extension metadata returns
+`FailedPrecondition` so the caller can finish materialization before retrying.
+
+Protocol version 9 streams native Arrow component record batches. A `DoGet`
+ticket contains one normalized epoch/shard request. Python builds an immutable
+`diskgraph.SamplingRequest`; the native client owns JSON ticket serialization and
+sets the request kind and protocol version. Typed constructor arguments reject
+invalid shards or resource settings before opening an RPC. Protocol metadata always
+lives in Flight `app_metadata`:
+
+- The initial empty schema carries `kind: "stream"`, `version`, `num_samples`
+  and `num_batches`.
+- Each minibatch starts with a metadata-only FlightData message carrying
+  `kind: "minibatch"`, `batch_id`, `num_components` and the full `layout`.
+  Its IPC message type is NONE and it has no Arrow data body.
+- Each component sends its own schema, optional dictionaries and record batch.
+  Those messages carry only `batch_id` and `component_id` in application metadata.
+
+Components can have different schemas and row counts. The client reads the
+minibatch layout before any component and yields only after all declared
+components arrive. Missing starts, repeated starts within a minibatch, invalid
+IDs and partial minibatches fail the stream. There are no binary wrapper columns,
+nested IPC streams, or iterator/end control record batches. Flight EOF completes
+the stream, and the client validates its received count and final RPC status.
+The server also verifies native counts.
+
+```text
+Empty schema + stream metadata
+Metadata-only minibatch 0 start + layout
+  Component 0 schema / dictionaries / record batch
+  Component 1 schema / dictionaries / record batch
+Metadata-only minibatch 1 start + layout
+  Component 0 schema / dictionaries / record batch
+  ...
+Successful EOF
+```
+
+Aligned columns share component record batches: seed IDs/positions, optional
+timestamps/weights, per-seed LP counts and appended target columns; sampled node
+IDs with their seed assignments per channel/type; edge source/destination
+indices per channel/relation; and ragged LP positive IDs/masks/optional weights
+and local indices. Negative masks, exclusion IDs and per-hop counts retain their
+own row domains. Feature batches retain their original columns and chunks.
+The empty negative mask remains an explicit component.
+
+Layout references select columns by position:
+
+```json
+{
+  "seed_node_ids": {"component": 0, "column": 2},
+  "target": {"component": 0, "columns": [5, 7]}
+}
+```
+
+`column` always returns an Arrow array. A half-open `columns` range always
+returns a record batch; no `array` or `as` marker is used. Whole feature batches
+use `[0, number_of_columns]`. Targets are appended as ordinary columns, without
+a struct wrapper. Original target field names, order, types, field metadata and
+schema metadata survive grouping and projection. Positional references support
+duplicate names. Invalid component IDs, column indices, ranges or row alignment
+fail rather than silently changing the trainer's data.
+
+`diskgraph.SamplingClient` uses Rust's low-level `FlightDataDecoder`, which
+supports schema changes and resets dictionaries between schemas. The ordinary
+PyArrow Flight reader assumes a fixed schema and cannot read this protocol.
+The native client validates and assembles a complete minibatch before exposing
+its arrays and record batches through the Arrow C Data Interface to Python.
+
+A one-slot native output queue reserves capacity before sampling each batch.
+Each logical minibatch has a cumulative 64 MiB wire-payload limit, including
+component schemas, dictionaries and application metadata. The Rust remote backend prefetches at most
+one native decoded batch per server through the remote backend. When the trainer pauses, client prefetch and gRPC
+buffers fill, Flight stops draining the output queue, and the native producer
+waits for capacity. There are no application credit messages. These queue limits
+exclude additional transport buffers and native sampling/encoding working memory.
+
+Cancelling a reader wakes a producer blocked on output capacity and releases its
+stream permit. Native sampling already in progress stops at the next batch
+boundary. Listener shutdown cancels all producers; a transport that is not being
+read may consume the existing five-second graceful-drain period before the
+listener force-closes it.
+New requests allow up to one second for a producer slot to become available,
+including release from a preceding cancelled stream, then report the server busy.
+
+`health` is the supported action. The earlier per-request
+graph/feature ticket API is replaced by complete task epochs. This service is
+intended for trusted peers on a restricted network.
+
+The Python extension exposes `diskgraph.SamplingServer`, a lifecycle handle for
+this same Rust implementation. It runs inside the Python/Ray actor process on
+Rust-owned threads. Constructor and close release the GIL; sampled batches never
+pass through a Python server handler. The core `SamplingService` wrapper owns
+this handle and advertises the worker's routable IP.
+
+Explicit close cancels producers and joins native serving threads. Context-manager
+exit closes the handle; dropping it signals asynchronous shutdown without
+blocking Python garbage collection. Native events use the extension's existing
+Python logging bridge. The CLI uses stderr tracing and handles SIGTERM/Ctrl-C.
+
+From `packages/diskgraph`, validate with:
+
+```bash
+cargo test -p diskgraph-gateway
+cargo clippy -p diskgraph-gateway --all-targets -- -D warnings
+```
+
+Python integration tests exercise the embedded listener and compare complete
+batches with the local iterator. The standalone executable remains available
+through Cargo; Python workers use the extension directly.
+
+
+The shared Python `diskgraph.RuntimeIndex` wrapper selects local sampling by
+default or the native `RemoteRuntimeIndex` backend with `sampling_endpoints`.
+Manifest resolution and option normalization are shared. Rust owns cached replica
+connections across extensions/stages, shard assignment, weighted-sidecar path
+relativization and deterministic round-robin merging. Ordered RHS catalog scans
+use only the first replica. Each native reader receives demand for one batch;
+a one-slot result channel retains it until consumed. No Python executor or
+per-replica Python batch export participates in prefetch.
+
+Closing an epoch cancels all its readers, joins their threads and discards queued
+batches. Closing an index cancels every epoch and releases connections; it is
+terminal. Dropping native handles signals cancellation without joining under the
+GIL. Explicit `close()` releases the GIL and performs deterministic shutdown.