From 1baf08c565a48b636e391d2d6950dda83b5780a5 Mon Sep 17 00:00:00 2001 From: mkhadkevich Date: Tue, 22 Sep 2026 06:38:48 +0000 Subject: [PATCH 1/3] test(sglang): align failover coverage with native pause routes Signed-off-by: mkhadkevich --- .../dynamo/sglang/tests/test_sglang_failover.py | 4 +++- .../test_sglang_memory_occupation_handlers.py | 17 ----------------- 2 files changed, 3 insertions(+), 18 deletions(-) diff --git a/components/src/dynamo/sglang/tests/test_sglang_failover.py b/components/src/dynamo/sglang/tests/test_sglang_failover.py index 30957eb227e2..d8a175778dd0 100644 --- a/components/src/dynamo/sglang/tests/test_sglang_failover.py +++ b/components/src/dynamo/sglang/tests/test_sglang_failover.py @@ -101,7 +101,9 @@ def unexpected_engine(**kwargs): monkeypatch.setattr(snapshot.sgl, "Engine", unexpected_engine) with pytest.raises(RuntimeError, match="writer-cohort fence"): - await snapshot.prepare_snapshot_engine(SimpleNamespace()) + await snapshot.prepare_snapshot_engine( + SimpleNamespace(server_args=SimpleNamespace()) + ) class _FakeEndpoint: diff --git a/components/src/dynamo/sglang/tests/test_sglang_memory_occupation_handlers.py b/components/src/dynamo/sglang/tests/test_sglang_memory_occupation_handlers.py index e524fa91557b..b6cb0e5103f0 100644 --- a/components/src/dynamo/sglang/tests/test_sglang_memory_occupation_handlers.py +++ b/components/src/dynamo/sglang/tests/test_sglang_memory_occupation_handlers.py @@ -238,20 +238,3 @@ async def test_clear_kv_blocks_reports_flush_exception(handler): chunks = [chunk async for chunk in handler.clear_kv_blocks({})] assert chunks == [{"status": "error", "message": "flush crashed"}] - - -@pytest.mark.parametrize("flag", ["handoff", "release_failover_lock"]) -@pytest.mark.asyncio -async def test_cooperative_handoff_fails_closed_without_mutating_worker(handler, flag): - lock = SimpleNamespace(release=AsyncMock()) - handler._gms_failover_lock = lock - - result = await handler.release_memory_occupation({flag: True}) - - assert result["status"] == "error" - assert "writer-cohort fencing" in result["message"] - lock.release.assert_not_awaited() - assert handler._gms_failover_lock is lock - handler.generate_endpoint.unregister_endpoint_instance.assert_not_awaited() - handler.engine.tokenizer_manager.pause_generation.assert_not_awaited() - handler.engine.tokenizer_manager.release_memory_occupation.assert_not_awaited() From 5355bb123ff445dc3107775c2e0fa1e53818c48e Mon Sep 17 00:00:00 2001 From: mkhadkevich Date: Tue, 22 Sep 2026 06:41:24 +0000 Subject: [PATCH 2/3] fix(gms): resolve default CUDA ordinal without recursion Signed-off-by: mkhadkevich --- .../tests/test_cuda_visible_devices.py | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/lib/gpu_memory_service/tests/test_cuda_visible_devices.py b/lib/gpu_memory_service/tests/test_cuda_visible_devices.py index 5ce442b03865..cb892de44a38 100644 --- a/lib/gpu_memory_service/tests/test_cuda_visible_devices.py +++ b/lib/gpu_memory_service/tests/test_cuda_visible_devices.py @@ -47,6 +47,21 @@ def test_socket_path_maps_cuda_visible_ordinal_to_nvml_device(monkeypatch, tmp_p assert calls == [("index", 3), ("index", 1)] +def test_socket_path_uses_device_ordinal_without_visibility_override( + monkeypatch, tmp_path +): + from gpu_memory_service.common import utils + + pynvml, calls = _fake_nvml() + monkeypatch.setitem(sys.modules, "pynvml", pynvml) + monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising=False) + monkeypatch.setenv("GMS_SOCKET_DIR", str(tmp_path)) + utils.invalidate_uuid_cache() + + assert utils.get_socket_path(2).endswith("gms_GPU-index:2_weights.sock") + assert calls == [("index", 2)] + + def test_socket_uuid_cache_tracks_visibility_mapping(monkeypatch, tmp_path): from gpu_memory_service.common import utils From eb02b9237c41b9451ee9d299b4fe1ec7eca83993 Mon Sep 17 00:00:00 2001 From: mkhadkevich Date: Tue, 22 Sep 2026 07:12:22 +0000 Subject: [PATCH 3/3] fix(vllm): match persistent KV context contract Signed-off-by: mkhadkevich --- .../integrations/vllm/worker.py | 1 - .../test_vllm_gms_worker_memory_accounting.py | 27 +++++++++++++++++++ 2 files changed, 27 insertions(+), 1 deletion(-) diff --git a/lib/gpu_memory_service/integrations/vllm/worker.py b/lib/gpu_memory_service/integrations/vllm/worker.py index 2e76bf7559a6..3cefb07e6fe4 100644 --- a/lib/gpu_memory_service/integrations/vllm/worker.py +++ b/lib/gpu_memory_service/integrations/vllm/worker.py @@ -527,6 +527,5 @@ def _maybe_get_memory_pool_context(self, tag: str): self._gms_kv_cache_config, self.vllm_config.model_config, torch.device(get_vmm_device_type().value, self._gms_device), - self.vllm_config.model_config, ) return super()._maybe_get_memory_pool_context(tag) diff --git a/lib/gpu_memory_service/tests/test_vllm_gms_worker_memory_accounting.py b/lib/gpu_memory_service/tests/test_vllm_gms_worker_memory_accounting.py index 22de822176a2..661995800c93 100644 --- a/lib/gpu_memory_service/tests/test_vllm_gms_worker_memory_accounting.py +++ b/lib/gpu_memory_service/tests/test_vllm_gms_worker_memory_accounting.py @@ -277,6 +277,33 @@ def test_vllm_kv_disable_flag_uses_base_worker_paths(monkeypatch): assert not hasattr(instance, "_gms_kv_manager") +def test_vllm_kv_context_passes_current_five_argument_contract(monkeypatch): + from gpu_memory_service.integrations.vllm import worker as worker_module + + captured = [] + sentinel = object() + manager = object() + kv_config = object() + model_config = object() + monkeypatch.setenv("GMS_VLLM_VMM_IPC_KV", "1") + monkeypatch.setattr( + worker_module, + "persistent_kv_allocation_context", + lambda *args: captured.append(args) or sentinel, + ) + instance = object.__new__(worker_module.GMSWorker) + instance._gms_kv_manager = manager + instance._gms_kv_engine_id = "engine" + instance._gms_kv_cache_config = kv_config + instance._gms_device = 0 + instance.vllm_config = SimpleNamespace(model_config=model_config) + + assert instance._maybe_get_memory_pool_context("kv_cache") is sentinel + assert len(captured) == 1 + assert len(captured[0]) == 5 + assert captured[0][:4] == (manager, "engine", kv_config, model_config) + + def test_vllm_kv_disable_flag_skips_sleep_wake_kv_lifecycle(monkeypatch): from gpu_memory_service.integrations.vllm import worker as worker_module