diff --git a/docs/docs/references/environment_variables.mdx b/docs/docs/references/environment_variables.mdx index 6a593e1f9..ea09a2fad 100644 --- a/docs/docs/references/environment_variables.mdx +++ b/docs/docs/references/environment_variables.mdx @@ -1068,6 +1068,11 @@ SGLang supports various environment variables that can be used to configure its Detect and report ranks that fall behind during collective ops. false + + SGLANG_ENABLE_RANK_CONSENSUS_CHECKER + Check for PP/TP-rank divergence. Kill the server when divergence occurs. Helpful for trouble-shooting server hangs issues. + False + SGLANG_FORCE_SHUTDOWN Force an immediate process-group shutdown on exit. diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 4cc5e5015..817b388ce 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -329,6 +329,7 @@ class Envs: SGLANG_LOG_REQUEST_HEADERS = EnvTuple(tuple()) SGLANG_LOG_SCHEDULER_STATUS_TARGET = EnvStr("") SGLANG_LOG_SCHEDULER_STATUS_INTERVAL = EnvFloat(60.0) + SGLANG_ENABLE_RANK_CONSENSUS_CHECKER = EnvBool(False) # =================================================================== # IPC, broadcasters, and ports diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index b617ed301..04dddd34b 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -309,6 +309,7 @@ from sglang.srt.utils import ( is_hip, is_mps, kill_itself_when_parent_died, + rank_consensus_checker, require_mlp_sync, set_gpu_proc_affinity, set_random_seed, @@ -655,6 +656,8 @@ class Scheduler( self.init_batch_result_processor() + self.init_rank_consensus_checker() + self.is_initializing = False self.init_startup_timing_summary() @@ -1688,6 +1691,8 @@ class Scheduler( if self.decode_offload_manager is not None: self.decode_offload_manager.release_host_resources() + rank_consensus_checker.shutdown() + def run_event_loop(self) -> None: """Run the scheduler's event loop. @@ -2114,6 +2119,16 @@ class Scheduler( get_running_batch=lambda: self.running_batch, ) + def init_rank_consensus_checker(self) -> None: + groups = [] + if self.attn_cp_group is not None and self.attn_tp_group is not None: + groups += [self.attn_cp_group, self.attn_tp_group] + else: + groups += [self.tp_group] + if self.pp_group is not None: + groups += [self.pp_group] + rank_consensus_checker.configure(groups) + def init_kv_events_publisher(self) -> None: self.kv_events_publisher = SchedulerKvEventsPublisher( kv_events_config=get_observability().kv_events_config, diff --git a/python/sglang/srt/mem_cache/unified_radix_cache.py b/python/sglang/srt/mem_cache/unified_radix_cache.py index 498be1329..c00b44b8f 100644 --- a/python/sglang/srt/mem_cache/unified_radix_cache.py +++ b/python/sglang/srt/mem_cache/unified_radix_cache.py @@ -95,6 +95,7 @@ if TYPE_CHECKING: from sglang.srt.mem_cache.memory_pool_host import PoolEntry from sglang.srt.server_args import ServerArgs +from sglang.srt.utils.rank_consensus_checker import rank_consensus T = TypeVar("T") @@ -491,6 +492,10 @@ class UnifiedRadixCache(BasePrefixCache): if self.host_pool_group is not None: self.host_pool_group.destroy() + @rank_consensus( + same_params=["params"], + same_results=["result.full_kv_hit_length", "result.swa_host_hit_length"], + ) def match_prefix(self, params: MatchPrefixParams) -> MatchResult: result = self.session.try_match_prefix(params) if result is not None: @@ -1700,6 +1705,7 @@ class UnifiedRadixCache(BasePrefixCache): operation_terminated = states[1].item() == 1 return can_terminate or operation_terminated + @rank_consensus(same_params=True, same_results=True) def check_prefetch_progress(self, req_id: str) -> bool: if req_id not in self.ongoing_prefetch: return True @@ -1917,6 +1923,7 @@ class UnifiedRadixCache(BasePrefixCache): return 0 return self.buffer_pipeline.staged_prefetch_swa_tokens(req_id) + @rank_consensus(same_params=True) def release_aborted_request(self, rid: str) -> None: self.prefetch_loaded_tokens_by_reqid.pop(rid, None) if ( diff --git a/python/sglang/srt/utils/rank_consensus_checker.py b/python/sglang/srt/utils/rank_consensus_checker.py new file mode 100644 index 000000000..953782dfc --- /dev/null +++ b/python/sglang/srt/utils/rank_consensus_checker.py @@ -0,0 +1,432 @@ +from __future__ import annotations + +import functools +import hashlib +import inspect +import logging +import os +import queue +import threading +from typing import TYPE_CHECKING, Any, Callable, List, Optional + +import torch +import torch.distributed as dist + +from sglang.srt.environ import envs + +if TYPE_CHECKING: + from sglang.srt.distributed.parallel_state import GroupCoordinator + +logger = logging.getLogger(__name__) + +_sync_groups: List[dist.ProcessGroup] = [] # Dedicated gloo groups (one per rank-set). +_q: Optional[queue.Queue[str]] = None +_worker_thread: Optional[threading.Thread] = None +_scheduler_thread: Optional[threading.Thread] = None + + +def rank_consensus(func=None, *, same_params=None, same_results=None, **kwargs): + """ + Mark a function that should be consensus in PP and TP ranks. Here consensus means, + the same order of calling, same parameters and return values optionally. + + The function must be called in the scheduler thread. + + Usages: + + * Assert that the function is called by all ranks. The parameters or results may not be same. + @rank_consensus + def foo(): + pass + + * Assert that all parameters are same in all ranks. + @rank_consensus(same_params = True) + def foo(a, b): + pass + + * Assert that some parameters are same in all ranks. + @rank_consensus(same_params = ["a", "c"]) + def foo(a, b, c): + pass + + * Assert that part of the parameters are same in all ranks. + @rank_consensus(same_params = ["a.req_id"]) + def foo(a): + pass + + * Assert that results are same in all ranks. + @rank_consensus(same_results = True) + def foo(): + return 1 + + * Assert for part of the results are same. + @rank_consensus(same_results = ["result.some_field"] + def foo(): + return SomeObject() + + @rank_consensus(same_results = ["result.field", "len(result.field2)"] + def foo(): + return SomeObject() + + * Assert the function is called by all ranks and all parameters and results are the same. + @rank_consensus(same_params = True, same_results = True) + def foo(): + return 1 + """ + if kwargs: + raise TypeError( + f"rank_consensus() got unexpected keyword argument(s): " f"{list(kwargs)}" + ) + + params_selector = _normalize_selector(same_params, "same_params") + results_selector = _normalize_selector(same_results, "same_results") + + def decorator(func: Callable) -> Callable: + # This decorator function called at import time. So it should be zero runtime overhead + # when the consensus checker is disabled. + if not envs.SGLANG_ENABLE_RANK_CONSENSUS_CHECKER.get(): + return func + + # Unwrap static/class-method descriptors so we always operate on the + # raw function. We remember the descriptor type so we can re-wrap the + # result and the class-body descriptor protocol keeps working. + if isinstance(func, (classmethod, staticmethod)): + raw_func = func.__func__ + descriptor_type = type(func) + else: + raw_func = func + descriptor_type = None + sig = inspect.signature(raw_func) + + # When calling class method or object method with "same_params=True", + # skip the first "cls" or "self", as the text format for that + # may include memory addresses, which are considered divergence. + skip_name: Optional[str] = None + if _is_method_with_receiver(func) and len(sig.parameters) > 0: + skip_name = next(iter(sig.parameters)) + + @functools.wraps(raw_func) + def wrapper(*args: Any, **kwargs: Any) -> Any: + params_payload = "" + if params_selector is not None: + # Bind once and apply defaults so that name-based selectors work + # regardless of whether the caller passed positionally or by kw. + bound = sig.bind(*args, **kwargs) + bound.apply_defaults() + arguments = dict(bound.arguments) + params_payload = _build_payload( + "call", params_selector, arguments, skip_name + ) + assert_same("%s called params=%s", raw_func.__name__, params_payload) + + result = raw_func(*args, **kwargs) + + result_payload = "" + if results_selector is not None: + result_scope = {"result": result} + result_payload = _build_payload( + "return", + results_selector, + result_scope, + ) + assert_same("%s returns result=%s", raw_func.__name__, result_payload) + return result + + # Re-wrap into the original descriptor type so class-body access + # (C.method / instance.method) still binds correctly. + if descriptor_type is staticmethod: + return staticmethod(wrapper) + if descriptor_type is classmethod: + return classmethod(wrapper) + return wrapper + + if func is not None: + # Bare `@rank_consensus` form. + return decorator(func) + else: + # `@rank_consensus(same_params=True, same_results=True)` form. + return decorator + + +def _normalize_selector( + value: None | bool | str | list[str], name: str +) -> None | bool | list[str]: + """Normalize a selector argument to one of: + ``None`` (skip), ``True`` (compare everything), or ``list[str]`` (the + expressions to evaluate). ``False`` is treated as ``None``. + """ + if value is None or value is False: + return None + if value is True: + return True + if isinstance(value, str): + return [value] + if isinstance(value, list) and all(isinstance(s, str) for s in value): + return list(value) + raise TypeError(f"{name} must be True / False / str / list[str], got {value!r}") + + +def _is_method_with_receiver(func: Any) -> bool: + """Return True iff ``func`` is a method whose first parameter is a + receiver (instance for instance-methods, class for class-methods) that + should be dropped from the ``same_params=True`` payload. + + Distinguishes: + * ``staticmethod`` object -> False (no receiver) + * ``classmethod`` object -> True (receiver is the class) + * plain ``def`` defined inside a class body (``__qualname__`` has a + dot before the final segment and is not a ```` closure) -> + True (instance method) + * anything else (module-level function, nested function, lambda) -> + False + """ + if isinstance(func, staticmethod): + return False + if isinstance(func, classmethod): + return True + if inspect.isfunction(func): + qualname = getattr(func, "__qualname__", "") + # ``C.m`` -> True; ``m`` -> False; ``outer..m`` -> False + # (closures aren't class-body methods). + if "." in qualname and "" not in qualname: + return True + return False + + +def _build_payload( + tag: str, + selector: bool | list[str], + scope: dict[str, Any], + skip_name: Optional[str] = None, +) -> str: + """Serialize the selected values into a single comparable string. + + ``skip_name`` only applies to the ``True`` (whole-scope) form and is used + to drop the receiver (``self`` / ``cls``) from method payloads; explicit + ``list[str]`` selectors honor exactly what the user listed. + """ + if selector is True: + # Whole scope is the payload. For the call checkpoint, the scope is + # the arguments dict; for the return checkpoint, the caller wrapped + # result into the scope, so we repr ``result`` directly. + if tag == "call": + if skip_name is not None: + scope = {k: v for k, v in scope.items() if k != skip_name} + return repr(scope) + return repr(scope["result"]) + parts: list[str] = [] + for expr in selector: + value = _eval_selector(expr, scope) + parts.append(f"{expr}={value!r}") + return " | ".join(parts) + + +def _eval_selector(expr: str, scope: dict[str, Any]) -> Any: + """Evaluate a selector expression in a restricted scope. + + Errors (unknown parameter name, missing attribute, bad syntax) propagate + -- they are caller bugs and must not be silently swallowed or confused + with cross-rank divergence. + """ + safe_builtins = { + "len": len, + "int": int, + "str": str, + "bool": bool, + "float": float, + "tuple": tuple, + "list": list, + "dict": dict, + "set": set, + "sorted": sorted, + "min": min, + "max": max, + "sum": sum, + } + return eval(expr, {"__builtins__": safe_builtins}, dict(scope)) + + +def enabled() -> bool: + """Test that the checker has been enabled and configure() is called.""" + return _q is not None + + +def assert_same(msg_fmt: str, *args: Any) -> None: + """Record a decision that every TP/PP rank must make identically. + + Must be called from the scheduler thread. If the env var is set and the + checker is configured, an assertion guards that the caller is on the + scheduler thread recorded at configure() time — events from other threads + would interleave out of order with peer ranks and corrupt the lock-step + drain. + + When the divergence checker is disabled, this is a zero-overhead no-op. + + Example: + assert_same("my decision: %s %d", "foo", 100) + + Prefer `@rank_consensus` over this function for code-cleanliness. + """ + if not enabled(): + return + # Sanity check: only the scheduler thread is allowed to enqueue. Other + # callers would race with the worker's min-length drain and desynchronize + # ranks, since their events would not exist on peer ranks. + if threading.current_thread() is not _scheduler_thread: + raise RuntimeError("rdc.assert_same must be called from the scheduler thread") + # Format eagerly: args may reference mutable state that mutates + # between now and when the worker thread drains the queue. + _q.put(msg_fmt % args) + + +def configure(groups: List[GroupCoordinator]) -> None: + """Initialize the checker. No-op if SGLANG_ENABLE_RANK_CONSENSUS_CHECKER is not set.""" + global _sync_groups, _q, _worker_thread, _scheduler_thread + + if not envs.SGLANG_ENABLE_RANK_CONSENSUS_CHECKER.get(): + return + + logger.warning( + "Rank consensus checker is enabled. The server will suicide if rank divergence detected." + ) + # Build a dedicated sync group. So our synchronization work will not affect + # the scheduler thread at all. + _sync_groups = _create_sync_groups(groups) + _q = queue.Queue() + # Assume the calling thread is the schedule thread. + # We will check assert_same() must be called by the scheduler thread. + _scheduler_thread = threading.current_thread() + _worker_thread = threading.Thread( + target=_worker_loop, name="rank_consensus_checker", daemon=True + ) + _worker_thread.start() + + +def _create_sync_groups( + groups: List[GroupCoordinator], +) -> List[dist.ProcessGroup]: + """Create duplicated groups, used for background thread""" + from sglang.srt.distributed.parallel_state import create_custom_parallel_group + + dedicated: List[dist.ProcessGroup] = [] + seen_rank_sets: set[tuple[int, ...]] = set() + for group in groups: + if group is None: + continue + # Skip single-rank groups: nothing to compare against. + if torch.distributed.get_world_size(group=group.cpu_group) == 1: + continue + group_ranks = tuple(torch.distributed.get_process_group_ranks(group.cpu_group)) + if group_ranks in seen_rank_sets: + continue + seen_rank_sets.add(group_ranks) + pg = create_custom_parallel_group(group_ranks=list(group_ranks), backend="gloo") + if pg is not None: + dedicated.append(pg) + return dedicated + + +def _destroy_dedicated_groups() -> None: + for pg in _sync_groups: + try: + torch.distributed.destroy_process_group(pg) + except Exception: + pass + + +def shutdown() -> None: + """Flush the queue, stop the worker thread, and disable assert_same.""" + global _q, _worker_thread, _sync_groups, _scheduler_thread + + q = _q + if q is None: + return + + # Put a sentinel value to wake the worker if it is blocked on _q.get(). + q.put(None) + if _worker_thread is not None: + _worker_thread.join() + _worker_thread = None + # Tear down the dedicated gloo groups BEFORE clearing _groups so the + # destroy helper can see them. Worker thread is already joined, so there + # is no concurrent all_reduce on these groups. + _destroy_dedicated_groups() + _q = None + _sync_groups = [] + _scheduler_thread = None + + +def _worker_loop() -> None: + """Consume events in lock-step with peer ranks via gloo all-reduce. + + Each iteration: + 1. Determine the items available in _q. + 2. Drain exactly the minimum number of items in all ranks. + 3. Compare all events are identical across ranks. + """ + while _q is not None: + # Drain first. Block waiting for the first event. + first = _q.get() + if first is None: + # shutdown() is called. + return + + # Drain more whenever available. + # Every rank should drain the same number. + count = _all_reduce_min_int(_q.qsize()) + events: List[str] = [first] + shutdown_signaled = False + for _ in range(count): + event = _q.get() + if event is None: + # shutdown() sentinel arrived mid-batch: stop draining but + # still check the events we already hold — they are real + # decisions every rank must agree on. Then exit, since the + # sentinel means shutdown() is waiting on worker_thread.join(). + shutdown_signaled = True + break + events.append(event) + + # Cross-rank check. + _check_for_consensus(events) + + if shutdown_signaled: + return + + +def _all_reduce_min_int(value: int) -> int: + """Reduce `value` to its global minimum across every configured group.""" + tensor = torch.tensor([value], dtype=torch.int64) + for group in _sync_groups: + dist.all_reduce(tensor, op=dist.ReduceOp.MIN, group=group) + return int(tensor.item()) + + +def _check_for_consensus(events: list[str]) -> None: + # Compute sha1 of concatenation of all msgs. + hasher = hashlib.sha1() + for msg in events: + hasher.update(msg.encode("utf-8")) + + # Determine if some rank has a different value. + value_bytes = hasher.digest() + min_value = torch.tensor(list(hasher.digest()), dtype=torch.uint8) + max_value = min_value.clone() + for group in _sync_groups: + dist.all_reduce(min_value, op=dist.ReduceOp.MIN, group=group) + dist.all_reduce(max_value, op=dist.ReduceOp.MAX, group=group) + if not torch.equal(min_value, max_value): + # When divergence, all rank should output the following log. + logger.critical( + f"Found rank divergence for {len(events)} events(s)! local hash: {value_bytes.hex()}, events = {events}" + ) + for handler in logger.handlers: + handler.flush() + + # os._exit instead of sys.exit: this runs in a background thread, where + # SystemExit would only kill the thread, not the process. os._exit tears + # down the whole scheduler process so a TP/PP mismatch can never + # silently keep serving. + os._exit(1) + + logger.debug(f"Consensus check passed for {len(events)} event(s).") diff --git a/test/registered/cpu/test_rank_consensus_checker.py b/test/registered/cpu/test_rank_consensus_checker.py new file mode 100644 index 000000000..edcdd0488 --- /dev/null +++ b/test/registered/cpu/test_rank_consensus_checker.py @@ -0,0 +1,586 @@ +import os +import queue +import threading +import traceback +import unittest +from multiprocessing import Process +from unittest.mock import patch + +import torch.distributed as dist +import torch.multiprocessing as mp + +from sglang.srt.distributed import parallel_state as ps +from sglang.srt.distributed.parallel_state import ( + get_pp_group, + get_tp_group, + init_distributed_environment, + initialize_model_parallel, +) +from sglang.srt.utils.rank_consensus_checker import ( + assert_same, + configure, + rank_consensus, + shutdown, +) +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase, find_available_port + +register_cpu_ci(est_time=30, suite="base-b-test-cpu") + + +def run_distributed_test( + rank: int, + world_size: int, + pp_size: int, + tp_size: int, + master_port: int, + fn, +) -> None: + """Child-process entry point: set up gloo, then run fn. + + Exit codes: + * 0 -> fn finished cleanly + * 1 -> rdc detected divergence and called os._exit(1) from its worker + * 2 -> fn raised (test setup/scenario bug) + """ + # CUDA_VISIBLE_DEVICES is set to "99" (a non-existent device) by the parent + # in _spawn() before this process starts, so by the time the test module + # (and torch) is re-imported here, is_cuda_alike() returns False and + # GroupCoordinator picks device="cpu". That keeps this test CPU-only and + # lets world_size exceed the host's physical GPU count. + + # The CUDA-only communicators (pynccl, custom allreduce) cannot be built + # without a GPU -- PyNcclCommunicator calls torch.cuda.device(device). + # initialize_model_parallel has no flag to disable pynccl, so patch + # init_model_parallel_group to force use_pynccl=False (and clear the + # module-level custom-allreduce default via its public setter). patch.object + # auto-restores on exit, including the os._exit(2) path below. + ps.set_custom_all_reduce(False) + + def _cpu_init_model_parallel_group( + *args, _orig=ps.init_model_parallel_group, **kwargs + ): + kwargs.setdefault("use_pynccl", False) + kwargs.setdefault("use_custom_allreduce", False) + return _orig(*args, **kwargs) + + with patch.object(ps, "init_model_parallel_group", _cpu_init_model_parallel_group): + try: + os.environ["RANK"] = str(rank) + os.environ["WORLD_SIZE"] = str(world_size) + os.environ["MASTER_ADDR"] = "localhost" + os.environ["MASTER_PORT"] = str(master_port) + os.environ["LOCAL_SIZE"] = str(world_size) + + init_distributed_environment( + world_size=world_size, + rank=rank, + distributed_init_method="env://", + local_rank=rank, + backend="gloo", + ) + + initialize_model_parallel( + tensor_model_parallel_size=tp_size, + pipeline_model_parallel_size=pp_size, + backend="gloo", + ) + + fn() + except Exception as e: + print(f"subprocess[{rank=}] has error: {e}", flush=True) + traceback.print_exc() + os._exit(2) + finally: + try: + if dist.is_initialized(): + dist.destroy_process_group() + except Exception: + pass + + +class _DummyClass: + def __init__(self, a: int = None, b: int = None): + self.a = a + self.b = b + + def __repr__(self) -> str: + return f"DummyClass(a={self.a}, b={self.b})" + + +class _MethodHost: + @rank_consensus(same_params=True) + def instance_method(obj, a, b): + return a + b + + @rank_consensus(same_params=True) + @classmethod + def class_method(klass, a): + return a + 1 + + @rank_consensus(same_params=True) + @staticmethod + def static_method(a, b): + return a * b + + +class RankConsensusCheckerTestCase(CustomTestCase): + def _spawn(self, fn, pp_size: int = 1, tp_size: int = 1, enable_env: bool = True): + """Run fn in world_size spawned gloo children. Returns True iff every + child exited with code 0. A detected divergence makes rdc call + os._exit(1) from its worker thread; an exception inside fn makes + run_distributed_test call os._exit(2). Either way _spawn returns + False for that child.""" + mp.set_start_method("spawn", force=True) + master_port = find_available_port(23456) + + old_env = os.getenv("SGLANG_ENABLE_RANK_CONSENSUS_CHECKER") + os.environ["SGLANG_ENABLE_RANK_CONSENSUS_CHECKER"] = str(enable_env) + + world_size = pp_size * tp_size + processes = [] + for rank in range(world_size): + p = Process( + target=run_distributed_test, + kwargs=dict( + rank=rank, + world_size=world_size, + pp_size=pp_size, + tp_size=tp_size, + master_port=master_port, + fn=fn, + ), + ) + p.start() + processes.append(p) + + for p in processes: + p.join() + + if old_env is None: + os.environ.pop("SGLANG_ENABLE_RANK_CONSENSUS_CHECKER") + else: + os.environ["SGLANG_ENABLE_RANK_CONSENSUS_CHECKER"] = old_env + + return all(p.exitcode == 0 for p in processes) + + +class TestAssertSame(RankConsensusCheckerTestCase): + @staticmethod + def same_fn(): + configure([get_tp_group()]) + assert_same("same %d", 10) + shutdown() + + def test_same(self): + """Same args on every rank -> no divergence, clean exit.""" + self.assertTrue(self._spawn(TestAssertSame.same_fn, tp_size=2)) + + @staticmethod + def divergence_fn(): + tp_group = get_tp_group() + configure([tp_group]) + assert_same("diverge %d", tp_group.rank_in_group) + shutdown() + + def test_divergence(self): + """Different args on different ranks -> rdc calls os._exit(1) -> child + exit code is 1 -> _spawn returns False.""" + self.assertFalse(self._spawn(TestAssertSame.divergence_fn, tp_size=2)) + + @staticmethod + def divergent_multi_group_fn(): + tp_group = get_tp_group() + pp_group = get_pp_group() + configure([tp_group, pp_group]) + assert_same("diverge %d", tp_group.rank_in_group) + shutdown() + + def test_divergence_detected_multi_group(self): + """Passing the same group twice must still surface the divergence.""" + self.assertFalse( + self._spawn( + TestAssertSame.divergent_multi_group_fn, + pp_size=2, + tp_size=2, + ) + ) + + @staticmethod + def wrong_thread_fn(): + tp_group = get_tp_group() + configure([tp_group]) + + err_box: queue.Queue = queue.Queue() + + def _other_thread(): + try: + assert_same("from other thread") + err_box.put(None) + except Exception as e: # noqa: BLE001 + err_box.put(e) + + t = threading.Thread(target=_other_thread) + t.start() + t.join() + + err = err_box.get() + shutdown() + assert isinstance( + err, RuntimeError + ), f"Expected RuntimeError from stray-thread assert_same, got {err!r}" + + def test_assert_same_rejects_non_scheduler_thread(self): + """Check that assert_same() must be called in the scheduler thread. Otherwise report error.""" + self.assertTrue(self._spawn(TestAssertSame.wrong_thread_fn, tp_size=2)) + + @staticmethod + def disabled_fn(): + tp_group = get_tp_group() + configure([tp_group]) + assert_same("diverge %d", tp_group.rank_in_group) + + def test_disabled_is_noop(self): + """Test that when SGLANG_ENABLE_RANK_CONSENSUS_CHECKER=false, assert_same is no-op.""" + self.assertTrue( + self._spawn(TestAssertSame.disabled_fn, tp_size=2, enable_env=False) + ) + + +class TestRankConsensusDecorator(RankConsensusCheckerTestCase): + @staticmethod + def consensus_bare_diverge_fn(): + @rank_consensus + def foo(a: int) -> int: + return a + + # Bare decorator only checks "was called", not args; even with diverging + # args this must exit clean (no rank divergence). + tp_group = get_tp_group() + configure([tp_group]) + foo(tp_group.rank_in_group) + shutdown() + + def test_bare_decorator_clean_with_diverging_args(self): + """Bare decorator only checks that every rank calls the function; + diverging args must NOT be flagged.""" + self.assertTrue( + self._spawn(TestRankConsensusDecorator.consensus_bare_diverge_fn, tp_size=2) + ) + + @staticmethod + def consensus_all_params_same_fn(): + @rank_consensus(same_params=True) + def foo(a: int, b: int) -> int: + return a + b + + configure([get_tp_group()]) + foo(1, 2) + shutdown() + + def test_all_params_same(self): + self.assertTrue( + self._spawn( + TestRankConsensusDecorator.consensus_all_params_same_fn, tp_size=2 + ) + ) + + @staticmethod + def consensus_all_params_diverge_fn(): + @rank_consensus(same_params=True) + def foo(a, b): + return a + b + + tp_group = get_tp_group() + configure([tp_group]) + # The second argument differs on rank. Expect divergence. + foo(1, tp_group.rank_in_group) + shutdown() + + def test_all_params_diverge(self): + self.assertFalse( + self._spawn( + TestRankConsensusDecorator.consensus_all_params_diverge_fn, tp_size=2 + ) + ) + + @staticmethod + def consensus_named_params_same_fn(): + @rank_consensus(same_params=["a", "c"]) + def foo(a: int, b: int, c: int) -> int: + return a + b + c + + tp_group = get_tp_group() + configure([tp_group]) + # b diverges but is NOT in the selector list. Expect good. + foo(1, tp_group.rank_in_group, 3) + shutdown() + + def test_named_params_ignores_unselected_divergence(self): + self.assertTrue( + self._spawn( + TestRankConsensusDecorator.consensus_named_params_same_fn, tp_size=2 + ) + ) + + @staticmethod + def consensus_named_params_diverge_fn(): + @rank_consensus(same_params=["a", "c"]) + def foo(a: int, b: int, c: int) -> int: + return a + b + c + + # c diverges and IS in the selector list. Expect divergence. + tp_group = get_tp_group() + configure([tp_group]) + foo(1, 2, tp_group.rank_in_group) + shutdown() + + def test_named_params_flags_selected_divergence(self): + self.assertFalse( + self._spawn( + TestRankConsensusDecorator.consensus_named_params_diverge_fn, + tp_size=2, + ) + ) + + @staticmethod + def consensus_dotted_param_same_fn(): + @rank_consensus(same_params=["_a.a"]) + def foo(_a: _DummyClass) -> None: + pass + + tp_group = get_tp_group() + configure([tp_group]) + dummy = _DummyClass(a=10, b=tp_group.rank_in_group) + foo(dummy) + shutdown() + + def test_dotted_param_same(self): + self.assertTrue( + self._spawn( + TestRankConsensusDecorator.consensus_dotted_param_same_fn, tp_size=2 + ) + ) + + @staticmethod + def consensus_dotted_param_diverge_fn(): + @rank_consensus(same_params=["_a.a"]) + def foo(_a: _DummyClass) -> None: + pass + + tp_group = get_tp_group() + configure([tp_group]) + dummy = _DummyClass(a=tp_group.rank_in_group, b=10) + foo(dummy) + shutdown() + + def test_dotted_param_diverge(self): + self.assertFalse( + self._spawn( + TestRankConsensusDecorator.consensus_dotted_param_diverge_fn, + tp_size=2, + ) + ) + + @staticmethod + def consensus_full_result_same_fn(): + @rank_consensus(same_results=True) + def foo(value: int) -> _DummyClass: + return _DummyClass(a=value, b=value * 2) + + configure([get_tp_group()]) + foo(5) + shutdown() + + def test_full_result_same(self): + self.assertTrue( + self._spawn( + TestRankConsensusDecorator.consensus_full_result_same_fn, tp_size=2 + ) + ) + + @staticmethod + def consensus_full_result_diverge_fn(): + @rank_consensus(same_results=True) + def foo(value: int) -> _DummyClass: + return _DummyClass(a=value, b=value * 2) + + tp_group = get_tp_group() + configure([tp_group]) + foo(tp_group.rank_in_group) + shutdown() + + def test_full_result_diverge(self): + self.assertFalse( + self._spawn( + TestRankConsensusDecorator.consensus_full_result_diverge_fn, + tp_size=2, + ) + ) + + @staticmethod + def consensus_partial_result_same_fn(): + @rank_consensus(same_results=["result.x", "len(result.y)"]) + def foo(x, y_list): + class _R: + pass + + r = _R() + r.x = x + r.y = y_list + return r + + tp_group = get_tp_group() + configure([tp_group]) + # x and len(y) both equal across ranks; y contents differ but are not selected. Expect good. + foo(x=3, y_list=[tp_group.rank_in_group] * 4) + shutdown() + + def test_partial_result_same(self): + self.assertTrue( + self._spawn( + TestRankConsensusDecorator.consensus_partial_result_same_fn, + tp_size=2, + ) + ) + + @staticmethod + def consensus_partial_result_diverge_fn(): + @rank_consensus(same_results=["result.x", "len(result.y)"]) + def foo(x, y_list): + class _R: + pass + + r = _R() + r.x = x + r.y = y_list + return r + + tp_group = get_tp_group() + configure([tp_group]) + # x diverges and IS selected. Expect divergence. + foo(x=tp_group.rank_in_group, y_list=[1, 2, 3]) + shutdown() + + def test_partial_result_diverge(self): + self.assertFalse( + self._spawn( + TestRankConsensusDecorator.consensus_partial_result_diverge_fn, + tp_size=2, + ) + ) + + @staticmethod + def consensus_both_same_fn(): + @rank_consensus(same_params=True, same_results=True) + def foo(a: int) -> int: + return a * 2 + + configure([get_tp_group()]) + foo(7) + shutdown() + + def test_both_same(self): + self.assertTrue( + self._spawn(TestRankConsensusDecorator.consensus_both_same_fn, tp_size=2) + ) + + @staticmethod + def consensus_both_diverge_fn(): + @rank_consensus(same_params=True, same_results=True) + def foo(a: int) -> int: + return a * 2 + + tp_group = get_tp_group() + configure([tp_group]) + foo(tp_group.rank_in_group) + shutdown() + + def test_both_diverge(self): + self.assertFalse( + self._spawn(TestRankConsensusDecorator.consensus_both_diverge_fn, tp_size=2) + ) + + @staticmethod + def consensus_instance_method_same_fn(): + configure([get_tp_group()]) + _MethodHost().instance_method(1, 2) + shutdown() + + def test_instance_method_receiver_dropped(self): + # Two ranks build two different _MethodHost instances; without the + # receiver-skip the per-rank address would diverge. Clean exit + # confirms the receiver is dropped. + self.assertTrue( + self._spawn( + TestRankConsensusDecorator.consensus_instance_method_same_fn, + tp_size=2, + ) + ) + + @staticmethod + def consensus_class_method_same_fn(): + configure([get_tp_group()]) + _MethodHost.class_method(5) + shutdown() + + def test_class_method_receiver_dropped(self): + # First param is named ``klass`` (not cls); detection must still work. + self.assertTrue( + self._spawn( + TestRankConsensusDecorator.consensus_class_method_same_fn, tp_size=2 + ) + ) + + @staticmethod + def consensus_class_method_via_instance_same_fn(): + configure([get_tp_group()]) + _MethodHost().class_method(5) + shutdown() + + def test_class_method_via_instance_receiver_dropped(self): + # Accessing the classmethod through an instance still binds the class + # as the receiver; verify it is still dropped. + self.assertTrue( + self._spawn( + TestRankConsensusDecorator.consensus_class_method_via_instance_same_fn, + tp_size=2, + ) + ) + + @staticmethod + def consensus_static_method_same_fn(): + configure([get_tp_group()]) + _MethodHost.static_method(3, 4) + shutdown() + + def test_static_method_no_receiver(self): + # Static method: no receiver, equal args -> clean. + self.assertTrue( + self._spawn( + TestRankConsensusDecorator.consensus_static_method_same_fn, + tp_size=2, + ) + ) + + @staticmethod + def consensus_static_method_diverge_fn(): + tp_group = get_tp_group() + configure([tp_group]) + # Static method: no receiver to drop, so a rank-dependent arg diverges. + _MethodHost.static_method(tp_group.rank_in_group, 4) + shutdown() + + def test_static_method_flags_diverging_arg(self): + # Static method: no receiver to drop, so a rank-dependent arg must + # still be flagged. Confirms we did not over-skip for static methods. + self.assertFalse( + self._spawn( + TestRankConsensusDecorator.consensus_static_method_diverge_fn, + tp_size=2, + ) + ) + + +if __name__ == "__main__": + unittest.main()