vrr/unsloth
2
0
Fork 0
mirror of https://github.com/unslothai/unsloth.git synced 2026-07-09 15:58:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 29

Studio: flush passthrough stream headers before upstream prefill stalls (#6835)

Browse source 
* Studio: flush passthrough stream headers before upstream prefill stalls

* Studio: clean up delayed passthrough send failures

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio: close passthrough preheader cleanup gaps

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio: retry delayed passthrough overflow truncation

* Studio: close completed passthrough send responses

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com>
This commit is contained in:
Rod Boev 2026-07-03 11:26:38 -04:00 • committed by GitHub
parent 026141a4a4
commit fbb5b0968c
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 669 additions and 10 deletions
Download patch file Download diff file Expand all files Collapse all files

132
studio/backend/routes/inference.py
View file

@ -784,6 +784,7 @@ def _set_stream_response_read_timeout(
_STREAM_DISCONNECT_POLL_TIMEOUT_S = 0.25
_OPENAI_PASSTHROUGH_PREHEADER_STATUS_WINDOW_S = 0.1
class _CompatSameTaskTimeout:
@ -10890,41 +10891,73 @@ async def _openai_passthrough_stream(
_cancel_keys = (payload.cancel_id, payload.session_id, completion_id)
_tracker = _TrackedCancel(cancel_event, *_cancel_keys)
_tracker.__enter__()
client = None
resp = None
send_task: Optional[asyncio.Task[Optional[httpx.Response]]] = None
async def _aclose_send_task(task: Optional[asyncio.Task[Optional[httpx.Response]]]) -> None:
if task is None:
return
if not task.done():
task.cancel()
try:
task_resp = await task
if task_resp is not None:
try:
await task_resp.aclose()
except Exception:
pass
except (asyncio.CancelledError, Exception):
pass
# Keep tracker cleanup paired if pre-header dispatch is cancelled.
try:
# Dispatch BEFORE returning StreamingResponse so transport errors and
# non-200 upstream statuses surface as real HTTP errors -- OpenAI SDKs
# rely on status codes to raise APIError/BadRequestError.
# Keep the pre-header window short so accepted SSE clients receive
# immediate headers in the common timeout-reduced stall.
client = httpx.AsyncClient(
timeout = _llama_streaming_generation_timeout(),
limits = httpx.Limits(max_keepalive_connections = 0),
trust_env = False,
)
resp = None
_truncate_budget = (
_OVERFLOW_TRUNCATE_MAX_RETRIES if _overflow_truncation_requested(payload) else 0
)
while True:
try:
req = client.build_request(
"POST", target_url, json = body, headers = {"Connection": "close"}
)
first_token_deadline = time.monotonic() + _DEFAULT_FIRST_TOKEN_TIMEOUT_S
resp = await _send_stream_with_preheader_cancel(
client, req, cancel_event, request = request
send_task = asyncio.create_task(
_send_stream_with_preheader_cancel(client, req, cancel_event, request = request)
)
done, _ = await asyncio.wait(
{send_task},
timeout = _OPENAI_PASSTHROUGH_PREHEADER_STATUS_WINDOW_S,
return_when = asyncio.FIRST_COMPLETED,
)
if send_task not in done:
break
# Dispatch returned quickly enough to preserve pre-header status.
resp = await send_task
send_task = None
except httpx.RequestError as e:
# llama-server subprocess crashed / starting / unreachable.
logger.error("openai passthrough stream: upstream unreachable: %s", e)
api_monitor.fail(monitor_id, _friendly_error(e))
await _aclose_send_task(send_task)
await _aclose_stream_resources(resp = resp, client = client)
raise HTTPException(
status_code = 502,
detail = _friendly_error(e),
)
if resp is None and send_task is not None and not send_task.done():
break
if resp is None:
api_monitor.finish(monitor_id, "cancelled")
await _aclose_send_task(send_task)
try:
await client.aclose()
except Exception:
@ -10969,6 +11002,8 @@ async def _openai_passthrough_stream(
api_monitor.fail(monitor_id, err_text[:500])
raise _openai_passthrough_error(upstream_status, err_text)
# Keep tracker cleanup paired if pre-header dispatch is cancelled after we
# have already committed headers.
async def _stream():
# Same httpx lifecycle pattern as _anthropic_passthrough_stream:
# save resp.aiter_lines() so the finally block can aclose() it on
@ -10976,10 +11011,10 @@ async def _openai_passthrough_stream(
lines_iter = None
# Watchers unblock aiter_lines() during prefill, before in-loop
# cancel/disconnect checks can run.
cancel_watcher = asyncio.create_task(_await_cancel_then_close(cancel_event, resp))
disconnect_watcher = asyncio.create_task(
_await_disconnect_then_close(request, resp, cancel_event)
)
cancel_watcher = None
disconnect_watcher = None
nonlocal resp, send_task, first_token_deadline, _truncate_budget
monitor_done = False
saw_finish_reason = False
saw_done = False
@ -11112,6 +11147,79 @@ async def _openai_passthrough_stream(
return lines
try:
while True:
if send_task is not None and not send_task.done():
try:
resp = await send_task
except httpx.RequestError as e:
logger.error("openai passthrough stream: upstream unreachable: %s", e)
api_monitor.fail(monitor_id, _friendly_error(e))
yield f"data: {json.dumps(_openai_stream_error_chunk(e))}\n\n"
return
send_task = None
elif send_task is not None:
try:
resp = send_task.result()
except httpx.RequestError as e:
logger.error("openai passthrough stream: upstream unreachable: %s", e)
api_monitor.fail(monitor_id, _friendly_error(e))
yield f"data: {json.dumps(_openai_stream_error_chunk(e))}\n\n"
return
send_task = None
if resp is None:
api_monitor.finish(monitor_id, "cancelled")
return
if resp.status_code == 200:
break
err_bytes = await resp.aread()
err_text = err_bytes.decode("utf-8", errors = "replace")
logger.error(
"openai passthrough upstream error: status=%s body=%s",
resp.status_code,
err_text[:500],
)
upstream_status = resp.status_code
try:
await resp.aclose()
except Exception:
pass
resp = None
if (
_truncate_budget > 0
and _classify_llama_generation_error(Exception(err_text))
and _apply_overflow_truncation(body, err_text)
):
_truncate_budget -= 1
req = client.build_request(
"POST", target_url, json = body, headers = {"Connection": "close"}
)
first_token_deadline = time.monotonic() + _DEFAULT_FIRST_TOKEN_TIMEOUT_S
send_task = asyncio.create_task(
_send_stream_with_preheader_cancel(
client, req, cancel_event, request = request
)
)
continue
upstream_error = _openai_passthrough_error(upstream_status, err_text)
error_payload = (
upstream_error.detail
if isinstance(upstream_error.detail, dict)
else openai_error_body(
str(upstream_error.detail),
status = upstream_status,
)
)
api_monitor.fail(monitor_id, err_text[:500])
yield f"data: {json.dumps(error_payload)}\n\n"
return
cancel_watcher = asyncio.create_task(_await_cancel_then_close(cancel_event, resp))
disconnect_watcher = asyncio.create_task(
_await_disconnect_then_close(request, resp, cancel_event)
)
lines_iter = resp.aiter_lines()
async for raw_line in _aiter_llama_stream_items(
lines_iter,
@ -11298,6 +11406,7 @@ async def _openai_passthrough_stream(
err = _openai_stream_error_chunk(e)
yield f"data: {json.dumps(err)}\n\n"
finally:
await _aclose_send_task(send_task)
await _aclose_stream_resources(
watchers = (cancel_watcher, disconnect_watcher),
iterator = lines_iter,
@ -11311,6 +11420,7 @@ async def _openai_passthrough_stream(
# finally never ran. Release the eagerly-opened upstream resp/client
# and the cancel-registry entry here; the watchers and line iterator
# are created inside _stream(), so there is nothing else to close.
await _aclose_send_task(send_task)
await _aclose_stream_resources(resp = resp, client = client)
_tracker.__exit__(None, None, None)
@ -11325,6 +11435,8 @@ async def _openai_passthrough_stream(
unstarted_cleanup = _unstarted_cleanup,
)
except BaseException:
await _aclose_send_task(send_task)
await _aclose_stream_resources(resp = resp, client = client)
_tracker.__exit__(None, None, None)
raise

547
studio/backend/tests/test_openai_tool_passthrough.py
View file

@ -1856,6 +1856,553 @@ class TestApiMonitorProviderAndCompletionStreams:
chunks = [chunk async for chunk in response.body_iterator]
return SimpleNamespace(chunks = chunks, body = "".join(chunks), monitor = monitor)
def test_passthrough_stream_preheader_dispatched_with_timeout(self, monkeypatch):
async def _run():
import routes.inference as inf_mod
gate = asyncio.Event()
async def fake_send(*_args, **_kwargs):
await gate.wait()
return httpx.Response(200, content = b"")
class Request:
async def is_disconnected(self):
return False
monitor = ApiMonitor(max_entries = 3)
monitor_id = monitor.start(
endpoint = "/v1/chat/completions",
method = "POST",
model = "gguf",
prompt = "hi",
)
monkeypatch.setattr(inf_mod, "api_monitor", monitor)
monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fake_send)
payload = ChatCompletionRequest(
model = "default",
messages = [ChatMessage(role = "user", content = "hi")],
stream = True,
)
response = await asyncio.wait_for(
_openai_passthrough_stream(
Request(),
threading.Event(),
SimpleNamespace(
base_url = "http://llama.test",
context_length = 4096,
_request_reasoning_kwargs = lambda *_args, **_kwargs: None,
),
payload,
"chatcmpl-test",
"chatcmpl-test",
monitor_id = monitor_id,
),
timeout = 0.2,
)
assert isinstance(response, _SameTaskStreamingResponse)
gate.set()
chunks = [
chunk.decode() if isinstance(chunk, bytes) else chunk
async for chunk in response.body_iterator
]
assert "data: [DONE]\n\n" in "".join(chunks)
asyncio.run(_run())
def test_passthrough_stream_preheader_non_200_in_window(self, monkeypatch):
async def _run():
import routes.inference as inf_mod
async def fake_send(*_args, **_kwargs):
return httpx.Response(400, content = b'{"error":"bad"}')
class Request:
async def is_disconnected(self):
return False
monitor = ApiMonitor(max_entries = 3)
monitor_id = monitor.start(
endpoint = "/v1/chat/completions",
method = "POST",
model = "gguf",
prompt = "hi",
)
monkeypatch.setattr(inf_mod, "api_monitor", monitor)
monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fake_send)
payload = ChatCompletionRequest(
model = "default",
messages = [ChatMessage(role = "user", content = "hi")],
stream = True,
)
with pytest.raises(HTTPException) as exc:
await _openai_passthrough_stream(
Request(),
threading.Event(),
SimpleNamespace(
base_url = "http://llama.test",
context_length = 4096,
_request_reasoning_kwargs = lambda *_args, **_kwargs: None,
),
payload,
"chatcmpl-test",
"chatcmpl-test",
monitor_id = monitor_id,
)
assert exc.value.status_code == 400
asyncio.run(_run())
def test_passthrough_stream_preheader_request_error_in_window(self, monkeypatch):
async def _run():
import routes.inference as inf_mod
async def fake_send(*_args, **_kwargs):
raise httpx.ConnectError("connectivity issue")
class Request:
async def is_disconnected(self):
return False
monitor = ApiMonitor(max_entries = 3)
monitor_id = monitor.start(
endpoint = "/v1/chat/completions",
method = "POST",
model = "gguf",
prompt = "hi",
)
monkeypatch.setattr(inf_mod, "api_monitor", monitor)
monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fake_send)
payload = ChatCompletionRequest(
model = "default",
messages = [ChatMessage(role = "user", content = "hi")],
stream = True,
)
with pytest.raises(HTTPException) as exc:
await _openai_passthrough_stream(
Request(),
threading.Event(),
SimpleNamespace(
base_url = "http://llama.test",
context_length = 4096,
_request_reasoning_kwargs = lambda *_args, **_kwargs: None,
),
payload,
"chatcmpl-test",
"chatcmpl-test",
monitor_id = monitor_id,
)
assert exc.value.status_code == 502
asyncio.run(_run())
def test_passthrough_stream_preheader_delayed_non_200_returns_sse_error(self, monkeypatch):
async def _run():
import routes.inference as inf_mod
gate = asyncio.Event()
async def fake_send(*_args, **_kwargs):
await gate.wait()
return httpx.Response(400, content = b'{"error":"bad"}')
class Request:
async def is_disconnected(self):
return False
monitor = ApiMonitor(max_entries = 3)
monitor_id = monitor.start(
endpoint = "/v1/chat/completions",
method = "POST",
model = "gguf",
prompt = "hi",
)
monkeypatch.setattr(inf_mod, "api_monitor", monitor)
monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fake_send)
payload = ChatCompletionRequest(
model = "default",
messages = [ChatMessage(role = "user", content = "hi")],
stream = True,
)
response = await asyncio.wait_for(
_openai_passthrough_stream(
Request(),
threading.Event(),
SimpleNamespace(
base_url = "http://llama.test",
context_length = 4096,
_request_reasoning_kwargs = lambda *_args, **_kwargs: None,
),
payload,
"chatcmpl-test",
"chatcmpl-test",
monitor_id = monitor_id,
),
timeout = 0.2,
)
assert isinstance(response, _SameTaskStreamingResponse)
gate.set()
chunks = [
chunk.decode() if isinstance(chunk, bytes) else chunk
async for chunk in response.body_iterator
]
body = "".join(chunks)
assert "data:" in body
assert '"error"' in body
[entry] = monitor.snapshot()
assert entry["status"] == "error"
assert "bad" in entry["error"]
asyncio.run(_run())
def test_passthrough_stream_preheader_delayed_context_error_keeps_error_envelope(
self, monkeypatch
):
async def _run():
import routes.inference as inf_mod
gate = asyncio.Event()
ctx_msg = "request (4096 tokens) exceeds the available context size (2048 tokens)"
async def fake_send(*_args, **_kwargs):
await gate.wait()
return httpx.Response(400, content = ctx_msg.encode("utf-8"))
class Request:
async def is_disconnected(self):
return False
monitor = ApiMonitor(max_entries = 3)
monitor_id = monitor.start(
endpoint = "/v1/chat/completions",
method = "POST",
model = "gguf",
prompt = "hi",
)
monkeypatch.setattr(inf_mod, "api_monitor", monitor)
monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fake_send)
payload = ChatCompletionRequest(
model = "default",
messages = [ChatMessage(role = "user", content = "hi")],
stream = True,
)
response = await asyncio.wait_for(
_openai_passthrough_stream(
Request(),
threading.Event(),
SimpleNamespace(
base_url = "http://llama.test",
context_length = 2048,
_request_reasoning_kwargs = lambda *_args, **_kwargs: None,
),
payload,
"chatcmpl-test",
"chatcmpl-test",
monitor_id = monitor_id,
),
timeout = 0.2,
)
assert isinstance(response, _SameTaskStreamingResponse)
gate.set()
chunks = [
chunk.decode() if isinstance(chunk, bytes) else chunk
async for chunk in response.body_iterator
]
body = "".join(chunks)
payload = json.loads(body.removeprefix("data: ").strip())
assert payload["error"]["code"] == "context_length_exceeded"
assert payload["error"]["param"] == "messages"
assert isinstance(payload["error"], dict)
asyncio.run(_run())
def test_passthrough_stream_preheader_delayed_context_error_retries_truncation(
self, monkeypatch
):
async def _run():
import routes.inference as inf_mod
gate = asyncio.Event()
calls = []
err_body = json.dumps(
{
"error": {
"message": "request (10000 tokens) exceeds the available context size (2048 tokens)",
"n_prompt_tokens": 10000,
"n_ctx": 2048,
}
}
).encode("utf-8")
async def fake_send(_client, req, *_args, **_kwargs):
calls.append(json.loads(req.content.decode("utf-8")))
if len(calls) == 1:
await gate.wait()
return httpx.Response(400, content = err_body)
return httpx.Response(200, content = b"")
class Request:
async def is_disconnected(self):
return False
monitor = ApiMonitor(max_entries = 3)
monitor_id = monitor.start(
endpoint = "/v1/chat/completions",
method = "POST",
model = "gguf",
prompt = "hi",
)
monkeypatch.setattr(inf_mod, "api_monitor", monitor)
monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fake_send)
messages = [
ChatMessage(role = "system", content = "system"),
*[
ChatMessage(role = "user", content = f"turn {idx} " + ("x" * 1000))
for idx in range(8)
],
]
payload = ChatCompletionRequest(
model = "default",
messages = messages,
stream = True,
context_overflow = "truncate_middle",
)
response = await asyncio.wait_for(
_openai_passthrough_stream(
Request(),
threading.Event(),
SimpleNamespace(
base_url = "http://llama.test",
context_length = 2048,
_request_reasoning_kwargs = lambda *_args, **_kwargs: None,
),
payload,
"chatcmpl-test",
"chatcmpl-test",
monitor_id = monitor_id,
),
timeout = 0.2,
)
assert isinstance(response, _SameTaskStreamingResponse)
gate.set()
chunks = [
chunk.decode() if isinstance(chunk, bytes) else chunk
async for chunk in response.body_iterator
]
assert "data: [DONE]\n\n" in "".join(chunks)
assert len(calls) == 2
assert len(calls[1]["messages"]) < len(calls[0]["messages"])
[entry] = monitor.snapshot()
assert entry["status"] == "completed"
asyncio.run(_run())
def test_passthrough_stream_preheader_delayed_request_error_cleans_up(self, monkeypatch):
async def _run():
import routes.inference as inf_mod
gate = asyncio.Event()
cancel_id = "delayed-request-error-cancel"
async def fake_send(*_args, **_kwargs):
await gate.wait()
raise httpx.ConnectError("delayed connectivity issue")
class Request:
async def is_disconnected(self):
return False
monitor = ApiMonitor(max_entries = 3)
monitor_id = monitor.start(
endpoint = "/v1/chat/completions",
method = "POST",
model = "gguf",
prompt = "hi",
)
monkeypatch.setattr(inf_mod, "api_monitor", monitor)
monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fake_send)
payload = ChatCompletionRequest(
model = "default",
messages = [ChatMessage(role = "user", content = "hi")],
stream = True,
cancel_id = cancel_id,
)
response = await asyncio.wait_for(
_openai_passthrough_stream(
Request(),
threading.Event(),
SimpleNamespace(
base_url = "http://llama.test",
context_length = 4096,
_request_reasoning_kwargs = lambda *_args, **_kwargs: None,
),
payload,
"chatcmpl-test",
"chatcmpl-test",
monitor_id = monitor_id,
),
timeout = 0.2,
)
assert isinstance(response, _SameTaskStreamingResponse)
assert cancel_id in inf_mod._CANCEL_REGISTRY
gate.set()
chunks = [
chunk.decode() if isinstance(chunk, bytes) else chunk
async for chunk in response.body_iterator
]
body = "".join(chunks)
assert "data:" in body
assert '"error"' in body
[entry] = monitor.snapshot()
assert entry["status"] == "error"
assert "Lost connection" in entry["error"]
assert cancel_id not in inf_mod._CANCEL_REGISTRY
asyncio.run(_run())
def test_passthrough_stream_preheader_cancel_cleans_pending_send(self, monkeypatch):
async def _run():
import routes.inference as inf_mod
entered = asyncio.Event()
cancelled = asyncio.Event()
cancel_id = "preheader-cancel-cleanup"
async def fake_send(*_args, **_kwargs):
entered.set()
try:
await asyncio.Event().wait()
except asyncio.CancelledError:
cancelled.set()
raise
class Request:
async def is_disconnected(self):
return False
monitor = ApiMonitor(max_entries = 3)
monitor_id = monitor.start(
endpoint = "/v1/chat/completions",
method = "POST",
model = "gguf",
prompt = "hi",
)
monkeypatch.setattr(inf_mod, "api_monitor", monitor)
monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fake_send)
payload = ChatCompletionRequest(
model = "default",
messages = [ChatMessage(role = "user", content = "hi")],
stream = True,
cancel_id = cancel_id,
)
task = asyncio.create_task(
_openai_passthrough_stream(
Request(),
threading.Event(),
SimpleNamespace(
base_url = "http://llama.test",
context_length = 4096,
_request_reasoning_kwargs = lambda *_args, **_kwargs: None,
),
payload,
"chatcmpl-test",
"chatcmpl-test",
monitor_id = monitor_id,
)
)
await asyncio.wait_for(entered.wait(), timeout = 0.2)
assert cancel_id in inf_mod._CANCEL_REGISTRY
task.cancel()
with pytest.raises(asyncio.CancelledError):
await task
await asyncio.wait_for(cancelled.wait(), timeout = 0.2)
assert cancel_id not in inf_mod._CANCEL_REGISTRY
asyncio.run(_run())
def test_passthrough_stream_unstarted_cleanup_closes_completed_send_response(self, monkeypatch):
async def _run():
import routes.inference as inf_mod
gate = asyncio.Event()
returned = asyncio.Event()
cancel_id = "unstarted-completed-send-cleanup"
class Stream(httpx.AsyncByteStream):
async def __aiter__(self):
if False:
yield b""
stream = Stream()
upstream_response = httpx.Response(200, stream = stream)
async def fake_send(*_args, **_kwargs):
await gate.wait()
returned.set()
return upstream_response
class Request:
async def is_disconnected(self):
return False
monitor = ApiMonitor(max_entries = 3)
monitor_id = monitor.start(
endpoint = "/v1/chat/completions",
method = "POST",
model = "gguf",
prompt = "hi",
)
monkeypatch.setattr(inf_mod, "api_monitor", monitor)
monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fake_send)
payload = ChatCompletionRequest(
model = "default",
messages = [ChatMessage(role = "user", content = "hi")],
stream = True,
cancel_id = cancel_id,
)
response = await asyncio.wait_for(
_openai_passthrough_stream(
Request(),
threading.Event(),
SimpleNamespace(
base_url = "http://llama.test",
context_length = 4096,
_request_reasoning_kwargs = lambda *_args, **_kwargs: None,
),
payload,
"chatcmpl-test",
"chatcmpl-test",
monitor_id = monitor_id,
),
timeout = 0.2,
)
assert isinstance(response, _SameTaskStreamingResponse)
assert cancel_id in inf_mod._CANCEL_REGISTRY
gate.set()
await asyncio.wait_for(returned.wait(), timeout = 0.2)
await asyncio.sleep(0)
await response._unstarted_cleanup()
assert upstream_response.is_closed
assert cancel_id not in inf_mod._CANCEL_REGISTRY
asyncio.run(_run())
def test_external_non_streaming_json_updates_monitor(self, monkeypatch):
async def _run():
import routes.inference as inf_mod