[AMOOOMA]

Split out of #40760, following review feedback to separate the generic pieces from the vLLM-specific ones. This PR is a small generic RunInference hook.

Why

Runners that autoscale on num_loaded_models_<model_identifier> (Dataflow's GPU autoscaling) look the metric up by the identifier RunInference annotates the inference step with: the model_identifier passed to RunInference, or a uuid it generates. ModelManager publishes under that identifier (and since #40714 keeps it stable across side-input model updates). But a model handler that publishes the metric itself, e.g. one that fronts an out-of-process model server, has no way to learn that value.

What

• ModelHandler.set_model_identifier(model_identifier): a no-op by default. _RunInferenceDoFn.setup() calls it on each worker, before the first load_model(), with the same tag it annotates the step with and gives ModelManager.
• Forwarded by the keyed, maybe-keyed, prebatched, preprocessing and postprocessing wrappers (KeyedModelHandler forwards to every per-key handler).

No behavior change for existing handlers.

Testing

base_test.py:

• An explicit model_identifier reaches the handler before load_model().
• Without one, the generated uuid equals the model_identifier annotation on RunInference/BeamML_RunInference.
• The identifier is forwarded through KeyedModelHandler and pre/post-processing wrappers.

---

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

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 79.82%. Comparing base (12bbb17) to head (c29a70f).
⚠️ Report is 10 commits behind head on master.

Additional details and impacted files

@@            Coverage Diff             @@
##           master   #40906      +/-   ##
==========================================
- Coverage   79.83%   79.82%   -0.02%     
==========================================
  Files         494      494              
  Lines       83575    83597      +22     
==========================================
+ Hits        66722    66731       +9     
- Misses      16853    16866      +13     

Flag	Coverage Δ
python	79.82% <100.00%> (-0.02%)	⬇️

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.