sgl-project--sglang/python/sglang/srt/entrypoints/grpc_bridge.py

"""Python-side bridge between the Rust gRPC server and TokenizerManager.

The RuntimeHandle exposes synchronous methods that Rust can call via PyO3
(with a brief GIL acquisition). Response chunks are pushed into Rust-side
channels via callback objects while all async work stays on the
TokenizerManager's event loop.
"""

import asyncio
import dataclasses
import json
import logging
from types import SimpleNamespace
from typing import Any, Awaitable, Callable, Dict, List, Optional

from pydantic import ValidationError

from sglang.srt.utils.msgspec_utils import msgspec_to_builtins

logger = logging.getLogger(__name__)


class _BadOpenAIRequest(ValueError):
    pass


class _CaseInsensitiveHeaders:
    __slots__ = ("_data",)

    def __init__(self, headers: Optional[Dict[str, str]] = None):
        self._data = {k.lower(): v for k, v in (headers or {}).items()}

    def get(self, name: str, default: Optional[str] = None) -> Optional[str]:
        return self._data.get(name.lower(), default)


class _GrpcRequest:
    """Small FastAPI Request shim used by OpenAIServing* and TokenizerManager."""

    def __init__(
        self,
        headers: Optional[Dict[str, str]] = None,
        is_disconnected_fn: Optional[Callable[[], bool]] = None,
    ):
        self.headers = _CaseInsensitiveHeaders(headers)
        self.state = SimpleNamespace()
        self._is_disconnected_fn = is_disconnected_fn

    async def is_disconnected(self) -> bool:
        if self._is_disconnected_fn is None:
            return False
        return bool(self._is_disconnected_fn())


class RuntimeHandle:
    """Thin Python handle that the Rust gRPC server calls into.

    Provides synchronous ``submit_*``, ``abort``, and info methods.
    Each submit method receives a ``chunk_callback`` (a Rust-side PyO3 object)
    that it invokes with ``(chunk_dict, finished, error)`` for each response
    chunk produced by TokenizerManager.
    """

    def __init__(
        self,
        tokenizer_manager,
        template_manager,
        server_args,
        scheduler_info: Optional[Dict] = None,
    ):
        self.tokenizer_manager = tokenizer_manager
        self.template_manager = template_manager
        self.server_args = server_args
        self.scheduler_info = scheduler_info or {}

        self._openai_serving_classes = None

        self.tokenizer_manager.auto_create_handle_loop()
        self._event_loop = self.tokenizer_manager.event_loop

    @property
    def _tm_loop(self):
        """Return the TokenizerManager loop used by communicator RPCs."""
        return self._event_loop

    def _safe_callback(self, chunk_callback, payload, **kwargs):
        """Invoke a Rust callback and return its ChunkSendStatus, if any."""
        try:
            return chunk_callback(payload, **kwargs)
        except Exception as e:
            logger.warning("gRPC chunk_callback failed: %s", e)
            return None

    def _send_native_error(self, chunk_callback, message: str):
        # ChunkCallback extracts the PyDict arg before reading error=.
        return self._safe_callback(chunk_callback, {}, finished=True, error=message)

    _BACKPRESSURE_TIMEOUT_S = 300.0

    @staticmethod
    def _is_pending_status(status) -> bool:
        return status is not None and status == type(status).Pending

    @staticmethod
    def _is_closed_status(status) -> bool:
        return status is not None and status == type(status).Closed

    def _abort_request_id(self, rid) -> None:
        if isinstance(rid, list):
            for single_rid in rid:
                self.tokenizer_manager.abort_request(rid=single_rid)
        else:
            self.tokenizer_manager.abort_request(rid=rid)

    async def _send_with_backpressure(
        self,
        chunk_callback,
        ready_event: Optional[asyncio.Event],
        payload,
        *,
        timeout_abort_rid=None,
        **kwargs,
    ) -> bool:
        status = self._safe_callback(chunk_callback, payload, **kwargs)
        if status is None or self._is_closed_status(status):
            return False
        if not self._is_pending_status(status):
            return True

        if kwargs.get("finished"):
            return True
        if ready_event is None:
            return True

        try:
            await asyncio.wait_for(
                ready_event.wait(), timeout=self._BACKPRESSURE_TIMEOUT_S
            )
        except asyncio.TimeoutError:
            if timeout_abort_rid is not None:
                self._abort_request_id(timeout_abort_rid)
                logger.warning(
                    "gRPC chunk backpressure wait timed out after %ss; aborted request",
                    self._BACKPRESSURE_TIMEOUT_S,
                )
            else:
                logger.warning(
                    "gRPC chunk backpressure wait timed out after %ss; closing stream",
                    self._BACKPRESSURE_TIMEOUT_S,
                )
            return False
        ready_event.clear()
        return True

    def _install_on_ready(self, chunk_callback) -> Optional[asyncio.Event]:
        set_on_ready = getattr(chunk_callback, "set_on_ready", None)
        if set_on_ready is None:
            return None
        ready_event = asyncio.Event()
        loop = self._tm_loop

        def _on_ready() -> None:
            loop.call_soon_threadsafe(ready_event.set)

        try:
            set_on_ready(_on_ready)
        except Exception as e:
            logger.warning("gRPC set_on_ready failed: %s", e)
            raise
        return ready_event

    @staticmethod
    def _uninstall_on_ready(chunk_callback) -> None:
        clear = getattr(chunk_callback, "clear_on_ready", None)
        if clear is None:
            return
        try:
            clear()
        except Exception as e:
            logger.warning("gRPC clear_on_ready failed: %s", e)

    def _submit_on_tm_loop(self, coro: Awaitable) -> None:
        future = asyncio.run_coroutine_threadsafe(coro, self._tm_loop)
        future.add_done_callback(self._log_unhandled_future_exception)

    @staticmethod
    def _log_unhandled_future_exception(future) -> None:
        try:
            future.result()
        except Exception as e:
            logger.error(
                "gRPC scheduled coroutine raised unhandled exception: %s",
                e,
                exc_info=True,
            )

    def _submit_json_unary(
        self,
        op_name: str,
        payload_coro_factory: Callable[[], Awaitable[Any]],
        chunk_callback,
        *,
        error_payload_fn: Optional[Callable[[Exception], Any]] = None,
    ) -> None:
        error_fn = error_payload_fn or (lambda e: {"error": {"message": str(e)}})

        async def _run() -> None:
            try:
                payload = await payload_coro_factory()
                self._safe_callback(
                    chunk_callback,
                    json.dumps(payload, default=str).encode("utf-8"),
                    finished=True,
                )
            except Exception as e:
                logger.error("gRPC %s error: %s", op_name, e)
                self._safe_callback(
                    chunk_callback,
                    json.dumps(error_fn(e), default=str).encode("utf-8"),
                    finished=True,
                    error=str(e),
                )

        self._submit_on_tm_loop(_run())

    def _get_openai_serving(self):
        """Lazily initialize OpenAI serving classes."""
        if self._openai_serving_classes is not None:
            return self._openai_serving_classes

        from sglang.srt.entrypoints.openai.serving_chat import OpenAIServingChat
        from sglang.srt.entrypoints.openai.serving_classify import (
            OpenAIServingClassify,
        )
        from sglang.srt.entrypoints.openai.serving_completions import (
            OpenAIServingCompletion,
        )
        from sglang.srt.entrypoints.openai.serving_embedding import (
            OpenAIServingEmbedding,
        )
        from sglang.srt.entrypoints.openai.serving_rerank import OpenAIServingRerank
        from sglang.srt.entrypoints.openai.serving_score import OpenAIServingScore

        self._openai_serving_classes = {
            "chat": OpenAIServingChat(self.tokenizer_manager, self.template_manager),
            "completion": OpenAIServingCompletion(
                self.tokenizer_manager, self.template_manager
            ),
            "embedding": OpenAIServingEmbedding(
                self.tokenizer_manager, self.template_manager
            ),
            "classify": OpenAIServingClassify(
                self.tokenizer_manager, self.template_manager
            ),
            "score": OpenAIServingScore(self.tokenizer_manager),
            "rerank": OpenAIServingRerank(
                self.tokenizer_manager, self.template_manager
            ),
        }
        return self._openai_serving_classes

    def submit_request(
        self,
        *,
        req_type: str,
        req_dict: dict,
        chunk_callback,
        is_disconnected_fn: Optional[Callable[[], bool]] = None,
    ):
        mock_request = (
            _GrpcRequest(is_disconnected_fn=is_disconnected_fn)
            if is_disconnected_fn is not None
            else None
        )
        if req_type == "generate":
            from sglang.srt.managers.io_struct import GenerateReqInput

            obj = GenerateReqInput(**req_dict)
            stream = req_dict.get("stream", False)
            self._submit_on_tm_loop(
                self._run_generate(obj, chunk_callback, stream, mock_request)
            )
        elif req_type == "embed":
            from sglang.srt.managers.io_struct import EmbeddingReqInput

            obj = EmbeddingReqInput(**req_dict)
            self._submit_on_tm_loop(self._run_embed(obj, chunk_callback, mock_request))
        else:
            raise ValueError(
                f"Unknown req_type: {req_type!r} (expected 'generate' or 'embed')"
            )

    async def _run_generate(self, obj, chunk_callback, stream: bool, request):
        ready_event = None
        try:
            ready_event = self._install_on_ready(chunk_callback) if stream else None
            gen = self.tokenizer_manager.generate_request(obj, request=request)
            if stream:
                async for chunk in gen:
                    finished = (
                        chunk.get("meta_info", {}).get("finish_reason") is not None
                    )
                    keep_going = await self._send_with_backpressure(
                        chunk_callback,
                        ready_event,
                        chunk,
                        finished=finished,
                        timeout_abort_rid=obj.rid,
                    )
                    if finished or not keep_going:
                        return
                # Defensive: generator exited without a finish_reason chunk.
                self._safe_callback(chunk_callback, {}, finished=True)
            else:
                result = await gen.__anext__()
                self._safe_callback(chunk_callback, result, finished=True)
        except StopAsyncIteration:
            self._safe_callback(chunk_callback, {}, finished=True)
        except Exception as e:
            logger.error("gRPC generate error for rid=%s: %s", obj.rid, e)
            self._send_native_error(chunk_callback, str(e))
        finally:
            if stream:
                self._uninstall_on_ready(chunk_callback)

    async def _run_embed(self, obj, chunk_callback, request):
        try:
            gen = self.tokenizer_manager.generate_request(obj, request=request)
            result = await gen.__anext__()
            self._safe_callback(chunk_callback, result, finished=True)
        except StopAsyncIteration:
            self._safe_callback(chunk_callback, {}, finished=True)
        except Exception as e:
            logger.error("gRPC embed error for rid=%s: %s", obj.rid, e)
            self._send_native_error(chunk_callback, str(e))

    # Bounded so a stuck TM loop can't deadlock the gRPC handler thread that
    # called abort. abort_request only enqueues a message on the ZMQ socket,
    # so a few seconds is generous; if we time out, log and drop — the client
    # will retry or give up.
    _ABORT_TIMEOUT_S = 5.0

    def abort(self, rid: str = "", abort_all: bool = False):
        """Abort a request by request ID or abort all active requests."""
        loop = self._tm_loop

        try:
            running_loop = asyncio.get_running_loop()
        except RuntimeError:
            running_loop = None

        if running_loop is loop:
            self.tokenizer_manager.abort_request(rid=rid, abort_all=abort_all)
            return

        future = asyncio.run_coroutine_threadsafe(
            self._abort_async(rid, abort_all),
            loop,
        )
        try:
            future.result(timeout=self._ABORT_TIMEOUT_S)
        except TimeoutError:
            future.cancel()
            logger.error(
                "gRPC abort timed out after %ss (rid=%r, abort_all=%s); "
                "tokenizer_manager loop appears stuck",
                self._ABORT_TIMEOUT_S,
                rid,
                abort_all,
            )

    async def _abort_async(self, rid: str, abort_all: bool) -> None:
        self.tokenizer_manager.abort_request(rid=rid, abort_all=abort_all)

    def get_model_info(self) -> str:
        model_config = self.tokenizer_manager.model_config
        result = {
            "model_path": self.tokenizer_manager.model_path,
            "tokenizer_path": self.server_args.tokenizer_path,
            "is_generation": self.tokenizer_manager.is_generation,
            "weight_version": self.server_args.weight_version,
            "model_type": getattr(model_config.hf_config, "model_type", None),
            "architectures": getattr(model_config.hf_config, "architectures", None),
        }
        return json.dumps(result, default=str)

    def get_server_info(self) -> str:
        result: Dict[str, Any] = dataclasses.asdict(self.server_args)
        result.update(self.scheduler_info)
        return json.dumps(msgspec_to_builtins(result), default=str)

    def health_check(self) -> bool:
        from sglang.srt.managers.tokenizer_manager import ServerStatus

        if self.tokenizer_manager.gracefully_exit:
            return False
        return self.tokenizer_manager.server_status not in (
            ServerStatus.Starting,
            ServerStatus.UnHealthy,
        )

    def tokenize(self, text: str, add_special_tokens: bool = True) -> str:
        tokenizer = self.tokenizer_manager.tokenizer
        tokens = tokenizer.encode(text, add_special_tokens=add_special_tokens)
        result = {
            "tokens": tokens,
            "count": len(tokens),
            "max_model_len": self.tokenizer_manager.model_config.context_len,
            "input_text": text,
        }
        return json.dumps(result)

    def detokenize(self, tokens: List[int]) -> str:
        tokenizer = self.tokenizer_manager.tokenizer
        text = tokenizer.decode(tokens)
        return json.dumps({"text": text})

    def list_models(self) -> str:
        served_model_name = self.tokenizer_manager.served_model_name
        models = [
            {
                "id": served_model_name,
                "root": served_model_name,
                "max_model_len": self.tokenizer_manager.model_config.context_len,
            }
        ]
        if self.server_args.enable_lora and hasattr(
            self.tokenizer_manager, "lora_registry"
        ):
            lora_registry = self.tokenizer_manager.lora_registry
            for _, lora_ref in lora_registry.get_all_adapters().items():
                models.append(
                    {
                        "id": lora_ref.lora_name,
                        "root": lora_ref.lora_path,
                        "parent": served_model_name,
                    }
                )
        return json.dumps(models)

    def get_load(self, chunk_callback, dp_rank: Optional[int] = None) -> None:
        async def _payload():
            result = await self.tokenizer_manager.get_loads(dp_rank=dp_rank)
            return [r.to_dict() for r in result]

        self._submit_json_unary("get_load", _payload, chunk_callback)

    def flush_cache(self, chunk_callback) -> None:
        async def _payload():
            ret = await self.tokenizer_manager.flush_cache()
            return {"success": ret.success, "message": "Cache flushed."}

        self._submit_json_unary(
            "flush_cache",
            _payload,
            chunk_callback,
            error_payload_fn=lambda e: {"success": False, "message": str(e)},
        )

    def pause_generation(self, mode: str, chunk_callback) -> None:
        async def _payload():
            from sglang.srt.managers.io_struct import PauseGenerationReqInput

            await self.tokenizer_manager.pause_generation(
                PauseGenerationReqInput(mode=mode)
            )
            return {"message": f"Generation paused (mode={mode})."}

        self._submit_json_unary("pause_generation", _payload, chunk_callback)

    def continue_generation(self, chunk_callback) -> None:
        async def _payload():
            from sglang.srt.managers.io_struct import ContinueGenerationReqInput

            await self.tokenizer_manager.continue_generation(
                ContinueGenerationReqInput()
            )
            return {"message": "Generation continued."}

        self._submit_json_unary("continue_generation", _payload, chunk_callback)

    def start_profile(self, output_dir: Optional[str], chunk_callback) -> None:
        async def _payload():
            from sglang.srt.managers.io_struct import ProfileReq

            req = ProfileReq(output_dir=output_dir) if output_dir else ProfileReq()
            await self.tokenizer_manager.start_profile(req)
            return {"message": "Profiling started."}

        self._submit_json_unary("start_profile", _payload, chunk_callback)

    def stop_profile(self, chunk_callback) -> None:
        async def _payload():
            await self.tokenizer_manager.stop_profile()
            return {"message": "Profiling stopped."}

        self._submit_json_unary("stop_profile", _payload, chunk_callback)

    def update_weights_from_disk(
        self, model_path: str, load_format: Optional[str], chunk_callback
    ) -> None:
        async def _payload():
            from sglang.srt.managers.io_struct import UpdateWeightFromDiskReqInput

            obj = UpdateWeightFromDiskReqInput(
                model_path=model_path, load_format=load_format
            )
            success, message, num_paused = (
                await self.tokenizer_manager.update_weights_from_disk(obj, request=None)
            )
            return {
                "success": success,
                "message": message,
                "num_paused_requests": num_paused,
            }

        self._submit_json_unary(
            "update_weights",
            _payload,
            chunk_callback,
            error_payload_fn=lambda e: {"success": False, "message": str(e)},
        )

    def _submit_openai(
        self,
        serving_key: str,
        streaming: bool,
        json_body: bytes,
        chunk_callback,
        trace_headers: Optional[Dict[str, str]],
        is_disconnected_fn: Optional[Callable[[], bool]],
    ) -> None:
        self._submit_on_tm_loop(
            self._run_openai_request(
                serving_key,
                json_body,
                chunk_callback,
                streaming=streaming,
                trace_headers=trace_headers,
                is_disconnected_fn=is_disconnected_fn,
            )
        )

    def submit_openai_chat(
        self,
        *,
        json_body: bytes,
        chunk_callback,
        trace_headers: Optional[Dict[str, str]] = None,
        is_disconnected_fn: Optional[Callable[[], bool]] = None,
    ) -> None:
        self._submit_openai(
            "chat", True, json_body, chunk_callback, trace_headers, is_disconnected_fn
        )

    def submit_openai_complete(
        self,
        *,
        json_body: bytes,
        chunk_callback,
        trace_headers: Optional[Dict[str, str]] = None,
        is_disconnected_fn: Optional[Callable[[], bool]] = None,
    ) -> None:
        self._submit_openai(
            "completion",
            True,
            json_body,
            chunk_callback,
            trace_headers,
            is_disconnected_fn,
        )

    def submit_openai_embed(
        self,
        *,
        json_body: bytes,
        chunk_callback,
        trace_headers: Optional[Dict[str, str]] = None,
        is_disconnected_fn: Optional[Callable[[], bool]] = None,
    ) -> None:
        self._submit_openai(
            "embedding",
            False,
            json_body,
            chunk_callback,
            trace_headers,
            is_disconnected_fn,
        )

    def submit_openai_classify(
        self,
        *,
        json_body: bytes,
        chunk_callback,
        trace_headers: Optional[Dict[str, str]] = None,
        is_disconnected_fn: Optional[Callable[[], bool]] = None,
    ) -> None:
        self._submit_openai(
            "classify",
            False,
            json_body,
            chunk_callback,
            trace_headers,
            is_disconnected_fn,
        )

    def submit_openai_score(
        self,
        *,
        json_body: bytes,
        chunk_callback,
        trace_headers: Optional[Dict[str, str]] = None,
        is_disconnected_fn: Optional[Callable[[], bool]] = None,
    ) -> None:
        self._submit_openai(
            "score", False, json_body, chunk_callback, trace_headers, is_disconnected_fn
        )

    def submit_openai_rerank(
        self,
        *,
        json_body: bytes,
        chunk_callback,
        trace_headers: Optional[Dict[str, str]] = None,
        is_disconnected_fn: Optional[Callable[[], bool]] = None,
    ) -> None:
        self._submit_openai(
            "rerank",
            False,
            json_body,
            chunk_callback,
            trace_headers,
            is_disconnected_fn,
        )

    def _get_openai_request_class(self, serving_key: str):
        """Return the Pydantic request class for a given serving key."""
        from sglang.srt.entrypoints.openai.protocol import (
            ChatCompletionRequest,
            ClassifyRequest,
            CompletionRequest,
            EmbeddingRequest,
            ScoringRequest,
            V1RerankReqInput,
        )

        return {
            "chat": ChatCompletionRequest,
            "completion": CompletionRequest,
            "embedding": EmbeddingRequest,
            "classify": ClassifyRequest,
            "score": ScoringRequest,
            "rerank": V1RerankReqInput,
        }[serving_key]

    async def _run_openai_request(
        self,
        serving_key: str,
        json_body: bytes,
        chunk_callback,
        streaming: bool,
        trace_headers: Optional[Dict[str, str]] = None,
        is_disconnected_fn: Optional[Callable[[], bool]] = None,
    ):
        try:
            serving = self._get_openai_serving()[serving_key]

            try:
                request_dict = json.loads(json_body)
                if not isinstance(request_dict, dict):
                    raise _BadOpenAIRequest(
                        f"Request body must be a JSON object, got {type(request_dict).__name__}"
                    )
                request_cls = self._get_openai_request_class(serving_key)
                request_obj = request_cls(**request_dict)
            except (json.JSONDecodeError, ValidationError, _BadOpenAIRequest) as e:
                error_body = json.dumps(
                    {"error": {"message": str(e), "type": "BadRequest"}}
                ).encode("utf-8")
                if streaming:
                    self._safe_callback(
                        chunk_callback, error_body, finished=True, error=str(e)
                    )
                else:
                    self._safe_callback(
                        chunk_callback, error_body, finished=True, status_code=400
                    )
                return

            mock_request = _GrpcRequest(
                headers=trace_headers,
                is_disconnected_fn=is_disconnected_fn,
            )

            result = await serving.handle_request(request_obj, mock_request)

            if hasattr(result, "body_iterator"):
                ready_event = self._install_on_ready(chunk_callback)
                data_buf: List[str] = []
                stream_closed = False

                async def _flush_event() -> bool:
                    """Flush buffered SSE data lines as one chunk. Returns False if Rust closed."""
                    if not data_buf:
                        return True
                    body = "\n".join(data_buf)
                    data_buf.clear()
                    if body == "[DONE]" or not body:
                        return True
                    return await self._send_with_backpressure(
                        chunk_callback,
                        ready_event,
                        body.encode("utf-8"),
                        finished=False,
                    )

                try:
                    async for raw_chunk in result.body_iterator:
                        if isinstance(raw_chunk, bytes):
                            raw_chunk = raw_chunk.decode("utf-8", errors="replace")
                        for line in raw_chunk.split("\n"):
                            line = line.rstrip("\r")
                            if not line:
                                if not await _flush_event():
                                    stream_closed = True
                                    break
                            elif line.startswith(":"):
                                continue  # SSE comment / heartbeat
                            elif line.startswith("data:"):
                                value = line[5:]
                                if value.startswith(" "):
                                    value = value[1:]
                                data_buf.append(value)
                            # event:, id:, retry:, unknown fields: ignored
                        if stream_closed:
                            break

                    if not stream_closed:
                        await _flush_event()
                        self._safe_callback(chunk_callback, b"", finished=True)
                finally:
                    self._uninstall_on_ready(chunk_callback)
            else:
                if hasattr(result, "model_dump"):
                    resp_bytes = json.dumps(result.model_dump()).encode("utf-8")
                elif hasattr(result, "body"):
                    resp_bytes = result.body
                elif isinstance(result, (dict, list)):
                    resp_bytes = json.dumps(result).encode("utf-8")
                else:
                    resp_bytes = str(result).encode("utf-8")
                status_code = int(
                    getattr(result, "status_code", None)
                    or getattr(result, "code", None)
                    or 200
                )
                self._safe_callback(
                    chunk_callback,
                    resp_bytes,
                    finished=True,
                    status_code=status_code,
                )

        except Exception as e:
            logger.error("gRPC OpenAI %s error: %s", serving_key, e)
            error_body = json.dumps({"error": {"message": str(e)}}).encode("utf-8")
            if streaming:
                self._safe_callback(
                    chunk_callback, error_body, finished=True, error=str(e)
                )
            else:
                self._safe_callback(
                    chunk_callback,
                    error_body,
                    finished=True,
                    status_code=int(getattr(e, "status_code", 500)),
                )