[Bug]: Requests aren't aborted when client disconnects if user provided a middleware

[tomeras91]

In #7111, @njhill did a great job to make sure requests are aborted if they are cancelled by the client. This works great and can really ease the load on the model engine in cases of timeouts / other client disconnects.

Unfortunately, this mechanism doesn't work if a BaseHTTPMiddleware is added to the server, which is the case when a use provides a function to add as middleware (via the --middleware cmd arg of vllm serve). In this case, some weird starlette behavior causes request.is_disconnected() to return False, even if the client disconnects. This is discussed in Kludex/starlette#2094.
One of the suggested solutions in that discussion (for example, in Kludex/starlette#2094 (comment) and Kludex/starlette#2094 (comment)) is to change the mechanism used to check if the request was actually disconnected. Copied verbatim from that discussion, use this helper instead of request.is_disconnected():

async def my_is_disconnected(request: Request) -> bool:
    assert hasattr(request, '_is_disconnected'), "private API in starlette changed"
    assert isinstance(request._is_disconnected, bool), "private API in starlette changed"

    if request._is_disconnected:
        return True

    message: Message = {}

    # this timeout may need to be tweaked.
    # Ideally request.receive() is non-blocking, in which case the timeout doesn't matter.
    # But the ASGI spec seems to imply it should be a blocking callable, as it doesn't specify
    # what should be returned in case of no new messages. In this situation you can return an
    # empty dict, but given that it's against the spec it seems inadvisable.
    with anyio.move_on_after(0.01):
        message = await request.receive()

    if message.get("type") == "http.disconnect":
        request._is_disconnected = True

    return request._is_disconnected

In vLLM, the is_cancelled callable is passed to iterate_with_cancellation or to merge_async_iterators. In all cases, request.is_disconnected is used as that callable. So it should be pretty easy to replace it with the helper function instead. The only issue I see is with the timeout in anyio.move_on_after.

So WDYT? Is it possible to tweak the is_cancelled callable so requests will be aborted upon client disconnect also when the user provides a middleware function?

[njhill]

Thanks @tomeras91! This solution seems a bit hacky, I guess it would be much better to fix it in starlette if possible. At minimum we should benchmark such a solution under heavy load. Perhaps this could be used as the is_disconnected function only in cases where the server is run with additional middleware.

[tomeras91]

I agree the solution seems hacky :/ (that's why I opened an issue to talk about it rather than open a PR to fix).
I also agree it's best if this would be solved in starlette. Unfortunately, I don't think it will in the near future. This weird behavior has been present since starlette version 0.21.0 which is already almost 3 years old. If it wasn't fixed by now, it probably won't ever be. Looking at the relevant issues in their repo, it seems they don't really treat it as a bug for some reason.

Do you have an idea on how to benchmark the possible overhead of the solution? I tried running benchmark_throughput.py but got pretty inconsistent results..

[tomeras91]

PR #9594 added the add_request_id middleware to the OpenAI compatible API server. Following the starlette issue mentioned above, now request cancelled by the client are not aborted from the LLM engine for all use-cases, even without a user-defined added middleware.

I believe we need to find a solution, as this aborting of requests is a very nice feature and I guess we don't want to give up on it