From 2a63ec70b65605dfcff1b0afc292ab0298459f20 Mon Sep 17 00:00:00 2001 From: Abhay Singh Date: Sun, 26 Jul 2026 20:03:47 +0530 Subject: [PATCH] fix(image): reuse image models instead of rebuilding them per request (#2513) (#2536) ## Description Fixes #2513. Image compression rebuilt its heavyweight models on every request: - `_compress_messages_worker` (`proxy/image_isolation.py`) created a new `ImageCompressor()` per call, and - `ImageCompressor.compress` (`image/compressor.py`) created a new `OnnxTechniqueRouter(use_siglip=...)` per image. Each `OnnxTechniqueRouter` loads native `ort.InferenceSession` models, and ONNX Runtime holds C++ memory that Python's GC does not eagerly reclaim. The image pool is a **persistent** single-worker `ProcessPoolExecutor`, so those sessions accumulated in the worker and RSS grew ~70 KB/request, reaching ~1.1 GB after ~15k image requests in a day (the log shows one `[RapidOCR] Using engine_name: onnxruntime` line per request, confirming reloads). ## Fix Load the models once and reuse them: - `ImageCompressor` caches the ONNX router on `self._onnx_router` (built lazily via `_get_onnx_router`) instead of building one per `compress()` call. - The isolation worker keeps a per-process `ImageCompressor` singleton (`_get_worker_compressor`) and reuses it across calls. - `_get_image_compressor()` (main process, used for the `has_images()` gate) returns a shared instance too. - Shared instances are marked `_is_singleton`, and `close()` is a no-op on them, so a caller's per-request `close()` no longer unloads the models the next request reuses. A non-singleton `close()` still releases the torch router and drops the cached ONNX router. RSS is now flat after the initial model load; behavior is otherwise unchanged. ## 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/image/compressor.py`: add `_onnx_router` cache + `_get_onnx_router`, use it in `compress()`, add the `_is_singleton` flag, and make `close()` a no-op on a singleton (drop the cached ONNX router on a real close). - `headroom/proxy/image_isolation.py`: reuse a per-worker `ImageCompressor` singleton in `_compress_messages_worker` instead of building/closing one per call. - `headroom/proxy/helpers.py`: `_get_image_compressor()` returns a shared singleton instance. - `tests/test_image_compressor_singleton_reuse.py` (new): the ONNX router is built once and cached, singleton `close()` is a no-op while non-singleton `close()` releases, and both `_get_image_compressor` and the worker helper return a shared singleton. - `tests/test_proxy_handler_helpers.py`: updated the two existing `_get_image_compressor` tests that pinned the old fresh-per-call behavior to assert the singleton reuse instead (and reset the new module global so they stay isolated). ## 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_image_compressor_singleton_reuse.py -q 5 passed # with the fix reverted, all five fail (router rebuilt per call, close() # unloads the shared models, helpers return fresh instances) $ uvx ruff@0.15.17 check headroom/image/compressor.py headroom/proxy/image_isolation.py headroom/proxy/helpers.py tests/test_image_compressor_singleton_reuse.py All checks passed! $ uvx mypy@1.20.2 --ignore-missing-imports headroom/image/compressor.py headroom/proxy/image_isolation.py headroom/proxy/helpers.py Success: no issues found in 3 source files ``` The pre-existing async tests in `tests/test_image_compression_isolation.py` (4 `@pytest.mark.asyncio` cases) fail identically on clean `main` in this environment because pytest-asyncio is not configured here; they are unrelated to this change and pass in CI. ## Real Behavior Proof - Environment: Windows 11, Python 3.12, project venv (`uv sync --extra proxy`), `uvx ruff@0.15.17` / `uvx mypy@1.20.2`, pytest in the venv. - Exact command / steps: with `OnnxTechniqueRouter` construction mocked, called `ImageCompressor._get_onnx_router()` twice and asserted a single construction; exercised `close()` on singleton vs non-singleton instances; and called `_get_image_compressor()` / `_get_worker_compressor()` twice each. Then reverted the three source files and re-ran. - Observed result: with the fix the ONNX router is constructed once and reused, singleton `close()` leaves `_router`/`_onnx_router` intact (no `release_models`), non-singleton `close()` releases and nulls them, and both helper accessors return the same `_is_singleton` instance; with the fix reverted every one of these fails (fresh construction / unconditional release / new instances). Ran against the actual modules. - Not tested: a live multi-hour image workload measuring RSS (the leak is inferred from the removed per-request model construction; the ONNX/torch model load itself is mocked here). ## 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 --- headroom/image/compressor.py | 40 +++++++-- headroom/proxy/helpers.py | 24 ++++-- headroom/proxy/image_isolation.py | 47 +++++++---- .../test_image_compressor_singleton_reuse.py | 82 +++++++++++++++++++ tests/test_proxy_handler_helpers.py | 12 ++- 5 files changed, 174 insertions(+), 31 deletions(-) create mode 100644 tests/test_image_compressor_singleton_reuse.py diff --git a/headroom/image/compressor.py b/headroom/image/compressor.py index f9ba35a78..5dc7baa97 100644 --- a/headroom/image/compressor.py +++ b/headroom/image/compressor.py @@ -138,8 +138,17 @@ class ImageCompressor: self.use_siglip = use_siglip self.device = device - # Lazy-loaded router + # Lazy-loaded routers. The ONNX router loads native ort.InferenceSession + # models that Python's GC does not eagerly reclaim, so building one per + # compress() call grew RSS unboundedly under image traffic (#2513). + # Cache it on the instance and reuse it. self._router: TrainedRouter | None = None + self._onnx_router: Any = None + + # Set on a process-wide shared instance (see the isolation worker and + # _get_image_compressor) so a per-request close() does not unload models + # the next request would just reload. + self._is_singleton = False # Last compression result (for metrics) self.last_result: CompressionResult | None = None @@ -163,13 +172,36 @@ class ImageCompressor: ) return self._router + def _get_onnx_router(self) -> Any: + """Lazy-load and cache the ONNX technique router. + + Building an ``OnnxTechniqueRouter`` loads native ``ort.InferenceSession`` + models; creating one per ``compress()`` call leaked native memory and + grew RSS unboundedly under image traffic (#2513). Cache one per instance. + """ + if self._onnx_router is None: + from .onnx_router import OnnxTechniqueRouter + + self._onnx_router = OnnxTechniqueRouter(use_siglip=self.use_siglip) + return self._onnx_router + def close(self, unload_models: bool = True) -> None: - """Release any router-held model state.""" + """Release any router-held model state. + + A process-wide shared instance (``_is_singleton``) keeps its models + loaded across requests, so a per-request ``close()`` must be a no-op + there; otherwise every image request would reload the ONNX/torch models + it just cached, reintroducing the #2513 leak. + """ + if self._is_singleton: + return if self._router is not None: # Only loaded routers hold heavyweight image models; plain has_images() # checks remain cheap and have nothing to release. self._router.release_models(unload_registry=unload_models) self._router = None + # Drop the cached ONNX router so its native sessions can be reclaimed. + self._onnx_router = None def has_images(self, messages: list[dict[str, Any]]) -> bool: """Check if messages contain images.""" @@ -675,9 +707,7 @@ class ImageCompressor: confidence = 0.0 else: try: - from .onnx_router import OnnxTechniqueRouter - - onnx_router = OnnxTechniqueRouter(use_siglip=self.use_siglip) + onnx_router = self._get_onnx_router() decision = onnx_router.classify(image_data, query) technique = decision.technique confidence = decision.confidence diff --git a/headroom/proxy/helpers.py b/headroom/proxy/helpers.py index 7dfc7d91b..562cb3385 100644 --- a/headroom/proxy/helpers.py +++ b/headroom/proxy/helpers.py @@ -929,24 +929,36 @@ async def request_with_transient_retry( # Image compression availability (do not retain a global compressor instance) _image_compressor_available: bool | None = None +_image_compressor_instance: Any = None def _get_image_compressor(): - """Create a short-lived image compressor on demand.""" - global _image_compressor_available + """Return the process-wide image compressor, or None if unavailable. + + The compressor caches heavyweight models; creating a new one per request + (and a new ONNX router per image) accumulated native memory and grew RSS + unboundedly (#2513). Reuse a single shared instance. It is marked a + singleton so a caller's per-request ``close()`` is a no-op and the models + stay loaded. The main-process handlers only call ``has_images()`` on it (the + heavy compression runs in the isolation worker), but sharing still avoids a + fresh object per request. + """ + global _image_compressor_available, _image_compressor_instance if _image_compressor_available is False: return None + if _image_compressor_instance is not None: + return _image_compressor_instance try: from headroom.image import ImageCompressor - # Callers own closing the compressor; this helper only memoizes whether - # the optional image stack is importable. - compressor = ImageCompressor() + instance = ImageCompressor() + instance._is_singleton = True if _image_compressor_available is None: logger.info("Image compression enabled (model: chopratejas/technique-router)") _image_compressor_available = True - return compressor + _image_compressor_instance = instance + return instance except ImportError as e: if _image_compressor_available is not False: logger.warning(f"Image compression not available: {e}") diff --git a/headroom/proxy/image_isolation.py b/headroom/proxy/image_isolation.py index b0176ef7a..2e2addf54 100644 --- a/headroom/proxy/image_isolation.py +++ b/headroom/proxy/image_isolation.py @@ -16,27 +16,42 @@ logger = logging.getLogger("headroom.proxy") _IMAGE_POOL_LOCK = threading.Lock() _IMAGE_POOL: ProcessPoolExecutor | None = None +# Per-worker-process cached compressor. The image pool is a persistent +# single-worker ProcessPoolExecutor, so building a fresh ImageCompressor (and, +# inside it, a fresh ONNX router loading native ort.InferenceSession models) on +# every call accumulated native memory in the worker and grew RSS to 1+ GB over +# a day (#2513). Load the models once per worker and reuse them. +_WORKER_COMPRESSOR: Any = None + + +def _get_worker_compressor() -> Any: + global _WORKER_COMPRESSOR + if _WORKER_COMPRESSOR is None: + from headroom.image import ImageCompressor + + instance = ImageCompressor() + # Shared across calls in this worker: don't let a per-call close() unload + # the models the next call reuses. + instance._is_singleton = True + _WORKER_COMPRESSOR = instance + return _WORKER_COMPRESSOR + def _compress_messages_worker( messages: list[dict[str, Any]], provider: str, ) -> tuple[list[dict[str, Any]], dict[str, Any] | None]: - from headroom.image import ImageCompressor - - compressor = ImageCompressor() - try: - compressed = compressor.compress(messages, provider=provider) - if compressor.last_result is None: - return compressed, None - return compressed, { - "technique": compressor.last_result.technique.value, - "original_tokens": compressor.last_result.original_tokens, - "compressed_tokens": compressor.last_result.compressed_tokens, - "confidence": compressor.last_result.confidence, - "savings_percent": compressor.last_result.savings_percent, - } - finally: - compressor.close() + compressor = _get_worker_compressor() + compressed = compressor.compress(messages, provider=provider) + if compressor.last_result is None: + return compressed, None + return compressed, { + "technique": compressor.last_result.technique.value, + "original_tokens": compressor.last_result.original_tokens, + "compressed_tokens": compressor.last_result.compressed_tokens, + "confidence": compressor.last_result.confidence, + "savings_percent": compressor.last_result.savings_percent, + } def _success_worker( diff --git a/tests/test_image_compressor_singleton_reuse.py b/tests/test_image_compressor_singleton_reuse.py new file mode 100644 index 000000000..4890780d8 --- /dev/null +++ b/tests/test_image_compressor_singleton_reuse.py @@ -0,0 +1,82 @@ +"""Image models must be loaded once and reused, not rebuilt per request (#2513). + +Every image request built a new ImageCompressor and a new OnnxTechniqueRouter, +each loading native ort.InferenceSession models that grew worker RSS to 1+ GB +over a day. These tests pin the caching / singleton behavior with the heavy +model construction mocked out. +""" + +from __future__ import annotations + +from unittest.mock import MagicMock, patch + +from headroom.image.compressor import ImageCompressor + + +def test_onnx_router_is_built_once_and_cached() -> None: + compressor = ImageCompressor() + fake_router = MagicMock(name="OnnxTechniqueRouter") + + with patch("headroom.image.onnx_router.OnnxTechniqueRouter", return_value=fake_router) as ctor: + first = compressor._get_onnx_router() + second = compressor._get_onnx_router() + + assert first is second is fake_router + assert ctor.call_count == 1 + + +def test_close_is_a_noop_on_a_singleton_instance() -> None: + compressor = ImageCompressor() + compressor._is_singleton = True + router = MagicMock() + compressor._router = router + compressor._onnx_router = MagicMock() + + compressor.close() + + # Models stay loaded so the next request reuses them. + assert compressor._router is router + assert compressor._onnx_router is not None + router.release_models.assert_not_called() + + +def test_close_releases_models_on_a_non_singleton_instance() -> None: + compressor = ImageCompressor() + router = MagicMock() + compressor._router = router + compressor._onnx_router = MagicMock() + + compressor.close() + + router.release_models.assert_called_once() + assert compressor._router is None + assert compressor._onnx_router is None + + +def test_get_image_compressor_returns_a_shared_singleton() -> None: + import headroom.proxy.helpers as helpers + + helpers._image_compressor_available = None + helpers._image_compressor_instance = None + try: + a = helpers._get_image_compressor() + b = helpers._get_image_compressor() + assert a is not None + assert a is b + assert a._is_singleton is True + finally: + helpers._image_compressor_available = None + helpers._image_compressor_instance = None + + +def test_worker_compressor_is_reused_across_calls() -> None: + import headroom.proxy.image_isolation as iso + + iso._WORKER_COMPRESSOR = None + try: + a = iso._get_worker_compressor() + b = iso._get_worker_compressor() + assert a is b + assert a._is_singleton is True + finally: + iso._WORKER_COMPRESSOR = None diff --git a/tests/test_proxy_handler_helpers.py b/tests/test_proxy_handler_helpers.py index 9292f9930..bb86076aa 100644 --- a/tests/test_proxy_handler_helpers.py +++ b/tests/test_proxy_handler_helpers.py @@ -721,10 +721,13 @@ def test_anthropic_image_compression_helper_only_rewrites_latest_eligible_turn() ) == [compressed] -def test_proxy_helper_creates_fresh_image_compressors(monkeypatch) -> None: +def test_proxy_helper_reuses_a_singleton_image_compressor(monkeypatch) -> None: + # #2513: the compressor caches heavyweight models, so it must be a + # process-wide singleton rather than a fresh instance per request. from headroom.proxy import helpers monkeypatch.setattr(helpers, "_image_compressor_available", None) + monkeypatch.setattr(helpers, "_image_compressor_instance", None) _FreshCompressor.instances = 0 with patch("headroom.image.ImageCompressor", _FreshCompressor): @@ -732,9 +735,9 @@ def test_proxy_helper_creates_fresh_image_compressors(monkeypatch) -> None: second = helpers._get_image_compressor() assert isinstance(first, _FreshCompressor) - assert isinstance(second, _FreshCompressor) - assert first is not second - assert _FreshCompressor.instances == 2 + assert first is second + assert first._is_singleton is True + assert _FreshCompressor.instances == 1 def test_proxy_helper_caches_image_stack_import_failure(monkeypatch) -> None: @@ -751,6 +754,7 @@ def test_proxy_helper_caches_image_stack_import_failure(monkeypatch) -> None: return real_import(name, *args, **kwargs) monkeypatch.setattr(helpers, "_image_compressor_available", None) + monkeypatch.setattr(helpers, "_image_compressor_instance", None) monkeypatch.setattr(builtins, "__import__", fake_import) assert helpers._get_image_compressor() is None