fix(image): decouple routing types from trained_router so importing the compressor doesn't import torch (#2513) (#2537)

## 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
This commit is contained in:
Abhay Singh 2026-08-12 10:52:56 +05:30 committed by GitHub
parent a540eb2c61
commit d7cf981093
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 110 additions and 35 deletions

View file

@ -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__)

View file

@ -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

View file

@ -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__)

View file

@ -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.

View file

@ -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"