diff --git a/python/sglang/srt/eplb/eplb_manager.py b/python/sglang/srt/eplb/eplb_manager.py index 0388dbee15e3..d454a1eaa61d 100644 --- a/python/sglang/srt/eplb/eplb_manager.py +++ b/python/sglang/srt/eplb/eplb_manager.py @@ -19,6 +19,7 @@ ) from sglang.srt.eplb.expert_location_updater import ExpertLocationUpdater from sglang.srt.runtime_context import get_exec, get_model, get_parallel +from sglang.srt.weight_cache.expert_layout import get_expert_layout_session if TYPE_CHECKING: from sglang.srt.configs.model_config import ModelConfig @@ -67,6 +68,13 @@ def __init__( f"[EPLBManager] system started, will rebalance per {self._rebalance_num_iterations} iterations." ) + # With the weight cache on, the expert weights are normally the + # daemon's copy, shared with every engine that maps it. A rebalance + # moves experts in place, so it has to keep the daemon's record of the + # expert layout in step (see weight_cache/expert_layout.py). + self._weight_cache_mode = _weight_cache_mode() + self._weight_cache_off_reason = None + self._main_generator = self._entrypoint() def on_forward_pass_end(self): @@ -112,6 +120,32 @@ def rebalance(self): ): return + # The daemon's expert layout record to move experts through, or None + # when the weights are this engine's own. + expert_layout = None + if self._weight_cache_mode != "off": + session = get_expert_layout_session() + mode, reason = ( + ("off", "this engine has no expert layout session") + if session is None + else session.rebalance_mode() + ) + if mode == "coordinated" and get_exec().moe.elastic_ep_backend is not None: + mode, reason = ( + "off", + "elastic EP also moves experts outside rebalance()", + ) + if mode == "off": + if reason != self._weight_cache_off_reason: + self._weight_cache_off_reason = reason + logger.warning( + "[EPLBManager] rebalancing is off with the weight cache: " + f"{reason}. The expert layout stays as it is." + ) + return + if mode == "coordinated": + expert_layout = session + logger.info("[EPLBManager] rebalance start") enable_timing = self._rebalance_layers_per_chunk is None @@ -152,6 +186,14 @@ def rebalance(self): for chunk_layer_ids in update_layer_ids_chunks: if len(update_layer_ids_chunks) > 1: yield + if expert_layout is not None: + reason = expert_layout.begin(chunk_layer_ids) + if reason is not None: + logger.info( + f"[EPLBManager] rebalance stopped before layers " + f"{chunk_layer_ids}: {reason}" + ) + return update_expert_location_with_recovery( expert_location_updater=self._get_expert_location_updater(), model=self._get_model(), @@ -174,6 +216,22 @@ def rebalance(self): # P2P waits only synchronize participating peers. Ranks without # moves must also install this chunk before NIXL resumes. dist.barrier() + if expert_layout is not None: + # Once committed, the daemon may serve these weights to the + # next engine, so the moves must have landed. + torch.get_device_module().synchronize() + reason = expert_layout.commit( + chunk_layer_ids, + get_global_expert_location_metadata().physical_to_logical_map_cpu, + ) + if reason is not None: + logger.error( + f"[EPLBManager] the weight cache daemons did not record the " + f"move of layers {chunk_layer_ids}: {reason}. Rebalancing " + f"is off; restart the daemons before another engine uses " + f"them." + ) + return self._log_rebalance_layout_after_update(update_layer_ids=all_update_layer_ids) @@ -351,3 +409,9 @@ def update_expert_location_with_recovery( def _chunk_list(items: List, chunk_size): for start_index in range(0, len(items), chunk_size): yield items[start_index : start_index + chunk_size] + + +def _weight_cache_mode() -> str: + # The model config's field. Inside EPLBManager.__init__, ``get_model`` names + # the model getter argument instead, and the model is not loaded yet there. + return get_model().weight_cache_mode diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index e646d198032f..0e030f7bfaac 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -255,6 +255,7 @@ from sglang.srt.utils.profile_utils import build_step_span_name from sglang.srt.utils.torch_memory_saver_adapter import TorchMemorySaverAdapter from sglang.srt.utils.weight_checker import WeightChecker +from sglang.srt.weight_cache.expert_layout import adopt_weight_cache_expert_layout _is_npu = is_npu() _is_cpu_amx_available = cpu_has_amx_support() @@ -751,9 +752,16 @@ def maybe_init_expert_location_metadata(self): else 0 ) set_global_expert_location_metadata( - compute_initial_expert_location_metadata( + # With the weight cache, route with the layout the daemon's + # weights are in, which an EPLB rebalance may have changed. + adopt_weight_cache_expert_layout( + compute_initial_expert_location_metadata( + model_config=self.model_config, + moe_ep_rank=expert_rank, + ), model_config=self.model_config, moe_ep_rank=expert_rank, + gpu_id=self.gpu_id, ) ) if ( diff --git a/python/sglang/srt/runtime_context.py b/python/sglang/srt/runtime_context.py index e17d01a4fe42..a7228a3bf5d2 100644 --- a/python/sglang/srt/runtime_context.py +++ b/python/sglang/srt/runtime_context.py @@ -621,6 +621,9 @@ class Resources(_FlagGroupBase): # (owning accessors live in sglang.srt.eplb). expert_distribution_recorder: Any = None expert_location_metadata: Any = None + # Weight cache: this engine's link to its daemon's record of the expert + # layout (owning accessors live in sglang.srt.weight_cache.expert_layout). + expert_layout_session: Any = None # LPLB: layer_id -> solver. lplb_solvers: dict = msgspec.field(default_factory=dict) # Named side streams (see RuntimeContext.get_stream): name -> stream. diff --git a/python/sglang/srt/weight_cache/daemon.py b/python/sglang/srt/weight_cache/daemon.py index 864e7b104e7e..a023f9c62b3e 100644 --- a/python/sglang/srt/weight_cache/daemon.py +++ b/python/sglang/srt/weight_cache/daemon.py @@ -59,6 +59,15 @@ spawn_world_rank, ) +from .expert_layout import ( + ATTACH, + LAYOUT_REFUSED, + OK, + REQUEST_TYPES, + ExpertLayoutPresence, + ExpertLayoutRegistry, + handle_expert_layout_request, +) from .protocol import ( CacheConfig, check_ipc_quant_support, @@ -188,6 +197,10 @@ def __init__( self.state_entries: Dict[str, Dict[str, Any]] = {} self.preloaded_weights_bytes = 0 self.transport_backend = None + # Where the routed experts of the served weights are (see + # expert_layout.py); tracked once load() builds the EPLB layout. + self.expert_layout = ExpertLayoutRegistry(None) + self.expert_layout_presence: Optional[ExpertLayoutPresence] = None def _init_distributed(self, server_args, model_config): """Initialize the distributed backend required for model loading. @@ -324,6 +337,7 @@ def load(self): # avoids maintaining a second copy of the model-parallel hierarchy. self._init_distributed(server_args, model_config) self._initialize_eplb_expert_location_metadata(model_config) + self._track_expert_layout() moe_dp_rank = get_parallel().moe_dp_rank moe_ep_rank = get_parallel().moe_ep_rank self.config = CacheConfig( @@ -488,6 +502,18 @@ def _initialize_eplb_expert_location_metadata(self, model_config) -> None: ) ) + def _track_expert_layout(self) -> None: + """Keep the record of the EPLB layout built above. An EPLB rebalance in + an engine moves experts in these weights and reports it here.""" + if not get_exec().moe.enable_eplb: + return + + from sglang.srt.eplb.expert_location import get_global_expert_location_metadata + + self.expert_layout = ExpertLayoutRegistry.from_metadata( + get_global_expert_location_metadata() + ) + def serve(self): """Block and serve IPC handles over Unix socket.""" # Do NOT unlink an existing socket here: stale-file cleanup is the launch @@ -523,6 +549,9 @@ def _signal_handler(signum, frame): signal.signal(signal.SIGTERM, _signal_handler) signal.signal(signal.SIGINT, _signal_handler) + self.expert_layout_presence = ExpertLayoutPresence(self.expert_layout) + self.expert_layout_presence.start() + try: while self._running: try: @@ -533,19 +562,11 @@ def _signal_handler(signum, frame): # (or dies mid-send) would block recv_msg forever and stall # every other engine rank. Bound each exchange instead. conn.settimeout(CLIENT_CONNECTION_TIMEOUT) - try: - self._handle_connection(conn) - except Exception as e: - logger.error( - f"[WeightCacheDaemon gpu={self.gpu_id}] " - f"Error handling connection: {e}", - exc_info=True, - ) - finally: - conn.close() + self._serve_connection(conn) except socket.timeout: continue finally: + self.expert_layout_presence.stop() sock.close() if os.path.exists(self.socket_path): os.unlink(self.socket_path) @@ -553,8 +574,24 @@ def _signal_handler(signum, frame): os.unlink(self.ready_path) logger.info(f"[WeightCacheDaemon gpu={self.gpu_id}] Shutdown complete") - def _handle_connection(self, conn: socket.socket): - """Handle a single client connection.""" + def _serve_connection(self, conn: socket.socket) -> None: + """Handle one accepted connection, then close it unless it became an + engine's expert layout attach connection.""" + keep_open = False + try: + keep_open = self._handle_connection(conn) + except Exception as e: + logger.error( + f"[WeightCacheDaemon gpu={self.gpu_id}] Error handling connection: {e}", + exc_info=True, + ) + finally: + if not keep_open: + conn.close() + + def _handle_connection(self, conn: socket.socket) -> bool: + """Handle a single client connection. Returns True when the connection + must stay open: an engine's expert layout attach connection.""" req = recv_msg(conn) if req.get("type") == "query_config": @@ -580,7 +617,17 @@ def _handle_connection(self, conn: socket.socket): send_msg( conn, {"status": "mismatch", "daemon_config": self.config.to_dict()} ) - return + return False + + # The engine must route with the expert layout these weights are in. + refusal = self.expert_layout.check_fetch(req.get("expert_layout")) + if refusal is not None: + logger.warning( + f"[WeightCacheDaemon gpu={self.gpu_id}] " + f"Refusing to serve the weights: {refusal}" + ) + send_msg(conn, {"status": LAYOUT_REFUSED, "message": refusal}) + return False logger.info( f"[WeightCacheDaemon gpu={self.gpu_id}] " @@ -601,6 +648,9 @@ def _handle_connection(self, conn: socket.socket): elif req.get("type") == "ping": send_msg(conn, {"status": "ok"}) + elif req.get("type") in REQUEST_TYPES: + return self._handle_expert_layout_request(conn, req) + else: send_msg( conn, @@ -609,6 +659,21 @@ def _handle_connection(self, conn: socket.socket): "message": f"Unknown request type: {req.get('type')}", }, ) + return False + + def _handle_expert_layout_request(self, conn: socket.socket, req) -> bool: + reply = handle_expert_layout_request(self.expert_layout, req) + if req["type"] != ATTACH or reply.get("status") != OK: + send_msg(conn, reply) + return False + # Attached: the connection now tells whether the engine is alive. + try: + send_msg(conn, reply) + except BaseException: + self.expert_layout.detach(req["engine_id"]) + raise + self.expert_layout_presence.watch(conn, req["engine_id"]) + return True def shutdown(self): """Release GPU memory and clean up.""" diff --git a/python/sglang/srt/weight_cache/expert_layout.py b/python/sglang/srt/weight_cache/expert_layout.py new file mode 100644 index 000000000000..1f1e40c0edb0 --- /dev/null +++ b/python/sglang/srt/weight_cache/expert_layout.py @@ -0,0 +1,933 @@ +# SPDX-License-Identifier: Apache-2.0 +"""The weight cache daemon's record of where its routed experts are. + +With the weight cache on, an engine's routed expert weights are the daemon's +tensors, mapped zero-copy. An EPLB rebalance moves experts in place, so it +moves them for the daemon and for every engine that maps it, while only the +rebalancing engine's ExpertLocationMetadata follows. The daemon therefore keeps +the layout of its weights -- the ``physical_to_logical_map`` -- with a version: + +* An engine attaches to its rank's daemon before it builds its expert location + metadata, and builds the metadata from the daemon's map. A restarted or a + second engine then routes with the layout the shared weights are in. +* A rebalance moves each chunk of layers in two phases: ``begin`` marks the + move pending on every rank's daemon, ``commit`` stores the new rows and bumps + the version after the weights have moved. An engine that disconnects while + its move is pending leaves the daemon dirty, and a dirty daemon serves no + weights. +* Only an engine that is alone on a daemon may move experts: a second engine + that maps the weights would keep routing with the old layout. + +The attach connection stays open for the engine's lifetime and is its +presence: the kernel closes it however the engine exits, so liveness needs no +PIDs, which another container cannot see. +""" + +import hashlib +import json +import logging +import os +import selectors +import socket +import stat +import threading +import time +import uuid +from dataclasses import dataclass +from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple + +from sglang.srt.runtime_context import get_exec, get_model, get_parallel, get_resources + +from .protocol import get_socket_path, recv_msg, send_msg + +logger = logging.getLogger(__name__) + +ATTACH = "expert_layout_attach" +QUERY = "expert_layout_query" +BEGIN = "expert_layout_begin" +COMMIT = "expert_layout_commit" +ABORT = "expert_layout_abort" +REQUEST_TYPES = frozenset((ATTACH, QUERY, BEGIN, COMMIT, ABORT)) + +# Reply statuses of the requests above. BUSY means not now (another engine +# maps the weights); REFUSED means that retrying will not help. +OK = "ok" +BUSY = "busy" +REFUSED = "refused" + +# The fetch_state status for weights the daemon refuses to serve because of +# their expert layout. +LAYOUT_REFUSED = "layout_refused" + +# What a rank found when it attached at startup. +TRACKED = "tracked" # the daemon tracks the layout of its weights +UNTRACKED = "untracked" # the daemon runs without --enable-eplb +LEGACY = "legacy" # the daemon predates these requests +ABSENT = "absent" # no daemon listens at the socket path + +# How long a starting engine waits for another engine's move to end. +PENDING_MOVE_TIMEOUT = 300.0 +_PENDING_MOVE_POLL_INTERVAL = 0.5 + + +def layout_digest(rows: Sequence[Sequence[int]]) -> str: + """SHA-256 of a physical_to_logical_map given as rows of ints.""" + text = json.dumps([[int(x) for x in row] for row in rows], separators=(",", ":")) + return hashlib.sha256(text.encode()).hexdigest() + + +def _busy(reason: str) -> Dict[str, Any]: + return {"status": BUSY, "reason": reason} + + +def _refused(reason: str) -> Dict[str, Any]: + return {"status": REFUSED, "reason": reason} + + +def _layers(layer_ids: Sequence[int]) -> str: + """Layer ids for messages, as runs: "layer 3", "layers 0-3,7".""" + runs: List[List[int]] = [] + for layer_id in sorted(layer_ids): + if runs and layer_id == runs[-1][1] + 1: + runs[-1][1] = layer_id + else: + runs.append([layer_id, layer_id]) + text = ",".join(str(a) if a == b else f"{a}-{b}" for a, b in runs) + return f"layer {text}" if len(layer_ids) == 1 else f"layers {text}" + + +def _is_int(value: Any) -> bool: + return isinstance(value, int) and not isinstance(value, bool) + + +_RESTART_HINT = ( + "Restart the weight cache daemons, and the engines that use them, so that " + "the weights are loaded from disk again." +) + + +# --------------------------------------------------------------------------- +# Daemon side +# --------------------------------------------------------------------------- + + +@dataclass +class _PendingMove: + engine_id: str + base_version: int + layer_ids: List[int] + started: float + + +class ExpertLayoutRegistry: + """Daemon side: the expert layout of the served weights, its version, the + pending move, and the attached engines. + + Every method takes the lock: the serve loop handles requests while the + presence thread reports engines that went away. + """ + + def __init__(self, rows: Optional[Sequence[Sequence[int]]]): + self._lock = threading.Lock() + self._rows = None if rows is None else [[int(x) for x in row] for row in rows] + self._digest = None if self._rows is None else layout_digest(self._rows) + self._version = 0 + self._pending: Optional[_PendingMove] = None + self._dirty: Optional[str] = None + self._attached: Dict[str, str] = {} # engine_id -> label + self._untracked_consumers = 0 + + @classmethod + def from_metadata(cls, metadata) -> "ExpertLayoutRegistry": + """Track the layout of the daemon's ExpertLocationMetadata; ``None`` + (no EPLB) tracks nothing.""" + if metadata is None: + return cls(None) + return cls(metadata.physical_to_logical_map_cpu.tolist()) + + @property + def tracked(self) -> bool: + return self._rows is not None + + def _label(self, engine_id: str) -> str: + return self._attached.get(engine_id, engine_id) + + def _snapshot_locked(self) -> Dict[str, Any]: + pending = self._pending + return { + "tracked": self.tracked, + "version": self._version, + "digest": self._digest, + "physical_to_logical_map": ( + None if self._rows is None else [list(row) for row in self._rows] + ), + "pending": ( + None + if pending is None + else { + "engine": self._label(pending.engine_id), + "base_version": pending.base_version, + "layer_ids": list(pending.layer_ids), + "seconds": time.monotonic() - pending.started, + } + ), + "dirty": self._dirty, + "attached": len(self._attached), + "untracked_consumers": self._untracked_consumers, + } + + def query(self) -> Dict[str, Any]: + with self._lock: + return {"status": OK, "layout": self._snapshot_locked()} + + def attach(self, engine_id: str, label: str) -> Dict[str, Any]: + with self._lock: + if engine_id in self._attached: + return _refused(f"engine {label} is attached already") + self._attached[engine_id] = label + logger.info( + f"[weight_cache] engine {label} attached " + f"({len(self._attached)} attached, layout version {self._version})" + ) + return {"status": OK, "layout": self._snapshot_locked()} + + def detach(self, engine_id: str) -> None: + """The engine's attach connection closed: it exited, however it did.""" + with self._lock: + label = self._label(engine_id) + self._attached.pop(engine_id, None) + pending = self._pending + if pending is not None and pending.engine_id == engine_id: + self._pending = None + self._dirty = ( + f"engine {label} went away while it moved the routed experts " + f"of {_layers(pending.layer_ids)} from layout version " + f"{pending.base_version}, so they may hold a mix of two layouts" + ) + logger.error( + f"[weight_cache] expert layout is now unknown: {self._dirty}" + ) + else: + logger.info( + f"[weight_cache] engine {label} detached " + f"({len(self._attached)} attached)" + ) + + def _check_mover_locked(self, engine_id: str) -> Optional[Dict[str, Any]]: + if not self.tracked: + return _refused( + "this weight cache daemon does not track the expert layout " + "(it runs without --enable-eplb)" + ) + if self._dirty is not None: + return _refused(f"the expert layout is unknown: {self._dirty}") + if engine_id not in self._attached: + return _refused("the engine is not attached to this weight cache daemon") + if self._untracked_consumers: + return _refused( + f"{self._untracked_consumers} engine(s) fetched these weights " + f"without following the expert layout and may still map them, so " + f"no expert may move until the weight cache daemon restarts" + ) + return None + + def _parse_layer_ids_locked(self, layer_ids: Any) -> Optional[List[int]]: + if not isinstance(layer_ids, (list, tuple)) or not layer_ids: + return None + if not all(_is_int(x) and 0 <= x < len(self._rows) for x in layer_ids): + return None + if len(set(layer_ids)) != len(layer_ids): + return None + return list(layer_ids) + + def begin( + self, engine_id: str, base_version: Any, layer_ids: Any + ) -> Dict[str, Any]: + """Mark a move of ``layer_ids`` pending, if this engine may move them now.""" + with self._lock: + refusal = self._check_mover_locked(engine_id) + if refusal is not None: + return refusal + if base_version != self._version: + return _refused( + f"the engine moves experts from layout version {base_version}, " + f"but the weights are in version {self._version}" + ) + layer_ids = self._parse_layer_ids_locked(layer_ids) + if layer_ids is None: + return _refused("the move names no valid layers") + if self._pending is not None: + if self._pending.engine_id == engine_id: + return _refused("the engine has a move pending already") + return _busy( + f"engine {self._label(self._pending.engine_id)} is moving experts" + ) + others = len(self._attached) - 1 + if others: + return _busy(f"{others} other engine(s) map these weights") + self._pending = _PendingMove( + engine_id, base_version, layer_ids, time.monotonic() + ) + return {"status": OK} + + def _pending_matches_locked( + self, engine_id: str, base_version: Any, layer_ids: Any + ) -> bool: + pending = self._pending + return ( + pending is not None + and pending.engine_id == engine_id + and pending.base_version == base_version + and isinstance(layer_ids, (list, tuple)) + and list(layer_ids) == pending.layer_ids + ) + + def commit( + self, engine_id: str, base_version: Any, layer_ids: Any, rows: Any + ) -> Dict[str, Any]: + """Store the new rows of the pending move: the weights have moved.""" + with self._lock: + if not self._pending_matches_locked(engine_id, base_version, layer_ids): + return _refused("no such move is pending") + num_physical_experts = len(self._rows[0]) if self._rows else 0 + if ( + not isinstance(rows, (list, tuple)) + or len(rows) != len(layer_ids) + or not all( + isinstance(row, (list, tuple)) + and len(row) == num_physical_experts + and all(_is_int(x) for x in row) + for row in rows + ) + ): + # The pending move stays: the weights moved, but to a layout + # the daemon cannot record. The engine going away makes it dirty. + return _refused( + f"the new rows do not fit {len(layer_ids)} layer(s) of " + f"{num_physical_experts} physical experts" + ) + for layer_id, row in zip(layer_ids, rows): + self._rows[layer_id] = list(row) + self._digest = layout_digest(self._rows) + self._version += 1 + self._pending = None + logger.info( + f"[weight_cache] expert layout version {self._version - 1} -> " + f"{self._version}: engine {self._label(engine_id)} moved the " + f"routed experts of {_layers(layer_ids)}" + ) + return {"status": OK, "version": self._version, "digest": self._digest} + + def abort( + self, engine_id: str, base_version: Any, layer_ids: Any + ) -> Dict[str, Any]: + """Drop the pending move: another rank cannot move, so nothing moved.""" + with self._lock: + if not self._pending_matches_locked(engine_id, base_version, layer_ids): + return _refused("no such move is pending") + self._pending = None + return {"status": OK} + + def check_fetch(self, layout: Optional[Dict[str, Any]]) -> Optional[str]: + """Why fetch_state must not serve the weights, or None to serve them. + + ``layout`` is the engine's ``expert_layout`` field: its engine_id, + version, and the digest of the map it routes with. Engines without the + field (older ones) route with their own initial layout. + """ + with self._lock: + if not self.tracked: + return None + if self._dirty is not None: + return ( + f"the expert layout of its weights is unknown: {self._dirty}. " + f"{_RESTART_HINT}" + ) + if self._pending is not None: + return ( + f"engine {self._label(self._pending.engine_id)} is moving the " + f"routed experts of {_layers(self._pending.layer_ids)}; retry " + f"when it is done" + ) + if layout is None: + if self._version: + return ( + f"EPLB moved the routed experts of its weights (expert layout " + f"version {self._version}), and an engine that does not " + f"follow the daemon's expert layout would route them with " + f"its initial layout. Upgrade the engine, or restart the " + f"weight cache daemon so that the weights are loaded from " + f"disk again." + ) + self._untracked_consumers += 1 + logger.warning( + "[weight_cache] an engine that does not follow the expert " + "layout fetched the weights; no expert may move until this " + "daemon restarts" + ) + return None + if layout.get("digest") != self._digest: + return ( + f"the expert layout the engine routes with (version " + f"{layout.get('version')}, digest {str(layout.get('digest'))[:12]}) " + f"is not the one its weights are in (version {self._version}, " + f"digest {self._digest[:12]})" + ) + if layout.get("engine_id") not in self._attached: + self._untracked_consumers += 1 + logger.warning( + "[weight_cache] an engine that is not attached fetched the " + "weights; no expert may move until this daemon restarts" + ) + return None + + +def handle_expert_layout_request( + registry: ExpertLayoutRegistry, req: Dict[str, Any] +) -> Dict[str, Any]: + """Daemon side: the reply to one of REQUEST_TYPES.""" + engine_id = req.get("engine_id") + if not isinstance(engine_id, str) or not engine_id: + return _refused("the request names no engine_id") + req_type = req.get("type") + if req_type == ATTACH: + return registry.attach(engine_id, str(req.get("label") or engine_id)) + if req_type == QUERY: + return registry.query() + if req_type == BEGIN: + return registry.begin(engine_id, req.get("base_version"), req.get("layer_ids")) + if req_type == COMMIT: + return registry.commit( + engine_id, req.get("base_version"), req.get("layer_ids"), req.get("rows") + ) + if req_type == ABORT: + return registry.abort(engine_id, req.get("base_version"), req.get("layer_ids")) + raise ValueError(f"not an expert layout request: {req_type!r}") + + +class ExpertLayoutPresence: + """Daemon side: holds the attach connections, and detaches an engine from + the registry when its connection closes.""" + + def __init__(self, registry: ExpertLayoutRegistry): + self._registry = registry + self._selector = selectors.DefaultSelector() + self._wakeup_recv, self._wakeup_send = socket.socketpair() + self._wakeup_recv.setblocking(False) + self._selector.register(self._wakeup_recv, selectors.EVENT_READ, None) + self._incoming: List[Tuple[socket.socket, str]] = [] + self._incoming_lock = threading.Lock() + self._stopping = threading.Event() + self._thread = threading.Thread( + target=self._run, name="weight-cache-expert-layout-presence", daemon=True + ) + + def start(self) -> None: + self._thread.start() + + def watch(self, conn: socket.socket, engine_id: str) -> None: + """Take over an attach connection until the engine closes it.""" + conn.setblocking(False) + with self._incoming_lock: + self._incoming.append((conn, engine_id)) + self._wakeup_send.send(b"\0") + + def stop(self) -> None: + if self._stopping.is_set(): + return + self._stopping.set() + self._wakeup_send.send(b"\0") + self._thread.join(timeout=5) + self._wakeup_send.close() + + def _run(self) -> None: + try: + while not self._stopping.is_set(): + for key, _ in self._selector.select(timeout=1.0): + if key.data is None: + self._drain_wakeups() + else: + self._check(key.fileobj, key.data) + with self._incoming_lock: + incoming, self._incoming = self._incoming, [] + for conn, engine_id in incoming: + self._selector.register(conn, selectors.EVENT_READ, engine_id) + finally: + for key in list(self._selector.get_map().values()): + key.fileobj.close() + self._selector.close() + + def _drain_wakeups(self) -> None: + try: + while self._wakeup_recv.recv(4096): + pass + except (BlockingIOError, InterruptedError): + pass + + def _check(self, conn: socket.socket, engine_id: str) -> None: + try: + data = conn.recv(1) + except (BlockingIOError, InterruptedError): + return + except OSError: + data = b"" + if data: + # An engine never writes on its attach connection after attaching. + logger.warning( + f"[weight_cache] engine {engine_id} wrote on its attach " + f"connection; treating it as gone" + ) + self._selector.unregister(conn) + conn.close() + self._registry.detach(engine_id) + + +# --------------------------------------------------------------------------- +# Engine side +# --------------------------------------------------------------------------- + + +def _connect(socket_path: str, timeout: float) -> Optional[socket.socket]: + """Connect to the daemon socket; None when no daemon listens there. + + Same rule as the weight loader: only a socket node owned by this user. + """ + try: + st = os.lstat(socket_path) + except FileNotFoundError: + return None + if not stat.S_ISSOCK(st.st_mode) or st.st_uid != os.getuid(): + raise RuntimeError( + f"[weight_cache] Refusing to connect: {socket_path} is not a socket " + f"owned by this user." + ) + sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) + sock.settimeout(timeout) + try: + sock.connect(socket_path) + except FileNotFoundError: + sock.close() + return None + except ConnectionRefusedError as e: + sock.close() + raise RuntimeError( + f"[weight_cache] The daemon socket exists at {socket_path} but refused " + f"the connection. The daemon may have crashed after creating the " + f"socket. Check daemon logs." + ) from e + except BaseException: + sock.close() + raise + return sock + + +class ExpertLayoutClient: + """Engine side: one rank's link to its daemon's expert layout record.""" + + def __init__(self, socket_path: str, label: str, timeout: float = 30.0): + self.socket_path = socket_path + self.label = label + self.timeout = timeout + self.engine_id = uuid.uuid4().hex + self._presence: Optional[socket.socket] = None + + @property + def attached(self) -> bool: + return self._presence is not None + + def attach(self) -> Tuple[str, Optional[Dict[str, Any]]]: + """Attach to the daemon and read its expert layout. + + Returns ("attached", layout), (LEGACY, None) for a daemon without the + expert layout requests, or (ABSENT, None) when no daemon listens. + """ + sock = _connect(self.socket_path, self.timeout) + if sock is None: + return ABSENT, None + try: + send_msg( + sock, {"type": ATTACH, "engine_id": self.engine_id, "label": self.label} + ) + reply = recv_msg(sock) + except BaseException: + sock.close() + raise + if reply.get("status") == OK: + # Kept open and never used again: the daemon sees it close when this + # process exits. + self._presence = sock + return "attached", reply["layout"] + sock.close() + if reply.get("status") == "error" and str(reply.get("message")).startswith( + "Unknown request type" + ): + return LEGACY, None + raise RuntimeError( + f"[weight_cache] The daemon at {self.socket_path} did not attach this " + f"engine: {reply.get('reason') or reply}" + ) + + def request(self, req_type: str, **fields: Any) -> Dict[str, Any]: + """One request on its own connection. A daemon that is gone, or does + not answer, reads as a refusal: callers are collective, so every rank + must get a reply to agree on.""" + try: + sock = _connect(self.socket_path, self.timeout) + if sock is None: + return _refused(f"no daemon listens at {self.socket_path} any more") + try: + send_msg( + sock, {"type": req_type, "engine_id": self.engine_id, **fields} + ) + return recv_msg(sock) + finally: + sock.close() + except Exception as e: + return _refused(f"the daemon at {self.socket_path} did not answer: {e}") + + def close(self) -> None: + if self._presence is not None: + self._presence.close() + self._presence = None + + +AllGather = Callable[[Any], List[Any]] + + +class ExpertLayoutSession: + """Engine side: what this rank agreed on with the others at startup, and + the two-phase move of experts in the daemon's weights. + + begin(), commit() and the first rebalance_mode() are collective: every rank + calls them at the same point, as EPLBManager.rebalance() does. + """ + + def __init__( + self, + *, + client: ExpertLayoutClient, + kind: str, + version: Optional[int], + all_gather: AllGather, + rebalance_off_reason: Optional[str], + ): + self.client = client + self.kind = kind + self.version = version + # Digest of the map this engine routes with; set once it is built. + self.digest: Optional[str] = None + # Set by the weight loader: True when the weights are the daemon's, + # False when they were loaded from disk instead. + self.weights_shared: Optional[bool] = None + self._all_gather = all_gather + self._off_reason = rebalance_off_reason + self._mode: Optional[Tuple[str, Optional[str]]] = None + + def fetch_fields(self) -> Dict[str, Any]: + """Extra fetch_state fields, so the daemon can check this engine's layout.""" + return { + "expert_layout": { + "engine_id": self.client.engine_id if self.client.attached else None, + "version": self.version, + "digest": self.digest, + } + } + + def rebalance_mode(self) -> Tuple[str, Optional[str]]: + """How a rebalance may move experts: ("private", None) when the weights + are this engine's own, ("coordinated", None) through begin() and + commit(), or ("off", reason).""" + if self._mode is None: + self._mode = self._agree_on_mode() + if self._mode[0] == "coordinated" and self._off_reason is not None: + return "off", self._off_reason + return self._mode + + def _agree_on_mode(self) -> Tuple[str, Optional[str]]: + if self.weights_shared is False: + local = ("private", None) + elif self.weights_shared is None: + local = ("off", "the weights are not loaded") + elif self._off_reason is not None: + local = ("off", self._off_reason) + else: + local = ("coordinated", None) + results = self._all_gather(local) + modes = {mode for mode, _ in results} + if len(modes) == 1: + return results[0] + reasons = [reason for mode, reason in results if mode == "off"] + return "off", ( + reasons[0] + if reasons + else "some ranks loaded the weights from disk and some from the daemon" + ) + + def begin(self, layer_ids: List[int]) -> Optional[str]: + """Mark a move of ``layer_ids`` pending on every rank's daemon. + + Returns None when every rank may move them now, else why not. A refusal + (not merely busy) turns rebalancing off for good. + """ + layer_ids = list(layer_ids) + reply = self.client.request( + BEGIN, base_version=self.version, layer_ids=layer_ids + ) + status = reply.get("status") + results = self._all_gather((status, reply.get("reason"))) + if all(s == OK for s, _ in results): + return None + if status == OK: + self.client.request(ABORT, base_version=self.version, layer_ids=layer_ids) + reason = "; ".join( + f"rank {rank}: {r}" for rank, (s, r) in enumerate(results) if s != OK + ) + if any(s not in (OK, BUSY) for s, _ in results): + self._off_reason = reason + return reason + + def commit(self, layer_ids: List[int], physical_to_logical_map) -> Optional[str]: + """Store the moved layers' rows on every rank's daemon. + + ``physical_to_logical_map`` is this engine's whole map after the move, + on CPU. Returns None when every daemon recorded the move and agrees + with this engine, else why not; rebalancing is then off for good. + """ + layer_ids = list(layer_ids) + expected_digest = layout_digest(physical_to_logical_map.tolist()) + reply = self.client.request( + COMMIT, + base_version=self.version, + layer_ids=layer_ids, + rows=physical_to_logical_map[layer_ids].tolist(), + ) + if reply.get("status") != OK: + local = (False, None, reply.get("reason")) + elif reply.get("digest") != expected_digest: + local = (False, reply.get("version"), "the daemon recorded another layout") + else: + local = (True, reply.get("version"), None) + results = self._all_gather(local) + if all(ok for ok, _, _ in results) and len({v for _, v, _ in results}) == 1: + self.version = reply["version"] + self.digest = expected_digest + return None + reason = "; ".join( + f"rank {rank}: {r or f'version {v}'}" + for rank, (ok, v, r) in enumerate(results) + if not ok or v != results[0][1] + ) + self._off_reason = reason + return reason + + +def decide_expert_layout( + summaries: Sequence[Dict[str, Any]], +) -> Tuple[bool, Optional[str]]: + """The startup decision, from every rank's attach summary in rank order. + + Returns whether to build the metadata from the daemons' layout, and why + rebalancing cannot go through the daemons (None if it can). Raises when + the layout of the shared weights is unknown or the ranks disagree on it. + """ + for rank, summary in enumerate(summaries): + if summary.get("error"): + raise RuntimeError( + f"[weight_cache] Rank {rank} could not read the expert layout of " + f"its weight cache daemon: {summary['error']}" + ) + if summary["dirty"]: + raise RuntimeError( + f"[weight_cache] The weight cache daemon of rank {rank} does not " + f"know the expert layout of its weights: {summary['dirty']}. " + f"{_RESTART_HINT}" + ) + tracked = [(rank, s) for rank, s in enumerate(summaries) if s["kind"] == TRACKED] + if len(tracked) == len(summaries): + if len({(s["version"], s["digest"]) for _, s in tracked}) > 1: + versions = ", ".join( + f"rank {rank}: version {s['version']}" for rank, s in tracked + ) + raise RuntimeError( + f"[weight_cache] The weight cache daemons disagree on the expert " + f"layout of their weights ({versions}), so it is unknown. " + f"{_RESTART_HINT}" + ) + return True, None + for rank, summary in tracked: + if summary["version"]: + raise RuntimeError( + f"[weight_cache] EPLB moved the routed experts in the weights of " + f"rank {rank}'s weight cache daemon (expert layout version " + f"{summary['version']}), but not every rank's daemon tracks the " + f"layout, so this engine cannot route with it. {_RESTART_HINT}" + ) + kinds = {s["kind"] for s in summaries} + if LEGACY in kinds: + return False, ( + "the weight cache daemon does not track the expert layout (an older " + "version), so a rebalance would move the experts under every engine " + "that maps them" + ) + if UNTRACKED in kinds: + return False, ( + "the weight cache daemon does not track the expert layout (start it " + "with --enable-eplb)" + ) + if kinds == {ABSENT}: + return False, "this engine did not attach to a weight cache daemon at startup" + return False, "not every rank found a weight cache daemon at startup" + + +def start_expert_layout_session( + *, + socket_path: str, + label: str, + all_gather: AllGather, + pending_move_timeout: float = PENDING_MOVE_TIMEOUT, +) -> Tuple[ExpertLayoutSession, Optional[List[List[int]]]]: + """Attach this rank to its daemon and agree with the other ranks. + + Collective. Returns the session and the daemons' physical_to_logical_map + if the engine must route with it, else None. + """ + client = ExpertLayoutClient(socket_path, label) + try: + try: + kind, layout = client.attach() + if layout is not None: + layout = _wait_for_pending_move(client, layout, pending_move_timeout) + kind = TRACKED if layout["tracked"] else UNTRACKED + summary = { + "kind": kind, + "dirty": None if layout is None else layout["dirty"], + "version": None if layout is None else layout["version"], + "digest": None if layout is None else layout["digest"], + } + except Exception as e: + # Still join the all_gather, so that every rank stops, not just this one. + summary = {"kind": None, "error": str(e), "dirty": None} + adopt, off_reason = decide_expert_layout(all_gather(summary)) + except BaseException: + client.close() + raise + session = ExpertLayoutSession( + client=client, + kind=kind, + version=summary["version"], + all_gather=all_gather, + rebalance_off_reason=off_reason, + ) + return session, layout["physical_to_logical_map"] if adopt else None + + +def _wait_for_pending_move( + client: ExpertLayoutClient, layout: Dict[str, Any], timeout: float +) -> Dict[str, Any]: + deadline = time.monotonic() + timeout + while layout["pending"] is not None and layout["dirty"] is None: + if time.monotonic() > deadline: + pending = layout["pending"] + raise RuntimeError( + f"[weight_cache] Engine {pending['engine']} has been moving the " + f"routed experts of {_layers(pending['layer_ids'])} " + f"in the weight cache daemon's weights for " + f"{pending['seconds']:.0f} s; this engine cannot learn their " + f"layout until the move ends." + ) + time.sleep(_PENDING_MOVE_POLL_INTERVAL) + reply = client.request(QUERY) + if reply.get("status") != OK: + raise RuntimeError( + f"[weight_cache] The daemon at {client.socket_path} did not report " + f"its expert layout: {reply.get('reason')}" + ) + layout = reply["layout"] + return layout + + +def _world_all_gather(obj: Any) -> List[Any]: + import torch.distributed as dist + + group = get_parallel().world_group.cpu_group + gathered = [None] * dist.get_world_size(group=group) + dist.all_gather_object(gathered, obj, group=group) + return gathered + + +def adopt_weight_cache_expert_layout( + metadata, *, model_config, moe_ep_rank: int, gpu_id: int +): + """The expert location metadata this engine must route with. + + With the weight cache on, attach to this rank's daemon and, when it tracks + the layout of its weights, build the metadata from the daemon's map: the + same path as ``--init-expert-location`` with that map. Collective; also + publishes the session for the weight loader and EPLBManager. + """ + if metadata is None or get_model().weight_cache_mode == "off": + return metadata + + from sglang.srt.eplb.expert_location import ExpertLocationMetadata + from sglang.srt.platforms import current_platform + + socket_path = get_model().weight_cache_socket or get_socket_path( + current_platform.get_device_uuid(gpu_id) + ) + session, rows = start_expert_layout_session( + socket_path=socket_path, + label=f"{socket.gethostname()}:{os.getpid()} tp_rank={get_parallel().tp_rank}", + all_gather=_world_all_gather, + ) + try: + if rows is not None and rows != metadata.physical_to_logical_map_cpu.tolist(): + _check_daemon_layout_fits(rows, metadata, version=session.version) + metadata = ExpertLocationMetadata.init_by_mapping( + model_config, physical_to_logical_map=rows, moe_ep_rank=moe_ep_rank + ) + logger.info( + f"[weight_cache] Routing with the weight cache daemon's expert " + f"layout (version {session.version}) instead of the initial one." + ) + except BaseException: + session.client.close() + raise + session.digest = layout_digest(metadata.physical_to_logical_map_cpu.tolist()) + set_expert_layout_session(session) + return metadata + + +def _check_daemon_layout_fits(rows, metadata, *, version: int) -> None: + shape = (len(rows), len(rows[0]) if rows else 0) + expected_shape = tuple(metadata.physical_to_logical_map_cpu.shape) + if shape != expected_shape: + raise RuntimeError( + f"[weight_cache] The weight cache daemon's expert layout has " + f"{shape[0]} layers x {shape[1]} physical experts, but this engine " + f"has {expected_shape[0]} x {expected_shape[1]}. Start the daemon " + f"and the engine with the same --ep-num-redundant-experts." + ) + if get_exec().moe.ep_dispatch_algorithm is None: + raise RuntimeError( + f"[weight_cache] The weight cache daemon's routed experts are in an " + f"expert layout (version {version}) other than this engine's initial " + f"one, and without --enable-eplb (or --init-expert-location) this " + f"engine routes as if every expert sat in its initial slot. Start " + f"the engine with --enable-eplb, or restart the weight cache daemon." + ) + + +def get_expert_layout_session() -> Optional[ExpertLayoutSession]: + return get_resources().expert_layout_session + + +def set_expert_layout_session(session: Optional[ExpertLayoutSession]) -> None: + get_resources().expert_layout_session = session + + +def mark_expert_weights_shared(shared: bool) -> None: + """Weight loader: record whether the weights are the daemon's (True) or + were loaded from disk (False).""" + session = get_expert_layout_session() + if session is not None: + session.weights_shared = shared diff --git a/python/sglang/srt/weight_cache/ipc_loader.py b/python/sglang/srt/weight_cache/ipc_loader.py index 6bf6a9305b4b..b2d25129cef8 100644 --- a/python/sglang/srt/weight_cache/ipc_loader.py +++ b/python/sglang/srt/weight_cache/ipc_loader.py @@ -25,6 +25,11 @@ from sglang.srt.platforms import current_platform from sglang.srt.runtime_context import get_exec, get_parallel +from .expert_layout import ( + LAYOUT_REFUSED, + get_expert_layout_session, + mark_expert_weights_shared, +) from .protocol import ( CacheConfig, check_ipc_quant_support, @@ -121,6 +126,7 @@ def load_model( "[IpcModelLoader] Weight cache not available or config mismatch, " "falling back to disk load" ) + mark_expert_weights_shared(False) return self._fallback_load(model_config, device_config) entries = cache_data["entries"] @@ -178,6 +184,7 @@ def load_model( f"total={time.perf_counter() - tic:.2f}s" ) + mark_expert_weights_shared(True) return model.eval() def _start_daemon_liveness_watchdog(self, daemon_pid: Optional[int]) -> None: @@ -555,9 +562,21 @@ def _fetch_from_cache(self, model_config, device_config) -> Optional[dict]: f"dtype={engine_config.dtype}" ) - send_msg(sock, {"type": "fetch_state", "config": engine_config.to_dict()}) + request = {"type": "fetch_state", "config": engine_config.to_dict()} + expert_layout_session = get_expert_layout_session() + if expert_layout_session is not None: + # Lets the daemon check that this engine routes with the expert + # layout its weights are in. + request.update(expert_layout_session.fetch_fields()) + send_msg(sock, request) result = recv_msg(sock) + if result.get("status") == LAYOUT_REFUSED: + raise RuntimeError( + f"[IpcModelLoader] The weight cache daemon at " + f"{self.socket_path} refused to serve its weights: " + f"{result.get('message')}" + ) if result.get("status") != "ok": daemon_config = result.get("daemon_config", {}) raise RuntimeError( diff --git a/test/registered/unit/eplb/test_eplb_weight_cache.py b/test/registered/unit/eplb/test_eplb_weight_cache.py new file mode 100644 index 000000000000..b82e654122ad --- /dev/null +++ b/test/registered/unit/eplb/test_eplb_weight_cache.py @@ -0,0 +1,177 @@ +import unittest +from types import SimpleNamespace +from unittest.mock import Mock, patch + +import torch + +from sglang.srt.eplb.eplb_manager import EPLBManager +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=2, suite="base-a-test-cpu") + +_MODULE = "sglang.srt.eplb.eplb_manager" + + +class _Session: + """Stands in for the engine's ExpertLayoutSession; records the moves.""" + + def __init__(self, calls, mode, reason=None, begin_reasons=()): + self._calls = calls + self._mode = (mode, reason) + self._begin_reasons = list(begin_reasons) + + def rebalance_mode(self): + return self._mode + + def begin(self, layer_ids): + self._calls.append(("begin", list(layer_ids))) + return self._begin_reasons.pop(0) if self._begin_reasons else None + + def commit(self, layer_ids, physical_to_logical_map): + self._calls.append( + ("commit", list(layer_ids), physical_to_logical_map.tolist()) + ) + return None + + +class TestEPLBWithWeightCache(unittest.TestCase): + def setUp(self): + self.calls = [] + + def _manager( + self, weight_cache_mode, session=None, *, busy_gpu=False, elastic_ep=None + ): + moe = SimpleNamespace( + eplb_rebalance_layers_per_chunk=1, + eplb_rebalance_num_iterations=2, + expert_distribution_recorder_buffer_size=2, + eplb_min_rebalancing_utilization_threshold=0.5, + ep_dispatch_algorithm="dynamic", + elastic_ep_backend=elastic_ep, + ) + self.recorder = Mock(recording=True) + # A busy GPU makes a rebalance stop right after reading the counts. + self.recorder.dump_record.return_value = { + "logical_count": None, + "average_utilization_rate_over_window": 1.0 if busy_gpu else 0.0, + } + self.metadata = SimpleNamespace( + physical_to_logical_map_cpu=torch.zeros(2, 4, dtype=torch.int64) + ) + + def move(*, update_layer_ids, **_): + # The engine's own map follows the moved layers. + for layer_id in update_layer_ids: + self.metadata.physical_to_logical_map_cpu[layer_id] = layer_id + 1 + self.calls.append(("move", list(update_layer_ids))) + + for name, replacement in ( + ("get_exec", lambda: SimpleNamespace(moe=moe)), + ("get_model", lambda: SimpleNamespace(weight_cache_mode=weight_cache_mode)), + ("get_parallel", lambda: SimpleNamespace(tp_rank=0)), + ("get_global_expert_distribution_recorder", lambda: self.recorder), + ("get_global_expert_location_metadata", lambda: self.metadata), + ("get_expert_layout_session", lambda: session), + ("update_expert_location_with_recovery", move), + ): + patcher = patch(f"{_MODULE}.{name}", replacement) + patcher.start() + self.addCleanup(patcher.stop) + # A commit waits for the moves to land. CPU CI has a CUDA build of + # torch and no driver, so a real device sync would raise there. + device = Mock() + device.synchronize.side_effect = lambda: self.calls.append(("sync",)) + patcher = patch.object(torch, "get_device_module", return_value=device) + patcher.start() + self.addCleanup(patcher.stop) + patcher = patch.object( + EPLBManager, "_compute_expert_location_metadata", return_value=None + ) + patcher.start() + self.addCleanup(patcher.stop) + return EPLBManager( + model_config=None, + get_model=lambda: SimpleNamespace( + routed_experts_weights_of_layer={0: [], 1: []} + ), + get_expert_location_updater=Mock(), + get_expert_backup_client=Mock(), + get_weight_updater=Mock(), + ) + + def test_rebalancing_stays_off_without_a_layout_record_to_follow(self): + for session, elastic_ep, reason in ( + (None, None, "no expert layout session"), + ( + _Session(self.calls, "off", "the daemon does not track it"), + None, + "the daemon does not track it", + ), + (_Session(self.calls, "coordinated"), "nixl", "elastic EP"), + ): + with self.subTest(reason=reason): + manager = self._manager("client", session, elastic_ep=elastic_ep) + with self.assertLogs(_MODULE, "WARNING") as logs: + self.assertEqual(list(manager.rebalance()), []) + self.assertIn( + "rebalancing is off with the weight cache", logs.output[0] + ) + self.assertIn(reason, logs.output[0]) + # Once per reason; elastic EP re-enabling rebalancing changes nothing. + manager.enable_rebalance() + with self.assertNoLogs(_MODULE, "WARNING"): + for _ in range(10): + manager.on_forward_pass_end() + self.recorder.dump_record.assert_not_called() + self.assertEqual(self.calls, []) + + def test_weights_loaded_from_disk_rebalance_as_usual(self): + manager = self._manager("client", _Session(self.calls, "private")) + self.assertEqual(list(manager.rebalance()), [None, None]) + self.assertEqual(self.calls, [("move", [0]), ("move", [1])]) + + def test_shared_weights_move_through_the_daemon(self): + manager = self._manager("daemon", _Session(self.calls, "coordinated")) + list(manager.rebalance()) + self.assertEqual( + self.calls, + [ + ("begin", [0]), + ("move", [0]), + ("sync",), + ("commit", [0], [[1] * 4, [0] * 4]), + ("begin", [1]), + ("move", [1]), + ("sync",), + ("commit", [1], [[1] * 4, [2] * 4]), + ], + ) + + def test_a_busy_daemon_stops_the_rebalance_before_the_next_move(self): + session = _Session( + self.calls, + "coordinated", + begin_reasons=[None, "rank 1: 1 other engine(s) map these weights"], + ) + manager = self._manager("client", session) + list(manager.rebalance()) + self.assertEqual( + self.calls, + [ + ("begin", [0]), + ("move", [0]), + ("sync",), + ("commit", [0], [[1] * 4, [0] * 4]), + ("begin", [1]), + ], + ) + + def test_rebalance_runs_without_the_weight_cache(self): + manager = self._manager("off", busy_gpu=True) + with self.assertNoLogs(_MODULE, "WARNING"): + self.assertEqual(list(manager.rebalance()), []) + self.recorder.dump_record.assert_called_once() + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/model_loader/test_weight_cache_expert_layout.py b/test/registered/unit/model_loader/test_weight_cache_expert_layout.py new file mode 100644 index 000000000000..39490ba4d12c --- /dev/null +++ b/test/registered/unit/model_loader/test_weight_cache_expert_layout.py @@ -0,0 +1,1052 @@ +""" +CPU-only unit tests for the weight cache daemon's record of its expert layout +(sglang.srt.weight_cache.expert_layout) and its users: + + - the record: the initial layout, begin/commit/abort and versions, the rule + that only an engine alone on the daemon moves experts, and the fetch_state + check + - the daemon's socket handling, including an engine killed while it moves + experts + - the startup agreement across ranks, with threads standing in for ranks + - building the expert location metadata from the daemon's map, which must + equal --init-expert-location with that map + - daemons and engines that predate the record + +No CUDA, no model, no daemon process. +""" + +import json +import os +import shutil +import socket +import subprocess +import sys +import tempfile +import threading +import time +import unittest +from types import SimpleNamespace +from unittest import mock + +import torch + +from sglang.srt.configs.load_config import LoadConfig, LoadFormat +from sglang.srt.eplb.expert_location import ( + ModelConfigForExpertLocation, + compute_initial_expert_location_metadata, +) +from sglang.srt.runtime_context import get_context, get_parallel, get_resources +from sglang.srt.weight_cache import expert_layout, ipc_loader +from sglang.srt.weight_cache.daemon import CLIENT_CONNECTION_TIMEOUT, WeightCacheDaemon +from sglang.srt.weight_cache.expert_layout import ( + ABSENT, + ATTACH, + BEGIN, + BUSY, + COMMIT, + LAYOUT_REFUSED, + LEGACY, + OK, + QUERY, + REFUSED, + TRACKED, + UNTRACKED, + ExpertLayoutClient, + ExpertLayoutPresence, + ExpertLayoutRegistry, + adopt_weight_cache_expert_layout, + decide_expert_layout, + get_expert_layout_session, + layout_digest, + mark_expert_weights_shared, + start_expert_layout_session, +) +from sglang.srt.weight_cache.protocol import CacheConfig, recv_msg, send_msg +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=20, suite="base-a-test-cpu") + +NUM_LAYERS = 3 +NUM_LOGICAL = 8 +NUM_PHYSICAL = 10 # two redundant experts over EP 2 + + +def _initial_rows(): + # The trivial layout: physical slot p holds logical expert p % 8. + return [[p % NUM_LOGICAL for p in range(NUM_PHYSICAL)] for _ in range(NUM_LAYERS)] + + +def _moved_row(layer_id, shift): + # Another layout that still holds every logical expert. + return [(3 * p + layer_id + shift) % NUM_LOGICAL for p in range(NUM_PHYSICAL)] + + +def _fetch_state(**fields): + """A fetch_state request; the test daemon matches every engine config.""" + config = CacheConfig( + model_path="/models/demo", + model_arch="Qwen3MoeForCausalLM", + tp_size=2, + tp_rank=0, + pp_size=1, + pp_rank=0, + dp_size=1, + ep_size=2, + moe_dp_size=1, + moe_dp_rank=0, + moe_ep_rank=0, + attn_dp_size=1, + enable_dp_lm_head=False, + attn_cp_size=1, + moe_dense_tp_size=None, + moe_a2a_backend="none", + quant_method="", + quant_config_hash="", + dtype="torch.bfloat16", + revision="", + device_capability="", + torch_version="", + ) + return {"type": "fetch_state", "config": config.to_dict(), **fields} + + +def _wait_for(predicate, timeout=10.0): + deadline = time.monotonic() + timeout + while not predicate(): + if time.monotonic() > deadline: + raise AssertionError("timed out waiting for the daemon") + time.sleep(0.02) + + +class _DaemonServer: + """A WeightCacheDaemon's request handling on a real Unix socket, served the + way serve() serves it, without a model, a ready file or signal handlers.""" + + def __init__(self, rows, transport_name="torch_ipc"): + self.dir = tempfile.mkdtemp(prefix="wcl") + self.path = os.path.join(self.dir, "d.sock") + daemon = WeightCacheDaemon.__new__(WeightCacheDaemon) + daemon.gpu_id = 0 + # Every engine config matches: these tests are about the expert layout. + daemon.config = mock.Mock() + daemon.config.matches.return_value = True + daemon.config.to_dict.return_value = {} + daemon.state_entries = {"w": {"shape": [1]}} + daemon.preloaded_weights_bytes = 0 + daemon.transport_backend = mock.Mock() + daemon.transport_backend.name = transport_name + daemon.transport_backend.send_fetch_state_response.side_effect = ( + lambda conn, **fields: send_msg( + conn, {"status": "ok", "transport_backend": transport_name, **fields} + ) + ) + daemon.expert_layout = ExpertLayoutRegistry(rows) + daemon.expert_layout_presence = ExpertLayoutPresence(daemon.expert_layout) + daemon.expert_layout_presence.start() + self.daemon = daemon + self._listener = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) + self._listener.bind(self.path) + self._listener.listen(8) + self._listener.settimeout(0.1) + self._stopping = threading.Event() + self._thread = threading.Thread(target=self._serve, daemon=True) + self._thread.start() + + @property + def registry(self) -> ExpertLayoutRegistry: + return self.daemon.expert_layout + + def layout(self): + return self.registry.query()["layout"] + + def _serve(self): + while not self._stopping.is_set(): + try: + conn, _ = self._listener.accept() + except socket.timeout: + continue + conn.settimeout(CLIENT_CONNECTION_TIMEOUT) + self.daemon._serve_connection(conn) + + def request(self, req): + with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as sock: + sock.connect(self.path) + send_msg(sock, req) + return recv_msg(sock) + + def close(self): + self._stopping.set() + self._thread.join(timeout=5) + self._listener.close() + self.daemon.expert_layout_presence.stop() + shutil.rmtree(self.dir, ignore_errors=True) + + +class _FakeWorld: + """all_gather over threads that stand in for the ranks of one engine.""" + + def __init__(self, size): + self._slots = [None] * size + self._barrier = threading.Barrier(size) + + def gather(self, rank): + def all_gather(obj): + self._slots[rank] = obj + self._barrier.wait(timeout=30) + gathered = list(self._slots) + self._barrier.wait(timeout=30) + return gathered + + return all_gather + + +def _run_ranks(*fns): + """Run one function per rank in threads; return (results, errors).""" + results = [None] * len(fns) + errors = [None] * len(fns) + + def run(rank): + try: + results[rank] = fns[rank]() + except BaseException as e: + errors[rank] = e + + threads = [threading.Thread(target=run, args=(rank,)) for rank in range(len(fns))] + for thread in threads: + thread.start() + for thread in threads: + thread.join(timeout=60) + return results, errors + + +def _move(path, moves, label="previous-engine"): + """An engine that attaches alone, moves (layer_id, row) pairs one by one, + and exits.""" + client = ExpertLayoutClient(path, label) + _, layout = client.attach() + version = layout["version"] + for layer_id, row in moves: + assert ( + client.request(BEGIN, base_version=version, layer_ids=[layer_id])["status"] + == OK + ) + reply = client.request( + COMMIT, base_version=version, layer_ids=[layer_id], rows=[row] + ) + assert reply["status"] == OK, reply + version = reply["version"] + client.close() + + +# A separate engine process: attaches, starts a move of layers 1-2, says so, +# and waits to be killed. +_DOOMED_ENGINE = f""" +import pickle, socket, struct, sys, time + +def send(sock, obj): + data = pickle.dumps(obj) + sock.sendall(struct.pack("!I", len(data)) + data) + +def recv(sock): + def exact(n): + buf = b"" + while len(buf) < n: + chunk = sock.recv(n - len(buf)) + if not chunk: + raise ConnectionError("closed") + buf += chunk + return buf + return pickle.loads(exact(struct.unpack("!I", exact(4))[0])) + +presence = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) +presence.connect(sys.argv[1]) +send(presence, {{"type": "{ATTACH}", "engine_id": "doomed", "label": "doomed-engine"}}) +assert recv(presence)["status"] == "{OK}" +move = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) +move.connect(sys.argv[1]) +send(move, {{"type": "{BEGIN}", "engine_id": "doomed", "base_version": 0, "layer_ids": [1, 2]}}) +print(recv(move)["status"], flush=True) +time.sleep(120) +""" + + +class TestExpertLayoutRegistry(CustomTestCase): + def _registry(self): + registry = ExpertLayoutRegistry(_initial_rows()) + registry.attach("a", "engine-a") + return registry + + def test_initial_layout(self): + self.assertFalse(ExpertLayoutRegistry.from_metadata(None).tracked) + metadata = SimpleNamespace( + physical_to_logical_map_cpu=torch.tensor(_initial_rows()) + ) + registry = ExpertLayoutRegistry.from_metadata(metadata) + layout = registry.query()["layout"] + self.assertTrue(layout["tracked"]) + self.assertEqual(layout["version"], 0) + self.assertEqual(layout["physical_to_logical_map"], _initial_rows()) + self.assertEqual(layout["digest"], layout_digest(_initial_rows())) + self.assertIsNone(layout["pending"]) + self.assertIsNone(layout["dirty"]) + # A snapshot is a copy. + layout["physical_to_logical_map"][0][0] = 7 + self.assertEqual( + registry.query()["layout"]["physical_to_logical_map"], _initial_rows() + ) + + def test_begin_commit_and_versions(self): + registry = self._registry() + self.assertEqual(registry.begin("a", 0, [0])["status"], OK) + pending = registry.query()["layout"]["pending"] + self.assertEqual( + (pending["engine"], pending["base_version"], pending["layer_ids"]), + ("engine-a", 0, [0]), + ) + expected = _initial_rows() + expected[0] = _moved_row(0, 1) + reply = registry.commit("a", 0, [0], [_moved_row(0, 1)]) + self.assertEqual((reply["status"], reply["version"]), (OK, 1)) + self.assertEqual(reply["digest"], layout_digest(expected)) + + # The next chunk starts from the new version. + self.assertEqual(registry.begin("a", 1, [1, 2])["status"], OK) + reply = registry.commit("a", 1, [1, 2], [_moved_row(1, 1), _moved_row(2, 1)]) + self.assertEqual(reply["version"], 2) + expected[1:] = [_moved_row(1, 1), _moved_row(2, 1)] + layout = registry.query()["layout"] + self.assertEqual(layout["physical_to_logical_map"], expected) + self.assertIsNone(layout["pending"]) + + # An aborted move leaves the version and the rows alone. + self.assertEqual(registry.begin("a", 2, [0])["status"], OK) + self.assertEqual(registry.abort("a", 2, [0])["status"], OK) + layout = registry.query()["layout"] + self.assertEqual((layout["version"], layout["pending"]), (2, None)) + self.assertEqual(layout["physical_to_logical_map"], expected) + + def test_malformed_moves_are_refused(self): + registry = self._registry() + for base_version, layer_ids in ( + (1, [0]), + (0, []), + (0, [3]), + (0, [0, 0]), + (0, [True]), + (0, "0"), + ): + with self.subTest(base_version=base_version, layer_ids=layer_ids): + reply = registry.begin("a", base_version, layer_ids) + self.assertEqual(reply["status"], REFUSED) + # Nothing is pending yet. + reply = registry.commit("a", 0, [0], [_moved_row(0, 1)]) + self.assertEqual(reply["status"], REFUSED) + + self.assertEqual(registry.begin("a", 0, [0, 1])["status"], OK) + # One move at a time; a commit must match it and fit its layers. + self.assertEqual(registry.begin("a", 0, [2])["status"], REFUSED) + for layer_ids, rows in ( + ([0], [_moved_row(0, 1)]), + ([0, 1], [_moved_row(0, 1)]), + ([0, 1], [_moved_row(0, 1), [0] * 4]), + ([0, 1], [_moved_row(0, 1), [0.5] * NUM_PHYSICAL]), + ): + with self.subTest(layer_ids=layer_ids, rows=rows): + self.assertEqual( + registry.commit("a", 0, layer_ids, rows)["status"], REFUSED + ) + # Still pending: the weights moved to a layout the daemon could not + # record, and the engine going away will leave the daemon dirty. + layout = registry.query()["layout"] + self.assertIsNotNone(layout["pending"]) + self.assertEqual(layout["version"], 0) + + def test_only_an_engine_alone_on_the_daemon_moves_experts(self): + registry = self._registry() + registry.attach("b", "engine-b") + self.assertEqual(registry.begin("a", 0, [0])["status"], BUSY) + self.assertEqual(registry.begin("b", 0, [0])["status"], BUSY) + registry.detach("b") + self.assertEqual(registry.begin("a", 0, [0])["status"], OK) + + # An engine attaching during a move sees it pending ... + reply = registry.attach("c", "engine-c") + self.assertEqual(reply["layout"]["pending"]["layer_ids"], [0]) + self.assertEqual(registry.begin("c", 0, [1])["status"], BUSY) + # ... the move ends, and then nothing moves while both are attached. + self.assertEqual(registry.commit("a", 0, [0], [_moved_row(0, 1)])["status"], OK) + self.assertEqual(registry.begin("a", 1, [1])["status"], BUSY) + self.assertEqual(registry.begin("c", 1, [1])["status"], BUSY) + + # Takeover: once the first engine is gone, the second one may move. + registry.detach("a") + self.assertIsNone(registry.query()["layout"]["dirty"]) + self.assertEqual(registry.begin("c", 1, [1])["status"], OK) + # An engine that is not attached never moves experts. + self.assertEqual(registry.begin("d", 1, [2])["status"], REFUSED) + + def test_engine_gone_mid_move_leaves_the_layout_unknown(self): + registry = self._registry() + self.assertEqual(registry.begin("a", 0, [1, 2])["status"], OK) + registry.detach("a") + layout = registry.query()["layout"] + self.assertIsNone(layout["pending"]) + self.assertIn("engine-a went away", layout["dirty"]) + self.assertIn("layers 1-2", layout["dirty"]) + + registry.attach("b", "engine-b") + self.assertEqual(registry.begin("b", 0, [0])["status"], REFUSED) + for fetch in ( + None, + {"engine_id": "b", "version": 0, "digest": layout["digest"]}, + ): + refusal = registry.check_fetch(fetch) + self.assertIn("is unknown", refusal) + self.assertIn("Restart the weight cache daemons", refusal) + + def test_fetch_check(self): + registry = self._registry() + current = { + "engine_id": "a", + "version": 0, + "digest": layout_digest(_initial_rows()), + } + self.assertIsNone(registry.check_fetch(current)) + self.assertIn( + "is not the one its weights are in", + registry.check_fetch(dict(current, digest="stale")), + ) + # While a move is in progress, no engine gets the weights. + registry.begin("a", 0, [0]) + self.assertIn("is moving", registry.check_fetch(current)) + registry.commit("a", 0, [0], [_moved_row(0, 1)]) + # An engine without the expert_layout field cannot follow a moved layout. + self.assertIn("Upgrade the engine", registry.check_fetch(None)) + # A daemon without EPLB tracks, and so checks, nothing. + self.assertIsNone(ExpertLayoutRegistry(None).check_fetch(None)) + + def test_engines_that_do_not_follow_the_layout_stop_all_moves(self): + unattached = { + "engine_id": "stranger", + "version": 0, + "digest": layout_digest(_initial_rows()), + } + for fetch in (None, unattached): + with self.subTest(fetch=fetch): + registry = self._registry() + self.assertIsNone(registry.check_fetch(fetch)) + reply = registry.begin("a", 0, [0]) + self.assertEqual(reply["status"], REFUSED) + self.assertIn("until the weight cache daemon restarts", reply["reason"]) + + +class TestDaemonExpertLayoutRequests(CustomTestCase): + def setUp(self): + self.server = _DaemonServer(_initial_rows()) + self.addCleanup(self.server.close) + + def _client(self, label="engine"): + client = ExpertLayoutClient(self.server.path, label) + self.addCleanup(client.close) + return client + + def test_attach_move_and_query_over_the_socket(self): + client = self._client() + kind, layout = client.attach() + self.assertEqual(kind, "attached") + self.assertEqual((layout["version"], layout["attached"]), (0, 1)) + reply = client.request(BEGIN, base_version=0, layer_ids=[2]) + self.assertEqual(reply["status"], OK) + reply = client.request( + COMMIT, base_version=0, layer_ids=[2], rows=[_moved_row(2, 5)] + ) + self.assertEqual((reply["status"], reply["version"]), (OK, 1)) + layout = client.request(QUERY)["layout"] + self.assertEqual(layout["physical_to_logical_map"][2], _moved_row(2, 5)) + + def test_closing_the_attach_connection_detaches(self): + client = self._client() + client.attach() + self.assertEqual(self.server.layout()["attached"], 1) + client.close() + _wait_for(lambda: self.server.layout()["attached"] == 0) + self.assertIsNone(self.server.layout()["dirty"]) + + def test_engine_killed_mid_move_leaves_the_daemon_dirty(self): + # A separate process, killed with SIGKILL while its move is pending: + # the kernel closing its attach connection is all the daemon sees. + child = subprocess.Popen( + [sys.executable, "-c", _DOOMED_ENGINE, self.server.path], + stdout=subprocess.PIPE, + text=True, + ) + try: + self.assertEqual(child.stdout.readline().strip(), OK) + self.assertIsNotNone(self.server.layout()["pending"]) + finally: + child.kill() + child.wait(timeout=10) + child.stdout.close() + _wait_for(lambda: self.server.layout()["dirty"] is not None) + layout = self.server.layout() + self.assertIn("doomed-engine went away", layout["dirty"]) + self.assertEqual(layout["version"], 0) + + # No engine gets the weights, old or new ... + for fields in ( + {}, + { + "expert_layout": { + "engine_id": "x", + "version": 0, + "digest": layout["digest"], + } + }, + ): + reply = self.server.request(_fetch_state(**fields)) + self.assertEqual(reply["status"], LAYOUT_REFUSED) + self.assertIn("Restart the weight cache daemons", reply["message"]) + self.server.daemon.transport_backend.send_fetch_state_response.assert_not_called() + # ... and a starting engine stops with the reason. + with self.assertRaisesRegex(RuntimeError, "doomed-engine went away"): + start_expert_layout_session( + socket_path=self.server.path, + label="next-engine", + all_gather=lambda obj: [obj], + ) + + def test_layout_check_comes_before_the_transport(self): + for transport_name in ("torch_ipc", "vmm_fd"): + with self.subTest(transport=transport_name): + server = _DaemonServer(_initial_rows(), transport_name=transport_name) + self.addCleanup(server.close) + client = ExpertLayoutClient(server.path, "engine") + self.addCleanup(client.close) + client.attach() + send = server.daemon.transport_backend.send_fetch_state_response + fields = { + "engine_id": client.engine_id, + "version": 0, + "digest": "stale", + } + reply = server.request(_fetch_state(expert_layout=fields)) + self.assertEqual(reply["status"], LAYOUT_REFUSED) + send.assert_not_called() + fields["digest"] = layout_digest(_initial_rows()) + reply = server.request(_fetch_state(expert_layout=fields)) + self.assertEqual( + (reply["status"], reply["transport_backend"]), + ("ok", transport_name), + ) + send.assert_called_once() + + +class TestStartupAgreement(CustomTestCase): + @staticmethod + def _tracked(version=0, digest="d", dirty=None): + return {"kind": TRACKED, "version": version, "digest": digest, "dirty": dirty} + + @staticmethod + def _other(kind): + return {"kind": kind, "version": None, "digest": None, "dirty": None} + + def test_decide(self): + tracked, other = self._tracked, self._other + self.assertEqual(decide_expert_layout([tracked(3), tracked(3)]), (True, None)) + for summaries, reason in ( + ([tracked(), other(LEGACY)], "older version"), + ([other(UNTRACKED), other(UNTRACKED)], "--enable-eplb"), + ([other(ABSENT), other(ABSENT)], "did not attach"), + ([tracked(), other(ABSENT)], "not every rank"), + ): + with self.subTest(kinds=[s["kind"] for s in summaries]): + adopt, off_reason = decide_expert_layout(summaries) + self.assertFalse(adopt) + self.assertIn(reason, off_reason) + for summaries, error in ( + ([tracked(3), tracked(4)], "disagree"), + ([tracked(3, "a"), tracked(3, "b")], "disagree"), + ([tracked(), tracked(dirty="engine x went away")], "engine x went away"), + ([tracked(2), other(UNTRACKED)], "not every rank's daemon"), + ): + with self.subTest(summaries=summaries): + with self.assertRaisesRegex(RuntimeError, error): + decide_expert_layout(summaries) + + def _start(self, servers, **kwargs): + world = _FakeWorld(len(servers)) + return _run_ranks( + *( + lambda rank=rank: start_expert_layout_session( + socket_path=servers[rank].path, + label=f"rank{rank}", + all_gather=world.gather(rank), + **kwargs, + ) + for rank in range(len(servers)) + ) + ) + + def _servers(self, count): + servers = [_DaemonServer(_initial_rows()) for _ in range(count)] + for server in servers: + self.addCleanup(server.close) + return servers + + def test_every_rank_adopts_the_daemons_layout(self): + servers = self._servers(2) + moved = [(layer_id, _moved_row(layer_id, 2)) for layer_id in range(NUM_LAYERS)] + for server in servers: + _move(server.path, moved) + results, errors = self._start(servers) + self.assertEqual(errors, [None, None]) + for session, rows in results: + self.addCleanup(session.client.close) + self.assertEqual(rows, [row for _, row in moved]) + self.assertEqual((session.kind, session.version), (TRACKED, NUM_LAYERS)) + + def test_daemons_at_different_versions_stop_every_rank(self): + servers = self._servers(2) + # Rank 1's daemon recorded a move that rank 0's did not, as after an + # engine died between the two commits. + _move(servers[1].path, [(0, _moved_row(0, 1))]) + _, errors = self._start(servers) + for error in errors: + self.assertIsInstance(error, RuntimeError) + self.assertIn("disagree", str(error)) + # The failed start leaves nothing attached. + for server in servers: + _wait_for(lambda server=server: server.layout()["attached"] == 0) + + def test_a_starting_engine_waits_for_a_move_in_progress(self): + (server,) = self._servers(1) + mover = ExpertLayoutClient(server.path, "mover") + self.addCleanup(mover.close) + mover.attach() + self.assertEqual( + mover.request(BEGIN, base_version=0, layer_ids=[0])["status"], OK + ) + commit = threading.Timer( + 0.5, + lambda: mover.request( + COMMIT, base_version=0, layer_ids=[0], rows=[_moved_row(0, 3)] + ), + ) + commit.start() + self.addCleanup(commit.cancel) + session, rows = start_expert_layout_session( + socket_path=server.path, label="next-engine", all_gather=lambda obj: [obj] + ) + self.addCleanup(session.client.close) + self.assertEqual(session.version, 1) + self.assertEqual(rows[0], _moved_row(0, 3)) + + def test_a_stuck_move_stops_the_start_with_the_reason(self): + (server,) = self._servers(1) + mover = ExpertLayoutClient(server.path, "stuck-engine") + self.addCleanup(mover.close) + mover.attach() + mover.request(BEGIN, base_version=0, layer_ids=[1]) + with self.assertRaisesRegex(RuntimeError, "stuck-engine has been moving"): + start_expert_layout_session( + socket_path=server.path, + label="next-engine", + all_gather=lambda obj: [obj], + pending_move_timeout=0.3, + ) + + def test_a_rank_that_cannot_read_its_daemon_stops_every_rank(self): + (server,) = self._servers(1) + # Rank 1's socket path holds a plain file, which an engine refuses. + not_a_socket = os.path.join(server.dir, "not-a-socket") + open(not_a_socket, "w").close() + world = _FakeWorld(2) + _, errors = _run_ranks( + lambda: start_expert_layout_session( + socket_path=server.path, label="rank0", all_gather=world.gather(0) + ), + lambda: start_expert_layout_session( + socket_path=not_a_socket, label="rank1", all_gather=world.gather(1) + ), + ) + for error in errors: + self.assertIsInstance(error, RuntimeError) + self.assertIn("Rank 1 could not read", str(error)) + _wait_for(lambda: server.layout()["attached"] == 0) + + +class TestCoordinatedMoves(CustomTestCase): + def setUp(self): + self.servers = [_DaemonServer(_initial_rows()) for _ in range(2)] + for server in self.servers: + self.addCleanup(server.close) + world = _FakeWorld(2) + results, errors = _run_ranks( + *( + lambda rank=rank: start_expert_layout_session( + socket_path=self.servers[rank].path, + label=f"rank{rank}", + all_gather=world.gather(rank), + ) + for rank in range(2) + ) + ) + self.assertEqual(errors, [None, None]) + self.sessions = [session for session, _ in results] + for session in self.sessions: + self.addCleanup(session.client.close) + session.digest = layout_digest(_initial_rows()) + session.weights_shared = True + + def _on_every_rank(self, fn): + results, errors = _run_ranks(*(lambda s=s: fn(s) for s in self.sessions)) + self.assertEqual(errors, [None, None]) + return results + + def test_a_move_is_recorded_by_every_rank_daemon(self): + new_map = torch.tensor(_initial_rows()) + new_map[1] = torch.tensor(_moved_row(1, 2)) + results = self._on_every_rank( + lambda s: (s.rebalance_mode(), s.begin([1]), s.commit([1], new_map)) + ) + for (mode, begin_reason, commit_reason), session, server in zip( + results, self.sessions, self.servers + ): + self.assertEqual(mode, ("coordinated", None)) + self.assertIsNone(begin_reason) + self.assertIsNone(commit_reason) + self.assertEqual(session.version, 1) + self.assertEqual(session.digest, layout_digest(new_map.tolist())) + layout = server.layout() + self.assertEqual(layout["version"], 1) + self.assertEqual(layout["physical_to_logical_map"], new_map.tolist()) + + def test_another_engine_on_one_rank_postpones_the_move_everywhere(self): + # A warm standby attaches to rank 1's daemon only. + standby = ExpertLayoutClient(self.servers[1].path, "standby") + self.addCleanup(standby.close) + standby.attach() + reasons = self._on_every_rank(lambda s: s.begin([0])) + for reason in reasons: + self.assertIn("rank 1: 1 other engine(s) map these weights", reason) + # Rank 0 began, then aborted: nothing is pending, nothing moved. + for server in self.servers: + self.assertIsNone(server.layout()["pending"]) + self.assertEqual(server.layout()["version"], 0) + # Busy only postpones; the move goes ahead once the standby is gone. + modes = self._on_every_rank(lambda s: s.rebalance_mode()) + self.assertEqual(modes, [("coordinated", None)] * 2) + standby.close() + _wait_for(lambda: self.servers[1].layout()["attached"] == 1) + self.assertEqual(self._on_every_rank(lambda s: s.begin([0])), [None, None]) + + def test_a_refusal_turns_rebalancing_off_for_good(self): + # An engine that does not follow the layout fetched rank 0's weights. + self.servers[0].request(_fetch_state()) + reasons = self._on_every_rank(lambda s: s.begin([0])) + for reason in reasons: + self.assertIn("rank 0: ", reason) + self.assertIn("until the weight cache daemon restarts", reason) + modes = self._on_every_rank(lambda s: s.rebalance_mode()) + self.assertEqual([mode for mode, _ in modes], ["off", "off"]) + + def test_a_commit_the_daemons_do_not_agree_with_turns_rebalancing_off(self): + self.assertEqual(self._on_every_rank(lambda s: s.begin([0])), [None, None]) + maps = [torch.tensor(_initial_rows()) for _ in range(2)] + for engine_map in maps: + engine_map[0] = torch.tensor(_moved_row(0, 1)) + # Rank 1's engine believes another layer moved too. + maps[1][2] = torch.tensor(_moved_row(2, 1)) + results, errors = _run_ranks( + *( + lambda rank=rank: self.sessions[rank].commit([0], maps[rank]) + for rank in range(2) + ) + ) + self.assertEqual(errors, [None, None]) + for reason in results: + self.assertIn("rank 1: the daemon recorded another layout", reason) + modes = self._on_every_rank(lambda s: s.rebalance_mode()) + self.assertEqual([mode for mode, _ in modes], ["off", "off"]) + + def test_rebalance_mode_is_agreed_by_every_rank(self): + self.sessions[0].weights_shared = False + modes = self._on_every_rank(lambda s: s.rebalance_mode()) + self.assertEqual( + modes, + [ + ( + "off", + "some ranks loaded the weights from disk and some from the daemon", + ) + ] + * 2, + ) + + def test_weights_loaded_from_disk_rebalance_without_the_daemons(self): + for session in self.sessions: + session.weights_shared = False + modes = self._on_every_rank(lambda s: s.rebalance_mode()) + self.assertEqual(modes, [("private", None)] * 2) + + +class TestAdoptTheDaemonLayout(CustomTestCase): + def setUp(self): + self.server = _DaemonServer(_initial_rows()) + self.addCleanup(self.server.close) + for patcher in ( + mock.patch.object( + ModelConfigForExpertLocation, + "from_model_config", + return_value=ModelConfigForExpertLocation( + num_layers=NUM_LAYERS, num_logical_experts=NUM_LOGICAL + ), + ), + # One rank. + mock.patch.object(expert_layout, "_world_all_gather", lambda obj: [obj]), + ): + patcher.start() + self.addCleanup(patcher.stop) + self.addCleanup(expert_layout.set_expert_layout_session, None) + + def _move_all_layers(self): + moved = [(layer_id, _moved_row(layer_id, 4)) for layer_id in range(NUM_LAYERS)] + _move(self.server.path, moved) + _wait_for(lambda: self.server.layout()["attached"] == 0) + return [row for _, row in moved] + + def _config(self, **overrides): + fields = dict( + device="cpu", + tp_size=2, + ep_size=2, + nnodes=1, + enable_eplb=True, + ep_num_redundant_experts=NUM_PHYSICAL - NUM_LOGICAL, + moe_a2a_backend="none", + ep_dispatch_algorithm="dynamic", + weight_cache_mode="client", + weight_cache_socket=self.server.path, + ) + fields.update(overrides) + return get_context().override_server_args(**fields) + + def _adopt(self, moe_ep_rank=0): + with get_parallel().override(tp_rank=moe_ep_rank): + initial = compute_initial_expert_location_metadata( + model_config=None, moe_ep_rank=moe_ep_rank + ) + adopted = adopt_weight_cache_expert_layout( + initial, model_config=None, moe_ep_rank=moe_ep_rank, gpu_id=0 + ) + session = get_expert_layout_session() + if session is not None: + self.addCleanup(session.client.close) + return initial, adopted, session + + def test_the_daemon_layout_equals_init_expert_location(self): + moved_rows = self._move_all_layers() + for a2a, dispatch in ( + ("none", "dynamic"), + ("deepep", "dynamic"), + ("deepep", "static"), + ): + for moe_ep_rank in (0, 1): + with self.subTest(a2a=a2a, dispatch=dispatch, moe_ep_rank=moe_ep_rank): + with self._config( + moe_a2a_backend=a2a, ep_dispatch_algorithm=dispatch + ): + _, adopted, session = self._adopt(moe_ep_rank) + with self._config( + moe_a2a_backend=a2a, + ep_dispatch_algorithm=dispatch, + weight_cache_mode="off", + init_expert_location=json.dumps( + {"physical_to_logical_map": moved_rows} + ), + ): + direct = compute_initial_expert_location_metadata( + model_config=None, moe_ep_rank=moe_ep_rank + ) + self.assertEqual( + adopted.physical_to_logical_map.tolist(), moved_rows + ) + for field in ( + "physical_to_logical_map", + "physical_to_logical_map_cpu", + "logical_to_all_physical_map", + "logical_to_all_physical_map_cpu", + "logical_to_all_physical_map_num_valid", + "logical_to_rank_dispatch_physical_map", + ): + expected = getattr(direct, field) + if expected is None: + self.assertIsNone(getattr(adopted, field), field) + else: + self.assertTrue( + torch.equal(getattr(adopted, field), expected), field + ) + self.assertEqual(adopted.ep_size, direct.ep_size) + self.assertEqual( + (session.kind, session.version), (TRACKED, NUM_LAYERS) + ) + self.assertEqual(session.digest, layout_digest(moved_rows)) + + def test_an_unmoved_layout_keeps_the_engine_metadata(self): + with self._config(): + initial, adopted, session = self._adopt() + self.assertIs(adopted, initial) + self.assertEqual(session.digest, layout_digest(_initial_rows())) + self.assertEqual(self.server.layout()["attached"], 1) + + def test_an_engine_without_eplb_cannot_follow_a_moved_layout(self): + self._move_all_layers() + with self._config(enable_eplb=False, ep_dispatch_algorithm=None): + with self.assertRaisesRegex(RuntimeError, "--enable-eplb"): + self._adopt() + # The engine that failed to start is not left attached. + _wait_for(lambda: self.server.layout()["attached"] == 0) + + def test_the_layout_shapes_must_match(self): + self._move_all_layers() + # 12 physical experts here, 10 in the daemon. + with self._config(ep_num_redundant_experts=4): + with self.assertRaisesRegex(RuntimeError, "--ep-num-redundant-experts"): + self._adopt() + _wait_for(lambda: self.server.layout()["attached"] == 0) + + def test_without_the_weight_cache_nothing_changes(self): + with self._config(weight_cache_mode="off"): + initial, adopted, session = self._adopt() + self.assertIs(adopted, initial) + self.assertIsNone(session) + self.assertEqual(self.server.layout()["attached"], 0) + + +class TestIpcLoaderSendsItsLayout(CustomTestCase): + def test_fetch_state_carries_the_layout_and_reports_a_refusal(self): + server = _DaemonServer(_initial_rows()) + self.addCleanup(server.close) + session, _ = start_expert_layout_session( + socket_path=server.path, label="engine", all_gather=lambda obj: [obj] + ) + self.addCleanup(session.client.close) + loader = ipc_loader.IpcModelLoader( + load_config=LoadConfig(load_format=LoadFormat.IPC_CACHE), + socket_path=server.path, + weight_cache_mode="client", + ) + parallel = SimpleNamespace( + tp_size=1, + tp_rank=0, + pp_size=1, + pp_rank=0, + moe_ep_size=1, + moe_dp_size=1, + moe_dp_rank=0, + moe_ep_rank=0, + dp_size=1, + attn_dp_size=1, + enable_dp_lm_head=False, + attn_cp_size=1, + moe_dense_tp_size=None, + ) + model_config = SimpleNamespace( + model_path="/models/demo", + hf_config=SimpleNamespace( + architectures=["Qwen3MoeForCausalLM"], quantization_config=None + ), + quantization=None, + revision=None, + dtype="torch.bfloat16", + ) + with ( + mock.patch.object(ipc_loader, "get_parallel", return_value=parallel), + mock.patch.object( + ipc_loader, + "get_exec", + return_value=SimpleNamespace( + moe=SimpleNamespace(moe_a2a_backend="none") + ), + ), + get_resources().override(expert_layout_session=session), + ): + session.digest = "stale" + with self.assertRaisesRegex( + RuntimeError, + "refused to serve its weights: the expert layout the engine routes " + "with .* is not the one its weights are in", + ): + loader._fetch_from_cache(model_config, SimpleNamespace(gpu_id=0)) + session.digest = layout_digest(_initial_rows()) + result = loader._fetch_from_cache(model_config, SimpleNamespace(gpu_id=0)) + self.assertEqual(result["status"], "ok") + + mark_expert_weights_shared(True) + self.assertTrue(session.weights_shared) + mark_expert_weights_shared(False) + self.assertFalse(session.weights_shared) + + +def _old_daemon_handle_connection(conn): + """How a daemon from before the expert layout record answers: fetch_state + reads only the config, and other request types get an error.""" + req = recv_msg(conn) + if req.get("type") == "fetch_state": + send_msg( + conn, {"status": "ok", "entries": {}, "transport_backend": "torch_ipc"} + ) + else: + send_msg( + conn, + {"status": "error", "message": f"Unknown request type: {req.get('type')}"}, + ) + + +class TestCompatibility(CustomTestCase): + def test_a_new_engine_with_an_old_daemon_keeps_its_layout(self): + server = _DaemonServer(_initial_rows()) + self.addCleanup(server.close) + server.daemon._handle_connection = _old_daemon_handle_connection + session, rows = start_expert_layout_session( + socket_path=server.path, label="engine", all_gather=lambda obj: [obj] + ) + self.assertEqual(session.kind, LEGACY) + self.assertIsNone(rows) + self.assertFalse(session.client.attached) + # Its fetch_state carries a field the old daemon does not read. + reply = server.request(_fetch_state(**session.fetch_fields())) + self.assertEqual(reply["status"], "ok") + # It does not rebalance the shared weights, as before. + session.weights_shared = True + mode, reason = session.rebalance_mode() + self.assertEqual(mode, "off") + self.assertIn("older version", reason) + + def test_an_old_engine_with_a_new_daemon(self): + server = _DaemonServer(_initial_rows()) + self.addCleanup(server.close) + # Served as before while every expert is where it started ... + reply = server.request(_fetch_state()) + self.assertEqual(reply["status"], "ok") + # ... after which no engine may move experts, since the old one would + # not follow them. + client = ExpertLayoutClient(server.path, "engine") + self.addCleanup(client.close) + client.attach() + self.assertEqual( + client.request(BEGIN, base_version=0, layer_ids=[0])["status"], REFUSED + ) + + def test_an_old_engine_is_refused_moved_experts(self): + server = _DaemonServer(_initial_rows()) + self.addCleanup(server.close) + _move(server.path, [(1, _moved_row(1, 1))]) + reply = server.request(_fetch_state()) + self.assertEqual(reply["status"], LAYOUT_REFUSED) + self.assertIn("Upgrade the engine", reply["message"]) + + +if __name__ == "__main__": + unittest.main()