From 1145805e7d5690246655f546f954a89761254838 Mon Sep 17 00:00:00 2001 From: psaab Date: Sun, 15 Mar 2026 14:29:10 -0700 Subject: [PATCH] Fix socket utilities and reserve_port for IPv6 dual-stack support (#20491) Co-authored-by: hnyls2002 --- python/sglang/srt/utils/common.py | 169 ++++++++++++++------- python/sglang/utils.py | 19 +-- test/registered/utils/test_socket_utils.py | 156 +++++++++++++++++++ 3 files changed, 274 insertions(+), 70 deletions(-) create mode 100644 test/registered/utils/test_socket_utils.py diff --git a/python/sglang/srt/utils/common.py b/python/sglang/srt/utils/common.py index 925552ede..6136442fc 100644 --- a/python/sglang/srt/utils/common.py +++ b/python/sglang/srt/utils/common.py @@ -79,7 +79,6 @@ import psutil import pybase64 import requests import torch -import torch.distributed import torch.distributed as dist import triton import zmq @@ -743,31 +742,84 @@ def wait_port_available( return False +def _get_addrinfos_for_bind(host=None, port=0): + """Return deduplicated addrinfo tuples for binding (one per address family). + + Args: + host: Bind address. None (with AI_PASSIVE) resolves to wildcard + addresses (0.0.0.0 / ::) suitable for accepting on all interfaces. + port: Port number. 0 lets the OS assign an available ephemeral port. + + Flags: + AI_ADDRCONFIG — only return families actually configured on this host. + AI_PASSIVE — return wildcard addresses suitable for bind(). + + Falls back to AF_INET if getaddrinfo fails (e.g. DNS misconfiguration). + """ + try: + infos = socket.getaddrinfo( + host, + port, + socket.AF_UNSPEC, + socket.SOCK_STREAM, + 0, + socket.AI_ADDRCONFIG | socket.AI_PASSIVE, + ) + seen = set() + return [i for i in infos if i[0] not in seen and not seen.add(i[0])] + except socket.gaierror: + fallback_host = "0.0.0.0" if host is None else host + return [(socket.AF_INET, socket.SOCK_STREAM, 0, "", (fallback_host, port))] + + +def try_bind_socket(host=None, port=0, *, reuse_addr=True, listen=False): + """Bind a TCP socket on the first available address family (IPv4/IPv6). + + Iterates over address families returned by _get_addrinfos_for_bind and + returns the first socket that successfully binds. + + Args: + host: Bind address. None binds to all interfaces (0.0.0.0 / ::). + port: Port number. 0 lets the OS assign an available ephemeral port; + use sock.getsockname()[1] to retrieve the assigned port. + reuse_addr: Set SO_REUSEADDR to allow quick port reuse after close. + listen: Call listen(1) after bind, making the socket ready to accept. + + Returns: + The bound socket. Caller is responsible for closing it. + + Raises: + OSError: If bind fails on all configured address families. + """ + for family, socktype, proto, _, sockaddr in _get_addrinfos_for_bind(host, port): + sock = socket.socket(family, socktype, proto) + try: + if reuse_addr: + sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + sock.bind(sockaddr) + if listen: + sock.listen(1) + return sock + except OSError: + sock.close() + raise OSError(f"Could not bind port {port} on any configured address family") + + def is_port_available(port): """Return whether a port is available.""" - with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: - try: - s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) - s.bind(("", port)) - s.listen(1) - return True - except socket.error: - return False - except OverflowError: - return False + try: + sock = try_bind_socket(port=port, listen=True) + sock.close() + return True + except (OSError, OverflowError): + return False def get_free_port(): - # try ipv4 - try: - with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: - s.bind(("", 0)) - return s.getsockname()[1] - except OSError: - # try ipv6 - with socket.socket(socket.AF_INET6, socket.SOCK_STREAM) as s: - s.bind(("", 0)) - return s.getsockname()[1] + sock = try_bind_socket() + port = sock.getsockname()[1] + sock.close() + return port def decode_video_base64(video_base64): @@ -1699,11 +1751,7 @@ def _get_fastapi_request_path(request) -> Tuple[str, bool]: def bind_port(port): """Bind to a specific port, assuming it's available.""" - sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) - sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) # Allows address reuse - sock.bind(("", port)) - sock.listen(1) - return sock + return try_bind_socket(port=port, listen=True) def get_amdgpu_memory_capacity(): @@ -2648,22 +2696,16 @@ def get_open_port() -> int: port = int(port) while True: try: - with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: - s.bind(("", port)) - return port + sock = try_bind_socket(port=port, reuse_addr=False) + sock.close() + return port except OSError: - port += 1 # Increment port number if already in use - logger.info("Port %d is already in use, trying port %d", port - 1, port) - # try ipv4 - try: - with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: - s.bind(("", 0)) - return s.getsockname()[1] - except OSError: - # try ipv6 - with socket.socket(socket.AF_INET6, socket.SOCK_STREAM) as s: - s.bind(("", 0)) - return s.getsockname()[1] + logger.info("Port %d is already in use, trying port %d", port, port + 1) + port += 1 + sock = try_bind_socket() + port = sock.getsockname()[1] + sock.close() + return port def is_valid_ipv6_address(address: str) -> bool: @@ -2934,31 +2976,40 @@ def get_local_ip_by_nic(interface: str = None) -> Optional[str]: def get_local_ip_by_remote() -> Optional[str]: - # try ipv4 - s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) - try: - s.connect(("8.8.8.8", 80)) # Doesn't need to be reachable - return s.getsockname()[0] - except Exception: - pass + # Google's public DNS servers, used to discover the local IP. + # UDP connect doesn't send packets; it just selects the right source address. + # https://developers.google.com/speed/public-dns/docs/using#addresses + # Try IPv4 first, then IPv6. getaddrinfo on a literal IP returns exactly + # one result, so we unpack directly instead of looping. + for dns_host, dns_port in [("8.8.8.8", 80), ("2001:4860:4860::8888", 80)]: + try: + family, socktype, proto, _, sockaddr = socket.getaddrinfo( + dns_host, + dns_port, + socket.AF_UNSPEC, + socket.SOCK_DGRAM, + 0, + socket.AI_ADDRCONFIG, + )[0] + with socket.socket(family, socktype, proto) as s: + s.connect(sockaddr) + return s.getsockname()[0] + except (socket.gaierror, OSError): + continue + # Fallback: resolve the local hostname to an IP address via /etc/hosts or DNS. + # Unreliable — many machines resolve hostname to 127.0.0.1, so we skip loopback. try: hostname = socket.gethostname() - ip = socket.gethostbyname(hostname) - if ip and ip != "127.0.0.1" and ip != "0.0.0.0": + ip = socket.getaddrinfo( + hostname, None, socket.AF_UNSPEC, 0, 0, socket.AI_ADDRCONFIG + )[0][4][0] + if ip and ip not in ("127.0.0.1", "0.0.0.0", "::1"): return ip except Exception: pass - # try ipv6 - try: - s = socket.socket(socket.AF_INET6, socket.SOCK_DGRAM) - # Google's public DNS server, see - # https://developers.google.com/speed/public-dns/docs/using#addresses - s.connect(("2001:4860:4860::8888", 80)) # Doesn't need to be reachable - return s.getsockname()[0] - except Exception: - logger.warning("Can not get local ip by remote") + logger.warning("Can not get local ip by remote") return None diff --git a/python/sglang/utils.py b/python/sglang/utils.py index 1a6187d91..c0fa8c06f 100644 --- a/python/sglang/utils.py +++ b/python/sglang/utils.py @@ -5,7 +5,6 @@ import json import logging import os import random -import socket import ssl import subprocess import sys @@ -124,6 +123,8 @@ def dump_state_text(filename: str, states: list, mode: str = "w"): def normalize_base_url(host: str, port: int) -> str: + from sglang.srt.utils.network import NetworkAddress + if host.startswith("http://") or host.startswith("https://"): warnings.warn( f"Including the scheme in --host ('{host}') is deprecated. " @@ -131,9 +132,8 @@ def normalize_base_url(host: str, port: int) -> str: DeprecationWarning, stacklevel=2, ) - else: - host = f"http://{host}" - return f"{host}:{port}" + return f"{host}:{port}" + return NetworkAddress(host, port).to_url() class HttpResponse: @@ -401,18 +401,15 @@ def reserve_port(host, start=30000, end=40000): Reserve an available port by trying to bind a socket. Returns a tuple (port, lock_socket) where `lock_socket` is kept open to hold the lock. """ + from sglang.srt.utils.common import try_bind_socket + candidates = list(range(start, end)) random.shuffle(candidates) - for port in candidates: - sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) - sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) try: - # Attempt to bind to the port on localhost - sock.bind((host, port)) + sock = try_bind_socket(host, port) return port, sock - except socket.error: - sock.close() # Failed to bind, try next port + except OSError: continue raise RuntimeError("No free port available.") diff --git a/test/registered/utils/test_socket_utils.py b/test/registered/utils/test_socket_utils.py new file mode 100644 index 000000000..62aa0f603 --- /dev/null +++ b/test/registered/utils/test_socket_utils.py @@ -0,0 +1,156 @@ +import os +import socket +import unittest +from unittest.mock import patch + +from sglang.srt.utils.common import ( + _get_addrinfos_for_bind, + bind_port, + get_free_port, + get_open_port, + is_port_available, + try_bind_socket, +) +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=1, suite="stage-a-cpu-only") + + +class TestTryBindSocket(CustomTestCase): + def test_bind_ephemeral_port(self): + """try_bind_socket() with port=0 should bind to an OS-assigned port.""" + sock = try_bind_socket() + try: + port = sock.getsockname()[1] + self.assertGreater(port, 0) + self.assertLessEqual(port, 65535) + finally: + sock.close() + + def test_bind_specific_port(self): + """try_bind_socket(port=N) should bind to that exact port.""" + port = get_free_port() + sock = try_bind_socket(port=port) + try: + self.assertEqual(sock.getsockname()[1], port) + finally: + sock.close() + + def test_bind_with_listen(self): + """try_bind_socket(listen=True) should return a listening socket.""" + sock = try_bind_socket(listen=True) + try: + # A listening socket has a valid bound address + port = sock.getsockname()[1] + self.assertGreater(port, 0) + finally: + sock.close() + + def test_bind_with_host(self): + """try_bind_socket(host='127.0.0.1') should bind to localhost.""" + sock = try_bind_socket(host="127.0.0.1") + try: + addr = sock.getsockname() + self.assertEqual(addr[0], "127.0.0.1") + finally: + sock.close() + + def test_bind_occupied_port_raises(self): + """try_bind_socket should raise OSError if port is occupied.""" + sock1 = try_bind_socket() + try: + port = sock1.getsockname()[1] + with self.assertRaises(OSError): + try_bind_socket(port=port, reuse_addr=False) + finally: + sock1.close() + + def test_returns_correct_family(self): + """Returned socket should be AF_INET or AF_INET6.""" + sock = try_bind_socket() + try: + self.assertIn(sock.family, (socket.AF_INET, socket.AF_INET6)) + finally: + sock.close() + + def test_gaierror_fallback(self): + """_get_addrinfos_for_bind should fall back to AF_INET on gaierror.""" + with patch( + "sglang.srt.utils.common.socket.getaddrinfo", + side_effect=socket.gaierror("mocked"), + ): + infos = _get_addrinfos_for_bind() + self.assertEqual(len(infos), 1) + family, socktype, _, _, sockaddr = infos[0] + self.assertEqual(family, socket.AF_INET) + self.assertEqual(sockaddr[0], "0.0.0.0") + + def test_gaierror_fallback_preserves_host(self): + """Fallback should use the provided host, not default to 0.0.0.0.""" + with patch( + "sglang.srt.utils.common.socket.getaddrinfo", + side_effect=socket.gaierror("mocked"), + ): + infos = _get_addrinfos_for_bind(host="10.0.0.1", port=8080) + self.assertEqual(infos[0][4], ("10.0.0.1", 8080)) + + +class TestSocketUtilities(CustomTestCase): + def test_is_port_available(self): + """is_port_available should return True for a free port.""" + port = get_free_port() + self.assertTrue(is_port_available(port)) + + def test_is_port_available_occupied(self): + """is_port_available should return False for an occupied port.""" + sock = bind_port(get_free_port()) + try: + port = sock.getsockname()[1] + self.assertFalse(is_port_available(port)) + finally: + sock.close() + + def test_get_free_port(self): + """get_free_port should return a valid port number.""" + port = get_free_port() + self.assertGreater(port, 0) + self.assertLessEqual(port, 65535) + + def test_bind_port(self): + """bind_port should return a listening socket.""" + port = get_free_port() + sock = bind_port(port) + try: + self.assertEqual(sock.getsockname()[1], port) + finally: + sock.close() + + def test_get_open_port(self): + """get_open_port should return a valid port number.""" + port = get_open_port() + self.assertGreater(port, 0) + self.assertLessEqual(port, 65535) + + def test_get_open_port_with_env_var(self): + """get_open_port should respect SGLANG_PORT env var.""" + free_port = get_free_port() + with patch.dict(os.environ, {"SGLANG_PORT": str(free_port)}): + port = get_open_port() + self.assertEqual(port, free_port) + + def test_get_open_port_env_var_occupied_increments(self): + """get_open_port should increment if SGLANG_PORT is occupied.""" + sock = bind_port(get_free_port()) + try: + occupied_port = sock.getsockname()[1] + with patch.dict(os.environ, {"SGLANG_PORT": str(occupied_port)}): + port = get_open_port() + # Should skip the occupied port and return a higher one + self.assertGreater(port, occupied_port) + finally: + sock.close() + + +if __name__ == "__main__": + unittest.main()