79 lines
2.3 KiB
Python
79 lines
2.3 KiB
Python
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)
|