fix(onnx): reduce retained cpu memory

headroom/image/onnx_router.py

@@ -20,6 +20,7 @@ from typing import Any
 import numpy as np
 
 from headroom.image.trained_router import ImageSignals, RouteDecision, Technique
+from headroom.onnx_runtime import create_cpu_session_options
 
 logger = logging.getLogger(__name__)
 
@@ -63,7 +64,9 @@ class OnnxTechniqueRouter:
 
         model_path = hf_hub_download(_TECHNIQUE_ROUTER_REPO, "model_quantized.onnx")
         self._classifier_session = ort.InferenceSession(
-            model_path, providers=["CPUExecutionProvider"]
+            model_path,
+            create_cpu_session_options(ort),
+            providers=["CPUExecutionProvider"],
         )
 
         tokenizer_path = hf_hub_download(_TECHNIQUE_ROUTER_REPO, "tokenizer.json")
@@ -95,7 +98,11 @@ class OnnxTechniqueRouter:
         logger.info("Loading SigLIP ONNX INT8 image encoder...")
 
         model_path = hf_hub_download(_SIGLIP_ENCODER_REPO, "image_encoder_int8.onnx")
-        self._siglip_session = ort.InferenceSession(model_path, providers=["CPUExecutionProvider"])
+        self._siglip_session = ort.InferenceSession(
+            model_path,
+            create_cpu_session_options(ort),
+            providers=["CPUExecutionProvider"],
+        )
 
         embeddings_path = hf_hub_download(_SIGLIP_ENCODER_REPO, "text_embeddings.npz")
         loaded = np.load(embeddings_path)

headroom/memory/adapters/embedders.py

@@ -19,6 +19,7 @@ from typing import TYPE_CHECKING, Any
 import numpy as np
 
 from headroom.models.config import ML_MODEL_DEFAULTS
+from headroom.onnx_runtime import create_cpu_session_options
 
 if TYPE_CHECKING:
     from sentence_transformers import SentenceTransformer
@@ -311,10 +312,14 @@ class OnnxLocalEmbedder:
         model_path = hf_hub_download(self.ONNX_REPO, "model.onnx")
         tok_path = hf_hub_download(self.ONNX_REPO, "tokenizer.json")
 
-        # Set thread count to avoid pthread_setaffinity_np errors in Docker containers
-        sess_options = ort.SessionOptions()
-        sess_options.intra_op_num_threads = 1
-        sess_options.inter_op_num_threads = 1
+        # Keep a small thread pool for Docker compatibility and disable ORT's
+        # CPU memory arena/pattern caches so long-running proxy workers do not
+        # retain large anonymous heaps after embedding bursts.
+        sess_options = create_cpu_session_options(
+            ort,
+            intra_op_num_threads=1,
+            inter_op_num_threads=1,
+        )
         self._session = ort.InferenceSession(
             model_path, sess_options, providers=["CPUExecutionProvider"]
         )

headroom/onnx_runtime.py

@@ -0,0 +1,51 @@
+"""ONNX Runtime helpers for long-running Headroom processes."""
+
+from __future__ import annotations
+
+import ctypes
+import sys
+from typing import Any
+
+
+def create_cpu_session_options(
+    ort: Any,
+    *,
+    intra_op_num_threads: int | None = None,
+    inter_op_num_threads: int | None = None,
+) -> Any:
+    """Create CPU-oriented ONNX Runtime session options.
+
+    Headroom runs as a long-lived proxy process, so we bias toward predictable
+    memory usage over peak ONNX throughput. Disabling ORT's CPU arena and memory
+    pattern caches reduces retained anonymous RSS after variable-size inference
+    workloads, which is especially important on small VMs.
+    """
+    sess_options = ort.SessionOptions()
+
+    if intra_op_num_threads is not None:
+        sess_options.intra_op_num_threads = intra_op_num_threads
+    if inter_op_num_threads is not None:
+        sess_options.inter_op_num_threads = inter_op_num_threads
+
+    if hasattr(sess_options, "enable_cpu_mem_arena"):
+        sess_options.enable_cpu_mem_arena = False
+    if hasattr(sess_options, "enable_mem_pattern"):
+        sess_options.enable_mem_pattern = False
+
+    return sess_options
+
+
+def trim_process_heap() -> bool:
+    """Ask glibc to return unused heap pages to the OS when available."""
+    if not sys.platform.startswith("linux"):
+        return False
+
+    try:
+        libc = ctypes.CDLL("libc.so.6")
+    except OSError:
+        return False
+
+    try:
+        return bool(libc.malloc_trim(0))
+    except Exception:
+        return False

headroom/transforms/kompress_compressor.py

@@ -14,12 +14,14 @@ Usage:
 
 from __future__ import annotations
 
+import gc
 import logging
 import threading
 from dataclasses import dataclass
 from typing import Any
 
 from ..config import TransformResult
+from ..onnx_runtime import create_cpu_session_options, trim_process_heap
 from ..tokenizer import Tokenizer
 from .base import Transform
 
@@ -174,7 +176,11 @@ def _load_kompress_onnx(model_id: str) -> tuple[Any, Any, str]:
         logger.info("Downloading Kompress ONNX model from %s ...", model_id)
         onnx_path = hf_hub_download(model_id, "onnx/kompress-int8.onnx")
 
-        session = ort.InferenceSession(onnx_path)
+        session = ort.InferenceSession(
+            onnx_path,
+            create_cpu_session_options(ort),
+            providers=["CPUExecutionProvider"],
+        )
         model = _OnnxModel(session)
         tokenizer = AutoTokenizer.from_pretrained("answerdotai/ModernBERT-base")
 
@@ -264,14 +270,17 @@ def unload_kompress_model(model_id: str | None = None) -> bool:
         else:
             return False
 
-        try:
-            import torch
+    try:
+        import torch
 
-            if torch.cuda.is_available():
-                torch.cuda.empty_cache()
-        except ImportError:
-            pass
-        return True
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+    except ImportError:
+        pass
+
+    gc.collect()
+    trim_process_heap()
+    return True
 
 
 # ── Compressor ────────────────────────────────────────────────────────

tests/test_onnx_runtime.py

@@ -0,0 +1,43 @@
+from headroom.onnx_runtime import create_cpu_session_options
+
+
+class _FakeSessionOptions:
+    def __init__(self):
+        self.intra_op_num_threads = None
+        self.inter_op_num_threads = None
+        self.enable_cpu_mem_arena = True
+        self.enable_mem_pattern = True
+
+
+class _FakeOrt:
+    SessionOptions = _FakeSessionOptions
+
+
+class _FakeSessionOptionsWithoutToggles:
+    def __init__(self):
+        self.intra_op_num_threads = None
+        self.inter_op_num_threads = None
+
+
+class _FakeOrtWithoutToggles:
+    SessionOptions = _FakeSessionOptionsWithoutToggles
+
+
+def test_create_cpu_session_options_disables_retention_features():
+    options = create_cpu_session_options(
+        _FakeOrt,
+        intra_op_num_threads=1,
+        inter_op_num_threads=2,
+    )
+
+    assert options.intra_op_num_threads == 1
+    assert options.inter_op_num_threads == 2
+    assert options.enable_cpu_mem_arena is False
+    assert options.enable_mem_pattern is False
+
+
+def test_create_cpu_session_options_handles_older_session_options():
+    options = create_cpu_session_options(_FakeOrtWithoutToggles)
+
+    assert options.intra_op_num_threads is None
+    assert options.inter_op_num_threads is None