[AMOOOMA]

Split out of #40760, following review feedback to separate the generic pieces from the vLLM-specific ones. This PR is the generic prefix-aware batching; vLLM support builds on it in a follow-up.

Why

BatchElements batches per worker, so an LLM model handler behind RunInference sends its server batches that mix unrelated prompt prefixes, and the server's prefix (KV) cache gets thrashed instead of reused.

What

BatchElements(key_fn=..., num_shards=...) (transforms/util.py)

• With key_fn set, a new WithBatchKey DoFn keys each element without the per-worker UUID that WithSharedKey / WithLengthBucketKey add. GroupIntoBatches then routes elements that share a key to the same worker, so every batch holds a single key.
• num_shards (an int, a per-key mapping, or a callable) spreads a hot key over a bounded number of random sub-keys, so one key isn't capped at one worker.
• Incomplete keyed batches flush after max_batch_duration_secs, or 1s when unset (GroupIntoBatches treats 0 as "no limit", which would stall low-traffic keys in streaming).
• Keyed batching works in streaming, unlike the unkeyed in-bundle path.

RunInference plumbing (ml/inference/base.py)

• ModelHandler(batch_key_fn=..., batch_key_num_shards=...) forwards to BatchElements.

apache_beam.ml.inference.prefix_batching (new, generic for LLM handlers)

• PrefixBatchingConfig, plus default extractors PromptPrefixKeyFn (completions) and ChatPrefixKeyFn (OpenAI-style chat messages; unknown message schemas raise rather than collapsing onto one key).
• Defaults to 8 shards per prefix. On a 5–40 RPS sine load (OpenOrca, T4, up to 50 workers), 4 shards left 18–33k requests queued behind hot prefixes at the end of the load. 8 shards wrote all 160k published requests and still scaled the pool down to 4 workers at low traffic.
• The generic BatchElements / ModelHandler defaults stay at 1 shard.

Testing

• util_test.py: homogeneous keyed batches, shard specs (int / mapping / callable), validation, default buffering.
• prefix_batching_test.py: extractors and config.
• base_test.py: batch_key_fn / num_shards reach BatchElements, and an end-to-end RunInference run produces prefix-pure batches.

---

Thank you for your contribution! Follow this checklist to help us incorporate your contribution quickly and easily:

• Replace these instructions with a description of your change and why you have made it.
• Mention the appropriate issue in your description (for example: addresses #123 or fixes #456).
• Update CHANGES.md with noteworthy changes.

[AMOOOMA]

R: @damccorm

[github-actions]

Stopping reviewer notifications for this pull request: review requested by someone other than the bot, ceding control. If you'd like to restart, comment assign set of reviewers

[codecov]

Codecov Report

❌ Patch coverage is 98.31933% with 2 lines in your changes missing coverage. Please review.
✅ Project coverage is 79.84%. Comparing base (12bbb17) to head (07aac31).
⚠️ Report is 11 commits behind head on master.

Files with missing lines	Patch %	Lines
sdks/python/apache_beam/transforms/util.py	96.36%	2 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master   #40905      +/-   ##
==========================================
+ Coverage   79.83%   79.84%   +0.01%     
==========================================
  Files         494      495       +1     
  Lines       83575    83699     +124     
==========================================
+ Hits        66722    66830     +108     
- Misses      16853    16869      +16     

Flag	Coverage Δ
python	79.84% <98.31%> (+0.01%)	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.