[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 how many requests each worker can handle, and sizes the pool roughly as ceil(backlog / reported). The estimate is the worker's sustained throughput in requests per second, so each worker is sized for about one second of it (scaled by multiplier). 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.

Usage: a model-server handler constructs the publisher directly, supplying only what it knows:

publisher = autoscaling_metrics.NumLoadedModelsPublisher(
    metrics_port_fn=lambda: server.port if server.is_up else None,
    completed_metric='vllm:request_success_total',
    running_metric='vllm:num_requests_running',
    model_identifier=model_identifier,
    model_name=model_name,
    config=autoscaling_config)  # An AutoscalingMetricConfig, or None.
publisher.start()  # Once the server is reachable; publisher.stop() at exit.

Components, each with one job:

• NumLoadedModelsPublisher: reads the server, estimates its capacity and publishes it from a background thread.
  • Each poll reads the port, scrapes /metrics, parses the two metrics, feeds them to its own CompletionRateCapacityEstimator and writes the estimate to num_loaded_models_<tag>. If config.load_signal_fn is set, it replaces the scrape and the estimate.
  • Metric tag: an explicit config.tag, then the RunInference model identifier verbatim, then a sanitized model name.
  • Resets the process-wide distribution before each update, so the runner sees the latest value rather than a lifetime mean.
  • Never injects a 0 on start, stop or a skipped poll, and a reloaded server's publisher supersedes the stale one.
• CompletionRateCapacityEstimator: turns a server's completed-requests counter and in-flight gauge 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.
  • Suits servers whose requests take long enough to be seen in flight (e.g. LLM servers); millisecond-latency servers can use a completion-rate mode or supply their own signal through config.load_signal_fn.
  • aggregation selects the estimate: 'concurrency' (default, above), or the raw completion rate without concurrency as 'window_mean' (trailing window), 'window_max' (peak interval in the window) or 'lifetime_max' (peak interval ever, never decays).
• AutoscalingMetricConfig: the user-facing knobs. Plain data.
• scrape_prometheus_metrics / parse_prometheus_metric: read Prometheus text exposition.

Pairs with the ModelHandler.set_model_identifier() hook (#40906), which supplies the model_identifier the publisher takes. No dependency between the two PRs.

Testing

autoscaling_metrics_test.py (100% line and branch coverage of the module):

• Prometheus scraping and parsing (scrape failure, summed series, exact names, timestamps, non-numeric values, escaped label values).
• Estimator math, running_percentile, counter reset, clock skew, negative gauge, warmup, idle ticks, window_secs=0, and each aggregation mode.
• Publisher: publishes values rounded, skips None and errors, resets the cell, supersession (a stopped publisher can't take its metric back), metric-tag precedence, config defaults.
• Publisher end to end against a fake server: server down, unreachable or incomplete scrapes, label filter, estimate, load_signal_fn override.

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

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

Additional details and impacted files

@@            Coverage Diff             @@
##           master   #40907      +/-   ##
==========================================
+ Coverage   79.83%   80.15%   +0.31%     
==========================================
  Files         494      495       +1     
  Lines       83575    83854     +279     
==========================================
+ Hits        66722    67213     +491     
+ Misses      16853    16641     -212     

Flag	Coverage Δ
python	80.15% <100.00%> (+0.31%)	⬆️

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.

[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

[damccorm]

What happens when total_completed == 0 while running > 0 (requests are in flight, but none have completed in the window yet) and now - self._first_active_ts >= self._warmup_timeout_secs?

Right now, update() falls through, computes mu = 0.0 / slot_secs = 0.0 (or 0.0 in window_mean / window_max), and returns max(float(self._min_value), 0.0) = 1.0. This causes the estimator to report the floor value 1 in two cases:

1. Slow initial LLM requests: If the first batch of LLM requests takes > 10s (DEFAULT_WARMUP_TIMEOUT_SECS = 10.0), at t = 10s (completed = 0, running > 0) the worker publishes 1 before a single request has finished.
2. First tick after an idle period >= window_secs: Because self._first_active_ts is set once for the lifetime of the estimator, if a worker is idle for >= window_secs (during which update() returns None and preserves the last published capacity) and then receives a new request (completed = 0, running > 0 on the first poll), all prior samples are pruned by cutoff = now - self._window_secs, now - self._first_active_ts >= warmup_timeout_secs is already true, and update() immediately returns 1.0—overwriting the worker's previously published capacity with 1.

Should we return None when total_completed <= 0 (except in LIFETIME_MAX mode when self._peak_rate > 0), and add a unit test covering this?

[AMOOOMA]

Good catch, fixed as suggested in e2ff8e3. With nothing completed in the window, update() returns None (keep the previous value) unless it's lifetime_max with a peak. I added a test for slow first requests and for the first poll after an idle period.

[damccorm]

In sdks/python/apache_beam/metrics/cells.pxd, MetricCell and DistributionCell are Cython cdef class extension types where _lock is a private cdef object _lock attribute and data is cdef readonly DistributionData data:

cdef class MetricCell(object):
  cdef object _lock
  ...
cdef class DistributionCell(MetricCell):
  cdef readonly DistributionData data

While unit tests run against pure-Python source, released Beam wheels and Dataflow worker containers compile cells.py with Cython—where accessing cell._lock from Python raises AttributeError: 'apache_beam.metrics.cells.DistributionCell' object has no attribute '_lock' and assigning cell.data = ... raises AttributeError: attribute 'data' of 'apache_beam.metrics.cells.DistributionCell' objects is not writable.

Could we either add a helper method on DistributionCell in cells.py (where self._lock and self.data are accessible inside the Cython class) or update cells.pxd so this works in Cython-compiled wheels?

[AMOOOMA]

Thanks, CI hit exactly this. I went with the helper (e2ff8e3): DistributionCell.reset_to(value) in cells.py swaps in a single-sample DistributionData under the cell's lock, so readers never see an empty distribution. _publish just calls it now. It's covered in cells_test.py, and I verified it with a Cython-compiled cells module as well as pure Python.