fix(atom): relay KV events from data-parallel ranks - #121
Draft
cquil11 wants to merge 2 commits into
Draft
Conversation
Signed-off-by: Cam Quilici <cjquilici@gmail.com>
Signed-off-by: Cam Quilici <cjquilici@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
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\n\n- relay KV events from every ATOM data-parallel EngineCore through one worker-local XSUB/XPUB proxy\n- preserve the one logical KV-event endpoint advertised to Infera for each ATOM worker\n- keep direct-bind behavior available for existing standalone hook callers\n- stop the relay with the worker and retain the existing ATOM process-group cleanup behavior\n\n## Root cause\n\nAn ATOM TP+DP worker creates one BlockManager in every data-parallel EngineCore process. The current hook gives every process the same externally advertised ZMQ bind endpoint. All DP ranks therefore race to bind one port; one wins and the remaining ranks fail during startup with zmq.error.ZMQError: Address already in use.\n\nThis reproduced simultaneously on all three workers of the MI355X DeepSeek-V4 2P/1D DPA validation in InferenceX run 31995111191, Slurm 39443. Every rank completed model load and CUDA-graph capture before failing at BlockManager.init -> publisher.ensure_bound() on the shared endpoint. The Infera health registry consequently remained at 0/2 prefill and 0/1 decode.\n\nThe fix binds one proxy endpoint in the Infera worker process and gives every spawned EngineCore a local IPC ingress to connect its PUB socket. ZMQ forwards all rank streams through the single advertised endpoint, preserving full DP cache visibility rather than publishing only rank zero or silently disabling KV-aware routing.\n\n## Validation\n\n- pytest -q tests/unit/engine tests/unit/router/test_kv_event_client.py: 50 passed\n- focused relay regression uses two independent EngineCore publishers and observes both streams through one advertised endpoint\n- worker wiring test verifies the child environment uses connect mode and the proxy is stopped during teardown\n- Ruff format/check: passed\n- applicable pre-commit hooks: passed\n- git diff --check: passed\n\nHardware revalidation will update the companion InferenceX PR after the current exact-head attempt finishes naturally. No running job was cancelled.