From 7f68ce34c4e1adf5fa771a78264e23acecd51903 Mon Sep 17 00:00:00 2001 From: 0z5a <0z5a@users.noreply.github.com> Date: Sun, 6 Sep 2026 20:51:47 +0000 Subject: [PATCH] fix: clear SGLang sleep state after empty-tag wakeup Keep Ray state consistent with acknowledged SGLang all-region wakeups. Cover CPU acknowledgement/failure contracts and real single-GPU sleep cycles with memory release, restoration, and deterministic output checks. The fixture configures TorchMemorySaver for Ray actors. AI-assisted implementation; human review and human-run tests remain required before requesting review. Signed-off-by: 0z5a <0z5a@users.noreply.github.com> Signed-off-by: 0z5a --- .../serve/engines/sglang/sglang_engine.py | 7 +- .../llm_tests/serve/test_llm_serve_sglang.py | 168 +++++++++++++++++- 2 files changed, 167 insertions(+), 8 deletions(-) diff --git a/python/ray/llm/_internal/serve/engines/sglang/sglang_engine.py b/python/ray/llm/_internal/serve/engines/sglang/sglang_engine.py index 2a7156ee2557..0e9226684148 100644 --- a/python/ray/llm/_internal/serve/engines/sglang/sglang_engine.py +++ b/python/ray/llm/_internal/serve/engines/sglang/sglang_engine.py @@ -69,7 +69,7 @@ class SGLangSleepConfig(BaseModel): - "kv_cache": Discard KV cache - "weights": Offload to CPU RAM - "cuda_graph": Discard CUDA graph - - None: Discard/Offload everything + - None or []: Discard/Offload everything """ @@ -81,7 +81,7 @@ class SGLangWakeupConfig(BaseModel): - "kv_cache": Restore KV cache only - "weights": Restore weights only - "cuda_graph": Restore CUDA graph only - - None: Restore everything + - None or []: Restore everything """ @@ -834,7 +834,8 @@ async def wakeup(self, **kwargs: Any) -> None: obj = ResumeMemoryOccupationReqInput(tags=config.tags) await self.engine.tokenizer_manager.resume_memory_occupation(obj, None) - if config.tags is None: + # SGLang treats both omitted tags and an empty list as all components. + if not config.tags: self._sleeping_tags.clear() else: self._sleeping_tags -= set(config.tags) diff --git a/release/llm_tests/serve/test_llm_serve_sglang.py b/release/llm_tests/serve/test_llm_serve_sglang.py index 3a52651e5d98..6a35329ccda3 100644 --- a/release/llm_tests/serve/test_llm_serve_sglang.py +++ b/release/llm_tests/serve/test_llm_serve_sglang.py @@ -1,6 +1,10 @@ +import asyncio import concurrent.futures +import os import re import sys +from types import SimpleNamespace +from unittest.mock import AsyncMock import httpx import pytest @@ -44,6 +48,12 @@ def _shutdown_and_wait_for_gpu_clear(baseline_mb: float) -> None: @pytest.fixture(scope="module") def sglang_client(): """Start an SGLang server once for all tests in this module.""" + from torch_memory_saver import configure_subprocess + + # Ray launches scheduler actors outside the engine's subprocess context. + with configure_subprocess(): + memory_saver_env = {"LD_PRELOAD": os.environ["LD_PRELOAD"]} + llm_config = LLMConfig( model_loading_config={ "model_id": RAY_MODEL_ID, @@ -56,10 +66,13 @@ def sglang_client(): } }, server_cls=SGLangServer, + runtime_env={"env_vars": memory_saver_env}, engine_kwargs={ "model_path": MODEL_ID, "tp_size": 1, "mem_fraction_static": 0.8, + "enable_memory_saver": True, + "enable_weights_cpu_backup": True, }, ) @@ -256,7 +269,11 @@ async def test_sglang_pause_resume_modes(sglang_client): @pytest.mark.asyncio -async def test_sglang_sleep_wakeup(sglang_client): +@pytest.mark.parametrize("sleep_kwargs", [{}, {"tags": None}, {"tags": []}]) +@pytest.mark.parametrize("wakeup_kwargs", [{}, {"tags": None}, {"tags": []}]) +async def test_sglang_sleep_wakeup( + sglang_client, sleep_kwargs, wakeup_kwargs, record_property +): """Verify sleep/wakeup lifecycle: GPU memory released then restored.""" handle = _get_llm_handle() @@ -267,18 +284,33 @@ async def test_sglang_sleep_wakeup(sglang_client): assert resp.choices[0].text is not None # Sleep (release all GPU memory). - await handle.sleep.remote() + awake_memory_mb = get_total_gpu_memory_mb() + record_property("awake_memory_mb", awake_memory_mb) + await handle.sleep.remote(**sleep_kwargs) assert await handle.is_sleeping.remote() is True + wait_for_condition( + lambda: get_total_gpu_memory_mb() < awake_memory_mb - 256, + timeout=30, + ) + sleeping_memory_mb = get_total_gpu_memory_mb() + record_property("sleeping_memory_mb", sleeping_memory_mb) # Wakeup and confirm state clears. - await handle.wakeup.remote() + await handle.wakeup.remote(**wakeup_kwargs) assert await handle.is_sleeping.remote() is False + wait_for_condition( + lambda: get_total_gpu_memory_mb() > sleeping_memory_mb + 256, + timeout=30, + ) + record_property("restored_memory_mb", get_total_gpu_memory_mb()) # Inference must work again after wakeup. - resp = sglang_client.completions.create( + resumed_resp = sglang_client.completions.create( model=RAY_MODEL_ID, prompt="Hello", max_tokens=4, temperature=0.0 ) - assert resp.choices[0].text is not None + record_property("before_sleep_text", resp.choices[0].text) + record_property("after_wakeup_text", resumed_resp.choices[0].text) + assert resumed_resp.choices[0].text == resp.choices[0].text @pytest.mark.asyncio @@ -682,5 +714,131 @@ def test_score_request_is_sglang_scoring_request(self): assert issubclass(ScoreRequest, ScoringRequest) +class TestSGLangSleepState: + """Ray-side state contracts with real SGLang request types, without a GPU. + + Keep these in the SGLang release suite, whose image installs SGLang. The + tokenizer manager is mocked; actual offload/restore is exercised above. + """ + + @pytest.fixture + def server(self): + server = SGLangServer.__new__(SGLangServer) + server._sleeping_tags = set() + server.engine = SimpleNamespace( + tokenizer_manager=SimpleNamespace( + release_memory_occupation=AsyncMock(), + resume_memory_occupation=AsyncMock(), + ) + ) + return server + + @pytest.mark.asyncio + @pytest.mark.parametrize("sleep_kwargs", [{}, {"tags": None}, {"tags": []}]) + @pytest.mark.parametrize("wakeup_kwargs", [{}, {"tags": None}, {"tags": []}]) + async def test_all_tags(self, server, sleep_kwargs, wakeup_kwargs): + assert await server.is_sleeping() is False + await server.sleep(**sleep_kwargs) + assert server._sleeping_tags == {"weights", "kv_cache", "cuda_graph"} + assert await server.is_sleeping() is True + + manager = server.engine.tokenizer_manager + manager.release_memory_occupation.assert_awaited_once() + release_args = manager.release_memory_occupation.await_args.args + assert release_args[0].tags == sleep_kwargs.get("tags") + assert release_args[1] is None + + await server.wakeup(**wakeup_kwargs) + manager.resume_memory_occupation.assert_awaited_once() + resume_args = manager.resume_memory_occupation.await_args.args + assert resume_args[0].tags == wakeup_kwargs.get("tags") + assert resume_args[1] is None + assert await server.is_sleeping() is False + assert server._sleeping_tags == set() + + @pytest.mark.asyncio + async def test_selective_wakeup_retains_other_tags(self, server): + await server.sleep(tags=["weights", "kv_cache"]) + await server.sleep(tags=["cuda_graph"]) + assert server._sleeping_tags == {"weights", "kv_cache", "cuda_graph"} + + for tag, remaining in [ + ("weights", {"kv_cache", "cuda_graph"}), + ("cuda_graph", {"kv_cache"}), + ("kv_cache", set()), + ]: + await server.wakeup(tags=[tag]) + assert server._sleeping_tags == remaining + assert await server.is_sleeping() is bool(remaining) + ( + request, + context, + ) = server.engine.tokenizer_manager.resume_memory_occupation.await_args.args + assert request.tags == [tag] + assert context is None + + @pytest.mark.asyncio + @pytest.mark.parametrize("operation", ["sleep", "wakeup"]) + @pytest.mark.parametrize("tags", [None, [], ["weights"]]) + @pytest.mark.parametrize("error_type", [RuntimeError, asyncio.CancelledError]) + async def test_failed_operation_preserves_state( + self, server, operation, tags, error_type + ): + initial = {"kv_cache"} if operation == "sleep" else {"weights", "kv_cache"} + server._sleeping_tags = initial.copy() + manager = server.engine.tokenizer_manager + rpc = ( + manager.release_memory_occupation + if operation == "sleep" + else manager.resume_memory_occupation + ) + rpc.side_effect = error_type("backend did not acknowledge") + + with pytest.raises(error_type, match="backend did not acknowledge"): + await getattr(server, operation)(tags=tags) + + rpc.assert_awaited_once() + assert server._sleeping_tags == initial + assert await server.is_sleeping() is True + + @pytest.mark.asyncio + @pytest.mark.parametrize("operation", ["sleep", "wakeup"]) + async def test_state_changes_only_after_acknowledgement(self, server, operation): + initial = ( + set() if operation == "sleep" else {"weights", "kv_cache", "cuda_graph"} + ) + expected = ( + {"weights", "kv_cache", "cuda_graph"} if operation == "sleep" else set() + ) + server._sleeping_tags = initial.copy() + entered = asyncio.Event() + acknowledged = asyncio.Event() + + async def control_rpc(*args): + entered.set() + await acknowledged.wait() + + manager = server.engine.tokenizer_manager + rpc = ( + manager.release_memory_occupation + if operation == "sleep" + else manager.resume_memory_occupation + ) + rpc.side_effect = control_rpc + task = asyncio.create_task(getattr(server, operation)(tags=[])) + try: + await asyncio.wait_for(entered.wait(), timeout=5) + assert server._sleeping_tags == initial + assert await server.is_sleeping() is bool(initial) + assert not task.done() + acknowledged.set() + await asyncio.wait_for(task, timeout=5) + assert server._sleeping_tags == expected + assert await server.is_sleeping() is bool(expected) + finally: + task.cancel() + await asyncio.gather(task, return_exceptions=True) + + if __name__ == "__main__": sys.exit(pytest.main(["-xvs", __file__]))