diff --git a/python/tokenspeed/runtime/engine/weight_transfer/manager.py b/python/tokenspeed/runtime/engine/weight_transfer/manager.py index 4b2398fb78..5d1dadeba7 100644 --- a/python/tokenspeed/runtime/engine/weight_transfer/manager.py +++ b/python/tokenspeed/runtime/engine/weight_transfer/manager.py @@ -95,6 +95,14 @@ def is_paused(self) -> bool: """Whether generation admission is currently paused.""" return self._async_llm.weight_transfer_admission_paused() + async def is_sleeping(self) -> bool: + """Whether the engine is in the sleep/wake-up suspended state. + + Mirrors :meth:`is_paused`; delegates to ``AsyncLLM.is_sleeping`` so the + sleep/wake-up state is queryable through the same RL-control surface. + """ + return await self._async_llm.is_sleeping() + def get_world_size(self, include_dp: bool = True) -> int: """Return the inference world size used to size the NCCL group. diff --git a/python/tokenspeed/runtime/entrypoints/control_server.py b/python/tokenspeed/runtime/entrypoints/control_server.py index 0fd8a246df..2ce93cdb7c 100644 --- a/python/tokenspeed/runtime/entrypoints/control_server.py +++ b/python/tokenspeed/runtime/entrypoints/control_server.py @@ -436,6 +436,11 @@ async def is_paused(request: Request): return await _proxy_to_rl_control(request) +@app.get("/is_sleeping") +async def is_sleeping(request: Request): + return await _proxy_to_rl_control(request) + + # --------------------------------------------------------------------------- # RL weight transfer — SGLang dialect, proxied to the same in-engine control app # diff --git a/python/tokenspeed/runtime/entrypoints/vllm_compat_http.py b/python/tokenspeed/runtime/entrypoints/vllm_compat_http.py index e62fc4336c..7016bdc0d2 100644 --- a/python/tokenspeed/runtime/entrypoints/vllm_compat_http.py +++ b/python/tokenspeed/runtime/entrypoints/vllm_compat_http.py @@ -230,6 +230,21 @@ async def is_paused(raw_request: Request) -> JSONResponse: return JSONResponse(content={"is_paused": paused}) +@router.get("/is_sleeping") +async def is_sleeping(raw_request: Request) -> JSONResponse: + try: + sleeping = await _manager(raw_request).is_sleeping() + except HTTPException: + raise + except Exception as e: # noqa: BLE001 - defensive + logger.exception("Failed to fetch sleep status") + return JSONResponse( + {"error": f"Failed to fetch sleep status: {e}"}, + status_code=HTTPStatus.INTERNAL_SERVER_ERROR.value, + ) + return JSONResponse(content={"is_sleeping": sleeping}) + + @router.get("/get_world_size") async def get_world_size( raw_request: Request, diff --git a/test/runtime/test_control_server.py b/test/runtime/test_control_server.py index 6404f5058c..d139bbd2c9 100644 --- a/test/runtime/test_control_server.py +++ b/test/runtime/test_control_server.py @@ -73,6 +73,25 @@ async def start_profile(): return mock +def _build_mock_rl_control() -> FastAPI: + """Mock of the in-engine RL-control app (``vllm_compat_http``). + + Only the status routes the sidecar proxies are needed; they return canned + booleans so the proxy-parity tests can assert byte-faithful relay. + """ + mock = FastAPI() + + @mock.get("/is_paused") + async def is_paused(): + return JSONResponse({"is_paused": False}) + + @mock.get("/is_sleeping") + async def is_sleeping(): + return JSONResponse({"is_sleeping": True}) + + return mock + + def _wait(port, path, timeout=15): deadline = time.time() + timeout while time.time() < deadline: @@ -222,6 +241,73 @@ def test_status_code_relayed(self): self.assertEqual(r.status_code, 404) +class TestRlControlProxy(unittest.TestCase): + """Proxy parity for the RL-control status routes (``/is_paused``, + ``/is_sleeping``). + + The sidecar mounts these as thin proxies to the in-engine RL-control app + (``vllm_compat_http``); they must relay the upstream ``{is_: bool}`` + body byte-faithfully, the same contract as the smg passthrough routes above. + """ + + RL_PORT = 28340 + SIDECAR_PORT = 28341 + + @classmethod + def setUpClass(cls): + from tokenspeed.runtime.entrypoints import control_server as hs + + cls.hs = hs + # Point the RL-control proxy at our mock; save/restore to keep the + # module global clean for the other test classes in this module. + cls._orig_rl_control_url = hs._rl_control_url + hs._rl_control_url = f"http://127.0.0.1:{cls.RL_PORT}" + hs._gateway_url = "http://127.0.0.1:1" # dead — no smg route hit here + hs._engine_grpc_addr = "127.0.0.1:1" + + cls._rl_server = uvicorn.Server( + uvicorn.Config( + _build_mock_rl_control(), + host="127.0.0.1", + port=cls.RL_PORT, + log_level="error", + ) + ) + cls._sidecar_server = uvicorn.Server( + uvicorn.Config( + hs.app, host="127.0.0.1", port=cls.SIDECAR_PORT, log_level="error" + ) + ) + cls._t_rl = threading.Thread(target=cls._rl_server.run, daemon=True) + cls._t_side = threading.Thread(target=cls._sidecar_server.run, daemon=True) + cls._t_rl.start() + cls._t_side.start() + assert _wait(cls.RL_PORT, "/is_paused"), "mock RL control failed to start" + assert _wait(cls.SIDECAR_PORT, "/is_paused"), "sidecar failed to start" + + @classmethod + def tearDownClass(cls): + cls._rl_server.should_exit = True + cls._sidecar_server.should_exit = True + cls.hs._rl_control_url = cls._orig_rl_control_url + + def _url(self, path): + return f"http://127.0.0.1:{self.SIDECAR_PORT}{path}" + + def test_is_paused_proxies_faithfully(self): + r = requests.get(self._url("/is_paused"), timeout=10) + self.assertEqual(r.status_code, 200) + self.assertEqual(r.json(), {"is_paused": False}) + + def test_is_sleeping_proxies_faithfully(self): + """``/is_sleeping`` must be mounted on the sidecar and relay the + upstream ``{is_sleeping: bool}`` body unchanged — parity with + ``/is_paused``.""" + r = requests.get(self._url("/is_sleeping"), timeout=10) + self.assertEqual(r.status_code, 200) + self.assertEqual(r.json(), {"is_sleeping": True}) + + class TestGrpcDirect(unittest.TestCase): """Unit tests for the gRPC-direct path (no live engine)."""