[AMOOOMA]

Split out of #40760, following review feedback to separate the generic pieces from the vLLM-specific ones and to split autoscaling into collection/publishing vs. vLLM consumption. This PR is the generic library; the vLLM handler uses it in a follow-up, and a Triton handler is planned to reuse it.

Why

Dataflow's GPU autoscaling reads num_loaded_models_<tag> in the BeamML_ModelManager namespace as the number of messages each worker can hold, and sizes the pool roughly as ceil(backlog / reported). A handler that serves inference through an out-of-process model server always has exactly one loaded copy per worker, so it needs to publish an estimate of the worker's capacity under that metric instead.

What (apache_beam.ml.inference.autoscaling_metrics, new)

Nothing in this module depends on a particular model server.

• NumLoadedModelsPublisher (publishing): polls a capacity signal on a background thread and publishes it under num_loaded_models_<tag>.
  • Resets the process-wide distribution before each update, so the runner sees the latest value rather than a lifetime mean.
  • Optional aggregation on top (last by default; also lifetime_max, max, pNN, mean).
  • Never publishes 0 on start/stop, and a reloaded server's publisher supersedes the stale one.
• PrometheusCounterGaugeSampler (collection): reads a server's completed-requests counter and in-flight gauge from its Prometheus /metrics endpoint. The two metric names are the only server-specific input.
• CompletionRateCapacityEstimator (collection): turns those samples into capacity = per-slot throughput × p90 concurrency × multiplier over a trailing window, with warmup and idle-tick handling.
  • Does no I/O: it's fed (completed_total, running) samples.
  • Normalizing by concurrency keeps the estimate steady when a worker is briefly starved after a rebalance, instead of cratering and triggering a downscale spiral.
• AutoscalingMetricConfig: user-facing knobs, metric-tag resolution (resolve_tag: an explicit tag, then the RunInference model identifier verbatim, then a sanitized model name), and factories for the estimator and publisher.

Pairs with the ModelHandler.set_model_identifier() hook (separate PR), which supplies the RunInference model identifier for resolve_tag. No dependency between the two PRs.

Testing

autoscaling_metrics_test.py:

• Prometheus parsing and sampler.
• Estimator math, counter reset, clock skew, warmup, idle ticks.
• Publisher aggregation, cell reset, supersession.
• Config tag precedence.

The estimator gives identical results to the version validated in #40760's staging runs (checked on 2,000 randomized traces).

---

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.

[codecov]

Codecov Report

❌ Patch coverage is 88.67188% with 29 lines in your changes missing coverage. Please review.
✅ Project coverage is 79.85%. Comparing base (12bbb17) to head (6fd3db5).
⚠️ Report is 10 commits behind head on master.

Files with missing lines	Patch %	Lines
...on/apache_beam/ml/inference/autoscaling_metrics.py	88.67%	29 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master   #40907      +/-   ##
==========================================
+ Coverage   79.83%   79.85%   +0.01%     
==========================================
  Files         494      495       +1     
  Lines       83575    83837     +262     
==========================================
+ Hits        66722    66944     +222     
- Misses      16853    16893      +40     

Flag	Coverage Δ
python	79.85% <88.67%> (+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.