chore: import upstream snapshot with attribution

python/sglang/test/scripted_runtime/tokenizer_recv_proxy.py

@@ -0,0 +1,78 @@
+from __future__ import annotations
+
+import time
+from collections import deque
+from typing import Any, Callable
+
+import zmq
+
+from sglang.srt.managers.io_struct import (
+    BatchTokenizedEmbeddingReqInput,
+    BatchTokenizedGenerateReqInput,
+    TokenizedEmbeddingReqInput,
+    TokenizedGenerateReqInput,
+    sock_recv,
+)
+
+_WORK_REQ_TYPES = (
+    TokenizedGenerateReqInput,
+    TokenizedEmbeddingReqInput,
+    BatchTokenizedGenerateReqInput,
+    BatchTokenizedEmbeddingReqInput,
+)
+
+
+class ScriptedTokenizerRecvProxy:
+
+    def __init__(self, *, underlying: zmq.Socket) -> None:
+        self._underlying = underlying
+        self._buffer: deque = deque()
+        self.work_reqs_seen: int = 0
+
+    def recv_pyobj(self, flags: int = 0) -> Any:
+        self._drain_underlying()
+
+        if self._buffer:
+            return self._buffer.popleft()
+
+        if flags & zmq.NOBLOCK:
+            raise zmq.ZMQError(zmq.EAGAIN, "Resource temporarily unavailable")
+        raise RuntimeError(
+            "ScriptedTokenizerRecvProxy.recv_pyobj: blocking recv is not supported"
+        )
+
+    def recv(self, flags: int = 0) -> bytes:
+        raise NotImplementedError(
+            "TODO: support ScriptedTokenizerRecvProxy.recv for msgpack IPC"
+        )
+
+    def wait_until_arrived(
+        self,
+        predicate: Callable[[Any], bool],
+        *,
+        timeout_s: float,
+        description: str = "matching object",
+    ) -> None:
+        start_len = len(self._buffer)
+        deadline = time.monotonic() + timeout_s
+        while True:
+            self._drain_underlying()
+            for i, obj in enumerate(self._buffer):
+                if i >= start_len and predicate(obj):
+                    return
+            if time.monotonic() >= deadline:
+                raise TimeoutError(
+                    f"ScriptedTokenizerRecvProxy: no {description} arrived on the "
+                    f"recv_from_tokenizer socket within {timeout_s}s"
+                )
+            time.sleep(0.005)
+
+    def _drain_underlying(self) -> None:
+        while True:
+            try:
+                req = sock_recv(self._underlying, zmq.NOBLOCK)
+            except zmq.ZMQError:
+                break
+            if isinstance(req, _WORK_REQ_TYPES):
+                self.work_reqs_seen += 1
+            self._buffer.append(req)