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: