fix(sglang): handle health signals from the tokenizer socket

grpc_servicer/smg_grpc_servicer/sglang/request_manager.py

@@ -597,18 +597,23 @@ def _dispatch_communicator_output(self, recv_obj) -> bool:
         return True
 
     async def _handle_tokenizer_loop(self):
-        """Process communicator responses from the scheduler's send_to_tokenizer path.
+        """Process outputs from the scheduler's send_to_tokenizer path.
 
         The scheduler sends communicator responses (ProfileReqOutput,
-        FlushCacheReqOutput, ...) directly to tokenizer_ipc_name, not through
-        the detokenizer. This loop only runs when recv_from_tokenizer is a
-        dedicated socket — in skip_tokenizer_init mode those responses arrive
-        on recv_from_scheduler and are dispatched by handle_loop instead.
+        FlushCacheReqOutput, ...) and busy-path health signals directly to
+        tokenizer_ipc_name, not through the detokenizer. This loop only runs
+        when recv_from_tokenizer is a dedicated socket — in skip_tokenizer_init
+        mode those responses arrive on recv_from_scheduler and are dispatched
+        by handle_loop instead.
         """
         while not self.gracefully_exit:
             try:
                 recv_obj = await self.recv_from_tokenizer.recv_pyobj()
-                if not self._dispatch_communicator_output(recv_obj):
+                if isinstance(recv_obj, HealthCheckOutput):
+                    self.last_receive_tstamp = real_time()
+                    if getattr(recv_obj, "rid", None):
+                        await self._handle_health_check_output(recv_obj)
+                elif not self._dispatch_communicator_output(recv_obj):
                     logger.warning(f"Unknown type on tokenizer socket: {type(recv_obj)}")
             except zmq.error.Again:
                 if self.gracefully_exit:

grpc_servicer/tests/test_sglang_tokenizer_health.py

@@ -0,0 +1,75 @@
+"""Engine-free coverage for SGLang tokenizer-socket health signals."""
+
+import ast
+import asyncio
+import logging
+from dataclasses import dataclass
+from pathlib import Path
+from types import SimpleNamespace
+from unittest.mock import AsyncMock, Mock
+
+import pytest
+
+_REQUEST_MANAGER = Path(__file__).parents[1] / "smg_grpc_servicer" / "sglang" / "request_manager.py"
+
+
+@dataclass
+class HealthCheckOutput:
+    rid: str | None = None
+
+
+def _tokenizer_loop():
+    tree = ast.parse(_REQUEST_MANAGER.read_text())
+    method = next(
+        node
+        for node in ast.walk(tree)
+        if isinstance(node, ast.AsyncFunctionDef) and node.name == "_handle_tokenizer_loop"
+    )
+    namespace = {
+        "HealthCheckOutput": HealthCheckOutput,
+        "get_exception_traceback": lambda: "",
+        "logger": logging.getLogger(__name__),
+        "real_time": Mock(return_value=42.0),
+        "zmq": SimpleNamespace(
+            error=SimpleNamespace(
+                Again=type("Again", (Exception,), {}),
+                ZMQError=type("ZMQError", (Exception,), {}),
+            )
+        ),
+    }
+    exec(
+        compile(ast.Module(body=[method], type_ignores=[]), str(_REQUEST_MANAGER), "exec"),
+        namespace,
+    )
+    return namespace["_handle_tokenizer_loop"], namespace["real_time"]
+
+
+@pytest.mark.parametrize(
+    ("rid", "handled"),
+    [(None, False), ("", False), ("HEALTH_CHECK_test", True)],
+)
+def test_tokenizer_health_signal_refreshes_liveness_without_spurious_request(rid, handled):
+    loop, real_time = _tokenizer_loop()
+    health_output = HealthCheckOutput(rid=rid)
+
+    async def receive():
+        manager.gracefully_exit = True
+        return health_output
+
+    manager = SimpleNamespace(
+        gracefully_exit=False,
+        last_receive_tstamp=0.0,
+        recv_from_tokenizer=SimpleNamespace(recv_pyobj=receive),
+        _dispatch_communicator_output=Mock(return_value=False),
+        _handle_health_check_output=AsyncMock(),
+    )
+
+    asyncio.run(loop(manager))
+
+    assert manager.last_receive_tstamp == 42.0
+    real_time.assert_called_once_with()
+    if handled:
+        manager._handle_health_check_output.assert_awaited_once_with(health_output)
+    else:
+        manager._handle_health_check_output.assert_not_awaited()
+    manager._dispatch_communicator_output.assert_not_called()