# 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.

torch's mode only instruments explicit synchronizing calls, so it misses
`non_blocking=True` CPU<->CUDA copies that the CUDA driver silently stages
through pageable memory (when the CPU tensor is not pinned or not densely
laid out). Those are caught by the wrappers on `Tensor.to`/`Tensor.cuda`/
`Tensor.copy_` installed by `_install_copy_checkers()`.
"""

import functools
import sys
import threading
import warnings
from contextlib import contextmanager, nullcontext
from contextvars import ContextVar

import torch

import vllm.envs as envs
from vllm.platforms import current_platform

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

# 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, 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()
    _install_copy_checkers()


_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 _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 _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


# ---------------------------------------------------------------------------
# Implicit sync detection in `non_blocking` CPU<->CUDA copies.
#
# A `non_blocking=True` copy between CPU and CUDA still blocks the host when
# the CPU side is not pinned or not densely laid out (gapped strided slices
# stage through a pageable temp): D2H copies into such memory wait for the
# stream, and H2D copies from it do too beyond a driver size threshold.
# None of that goes through an API torch's sync debug mode instruments, so
# wrap the Python-level transfer entry points and flag such copies on threads
# in a checked region.

IMPLICIT_SYNC_ERROR_MESSAGE = (
    "Implicit GPU<->CPU sync detected - use a pinned, densely laid out CPU "
    "tensor, or wrap with gpu_sync_allowed()"
)

# Captured at import; `_install_copy_checkers` replaces these attributes.
_torch_to = torch.Tensor.to
_torch_cuda = torch.Tensor.cuda
_torch_copy_ = torch.Tensor.copy_

# torch's own parser for `Tensor.to`'s overloaded signature, returning
# (device, dtype, non_blocking, memory_format) (also used by torch._dynamo).
_parse_to = getattr(torch._C._nn, "_parse_to", None)


def _active_check_mode() -> str | None:
    """The mode the calling thread is being checked in, if any."""
    if not _sync_check_enabled or _allow_depth.get() or torch.compiler.is_compiling():
        return None
    return _checking.get()


def _is_dense(t: torch.Tensor) -> bool:
    """Whether `t`'s elements fill a contiguous storage range, possibly
    permuted (e.g. a transposed matrix). Such layouts can be copied with
    pitched cudaMemcpy2D/3DAsync directly from/to pinned memory, whereas
    gapped layouts (e.g. strided slices) stage through a pageable temp."""
    if t.is_contiguous():
        # The common case, answered by a memoized TensorImpl flag.
        return True
    expected_stride = 1
    for stride, size in sorted(zip(t.stride(), t.shape)):
        if size <= 1:
            continue
        if stride != expected_stride:
            return False
        expected_stride *= size
    return True


def _cpu_copy_stall_reason(
    cpu_tensor: torch.Tensor, device: torch.device | None
) -> str | None:
    """Why a `non_blocking` copy between `cpu_tensor` and a CUDA tensor would
    silently block the host, or None if it is genuinely asynchronous.

    `device` describes the other side of the copy; a non-CUDA `device` means
    the copy does not involve the GPU at all.
    """
    if device is None or device.type != "cuda" or cpu_tensor.device.type != "cpu":
        return None
    if cpu_tensor.numel() == 0:
        return None  # Nothing to copy; torch issues no CUDA call at all.
    if not _is_dense(cpu_tensor):
        return "the CPU tensor is not densely laid out"
    if not cpu_tensor.is_pinned():
        return "the CPU tensor is not pinned"
    return None


def _report_implicit_sync(op: str, direction: str, reason: str, mode: str) -> None:
    msg = (
        f"Tensor.{op}: a non_blocking {direction} copy that may be silently "
        f"synchronous ({reason}). {IMPLICIT_SYNC_ERROR_MESSAGE}"
    )
    if mode == "error":
        raise RuntimeError(msg)
    warnings.warn(msg, stacklevel=3)


def _checked_to(self, *args, **kwargs):
    # Only H2D is checked: D2H `to` allocates a pinned destination in
    # supported torch versions, so it is genuinely asynchronous.
    if (mode := _active_check_mode()) is not None and _parse_to is not None:
        try:
            device, _, non_blocking, _ = _parse_to(*args, **kwargs)
        except Exception:
            pass  # Invalid args; let the original call raise.
        else:
            if non_blocking and (reason := _cpu_copy_stall_reason(self, device)):
                _report_implicit_sync("to", "H2D", reason, mode)
    return _torch_to(self, *args, **kwargs)


def _checked_cuda(self, *args, **kwargs):
    # Signature: cuda(device=None, non_blocking=False, memory_format=None).
    if (mode := _active_check_mode()) is not None:
        device = args[0] if args else kwargs.get("device")
        non_blocking = args[1] if len(args) > 1 else kwargs.get("non_blocking", False)
        if non_blocking:
            if device is None:
                device = torch.device("cuda")
            elif isinstance(device, int):
                device = torch.device("cuda", device)
            elif not isinstance(device, torch.device):
                device = torch.device(device)
            if reason := _cpu_copy_stall_reason(self, device):
                _report_implicit_sync("cuda", "H2D", reason, mode)
    return _torch_cuda(self, *args, **kwargs)


def _checked_copy_(self, *args, **kwargs):
    # Signature: copy_(src, non_blocking=False).
    if (
        (mode := _active_check_mode()) is not None
        and args
        and isinstance(src := args[0], torch.Tensor)
    ):
        non_blocking = args[1] if len(args) > 1 else kwargs.get("non_blocking", False)
        if non_blocking:
            if self.device.type == "cuda":
                reason = _cpu_copy_stall_reason(src, self.device)
                direction = "H2D"
            else:
                reason = _cpu_copy_stall_reason(self, src.device)
                direction = "D2H"
            if reason:
                _report_implicit_sync("copy_", direction, reason, mode)
    return _torch_copy_(self, *args, **kwargs)


_copy_checkers_installed: bool = False


def _install_copy_checkers() -> None:
    """Wrap the Python-level CPU<->CUDA transfer entry points with the checks
    above.

    The patch is process-global, but it only reports on threads inside a
    checked region, matching the scoping of the torch-level check. Copies
    issued from C++ (e.g. inside compiled graphs) bypass it.
    """
    global _copy_checkers_installed
    if _copy_checkers_installed:
        return
    _copy_checkers_installed = True
    torch.Tensor.to = _checked_to  # type: ignore[method-assign]
    torch.Tensor.cuda = _checked_cuda  # type: ignore[method-assign]
    torch.Tensor.copy_ = _checked_copy_  # type: ignore[method-assign]


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