## Description
Codecov flagged 9 uncovered lines on #2716 after it merged:
`hf_entry_known_absent`'s own body
in `headroom/onnx_runtime.py` was only ever exercised indirectly (every
existing test in
`tests/test_transforms/test_kompress_compressor.py` monkeypatched it
away rather than calling the
real implementation), and `_load_pytorch_weights` /
`_load_kompress_pytorch` in
`headroom/transforms/kompress_compressor.py` had three untested
branches: the double cache-miss
under `allow_download=False` (merged.pt confirmed absent AND the plain
fallback also not cached),
a genuine non-404 download failure propagating instead of silently
falling back, and the
already-cached fast path in `_load_kompress_pytorch`.
## Type of Change
- [ ] Bug fix
- [ ] New feature
- [x] Test coverage improvement, no production code change
## Changes Made
- `tests/test_onnx_runtime.py`: added `_write_fake_hf_cache` (builds a
minimal on-disk HF hub
cache layout, including the `.no_exist/<hash>/<filename>` marker
huggingface_hub writes after a
real 404) and three direct tests of `hf_entry_known_absent` against the
real
`huggingface_hub.try_to_load_from_cache`, not a mock of it.
- `tests/test_transforms/test_kompress_compressor.py`: added
`test_cache_only_raises_when_confirmed_absent_but_plain_also_missing`,
`test_genuine_download_failure_propagates_instead_of_falling_back`, and
a new
`TestLoadKompressPytorchCaching` class covering the already-cached fast
path.
## Testing
```text
$ .venv/bin/python3 -m pytest tests/test_onnx_runtime.py tests/test_transforms/test_kompress_compressor.py -q
51 passed
$ .venv/bin/python3 -m pytest tests/ -k "kompress or onnx_runtime" -q --cov=headroom.transforms.kompress_compressor --cov=headroom.onnx_runtime --cov-report=term-missing
# before: onnx_runtime.py Missing includes 132-136 (hf_entry_known_absent's entire body);
# kompress_compressor.py Missing includes 805-806, 818, 836
# after: none of those lines appear in Missing anymore
191 passed, 7 skipped
$ .venv/bin/python3 -m ruff format --check tests/test_onnx_runtime.py tests/test_transforms/test_kompress_compressor.py
2 files already formatted
$ .venv/bin/python3 -m ruff check tests/test_onnx_runtime.py tests/test_transforms/test_kompress_compressor.py
All checks passed!
```
## Review Readiness
- Test-only, additive diff (113 insertions, 0 deletions, 0 lines touched
outside the two test
files). No behavior change possible.
## Checklist
- [x] My code follows the project's style guidelines
- [x] I have performed a self-review of my code
- [x] New and existing unit tests pass locally with my changes
- [x] I did not edit `CHANGELOG.md`
## Additional Notes
Not closing: the remaining branch-partial on the `device == "auto"`
cuda/mps/cpu selection in
`_load_kompress_pytorch` (would need mocking `torch.cuda.is_available()`
/
`torch.backends.mps.is_available()` for marginal benefit); left as-is.
# PR draft: fix(kompress): load merged.pt for the v2 checkpoint instead
of the unmerged PEFT safetensors
Branch: `rnoz/fix-kompress-merged-checkpoint` (off `upstream/main`). Two
commits.
Issue: https://github.com/headroomlabs-ai/headroom/issues/2714 (filed,
open).
---
## Description
`_load_kompress_pytorch` in `headroom/transforms/kompress_compressor.py`
downloaded `model.safetensors` from the default model repo
`chopratejas/kompress-v2-base` and loaded it with `strict=False`,
discarding the missing/unexpected key report. That file is the unmerged
PEFT checkpoint (encoder keys prefixed `encoder.base_model.model...`),
which never matches `HeadroomCompressorModel`'s plain `encoder.*` keys.
The LoRA-adapted encoder weights were silently dropped while
`token_head`/`span_conv` happened to match and loaded fine, so the model
ran with a stock, non-adapted `answerdotai/ModernBERT-base` encoder
feeding correctly trained decision heads, with no error and a healthy
status reported everywhere.
`scripts/export_kompress_v2_onnx.py` already documents this exact
mismatch and loads the correct `merged.pt` sub-state-dicts for its own
export path. This PR mirrors that same loading logic into the runtime
PyTorch loader, with a fallback to the plain `model.safetensors` format
for repos that never shipped a `merged.pt` (verified against the v1
`chopratejas/kompress-base` repo via the public HF API, which has no
`merged.pt`). Both paths now check the missing/unexpected key report and
raise instead of silently proceeding on a mismatch.
A second commit fixes a gap an adversarial review caught in the first:
the cache-only (`allow_download=False`, startup preload) path could not
tell "this repo genuinely has no merged.pt" apart from "merged.pt exists
but is not downloaded yet", so it would have fallen back to a stale
`model.safetensors` left over from before this fix on exactly the
upgrade path this PR is meant to close. It now uses `huggingface_hub`'s
own `.no_exist` cache marker (via a new `hf_entry_known_absent()` helper
in `headroom/onnx_runtime.py`) to make that distinction without a
network call, and only falls back when absence is confirmed; otherwise
it raises `KompressModelNotCached` so the caller defers instead of
guessing.
Closes#2714
## Type of Change
- [x] Bug fix (non-breaking change that fixes an issue)
## Changes Made
- `headroom/transforms/kompress_compressor.py`: added
`_load_merged_state_dict`, `_load_plain_state_dict`, and
`_load_pytorch_weights`, replacing the inline `model.safetensors`
download + `load_state_dict(strict=False)` call in
`_load_kompress_pytorch`. `merged.pt` is tried first; the plain format
is only used when its absence is confirmed.
- `headroom/onnx_runtime.py`: added `hf_entry_known_absent()`, a thin
wrapper around `huggingface_hub.try_to_load_from_cache()` that reads the
on-disk `.no_exist` marker HF writes after a real 404, so cache-only
code can distinguish "confirmed absent" from "never checked" without
hitting the network.
- `tests/test_transforms/test_kompress_compressor.py`: added
`TestPytorchWeightLoading` (8 tests) covering the merged-checkpoint
happy path, missing-section and key-mismatch failures, the plain-format
fallback for repos without `merged.pt`, and the cache-only ambiguity fix
(confirmed-absent vs unconfirmed).
## Testing
- [x] Unit tests pass (`pytest`)
- [x] Linting passes (`ruff check`)
- [x] Type checking passes (`mypy`)
- [x] Manual testing performed (real download against the live
`chopratejas/kompress-v2-base` repo, not just mocks)
### Test Output
```text
$ .venv/bin/python3 -m pytest tests/ -k kompress -q
178 passed, 7 skipped in 22.39s
$ .venv/bin/python3 -m ruff check headroom/onnx_runtime.py headroom/transforms/kompress_compressor.py tests/test_transforms/test_kompress_compressor.py
All checks passed!
$ .venv/bin/python3 -m ruff format --check headroom/onnx_runtime.py headroom/transforms/kompress_compressor.py tests/test_transforms/test_kompress_compressor.py
3 files already formatted
$ .venv/bin/python3 -m mypy headroom/onnx_runtime.py headroom/transforms/kompress_compressor.py
Success: no issues found in 2 source files
```
## Real Behavior Proof
Ran the actual fixed loader against the live
`chopratejas/kompress-v2-base` HF repo (not a mock), before and after
each fix:
```text
# BEFORE (parsed the real cached model.safetensors header by hand, no safetensors lib needed):
total tensors: 316
encoder.base_model.model.embeddings.norm.weight
encoder.base_model.model.embeddings.tok_embeddings.weight
...(all 310 encoder tensors share this prefix)...
span_conv.0.bias / span_conv.0.weight / span_conv.2.bias / span_conv.2.weight
token_head.bias / token_head.weight
exact prefix match count with plain "encoder.<rest>": 0
# This confirms the pre-fix code's model.load_state_dict(state_dict, strict=False)
# silently dropped every encoder weight (0 keys match HeadroomCompressorModel.encoder),
# while token_head/span_conv happened to match and loaded.
# AFTER (commit 1, real merged.pt download + load):
$ .venv/bin/python3 -c "
import headroom.transforms.kompress_compressor as kmod
model = kmod._get_model_class()()
kmod._load_pytorch_weights(model, 'chopratejas/kompress-v2-base', allow_download=True)
print('SUCCESS: 0 missing/unexpected keys across all three sections')
"
SUCCESS: 0 missing/unexpected keys across all three sections
# AFTER (full pipeline, real end-to-end compression through the public API):
$ .venv/bin/python3 -c "
import headroom.transforms.kompress_compressor as kmod
compressor = kmod.KompressCompressor()
result = compressor.compress(sample_traceback_plus_boilerplate_text)
print(result.original_tokens, result.compressed_tokens, result.tokens_saved)
print('ValueError: bad input' in result.compressed)
"
497 454 43
True # must-keep line (the actual error) survived compression
# AFTER (commit 2, cache-only ambiguity): unit tests
# test_cache_only_defers_instead_of_using_stale_plain_checkpoint: PASSED
# test_cache_only_uses_plain_checkpoint_when_merged_pt_confirmed_absent: PASSED
```
## Review Readiness
- [x] I have performed a self-review
- [x] An independent adversarial review pass was run on both commits
before this PR was opened; its one finding (the cache-only ambiguity) is
fixed in commit 2, verified with new regression tests
- [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 (N/A:
internal loader behavior, no public API or config surface changed)
- [x] My changes generate no new warnings
- [x] I have added tests that prove my fix is effective (regression
tests for both the original silent-drop bug and the cache-only ambiguity
found in review)
- [x] New and existing unit tests pass locally with my changes
- [x] I did not edit `CHANGELOG.md`
## Additional Notes
- Both `merged.pt` and the plain `model.safetensors` fallback now raise
loudly on any state-dict mismatch instead of proceeding with
partially-loaded weights, closing the general silent-failure class this
bug belonged to, not just this one instance of it.
- No other call sites of `_load_kompress_pytorch` or its removed inline
code exist; its public signature is unchanged.
## Description
With `HEADROOM_KOMPRESS_BACKEND=onnx_coreml`, every Kompress compression
call and the startup canary crash with `'_OnnxModel' object has no
attribute 'parameters'`, so Kompress silently degrades to passthrough
and `/health` reports `kompress: unhealthy, backend: null`.
Root cause: `headroom/transforms/kompress_compressor.py` gated the
ONNX-vs-PyTorch branch with an exact string match `backend == "onnx"`.
But `_load_kompress_onnx` returns `onnx_coreml` (CoreML) or `onnx_cpu` —
never the bare string `onnx`. So under `onnx_coreml` the code built
PyTorch tensors and dispatched to a device via
`next(model.parameters())`, which the `_OnnxModel` wrapper doesn't
implement. This is the accelerated backend Apple Silicon users reach
for, so the fast path is exactly the broken one.
Fixes#2442
## 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
- Change the four exact-match `backend == "onnx"` sites in
`headroom/transforms/kompress_compressor.py` to
`backend.startswith("onnx")`, matching the convention already used by
`_model_device_type`: `_timed_canary`, `compress`, `compress_batch`, and
the batch-parallelism guard in `_should_use_sequential_fallback`.
- Update the guard comment ("ONNX CPU provider" → "ONNX EPs") since it
now covers all ONNX execution providers.
- Add regression tests exercising `_timed_canary` on `onnx_coreml` (must
take the numpy path and never touch `.parameters()`) with a negative
control proving the PyTorch branch still dispatches to a device.
- Leave `CHANGELOG.md` untouched — release-please generates it from
conventional commits.
- Out of scope: the secondary `/health` under-reporting the issue flags
as informational (deferred-preload warmup object never flips to
`loaded`).
## Testing
- [x] Unit tests pass (`python -m pytest
tests/test_transforms/test_kompress_compressor.py::TestOnnxBackendPrefixGating
-q`)
- [x] Linting passes (`ruff check`, `ruff format --check` on the two
changed files)
- [x] Type checking passes (`mypy
headroom/transforms/kompress_compressor.py --ignore-missing-imports`)
- [x] New tests added for new functionality
### Test Output
```text
$ python -m pytest tests/test_transforms/test_kompress_compressor.py::TestOnnxBackendPrefixGating -q
collected 2 items
tests\test_transforms\test_kompress_compressor.py .. [100%]
2 passed in 2.20s
$ ruff check headroom/transforms/kompress_compressor.py tests/test_transforms/test_kompress_compressor.py
All checks passed!
```
## Real Behavior Proof
- Environment: Windows 11, Python 3.13, local dev checkout on a branch
off upstream/main (no Apple Silicon / CoreML hardware available)
- Exact command / steps: Ran the new `TestOnnxBackendPrefixGating`
regression; then temporarily reverted one site back to `backend ==
"onnx"` and re-ran to confirm the test discriminates.
- Observed result: With the fix, `_timed_canary(model, tokenizer,
"onnx_coreml")` returns a float and never touches `.parameters()`.
Reverting one site makes the onnx_coreml test fail (it takes the `pt`
tensor path and hits the paramless model), proving the test catches the
exact bug. The issue reporter separately verified the fix on real Apple
Silicon hardware (onnxruntime 1.27.0, CoreMLExecutionProvider): zero
occurrences of the error afterward and compression completing on the
CoreML session.
- Not tested: End-to-end run on real CoreML hardware from this
environment — reproduced via the unit-level device-dispatch seam
instead; hardware confirmation is in the issue.
## Review Readiness
- [x] I have performed a self-review
- [x] This PR is ready for human review
Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
## Description
`ContentRouter.eager_load_compressors()` runs a network
`hf_hub_download` of the Kompress ONNX model on the **blocking
startup/lifespan path**, before the proxy binds its port. On a cold
cache this is unsafe:
- the download can hang long enough to blow the supervisor's bind
timeout, or
- a native crash in the download/ML stack (an **uncatchable `Fatal
Python error: Aborted` / SIGABRT**) kills the interpreter before it ever
`listen()`s.
Either way the supervisor sees "proxy never opened its port" and gives
up. We observed this in the field from the desktop app (process aborted
during `eager_load_compressors -> _load_kompress_onnx ->
hf_hub_download` of `onnx/kompress-int8.onnx`, while the only Python
thread was parked in the HuggingFace download file-lock; the abort came
from a native thread, so `try/except` at the call site cannot catch it).
The eager preload is a latency optimization and must never be able to
block — or kill — startup. This change makes startup preload
**cache-only**: if the model isn't already cached, we defer the download
to first use (off the startup path) and bind the port normally. Warm
starts are 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
- `onnx_runtime.hf_hub_download_local_first(...)`: added `allow_network`
(default `True`). When `False`, a cache miss re-raises the local-lookup
error instead of falling back to a network download.
- `kompress_compressor`: added `allow_download` (default `True`)
threaded through `preload()` -> `_load_kompress()` ->
`_load_kompress_onnx()` / `_load_kompress_pytorch()` and the ModernBERT
tokenizer load. Added `KompressModelNotCached`, raised when a cache-only
load misses. Auto-mode no longer falls back to a PyTorch network
download on a cache-only miss — it propagates so the caller can defer.
- `content_router.eager_load_compressors()`: calls
`preload(allow_download=False)`. On `KompressModelNotCached` it logs and
reports the component as `"deferred"` (a status
`warmup.merge_transform_status` already handles gracefully) instead of
letting a cold download run on the startup path.
Default (first-request) loading behavior and warm-start preload are
unchanged.
## 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
New tests in `tests/test_kompress_preload_deferral.py` cover: cache-only
`hf_hub_download_local_first` never hits the network; default still
falls back; cache-only ONNX load raises `KompressModelNotCached`;
auto-mode does **not** trigger a PyTorch download on a cache-only miss;
and `eager_load_compressors` reports `deferred` (cold) / `enabled`
(warm). Existing `_load_kompress` dispatch tests updated for the new
keyword-only param.
> Note on environment: I do not have a clean reproduction of the native
SIGABRT itself (it depends on a specific machine's HF download/ML native
stack), so the "Manual testing performed" box is left unchecked. The
tests target the structural fix — that startup preload can no longer
perform a network download — which is the precondition for the crash.
## Test Output
```
$ uv run pytest -v tests/test_kompress_preload_deferral.py
tests/test_kompress_preload_deferral.py::test_local_first_no_network_when_disallowed PASSED
tests/test_kompress_preload_deferral.py::test_local_first_falls_back_to_network_by_default PASSED
tests/test_kompress_preload_deferral.py::test_load_kompress_onnx_cache_miss_raises_not_cached PASSED
tests/test_kompress_preload_deferral.py::test_load_kompress_auto_does_not_pytorch_download_on_cache_miss PASSED
tests/test_kompress_preload_deferral.py::test_eager_load_defers_when_model_not_cached PASSED
tests/test_kompress_preload_deferral.py::test_eager_load_enabled_when_model_cached PASSED
6 passed in 4.82s
$ uv run pytest tests/test_transforms/test_kompress_compressor.py tests/test_transforms_content_router.py tests/test_onnx_runtime.py tests/test_proxy_warmup.py
63 passed
$ uv run ruff check <changed files> # All checks passed!
$ uv run mypy headroom/onnx_runtime.py headroom/transforms/kompress_compressor.py headroom/transforms/content_router.py
Success: no issues found
```
## 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 (auto-generated from
conventional commits)
## Additional Notes
This contains the cold-start case. A native crash in onnxruntime
*session init* (as opposed to the download) on first request would still
be a separate issue; it is not what was observed here (the abort was
during the HF download), and isolating it would be a larger, separate
change.
---------
Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
## Summary
This PR was originally \"HEADROOM_KOMPRESS_BACKEND env + GPU/MPS
auto-detect\" (for #202). While it sat, main independently shipped the
backend-selection env var in a2ea9648 (\"fix: add Kompress backend and
thread controls\") with a richer backend set (`auto` / `onnx` /
`onnx_cpu` / `onnx_coreml` / `pytorch` / `pytorch_mps` + shorthand
aliases) and an explicit design decision to keep `auto` on the
ONNX-CPU-first path rather than auto-preferring accelerators. Rather
than re-litigate that, this PR has been rebased onto latest main and
rescoped to the two pieces main still lacks:
1. **Warn on unrecognized `HEADROOM_KOMPRESS_BACKEND` values** —
previously typos (`gpu`, `cudaa`, …) silently mapped to `auto`,
indistinguishable from the default. Now a warning names the offending
value and the accepted set; behavior still falls back to `auto`.
2. **Documentation** — the env var and its six backends/aliases were
undocumented outside the source. Added a \"Kompress backend selection\"
section to `wiki/configuration.md` and a CHANGELOG entry.
## Testing
- `pytest tests/test_transforms/test_kompress_compressor.py` — 28 passed
(includes 2 new tests: warning fires on unrecognized value; valid values
and unset stay silent)
- `ruff check` / `ruff format` clean on touched files
- No behavior change beyond the new warning, so no GPU/MPS hardware
validation is required for this scope.
Refs #202🤖 Generated with [Claude Code](https://claude.com/claude-code)
---------
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
KompressConfig now accepts model_id, chunk_words, and score_threshold so
domain-specific models (e.g. kompress-finance with 50-word chunks) can be
used without forking the compressor. Model cache is keyed by model_id,
allowing multiple models to coexist. All defaults match prior behavior.
Also fix mypy errors in memory/sync.py from recent merge.
Implements compress_batch() for issue #151. Compresses N texts with
batched forward passes on GPU and falls back to sequential compress()
on CPU where batching doesn't help.
Measured performance (RTX 3080 Ti, 1000-word / ~6K-char inputs):
GPU (PyTorch + CUDA):
N=1: 2.68x speedup (multi-chunk text batches within single call)
N=5: 2.75x speedup
N=12: 2.49x speedup
CPU (ONNX): fallback to sequential — parity with compress() in loop
ONNX Runtime's CPU execution provider does not parallelize across the
batch dimension for this model architecture; verified across default,
physical-cores-only, and single-thread configurations. The fallback
keeps the API useful while that limitation exists.
Features:
- Per-item target_ratio: scalar applies to all, list allows per-text
- Input order preserved in output
- Passthrough parity with compress() on short texts / errors
- Configurable batch_size (default 32)
Tests: 8 new (TestKompressCompressorBatch), 21 total pass.
Closes#151
KompressCompressor now tries ONNX Runtime first (156MB INT8 model),
falls back to PyTorch only if ONNX unavailable. No torch needed for
text compression — just onnxruntime (~50MB) + transformers (tokenizer).
Changes:
- Add onnxruntime + transformers to [proxy] extra in pyproject.toml
- Add _OnnxModel wrapper with get_scores/get_keep_mask interface
- _load_kompress() tries ONNX first, falls back to PyTorch
- is_kompress_available() returns True if EITHER backend available
- compress() handles both numpy (ONNX) and tensor (PyTorch) outputs
Dependency impact:
Before: pip install headroom-ai[proxy] → no text compression
After: pip install headroom-ai[proxy] → Kompress ONNX INT8 (156MB)
[ml] extra still available for full PyTorch (600MB, GPU support)
The proxy startup crashed with `ModuleNotFoundError: No module named
'torch'` when installed with just `[proxy]` extras because
kompress_compressor.py had unconditional top-level torch imports.
Moved torch/transformers imports to be lazy so the module is safely
importable without the [ml] extra. Added tests for import safety.
Bumped version to 0.4.5
