distributed: bind only to 127.0.0.1 for local-only usage (#1061)

PygmalionAI · Jan 1, 2025 · f81e7d7 · f81e7d7
1 parent 58aff37
commit f81e7d7
Showing 1 changed file with 11 additions and 6 deletions.
diff --git a/aphrodite/distributed/device_communicators/shm_broadcast.py b/aphrodite/distributed/device_communicators/shm_broadcast.py
@@ -194,7 +194,9 @@ def __init__(
             # see http://api.zeromq.org/3-3:zmq-setsockopt for more details
             self.local_socket.setsockopt(XPUB_VERBOSE, True)
             local_subscribe_port = get_open_port()
-            self.local_socket.bind(f"tcp://*:{local_subscribe_port}")
+            socket_addr = f"tcp://127.0.0.1:{local_subscribe_port}"
+            logger.debug(f"Binding to {socket_addr}")
+            self.local_socket.bind(socket_addr)
 
             self.current_idx = 0
 
@@ -210,7 +212,8 @@ def __init__(
             self.remote_socket = context.socket(XPUB)
             self.remote_socket.setsockopt(XPUB_VERBOSE, True)
             remote_subscribe_port = get_open_port()
-            self.remote_socket.bind(f"tcp://*:{remote_subscribe_port}")
+            socket_addr = f"tcp://*:{remote_subscribe_port}"
+            self.remote_socket.bind(socket_addr)
 
         else:
             remote_subscribe_port = None
@@ -254,8 +257,9 @@ def create_from_handle(handle: Handle, rank) -> "MessageQueue":
 
             self.local_socket = context.socket(SUB)
             self.local_socket.setsockopt_string(SUBSCRIBE, "")
-            self.local_socket.connect(
-                f"tcp://{handle.connect_ip}:{handle.local_subscribe_port}")
+            socket_addr = f"tcp://127.0.0.1:{handle.local_subscribe_port}"
+            logger.debug(f"Connecting to {socket_addr}")
+            self.local_socket.connect(socket_addr)
 
             self.remote_socket = None
         else:
@@ -269,8 +273,9 @@ def create_from_handle(handle: Handle, rank) -> "MessageQueue":
 
             self.remote_socket = context.socket(SUB)
             self.remote_socket.setsockopt_string(SUBSCRIBE, "")
-            self.remote_socket.connect(
-                f"tcp://{handle.connect_ip}:{handle.remote_subscribe_port}")
+            socket_addr = f"tcp://{handle.connect_ip}:{handle.remote_subscribe_port}"
+            logger.debug(f"Connecting to {socket_addr}")
+            self.remote_socket.connect(socket_addr)
 
         return self