chore: import upstream snapshot with attribution

python/sglang/srt/managers/load_snapshot.py

@@ -0,0 +1,654 @@
+"""Load snapshot: publish scheduler load metrics for DP balancing and /v1/loads.
+
+Architecture
+------------
+
+Each scheduler periodically publishes a ``LoadSnapshot`` containing its
+current load metrics (running reqs, tokens, throughput, ...).  Two
+transport backends are supported:
+
+**SHM mode** (single-node, default)::
+
+    Scheduler  ──ShmLoadSnapshotWriter──▶  /dev/shm mmap file
+                                               ▲
+    TokenizerManager  ──ShmLoadSnapshotReader───┘  (for /v1/loads)
+    DataParallelController  ──ShmLoadSnapshotReader─┘  (for dispatch)
+
+**ZMQ mode** (multi-node DP attention, or ``SGLANG_LOAD_SNAPSHOT_USE_ZMQ=1``)::
+
+    Scheduler (any node)  ──ZmqLoadSnapshotWriter (PUSH)──▶  network
+                                                               │
+    ZmqShmLoadSnapshotReader (PULL, node 0)  ◀─────────────────┘
+        │  drains zmq, writes to SHM
+        ▼
+    /dev/shm mmap file (node 0)
+        ▲
+    TokenizerManager / DataParallelController  ──ShmLoadSnapshotReader──┘
+
+Shared memory does not work across nodes, so multi-node DP attention
+requires the ZMQ transport.  The ``ZmqShmLoadSnapshotReader`` on node 0
+receives snapshots from all schedulers via zmq PUSH/PULL and writes them
+into the local SHM file.  All readers (TokenizerManager,
+DataParallelController) on
+node 0 then read from SHM.
+
+``zmq_reader_owner()`` decides which process on node 0 binds the zmq
+PULL socket (only one can bind); the other reads plain SHM.
+"""
+
+from __future__ import annotations
+
+import fcntl
+import hashlib
+import logging
+import mmap
+import os
+import struct
+from contextlib import contextmanager
+from typing import Optional
+
+import msgspec
+import msgspec.msgpack
+import msgspec.structs
+
+from sglang.srt.environ import envs
+from sglang.srt.utils.network import is_zmq_endpoint_ipv6
+
+logger = logging.getLogger(__name__)
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+
+def should_use_zmq(server_args) -> bool:
+    """Whether to use zmq PUSH/PULL instead of shared memory for load snapshots.
+
+    Shared memory (mmap) only works within a single node.  When schedulers
+    run on multiple nodes (multi-node DP attention), they cannot write to
+    the SHM file on node 0, so we fall back to zmq transport.  The env var
+    ``SGLANG_LOAD_SNAPSHOT_USE_ZMQ`` forces zmq mode for testing.
+    """
+    return (
+        server_args.enable_dp_attention and server_args.nnodes > 1
+    ) or envs.SGLANG_LOAD_SNAPSHOT_USE_ZMQ.get()
+
+
+_LOAD_AWARE_METHODS = frozenset({"total_requests", "total_tokens"})
+
+
+def _tokenizer_load_snapshot_owner_caller(server_args) -> str:
+    """The caller that plays the tokenizer-side zmq owner role.
+
+    In multi-tokenizer mode (``tokenizer_worker_num > 1``) there are N
+    independent ``TokenizerWorker`` processes that would all try to bind the
+    same zmq PULL endpoint.  Instead, the single ``MultiTokenizerRouter``
+    process owns the socket (polls zmq -> SHM) and every worker reads SHM.
+    """
+    if server_args.tokenizer_worker_num > 1:
+        return "MultiTokenizerRouter"
+    return "TokenizerManager"
+
+
+def zmq_reader_owner(server_args, caller: str) -> bool:
+    """Decide which process owns the zmq PULL socket.
+
+    Exactly one of ``"DataParallelController"``, ``"TokenizerManager"``, or
+    ``"MultiTokenizerRouter"`` must return True when zmq mode is active.  The
+    owner polls zmq -> SHM; the others read SHM.
+
+    Rules:
+      - Non-zero node_rank: no TokenizerManager, DataParallelController only
+        launches schedulers and waits -> nobody owns it.
+      - dp_size == 1: no DataParallelController exists -> tokenizer-side owner
+        owns it.
+      - dp_size > 1, load-aware method: DataParallelController polls on every
+        dispatch via refresh_load_budget() -> DataParallelController owns it.
+      - dp_size > 1, round-robin / other: DataParallelController never reads
+        load data -> tokenizer-side owner owns it (polls on /v1/loads calls).
+
+    The tokenizer-side owner is the ``"MultiTokenizerRouter"`` caller in
+    multi-tokenizer mode, otherwise the ``"TokenizerManager"`` caller.
+    """
+    if not should_use_zmq(server_args):
+        return False
+    if server_args.node_rank != 0:
+        return False
+    tokenizer_owner = _tokenizer_load_snapshot_owner_caller(server_args)
+    if server_args.dp_size == 1:
+        return caller == tokenizer_owner
+    if server_args.load_balance_method.lower() in _LOAD_AWARE_METHODS:
+        return caller == "DataParallelController"
+    return caller == tokenizer_owner
+
+
+# ---------------------------------------------------------------------------
+# LoadSnapshot data class
+# ---------------------------------------------------------------------------
+
+
+class MemoryMetrics(msgspec.Struct, array_like=True):
+    """Memory breakdown metrics."""
+
+    weight_gb: float
+    kv_cache_gb: float
+    graph_gb: float
+    token_capacity: int
+
+
+class SpeculativeMetrics(msgspec.Struct, array_like=True):
+    """Speculative decoding metrics."""
+
+    accept_length: float
+    accept_rate: float
+
+
+class LoRAMetrics(msgspec.Struct, array_like=True):
+    """LoRA adapter pool metrics."""
+
+    slots_used: int
+    slots_total: int
+    utilization: float
+
+
+class DisaggregationMetrics(msgspec.Struct, array_like=True):
+    """PD disaggregation metrics."""
+
+    mode: str  # "prefill", "decode", or "null"
+    prefill_bootstrap_queue_reqs: int = 0
+    prefill_inflight_queue_reqs: int = 0
+    decode_prealloc_queue_reqs: int = 0
+    decode_transfer_queue_reqs: int = 0
+    decode_retracted_queue_reqs: int = 0
+    kv_transfer_speed_gb_s: float = 0.0
+    kv_transfer_latency_ms: float = 0.0
+
+
+class QueueMetrics(msgspec.Struct, array_like=True):
+    """Detailed queue info breakdown."""
+
+    waiting: int
+    grammar: int
+    paused: int
+    retracted: int
+
+
+_CORE_KEYS = (
+    "timestamp",
+    "dp_rank",
+    "num_running_reqs",
+    "num_waiting_reqs",
+    "num_waiting_uncached_tokens",
+    "num_used_tokens",
+    "num_total_tokens",
+    "max_total_num_tokens",
+    "max_running_requests",
+    "token_usage",
+    "gen_throughput",
+    "cache_hit_rate",
+    "utilization",
+)
+
+
+class LoadSnapshot(msgspec.Struct, omit_defaults=True):
+    """Per-DP-rank load metrics: the SHM/zmq wire format and the /v1/loads source."""
+
+    timestamp: float = 0.0
+    dp_rank: int = 0
+    num_running_reqs: int = 0
+    num_waiting_reqs: int = 0
+    num_waiting_uncached_tokens: int = 0
+    num_used_tokens: int = 0
+    num_total_tokens: int = 0
+    max_total_num_tokens: int = 0
+    max_running_requests: int = 0
+    token_usage: float = 0.0
+    gen_throughput: float = 0.0
+    cache_hit_rate: float = 0.0
+    utilization: float = 0.0
+
+    memory: Optional[MemoryMetrics] = None
+    speculative: Optional[SpeculativeMetrics] = None
+    lora: Optional[LoRAMetrics] = None
+    disaggregation: Optional[DisaggregationMetrics] = None
+    queues: Optional[QueueMetrics] = None
+
+    VALID_SECTIONS = frozenset(
+        {"core", "memory", "spec", "lora", "disagg", "queues", "all"}
+    )
+
+    def to_dict(self, include: Optional[set[str]] = None) -> dict:
+        load = {key: getattr(self, key) for key in _CORE_KEYS}
+
+        if include is None or "all" in include:
+            include_all = True
+        else:
+            if not (include <= self.VALID_SECTIONS):
+                raise ValueError(
+                    f"Invalid include sections: {include - self.VALID_SECTIONS}. "
+                    f"Valid options: {sorted(self.VALID_SECTIONS)}"
+                )
+            if include == {"core"}:
+                return load
+            include_all = False
+
+        for field, include_name, section in (
+            ("memory", "memory", self.memory),
+            ("speculative", "spec", self.speculative),
+            ("lora", "lora", self.lora),
+            ("disaggregation", "disagg", self.disaggregation),
+            ("queues", "queues", self.queues),
+        ):
+            if section is None or (not include_all and include_name not in include):
+                continue
+            load[field] = msgspec.structs.asdict(section)
+
+        return load
+
+
+def _enc_hook(obj):
+    """Coerce numpy scalars to native Python; msgpack has no numpy types."""
+    to_item = getattr(obj, "item", None)
+    if to_item is not None:
+        return to_item()
+    raise NotImplementedError(f"cannot encode {type(obj).__name__} in load snapshot")
+
+
+snapshot_encoder = msgspec.msgpack.Encoder(enc_hook=_enc_hook)
+snapshot_decoder = msgspec.msgpack.Decoder(LoadSnapshot)
+
+
+# ---------------------------------------------------------------------------
+# SHM file layout utilities
+# ---------------------------------------------------------------------------
+
+MAGIC = b"SLNS"
+VERSION = 2
+HEADER_STRUCT = struct.Struct("<4sHHI")
+SLOT_LEN_STRUCT = struct.Struct("<I")
+SLOT_SIZE = 16 * 1024
+
+
+@contextmanager
+def file_lock(fd: int, lock_type: int):
+    fcntl.flock(fd, lock_type)
+    try:
+        yield
+    finally:
+        fcntl.flock(fd, fcntl.LOCK_UN)
+
+
+def shm_path_for(ipc_name: str) -> str:
+    name = os.path.basename(ipc_name.rstrip("/")) or "default"
+    safe_name = "".join(c if c.isalnum() or c in "._-" else "_" for c in name)
+    digest = hashlib.blake2s(ipc_name.encode(), digest_size=4).hexdigest()
+    return f"/dev/shm/sglang_loads_{safe_name}_{digest}.shm"
+
+
+def file_size(dp_size: int, slot_size: int = SLOT_SIZE) -> int:
+    return HEADER_STRUCT.size + dp_size * slot_size
+
+
+def slot_offset(dp_rank: int, slot_size: int = SLOT_SIZE) -> int:
+    return HEADER_STRUCT.size + dp_rank * slot_size
+
+
+# ---------------------------------------------------------------------------
+# Writers
+# ---------------------------------------------------------------------------
+
+
+class ShmLoadSnapshotWriter:
+    def __init__(
+        self, path: str, dp_size: int, dp_rank: int, publish_interval: int = 1
+    ):
+        if dp_rank < 0 or dp_rank >= dp_size:
+            raise ValueError(f"invalid dp_rank={dp_rank} for dp_size={dp_size}")
+        self.publish_interval = max(1, publish_interval)
+        self.publish_counter = 0
+
+        self.path = path
+        self.dp_size = dp_size
+        self.dp_rank = dp_rank
+        self.slot_size = SLOT_SIZE
+        self.fd = -1
+        size = file_size(dp_size, self.slot_size)
+
+        self.fd = os.open(path, os.O_CREAT | os.O_RDWR, 0o600)
+        try:
+            with file_lock(self.fd, fcntl.LOCK_EX):
+                os.ftruncate(self.fd, size)
+                self.mmap = mmap.mmap(self.fd, size, access=mmap.ACCESS_WRITE)
+                HEADER_STRUCT.pack_into(
+                    self.mmap, 0, MAGIC, VERSION, dp_size, self.slot_size
+                )
+                self._write_payload(LoadSnapshot(dp_rank=dp_rank))
+        except Exception:
+            if self.fd >= 0:
+                os.close(self.fd)
+            raise
+
+    def write(self, snapshot: LoadSnapshot) -> None:
+        if snapshot.dp_rank != self.dp_rank:
+            raise ValueError(
+                f"snapshot dp_rank={snapshot.dp_rank} does not match writer dp_rank={self.dp_rank}"
+            )
+
+        with file_lock(self.fd, fcntl.LOCK_EX):
+            self._write_payload(snapshot)
+
+    def _write_payload(self, snapshot: LoadSnapshot) -> None:
+        payload = snapshot_encoder.encode(snapshot)
+        max_payload_size = self.slot_size - SLOT_LEN_STRUCT.size
+        if len(payload) > max_payload_size:
+            raise ValueError(
+                f"load snapshot payload size {len(payload)} exceeds slot payload "
+                f"capacity {max_payload_size}"
+            )
+
+        offset = slot_offset(self.dp_rank, self.slot_size)
+        payload_start = offset + SLOT_LEN_STRUCT.size
+        payload_end = payload_start + len(payload)
+        slot_end = offset + self.slot_size
+
+        SLOT_LEN_STRUCT.pack_into(self.mmap, offset, 0)
+        self.mmap[payload_start:payload_end] = payload
+        self.mmap[payload_end:slot_end] = b"\0" * (slot_end - payload_end)
+        SLOT_LEN_STRUCT.pack_into(self.mmap, offset, len(payload))
+
+    def close(self) -> None:
+        self.mmap.close()
+        os.close(self.fd)
+
+
+class ZmqLoadSnapshotWriter:
+    """Sends load snapshots via zmq PUSH to a ZmqShmLoadSnapshotReader.
+
+    CONFLATE is set so only the latest message is kept in the send
+    buffer when the reader is slower than the writer.
+    """
+
+    def __init__(
+        self, endpoint: str, dp_size: int, dp_rank: int, publish_interval: int = 1
+    ):
+        import zmq as _zmq
+
+        if dp_rank < 0 or dp_rank >= dp_size:
+            raise ValueError(f"invalid dp_rank={dp_rank} for dp_size={dp_size}")
+        self.publish_interval = max(1, publish_interval)
+        self.publish_counter = 0
+        self.dp_size = dp_size
+        self.dp_rank = dp_rank
+
+        self._zmq = _zmq
+        self._ctx = _zmq.Context.instance()
+        self._socket = self._ctx.socket(_zmq.PUSH)
+        if is_zmq_endpoint_ipv6(endpoint):
+            self._socket.setsockopt(_zmq.IPV6, 1)
+        self._socket.setsockopt(_zmq.LINGER, 0)
+        self._socket.setsockopt(_zmq.CONFLATE, 1)
+        self._socket.connect(endpoint)
+
+    def write(self, snapshot: LoadSnapshot) -> None:
+        if snapshot.dp_rank != self.dp_rank:
+            raise ValueError(
+                f"snapshot dp_rank={snapshot.dp_rank} does not match "
+                f"writer dp_rank={self.dp_rank}"
+            )
+        try:
+            self._socket.send(snapshot_encoder.encode(snapshot), self._zmq.NOBLOCK)
+        except self._zmq.Again:
+            pass
+
+    def close(self) -> None:
+        self._socket.close()
+
+
+# ---------------------------------------------------------------------------
+# Readers
+# ---------------------------------------------------------------------------
+
+
+class ShmLoadSnapshotReader:
+    def __init__(self, path: str, dp_size: int):
+        self.path = path
+        self.dp_size = dp_size
+        self.mmap: Optional[mmap.mmap] = None
+        self.fd: Optional[int] = None
+        self.slot_size = SLOT_SIZE
+        self._header_warning_logged = False
+        self._attach()
+
+    def _attach(self) -> bool:
+        if self.mmap is not None:
+            return True
+
+        try:
+            fd = os.open(self.path, os.O_RDONLY)
+        except FileNotFoundError:
+            return False
+
+        size = os.fstat(fd).st_size
+        if size < HEADER_STRUCT.size:
+            os.close(fd)
+            return False
+
+        try:
+            with file_lock(fd, fcntl.LOCK_SH):
+                mapped = mmap.mmap(fd, size, access=mmap.ACCESS_READ)
+                magic, version, dp_size, slot_size = HEADER_STRUCT.unpack_from(
+                    mapped, 0
+                )
+        except (OSError, ValueError):
+            os.close(fd)
+            return False
+
+        if (
+            magic != MAGIC
+            or version != VERSION
+            or dp_size != self.dp_size
+            or slot_size < SLOT_LEN_STRUCT.size
+            or size < file_size(self.dp_size, slot_size)
+        ):
+            mapped.close()
+            os.close(fd)
+            if not self._header_warning_logged:
+                logger.warning("load shm header mismatch at %s", self.path)
+                self._header_warning_logged = True
+            return False
+
+        self.mmap = mapped
+        self.fd = fd
+        self.slot_size = slot_size
+        return True
+
+    def read(self, dp_rank: int) -> Optional[LoadSnapshot]:
+        if dp_rank < 0 or dp_rank >= self.dp_size:
+            return None
+        if not self._attach():
+            return None
+
+        assert self.fd is not None
+        with file_lock(self.fd, fcntl.LOCK_SH):
+            return self._read_slot(dp_rank)
+
+    def _read_slot(self, dp_rank: int) -> Optional[LoadSnapshot]:
+        assert self.mmap is not None
+        offset = slot_offset(dp_rank, self.slot_size)
+        (payload_len,) = SLOT_LEN_STRUCT.unpack_from(self.mmap, offset)
+        max_payload_size = self.slot_size - SLOT_LEN_STRUCT.size
+        if payload_len == 0 or payload_len > max_payload_size:
+            return None
+
+        payload_start = offset + SLOT_LEN_STRUCT.size
+        payload_end = payload_start + payload_len
+        try:
+            return snapshot_decoder.decode(self.mmap[payload_start:payload_end])
+        except Exception as e:
+            logger.debug("load snapshot decode failed for rank %s: %s", dp_rank, e)
+            return None
+
+    def read_all(self) -> list[LoadSnapshot]:
+        if not self._attach():
+            return []
+
+        assert self.fd is not None
+        with file_lock(self.fd, fcntl.LOCK_SH):
+            loads = []
+            for r in range(self.dp_size):
+                load = self._read_slot(r)
+                if load is not None:
+                    loads.append(load)
+            return loads
+
+    def close(self) -> None:
+        if self.mmap is not None:
+            self.mmap.close()
+            self.mmap = None
+        if self.fd is not None:
+            os.close(self.fd)
+            self.fd = None
+
+
+class ZmqShmLoadSnapshotReader:
+    """Receives snapshots via zmq PULL from writers, writes to SHM, reads from SHM.
+
+    Transparently wraps a ShmLoadSnapshotReader.  Every read() / read_all()
+    first drains the PULL socket into SHM so callers always see fresh data.
+    """
+
+    def __init__(self, endpoint: str, shm_path: str, dp_size: int):
+        import zmq as _zmq
+
+        self._zmq = _zmq
+        self._ctx = _zmq.Context.instance()
+        self._socket = self._ctx.socket(_zmq.PULL)
+        if is_zmq_endpoint_ipv6(endpoint):
+            self._socket.setsockopt(_zmq.IPV6, 1)
+        self._socket.setsockopt(_zmq.LINGER, 0)
+        self._socket.setsockopt(_zmq.CONFLATE, 1)
+        self._socket.bind(endpoint)
+
+        self._endpoint = endpoint
+        self._shm_path = shm_path
+        self.dp_size = dp_size
+        self._shm_reader = ShmLoadSnapshotReader(shm_path, dp_size)
+        self._shm_writers: dict[int, ShmLoadSnapshotWriter] = {}
+
+    def _poll(self) -> None:
+        """Drain zmq messages and write latest per dp_rank to SHM."""
+        latest: dict[int, LoadSnapshot] = {}
+        while True:
+            try:
+                data = self._socket.recv(self._zmq.NOBLOCK)
+            except self._zmq.Again:
+                break
+            try:
+                snapshot = snapshot_decoder.decode(data)
+                if 0 <= snapshot.dp_rank < self.dp_size:
+                    latest[snapshot.dp_rank] = snapshot
+            except Exception as e:
+                logger.warning("load snapshot zmq decode failed: %s", e)
+
+        for dp_rank, snapshot in latest.items():
+            if dp_rank not in self._shm_writers:
+                self._shm_writers[dp_rank] = ShmLoadSnapshotWriter(
+                    self._shm_path, self.dp_size, dp_rank
+                )
+            try:
+                self._shm_writers[dp_rank].write(snapshot)
+            except Exception as e:
+                logger.warning(
+                    "load snapshot shm write failed for rank %d: %s", dp_rank, e
+                )
+
+    def fileno(self) -> int:
+        """Edge-triggered fd that becomes readable when zmq messages arrive.
+
+        Lets an owner process register the reader with an event loop and drain
+        it via ``poll()`` instead of polling on a timer.
+        """
+        return self._socket.getsockopt(self._zmq.FD)
+
+    def poll(self) -> None:
+        """Drain the zmq PULL socket into SHM.
+
+        Public entry point so an owner process (e.g. MultiTokenizerRouter) can
+        keep SHM fresh without touching internals.
+        """
+        self._poll()
+
+    def read(self, dp_rank: int) -> Optional[LoadSnapshot]:
+        self._poll()
+        return self._shm_reader.read(dp_rank)
+
+    def read_all(self) -> list[LoadSnapshot]:
+        self._poll()
+        return self._shm_reader.read_all()
+
+    def close(self) -> None:
+        for w in self._shm_writers.values():
+            w.close()
+        self._shm_writers.clear()
+        self._shm_reader.close()
+        self._socket.close()
+        if self._endpoint.startswith("ipc://"):
+            try:
+                os.unlink(self._endpoint[len("ipc://") :])
+            except OSError:
+                pass
+
+
+# ---------------------------------------------------------------------------
+# Factory functions
+# ---------------------------------------------------------------------------
+
+
+def _zmq_addr_for(port_args) -> str:
+    """Return the zmq PUSH/PULL address from PortArgs.
+
+    For dp_attention (TCP mode), uses the ``load_collector_ipc_name`` field
+    stored in PortArgs.  For single-node IPC (env-var override), derives
+    a deterministic IPC path from ``instance_id``.
+    """
+    ipc_name = getattr(port_args, "load_collector_ipc_name", "")
+    if ipc_name:
+        return ipc_name
+    safe = "".join(
+        c if c.isalnum() or c in "._-" else "_" for c in port_args.instance_id
+    )
+    digest = hashlib.blake2s(port_args.instance_id.encode(), digest_size=4).hexdigest()
+    return f"ipc:///tmp/sglang_load_collector_{safe}_{digest}.sock"
+
+
+def create_load_snapshot_writer(
+    server_args,
+    port_args,
+    dp_size: int,
+    dp_rank: int,
+    publish_interval: int = 1,
+):
+    """Return a SHM or ZMQ writer based on server configuration."""
+    if should_use_zmq(server_args):
+        return ZmqLoadSnapshotWriter(
+            _zmq_addr_for(port_args), dp_size, dp_rank, publish_interval
+        )
+    return ShmLoadSnapshotWriter(
+        shm_path_for(port_args.instance_id), dp_size, dp_rank, publish_interval
+    )
+
+
+def create_load_snapshot_reader(server_args, port_args, caller: str):
+    """Create a load snapshot reader.
+
+    Args:
+        caller: ``"DataParallelController"``, ``"TokenizerManager"``, or
+            ``"MultiTokenizerRouter"`` -- determines who binds the zmq PULL
+            socket when zmq mode is active.
+    """
+    dp_size = server_args.dp_size
+    if zmq_reader_owner(server_args, caller):
+        return ZmqShmLoadSnapshotReader(
+            _zmq_addr_for(port_args), shm_path_for(port_args.instance_id), dp_size
+        )
+    return ShmLoadSnapshotReader(shm_path_for(port_args.instance_id), dp_size)