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: if TYPE_CHECKING:
from .trained_router import TrainedRouter 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__) 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 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 from headroom.onnx_runtime import create_cpu_session_options, hf_hub_download_local_first
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)

View file

@ -12,8 +12,6 @@ from __future__ import annotations
import gc import gc
import io import io
from dataclasses import dataclass
from enum import Enum
from pathlib import Path from pathlib import Path
from typing import Any from typing import Any
@ -28,6 +26,11 @@ except ImportError:
from headroom.models.config import ML_MODEL_DEFAULTS 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: def _extract_tensor(output: torch.Tensor | BaseModelOutputWithPooling) -> torch.Tensor:
"""Extract tensor from model output. """Extract tensor from model output.
@ -55,37 +58,6 @@ def _extract_tensor(output: torch.Tensor | BaseModelOutputWithPooling) -> torch.
return output 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: class TrainedRouter:
"""Router using trained MiniLM classifier + SigLIP image analysis. """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"