sglang 0.5.3rc0__py3-none-any.whl → 0.5.4__py3-none-any.whl

sglang/srt/entrypoints/grpc_request_manager.py

@@ -1,580 +0,0 @@
-"""
-gRPC Request Manager - Orchestrates request lifecycle without tokenization.
-Mimics TokenizerManager's state management and ZMQ communication patterns.
-"""
-
-import asyncio
-import dataclasses
-import logging
-import os
-import signal
-import sys
-import threading
-import time
-from typing import Any, Dict, List, Optional, Union
-
-import grpc
-import zmq
-import zmq.asyncio
-
-from sglang.srt.managers.io_struct import (
-    AbortReq,
-    BatchEmbeddingOut,
-    BatchTokenIDOut,
-    HealthCheckOutput,
-    TokenizedEmbeddingReqInput,
-    TokenizedGenerateReqInput,
-)
-from sglang.srt.server_args import PortArgs, ServerArgs
-from sglang.srt.utils import get_zmq_socket, kill_process_tree
-from sglang.utils import get_exception_traceback
-
-logger = logging.getLogger(__name__)
-
-
-class GrpcSignalHandler:
-    """Minimal signal handler for gRPC server - delegates real crash handling to scheduler."""
-
-    def __init__(self, grpc_manager):
-        self.grpc_manager = grpc_manager
-
-    def sigterm_handler(self, signum=None, frame=None):
-        """Handle SIGTERM by gracefully shutting down gRPC server."""
-        logger.warning(
-            f"SIGTERM received. {signum=} {frame=}. Shutting down gRPC server..."
-        )
-        self.grpc_manager.gracefully_exit = True
-
-    def running_phase_sigquit_handler(self, signum=None, frame=None):
-        """Handle SIGQUIT from failed scheduler process."""
-        logger.error(
-            "Received SIGQUIT from scheduler process. Scheduler failed, shutting down gRPC server."
-        )
-        logger.info(
-            "Note: Crash dumps are handled by the scheduler process, not the gRPC server."
-        )
-        # Just exit cleanly - the scheduler handles crash dumps
-        kill_process_tree(os.getpid(), include_parent=True)
-
-
-@dataclasses.dataclass
-class GrpcReqState:
-    """State tracking for a gRPC request."""
-
-    # Request identification
-    request_id: str
-    grpc_context: Optional[grpc.aio.ServicerContext]
-
-    # Communication
-    out_queue: asyncio.Queue
-    finished: bool
-    event: asyncio.Event
-    obj: Union[TokenizedGenerateReqInput, TokenizedEmbeddingReqInput]
-
-    # Metrics (same as TokenizerManager's ReqState)
-    created_time: float
-    finished_time: float = 0.0
-    first_token_time: float = 0.0
-    last_time: float = 0.0
-    last_completion_tokens: int = 1
-
-    # Streaming state
-    last_output_offset: int = 0
-    stream_finished: bool = False
-
-    # Output accumulation
-    text: str = ""
-    output_ids: List[int] = dataclasses.field(default_factory=list)
-    input_token_logprobs_val: List[float] = dataclasses.field(default_factory=list)
-    input_token_logprobs_idx: List[int] = dataclasses.field(default_factory=list)
-    output_token_logprobs_val: List[float] = dataclasses.field(default_factory=list)
-    output_token_logprobs_idx: List[int] = dataclasses.field(default_factory=list)
-    input_top_logprobs_val: List[List[float]] = dataclasses.field(default_factory=list)
-    input_top_logprobs_idx: List[List[int]] = dataclasses.field(default_factory=list)
-    output_top_logprobs_val: List[List[float]] = dataclasses.field(default_factory=list)
-    output_top_logprobs_idx: List[List[int]] = dataclasses.field(default_factory=list)
-
-    # Session state
-    session_id: Optional[str] = None
-    is_session_request: bool = False
-
-
-class GrpcRequestManager:
-    """
-    Manages gRPC request lifecycle, mimicking TokenizerManager's orchestration
-    behaviors without tokenization.
-    """
-
-    def __init__(
-        self,
-        server_args: ServerArgs,
-        port_args: PortArgs,
-    ):
-        """Initialize the gRPC request manager."""
-        self.server_args = server_args
-        self.port_args = port_args
-
-        # ZMQ Communication Setup (same pattern as TokenizerManager)
-        context = zmq.asyncio.Context(2)
-
-        # Socket for receiving outputs from scheduler
-        self.recv_from_scheduler = get_zmq_socket(
-            context, zmq.PULL, port_args.detokenizer_ipc_name, bind=True
-        )
-
-        # Socket for sending requests to scheduler
-        self.send_to_scheduler = get_zmq_socket(
-            context, zmq.PUSH, port_args.scheduler_input_ipc_name, bind=True
-        )
-
-        # State Management (from TokenizerManager)
-        self.rid_to_state: Dict[str, GrpcReqState] = {}
-        self.asyncio_tasks: set = set()
-        self.gracefully_exit = False
-        self.no_create_loop = False
-        self.event_loop = None
-
-        # Pause/Resume Control
-        self.is_pause = False
-        self.is_pause_cond = asyncio.Condition()
-
-        # Metrics
-        self.request_counter = 0
-        self.request_counter_lock = asyncio.Lock()
-        self.last_receive_tstamp = time.time()
-
-        # Crash dump for debugging
-        self.crash_dump_request_list = []
-        self.crash_dump_performed = False
-
-        logger.info(
-            f"GrpcRequestManager initialized with ZMQ IPC: "
-            f"recv={port_args.detokenizer_ipc_name}, "
-            f"send={port_args.scheduler_input_ipc_name}"
-        )
-
-    async def generate_request(
-        self,
-        obj: TokenizedGenerateReqInput,
-        request_id: Optional[str] = None,
-        grpc_context: Optional[grpc.aio.ServicerContext] = None,
-    ) -> asyncio.Queue:
-        """
-        Submit a generation request to the scheduler.
-        Returns a queue for streaming outputs.
-        """
-        # Generate request ID if not provided
-        if request_id is None:
-            async with self.request_counter_lock:
-                request_id = f"grpc-{self.request_counter}"
-                self.request_counter += 1
-
-        obj.rid = request_id
-
-        # TODO: support log_request
-
-        # Create request state
-        state = GrpcReqState(
-            request_id=request_id,
-            grpc_context=grpc_context,
-            out_queue=asyncio.Queue(),
-            finished=False,
-            event=asyncio.Event(),
-            obj=obj,
-            created_time=time.time(),
-        )
-
-        # Track session if needed
-        if hasattr(obj, "session_params") and obj.session_params:
-            state.session_id = obj.session_params.session_id
-            state.is_session_request = True
-
-        # Register state
-        self.rid_to_state[request_id] = state
-        self.record_request_for_crash_dump(obj)
-
-        # Send to scheduler via ZMQ
-        try:
-            await self._send_to_scheduler(obj)
-        except Exception as e:
-            # Clean up on failure
-            del self.rid_to_state[request_id]
-            raise RuntimeError(f"Failed to send request to scheduler: {e}")
-
-        return state.out_queue
-
-    async def embedding_request(
-        self,
-        obj: TokenizedEmbeddingReqInput,
-        request_id: Optional[str] = None,
-    ) -> asyncio.Future:
-        """
-        Submit an embedding request to the scheduler.
-        Returns a future that will contain the embedding result.
-        """
-        # Generate request ID if not provided
-        if request_id is None:
-            async with self.request_counter_lock:
-                request_id = f"grpc-embed-{self.request_counter}"
-                self.request_counter += 1
-
-        obj.rid = request_id
-
-        # Create request state
-        state = GrpcReqState(
-            request_id=request_id,
-            grpc_context=None,
-            out_queue=asyncio.Queue(),
-            finished=False,
-            event=asyncio.Event(),
-            obj=obj,
-            created_time=time.time(),
-        )
-
-        # Register state
-        self.rid_to_state[request_id] = state
-
-        # Create future for result
-        future = asyncio.Future()
-
-        # Send to scheduler
-        try:
-            await self._send_to_scheduler(obj)
-        except Exception as e:
-            del self.rid_to_state[request_id]
-            future.set_exception(e)
-            return future
-
-        # Wait for result in background
-        async def wait_for_result():
-            try:
-                # Wait for completion
-                await state.event.wait()
-                # Get result from queue
-                result = await state.out_queue.get()
-                future.set_result(result)
-            except Exception as e:
-                future.set_exception(e)
-            finally:
-                # Clean up
-                if request_id in self.rid_to_state:
-                    del self.rid_to_state[request_id]
-
-        asyncio.create_task(wait_for_result())
-        return future
-
-    async def abort_request(self, request_id: str) -> bool:
-        """Abort a running request."""
-        if request_id not in self.rid_to_state:
-            return False
-
-        # Send abort to scheduler
-        abort_req = AbortReq(rid=request_id)
-        try:
-            await self._send_to_scheduler(abort_req)
-        except Exception as e:
-            logger.error(f"Failed to send abort request: {e}")
-            return False
-
-        # Mark as finished
-        state = self.rid_to_state.get(request_id)
-        if state:
-            state.finished = True
-            state.stream_finished = True
-            state.event.set()
-
-            # Send abort notification to output queue
-            await state.out_queue.put({"error": "Request aborted", "abort": True})
-
-        return True
-
-    async def pause_generation(self):
-        """Pause generation processing."""
-        async with self.is_pause_cond:
-            self.is_pause = True
-            logger.info("Generation paused")
-
-    async def resume_generation(self):
-        """Resume generation processing."""
-        async with self.is_pause_cond:
-            self.is_pause = False
-            self.is_pause_cond.notify_all()
-            logger.info("Generation resumed")
-
-    async def handle_loop(self):
-        """
-        Main event loop - processes outputs from scheduler.
-        Mimics TokenizerManager's handle_loop.
-        """
-        while not self.gracefully_exit:
-            try:
-                # Receive from scheduler
-                recv_obj = await self.recv_from_scheduler.recv_pyobj()
-                self.last_receive_tstamp = time.time()
-
-                # Check for pause
-                async with self.is_pause_cond:
-                    while self.is_pause:
-                        await self.is_pause_cond.wait()
-
-                # Handle different output types
-                if isinstance(recv_obj, BatchTokenIDOut):
-                    await self._handle_batch_output(recv_obj)
-                elif isinstance(recv_obj, BatchEmbeddingOut):
-                    await self._handle_embedding_output(recv_obj)
-                elif isinstance(recv_obj, HealthCheckOutput):
-                    await self._handle_health_check_output(recv_obj)
-                else:
-                    logger.warning(f"Unknown output type: {type(recv_obj)}")
-
-            except zmq.error.Again:
-                # Timeout, check if we should exit
-                if self.gracefully_exit:
-                    break
-                continue
-            except Exception as e:
-                logger.error(f"Handle loop error: {e}\n{get_exception_traceback()}")
-                if self.gracefully_exit:
-                    break
-
-    async def _handle_batch_output(self, batch_out: BatchTokenIDOut):
-        """Handle batch generation output from scheduler."""
-        # Process each request in the batch
-        for i, rid in enumerate(batch_out.rids):
-            if rid not in self.rid_to_state:
-                continue
-
-            state = self.rid_to_state[rid]
-
-            # Update metrics
-            now = time.time()
-            if state.first_token_time == 0.0:
-                state.first_token_time = now
-            state.last_time = now
-
-            # Extract output for this request
-            output_data = {
-                "request_id": rid,
-                "text": batch_out.decoded_texts[i] if batch_out.decoded_texts else "",
-                "token_ids": batch_out.output_ids[i] if batch_out.output_ids else [],
-                "finished": batch_out.finished_reasons[i] is not None,
-                "meta_info": {
-                    "prompt_tokens": (
-                        batch_out.prompt_tokens[i] if batch_out.prompt_tokens else 0
-                    ),
-                    "completion_tokens": (
-                        batch_out.completion_tokens[i]
-                        if batch_out.completion_tokens
-                        else 0
-                    ),
-                    "finish_reason": (
-                        str(batch_out.finished_reasons[i])
-                        if batch_out.finished_reasons[i]
-                        else None
-                    ),
-                },
-            }
-
-            # Add logprobs if available
-            if batch_out.output_token_logprobs_val and i < len(
-                batch_out.output_token_logprobs_val
-            ):
-                output_data["logprobs"] = {
-                    "tokens": batch_out.output_token_logprobs_val[i],
-                    "top_logprobs": (
-                        batch_out.output_top_logprobs_val[i]
-                        if batch_out.output_top_logprobs_val
-                        and i < len(batch_out.output_top_logprobs_val)
-                        else None
-                    ),
-                }
-
-            # Update state
-            if output_data["text"]:
-                state.text += output_data["text"][state.last_output_offset :]
-                state.last_output_offset = len(output_data["text"])
-
-            if output_data["token_ids"]:
-                state.output_ids.extend(output_data["token_ids"])
-
-            # Send to output queue
-            await state.out_queue.put(output_data)
-
-            # Handle completion
-            if output_data["finished"]:
-                state.finished = True
-                state.finished_time = now
-                state.stream_finished = True
-                state.event.set()
-
-                # Remove from tracking after a delay
-                async def cleanup():
-                    await asyncio.sleep(5.0)
-                    if rid in self.rid_to_state:
-                        del self.rid_to_state[rid]
-
-                asyncio.create_task(cleanup())
-
-    async def _handle_embedding_output(self, batch_out: BatchEmbeddingOut):
-        """Handle batch embedding output from scheduler."""
-        for i, rid in enumerate(batch_out.rids):
-            if rid not in self.rid_to_state:
-                continue
-
-            state = self.rid_to_state[rid]
-
-            # Create result
-            result = {
-                "request_id": rid,
-                "embedding": batch_out.embeddings[i],
-                "prompt_tokens": (
-                    batch_out.prompt_tokens[i] if batch_out.prompt_tokens else 0
-                ),
-                "finish_reason": (
-                    batch_out.finish_reason[i] if batch_out.finish_reason else None
-                ),
-            }
-
-            # Send result
-            await state.out_queue.put(result)
-
-            # Mark as finished
-            state.finished = True
-            state.finished_time = time.time()
-            state.event.set()
-
-    async def _handle_health_check_output(self, health_out: HealthCheckOutput):
-        """Handle health check output from scheduler."""
-        rid = health_out.rid
-
-        if rid not in self.rid_to_state:
-            logger.warning(f"Health check output for unknown request: {rid}")
-            return
-
-        state = self.rid_to_state[rid]
-
-        # Create health check result
-        result = {
-            "request_id": rid,
-            "healthy": True,  # If we got a response, scheduler is healthy
-            "output_text": (
-                health_out.output_str if hasattr(health_out, "output_str") else ""
-            ),
-            "finish_reason": (
-                health_out.finish_reason
-                if hasattr(health_out, "finish_reason")
-                else "stop"
-            ),
-        }
-
-        # Send result
-        await state.out_queue.put(result)
-
-        # Mark as finished
-        state.finished = True
-        state.finished_time = time.time()
-        state.event.set()
-
-    async def _send_to_scheduler(self, obj):
-        """Send an object to the scheduler via ZMQ."""
-        try:
-            self.send_to_scheduler.send_pyobj(obj)
-        except Exception as e:
-            logger.error(f"Failed to send to scheduler: {e}")
-            raise
-
-    def record_request_for_crash_dump(self, obj):
-        """Record request for potential crash dump."""
-        if len(self.crash_dump_request_list) < 100:
-            self.crash_dump_request_list.append(
-                {
-                    "time": time.time(),
-                    "request_id": getattr(obj, "rid", "unknown"),
-                    "type": type(obj).__name__,
-                }
-            )
-
-    async def shutdown(self):
-        """Gracefully shutdown the request manager."""
-        logger.info("Shutting down GrpcRequestManager")
-        self.gracefully_exit = True
-
-        # Cancel all pending requests
-        for rid, state in self.rid_to_state.items():
-            if not state.finished:
-                await state.out_queue.put(
-                    {"error": "Server shutting down", "shutdown": True}
-                )
-                state.finished = True
-                state.event.set()
-
-        # Wait for tasks to complete
-        if self.asyncio_tasks:
-            await asyncio.gather(*list(self.asyncio_tasks), return_exceptions=True)
-
-        # Close ZMQ sockets
-        self.recv_from_scheduler.close()
-        self.send_to_scheduler.close()
-
-        logger.info("GrpcRequestManager shutdown complete")
-
-    def get_server_info(self) -> Dict[str, Any]:
-        """Get server information for health checks."""
-        return {
-            "active_requests": len(self.rid_to_state),
-            "paused": self.is_pause,
-            "last_receive_time": self.last_receive_tstamp,
-        }
-
-    def auto_create_handle_loop(self):
-        """Automatically create and start the handle_loop task, matching TokenizerManager pattern."""
-        if self.no_create_loop:
-            return
-
-        self.no_create_loop = True
-        loop = asyncio.get_event_loop()
-        self.asyncio_tasks.add(
-            loop.create_task(print_exception_wrapper(self.handle_loop))
-        )
-
-        self.event_loop = loop
-
-        # We cannot add signal handler when the grpc manager is not in
-        # the main thread due to the CPython limitation.
-        if threading.current_thread() is threading.main_thread():
-            signal_handler = GrpcSignalHandler(self)
-            loop.add_signal_handler(signal.SIGTERM, signal_handler.sigterm_handler)
-            # Update the signal handler for the process. It overrides the sigquit handler in the launch phase.
-            loop.add_signal_handler(
-                signal.SIGQUIT, signal_handler.running_phase_sigquit_handler
-            )
-        else:
-            logger.warning(
-                "Signal handler is not added because the grpc request manager is "
-                "not in the main thread. This disables graceful shutdown of the "
-                "grpc request manager when SIGTERM is received."
-            )
-        self.asyncio_tasks.add(
-            loop.create_task(print_exception_wrapper(self.sigterm_watchdog))
-        )
-
-    async def sigterm_watchdog(self):
-        """Watchdog to handle SIGTERM gracefully, matching TokenizerManager pattern."""
-        while not self.gracefully_exit:
-            await asyncio.sleep(1.0)
-
-
-async def print_exception_wrapper(func):
-    """
-    Sometimes an asyncio function does not print exception.
-    We do another wrapper to handle the exception.
-    """
-    try:
-        await func()
-    except Exception:
-        traceback = get_exception_traceback()
-        logger.error(f"GrpcRequestManager hit an exception: {traceback}")
-        if hasattr(func, "__self__") and isinstance(func.__self__, GrpcRequestManager):
-            func.__self__.dump_requests_before_crash()
-        kill_process_tree(os.getpid(), include_parent=True)
-        sys.exit(1)

sglang/srt/layers/quantization/deep_gemm_wrapper/configurer.py

@@ -1,32 +0,0 @@
-import logging
-
-import torch
-
-from sglang.srt.utils import get_bool_env_var, get_device_sm
-
-logger = logging.getLogger(__name__)
-
-
-def _compute_enable_deep_gemm():
-    sm_version = get_device_sm()
-    if sm_version < 90:
-        return False
-
-    try:
-        import deep_gemm
-    except ImportError:
-        logger.warning("Failed to import deep_gemm, disable ENABLE_JIT_DEEPGEMM.")
-        return False
-
-    return get_bool_env_var("SGL_ENABLE_JIT_DEEPGEMM", default="true")
-
-
-def _is_blackwell_arch() -> bool:
-    major, minor = torch.cuda.get_device_capability(torch.cuda.current_device())
-    return major == 10
-
-
-ENABLE_JIT_DEEPGEMM = _compute_enable_deep_gemm()
-
-DEEPGEMM_BLACKWELL = ENABLE_JIT_DEEPGEMM and _is_blackwell_arch()
-DEEPGEMM_SCALE_UE8M0 = DEEPGEMM_BLACKWELL