From d7cf981093cf505192a3736dadd0254a120830a1 Mon Sep 17 00:00:00 2001 From: Abhay Singh Date: Wed, 12 Aug 2026 10:52:56 +0530 Subject: [PATCH] fix(image): decouple routing types from trained_router so importing the compressor doesn't import torch (#2513) (#2537) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Description Addresses the secondary crash in #2513. `image/compressor.py` did `from .trained_router import Technique` at module scope, and `image/onnx_router.py` imported `ImageSignals` / `RouteDecision` / `Technique` from `trained_router` the same way. `trained_router` imports `torch` and `transformers` at module scope, so merely importing the image compressor eagerly pulled in the heavy ML stack. On Python 3.13+ that eager import crashed the first image request with: ``` AttributeError: module 'torch' has no attribute 'compiler' ``` because `transformers` touches `torch.compiler` during its own import, before torch has finished initializing inside the proxy process. ## Fix Move the dependency-free routing types — the `Technique` enum and the `ImageSignals` / `RouteDecision` dataclasses — into a new `headroom/image/image_types.py` (no torch / transformers / onnx imports). `compressor.py` and `onnx_router.py` import them from there; `trained_router` re-exports them so existing `from .trained_router import Technique` imports keep working. Importing the compressor or the ONNX router no longer imports `trained_router`, so the torch/transformers stack is only loaded when the PyTorch router is actually used (lazily, inside `_get_router`). ## 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/image_types.py` (new): `Technique`, `ImageSignals`, `RouteDecision` — pure enum/dataclasses. - `headroom/image/trained_router.py`: import and re-export those types from `image_types` (drop the local definitions and the now-unused `dataclass` / `Enum` imports). - `headroom/image/compressor.py`, `headroom/image/onnx_router.py`: import the routing types from `image_types`. - `tests/test_image_types_torch_decoupling.py` (new): subprocess checks that importing the compressor / ONNX router does not import `trained_router`, that `image_types` imports no torch, and that all three re-export paths resolve to the same objects. ## 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_types_torch_decoupling.py -q 4 passed # with the compressor import reverted, the "does not import trained_router" # check fails $ uvx ruff@0.15.17 check headroom/image/image_types.py headroom/image/trained_router.py headroom/image/compressor.py headroom/image/onnx_router.py tests/test_image_types_torch_decoupling.py All checks passed! $ uvx mypy@1.20.2 --ignore-missing-imports headroom/image/image_types.py headroom/image/trained_router.py headroom/image/compressor.py headroom/image/onnx_router.py Success: no issues found in 4 source files ``` `tests/test_image_compression.py::TestOnnxRouter::test_full_classify_with_image` fails identically on clean `main` in this environment (it needs real ONNX model weights that aren't available locally); it is unrelated to this change. ## Real Behavior Proof - Environment: Windows 11, Python 3.12, project venv (`uv sync --extra proxy`, torch not installed here), `uvx ruff@0.15.17` / `uvx mypy@1.20.2`, pytest in the venv. - Exact command / steps: in a fresh subprocess, imported `headroom.image.compressor` / `headroom.image.onnx_router` / `headroom.image.image_types` and checked `sys.modules`; also asserted `headroom.image.Technique`, `trained_router.Technique`, and `image_types.Technique` are the same object. Then reverted the compressor import and re-ran. - Observed result: with the fix, importing the compressor and the ONNX router leaves `headroom.image.trained_router` out of `sys.modules`, `image_types` pulls in no `torch`, and all re-export paths are identical objects; with the fix reverted, importing the compressor pulls `trained_router` back in (the eager path that triggers the torch import). Ran against the actual modules. - Not tested: the Python 3.13 `torch.compiler` crash itself (this environment is 3.12 without torch); the fix removes the eager import that causes it. ## 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 | 5 +- headroom/image/image_types.py | 45 ++++++++++++++++++ headroom/image/onnx_router.py | 2 +- headroom/image/trained_router.py | 38 ++------------- tests/test_image_types_torch_decoupling.py | 55 ++++++++++++++++++++++ 5 files changed, 110 insertions(+), 35 deletions(-) create mode 100644 headroom/image/image_types.py create mode 100644 tests/test_image_types_torch_decoupling.py diff --git a/headroom/image/compressor.py b/headroom/image/compressor.py index 5dc7baa97..448d52be4 100644 --- a/headroom/image/compressor.py +++ b/headroom/image/compressor.py @@ -31,7 +31,10 @@ from typing import TYPE_CHECKING, Any if TYPE_CHECKING: from .trained_router import TrainedRouter -from .trained_router import Technique +# Import the enum from the dependency-free module so importing the compressor +# does not eagerly import trained_router (and thus torch/transformers) — that +# eager import crashed on Python 3.13+ (#2513). +from .image_types import Technique logger = logging.getLogger(__name__) diff --git a/headroom/image/image_types.py b/headroom/image/image_types.py new file mode 100644 index 000000000..572fcf7f3 --- /dev/null +++ b/headroom/image/image_types.py @@ -0,0 +1,45 @@ +"""Lightweight image-routing types shared across the image stack. + +Kept dependency-free (pure enum + dataclasses, no torch / transformers / onnx) +so importing the image compressor or the ONNX router does not eagerly import +the heavy ML stack via ``trained_router``. On Python 3.13+ that eager import +crashed with ``AttributeError: module 'torch' has no attribute 'compiler'`` +because ``transformers`` touched ``torch.compiler`` before torch finished +initializing inside the proxy process (#2513). +""" + +from __future__ import annotations + +from dataclasses import dataclass +from enum import Enum + + +class Technique(Enum): + """Image optimization techniques.""" + + TRANSCODE = "transcode" # Convert to text description (99% savings) + CROP = "crop" # Extract relevant region (50-90% savings) + PRESERVE = "preserve" # Keep full quality (0% savings) + FULL_LOW = "full_low" # Full image, lower quality (87% savings) + + +@dataclass +class ImageSignals: + """Signals extracted from image analysis.""" + + has_text: float + is_document: float + is_complex: float + has_small_details: float + + +@dataclass +class RouteDecision: + """Result of routing decision.""" + + technique: Technique + confidence: float + reason: str + image_signals: ImageSignals | None = None + query_prediction: str | None = None + query_confidence: float | None = None diff --git a/headroom/image/onnx_router.py b/headroom/image/onnx_router.py index 64ceeefa6..509303bac 100644 --- a/headroom/image/onnx_router.py +++ b/headroom/image/onnx_router.py @@ -19,7 +19,7 @@ from typing import Any import numpy as np -from headroom.image.trained_router import ImageSignals, RouteDecision, Technique +from headroom.image.image_types import ImageSignals, RouteDecision, Technique from headroom.onnx_runtime import create_cpu_session_options, hf_hub_download_local_first logger = logging.getLogger(__name__) diff --git a/headroom/image/trained_router.py b/headroom/image/trained_router.py index b53c30d19..dfa16ef70 100644 --- a/headroom/image/trained_router.py +++ b/headroom/image/trained_router.py @@ -12,8 +12,6 @@ from __future__ import annotations import gc import io -from dataclasses import dataclass -from enum import Enum from pathlib import Path from typing import Any @@ -28,6 +26,11 @@ except ImportError: from headroom.models.config import ML_MODEL_DEFAULTS +# Re-exported from the dependency-free module so existing +# ``from .trained_router import Technique`` imports keep working without this +# module (which imports torch/transformers) being needed just for the types. +from .image_types import ImageSignals, RouteDecision, Technique + def _extract_tensor(output: torch.Tensor | BaseModelOutputWithPooling) -> torch.Tensor: """Extract tensor from model output. @@ -55,37 +58,6 @@ def _extract_tensor(output: torch.Tensor | BaseModelOutputWithPooling) -> torch. return output -class Technique(Enum): - """Image optimization techniques.""" - - TRANSCODE = "transcode" # Convert to text description (99% savings) - CROP = "crop" # Extract relevant region (50-90% savings) - PRESERVE = "preserve" # Keep full quality (0% savings) - FULL_LOW = "full_low" # Full image, lower quality (87% savings) - - -@dataclass -class ImageSignals: - """Signals extracted from image analysis.""" - - has_text: float - is_document: float - is_complex: float - has_small_details: float - - -@dataclass -class RouteDecision: - """Result of routing decision.""" - - technique: Technique - confidence: float - reason: str - image_signals: ImageSignals | None = None - query_prediction: str | None = None - query_confidence: float | None = None - - class TrainedRouter: """Router using trained MiniLM classifier + SigLIP image analysis. diff --git a/tests/test_image_types_torch_decoupling.py b/tests/test_image_types_torch_decoupling.py new file mode 100644 index 000000000..4c7f37426 --- /dev/null +++ b/tests/test_image_types_torch_decoupling.py @@ -0,0 +1,55 @@ +"""Importing the image compressor must not eagerly import trained_router (torch). + +#2513 (secondary): ``from .trained_router import Technique`` at module scope +pulled in torch / transformers, which on Python 3.13+ crashed with +``AttributeError: module 'torch' has no attribute 'compiler'`` before torch had +finished initializing. The routing types now live in a dependency-free module. + +The "does not import" checks run in a fresh subprocess so module caching from +other tests can't mask the import graph. +""" + +from __future__ import annotations + +import subprocess +import sys + + +def _assert_import_excludes(module_name: str, excluded: str) -> None: + code = ( + "import sys, importlib\n" + f"importlib.import_module({module_name!r})\n" + f"assert {excluded!r} not in sys.modules, {excluded!r} + ' was imported'\n" + "print('ok')\n" + ) + result = subprocess.run( + [sys.executable, "-c", code], capture_output=True, text=True, timeout=120 + ) + assert result.returncode == 0, result.stderr + assert result.stdout.strip() == "ok" + + +def test_importing_compressor_does_not_import_trained_router() -> None: + _assert_import_excludes("headroom.image.compressor", "headroom.image.trained_router") + + +def test_importing_onnx_router_does_not_import_trained_router() -> None: + _assert_import_excludes("headroom.image.onnx_router", "headroom.image.trained_router") + + +def test_image_types_module_does_not_import_torch() -> None: + _assert_import_excludes("headroom.image.image_types", "torch") + + +def test_technique_reexports_are_the_same_object() -> None: + from headroom.image import Technique as via_package + from headroom.image.image_types import ImageSignals, RouteDecision, Technique + from headroom.image.trained_router import ImageSignals as tr_signals + from headroom.image.trained_router import RouteDecision as tr_decision + from headroom.image.trained_router import Technique as tr_technique + + assert via_package is Technique + assert tr_technique is Technique + assert tr_signals is ImageSignals + assert tr_decision is RouteDecision + assert Technique.PRESERVE.value == "preserve"