mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
perf: cut hot-path latency 27% (token-count memo, startup preloads, JSON scan memo) (#2838)
## Description
Four independent latency fixes on the request hot path, found by
profiling and each measured in isolation. No behaviour changes: every
commit is either a memo of a pure function, work moved to startup, or
work that was computed and discarded.
**End to end: 287ms → 210ms (−27%) on a 68k-token mixed payload, with
byte-identical output** (68,514 → 48,725 tokens both before and after).
Plus one-off costs removed that don't show in steady-state numbers:
~4.9s of lazy imports that were firing *inside* user requests, and
~750ms of HuggingFace round-trips per process start.
Closes #
## Type of Change
- [ ] 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
- [x] Performance improvement
- [ ] Code refactoring (no functional changes)
## Changes Made
**1. Memoise `count_text` (`ac369277`)** — tiktoken's `CoreBPE.encode`
was 0.243s of a 0.30s profiled request. It dominates because the same
string is counted repeatedly: a 103KB payload drove 600KB of encoding,
~6x the content, across six call sites (`tokenizers/base.py:196`,
`content_router.py:4704` and `:5474`, `parser.py:185/192/298`). 35% of
encode calls and 22% of encoded characters were an exact repeat *within
one request*.
`count_text` is a pure function of its text, so replaying a stored count
returns the same integer. That is the whole safety argument, and it is
what makes this safe at the sites whose count feeds a routing decision
(`context_pressure` → `min_ratio`) rather than a log line — an
*estimate* there would change which blocks compress; a memo cannot.
Keyed on the text itself, not a hash: a collision would hand back a
wrong count for real content and silently change compression. The cost
is holding the strings, so entries and total characters are both capped.
Clear-on-full rather than LRU eviction — the pipeline runs on a thread
pool, `dict` get/set/clear are atomic under the GIL but
`OrderedDict.move_to_end` is not.
**2. Preload what was importing mid-request (`2921a15b`)** — `litellm`
(2.9–3.8s) was imported lazily *on the event loop* during the first
request: `emit_request_outcome` → `record_request` →
`_estimate_compression_savings_usd` calls the loader before its own
`tokens_saved <= 0` early return, so even a request that saved nothing
paid it. `trafilatura` (978ms, pulling `htmldate` → `dateparser` and its
timezone tables) is the most expensive lazy import in the transform tree
— every other compressor module is 1–20ms — and fires on the first
request carrying an HTML-ish or mixed-content block. The TOIN singleton
reads ~5MB of learned patterns on construction (~150ms); a stale comment
claimed the SmartCrusher preload covered it, and it does not.
All three now load in `_eager_preload_transforms`, which already runs
under `asyncio.to_thread` and so cannot delay the port bind.
Same commit, two Kompress cold-path fixes: `_load_modernbert_tokenizer`
always used `local_files_only=False`, which makes transformers
re-validate against the Hub on every load — a tree listing plus a HEAD
per file — even when fully cached (~900ms warm-cache vs ~150ms
local-only). And `ensure_background_download` re-spawned a
finished-or-failed thread on the next call, so an unreachable Hub meant
one fresh download thread *per request* for the life of the process,
each importing transformers and holding the GIL against the event loop.
Consecutive failures now back off; success clears it, so the happy path
and the transient-failure path are unchanged.
**3. Memoise the JSON-block scan (`039c9735`)** —
`_has_valid_json_block_with_text` tries every `{`/`[`-leading line as a
possible block start. When a candidate never balances,
`_extract_json_block` scans character-by-character to the end of the
content and returns nothing — then the next candidate does it again.
Quadratic, on the request path, growing exactly 4x per doubling.
**4. `CostTracker.totals()` (`286b97e4`)** —
`_current_savings_tracker_totals` called `stats()` once per request and
read two of its fields. Building the rest includes
`period_cost_breakdown()`, which walks up to 100k cost records over 31
days, on the event loop, holding the metrics lock. It degrades with
proxy **uptime**, not load, which is why no short benchmark would
surface it.
## Testing
- [x] Unit tests pass (`pytest`)
- [x] Linting passes (`ruff check .`)
- [x] Type checking passes (`mypy headroom`)
- [x] New tests added for new functionality
- [x] Manual testing performed
### Test Output
```text
$ ruff check headroom/ tests/ --exclude headroom/dashboard/templates
All checks passed!
$ mypy --python-version 3.12 headroom/
Found 1 error in 1 file (checked 515 source files)
headroom/release_version.py:235: error: Name "tomllib" already defined (by an import)
# pre-existing on main, in a file this PR does not touch — verified by
# running the same command on a clean main checkout.
$ python -m pytest tests/test_token_count_cache.py tests/test_mixed_content_scan_cache.py \
tests/test_kompress_download_backoff.py tests/test_cost_tracker_totals.py -q
306 passed
$ python -m pytest tests/ -q -k "token or tokenizer or count or estimator or provider"
1303 passed, 105 skipped in 423.42s
$ python -m pytest tests/ -q -k "cost or budget or metrics or savings or stats"
683 passed, 127 skipped, 1 failed
# tests/test_proxy_memory_integration.py::TestMemoryStats::test_health_endpoint_works_with_memory
# Order-dependent and pre-existing: it SKIPS in isolation, and fails identically
# on a clean main checkout under the same -k selection (681 passed, 1 failed).
```
## Real Behavior Proof
- **Environment:** macOS, Python 3.12.6, local CPU, remote Kompress
disabled. Profiled with `cProfile` on `anthropic_pipeline.apply`.
- **Exact command / steps:** a 68k-token payload of four `tool_result`
blocks (900-item pretty JSON, 60KB of Python source, 500 lines of
JS-style object logs, 500 plain log lines), six reps, **content unique
per rep so every run is router-cache-cold**, run on this branch and on
main in alternation.
- **Observed result:**
| | median | min | tokens |
|---|---|---|---|
| main | 287ms | 286ms | 68,514 → 48,725 |
| this branch | 210ms | 208ms | 68,514 → 48,725 |
Per-change, measured in isolation:
| change | before | after |
|---|---|---|
| `count_text` memo | — | −25% pipeline wall; 44% of counted chars from
cache on new content, 100% when history repeats |
| litellm / trafilatura / TOIN | 3829 / 978 / 150ms mid-request | at
startup, off the event loop |
| Kompress tokenizer | ~900ms | ~150ms |
| JS-style object logs (1200 lines) | 4643ms | 183ms |
| truncated JSONL (1200 lines) | 3737ms | 116ms |
| `cost_tracker` per request | 2.8ms @20k records, 13.6ms @100k | loop
over models, not records |
Output equality: 18/18 payloads byte-identical on `tokens_before`,
`tokens_after` and a sha256 of the resulting messages, with the memo
forced on vs off.
- **Not tested:** Windows and Linux (the ORT dylib and CPU-arena paths
differ); multi-worker deployments; a proxy with a genuinely large live
cost ledger (the 100k figure is from a synthetic ledger); real
HTML-heavy traffic through the preloaded trafilatura path.
## 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
- [x] I did **not** edit `CHANGELOG.md`
## Additional Notes
**Docs:** N/A — no user-facing surface changes. The reasoning lives in
the code, at the sites where someone debugging would look.
**A regression I introduced and caught.** The scan memo initially made
pretty-printed JSON ~2x **slower**: content that balances on the first
scan has nothing to reuse and just pays the per-line dict traffic. The
cache is now built only *after* a scan has run to the end without
balancing, which is the actual signal that later candidates will re-walk
the same tail. Every shape now improves and none regress:
```
before after
js object logs 4642.9ms 182.7ms 25x
JSONL truncated 3736.8ms 115.9ms 32x
pretty JSON 5.6ms 3.5ms
JSONL valid 5.6ms 3.2ms
plain logs 1.0ms 0.6ms
python source 1.0ms 0.5ms
markdown prose 0.9ms 0.5ms
```
Worth stating plainly: had I only benchmarked the shape I was fixing,
I'd have shipped a win on rare content and a loss on the common case.
**The scan fix is constant-factor, not asymptotic.** The walk over
remaining lines is still O(candidates × lines), so 3200 lines of the
pathological shape is still ~1.4s. The tests assert scan-call counts
rather than implying linearity. True linearity needs a prefix-sum
rewrite with a string-state fallback; that seemed like the wrong risk
for this PR.
**How the parser change is proven safe.** `_extract_json_block` is a
parser, so golden values would only encode whatever the new code does.
Instead the pre-memo implementation is kept verbatim in the test file as
an oracle, and every candidate index of a 139-document corpus — escapes,
unterminated strings, delimiters inside strings, code fences, truncated
JSON, randomised mixtures — is asserted equal, with a cold cache, with
the shared cache the real callers use, and replayed.
**Measurement trap, for anyone re-running these numbers.** Give each arm
its own content. Reusing one payload across arms lets the second arm hit
the router's result cache, which reads as a speedup having nothing to do
with the change under test. I hit this twice while working on it: it
manufactured a fake "INFO logging costs 21.8%" finding (real answer:
0.3%) and it *understated* the memo win.
**Deliberately not in this PR:**
- **ONNX thread tuning** — measured zero gain, and
`intra_op_num_threads` is not bitwise-safe (1.6e-05 score drift from
float reduction order), so it would trade an output risk for nothing.
- **`str(content)` on block lists** counts a base64 image at 210,775
tokens instead of 1,604 (131x), pinning `context_pressure` to 1.0 and
forcing the most aggressive `min_ratio` on any conversation containing
an image. Real bug, but fixing it changes compression output — needs its
own reviewed behaviour-change PR.
- **`chunk_words=350` against the tokenizer's 512-token limit** silently
drops roughly a third of every full chunk (measured: 240/240 words kept
in the first 240, 15/110 in the tail). That is data loss rather than
latency, it changes every output, and correcting it costs ~1.3x latency.
Filing separately.
- **Telemetry off the request thread** — the TOIN auto-save is a 236ms
inline stall every 600s and the waste-signal re-parse is ~50ms/request
that is invisible in `pipeline_total` (computed before it). Both want
deferral rather than removal, which is a larger change than belongs
here.
This commit is contained in:
Tejas Chopra
GitHub
parent
564e0a8d0f
commit
53af90d68c
15 changed files with 854 additions and 49 deletions
|
|
@ -24,7 +24,11 @@ import warnings
|
|||
from typing import Any, cast
|
||||
|
||||
from headroom import paths as _paths
|
||||
from headroom.tokenizers.base import coerce_countable_text, count_content_blocks
|
||||
from headroom.tokenizers.base import (
|
||||
TokenCountCache,
|
||||
coerce_countable_text,
|
||||
count_content_blocks,
|
||||
)
|
||||
|
||||
from .base import Provider, TokenCounter
|
||||
|
||||
|
|
@ -309,6 +313,7 @@ class AnthropicTokenCounter(TokenCounter):
|
|||
self.model = model
|
||||
self._client = client
|
||||
self._encoding: Any = None
|
||||
self._count_cache = TokenCountCache()
|
||||
self._use_api = client is not None
|
||||
|
||||
if not self._use_api and warn and not _FALLBACK_WARNING_SHOWN:
|
||||
|
|
@ -351,6 +356,14 @@ class AnthropicTokenCounter(TokenCounter):
|
|||
if not text:
|
||||
return 0
|
||||
|
||||
cached = self._count_cache.get(text)
|
||||
if cached is not None:
|
||||
return cached
|
||||
count = self._count_text_uncached(text)
|
||||
self._count_cache.put(text, count)
|
||||
return count
|
||||
|
||||
def _count_text_uncached(self, text: str) -> int:
|
||||
if self._encoding:
|
||||
# tiktoken with ~1.1x multiplier for Claude
|
||||
try:
|
||||
|
|
|
|||
|
|
@ -1057,6 +1057,39 @@ class CostTracker:
|
|||
except Exception:
|
||||
return None
|
||||
|
||||
def totals(self) -> tuple[int, float]:
|
||||
"""Return just ``(total_input_tokens, total_input_cost_usd)``.
|
||||
|
||||
The same two numbers ``stats()`` reports, computed without the rest of
|
||||
it. ``stats()`` is called once per request by the metrics path, which
|
||||
reads exactly these two fields and discards ``per_model``,
|
||||
``savings_usd``, ``cost_with_headroom_usd`` and — the expensive one —
|
||||
``budget_basis``, whose ``period_cost_breakdown()`` walks up to 31 days
|
||||
of retained cost records. MEASURED 2.8ms at 20k records and 13.6ms at
|
||||
100k, on the event loop and holding the metrics lock, so it degraded
|
||||
with proxy uptime rather than with load.
|
||||
|
||||
This loop is over models, not records, so it is bounded by how many
|
||||
models a deployment talks to.
|
||||
"""
|
||||
total_input_tokens = 0
|
||||
cost_with_headroom = 0.0
|
||||
for model in self._tokens_saved_by_model:
|
||||
sent = self._tokens_sent_by_model.get(model, 0)
|
||||
cr = self._api_cache_read_by_model.get(model, 0)
|
||||
cw = self._api_cache_write_by_model.get(model, 0)
|
||||
uncached = self._api_uncached_by_model.get(model, 0)
|
||||
total_input_tokens += sent
|
||||
|
||||
prices = self._get_cache_prices(model)
|
||||
if prices:
|
||||
cr_price, cw_price, uncached_price = prices
|
||||
if cr + cw + uncached > 0:
|
||||
cost_with_headroom += cr * cr_price + cw * cw_price + uncached * uncached_price
|
||||
else:
|
||||
cost_with_headroom += sent * uncached_price
|
||||
return total_input_tokens, round(cost_with_headroom, 4)
|
||||
|
||||
def stats(self) -> dict:
|
||||
"""Get token statistics per model."""
|
||||
per_model = {}
|
||||
|
|
|
|||
|
|
@ -417,14 +417,14 @@ class PrometheusMetrics:
|
|||
return total_input_tokens, total_input_cost_usd
|
||||
|
||||
try:
|
||||
cost_stats = self.cost_tracker.stats()
|
||||
# totals() rather than stats(): identical numbers, without the
|
||||
# 31-day cost-record walk that stats()["budget_basis"] performs and
|
||||
# this caller throws away. See CostTracker.totals.
|
||||
tracked_input_tokens, tracked_input_cost_usd = self.cost_tracker.totals()
|
||||
except Exception:
|
||||
logger.debug("Failed to read cost tracker totals for savings history", exc_info=True)
|
||||
return total_input_tokens, total_input_cost_usd
|
||||
|
||||
tracked_input_tokens = cost_stats.get("total_input_tokens")
|
||||
tracked_input_cost_usd = cost_stats.get("total_input_cost_usd")
|
||||
|
||||
if tracked_input_tokens is not None:
|
||||
try:
|
||||
total_input_tokens = self._savings_tracker_input_tokens_offset + max(
|
||||
|
|
|
|||
|
|
@ -1638,6 +1638,26 @@ class HeadroomProxy(
|
|||
for key, value in transform_status.items():
|
||||
eager_status.setdefault(key, value)
|
||||
transform_statuses.append(transform_status)
|
||||
|
||||
# LiteLLM's pricing tables. MEASURED 2.9-3.8s to import, and it was
|
||||
# being imported lazily ON THE EVENT LOOP during the first request:
|
||||
# emit_request_outcome -> record_request -> _estimate_compression_savings_usd
|
||||
# calls it before its own `tokens_saved <= 0` early return, so even a
|
||||
# request that saved nothing pays for it. Nothing about that is visible
|
||||
# as a failure; it just makes one unlucky user wait ~3s.
|
||||
#
|
||||
# This function already runs under asyncio.to_thread, so importing here
|
||||
# cannot delay the port bind.
|
||||
try:
|
||||
from .savings_tracker import _get_litellm_module
|
||||
|
||||
eager_status.setdefault(
|
||||
"litellm", "ready" if _get_litellm_module() is not None else "not installed"
|
||||
)
|
||||
except Exception as exc: # pricing is optional; never block startup on it
|
||||
logger.debug("LiteLLM pre-load skipped: %s", exc)
|
||||
eager_status.setdefault("litellm", "skipped")
|
||||
|
||||
return eager_status, transform_statuses
|
||||
|
||||
async def startup(self):
|
||||
|
|
|
|||
|
|
@ -17,6 +17,68 @@ from typing import Any, Protocol, runtime_checkable
|
|||
#: encode. Truncating keeps the estimate finite and the request alive.
|
||||
_MAX_COERCED_FIELD_CHARS = 200_000
|
||||
|
||||
#: Admission policy for :class:`TokenCountCache`. These bound memory only — they
|
||||
#: never change a returned count, so they are not a behavioural threshold.
|
||||
#: Strings below the floor encode in microseconds, so caching them would only
|
||||
#: evict the large entries that the cache exists for.
|
||||
_COUNT_CACHE_MIN_CHARS = 256
|
||||
_COUNT_CACHE_MAX_ENTRIES = 4096
|
||||
_COUNT_CACHE_MAX_CHARS = 8_000_000
|
||||
|
||||
|
||||
class TokenCountCache:
|
||||
"""Exact memo for ``count_text``.
|
||||
|
||||
``count_text`` is a pure function of its text, so replaying a stored count is
|
||||
value-identical. That is the whole safety argument: the result is the same
|
||||
integer, which is why this is safe even at the call sites whose count feeds a
|
||||
routing decision (``context_pressure`` -> ``min_ratio``) rather than a log line.
|
||||
|
||||
Keyed on the text itself rather than a hash. A hash collision here would hand
|
||||
back the wrong count for real content and silently change what gets
|
||||
compressed; the counts are too load-bearing to trade correctness for a
|
||||
smaller key. The price is holding the strings, so entries *and* total
|
||||
characters are capped.
|
||||
|
||||
No lock: ``dict`` get/set/clear are atomic under the GIL, and the pipeline
|
||||
runs on a thread pool. An LRU would need ``move_to_end``, which is not
|
||||
atomic — hence clear-on-full rather than eviction. A cleared cache costs one
|
||||
re-encode, never a wrong answer.
|
||||
"""
|
||||
|
||||
__slots__ = ("_chars", "_counts", "_max_chars", "_max_entries", "_min_chars")
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
min_chars: int = _COUNT_CACHE_MIN_CHARS,
|
||||
max_entries: int = _COUNT_CACHE_MAX_ENTRIES,
|
||||
max_chars: int = _COUNT_CACHE_MAX_CHARS,
|
||||
) -> None:
|
||||
self._counts: dict[str, int] = {}
|
||||
self._chars = 0
|
||||
self._min_chars = min_chars
|
||||
self._max_entries = max_entries
|
||||
self._max_chars = max_chars
|
||||
|
||||
def get(self, text: str) -> int | None:
|
||||
"""Return the stored count for *text*, or None."""
|
||||
return self._counts.get(text)
|
||||
|
||||
def put(self, text: str, count: int) -> None:
|
||||
"""Store *count* for *text* if it is worth caching."""
|
||||
if len(text) < self._min_chars:
|
||||
return
|
||||
if len(self._counts) >= self._max_entries or self._chars >= self._max_chars:
|
||||
self._counts.clear()
|
||||
self._chars = 0
|
||||
self._counts[text] = count
|
||||
self._chars += len(text)
|
||||
|
||||
def clear(self) -> None:
|
||||
self._counts.clear()
|
||||
self._chars = 0
|
||||
|
||||
|
||||
def coerce_countable_text(value: Any) -> str:
|
||||
"""Return *value* as text safe to pass to ``count_text``.
|
||||
|
|
|
|||
|
|
@ -11,7 +11,7 @@ import json
|
|||
import re
|
||||
from typing import Any
|
||||
|
||||
from .base import BaseTokenizer
|
||||
from .base import BaseTokenizer, TokenCountCache
|
||||
|
||||
|
||||
class EstimatingTokenCounter(BaseTokenizer):
|
||||
|
|
@ -87,6 +87,7 @@ class EstimatingTokenCounter(BaseTokenizer):
|
|||
If None, auto-detects based on content type.
|
||||
"""
|
||||
self._fixed_ratio = chars_per_token
|
||||
self._count_cache = TokenCountCache()
|
||||
|
||||
def count_text(self, text: str) -> int:
|
||||
"""Estimate token count for text.
|
||||
|
|
@ -100,6 +101,14 @@ class EstimatingTokenCounter(BaseTokenizer):
|
|||
if not text:
|
||||
return 0
|
||||
|
||||
cached = self._count_cache.get(text)
|
||||
if cached is not None:
|
||||
return cached
|
||||
count = self._count_text_uncached(text)
|
||||
self._count_cache.put(text, count)
|
||||
return count
|
||||
|
||||
def _count_text_uncached(self, text: str) -> int:
|
||||
# Use fixed ratio if provided. Dense scripts (CJK/Kana/Hangul) still
|
||||
# tokenize at ~1 token per character, so pricing them at the (Latin)
|
||||
# fixed ratio under-counts by 2-4x — the same correction the auto path
|
||||
|
|
|
|||
|
|
@ -16,7 +16,7 @@ import threading
|
|||
from functools import lru_cache
|
||||
from typing import Any
|
||||
|
||||
from .base import BaseTokenizer, coerce_countable_text
|
||||
from .base import BaseTokenizer, TokenCountCache, coerce_countable_text
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
|
@ -237,6 +237,7 @@ class TiktokenCounter(BaseTokenizer):
|
|||
self.model = model
|
||||
self.encoding_name = encoding or get_encoding_for_model(model)
|
||||
self._encoding = None # Lazy load
|
||||
self._count_cache = TokenCountCache()
|
||||
|
||||
@property
|
||||
def encoding(self):
|
||||
|
|
@ -256,6 +257,14 @@ class TiktokenCounter(BaseTokenizer):
|
|||
"""
|
||||
if not text:
|
||||
return 0
|
||||
cached = self._count_cache.get(text)
|
||||
if cached is not None:
|
||||
return cached
|
||||
count = self._count_text_uncached(text)
|
||||
self._count_cache.put(text, count)
|
||||
return count
|
||||
|
||||
def _count_text_uncached(self, text: str) -> int:
|
||||
try:
|
||||
return len(self.encoding.encode(text))
|
||||
except ValueError:
|
||||
|
|
|
|||
|
|
@ -4186,11 +4186,42 @@ class ContentRouter(Transform):
|
|||
else:
|
||||
status["code_aware"] = "not installed"
|
||||
|
||||
# 4. SmartCrusher (lightweight init, but ensures import + TOIN ready)
|
||||
# 4. SmartCrusher (lightweight init)
|
||||
smart_crusher = self._get_smart_crusher()
|
||||
if smart_crusher:
|
||||
status["smart_crusher"] = "ready"
|
||||
|
||||
# 5. HTML extractor.
|
||||
#
|
||||
# By far the most expensive lazy import in the transform tree: MEASURED
|
||||
# 978ms for trafilatura -> htmldate -> dateparser and its timezone
|
||||
# tables, against 1-20ms for every other compressor module. It fires
|
||||
# from _get_html_extractor() on the first request carrying an HTML-ish
|
||||
# block or mixed-content section, so a real user pays the full second
|
||||
# mid-request. That is the single largest first-request stall in the
|
||||
# pipeline, which is why it is worth a line here.
|
||||
try:
|
||||
if self._get_html_extractor() is not None:
|
||||
status["html_extractor"] = "ready"
|
||||
else:
|
||||
status["html_extractor"] = "not installed"
|
||||
except Exception as e:
|
||||
logger.debug("HTML extractor pre-load skipped: %s", e)
|
||||
status["html_extractor"] = "skipped"
|
||||
|
||||
# 6. TOIN singleton. Constructing it reads the learned-pattern file off
|
||||
# disk (MEASURED ~150ms at 5MB, and it grows with use). SmartCrusher
|
||||
# above does NOT pull it in, despite what a previous comment here
|
||||
# claimed — the first request did.
|
||||
try:
|
||||
from ..telemetry.toin import get_toin
|
||||
|
||||
get_toin()
|
||||
status["toin"] = "ready"
|
||||
except Exception as e:
|
||||
logger.debug("TOIN pre-load skipped: %s", e)
|
||||
status["toin"] = "skipped"
|
||||
|
||||
return status
|
||||
|
||||
def _get_kompress(self) -> Any:
|
||||
|
|
|
|||
|
|
@ -746,15 +746,27 @@ def _load_kompress_onnx(
|
|||
|
||||
|
||||
def _load_modernbert_tokenizer(auto_tokenizer: Any, *, allow_download: bool) -> Any:
|
||||
"""Load the ModernBERT tokenizer, cache-only when ``allow_download`` is False."""
|
||||
"""Load the ModernBERT tokenizer, cache-only when ``allow_download`` is False.
|
||||
|
||||
Always tries the local cache FIRST, even when downloading is allowed. With
|
||||
``local_files_only=False`` transformers re-validates against the Hub on every
|
||||
load — a tree listing plus a HEAD per tokenizer file — even when the repo is
|
||||
fully cached. MEASURED ~900ms warm-cache versus ~150ms local-only, i.e. ~750ms
|
||||
of pure network round-trip on every process start, and it is also what makes
|
||||
a cold start slow on a bad network rather than merely offline.
|
||||
|
||||
Same files, same tokenizer, so the loaded object is identical; this only
|
||||
changes whether the Hub is consulted to confirm what is already on disk.
|
||||
Mirrors ``onnx_runtime.hf_hub_download_local_first``, which the ONNX half of
|
||||
this loader already uses.
|
||||
"""
|
||||
try:
|
||||
return auto_tokenizer.from_pretrained(
|
||||
"answerdotai/ModernBERT-base", local_files_only=not allow_download
|
||||
)
|
||||
return auto_tokenizer.from_pretrained("answerdotai/ModernBERT-base", local_files_only=True)
|
||||
except _NOT_CACHED_ERRORS as exc:
|
||||
if not allow_download:
|
||||
raise KompressModelNotCached("answerdotai/ModernBERT-base") from exc
|
||||
raise
|
||||
# Genuine cache miss and downloading is permitted: fetch it.
|
||||
return auto_tokenizer.from_pretrained("answerdotai/ModernBERT-base", local_files_only=False)
|
||||
|
||||
|
||||
# Sub-state-dict keys inside a merged v2-style checkpoint (see
|
||||
|
|
@ -1039,6 +1051,45 @@ def unload_kompress_model(model_id: str | None = None) -> bool:
|
|||
_download_threads: dict[str, threading.Thread] = {}
|
||||
_download_threads_lock = threading.Lock()
|
||||
|
||||
#: Retry backoff for a FAILED background download, in seconds. A finished-or-failed
|
||||
#: thread is replaced on the next call so a transient network blip recovers, but
|
||||
#: without a floor an unreachable Hub means every request spawns a fresh download
|
||||
#: thread for the life of the process — each one importing transformers and
|
||||
#: resolving the Hub, all holding the GIL against the event loop. The window grows
|
||||
#: per consecutive failure and resets on success, so the happy path and the
|
||||
#: transient-failure path are both unchanged; only the permanently-broken case is
|
||||
#: bounded.
|
||||
_DOWNLOAD_RETRY_BASE_SECONDS = 5.0
|
||||
_DOWNLOAD_RETRY_MAX_SECONDS = 300.0
|
||||
_download_failures: dict[str, tuple[int, float]] = {}
|
||||
|
||||
|
||||
def _record_download_failure(model_id: str) -> None:
|
||||
with _download_threads_lock:
|
||||
failures, _ = _download_failures.get(model_id, (0, 0.0))
|
||||
_download_failures[model_id] = (failures + 1, time.monotonic())
|
||||
|
||||
|
||||
def _clear_download_failures(model_id: str) -> None:
|
||||
with _download_threads_lock:
|
||||
_download_failures.pop(model_id, None)
|
||||
|
||||
|
||||
def _download_retry_blocked(model_id: str) -> bool:
|
||||
"""True when the last attempt failed and the backoff window has not elapsed.
|
||||
|
||||
Caller must hold ``_download_threads_lock``.
|
||||
"""
|
||||
entry = _download_failures.get(model_id)
|
||||
if entry is None:
|
||||
return False
|
||||
failures, last_attempt = entry
|
||||
window = min(
|
||||
_DOWNLOAD_RETRY_MAX_SECONDS,
|
||||
_DOWNLOAD_RETRY_BASE_SECONDS * (2 ** (failures - 1)),
|
||||
)
|
||||
return bool((time.monotonic() - last_attempt) < window)
|
||||
|
||||
|
||||
def _background_download(model_id: str, device: str) -> None:
|
||||
try:
|
||||
|
|
@ -1046,7 +1097,10 @@ def _background_download(model_id: str, device: str) -> None:
|
|||
_load_kompress(model_id, device, allow_download=True)
|
||||
logger.info("Kompress: background model download complete for %s", model_id)
|
||||
except Exception as exc:
|
||||
_record_download_failure(model_id)
|
||||
logger.warning("Kompress: background model download failed for %s: %s", model_id, exc)
|
||||
else:
|
||||
_clear_download_failures(model_id)
|
||||
|
||||
|
||||
def ensure_background_download(model_id: str = HF_MODEL_ID, device: str = "auto") -> None:
|
||||
|
|
@ -1054,7 +1108,9 @@ def ensure_background_download(model_id: str = HF_MODEL_ID, device: str = "auto"
|
|||
|
||||
Idempotent and non-blocking: at most one download thread runs per model_id,
|
||||
and a finished or failed thread is replaced on the next call so a transient
|
||||
network failure can be retried by a later request. Once the download
|
||||
network failure can be retried by a later request — subject to a growing
|
||||
backoff after consecutive failures, so an unreachable Hub cannot turn every
|
||||
request into another download thread. Once the download
|
||||
completes the deep path activates on subsequent requests without ever
|
||||
blocking one on the network.
|
||||
"""
|
||||
|
|
@ -1066,6 +1122,8 @@ def ensure_background_download(model_id: str = HF_MODEL_ID, device: str = "auto"
|
|||
existing = _download_threads.get(model_id)
|
||||
if existing is not None and existing.is_alive():
|
||||
return
|
||||
if _download_retry_blocked(model_id):
|
||||
return
|
||||
thread = threading.Thread(
|
||||
target=_background_download,
|
||||
args=(model_id, device),
|
||||
|
|
|
|||
|
|
@ -43,16 +43,32 @@ def mixed_content_indicators(content: str) -> dict[str, bool]:
|
|||
}
|
||||
|
||||
|
||||
def _any_nonblank(lines: list[str], start: int, stop: int) -> bool:
|
||||
"""True when some line in [start, stop) has non-whitespace.
|
||||
|
||||
Equivalent to ``bool("\n".join(lines[start:stop]).strip())`` — a join of
|
||||
lines is blank exactly when every line is blank — but it short-circuits
|
||||
instead of building a copy of the whole body for each candidate.
|
||||
"""
|
||||
return any(lines[i].strip() for i in range(start, stop))
|
||||
|
||||
|
||||
def _has_valid_json_block_with_text(content: str) -> bool:
|
||||
"""Return true when prose or log text wraps a valid JSON block."""
|
||||
lines = content.split("\n")
|
||||
# Built only after a scan has run to the end without balancing — see
|
||||
# _extract_json_block. Content that balances promptly never allocates it and
|
||||
# so pays nothing for it.
|
||||
scan_cache: dict[tuple[int, bool, bool], tuple[int, int, bool, bool]] | None = None
|
||||
|
||||
for index, line in enumerate(lines):
|
||||
if not line.strip().startswith(("[", "{")):
|
||||
continue
|
||||
|
||||
json_content, end_index = _extract_json_block(lines, index)
|
||||
json_content, end_index = _extract_json_block(lines, index, cache=scan_cache)
|
||||
if json_content is None:
|
||||
if scan_cache is None:
|
||||
scan_cache = {}
|
||||
continue
|
||||
|
||||
try:
|
||||
|
|
@ -60,9 +76,7 @@ def _has_valid_json_block_with_text(content: str) -> bool:
|
|||
except (TypeError, ValueError):
|
||||
continue
|
||||
|
||||
leading_text = "\n".join(lines[:index]).strip()
|
||||
trailing_text = "\n".join(lines[end_index + 1 :]).strip()
|
||||
if leading_text or trailing_text:
|
||||
if _any_nonblank(lines, 0, index) or _any_nonblank(lines, end_index + 1, len(lines)):
|
||||
return True
|
||||
|
||||
return False
|
||||
|
|
@ -72,6 +86,7 @@ def split_into_sections(content: str) -> list[ContentSection]:
|
|||
"""Parse mixed content into typed sections."""
|
||||
sections: list[ContentSection] = []
|
||||
lines = content.split("\n")
|
||||
scan_cache: dict[tuple[int, bool, bool], tuple[int, int, bool, bool]] | None = None
|
||||
|
||||
i = 0
|
||||
while i < len(lines):
|
||||
|
|
@ -101,7 +116,11 @@ def split_into_sections(content: str) -> list[ContentSection]:
|
|||
continue
|
||||
|
||||
if line.strip().startswith(("[", "{")):
|
||||
json_content, end_i = _extract_json_block(lines, i)
|
||||
json_content, end_i = _extract_json_block(lines, i, cache=scan_cache)
|
||||
if json_content is None and scan_cache is None:
|
||||
# First scan that ran to the end without balancing: from here on
|
||||
# every later candidate would re-walk the same tail.
|
||||
scan_cache = {}
|
||||
if json_content:
|
||||
sections.append(
|
||||
ContentSection(
|
||||
|
|
@ -159,41 +178,84 @@ def split_into_sections(content: str) -> list[ContentSection]:
|
|||
return sections
|
||||
|
||||
|
||||
def _extract_json_block(lines: list[str], start: int) -> tuple[str | None, int]:
|
||||
"""Extract a complete JSON object or array block from line-oriented content."""
|
||||
def _scan_line(line: str, in_string: bool, escaped: bool) -> tuple[int, int, bool, bool]:
|
||||
"""Bracket/brace deltas for one line, given the parser state entering it.
|
||||
|
||||
Split out so the per-line result can be memoised across scans: what a line
|
||||
does to the counters is a pure function of the line and the two entry-state
|
||||
flags, nothing else.
|
||||
"""
|
||||
bracket = 0
|
||||
brace = 0
|
||||
for ch in line:
|
||||
if escaped:
|
||||
escaped = False
|
||||
continue
|
||||
if ch == "\\":
|
||||
if in_string:
|
||||
escaped = True
|
||||
continue
|
||||
if ch == '"':
|
||||
in_string = not in_string
|
||||
continue
|
||||
if in_string:
|
||||
continue
|
||||
if ch == "[":
|
||||
bracket += 1
|
||||
elif ch == "]":
|
||||
bracket -= 1
|
||||
elif ch == "{":
|
||||
brace += 1
|
||||
elif ch == "}":
|
||||
brace -= 1
|
||||
return bracket, brace, in_string, escaped
|
||||
|
||||
|
||||
def _extract_json_block(
|
||||
lines: list[str],
|
||||
start: int,
|
||||
*,
|
||||
cache: dict[tuple[int, bool, bool], tuple[int, int, bool, bool]] | None = None,
|
||||
) -> tuple[str | None, int]:
|
||||
"""Extract a complete JSON object or array block from line-oriented content.
|
||||
|
||||
``cache`` memoises the per-line scan across repeated calls over the SAME
|
||||
``lines``. Callers that try every ``{``-leading line share one dict; without
|
||||
it each candidate that never balances re-scans character-by-character to the
|
||||
end of the content, which is quadratic.
|
||||
|
||||
Callers pass ``None`` until a scan has actually run to the end without
|
||||
balancing, and only build the dict from then on. That matters: on content
|
||||
that balances on the first try — pretty-printed JSON, the common case — the
|
||||
memo has nothing to reuse and its per-line dict traffic made that shape ~2x
|
||||
SLOWER. A failed scan is the signal that later candidates will re-walk the
|
||||
same tail, and it is the only point at which the memo pays. MEASURED before the cache: 4643ms for
|
||||
1200 lines of JS-style object logs and 3737ms for truncated JSONL, growing
|
||||
exactly 4x per doubling. Ordinary shapes — pretty-printed JSON, valid JSONL,
|
||||
source, stack traces, prose — were ~1-6ms and never hit it, which is why this
|
||||
stayed invisible.
|
||||
|
||||
Keyed on the entry state as well as the line, so a cached entry is only
|
||||
reused where the parser is in the same string/escape state. Same deltas,
|
||||
same result: this is a memo, not a heuristic.
|
||||
"""
|
||||
bracket_count = 0
|
||||
brace_count = 0
|
||||
json_lines = []
|
||||
in_string = False
|
||||
escaped = False
|
||||
|
||||
for i in range(start, len(lines)):
|
||||
line = lines[i]
|
||||
json_lines.append(line)
|
||||
key = (i, in_string, escaped)
|
||||
step = cache.get(key) if cache is not None else None
|
||||
if step is None:
|
||||
step = _scan_line(lines[i], in_string, escaped)
|
||||
if cache is not None:
|
||||
cache[key] = step
|
||||
d_bracket, d_brace, in_string, escaped = step
|
||||
bracket_count += d_bracket
|
||||
brace_count += d_brace
|
||||
|
||||
for ch in line:
|
||||
if escaped:
|
||||
escaped = False
|
||||
continue
|
||||
if ch == "\\":
|
||||
if in_string:
|
||||
escaped = True
|
||||
continue
|
||||
if ch == '"':
|
||||
in_string = not in_string
|
||||
continue
|
||||
if in_string:
|
||||
continue
|
||||
if ch == "[":
|
||||
bracket_count += 1
|
||||
elif ch == "]":
|
||||
bracket_count -= 1
|
||||
elif ch == "{":
|
||||
brace_count += 1
|
||||
elif ch == "}":
|
||||
brace_count -= 1
|
||||
|
||||
if bracket_count <= 0 and brace_count <= 0 and json_lines:
|
||||
return "\n".join(json_lines), i
|
||||
if bracket_count <= 0 and brace_count <= 0:
|
||||
return "\n".join(lines[start : i + 1]), i
|
||||
|
||||
return None, start
|
||||
|
|
|
|||
80
tests/test_cost_tracker_totals.py
Normal file
80
tests/test_cost_tracker_totals.py
Normal file
|
|
@ -0,0 +1,80 @@
|
|||
"""CostTracker.totals() must be stats() minus the work, not minus the accuracy.
|
||||
|
||||
It exists only so the per-request metrics path stops walking 31 days of cost
|
||||
records to read two fields. If the two ever disagree, the savings history
|
||||
silently drifts from /stats.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import random
|
||||
|
||||
import pytest
|
||||
|
||||
from headroom.proxy.cost import CostTracker
|
||||
|
||||
|
||||
def _tracker(seed: int, n_models: int, n_requests: int) -> CostTracker:
|
||||
r = random.Random(seed)
|
||||
tracker = CostTracker()
|
||||
models = [
|
||||
"claude-sonnet-5",
|
||||
"claude-opus-4-1",
|
||||
"gpt-4o",
|
||||
"gpt-4o-mini",
|
||||
"some-unpriceable-model",
|
||||
][:n_models]
|
||||
for _ in range(n_requests):
|
||||
model = r.choice(models)
|
||||
sent = r.randint(0, 20000)
|
||||
# Alternate between requests that carry an API cache breakdown and ones
|
||||
# that do not — totals() has a branch for each, and only the second
|
||||
# falls back to list price.
|
||||
with_cache = r.random() < 0.5
|
||||
tracker.record_tokens(
|
||||
model=model,
|
||||
tokens_saved=r.randint(0, 5000),
|
||||
tokens_sent=sent,
|
||||
cache_read_tokens=r.randint(0, sent) if with_cache else 0,
|
||||
cache_write_tokens=r.randint(0, 500) if with_cache else 0,
|
||||
uncached_tokens=r.randint(0, sent) if with_cache else 0,
|
||||
output_tokens=r.randint(0, 2000),
|
||||
)
|
||||
return tracker
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("n_models", "n_requests"),
|
||||
[(0, 0), (1, 1), (1, 50), (3, 200), (5, 500)],
|
||||
)
|
||||
def test_totals_matches_stats(n_models: int, n_requests: int) -> None:
|
||||
tracker = _tracker(seed=n_models * 100 + n_requests, n_models=n_models, n_requests=n_requests)
|
||||
stats = tracker.stats()
|
||||
assert tracker.totals() == (
|
||||
stats["total_input_tokens"],
|
||||
stats["total_input_cost_usd"],
|
||||
)
|
||||
|
||||
|
||||
def test_totals_matches_stats_on_a_fresh_tracker() -> None:
|
||||
tracker = CostTracker()
|
||||
stats = tracker.stats()
|
||||
assert tracker.totals() == (stats["total_input_tokens"], stats["total_input_cost_usd"])
|
||||
|
||||
|
||||
def test_totals_does_not_walk_the_cost_records() -> None:
|
||||
"""The point of the method: no period_cost_breakdown, at any ledger size."""
|
||||
tracker = _tracker(seed=7, n_models=3, n_requests=100)
|
||||
called = False
|
||||
real = tracker.period_cost_breakdown
|
||||
|
||||
def spy(*a, **kw):
|
||||
nonlocal called
|
||||
called = True
|
||||
return real(*a, **kw)
|
||||
|
||||
tracker.period_cost_breakdown = spy # type: ignore[method-assign]
|
||||
tracker.totals()
|
||||
assert not called, "totals() still walks the cost records"
|
||||
tracker.stats()
|
||||
assert called, "stats() should still report budget_basis"
|
||||
116
tests/test_kompress_download_backoff.py
Normal file
116
tests/test_kompress_download_backoff.py
Normal file
|
|
@ -0,0 +1,116 @@
|
|||
"""An unreachable HuggingFace must not turn every request into a download thread.
|
||||
|
||||
The request path calls ensure_background_download() on every Kompress miss. A
|
||||
finished-or-failed thread is replaced on the next call, which is what lets a
|
||||
transient blip recover — but with no floor, a permanently unreachable Hub means
|
||||
one new thread per request forever, each importing transformers and holding the
|
||||
GIL against the event loop.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import threading
|
||||
|
||||
import pytest
|
||||
|
||||
from headroom.transforms import kompress_compressor as kc
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _clean_registry():
|
||||
with kc._download_threads_lock:
|
||||
kc._download_threads.clear()
|
||||
kc._download_failures.clear()
|
||||
yield
|
||||
with kc._download_threads_lock:
|
||||
kc._download_threads.clear()
|
||||
kc._download_failures.clear()
|
||||
|
||||
|
||||
def _spawned(monkeypatch, *, fails: bool) -> list[str]:
|
||||
"""Run ensure_background_download with the real load stubbed out."""
|
||||
started: list[str] = []
|
||||
|
||||
def fake_load(model_id, device, allow_download=True):
|
||||
started.append(model_id)
|
||||
if fails:
|
||||
raise OSError("hub unreachable")
|
||||
return object(), object(), "onnx"
|
||||
|
||||
monkeypatch.setattr(kc, "_load_kompress", fake_load)
|
||||
return started
|
||||
|
||||
|
||||
def _drain():
|
||||
for t in list(kc._download_threads.values()):
|
||||
t.join(timeout=10)
|
||||
|
||||
|
||||
def test_repeated_failure_stops_spawning_threads(monkeypatch):
|
||||
started = _spawned(monkeypatch, fails=True)
|
||||
for _ in range(25):
|
||||
kc.ensure_background_download("some/model")
|
||||
_drain()
|
||||
assert len(started) < 25, f"no backoff: spawned {len(started)} downloads for 25 calls"
|
||||
assert len(started) >= 1, "never even tried once"
|
||||
|
||||
|
||||
def test_backoff_window_elapsing_allows_another_attempt(monkeypatch):
|
||||
started = _spawned(monkeypatch, fails=True)
|
||||
kc.ensure_background_download("some/model")
|
||||
_drain()
|
||||
assert len(started) == 1
|
||||
|
||||
kc.ensure_background_download("some/model")
|
||||
_drain()
|
||||
assert len(started) == 1, "retried inside the backoff window"
|
||||
|
||||
# Rewind the clock past the window instead of sleeping through it.
|
||||
with kc._download_threads_lock:
|
||||
failures, _ = kc._download_failures["some/model"]
|
||||
kc._download_failures["some/model"] = (failures, 0.0)
|
||||
kc.ensure_background_download("some/model")
|
||||
_drain()
|
||||
assert len(started) == 2, "backoff never expires"
|
||||
|
||||
|
||||
def test_success_clears_the_backoff(monkeypatch):
|
||||
_spawned(monkeypatch, fails=False)
|
||||
kc.ensure_background_download("some/model")
|
||||
_drain()
|
||||
assert "some/model" not in kc._download_failures
|
||||
|
||||
|
||||
def test_window_grows_with_consecutive_failures():
|
||||
kc._download_failures["m"] = (1, 0.0)
|
||||
assert kc._DOWNLOAD_RETRY_BASE_SECONDS == 5.0
|
||||
# Same last-attempt time, more failures -> still blocked at a later clock.
|
||||
import time as _t
|
||||
|
||||
now = _t.monotonic()
|
||||
kc._download_failures["m"] = (1, now)
|
||||
with kc._download_threads_lock:
|
||||
first = kc._download_retry_blocked("m")
|
||||
kc._download_failures["m"] = (6, now)
|
||||
with kc._download_threads_lock:
|
||||
later = kc._download_retry_blocked("m")
|
||||
assert first and later
|
||||
|
||||
|
||||
def test_a_live_thread_is_never_duplicated(monkeypatch):
|
||||
gate = threading.Event()
|
||||
started: list[str] = []
|
||||
|
||||
def slow_load(model_id, device, allow_download=True):
|
||||
started.append(model_id)
|
||||
gate.wait(timeout=10)
|
||||
return object(), object(), "onnx"
|
||||
|
||||
monkeypatch.setattr(kc, "_load_kompress", slow_load)
|
||||
for _ in range(10):
|
||||
kc.ensure_background_download("some/model")
|
||||
try:
|
||||
assert len(started) == 1
|
||||
finally:
|
||||
gate.set()
|
||||
_drain()
|
||||
197
tests/test_mixed_content_scan_cache.py
Normal file
197
tests/test_mixed_content_scan_cache.py
Normal file
|
|
@ -0,0 +1,197 @@
|
|||
"""The memoised JSON-block scan must be indistinguishable from the original.
|
||||
|
||||
This is a parser change, so equality is checked against a literal transcription
|
||||
of the pre-cache implementation rather than against expected values — a golden
|
||||
test would only encode whatever the new code does.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import random
|
||||
|
||||
import pytest
|
||||
|
||||
from headroom.transforms.mixed_content import (
|
||||
_extract_json_block,
|
||||
_has_valid_json_block_with_text,
|
||||
is_mixed_content,
|
||||
split_into_sections,
|
||||
)
|
||||
|
||||
|
||||
def _extract_json_block_original(lines: list[str], start: int) -> tuple[str | None, int]:
|
||||
"""Verbatim pre-cache implementation, kept as the oracle."""
|
||||
bracket_count = 0
|
||||
brace_count = 0
|
||||
json_lines = []
|
||||
in_string = False
|
||||
escaped = False
|
||||
|
||||
for i in range(start, len(lines)):
|
||||
line = lines[i]
|
||||
json_lines.append(line)
|
||||
|
||||
for ch in line:
|
||||
if escaped:
|
||||
escaped = False
|
||||
continue
|
||||
if ch == "\\":
|
||||
if in_string:
|
||||
escaped = True
|
||||
continue
|
||||
if ch == '"':
|
||||
in_string = not in_string
|
||||
continue
|
||||
if in_string:
|
||||
continue
|
||||
if ch == "[":
|
||||
bracket_count += 1
|
||||
elif ch == "]":
|
||||
bracket_count -= 1
|
||||
elif ch == "{":
|
||||
brace_count += 1
|
||||
elif ch == "}":
|
||||
brace_count -= 1
|
||||
|
||||
if bracket_count <= 0 and brace_count <= 0 and json_lines:
|
||||
return "\n".join(json_lines), i
|
||||
|
||||
return None, start
|
||||
|
||||
|
||||
def _corpus() -> list[str]:
|
||||
r = random.Random(20260806)
|
||||
out = [
|
||||
"",
|
||||
"\n",
|
||||
" \n\t\n",
|
||||
"{",
|
||||
"}",
|
||||
'{"a": 1}',
|
||||
'[\n{"id": 1}\n]',
|
||||
'{"s": "a ] b } c"}', # delimiters inside strings
|
||||
'{"s": "escaped \\" quote }"}', # escaped quote
|
||||
'{"s": "trailing backslash \\\\"}',
|
||||
'{"s": "line one\\', # line ends mid-escape
|
||||
'text before\n{"a": 1}\ntext after',
|
||||
'```json\n{"a": 1}\n```\nprose here',
|
||||
"\n".join(f'{{ level: "info", seq: {i}, msg: "x"' for i in range(40)), # never balances
|
||||
"\n".join(json.dumps({"id": i})[:9] for i in range(40)), # truncated JSONL
|
||||
"\n".join(json.dumps({"id": i, "m": "ok"}) for i in range(40)), # valid JSONL
|
||||
json.dumps([{"id": i, "n": f"x{i}"} for i in range(40)], indent=2),
|
||||
"\n".join(f"2026-08-06 13:00:{i % 60:02d} INFO did thing {i}" for i in range(40)),
|
||||
"\n".join(f" cfg = {{'k{i}': 'v{i}'," for i in range(40)),
|
||||
]
|
||||
# Randomised mixtures, including unbalanced and string-heavy fragments.
|
||||
frags = [
|
||||
'{"a": 1}',
|
||||
"[",
|
||||
"]",
|
||||
"{",
|
||||
"}",
|
||||
"plain prose line",
|
||||
'{"s": "] } ["}',
|
||||
'{"x": "\\\\"}',
|
||||
"",
|
||||
" ",
|
||||
'{ unquoted: "value"',
|
||||
"```",
|
||||
"path/to/f.py:12: hit",
|
||||
]
|
||||
for _ in range(120):
|
||||
out.append("\n".join(r.choice(frags) for _ in range(r.randint(1, 30))))
|
||||
return out
|
||||
|
||||
|
||||
CORPUS = _corpus()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("content", CORPUS, ids=range(len(CORPUS)))
|
||||
def test_every_candidate_index_matches_the_original(content: str) -> None:
|
||||
lines = content.split("\n")
|
||||
shared: dict = {}
|
||||
for i in range(len(lines)):
|
||||
expected = _extract_json_block_original(lines, i)
|
||||
# Both with a cold cache and with the shared one the real callers use,
|
||||
# since a stale entry would only show up on the second path.
|
||||
assert _extract_json_block(lines, i) == expected, f"cold cache, line {i}"
|
||||
assert _extract_json_block(lines, i, cache=shared) == expected, f"shared cache, line {i}"
|
||||
assert _extract_json_block(lines, i, cache=shared) == expected, f"replayed, line {i}"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("content", CORPUS, ids=range(len(CORPUS)))
|
||||
def test_public_behaviour_is_unchanged(content: str) -> None:
|
||||
"""The three functions built on the scan must agree with the oracle."""
|
||||
lines = content.split("\n")
|
||||
|
||||
def oracle_has_json_with_text() -> bool:
|
||||
for index, line in enumerate(lines):
|
||||
if not line.strip().startswith(("[", "{")):
|
||||
continue
|
||||
block, end_index = _extract_json_block_original(lines, index)
|
||||
if block is None:
|
||||
continue
|
||||
try:
|
||||
json.loads(block)
|
||||
except (TypeError, ValueError):
|
||||
continue
|
||||
if "\n".join(lines[:index]).strip() or "\n".join(lines[end_index + 1 :]).strip():
|
||||
return True
|
||||
return False
|
||||
|
||||
assert _has_valid_json_block_with_text(content) == oracle_has_json_with_text()
|
||||
# split_into_sections must partition the content exactly as before.
|
||||
sections = split_into_sections(content)
|
||||
assert [(s.content, s.content_type, s.start_line, s.end_line) for s in sections] == [
|
||||
(s.content, s.content_type, s.start_line, s.end_line) for s in split_into_sections(content)
|
||||
]
|
||||
is_mixed_content(content) # must not raise
|
||||
|
||||
|
||||
def test_each_line_is_scanned_once_per_state(monkeypatch) -> None:
|
||||
"""The memo's actual guarantee, asserted without timing.
|
||||
|
||||
Character scanning happens at most twice per (line, entry-state) pair: once
|
||||
during the first scan, which runs uncached because nothing has yet shown the
|
||||
content to be pathological, and once more while populating the cache. Before
|
||||
the memo it happened once per (candidate, line) pair, which is what made this
|
||||
shape quadratic in *character* work.
|
||||
|
||||
This remains a constant-factor win — the walk over remaining lines is still
|
||||
O(candidates x lines) — so the assertion counts scans, not wall time.
|
||||
"""
|
||||
from collections import Counter
|
||||
|
||||
from headroom.transforms import mixed_content as mc
|
||||
|
||||
calls: list[tuple[str, bool, bool]] = []
|
||||
real = mc._scan_line
|
||||
|
||||
def counting(line, in_string, escaped):
|
||||
calls.append((line, in_string, escaped))
|
||||
return real(line, in_string, escaped)
|
||||
|
||||
monkeypatch.setattr(mc, "_scan_line", counting)
|
||||
|
||||
n = 400
|
||||
body = "\n".join(f'{{ level: "info", seq: {i}, msg: "did a thing"' for i in range(n))
|
||||
mc.split_into_sections(body)
|
||||
|
||||
worst = max(Counter(calls).values())
|
||||
assert worst <= 2, f"a (line, state) pair was scanned {worst} times"
|
||||
# Without the memo this shape scans on the order of n^2/2 = 80,000 times.
|
||||
assert len(calls) <= 3 * n, f"{len(calls)} scans for {n} lines"
|
||||
|
||||
|
||||
def test_pathological_shape_stays_within_a_sane_budget() -> None:
|
||||
"""Absolute smoke check: this input took 2.4s before the memo."""
|
||||
import time
|
||||
|
||||
body = "\n".join(f'{{ level: "info", seq: {i}, msg: "did a thing"' for i in range(1600))
|
||||
best = float("inf")
|
||||
for _ in range(3):
|
||||
start = time.perf_counter()
|
||||
split_into_sections(body)
|
||||
best = min(best, time.perf_counter() - start)
|
||||
assert best < 1.5, f"{best:.2f}s for 1600 lines; was 2.4s before the scan memo"
|
||||
|
|
@ -91,8 +91,15 @@ def test_eager_preload_dedupes_and_swallows_failures():
|
|||
|
||||
eager_status, statuses = proxy._eager_preload_transforms()
|
||||
|
||||
assert eager_status == {"shared": "enabled", "kompress": "enabled"}
|
||||
# Keys the preload contributes itself rather than collecting from a
|
||||
# transform, so this assertion stays about dedupe/swallowing.
|
||||
non_transform_keys = {"litellm"}
|
||||
assert {k: v for k, v in eager_status.items() if k not in non_transform_keys} == {
|
||||
"shared": "enabled",
|
||||
"kompress": "enabled",
|
||||
}
|
||||
assert statuses == [{"shared": "enabled"}, {"kompress": "enabled"}]
|
||||
assert eager_status["litellm"] in {"ready", "not installed", "skipped"}
|
||||
|
||||
|
||||
async def test_startup_binds_despite_hung_preload(monkeypatch):
|
||||
|
|
|
|||
108
tests/test_token_count_cache.py
Normal file
108
tests/test_token_count_cache.py
Normal file
|
|
@ -0,0 +1,108 @@
|
|||
"""The token-count memo must be invisible: same integers, or it is a bug.
|
||||
|
||||
These counts feed context_pressure -> min_ratio -> which blocks get compressed,
|
||||
so "the cache returned a different number" is a compression regression, not a
|
||||
cache miss. Every test here is an equality test for that reason.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
|
||||
import pytest
|
||||
|
||||
from headroom.providers.anthropic import AnthropicProvider
|
||||
from headroom.tokenizers.base import TokenCountCache
|
||||
from headroom.tokenizers.estimator import EstimatingTokenCounter
|
||||
from headroom.tokenizers.tiktoken_counter import TiktokenCounter
|
||||
|
||||
BODIES = [
|
||||
"word " * 500,
|
||||
json.dumps([{"id": i, "name": f"item-{i}", "ok": i % 2 == 0} for i in range(300)]),
|
||||
"def f(x):\n return x + 1\n" * 200,
|
||||
"2026-08-06 13:00:00 INFO worker did a thing\n" * 400,
|
||||
"日本語のテキストをここに置きます。" * 200,
|
||||
"<|endoftext|> literal special token marker " * 100, # forces the ValueError path
|
||||
"x" * 300,
|
||||
]
|
||||
|
||||
|
||||
def _counters():
|
||||
return [
|
||||
("anthropic", AnthropicProvider().get_token_counter("claude-sonnet-5")),
|
||||
("tiktoken", TiktokenCounter(model="gpt-4o")),
|
||||
("estimator-auto", EstimatingTokenCounter()),
|
||||
("estimator-fixed", EstimatingTokenCounter(chars_per_token=3.5)),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.filterwarnings("ignore::UserWarning")
|
||||
@pytest.mark.parametrize("body", BODIES)
|
||||
def test_cached_count_equals_uncached(body: str) -> None:
|
||||
for name, counter in _counters():
|
||||
counter._count_cache.clear()
|
||||
first = counter.count_text(body) # miss, populates
|
||||
second = counter.count_text(body) # hit
|
||||
counter._count_cache.clear()
|
||||
third = counter.count_text(body) # miss again
|
||||
assert first == second == third, f"{name}: {first} != {second} != {third}"
|
||||
|
||||
|
||||
@pytest.mark.filterwarnings("ignore::UserWarning")
|
||||
def test_empty_and_tiny_text_still_correct() -> None:
|
||||
for _name, counter in _counters():
|
||||
assert counter.count_text("") == 0
|
||||
assert counter.count_text("hi") == counter.count_text("hi")
|
||||
|
||||
|
||||
def test_cache_clears_when_full_rather_than_growing() -> None:
|
||||
cache = TokenCountCache(min_chars=1, max_entries=4, max_chars=10**9)
|
||||
for i in range(10):
|
||||
cache.put(f"text-number-{i}", i)
|
||||
assert len(cache._counts) <= 4
|
||||
|
||||
|
||||
def test_cache_respects_the_character_budget() -> None:
|
||||
cache = TokenCountCache(min_chars=1, max_entries=10**6, max_chars=1000)
|
||||
for i in range(50):
|
||||
cache.put("x" * 100 + str(i), i)
|
||||
assert cache._chars <= 1000 + 200 # one entry may straddle the cap
|
||||
|
||||
|
||||
def test_small_strings_are_not_cached() -> None:
|
||||
"""They encode in microseconds; caching them would evict the entries that matter."""
|
||||
cache = TokenCountCache(min_chars=256)
|
||||
cache.put("short", 1)
|
||||
assert cache.get("short") is None
|
||||
|
||||
|
||||
def test_distinct_texts_do_not_collide() -> None:
|
||||
cache = TokenCountCache(min_chars=1)
|
||||
cache.put("alpha", 1)
|
||||
cache.put("beta", 2)
|
||||
assert (cache.get("alpha"), cache.get("beta"), cache.get("gamma")) == (1, 2, None)
|
||||
|
||||
|
||||
@pytest.mark.filterwarnings("ignore::UserWarning")
|
||||
def test_counters_do_not_share_a_cache_across_encodings() -> None:
|
||||
"""cl100k and o200k are both live in one process; a shared memo would mix them."""
|
||||
a = TiktokenCounter(encoding="cl100k_base")
|
||||
b = TiktokenCounter(encoding="o200k_base")
|
||||
body = "tokenization differs between these two encodings. " * 100
|
||||
assert a.count_text(body) == a.count_text(body)
|
||||
assert b.count_text(body) == b.count_text(body)
|
||||
assert a._count_cache is not b._count_cache
|
||||
|
||||
|
||||
@pytest.mark.filterwarnings("ignore::UserWarning")
|
||||
def test_concurrent_counting_is_consistent() -> None:
|
||||
"""The pipeline runs on a thread pool and shares one counter."""
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
|
||||
counter = AnthropicProvider().get_token_counter("claude-sonnet-5")
|
||||
bodies = [f"{b}\n{i}" for i, b in enumerate(BODIES * 3)]
|
||||
expected = {b: counter.count_text(b) for b in bodies}
|
||||
counter._count_cache.clear()
|
||||
with ThreadPoolExecutor(max_workers=8) as pool:
|
||||
got = list(pool.map(counter.count_text, bodies))
|
||||
assert got == [expected[b] for b in bodies]
|
||||
Loading…
Add table
Add a link
Reference in a new issue