Make batch inference resumable and leak-free

[RickyVishwakarma]

What

Hardens the SGLang batch runner (infer.py) for long-horizon documents.

• Resume: skip pages whose output already exists; --overwrite forces a re-run.
• Atomic writes: each page streams to a .part file, renamed only on completion —
  a present output file always means a fully parsed page, so resume is safe.
• Fix temp-dir leak: rendered PDF pages now go to a temp dir that is cleaned up
  after the run; --keep_temp retains it.
• Run manifest: writes <output_dir>/manifest.json with per-page status
  (ok/skipped/failed), tokens, decode time, and retry attempts.
• Ctrl-C: cancels pending jobs, keeps finished pages, still writes the manifest.

Tests

Adds tests/test_infer.py (mocks the SGLang endpoint — runs CPU-only, no GPU/model)
and a GitHub Actions CI workflow, per CONTRIBUTING.md. All 11 tests pass locally.

Why

A long PDF that fails or is interrupted partway currently has to be re-run from page 1.
This makes batch runs idempotent and auditable.

[kushdab]

Well-designed PR — the atomic-write + resume + manifest combination is exactly what batch OCR pipelines need for production reliability. The implementation is clean throughout. A few observations:

What's particularly strong

• .part → os.replace() atomicity is the right pattern. A present .md file always means a fully parsed page, so resume is provably safe. remove_partial() on failure correctly cleans the .part sibling.
• is_completed() checking both existence and non-empty catches the truncation edge case where a previous run created the file but wrote nothing.
• executor.shutdown(wait=False, cancel_futures=True) on Ctrl-C cancels pending (not yet started) futures. Futures already in-flight run to completion (or raise) before shutdown(wait=True) returns — so the .part cleanup path in infer_one is always reached, and the manifest captures their final status. This is correct.
• avg_decode is now filtered on status == "ok" rather than tokens > 0, which implicitly fixes the same averaging bug that PR Fix avg_decode calculation when decode_time is 0 for single tokens #13 targets — skipped requests contribute zero tokens but are excluded from the decode average by construction.

One functional gap: missing --image_mode validation for PDF mode

gundam uses crop-mode tiling (calling infer() per tile), which is designed for a single image. When --pdf is passed, infer_multi() is used under the hood — gundam silently produces wrong output there. PRs #26 and #21 both add this guard; it would be worth including here too:

# In parse_args(), after args = parser.parse_args():
if args.pdf and args.image_mode != "base":
    parser.error("--pdf requires --image_mode base (gundam is for single images only)")

Minor: manifest.json is a fixed name in output_dir

If two runs share an output_dir (e.g. resuming a PDF batch mid-way through, then running a second PDF into the same dir), the second run overwrites the first manifest. Could accept --manifest as an optional path that defaults to <output_dir>/manifest.json, or include the source filename in the manifest name (<prefix>_manifest.json). Low priority, but easy to add.

Overlap note

This PR covers the same ground as:

• PR infer.py: CLI flags for ngram params, --resume, --results_jsonl, --max_pages/images, mode validation, tempdir cleanup #21 — --resume, structured per-request results, --max_pages. Different implementation: infer.py: CLI flags for ngram params, --resume, --results_jsonl, --max_pages/images, mode validation, tempdir cleanup #21 uses a write_output flag; Make batch inference resumable and leak-free #36 uses .part files + os.replace. The atomic-write approach here is more robust against partial writes.
• PR fix: clean up PDF tmpdir and validate --image_mode in PDF mode #26 — PDF tmpdir cleanup. Different mechanism: fix: clean up PDF tmpdir and validate --image_mode in PDF mode #26 keeps tmpdir management inside pdf_to_images() (as a returned TemporaryDirectory); Make batch inference resumable and leak-free #36 creates it in build_jobs() and tears it down in run()'s finally. Both are correct; the finally path here is slightly easier to reason about.
• PR fix(infer+docs): PDF tmpdir cleanup + remove broken kernels install from README #34 — PDF tmpdir (via atexit) + README kernels fix.

Of the tmpdir approaches across these three PRs, finally is the most reliable — it runs even if run() raises, which atexit also guarantees, but finally makes the scope explicit. The missing piece relative to PR #26 is the --image_mode validation noted above.

Overall this is ready to merge pending that one guard. ✅

[RickyVishwakarma]

Added the --image_mode base guard for PDF mode in parse_args().

[rajpratham1]

This is a substantial and well-thought-out enhancement rather than just a bug fix. It makes long-running OCR jobs resumable, prevents partially written output files from being treated as completed work, adds interruption handling, introduces a manifest for tracking progress, and backs everything with comprehensive unit tests. The implementation also updates the documentation and CI, making the feature much easier to maintain.

[RickyVishwakarma]

This is a substantial and well-thought-out enhancement rather than just a bug fix. It makes long-running OCR jobs resumable, prevents partially written output files from being treated as completed work, adds interruption handling, introduces a manifest for tracking progress, and backs everything with comprehensive unit tests. The implementation also updates the documentation and CI, making the feature much easier to maintain.

Not in progress — it's complete. Added the --image_mode base guard kushdab requested, with tests and a README fix. Ready to merge. Thanks for the review!