🐛 fixup cancel for v0.6.5

Signed-off-by: Joe Runde <[email protected]>

Author: joerunde

src/vllm_tgis_adapter/grpc/grpc_server.py

@@ -23,7 +23,6 @@
     extract_trace_headers,
     log_tracing_disabled_warning,
 )
-from vllm.utils import iterate_with_cancellation
 
 from vllm_tgis_adapter.logging import init_logger
 from vllm_tgis_adapter.tgis_utils import logs
@@ -262,11 +261,8 @@ async def Generate(
                 ),
             )
 
-        async def is_cancelled() -> bool:
-            return context.cancelled()
-
         result_generator: AsyncIterator[tuple[int, RequestOutput]] = (
-            merge_async_iterators(*generators, is_cancelled=is_cancelled)
+            merge_async_iterators(*generators)
         )
 
         resp_options = request.params.response
@@ -357,11 +353,6 @@ async def GenerateStream(  # noqa: PLR0915, C901
             **kwargs,
         )
 
-        async def is_cancelled() -> bool:
-            return context.cancelled()
-
-        result_generator = iterate_with_cancellation(result_generator, is_cancelled)
-
         resp_options = request.params.response
 
         first_response = None

src/vllm_tgis_adapter/tgis_utils/logs.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+import asyncio
 import functools
 import logging
 import time
@@ -86,6 +87,12 @@ async def generate_with_logging(*args, **kwargs) -> AsyncGenerator[RequestOutput
             async for response in old_generate_fn(*args, **kwargs):
                 last = response
                 yield response
+        except asyncio.CancelledError:
+            _log_cancellation(
+                request_id=request_id,
+                correlation_id=correlation_id,
+            )
+            raise
         except BaseException as e:
             # Log any error
             _log_error(
@@ -118,6 +125,14 @@ def _log_error(request_id: str, correlation_id: str, exception_str: str) -> None
     )
 
 
+def _log_cancellation(request_id: str, correlation_id: str) -> None:
+    logger.info(
+        "Request cancelled: request_id=%s correlation_id=%s",
+        request_id,
+        correlation_id,
+    )
+
+
 def _log_request(
     request_id: str,
     params: SamplingParams,