1 commit

Author	SHA1	Message	Date
Abhay Singh	2a63ec70b6	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	2026-07-26 07:33:47 -07:00