Skip to content

GitLab

Unverified Commit b78ec99a authored by Schwinn Saereesitthipitak's avatar Schwinn Saereesitthipitak Committed by GitHub
Browse files

test: reorganize and prune GPU Memory Service tests (#7828) 


Signed-off-by: default avatarSchwinn Saereesitthipitak <schwinns@nvidia.com>
parent 11bb8498
Expand all Hide whitespace changes
Inline Side-by-side
Showing with 1264 additions and 929 deletions
tests/gms/integration/test_gms_shadow_failover.py deleted 100644 → 0
View file @ 11bb8498
This diff is collapsed.
tests/gms/integration/test_gms_sleep_wake.py deleted 100644 → 0
View file @ 11bb8498
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
from __future__ import annotations
import logging
import time
from contextlib import ExitStack
from typing import Callable
import pytest
from gpu_memory_service.server.fsm import ServerState
from tests.utils.constants import FAULT_TOLERANCE_MODEL_NAME
from tests.utils.managed_process import DynamoFrontendProcess, ManagedProcess
from ..harness.gms import GMSServerProcess
from ..harness.runtime import (
MIN_EXPECTED_MEMORY_RETURN_FRACTION,
get_gpu_memory_used,
send_completion,
)
from ..harness.sglang import SGLangWithGMSProcess
from ..harness.vllm import VLLMWithGMSProcess
pytestmark = [pytest.mark.nightly]
# Event flow under test:
# 1. Weights are published once as a committed layout.
# 2. KV cache starts as a live RW layout build.
# 3. Sleep keeps weights committed but aborts and clears the KV layout.
# 4. Wake reconnects weights as RO to the same committed layout.
# 5. Wake recreates KV cache in a fresh RW layout after the old one was cleared.
logger = logging.getLogger(__name__)
def _run_sleep_wake_test(
request,
ports: dict,
make_engine: Callable[[], ManagedProcess],
) -> None:
with ExitStack() as stack:
weights_gms = stack.enter_context(
GMSServerProcess(request, device=0, tag="weights")
)
kv_cache_gms = stack.enter_context(
GMSServerProcess(request, device=0, tag="kv_cache")
)
stack.enter_context(
DynamoFrontendProcess(request, frontend_port=ports["frontend"])
)
with make_engine() as engine:
result = send_completion(ports["frontend"])
logger.info("Initial inference result: %s", result)
assert result["choices"]
# Before sleep, weights must already be published and visible to RO
# readers while KV cache remains a live RW layout owned by the engine.
deadline = time.monotonic() + 30.0
while True:
weights_before_sleep = weights_gms.get_runtime_state()
kv_before_sleep = kv_cache_gms.get_runtime_state()
if (
weights_before_sleep.state == ServerState.RO
and weights_before_sleep.allocation_count > 0
and weights_before_sleep.memory_layout_hash
and kv_before_sleep.state == ServerState.RW
and kv_before_sleep.allocation_count > 0
):
break
if time.monotonic() > deadline:
raise TimeoutError("initial GMS state did not stabilize")
time.sleep(0.1)
mem_before = get_gpu_memory_used()
logger.info("Memory before sleep: %.0f MB", mem_before / (1 << 20))
sleep_result = engine.sleep()
assert sleep_result["status"] == "ok"
mem_after_sleep = get_gpu_memory_used()
released_bytes = mem_before - mem_after_sleep
logger.info("Memory after sleep: %.0f MB", mem_after_sleep / (1 << 20))
assert mem_after_sleep < mem_before, "Sleep should reduce memory"
assert released_bytes > 0
# Sleep preserves the committed weights layout but aborts and clears the
# mutable KV-cache layout, which is what should release GPU memory.
deadline = time.monotonic() + 30.0
while True:
weights_after_sleep = weights_gms.get_runtime_state()
kv_after_sleep = kv_cache_gms.get_runtime_state()
if (
weights_after_sleep.state == ServerState.COMMITTED
and weights_after_sleep.allocation_count
== weights_before_sleep.allocation_count
and weights_after_sleep.memory_layout_hash
== weights_before_sleep.memory_layout_hash
and kv_after_sleep.state == ServerState.EMPTY
and kv_after_sleep.allocation_count == 0
):
break
if time.monotonic() > deadline:
raise TimeoutError(
"sleep did not drive GMS into the expected state"
)
time.sleep(0.1)
# Weights are immutable across sleep/wake, so their event history should
# still be the original publish: connect once, commit once.
weights_events = weights_gms.get_event_history().events
assert [event.kind for event in weights_events] == [
"rw_connected",
"committed",
]
# KV cache is different: sleep must abort the old RW layout and clear its
# server-owned allocations before wake can start a new RW layout.
kv_events = kv_cache_gms.get_event_history().events
assert [event.kind for event in kv_events] == [
"rw_connected",
"rw_aborted",
"allocations_cleared",
]
assert kv_events[-1].allocation_count > 0
wake_result = engine.wake()
assert wake_result["status"] == "ok"
mem_after_wake = get_gpu_memory_used()
reacquired_bytes = mem_after_wake - mem_after_sleep
logger.info("Memory after wake: %.0f MB", mem_after_wake / (1 << 20))
assert mem_after_wake > mem_after_sleep, "Wake should reacquire memory"
assert (
reacquired_bytes
) >= released_bytes * MIN_EXPECTED_MEMORY_RETURN_FRACTION
# Wake reconnects weights as RO to the same committed layout, but KV cache
# must come back as a fresh RW layout with new allocations.
deadline = time.monotonic() + 30.0
while True:
weights_after_wake = weights_gms.get_runtime_state()
kv_after_wake = kv_cache_gms.get_runtime_state()
if (
weights_after_wake.state == ServerState.RO
and weights_after_wake.allocation_count
== weights_before_sleep.allocation_count
and weights_after_wake.memory_layout_hash
== weights_before_sleep.memory_layout_hash
and kv_after_wake.state == ServerState.RW
and kv_after_wake.allocation_count > 0
):
break
if time.monotonic() > deadline:
raise TimeoutError("wake did not restore the expected GMS state")
time.sleep(0.1)
weights_events_after_wake = weights_gms.get_event_history().events
assert [event.kind for event in weights_events_after_wake] == [
"rw_connected",
"committed",
]
# The wake history should therefore extend the old KV sequence with one
# new RW connect after the previous layout was fully cleared.
kv_events_after_wake = kv_cache_gms.get_event_history().events
assert [event.kind for event in kv_events_after_wake] == [
"rw_connected",
"rw_aborted",
"allocations_cleared",
"rw_connected",
]
assert kv_events_after_wake[2].allocation_count > 0
result = send_completion(ports["frontend"], "Goodbye")
logger.info("Post-wake inference result: %s", result)
assert result["choices"]
logger.info(
"Memory freed: %.0f MB", (mem_before - mem_after_sleep) / (1 << 20)
)
@pytest.mark.vllm
@pytest.mark.e2e
@pytest.mark.gpu_1
@pytest.mark.model(FAULT_TOLERANCE_MODEL_NAME)
@pytest.mark.timeout(300)
def test_gms_basic_sleep_wake_vllm(
request,
runtime_services_dynamic_ports,
gms_ports,
predownload_models,
):
ports = gms_ports
_run_sleep_wake_test(
request,
ports,
make_engine=lambda: VLLMWithGMSProcess(
request,
"engine",
ports["shadow_system"],
ports["shadow_kv_event"],
ports["shadow_nixl"],
ports["frontend"],
),
)
@pytest.mark.skip(reason="Nightly CI failure: https://linear.app/nvidia/issue/DYN-2567")
@pytest.mark.sglang
@pytest.mark.e2e
@pytest.mark.gpu_1
@pytest.mark.model(FAULT_TOLERANCE_MODEL_NAME)
@pytest.mark.timeout(300)
def test_gms_basic_sleep_wake_sglang(
request,
runtime_services_dynamic_ports,
gms_ports,
predownload_models,
):
ports = gms_ports
_run_sleep_wake_test(
request,
ports,
make_engine=lambda: SGLangWithGMSProcess(
request,
"engine",
ports["shadow_system"],
ports["shadow_sglang"],
ports["frontend"],
),
)
tests/gms/integration/test_gms_vllm_worker.py deleted 100644 → 0
View file @ 11bb8498
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
from __future__ import annotations
from contextlib import contextmanager
from types import SimpleNamespace
import pytest
pytestmark = [
pytest.mark.pre_merge,
pytest.mark.unit,
pytest.mark.gpu_0,
pytest.mark.vllm,
]
class _FakeManager:
def __init__(self, *, is_unmapped: bool = False) -> None:
self.is_unmapped = is_unmapped
self.calls: list[object] = []
def unmap_all_vas(self) -> None:
self.calls.append("unmap_all_vas")
self.is_unmapped = True
def abort(self) -> None:
self.calls.append("abort")
def connect(self, lock_type, timeout_ms=None) -> None:
self.calls.append(("connect", lock_type.value))
self.is_unmapped = False
def reallocate_all_handles(self, *, tag: str) -> None:
self.calls.append(("reallocate_all_handles", tag))
def remap_all_vas(self) -> None:
self.calls.append("remap_all_vas")
self.is_unmapped = False
def test_initialize_from_config_uses_kv_cache_gms_tag(monkeypatch):
import gpu_memory_service.integrations.vllm.worker as worker_module
import vllm.distributed.kv_transfer as kv_transfer
from gpu_memory_service.integrations.vllm.worker import GMSWorker
create_calls: list[tuple[object, ...]] = []
pool_calls: list[tuple[str, str]] = []
kv_transfer_calls: list[object] = []
kv_init_calls: list[object] = []
@contextmanager
def fake_use_mem_pool(tag, device):
pool_calls.append((tag, str(device)))
yield
def fake_get_or_create(socket_path, device, mode, *, tag, timeout_ms=None):
create_calls.append((socket_path, device, mode.value, tag, timeout_ms))
return object()
monkeypatch.setattr(worker_module, "gms_use_mem_pool", fake_use_mem_pool)
monkeypatch.setattr(
worker_module,
"get_or_create_gms_client_memory_manager",
fake_get_or_create,
)
monkeypatch.setattr(
worker_module,
"get_socket_path",
lambda device, tag: f"/tmp/{tag}-{device}.sock",
)
monkeypatch.setattr(
kv_transfer,
"ensure_kv_transfer_initialized",
lambda vllm_config, kv_cache_config: kv_transfer_calls.append(kv_cache_config),
)
worker = object.__new__(GMSWorker)
worker.local_rank = 3
worker.vllm_config = SimpleNamespace(
model_config=SimpleNamespace(enable_sleep_mode=True)
)
worker.model_runner = SimpleNamespace(
initialize_kv_cache=lambda kv_cache_config: kv_init_calls.append(
kv_cache_config
)
)
worker.initialize_from_config("kv-config")
assert create_calls == [("/tmp/kv_cache-3.sock", 3, "rw", "kv_cache", None)]
assert pool_calls == [("kv_cache", "cuda:3")]
assert kv_transfer_calls == ["kv-config"]
assert kv_init_calls == ["kv-config"]
def test_sleep_level_2_unmaps_weights_and_kv_cache(monkeypatch):
import gpu_memory_service.integrations.vllm.worker as worker_module
from gpu_memory_service.integrations.vllm.worker import GMSWorker
weights = _FakeManager()
kv_cache = _FakeManager()
monkeypatch.setattr(
worker_module,
"get_gms_client_memory_manager",
lambda tag: weights if tag == "weights" else kv_cache,
)
monkeypatch.setattr(
worker_module.torch.cuda,
"mem_get_info",
lambda: (2 << 30, 8 << 30),
)
worker = object.__new__(GMSWorker)
worker.sleep(level=2)
assert weights.calls == ["unmap_all_vas", "abort"]
assert kv_cache.calls == ["unmap_all_vas", "abort"]
def test_wake_up_remaps_weights_and_reallocates_kv_cache(monkeypatch):
import gpu_memory_service.integrations.vllm.worker as worker_module
from gpu_memory_service.integrations.vllm.worker import GMSWorker
weights = _FakeManager(is_unmapped=True)
kv_cache = _FakeManager(is_unmapped=True)
fp8_calls: list[str] = []
monkeypatch.setattr(
worker_module,
"get_gms_client_memory_manager",
lambda tag: weights if tag == "weights" else kv_cache,
)
worker = object.__new__(GMSWorker)
worker.local_rank = 0
worker.cache_config = SimpleNamespace(cache_dtype="fp8_e4m3")
worker.model_runner = SimpleNamespace(
kv_caches={"layer_0": True},
init_fp8_kv_scales=lambda: fp8_calls.append("fp8"),
)
worker.wake_up(["weights", "kv_cache"])
assert weights.calls == [
("connect", "ro"),
"remap_all_vas",
]
assert kv_cache.calls == [
("connect", "rw"),
("reallocate_all_handles", "kv_cache"),
"remap_all_vas",
]
assert fp8_calls == ["fp8"]
def test_maybe_get_memory_pool_context_routes_tags(monkeypatch):
import gpu_memory_service.integrations.vllm.worker as worker_module
from gpu_memory_service.integrations.vllm.worker import GMSWorker, Worker
kv_cache_context = object()
super_calls: list[str] = []
mem_pool_calls: list[tuple[str, str]] = []
def fake_use_mem_pool(tag, device):
mem_pool_calls.append((tag, str(device)))
return kv_cache_context
def fake_super_context(self, tag):
del self
super_calls.append(tag)
return f"super:{tag}"
monkeypatch.setattr(worker_module, "gms_use_mem_pool", fake_use_mem_pool)
monkeypatch.setattr(Worker, "_maybe_get_memory_pool_context", fake_super_context)
worker = object.__new__(GMSWorker)
worker.local_rank = 2
weights_context = worker._maybe_get_memory_pool_context("weights")
with weights_context:
pass
assert mem_pool_calls == []
assert super_calls == []
assert worker._maybe_get_memory_pool_context("kv_cache") is kv_cache_context
assert mem_pool_calls == [("kv_cache", "cuda:2")]
assert super_calls == []
assert worker._maybe_get_memory_pool_context("other") == "super:other"
assert super_calls == ["other"]
tests/gpu_memory_service/common/gms.py 0 → 100644
View file @ b78ec99a
This diff is collapsed.
tests/gpu_memory_service/common/runtime.py 0 → 100644
View file @ b78ec99a
This diff is collapsed.
tests/gpu_memory_service/flow_assertions.py 0 → 100644
View file @ b78ec99a
This diff is collapsed.
tests/gpu_memory_service/test_quiesce_resume.py 0 → 100644
View file @ b78ec99a
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
from __future__ import annotations
import logging
import pytest
from tests.gpu_memory_service.common.runtime import (
GMSProcessManager,
SGLangWithGMSProcess,
VLLMWithGMSProcess,
get_gpu_memory_used,
)
from tests.gpu_memory_service.flow_assertions import (
assert_completion_ok,
assert_kv_history,
assert_memory_restored_after_quiesce,
assert_weights_published_once,
quiesce_engine,
wait_for_resumed_layout,
)
from tests.utils.constants import FAULT_TOLERANCE_MODEL_NAME
pytestmark = [pytest.mark.nightly, pytest.mark.fault_tolerance]
# Event flow under test:
# 1. Weights are published once as a committed layout.
# 2. KV cache starts as a live RW layout build.
# 3. Quiesce keeps weights committed but aborts and clears the KV layout.
# 4. Resume reconnects weights as RO to the same committed layout.
# 5. Resume recreates KV cache in a fresh RW layout after the old one was cleared.
logger = logging.getLogger(__name__)
def _run_quiesce_resume_test(
request,
engine_cls,
) -> None:
with GMSProcessManager(request, engine_cls) as manager:
frontend_port = manager.frontend_port
weights_gms = manager.weights_gms
kv_cache_gms = manager.kv_cache_gms
engine = manager.start_engine("engine")
assert_completion_ok(
frontend_port,
"Hello",
failure_message="Initial inference failed",
success_message="Initial inference result",
)
# Before quiesce, weights must already be published and visible to RO
# readers while KV cache remains a live RW layout owned by the engine.
weights_before_quiesce, released_bytes, mem_after_quiesce = quiesce_engine(
weights_gms,
kv_cache_gms,
engine,
quiesce_label="Engine quiesce",
)
# Weights are immutable across quiesce/resume, so their event history should
# still be the original publish: connect once, commit once.
weights_events = weights_gms.get_event_history().events
assert_weights_published_once(weights_events)
# KV cache is different: quiesce must abort the old RW layout and clear
# its server-owned allocations before resume can start a new RW layout.
kv_events = kv_cache_gms.get_event_history().events
assert_kv_history(kv_events, cleared_layouts=1)
assert kv_events[-1].allocation_count > 0
resume_result = engine.resume()
assert resume_result["status"] == "ok"
mem_after_resume = get_gpu_memory_used()
assert_memory_restored_after_quiesce(
"Memory after resume",
mem_after_quiesce,
mem_after_resume,
released_bytes,
)
# Resume reconnects weights as RO to the same committed layout, but KV cache
# must come back as a fresh RW layout with new allocations.
wait_for_resumed_layout(
weights_gms,
kv_cache_gms,
weights_before_quiesce,
)
weights_events_after_resume = weights_gms.get_event_history().events
assert_weights_published_once(weights_events_after_resume)
# The resume history should therefore extend the old KV sequence with one
# new RW connect after the previous layout was fully cleared.
kv_events_after_resume = kv_cache_gms.get_event_history().events
assert_kv_history(
kv_events_after_resume,
cleared_layouts=1,
suffix=["rw_connected"],
)
assert kv_events_after_resume[2].allocation_count > 0
assert_completion_ok(
frontend_port,
"Goodbye",
failure_message="Post-resume inference failed",
success_message="Post-resume inference result",
)
logger.info("Memory freed: %.0f MB", released_bytes / (1 << 20))
@pytest.mark.e2e
@pytest.mark.gpu_1
@pytest.mark.model(FAULT_TOLERANCE_MODEL_NAME)
@pytest.mark.timeout(300)
@pytest.mark.vllm
def test_gms_basic_quiesce_resume_vllm(
request,
runtime_services_dynamic_ports,
predownload_models,
):
_run_quiesce_resume_test(request, VLLMWithGMSProcess)
@pytest.mark.e2e
@pytest.mark.gpu_1
@pytest.mark.model(FAULT_TOLERANCE_MODEL_NAME)
@pytest.mark.timeout(300)
@pytest.mark.sglang
def test_gms_basic_quiesce_resume_sglang(
request,
runtime_services_dynamic_ports,
predownload_models,
):
_run_quiesce_resume_test(request, SGLangWithGMSProcess)
tests/gpu_memory_service/test_shadow_failover.py 0 → 100644
View file @ b78ec99a
This diff is collapsed.
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment