[P/D] [NixlConnector] zmq handshake abstraction #26527
Conversation
wseaton
changed the title
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request refactors the NIXL connector by abstracting the ZMQ handshake logic into a separate HandshakeStrategy. This is a good architectural improvement that enhances modularity and prepares for supporting other communication protocols. The changes are well-structured, with new files for lazy imports and the handshake strategy itself. However, I've identified a critical issue where a new method will cause an AttributeError due to accessing a non-existent attribute. I've also noted a potential performance issue in the new handshake logic that appears to perform an unnecessary network operation.
vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector/__init__.py
Outdated
Show resolved
Hide resolved
| # Handshake with remote agent-rank0 first to get the tp_size of remote | ||
| path = make_zmq_path("tcp", host, port) | ||
| logger.debug("Querying master rank metadata on path: %s", path) | ||
| metadata, agent_name_0 = handshake(path, 0) |
There was a problem hiding this comment.
The logic here initiates a handshake with the remote rank 0, but the agent for rank 0 appears to be unused if the target remote rank (p_remote_rank) is not 0. The _read_blocks method only uses the agent corresponding to p_remote_rank. This initial handshake with rank 0 seems to introduce an unnecessary network round trip and agent registration. The previous implementation only connected to the required p_remote_rank. Could you clarify the need for this initial handshake with rank 0 or remove it if it's redundant? The comment "to get the tp_size of remote" is also misleading as remote_tp_size is an argument.
There was a problem hiding this comment.
cc @NickLucche, I think we were saying the handshake through rank 0 was desirable in certain cases?
There was a problem hiding this comment.
Ended up changing to align with current behavior for now
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector/handshake.py
Show resolved
Hide resolved
vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector/__init__.py
Outdated
Show resolved
Hide resolved
wseaton
changed the title
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
wseaton
force-pushed
the
handshake-abstraction
branch
from
9790a03 to
fc48b3f
Compare
|
Need to align the interface with #26338, since they interact |
Signed-off-by: Will Eaton <[email protected]>
Signed-off-by: Will Eaton <[email protected]>
Signed-off-by: Will Eaton <[email protected]>
Signed-off-by: Will Eaton <[email protected]>
Signed-off-by: Will Eaton <[email protected]>
wseaton
force-pushed
the
handshake-abstraction
branch
from
5c3e7f5 to
3251cdd
Compare
|
This pull request has merge conflicts that must be resolved before it can be |
1 participant
Just the ZMQ and base class changes split from: #22274