Ec connector nixl transfer#9
Open
omerpaz95 wants to merge 22 commits into
Open
Conversation
Signed-off-by: omerpaz95 <omerpaz95@gmail.com>
Signed-off-by: omerpaz95 <omerpaz95@gmail.com>
Signed-off-by: omerpaz95 <omerpaz95@gmail.com>
…nd pass to EC CPU Connector Signed-off-by: omerpaz95 <omerpaz95@gmail.com>
…from config for cleaner implementation Signed-off-by: omerpaz95 <omerpaz95@gmail.com>
Signed-off-by: omerpaz95 <omerpaz95@gmail.com>
The block free-list belongs with the memory it tracks. Scheduler-only usage (workers never call alloc/free) so no locking is needed. alloc() returns block indices (not tensor views) because that's what downstream consumers — ZMQ payloads, connector metadata, NIXL descriptor offset math — actually need. Signed-off-by: omerpaz95 <omerpaz95@gmail.com>
Producer ROUTER socket binds to these; consumer DEALERs connect to the host:port advertised by each E node. Default 5601 (one above NIXL) so both connectors can coexist on one host in tests. Signed-off-by: omerpaz95 <omerpaz95@gmail.com>
Enables unconditional ``connector.shutdown()`` from teardown paths without subclass checks. Subclasses override to release resources. Signed-off-by: omerpaz95 <omerpaz95@gmail.com>
XferReq/XferAck are the sole messages on the scheduler↔scheduler ZMQ side channel. Compatibility hash guards against silently talking to a peer built from a drifted codebase. Signed-off-by: omerpaz95 <omerpaz95@gmail.com>
Implements the NixlWrapper subset the connector uses (register_memory, add_remote_agent, prep_xfer_dlist, make_prepped_xfer, transfer, check_xfer_state, ...). WRITE xfers synchronously memmove bytes between registered regions via ctypes, so tests can assert real data arrival without a real NIXL runtime. Signed-off-by: omerpaz95 <omerpaz95@gmail.com>
Opens the same mmap workers use, creates a NIXL agent, registers the region, and builds an xfer_descs handle + serialized peer descriptor bytes for inclusion in XferReq messages. State dicts (_encoding_map, _ready, _dealers) are initialized empty. Producer scaffolding is added but the listener thread arrives in the next task. Signed-off-by: omerpaz95 <omerpaz95@gmail.com>
Listener decodes XferReq messages, validates the compatibility hash, and either dispatches (when the encoding is locally available) or parks in _waiting for the future save_caches path to drain. NIXL transfer itself is stubbed; implemented in the next task. Signed-off-by: omerpaz95 <omerpaz95@gmail.com>
Per-transfer lifecycle: register source tensor, add remote agent, prep local/remote dlists, make_prepped_xfer + transfer, poll for DONE, release everything on return. Consumer's memory descriptor arrives as a msgpack-encoded list[tuple[int,int,int]] and is rebuilt on our side via get_xfer_descs — no opaque NIXL object crosses the wire. Signed-off-by: omerpaz95 <omerpaz95@gmail.com>
has_cache_item is the engine's readiness probe: True if the xfer completed, None if still in flight, False if we've never heard of the hash. _drain_acks polls every DEALER non-blocking and routes each ack to _complete. Success adds to _ready; failure frees the blocks. Signed-off-by: omerpaz95 <omerpaz95@gmail.com>
Allocates N blocks from the region, records the mapping in _encoding_map, sends XferReq to the E node. Encoder node address is resolved from kv_transfer_params (future PR) or falls back to ec_connector_extra_config.default_encoder_node. DEALER sockets are cached per (host,port) so we don't churn connections. Signed-off-by: omerpaz95 <omerpaz95@gmail.com>
Drains any pending acks, then pops every ready mm_hash from _encoding_map into the ECCPUConnectorMetadata the worker receives. Worker's start_load_caches CPU→GPU-copies from the block indices. Hand-off-and-forget: clears _ready each step. Updates existing worker-path tests to match the new _ready-gated semantics. Signed-off-by: omerpaz95 <omerpaz95@gmail.com>
Future save_caches populates _encodings[mm_hash] then calls _drain_waiting(mm_hash) so any XferReqs that arrived early get served. Tested via a direct call until save_caches lands in a follow-up PR. Signed-off-by: omerpaz95 <omerpaz95@gmail.com>
Two connectors (producer + consumer) wired over real ZMQ sockets with FakeNixlWrapper standing in for NIXL's data plane. Asserts the consumer's mmap actually receives the producer's tensor bytes, and that the wait-then-drain path works for requests that arrive before the encoding is ready. Signed-off-by: omerpaz95 <omerpaz95@gmail.com>
Spec covers the scheduler-side NIXL transport, scheduler↔scheduler ZMQ control plane, transient NIXL agent metadata per transfer, and the single-dict + ready-set state model on the consumer. Plan decomposes the work into 13 TDD-structured tasks. Signed-off-by: omerpaz95 <omerpaz95@gmail.com>
Signed-off-by: omerpaz95 <omerpaz95@gmail.com>
Signed-off-by: omerpaz95 <omerpaz95@gmail.com>
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
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
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.
Summary
Adds a remote-transport mechanism on top of the existing
ECCPUConnector: aconsumer (PD) scheduler can fetch an encoder output by
mm_hashfrom a remoteproducer (E) scheduler over ZMQ + NIXL, landing the bytes directly in the
consumer's CPU encoder-cache mmap. Scheduler-to-scheduler control plane,
scheduler-side NIXL data plane; worker processes are untouched.
over ZMQ; one listener thread on the producer; consumer drains acks on demand.
WRITEfrom a producerdict[mm_hash, Tensor]into the consumer's pre-registered
ECSharedRegionmmap._encoding_map(mm_hash → block indices, populated at alloctime) +
_ready(hashes whose xfer completed)._encodings(populated by a futuresave_cachesPR; thisPR's tests fixture-inject) +
_waiting(parked requests).ECSharedRegiongainsalloc(n)/free(indices)so the free-list liveswith the memory it tracks.
messages.pymodule.See
docs/superpowers/specs/2026-05-04-ec-nixl-transfer-design.mdfor thefull design and
docs/superpowers/plans/2026-05-04-ec-nixl-transfer.mdforthe TDD plan.
Out of scope (future PRs)
save_caches(GPU→CPU staging into_encodings).encoder_node_addressontokv_transfer_params— this PR usesec_connector_extra_config.default_encoder_nodeas an MVP fallback.Test plan
tests/v1/ec_connector/unit/test_ec_shared_region_alloc.py— 5/5tests/v1/ec_connector/unit/test_ec_messages.py— 4/4tests/v1/ec_connector/unit/test_ec_consumer_scheduler.py— state machinetests/v1/ec_connector/unit/test_ec_producer_scheduler.py— listener + NIXL xfer + draintests/v1/ec_connector/unit/test_ec_cpu_connector.py— existing worker tests, no regressiontests/v1/ec_connector/integration/test_ec_zmq_loopback.py— end-to-end over real ZMQ, asserts bytes land correctly in consumer mmaptests/v1/ec_connector/: 59/59pre-commit runclean on every commitReal NIXL runtime not exercised here (no NIXL install in dev env);
FakeNixlWrappermirrors theNixlWrappermethod surface and synchronouslyctypes.memmoves bytes between registered regions onWRITE.