Split oversized result uploads so wide-taxonomy batches don't get rejected - #149

sequenceDiagram
  participant Worker
  participant ResultPoster
  participant post_batch_results
  participant chunk_results_by_size
  participant HTTP
  Worker->>ResultPoster: __init__(max_post_bytes=settings.antenna_result_post_max_bytes)
  Note over ResultPoster: stores self.max_post_bytes
  Worker->>ResultPoster: post_batch(results)
  ResultPoster->>post_batch_results: post_batch_results(..., max_bytes=self.max_post_bytes)
  post_batch_results->>post_batch_results: serialize to JSON dicts
  post_batch_results->>chunk_results_by_size: JSON dicts, max_bytes
  chunk_results_by_size-->>post_batch_results: list of chunks
  loop for each chunk
    post_batch_results->>HTTP: POST {results: chunk}
    HTTP-->>post_batch_results: 200/validation error
    post_batch_results->>post_batch_results: log per-chunk status
  end
  post_batch_results-->>ResultPoster: all_ok flag
  ResultPoster-->>Worker: success status

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

_{Comment @coderabbitai help to get the list of available commands.}

Claude says: This PR had been open a couple of months collecting small fixes, so before another round of line-by-line I did a structural pass: what the change is really deciding, and what might be wrong in code the diff never touches. Most of what came out is already fixed on the branch. The rest are decisions worth making on purpose rather than by default.

Fixed on the branch

1. The size cap was not actually capping. (fa120c7)

This is the one worth reading. The packer measured each result with a compact JSON encoding, but requests serializes a json= argument using json.dumps defaults, which put a space after every comma and colon. On the long numeric scores and logits arrays this PR exists to handle, that whitespace is about a quarter of the body. So the cap was an undercount of the thing it was capping. At the shipped 25 MB default with the ~29k-class model, a batch packed into a 25.06 MB request:

29000 classes, cap 25.0 MB | before: max body 25.06 MB (100%)  <-- over cap
                           | after:  max body 23.50 MB (94%)

The overshoot is small at the current default and the proxy limit leaves room, so this was very unlikely to bite in practice. It mattered because the guarantee was not sound: the margin came from headroom rather than from the cap. Sizes are now measured the way requests encodes them, and the envelope and separator costs are derived from the encoder rather than written out, so the two cannot drift apart again.

2. A ValidationError could strand every upload behind it. (fa120c7)

CodeRabbit flagged this and was half right; details and the empirical check are in that thread. Short version: model_validate() raises Pydantic's ValidationError, which is not a requests exception and escaped the handler. Before chunking that failed one POST. After chunking it aborts the loop, so uploads queued behind the bad response are never attempted and their images wait on redelivery for nothing. The handler now catches ValueError as well, which covers it. The other half of the report, JSON decode errors, was already covered, since requests.exceptions.JSONDecodeError subclasses RequestException.

3. The five failing test_worker.py tests. (fa120c7)

These were not a product regression, which surprised me. The worker tests build settings as a bare MagicMock, so the new antenna_result_post_max_bytes arrived as a MagicMock and raised TypeError the moment the packer compared it against an int. That got swallowed by the broad except Exception in _post_results_sync and reported as a failed post, so the symptom was assert 0 == 2: results that were simply never posted. Production always gets a real int from Pydantic, so the impact was confined to the tests.

Both settings helpers now supply a real number. Worth noting the same trap had already been hit once before, which is why localization_batch_size carries a # Real integer comment. A MagicMock settings object silently invents a value for every new setting the worker reads, so each one of these lands as a mysterious product failure. If it happens a third time, building the double from a real Settings would end the category.

4. benchmark.py ignored the new setting. (da59959)

The kind of thing a diff review will not catch, since the file is not in the diff. It built its ResultPoster without the cap, so it fell back to the module default and quietly disregarded AMI_ANTENNA_RESULT_POST_MAX_BYTES. A benchmark against a tuned deployment would report a different request count than the worker actually makes. It now takes the setting, like the worker does. The ResultPoster docstring also documents max_post_bytes and future_timeout, per the other CodeRabbit note.

Worth a decision before merge

5. Chunking over gzip, made explicit. The body argued gzip is "the highest-leverage lever" and that payloads will grow as embeddings arrive, then shipped chunking. I think that is the right call: chunking is client-side only and unblocks affected jobs now, whereas gzip needs the API server to decompress request bodies and the proxy to pass them through. They also compose rather than compete. But it was being made by default rather than on purpose, so the body now says so. Happy to be overruled on the ordering.

6. The success boolean is currently decorative. post_batch_results() returns True only when every upload lands, and that is now pinned by tests. But the value feeds metrics and logging only: _process_job returns did_work regardless, so a batch with a failed upload does not fail the job. In practice recovery comes from redelivery of the unacknowledged images, which is per-image and works. I lean towards leaving it, since redelivery is the more granular mechanism and failing the job would be coarser. Flagging it because "we rely on redelivery, not on the return value" is a design position worth holding deliberately rather than discovering later.

7. Still no live-worker run. Everything here is verified against the serialization path, not a GPU job. The remaining unknown is the one the PR already names: an image dense enough to exceed the cap on its own still goes out in a single upload. A real large-taxonomy job, checking each upload lands under the proxy limit, is the last gap. That is a "confirm after merge" item to me rather than a blocker, given the cap now measures honestly.

Smaller cleanups (low priority)

• The body claimed 7 tests covering "all-chunks-must-succeed semantics", but no test actually exercised a failing chunk. That claim is now true: TestPostBatchResultsFailureSemantics pins that a failed upload reports failure, does not abort the uploads after it, and that an off-schema or undecodable response fails only its own upload. 13 tests in that file now.
• Sending compact JSON would cut roughly 20% off every request for free, by encoding the body in the client instead of handing requests a dict. Left out here because the shared test HTTP shim forwards only the json= argument and would need updating with it, and that file is being edited on another branch. Noted as a follow-up.

Where it stands

trapdata/antenna is 26/26 locally and 0 failures in CI. The full CI run is 49 passed, 3 failed on both 3.10 and 3.12, and those three are test_api.py::test_config_num_classification_predictions, test_api.py::test_logits_in_classification_response, and test_models.py::test_url. All three are environmental, unrelated to this branch, and being fixed separately: PR #163 fails the identical three on its current run. None of the five test_worker.py failures remain.

No history was rewritten, both pushes were fast-forwards, so anything stacked on this branch is unaffected.