akadmin/headroom
1
0
Fork 0
mirror of https://github.com/headroomlabs-ai/headroom.git synced 2026-08-27 14:17:10 -04:00
Code Issues Projects Releases Packages Wiki Activity Actions

fix(kompress): let orgs run Kompress on their own inference stack (#2736)

Browse source 
## What this enables

An org pulls the Kompress weights from HuggingFace, serves them on their
own infrastructure, and points Headroom at it:

```bash
HEADROOM_KOMPRESS_ENDPOINT=https://ml.internal.acme.com
```

No credential needed, no local ML dependencies, and original content
never leaves their network (the CCR store stays proxy-local, so
`headroom_retrieve` keeps working).

## The one thing that was actually broken

Almost all of this already worked. The blocker was a hardcoded path:

```python
self._url = endpoint.rstrip("/") + "/compress"
```

Real inference servers don't serve at `/compress`:

| Stack | Path |
|---|---|
| TorchServe | `/predictions/kompress` |
| KServe / Seldon | `/v1/models/kompress:predict` |
| SageMaker | `/invocations` |

Appending `/compress` to those 404s. And because remote Kompress **fails
open**, that 404 is invisible — compression silently stops instead of
erroring. The only workaround was standing up a reverse proxy purely to
rename a path.

## Two new env vars, both defaulting to current behaviour

| Var | Default | Purpose |
|---|---|---|
| `HEADROOM_KOMPRESS_ENDPOINT_PATH` | `/compress` | Set empty to use the
endpoint URL verbatim |
| `HEADROOM_KOMPRESS_ENDPOINT_HEADERS` | *(none)* | `k=v,k2=v2`, merged
last so it can replace `Authorization` |

Headers are applied after the token deliberately, so a gateway wanting
`x-api-key` or `X-Tenant-Id` needs no separate auth-scheme setting.

## No regression

With only `HEADROOM_KOMPRESS_ENDPOINT` set, the request is
**byte-identical** to before — `POST <endpoint>/compress` with an
optional Bearer token. Existing Modal deployments need no change.

`os.environ.get` with a default distinguishes "unset" (use `/compress`)
from an explicit empty value (endpoint is a complete URL), so the escape
hatch can't fire by accident. The regression cases are deliberately the
*first* tests in the new file.

Verified through the real router wiring:

```
modal (today's config)           -> https://acme--kompress.modal.run/compress
modal + token                    -> …/compress  {'authorization': 'Bearer tok'}
self-hosted KServe (full URL)    -> https://ml.acme.com/v1/models/kompress:predict
self-hosted TorchServe (path)    -> https://ts.acme.com/predictions/kompress
self-hosted, x-api-key, no token -> …/compress  {'x-api-key': 'k', 'x-tenant-id': 'acme'}
```

## Documents the HTTP contract

The endpoint contract was only discoverable by reading the source. Now
in the module docstring:

```
request   {"content": "<text>", "target_ratio": 0.5 | null}
response  {"compressed": "<text>",       # REQUIRED
           "original_tokens": int,        # optional, derived if absent
           "compressed_tokens": int,      # optional
           "compression_ratio": float,    # optional
           "model_used": str}             # optional
```

`compressed` is the only required field, so a shim in front of an
existing inference server is a few lines.

Also logs the **resolved** URL at startup — with fail-open, a mistyped
path otherwise manifests as nothing happening at all.

## Notes

- `parse_endpoint_headers` reimplements the
`HEADROOM_OTEL_METRICS_HEADERS` format rather than importing it:
`observability.metrics` imports opentelemetry at module scope, and
remote Kompress exists precisely so a proxy can run without heavy
optional deps.
- 27 new tests. Pre-existing unrelated flake in
`test_content_router_single_item_deadline.py` (fails 3/3 on clean main).

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
This commit is contained in:
Tejas Chopra 2026-08-03 07:26:06 -07:00 • committed by GitHub
parent 7c9b046595
commit 3d23d76248
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 248 additions and 5 deletions
Download patch file Download diff file Expand all files Collapse all files

16
headroom/transforms/content_router.py
View file

@ -4230,14 +4230,26 @@ class ContentRouter(Transform):
return None
if getattr(self, "_kompress_remote", None) is None:
from .kompress_compressor import KompressConfig
from .kompress_remote import RemoteKompressCompressor
from .kompress_remote import (
DEFAULT_ENDPOINT_PATH,
RemoteKompressCompressor,
parse_endpoint_headers,
)
# Defaults reproduce the previous behaviour exactly, so existing
# (Modal) deployments are unaffected: os.environ.get with a default
# distinguishes "unset" (use /compress) from an explicit empty value
# (the operator's endpoint is already a complete URL).
self._kompress_remote = RemoteKompressCompressor(
endpoint=endpoint,
token=os.environ.get("HEADROOM_KOMPRESS_ENDPOINT_TOKEN") or None,
config=KompressConfig(enable_ccr=self.config.ccr_inject_marker),
path=os.environ.get("HEADROOM_KOMPRESS_ENDPOINT_PATH", DEFAULT_ENDPOINT_PATH),
headers=parse_endpoint_headers(
os.environ.get("HEADROOM_KOMPRESS_ENDPOINT_HEADERS")
),
)
logger.info("Kompress: using remote endpoint %s", endpoint)
logger.info("Kompress: using remote endpoint %s", self._kompress_remote.url)
return self._kompress_remote
def _get_image_optimizer(self) -> Any:

96
headroom/transforms/kompress_remote.py
View file

@ -10,8 +10,45 @@ Only the model inference is remote. The CCR store + retrieval marker stay
proxy-local (the endpoint is stateless, ``enable_ccr=False``), so
``headroom_retrieve`` keeps working and original content never persists off-box.
Enabled by ``HEADROOM_KOMPRESS_ENDPOINT`` (+ optional
``HEADROOM_KOMPRESS_ENDPOINT_TOKEN``) — see ``ContentRouter._get_kompress``.
Enabled by ``HEADROOM_KOMPRESS_ENDPOINT`` — see ``ContentRouter._get_kompress``.
# Bring-your-own deployment
The endpoint does not have to be Headroom Labs'. An org can pull the Kompress
weights from HuggingFace, serve them on its own stack (vLLM, TorchServe,
SageMaker, KServe, a bare FastAPI box) and point Headroom at it. Nothing about
this class is Modal-specific, and no credential is required — auth is whatever
the operator's own infrastructure expects, including none at all:
HEADROOM_KOMPRESS_ENDPOINT https://ml.internal.acme.com
HEADROOM_KOMPRESS_ENDPOINT_PATH /compress (default; set empty to use
the endpoint URL verbatim)
HEADROOM_KOMPRESS_ENDPOINT_TOKEN optional; sent as `Authorization: Bearer`
HEADROOM_KOMPRESS_ENDPOINT_HEADERS optional; `k=v,k2=v2`, applied last so it
can replace the Authorization header for
stacks that want `x-api-key` or similar
Both new knobs default to today's behaviour: with only
``HEADROOM_KOMPRESS_ENDPOINT`` set, the request is byte-identical to before —
``POST <endpoint>/compress`` with an optional Bearer token. Existing Modal
deployments need no change.
# The HTTP contract
Deliberately small, so a shim in front of an existing inference server is a few
lines. ``POST <endpoint><path>``:
request {"content": "<text>", "target_ratio": 0.5 | null}
response {"compressed": "<text>", # REQUIRED, must be a string
"original_tokens": int, # optional, defaults to word count
"compressed_tokens": int, # optional, defaults to word count
"compression_ratio": float, # optional, defaults to 1.0
"model_used": str} # optional
``compressed`` is the only required field; every other value is derived if
absent. Any non-2xx, timeout, malformed field, or missing ``compressed`` makes
this pass the content through verbatim — a broken endpoint costs compression,
never correctness.
"""
from __future__ import annotations
@ -24,6 +61,33 @@ from .kompress_compressor import KompressConfig, KompressResult, store_kompress_
logger = logging.getLogger(__name__)
# Appended to the configured endpoint unless overridden. An operator whose stack
# already serves a full path (``/v1/models/kompress:predict``) sets
# HEADROOM_KOMPRESS_ENDPOINT_PATH="" and gives the complete URL instead.
DEFAULT_ENDPOINT_PATH = "/compress"
def parse_endpoint_headers(raw: str | None) -> dict[str, str]:
"""Parse ``k=v,k2=v2`` into a header dict.
Same format as ``HEADROOM_OTEL_METRICS_HEADERS`` so operators meet one
convention. Reimplemented rather than imported from
:mod:`headroom.observability.metrics`, which pulls in opentelemetry at module
scope — remote Kompress exists precisely so a proxy can run without heavy
optional deps, so it must not drag one in through a parsing helper.
"""
pairs: dict[str, str] = {}
for item in (raw or "").split(","):
part = item.strip()
if not part or "=" not in part:
continue
key, _, value = part.partition("=")
key, value = key.strip(), value.strip()
if key and value:
pairs[key] = value
return pairs
# Below this word count local Kompress passes through verbatim (KompressCompressor
# .compress); mirror it so we never pay a round-trip on a trivially small block.
_MIN_WORDS = 10
@ -48,15 +112,41 @@ class RemoteKompressCompressor:
token: str | None = None,
config: KompressConfig | None = None,
timeout: float = 20.0,
path: str | None = DEFAULT_ENDPOINT_PATH,
headers: dict[str, str] | None = None,
) -> None:
self.config = config or KompressConfig()
self._url = endpoint.rstrip("/") + "/compress"
# Default keeps the pre-existing behaviour exactly: <endpoint>/compress.
# An empty (or None) path means the caller has supplied a complete URL —
# needed because real inference servers do not serve at /compress
# (TorchServe /predictions/<model>, KServe /v1/models/<name>:predict,
# SageMaker /invocations), and appending to those yields a 404.
if path:
suffix = path if path.startswith("/") else "/" + path
self._url = endpoint.rstrip("/") + suffix
else:
self._url = endpoint
self._headers = {"content-type": "application/json"}
if token:
self._headers["authorization"] = f"Bearer {token}"
# Applied last on purpose: lets an operator replace `authorization` with
# whatever their gateway wants (x-api-key, a signed header, a tenant id)
# without needing a separate auth-scheme setting.
if headers:
self._headers.update(headers)
# httpx.Client is safe to share across the proxy's worker threads.
self._client = httpx.Client(timeout=timeout)
@property
def url(self) -> str:
"""The resolved POST target.
Logged when the router builds this, so a wrong ``_PATH`` shows up as a
visibly odd URL at startup rather than as silent pass-through later —
the fail-open contract means a 404 never surfaces as an error.
"""
return self._url
# Nothing to load locally; short-circuit the router straight to compress().
def is_ready(self) -> bool:
return True

141
tests/test_kompress_remote_endpoint.py Normal file
View file

@ -0,0 +1,141 @@
"""Endpoint resolution for bring-your-own Kompress deployments.
The load-bearing test here is the first one: an operator with only
``HEADROOM_KOMPRESS_ENDPOINT`` set must get the exact same request as before
these knobs existed. Everything else is additive.
"""
from __future__ import annotations
import pytest
from headroom.transforms.kompress_remote import (
DEFAULT_ENDPOINT_PATH,
RemoteKompressCompressor,
parse_endpoint_headers,
)
class TestNoRegressionForExistingDeployments:
"""Modal users set one env var and must be unaffected."""
def test_default_appends_compress(self):
c = RemoteKompressCompressor(endpoint="https://acme--kompress.modal.run")
assert c.url == "https://acme--kompress.modal.run/compress"
def test_trailing_slash_does_not_double_up(self):
c = RemoteKompressCompressor(endpoint="https://acme--kompress.modal.run/")
assert c.url == "https://acme--kompress.modal.run/compress"
def test_token_still_sent_as_bearer(self):
c = RemoteKompressCompressor(endpoint="https://x.modal.run", token="secret")
assert c._headers["authorization"] == "Bearer secret"
assert c._headers["content-type"] == "application/json"
def test_no_token_means_no_auth_header(self):
"""Self-hosted stacks frequently need no credential at all."""
c = RemoteKompressCompressor(endpoint="https://ml.internal")
assert "authorization" not in c._headers
def test_default_path_constant_is_the_historical_value(self):
assert DEFAULT_ENDPOINT_PATH == "/compress"
class TestSelfHostedPaths:
"""Real inference servers do not serve at /compress."""
@pytest.mark.parametrize(
"endpoint,path,expected",
[
# KServe / Seldon
(
"https://ml.acme.com",
"/v1/models/kompress:predict",
"https://ml.acme.com/v1/models/kompress:predict",
),
# TorchServe
(
"https://torchserve.acme.com",
"/predictions/kompress",
"https://torchserve.acme.com/predictions/kompress",
),
# SageMaker
(
"https://runtime.sagemaker.internal",
"/invocations",
"https://runtime.sagemaker.internal/invocations",
),
# A leading slash is optional in the env var.
("https://ml.acme.com", "invocations", "https://ml.acme.com/invocations"),
# Endpoint with its own base path, plus a suffix.
(
"https://gw.acme.com/kompress",
"/compress",
"https://gw.acme.com/kompress/compress",
),
],
)
def test_path_override(self, endpoint, path, expected):
assert RemoteKompressCompressor(endpoint=endpoint, path=path).url == expected
@pytest.mark.parametrize("empty", ["", None])
def test_empty_path_uses_the_url_verbatim(self, empty):
"""The escape hatch: the operator supplies a complete URL.
Without this, an endpoint that is already a full path gets /compress
appended and 404s — and because remote Kompress fails open, that 404 is
invisible: compression silently stops instead of erroring.
"""
url = "https://ml.acme.com/v1/models/kompress:predict"
assert RemoteKompressCompressor(endpoint=url, path=empty).url == url
def test_verbatim_url_keeps_its_trailing_slash_untouched(self):
url = "https://ml.acme.com/predict/"
assert RemoteKompressCompressor(endpoint=url, path="").url == url
class TestCustomHeaders:
def test_extra_headers_are_merged(self):
c = RemoteKompressCompressor(
endpoint="https://ml.acme.com",
headers={"x-tenant-id": "acme", "x-env": "prod"},
)
assert c._headers["x-tenant-id"] == "acme"
assert c._headers["x-env"] == "prod"
assert c._headers["content-type"] == "application/json"
def test_headers_can_replace_the_bearer_scheme(self):
"""A gateway wanting x-api-key should not need a new setting."""
c = RemoteKompressCompressor(
endpoint="https://ml.acme.com",
token="ignored",
headers={"authorization": "Token abc123"},
)
assert c._headers["authorization"] == "Token abc123"
def test_api_key_header_without_any_token(self):
c = RemoteKompressCompressor(endpoint="https://ml.acme.com", headers={"x-api-key": "k"})
assert c._headers["x-api-key"] == "k"
assert "authorization" not in c._headers
class TestHeaderParsing:
@pytest.mark.parametrize(
"raw,expected",
[
(None, {}),
("", {}),
(" ", {}),
("x-api-key=abc", {"x-api-key": "abc"}),
("a=1,b=2", {"a": "1", "b": "2"}),
(" a = 1 , b = 2 ", {"a": "1", "b": "2"}),
("malformed", {}),
("a=1,malformed,b=2", {"a": "1", "b": "2"}),
("a=", {}),
("=1", {}),
# A value containing '=' (e.g. base64) must survive intact.
("authorization=Basic dXNlcjpwYXNz==", {"authorization": "Basic dXNlcjpwYXNz=="}),
],
)
def test_parse(self, raw, expected):
assert parse_endpoint_headers(raw) == expected