[Bug Fix] Fix non-streaming request abort failure when --enable-metrics is enabled (#20625)
This commit is contained in:
@@ -1344,6 +1344,12 @@ def add_prometheus_track_response_middleware(app):
|
|||||||
)
|
)
|
||||||
)
|
)
|
||||||
|
|
||||||
|
# Fix: replace BaseHTTPMiddleware's call_next with a pure ASGI version
|
||||||
|
# that passes `receive` through, so request.is_disconnected() keeps working.
|
||||||
|
from sglang.srt.utils.http_middleware_patch import patch_app_http_middleware
|
||||||
|
|
||||||
|
patch_app_http_middleware(app)
|
||||||
|
|
||||||
@app.middleware("http")
|
@app.middleware("http")
|
||||||
async def track_http_status_code(request, call_next):
|
async def track_http_status_code(request, call_next):
|
||||||
# With recording all requests, we have the risk of high cardinality if requests have arbitrary unhandled paths.
|
# With recording all requests, we have the risk of high cardinality if requests have arbitrary unhandled paths.
|
||||||
|
|||||||
@@ -0,0 +1,67 @@
|
|||||||
|
"""
|
||||||
|
Fix @app.middleware("http") whose BaseHTTPMiddleware call_next replaces
|
||||||
|
ASGI ``receive``, breaking request.is_disconnected() and preventing
|
||||||
|
non-streaming request abort on client disconnect.
|
||||||
|
|
||||||
|
patch_app_http_middleware(app) replaces @app.middleware("http") with a
|
||||||
|
version whose call_next passes ``receive`` through untouched.
|
||||||
|
"""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
from starlette.requests import Request
|
||||||
|
|
||||||
|
|
||||||
|
class _SentResponse:
|
||||||
|
"""Response proxy returned after the real response was already sent."""
|
||||||
|
|
||||||
|
def __init__(self, status_code: int):
|
||||||
|
self.status_code = status_code
|
||||||
|
|
||||||
|
|
||||||
|
class _PureASGIDispatch:
|
||||||
|
"""Pure ASGI middleware providing a fixed call_next that passes
|
||||||
|
``receive`` through untouched (unlike BaseHTTPMiddleware)."""
|
||||||
|
|
||||||
|
def __init__(self, app, dispatch):
|
||||||
|
self.app = app
|
||||||
|
self.dispatch = dispatch
|
||||||
|
|
||||||
|
async def __call__(self, scope, receive, send):
|
||||||
|
if scope["type"] != "http":
|
||||||
|
await self.app(scope, receive, send)
|
||||||
|
return
|
||||||
|
|
||||||
|
request = Request(scope, receive)
|
||||||
|
status_code = 500
|
||||||
|
|
||||||
|
async def call_next(_request):
|
||||||
|
nonlocal status_code
|
||||||
|
|
||||||
|
async def send_and_capture(message):
|
||||||
|
nonlocal status_code
|
||||||
|
if message["type"] == "http.response.start":
|
||||||
|
status_code = message["status"]
|
||||||
|
await send(message)
|
||||||
|
|
||||||
|
await self.app(scope, receive, send_and_capture)
|
||||||
|
return _SentResponse(status_code)
|
||||||
|
|
||||||
|
await self.dispatch(request, call_next)
|
||||||
|
|
||||||
|
|
||||||
|
def patch_app_http_middleware(app):
|
||||||
|
"""Replace @app.middleware("http") with a fixed-call_next version."""
|
||||||
|
_orig = app.middleware
|
||||||
|
|
||||||
|
def _fixed(middleware_type):
|
||||||
|
if middleware_type == "http":
|
||||||
|
|
||||||
|
def decorator(fn):
|
||||||
|
app.add_middleware(_PureASGIDispatch, dispatch=fn)
|
||||||
|
return fn
|
||||||
|
|
||||||
|
return decorator
|
||||||
|
return _orig(middleware_type)
|
||||||
|
|
||||||
|
app.middleware = _fixed
|
||||||
@@ -0,0 +1,78 @@
|
|||||||
|
"""
|
||||||
|
Unit test for _PureASGIDispatch: verify that the ASGI ``receive`` callable
|
||||||
|
is passed through untouched so that request.is_disconnected() works.
|
||||||
|
|
||||||
|
Background: @app.middleware("http") wraps handlers with BaseHTTPMiddleware
|
||||||
|
whose call_next() replaces the ASGI ``receive``, breaking
|
||||||
|
request.is_disconnected() and preventing non-streaming abort on client
|
||||||
|
disconnect. _PureASGIDispatch fixes this. The existing test_abort.py
|
||||||
|
already covers the full e2e abort flow.
|
||||||
|
"""
|
||||||
|
|
||||||
|
import asyncio
|
||||||
|
import unittest
|
||||||
|
|
||||||
|
from starlette.requests import Request
|
||||||
|
|
||||||
|
from sglang.srt.utils.http_middleware_patch import _PureASGIDispatch
|
||||||
|
from sglang.test.ci.ci_register import register_cuda_ci
|
||||||
|
from sglang.test.test_utils import CustomTestCase
|
||||||
|
|
||||||
|
register_cuda_ci(est_time=10, suite="stage-a-cpu-only")
|
||||||
|
|
||||||
|
_HTTP_SCOPE = {
|
||||||
|
"type": "http",
|
||||||
|
"asgi": {"version": "3.0"},
|
||||||
|
"http_version": "1.1",
|
||||||
|
"method": "POST",
|
||||||
|
"path": "/test",
|
||||||
|
"query_string": b"",
|
||||||
|
"root_path": "",
|
||||||
|
"headers": [],
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
class TestPureASGIDispatchReceivePassthrough(CustomTestCase):
|
||||||
|
"""Verify _PureASGIDispatch passes ``receive`` through untouched."""
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
async def _run_with_receive(receive_msg):
|
||||||
|
"""Invoke _PureASGIDispatch and return request.is_disconnected()."""
|
||||||
|
result = {}
|
||||||
|
|
||||||
|
async def dispatch(request: Request, call_next):
|
||||||
|
result["disconnected"] = await request.is_disconnected()
|
||||||
|
await call_next(request)
|
||||||
|
|
||||||
|
async def inner_app(scope, receive, send):
|
||||||
|
await send({"type": "http.response.start", "status": 200, "headers": []})
|
||||||
|
await send({"type": "http.response.body", "body": b""})
|
||||||
|
|
||||||
|
middleware = _PureASGIDispatch(inner_app, dispatch=dispatch)
|
||||||
|
|
||||||
|
async def receive():
|
||||||
|
return receive_msg
|
||||||
|
|
||||||
|
sent = []
|
||||||
|
|
||||||
|
async def send(msg):
|
||||||
|
sent.append(msg)
|
||||||
|
|
||||||
|
await middleware(_HTTP_SCOPE, receive, send)
|
||||||
|
return result["disconnected"]
|
||||||
|
|
||||||
|
def test_is_disconnected_on_client_disconnect(self):
|
||||||
|
"""receive() -> http.disconnect: is_disconnected() must return True."""
|
||||||
|
self.assertTrue(
|
||||||
|
asyncio.run(self._run_with_receive({"type": "http.disconnect"}))
|
||||||
|
)
|
||||||
|
|
||||||
|
def test_not_disconnected_when_connected(self):
|
||||||
|
"""receive() -> http.request: is_disconnected() must return False."""
|
||||||
|
self.assertFalse(
|
||||||
|
asyncio.run(self._run_with_receive({"type": "http.request", "body": b""}))
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
unittest.main()
|
||||||
Reference in New Issue
Block a user