mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
4 commits
| Author | SHA1 | Message | Date | |
|---|---|---|---|---|
Abhay Singh
|
5c561bd913
|
fix(onnx): stop ONNX thread pools from spinning idle cores (#2495) (#2540)
## Description Fixes #2495 (tokensave / the proxy using ~100% of all cores). ONNX Runtime's intra-op (and inter-op) thread pools **spin-wait on every core between inferences** by default. Headroom is a long-lived process that keeps ONNX models loaded — the kompress code compressor ("tokensave"), the image technique/SigLIP routers, and the memory embedder — so once a model is loaded, its idle thread pool keeps every core busy even when no compression is running. That matches the report exactly: CPU climbs to ~100% of all cores "after a period of time" and the whole machine slows down, with no obvious trigger. `create_cpu_session_options` (the shared factory every CPU ONNX session goes through) configured threads and the memory arena but never touched spinning, so ORT's default (spin enabled) was in effect everywhere. ## Fix Disable intra-op and inter-op thread spinning in `create_cpu_session_options` so idle ORT threads block instead of spin-waiting. This applies to every ONNX session built through the factory (kompress + the image routers). It: - is **best-effort per key** (wrapped in try/except) so an older ORT build that doesn't recognize a config key still creates a session; - is **overridable** via `HEADROOM_ONNX_ALLOW_SPINNING=1` for a dedicated/batch box that wants ORT's peak-throughput spinning; - does not change active-inference throughput meaningfully — blocking threads wake on new work with only microsecond-scale latency, which is the recommended setting for a server/proxy with idle periods. The memory embedder already builds its own options with `intra_op_num_threads=1`; this change is orthogonal and additionally quiets its idle spinning if it were ever routed through the factory. ## Type of Change - [x] Bug fix (non-breaking change that fixes an issue) - [ ] New feature (non-breaking change that adds functionality) - [ ] Breaking change (fix or feature that would cause existing functionality to change) - [ ] Documentation update - [ ] Performance improvement - [ ] Code refactoring (no functional changes) ## Changes Made - `headroom/onnx_runtime.py`: add `ONNX_ALLOW_SPINNING_ENV` + `onnx_thread_spinning_enabled()`; disable `session.intra_op.allow_spinning` / `session.inter_op.allow_spinning` in `create_cpu_session_options` unless spinning is explicitly re-enabled. - `tests/test_onnx_runtime.py`: spinning is disabled by default (both keys), `HEADROOM_ONNX_ALLOW_SPINNING=1` re-enables it, an explicit `0` disables it, and a config key an older ORT rejects doesn't break session creation. ## Testing - [x] Unit tests pass (`pytest`) - [x] Linting passes (`ruff check .`) - [x] Type checking passes (`mypy headroom`) - [x] New tests added for new functionality - [ ] Manual testing performed ### Test Output ```text $ python -m pytest tests/test_onnx_runtime.py -q 11 passed # with the fix reverted the new symbols don't exist, so the spinning tests # fail at import — the pre-fix factory left ORT's spinning at its (enabled) default $ uvx ruff@0.15.17 check headroom/onnx_runtime.py tests/test_onnx_runtime.py All checks passed! $ uvx mypy@1.20.2 --ignore-missing-imports headroom/onnx_runtime.py Success: no issues found in 1 source file ``` ## Real Behavior Proof - Environment: Windows 11, Python 3.12, project venv (`uv sync --extra proxy`, onnxruntime 1.23.2 installed), `uvx ruff@0.15.17` / `uvx mypy@1.20.2`, pytest in the venv. - Exact command / steps: built a real `onnxruntime.SessionOptions` via `create_cpu_session_options(ort)` and read back `session.intra_op.allow_spinning` / `session.inter_op.allow_spinning`; repeated with `HEADROOM_ONNX_ALLOW_SPINNING=1`. - Observed result: by default both keys read back `"0"` (spinning disabled); with `HEADROOM_ONNX_ALLOW_SPINNING=1` neither key is set (ORT's default spinning restored). Against a real ORT the pre-fix factory set neither key, so ORT's default (spinning enabled) applied — the idle all-cores burn. Ran against the actual module and real onnxruntime. - Not tested: a live multi-hour VS Code + Claude session measuring CPU before/after (the spinning-disable is the documented ORT remedy for idle-CPU in a long-lived process; the config change itself is verified end to end against real ORT). ## Review Readiness - [x] I have performed a self-review - [x] This PR is ready for human review ## Checklist - [x] My code follows the project's style guidelines - [x] I have performed a self-review of my code - [x] I have commented my code, particularly in hard-to-understand areas - [ ] I have made corresponding changes to the documentation - [x] My changes generate no new warnings - [x] I have added tests that prove my fix is effective or that my feature works - [x] New and existing unit tests pass locally with my changes - [ ] I have updated the CHANGELOG.md if applicable --------- Co-authored-by: JD Davis <jd@jds-macbook-air.tail2a279.ts.net> |
||
|
Raúl
|
1a2688b57f
|
test(kompress): close patch-coverage gaps from #2716 (#2721)
## Description Codecov flagged 9 uncovered lines on #2716 after it merged: `hf_entry_known_absent`'s own body in `headroom/onnx_runtime.py` was only ever exercised indirectly (every existing test in `tests/test_transforms/test_kompress_compressor.py` monkeypatched it away rather than calling the real implementation), and `_load_pytorch_weights` / `_load_kompress_pytorch` in `headroom/transforms/kompress_compressor.py` had three untested branches: the double cache-miss under `allow_download=False` (merged.pt confirmed absent AND the plain fallback also not cached), a genuine non-404 download failure propagating instead of silently falling back, and the already-cached fast path in `_load_kompress_pytorch`. ## Type of Change - [ ] Bug fix - [ ] New feature - [x] Test coverage improvement, no production code change ## Changes Made - `tests/test_onnx_runtime.py`: added `_write_fake_hf_cache` (builds a minimal on-disk HF hub cache layout, including the `.no_exist/<hash>/<filename>` marker huggingface_hub writes after a real 404) and three direct tests of `hf_entry_known_absent` against the real `huggingface_hub.try_to_load_from_cache`, not a mock of it. - `tests/test_transforms/test_kompress_compressor.py`: added `test_cache_only_raises_when_confirmed_absent_but_plain_also_missing`, `test_genuine_download_failure_propagates_instead_of_falling_back`, and a new `TestLoadKompressPytorchCaching` class covering the already-cached fast path. ## Testing ```text $ .venv/bin/python3 -m pytest tests/test_onnx_runtime.py tests/test_transforms/test_kompress_compressor.py -q 51 passed $ .venv/bin/python3 -m pytest tests/ -k "kompress or onnx_runtime" -q --cov=headroom.transforms.kompress_compressor --cov=headroom.onnx_runtime --cov-report=term-missing # before: onnx_runtime.py Missing includes 132-136 (hf_entry_known_absent's entire body); # kompress_compressor.py Missing includes 805-806, 818, 836 # after: none of those lines appear in Missing anymore 191 passed, 7 skipped $ .venv/bin/python3 -m ruff format --check tests/test_onnx_runtime.py tests/test_transforms/test_kompress_compressor.py 2 files already formatted $ .venv/bin/python3 -m ruff check tests/test_onnx_runtime.py tests/test_transforms/test_kompress_compressor.py All checks passed! ``` ## Review Readiness - Test-only, additive diff (113 insertions, 0 deletions, 0 lines touched outside the two test files). No behavior change possible. ## Checklist - [x] My code follows the project's style guidelines - [x] I have performed a self-review of my code - [x] New and existing unit tests pass locally with my changes - [x] I did not edit `CHANGELOG.md` ## Additional Notes Not closing: the remaining branch-partial on the `device == "auto"` cuda/mps/cpu selection in `_load_kompress_pytorch` (would need mocking `torch.cuda.is_available()` / `torch.backends.mps.is_available()` for marginal benefit); left as-is. |
||
|
Tejas Chopra
|
36202f4d0b
|
fix(windows): unwedge compression on degraded ONNX runtimes (every request timing out at 30s, 0% savings) (#822)
## Summary Multiple Windows users reported (via Discord, on v0.23.0, `pip install "headroom-ai[all]"`) that the proxy delivers **zero compression** and adds **+30s latency to every request**: `Optimization failed: TimeoutError:` with `compression_first_stage ≈ 30000ms` on every optimization attempt, for the lifetime of the process. Log analysis showed the wedge starts at the **first message eligible for real compression** (earlier requests succeed because everything is skipped/excluded) — and never recovers, even though the Kompress model loaded successfully at startup. ### Root cause chain 1. `create_cpu_session_options` disabled ONNX Runtime's CPU memory arena on **all** platforms. On Windows this is catastrophic: every `Run()` falls back to per-node `VirtualAlloc`/free, slowing ModernBERT inference by 2–3 orders of magnitude (onnxruntime#11627). One reporter's perf summary showed max optimization overhead of **200,369ms** (~13 chunks × ~15s) — slow, not deadlocked. 2. The first slow inference outlives the proxy's 30s compression-stage timeout. `asyncio.wait_for` abandons the future but **cannot kill the executor thread**, which keeps holding the Kompress `BoundedSemaphore(1)`. 3. Every later compression blocks on an **unbounded** `semaphore.acquire()`, times out at exactly 30s, and leaks another thread — permanently wedging the proxy until restart. Two adjacent Windows bugs found in the same logs are fixed too: `subprocess.run(text=True)` without `encoding=` decodes child output with cp1252, so rtk's emoji output killed reader threads (`UnicodeDecodeError: 'charmap' codec can't decode byte 0x8f`); and the OpenAI handler logged `Optimization failed: ` with an empty message because `str(asyncio.TimeoutError())` is empty. ### Fixes - **`onnx_runtime.py`** — keep the CPU arena at ORT's default on Windows; Linux/macOS keep the legacy low-RSS behavior (arena disabled) bit-for-bit. New `HEADROOM_ONNX_CPU_ARENA` env overrides either way. All ONNX sessions (Kompress, image router, memory embedders) share this helper, so one fix covers them all. - **`kompress_compressor.py`** — three layers of wedge-proofing, each fail-safing to passthrough instead of blocking: - bounded semaphore acquire (`HEADROOM_KOMPRESS_ACQUIRE_TIMEOUT_SECONDS`, default 5s) - wall-clock budget per compress/compress_batch call (`HEADROOM_KOMPRESS_TIME_BUDGET_SECONDS`, default 20s — under the 30s stage timeout, so Kompress gives up before the request is abandoned). Batch bail never emits a partially-covered text. - preload canary (`HEADROOM_KOMPRESS_CANARY_SECONDS`, default 5s, one retry to forgive cold-start warmup): machines that can never finish inference inside the stage timeout get ML compression disabled up front with one actionable warning, instead of a guaranteed 30s timeout per request. - Setting any knob `<= 0` disables that guard (restores legacy behavior). First give-up logs at WARNING with remediation hints; repeats drop to DEBUG. - **`proxy/helpers.py`, `interceptors/astgrep.py`** — `encoding="utf-8", errors="replace"` on rtk/lean-ctx/ast-grep subprocess calls. - **`handlers/openai.py`** — failure log now includes request id + exception type, matching the Anthropic handler. ### Non-Windows perf - Session options on Linux/macOS are unchanged (pinned by tests). - The only new hot-path cost is one `time.monotonic()` + a bounded acquire per chunk: micro-benchmarked at sub-microsecond (bounded acquire measured marginally *faster* than the old context-manager acquire), vs 50–500ms of inference per chunk. - Real-model smoke run on macOS: identical compression output (ratio 0.262 on a 1020-word sample), canary passes, budget/acquire give-up paths verified against the real ONNX stack by forcing tiny env values. Related (same symptom, different root cause — **not** addressed here): #810 tracks the blocked-tiktoken-download hang, which produces the same per-request 30s `TimeoutError` signature. The bounded-acquire/budget changes in this PR limit the blast radius of Kompress-side slowness only. ## Validation - `.venv/bin/ruff check headroom/ tests/...` — clean - `.venv/bin/ruff format --check` — clean (355 files) - `.venv/bin/mypy` on all five changed source files — no issues - `python -m pytest tests/test_onnx_runtime.py tests/test_kompress_failsafe.py tests/test_subprocess_encoding.py` — 25 passed (new coverage: arena platform matrix + env overrides, stuck-semaphore passthrough for compress and batch, budget bail incl. mid-batch no-data-loss, canary trip/pass/retry/disable/error-safety, UTF-8 subprocess kwargs) - `python -m pytest tests/test_transforms_content_router.py tests/test_proxy_handler_helpers.py tests/test_codex_ws_compression_scheduler.py tests/test_proxy_warmup.py tests/test_proxy_pipeline_lifecycle.py` — 52 passed (existing suites for touched areas) <!-- headroom-maintainer-template-completion:start --> ## Description This PR prepares `fix(windows): unwedge compression on degraded ONNX runtimes (every request timing out at 30s, 0% savings)` for review by documenting the intended change, validation evidence, and remaining merge-readiness context. Linked issues: None declared. ## Type of Change - [x] Bug fix - [ ] New feature - [ ] Documentation - [ ] Refactor - [ ] Tests only ## Changes Made - Commit: fix(windows): unwedge compression on degraded ONNX runtimes - Commit: fix(kompress): run preload canary off the startup path - Touches `headroom/onnx_runtime.py` - Touches `headroom/proxy/handlers/openai.py` - Touches `headroom/proxy/helpers.py` - Touches `headroom/proxy/interceptors/astgrep.py` - Touches `headroom/transforms/kompress_compressor.py` - Touches `tests/test_kompress_failsafe.py` ## Testing - [x] GitHub checks reviewed - [x] Metadata/template validation - [ ] Local functional testing ### Test Output ```text gh pr view 822 --repo chopratejas/headroom --json statusCheckRollup - CI / changes: SUCCESS - CodeQL / Analyze (actions): SUCCESS - Evaluation Suite / smoke-test: SUCCESS - Init E2E / docker-init-e2e: SUCCESS - PR Governance / label: SUCCESS - Wrap E2E / docker-wrap-e2e: SUCCESS - CodeQL / Analyze (c-cpp): SUCCESS - CodeQL / Analyze (javascript-typescript): SUCCESS - CodeQL / Analyze (python): SUCCESS - CodeQL / Analyze (rust): SUCCESS - Evaluation Suite / weekly-suite: SKIPPED - CI / commitlint: SUCCESS ``` ## Real Behavior Proof - Environment: GitHub PR metadata and checks for `chopratejas/headroom` PR #822. - Exact command / steps: Reviewed PR title, commits, changed files, linked issues, labels, and check rollup; appended this maintainer template completion block without replacing the author's original description. - Observed result: PR body now contains all required governance sections, checked readiness fields, and a non-placeholder validation evidence block. - Not tested: This pass updated PR metadata only; code validation remains represented by the linked GitHub checks and any author-provided evidence above. ## Review Readiness - [x] I have performed a self-review - [x] This PR is ready for human review <!-- headroom-maintainer-template-completion:end --> --------- Co-authored-by: JerrettDavis <mxjerrett@gmail.com> |
||
|
Kayzo
|
fbbde51db1 | fix(onnx): reduce retained cpu memory |