Repository navigation
Fix potential race condition in shm_broadcast.py and unify socket cre… - #15903
Baidu-AIAK wants to merge 2 commits into
Conversation
Summary of ChangesHello @Baidu-AIAK, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request resolves a potential race condition in the distributed communication layer by enhancing the port binding mechanism. It transitions from a non-atomic Highlights
🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console. Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
![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 aims to fix a potential race condition by refactoring socket creation to use ZeroMQ's atomic bind_to_random_port method. The approach is sound, but the implementation has a critical bug in python/sglang/srt/utils/common.py where a missing f-string would cause bind_to_random_port to fail. Additionally, the handling of IPv6 addresses is incorrect, as the ZMQ_IPV6 socket option is set after the socket is bound. I've provided a suggestion in common.py to fix both issues by also moving the IPv6 logic into the get_zmq_socket helper function. This also allows for simplifying the code in shm_broadcast.py by removing the now-redundant IPv6 check.
| config_socket(socket, socket_type) | ||
| port = socket.bind_to_random_port("tcp://*") | ||
| port = socket.bind_to_random_port("tcp://{host}") | ||
| return port, socket |
There was a problem hiding this comment.
This block contains a critical bug and a logic error. First, the call to bind_to_random_port is missing an f for the f-string, so it will try to bind to the literal address tcp://{host}. Second, the ZMQ_IPV6 socket option must be set before binding for IPv6 hosts. The calling code in shm_broadcast.py attempts this after binding, which is incorrect. Let's fix both issues here by moving the IPv6 check into this function and correcting the f-string.
| config_socket(socket, socket_type) | |
| port = socket.bind_to_random_port("tcp://*") | |
| port = socket.bind_to_random_port("tcp://{host}") | |
| return port, socket | |
| config_socket(socket, socket_type) | |
| if is_valid_ipv6_address(host): | |
| socket.setsockopt(zmq.IPV6, 1) | |
| port = socket.bind_to_random_port(f"tcp://{host}") | |
| return port, socket |
There was a problem hiding this comment.
I've fixed the missing f-string prefix. I also followed your suggestion to move the IPv6 configuration logic inside get_zmq_socket to ensure ZMQ_IPV6 is set before binding. The changes are pushed.
| if is_valid_ipv6_address(connect_ip): | ||
| self.remote_socket.setsockopt(IPV6, 1) |
There was a problem hiding this comment.
This IPv6 check is performed after the socket is already bound within get_zmq_socket, which is too late for the ZMQ_IPV6 option to take effect. With the suggested change in common.py to handle IPv6 configuration inside get_zmq_socket, this block becomes both incorrect and redundant and can be removed.
3c1ea64 to
2cdfc80
Compare
|
Hi maintainers, |
|
/gemini review |
|
Warning You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again! |
|
/tag-and-rerun-ci |
1994949 to
02183b6
Compare
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request effectively addresses a potential race condition in shm_broadcast.py by refactoring the socket creation to use ZeroMQ's atomic bind_to_random_port method. The unification of get_zmq_socket and the removal of get_zmq_socket_on_host improve code clarity and maintainability. The changes also enhance IPv6 support and add better error handling. Overall, this is a solid improvement.
I've found one critical issue where the refactored get_zmq_socket function is being used incorrectly in python/sglang/srt/disaggregation/common/conn.py, which would lead to failures when using IPv6. I've provided a comment with details on how to resolve this.
| self.rank_port, self.server_socket = get_zmq_socket( | ||
| context, zmq.PULL, host=zmq_bind_host | ||
| ) |
There was a problem hiding this comment.
While this change correctly replaces the deprecated function, it exposes a bug due to how the host parameter is prepared on the preceding line. The new get_zmq_socket function expects a raw IP address for the host parameter and handles IPv6 address wrapping internally. However, on line 75, maybe_wrap_ipv6_address is called, which passes a pre-wrapped address (e.g., [::1]) to get_zmq_socket. This will cause is_valid_ipv6_address inside get_zmq_socket to fail and prevent the ZMQ_IPV6 socket option from being set correctly, leading to errors when binding to an IPv6 address.
To fix this, you should remove the call to maybe_wrap_ipv6_address on line 75:
zmq_bind_host = self.local_ipYou may also want to adjust the logging on line 79 to ensure the address is displayed correctly for IPv6.
There was a problem hiding this comment.
Good catch. I didn't realize get_zmq_socket handles IPv6 wrapping internally. I've removed the call to maybe_wrap_ipv6_address on line 75 to pass the raw IP. I also updated the logging statement on line 79 as suggested.
2ebaa92 to
dea9d77
Compare
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request effectively addresses a potential race condition in shm_broadcast.py by replacing the non-atomic get_open_port() and socket.bind() sequence with a unified get_zmq_socket helper that leverages ZeroMQ's atomic bind_to_random_port. The refactoring of get_zmq_socket is well-executed, improving IPv6 handling and adding robust error handling with resource cleanup. The consolidation of get_zmq_socket_on_host into the main helper function is a good application of the DRY principle, simplifying the codebase. Overall, these changes significantly enhance the robustness and maintainability of the socket creation logic. I have a couple of minor suggestions to further improve the new utility function.
| local_subscribe_port, self.local_socket = get_zmq_socket( | ||
| context, XPUB, endpoint=None, host="127.0.0.1" | ||
| ) | ||
| # set the verbose option so that we can receive every subscription | ||
| # message. otherwise, we will only receive the first subscription | ||
| # see http://api.zeromq.org/3-3:zmq-setsockopt for more details | ||
| self.local_socket.setsockopt(XPUB_VERBOSE, True) |
There was a problem hiding this comment.
I see that you're setting XPUB_VERBOSE on the socket after it's created by get_zmq_socket. This pattern is repeated for the remote socket as well. To improve encapsulation, you could consider extending get_zmq_socket to accept an optional dictionary of socket options. This would centralize socket configuration and make its usage cleaner. For example:
# In common.py
def get_zmq_socket(..., sockopts: Optional[Dict[int, Any]] = None):
# ...
if sockopts:
for opt, val in sockopts.items():
socket.setsockopt(opt, val)
# ...
# In shm_broadcast.py
local_subscribe_port, self.local_socket = get_zmq_socket(
context, XPUB, endpoint=None, host="127.0.0.1",
sockopts={XPUB_VERBOSE: True}
)This is a suggestion for a future improvement and not critical for this PR.
| @@ -1345,54 +1346,49 @@ def get_zmq_socket( | |||
| socket_type: Type of ZeroMQ socket to create. | |||
| endpoint: Optional endpoint to bind/connect to. If None, binds to a random TCP port. | |||
| bind: Whether to bind (True) or connect (False) to the endpoint. Ignored if endpoint is None. | |||
| host: Host to bind/connect to, without "tcp://" prefix. | |||
There was a problem hiding this comment.
The docstring for the host parameter is slightly misleading. It's only used for binding to a random port when endpoint is None, not for connecting. Clarifying this would improve readability and prevent potential misuse.
| host: Host to bind/connect to, without "tcp://" prefix. | |
| host: Host to bind to when `endpoint` is None. Defaults to "*" (all interfaces). |
cc2cd18 to
b00cf9d
Compare
|
/rerun-and-failed-ci |
fc7cf25 to
bb3edda
Compare
|
f00027c to
b5695e6
Compare
…ation Change-Id: I6f7037e98b04f06e23e20a7b817c5e2e9621a885
b5695e6 to
004f0c7
Compare
This branch has not been deployed
4 participants
Motivation
Fixes #15885
This PR addresses a potential race condition (TOCTOU) in
python/sglang/srt/distributed/device_communicators/shm_broadcast.py.Currently, the code uses get_open_port() to find an available port and subsequently binds to it. This approach is not atomic: another process could claim the port between the time it is returned by get_open_port() and the time socket.bind() is called.
To resolve this, I have refactored the code to use the existing get_zmq_socket helper function, which leverages ZeroMQ's native bind_to_random_port method. This method is atomic and safer.
Additionally, to maintain the original behavior of binding strictly to localhost (127.0.0.1) rather than all interfaces (0.0.0.0), I extended the get_zmq_socket helper to support a customizable host parameter.
Modifications
Refactored get_zmq_socket to fix critical logic errors:
Removed the redundant get_zmq_socket_on_host function. Its functionality is now fully covered by get_zmq_socket (following DRY principle).
Accuracy Tests
This is a refactoring of the networking/infrastructure layer (socket initialization). It does not affect the model kernel, forward pass, or logical output. Therefore, accuracy remains unchanged.
Benchmarking and Profiling
This change only affects the initialization phase of the router. It does not impact the inference critical path or runtime throughput. No performance regression is expected.
Checklist