feat(kompress): warn on unrecognized HEADROOM_KOMPRESS_BACKEND + document backend selection (#204)

## 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>
This commit is contained in:
mbachaud 2026-06-10 18:13:17 -07:00 committed by GitHub
parent 35b46d6e84
commit 6367d0b722
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 79 additions and 1 deletions

View file

@ -35,6 +35,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
### Added
* **kompress:** warn when `HEADROOM_KOMPRESS_BACKEND` is set to an unrecognized
value instead of silently falling back to `auto`, and document the backend
selection env var (`auto` / `onnx` / `onnx_cpu` / `onnx_coreml` / `pytorch` /
`pytorch_mps` plus shorthand aliases) in `wiki/configuration.md` (issue
[#202](https://github.com/chopratejas/headroom/issues/202), PR
[#204](https://github.com/chopratejas/headroom/pull/204)).
* **proxy:** per-provider attribution in the savings history rollups. Each `/stats-history` bucket (hourly/daily/weekly/monthly) now carries a `by_provider` map breaking down `tokens_saved`, `compression_savings_usd_delta`, `total_input_tokens_delta`, and `total_input_cost_usd_delta` per provider, so consumers can show how savings and spend are distributed across providers within a time period. Providers only appear in a bucket where they moved a counter; legacy history checkpoints with no provider collapse into `"unknown"`. Affected files: `headroom/proxy/savings_tracker.py`, `headroom/proxy/prometheus_metrics.py`.
### Changed

View file

@ -90,7 +90,16 @@ def _selected_backend() -> KompressBackend:
"pytorch_mps": "pytorch_mps",
"auto": "auto",
}
return aliases.get(raw, "auto") # type: ignore[return-value]
backend = aliases.get(raw)
if backend is None:
logger.warning(
"%s has unrecognized value %r; falling back to 'auto'. Valid values: %s",
KOMPRESS_BACKEND_ENV,
os.environ.get(KOMPRESS_BACKEND_ENV, ""),
", ".join(sorted(set(aliases.values()))),
)
return "auto"
return backend # type: ignore[return-value]
def _env_int(name: str) -> int | None:

View file

@ -8,6 +8,7 @@ Covers:
- Transform interface: apply() method
"""
import logging
from types import SimpleNamespace
from unittest.mock import MagicMock, patch
@ -83,6 +84,30 @@ class TestKompressBackendSelection:
monkeypatch.setenv("HEADROOM_KOMPRESS_BACKEND", "unknown")
assert kmod._selected_backend() == "auto"
def test_unrecognized_backend_warns_and_falls_back_to_auto(self, monkeypatch, caplog) -> None:
import headroom.transforms.kompress_compressor as kmod
monkeypatch.setenv("HEADROOM_KOMPRESS_BACKEND", "tpu")
with caplog.at_level(logging.WARNING, logger=kmod.logger.name):
assert kmod._selected_backend() == "auto"
assert any(
"unrecognized" in record.getMessage() and "tpu" in record.getMessage()
for record in caplog.records
)
def test_valid_backend_values_do_not_warn(self, monkeypatch, caplog) -> None:
import headroom.transforms.kompress_compressor as kmod
with caplog.at_level(logging.WARNING, logger=kmod.logger.name):
for value in ("auto", "onnx", "cpu", "coreml", "mps", "torch", "ONNX-CPU"):
monkeypatch.setenv("HEADROOM_KOMPRESS_BACKEND", value)
kmod._selected_backend()
monkeypatch.delenv("HEADROOM_KOMPRESS_BACKEND", raising=False)
kmod._selected_backend()
assert not caplog.records
def test_forced_pytorch_mps_backend_uses_mps_device(self, monkeypatch) -> None:
import headroom.transforms.kompress_compressor as kmod

View file

@ -70,6 +70,44 @@ headroom proxy --no-ccr-expansion
headroom proxy --help
```
### Kompress backend selection
Kompress (the model-based compressor) can run on two engines:
- **ONNX Runtime** — lightweight, CPU-first. Installed with
`pip install headroom-ai[proxy]`. Optionally uses the CoreML execution
provider on macOS.
- **PyTorch** — heavier, supports CUDA and Apple-Silicon MPS
acceleration. Installed with `pip install headroom-ai[ml]`. With
`device=auto` it selects `cuda`, then `mps`, then `cpu`.
Select the backend via the `HEADROOM_KOMPRESS_BACKEND` environment
variable:
| Value | Behavior |
|---------------------|------------------------------------------------------------------------|
| `auto` | Default. ONNX CPU first (stable, lightweight), PyTorch as fallback. |
| `onnx` / `onnx_cpu` | Force ONNX Runtime on CPU. |
| `onnx_coreml` | Force ONNX Runtime with the CoreML provider (CPU fallback). |
| `pytorch` | Force PyTorch with automatic device selection (CUDA → MPS → CPU). |
| `pytorch_mps` | Force PyTorch on Apple-Silicon MPS; falls back to ONNX CPU on failure. |
Values are case-insensitive and hyphens are accepted (`onnx-cpu` ==
`onnx_cpu`). Shorthand aliases: `cpu``onnx_cpu`, `coreml`
`onnx_coreml`, `mps` / `torch_mps``pytorch_mps`, `torch`
`pytorch`. Unrecognized values log a warning and fall back to `auto`.
Example — opt in to MPS on an Apple-Silicon machine:
```bash
export HEADROOM_KOMPRESS_BACKEND=mps
headroom proxy ...
```
The default deliberately stays on ONNX CPU so existing installs keep
their compression quality and performance characteristics; accelerator
backends are opt-in.
## Per-Request Overrides
Override configuration for specific requests: