feat: request migration for trtllm

components/src/dynamo/sglang/main.py

@@ -4,7 +4,6 @@
 import asyncio
 import logging
 import os
-import signal
 import sys
 
 import sglang as sgl
@@ -79,31 +78,26 @@ async def worker():
         config.dynamo_args.use_kv_events,
     )
 
-    def signal_handler():
-        asyncio.create_task(graceful_shutdown(runtime))
-
-    for sig in (signal.SIGTERM, signal.SIGINT):
-        loop.add_signal_handler(sig, signal_handler)
-
-    logging.info("Signal handlers will trigger a graceful shutdown of the runtime")
-
-    if config.dynamo_args.embedding_worker:
-        await init_embedding(runtime, config)
-    elif config.dynamo_args.multimodal_processor:
-        await init_multimodal_processor(runtime, config)
-    elif config.dynamo_args.multimodal_encode_worker:
-        await init_multimodal_encode_worker(runtime, config)
-    elif config.dynamo_args.multimodal_worker:
-        if config.serving_mode != DisaggregationMode.PREFILL:
-            await init_multimodal_worker(runtime, config)
+    try:
+        if config.dynamo_args.embedding_worker:
+            await init_embedding(runtime, config)
+        elif config.dynamo_args.multimodal_processor:
+            await init_multimodal_processor(runtime, config)
+        elif config.dynamo_args.multimodal_encode_worker:
+            await init_multimodal_encode_worker(runtime, config)
+        elif config.dynamo_args.multimodal_worker:
+            if config.serving_mode != DisaggregationMode.PREFILL:
+                await init_multimodal_worker(runtime, config)
+            else:
+                await init_multimodal_prefill_worker(runtime, config)
+        elif config.dynamo_args.diffusion_worker:
+            await init_diffusion(runtime, config)
+        elif config.serving_mode != DisaggregationMode.PREFILL:
+            await init(runtime, config)
         else:
-            await init_multimodal_prefill_worker(runtime, config)
-    elif config.dynamo_args.diffusion_worker:
-        await init_diffusion(runtime, config)
-    elif config.serving_mode != DisaggregationMode.PREFILL:
-        await init(runtime, config)
-    else:
-        await init_prefill(runtime, config)
+            await init_prefill(runtime, config)
+    finally:
+        runtime.shutdown()
 
 
 async def init(runtime: DistributedRuntime, config: Config):
@@ -184,6 +178,7 @@ async def init(runtime: DistributedRuntime, config: Config):
         logging.error(f"Failed to serve endpoints: {e}")
         raise
     finally:
+        logging.info("Finally block reached")
         metrics_task.cancel()
         try:
             await metrics_task
@@ -640,12 +635,6 @@ async def _do_warmup():
         logging.warning(f"Prefill warmup failed: {e}")
 
 
-async def graceful_shutdown(runtime):
-    logging.info("Received shutdown signal, shutting down DistributedRuntime")
-    runtime.shutdown()
-    logging.info("DistributedRuntime shutdown complete")
-
-
 def main():
     uvloop.run(worker())
 

components/src/dynamo/sglang/request_handlers/embedding/embedding_handler.py

@@ -29,7 +29,7 @@ def cleanup(self):
         logging.info("Engine shutdown")
         super().cleanup()
 
-    async def generate(self, request: dict, context: Context):
+    async def _internal_generate(self, request: dict, context: Context):
         """
         Generate embeddings for the given input.
 

components/src/dynamo/sglang/request_handlers/handler_base.py

@@ -169,8 +169,30 @@ def register_engine_routes(self, runtime) -> None:
             "resume_memory_occupation", self.resume_memory_occupation
         )
 
-    @abstractmethod
     async def generate(self, request: Dict[str, Any], context: Context):
+        """Generate response from request with cancellation handling.
+
+        This is a wrapper around _internal_generate that provides cancellation
+        handling for the generation process. It monitors for local cancellation
+        events (e.g., SIGINT, SIGTERM) that cause the local process to be terminated,
+        not remote client request cancellations.
+
+        When a local cancellation is detected via the context, it logs the cancellation
+        and raises GeneratorExit to gracefully terminate the generation.
+
+
+        Raises:
+            GeneratorExit: When the request is cancelled due to local cancellation.
+        """
+        try:
+            async for item in self._internal_generate(request, context):
+                yield item
+        except asyncio.CancelledError:
+            logging.info(f"Request {context.id()} cancelled")
+            raise GeneratorExit("Request cancelled")
+
+    @abstractmethod
+    async def _internal_generate(self, request: Dict[str, Any], context: Context):
         """Generate response from request.
 
         Args:

components/src/dynamo/sglang/request_handlers/llm/decode_handler.py

@@ -86,7 +86,7 @@ def _build_sampling_params(self, request: Dict[str, Any]) -> Dict[str, Any]:
 
         return {k: v for k, v in param_mapping.items() if v is not None}
 
-    async def generate(
+    async def _internal_generate(
         self, request: Dict[str, Any], context: Context
     ) -> AsyncGenerator[Dict[str, Any], None]:
         """Generate response in aggregated or disaggregated mode.

components/src/dynamo/sglang/request_handlers/llm/prefill_handler.py

@@ -53,7 +53,7 @@ def cleanup(self) -> None:
         logging.info("Prefill engine shutdown")
         super().cleanup()
 
-    async def generate(
+    async def _internal_generate(
         self, request: Dict[str, Any], context: Context
     ) -> AsyncGenerator[Dict[str, Any], None]:
         """Generate prefill output and provide bootstrap info for decode worker.

components/src/dynamo/sglang/request_handlers/multimodal/encode_worker_handler.py

@@ -91,7 +91,7 @@ def __init__(
     def cleanup(self):
         pass
 
-    async def generate(
+    async def _internal_generate(
         self, request: SglangMultimodalRequest, context: Context
     ) -> AsyncIterator[str]:
         """

components/src/dynamo/sglang/request_handlers/multimodal/processor_handler.py

@@ -54,7 +54,9 @@ def __init__(
     def cleanup(self):
         pass
 
-    async def generate(self, raw_request: MultiModalRequest, context: Context):
+    async def _internal_generate(
+        self, raw_request: MultiModalRequest, context: Context
+    ):
         """
         Process multimodal request and forward to encode worker.
 

components/src/dynamo/sglang/request_handlers/multimodal/worker_handler.py

@@ -269,7 +269,7 @@ def _validate_and_parse_request(self, request) -> SglangMultimodalRequest:
                 request = SglangMultimodalRequest.model_validate(request)
         return request
 
-    async def generate(
+    async def _internal_generate(
         self, request: SglangMultimodalRequest, context: Context
     ) -> AsyncIterator[str]:
         """
@@ -428,7 +428,7 @@ async def async_init(self):
         """Initialize async components like connector"""
         await self.embeddings_processor.initialize()
 
-    async def generate(
+    async def _internal_generate(
         self, disagg_request: DisaggSglangMultimodalRequest, context: Context
     ) -> AsyncIterator[str]:
         """

components/src/dynamo/trtllm/main.py

@@ -71,8 +71,9 @@
 configure_dynamo_logging()
 
 
-async def graceful_shutdown(runtime):
+async def graceful_shutdown(runtime, shutdown_event):
     logging.info("Received shutdown signal, shutting down DistributedRuntime")
+    shutdown_event.set()
     runtime.shutdown()
     logging.info("DistributedRuntime shutdown complete")
 
@@ -128,6 +129,9 @@ async def worker():
     config = cmd_line_args()
 
     loop = asyncio.get_running_loop()
+    # Create shutdown event
+    shutdown_event = asyncio.Event()
+
     # Enable NATS based on use_kv_events flag (derived from publish_events_and_metrics)
     runtime = DistributedRuntime(
         loop, config.store_kv, config.request_plane, config.use_kv_events
@@ -136,17 +140,19 @@ async def worker():
     # Set up signal handler for graceful shutdown
     def signal_handler():
         # Schedule the shutdown coroutine instead of calling it directly
-        asyncio.create_task(graceful_shutdown(runtime))
+        asyncio.create_task(graceful_shutdown(runtime, shutdown_event))
 
     for sig in (signal.SIGTERM, signal.SIGINT):
         loop.add_signal_handler(sig, signal_handler)
 
     logging.info("Signal handlers set up for graceful shutdown")
 
-    await init(runtime, config)
+    await init(runtime, config, shutdown_event)
 
 
-async def init(runtime: DistributedRuntime, config: Config):
+async def init(
+    runtime: DistributedRuntime, config: Config, shutdown_event: asyncio.Event
+):
     """
     Instantiate and serve
     """
@@ -425,6 +431,7 @@ async def init(runtime: DistributedRuntime, config: Config):
             runtime=runtime,  # Pass runtime for graceful shutdown
             metrics_collector=metrics_collector,
             kv_block_size=config.kv_block_size,
+            shutdown_event=shutdown_event,
         )
 
         # Register the model with runtime config

components/src/dynamo/trtllm/request_handlers/handler_base.py

@@ -67,6 +67,7 @@ class RequestHandlerConfig:
     ] = None  # DistributedRuntime reference for graceful shutdown
     metrics_collector: Optional[Any] = None  # TensorRT-LLM MetricsCollector
     kv_block_size: int = 32
+    shutdown_event: Optional[asyncio.Event] = None
 
 
 class HandlerBase:
@@ -88,6 +89,7 @@ def __init__(self, config: RequestHandlerConfig):
         # Store runtime reference for graceful shutdown
         self.runtime = config.runtime
         self.kv_block_size: int = config.kv_block_size
+        self.shutdown_event = config.shutdown_event
 
     def check_error(self, result: dict):
         """
@@ -170,47 +172,101 @@ def _extract_logprobs(
 
         return log_probs if log_probs else None, top_logprobs if top_logprobs else None
 
-    async def _handle_cancellation(
+    async def _handle_cancellation_and_shutdown(
         self, generation_result: GenerationResult, context: Context
     ):
-        """Background task to handle cancellation by monitoring context state."""
+        """
+        Background task to handle cancellation and shutdown by monitoring both signals.
+        Returns 'shutdown' if shutdown was triggered, 'cancelled' if cancelled, None otherwise.
+        """
         try:
-            # Wait asynchronously for cancellation signal instead of polling
-            await context.async_killed_or_stopped()
-            # Abort the generation
-            generation_result.abort()
-            logging.debug(f"Aborted Request ID: {context.id()}")
+            # Get the cancellation future (already a Future, not a coroutine)
+            cancellation_future = context.async_killed_or_stopped()
+
+            if self.shutdown_event:
+                # Create task for shutdown monitoring
+                shutdown_task = asyncio.create_task(self.shutdown_event.wait())
+                # Wait for whichever happens first
+                # asyncio.wait() accepts both Futures and Tasks
+                done, pending = await asyncio.wait(
+                    [cancellation_future, shutdown_task],
+                    return_when=asyncio.FIRST_COMPLETED,
+                )
+
+                # Cancel the pending task/future
+                for task in pending:
+                    task.cancel()
+                    try:
+                        await task
+                    except asyncio.CancelledError:
+                        pass
+
+                # Abort the generation
+                logging.info(f"Aborting request {context.id()}")
+                generation_result.abort()
+                logging.debug(f"Aborted Request ID: {context.id()}")
+
+                # Check which event triggered and return the reason
+                if shutdown_task in done:
+                    logging.info(f"Shutdown event detected for request {context.id()}")
+                    return "shutdown"
+                else:
+                    return "cancelled"
+            else:
+                # No shutdown event, just wait for cancellation
+                await cancellation_future
+                logging.info(f"Aborting request {context.id()}")
+                generation_result.abort()
+                logging.debug(f"Aborted Request ID: {context.id()}")
+                return "cancelled"
         except asyncio.CancelledError:
-            # Task was cancelled, which is expected when generation completes
-            pass
+            # Task was cancelled, which is expected when generation completes normally
+            return None
 
     @asynccontextmanager
     async def _cancellation_monitor(
         self, generation_result: GenerationResult, context: Context
     ) -> AsyncGenerator[asyncio.Task, None]:
         """
-        Context manager for monitoring request cancellation.
+        Context manager for monitoring request cancellation and shutdown.
 
-        Automatically creates a background task to monitor for cancellation and
-        cleans it up when the context exits.
+        Automatically creates a background task to monitor for cancellation
+        and shutdown events, cleaning it up when the context exits.
+
+        If shutdown event was triggered, raises GeneratorExit on exit.
 
         Yields:
-            asyncio.Task: The cancellation monitoring task
+            asyncio.Task: The monitoring task
         """
-        cancellation_task = asyncio.create_task(
-            self._handle_cancellation(generation_result, context)
+        monitor_task = asyncio.create_task(
+            self._handle_cancellation_and_shutdown(generation_result, context)
         )
 
         try:
-            yield cancellation_task
+            yield monitor_task
         finally:
-            # Clean up the background cancellation task
-            if not cancellation_task.done():
-                cancellation_task.cancel()
+            # Clean up the background monitoring task
+            shutdown_triggered = False
+            if not monitor_task.done():
+                monitor_task.cancel()
                 try:
-                    await cancellation_task
+                    await monitor_task
                 except asyncio.CancelledError:
                     pass
+            else:
+                # Task completed, check if it was due to shutdown
+                try:
+                    result = monitor_task.result()
+                    if result == "shutdown":
+                        shutdown_triggered = True
+                except Exception:
+                    pass
+
+            # Raise GeneratorExit if shutdown was triggered
+            if shutdown_triggered:
+                raise GeneratorExit(
+                    "Decode engine was shut down during token generation"
+                )
 
     def _decode_disaggregated_params_from_prefill(
         self, prefill_result: dict
@@ -653,7 +709,7 @@ async def generate_locally(
                 trace_headers=trace_headers,
             )
 
-            # Use the context manager to handle cancellation monitoring
+            # Use the context manager to handle cancellation and shutdown monitoring
             async with self._cancellation_monitor(generation_result, context):
                 async for res in generation_result:
                     # TRTLLM engine needs to start generating tokens first before stats