wehub-resource-sync
62 lines
1.9 KiB
Python
62 lines
1.9 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
import pytest
|
|
import torch
|
|
|
|
import vllm.utils.gpu_sync_debug as gsd
|
|
from vllm.utils.gpu_sync_debug import (
|
|
SYNC_ERROR_MESSAGE,
|
|
gpu_sync_allowed,
|
|
with_gpu_sync_check,
|
|
)
|
|
|
|
from ..utils import create_new_process_for_each_test
|
|
|
|
pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
|
|
|
|
|
|
def _no_sync():
|
|
# Pure on-GPU compute, no implicit CPU sync...
|
|
x = torch.ones(4, device="cuda") + 1
|
|
# ...plus a sync that we explicitly allow.
|
|
with gpu_sync_allowed():
|
|
return x.cpu()
|
|
|
|
|
|
def _causes_sync():
|
|
x = torch.ones(4, device="cuda")
|
|
# An allowed sync (suppressed)...
|
|
with gpu_sync_allowed():
|
|
x.cpu()
|
|
# ...then an un-allowed sync that should trip the check.
|
|
return x.cpu()
|
|
|
|
|
|
@pytest.mark.parametrize("mode", ["warn", "error"])
|
|
@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)
|
|
monkeypatch.setattr(gsd, "_sync_check_enabled", True)
|
|
|
|
# Guarded syncs always pass.
|
|
with_gpu_sync_check(_no_sync)()
|
|
|
|
if mode == "error":
|
|
# "error" mode turns the stray sync into a RuntimeError.
|
|
with pytest.raises(RuntimeError, match=SYNC_ERROR_MESSAGE):
|
|
with_gpu_sync_check(_causes_sync)()
|
|
else:
|
|
# "warn" mode only warns, so the call still succeeds.
|
|
with_gpu_sync_check(_causes_sync)()
|
|
|
|
|
|
@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_enabled", True)
|
|
|
|
with_gpu_sync_check(_no_sync)()
|
|
with_gpu_sync_check(_causes_sync)()
|