diff --git a/tests/utils_/test_gpu_sync_debug.py b/tests/utils_/test_gpu_sync_debug.py index ea9b76f34bc..12001a026a9 100644 --- a/tests/utils_/test_gpu_sync_debug.py +++ b/tests/utils_/test_gpu_sync_debug.py @@ -1,5 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import threading +import warnings + import pytest import torch @@ -36,7 +39,9 @@ def _causes_sync(): @create_new_process_for_each_test() def test_with_env_set(monkeypatch, mode): # Env set + gate flipped on: the unguarded sync is detected. - monkeypatch.setenv("VLLM_GPU_SYNC_CHECK", mode) + # `_SYNC_CHECK_MODE` is read from the env once at import, so patch the + # module attribute rather than the environment. + monkeypatch.setattr(gsd, "_SYNC_CHECK_MODE", mode) monkeypatch.setattr(gsd, "_sync_check_enabled", True) # Guarded syncs always pass. @@ -51,10 +56,104 @@ def test_with_env_set(monkeypatch, mode): with_gpu_sync_check(_causes_sync)() +@create_new_process_for_each_test() +def test_other_threads_are_not_policed(monkeypatch): + """A background thread that syncs deliberately must not be broken by the + check being armed on the thread running the decorated function.""" + monkeypatch.setattr(gsd, "_SYNC_CHECK_MODE", "error") + monkeypatch.setattr(gsd, "_sync_check_enabled", True) + + def sync_on_worker(): + failure: list[BaseException] = [] + + def worker(): + try: + torch.ones(4, device="cuda").cpu() + except BaseException as exc: # pragma: no cover - failure path + failure.append(exc) + + thread = threading.Thread(target=worker) + thread.start() + thread.join() + assert not failure, f"background thread raised: {failure[0]!r}" + + with_gpu_sync_check(sync_on_worker)() + + +@create_new_process_for_each_test() +def test_allow_on_other_thread_does_not_disarm(monkeypatch): + """`gpu_sync_allowed()` on one thread must not suppress the check on + another. It is scoped by ContextVar rather than torch's process-global + sync debug mode, which a previous implementation mutated.""" + monkeypatch.setattr(gsd, "_SYNC_CHECK_MODE", "error") + monkeypatch.setattr(gsd, "_sync_check_enabled", True) + + def main_syncs_while_worker_allows(): + stop = threading.Event() + + def worker(): + with gpu_sync_allowed(): + while not stop.is_set(): + torch.ones(4, device="cuda").cpu() + + thread = threading.Thread(target=worker) + thread.start() + try: + # Must still be reported despite the worker's open allow region. + torch.ones(4, device="cuda").cpu() + finally: + stop.set() + thread.join() + + with pytest.raises(RuntimeError, match=SYNC_ERROR_MESSAGE): + with_gpu_sync_check(main_syncs_while_worker_allows)() + + +@create_new_process_for_each_test() +def test_suppressing_works_while_compiling(monkeypatch): + """`_suppressing` wraps torch compile entry points, which run with + `torch.compiler.is_compiling()` true. `gpu_sync_allowed()` deliberately + no-ops in that state, so `_suppressing` must not route through it.""" + monkeypatch.setattr(gsd, "_SYNC_CHECK_MODE", "error") + monkeypatch.setattr(gsd, "_sync_check_enabled", True) + # Emulate being inside a torch compile, as inductor passes are. + monkeypatch.setattr(torch.compiler, "is_compiling", lambda: True) + + suppressed = gsd._suppressing(lambda: torch.ones(4, device="cuda").cpu()) + with_gpu_sync_check(suppressed)() + + +@create_new_process_for_each_test() +def test_sync_debug_mode_restored_after_checked_call(monkeypatch): + """The mode is armed only for the duration of a checked call. Leaving it + on process-wide made every sync outside a checked region emit a + `UserWarning` whenever our handler was not the installed one.""" + monkeypatch.setattr(gsd, "_SYNC_CHECK_MODE", "error") + monkeypatch.setattr(gsd, "_sync_check_enabled", True) + + before = torch.cuda.get_sync_debug_mode() + + def nested(): + # `execute_model` -> `sample_tokens` both carry the decorator. + assert torch.cuda.get_sync_debug_mode() != 0, "armed inside" + with_gpu_sync_check(lambda: None)() + assert torch.cuda.get_sync_debug_mode() != 0, "still armed after inner" + + with_gpu_sync_check(nested)() + assert torch.cuda.get_sync_debug_mode() == before + + # With the mode back to its original value, torch emits nothing for a + # sync outside a checked region. + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + torch.ones(4, device="cuda").cpu() + assert not [r for r in caught if gsd._TORCH_SYNC_WARNING in str(r.message)] + + @create_new_process_for_each_test() def test_without_env_set(monkeypatch): # Env unset: the decorator is a pass-through, no sync is detected. - monkeypatch.delenv("VLLM_GPU_SYNC_CHECK", raising=False) + monkeypatch.setattr(gsd, "_SYNC_CHECK_MODE", None) monkeypatch.setattr(gsd, "_sync_check_enabled", True) with_gpu_sync_check(_no_sync)() diff --git a/vllm/utils/gpu_sync_debug.py b/vllm/utils/gpu_sync_debug.py index 1e2114f3b5b..758b73e8422 100644 --- a/vllm/utils/gpu_sync_debug.py +++ b/vllm/utils/gpu_sync_debug.py @@ -1,8 +1,20 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Detect unintended GPU<->CPU syncs in the hot path. + +`torch.cuda.set_sync_debug_mode` is process-global, so we arm it at "warn" +(which never raises) and decide in `_sync_warning_hook` whether a given sync +is a failure. The scoping lives in ContextVars, which are per-thread and +per-asyncio-task, so an allow region opened on one thread is invisible to +every other one. +""" + import functools import sys -from contextlib import contextmanager +import threading +import warnings +from contextlib import contextmanager, nullcontext +from contextvars import ContextVar import torch @@ -13,153 +25,210 @@ SYNC_ERROR_MESSAGE = ( "GPU<->CPU sync detected - avoid it or wrap with gpu_sync_allowed()" ) +# Warning text torch emits in "warn" mode. +_TORCH_SYNC_WARNING = "called a synchronizing CUDA operation" + _GPU_SYNC_ALLOWED_FIRST_SEEN: set[tuple[str, int]] = set() -# Global sync-check gate. Off during engine setup (model load, KV cache -# init, warmup/compile) so first-compile and lazy-init syncs pass through; -# flipped on by `enable_gpu_sync_check()` at the end of -# `GPUWorker.compile_or_warm_up_model`, after which `with_gpu_sync_check`- -# decorated functions activate the configured debug mode. +# Read once: `envs.__getattr__` re-evaluates the env lambda on every access, +# which dominates the cost of a disabled `gpu_sync_allowed()`. Tests toggling +# the env var must patch this attribute instead. +_SYNC_CHECK_MODE: str | None = ( + envs.VLLM_GPU_SYNC_CHECK if current_platform.is_cuda_alike() else None +) + +# Configured mode while a thread is inside a checked call, else None. +_checking: ContextVar[str | None] = ContextVar("vllm_gpu_sync_checking", default=None) +# Depth of nested `gpu_sync_allowed()` regions on this thread. +_allow_depth: ContextVar[int] = ContextVar("vllm_gpu_sync_allow_depth", default=0) + +# Off during engine setup so model load, KV cache init and warmup may sync +# freely; flipped on by `enable_gpu_sync_check()` once warmup completes. _sync_check_enabled: bool = False def enable_gpu_sync_check() -> None: - """Flip the sync-check gate on. Call once per worker, after warmup / - first-compile is complete. No-op unless `VLLM_GPU_SYNC_CHECK` is set.""" - if envs.VLLM_GPU_SYNC_CHECK is None: + """Flip the sync-check gate on, once per worker, after warmup.""" + if _SYNC_CHECK_MODE is None: return global _sync_check_enabled _sync_check_enabled = True _install_compile_time_sync_suppressors() +_arm_lock = threading.Lock() +_arm_count: int = 0 +_saved_sync_debug_mode: int = 0 +_prev_showwarning = warnings.showwarning +_sync_filter_head: tuple | None = None + + +def _sync_warning_hook(message, category, filename, lineno, file=None, line=None): + """Turn torch's sync warning into an error, but only on a thread that is + being checked and is outside any allow region.""" + if _TORCH_SYNC_WARNING in str(message): + mode = _checking.get() + if mode is None or _allow_depth.get(): + return None + if mode == "error": + raise RuntimeError(SYNC_ERROR_MESSAGE) + return _prev_showwarning(message, category, filename, lineno, file, line) + + +def _install_warning_hook() -> None: + """(Re)install the hook and a filter that lets torch's warning reach it. + + Done per checked call because pytest runs each test inside + `warnings.catch_warnings()`, which restores both `showwarning` and + `filters`. The hook is left in place afterwards: outside a checked call + the debug mode is disarmed, so torch emits nothing for it to see. + """ + global _prev_showwarning, _sync_filter_head + if warnings.showwarning is not _sync_warning_hook: + _prev_showwarning = warnings.showwarning + warnings.showwarning = _sync_warning_hook + # "always" so the warning survives filtering and isn't deduplicated by + # `__warningregistry__`. `filterwarnings` prepends, so only re-assert it + # once ours is no longer in front. + if warnings.filters[:1] != [_sync_filter_head]: + warnings.filterwarnings( + "always", message=_TORCH_SYNC_WARNING, category=UserWarning + ) + _sync_filter_head = warnings.filters[0] + + +@contextmanager +def _checked_region(mode: str): + """Police syncs on this thread for the duration of the block. + + The debug mode is armed per call so that syncs outside a checked region + emit nothing, and refcounted because `execute_model` and `sample_tokens` + nest. "warn" rather than "error" because torch's error mode raises on + whichever thread synced, with no way to exempt one. + """ + global _arm_count, _saved_sync_debug_mode + _install_warning_hook() + with _arm_lock: + if _arm_count == 0: + _saved_sync_debug_mode = torch.cuda.get_sync_debug_mode() + torch.cuda.set_sync_debug_mode("warn") + _arm_count += 1 + token = _checking.set(mode) + try: + yield + finally: + _checking.reset(token) + with _arm_lock: + _arm_count -= 1 + if _arm_count == 0: + torch.cuda.set_sync_debug_mode(_saved_sync_debug_mode) + + +@contextmanager +def _allow_syncs(): + token = _allow_depth.set(_allow_depth.get() + 1) + try: + yield + finally: + _allow_depth.reset(token) + + +# Shared, since `nullcontext` is stateless: reusing one instance avoids an +# allocation per call, which is most of the cost when the check is disabled. +_NOOP_CM = nullcontext() + _compile_time_suppressors_installed: bool = False -def _install_compile_time_sync_suppressors() -> None: - """Wrap torch inductor/aot_autograd compile entry points so the - synchronizing ops those passes perform don't trip the - sync-check mode we set around `execute_model` / `sample_tokens`. +def _suppressing(fn): + """Allow the syncs `fn` performs on its calling thread.""" - Warmup-time compiles already run under the gate (before - `enable_gpu_sync_check`), but post-warmup compiles fire inside - `execute_model` and we want to avoid this tripping the sync check. + @functools.wraps(fn) + def wrapper(*args, **kwargs): + # Not `gpu_sync_allowed()`, which no-ops while + # `torch.compiler.is_compiling()` -- exactly when these run. + with _allow_syncs(): + return fn(*args, **kwargs) + + return wrapper + + +def _install_compile_time_sync_suppressors() -> None: + """Allow the syncs torch's compile passes perform. + + Warmup-time compiles run before the gate flips, but lazy ones fire inside + `execute_model`. """ global _compile_time_suppressors_installed if _compile_time_suppressors_installed: return _compile_time_suppressors_installed = True - try: # noqa: BLE001 + try: from torch._inductor.fx_passes import joint_graph as _jg - _orig_joint = _jg.joint_graph_passes - - @functools.wraps(_orig_joint) - def _wrapped_joint(*args, **kwargs): - prev_mode = torch.cuda.get_sync_debug_mode() - if not prev_mode: - return _orig_joint(*args, **kwargs) - torch.cuda.set_sync_debug_mode(0) - try: - return _orig_joint(*args, **kwargs) - finally: - torch.cuda.set_sync_debug_mode(prev_mode) - - # `compile_fx` does `from .fx_passes.joint_graph import - # joint_graph_passes`, which binds the *function object* at import - # time. Patching just the module attribute won't update that rebind, - # so patch every already-imported reference we can find. Restrict - # the scan to torch's compile-time modules. - import sys as _sys - - setattr(_jg, "joint_graph_passes", _wrapped_joint) # noqa: B010 - for _name, _mod in list(_sys.modules.items()): - if _mod is None: - continue - if not ( - _name.startswith("torch._inductor") - or _name.startswith("torch._functorch") - or _name.startswith("torch._dynamo") + orig = _jg.joint_graph_passes + wrapped = _suppressing(orig) + # `compile_fx` imports this by value, so patching the defining module + # alone misses that rebind; patch every compile-time module still + # holding the original. + _jg.joint_graph_passes = wrapped + for name, mod in list(sys.modules.items()): + if ( + mod is not None + and name.startswith( + ("torch._inductor", "torch._functorch", "torch._dynamo") + ) + and getattr(mod, "joint_graph_passes", None) is orig ): - continue - if getattr(_mod, "joint_graph_passes", None) is _orig_joint: - setattr(_mod, "joint_graph_passes", _wrapped_joint) # noqa: B010 + setattr(mod, "joint_graph_passes", wrapped) # noqa: B010 + except Exception: # pragma: no cover + pass + + try: + # Inductor builds its cudagraph tree lazily, so `deferred_cudagraphify` + # and the `capture_begin` sync inside it can fire during + # `execute_model`. It resolves `cudagraphify` as a module global at + # call time, so patching the attribute is enough. + from torch._inductor import cudagraph_trees as _ct + + _ct.cudagraphify = _suppressing(_ct.cudagraphify) except Exception: # pragma: no cover pass -@contextmanager -def _suppress_gpu_sync_check(prev_mode: int): - torch.cuda.set_sync_debug_mode(0) - try: - yield - finally: - torch.cuda.set_sync_debug_mode(prev_mode) +def gpu_sync_allowed(first_only: bool = False): + """Allow GPU<->CPU syncs inside the `with` block, on this thread only. + + With `first_only`, only the first entry from a given call site + (filename, lineno) is allowed, so later syncs there are still reported. + """ + if _SYNC_CHECK_MODE is None or torch.compiler.is_compiling(): + return _NOOP_CM + if first_only: + frame = sys._getframe(1) + key = (frame.f_code.co_filename, frame.f_lineno) + if key in _GPU_SYNC_ALLOWED_FIRST_SEEN: + return _NOOP_CM + _GPU_SYNC_ALLOWED_FIRST_SEEN.add(key) + return _allow_syncs() -@contextmanager -def _noop_cm(): - yield +def with_gpu_sync_check(fn): + """Report GPU<->CPU syncs performed by `fn` on its calling thread. - -if current_platform.is_cuda_alike(): - - def gpu_sync_allowed(first_only: bool = False): - """Context manager that suppresses `torch.cuda.set_sync_debug_mode` for the - duration of the `with` block. - - If `first_only` is True, only the first entry from this call site - suppresses the sync check; subsequent entries from the same site are - no-ops so any further GPU syncs will be reported. The "site" is the - caller's (filename, lineno), so different - `with gpu_sync_allowed(first_only=True):` lines track independently. - """ - if envs.VLLM_GPU_SYNC_CHECK is None or torch.compiler.is_compiling(): - return _noop_cm() - prev_mode = torch.cuda.get_sync_debug_mode() - if not prev_mode: - return _noop_cm() - if first_only: - frame = sys._getframe(1) - key = (frame.f_code.co_filename, frame.f_lineno) - if key in _GPU_SYNC_ALLOWED_FIRST_SEEN: - return _noop_cm() - _GPU_SYNC_ALLOWED_FIRST_SEEN.add(key) - return _suppress_gpu_sync_check(prev_mode) - - def with_gpu_sync_check(fn): - """Decorator that enables `torch.cuda.set_sync_debug_mode` around `fn` - when `VLLM_GPU_SYNC_CHECK` is set *and* the gate has been flipped by - `enable_gpu_sync_check()`. Before the gate flips (i.e. during - engine setup / warmup) the decorated function runs as-is. - """ - mode = envs.VLLM_GPU_SYNC_CHECK - if mode is None: - return fn - - @functools.wraps(fn) - def wrapper(*args, **kwargs): - if not _sync_check_enabled: - return fn(*args, **kwargs) - prev_mode = torch.cuda.get_sync_debug_mode() - torch.cuda.set_sync_debug_mode(mode) - try: - return fn(*args, **kwargs) - except RuntimeError as re: - if str(re) == "called a synchronizing CUDA operation": - raise RuntimeError(SYNC_ERROR_MESSAGE) from re - raise re - finally: - torch.cuda.set_sync_debug_mode(prev_mode) - - return wrapper - -else: - # No-op the methods in non-CUDA cases. - - def gpu_sync_allowed(first_only: bool = False): - return _noop_cm() - - def with_gpu_sync_check(fn): + Active only once `enable_gpu_sync_check()` has flipped the gate. Other + threads are never policed, so deliberate syncs there (e.g. the EPLB + transfer worker) are unaffected. + """ + if (mode := _SYNC_CHECK_MODE) is None: return fn + + @functools.wraps(fn) + def wrapper(*args, **kwargs): + if not _sync_check_enabled: + return fn(*args, **kwargs) + with _checked_region(mode): + return fn(*args, **kwargs) + + return wrapper