mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
fix(proxy/batch): don't crash an OpenAI batch on a valid-JSON non-object line (#2316)
## Description
A single malformed line can abort compression for an entire OpenAI batch
upload.
`_compress_batch_jsonl` parses each JSONL line and immediately reads the
request body:
```python
request_obj = json.loads(line)
body = request_obj.get("body", {})
messages = body.get("messages", [])
...
except json.JSONDecodeError as e:
...
compressed_lines.append(line) # keep original on error
```
`json.loads` returns a valid JSON *value*, which isn't necessarily an
object. A line like `[1, 2, 3]`, `"hello"`, or `null` parses fine, but
`request_obj.get(...)` on a list/str/None raises `AttributeError`.
Likewise a request object whose `body` is present but not a dict
(`{"body": "..."}`) makes `body.get("messages", ...)` raise. The
surrounding `except json.JSONDecodeError` doesn't catch
`AttributeError`, so the exception propagates out of
`_compress_batch_jsonl` and the whole batch-create request fails.
This is the OpenAI batch upload path; the file is user-supplied, so a
single stray non-object line takes the batch down instead of just being
passed through like the other non-compressible cases already are.
## Fix
Guard for non-object shapes and pass them through unchanged:
```python
request_obj = json.loads(line)
if not isinstance(request_obj, dict):
compressed_lines.append(line)
total_requests += 1
continue
body = request_obj.get("body", {})
if not isinstance(body, dict):
body = {}
```
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/proxy/handlers/batch.py`: in `_compress_batch_jsonl`, pass
through a non-dict parsed line and coalesce a non-dict `body` to `{}`.
- `tests/test_proxy_handlers_batch.py`: new test that array / string /
null lines and a non-dict `body` pass through without crashing and are
preserved.
- `CHANGELOG.md`: Bug Fixes entry.
## Testing
- [ ] 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
### Test Output
```text
$ uvx ruff@0.15.17 check headroom/proxy/handlers/batch.py tests/test_proxy_handlers_batch.py
All checks passed!
$ uvx mypy@1.20.2 --ignore-missing-imports headroom/proxy/handlers/batch.py
Success: no issues found in 1 source file
```
## Real Behavior Proof
- Environment: Windows 11, Python 3.12, `uvx ruff@0.15.17` / `uvx
mypy@1.20.2`. A full `pytest` OOM-kills this box (ML stack import), so I
reproduced the per-line handling with a dependency-free script and left
the full pytest to CI.
- Exact command / steps: ran lines `[1,2,3]`, `"hello"`, `null`, and
`{"body": "not-a-dict"}` through the OLD (bare `.get`) and NEW
(isinstance-guarded) logic, plus a normal request and a `not-json` line
as controls.
- Observed result: OLD raises `AttributeError` on each non-object line
and on the non-dict body; NEW passes the non-object lines through
unchanged, coalesces the non-dict body to `{}`, still processes a normal
request, and still passes `not-json` through as a JSON-decode error.
- Not tested: a live OpenAI batch upload end-to-end; full local `pytest`
deferred to CI (OOM).
## 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" box is unchecked because the full suite
imports the ML stack, which I can't run here. The new test reuses the
existing `DummyBatchHandler` / `install_batch_support_modules` harness
in `tests/test_proxy_handlers_batch.py` (the same one the neighbouring
invalid-line test uses), so it runs under the normal CI pytest job;
behaviour is additionally verified by the standalone proof above.
---------
Co-authored-by: JerrettDavis <mxjerrett@gmail.com>
This commit is contained in:
parent
c471800e8e
commit
1f2c681c0b
2 changed files with 45 additions and 0 deletions
|
|
@ -1104,7 +1104,18 @@ class BatchHandlerMixin:
|
|||
|
||||
try:
|
||||
request_obj = json.loads(line)
|
||||
# A JSONL line can be valid JSON that isn't a request object
|
||||
# (e.g. a stray array, string, or null). `.get` on it would raise
|
||||
# AttributeError, which the outer `except json.JSONDecodeError`
|
||||
# below does not catch — crashing the whole batch. Pass such lines
|
||||
# through unchanged, like the other non-compressible cases.
|
||||
if not isinstance(request_obj, dict):
|
||||
compressed_lines.append(line)
|
||||
total_requests += 1
|
||||
continue
|
||||
body = request_obj.get("body", {})
|
||||
if not isinstance(body, dict):
|
||||
body = {}
|
||||
messages = body.get("messages", [])
|
||||
model = body.get("model", "gpt-4")
|
||||
|
||||
|
|
|
|||
|
|
@ -651,6 +651,40 @@ async def test_compress_batch_jsonl_without_optimization_handles_invalid_lines(
|
|||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_compress_batch_jsonl_handles_non_object_lines(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
# A JSONL line that is valid JSON but not a request object (array/string/
|
||||
# null), or a request whose `body` isn't a dict, must pass through instead
|
||||
# of crashing the whole batch (`.get` on a non-dict raises AttributeError,
|
||||
# which the JSONDecodeError guard does not catch).
|
||||
install_batch_support_modules(monkeypatch, tokenizer_count=12)
|
||||
handler = DummyBatchHandler()
|
||||
content = "\n".join(
|
||||
[
|
||||
json.dumps(
|
||||
{"body": {"model": "gpt-4o", "messages": [{"role": "user", "content": "hi"}]}}
|
||||
),
|
||||
json.dumps([1, 2, 3]),
|
||||
json.dumps("hello"),
|
||||
"null",
|
||||
json.dumps({"body": "not-a-dict"}),
|
||||
]
|
||||
)
|
||||
|
||||
lines, stats = await handler._compress_batch_jsonl(content, "req-1")
|
||||
|
||||
assert len(lines) == 5
|
||||
assert json.loads(lines[1]) == [1, 2, 3]
|
||||
assert json.loads(lines[2]) == "hello"
|
||||
assert json.loads(lines[3]) is None
|
||||
assert json.loads(lines[4]) == {"body": "not-a-dict"}
|
||||
assert stats["total_requests"] == 5
|
||||
# None of these are JSON decode errors, so the error counter stays at 0.
|
||||
assert stats["errors"] == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_compress_batch_jsonl_uses_pipeline_and_ccr_injection(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue