diffusion: reduce effort of supporting new model (#12982)

python/sglang/multimodal_gen/runtime/utils/common.py

@@ -126,8 +126,26 @@ def is_port_available(port):
 
 
 def get_zmq_socket(
-    context: zmq.Context, socket_type: zmq.SocketType, endpoint: str, bind: bool
-) -> zmq.Socket:
+    context: zmq.Context,
+    socket_type: zmq.SocketType,
+    endpoint: str,
+    bind: bool,
+    max_bind_retries: int = 10,
+) -> tuple[zmq.Socket, str]:
+    """
+    Create and configure a ZMQ socket.
+
+    Args:
+        context: ZMQ context
+        socket_type: Type of ZMQ socket
+        endpoint: Endpoint string (e.g., "tcp://localhost:5555")
+        bind: Whether to bind (True) or connect (False)
+        max_bind_retries: Maximum number of retries if bind fails due to address already in use
+
+    Returns:
+        A tuple of (socket, actual_endpoint). The actual_endpoint may differ from the
+        requested endpoint if bind retry was needed.
+    """
     mem = psutil.virtual_memory()
     total_mem = mem.total / 1024**3
     available_mem = mem.available / 1024**3
@@ -165,11 +183,67 @@ def get_zmq_socket(
         raise ValueError(f"Unsupported socket type: {socket_type}")
 
     if bind:
-        socket.bind(endpoint)
+        # Parse port from endpoint for retry logic
+        import re
+
+        port_match = re.search(r":(\d+)$", endpoint)
+
+        if port_match and max_bind_retries > 1:
+            original_port = int(port_match.group(1))
+            last_exception = None
+
+            for attempt in range(max_bind_retries):
+                try:
+                    current_endpoint = endpoint
+                    if attempt > 0:
+                        # Try next port (increment by 42 to match settle_port logic)
+                        current_port = original_port + attempt * 42
+                        current_endpoint = re.sub(
+                            r":(\d+)$", f":{current_port}", endpoint
+                        )
+                        logger.info(
+                            f"ZMQ bind failed for port {original_port + (attempt - 1) * 42}, "
+                            f"retrying with port {current_port} (attempt {attempt + 1}/{max_bind_retries})"
+                        )
+
+                    socket.bind(current_endpoint)
+
+                    if attempt > 0:
+                        logger.warning(
+                            f"Successfully bound ZMQ socket to {current_endpoint} after {attempt + 1} attempts. "
+                            f"Original port {original_port} was unavailable."
+                        )
+
+                    return socket, current_endpoint
+
+                except zmq.ZMQError as e:
+                    last_exception = e
+                    if e.errno == zmq.EADDRINUSE and attempt < max_bind_retries - 1:
+                        # Address already in use, try next port
+                        continue
+                    elif attempt == max_bind_retries - 1:
+                        # Last attempt failed
+                        logger.error(
+                            f"Failed to bind ZMQ socket after {max_bind_retries} attempts. "
+                            f"Original endpoint: {endpoint}, Last tried port: {original_port + attempt * 42}"
+                        )
+                        raise
+                    else:
+                        # Different error, raise immediately
+                        raise
+
+            # Should not reach here, but just in case
+            if last_exception:
+                raise last_exception
+        else:
+            # No retry logic needed (either no port in endpoint or max_bind_retries == 1)
+            socket.bind(endpoint)
+            return socket, endpoint
     else:
         socket.connect(endpoint)
+        return socket, endpoint
 
-    return socket
+    return socket, endpoint
 
 
 # https://pytorch.org/docs/stable/notes/hip.html#checking-for-hip

python/sglang/multimodal_gen/runtime/utils/logging_utils.py

@@ -34,7 +34,7 @@ _warned_main_process = False
 
 _FORMAT = (
     f"{SGL_DIFFUSION_LOGGING_PREFIX}%(levelname)s %(asctime)s "
-    "[%(filename)s:%(lineno)d] %(message)s"
+    "[%(filename)s: %(lineno)d] %(message)s"
 )
 
 # _FORMAT = "[%(asctime)s] %(message)s"