vllm.utils.gpu_sync_debug ¶

@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)

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:
        from torch._inductor.fx_passes import joint_graph as _jg

        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
            ):
                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

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]

def _suppressing(fn):
    """Allow the syncs `fn` performs on its calling thread."""

    @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 _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 enable_gpu_sync_check() -> 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()

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()

def with_gpu_sync_check(fn):
    """Report GPU<->CPU syncs performed by `fn` on its calling thread.

    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