fix(gms): align failover tests and runtime contracts

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:

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

lib/gpu_memory_service/integrations/vllm/worker.py

@@ -481,6 +481,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)

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
 

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