sglang 0.5.3rc0__py3-none-any.whl → 0.5.3rc2__py3-none-any.whl

sglang/srt/entrypoints/grpc_request_manager.py

@@ -4,6 +4,7 @@ Mimics TokenizerManager's state management and ZMQ communication patterns.
 """
 
 import asyncio
+import copy
 import dataclasses
 import logging
 import os
@@ -11,7 +12,8 @@ import signal
 import sys
 import threading
 import time
-from typing import Any, Dict, List, Optional, Union
+import uuid
+from typing import Any, AsyncGenerator, Dict, List, Optional, Union
 
 import grpc
 import zmq
@@ -19,8 +21,8 @@ import zmq.asyncio
 
 from sglang.srt.managers.io_struct import (
     AbortReq,
-    BatchEmbeddingOut,
-    BatchTokenIDOut,
+    BatchEmbeddingOutput,
+    BatchTokenIDOutput,
     HealthCheckOutput,
     TokenizedEmbeddingReqInput,
     TokenizedGenerateReqInput,
@@ -79,11 +81,10 @@ class GrpcReqState:
     last_completion_tokens: int = 1
 
     # Streaming state
-    last_output_offset: int = 0
     stream_finished: bool = False
+    input_logprobs_sent: bool = False  # Track if input logprobs were sent in streaming
 
-    # Output accumulation
-    text: str = ""
+    # Token accumulation (for non-streaming)
     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)
@@ -109,22 +110,23 @@ class GrpcRequestManager:
         self,
         server_args: ServerArgs,
         port_args: PortArgs,
+        bootstrap_server=None,
     ):
         """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)
+        self.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
+            self.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
+            self.context, zmq.PUSH, port_args.scheduler_input_ipc_name, bind=True
         )
 
         # State Management (from TokenizerManager)
@@ -139,41 +141,158 @@ class GrpcRequestManager:
         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
 
+        # Bootstrap server (passed from serve_grpc, not started here)
+        self.bootstrap_server = bootstrap_server
+
         logger.info(
             f"GrpcRequestManager initialized with ZMQ IPC: "
             f"recv={port_args.detokenizer_ipc_name}, "
             f"send={port_args.scheduler_input_ipc_name}"
         )
+        if self.bootstrap_server:
+            logger.info(
+                f"Bootstrap server initialized for disaggregation mode: "
+                f"{server_args.disaggregation_mode}"
+            )
 
     async def generate_request(
         self,
         obj: TokenizedGenerateReqInput,
         request_id: Optional[str] = None,
         grpc_context: Optional[grpc.aio.ServicerContext] = None,
-    ) -> asyncio.Queue:
+    ) -> AsyncGenerator[Union[Dict, List[Dict]], None]:
         """
-        Submit a generation request to the scheduler.
-        Returns a queue for streaming outputs.
+        Submit a generation request to the scheduler with n>1 parallel sampling support.
+
+        This method implements the same two-phase approach as tokenizer_manager.py:
+        1. Phase 1: Send prefix caching request (max_new_tokens=0)
+        2. Phase 2: Send n generation requests that reuse the cached prefix
+
+        Yields individual responses for streaming, or aggregated responses for non-streaming.
         """
+        n = getattr(obj.sampling_params, "n", 1)
+
+        if n <= 1:
+            async for response in self._handle_single_request(
+                obj, request_id, grpc_context
+            ):
+                yield response
+            return
+
+        # N>1 handling - two-phase approach
+        logger.debug(f"Multiple sampling request (n={n}), using two-phase approach")
+
+        # Generate base request ID if not provided
+        if request_id is None:
+            base_request_id = f"grpc-{uuid.uuid4().hex}"
+        else:
+            base_request_id = request_id
+
+        # Phase 1: Cache the common prefix
+        logger.debug(f"Phase 1: Caching prefix for request {base_request_id}")
+        prefix_obj = copy.copy(obj)
+        prefix_obj.sampling_params = copy.copy(obj.sampling_params)
+        prefix_obj.sampling_params.max_new_tokens = 0  # Prefill-only
+        prefix_obj.sampling_params.n = 1  # Don't replicate prefix request
+
+        # Send prefix caching request and consume response
+        async for _ in self._handle_single_request(
+            prefix_obj, f"{base_request_id}-prefix", grpc_context
+        ):
+            # Consume prefix response (usually just one chunk with finish_reason)
+            pass
+
+        logger.debug(f"Phase 1 completed: Prefix cached for {base_request_id}")
+
+        # Phase 2: Generate n parallel requests
+        logger.debug(f"Phase 2: Generating {n} parallel requests")
+        generators = []
+        request_ids = []
+
+        for i in range(n):
+            # Create individual generation request
+            gen_obj = copy.copy(obj)
+            gen_obj.sampling_params = copy.copy(obj.sampling_params)
+            gen_obj.sampling_params.n = 1  # Each request generates 1 response
+
+            gen_request_id = f"{base_request_id}-{i}"
+            request_ids.append(gen_request_id)
+
+            # Start generation request
+            generators.append(
+                self._handle_single_request(gen_obj, gen_request_id, grpc_context)
+            )
+
+        # Handle response aggregation
+        is_stream = getattr(obj, "stream", False)
+
+        if not is_stream:
+            # Non-streaming: collect all responses and return as batch
+            logger.debug(f"Non-streaming mode: collecting {n} responses")
+            responses = []
+            for generator in generators:
+                async for response in generator:
+                    responses.append(response)
+            yield responses  # Return all responses as a batch
+        else:
+            # Streaming mode: multiplex responses with index for ordering
+            logger.debug(f"Streaming mode: multiplexing {n} streams")
+            rid_to_index = {rid: i for i, rid in enumerate(request_ids)}
+
+            # Create async tasks for all generators
+            task_map = {}
+            for generator in generators:
+                task = asyncio.create_task(generator.__anext__())
+                task_map[task] = generator
+
+            # Process responses as they arrive
+            while task_map:
+                done, _ = await asyncio.wait(
+                    task_map.keys(), return_when=asyncio.FIRST_COMPLETED
+                )
+
+                for task in done:
+                    generator = task_map.pop(task)
+                    try:
+                        response = await task
+
+                        # Add index for client-side ordering
+                        if isinstance(response, dict) and "meta_info" in response:
+                            response_rid = response["meta_info"].get("id", "")
+                            if response_rid in rid_to_index:
+                                response["index"] = rid_to_index[response_rid]
+
+                        yield response
+
+                        # Create next task for this generator
+                        next_task = asyncio.create_task(generator.__anext__())
+                        task_map[next_task] = generator
+
+                    except StopAsyncIteration:
+                        # This generator is finished
+                        pass
+
+    async def _handle_single_request(
+        self,
+        obj: TokenizedGenerateReqInput,
+        request_id: Optional[str] = None,
+        grpc_context: Optional[grpc.aio.ServicerContext] = None,
+    ):
+        """Handle a single request - core implementation without n>1 logic."""
         # 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
+            request_id = f"grpc-{uuid.uuid4().hex}"
 
         obj.rid = request_id
 
+        # Create and register request state
         # TODO: support log_request
-
-        # Create request state
         state = GrpcReqState(
             request_id=request_id,
             grpc_context=grpc_context,
@@ -189,19 +308,51 @@ class GrpcRequestManager:
             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:
+            # Send to scheduler - let exceptions bubble up to grpc_server.py
             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
+            is_stream = getattr(obj, "stream", False)
+
+            while True:
+                # Client cancelled - notify scheduler and exit
+                if grpc_context and grpc_context.cancelled():
+                    await self.abort_request(request_id)
+                    return
+
+                try:
+                    response = await asyncio.wait_for(state.out_queue.get(), timeout=4)
+
+                    if is_stream:
+                        yield response
+
+                    # Non-streaming: yield final response with accumulated tokens from state
+                    if isinstance(response, dict) and response.get("finished", False):
+                        if not is_stream:
+                            final_response = response.copy()
+                            final_response["token_ids"] = state.output_ids
+                            yield final_response
+                        break
+
+                except asyncio.TimeoutError:
+                    # Timeout waiting for response - abort and cleanup
+                    logger.warning(
+                        f"Timeout waiting for response for request {request_id}"
+                    )
+                    await self.abort_request(request_id)
+                    return
+
+        finally:
+            # Always clean up request state when exiting
+            self._cleanup_request_state(request_id)
+
+    def _cleanup_request_state(self, request_id: str):
+        """Clean up local request state (does not notify scheduler)."""
+        if request_id in self.rid_to_state:
+            del self.rid_to_state[request_id]
 
     async def embedding_request(
         self,
@@ -214,9 +365,7 @@ class GrpcRequestManager:
         """
         # 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
+            request_id = f"grpc-embed-{uuid.uuid4().hex}"
 
         obj.rid = request_id
 
@@ -318,9 +467,9 @@ class GrpcRequestManager:
                         await self.is_pause_cond.wait()
 
                 # Handle different output types
-                if isinstance(recv_obj, BatchTokenIDOut):
+                if isinstance(recv_obj, BatchTokenIDOutput):
                     await self._handle_batch_output(recv_obj)
-                elif isinstance(recv_obj, BatchEmbeddingOut):
+                elif isinstance(recv_obj, BatchEmbeddingOutput):
                     await self._handle_embedding_output(recv_obj)
                 elif isinstance(recv_obj, HealthCheckOutput):
                     await self._handle_health_check_output(recv_obj)
@@ -332,12 +481,71 @@ class GrpcRequestManager:
                 if self.gracefully_exit:
                     break
                 continue
+            except zmq.error.ZMQError as e:
+                # Socket closed or other ZMQ error - exit cleanly if shutting down
+                if self.gracefully_exit:
+                    logger.debug(f"ZMQ recv interrupted during shutdown: {e}")
+                    break
+                logger.error(
+                    f"ZMQ error in handle loop: {e}\n{get_exception_traceback()}"
+                )
+                break
             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):
+    def _convert_logprob_style(
+        self,
+        state: GrpcReqState,
+        batch_out: BatchTokenIDOutput,
+        batch_index: int,
+    ):
+        """
+        Convert and accumulate logprobs from batch output to state.
+        Follows the same logic as tokenizer_manager.convert_logprob_style.
+        """
+        # Early exit if no input logprobs at all
+        if batch_out.input_token_logprobs_val is None:
+            return
+
+        # Accumulate input token logprobs (only if list is non-empty)
+        if len(batch_out.input_token_logprobs_val) > 0:
+            state.input_token_logprobs_val.extend(
+                batch_out.input_token_logprobs_val[batch_index]
+            )
+            state.input_token_logprobs_idx.extend(
+                batch_out.input_token_logprobs_idx[batch_index]
+            )
+
+        # Always accumulate output token logprobs
+        state.output_token_logprobs_val.extend(
+            batch_out.output_token_logprobs_val[batch_index]
+        )
+        state.output_token_logprobs_idx.extend(
+            batch_out.output_token_logprobs_idx[batch_index]
+        )
+
+        # Handle top logprobs if requested
+        if state.obj.top_logprobs_num > 0:
+            # Accumulate input top logprobs (only if list is non-empty)
+            if len(batch_out.input_top_logprobs_val) > 0:
+                state.input_top_logprobs_val.extend(
+                    batch_out.input_top_logprobs_val[batch_index]
+                )
+                state.input_top_logprobs_idx.extend(
+                    batch_out.input_top_logprobs_idx[batch_index]
+                )
+
+            # Always accumulate output top logprobs
+            state.output_top_logprobs_val.extend(
+                batch_out.output_top_logprobs_val[batch_index]
+            )
+            state.output_top_logprobs_idx.extend(
+                batch_out.output_top_logprobs_idx[batch_index]
+            )
+
+    async def _handle_batch_output(self, batch_out: BatchTokenIDOutput):
         """Handle batch generation output from scheduler."""
         # Process each request in the batch
         for i, rid in enumerate(batch_out.rids):
@@ -355,7 +563,6 @@ class GrpcRequestManager:
             # 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": {
@@ -367,37 +574,81 @@ class GrpcRequestManager:
                         if batch_out.completion_tokens
                         else 0
                     ),
+                    "cached_tokens": (
+                        batch_out.cached_tokens[i] if batch_out.cached_tokens else 0
+                    ),
                     "finish_reason": (
-                        str(batch_out.finished_reasons[i])
+                        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"])
+            # Accumulate logprobs (following tokenizer_manager pattern)
+            if state.obj.return_logprob:
+                self._convert_logprob_style(state, batch_out, i)
 
+            # Send input logprobs based if available
+            if (
+                state.obj.return_logprob
+                and state.obj.logprob_start_len >= 0
+                and state.input_token_logprobs_val
+            ):
+                if state.obj.stream and not state.input_logprobs_sent:
+                    # Streaming: send input logprobs once in first chunk that has them
+                    output_data["input_logprobs"] = {
+                        "token_logprobs_val": state.input_token_logprobs_val,
+                        "token_logprobs_idx": state.input_token_logprobs_idx,
+                        "top_logprobs_val": state.input_top_logprobs_val,
+                        "top_logprobs_idx": state.input_top_logprobs_idx,
+                    }
+                    state.input_logprobs_sent = True
+                elif not state.obj.stream and output_data["finished"]:
+                    # Non-streaming: send input logprobs in final chunk
+                    output_data["input_logprobs"] = {
+                        "token_logprobs_val": state.input_token_logprobs_val,
+                        "token_logprobs_idx": state.input_token_logprobs_idx,
+                        "top_logprobs_val": state.input_top_logprobs_val,
+                        "top_logprobs_idx": state.input_top_logprobs_idx,
+                    }
+
+            # Send output logprobs if available
+            if (
+                state.obj.return_logprob
+                and batch_out.output_token_logprobs_val
+                and i < len(batch_out.output_token_logprobs_val)
+            ):
+                if state.obj.stream:
+                    # For streaming: send incremental logprobs (only new tokens in this chunk)
+                    # NOTE: this is different than TokenizerManager, which always accumulates
+                    def get_part(attr_name):
+                        source_list = getattr(batch_out, attr_name, None)
+                        return (
+                            source_list[i]
+                            if source_list and i < len(source_list)
+                            else []
+                        )
+
+                    output_data["output_logprobs"] = {
+                        "token_logprobs_val": batch_out.output_token_logprobs_val[i],
+                        "token_logprobs_idx": get_part("output_token_logprobs_idx"),
+                        "top_logprobs_val": get_part("output_top_logprobs_val"),
+                        "top_logprobs_idx": get_part("output_top_logprobs_idx"),
+                    }
+                elif output_data["finished"]:
+                    # Non-streaming: send cumulative output logprobs in final chunk
+                    output_data["output_logprobs"] = {
+                        "token_logprobs_val": state.output_token_logprobs_val,
+                        "token_logprobs_idx": state.output_token_logprobs_idx,
+                        "top_logprobs_val": state.output_top_logprobs_val,
+                        "top_logprobs_idx": state.output_top_logprobs_idx,
+                    }
+
+            # Update state for accumulation
             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
@@ -415,7 +666,7 @@ class GrpcRequestManager:
 
                 asyncio.create_task(cleanup())
 
-    async def _handle_embedding_output(self, batch_out: BatchEmbeddingOut):
+    async def _handle_embedding_output(self, batch_out: BatchEmbeddingOutput):
         """Handle batch embedding output from scheduler."""
         for i, rid in enumerate(batch_out.rids):
             if rid not in self.rid_to_state:
@@ -499,8 +750,17 @@ class GrpcRequestManager:
         logger.info("Shutting down GrpcRequestManager")
         self.gracefully_exit = True
 
+        # Cancel all asyncio tasks FIRST - this will interrupt blocked recv() calls
+        for task in list(self.asyncio_tasks):
+            if not task.done():
+                task.cancel()
+
+        # Give tasks a moment to process cancellation
+        if self.asyncio_tasks:
+            await asyncio.gather(*list(self.asyncio_tasks), return_exceptions=True)
+
         # Cancel all pending requests
-        for rid, state in self.rid_to_state.items():
+        for rid, state in list(self.rid_to_state.items()):
             if not state.finished:
                 await state.out_queue.put(
                     {"error": "Server shutting down", "shutdown": True}
@@ -512,10 +772,25 @@ class GrpcRequestManager:
         if self.asyncio_tasks:
             await asyncio.gather(*list(self.asyncio_tasks), return_exceptions=True)
 
+        # Shutdown bootstrap server if running
+        if self.bootstrap_server:
+            logger.info("Shutting down bootstrap server")
+            try:
+                if hasattr(self.bootstrap_server, "shutdown"):
+                    if asyncio.iscoroutinefunction(self.bootstrap_server.shutdown):
+                        await self.bootstrap_server.shutdown()
+                    else:
+                        self.bootstrap_server.shutdown()
+            except Exception as e:
+                logger.warning(f"Error shutting down bootstrap server: {e}")
+
         # Close ZMQ sockets
         self.recv_from_scheduler.close()
         self.send_to_scheduler.close()
 
+        # Terminate the ZMQ context - this is critical for asyncio loop to exit cleanly
+        self.context.term()
+
         logger.info("GrpcRequestManager shutdown complete")
 
     def get_server_info(self) -> Dict[str, Any]: