mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
0cddac632d
## Description
The usage reporter permanently drops a reporting window's usage whenever
the send to the cloud fails.
`UsageReporter._report_usage` computes usage as a **delta** against the
last snapshot, POSTs it, and then rebases the baseline:
```python
try:
resp = await client.post(f"{self._cloud_url}/v1/license/usage", json=payload, timeout=10.0)
if resp.status_code == 200:
...
else:
logger.warning("Usage report returned status %d", resp.status_code)
except Exception:
logger.warning("Failed to send usage report", exc_info=True)
# Update snapshot
self._snapshot_metrics() # runs on success, non-200, AND exception
self._last_report_time = now
```
`_snapshot_metrics()` rebases `_last_tokens_saved_by_model` /
`_last_tokens_sent_by_model` / `_last_requests_by_model` to the current
cumulative counters. Because it runs unconditionally after the POST, a
report that fails to send — non-200 or a raised exception — still
advances the baseline. The module is explicitly built to tolerate a
briefly-unreachable cloud (7-day grace, cached license), so this is a
normal, recurring situation.
The consequence: the failed window's requests and tokens are never
re-included. The next report is a delta from the advanced baseline, so
that window is silently and permanently lost from usage-based billing /
quota. Every transient network blip under-counts usage. (The
`total_requests == 0` early-return already gets this right — it advances
only `_last_report_time`, without snapshotting, since there's nothing to
lose.)
## Fix
Advance the baseline (`_snapshot_metrics()` and `_last_report_time`)
only inside the `resp.status_code == 200` branch. On a non-200 or an
exception, both baselines are left intact, so the next report covers the
full period since the last successful send and re-includes the
previously-failed window.
Closes #
## Type of Change
- [x] Bug fix (non-breaking change that fixes an issue)
- [ ] New feature (non-breaking change that adds functionality)
- [ ] Breaking change (fix or feature that would cause existing
functionality to change)
- [ ] Documentation update
- [ ] Performance improvement
- [ ] Code refactoring (no functional changes)
## Changes Made
- `headroom/telemetry/reporter.py`: move `_snapshot_metrics()` +
`_last_report_time = now` into the 200 branch of `_report_usage`.
- `tests/test_usage_reporter_snapshot.py`: new tests — baseline advances
on 200, and stays intact on a non-200 and on an exception (window
preserved).
- `CHANGELOG.md`: Bug Fixes entry.
## Testing
- [ ] Unit tests pass (`pytest`)
- [x] Linting passes (`ruff check .`)
- [ ] Type checking passes (`mypy headroom`)
- [x] New tests added for new functionality
- [ ] Manual testing performed
### Test Output
```text
$ uvx ruff@0.15.17 check headroom/telemetry/reporter.py tests/test_usage_reporter_snapshot.py
All checks passed!
$ python -m py_compile headroom/telemetry/reporter.py tests/test_usage_reporter_snapshot.py
OK
```
## Real Behavior Proof
- Environment: Windows 11, Python 3.12, `uvx ruff@0.15.17`. Importing
`headroom` pulls in the torch/transformers stack and a full `pytest`
gets OOM-killed on this box, so I verified the delta accounting with a
dependency-free script that models two windows across a failed then
successful send, and left the full pytest to CI.
- Exact command / steps: window 1 saves 100 tokens and the send fails;
window 2 saves another 50 (cumulative 150) and the send succeeds. Ran
under the old (unconditional snapshot) and new (snapshot-on-200) logic.
- Observed result: old delivers only 50 tokens total (window 1's 100
dropped when the baseline advanced on the failed send); new delivers the
full 150 (window 2's delta re-includes window 1). The new tests assert
the baseline advances on 200 and is untouched on a non-200 / exception,
driving the real `_report_usage` with a fake proxy + client.
- Not tested: a live cloud round-trip; full local `pytest` deferred to
CI (OOM, per above).
## 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
- [ ] New and existing unit tests pass locally with my changes
- [x] I have updated the CHANGELOG.md if applicable
## Additional Notes
The "unit tests pass locally" and "type checking" boxes are unchecked
because the full suite imports the ML stack, which I can't run in this
environment; the change moves two lines into the success branch,
verified by the standalone delta-accounting proof and new tests that
drive the real `_report_usage` via `object.__new__` with a fake proxy
and HTTP client (200, non-200, and exception).
Co-authored-by: JerrettDavis <mxjerrett@gmail.com>
Co-authored-by: Tejas Chopra <chopratejas@gmail.com>
72 lines
2.1 KiB
Python
72 lines
2.1 KiB
Python
"""UsageReporter must only advance its delta baseline after a confirmed 200.
|
|
|
|
Snapshotting on a failed send permanently drops that window's usage from the
|
|
delta-based usage report (billing/quota under-count)."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from types import SimpleNamespace
|
|
|
|
import anyio
|
|
|
|
from headroom.telemetry.reporter import UsageReporter
|
|
|
|
|
|
class _Resp:
|
|
def __init__(self, status_code: int) -> None:
|
|
self.status_code = status_code
|
|
|
|
def json(self) -> dict:
|
|
return {}
|
|
|
|
|
|
def _make_reporter(post_outcome) -> UsageReporter:
|
|
r = object.__new__(UsageReporter)
|
|
ct = SimpleNamespace(
|
|
_tokens_saved_by_model={"m": 100},
|
|
_tokens_sent_by_model={"m": 400},
|
|
_requests_by_model={"m": 5},
|
|
)
|
|
r._proxy = SimpleNamespace(cost_tracker=ct)
|
|
r._last_report_time = None
|
|
r._last_tokens_saved_by_model = {}
|
|
r._last_tokens_sent_by_model = {}
|
|
r._last_requests_by_model = {}
|
|
r._license_key = "k"
|
|
r._cloud_url = "https://cloud.example"
|
|
r._license_info = None
|
|
|
|
class _Client:
|
|
async def post(self, *args, **kwargs):
|
|
if isinstance(post_outcome, Exception):
|
|
raise post_outcome
|
|
return _Resp(post_outcome)
|
|
|
|
async def _get_client():
|
|
return _Client()
|
|
|
|
r._get_client = _get_client
|
|
return r
|
|
|
|
|
|
def test_baseline_advances_only_on_success():
|
|
r = _make_reporter(200)
|
|
anyio.run(r._report_usage)
|
|
# Success -> baseline rebased to the current cumulative counters.
|
|
assert r._last_tokens_saved_by_model == {"m": 100}
|
|
assert r._last_requests_by_model == {"m": 5}
|
|
|
|
|
|
def test_baseline_not_advanced_on_non_200():
|
|
r = _make_reporter(500)
|
|
anyio.run(r._report_usage)
|
|
# Failed send -> baseline untouched so the window is retried next report.
|
|
assert r._last_tokens_saved_by_model == {}
|
|
assert r._last_requests_by_model == {}
|
|
|
|
|
|
def test_baseline_not_advanced_on_exception():
|
|
r = _make_reporter(RuntimeError("network down"))
|
|
anyio.run(r._report_usage)
|
|
assert r._last_tokens_saved_by_model == {}
|
|
assert r._last_requests_by_model == {}
|