fix(studio): Windows GGUF cancel hang + CPU spinlock overhead (#5692)#5749
Conversation
…hai#5692) Two fixes for Windows-native GGUF inference via llama-server: **Issue 1 — GPU/CUDA Hang on Stream Cancellation:** - Add `Connection: close` header to all httpx requests proxying to llama-server, preventing Keep-Alive from masking downstream socket closure. - Introduce `_await_disconnect_then_close` background watcher that polls `request.is_disconnected()` every 100ms and calls `resp.aclose()` immediately when the client disconnects. This runs alongside the existing cancel-POST watcher and covers client aborts that never reach the /cancel endpoint (tab close, proxy aborts, Colab, mobile navigation, etc.). - Change all StreamingResponse `Connection: keep-alive` headers to `Connection: close`. **Issue 2 — High CPU Spinlock & KV Cache Backup Overhead:** - Set OMP_WAIT_POLICY=PASSIVE and OMP_NUM_THREADS=2 in the llama-server subprocess environment on Windows to prevent OpenMP from spin-waiting on all logical cores while the GPU decodes. - Limit `--threads` to 2 on Windows when the model is fully GPU-offloaded (`-ngl -1`). Auto-detect otherwise. - Pass `--cache-ram 0 --ctx-checkpoints 0 --no-cache-prompt --checkpoint-every-n-tokens -1` on Windows to disable prompt-cache snapshots that copy KV cache to system RAM over the WDDM/PCI-E bus. Closes unslothai#5692.
for more information, see https://pre-commit.ci
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces several optimizations and resource management improvements for inference, particularly targeting Windows environments. Key changes include limiting CPU threads and setting passive OpenMP wait policies when models are fully GPU-offloaded, disabling KV-cache checkpoints to system RAM to reduce PCI-E overhead, and implementing a background watcher to close response streams immediately upon client disconnect. Feedback from the review suggests simplifying the initialization logic for the GPU offload flag and recommends logging exceptions instead of using silent pass blocks when closing responses to improve debuggability.
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 881eb11e18
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| while not await request.is_disconnected(): | ||
| await asyncio.sleep(0.1) | ||
| try: | ||
| await resp.aclose() |
There was a problem hiding this comment.
When _await_disconnect_then_close closes resp on client disconnect, it never sets cancel_event, but both passthrough streamers treat RemoteProtocolError/ReadError/CloseError as ignorable only when cancel_event.is_set(). In the exact prefill-blocked case this watcher was added for, the loop never reaches its own request.is_disconnected() check, so the close is interpreted as an upstream error and re-raised instead of a normal disconnect path.
Useful? React with 👍 / 👎.
for more information, see https://pre-commit.ci
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 8addf2ac83
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
- Simplify _fully_gpu_offloaded init: default to False, only set True in the gpu_indices branch, drop redundant else. - Log exceptions in _await_disconnect_then_close at debug level instead of silent pass, per review suggestion.
|
Reviewed against 1.
|
|
Quick follow-up with a git-am-able patch and additional validation results. Copy-pasteable patchThe three review fixes as a single commit applied against PR head From 01e78eb86df3ff9f832b83e194a929469a35d2bd Mon Sep 17 00:00:00 2001
Subject: review: address codex-bot findings on PR 5749
---
studio/backend/core/inference/llama_cpp.py | 22 ++++++++++++---------
studio/backend/routes/inference.py | 13 +++++++++---
2 files changed, 23 insertions(+), 12 deletions(-)
diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py
@@ -3118,12 +3118,13 @@ class LlamaCppBackend:
else:
self._api_key = None
- # On Windows, llama-server's default KV-cache checkpoints write
- # prompt-cache snapshots to system RAM over the WDDM driver /
- # PCI-E bus. This adds substantial overhead even when the model
- # is fully GPU-offloaded. Disable all checkpoint machinery to
- # keep the cache entirely in VRAM. #5692.
- if _sys.platform == "win32":
+ # On Windows with full GPU offload, llama-server's default
+ # KV-cache checkpoints write prompt-cache snapshots to system
+ # RAM over the WDDM driver / PCI-E bus. Disable all checkpoint
+ # machinery to keep the cache entirely in VRAM. Scoped to
+ # full-offload only: CPU-only / partial-offload Windows runs
+ # still benefit from prompt caching across turns. #5692.
+ if _sys.platform == "win32" and _fully_gpu_offloaded:
cmd.extend(
[
"--cache-ram",
@@ -3177,9 +3178,12 @@ class LlamaCppBackend:
# PASSIVE tells the runtime to yield instead of spinning;
# limiting to 2 threads keeps the few remaining CPU-side
# operations (tokenisation, copy) from monopolising the
- # scheduler. #5692.
- env.setdefault("OMP_WAIT_POLICY", "PASSIVE")
- env.setdefault("OMP_NUM_THREADS", "2")
+ # scheduler. Scoped to full-offload only so CPU-only and
+ # partial-offload runs keep their default OpenMP parallelism.
+ # #5692.
+ if _fully_gpu_offloaded:
+ env.setdefault("OMP_WAIT_POLICY", "PASSIVE")
+ env.setdefault("OMP_NUM_THREADS", "2")
else:
# Linux: set LD_LIBRARY_PATH for shared libs next to the binary
# and CUDA runtime libs (libcudart, libcublas, etc.)
diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py
@@ -419,7 +419,7 @@ async def _await_cancel_then_close(cancel_event, resp) -> None:
return
-async def _await_disconnect_then_close(request, resp) -> None:
+async def _await_disconnect_then_close(request, resp, cancel_event) -> None:
"""Watch ``request.is_disconnected()`` and close ``resp`` when the client
disconnects (tab closed, navigation away, fetch abort). This prevents
llama-server from continuing to decode after the client is gone — a
@@ -432,10 +432,17 @@ async def _await_disconnect_then_close(request, resp) -> None:
Starlette Request directly, so it covers client-initiated aborts that
never reach the /cancel POST path (proxied fetches, Colab, mobile
tab-close, etc.).
+
+ Sets ``cancel_event`` before closing ``resp`` so the streamer's
+ ``except (RemoteProtocolError, ReadError, CloseError)`` clause -- which
+ only suppresses the re-raise when ``cancel_event.is_set()`` -- treats
+ the watcher-driven close as a normal client cancellation rather than
+ an upstream error. Mirrors the in-loop fallback path.
"""
try:
while not await request.is_disconnected():
await asyncio.sleep(0.1)
+ cancel_event.set()
try:
await resp.aclose()
except Exception as e:
@@ -5065,7 +5072,7 @@ async def _anthropic_passthrough_stream(
_await_cancel_then_close(cancel_event, resp)
)
disconnect_watcher = asyncio.create_task(
- _await_disconnect_then_close(request, resp)
+ _await_disconnect_then_close(request, resp, cancel_event)
)
lines_iter = resp.aiter_lines()
async for raw_line in lines_iter:
@@ -5491,7 +5498,7 @@ async def _openai_passthrough_stream(
_await_cancel_then_close(cancel_event, resp)
)
disconnect_watcher = asyncio.create_task(
- _await_disconnect_then_close(request, resp)
+ _await_disconnect_then_close(request, resp, cancel_event)
)
try:
lines_iter = resp.aiter_lines()Additional validation since the previous comment
The race test stresses the failure mode the disconnect-watcher's new The Anthropic-path |
- _await_disconnect_then_close: set cancel_event before resp.aclose() so the streamer's RemoteProtocolError handler treats the watcher-driven close as cancellation, not an upstream error. Both call sites pass cancel_event through. - Windows --cache-ram / --no-cache-prompt / --ctx-checkpoints block: gate on _fully_gpu_offloaded so CPU and partial-offload Windows runs keep prompt-cache reuse across turns. - Windows OMP_WAIT_POLICY / OMP_NUM_THREADS env: same gate so CPU and partial-offload Windows runs keep default OpenMP parallelism.
|
Pushed commit Three changes, all behaviour-preserving outside the targeted scenarios:
Local + cross-platform validation:
Let me know if you want any of these test files added to |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 38903b6402
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
- Drop the function-local `import sys as _sys` introduced as an F823 workaround; remove the redundant in-function `import os`/`import sys` block so module-level imports resolve sys/os instead. F823 no longer triggers because no shadowing import remains inside load_model. - Rename `_fully_gpu_offloaded` and `_t` to `fully_gpu_offloaded` and `threads_arg`. Underscore-prefixed names usually mean private/module- level; plain locals match Python style for in-function temporaries. No behavior change. ruff clean, py_compile clean, 35 studio cancel- infra tests + 13 launch-gating AST locks + 6 disconnect-watcher locks + 4 spoof live-import tests all pass.
|
Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits. |
The Windows-launch-gating AST regression locks pinned the exact variable names from the upstream PR head (`_fully_gpu_offloaded`, `_sys`, `_t`). The follow-up cleanup commit on PR unslothai#5749 drops the underscore prefixes, which would otherwise break the locks on the next upstream sync. Accept either form so the lock survives the rename. 19 tests still pass on the current staging-branch source.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: d992261380
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: fc83bb72e5
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
CI passes on fork. |
for more information, see https://pre-commit.ci
4 participants
Summary
Two fixes for Windows-native GGUF inference via llama-server (unsloth Studio).
Issue 1 — GPU/CUDA Hang on Stream Cancellation (Socket Leak)
When a user cancels a streaming chat completion midway (stops generation, closes tab), llama-server continues decoding in the background because the httpx Keep-Alive connection masks the downstream socket closure.
Changes in
studio/backend/routes/inference.py:Connection: closeheader to all httpx requests proxying to llama-server, preventing Keep-Alive from masking downstream socket closure._await_disconnect_then_closebackground watcher that pollsrequest.is_disconnected()every 100ms and callsresp.aclose()immediately when the client disconnects. This runs alongside the existing cancel-POST watcher and covers client aborts that never reach the /cancel endpoint (tab close, proxy aborts, Colab, mobile navigation, etc.).Connection: keep-aliveheaders toConnection: close.Issue 2 — High CPU Spinlock & KV Cache Backup Overhead
During GGUF inference on Windows, CPU usage spikes to 100% even when the model is fully GPU-offloaded (
-ngl -1), causing thermal throttling and fan noise. Two root causes:Changes in
studio/backend/core/inference/llama_cpp.py:OMP_WAIT_POLICY=PASSIVEandOMP_NUM_THREADS=2in the llama-server subprocess environment on Windows to prevent OpenMP from spin-waiting on all logical cores while the GPU decodes.--threadsto 2 on Windows when the model is fully GPU-offloaded (-ngl -1). Auto-detect otherwise.--cache-ram 0 --ctx-checkpoints 0 --no-cache-prompt --checkpoint-every-n-tokens -1on Windows to disable prompt-cache snapshots that copy KV cache to system RAM over the WDDM/PCI-E bus.Fixes #5692.