264 lines
9.8 KiB
Python
264 lines
9.8 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
"""Characterization tests for ``vllm_omni.worker.base.OmniGPUWorkerBase``.
|
|
|
|
Pins the CURRENT behaviour of ``determine_available_memory`` (device-level
|
|
profiling is the only path — the NVML / process-scoped arm was removed with
|
|
parallel stage init, which coordinates concurrent same-device measurement via
|
|
admission + SH/EX device locks instead), plus the
|
|
memory-pool / sleep / wake-up plumbing that a later change may touch.
|
|
|
|
Workers are constructed via ``object.__new__`` with only the attributes each
|
|
method reads; module-level collaborators are monkeypatched. Pure CPU.
|
|
"""
|
|
|
|
from contextlib import contextmanager, nullcontext
|
|
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
import vllm_omni.worker.base as base
|
|
from vllm_omni.worker.base import OmniGPUWorkerBase
|
|
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
|
|
|
GIB = 1024**3
|
|
|
|
|
|
def _fake_memory_profiling(*, non_torch: int, torch_peak: int):
|
|
"""A stand-in for ``vllm.utils.mem_utils.memory_profiling`` context manager."""
|
|
|
|
@contextmanager
|
|
def _mp(snapshot, weights_memory): # noqa: ARG001 - signature parity only
|
|
# total_consumed mirrors upstream vllm.utils.mem_utils.MemoryProfilingResult
|
|
# (added by upstream 58b2012aa2) and is read by OmniGPUWorkerBase.
|
|
yield SimpleNamespace(
|
|
non_torch_increase=non_torch,
|
|
torch_peak_increase=torch_peak,
|
|
total_consumed=torch_peak + non_torch,
|
|
)
|
|
|
|
return _mp
|
|
|
|
|
|
def _make_worker(*, requested_memory: int, kv_cache_memory_bytes: int = 0, model_memory_usage: int = 0):
|
|
worker = object.__new__(OmniGPUWorkerBase)
|
|
worker.cache_config = SimpleNamespace(kv_cache_memory_bytes=kv_cache_memory_bytes)
|
|
worker.model_runner = SimpleNamespace(
|
|
profile_run=lambda: None,
|
|
model_memory_usage=model_memory_usage,
|
|
)
|
|
worker.init_snapshot = object() # unused once memory_profiling is faked
|
|
worker.requested_memory = requested_memory
|
|
worker.local_rank = 0
|
|
return worker
|
|
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# determine_available_memory #
|
|
# --------------------------------------------------------------------------- #
|
|
def test_no_process_scoped_collaborators_remain():
|
|
"""The NVML arm is gone from base.py, so the method CANNOT take a
|
|
process-scoped path on any host: its collaborators are not even imported.
|
|
(The removal itself happened in the parallel-stage-init feature commit;
|
|
tests/engine/test_parallel_stage_init.py pins the same invariant.)"""
|
|
assert not hasattr(base, "is_process_scoped_memory_available")
|
|
assert not hasattr(base, "detect_pid_host")
|
|
assert not hasattr(base, "get_process_gpu_memory")
|
|
|
|
|
|
def test_determine_available_memory_profiling_path(monkeypatch):
|
|
"""available = requested - (weights + peak + non_torch); the only path."""
|
|
worker = _make_worker(requested_memory=30 * GIB, model_memory_usage=10 * GIB)
|
|
monkeypatch.setattr(base, "memory_profiling", _fake_memory_profiling(non_torch=1 * GIB, torch_peak=2 * GIB))
|
|
|
|
out = OmniGPUWorkerBase.determine_available_memory(worker)
|
|
|
|
profiled = 10 * GIB + 2 * GIB + 1 * GIB
|
|
assert out == 30 * GIB - profiled
|
|
|
|
|
|
def test_determine_available_memory_populates_total_consumed(monkeypatch):
|
|
"""total_consumed mirrors upstream's MemoryProfilingResult field."""
|
|
worker = _make_worker(requested_memory=30 * GIB, model_memory_usage=10 * GIB)
|
|
monkeypatch.setattr(base, "memory_profiling", _fake_memory_profiling(non_torch=1 * GIB, torch_peak=2 * GIB))
|
|
|
|
OmniGPUWorkerBase.determine_available_memory(worker)
|
|
|
|
assert worker.total_consumed == 3 * GIB
|
|
|
|
|
|
def test_determine_available_memory_clamps_to_zero(monkeypatch):
|
|
"""Over-subscription clamps the KV budget to 0, never negative."""
|
|
worker = _make_worker(requested_memory=5 * GIB, model_memory_usage=8 * GIB)
|
|
monkeypatch.setattr(base, "memory_profiling", _fake_memory_profiling(non_torch=0, torch_peak=0))
|
|
|
|
out = OmniGPUWorkerBase.determine_available_memory(worker)
|
|
|
|
assert out == 0
|
|
|
|
|
|
def test_determine_available_memory_kv_cache_short_circuit(monkeypatch):
|
|
"""A pre-set kv_cache_memory_bytes short-circuits profiling and is returned."""
|
|
worker = _make_worker(requested_memory=30 * GIB, kv_cache_memory_bytes=7 * GIB)
|
|
profile_calls = []
|
|
worker.model_runner.profile_run = lambda: profile_calls.append(True)
|
|
# is_rocm gates only an extra synchronize on the short-circuit path.
|
|
monkeypatch.setattr(base, "current_omni_platform", SimpleNamespace(is_rocm=lambda: False))
|
|
|
|
out = OmniGPUWorkerBase.determine_available_memory(worker)
|
|
|
|
assert out == 7 * GIB
|
|
assert profile_calls == [True] # profile_run still runs before returning
|
|
|
|
|
|
def test_determine_available_memory_kv_short_circuit_syncs_on_rocm(monkeypatch):
|
|
"""On ROCm the short-circuit adds an accelerator synchronize (a later change may remove it)."""
|
|
worker = _make_worker(requested_memory=30 * GIB, kv_cache_memory_bytes=7 * GIB)
|
|
monkeypatch.setattr(base, "current_omni_platform", SimpleNamespace(is_rocm=lambda: True))
|
|
synced = []
|
|
monkeypatch.setattr(torch.accelerator, "synchronize", lambda: synced.append(True))
|
|
|
|
out = OmniGPUWorkerBase.determine_available_memory(worker)
|
|
|
|
assert out == 7 * GIB
|
|
assert synced == [True] # ROCm path synchronizes; non-ROCm path (above) does not
|
|
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# _maybe_get_memory_pool_context #
|
|
# --------------------------------------------------------------------------- #
|
|
def test_memory_pool_context_disabled_returns_nullcontext():
|
|
worker = object.__new__(OmniGPUWorkerBase)
|
|
worker.rank = 0
|
|
worker.cache_config = SimpleNamespace(enable_sleep_mode=False)
|
|
# No vllm_config attribute -> v1_config_enabled stays False.
|
|
|
|
ctx = OmniGPUWorkerBase._maybe_get_memory_pool_context(worker, "weights")
|
|
|
|
assert isinstance(ctx, type(nullcontext()))
|
|
|
|
|
|
def test_memory_pool_context_enabled_uses_cumem_pool_with_tag(monkeypatch):
|
|
import vllm.device_allocator.cumem as cumem_mod
|
|
|
|
worker = object.__new__(OmniGPUWorkerBase)
|
|
worker.rank = 0
|
|
worker.cache_config = SimpleNamespace(enable_sleep_mode=True)
|
|
monkeypatch.setattr(base, "current_omni_platform", SimpleNamespace(synchronize=lambda: None))
|
|
|
|
recorded = {}
|
|
|
|
class FakeAllocator:
|
|
@classmethod
|
|
def get_instance(cls):
|
|
return cls()
|
|
|
|
def use_memory_pool(self, tag):
|
|
recorded["tag"] = tag
|
|
return "POOL_CTX"
|
|
|
|
monkeypatch.setattr(cumem_mod, "CuMemAllocator", FakeAllocator)
|
|
|
|
ctx = OmniGPUWorkerBase._maybe_get_memory_pool_context(worker, "weights")
|
|
|
|
assert ctx == "POOL_CTX"
|
|
assert recorded["tag"] == "weights"
|
|
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# sleep / wake_up #
|
|
# --------------------------------------------------------------------------- #
|
|
def _fake_platform():
|
|
return SimpleNamespace(
|
|
get_current_memory_usage=lambda device: 0,
|
|
empty_cache=lambda: None,
|
|
synchronize=lambda: None,
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("level", "expected_tags"),
|
|
[(1, ("weights",)), (2, tuple())],
|
|
)
|
|
def test_sleep_offload_tags_by_level(monkeypatch, level, expected_tags):
|
|
import vllm.device_allocator.cumem as cumem_mod
|
|
|
|
worker = object.__new__(OmniGPUWorkerBase)
|
|
worker.rank = 0
|
|
worker.device = "cuda:0"
|
|
order: list[str] = []
|
|
platform = SimpleNamespace(
|
|
get_current_memory_usage=lambda device: 0,
|
|
empty_cache=lambda: order.append("empty_cache"),
|
|
synchronize=lambda: order.append("sync"),
|
|
collect=lambda: None,
|
|
)
|
|
monkeypatch.setattr(base, "current_omni_platform", platform)
|
|
|
|
calls = {}
|
|
|
|
class FakeAllocator:
|
|
@classmethod
|
|
def get_instance(cls):
|
|
return cls()
|
|
|
|
def sleep(self, offload_tags):
|
|
order.append("allocator_sleep")
|
|
calls["offload_tags"] = offload_tags
|
|
|
|
monkeypatch.setattr(cumem_mod, "CuMemAllocator", FakeAllocator)
|
|
|
|
assert OmniGPUWorkerBase.sleep(worker, level=level) is True
|
|
assert calls["offload_tags"] == expected_tags
|
|
assert order[0] == "sync"
|
|
assert "allocator_sleep" in order
|
|
assert order.index("sync") < order.index("allocator_sleep")
|
|
assert "empty_cache" not in order
|
|
|
|
|
|
def test_wake_up_forwards_tags_to_allocator(monkeypatch):
|
|
import vllm.device_allocator.cumem as cumem_mod
|
|
|
|
worker = object.__new__(OmniGPUWorkerBase)
|
|
worker.rank = 0
|
|
monkeypatch.setattr(base, "current_omni_platform", SimpleNamespace(synchronize=lambda: None))
|
|
|
|
calls = {}
|
|
|
|
class FakeAllocator:
|
|
@classmethod
|
|
def get_instance(cls):
|
|
return cls()
|
|
|
|
def wake_up(self, tags):
|
|
calls["tags"] = tags
|
|
|
|
monkeypatch.setattr(cumem_mod, "CuMemAllocator", FakeAllocator)
|
|
|
|
assert OmniGPUWorkerBase.wake_up(worker, tags=["weights"]) is True
|
|
assert calls["tags"] == ["weights"]
|
|
|
|
|
|
def test_cuda_profiler_is_wired_for_omni_worker(monkeypatch):
|
|
import vllm.profiler.wrapper as profiler_wrapper
|
|
|
|
def fake_gpu_worker_init(worker, *args, **kwargs):
|
|
worker.vllm_config = SimpleNamespace(profiler_config=SimpleNamespace(profiler="cuda"))
|
|
worker.local_rank = 0
|
|
worker.rank = 0
|
|
worker.profiler = None
|
|
|
|
class FakeCudaProfiler:
|
|
def __init__(self, profiler_config):
|
|
self.profiler_config = profiler_config
|
|
|
|
monkeypatch.setattr(base.GPUWorker, "__init__", fake_gpu_worker_init)
|
|
monkeypatch.setattr(profiler_wrapper, "CudaProfilerWrapper", FakeCudaProfiler)
|
|
|
|
worker = OmniGPUWorkerBase()
|
|
|
|
assert isinstance(worker.profiler, FakeCudaProfiler)
|
|
assert worker.profiler.profiler_config.profiler == "cuda"
|