Repository: huggingface/transformers
Pull request: feat: support request cancellation

Support client-side request cancellation, so server avoids continuing to generate tokens needlessly, saving resources (also helps server exit).

Will rebase when #40479 is merged to get clearer diff.

Changed files:
filename: src/transformers/commands/serving.py
patch: @@ -11,6 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+import asyncio
 import base64
 import copy
 import datetime
@@ -24,7 +25,7 @@
 import threading
 import time
 from argparse import ArgumentParser, Namespace
-from collections.abc import Generator, Iterable
+from collections.abc import AsyncGenerator, Generator, Iterable
 from contextlib import asynccontextmanager
 from dataclasses import dataclass, field
 from io import BytesIO
@@ -127,10 +128,11 @@ class TransformersResponseCreateParamsStreaming(ResponseCreateParamsStreaming, t
 
     class TransformersCompletionCreateParamsStreaming(CompletionCreateParamsStreaming, total=False):
         """
-        OpenAI's CompletionCreateParamsStreaming with an additional field for the generation config (as a json string).
+        OpenAI's CompletionCreateParamsStreaming with additional fields for the generation config (as a json string) and passing the request_id
         """
 
         generation_config: str
+        request_id: str
 
     class TransformersTranscriptionCreateParams(TranscriptionCreateParamsBase, total=False):
         """
@@ -784,7 +786,7 @@ def get_gen_models(self) -> list[dict[str, any]]:
                 for model in model_infos
             ]
 
-    def continuous_batching_chat_completion(self, req: dict) -> Generator[str, None, None]:
+    def continuous_batching_chat_completion(self, req: dict) -> AsyncGenerator[str, None]:
         """
         Generates an OpenAI Chat Completion using continuous batching.
 
@@ -832,13 +834,8 @@ def continuous_batching_chat_completion(self, req: dict) -> Generator[str, None,
             model.device
         )
 
-        def stream_chat_completion(_inputs):
+        def stream_chat_completion(request_id, decode_stream):
             try:
-                decode_stream = DecodeStream(_inputs.tolist(), False)
-                request_id = self.running_continuous_batching_manager.add_request(
-                    _inputs, request_id=req.get("request_id"), max_new_tokens=generation_config.max_new_tokens
-                )
-
                 # Emit the assistant role to start the stream. Other chunks won't have a role, as it is implicit
                 # they come from the assistant.
                 yield self.build_chat_completion_chunk(request_id, role="assistant", model=model_id_and_revision)
@@ -862,9 +859,25 @@ def stream_chat_completion(_inputs):
 
             except Exception as e:
                 logger.error(str(e))
+                self.running_continuous_batching_manager.cancel_request(request_id)
                 yield f'data: {{"error": "{str(e)}"}}'
 
-        return stream_chat_completion(inputs[0])
+        async def cancellation_wrapper(_inputs):
+            request_id = None
+            try:
+                decode_stream = DecodeStream(_inputs.tolist(), False)
+                request_id = self.running_continuous_batching_manager.add_request(
+                    _inputs, request_id=req.get("request_id"), max_new_tokens=generation_config.max_new_tokens
+                )
+                for chunk in stream_chat_completion(request_id, decode_stream):
+                    yield chunk
+                    await asyncio.sleep(0)  # Yield control to the event loop to check for cancellations
+            except asyncio.CancelledError:
+                if request_id is not None:
+                    self.running_continuous_batching_manager.cancel_request(request_id)
+                    logger.warning(f"Request {request_id} was cancelled.")
+
+        return cancellation_wrapper(inputs[0])
 
     @staticmethod
     def get_model_modality(model: "PreTrainedModel") -> Modality:

filename: src/transformers/generation/continuous_batching/continuous_api.py
patch: @@ -226,6 +226,7 @@ def prepare_next_batch(self) -> bool:
         """Prepare tensors and metadata for the next model forward pass."""
         # Get new requests from the queue
         self._get_new_requests()
+        self.scheduler.clear_cancelled_requests()
         if not self.scheduler.has_pending_requests():
             return False
 
@@ -547,6 +548,15 @@ def add_requests(self, inputs: list[list[int]], **kwargs):
         for input_ids in inputs:
             self.add_request(input_ids, **kwargs)
 
+    def cancel_request(self, request_id: str):
+        """Cancel a request by its ID.
+
+        Args:
+            request_id: The ID of the request to cancel
+        """
+        if self.batch_processor is not None:
+            self.batch_processor.scheduler.set_request_cancellation(request_id)
+
     def get_result(self, request_id=None, timeout=None) -> Optional[GenerationOutput]:
         """Retrieve one result from the output queue.
 
@@ -577,10 +587,13 @@ def __iter__(self):
 
     def request_id_iter(self, request_id):
         """Iterate over results matching a specific request id as they become available."""
-        while self._generation_thread is not None and self._generation_thread.is_alive():
+        request_cancelled = False
+        while self._generation_thread is not None and self._generation_thread.is_alive() and not request_cancelled:
             result = self.get_result(request_id=request_id, timeout=0.1)
             if result is not None:
                 yield result
+            if self.batch_processor is not None:
+                request_cancelled = self.batch_processor.scheduler.request_is_cancelled(request_id)
 
     @traced
     def warmup(self, batch_processor):

filename: src/transformers/generation/continuous_batching/scheduler.py
patch: @@ -12,6 +12,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+import threading
 from abc import ABC, abstractmethod
 from collections import deque
 
@@ -32,6 +33,8 @@ def __init__(self, cache: PagedAttentionCache, retain_cache_on_finish: bool = Fa
         self.waiting_requests_order: deque[str] = deque()
         self.cache = cache
         self.retain_cache_on_finish = retain_cache_on_finish
+        self._cancellation_lock = threading.Lock()
+        self._requests_to_cancel: set[str] = set()
 
     @abstractmethod
     def add_waiting_request(self, state: RequestState):
@@ -58,6 +61,30 @@ def get_active_request_static_outputs(self, request_id: str) -> list[int]:
             return self.active_requests[request_id].static_outputs
         return []
 
+    @traced
+    def set_request_cancellation(self, request_id: str):
+        with self._cancellation_lock:
+            self._requests_to_cancel.add(request_id)
+
+    @traced
+    def clear_cancelled_requests(self):
+        with self._cancellation_lock:
+            for request_id in self._requests_to_cancel:
+                if request_id in self.active_requests:
+                    del self.active_requests[request_id]
+                if request_id in self.waiting_requests:
+                    del self.waiting_requests[request_id]
+                if request_id in self.waiting_requests_order:
+                    self.waiting_requests_order.remove(request_id)
+                self.cache.free_blocks(request_id)
+            self._requests_to_cancel = set()
+
+    @traced
+    def request_is_cancelled(self, request_id: str) -> bool:
+        return request_id in self._requests_to_cancel or (
+            request_id not in self.active_requests and request_id not in self.waiting_requests
+        )
+
 
 @attach_tracer()
 class FIFOScheduler(Scheduler):

filename: tests/commands/test_serving.py
patch: @@ -19,6 +19,7 @@
 from unittest.mock import patch
 
 import aiohttp.client_exceptions
+import requests
 from huggingface_hub import AsyncInferenceClient, ChatCompletionStreamOutput
 from parameterized import parameterized
 
@@ -492,6 +493,37 @@ def test_tool_call(self):
         self.assertTrue(all(reason is None for reason in finish_reasons[:-1]))
 
 
+def _get_scheduler(serve_command):
+    # Defensive navigation in case any layer is renamed in the future
+    cbm = getattr(serve_command, "running_continuous_batching_manager", None)
+    assert cbm is not None, "ServeCommand has no running_continuous_batching_manager"
+    bp = getattr(cbm, "batch_processor", None)
+    assert bp is not None, "CBM has no batch_processor"
+    sched = getattr(bp, "scheduler", None)
+    assert sched is not None, "batch_processor has no scheduler"
+    return sched
+
+
+def _open_stream_and_cancel(base_url: str, request_id: str):
+    with requests.Session() as s:
+        with s.post(
+            f"{base_url}/v1/chat/completions",
+            json={
+                "model": "Qwen/Qwen2.5-0.5B-Instruct",
+                "stream": True,
+                "messages": [{"role": "user", "content": "Count slowly so I can cancel you."}],
+                "request_id": request_id,
+            },
+            stream=True,
+            timeout=30,
+        ) as resp:
+            assert resp.status_code == 200
+
+            for _ in resp.iter_content(chunk_size=None):
+                resp.close()
+                break
+
+
 @slow  # server startup time is slow on our push CI
 @require_openai
 class ServeCompletionsContinuousBatchingIntegrationTest(ServeCompletionsMixin, unittest.TestCase):
@@ -560,6 +592,33 @@ def test_max_tokens_not_set_in_req(self):
             )
         )
 
+    def test_request_cancellation(self):
+        """Tests that a request can be cancelled."""
+
+        base_url = f"http://127.0.0.1:{self.port}"
+        request_id = "test-cancel"
+
+        _open_stream_and_cancel(base_url, request_id)
+
+        scheduler = _get_scheduler(self.serve_command)
+
+        # Because cancellation is non-blocking, poll for a short, bounded time.
+        deadline = time.time() + 8.0  # generous but still CI-friendly
+        last_seen = None
+        while time.time() < deadline:
+            is_cancelled = scheduler.request_is_cancelled(request_id)
+            if is_cancelled:
+                break
+            last_seen = time.time()
+            time.sleep(0.1)  # don't spin the CPU
+
+        is_cancelled = scheduler.request_is_cancelled(request_id)
+        self.assertTrue(
+            is_cancelled,
+            f"Request {request_id} still present in scheduler after cancellation "
+            f"(last seen at {last_seen}). Check cancellation propagation.",
+        )
+
 
 @require_openai
 class ServeResponsesMixin: