mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
2a63ec70b6
## 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
82 lines
2.5 KiB
Python
82 lines
2.5 KiB
Python
"""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
|