From 5c561bd913ea60fad2c3c53f4b65e679e7d248d0 Mon Sep 17 00:00:00 2001 From: Abhay Singh Date: Sat, 8 Aug 2026 12:03:57 +0530 Subject: [PATCH] fix(onnx): stop ONNX thread pools from spinning idle cores (#2495) (#2540) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## 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 --- headroom/onnx_runtime.py | 31 +++++++++++++++++++++ tests/test_onnx_runtime.py | 57 ++++++++++++++++++++++++++++++++++++++ 2 files changed, 88 insertions(+) diff --git a/headroom/onnx_runtime.py b/headroom/onnx_runtime.py index aa3e7f7b2..8157c2202 100644 --- a/headroom/onnx_runtime.py +++ b/headroom/onnx_runtime.py @@ -13,6 +13,7 @@ logger = logging.getLogger(__name__) # Override for the CPU memory-arena default below: "1"/"true" forces the # arena ON, "0"/"false" forces it OFF, unset/"auto" uses the platform default. ONNX_CPU_ARENA_ENV = "HEADROOM_ONNX_CPU_ARENA" +ONNX_ALLOW_SPINNING_ENV = "HEADROOM_ONNX_ALLOW_SPINNING" _TRUTHY = frozenset({"1", "true", "yes", "on"}) _FALSY = frozenset({"0", "false", "no", "off"}) @@ -46,6 +47,22 @@ def cpu_arena_enabled() -> bool: return sys.platform == "win32" +def onnx_thread_spinning_enabled() -> bool: + """Whether ONNX Runtime intra/inter-op thread pools may spin-wait when idle. + + ORT's thread pools spin-wait on every core between inferences by default, so + a long-lived proxy that keeps compression/embedding models loaded pegs all + cores even while completely idle — the machine slows to a crawl after a + while (#2495). Default to blocking idle threads (spinning off). Set + ``HEADROOM_ONNX_ALLOW_SPINNING=1`` to restore ORT's spinning for peak + throughput on a dedicated/batch box. + """ + override = _env_flag(ONNX_ALLOW_SPINNING_ENV) + if override is not None: + return override + return False + + # Pin model artifacts to immutable commit SHAs so a changed or compromised # upstream HuggingFace repo cannot be pulled silently (supply-chain integrity). # Repos not listed here fall back to the floating default ref. Set @@ -159,6 +176,20 @@ def create_cpu_session_options( if inter_op_num_threads is not None: sess_options.inter_op_num_threads = inter_op_num_threads + if not onnx_thread_spinning_enabled(): + # ORT's thread pools spin-wait on all cores between inferences by + # default, so idle-but-loaded models peg every core in a long-lived + # proxy (#2495). Make idle threads block instead. Best-effort: older ORT + # builds may not recognize a key. + for spin_key in ( + "session.intra_op.allow_spinning", + "session.inter_op.allow_spinning", + ): + try: + sess_options.add_session_config_entry(spin_key, "0") + except Exception: + pass + if not cpu_arena_enabled(): if hasattr(sess_options, "enable_cpu_mem_arena"): sess_options.enable_cpu_mem_arena = False diff --git a/tests/test_onnx_runtime.py b/tests/test_onnx_runtime.py index 63d29abe6..3e7f6f949 100644 --- a/tests/test_onnx_runtime.py +++ b/tests/test_onnx_runtime.py @@ -2,10 +2,12 @@ import os import sys from headroom.onnx_runtime import ( + ONNX_ALLOW_SPINNING_ENV, ONNX_CPU_ARENA_ENV, cpu_arena_enabled, create_cpu_session_options, hf_entry_known_absent, + onnx_thread_spinning_enabled, ) @@ -15,6 +17,10 @@ class _FakeSessionOptions: self.inter_op_num_threads = None self.enable_cpu_mem_arena = True self.enable_mem_pattern = True + self.config_entries: dict[str, str] = {} + + def add_session_config_entry(self, key: str, value: str) -> None: + self.config_entries[key] = value class _FakeOrt: @@ -26,6 +32,10 @@ class _FakeSessionOptionsWithoutToggles: self.intra_op_num_threads = None self.inter_op_num_threads = None + def add_session_config_entry(self, key: str, value: str) -> None: + # No config storage on this stand-in; ORT here just accepts the call. + return None + class _FakeOrtWithoutToggles: SessionOptions = _FakeSessionOptionsWithoutToggles @@ -108,6 +118,53 @@ def test_create_cpu_session_options_handles_older_session_options(monkeypatch): assert options.inter_op_num_threads is None +def test_thread_spinning_disabled_by_default(monkeypatch): + # #2495: ORT thread pools spin-wait on all cores between inferences, so a + # long-lived proxy pegs every core while idle. Disable spinning by default. + monkeypatch.delenv(ONNX_ALLOW_SPINNING_ENV, raising=False) + monkeypatch.delenv(ONNX_CPU_ARENA_ENV, raising=False) + + assert onnx_thread_spinning_enabled() is False + options = create_cpu_session_options(_FakeOrt) + assert options.config_entries.get("session.intra_op.allow_spinning") == "0" + assert options.config_entries.get("session.inter_op.allow_spinning") == "0" + + +def test_thread_spinning_env_can_reenable(monkeypatch): + monkeypatch.setenv(ONNX_ALLOW_SPINNING_ENV, "1") + monkeypatch.delenv(ONNX_CPU_ARENA_ENV, raising=False) + + assert onnx_thread_spinning_enabled() is True + options = create_cpu_session_options(_FakeOrt) + assert "session.intra_op.allow_spinning" not in options.config_entries + assert "session.inter_op.allow_spinning" not in options.config_entries + + +def test_thread_spinning_env_explicit_off(monkeypatch): + monkeypatch.setenv(ONNX_ALLOW_SPINNING_ENV, "0") + + assert onnx_thread_spinning_enabled() is False + options = create_cpu_session_options(_FakeOrt) + assert options.config_entries.get("session.intra_op.allow_spinning") == "0" + + +def test_spinning_disable_is_best_effort_on_older_ort(monkeypatch): + # An ORT build that rejects the config key must not break session creation. + monkeypatch.delenv(ONNX_ALLOW_SPINNING_ENV, raising=False) + monkeypatch.setattr(sys, "platform", "linux") + + class _RejectingSessionOptions(_FakeSessionOptions): + def add_session_config_entry(self, key: str, value: str) -> None: + raise RuntimeError(f"unknown config key: {key}") + + class _RejectingOrt: + SessionOptions = _RejectingSessionOptions + + # Must not raise. + options = create_cpu_session_options(_RejectingOrt) + assert options.enable_cpu_mem_arena is False + + def _write_fake_hf_cache( root: str, repo_id: str, revision: str, *, no_exist_files: list[str] ) -> None: