Skip to content

[None][fix] Use decorator for request cancelation and handle CancelledError #5559

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
pathorn wants to merge 1 commit into from
Closed

[None][fix] Use decorator for request cancelation and handle CancelledError #5559

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
109 changes: 103 additions & 6 deletions tensorrt_llm/serve/openai_server.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,17 +7,20 @@
from collections import deque
from contextlib import asynccontextmanager
from datetime import datetime
from functools import wraps
from http import HTTPStatus
from pathlib import Path
from typing import (Annotated, Any, AsyncGenerator, AsyncIterator, List,
Optional, Union)

import uvicorn
from fastapi import Body, FastAPI, Request
from fastapi.exceptions import RequestValidationError
from fastapi.exceptions import RequestValidationError, HTTPException
from fastapi.responses import JSONResponse, Response, StreamingResponse
from starlette.routing import Mount
from transformers import AutoConfig, AutoProcessor
from openai.types.chat import ChatCompletionMessageParam
from pydantic import BaseModel

from tensorrt_llm._tensorrt_engine import LLM
# yapf: disable
Expand Down Expand Up @@ -69,6 +72,72 @@
TIMEOUT_KEEP_ALIVE = 5 # seconds.


async def disconnect_poller(request: Request, result: Any):
"""
Poll for a disconnect.
If the request disconnects, stop polling and return.
"""
try:
while True:
message = await request.receive()
if message["type"] == "http.disconnect":
break

print("Request disconnected")

return result
except asyncio.CancelledError:
print("Stopping polling loop")


def cancel_on_disconnect(model_type: Type[BaseModel]):
"""
Decorator that will check if the client disconnects,
and cancel the task if required.
"""

def cancel_on_disconnect_inner(handler: Callable):

@wraps(handler)
async def cancel_on_disconnect_decorator(self, request: model_type, raw_request: Request):
sentinel = object()

# Create two tasks, one to poll the request and check if the
# client disconnected, and another which is the request handler
poller_task = asyncio.ensure_future(disconnect_poller(raw_request, sentinel))
handler_task = asyncio.ensure_future(handler(self, request=request, raw_request=raw_request))

done, pending = await asyncio.wait(
[poller_task, handler_task], return_when=asyncio.FIRST_COMPLETED
)

Comment on lines +105 to +113

@coderabbitai coderabbitai Bot Oct 3, 2025

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🔴 Critical

Guard decorator against raw_request being None

Routes such as health_generate invoke openai_chat with raw_request=None. The new decorator always schedules disconnect_poller(raw_request, ...), so these calls now raise AttributeError (None has no receive) before reaching the handler. This regresses the health endpoint and any other internal callers without a raw Request. Add a fast path that skips the poller when raw_request is missing.

-            poller_task = asyncio.ensure_future(disconnect_poller(raw_request, sentinel))
-            handler_task = asyncio.ensure_future(handler(self, request=request, raw_request=raw_request))
+            if raw_request is None:
+                return await handler(self, request=request, raw_request=raw_request)
+
+            poller_task = asyncio.ensure_future(disconnect_poller(raw_request, sentinel))
+            handler_task = asyncio.ensure_future(handler(self, request=request, raw_request=raw_request))
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
# Create two tasks, one to poll the request and check if the
# client disconnected, and another which is the request handler
poller_task = asyncio.ensure_future(disconnect_poller(raw_request, sentinel))
handler_task = asyncio.ensure_future(handler(self, request=request, raw_request=raw_request))
done, pending = await asyncio.wait(
[poller_task, handler_task], return_when=asyncio.FIRST_COMPLETED
)
# Create two tasks, one to poll the request and check if the
# client disconnected, and another which is the request handler
if raw_request is None:
return await handler(self, request=request, raw_request=raw_request)
poller_task = asyncio.ensure_future(disconnect_poller(raw_request, sentinel))
handler_task = asyncio.ensure_future(handler(self, request=request, raw_request=raw_request))
done, pending = await asyncio.wait(
[poller_task, handler_task], return_when=asyncio.FIRST_COMPLETED
)
🤖 Prompt for AI Agents 
In tensorrt_llm/serve/openai_server.py around lines 102 to 110, the decorator
always schedules disconnect_poller(raw_request, ...) even when raw_request is
None (used by internal routes like health_generate), causing AttributeError;
modify the decorator to fast-path when raw_request is falsy by not creating or
scheduling the poller and only awaiting the handler task (or awaiting handler
immediately), otherwise create poller_task as before; ensure sentinel logic and
cleanup still occur when poller is skipped so handler cancellation/cleanup
remains correct.

# Cancel any outstanding tasks
for t in pending:
t.cancel()

try:
await t
except asyncio.CancelledError:
print(f"{t} was cancelled")
except Exception as exc:
print(f"{t} raised {exc} when being cancelled")

# Return the result if the handler finished first
if handler_task in done:
return await handler_task

# Otherwise, raise an exception
# This is not exactly needed, but it will prevent
# validation errors if your request handler is supposed
# to return something.
print("Raising an HTTP error because I was disconnected!!")

raise HTTPException(503)

return cancel_on_disconnect_decorator

return cancel_on_disconnect_inner

class OpenAIServer:

def __init__(self,
Expand Down Expand Up @@ -407,8 +476,12 @@ async def _extract_metrics(self, res: RequestOutput, raw_request: Request):
async with self.perf_metrics_lock:
self.perf_metrics.append(item)


@cancel_on_disconnect(ChatCompletionRequest)
async def openai_chat(self, request: ChatCompletionRequest, raw_request: Request) -> Response:

did_complete = False

def get_role() -> str:
if request.add_generation_prompt:
role = "assistant"
Expand All @@ -418,6 +491,7 @@ def get_role() -> str:

async def chat_stream_generator(
promise: RequestOutput, postproc_params: PostprocParams) -> AsyncGenerator[str, None]:
nonlocal did_complete
try:
if not self.postproc_worker_enabled:
post_processor, args = postproc_params.post_processor, postproc_params.postproc_args
Expand All @@ -431,13 +505,21 @@ async def chat_stream_generator(
async for res in promise:
pp_results = res.outputs[0]._postprocess_result if self.postproc_worker_enabled else post_processor(res, args)
for pp_res in pp_results:
yield pp_res
for choice in pp_res.choices:
if choice.finish_reason is not None:
did_complete = True

pp_res_json = pp_res.model_dump_json(exclude_unset=True)
yield f"data: {pp_res_json}\n\n"
yield "data: [DONE]\n\n"
await self._extract_metrics(res, raw_request)
nvtx_mark("generation ends")
except:
logger.error(traceback.format_exc())
raise
finally:
if not did_complete:
promise.abort()

async def create_chat_response(
promise: RequestOutput, postproc_params: PostprocParams, disaggregated_params: Optional[LlmDisaggregatedParams] = None) -> ChatCompletionResponse:
Expand All @@ -455,6 +537,7 @@ async def create_chat_response(
await self._extract_metrics(promise, raw_request)
return chat_response

promise: Optional[RequestOutput] = None
try:
check_multiple_response(request.n, self.llm.args.backend)
conversation: List[ConversationMessage] = []
Expand Down Expand Up @@ -512,7 +595,6 @@ async def create_chat_response(
disaggregated_params=disaggregated_params,
cache_salt=request.cache_salt,
)
asyncio.create_task(self.await_disconnected(raw_request, promise))
if not self.postproc_worker_enabled:
postproc_args.tokenizer = self.tokenizer
postproc_args.num_prompt_tokens = len(promise.prompt_token_ids)
Expand All @@ -528,6 +610,10 @@ async def create_chat_response(
logger.error(traceback.format_exc())
# If internal executor error is raised, shutdown the server
signal.raise_signal(signal.SIGINT)
except asyncio.CancelledError:
if promise is not None:
promise.abort()
return self.create_error_response("cancelled")
except Exception as e:
logger.error(traceback.format_exc())
return self.create_error_response(str(e))
Expand Down Expand Up @@ -608,6 +694,7 @@ async def create_mm_embedding_response(promise: RequestOutput):
logger.error(traceback.format_exc())
return self.create_error_response(str(e))

@cancel_on_disconnect(CompletionRequest)
async def openai_completion(self, request: CompletionRequest, raw_request: Request) -> Response:

async def completion_response(promise: RequestOutput,
Expand Down Expand Up @@ -652,6 +739,7 @@ def merge_completion_responses(responses: List[CompletionResponse]) -> Completio
return merged_rsp

async def completion_generator(promise: RequestOutput, params: Optional[PostprocParams]):
did_complete = False
try:
async for output in promise:
if not self.postproc_worker_enabled:
Expand All @@ -660,12 +748,18 @@ async def completion_generator(promise: RequestOutput, params: Optional[Postproc
else:
pp_result = output.outputs[0]._postprocess_result
for pp_res in pp_result:
yield pp_res
for choice in pp_res.choices:
if choice.finish_reason is not None:
did_complete = True
pp_res_json = pp_res.model_dump_json(exclude_unset=True)
yield f"data: {pp_res_json}\n\n"
await self._extract_metrics(output, raw_request)
except:
logger.error(traceback.format_exc())
raise

finally:
if not did_complete:
promise.abort()

async def merge_generators(generators: List[AsyncIterator[Any]]):
result_queue = asyncio.Queue()
Expand Down Expand Up @@ -737,7 +831,6 @@ async def generator_wrapper(generator: AsyncIterator[Any]):
lora_request=request.lora_request,
disaggregated_params=disaggregated_params
)
asyncio.create_task(self.await_disconnected(raw_request, promise))
if not self.postproc_worker_enabled:
postproc_args.tokenizer = self.tokenizer
postproc_args.num_prompt_tokens = len(promise.prompt_token_ids)
Expand All @@ -759,6 +852,10 @@ async def generator_wrapper(generator: AsyncIterator[Any]):
logger.error(traceback.format_exc())
# If internal executor error is raised, shutdown the server
signal.raise_signal(signal.SIGINT)
except asyncio.CancelledError:
for promise in promises:
promise.abort()
return self.create_error_response("cancelled")
except Exception as e:
logger.error(traceback.format_exc())
return self.create_error_response(str(e))
Expand Down
48 changes: 26 additions & 22 deletions tensorrt_llm/serve/postprocess_handlers.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -123,8 +123,9 @@ def chat_stream_post_processor(rsp: GenerationResultBase,

def yield_first_chat(num_tokens: int,
idx: int,
rsp: GenerationResultBase,
role: str = None,
content: str = None):
content: str = None) -> ChatCompletionStreamResponse:
choice_data = ChatCompletionResponseStreamChoice(index=idx,
delta=DeltaMessage(
role=role,
Expand All @@ -135,11 +136,12 @@ def yield_first_chat(num_tokens: int,
if include_continuous_usage:
chunk.usage = UsageInfo(prompt_tokens=num_tokens,
total_tokens=num_tokens,
completion_tokens=0)
data = chunk.model_dump_json(exclude_none=True)
return data
completion_tokens=0,
prompt_tokens_details=PromptTokensDetails(cached_tokens=min(num_tokens, (getattr(rsp, 'num_reused_blocks', 0) or 0) * 32))
)
return chunk

res: List[str] = []
res: List[ChatCompletionStreamResponse] = []
finish_reason_sent = [False] * args.num_choices
prompt_tokens = args.num_prompt_tokens
if stream_option := args.stream_options:
Expand All @@ -151,11 +153,11 @@ def yield_first_chat(num_tokens: int,
if args.first_iteration:
for i in range(args.num_choices):
res.append(
f"data: {yield_first_chat(prompt_tokens, i, role=args.role)} \n\n"
yield_first_chat(prompt_tokens, i, rsp, role=args.role)
)
if args.echo and args.last_message_content:
res.append(
f"data: {yield_first_chat(prompt_tokens, i, content=args.last_message_content)} \n\n"
yield_first_chat(prompt_tokens, i, rsp, content=args.last_message_content)
)
args.first_iteration = False

Expand Down Expand Up @@ -204,23 +206,24 @@ def yield_first_chat(num_tokens: int,
if include_continuous_usage:
chunk.usage = UsageInfo(prompt_tokens=prompt_tokens,
completion_tokens=output.length,
total_tokens=output.length + prompt_tokens)
data = chunk.model_dump_json(exclude_none=True)
res.append(f"data: {data}\n\n")
total_tokens=output.length + prompt_tokens,
prompt_tokens_details=PromptTokensDetails(cached_tokens=min(prompt_tokens, (getattr(rsp, 'num_reused_blocks', 0) or 0) * 32))
)
res.append(chunk)

if include_usage and rsp._done:
completion_tokens = sum(output.length for output in rsp.outputs)
final_usage = UsageInfo(
prompt_tokens=prompt_tokens,
completion_tokens=completion_tokens,
total_tokens=prompt_tokens + completion_tokens,
prompt_tokens_details=PromptTokensDetails(cached_tokens=min(prompt_tokens, (getattr(rsp, 'num_reused_blocks', 0) or 0) * 32))
)

final_usage_chunk = ChatCompletionStreamResponse(choices=[],
model=args.model,
usage=final_usage)
final_usage_data = final_usage_chunk.model_dump_json()
res.append(f"data: {final_usage_data}\n\n")
res.append(final_usage_chunk)
return res


Expand Down Expand Up @@ -279,6 +282,7 @@ def chat_response_post_processor(
prompt_tokens=num_prompt_tokens,
completion_tokens=num_generated_tokens,
total_tokens=num_prompt_tokens + num_generated_tokens,
prompt_tokens_details=PromptTokensDetails(cached_tokens=min(num_prompt_tokens, (getattr(rsp, 'num_reused_blocks', 0) or 0) * 32))
)
response = ChatCompletionResponse(
model=args.model,
Expand Down Expand Up @@ -339,23 +343,24 @@ def completion_stream_post_processor(rsp: DetokenizedGenerationResultBase,
if include_continuous_usage:
chunk.usage = UsageInfo(prompt_tokens=prompt_tokens,
completion_tokens=output.length,
total_tokens=output.length + prompt_tokens)
data = chunk.model_dump_json(exclude_unset=False)
res.append(f"data: {data}\n\n")
total_tokens=output.length + prompt_tokens,
prompt_tokens_details=PromptTokensDetails(cached_tokens=min(prompt_tokens, (getattr(rsp, 'num_reused_blocks', 0) or 0) * 32))
)
res.append(chunk)

if include_usage and rsp._done:
completion_tokens = sum(output.length for output in rsp.outputs)
final_usage = UsageInfo(
prompt_tokens=prompt_tokens,
completion_tokens=completion_tokens,
total_tokens=prompt_tokens + completion_tokens,
prompt_tokens_details=PromptTokensDetails(cached_tokens=min(prompt_tokens, (getattr(rsp, 'num_reused_blocks', 0) or 0) * 32))
)

final_usage_chunk = ChatCompletionStreamResponse(choices=[],
final_usage_chunk = CompletionStreamResponse(choices=[],
model=args.model,
usage=final_usage)
final_usage_data = final_usage_chunk.model_dump_json()
res.append(f"data: {final_usage_data}\n\n")
res.append(final_usage_chunk)
args.first_iteration = False
return res

Expand Down Expand Up @@ -392,10 +397,9 @@ def completion_response_post_processor(

usage = UsageInfo(prompt_tokens=prompt_tokens,
completion_tokens=completion_tokens,
total_tokens=completion_tokens + prompt_tokens)
response = CompletionResponse(choices=choices,
model=args.model,
usage=usage)
total_tokens=completion_tokens + prompt_tokens,
prompt_tokens_details=PromptTokensDetails(cached_tokens=min(prompt_tokens, (getattr(rsp, 'num_reused_blocks', 0) or 0) * 32)))
response = CompletionResponse(choices=choices, model=args.model, usage=usage)
return response


Expand Down
Loading