chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

[Utils] Add NetworkAddress abstraction for IPv6-safe address handling (#20306)

Browse Source
This commit is contained in:
Liangsheng Yin
2026-03-11 00:07:37 -07:00
committed by GitHub
parent 55e6acf834
commit 61cad15d28
12 changed files with 539 additions and 270 deletions
Download Patch File Download Diff File Expand all files Collapse all files

30
python/sglang/srt/disaggregation/common/conn.py
View File

@@ -3,7 +3,6 @@ from __future__ import annotations
import asyncio
import dataclasses
import logging
import socket
import threading
import time
from collections import defaultdict
@@ -37,12 +36,10 @@ from sglang.srt.layers.dp_attention import (
)
from sglang.srt.server_args import ServerArgs
from sglang.srt.utils import (
format_tcp_address,
get_local_ip_auto,
get_zmq_socket_on_host,
is_valid_ipv6_address,
maybe_wrap_ipv6_address,
)
from sglang.srt.utils.network import NetworkAddress
logger = logging.getLogger(__name__)
@@ -116,11 +113,10 @@ class CommonKVManager(BaseKVManager):
# bind zmq socket
context = zmq.Context()
zmq_bind_host = maybe_wrap_ipv6_address(self.local_ip)
self.rank_port, self.server_socket = get_zmq_socket_on_host(
context, zmq.PULL, host=zmq_bind_host
context, zmq.PULL, host=self.local_ip
)
logger.debug(f"kv manager bind to {zmq_bind_host}:{self.rank_port}")
logger.debug(f"kv manager bind to {self.local_ip}:{self.rank_port}")
self.request_status: Dict[int, KVPoll] = {}
self.failure_records: Dict[int, str] = {}
@@ -255,20 +251,14 @@ class CommonKVManager(BaseKVManager):
"""Register prefill server info to bootstrap server via HTTP POST."""
if self.dist_init_addr:
# Multi-node case: bootstrap server's host is dist_init_addr
if self.dist_init_addr.startswith("["): # [ipv6]:port or [ipv6]
if self.dist_init_addr.endswith("]"):
host = self.dist_init_addr
else:
host, _ = self.dist_init_addr.rsplit(":", 1)
else:
host = socket.gethostbyname(self.dist_init_addr.rsplit(":", 1)[0])
host = NetworkAddress.parse(self.dist_init_addr).host
else:
# Single-node case: bootstrap server's host is the same as http server's host
host = self.bootstrap_host
host = maybe_wrap_ipv6_address(host)
bootstrap_server_url = f"{host}:{self.bootstrap_port}"
url = f"http://{bootstrap_server_url}/route"
bootstrap_na = NetworkAddress(host, self.bootstrap_port)
bootstrap_server_url = bootstrap_na.to_host_port_str()
url = f"{bootstrap_na.to_url()}/route"
payload = {
"attn_tp_size": self.attn_tp_size,
"attn_tp_rank": self.attn_tp_rank,
@@ -656,10 +646,8 @@ class CommonKVReceiver(BaseKVReceiver):
def _connect_to_bootstrap_server(cls, bootstrap_info: dict):
ip_address = bootstrap_info["rank_ip"]
port = bootstrap_info["rank_port"]
is_ipv6_address = is_valid_ipv6_address(ip_address)
sock, lock = cls._connect(
format_tcp_address(ip_address, port), is_ipv6=is_ipv6_address
)
na = NetworkAddress(ip_address, port)
sock, lock = cls._connect(na.to_tcp(), is_ipv6=na.is_ipv6)
return sock, lock
def _register_kv_args(self):

12
python/sglang/srt/disaggregation/mooncake/conn.py
View File

@@ -35,7 +35,7 @@ from sglang.srt.disaggregation.utils import (
from sglang.srt.distributed.parallel_state import get_mooncake_transfer_engine
from sglang.srt.environ import envs
from sglang.srt.server_args import ServerArgs
from sglang.srt.utils import format_tcp_address, is_valid_ipv6_address
from sglang.srt.utils.network import NetworkAddress
logger = logging.getLogger(__name__)
@@ -552,9 +552,8 @@ class MooncakeKVManager(CommonKVManager):
aux_index: int,
data: bytes,
):
socket = self._connect(
format_tcp_address(remote, dst_port), is_ipv6=is_valid_ipv6_address(remote)
)
na = NetworkAddress(remote, dst_port)
socket = self._connect(na.to_tcp(), is_ipv6=na.is_ipv6)
socket.send_multipart(
[
@@ -759,9 +758,8 @@ class MooncakeKVManager(CommonKVManager):
def sync_status_to_decode_endpoint(
self, remote: str, dst_port: int, room: int, status: int, prefill_rank: int
):
self._connect(
format_tcp_address(remote, dst_port), is_ipv6=is_valid_ipv6_address(remote)
).send_multipart(
na = NetworkAddress(remote, dst_port)
self._connect(na.to_tcp(), is_ipv6=na.is_ipv6).send_multipart(
[
str(room).encode("ascii"),
str(status).encode("ascii"),

14
python/sglang/srt/disaggregation/mori/conn.py
View File

@@ -38,11 +38,10 @@ from sglang.srt.disaggregation.utils import (
)
from sglang.srt.server_args import ServerArgs
from sglang.srt.utils.common import (
format_tcp_address,
get_int_env_var,
get_local_ip_auto,
is_valid_ipv6_address,
)
from sglang.srt.utils.network import NetworkAddress
logger = logging.getLogger(__name__)
MORI_GUARD = b"MoriMsgGuard"
@@ -417,10 +416,8 @@ class MoriKVManager(CommonKVManager):
]
for info in infos:
try:
endpoint = format_tcp_address(info.endpoint, info.dst_port)
socket = self._connect(
endpoint, is_ipv6=is_valid_ipv6_address(info.endpoint)
)
na = NetworkAddress(info.endpoint, info.dst_port)
socket = self._connect(na.to_tcp(), is_ipv6=na.is_ipv6)
socket.send_multipart(payload)
except Exception:
logger.exception(
@@ -754,9 +751,8 @@ class MoriKVManager(CommonKVManager):
aux_index: int,
data: bytes,
):
socket = self._connect(
format_tcp_address(remote, dst_port), is_ipv6=is_valid_ipv6_address(remote)
)
na = NetworkAddress(remote, dst_port)
socket = self._connect(na.to_tcp(), is_ipv6=na.is_ipv6)
socket.send_multipart(
[

9
python/sglang/srt/distributed/device_communicators/mooncake_transfer_engine.py
View File

@@ -4,7 +4,8 @@ import os
from typing import List, Optional
from sglang.srt.environ import envs
from sglang.srt.utils import get_free_port, maybe_wrap_ipv6_address
from sglang.srt.utils import get_free_port
from sglang.srt.utils.network import NetworkAddress
logger = logging.getLogger(__name__)
@@ -117,9 +118,9 @@ class MooncakeTransferEngine:
hostname=self.hostname,
device_name=self.ib_device,
)
self.session_id = (
f"{maybe_wrap_ipv6_address(self.hostname)}:{self.engine.get_rpc_port()}"
)
self.session_id = NetworkAddress(
self.hostname, self.engine.get_rpc_port()
).to_host_port_str()
def register(self, ptr, length):
try:

15
python/sglang/srt/distributed/device_communicators/shm_broadcast.py
View File

@@ -17,11 +17,10 @@ from zmq import IPV6 # type: ignore
from zmq import SUB, SUBSCRIBE, XPUB, XPUB_VERBOSE, Context # type: ignore
from sglang.srt.utils import (
format_tcp_address,
get_local_ip_auto,
get_open_port,
is_valid_ipv6_address,
)
from sglang.srt.utils.network import NetworkAddress
# SGLANG_RINGBUFFER_WARNING_INTERVAL can be set to 60
SGLANG_RINGBUFFER_WARNING_INTERVAL = int(
@@ -229,9 +228,10 @@ class MessageQueue:
self.remote_socket = context.socket(XPUB)
self.remote_socket.setsockopt(XPUB_VERBOSE, True)
remote_subscribe_port = get_open_port()
if is_valid_ipv6_address(connect_ip):
na = NetworkAddress(connect_ip, remote_subscribe_port)
if na.is_ipv6:
self.remote_socket.setsockopt(IPV6, 1)
address = format_tcp_address(connect_ip, remote_subscribe_port)
address = na.to_tcp()
logger.debug(f"class MessageQueue: Binding remote socket to {address=}")
self.remote_socket.bind(address)
@@ -292,11 +292,10 @@ class MessageQueue:
self.remote_socket = context.socket(SUB)
self.remote_socket.setsockopt_string(SUBSCRIBE, "")
if is_valid_ipv6_address(handle.connect_ip):
na = NetworkAddress(handle.connect_ip, handle.remote_subscribe_port)
if na.is_ipv6:
self.remote_socket.setsockopt(IPV6, 1)
socket_addr = format_tcp_address(
handle.connect_ip, handle.remote_subscribe_port
)
socket_addr = na.to_tcp()
logger.debug("Connecting to %s", socket_addr)
self.remote_socket.connect(socket_addr)

14
python/sglang/srt/managers/data_parallel_controller.py
View File

@@ -48,12 +48,12 @@ from sglang.srt.server_args import (
from sglang.srt.utils import numa_utils
from sglang.srt.utils.common import (
bind_port,
configure_ipv6,
configure_logger,
get_zmq_socket,
kill_itself_when_parent_died,
maybe_reindex_device_id,
)
from sglang.srt.utils.network import NetworkAddress
from sglang.srt.utils.torch_memory_saver_adapter import TorchMemorySaverAdapter
from sglang.srt.utils.watchdog import Watchdog
from sglang.utils import TypeBasedDispatcher, get_exception_traceback
@@ -289,13 +289,13 @@ class DataParallelController:
"""
# Determine the endpoint for inter-node communication
if server_args.dist_init_addr is None:
endpoint = f"tcp://127.0.0.1:{server_args.port + DP_ATTENTION_HANDSHAKE_PORT_DELTA}"
elif server_args.dist_init_addr.startswith("["): # ipv6 address
port, host = configure_ipv6(server_args.dist_init_addr)
endpoint = f"tcp://{host}:{int(port) + DP_ATTENTION_HANDSHAKE_PORT_DELTA}"
na = NetworkAddress(
"127.0.0.1", server_args.port + DP_ATTENTION_HANDSHAKE_PORT_DELTA
)
else:
host, port = server_args.dist_init_addr.split(":")
endpoint = f"tcp://{host}:{int(port) + DP_ATTENTION_HANDSHAKE_PORT_DELTA}"
na = NetworkAddress.parse(server_args.dist_init_addr)
na = NetworkAddress(na.host, na.port + DP_ATTENTION_HANDSHAKE_PORT_DELTA)
endpoint = na.to_tcp()
if server_args.node_rank == 0:
# Node 0: Broadcast worker ports to all other nodes

41
python/sglang/srt/server_args.py
View File

@@ -35,7 +35,6 @@ from sglang.srt.parser.reasoning_parser import ReasoningParser
from sglang.srt.utils.common import (
LORA_TARGET_ALL_MODULES,
SUPPORTED_LORA_TARGET_MODULES,
configure_ipv6,
cpu_has_amx_support,
get_bool_env_var,
get_device,
@@ -60,7 +59,6 @@ from sglang.srt.utils.common import (
is_sm100_supported,
is_sm120_supported,
is_triton_kernels_available,
is_valid_ipv6_address,
json_list_type,
nullable_str,
parse_connector_type,
@@ -69,6 +67,7 @@ from sglang.srt.utils.common import (
xpu_has_xmx_support,
)
from sglang.srt.utils.hf_transformers_utils import check_gguf_file
from sglang.srt.utils.network import NetworkAddress
from sglang.utils import is_in_ci
logger = logging.getLogger(__name__)
@@ -5582,10 +5581,7 @@ class ServerArgs:
host = "127.0.0.1"
elif host == "::":
host = "::1"
if is_valid_ipv6_address(host):
return f"{scheme}://[{host}]:{self.port}"
else:
return f"{scheme}://{host}:{self.port}"
return NetworkAddress(host, self.port).to_url(scheme)
def ssl_verify(self):
"""Return the value for the requests library's ``verify=`` parameter.
@@ -6186,23 +6182,16 @@ class PortArgs:
else:
# DP attention. Use TCP + port to handle both single-node and multi-node.
if server_args.nnodes == 1 and server_args.dist_init_addr is None:
dist_init_addr = ("127.0.0.1", server_args.port + ZMQ_TCP_PORT_DELTA)
elif server_args.dist_init_addr.startswith("["): # ipv6 address
port_num, host = configure_ipv6(server_args.dist_init_addr)
dist_init_addr = (host, str(port_num))
na = NetworkAddress("127.0.0.1", server_args.port + ZMQ_TCP_PORT_DELTA)
else:
dist_init_addr = server_args.dist_init_addr.split(":")
na = NetworkAddress.parse(server_args.dist_init_addr)
assert (
len(dist_init_addr) == 2
), "please provide --dist-init-addr as host:port of head node"
dist_init_host, dist_init_port = dist_init_addr
dist_init_port = int(dist_init_port)
dist_init_host = na.host
dist_init_port = na.port
port_base = dist_init_port + 1
detokenizer_port = port_base + 1
rpc_port = port_base + 2
metrics_ipc_name = port_base + 3
metrics_port = port_base + 3
if dp_rank is None:
# TokenizerManager to DataParallelController
scheduler_input_port = port_base + 4
@@ -6217,7 +6206,7 @@ class PortArgs:
wait_port_available(detokenizer_port, "detokenizer_port")
wait_port_available(nccl_port, "nccl_port")
wait_port_available(rpc_port, "rpc_port")
wait_port_available(metrics_ipc_name, "metrics_ipc_name")
wait_port_available(metrics_port, "metrics_port")
# Check scheduler_input_port only for dp.
# Skip check when using worker_ports since the port is already bound by our ZMQ socket
if dp_rank is None or worker_ports is None:
@@ -6229,12 +6218,16 @@ class PortArgs:
raise
return PortArgs(
tokenizer_ipc_name=f"tcp://{dist_init_host}:{port_base}",
scheduler_input_ipc_name=f"tcp://{dist_init_host}:{scheduler_input_port}",
detokenizer_ipc_name=f"tcp://{dist_init_host}:{detokenizer_port}",
tokenizer_ipc_name=NetworkAddress(dist_init_host, port_base).to_tcp(),
scheduler_input_ipc_name=NetworkAddress(
dist_init_host, scheduler_input_port
).to_tcp(),
detokenizer_ipc_name=NetworkAddress(
dist_init_host, detokenizer_port
).to_tcp(),
nccl_port=nccl_port,
rpc_ipc_name=f"tcp://{dist_init_host}:{rpc_port}",
metrics_ipc_name=f"tcp://{dist_init_host}:{metrics_ipc_name}",
rpc_ipc_name=NetworkAddress(dist_init_host, rpc_port).to_tcp(),
metrics_ipc_name=NetworkAddress(dist_init_host, metrics_port).to_tcp(),
tokenizer_worker_ipc_name=tokenizer_worker_ipc_name,
)

53
python/sglang/srt/utils/common.py
View File

@@ -1487,9 +1487,16 @@ def get_zmq_socket_on_host(
Tuple of (port, socket) where port is the randomly assigned TCP port.
"""
socket = context.socket(socket_type)
# Bind to random TCP port
# Bind to random TCP port, auto-wrapping IPv6 and setting zmq.IPV6 flag
config_socket(socket, socket_type)
bind_host = f"tcp://{host}" if host else "tcp://*"
if host:
if is_valid_ipv6_address(host):
socket.setsockopt(zmq.IPV6, 1)
bind_host = f"tcp://[{host}]"
else:
bind_host = f"tcp://{host}"
else:
bind_host = "tcp://*"
port = socket.bind_to_random_port(bind_host)
return port, socket
@@ -2667,48 +2674,6 @@ def is_valid_ipv6_address(address: str) -> bool:
return False
def maybe_wrap_ipv6_address(address: str) -> str:
if is_valid_ipv6_address(address):
return f"[{address}]"
return address
def format_tcp_address(ip: str, port: int) -> str:
return f"tcp://{maybe_wrap_ipv6_address(ip)}:{port}"
def configure_ipv6(dist_init_addr):
addr = dist_init_addr
end = addr.find("]")
if end == -1:
raise ValueError("invalid IPv6 address format: missing ']'")
host = addr[: end + 1]
# this only validates the address without brackets: we still need the below checks.
# if it's invalid, immediately raise an error so we know it's not formatting issues.
if not is_valid_ipv6_address(host[1:end]):
raise ValueError(f"invalid IPv6 address: {host}")
port_str = None
if len(addr) > end + 1:
if addr[end + 1] == ":":
port_str = addr[end + 2 :]
else:
raise ValueError("received IPv6 address format: expected ':' after ']'")
if not port_str:
raise ValueError(
"a port must be specified in IPv6 address (format: [ipv6]:port)"
)
try:
port = int(port_str)
except ValueError:
raise ValueError(f"invalid port in IPv6 address: '{port_str}'")
return port, host
def launch_dummy_health_check_server(host, port, enable_metrics):
import asyncio

129
python/sglang/srt/utils/network.py Normal file
View File

@@ -0,0 +1,129 @@
from __future__ import annotations
import ipaddress
import socket
from dataclasses import dataclass
from typing import Tuple
def _is_ipv6(host: str) -> bool:
"""Check whether *host* is a valid IPv6 address (without brackets)."""
try:
ipaddress.IPv6Address(host)
return True
except ValueError:
return False
def _wrap(host: str) -> str:
"""Wrap an IPv6 address in brackets; pass IPv4/hostname through."""
return f"[{host}]" if _is_ipv6(host) else host
def _parse_port(s: str) -> int:
try:
port = int(s)
except ValueError:
raise ValueError(f"Invalid port number: {s!r}")
if not (0 <= port <= 65535):
raise ValueError(f"Port out of range (0-65535): {port}")
return port
@dataclass(frozen=True)
class NetworkAddress:
host: str
port: int
@property
def is_ipv6(self) -> bool:
return _is_ipv6(self.host)
@property
def family(self) -> socket.AddressFamily:
return socket.AF_INET6 if self.is_ipv6 else socket.AF_INET
def to_url(self, scheme: str = "http") -> str:
"""``http://127.0.0.1:30000`` or ``http://[::1]:30000``."""
return f"{scheme}://{_wrap(self.host)}:{self.port}"
def to_tcp(self) -> str:
"""``tcp://`` endpoint for ZMQ / torch distributed."""
return self.to_url("tcp")
def to_host_port_str(self) -> str:
"""``host:port`` string for gRPC listen address, session IDs, logs."""
return f"{_wrap(self.host)}:{self.port}"
def to_bind_tuple(self) -> Tuple[str, int]:
"""Raw ``(host, port)`` tuple for ``socket.bind()`` / ``socket.connect()``.
Returns the *unwrapped* host — sockets need the raw address, not
the bracketed form.
"""
return (self.host, self.port)
@staticmethod
def parse(addr: str) -> NetworkAddress:
"""Parse a ``host:port`` string into a ``NetworkAddress``.
Accepted formats::
[::1]:8000 → NetworkAddress("::1", 8000)
127.0.0.1:8000 → NetworkAddress("127.0.0.1", 8000)
my-hostname:8000 → NetworkAddress("my-hostname", 8000)
IPv6 addresses **must** be bracketed. Bare ``::1:8000`` is
ambiguous and will raise ``ValueError``.
Raises:
ValueError: If the string cannot be unambiguously parsed.
"""
if not addr:
raise ValueError("Empty address string")
# --- Bracketed IPv6: [addr]:port ---
if addr.startswith("["):
close = addr.find("]")
if close == -1:
raise ValueError(f"Missing closing bracket in IPv6 address: {addr!r}")
host = addr[1:close]
if not _is_ipv6(host):
raise ValueError(f"Invalid IPv6 address inside brackets: {host!r}")
rest = addr[close + 1 :]
if not rest.startswith(":") or len(rest) < 2:
raise ValueError(
f"Expected ':port' after closing bracket, got: {rest!r}"
)
return NetworkAddress(host, _parse_port(rest[1:]))
# --- Plain host:port (IPv4 / hostname) ---
if ":" not in addr:
raise ValueError(f"Missing port in address (expected host:port): {addr!r}")
host, port_str = addr.rsplit(":", 1)
if not host:
raise ValueError(f"Empty host in address: {addr!r}")
# Guard against bare IPv6 slipping through
if ":" in host and _is_ipv6(host):
raise ValueError(
f"Bare IPv6 address without brackets is ambiguous: {addr!r}. "
f"Use [{host}]:{port_str} instead."
)
return NetworkAddress(host, _parse_port(port_str))
@staticmethod
def from_parts(host: str, port: int) -> NetworkAddress:
"""Create from separate host and port, stripping brackets if present.
Useful when the host may come from user input that already has
brackets (e.g. ``[::1]``).
"""
if host.startswith("[") and host.endswith("]"):
host = host[1:-1]
return NetworkAddress(host, port)
def __str__(self) -> str:
return self.to_host_port_str()
def __repr__(self) -> str:
return f"NetworkAddress({self.host!r}, {self.port})"

24
python/sglang/utils.py
View File

@@ -432,31 +432,23 @@ def _prebind_listening_socket(host: str, port: int) -> socket.socket:
"""
import errno
from sglang.srt.utils.common import is_valid_ipv6_address
from sglang.srt.utils.network import NetworkAddress
# Determine socket family based on address type
if host and is_valid_ipv6_address(host):
family = socket.AF_INET6
else:
family = socket.AF_INET
sock = socket.socket(family=family, type=socket.SOCK_STREAM)
# Set socket options to allow port reuse
na = NetworkAddress(host or "0.0.0.0", port)
sock = socket.socket(family=na.family, type=socket.SOCK_STREAM)
# SO_REUSEADDR: Allows binding to a port in TIME_WAIT state
sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
try:
sock.bind((host or "", port))
sock.bind(na.to_bind_tuple())
sock.set_inheritable(True)
logger.info(
f"Successfully reserved port {port} on host '{host or ('::' if family == socket.AF_INET6 else '0.0.0.0')}'"
)
logger.info(f"Successfully reserved port {port} on host '{na.host}'")
except OSError as e:
sock.close()
if e.errno == errno.EADDRINUSE:
raise RuntimeError(
f"Port {port} is already in use on host '{host or ('::' if family == socket.AF_INET6 else '0.0.0.0')}'. "
f"Port {port} is already in use on host '{na.host}'. "
f"Please choose a different port with --port or stop the process using this port. "
f"You can find the process using: lsof -i :{port} or netstat -tlnp | grep {port}"
) from e
@@ -467,9 +459,7 @@ def _prebind_listening_socket(host: str, port: int) -> socket.socket:
f"Consider using a port >= 1024 (e.g., --port 8000) or run with sudo."
) from e
else:
raise RuntimeError(
f"Failed to bind to {host or ('::' if family == socket.AF_INET6 else '0.0.0.0')}:{port}: {e}"
) from e
raise RuntimeError(f"Failed to bind to {na}: {e}") from e
return sock