fix(inference): only retry batch-size probing on genuine OOM errors

[rtmalikian]

Summary

Fixes #1607: batch-size probing loops in the inference paths caught every RuntimeError and retried with batch_size * 3 // 4 with no lower bound or termination. A non-OOM failure was silently swallowed, batch size 1 became 0 on the first retry, and a persistent error looped at zero forever. (Every probe site also carried a dead duplicate except torch.cuda.OutOfMemoryError branch, unreachable because it subclasses RuntimeError.)

Changes

• New FlagEmbedding/inference/utils.py (dependency-free) with three helpers:
  • is_oom_error(e) — true for genuine OOM failures (CUDA / CPU allocator messages), covering torch.cuda.OutOfMemoryError
  • shrink_batch_size_on_oom(batch_size, e) — returns batch_size * 3 // 4, re-raising the original OOM error when the result would drop below 1
  • validate_batch_size(batch_size) — rejects non-positive caller-supplied batch sizes early with ValueError (None still allowed for the default fallback)
• 9 files, 10 probe sites — same surgical change in every handler: non-OOM RuntimeErrors now propagate immediately with their original traceback; only genuine OOM shrinks and retries; validate_batch_size at the top of each probing function. Files: encoder_only/base.py, encoder_only/m3.py, decoder_only/base.py, decoder_only/icl.py (2 sites), decoder_only/pseudo_moe.py, reranker/encoder_only/base.py, reranker/decoder_only/{base,layerwise,lightweight}.py
• New regression test tests/test_inference_batch_probing.py (13 tests, torch-free) covering helper behavior, fixed-loop semantics, and a structural sweep asserting every probe loop in the inference tree uses the fixed policy

Verification

• Reproduced the bug on the unmodified tree with the issue's stdlib-only AST repro: non-OOM RuntimeError at batch size 1 → probe sizes [1, 0, 0, 0, 0, 0, 0, 0] (swallowed, retried forever at 0)
• After the fix: non-OOM raises immediately after [1]; persistent OOM attempts [2, 1] then re-raises the original OOM (never 0); transient OOM at 32 recovers at 6 ([32, 24, 18, 13, 9, 6]); batch_size=0 → ValueError
• python3 -m unittest tests.test_inference_batch_probing: 13/13 OK; 3 structural tests verified to FAIL on the pre-fix tree (via git stash) and pass after
• python3 -m py_compile clean on all changed files
• Repo test suite: existing tests all require torch/transformers + HF model downloads, which are unavailable in this environment, so they could not be run here

Author

• Raphael Malikian (rtmalikian@gmail.com, GitHub: rtmalikian)
• Happy to collaborate on this or similar retrieval/inference work; open to contract or full-time roles in clinical AI

This PR was prepared with AI assistance and reviewed before submission.