From 6fd3db54913804990716e083ea7828dc0b437cbc Mon Sep 17 00:00:00 2001 From: AMOOOMA Date: Fri, 30 Oct 2026 16:37:46 +0000 Subject: [PATCH 01/14] [RunInference] Add a capacity-based autoscaling signal for model-server handlers Dataflow's GPU autoscaling reads num_loaded_models_ in the BeamML_ModelManager namespace as the number of messages each worker can hold. A model handler that serves inference through an out-of-process server (vLLM, Triton, ...) always has exactly one loaded copy per worker, so it needs to publish a capacity estimate under that metric instead. Add a generic library for that in apache_beam.ml.inference.autoscaling_metrics: - NumLoadedModelsPublisher publishes a capacity signal on a background thread, resetting the process-wide distribution so the latest value (not a lifetime mean) is reported, with optional aggregation policies on top. - PrometheusCounterGaugeSampler reads a server's completed-requests counter and in-flight gauge from its Prometheus /metrics endpoint; the metric names are the only server-specific input. - CompletionRateCapacityEstimator turns those samples into a per-worker capacity estimate: per-slot throughput times p90 concurrency over a trailing window. It does no I/O. - AutoscalingMetricConfig holds the user-facing knobs, resolves the metric tag (explicit tag, then the RunInference model identifier verbatim, then a sanitized model name), and builds the estimator and publisher. --- .../ml/inference/autoscaling_metrics.py | 632 ++++++++++++++++++ .../ml/inference/autoscaling_metrics_test.py | 264 ++++++++ 2 files changed, 896 insertions(+) create mode 100644 sdks/python/apache_beam/ml/inference/autoscaling_metrics.py create mode 100644 sdks/python/apache_beam/ml/inference/autoscaling_metrics_test.py diff --git a/sdks/python/apache_beam/ml/inference/autoscaling_metrics.py b/sdks/python/apache_beam/ml/inference/autoscaling_metrics.py new file mode 100644 index 000000000000..589aa697c834 --- /dev/null +++ b/sdks/python/apache_beam/ml/inference/autoscaling_metrics.py @@ -0,0 +1,632 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +"""Capacity-based autoscaling signal for model handlers that run a model server. + +Generic across model servers (vLLM, Triton, ...): a model handler that serves +inference through an out-of-process server uses these pieces to tell the runner +how much work each worker can absorb. Nothing here depends on a particular +server. + +Dataflow's GPU autoscaling reads the process-wide distribution +``num_loaded_models_`` in the ``BeamML_ModelManager`` namespace (the same +metric ``model_manager.py`` publishes) as *the number of messages each worker +can hold*, and sizes the pool roughly as ``ceil(backlog / reported)``. A model +server always has exactly one loaded copy per worker, so the handler instead +publishes that metric with an estimate of the worker's message capacity. + +Building blocks, from the bottom up: + +* :func:`scrape_prometheus_metrics` / :func:`parse_prometheus_metric` -- fetch + and read Prometheus text exposition. +* :class:`PrometheusCounterGaugeSampler` -- reads a server's completed-requests + counter and in-flight gauge (the only server-specific inputs are the two + metric names). +* :class:`CompletionRateCapacityEstimator` -- turns those samples into a + per-worker capacity estimate: per-slot throughput x typical concurrency over a + trailing window. Pure computation; no I/O. +* :class:`NumLoadedModelsPublisher` -- polls a capacity signal on a background + thread and publishes it under ``num_loaded_models_``. +* :class:`AutoscalingMetricConfig` -- the user-facing knobs, plus the tag + resolution and factories a handler uses to build the pieces above. +""" + +# pytype: skip-file + +import dataclasses +import logging +import threading +import time +import urllib.request +from collections.abc import Callable +from collections.abc import Mapping +from dataclasses import dataclass +from typing import Optional + +from apache_beam.metrics.execution import MetricsEnvironment +from apache_beam.metrics.metricbase import MetricName + +__all__ = [ + 'AutoscalingMetricConfig', + 'CompletionRateCapacityEstimator', + 'NumLoadedModelsPublisher', + 'PrometheusCounterGaugeSampler', + 'parse_prometheus_metric', + 'sanitize_metric_tag', + 'scrape_prometheus_metrics', +] + +MODEL_MANAGER_METRICS_NAMESPACE = 'BeamML_ModelManager' +NUM_LOADED_MODELS_METRIC_PREFIX = 'num_loaded_models_' + +DEFAULT_POLL_INTERVAL_SECS = 5.0 +# Dimensionless multiplier applied to the capacity estimate before it is +# published. The runner sizes the pool roughly as ``ceil(backlog / reported)``, +# so a larger value lets each worker hold more backlog (fewer workers) and a +# smaller one scales out more aggressively. +DEFAULT_MULTIPLIER = 1.0 +# Trailing window and concurrency percentile used by +# CompletionRateCapacityEstimator. +DEFAULT_SIGNAL_WINDOW_SECS = 1800.0 # 30 minutes. +DEFAULT_RUNNING_PERCENTILE = 90.0 +DEFAULT_MIN_REQUESTS = 3 +DEFAULT_WARMUP_TIMEOUT_SECS = 10.0 +# The capacity estimate is already smoothed, so the publisher forwards its latest +# value unchanged by default (no second smoothing layer). 'lifetime_max' etc. +# remain available as an explicit policy on top. +DEFAULT_AGGREGATION = 'last' +DEFAULT_AGGREGATION_WINDOW_SECS = 1800.0 # 30 minutes. + + +def sanitize_metric_tag(name: str) -> str: + """Turns a model name into a metric-tag-safe token for the autoscaler. + + Args: + name: The model name (may contain ``/``, ``:`` etc.). + + Returns: + ``name`` with characters outside ``[A-Za-z0-9_]`` replaced by ``_``. + """ + return ''.join(c if (c.isalnum() or c == '_') else '_' for c in name) + + +def scrape_prometheus_metrics( + port: int, + timeout_secs: float = 5.0, + path: str = '/metrics') -> Optional[str]: + """Fetches ``http://localhost:``; None on any failure.""" + url = f'http://localhost:{port}{path}' + try: + with urllib.request.urlopen(url, timeout=timeout_secs) as response: + return response.read().decode('utf-8', errors='replace') + except Exception: # pylint: disable=broad-except + return None + + +def _parse_prometheus_labels(label_str: str) -> dict[str, str]: + """Parses the ``{k="v",...}`` label section of a Prometheus sample line.""" + labels: dict[str, str] = {} + i = 0 + n = len(label_str) + while i < n: + eq = label_str.find('=', i) + if eq == -1: + break + name = label_str[i:eq].strip() + j = eq + 1 + if j < n and label_str[j] == '"': + j += 1 + value_chars = [] + while j < n and label_str[j] != '"': + if label_str[j] == '\\' and j + 1 < n: + j += 1 + value_chars.append(label_str[j]) + j += 1 + value = ''.join(value_chars) + j += 1 # skip closing quote + else: + comma = label_str.find(',', j) + comma = n if comma == -1 else comma + value = label_str[j:comma].strip() + j = comma + labels[name] = value + comma = label_str.find(',', j) + if comma == -1: + break + i = comma + 1 + return labels + + +def parse_prometheus_metric( + metrics_text: str, + metric_name: str, + label_filter: Optional[Mapping[str, str]] = None, +) -> Optional[float]: + """Sums the values of a metric from Prometheus text-format exposition. + + Args: + metrics_text: The full body returned by the server's ``/metrics`` endpoint. + metric_name: The metric family name to match (e.g. + ``vllm:num_requests_running``). + label_filter: If provided, only samples whose labels contain all of these + key/value pairs are summed. + + Returns: + The sum of the matching sample values, or ``None`` if the metric is absent. + """ + total: Optional[float] = None + for raw_line in metrics_text.splitlines(): + line = raw_line.strip() + if not line or line.startswith('#'): + continue + if not line.startswith(metric_name): + continue + remainder = line[len(metric_name):] + if remainder and remainder[0] not in (' ', '\t', '{'): + continue + labels: Mapping[str, str] = {} + if remainder.startswith('{'): + close = remainder.find('}') + if close == -1: + continue + labels = _parse_prometheus_labels(remainder[1:close]) + remainder = remainder[close + 1:] + parts = remainder.split() + if not parts: + continue + try: + value = float(parts[0]) + except ValueError: + continue + if label_filter and any(labels.get(k) != v + for k, v in label_filter.items()): + continue + total = value if total is None else total + value + return total + + +def _nearest_rank_percentile(values: list[float], q: float) -> float: + """Nearest-rank percentile (``q`` in [0, 100]) of a non-empty list.""" + ordered = sorted(values) + q = min(max(q, 0.0), 100.0) + rank = int((q * len(ordered) + 99) // 100) # ceil(q / 100 * n). + idx = min(max(rank - 1, 0), len(ordered) - 1) + return ordered[idx] + + +class NumLoadedModelsPublisher: + """Publishes the ``num_loaded_models_`` autoscaling metric on a timer. + + ``signal_fn`` produces the raw load sample (e.g. requests running in + parallel, or an estimated per-worker capacity); it is polled on a background + daemon thread every ``poll_interval_secs``. Returning ``None`` (e.g. a + transient scrape failure) skips that tick, leaving the running state and last + published value unchanged. + + The published value is an aggregate of the observed samples, chosen so the + pool does not scale *down* the instant an input burst clears (which would + thrash and then be unable to absorb the next burst): + + * ``'lifetime_max'`` -- the maximum sample observed over the whole + pipeline lifetime (monotonic; never decays). Sizes the + fleet to the worst load ever seen and holds it. + * ``'max'`` -- peak sample in the trailing ``aggregation_window_secs`` + window; decays that long after load drops. + * ``'p'`` -- a windowed percentile, e.g. ``'p90'``. Smoother than max: + ignores brief spikes and decays as low samples fill it. + * ``'mean'`` -- the average over the window. + * ``'last'`` -- the raw last sample (no smoothing). Use this when + ``signal_fn`` already smooths its output. + """ + # The publisher currently owning each metric name in this process. A model + # update or server restart creates a new server (and publisher) for the same + # identifier; starting it supersedes the stale one so exactly one thread + # publishes per metric name. + _active: dict[str, 'NumLoadedModelsPublisher'] = {} + _active_lock = threading.Lock() + + def __init__( + self, + tag: str, + signal_fn: Callable[[], Optional[float]], + poll_interval_secs: float = DEFAULT_POLL_INTERVAL_SECS, + aggregation: str = DEFAULT_AGGREGATION, + aggregation_window_secs: float = DEFAULT_AGGREGATION_WINDOW_SECS, + multiplier: float = DEFAULT_MULTIPLIER, + ): + if not tag: + raise ValueError('num_loaded_models tag must be a non-empty string.') + self._metric_name = NUM_LOADED_MODELS_METRIC_PREFIX + tag + self._signal_fn = signal_fn + self._poll_interval_secs = poll_interval_secs + self._aggregation = (aggregation or DEFAULT_AGGREGATION).lower() + self._aggregation_window_secs = max(float(aggregation_window_secs), 0.0) + self._multiplier = float(multiplier) + # Trailing (timestamp, value) samples for the windowed aggregation modes. + self._samples: list[tuple[float, float]] = [] + # Monotonic peak for the 'lifetime_max' mode (no sample retention). + self._lifetime_max: Optional[float] = None + self._stop_event = threading.Event() + self._thread: Optional[threading.Thread] = None + + def _publish(self, value: float) -> None: + ivalue = int(round(value * self._multiplier)) + # Process-wide DistributionCells are read without resetting by the Python + # SDK harness (_request_harness_monitoring_infos), and Dataflow's worker + # harness reports their cumulative mean (sum / count). Resetting before + # updating keeps (sum=ivalue, count=1) so the latest capacity is reported + # instead of an unweighted lifetime average since worker boot. + cell = MetricsEnvironment.process_wide_container().get_distribution( + MetricName(MODEL_MANAGER_METRICS_NAMESPACE, self._metric_name)) + cumulative = cell.get_cumulative() + if cumulative.count != 1 or cumulative.sum != ivalue: + cell.reset() + cell.update(ivalue) + + def _aggregate(self, now: float) -> Optional[float]: + """Aggregates in-window samples per ``self._aggregation``. + + Evicts samples older than ``aggregation_window_secs`` first, then reduces + the remainder. + + Args: + now: Current timestamp used for window eviction. + + Returns: + The aggregated value, or ``None`` if the window is empty. + """ + if self._aggregation_window_secs > 0: + cutoff = now - self._aggregation_window_secs + self._samples = [(t, v) for (t, v) in self._samples if t >= cutoff] + values = [v for (_, v) in self._samples] + if not values: + return None + mode = self._aggregation + if mode == 'last': + return values[-1] + if mode == 'mean': + return sum(values) / len(values) + if mode.startswith('p'): + try: + q = float(int(mode[1:])) + except ValueError: + return max(values) + return _nearest_rank_percentile(values, q) + return max(values) # 'max' and unknown modes. + + def _run(self) -> None: + while not self._stop_event.is_set(): + try: + value = self._signal_fn() + if value is not None: + if self._aggregation == 'lifetime_max': + # Monotonic running max over the whole pipeline lifetime; no + # per-sample retention and never decays. + self._lifetime_max = ( + float(value) if self._lifetime_max is None else max( + self._lifetime_max, float(value))) + aggregate = self._lifetime_max + else: + now = time.time() + self._samples.append((now, float(value))) + aggregate = self._aggregate(now) + if aggregate is not None: + self._publish(aggregate) + except Exception: # pylint: disable=broad-except + logging.exception( + 'Failed to publish autoscaling metric %s; will retry.', + self._metric_name, + ) + self._stop_event.wait(self._poll_interval_secs) + + def start(self) -> None: + """Starts the background polling thread. + + Deliberately publishes no initial 0: reporting 0 would tell the autoscaler + this replica has no capacity and could trigger a spurious downscale before + the first real sample lands. The metric stays absent until the first + aggregate is published. + """ + if self._thread is not None: + return + with NumLoadedModelsPublisher._active_lock: + previous = NumLoadedModelsPublisher._active.get(self._metric_name) + NumLoadedModelsPublisher._active[self._metric_name] = self + if previous is not None and previous is not self: + logging.info( + 'Replacing the active publisher for %s (model reload or server ' + 'restart).', + self._metric_name) + previous.stop() + self._thread = threading.Thread( + target=self._run, name='num-loaded-models-publisher', daemon=True) + self._thread.start() + + def stop(self) -> None: + """Stops polling. + + Deliberately publishes no final 0. When a worker is drained during a + downscale, reporting 0 shrinks the aggregate capacity and can drive a + feedback loop of further downscales that tears down healthy replicas. The + drained worker's contribution ages out via metric staleness instead. + """ + self._stop_event.set() + with NumLoadedModelsPublisher._active_lock: + if NumLoadedModelsPublisher._active.get(self._metric_name) is self: + del NumLoadedModelsPublisher._active[self._metric_name] + thread = self._thread + if thread is not None: + thread.join(timeout=self._poll_interval_secs + 5) + self._thread = None + + +class PrometheusCounterGaugeSampler: + """Reads a completed-requests counter and an in-flight gauge from a server. + + Calling the sampler with the server's metrics port scrapes + ``http://localhost:/metrics`` and returns ``(completed_total, + running)``: the value of the monotonically increasing ``completed_metric`` + counter and of the ``running_metric`` gauge, each summed over the samples that + match ``label_filter``. It returns ``None`` when the endpoint is unreachable + or either metric is missing. + + The metric names are the only server-specific inputs, e.g. + ``('vllm:request_success_total', 'vllm:num_requests_running')`` for vLLM. + """ + def __init__( + self, + completed_metric: str, + running_metric: str, + label_filter: Optional[Mapping[str, str]] = None, + timeout_secs: float = 5.0, + ): + self._completed_metric = completed_metric + self._running_metric = running_metric + self._label_filter = dict(label_filter) if label_filter else None + self._timeout_secs = timeout_secs + + def __call__(self, port: int) -> Optional[tuple[float, float]]: + text = scrape_prometheus_metrics(port, timeout_secs=self._timeout_secs) + if text is None: + return None + completed = parse_prometheus_metric( + text, self._completed_metric, label_filter=self._label_filter) + running = parse_prometheus_metric( + text, self._running_metric, label_filter=self._label_filter) + if completed is None or running is None: + return None + return completed, running + + +class CompletionRateCapacityEstimator: + """Estimates a worker's message capacity from completion and concurrency. + + Feed it periodic samples of a model server's monotonically increasing + *completed requests* counter and its *in-flight requests* gauge (for example + from :class:`PrometheusCounterGaugeSampler`) via :meth:`update`. Over a + trailing ``window_secs`` window it derives: + + * ``mu`` -- per-slot throughput, + ``sum(delta_completed) / sum(running * dt)`` + (a ratio of windowed sums, so discrete completions don't make it bursty). + This approximates ``1 / avg_latency`` and is roughly independent of the + offered load. + * ``running_ref`` -- the ``running_percentile`` (default p90) of observed + in-flight concurrency. p90 rather than max avoids latching onto rare, + heavily backlogged spikes that overstate normal-load capacity. + + and estimates a per-worker message capacity:: + + capacity = max(min_value, mu * running_ref * multiplier) + + i.e. the worker's sustained throughput scaled by a dimensionless tuning + ``multiplier``. Normalizing by ``running`` makes the estimate + load-independent: right after an autoscaling event, when the runner + rebalances work and the worker is briefly starved, the raw completion rate + collapses but ``mu`` holds and ``running_ref`` still reflects the concurrency + the worker normally serves, so the estimate stays steady instead of cratering + and triggering a downscale spiral. + + The estimate is already windowed; feed it to ``NumLoadedModelsPublisher`` + with ``aggregation='last'`` unless you deliberately want a second policy + layer (e.g. ``'lifetime_max'`` to hold the peak estimate). + + This class does no I/O, so it is independent of how the samples are read. + """ + def __init__( + self, + multiplier: float = DEFAULT_MULTIPLIER, + min_value: int = 1, + window_secs: float = DEFAULT_SIGNAL_WINDOW_SECS, + running_percentile: float = DEFAULT_RUNNING_PERCENTILE, + min_requests: int = DEFAULT_MIN_REQUESTS, + warmup_timeout_secs: float = DEFAULT_WARMUP_TIMEOUT_SECS, + ): + self._multiplier = max(float(multiplier), 0.0) + self._min_value = int(min_value) + self._window_secs = max(float(window_secs), 0.0) + self._running_percentile = min(max(float(running_percentile), 0.0), 100.0) + self._min_requests = max(int(min_requests), 0) + self._warmup_timeout_secs = max(float(warmup_timeout_secs), 0.0) + self._last_count: Optional[float] = None + self._last_ts: Optional[float] = None + self._first_active_ts: Optional[float] = None + # Trailing (timestamp, delta_completed, running, dt) samples. + self._samples: list[tuple[float, float, float, float]] = [] + + def update( + self, + completed_total: float, + running: float, + now: Optional[float] = None) -> Optional[float]: + """Records one sample and returns the current estimate. + + Args: + completed_total: Current value of the completed-requests counter. + running: Current number of in-flight requests. + now: Sample timestamp in seconds; defaults to ``time.time()``. + + Returns: + The capacity estimate, or ``None`` when there is not enough data yet + (first sample, counter reset, idle window, or warmup), in which case the + caller should keep its previous value. + """ + if now is None: + now = time.time() + prev_count, prev_ts = self._last_count, self._last_ts + self._last_count, self._last_ts = completed_total, now + if prev_count is None or prev_ts is None: + # First sample only establishes a baseline; a rate needs two samples. + return None + dt = now - prev_ts + delta = completed_total - prev_count + if dt <= 0 or delta < 0: + # Clock skew or counter reset (server restart); skip this tick. + return None + if running < 0: + return None + # Evict samples older than the window, and skip completely idle intervals so + # trough/idle periods do not dilute running_ref (p90 concurrency) to 0. + if self._window_secs > 0: + cutoff = now - self._window_secs + self._samples = [s for s in self._samples if s[0] >= cutoff] + if delta == 0 and running == 0: + return None + self._samples.append((now, float(delta), float(running), float(dt))) + # Per-slot throughput as a ratio of windowed sums (stable vs. per-tick): + # total completions divided by total slot-seconds of in-flight work. + total_completed = sum(d for (_, d, _, _) in self._samples) + total_slot_secs = sum(r * dts for (_, _, r, dts) in self._samples) + if total_slot_secs <= 0: + # Window is entirely idle (no in-flight requests); nothing to report. + return None + if self._first_active_ts is None: + self._first_active_ts = prev_ts + if (total_completed < self._min_requests and + (now - self._first_active_ts) < self._warmup_timeout_secs): + # Hold off publishing until a newly started worker completes at least + # min_requests or warmup_timeout_secs elapses, so an uncalibrated cold + # worker does not briefly overwrite the job-wide metric with 1. + return None + mu = total_completed / total_slot_secs + running_ref = _nearest_rank_percentile( + [r for (_, _, r, _) in self._samples if r > 0], + self._running_percentile) + estimate = mu * running_ref * self._multiplier + return max(float(self._min_value), estimate) + + +@dataclass(frozen=True) +class AutoscalingMetricConfig: + """User-facing knobs for the ``num_loaded_models_`` autoscaling signal. + + Model handlers that publish this signal take an instance (typically as an + ``autoscaling_config`` argument) to override the defaults; leave it unset for + the defaults. Each field: + + * ``tag`` -- metric tag to publish under. Leave unset (the default) so the + handler fills it in via :meth:`resolve_tag` from the identifier RunInference + reports for the step (``RunInference(model_identifier=...)``, or the random + id it generates), which is what Dataflow's autoscaler looks up and which + stays stable across model updates. Only set it explicitly if you also pass + the same value as ``model_identifier``. + * ``poll_interval_secs`` -- how often to sample and publish (default 5). + * ``multiplier`` -- dimensionless factor applied to the capacity estimate + before publishing; larger values mean fewer workers, smaller values scale + out more aggressively (default 1.0). + * ``signal_window_secs`` -- trailing window over which + :class:`CompletionRateCapacityEstimator` estimates throughput and p90 + concurrency (default 1800). Ignored when ``load_signal_fn`` is set. + * ``min_requests`` -- minimum completed requests before a newly started + worker begins publishing, unless ``warmup_timeout_secs`` elapses first + (default 3). + * ``warmup_timeout_secs`` -- maximum time (in seconds) since a worker first + sees active requests to wait for ``min_requests`` before publishing anyway + (default 10). + * ``aggregation`` -- optional policy the publisher applies on top of the + (already smoothed) signal: ``'last'`` (default; forward the latest + estimate), ``'lifetime_max'`` (hold the peak estimate for the whole run, + never decays), or ``'max'``/``'p90'``/``'p95'``/``'mean'`` over + ``aggregation_window_secs``. Only change this from ``'last'`` when you + deliberately want a second smoothing layer, e.g. to stop the pool scaling + down the instant a burst clears. + * ``aggregation_window_secs`` -- trailing window for the windowed + ``aggregation`` modes; unused by ``'last'`` and ``'lifetime_max'`` + (default 1800). + * ``load_signal_fn`` -- optional override that, given the server's metrics + port, returns the per-worker capacity value directly, bypassing the + built-in completion-rate estimate. + + A handler calls :meth:`resolve_tag`, builds its capacity signal from a + server-specific sampler and :meth:`make_capacity_estimator` (or uses + ``load_signal_fn``), and passes that signal to :meth:`make_publisher`. + """ + tag: Optional[str] = None + poll_interval_secs: float = DEFAULT_POLL_INTERVAL_SECS + multiplier: float = DEFAULT_MULTIPLIER + signal_window_secs: float = DEFAULT_SIGNAL_WINDOW_SECS + min_requests: int = DEFAULT_MIN_REQUESTS + warmup_timeout_secs: float = DEFAULT_WARMUP_TIMEOUT_SECS + aggregation: str = DEFAULT_AGGREGATION + aggregation_window_secs: float = DEFAULT_AGGREGATION_WINDOW_SECS + load_signal_fn: Optional[Callable[[int], Optional[float]]] = None + + def resolve_tag( + self, + model_identifier: Optional[str], + fallback_name: str = '') -> 'AutoscalingMetricConfig': + """Returns a copy whose ``tag`` is set. + + Precedence: an explicit ``tag``; else ``model_identifier``, used verbatim + because it must match the identifier RunInference annotates the step with + (see ``ModelHandler.set_model_identifier``); else ``fallback_name`` + sanitized with :func:`sanitize_metric_tag`, for handlers used outside + RunInference. + + Raises: + ValueError: if no tag can be derived. + """ + if self.tag: + return self + tag = model_identifier or sanitize_metric_tag(fallback_name or '') + if not tag: + raise ValueError( + 'Cannot derive a num_loaded_models tag: pass a model_identifier, a ' + 'fallback_name, or set AutoscalingMetricConfig.tag.') + return dataclasses.replace(self, tag=tag) + + def make_capacity_estimator(self) -> CompletionRateCapacityEstimator: + """Builds a :class:`CompletionRateCapacityEstimator` from these knobs.""" + return CompletionRateCapacityEstimator( + multiplier=self.multiplier, + window_secs=self.signal_window_secs, + min_requests=self.min_requests, + warmup_timeout_secs=self.warmup_timeout_secs) + + def make_publisher( + self, signal_fn: Callable[[], + Optional[float]]) -> NumLoadedModelsPublisher: + if not self.tag: + raise ValueError( + 'AutoscalingMetricConfig.tag is unset; call resolve_tag() first.') + return NumLoadedModelsPublisher( + tag=self.tag, + signal_fn=signal_fn, + poll_interval_secs=self.poll_interval_secs, + aggregation=self.aggregation, + aggregation_window_secs=self.aggregation_window_secs, + ) diff --git a/sdks/python/apache_beam/ml/inference/autoscaling_metrics_test.py b/sdks/python/apache_beam/ml/inference/autoscaling_metrics_test.py new file mode 100644 index 000000000000..c7c4f4887073 --- /dev/null +++ b/sdks/python/apache_beam/ml/inference/autoscaling_metrics_test.py @@ -0,0 +1,264 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +import io +import unittest +from unittest import mock + +from apache_beam.ml.inference import autoscaling_metrics + + +class AutoscalingMetricsTest(unittest.TestCase): + def test_sanitize_metric_tag(self): + self.assertEqual( + autoscaling_metrics.sanitize_metric_tag('Qwen/Qwen2.5-7B:v1'), + 'Qwen_Qwen2_5_7B_v1') + self.assertEqual( + autoscaling_metrics.sanitize_metric_tag('normal_tag_1'), 'normal_tag_1') + + def test_parse_prometheus_metric(self): + prom_text = ( + '# HELP vllm:num_requests_running Running requests\n' + '# TYPE vllm:num_requests_running gauge\n' + 'vllm:num_requests_running{model_name="qwen"} 4.0\n' + 'vllm:num_requests_running{model_name="other"} 6.0\n') + self.assertEqual( + autoscaling_metrics.parse_prometheus_metric( + prom_text, + 'vllm:num_requests_running', + label_filter={'model_name': 'qwen'}), + 4.0) + self.assertEqual( + autoscaling_metrics.parse_prometheus_metric( + prom_text, 'vllm:num_requests_running'), + 10.0) + self.assertIsNone( + autoscaling_metrics.parse_prometheus_metric(prom_text, 'vllm:missing')) + + def test_num_loaded_models_publisher_aggregation(self): + samples = [(100.0 + i * 10, float(i + 1)) for i in range(10)] + pub = autoscaling_metrics.NumLoadedModelsPublisher( + tag='test', + signal_fn=lambda: 1.0, + aggregation='p90', + aggregation_window_secs=1800.0) + pub._samples = list(samples) + self.assertEqual(pub._aggregate(now=200.0), 9.0) + + pub._samples = [(100.0, 2.0), (1900.0, 8.0), (1950.0, 10.0)] + self.assertEqual(pub._aggregate(now=2000.0), 10.0) + + pub._samples = [(100.0, 2.0), (102.0, 5.0), (104.0, 3.0)] + pub._aggregation_window_secs = 60.0 + pub._aggregation = 'max' + self.assertEqual(pub._aggregate(now=105.0), 5.0) + pub._aggregation = 'last' + self.assertEqual(pub._aggregate(now=105.0), 3.0) + pub._aggregation = 'mean' + self.assertAlmostEqual(pub._aggregate(now=105.0), 3.3333333333333335) + + def test_publish_resets_process_wide_distribution_cell(self): + pub = autoscaling_metrics.NumLoadedModelsPublisher( + tag='reset_test', signal_fn=lambda: 1.0, multiplier=1.0) + cell = ( + autoscaling_metrics.MetricsEnvironment.process_wide_container(). + get_distribution( + autoscaling_metrics.MetricName( + autoscaling_metrics.MODEL_MANAGER_METRICS_NAMESPACE, + 'num_loaded_models_reset_test'))) + cell.reset() + self.addCleanup(cell.reset) + for _ in range(10): + pub._publish(2.0) + self.assertEqual(cell.get_cumulative().count, 1) + self.assertEqual(cell.get_cumulative().sum, 2) + pub._publish(5.0) + self.assertEqual(cell.get_cumulative().count, 1) + self.assertEqual(cell.get_cumulative().sum, 5) + + def test_publisher_for_same_tag_supersedes_previous(self): + """A reloaded model's publisher takes over its tag from the stale one.""" + active = autoscaling_metrics.NumLoadedModelsPublisher._active + first = autoscaling_metrics.NumLoadedModelsPublisher( + tag='reload', signal_fn=lambda: None, poll_interval_secs=0.01) + second = autoscaling_metrics.NumLoadedModelsPublisher( + tag='reload', signal_fn=lambda: None, poll_interval_secs=0.01) + other = autoscaling_metrics.NumLoadedModelsPublisher( + tag='other', signal_fn=lambda: None, poll_interval_secs=0.01) + try: + first.start() + other.start() + self.assertTrue(first._thread.is_alive()) + second.start() + # first is stopped and unregistered; second owns the tag; other is + # untouched. + self.assertTrue(first._stop_event.is_set()) + self.assertIsNone(first._thread) + self.assertTrue(second._thread.is_alive()) + self.assertTrue(other._thread.is_alive()) + self.assertIs(active['num_loaded_models_reload'], second) + # Stopping the stale one again must not unregister the new owner. + first.stop() + self.assertIs(active['num_loaded_models_reload'], second) + finally: + first.stop() + second.stop() + other.stop() + self.assertNotIn('num_loaded_models_reload', active) + + def test_prometheus_counter_gauge_sampler(self): + sampler = autoscaling_metrics.PrometheusCounterGaugeSampler( + completed_metric='vllm:request_success_total', + running_metric='vllm:num_requests_running') + responses = [ + 'vllm:request_success_total 10.0\nvllm:num_requests_running 4.0\n', + # Missing the gauge -> None. + 'vllm:request_success_total 11.0\n', + ] + + def fake_urlopen(url, timeout=5.0): + self.assertEqual(url, 'http://localhost:8000/metrics') + del timeout + return io.BytesIO(responses.pop(0).encode('utf-8')) + + with mock.patch('urllib.request.urlopen', side_effect=fake_urlopen): + self.assertEqual(sampler(8000), (10.0, 4.0)) + self.assertIsNone(sampler(8000)) + # Unreachable endpoint -> None. + with mock.patch('urllib.request.urlopen', side_effect=OSError('refused')): + self.assertIsNone(sampler(8000)) + + def test_prometheus_counter_gauge_sampler_label_filter(self): + sampler = autoscaling_metrics.PrometheusCounterGaugeSampler( + completed_metric='nv_inference_request_success', + running_metric='nv_inference_pending_request_count', + label_filter={'model': 'a'}) + body = ( + 'nv_inference_request_success{model="a"} 10\n' + 'nv_inference_request_success{model="b"} 99\n' + 'nv_inference_pending_request_count{model="a"} 2\n' + 'nv_inference_pending_request_count{model="b"} 50\n') + with mock.patch('urllib.request.urlopen', + return_value=io.BytesIO(body.encode('utf-8'))): + # Only model "a" counts. + self.assertEqual(sampler(8002), (10.0, 2.0)) + + def test_capacity_estimator(self): + estimator = autoscaling_metrics.CompletionRateCapacityEstimator( + multiplier=2.0, min_value=1, window_secs=60.0, running_percentile=90.0) + # First sample establishes baseline -> None. + self.assertIsNone(estimator.update(10.0, 4.0, now=100.0)) + # delta=20 over dt=5s with running=4 -> mu = 20 / (4 * 5) = 1.0, + # running_ref = 4.0, multiplier = 2.0 -> 8.0. + self.assertAlmostEqual(estimator.update(30.0, 4.0, now=105.0), 8.0) + + def test_capacity_estimator_skips_counter_reset_and_clock_skew(self): + estimator = autoscaling_metrics.CompletionRateCapacityEstimator( + window_secs=60.0, min_requests=0) + self.assertIsNone(estimator.update(50.0, 2.0, now=100.0)) + # Counter went down (server restart) -> skipped, new baseline. + self.assertIsNone(estimator.update(5.0, 2.0, now=105.0)) + # Non-increasing timestamp -> skipped. + self.assertIsNone(estimator.update(9.0, 2.0, now=105.0)) + # delta=4 over 5s with running=2 -> mu=0.4, p90=2 -> 0.8 -> min_value 1. + self.assertAlmostEqual(estimator.update(13.0, 2.0, now=110.0), 1.0) + + def test_capacity_estimator_warmup_min_requests_and_timeout(self): + # Requires >= 3 completed requests before the 10s warmup timeout. + estimator = autoscaling_metrics.CompletionRateCapacityEstimator( + min_requests=3, warmup_timeout_secs=10.0) + self.assertIsNone(estimator.update(0.0, 2.0, now=100.0)) # Baseline. + # 1 completion (< 3) at elapsed=5s (< 10s) -> None. + self.assertIsNone(estimator.update(1.0, 2.0, now=105.0)) + # 4 completions (>= 3) at elapsed=8s -> publishes + # (4 / (2*5 + 2*3)) * 2 = 0.5 -> clamped to min_value=1.0. + self.assertAlmostEqual(estimator.update(4.0, 2.0, now=108.0), 1.0) + + slow = autoscaling_metrics.CompletionRateCapacityEstimator( + min_requests=3, warmup_timeout_secs=10.0) + self.assertIsNone(slow.update(0.0, 1.0, now=100.0)) # Baseline. + # 1 completion (< 3) at elapsed=5s -> None. + self.assertIsNone(slow.update(1.0, 1.0, now=105.0)) + # 2 completions (< 3) but elapsed=10s (>= timeout) -> publishes. + self.assertAlmostEqual(slow.update(2.0, 1.0, now=110.0), 1.0) + + def test_capacity_estimator_skips_idle_ticks_without_diluting_p90(self): + estimator = autoscaling_metrics.CompletionRateCapacityEstimator( + multiplier=1.0, + min_value=1, + window_secs=60.0, + running_percentile=90.0, + min_requests=1) + self.assertIsNone(estimator.update(0.0, 5.0, now=100.0)) # Baseline. + # 25 completions over 5s with running=5 -> mu=1.0, p90=5 -> 5.0. + self.assertAlmostEqual(estimator.update(25.0, 5.0, now=105.0), 5.0) + # Three completely idle ticks (delta=0, running=0) -> None, without adding + # r=0 samples that would dilute p90 concurrency. + for t in (110.0, 115.0, 120.0): + self.assertIsNone(estimator.update(25.0, 0.0, now=t)) + # 5 completions that already finished before the scrape (delta=5, + # running=0): total_completed=30, total_slot_secs=25 -> mu=1.2, p90=5 -> 6. + self.assertAlmostEqual(estimator.update(30.0, 0.0, now=125.0), 6.0) + + def test_capacity_estimator_defaults_to_wall_clock(self): + estimator = autoscaling_metrics.CompletionRateCapacityEstimator( + min_requests=0) + with mock.patch.object(autoscaling_metrics.time, + 'time', + side_effect=[100.0, 105.0]): + self.assertIsNone(estimator.update(0.0, 1.0)) + self.assertAlmostEqual(estimator.update(10.0, 1.0), 2.0) + + def test_autoscaling_metric_config_resolve_tag(self): + cfg = autoscaling_metrics.AutoscalingMetricConfig( + poll_interval_secs=7.0, + multiplier=2.0, + signal_window_secs=30.0, + min_requests=5, + warmup_timeout_secs=15.0, + aggregation='lifetime_max', + aggregation_window_secs=40.0) + # The RunInference identifier is used verbatim (it has to match the step + # annotation), even if it contains characters sanitize would change. + self.assertEqual( + cfg.resolve_tag('my-step/id', 'org/model:v1').tag, 'my-step/id') + # Outside RunInference: the sanitized fallback name. + self.assertEqual(cfg.resolve_tag(None, 'org/model:v1').tag, 'org_model_v1') + # An explicit tag always wins. + self.assertEqual( + autoscaling_metrics.AutoscalingMetricConfig(tag='custom').resolve_tag( + 'step-id', 'model').tag, + 'custom') + with self.assertRaises(ValueError): + autoscaling_metrics.AutoscalingMetricConfig().resolve_tag(None, '') + with self.assertRaises(ValueError): + autoscaling_metrics.AutoscalingMetricConfig().make_publisher(lambda: None) + + estimator = cfg.make_capacity_estimator() + self.assertEqual(estimator._multiplier, 2.0) + self.assertEqual(estimator._window_secs, 30.0) + self.assertEqual(estimator._min_requests, 5) + self.assertEqual(estimator._warmup_timeout_secs, 15.0) + publisher = cfg.resolve_tag('step-id').make_publisher(lambda: None) + self.assertEqual(publisher._metric_name, 'num_loaded_models_step-id') + self.assertEqual(publisher._aggregation, 'lifetime_max') + self.assertEqual(publisher._poll_interval_secs, 7.0) + # The multiplier is applied by the estimator, never a second time here. + self.assertEqual(publisher._multiplier, 1.0) + + +if __name__ == '__main__': + unittest.main() From 3610dd19ebf2c9a47c3464cc58e861eaa3b76df7 Mon Sep 17 00:00:00 2001 From: AMOOOMA Date: Mon, 2 Nov 2026 23:57:59 +0000 Subject: [PATCH 02/14] Clean up docstrings and readability in autoscaling_metrics - Docstrings: Google-style Args/Attributes/Raises throughout; module docstring shows how a handler wires the pieces together; the publisher docstring now reflects that 'last' is the default policy; the estimator docstring states which servers it suits (latencies long enough to be seen in flight) and shows the formula in one block. - Parse Prometheus sample and label lines with regexes instead of a hand-rolled index parser. Quoted label values may now contain commas, braces and escapes. - Validate the publisher's aggregation policy up front (unknown policies raise ValueError instead of silently behaving like 'max') and split ingestion into _record(), which also stops 'last' from retaining samples. - Store estimator samples as a NamedTuple instead of positional tuples. Outputs are unchanged for the estimator, every valid aggregation policy, and valid Prometheus exposition (checked against the previous commit on randomized inputs). --- .../ml/inference/autoscaling_metrics.py | 645 +++++++++--------- .../ml/inference/autoscaling_metrics_test.py | 72 +- 2 files changed, 377 insertions(+), 340 deletions(-) diff --git a/sdks/python/apache_beam/ml/inference/autoscaling_metrics.py b/sdks/python/apache_beam/ml/inference/autoscaling_metrics.py index 589aa697c834..1752fa5ff11f 100644 --- a/sdks/python/apache_beam/ml/inference/autoscaling_metrics.py +++ b/sdks/python/apache_beam/ml/inference/autoscaling_metrics.py @@ -17,44 +17,51 @@ """Capacity-based autoscaling signal for model handlers that run a model server. -Generic across model servers (vLLM, Triton, ...): a model handler that serves -inference through an out-of-process server uses these pieces to tell the runner -how much work each worker can absorb. Nothing here depends on a particular -server. - Dataflow's GPU autoscaling reads the process-wide distribution -``num_loaded_models_`` in the ``BeamML_ModelManager`` namespace (the same -metric ``model_manager.py`` publishes) as *the number of messages each worker -can hold*, and sizes the pool roughly as ``ceil(backlog / reported)``. A model -server always has exactly one loaded copy per worker, so the handler instead -publishes that metric with an estimate of the worker's message capacity. - -Building blocks, from the bottom up: - -* :func:`scrape_prometheus_metrics` / :func:`parse_prometheus_metric` -- fetch - and read Prometheus text exposition. -* :class:`PrometheusCounterGaugeSampler` -- reads a server's completed-requests - counter and in-flight gauge (the only server-specific inputs are the two - metric names). -* :class:`CompletionRateCapacityEstimator` -- turns those samples into a - per-worker capacity estimate: per-slot throughput x typical concurrency over a - trailing window. Pure computation; no I/O. -* :class:`NumLoadedModelsPublisher` -- polls a capacity signal on a background - thread and publishes it under ``num_loaded_models_``. -* :class:`AutoscalingMetricConfig` -- the user-facing knobs, plus the tag - resolution and factories a handler uses to build the pieces above. +``num_loaded_models_`` in the ``BeamML_ModelManager`` namespace as the +number of messages each worker can hold, and sizes the pool roughly as +``ceil(backlog / value)``. ``ModelManager`` publishes the number of loaded +model copies there. A handler that serves inference through an out-of-process +model server (vLLM, Triton, ...) always has exactly one copy per worker, so it +publishes an estimate of the worker's capacity under that metric instead. + +The pieces, none of which depend on a particular model server: + +* :class:`PrometheusCounterGaugeSampler` reads a server's completed-requests + counter and in-flight gauge from its Prometheus ``/metrics`` endpoint. +* :class:`CompletionRateCapacityEstimator` turns those samples into a + per-worker capacity estimate. It does no I/O. +* :class:`NumLoadedModelsPublisher` publishes a capacity signal under + ``num_loaded_models_`` from a background thread. +* :class:`AutoscalingMetricConfig` holds the user-facing knobs, resolves the + metric tag, and builds the estimator and publisher. + +A handler typically wires them as:: + + config = config.resolve_tag(model_identifier, model_name) + sampler = PrometheusCounterGaugeSampler(completed_metric, running_metric) + estimator = config.make_capacity_estimator() + + def signal(): + sample = sampler(metrics_port) + return None if sample is None else estimator.update(*sample) + + publisher = config.make_publisher(signal) + publisher.start() # Once the server is reachable; publisher.stop() at exit. """ # pytype: skip-file import dataclasses import logging +import re import threading import time import urllib.request from collections.abc import Callable from collections.abc import Mapping from dataclasses import dataclass +from typing import NamedTuple from typing import Optional from apache_beam.metrics.execution import MetricsEnvironment @@ -70,30 +77,45 @@ 'scrape_prometheus_metrics', ] +# The metric the runner reads (shared with model_manager.py). MODEL_MANAGER_METRICS_NAMESPACE = 'BeamML_ModelManager' NUM_LOADED_MODELS_METRIC_PREFIX = 'num_loaded_models_' +# Publisher defaults. The capacity estimate is already smoothed, so by default +# the publisher forwards the latest value without a second smoothing layer. DEFAULT_POLL_INTERVAL_SECS = 5.0 -# Dimensionless multiplier applied to the capacity estimate before it is -# published. The runner sizes the pool roughly as ``ceil(backlog / reported)``, -# so a larger value lets each worker hold more backlog (fewer workers) and a +DEFAULT_AGGREGATION = 'last' +DEFAULT_AGGREGATION_WINDOW_SECS = 1800.0 # 30 minutes. + +# Estimator defaults. +# A larger multiplier lets each worker hold more backlog (fewer workers); a # smaller one scales out more aggressively. DEFAULT_MULTIPLIER = 1.0 -# Trailing window and concurrency percentile used by -# CompletionRateCapacityEstimator. DEFAULT_SIGNAL_WINDOW_SECS = 1800.0 # 30 minutes. DEFAULT_RUNNING_PERCENTILE = 90.0 DEFAULT_MIN_REQUESTS = 3 DEFAULT_WARMUP_TIMEOUT_SECS = 10.0 -# The capacity estimate is already smoothed, so the publisher forwards its latest -# value unchanged by default (no second smoothing layer). 'lifetime_max' etc. -# remain available as an explicit policy on top. -DEFAULT_AGGREGATION = 'last' -DEFAULT_AGGREGATION_WINDOW_SECS = 1800.0 # 30 minutes. + +# One sample line of the Prometheus text format: +# metric_name{label="value",...} value [timestamp] +_PROMETHEUS_SAMPLE_RE = re.compile( + r'(?P[a-zA-Z_:][a-zA-Z0-9_:]*)' + r'(?:\{(?P(?:[^"}]|"(?:[^"\\]|\\.)*")*)\})?' + r'\s+(?P\S+)') +# One label pair; values are quoted and may contain \\, \" and \n escapes. +_PROMETHEUS_LABEL_RE = re.compile( + r'(?P[a-zA-Z_][a-zA-Z0-9_]*)\s*=\s*"(?P(?:[^"\\]|\\.)*)"') +_PROMETHEUS_ESCAPE_RE = re.compile(r'\\(.)') + +_WINDOWED_AGGREGATIONS = ('max', 'mean') +_PERCENTILE_AGGREGATION_RE = re.compile(r'p(\d+)') def sanitize_metric_tag(name: str) -> str: - """Turns a model name into a metric-tag-safe token for the autoscaler. + """Turns a model name into a metric-tag-safe token. + + Only used for the fallback tag of handlers running outside RunInference; the + RunInference model identifier is always used verbatim. Args: name: The model name (may contain ``/``, ``:`` etc.). @@ -108,7 +130,17 @@ def scrape_prometheus_metrics( port: int, timeout_secs: float = 5.0, path: str = '/metrics') -> Optional[str]: - """Fetches ``http://localhost:``; None on any failure.""" + """Fetches a local Prometheus endpoint. + + Args: + port: Port of the server on ``localhost``. + timeout_secs: Request timeout. + path: Endpoint path. + + Returns: + The response body, or ``None`` on any failure (e.g. the server is still + starting or restarting), so callers can simply skip that sample. + """ url = f'http://localhost:{port}{path}' try: with urllib.request.urlopen(url, timeout=timeout_secs) as response: @@ -117,38 +149,18 @@ def scrape_prometheus_metrics( return None +def _unescape_label_value(value: str) -> str: + """Undoes the backslash escapes in a Prometheus label value.""" + return _PROMETHEUS_ESCAPE_RE.sub( + lambda m: '\n' if m.group(1) == 'n' else m.group(1), value) + + def _parse_prometheus_labels(label_str: str) -> dict[str, str]: - """Parses the ``{k="v",...}`` label section of a Prometheus sample line.""" - labels: dict[str, str] = {} - i = 0 - n = len(label_str) - while i < n: - eq = label_str.find('=', i) - if eq == -1: - break - name = label_str[i:eq].strip() - j = eq + 1 - if j < n and label_str[j] == '"': - j += 1 - value_chars = [] - while j < n and label_str[j] != '"': - if label_str[j] == '\\' and j + 1 < n: - j += 1 - value_chars.append(label_str[j]) - j += 1 - value = ''.join(value_chars) - j += 1 # skip closing quote - else: - comma = label_str.find(',', j) - comma = n if comma == -1 else comma - value = label_str[j:comma].strip() - j = comma - labels[name] = value - comma = label_str.find(',', j) - if comma == -1: - break - i = comma + 1 - return labels + """Parses the inside of the ``{k="v",...}`` section of a sample line.""" + return { + match['name']: _unescape_label_value(match['value']) + for match in _PROMETHEUS_LABEL_RE.finditer(label_str) + } def parse_prometheus_metric( @@ -156,45 +168,32 @@ def parse_prometheus_metric( metric_name: str, label_filter: Optional[Mapping[str, str]] = None, ) -> Optional[float]: - """Sums the values of a metric from Prometheus text-format exposition. + """Sums the values of one metric from Prometheus text-format exposition. Args: metrics_text: The full body returned by the server's ``/metrics`` endpoint. - metric_name: The metric family name to match (e.g. + metric_name: The exact sample name to match (e.g. ``vllm:num_requests_running``). label_filter: If provided, only samples whose labels contain all of these key/value pairs are summed. Returns: - The sum of the matching sample values, or ``None`` if the metric is absent. + The sum of the matching sample values, or ``None`` if there are none. """ total: Optional[float] = None - for raw_line in metrics_text.splitlines(): - line = raw_line.strip() - if not line or line.startswith('#'): - continue - if not line.startswith(metric_name): - continue - remainder = line[len(metric_name):] - if remainder and remainder[0] not in (' ', '\t', '{'): + for line in metrics_text.splitlines(): + match = _PROMETHEUS_SAMPLE_RE.match(line.strip()) + # Comment, blank and other-metric lines don't match. + if match is None or match['name'] != metric_name: continue - labels: Mapping[str, str] = {} - if remainder.startswith('{'): - close = remainder.find('}') - if close == -1: + if label_filter: + labels = _parse_prometheus_labels(match['labels'] or '') + if any(labels.get(k) != v for k, v in label_filter.items()): continue - labels = _parse_prometheus_labels(remainder[1:close]) - remainder = remainder[close + 1:] - parts = remainder.split() - if not parts: - continue try: - value = float(parts[0]) + value = float(match['value']) except ValueError: continue - if label_filter and any(labels.get(k) != v - for k, v in label_filter.items()): - continue total = value if total is None else total + value return total @@ -203,39 +202,66 @@ def _nearest_rank_percentile(values: list[float], q: float) -> float: """Nearest-rank percentile (``q`` in [0, 100]) of a non-empty list.""" ordered = sorted(values) q = min(max(q, 0.0), 100.0) - rank = int((q * len(ordered) + 99) // 100) # ceil(q / 100 * n). - idx = min(max(rank - 1, 0), len(ordered) - 1) - return ordered[idx] + # Nearest rank ceil(q / 100 * n), in arithmetic that avoids float rounding. + rank = int((q * len(ordered) + 99) // 100) + return ordered[min(max(rank - 1, 0), len(ordered) - 1)] + + +def _parse_aggregation(aggregation: str) -> tuple[str, Optional[float]]: + """Validates an aggregation policy; returns ``(mode, percentile)``.""" + mode = (aggregation or DEFAULT_AGGREGATION).lower() + if mode in ('last', 'lifetime_max') + _WINDOWED_AGGREGATIONS: + return mode, None + percentile = _PERCENTILE_AGGREGATION_RE.fullmatch(mode) + if percentile and 0 <= int(percentile.group(1)) <= 100: + return mode, float(percentile.group(1)) + raise ValueError( + f'Unknown aggregation {aggregation!r}: expected one of last, max, mean, ' + 'lifetime_max, or pNN with NN in [0, 100].') class NumLoadedModelsPublisher: - """Publishes the ``num_loaded_models_`` autoscaling metric on a timer. - - ``signal_fn`` produces the raw load sample (e.g. requests running in - parallel, or an estimated per-worker capacity); it is polled on a background - daemon thread every ``poll_interval_secs``. Returning ``None`` (e.g. a - transient scrape failure) skips that tick, leaving the running state and last - published value unchanged. - - The published value is an aggregate of the observed samples, chosen so the - pool does not scale *down* the instant an input burst clears (which would - thrash and then be unable to absorb the next burst): - - * ``'lifetime_max'`` -- the maximum sample observed over the whole - pipeline lifetime (monotonic; never decays). Sizes the - fleet to the worst load ever seen and holds it. - * ``'max'`` -- peak sample in the trailing ``aggregation_window_secs`` - window; decays that long after load drops. - * ``'p'`` -- a windowed percentile, e.g. ``'p90'``. Smoother than max: - ignores brief spikes and decays as low samples fill it. - * ``'mean'`` -- the average over the window. - * ``'last'`` -- the raw last sample (no smoothing). Use this when - ``signal_fn`` already smooths its output. + """Publishes a capacity signal as the ``num_loaded_models_`` metric. + + A daemon thread calls ``signal_fn`` every ``poll_interval_secs``, applies the + ``aggregation`` policy, scales the result by ``multiplier`` and publishes it + to the process-wide distribution the runner reads. A ``signal_fn`` result of + ``None`` (e.g. the server is restarting) skips the tick and leaves the last + published value in place. + + Aggregation policies: + + * ``'last'`` (default): the latest sample. Use it when ``signal_fn`` already + smooths its output, as :class:`CompletionRateCapacityEstimator` does. + * ``'max'``: the peak over the trailing ``aggregation_window_secs``. + * ``'pNN'`` (e.g. ``'p90'``): a percentile over that window. Smoother than + ``'max'``: it ignores brief spikes. + * ``'mean'``: the average over that window. + * ``'lifetime_max'``: the peak since the publisher started. It never decays, + so the pool stays sized for the worst load seen. + + The publisher never reports 0, whether on start, on stop or on a skipped + tick: a 0 tells the runner the worker has no capacity and can trigger a + spurious downscale. Only one publisher per metric name is active in a + process; starting a new one (e.g. after a model reload) stops the previous + one. + + Args: + tag: Metric tag. It must equal the RunInference model identifier the runner + looks up (see :meth:`AutoscalingMetricConfig.resolve_tag`). + signal_fn: Returns the current capacity sample, or ``None`` to skip. + poll_interval_secs: Seconds between samples. + aggregation: One of the policies above. + aggregation_window_secs: Trailing window for ``'max'``, ``'pNN'`` and + ``'mean'``; 0 keeps every sample. + multiplier: Factor applied to every published value. Leave it at 1.0 when + ``signal_fn`` already applies one, as the estimator does. + + Raises: + ValueError: if ``tag`` is empty or ``aggregation`` is not a known policy. """ - # The publisher currently owning each metric name in this process. A model - # update or server restart creates a new server (and publisher) for the same - # identifier; starting it supersedes the stale one so exactly one thread - # publishes per metric name. + # The publisher owning each metric name in this process, so a reloaded + # model's publisher can supersede the stale one. _active: dict[str, 'NumLoadedModelsPublisher'] = {} _active_lock = threading.Lock() @@ -253,93 +279,19 @@ def __init__( self._metric_name = NUM_LOADED_MODELS_METRIC_PREFIX + tag self._signal_fn = signal_fn self._poll_interval_secs = poll_interval_secs - self._aggregation = (aggregation or DEFAULT_AGGREGATION).lower() + self._aggregation, self._percentile = _parse_aggregation(aggregation) self._aggregation_window_secs = max(float(aggregation_window_secs), 0.0) self._multiplier = float(multiplier) - # Trailing (timestamp, value) samples for the windowed aggregation modes. + # (timestamp, value) samples kept for the windowed policies. self._samples: list[tuple[float, float]] = [] - # Monotonic peak for the 'lifetime_max' mode (no sample retention). self._lifetime_max: Optional[float] = None self._stop_event = threading.Event() self._thread: Optional[threading.Thread] = None - def _publish(self, value: float) -> None: - ivalue = int(round(value * self._multiplier)) - # Process-wide DistributionCells are read without resetting by the Python - # SDK harness (_request_harness_monitoring_infos), and Dataflow's worker - # harness reports their cumulative mean (sum / count). Resetting before - # updating keeps (sum=ivalue, count=1) so the latest capacity is reported - # instead of an unweighted lifetime average since worker boot. - cell = MetricsEnvironment.process_wide_container().get_distribution( - MetricName(MODEL_MANAGER_METRICS_NAMESPACE, self._metric_name)) - cumulative = cell.get_cumulative() - if cumulative.count != 1 or cumulative.sum != ivalue: - cell.reset() - cell.update(ivalue) - - def _aggregate(self, now: float) -> Optional[float]: - """Aggregates in-window samples per ``self._aggregation``. - - Evicts samples older than ``aggregation_window_secs`` first, then reduces - the remainder. - - Args: - now: Current timestamp used for window eviction. - - Returns: - The aggregated value, or ``None`` if the window is empty. - """ - if self._aggregation_window_secs > 0: - cutoff = now - self._aggregation_window_secs - self._samples = [(t, v) for (t, v) in self._samples if t >= cutoff] - values = [v for (_, v) in self._samples] - if not values: - return None - mode = self._aggregation - if mode == 'last': - return values[-1] - if mode == 'mean': - return sum(values) / len(values) - if mode.startswith('p'): - try: - q = float(int(mode[1:])) - except ValueError: - return max(values) - return _nearest_rank_percentile(values, q) - return max(values) # 'max' and unknown modes. - - def _run(self) -> None: - while not self._stop_event.is_set(): - try: - value = self._signal_fn() - if value is not None: - if self._aggregation == 'lifetime_max': - # Monotonic running max over the whole pipeline lifetime; no - # per-sample retention and never decays. - self._lifetime_max = ( - float(value) if self._lifetime_max is None else max( - self._lifetime_max, float(value))) - aggregate = self._lifetime_max - else: - now = time.time() - self._samples.append((now, float(value))) - aggregate = self._aggregate(now) - if aggregate is not None: - self._publish(aggregate) - except Exception: # pylint: disable=broad-except - logging.exception( - 'Failed to publish autoscaling metric %s; will retry.', - self._metric_name, - ) - self._stop_event.wait(self._poll_interval_secs) - def start(self) -> None: """Starts the background polling thread. - Deliberately publishes no initial 0: reporting 0 would tell the autoscaler - this replica has no capacity and could trigger a spurious downscale before - the first real sample lands. The metric stays absent until the first - aggregate is published. + The metric stays absent until the first sample is published. """ if self._thread is not None: return @@ -359,10 +311,9 @@ def start(self) -> None: def stop(self) -> None: """Stops polling. - Deliberately publishes no final 0. When a worker is drained during a - downscale, reporting 0 shrinks the aggregate capacity and can drive a - feedback loop of further downscales that tears down healthy replicas. The - drained worker's contribution ages out via metric staleness instead. + Publishes no final 0: when a worker is drained during a downscale, a 0 + would shrink the aggregate capacity and could drive a feedback loop of + further downscales. The drained worker's value ages out instead. """ self._stop_event.set() with NumLoadedModelsPublisher._active_lock: @@ -373,19 +324,67 @@ def stop(self) -> None: thread.join(timeout=self._poll_interval_secs + 5) self._thread = None + def _run(self) -> None: + while not self._stop_event.is_set(): + try: + sample = self._signal_fn() + if sample is not None: + self._publish(self._record(float(sample), time.time())) + except Exception: # pylint: disable=broad-except + logging.exception( + 'Failed to publish autoscaling metric %s; will retry.', + self._metric_name) + self._stop_event.wait(self._poll_interval_secs) + + def _record(self, sample: float, now: float) -> float: + """Adds a sample and returns the value to publish under the policy.""" + if self._aggregation == 'last': + return sample + if self._aggregation == 'lifetime_max': + if self._lifetime_max is None or sample > self._lifetime_max: + self._lifetime_max = sample + return self._lifetime_max + self._samples.append((now, sample)) + if self._aggregation_window_secs > 0: + cutoff = now - self._aggregation_window_secs + self._samples = [(t, v) for (t, v) in self._samples if t >= cutoff] + values = [v for (_, v) in self._samples] + if self._aggregation == 'max': + return max(values) + if self._aggregation == 'mean': + return sum(values) / len(values) + return _nearest_rank_percentile(values, self._percentile) + + def _publish(self, value: float) -> None: + published = int(round(value * self._multiplier)) + # The SDK harness reports process-wide distributions without resetting + # them, and Dataflow reads their cumulative mean (sum / count). Resetting + # before each update keeps the cell at (sum=published, count=1), so the + # runner sees the latest value rather than an average since worker start. + cell = MetricsEnvironment.process_wide_container().get_distribution( + MetricName(MODEL_MANAGER_METRICS_NAMESPACE, self._metric_name)) + current = cell.get_cumulative() + if current.count != 1 or current.sum != published: + cell.reset() + cell.update(published) + class PrometheusCounterGaugeSampler: """Reads a completed-requests counter and an in-flight gauge from a server. Calling the sampler with the server's metrics port scrapes ``http://localhost:/metrics`` and returns ``(completed_total, - running)``: the value of the monotonically increasing ``completed_metric`` - counter and of the ``running_metric`` gauge, each summed over the samples that - match ``label_filter``. It returns ``None`` when the endpoint is unreachable - or either metric is missing. + running)``, each summed over the samples that match ``label_filter``. The two + metric names are the only server-specific inputs, e.g. + ``vllm:request_success_total`` and ``vllm:num_requests_running`` for vLLM. - The metric names are the only server-specific inputs, e.g. - ``('vllm:request_success_total', 'vllm:num_requests_running')`` for vLLM. + Args: + completed_metric: Name of the monotonically increasing counter of completed + requests. + running_metric: Name of the gauge of requests currently in flight. + label_filter: Optional labels a sample must have to be counted, e.g. to + select one model on a multi-model server. + timeout_secs: Scrape timeout. """ def __init__( self, @@ -400,6 +399,7 @@ def __init__( self._timeout_secs = timeout_secs def __call__(self, port: int) -> Optional[tuple[float, float]]: + """Returns ``(completed_total, running)``, or ``None`` if unavailable.""" text = scrape_prometheus_metrics(port, timeout_secs=self._timeout_secs) if text is None: return None @@ -412,40 +412,56 @@ def __call__(self, port: int) -> Optional[tuple[float, float]]: return completed, running +class _Interval(NamedTuple): + """One sampling interval in the estimator's trailing window.""" + end_ts: float + completed: float # Requests completed during the interval. + running: float # Requests in flight at the end of the interval. + secs: float # Length of the interval. + + class CompletionRateCapacityEstimator: - """Estimates a worker's message capacity from completion and concurrency. - - Feed it periodic samples of a model server's monotonically increasing - *completed requests* counter and its *in-flight requests* gauge (for example - from :class:`PrometheusCounterGaugeSampler`) via :meth:`update`. Over a - trailing ``window_secs`` window it derives: - - * ``mu`` -- per-slot throughput, - ``sum(delta_completed) / sum(running * dt)`` - (a ratio of windowed sums, so discrete completions don't make it bursty). - This approximates ``1 / avg_latency`` and is roughly independent of the - offered load. - * ``running_ref`` -- the ``running_percentile`` (default p90) of observed - in-flight concurrency. p90 rather than max avoids latching onto rare, - heavily backlogged spikes that overstate normal-load capacity. - - and estimates a per-worker message capacity:: - - capacity = max(min_value, mu * running_ref * multiplier) - - i.e. the worker's sustained throughput scaled by a dimensionless tuning - ``multiplier``. Normalizing by ``running`` makes the estimate - load-independent: right after an autoscaling event, when the runner - rebalances work and the worker is briefly starved, the raw completion rate - collapses but ``mu`` holds and ``running_ref`` still reflects the concurrency - the worker normally serves, so the estimate stays steady instead of cratering - and triggering a downscale spiral. - - The estimate is already windowed; feed it to ``NumLoadedModelsPublisher`` - with ``aggregation='last'`` unless you deliberately want a second policy - layer (e.g. ``'lifetime_max'`` to hold the peak estimate). - - This class does no I/O, so it is independent of how the samples are read. + """Estimates how many messages a worker can absorb. + + Feed it periodic samples of a model server's completed-requests counter and + in-flight gauge (e.g. from :class:`PrometheusCounterGaugeSampler`) through + :meth:`update`. Over a trailing ``window_secs`` window it computes:: + + mu = sum(completed) / sum(running * secs) + running_ref = percentile(running, running_percentile) + capacity = max(min_value, mu * running_ref * multiplier) + + ``mu`` is the throughput of one in-flight slot (roughly ``1 / latency``), + computed as a ratio of windowed sums so discrete completions don't make it + jumpy. ``running_ref`` is the concurrency the worker typically sustains; the + p90 rather than the max avoids latching onto rare backlog spikes. Their + product is the worker's sustained throughput. + + Normalizing by concurrency keeps the estimate steady when load changes: right + after a scaling event, when the runner rebalances work and a worker is + briefly starved, its raw completion rate collapses, but ``mu`` and + ``running_ref`` hold, so the estimate does not crater and trigger a downscale + spiral. + + The estimator suits servers whose requests take long enough to be seen in + flight by periodic sampling, such as LLM servers with latencies of a second + or more. For servers with millisecond latencies the in-flight gauge reads + almost always 0; pass the publisher a different signal instead. + + The estimate is already smoothed, so publish it with ``aggregation='last'``. + This class does no I/O. + + Args: + multiplier: Dimensionless factor on the estimate; larger means fewer + workers. + min_value: Floor of the estimate. + window_secs: Trailing window; 0 keeps every interval. + running_percentile: Percentile of in-flight concurrency used as + ``running_ref``. + min_requests: Completions a newly started worker needs before it reports, + unless ``warmup_timeout_secs`` passes first. + warmup_timeout_secs: Longest time after a worker first sees work that it + waits for ``min_requests``. """ def __init__( self, @@ -465,8 +481,7 @@ def __init__( self._last_count: Optional[float] = None self._last_ts: Optional[float] = None self._first_active_ts: Optional[float] = None - # Trailing (timestamp, delta_completed, running, dt) samples. - self._samples: list[tuple[float, float, float, float]] = [] + self._samples: list[_Interval] = [] def update( self, @@ -481,99 +496,79 @@ def update( now: Sample timestamp in seconds; defaults to ``time.time()``. Returns: - The capacity estimate, or ``None`` when there is not enough data yet - (first sample, counter reset, idle window, or warmup), in which case the - caller should keep its previous value. + The capacity estimate, or ``None`` when there is not enough data (first + sample, counter reset, idle window or warmup); the caller should then + keep its previous value. """ if now is None: now = time.time() prev_count, prev_ts = self._last_count, self._last_ts self._last_count, self._last_ts = completed_total, now if prev_count is None or prev_ts is None: - # First sample only establishes a baseline; a rate needs two samples. - return None - dt = now - prev_ts - delta = completed_total - prev_count - if dt <= 0 or delta < 0: - # Clock skew or counter reset (server restart); skip this tick. - return None - if running < 0: - return None - # Evict samples older than the window, and skip completely idle intervals so - # trough/idle periods do not dilute running_ref (p90 concurrency) to 0. + return None # The first sample only sets the baseline. + secs = now - prev_ts + completed = completed_total - prev_count + if secs <= 0 or completed < 0 or running < 0: + return None # Clock skew, counter reset (server restart) or bad gauge. + if self._window_secs > 0: cutoff = now - self._window_secs - self._samples = [s for s in self._samples if s[0] >= cutoff] - if delta == 0 and running == 0: - return None - self._samples.append((now, float(delta), float(running), float(dt))) - # Per-slot throughput as a ratio of windowed sums (stable vs. per-tick): - # total completions divided by total slot-seconds of in-flight work. - total_completed = sum(d for (_, d, _, _) in self._samples) - total_slot_secs = sum(r * dts for (_, _, r, dts) in self._samples) - if total_slot_secs <= 0: - # Window is entirely idle (no in-flight requests); nothing to report. + self._samples = [s for s in self._samples if s.end_ts >= cutoff] + if completed == 0 and running == 0: + # Skip fully idle intervals so troughs don't drag running_ref to 0. return None + self._samples.append( + _Interval(now, float(completed), float(running), float(secs))) + + total_completed = sum(s.completed for s in self._samples) + slot_secs = sum(s.running * s.secs for s in self._samples) + if slot_secs <= 0: + return None # Nothing was in flight during the window. if self._first_active_ts is None: self._first_active_ts = prev_ts if (total_completed < self._min_requests and - (now - self._first_active_ts) < self._warmup_timeout_secs): - # Hold off publishing until a newly started worker completes at least - # min_requests or warmup_timeout_secs elapses, so an uncalibrated cold - # worker does not briefly overwrite the job-wide metric with 1. + now - self._first_active_ts < self._warmup_timeout_secs): + # Wait until a new worker is calibrated, so it doesn't briefly report + # the floor value for the whole job. return None - mu = total_completed / total_slot_secs + + mu = total_completed / slot_secs running_ref = _nearest_rank_percentile( - [r for (_, _, r, _) in self._samples if r > 0], + [s.running for s in self._samples if s.running > 0], self._running_percentile) - estimate = mu * running_ref * self._multiplier - return max(float(self._min_value), estimate) + return max(float(self._min_value), mu * running_ref * self._multiplier) @dataclass(frozen=True) class AutoscalingMetricConfig: - """User-facing knobs for the ``num_loaded_models_`` autoscaling signal. - - Model handlers that publish this signal take an instance (typically as an - ``autoscaling_config`` argument) to override the defaults; leave it unset for - the defaults. Each field: - - * ``tag`` -- metric tag to publish under. Leave unset (the default) so the - handler fills it in via :meth:`resolve_tag` from the identifier RunInference - reports for the step (``RunInference(model_identifier=...)``, or the random - id it generates), which is what Dataflow's autoscaler looks up and which - stays stable across model updates. Only set it explicitly if you also pass - the same value as ``model_identifier``. - * ``poll_interval_secs`` -- how often to sample and publish (default 5). - * ``multiplier`` -- dimensionless factor applied to the capacity estimate - before publishing; larger values mean fewer workers, smaller values scale - out more aggressively (default 1.0). - * ``signal_window_secs`` -- trailing window over which - :class:`CompletionRateCapacityEstimator` estimates throughput and p90 - concurrency (default 1800). Ignored when ``load_signal_fn`` is set. - * ``min_requests`` -- minimum completed requests before a newly started - worker begins publishing, unless ``warmup_timeout_secs`` elapses first - (default 3). - * ``warmup_timeout_secs`` -- maximum time (in seconds) since a worker first - sees active requests to wait for ``min_requests`` before publishing anyway - (default 10). - * ``aggregation`` -- optional policy the publisher applies on top of the - (already smoothed) signal: ``'last'`` (default; forward the latest - estimate), ``'lifetime_max'`` (hold the peak estimate for the whole run, - never decays), or ``'max'``/``'p90'``/``'p95'``/``'mean'`` over - ``aggregation_window_secs``. Only change this from ``'last'`` when you - deliberately want a second smoothing layer, e.g. to stop the pool scaling - down the instant a burst clears. - * ``aggregation_window_secs`` -- trailing window for the windowed - ``aggregation`` modes; unused by ``'last'`` and ``'lifetime_max'`` - (default 1800). - * ``load_signal_fn`` -- optional override that, given the server's metrics - port, returns the per-worker capacity value directly, bypassing the - built-in completion-rate estimate. - - A handler calls :meth:`resolve_tag`, builds its capacity signal from a - server-specific sampler and :meth:`make_capacity_estimator` (or uses - ``load_signal_fn``), and passes that signal to :meth:`make_publisher`. + """User-facing knobs for the capacity-based autoscaling signal. + + Model handlers that publish the signal take an instance, typically as an + ``autoscaling_config`` argument; leave it unset for the defaults. A handler + calls :meth:`resolve_tag`, builds its signal from a server-specific sampler + and :meth:`make_capacity_estimator` (or from ``load_signal_fn``), and passes + the signal to :meth:`make_publisher`. + + Attributes: + tag: Metric tag to publish under. Leave unset so the handler uses the + RunInference model identifier (see :meth:`resolve_tag`), which is what + the runner looks up and which stays stable across model updates. Only + set it if you pass the same value as ``RunInference(model_identifier=)``. + poll_interval_secs: How often to sample and publish. + multiplier: Factor on the capacity estimate; larger means fewer workers, + smaller scales out more aggressively. + signal_window_secs: Trailing window of the capacity estimator. + min_requests: Completions a newly started worker needs before it + publishes, unless ``warmup_timeout_secs`` passes first. + warmup_timeout_secs: Longest time a new worker waits for ``min_requests``. + aggregation: Publisher policy on top of the already-smoothed estimate (see + :class:`NumLoadedModelsPublisher`). Keep ``'last'`` unless you want a + second smoothing layer, e.g. ``'lifetime_max'`` to never scale down. + aggregation_window_secs: Window of the ``'max'``, ``'pNN'`` and ``'mean'`` + policies. + load_signal_fn: Optional replacement for the estimator: given the server's + metrics port, returns the capacity directly. Handlers that support it + call it instead of the estimator. """ tag: Optional[str] = None poll_interval_secs: float = DEFAULT_POLL_INTERVAL_SECS @@ -591,11 +586,10 @@ def resolve_tag( fallback_name: str = '') -> 'AutoscalingMetricConfig': """Returns a copy whose ``tag`` is set. - Precedence: an explicit ``tag``; else ``model_identifier``, used verbatim - because it must match the identifier RunInference annotates the step with - (see ``ModelHandler.set_model_identifier``); else ``fallback_name`` - sanitized with :func:`sanitize_metric_tag`, for handlers used outside - RunInference. + Precedence: an explicit ``tag``; then ``model_identifier``, used verbatim + because it must equal the identifier RunInference annotates the step with + (see ``ModelHandler.set_model_identifier``); then ``fallback_name`` passed + through :func:`sanitize_metric_tag`, for handlers used outside RunInference. Raises: ValueError: if no tag can be derived. @@ -620,6 +614,13 @@ def make_capacity_estimator(self) -> CompletionRateCapacityEstimator: def make_publisher( self, signal_fn: Callable[[], Optional[float]]) -> NumLoadedModelsPublisher: + """Builds a :class:`NumLoadedModelsPublisher` for ``signal_fn``. + + The multiplier is not applied again here; the estimator already applies it. + + Raises: + ValueError: if ``tag`` is unset; call :meth:`resolve_tag` first. + """ if not self.tag: raise ValueError( 'AutoscalingMetricConfig.tag is unset; call resolve_tag() first.') diff --git a/sdks/python/apache_beam/ml/inference/autoscaling_metrics_test.py b/sdks/python/apache_beam/ml/inference/autoscaling_metrics_test.py index c7c4f4887073..819cce3bab13 100644 --- a/sdks/python/apache_beam/ml/inference/autoscaling_metrics_test.py +++ b/sdks/python/apache_beam/ml/inference/autoscaling_metrics_test.py @@ -48,27 +48,63 @@ def test_parse_prometheus_metric(self): self.assertIsNone( autoscaling_metrics.parse_prometheus_metric(prom_text, 'vllm:missing')) + def test_parse_prometheus_metric_exposition_details(self): + prom_text = ( + # Series of one counter that differ only by label are summed. + 'vllm:request_success_total{finished_reason="stop"} 10.0\n' + 'vllm:request_success_total{finished_reason="length"} 5.0\n' + # A metric whose name merely starts with the requested one is ignored. + 'vllm:request_success_total_created 1.7e9\n' + # Optional trailing timestamps are ignored. + 'nv_inference_request_success{model="a"} 3 1700000000000\n' + # Quoted label values may contain commas, braces and escapes. + 'labeled{path="a,b}c",quote="say \\"hi\\"",n="x\\ny"} 7\n' + ' \n' + 'not a sample line\n') + parse = autoscaling_metrics.parse_prometheus_metric + self.assertEqual(parse(prom_text, 'vllm:request_success_total'), 15.0) + self.assertEqual(parse(prom_text, 'nv_inference_request_success'), 3.0) + self.assertEqual( + parse(prom_text, 'labeled', label_filter={'path': 'a,b}c'}), 7.0) + self.assertEqual( + parse(prom_text, 'labeled', label_filter={'quote': 'say "hi"'}), 7.0) + self.assertEqual( + parse(prom_text, 'labeled', label_filter={'n': 'x\ny'}), 7.0) + self.assertIsNone( + parse(prom_text, 'labeled', label_filter={'path': 'other'})) + def test_num_loaded_models_publisher_aggregation(self): - samples = [(100.0 + i * 10, float(i + 1)) for i in range(10)] - pub = autoscaling_metrics.NumLoadedModelsPublisher( - tag='test', - signal_fn=lambda: 1.0, - aggregation='p90', - aggregation_window_secs=1800.0) - pub._samples = list(samples) - self.assertEqual(pub._aggregate(now=200.0), 9.0) + def published(aggregation, samples, window_secs=60.0): + """Feeds (timestamp, value) samples; returns the value after each.""" + pub = autoscaling_metrics.NumLoadedModelsPublisher( + tag='test', + signal_fn=lambda: None, + aggregation=aggregation, + aggregation_window_secs=window_secs) + return [pub._record(value, now) for (now, value) in samples] - pub._samples = [(100.0, 2.0), (1900.0, 8.0), (1950.0, 10.0)] - self.assertEqual(pub._aggregate(now=2000.0), 10.0) + samples = [(100.0, 2.0), (102.0, 5.0), (104.0, 3.0)] + self.assertEqual(published('last', samples), [2.0, 5.0, 3.0]) + self.assertEqual(published('max', samples), [2.0, 5.0, 5.0]) + self.assertEqual(published('mean', samples), [2.0, 3.5, 10.0 / 3]) + self.assertEqual(published('lifetime_max', samples), [2.0, 5.0, 5.0]) + # Windowed policies forget samples older than the window. + self.assertEqual( + published('max', [(100.0, 9.0), (130.0, 4.0), (170.0, 3.0)]), + [9.0, 9.0, 4.0]) + # lifetime_max never decays. + self.assertEqual( + published('lifetime_max', [(100.0, 9.0), (1000.0, 1.0)]), [9.0, 9.0]) + # p90 over 1..10 is 9 (nearest rank). + ten = [(100.0 + i, float(i + 1)) for i in range(10)] + self.assertEqual(published('p90', ten, window_secs=1800.0)[-1], 9.0) + self.assertEqual(published('P50', ten, window_secs=1800.0)[-1], 5.0) - pub._samples = [(100.0, 2.0), (102.0, 5.0), (104.0, 3.0)] - pub._aggregation_window_secs = 60.0 - pub._aggregation = 'max' - self.assertEqual(pub._aggregate(now=105.0), 5.0) - pub._aggregation = 'last' - self.assertEqual(pub._aggregate(now=105.0), 3.0) - pub._aggregation = 'mean' - self.assertAlmostEqual(pub._aggregate(now=105.0), 3.3333333333333335) + def test_num_loaded_models_publisher_rejects_unknown_aggregation(self): + for aggregation in ('median', 'p', 'p101', 'p9.5'): + with self.assertRaisesRegex(ValueError, 'Unknown aggregation'): + autoscaling_metrics.NumLoadedModelsPublisher( + tag='test', signal_fn=lambda: None, aggregation=aggregation) def test_publish_resets_process_wide_distribution_cell(self): pub = autoscaling_metrics.NumLoadedModelsPublisher( From f84e5f4f3303b9aaefe1527dc466aa8104060fea Mon Sep 17 00:00:00 2001 From: AMOOOMA Date: Tue, 3 Nov 2026 00:21:23 +0000 Subject: [PATCH 03/14] Make NumLoadedModelsPublisher a plain sink Each component now has one job: the sampler reads, the estimator computes (the only window and the only multiplier), the publisher publishes, and the config wires them together. Remove the publisher's aggregation policies, its aggregation window and its multiplier, and the matching AutoscalingMetricConfig fields. The capacity estimate is already smoothed and its multiplier already applied, so for the estimator these were a no-op ('last', 1.0) that read like a second smoothing layer. A handler whose signal needs extra smoothing can apply it in its signal_fn. The estimator, sampler and Prometheus parsing are unchanged. --- .../ml/inference/autoscaling_metrics.py | 136 ++++-------------- .../ml/inference/autoscaling_metrics_test.py | 67 ++++----- 2 files changed, 61 insertions(+), 142 deletions(-) diff --git a/sdks/python/apache_beam/ml/inference/autoscaling_metrics.py b/sdks/python/apache_beam/ml/inference/autoscaling_metrics.py index 1752fa5ff11f..37d447a2e9e6 100644 --- a/sdks/python/apache_beam/ml/inference/autoscaling_metrics.py +++ b/sdks/python/apache_beam/ml/inference/autoscaling_metrics.py @@ -27,14 +27,15 @@ The pieces, none of which depend on a particular model server: -* :class:`PrometheusCounterGaugeSampler` reads a server's completed-requests - counter and in-flight gauge from its Prometheus ``/metrics`` endpoint. -* :class:`CompletionRateCapacityEstimator` turns those samples into a - per-worker capacity estimate. It does no I/O. -* :class:`NumLoadedModelsPublisher` publishes a capacity signal under - ``num_loaded_models_`` from a background thread. -* :class:`AutoscalingMetricConfig` holds the user-facing knobs, resolves the - metric tag, and builds the estimator and publisher. +* :class:`PrometheusCounterGaugeSampler` *reads*: it scrapes a server's + completed-requests counter and in-flight gauge from its Prometheus + ``/metrics`` endpoint. +* :class:`CompletionRateCapacityEstimator` *computes*: it turns those samples + into a smoothed per-worker capacity estimate. It does no I/O. +* :class:`NumLoadedModelsPublisher` *publishes*: a background thread polls a + capacity signal and writes each value to ``num_loaded_models_`` as is. +* :class:`AutoscalingMetricConfig` *configures*: it holds the user-facing + knobs, resolves the metric tag, and builds the estimator and publisher. A handler typically wires them as:: @@ -81,11 +82,7 @@ def signal(): MODEL_MANAGER_METRICS_NAMESPACE = 'BeamML_ModelManager' NUM_LOADED_MODELS_METRIC_PREFIX = 'num_loaded_models_' -# Publisher defaults. The capacity estimate is already smoothed, so by default -# the publisher forwards the latest value without a second smoothing layer. DEFAULT_POLL_INTERVAL_SECS = 5.0 -DEFAULT_AGGREGATION = 'last' -DEFAULT_AGGREGATION_WINDOW_SECS = 1800.0 # 30 minutes. # Estimator defaults. # A larger multiplier lets each worker hold more backlog (fewer workers); a @@ -107,9 +104,6 @@ def signal(): r'(?P[a-zA-Z_][a-zA-Z0-9_]*)\s*=\s*"(?P(?:[^"\\]|\\.)*)"') _PROMETHEUS_ESCAPE_RE = re.compile(r'\\(.)') -_WINDOWED_AGGREGATIONS = ('max', 'mean') -_PERCENTILE_AGGREGATION_RE = re.compile(r'p(\d+)') - def sanitize_metric_tag(name: str) -> str: """Turns a model name into a metric-tag-safe token. @@ -207,58 +201,30 @@ def _nearest_rank_percentile(values: list[float], q: float) -> float: return ordered[min(max(rank - 1, 0), len(ordered) - 1)] -def _parse_aggregation(aggregation: str) -> tuple[str, Optional[float]]: - """Validates an aggregation policy; returns ``(mode, percentile)``.""" - mode = (aggregation or DEFAULT_AGGREGATION).lower() - if mode in ('last', 'lifetime_max') + _WINDOWED_AGGREGATIONS: - return mode, None - percentile = _PERCENTILE_AGGREGATION_RE.fullmatch(mode) - if percentile and 0 <= int(percentile.group(1)) <= 100: - return mode, float(percentile.group(1)) - raise ValueError( - f'Unknown aggregation {aggregation!r}: expected one of last, max, mean, ' - 'lifetime_max, or pNN with NN in [0, 100].') - - class NumLoadedModelsPublisher: """Publishes a capacity signal as the ``num_loaded_models_`` metric. - A daemon thread calls ``signal_fn`` every ``poll_interval_secs``, applies the - ``aggregation`` policy, scales the result by ``multiplier`` and publishes it - to the process-wide distribution the runner reads. A ``signal_fn`` result of - ``None`` (e.g. the server is restarting) skips the tick and leaves the last - published value in place. - - Aggregation policies: - - * ``'last'`` (default): the latest sample. Use it when ``signal_fn`` already - smooths its output, as :class:`CompletionRateCapacityEstimator` does. - * ``'max'``: the peak over the trailing ``aggregation_window_secs``. - * ``'pNN'`` (e.g. ``'p90'``): a percentile over that window. Smoother than - ``'max'``: it ignores brief spikes. - * ``'mean'``: the average over that window. - * ``'lifetime_max'``: the peak since the publisher started. It never decays, - so the pool stays sized for the worst load seen. - - The publisher never reports 0, whether on start, on stop or on a skipped - tick: a 0 tells the runner the worker has no capacity and can trigger a - spurious downscale. Only one publisher per metric name is active in a - process; starting a new one (e.g. after a model reload) stops the previous - one. + A daemon thread calls ``signal_fn`` every ``poll_interval_secs`` and writes + each value, rounded to an int, to the process-wide distribution the runner + reads. The publisher does not transform values: any smoothing belongs in + ``signal_fn`` (e.g. :class:`CompletionRateCapacityEstimator`, which already + smooths and never returns less than its ``min_value``). + + A ``signal_fn`` result of ``None`` (e.g. the server is restarting) skips the + tick and leaves the last published value in place. The publisher never + injects a 0 of its own, on start, on stop or on a skipped tick: a 0 tells the + runner the worker has no capacity and can trigger a spurious downscale. Only + one publisher per metric name is active in a process; starting a new one + (e.g. after a model reload) stops the previous one. Args: tag: Metric tag. It must equal the RunInference model identifier the runner looks up (see :meth:`AutoscalingMetricConfig.resolve_tag`). - signal_fn: Returns the current capacity sample, or ``None`` to skip. - poll_interval_secs: Seconds between samples. - aggregation: One of the policies above. - aggregation_window_secs: Trailing window for ``'max'``, ``'pNN'`` and - ``'mean'``; 0 keeps every sample. - multiplier: Factor applied to every published value. Leave it at 1.0 when - ``signal_fn`` already applies one, as the estimator does. + signal_fn: Returns the current capacity, or ``None`` to skip the tick. + poll_interval_secs: Seconds between polls. Raises: - ValueError: if ``tag`` is empty or ``aggregation`` is not a known policy. + ValueError: if ``tag`` is empty. """ # The publisher owning each metric name in this process, so a reloaded # model's publisher can supersede the stale one. @@ -270,21 +236,12 @@ def __init__( tag: str, signal_fn: Callable[[], Optional[float]], poll_interval_secs: float = DEFAULT_POLL_INTERVAL_SECS, - aggregation: str = DEFAULT_AGGREGATION, - aggregation_window_secs: float = DEFAULT_AGGREGATION_WINDOW_SECS, - multiplier: float = DEFAULT_MULTIPLIER, ): if not tag: raise ValueError('num_loaded_models tag must be a non-empty string.') self._metric_name = NUM_LOADED_MODELS_METRIC_PREFIX + tag self._signal_fn = signal_fn self._poll_interval_secs = poll_interval_secs - self._aggregation, self._percentile = _parse_aggregation(aggregation) - self._aggregation_window_secs = max(float(aggregation_window_secs), 0.0) - self._multiplier = float(multiplier) - # (timestamp, value) samples kept for the windowed policies. - self._samples: list[tuple[float, float]] = [] - self._lifetime_max: Optional[float] = None self._stop_event = threading.Event() self._thread: Optional[threading.Thread] = None @@ -312,7 +269,7 @@ def stop(self) -> None: """Stops polling. Publishes no final 0: when a worker is drained during a downscale, a 0 - would shrink the aggregate capacity and could drive a feedback loop of + would shrink the pool's total capacity and could drive a feedback loop of further downscales. The drained worker's value ages out instead. """ self._stop_event.set() @@ -327,36 +284,17 @@ def stop(self) -> None: def _run(self) -> None: while not self._stop_event.is_set(): try: - sample = self._signal_fn() - if sample is not None: - self._publish(self._record(float(sample), time.time())) + value = self._signal_fn() + if value is not None: + self._publish(value) except Exception: # pylint: disable=broad-except logging.exception( 'Failed to publish autoscaling metric %s; will retry.', self._metric_name) self._stop_event.wait(self._poll_interval_secs) - def _record(self, sample: float, now: float) -> float: - """Adds a sample and returns the value to publish under the policy.""" - if self._aggregation == 'last': - return sample - if self._aggregation == 'lifetime_max': - if self._lifetime_max is None or sample > self._lifetime_max: - self._lifetime_max = sample - return self._lifetime_max - self._samples.append((now, sample)) - if self._aggregation_window_secs > 0: - cutoff = now - self._aggregation_window_secs - self._samples = [(t, v) for (t, v) in self._samples if t >= cutoff] - values = [v for (_, v) in self._samples] - if self._aggregation == 'max': - return max(values) - if self._aggregation == 'mean': - return sum(values) / len(values) - return _nearest_rank_percentile(values, self._percentile) - def _publish(self, value: float) -> None: - published = int(round(value * self._multiplier)) + published = int(round(value)) # The SDK harness reports process-wide distributions without resetting # them, and Dataflow reads their cumulative mean (sum / count). Resetting # before each update keeps the cell at (sum=published, count=1), so the @@ -448,8 +386,7 @@ class CompletionRateCapacityEstimator: or more. For servers with millisecond latencies the in-flight gauge reads almost always 0; pass the publisher a different signal instead. - The estimate is already smoothed, so publish it with ``aggregation='last'``. - This class does no I/O. + The estimate is already smoothed; publish it as is. This class does no I/O. Args: multiplier: Dimensionless factor on the estimate; larger means fewer @@ -556,16 +493,11 @@ class AutoscalingMetricConfig: set it if you pass the same value as ``RunInference(model_identifier=)``. poll_interval_secs: How often to sample and publish. multiplier: Factor on the capacity estimate; larger means fewer workers, - smaller scales out more aggressively. + smaller scales out more aggressively. Applied once, by the estimator. signal_window_secs: Trailing window of the capacity estimator. min_requests: Completions a newly started worker needs before it publishes, unless ``warmup_timeout_secs`` passes first. warmup_timeout_secs: Longest time a new worker waits for ``min_requests``. - aggregation: Publisher policy on top of the already-smoothed estimate (see - :class:`NumLoadedModelsPublisher`). Keep ``'last'`` unless you want a - second smoothing layer, e.g. ``'lifetime_max'`` to never scale down. - aggregation_window_secs: Window of the ``'max'``, ``'pNN'`` and ``'mean'`` - policies. load_signal_fn: Optional replacement for the estimator: given the server's metrics port, returns the capacity directly. Handlers that support it call it instead of the estimator. @@ -576,8 +508,6 @@ class AutoscalingMetricConfig: signal_window_secs: float = DEFAULT_SIGNAL_WINDOW_SECS min_requests: int = DEFAULT_MIN_REQUESTS warmup_timeout_secs: float = DEFAULT_WARMUP_TIMEOUT_SECS - aggregation: str = DEFAULT_AGGREGATION - aggregation_window_secs: float = DEFAULT_AGGREGATION_WINDOW_SECS load_signal_fn: Optional[Callable[[int], Optional[float]]] = None def resolve_tag( @@ -616,8 +546,6 @@ def make_publisher( Optional[float]]) -> NumLoadedModelsPublisher: """Builds a :class:`NumLoadedModelsPublisher` for ``signal_fn``. - The multiplier is not applied again here; the estimator already applies it. - Raises: ValueError: if ``tag`` is unset; call :meth:`resolve_tag` first. """ @@ -628,6 +556,4 @@ def make_publisher( tag=self.tag, signal_fn=signal_fn, poll_interval_secs=self.poll_interval_secs, - aggregation=self.aggregation, - aggregation_window_secs=self.aggregation_window_secs, ) diff --git a/sdks/python/apache_beam/ml/inference/autoscaling_metrics_test.py b/sdks/python/apache_beam/ml/inference/autoscaling_metrics_test.py index 819cce3bab13..f837f8e6642e 100644 --- a/sdks/python/apache_beam/ml/inference/autoscaling_metrics_test.py +++ b/sdks/python/apache_beam/ml/inference/autoscaling_metrics_test.py @@ -73,42 +73,40 @@ def test_parse_prometheus_metric_exposition_details(self): self.assertIsNone( parse(prom_text, 'labeled', label_filter={'path': 'other'})) - def test_num_loaded_models_publisher_aggregation(self): - def published(aggregation, samples, window_secs=60.0): - """Feeds (timestamp, value) samples; returns the value after each.""" - pub = autoscaling_metrics.NumLoadedModelsPublisher( - tag='test', - signal_fn=lambda: None, - aggregation=aggregation, - aggregation_window_secs=window_secs) - return [pub._record(value, now) for (now, value) in samples] + def test_publisher_publishes_signal_values_as_is(self): + """Each signal value is published rounded; None skips the tick.""" + cell = ( + autoscaling_metrics.MetricsEnvironment.process_wide_container(). + get_distribution( + autoscaling_metrics.MetricName( + autoscaling_metrics.MODEL_MANAGER_METRICS_NAMESPACE, + 'num_loaded_models_sink_test'))) + cell.reset() + self.addCleanup(cell.reset) + published = [] + signal_values = [3.4, None, 7.6] - samples = [(100.0, 2.0), (102.0, 5.0), (104.0, 3.0)] - self.assertEqual(published('last', samples), [2.0, 5.0, 3.0]) - self.assertEqual(published('max', samples), [2.0, 5.0, 5.0]) - self.assertEqual(published('mean', samples), [2.0, 3.5, 10.0 / 3]) - self.assertEqual(published('lifetime_max', samples), [2.0, 5.0, 5.0]) - # Windowed policies forget samples older than the window. - self.assertEqual( - published('max', [(100.0, 9.0), (130.0, 4.0), (170.0, 3.0)]), - [9.0, 9.0, 4.0]) - # lifetime_max never decays. - self.assertEqual( - published('lifetime_max', [(100.0, 9.0), (1000.0, 1.0)]), [9.0, 9.0]) - # p90 over 1..10 is 9 (nearest rank). - ten = [(100.0 + i, float(i + 1)) for i in range(10)] - self.assertEqual(published('p90', ten, window_secs=1800.0)[-1], 9.0) - self.assertEqual(published('P50', ten, window_secs=1800.0)[-1], 5.0) + def signal_fn(): + if published: + # Record what the previous tick left in the cell. + published[-1] = cell.get_cumulative().sum + if not signal_values: + pub._stop_event.set() # pylint: disable=protected-access + return None + published.append(None) + return signal_values.pop(0) - def test_num_loaded_models_publisher_rejects_unknown_aggregation(self): - for aggregation in ('median', 'p', 'p101', 'p9.5'): - with self.assertRaisesRegex(ValueError, 'Unknown aggregation'): - autoscaling_metrics.NumLoadedModelsPublisher( - tag='test', signal_fn=lambda: None, aggregation=aggregation) + pub = autoscaling_metrics.NumLoadedModelsPublisher( + tag='sink_test', signal_fn=signal_fn, poll_interval_secs=0) + # Run the polling loop on this thread so the test is deterministic. + pub._run() # pylint: disable=protected-access + # 3.4 -> 3; None leaves 3 in place; 7.6 -> 8. + self.assertEqual(published, [3, 3, 8]) + self.assertEqual(cell.get_cumulative().count, 1) def test_publish_resets_process_wide_distribution_cell(self): pub = autoscaling_metrics.NumLoadedModelsPublisher( - tag='reset_test', signal_fn=lambda: 1.0, multiplier=1.0) + tag='reset_test', signal_fn=lambda: 1.0) cell = ( autoscaling_metrics.MetricsEnvironment.process_wide_container(). get_distribution( @@ -264,9 +262,7 @@ def test_autoscaling_metric_config_resolve_tag(self): multiplier=2.0, signal_window_secs=30.0, min_requests=5, - warmup_timeout_secs=15.0, - aggregation='lifetime_max', - aggregation_window_secs=40.0) + warmup_timeout_secs=15.0) # The RunInference identifier is used verbatim (it has to match the step # annotation), even if it contains characters sanitize would change. self.assertEqual( @@ -290,10 +286,7 @@ def test_autoscaling_metric_config_resolve_tag(self): self.assertEqual(estimator._warmup_timeout_secs, 15.0) publisher = cfg.resolve_tag('step-id').make_publisher(lambda: None) self.assertEqual(publisher._metric_name, 'num_loaded_models_step-id') - self.assertEqual(publisher._aggregation, 'lifetime_max') self.assertEqual(publisher._poll_interval_secs, 7.0) - # The multiplier is applied by the estimator, never a second time here. - self.assertEqual(publisher._multiplier, 1.0) if __name__ == '__main__': From 4345b29ba53e9f5712a5d0477ebd42b7d56d26ed Mon Sep 17 00:00:00 2001 From: AMOOOMA Date: Tue, 3 Nov 2026 00:34:49 +0000 Subject: [PATCH 04/14] Package the model-server capacity signal behind one call Add AutoscalingMetricConfig.make_model_server_publisher(metrics_port_fn, completed_metric, running_metric). A handler for a model server that exposes a completed-requests counter and an in-flight gauge now only supplies those two metric names and how to find the server's port; the library builds the sampler and estimator, wires them together, honors load_signal_fn, skips ticks while the server is down, and returns a ready publisher. The sampler, estimator and make_publisher stay public for handlers that need a different signal. --- .../ml/inference/autoscaling_metrics.py | 90 ++++++++++++++++--- .../ml/inference/autoscaling_metrics_test.py | 51 +++++++++++ 2 files changed, 127 insertions(+), 14 deletions(-) diff --git a/sdks/python/apache_beam/ml/inference/autoscaling_metrics.py b/sdks/python/apache_beam/ml/inference/autoscaling_metrics.py index 37d447a2e9e6..797746f8e015 100644 --- a/sdks/python/apache_beam/ml/inference/autoscaling_metrics.py +++ b/sdks/python/apache_beam/ml/inference/autoscaling_metrics.py @@ -37,18 +37,19 @@ * :class:`AutoscalingMetricConfig` *configures*: it holds the user-facing knobs, resolves the metric tag, and builds the estimator and publisher. -A handler typically wires them as:: +A handler for a model server that exposes a completed-requests counter and an +in-flight gauge only supplies the server-specific parts, the two metric names +and how to find the server's port, and gets a ready publisher:: config = config.resolve_tag(model_identifier, model_name) - sampler = PrometheusCounterGaugeSampler(completed_metric, running_metric) - estimator = config.make_capacity_estimator() - - def signal(): - sample = sampler(metrics_port) - return None if sample is None else estimator.update(*sample) - - publisher = config.make_publisher(signal) + publisher = config.make_model_server_publisher( + metrics_port_fn=lambda: server.port if server.is_up else None, + completed_metric='vllm:request_success_total', + running_metric='vllm:num_requests_running') publisher.start() # Once the server is reachable; publisher.stop() at exit. + +Handlers whose server needs a different capacity signal build it themselves +and pass it to :meth:`AutoscalingMetricConfig.make_publisher`. """ # pytype: skip-file @@ -476,15 +477,41 @@ def update( return max(float(self._min_value), mu * running_ref * self._multiplier) +class _ModelServerCapacitySignal: + """Capacity signal of a model server: scrapes it, then estimates capacity.""" + def __init__( + self, + metrics_port_fn: Callable[[], Optional[int]], + sampler: PrometheusCounterGaugeSampler, + estimator: CompletionRateCapacityEstimator, + load_signal_fn: Optional[Callable[[int], Optional[float]]] = None, + ): + self._metrics_port_fn = metrics_port_fn + self._sampler = sampler + self._estimator = estimator + self._load_signal_fn = load_signal_fn + + def __call__(self) -> Optional[float]: + port = self._metrics_port_fn() + if port is None: + return None # The server is not up; skip the tick. + if self._load_signal_fn is not None: + return self._load_signal_fn(port) + sample = self._sampler(port) + if sample is None: + return None + return self._estimator.update(*sample) + + @dataclass(frozen=True) class AutoscalingMetricConfig: """User-facing knobs for the capacity-based autoscaling signal. Model handlers that publish the signal take an instance, typically as an ``autoscaling_config`` argument; leave it unset for the defaults. A handler - calls :meth:`resolve_tag`, builds its signal from a server-specific sampler - and :meth:`make_capacity_estimator` (or from ``load_signal_fn``), and passes - the signal to :meth:`make_publisher`. + calls :meth:`resolve_tag` and then :meth:`make_model_server_publisher` with + its server's metric names, or builds its own signal and passes it to + :meth:`make_publisher`. Attributes: tag: Metric tag to publish under. Leave unset so the handler uses the @@ -499,8 +526,8 @@ class AutoscalingMetricConfig: publishes, unless ``warmup_timeout_secs`` passes first. warmup_timeout_secs: Longest time a new worker waits for ``min_requests``. load_signal_fn: Optional replacement for the estimator: given the server's - metrics port, returns the capacity directly. Handlers that support it - call it instead of the estimator. + metrics port, returns the capacity directly. Used by + :meth:`make_model_server_publisher`. """ tag: Optional[str] = None poll_interval_secs: float = DEFAULT_POLL_INTERVAL_SECS @@ -541,6 +568,41 @@ def make_capacity_estimator(self) -> CompletionRateCapacityEstimator: min_requests=self.min_requests, warmup_timeout_secs=self.warmup_timeout_secs) + def make_model_server_publisher( + self, + metrics_port_fn: Callable[[], Optional[int]], + completed_metric: str, + running_metric: str, + label_filter: Optional[Mapping[str, str]] = None, + ) -> NumLoadedModelsPublisher: + """Builds a publisher of a model server's capacity. + + The publisher scrapes the server's Prometheus ``/metrics`` endpoint with a + :class:`PrometheusCounterGaugeSampler`, feeds the samples to a + :class:`CompletionRateCapacityEstimator` built from these knobs, and + publishes the estimate. If ``load_signal_fn`` is set, it replaces the + scrape and the estimate. + + Args: + metrics_port_fn: Returns the port serving ``/metrics`` while the server + is up, or ``None`` while it is not (the publisher then skips the tick). + completed_metric: Name of the server's counter of completed requests, + e.g. ``vllm:request_success_total``. + running_metric: Name of the server's gauge of in-flight requests, e.g. + ``vllm:num_requests_running``. + label_filter: Optional labels a sample must have to be counted. + + Raises: + ValueError: if ``tag`` is unset; call :meth:`resolve_tag` first. + """ + return self.make_publisher( + _ModelServerCapacitySignal( + metrics_port_fn, + PrometheusCounterGaugeSampler( + completed_metric, running_metric, label_filter=label_filter), + self.make_capacity_estimator(), + self.load_signal_fn)) + def make_publisher( self, signal_fn: Callable[[], Optional[float]]) -> NumLoadedModelsPublisher: diff --git a/sdks/python/apache_beam/ml/inference/autoscaling_metrics_test.py b/sdks/python/apache_beam/ml/inference/autoscaling_metrics_test.py index f837f8e6642e..123f51bded53 100644 --- a/sdks/python/apache_beam/ml/inference/autoscaling_metrics_test.py +++ b/sdks/python/apache_beam/ml/inference/autoscaling_metrics_test.py @@ -288,6 +288,57 @@ def test_autoscaling_metric_config_resolve_tag(self): self.assertEqual(publisher._metric_name, 'num_loaded_models_step-id') self.assertEqual(publisher._poll_interval_secs, 7.0) + def test_make_model_server_publisher(self): + """The handler supplies metric names and a port; the rest is packaged.""" + config = autoscaling_metrics.AutoscalingMetricConfig( + multiplier=2.0, poll_interval_secs=7.0).resolve_tag('step-id') + port = [None] + publisher = config.make_model_server_publisher( + metrics_port_fn=lambda: port[0], + completed_metric='done_total', + running_metric='running') + self.assertEqual(publisher._metric_name, 'num_loaded_models_step-id') + self.assertEqual(publisher._poll_interval_secs, 7.0) + signal = publisher._signal_fn # What the publisher's thread polls. + + # Server not up yet -> skip the tick without scraping. + with mock.patch('urllib.request.urlopen') as urlopen: + self.assertIsNone(signal()) + urlopen.assert_not_called() + + port[0] = 8000 + responses = [ + 'done_total 10\nrunning 4\n', + 'done_total 30\nrunning 4\n', + ] + + def fake_urlopen(url, timeout=5.0): + self.assertEqual(url, 'http://localhost:8000/metrics') + del timeout + return io.BytesIO(responses.pop(0).encode('utf-8')) + + with mock.patch('urllib.request.urlopen', side_effect=fake_urlopen): + with mock.patch.object(autoscaling_metrics.time, + 'time', + side_effect=[100.0, 105.0]): + self.assertIsNone(signal()) # Baseline sample. + # delta=20 over 5s with running=4 -> mu=1, p90=4, multiplier 2 -> 8. + self.assertAlmostEqual(signal(), 8.0) + + def test_make_model_server_publisher_load_signal_fn_override(self): + config = autoscaling_metrics.AutoscalingMetricConfig( + load_signal_fn=lambda port: port / 1000.0).resolve_tag('step-id') + publisher = config.make_model_server_publisher( + metrics_port_fn=lambda: 8000, + completed_metric='done_total', + running_metric='running') + with mock.patch('urllib.request.urlopen') as urlopen: + self.assertEqual(publisher._signal_fn(), 8.0) + urlopen.assert_not_called() + with self.assertRaises(ValueError): + autoscaling_metrics.AutoscalingMetricConfig().make_model_server_publisher( + lambda: 8000, 'done_total', 'running') + if __name__ == '__main__': unittest.main() From 833afa78493f4df71bbb4d14af52eea7dd1b4729 Mon Sep 17 00:00:00 2001 From: AMOOOMA Date: Tue, 3 Nov 2026 00:42:46 +0000 Subject: [PATCH 05/14] Fold the sampler and signal glue into make_model_server_publisher Producing one capacity value crossed five classes (publisher -> signal -> sampler -> estimator -> publisher), two of which were only glue. Write the poll as a single function inside make_model_server_publisher (port, scrape, parse the two metrics, estimate) and remove PrometheusCounterGaugeSampler, _ModelServerCapacitySignal and AutoscalingMetricConfig.make_capacity_estimator. What remains is three classes with one job each (estimator computes, publisher publishes, config configures) plus the scrape/parse functions. Published values are unchanged. --- .../ml/inference/autoscaling_metrics.py | 152 +++++------------- .../ml/inference/autoscaling_metrics_test.py | 72 ++++----- 2 files changed, 74 insertions(+), 150 deletions(-) diff --git a/sdks/python/apache_beam/ml/inference/autoscaling_metrics.py b/sdks/python/apache_beam/ml/inference/autoscaling_metrics.py index 797746f8e015..e75e354aa8ab 100644 --- a/sdks/python/apache_beam/ml/inference/autoscaling_metrics.py +++ b/sdks/python/apache_beam/ml/inference/autoscaling_metrics.py @@ -25,17 +25,19 @@ model server (vLLM, Triton, ...) always has exactly one copy per worker, so it publishes an estimate of the worker's capacity under that metric instead. -The pieces, none of which depend on a particular model server: +Nothing here depends on a particular model server. There are three classes, +each with one job: -* :class:`PrometheusCounterGaugeSampler` *reads*: it scrapes a server's - completed-requests counter and in-flight gauge from its Prometheus - ``/metrics`` endpoint. -* :class:`CompletionRateCapacityEstimator` *computes*: it turns those samples - into a smoothed per-worker capacity estimate. It does no I/O. +* :class:`CompletionRateCapacityEstimator` *computes* a smoothed per-worker + capacity from a server's completed-requests counter and in-flight gauge. It + does no I/O. * :class:`NumLoadedModelsPublisher` *publishes*: a background thread polls a capacity signal and writes each value to ``num_loaded_models_`` as is. * :class:`AutoscalingMetricConfig` *configures*: it holds the user-facing - knobs, resolves the metric tag, and builds the estimator and publisher. + knobs, resolves the metric tag, and builds the publisher. + +plus :func:`scrape_prometheus_metrics` and :func:`parse_prometheus_metric` to +read a server's Prometheus ``/metrics`` endpoint. A handler for a model server that exposes a completed-requests counter and an in-flight gauge only supplies the server-specific parts, the two metric names @@ -73,7 +75,6 @@ 'AutoscalingMetricConfig', 'CompletionRateCapacityEstimator', 'NumLoadedModelsPublisher', - 'PrometheusCounterGaugeSampler', 'parse_prometheus_metric', 'sanitize_metric_tag', 'scrape_prometheus_metrics', @@ -144,16 +145,13 @@ def scrape_prometheus_metrics( return None -def _unescape_label_value(value: str) -> str: - """Undoes the backslash escapes in a Prometheus label value.""" - return _PROMETHEUS_ESCAPE_RE.sub( - lambda m: '\n' if m.group(1) == 'n' else m.group(1), value) - - def _parse_prometheus_labels(label_str: str) -> dict[str, str]: - """Parses the inside of the ``{k="v",...}`` section of a sample line.""" + """Parses the inside of ``{k="v",...}``, undoing backslash escapes.""" + def unescape(match: re.Match) -> str: + return '\n' if match.group(1) == 'n' else match.group(1) + return { - match['name']: _unescape_label_value(match['value']) + match['name']: _PROMETHEUS_ESCAPE_RE.sub(unescape, match['value']) for match in _PROMETHEUS_LABEL_RE.finditer(label_str) } @@ -308,49 +306,6 @@ def _publish(self, value: float) -> None: cell.update(published) -class PrometheusCounterGaugeSampler: - """Reads a completed-requests counter and an in-flight gauge from a server. - - Calling the sampler with the server's metrics port scrapes - ``http://localhost:/metrics`` and returns ``(completed_total, - running)``, each summed over the samples that match ``label_filter``. The two - metric names are the only server-specific inputs, e.g. - ``vllm:request_success_total`` and ``vllm:num_requests_running`` for vLLM. - - Args: - completed_metric: Name of the monotonically increasing counter of completed - requests. - running_metric: Name of the gauge of requests currently in flight. - label_filter: Optional labels a sample must have to be counted, e.g. to - select one model on a multi-model server. - timeout_secs: Scrape timeout. - """ - def __init__( - self, - completed_metric: str, - running_metric: str, - label_filter: Optional[Mapping[str, str]] = None, - timeout_secs: float = 5.0, - ): - self._completed_metric = completed_metric - self._running_metric = running_metric - self._label_filter = dict(label_filter) if label_filter else None - self._timeout_secs = timeout_secs - - def __call__(self, port: int) -> Optional[tuple[float, float]]: - """Returns ``(completed_total, running)``, or ``None`` if unavailable.""" - text = scrape_prometheus_metrics(port, timeout_secs=self._timeout_secs) - if text is None: - return None - completed = parse_prometheus_metric( - text, self._completed_metric, label_filter=self._label_filter) - running = parse_prometheus_metric( - text, self._running_metric, label_filter=self._label_filter) - if completed is None or running is None: - return None - return completed, running - - class _Interval(NamedTuple): """One sampling interval in the estimator's trailing window.""" end_ts: float @@ -363,8 +318,7 @@ class CompletionRateCapacityEstimator: """Estimates how many messages a worker can absorb. Feed it periodic samples of a model server's completed-requests counter and - in-flight gauge (e.g. from :class:`PrometheusCounterGaugeSampler`) through - :meth:`update`. Over a trailing ``window_secs`` window it computes:: + in-flight gauge through :meth:`update`. Over a trailing ``window_secs`` window it computes:: mu = sum(completed) / sum(running * secs) running_ref = percentile(running, running_percentile) @@ -477,32 +431,6 @@ def update( return max(float(self._min_value), mu * running_ref * self._multiplier) -class _ModelServerCapacitySignal: - """Capacity signal of a model server: scrapes it, then estimates capacity.""" - def __init__( - self, - metrics_port_fn: Callable[[], Optional[int]], - sampler: PrometheusCounterGaugeSampler, - estimator: CompletionRateCapacityEstimator, - load_signal_fn: Optional[Callable[[int], Optional[float]]] = None, - ): - self._metrics_port_fn = metrics_port_fn - self._sampler = sampler - self._estimator = estimator - self._load_signal_fn = load_signal_fn - - def __call__(self) -> Optional[float]: - port = self._metrics_port_fn() - if port is None: - return None # The server is not up; skip the tick. - if self._load_signal_fn is not None: - return self._load_signal_fn(port) - sample = self._sampler(port) - if sample is None: - return None - return self._estimator.update(*sample) - - @dataclass(frozen=True) class AutoscalingMetricConfig: """User-facing knobs for the capacity-based autoscaling signal. @@ -560,14 +488,6 @@ def resolve_tag( 'fallback_name, or set AutoscalingMetricConfig.tag.') return dataclasses.replace(self, tag=tag) - def make_capacity_estimator(self) -> CompletionRateCapacityEstimator: - """Builds a :class:`CompletionRateCapacityEstimator` from these knobs.""" - return CompletionRateCapacityEstimator( - multiplier=self.multiplier, - window_secs=self.signal_window_secs, - min_requests=self.min_requests, - warmup_timeout_secs=self.warmup_timeout_secs) - def make_model_server_publisher( self, metrics_port_fn: Callable[[], Optional[int]], @@ -577,11 +497,10 @@ def make_model_server_publisher( ) -> NumLoadedModelsPublisher: """Builds a publisher of a model server's capacity. - The publisher scrapes the server's Prometheus ``/metrics`` endpoint with a - :class:`PrometheusCounterGaugeSampler`, feeds the samples to a - :class:`CompletionRateCapacityEstimator` built from these knobs, and - publishes the estimate. If ``load_signal_fn`` is set, it replaces the - scrape and the estimate. + Every poll scrapes the server's Prometheus ``/metrics`` endpoint, reads the + two metrics, feeds them to a :class:`CompletionRateCapacityEstimator` built + from these knobs, and publishes the estimate. If ``load_signal_fn`` is set, + it replaces the scrape and the estimate. Args: metrics_port_fn: Returns the port serving ``/metrics`` while the server @@ -590,18 +509,35 @@ def make_model_server_publisher( e.g. ``vllm:request_success_total``. running_metric: Name of the server's gauge of in-flight requests, e.g. ``vllm:num_requests_running``. - label_filter: Optional labels a sample must have to be counted. + label_filter: Optional labels a sample must have to be counted, e.g. to + select one model on a multi-model server. Raises: ValueError: if ``tag`` is unset; call :meth:`resolve_tag` first. """ - return self.make_publisher( - _ModelServerCapacitySignal( - metrics_port_fn, - PrometheusCounterGaugeSampler( - completed_metric, running_metric, label_filter=label_filter), - self.make_capacity_estimator(), - self.load_signal_fn)) + estimator = CompletionRateCapacityEstimator( + multiplier=self.multiplier, + window_secs=self.signal_window_secs, + min_requests=self.min_requests, + warmup_timeout_secs=self.warmup_timeout_secs) + + def capacity() -> Optional[float]: + """One poll: returns the current capacity, or None to skip the tick.""" + port = metrics_port_fn() + if port is None: + return None # The server is not up. + if self.load_signal_fn is not None: + return self.load_signal_fn(port) + text = scrape_prometheus_metrics(port) + if text is None: + return None + completed = parse_prometheus_metric(text, completed_metric, label_filter) + running = parse_prometheus_metric(text, running_metric, label_filter) + if completed is None or running is None: + return None + return estimator.update(completed, running) + + return self.make_publisher(capacity) def make_publisher( self, signal_fn: Callable[[], diff --git a/sdks/python/apache_beam/ml/inference/autoscaling_metrics_test.py b/sdks/python/apache_beam/ml/inference/autoscaling_metrics_test.py index 123f51bded53..5d02156dd5fe 100644 --- a/sdks/python/apache_beam/ml/inference/autoscaling_metrics_test.py +++ b/sdks/python/apache_beam/ml/inference/autoscaling_metrics_test.py @@ -153,43 +153,6 @@ def test_publisher_for_same_tag_supersedes_previous(self): other.stop() self.assertNotIn('num_loaded_models_reload', active) - def test_prometheus_counter_gauge_sampler(self): - sampler = autoscaling_metrics.PrometheusCounterGaugeSampler( - completed_metric='vllm:request_success_total', - running_metric='vllm:num_requests_running') - responses = [ - 'vllm:request_success_total 10.0\nvllm:num_requests_running 4.0\n', - # Missing the gauge -> None. - 'vllm:request_success_total 11.0\n', - ] - - def fake_urlopen(url, timeout=5.0): - self.assertEqual(url, 'http://localhost:8000/metrics') - del timeout - return io.BytesIO(responses.pop(0).encode('utf-8')) - - with mock.patch('urllib.request.urlopen', side_effect=fake_urlopen): - self.assertEqual(sampler(8000), (10.0, 4.0)) - self.assertIsNone(sampler(8000)) - # Unreachable endpoint -> None. - with mock.patch('urllib.request.urlopen', side_effect=OSError('refused')): - self.assertIsNone(sampler(8000)) - - def test_prometheus_counter_gauge_sampler_label_filter(self): - sampler = autoscaling_metrics.PrometheusCounterGaugeSampler( - completed_metric='nv_inference_request_success', - running_metric='nv_inference_pending_request_count', - label_filter={'model': 'a'}) - body = ( - 'nv_inference_request_success{model="a"} 10\n' - 'nv_inference_request_success{model="b"} 99\n' - 'nv_inference_pending_request_count{model="a"} 2\n' - 'nv_inference_pending_request_count{model="b"} 50\n') - with mock.patch('urllib.request.urlopen', - return_value=io.BytesIO(body.encode('utf-8'))): - # Only model "a" counts. - self.assertEqual(sampler(8002), (10.0, 2.0)) - def test_capacity_estimator(self): estimator = autoscaling_metrics.CompletionRateCapacityEstimator( multiplier=2.0, min_value=1, window_secs=60.0, running_percentile=90.0) @@ -279,11 +242,6 @@ def test_autoscaling_metric_config_resolve_tag(self): with self.assertRaises(ValueError): autoscaling_metrics.AutoscalingMetricConfig().make_publisher(lambda: None) - estimator = cfg.make_capacity_estimator() - self.assertEqual(estimator._multiplier, 2.0) - self.assertEqual(estimator._window_secs, 30.0) - self.assertEqual(estimator._min_requests, 5) - self.assertEqual(estimator._warmup_timeout_secs, 15.0) publisher = cfg.resolve_tag('step-id').make_publisher(lambda: None) self.assertEqual(publisher._metric_name, 'num_loaded_models_step-id') self.assertEqual(publisher._poll_interval_secs, 7.0) @@ -325,6 +283,36 @@ def fake_urlopen(url, timeout=5.0): # delta=20 over 5s with running=4 -> mu=1, p90=4, multiplier 2 -> 8. self.assertAlmostEqual(signal(), 8.0) + def test_make_model_server_publisher_skips_unreadable_scrapes(self): + config = autoscaling_metrics.AutoscalingMetricConfig( + min_requests=0).resolve_tag('step-id') + publisher = config.make_model_server_publisher( + metrics_port_fn=lambda: 8000, + completed_metric='done_total', + running_metric='running', + label_filter={'model': 'a'}) + signal = publisher._signal_fn + scrapes = [ + None, # Endpoint unreachable. + 'done_total{model="a"} 1\n', # Gauge missing. + # Only model "a" counts. + 'done_total{model="a"} 10\ndone_total{model="b"} 99\n' + 'running{model="a"} 2\nrunning{model="b"} 50\n', + 'done_total{model="a"} 20\ndone_total{model="b"} 999\n' + 'running{model="a"} 2\nrunning{model="b"} 50\n', + ] + with mock.patch.object(autoscaling_metrics, + 'scrape_prometheus_metrics', + side_effect=lambda port: scrapes.pop(0)): + with mock.patch.object(autoscaling_metrics.time, + 'time', + side_effect=[100.0, 110.0]): + self.assertIsNone(signal()) + self.assertIsNone(signal()) + self.assertIsNone(signal()) # Baseline sample. + # delta=10 over 10s with running=2 -> mu=0.5, p90=2 -> 1.0. + self.assertAlmostEqual(signal(), 1.0) + def test_make_model_server_publisher_load_signal_fn_override(self): config = autoscaling_metrics.AutoscalingMetricConfig( load_signal_fn=lambda port: port / 1000.0).resolve_tag('step-id') From 5085d0d278c4171f1401567d56d9e0f888460118 Mon Sep 17 00:00:00 2001 From: AMOOOMA Date: Tue, 3 Nov 2026 00:50:42 +0000 Subject: [PATCH 06/14] Make AutoscalingMetricConfig plain data The config class mixed knobs with construction: resolve_tag() returned a copy with the tag filled in, and make_model_server_publisher() held the polling logic. Move construction to module-level functions: - make_model_server_publisher(...): the one call a model-server handler makes (metric names, port, model identifier, optional config). - make_publisher(...): for a handler's own capacity signal. Both resolve the metric tag internally, so handlers no longer call a separate resolve_tag step, and both take keyword-only arguments so the identifier, model name and metric names can't be swapped. AutoscalingMetricConfig is now just the knobs. Published values are unchanged. --- .../ml/inference/autoscaling_metrics.py | 264 +++++++++--------- .../ml/inference/autoscaling_metrics_test.py | 72 ++--- 2 files changed, 179 insertions(+), 157 deletions(-) diff --git a/sdks/python/apache_beam/ml/inference/autoscaling_metrics.py b/sdks/python/apache_beam/ml/inference/autoscaling_metrics.py index e75e354aa8ab..82e83b1393d3 100644 --- a/sdks/python/apache_beam/ml/inference/autoscaling_metrics.py +++ b/sdks/python/apache_beam/ml/inference/autoscaling_metrics.py @@ -25,38 +25,36 @@ model server (vLLM, Triton, ...) always has exactly one copy per worker, so it publishes an estimate of the worker's capacity under that metric instead. -Nothing here depends on a particular model server. There are three classes, -each with one job: - -* :class:`CompletionRateCapacityEstimator` *computes* a smoothed per-worker - capacity from a server's completed-requests counter and in-flight gauge. It - does no I/O. -* :class:`NumLoadedModelsPublisher` *publishes*: a background thread polls a - capacity signal and writes each value to ``num_loaded_models_`` as is. -* :class:`AutoscalingMetricConfig` *configures*: it holds the user-facing - knobs, resolves the metric tag, and builds the publisher. - -plus :func:`scrape_prometheus_metrics` and :func:`parse_prometheus_metric` to -read a server's Prometheus ``/metrics`` endpoint. - -A handler for a model server that exposes a completed-requests counter and an -in-flight gauge only supplies the server-specific parts, the two metric names -and how to find the server's port, and gets a ready publisher:: - - config = config.resolve_tag(model_identifier, model_name) - publisher = config.make_model_server_publisher( +Nothing here depends on a particular model server. A handler for a model server +that exposes a completed-requests counter and an in-flight gauge makes one call, +supplying only what it knows: the two metric names, how to find the server's +port, and the RunInference model identifier:: + + publisher = make_model_server_publisher( + config=autoscaling_config, # An AutoscalingMetricConfig, or None. + model_identifier=model_identifier, + model_name=model_name, metrics_port_fn=lambda: server.port if server.is_up else None, completed_metric='vllm:request_success_total', running_metric='vllm:num_requests_running') publisher.start() # Once the server is reachable; publisher.stop() at exit. -Handlers whose server needs a different capacity signal build it themselves -and pass it to :meth:`AutoscalingMetricConfig.make_publisher`. +Handlers whose server needs a different capacity signal pass it to +:func:`make_publisher` instead. + +Inside, each component has one job: + +* :class:`AutoscalingMetricConfig` holds the user-facing knobs. Plain data. +* :class:`CompletionRateCapacityEstimator` computes a worker's capacity from + the counter and the gauge. It does no I/O. +* :class:`NumLoadedModelsPublisher` publishes: a background thread polls a + capacity signal and writes each value to ``num_loaded_models_``. +* :func:`scrape_prometheus_metrics` and :func:`parse_prometheus_metric` read a + server's Prometheus ``/metrics`` endpoint. """ # pytype: skip-file -import dataclasses import logging import re import threading @@ -75,6 +73,8 @@ 'AutoscalingMetricConfig', 'CompletionRateCapacityEstimator', 'NumLoadedModelsPublisher', + 'make_model_server_publisher', + 'make_publisher', 'parse_prometheus_metric', 'sanitize_metric_tag', 'scrape_prometheus_metrics', @@ -218,7 +218,7 @@ class NumLoadedModelsPublisher: Args: tag: Metric tag. It must equal the RunInference model identifier the runner - looks up (see :meth:`AutoscalingMetricConfig.resolve_tag`). + looks up; :func:`make_publisher` derives it. signal_fn: Returns the current capacity, or ``None`` to skip the tick. poll_interval_secs: Seconds between polls. @@ -436,122 +436,136 @@ class AutoscalingMetricConfig: """User-facing knobs for the capacity-based autoscaling signal. Model handlers that publish the signal take an instance, typically as an - ``autoscaling_config`` argument; leave it unset for the defaults. A handler - calls :meth:`resolve_tag` and then :meth:`make_model_server_publisher` with - its server's metric names, or builds its own signal and passes it to - :meth:`make_publisher`. + ``autoscaling_config`` argument; leave it unset for the defaults. Attributes: - tag: Metric tag to publish under. Leave unset so the handler uses the - RunInference model identifier (see :meth:`resolve_tag`), which is what - the runner looks up and which stays stable across model updates. Only - set it if you pass the same value as ``RunInference(model_identifier=)``. - poll_interval_secs: How often to sample and publish. + poll_interval_secs: How often to read the server and publish. multiplier: Factor on the capacity estimate; larger means fewer workers, - smaller scales out more aggressively. Applied once, by the estimator. - signal_window_secs: Trailing window of the capacity estimator. + smaller scales out more aggressively. + signal_window_secs: Trailing window of the capacity estimate. min_requests: Completions a newly started worker needs before it publishes, unless ``warmup_timeout_secs`` passes first. warmup_timeout_secs: Longest time a new worker waits for ``min_requests``. - load_signal_fn: Optional replacement for the estimator: given the server's - metrics port, returns the capacity directly. Used by - :meth:`make_model_server_publisher`. + load_signal_fn: Optional replacement for the built-in estimate: given the + server's metrics port, returns the capacity directly. + tag: Optional metric tag override. Leave unset so the RunInference model + identifier is used, which is what the runner looks up and which stays + stable across model updates. Only set it if you pass the same value as + ``RunInference(model_identifier=...)``. """ - tag: Optional[str] = None poll_interval_secs: float = DEFAULT_POLL_INTERVAL_SECS multiplier: float = DEFAULT_MULTIPLIER signal_window_secs: float = DEFAULT_SIGNAL_WINDOW_SECS min_requests: int = DEFAULT_MIN_REQUESTS warmup_timeout_secs: float = DEFAULT_WARMUP_TIMEOUT_SECS load_signal_fn: Optional[Callable[[int], Optional[float]]] = None + tag: Optional[str] = None - def resolve_tag( - self, - model_identifier: Optional[str], - fallback_name: str = '') -> 'AutoscalingMetricConfig': - """Returns a copy whose ``tag`` is set. - Precedence: an explicit ``tag``; then ``model_identifier``, used verbatim - because it must equal the identifier RunInference annotates the step with - (see ``ModelHandler.set_model_identifier``); then ``fallback_name`` passed - through :func:`sanitize_metric_tag`, for handlers used outside RunInference. +def _metric_tag( + config: AutoscalingMetricConfig, + model_identifier: Optional[str], + model_name: str) -> str: + """Picks the ``num_loaded_models`` tag. - Raises: - ValueError: if no tag can be derived. - """ - if self.tag: - return self - tag = model_identifier or sanitize_metric_tag(fallback_name or '') - if not tag: - raise ValueError( - 'Cannot derive a num_loaded_models tag: pass a model_identifier, a ' - 'fallback_name, or set AutoscalingMetricConfig.tag.') - return dataclasses.replace(self, tag=tag) + Precedence: ``config.tag``; then ``model_identifier``, used verbatim because + it must equal the identifier RunInference annotates the step with (see + ``ModelHandler.set_model_identifier``); then ``model_name`` passed through + :func:`sanitize_metric_tag`, for handlers used outside RunInference. + """ + tag = config.tag or model_identifier or sanitize_metric_tag(model_name) + if not tag: + raise ValueError( + 'Cannot derive a num_loaded_models tag: pass a model_identifier, a ' + 'model_name, or set AutoscalingMetricConfig.tag.') + return tag + + +def make_publisher( + *, + signal_fn: Callable[[], Optional[float]], + model_identifier: Optional[str], + model_name: str = '', + config: Optional[AutoscalingMetricConfig] = None, +) -> NumLoadedModelsPublisher: + """Builds a publisher for a handler's own capacity signal. - def make_model_server_publisher( - self, - metrics_port_fn: Callable[[], Optional[int]], - completed_metric: str, - running_metric: str, - label_filter: Optional[Mapping[str, str]] = None, - ) -> NumLoadedModelsPublisher: - """Builds a publisher of a model server's capacity. - - Every poll scrapes the server's Prometheus ``/metrics`` endpoint, reads the - two metrics, feeds them to a :class:`CompletionRateCapacityEstimator` built - from these knobs, and publishes the estimate. If ``load_signal_fn`` is set, - it replaces the scrape and the estimate. + Args: + signal_fn: Returns the current capacity, or ``None`` to skip the tick. + model_identifier: The identifier RunInference reported through + ``ModelHandler.set_model_identifier``, or ``None`` outside RunInference. + model_name: Fallback for the metric tag outside RunInference. + config: Knobs; ``None`` uses the defaults. - Args: - metrics_port_fn: Returns the port serving ``/metrics`` while the server - is up, or ``None`` while it is not (the publisher then skips the tick). - completed_metric: Name of the server's counter of completed requests, - e.g. ``vllm:request_success_total``. - running_metric: Name of the server's gauge of in-flight requests, e.g. - ``vllm:num_requests_running``. - label_filter: Optional labels a sample must have to be counted, e.g. to - select one model on a multi-model server. - - Raises: - ValueError: if ``tag`` is unset; call :meth:`resolve_tag` first. - """ - estimator = CompletionRateCapacityEstimator( - multiplier=self.multiplier, - window_secs=self.signal_window_secs, - min_requests=self.min_requests, - warmup_timeout_secs=self.warmup_timeout_secs) - - def capacity() -> Optional[float]: - """One poll: returns the current capacity, or None to skip the tick.""" - port = metrics_port_fn() - if port is None: - return None # The server is not up. - if self.load_signal_fn is not None: - return self.load_signal_fn(port) - text = scrape_prometheus_metrics(port) - if text is None: - return None - completed = parse_prometheus_metric(text, completed_metric, label_filter) - running = parse_prometheus_metric(text, running_metric, label_filter) - if completed is None or running is None: - return None - return estimator.update(completed, running) - - return self.make_publisher(capacity) - - def make_publisher( - self, signal_fn: Callable[[], - Optional[float]]) -> NumLoadedModelsPublisher: - """Builds a :class:`NumLoadedModelsPublisher` for ``signal_fn``. - - Raises: - ValueError: if ``tag`` is unset; call :meth:`resolve_tag` first. - """ - if not self.tag: - raise ValueError( - 'AutoscalingMetricConfig.tag is unset; call resolve_tag() first.') - return NumLoadedModelsPublisher( - tag=self.tag, - signal_fn=signal_fn, - poll_interval_secs=self.poll_interval_secs, - ) + Raises: + ValueError: if no metric tag can be derived. + """ + config = config or AutoscalingMetricConfig() + return NumLoadedModelsPublisher( + tag=_metric_tag(config, model_identifier, model_name), + signal_fn=signal_fn, + poll_interval_secs=config.poll_interval_secs) + + +def make_model_server_publisher( + *, + metrics_port_fn: Callable[[], Optional[int]], + completed_metric: str, + running_metric: str, + model_identifier: Optional[str], + model_name: str = '', + config: Optional[AutoscalingMetricConfig] = None, + label_filter: Optional[Mapping[str, str]] = None, +) -> NumLoadedModelsPublisher: + """Builds a publisher of a model server's capacity. + + Every poll reads the server's Prometheus ``/metrics`` endpoint, parses the + two metrics, feeds them to a :class:`CompletionRateCapacityEstimator` and + publishes the estimate. If ``config.load_signal_fn`` is set, it replaces the + scrape and the estimate. + + Args: + metrics_port_fn: Returns the port serving ``/metrics`` while the server is + up, or ``None`` while it is not (the publisher then skips the tick). + completed_metric: Name of the server's counter of completed requests, e.g. + ``vllm:request_success_total``. + running_metric: Name of the server's gauge of in-flight requests, e.g. + ``vllm:num_requests_running``. + model_identifier: The identifier RunInference reported through + ``ModelHandler.set_model_identifier``, or ``None`` outside RunInference. + model_name: Fallback for the metric tag outside RunInference. + config: Knobs; ``None`` uses the defaults. + label_filter: Optional labels a sample must have to be counted, e.g. to + select one model on a multi-model server. + + Raises: + ValueError: if no metric tag can be derived. + """ + config = config or AutoscalingMetricConfig() + estimator = CompletionRateCapacityEstimator( + multiplier=config.multiplier, + window_secs=config.signal_window_secs, + min_requests=config.min_requests, + warmup_timeout_secs=config.warmup_timeout_secs) + + def read_capacity() -> Optional[float]: + """Reads the server once; returns its capacity, or None to skip.""" + port = metrics_port_fn() + if port is None: + return None # The server is not up. + if config.load_signal_fn is not None: + return config.load_signal_fn(port) + text = scrape_prometheus_metrics(port) + if text is None: + return None + completed = parse_prometheus_metric(text, completed_metric, label_filter) + running = parse_prometheus_metric(text, running_metric, label_filter) + if completed is None or running is None: + return None + return estimator.update(completed, running) + + return make_publisher( + signal_fn=read_capacity, + model_identifier=model_identifier, + model_name=model_name, + config=config) diff --git a/sdks/python/apache_beam/ml/inference/autoscaling_metrics_test.py b/sdks/python/apache_beam/ml/inference/autoscaling_metrics_test.py index 5d02156dd5fe..b95e6f324d05 100644 --- a/sdks/python/apache_beam/ml/inference/autoscaling_metrics_test.py +++ b/sdks/python/apache_beam/ml/inference/autoscaling_metrics_test.py @@ -219,42 +219,52 @@ def test_capacity_estimator_defaults_to_wall_clock(self): self.assertIsNone(estimator.update(0.0, 1.0)) self.assertAlmostEqual(estimator.update(10.0, 1.0), 2.0) - def test_autoscaling_metric_config_resolve_tag(self): - cfg = autoscaling_metrics.AutoscalingMetricConfig( - poll_interval_secs=7.0, - multiplier=2.0, - signal_window_secs=30.0, - min_requests=5, - warmup_timeout_secs=15.0) + def test_make_publisher_metric_tag(self): + def tag(model_identifier, model_name='', **config): + publisher = autoscaling_metrics.make_publisher( + signal_fn=lambda: None, + model_identifier=model_identifier, + model_name=model_name, + config=autoscaling_metrics.AutoscalingMetricConfig(**config)) + return publisher._metric_name + # The RunInference identifier is used verbatim (it has to match the step # annotation), even if it contains characters sanitize would change. self.assertEqual( - cfg.resolve_tag('my-step/id', 'org/model:v1').tag, 'my-step/id') - # Outside RunInference: the sanitized fallback name. - self.assertEqual(cfg.resolve_tag(None, 'org/model:v1').tag, 'org_model_v1') + tag('my-step/id', 'org/model:v1'), 'num_loaded_models_my-step/id') + # Outside RunInference: the sanitized model name. + self.assertEqual( + tag(None, 'org/model:v1'), 'num_loaded_models_org_model_v1') # An explicit tag always wins. self.assertEqual( - autoscaling_metrics.AutoscalingMetricConfig(tag='custom').resolve_tag( - 'step-id', 'model').tag, - 'custom') + tag('step-id', 'model', tag='custom'), 'num_loaded_models_custom') with self.assertRaises(ValueError): - autoscaling_metrics.AutoscalingMetricConfig().resolve_tag(None, '') - with self.assertRaises(ValueError): - autoscaling_metrics.AutoscalingMetricConfig().make_publisher(lambda: None) + tag(None, '') - publisher = cfg.resolve_tag('step-id').make_publisher(lambda: None) - self.assertEqual(publisher._metric_name, 'num_loaded_models_step-id') + def test_make_publisher_uses_config_and_defaults(self): + publisher = autoscaling_metrics.make_publisher( + signal_fn=lambda: None, + model_identifier='step-id', + config=autoscaling_metrics.AutoscalingMetricConfig( + poll_interval_secs=7.0)) self.assertEqual(publisher._poll_interval_secs, 7.0) + # No config -> defaults. + default = autoscaling_metrics.make_publisher( + signal_fn=lambda: None, model_identifier='step-id') + self.assertEqual( + default._poll_interval_secs, + autoscaling_metrics.DEFAULT_POLL_INTERVAL_SECS) def test_make_model_server_publisher(self): """The handler supplies metric names and a port; the rest is packaged.""" - config = autoscaling_metrics.AutoscalingMetricConfig( - multiplier=2.0, poll_interval_secs=7.0).resolve_tag('step-id') port = [None] - publisher = config.make_model_server_publisher( + publisher = autoscaling_metrics.make_model_server_publisher( metrics_port_fn=lambda: port[0], completed_metric='done_total', - running_metric='running') + running_metric='running', + model_identifier='step-id', + config=autoscaling_metrics.AutoscalingMetricConfig( + multiplier=2.0, poll_interval_secs=7.0)) self.assertEqual(publisher._metric_name, 'num_loaded_models_step-id') self.assertEqual(publisher._poll_interval_secs, 7.0) signal = publisher._signal_fn # What the publisher's thread polls. @@ -284,12 +294,12 @@ def fake_urlopen(url, timeout=5.0): self.assertAlmostEqual(signal(), 8.0) def test_make_model_server_publisher_skips_unreadable_scrapes(self): - config = autoscaling_metrics.AutoscalingMetricConfig( - min_requests=0).resolve_tag('step-id') - publisher = config.make_model_server_publisher( + publisher = autoscaling_metrics.make_model_server_publisher( metrics_port_fn=lambda: 8000, completed_metric='done_total', running_metric='running', + model_identifier='step-id', + config=autoscaling_metrics.AutoscalingMetricConfig(min_requests=0), label_filter={'model': 'a'}) signal = publisher._signal_fn scrapes = [ @@ -314,18 +324,16 @@ def test_make_model_server_publisher_skips_unreadable_scrapes(self): self.assertAlmostEqual(signal(), 1.0) def test_make_model_server_publisher_load_signal_fn_override(self): - config = autoscaling_metrics.AutoscalingMetricConfig( - load_signal_fn=lambda port: port / 1000.0).resolve_tag('step-id') - publisher = config.make_model_server_publisher( + publisher = autoscaling_metrics.make_model_server_publisher( metrics_port_fn=lambda: 8000, completed_metric='done_total', - running_metric='running') + running_metric='running', + model_identifier='step-id', + config=autoscaling_metrics.AutoscalingMetricConfig( + load_signal_fn=lambda port: port / 1000.0)) with mock.patch('urllib.request.urlopen') as urlopen: self.assertEqual(publisher._signal_fn(), 8.0) urlopen.assert_not_called() - with self.assertRaises(ValueError): - autoscaling_metrics.AutoscalingMetricConfig().make_model_server_publisher( - lambda: 8000, 'done_total', 'running') if __name__ == '__main__': From 746051e600b260f35e63effaf0f5f752775d201d Mon Sep 17 00:00:00 2001 From: AMOOOMA Date: Tue, 3 Nov 2026 01:02:44 +0000 Subject: [PATCH 07/14] Let NumLoadedModelsPublisher read the server itself Remove the make_publisher and make_model_server_publisher wrappers. A handler now constructs NumLoadedModelsPublisher directly with its server's metric names, a metrics_port_fn, the model identifier and an optional AutoscalingMetricConfig; the publisher resolves the metric tag, owns its CompletionRateCapacityEstimator, and on each poll reads the port, scrapes /metrics (or calls config.load_signal_fn), updates the estimate and publishes it. Behavior is unchanged. --- .../ml/inference/autoscaling_metrics.py | 347 +++++++----------- .../ml/inference/autoscaling_metrics_test.py | 219 ++++++----- 2 files changed, 250 insertions(+), 316 deletions(-) diff --git a/sdks/python/apache_beam/ml/inference/autoscaling_metrics.py b/sdks/python/apache_beam/ml/inference/autoscaling_metrics.py index 82e83b1393d3..df332d829cdd 100644 --- a/sdks/python/apache_beam/ml/inference/autoscaling_metrics.py +++ b/sdks/python/apache_beam/ml/inference/autoscaling_metrics.py @@ -26,29 +26,26 @@ publishes an estimate of the worker's capacity under that metric instead. Nothing here depends on a particular model server. A handler for a model server -that exposes a completed-requests counter and an in-flight gauge makes one call, -supplying only what it knows: the two metric names, how to find the server's -port, and the RunInference model identifier:: +that exposes a completed-requests counter and an in-flight gauge creates a +:class:`NumLoadedModelsPublisher`, supplying only what it knows: the two metric +names, how to find the server's port, and the RunInference model identifier:: - publisher = make_model_server_publisher( - config=autoscaling_config, # An AutoscalingMetricConfig, or None. - model_identifier=model_identifier, - model_name=model_name, + publisher = 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') + 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. -Handlers whose server needs a different capacity signal pass it to -:func:`make_publisher` instead. - -Inside, each component has one job: +The pieces: +* :class:`NumLoadedModelsPublisher` reads the server, estimates its capacity + and publishes it from a background thread. +* :class:`CompletionRateCapacityEstimator` is the estimate itself: pure + computation on the counter and the gauge, no I/O. * :class:`AutoscalingMetricConfig` holds the user-facing knobs. Plain data. -* :class:`CompletionRateCapacityEstimator` computes a worker's capacity from - the counter and the gauge. It does no I/O. -* :class:`NumLoadedModelsPublisher` publishes: a background thread polls a - capacity signal and writes each value to ``num_loaded_models_``. * :func:`scrape_prometheus_metrics` and :func:`parse_prometheus_metric` read a server's Prometheus ``/metrics`` endpoint. """ @@ -73,8 +70,6 @@ 'AutoscalingMetricConfig', 'CompletionRateCapacityEstimator', 'NumLoadedModelsPublisher', - 'make_model_server_publisher', - 'make_publisher', 'parse_prometheus_metric', 'sanitize_metric_tag', 'scrape_prometheus_metrics', @@ -200,112 +195,6 @@ def _nearest_rank_percentile(values: list[float], q: float) -> float: return ordered[min(max(rank - 1, 0), len(ordered) - 1)] -class NumLoadedModelsPublisher: - """Publishes a capacity signal as the ``num_loaded_models_`` metric. - - A daemon thread calls ``signal_fn`` every ``poll_interval_secs`` and writes - each value, rounded to an int, to the process-wide distribution the runner - reads. The publisher does not transform values: any smoothing belongs in - ``signal_fn`` (e.g. :class:`CompletionRateCapacityEstimator`, which already - smooths and never returns less than its ``min_value``). - - A ``signal_fn`` result of ``None`` (e.g. the server is restarting) skips the - tick and leaves the last published value in place. The publisher never - injects a 0 of its own, on start, on stop or on a skipped tick: a 0 tells the - runner the worker has no capacity and can trigger a spurious downscale. Only - one publisher per metric name is active in a process; starting a new one - (e.g. after a model reload) stops the previous one. - - Args: - tag: Metric tag. It must equal the RunInference model identifier the runner - looks up; :func:`make_publisher` derives it. - signal_fn: Returns the current capacity, or ``None`` to skip the tick. - poll_interval_secs: Seconds between polls. - - Raises: - ValueError: if ``tag`` is empty. - """ - # The publisher owning each metric name in this process, so a reloaded - # model's publisher can supersede the stale one. - _active: dict[str, 'NumLoadedModelsPublisher'] = {} - _active_lock = threading.Lock() - - def __init__( - self, - tag: str, - signal_fn: Callable[[], Optional[float]], - poll_interval_secs: float = DEFAULT_POLL_INTERVAL_SECS, - ): - if not tag: - raise ValueError('num_loaded_models tag must be a non-empty string.') - self._metric_name = NUM_LOADED_MODELS_METRIC_PREFIX + tag - self._signal_fn = signal_fn - self._poll_interval_secs = poll_interval_secs - self._stop_event = threading.Event() - self._thread: Optional[threading.Thread] = None - - def start(self) -> None: - """Starts the background polling thread. - - The metric stays absent until the first sample is published. - """ - if self._thread is not None: - return - with NumLoadedModelsPublisher._active_lock: - previous = NumLoadedModelsPublisher._active.get(self._metric_name) - NumLoadedModelsPublisher._active[self._metric_name] = self - if previous is not None and previous is not self: - logging.info( - 'Replacing the active publisher for %s (model reload or server ' - 'restart).', - self._metric_name) - previous.stop() - self._thread = threading.Thread( - target=self._run, name='num-loaded-models-publisher', daemon=True) - self._thread.start() - - def stop(self) -> None: - """Stops polling. - - Publishes no final 0: when a worker is drained during a downscale, a 0 - would shrink the pool's total capacity and could drive a feedback loop of - further downscales. The drained worker's value ages out instead. - """ - self._stop_event.set() - with NumLoadedModelsPublisher._active_lock: - if NumLoadedModelsPublisher._active.get(self._metric_name) is self: - del NumLoadedModelsPublisher._active[self._metric_name] - thread = self._thread - if thread is not None: - thread.join(timeout=self._poll_interval_secs + 5) - self._thread = None - - def _run(self) -> None: - while not self._stop_event.is_set(): - try: - value = self._signal_fn() - if value is not None: - self._publish(value) - except Exception: # pylint: disable=broad-except - logging.exception( - 'Failed to publish autoscaling metric %s; will retry.', - self._metric_name) - self._stop_event.wait(self._poll_interval_secs) - - def _publish(self, value: float) -> None: - published = int(round(value)) - # The SDK harness reports process-wide distributions without resetting - # them, and Dataflow reads their cumulative mean (sum / count). Resetting - # before each update keeps the cell at (sum=published, count=1), so the - # runner sees the latest value rather than an average since worker start. - cell = MetricsEnvironment.process_wide_container().get_distribution( - MetricName(MODEL_MANAGER_METRICS_NAMESPACE, self._metric_name)) - current = cell.get_cumulative() - if current.count != 1 or current.sum != published: - cell.reset() - cell.update(published) - - class _Interval(NamedTuple): """One sampling interval in the estimator's trailing window.""" end_ts: float @@ -462,71 +351,32 @@ class AutoscalingMetricConfig: tag: Optional[str] = None -def _metric_tag( - config: AutoscalingMetricConfig, - model_identifier: Optional[str], - model_name: str) -> str: - """Picks the ``num_loaded_models`` tag. - - Precedence: ``config.tag``; then ``model_identifier``, used verbatim because - it must equal the identifier RunInference annotates the step with (see - ``ModelHandler.set_model_identifier``); then ``model_name`` passed through - :func:`sanitize_metric_tag`, for handlers used outside RunInference. - """ - tag = config.tag or model_identifier or sanitize_metric_tag(model_name) - if not tag: - raise ValueError( - 'Cannot derive a num_loaded_models tag: pass a model_identifier, a ' - 'model_name, or set AutoscalingMetricConfig.tag.') - return tag - - -def make_publisher( - *, - signal_fn: Callable[[], Optional[float]], - model_identifier: Optional[str], - model_name: str = '', - config: Optional[AutoscalingMetricConfig] = None, -) -> NumLoadedModelsPublisher: - """Builds a publisher for a handler's own capacity signal. - - Args: - signal_fn: Returns the current capacity, or ``None`` to skip the tick. - model_identifier: The identifier RunInference reported through - ``ModelHandler.set_model_identifier``, or ``None`` outside RunInference. - model_name: Fallback for the metric tag outside RunInference. - config: Knobs; ``None`` uses the defaults. - - Raises: - ValueError: if no metric tag can be derived. - """ - config = config or AutoscalingMetricConfig() - return NumLoadedModelsPublisher( - tag=_metric_tag(config, model_identifier, model_name), - signal_fn=signal_fn, - poll_interval_secs=config.poll_interval_secs) - - -def make_model_server_publisher( - *, - metrics_port_fn: Callable[[], Optional[int]], - completed_metric: str, - running_metric: str, - model_identifier: Optional[str], - model_name: str = '', - config: Optional[AutoscalingMetricConfig] = None, - label_filter: Optional[Mapping[str, str]] = None, -) -> NumLoadedModelsPublisher: - """Builds a publisher of a model server's capacity. - - Every poll reads the server's Prometheus ``/metrics`` endpoint, parses the - two metrics, feeds them to a :class:`CompletionRateCapacityEstimator` and - publishes the estimate. If ``config.load_signal_fn`` is set, it replaces the - scrape and the estimate. +class NumLoadedModelsPublisher: + """Publishes a model server's capacity as ``num_loaded_models_``. + + Every ``poll_interval_secs`` a daemon thread reads the server's Prometheus + ``/metrics`` endpoint, feeds the two metrics to a + :class:`CompletionRateCapacityEstimator`, and writes the estimate, rounded + to an int, to the process-wide distribution the runner reads. If + ``config.load_signal_fn`` is set, it replaces the scrape and the estimate. + + A poll is skipped, leaving the last published value in place, while the + server is down, when the scrape fails, or while the estimator has no + estimate yet. The publisher never injects a 0 of its own, on start, on stop + or on a skipped poll: a 0 tells the runner the worker has no capacity and can + trigger a spurious downscale. Only one publisher per metric name is active in + a process; starting a new one (e.g. after a model reload) stops the previous + one. + + The metric tag is ``config.tag`` if set; else ``model_identifier``, used + verbatim because it must equal the identifier RunInference annotates the + step with (see ``ModelHandler.set_model_identifier``); else ``model_name`` + passed through :func:`sanitize_metric_tag`, for handlers used outside + RunInference. Args: metrics_port_fn: Returns the port serving ``/metrics`` while the server is - up, or ``None`` while it is not (the publisher then skips the tick). + up, or ``None`` while it is not. completed_metric: Name of the server's counter of completed requests, e.g. ``vllm:request_success_total``. running_metric: Name of the server's gauge of in-flight requests, e.g. @@ -541,31 +391,118 @@ def make_model_server_publisher( Raises: ValueError: if no metric tag can be derived. """ - config = config or AutoscalingMetricConfig() - estimator = CompletionRateCapacityEstimator( - multiplier=config.multiplier, - window_secs=config.signal_window_secs, - min_requests=config.min_requests, - warmup_timeout_secs=config.warmup_timeout_secs) - - def read_capacity() -> Optional[float]: + # The publisher owning each metric name in this process, so a reloaded + # model's publisher can supersede the stale one. + _active: dict[str, 'NumLoadedModelsPublisher'] = {} + _active_lock = threading.Lock() + + def __init__( + self, + *, + metrics_port_fn: Callable[[], Optional[int]], + completed_metric: str, + running_metric: str, + model_identifier: Optional[str], + model_name: str = '', + config: Optional[AutoscalingMetricConfig] = None, + label_filter: Optional[Mapping[str, str]] = None, + ): + config = config or AutoscalingMetricConfig() + tag = config.tag or model_identifier or sanitize_metric_tag(model_name) + if not tag: + raise ValueError( + 'Cannot derive a num_loaded_models tag: pass a model_identifier, a ' + 'model_name, or set AutoscalingMetricConfig.tag.') + self._metric_name = NUM_LOADED_MODELS_METRIC_PREFIX + tag + self._metrics_port_fn = metrics_port_fn + self._completed_metric = completed_metric + self._running_metric = running_metric + self._label_filter = dict(label_filter) if label_filter else None + self._load_signal_fn = config.load_signal_fn + self._poll_interval_secs = config.poll_interval_secs + self._estimator = CompletionRateCapacityEstimator( + multiplier=config.multiplier, + window_secs=config.signal_window_secs, + min_requests=config.min_requests, + warmup_timeout_secs=config.warmup_timeout_secs) + self._stop_event = threading.Event() + self._thread: Optional[threading.Thread] = None + + def start(self) -> None: + """Starts the background polling thread. + + The metric stays absent until the first sample is published. + """ + if self._thread is not None: + return + with NumLoadedModelsPublisher._active_lock: + previous = NumLoadedModelsPublisher._active.get(self._metric_name) + NumLoadedModelsPublisher._active[self._metric_name] = self + if previous is not None and previous is not self: + logging.info( + 'Replacing the active publisher for %s (model reload or server ' + 'restart).', + self._metric_name) + previous.stop() + self._thread = threading.Thread( + target=self._run, name='num-loaded-models-publisher', daemon=True) + self._thread.start() + + def stop(self) -> None: + """Stops polling. + + Publishes no final 0: when a worker is drained during a downscale, a 0 + would shrink the pool's total capacity and could drive a feedback loop of + further downscales. The drained worker's value ages out instead. + """ + self._stop_event.set() + with NumLoadedModelsPublisher._active_lock: + if NumLoadedModelsPublisher._active.get(self._metric_name) is self: + del NumLoadedModelsPublisher._active[self._metric_name] + thread = self._thread + if thread is not None: + thread.join(timeout=self._poll_interval_secs + 5) + self._thread = None + + def _run(self) -> None: + while not self._stop_event.is_set(): + try: + capacity = self._read_capacity() + if capacity is not None: + self._publish(capacity) + except Exception: # pylint: disable=broad-except + logging.exception( + 'Failed to publish autoscaling metric %s; will retry.', + self._metric_name) + self._stop_event.wait(self._poll_interval_secs) + + def _read_capacity(self) -> Optional[float]: """Reads the server once; returns its capacity, or None to skip.""" - port = metrics_port_fn() + port = self._metrics_port_fn() if port is None: return None # The server is not up. - if config.load_signal_fn is not None: - return config.load_signal_fn(port) + if self._load_signal_fn is not None: + return self._load_signal_fn(port) text = scrape_prometheus_metrics(port) if text is None: return None - completed = parse_prometheus_metric(text, completed_metric, label_filter) - running = parse_prometheus_metric(text, running_metric, label_filter) + completed = parse_prometheus_metric( + text, self._completed_metric, self._label_filter) + running = parse_prometheus_metric( + text, self._running_metric, self._label_filter) if completed is None or running is None: return None - return estimator.update(completed, running) + return self._estimator.update(completed, running) - return make_publisher( - signal_fn=read_capacity, - model_identifier=model_identifier, - model_name=model_name, - config=config) + def _publish(self, value: float) -> None: + published = int(round(value)) + # The SDK harness reports process-wide distributions without resetting + # them, and Dataflow reads their cumulative mean (sum / count). Resetting + # before each update keeps the cell at (sum=published, count=1), so the + # runner sees the latest value rather than an average since worker start. + cell = MetricsEnvironment.process_wide_container().get_distribution( + MetricName(MODEL_MANAGER_METRICS_NAMESPACE, self._metric_name)) + current = cell.get_cumulative() + if current.count != 1 or current.sum != published: + cell.reset() + cell.update(published) diff --git a/sdks/python/apache_beam/ml/inference/autoscaling_metrics_test.py b/sdks/python/apache_beam/ml/inference/autoscaling_metrics_test.py index b95e6f324d05..39c403b3dd27 100644 --- a/sdks/python/apache_beam/ml/inference/autoscaling_metrics_test.py +++ b/sdks/python/apache_beam/ml/inference/autoscaling_metrics_test.py @@ -21,6 +21,16 @@ from apache_beam.ml.inference import autoscaling_metrics +def _publisher(model_identifier='step-id', port=None, **kwargs): + """A publisher for a fake server exposing ``done_total`` and ``running``.""" + return autoscaling_metrics.NumLoadedModelsPublisher( + metrics_port_fn=lambda: port, + completed_metric='done_total', + running_metric='running', + model_identifier=model_identifier, + **kwargs) + + class AutoscalingMetricsTest(unittest.TestCase): def test_sanitize_metric_tag(self): self.assertEqual( @@ -73,8 +83,74 @@ def test_parse_prometheus_metric_exposition_details(self): self.assertIsNone( parse(prom_text, 'labeled', label_filter={'path': 'other'})) + def test_capacity_estimator(self): + estimator = autoscaling_metrics.CompletionRateCapacityEstimator( + multiplier=2.0, min_value=1, window_secs=60.0, running_percentile=90.0) + # First sample establishes baseline -> None. + self.assertIsNone(estimator.update(10.0, 4.0, now=100.0)) + # delta=20 over dt=5s with running=4 -> mu = 20 / (4 * 5) = 1.0, + # running_ref = 4.0, multiplier = 2.0 -> 8.0. + self.assertAlmostEqual(estimator.update(30.0, 4.0, now=105.0), 8.0) + + def test_capacity_estimator_skips_counter_reset_and_clock_skew(self): + estimator = autoscaling_metrics.CompletionRateCapacityEstimator( + window_secs=60.0, min_requests=0) + self.assertIsNone(estimator.update(50.0, 2.0, now=100.0)) + # Counter went down (server restart) -> skipped, new baseline. + self.assertIsNone(estimator.update(5.0, 2.0, now=105.0)) + # Non-increasing timestamp -> skipped. + self.assertIsNone(estimator.update(9.0, 2.0, now=105.0)) + # delta=4 over 5s with running=2 -> mu=0.4, p90=2 -> 0.8 -> min_value 1. + self.assertAlmostEqual(estimator.update(13.0, 2.0, now=110.0), 1.0) + + def test_capacity_estimator_warmup_min_requests_and_timeout(self): + # Requires >= 3 completed requests before the 10s warmup timeout. + estimator = autoscaling_metrics.CompletionRateCapacityEstimator( + min_requests=3, warmup_timeout_secs=10.0) + self.assertIsNone(estimator.update(0.0, 2.0, now=100.0)) # Baseline. + # 1 completion (< 3) at elapsed=5s (< 10s) -> None. + self.assertIsNone(estimator.update(1.0, 2.0, now=105.0)) + # 4 completions (>= 3) at elapsed=8s -> publishes + # (4 / (2*5 + 2*3)) * 2 = 0.5 -> clamped to min_value=1.0. + self.assertAlmostEqual(estimator.update(4.0, 2.0, now=108.0), 1.0) + + slow = autoscaling_metrics.CompletionRateCapacityEstimator( + min_requests=3, warmup_timeout_secs=10.0) + self.assertIsNone(slow.update(0.0, 1.0, now=100.0)) # Baseline. + # 1 completion (< 3) at elapsed=5s -> None. + self.assertIsNone(slow.update(1.0, 1.0, now=105.0)) + # 2 completions (< 3) but elapsed=10s (>= timeout) -> publishes. + self.assertAlmostEqual(slow.update(2.0, 1.0, now=110.0), 1.0) + + def test_capacity_estimator_skips_idle_ticks_without_diluting_p90(self): + estimator = autoscaling_metrics.CompletionRateCapacityEstimator( + multiplier=1.0, + min_value=1, + window_secs=60.0, + running_percentile=90.0, + min_requests=1) + self.assertIsNone(estimator.update(0.0, 5.0, now=100.0)) # Baseline. + # 25 completions over 5s with running=5 -> mu=1.0, p90=5 -> 5.0. + self.assertAlmostEqual(estimator.update(25.0, 5.0, now=105.0), 5.0) + # Three completely idle ticks (delta=0, running=0) -> None, without adding + # r=0 samples that would dilute p90 concurrency. + for t in (110.0, 115.0, 120.0): + self.assertIsNone(estimator.update(25.0, 0.0, now=t)) + # 5 completions that already finished before the scrape (delta=5, + # running=0): total_completed=30, total_slot_secs=25 -> mu=1.2, p90=5 -> 6. + self.assertAlmostEqual(estimator.update(30.0, 0.0, now=125.0), 6.0) + + def test_capacity_estimator_defaults_to_wall_clock(self): + estimator = autoscaling_metrics.CompletionRateCapacityEstimator( + min_requests=0) + with mock.patch.object(autoscaling_metrics.time, + 'time', + side_effect=[100.0, 105.0]): + self.assertIsNone(estimator.update(0.0, 1.0)) + self.assertAlmostEqual(estimator.update(10.0, 1.0), 2.0) + def test_publisher_publishes_signal_values_as_is(self): - """Each signal value is published rounded; None skips the tick.""" + """Each capacity is published rounded; None skips the tick.""" cell = ( autoscaling_metrics.MetricsEnvironment.process_wide_container(). get_distribution( @@ -86,7 +162,8 @@ def test_publisher_publishes_signal_values_as_is(self): published = [] signal_values = [3.4, None, 7.6] - def signal_fn(): + def load_signal_fn(port): + del port if published: # Record what the previous tick left in the cell. published[-1] = cell.get_cumulative().sum @@ -96,8 +173,11 @@ def signal_fn(): published.append(None) return signal_values.pop(0) - pub = autoscaling_metrics.NumLoadedModelsPublisher( - tag='sink_test', signal_fn=signal_fn, poll_interval_secs=0) + pub = _publisher( + 'sink_test', + port=8000, + config=autoscaling_metrics.AutoscalingMetricConfig( + poll_interval_secs=0, load_signal_fn=load_signal_fn)) # Run the polling loop on this thread so the test is deterministic. pub._run() # pylint: disable=protected-access # 3.4 -> 3; None leaves 3 in place; 7.6 -> 8. @@ -105,8 +185,7 @@ def signal_fn(): self.assertEqual(cell.get_cumulative().count, 1) def test_publish_resets_process_wide_distribution_cell(self): - pub = autoscaling_metrics.NumLoadedModelsPublisher( - tag='reset_test', signal_fn=lambda: 1.0) + pub = _publisher('reset_test') cell = ( autoscaling_metrics.MetricsEnvironment.process_wide_container(). get_distribution( @@ -126,12 +205,10 @@ def test_publish_resets_process_wide_distribution_cell(self): def test_publisher_for_same_tag_supersedes_previous(self): """A reloaded model's publisher takes over its tag from the stale one.""" active = autoscaling_metrics.NumLoadedModelsPublisher._active - first = autoscaling_metrics.NumLoadedModelsPublisher( - tag='reload', signal_fn=lambda: None, poll_interval_secs=0.01) - second = autoscaling_metrics.NumLoadedModelsPublisher( - tag='reload', signal_fn=lambda: None, poll_interval_secs=0.01) - other = autoscaling_metrics.NumLoadedModelsPublisher( - tag='other', signal_fn=lambda: None, poll_interval_secs=0.01) + fast = autoscaling_metrics.AutoscalingMetricConfig(poll_interval_secs=0.01) + first = _publisher('reload', config=fast) + second = _publisher('reload', config=fast) + other = _publisher('other', config=fast) try: first.start() other.start() @@ -153,77 +230,10 @@ def test_publisher_for_same_tag_supersedes_previous(self): other.stop() self.assertNotIn('num_loaded_models_reload', active) - def test_capacity_estimator(self): - estimator = autoscaling_metrics.CompletionRateCapacityEstimator( - multiplier=2.0, min_value=1, window_secs=60.0, running_percentile=90.0) - # First sample establishes baseline -> None. - self.assertIsNone(estimator.update(10.0, 4.0, now=100.0)) - # delta=20 over dt=5s with running=4 -> mu = 20 / (4 * 5) = 1.0, - # running_ref = 4.0, multiplier = 2.0 -> 8.0. - self.assertAlmostEqual(estimator.update(30.0, 4.0, now=105.0), 8.0) - - def test_capacity_estimator_skips_counter_reset_and_clock_skew(self): - estimator = autoscaling_metrics.CompletionRateCapacityEstimator( - window_secs=60.0, min_requests=0) - self.assertIsNone(estimator.update(50.0, 2.0, now=100.0)) - # Counter went down (server restart) -> skipped, new baseline. - self.assertIsNone(estimator.update(5.0, 2.0, now=105.0)) - # Non-increasing timestamp -> skipped. - self.assertIsNone(estimator.update(9.0, 2.0, now=105.0)) - # delta=4 over 5s with running=2 -> mu=0.4, p90=2 -> 0.8 -> min_value 1. - self.assertAlmostEqual(estimator.update(13.0, 2.0, now=110.0), 1.0) - - def test_capacity_estimator_warmup_min_requests_and_timeout(self): - # Requires >= 3 completed requests before the 10s warmup timeout. - estimator = autoscaling_metrics.CompletionRateCapacityEstimator( - min_requests=3, warmup_timeout_secs=10.0) - self.assertIsNone(estimator.update(0.0, 2.0, now=100.0)) # Baseline. - # 1 completion (< 3) at elapsed=5s (< 10s) -> None. - self.assertIsNone(estimator.update(1.0, 2.0, now=105.0)) - # 4 completions (>= 3) at elapsed=8s -> publishes - # (4 / (2*5 + 2*3)) * 2 = 0.5 -> clamped to min_value=1.0. - self.assertAlmostEqual(estimator.update(4.0, 2.0, now=108.0), 1.0) - - slow = autoscaling_metrics.CompletionRateCapacityEstimator( - min_requests=3, warmup_timeout_secs=10.0) - self.assertIsNone(slow.update(0.0, 1.0, now=100.0)) # Baseline. - # 1 completion (< 3) at elapsed=5s -> None. - self.assertIsNone(slow.update(1.0, 1.0, now=105.0)) - # 2 completions (< 3) but elapsed=10s (>= timeout) -> publishes. - self.assertAlmostEqual(slow.update(2.0, 1.0, now=110.0), 1.0) - - def test_capacity_estimator_skips_idle_ticks_without_diluting_p90(self): - estimator = autoscaling_metrics.CompletionRateCapacityEstimator( - multiplier=1.0, - min_value=1, - window_secs=60.0, - running_percentile=90.0, - min_requests=1) - self.assertIsNone(estimator.update(0.0, 5.0, now=100.0)) # Baseline. - # 25 completions over 5s with running=5 -> mu=1.0, p90=5 -> 5.0. - self.assertAlmostEqual(estimator.update(25.0, 5.0, now=105.0), 5.0) - # Three completely idle ticks (delta=0, running=0) -> None, without adding - # r=0 samples that would dilute p90 concurrency. - for t in (110.0, 115.0, 120.0): - self.assertIsNone(estimator.update(25.0, 0.0, now=t)) - # 5 completions that already finished before the scrape (delta=5, - # running=0): total_completed=30, total_slot_secs=25 -> mu=1.2, p90=5 -> 6. - self.assertAlmostEqual(estimator.update(30.0, 0.0, now=125.0), 6.0) - - def test_capacity_estimator_defaults_to_wall_clock(self): - estimator = autoscaling_metrics.CompletionRateCapacityEstimator( - min_requests=0) - with mock.patch.object(autoscaling_metrics.time, - 'time', - side_effect=[100.0, 105.0]): - self.assertIsNone(estimator.update(0.0, 1.0)) - self.assertAlmostEqual(estimator.update(10.0, 1.0), 2.0) - - def test_make_publisher_metric_tag(self): + def test_publisher_metric_tag(self): def tag(model_identifier, model_name='', **config): - publisher = autoscaling_metrics.make_publisher( - signal_fn=lambda: None, - model_identifier=model_identifier, + publisher = _publisher( + model_identifier, model_name=model_name, config=autoscaling_metrics.AutoscalingMetricConfig(**config)) return publisher._metric_name @@ -241,33 +251,26 @@ def tag(model_identifier, model_name='', **config): with self.assertRaises(ValueError): tag(None, '') - def test_make_publisher_uses_config_and_defaults(self): - publisher = autoscaling_metrics.make_publisher( - signal_fn=lambda: None, - model_identifier='step-id', + def test_publisher_config_and_defaults(self): + publisher = _publisher( config=autoscaling_metrics.AutoscalingMetricConfig( poll_interval_secs=7.0)) self.assertEqual(publisher._poll_interval_secs, 7.0) # No config -> defaults. - default = autoscaling_metrics.make_publisher( - signal_fn=lambda: None, model_identifier='step-id') self.assertEqual( - default._poll_interval_secs, + _publisher()._poll_interval_secs, autoscaling_metrics.DEFAULT_POLL_INTERVAL_SECS) - def test_make_model_server_publisher(self): - """The handler supplies metric names and a port; the rest is packaged.""" + def test_publisher_reads_server_capacity(self): + """The handler supplies metric names and a port; the rest is built in.""" port = [None] - publisher = autoscaling_metrics.make_model_server_publisher( + publisher = autoscaling_metrics.NumLoadedModelsPublisher( metrics_port_fn=lambda: port[0], completed_metric='done_total', running_metric='running', model_identifier='step-id', - config=autoscaling_metrics.AutoscalingMetricConfig( - multiplier=2.0, poll_interval_secs=7.0)) - self.assertEqual(publisher._metric_name, 'num_loaded_models_step-id') - self.assertEqual(publisher._poll_interval_secs, 7.0) - signal = publisher._signal_fn # What the publisher's thread polls. + config=autoscaling_metrics.AutoscalingMetricConfig(multiplier=2.0)) + signal = publisher._read_capacity # What the publisher's thread polls. # Server not up yet -> skip the tick without scraping. with mock.patch('urllib.request.urlopen') as urlopen: @@ -293,15 +296,12 @@ def fake_urlopen(url, timeout=5.0): # delta=20 over 5s with running=4 -> mu=1, p90=4, multiplier 2 -> 8. self.assertAlmostEqual(signal(), 8.0) - def test_make_model_server_publisher_skips_unreadable_scrapes(self): - publisher = autoscaling_metrics.make_model_server_publisher( - metrics_port_fn=lambda: 8000, - completed_metric='done_total', - running_metric='running', - model_identifier='step-id', + def test_publisher_skips_unreadable_scrapes(self): + publisher = _publisher( + port=8000, config=autoscaling_metrics.AutoscalingMetricConfig(min_requests=0), label_filter={'model': 'a'}) - signal = publisher._signal_fn + signal = publisher._read_capacity scrapes = [ None, # Endpoint unreachable. 'done_total{model="a"} 1\n', # Gauge missing. @@ -323,16 +323,13 @@ def test_make_model_server_publisher_skips_unreadable_scrapes(self): # delta=10 over 10s with running=2 -> mu=0.5, p90=2 -> 1.0. self.assertAlmostEqual(signal(), 1.0) - def test_make_model_server_publisher_load_signal_fn_override(self): - publisher = autoscaling_metrics.make_model_server_publisher( - metrics_port_fn=lambda: 8000, - completed_metric='done_total', - running_metric='running', - model_identifier='step-id', + def test_publisher_load_signal_fn_override(self): + publisher = _publisher( + port=8000, config=autoscaling_metrics.AutoscalingMetricConfig( load_signal_fn=lambda port: port / 1000.0)) with mock.patch('urllib.request.urlopen') as urlopen: - self.assertEqual(publisher._signal_fn(), 8.0) + self.assertEqual(publisher._read_capacity(), 8.0) urlopen.assert_not_called() From 098fd2557018a601c16b31a83c62bceb17dce360 Mon Sep 17 00:00:00 2001 From: AMOOOMA Date: Tue, 3 Nov 2026 01:08:41 +0000 Subject: [PATCH 08/14] Describe the capacity signal in requests and state its units --- .../ml/inference/autoscaling_metrics.py | 18 +++++++++++++----- 1 file changed, 13 insertions(+), 5 deletions(-) diff --git a/sdks/python/apache_beam/ml/inference/autoscaling_metrics.py b/sdks/python/apache_beam/ml/inference/autoscaling_metrics.py index df332d829cdd..ea4f662bf210 100644 --- a/sdks/python/apache_beam/ml/inference/autoscaling_metrics.py +++ b/sdks/python/apache_beam/ml/inference/autoscaling_metrics.py @@ -18,8 +18,8 @@ """Capacity-based autoscaling signal for model handlers that run a model server. Dataflow's GPU autoscaling reads the process-wide distribution -``num_loaded_models_`` in the ``BeamML_ModelManager`` namespace as the -number of messages each worker can hold, and sizes the pool roughly as +``num_loaded_models_`` in the ``BeamML_ModelManager`` namespace as how +many requests each worker can handle, and sizes the pool roughly as ``ceil(backlog / value)``. ``ModelManager`` publishes the number of loaded model copies there. A handler that serves inference through an out-of-process model server (vLLM, Triton, ...) always has exactly one copy per worker, so it @@ -204,10 +204,11 @@ class _Interval(NamedTuple): class CompletionRateCapacityEstimator: - """Estimates how many messages a worker can absorb. + """Estimates how many requests a worker can handle. Feed it periodic samples of a model server's completed-requests counter and - in-flight gauge through :meth:`update`. Over a trailing ``window_secs`` window it computes:: + in-flight gauge through :meth:`update`. Over a trailing ``window_secs`` + window it computes:: mu = sum(completed) / sum(running * secs) running_ref = percentile(running, running_percentile) @@ -219,6 +220,12 @@ class CompletionRateCapacityEstimator: p90 rather than the max avoids latching onto rare backlog spikes. Their product is the worker's sustained throughput. + Units: ``mu`` is completions per slot-second, so ``capacity`` is in requests + per second. The runner reads the published value as how many requests one + worker can handle, so each worker is sized for about one second of its + sustained throughput; ``multiplier`` scales that, e.g. 2.0 sizes each worker + for two seconds' worth of requests. + Normalizing by concurrency keeps the estimate steady when load changes: right after a scaling event, when the runner rebalances work and a worker is briefly starved, its raw completion rate collapses, but ``mu`` and @@ -228,7 +235,8 @@ class CompletionRateCapacityEstimator: The estimator suits servers whose requests take long enough to be seen in flight by periodic sampling, such as LLM servers with latencies of a second or more. For servers with millisecond latencies the in-flight gauge reads - almost always 0; pass the publisher a different signal instead. + almost always 0; supply a different signal through + ``AutoscalingMetricConfig.load_signal_fn`` instead. The estimate is already smoothed; publish it as is. This class does no I/O. From 3a2d34538c11af09dfa5777a8c1e0af45f10d607 Mon Sep 17 00:00:00 2001 From: AMOOOMA Date: Tue, 3 Nov 2026 01:20:19 +0000 Subject: [PATCH 09/14] Add completion-rate aggregation modes to the capacity estimator CompletionRateCapacityEstimator and AutoscalingMetricConfig take an aggregation: 'concurrency' (the default and the existing estimate), 'window_mean', 'window_max' or 'lifetime_max'. The last three use the raw completion rate and ignore concurrency, so they also work for servers whose in-flight gauge is almost always 0. The default is unchanged. --- .../ml/inference/autoscaling_metrics.py | 99 ++++++++++++++----- .../ml/inference/autoscaling_metrics_test.py | 57 +++++++++++ 2 files changed, 132 insertions(+), 24 deletions(-) diff --git a/sdks/python/apache_beam/ml/inference/autoscaling_metrics.py b/sdks/python/apache_beam/ml/inference/autoscaling_metrics.py index ea4f662bf210..3eb81d22b75c 100644 --- a/sdks/python/apache_beam/ml/inference/autoscaling_metrics.py +++ b/sdks/python/apache_beam/ml/inference/autoscaling_metrics.py @@ -90,6 +90,13 @@ DEFAULT_MIN_REQUESTS = 3 DEFAULT_WARMUP_TIMEOUT_SECS = 10.0 +# How CompletionRateCapacityEstimator turns its samples into an estimate. +CONCURRENCY = 'concurrency' +WINDOW_MEAN = 'window_mean' +WINDOW_MAX = 'window_max' +LIFETIME_MAX = 'lifetime_max' +_AGGREGATIONS = (CONCURRENCY, WINDOW_MEAN, WINDOW_MAX, LIFETIME_MAX) + # One sample line of the Prometheus text format: # metric_name{label="value",...} value [timestamp] _PROMETHEUS_SAMPLE_RE = re.compile( @@ -207,21 +214,34 @@ class CompletionRateCapacityEstimator: """Estimates how many requests a worker can handle. Feed it periodic samples of a model server's completed-requests counter and - in-flight gauge through :meth:`update`. Over a trailing ``window_secs`` - window it computes:: + in-flight gauge through :meth:`update`. Each sample closes an interval; + ``aggregation`` picks how the intervals become an estimate. Every mode + returns ``max(min_value, value * multiplier)``. + + ``'concurrency'`` (the default) computes, over a trailing ``window_secs`` + window:: mu = sum(completed) / sum(running * secs) - running_ref = percentile(running, running_percentile) - capacity = max(min_value, mu * running_ref * multiplier) + value = mu * percentile(running, running_percentile) ``mu`` is the throughput of one in-flight slot (roughly ``1 / latency``), computed as a ratio of windowed sums so discrete completions don't make it - jumpy. ``running_ref`` is the concurrency the worker typically sustains; the - p90 rather than the max avoids latching onto rare backlog spikes. Their - product is the worker's sustained throughput. + jumpy. The percentile (``running_ref``) is the concurrency the worker + typically sustains; the p90 rather than the max avoids latching onto rare + backlog spikes. Their product is the worker's sustained throughput. + + The other modes use the raw completion rate and ignore concurrency: + + * ``'window_mean'``: completions per second over the trailing window. + * ``'window_max'``: the highest single-interval completion rate in the + trailing window. + * ``'lifetime_max'``: the highest single-interval completion rate seen since + the estimator was created. Never decays. + + Single-interval rates are noisy when requests take seconds, since a few + completions can land in one interval; the max modes latch onto such bursts. - Units: ``mu`` is completions per slot-second, so ``capacity`` is in requests - per second. The runner reads the published value as how many requests one + Units: every mode is in requests per second. The runner reads the published value as how many requests one worker can handle, so each worker is sized for about one second of its sustained throughput; ``multiplier`` scales that, e.g. 2.0 sizes each worker for two seconds' worth of requests. @@ -230,13 +250,17 @@ class CompletionRateCapacityEstimator: after a scaling event, when the runner rebalances work and a worker is briefly starved, its raw completion rate collapses, but ``mu`` and ``running_ref`` hold, so the estimate does not crater and trigger a downscale - spiral. + spiral. ``'window_mean'`` has no such protection; the max modes avoid it by + holding the peak. - The estimator suits servers whose requests take long enough to be seen in - flight by periodic sampling, such as LLM servers with latencies of a second - or more. For servers with millisecond latencies the in-flight gauge reads - almost always 0; supply a different signal through - ``AutoscalingMetricConfig.load_signal_fn`` instead. + ``'concurrency'`` suits servers whose requests take long enough to be seen + in flight by periodic sampling, such as LLM servers with latencies of a + second or more. For servers with millisecond latencies the in-flight gauge + reads almost always 0; use a completion-rate mode, or supply a different + signal through ``AutoscalingMetricConfig.load_signal_fn``. + + In every mode, intervals with no completions and nothing in flight are + skipped, so idle periods don't drag the estimate down. The estimate is already smoothed; publish it as is. This class does no I/O. @@ -244,13 +268,19 @@ class CompletionRateCapacityEstimator: multiplier: Dimensionless factor on the estimate; larger means fewer workers. min_value: Floor of the estimate. - window_secs: Trailing window; 0 keeps every interval. + window_secs: Trailing window; 0 keeps every interval. ``'lifetime_max'`` + uses it only to count completions for ``min_requests``. running_percentile: Percentile of in-flight concurrency used as - ``running_ref``. + ``running_ref``; ``'concurrency'`` mode only. min_requests: Completions a newly started worker needs before it reports, unless ``warmup_timeout_secs`` passes first. warmup_timeout_secs: Longest time after a worker first sees work that it waits for ``min_requests``. + aggregation: ``'concurrency'``, ``'window_mean'``, ``'window_max'`` or + ``'lifetime_max'``. + + Raises: + ValueError: if ``aggregation`` is not one of the above. """ def __init__( self, @@ -260,7 +290,13 @@ def __init__( running_percentile: float = DEFAULT_RUNNING_PERCENTILE, min_requests: int = DEFAULT_MIN_REQUESTS, warmup_timeout_secs: float = DEFAULT_WARMUP_TIMEOUT_SECS, + aggregation: str = CONCURRENCY, ): + if aggregation not in _AGGREGATIONS: + raise ValueError( + f'Unknown aggregation {aggregation!r}; expected one of ' + f'{", ".join(_AGGREGATIONS)}.') + self._aggregation = aggregation self._multiplier = max(float(multiplier), 0.0) self._min_value = int(min_value) self._window_secs = max(float(window_secs), 0.0) @@ -271,6 +307,7 @@ def __init__( self._last_ts: Optional[float] = None self._first_active_ts: Optional[float] = None self._samples: list[_Interval] = [] + self._peak_rate = 0.0 # Highest single-interval rate, for lifetime_max. def update( self, @@ -308,10 +345,11 @@ def update( return None self._samples.append( _Interval(now, float(completed), float(running), float(secs))) + self._peak_rate = max(self._peak_rate, completed / secs) total_completed = sum(s.completed for s in self._samples) slot_secs = sum(s.running * s.secs for s in self._samples) - if slot_secs <= 0: + if self._aggregation == CONCURRENCY and slot_secs <= 0: return None # Nothing was in flight during the window. if self._first_active_ts is None: self._first_active_ts = prev_ts @@ -321,11 +359,19 @@ def update( # the floor value for the whole job. return None - mu = total_completed / slot_secs - running_ref = _nearest_rank_percentile( - [s.running for s in self._samples if s.running > 0], - self._running_percentile) - return max(float(self._min_value), mu * running_ref * self._multiplier) + if self._aggregation == WINDOW_MEAN: + value = total_completed / sum(s.secs for s in self._samples) + elif self._aggregation == WINDOW_MAX: + value = max(s.completed / s.secs for s in self._samples) + elif self._aggregation == LIFETIME_MAX: + value = self._peak_rate + else: + mu = total_completed / slot_secs + running_ref = _nearest_rank_percentile( + [s.running for s in self._samples if s.running > 0], + self._running_percentile) + value = mu * running_ref + return max(float(self._min_value), value * self._multiplier) @dataclass(frozen=True) @@ -343,6 +389,9 @@ class AutoscalingMetricConfig: min_requests: Completions a newly started worker needs before it publishes, unless ``warmup_timeout_secs`` passes first. warmup_timeout_secs: Longest time a new worker waits for ``min_requests``. + aggregation: How the estimate is computed: ``'concurrency'`` (default), + ``'window_mean'``, ``'window_max'`` or ``'lifetime_max'``. See + :class:`CompletionRateCapacityEstimator`. load_signal_fn: Optional replacement for the built-in estimate: given the server's metrics port, returns the capacity directly. tag: Optional metric tag override. Leave unset so the RunInference model @@ -355,6 +404,7 @@ class AutoscalingMetricConfig: signal_window_secs: float = DEFAULT_SIGNAL_WINDOW_SECS min_requests: int = DEFAULT_MIN_REQUESTS warmup_timeout_secs: float = DEFAULT_WARMUP_TIMEOUT_SECS + aggregation: str = CONCURRENCY load_signal_fn: Optional[Callable[[int], Optional[float]]] = None tag: Optional[str] = None @@ -432,7 +482,8 @@ def __init__( multiplier=config.multiplier, window_secs=config.signal_window_secs, min_requests=config.min_requests, - warmup_timeout_secs=config.warmup_timeout_secs) + warmup_timeout_secs=config.warmup_timeout_secs, + aggregation=config.aggregation) self._stop_event = threading.Event() self._thread: Optional[threading.Thread] = None diff --git a/sdks/python/apache_beam/ml/inference/autoscaling_metrics_test.py b/sdks/python/apache_beam/ml/inference/autoscaling_metrics_test.py index 39c403b3dd27..3be4ba41e9b2 100644 --- a/sdks/python/apache_beam/ml/inference/autoscaling_metrics_test.py +++ b/sdks/python/apache_beam/ml/inference/autoscaling_metrics_test.py @@ -149,6 +149,58 @@ def test_capacity_estimator_defaults_to_wall_clock(self): self.assertIsNone(estimator.update(0.0, 1.0)) self.assertAlmostEqual(estimator.update(10.0, 1.0), 2.0) + def test_capacity_estimator_completion_rate_aggregations(self): + """The completion-rate modes ignore concurrency.""" + def run(aggregation, **kwargs): + estimator = autoscaling_metrics.CompletionRateCapacityEstimator( + min_value=0, + min_requests=0, + window_secs=60.0, + aggregation=aggregation, + **kwargs) + # (completed_total, running, now): rates of 2, 6, idle, 1 req/s. + samples = [(0, 3, 0.0), (10, 3, 5.0), (40, 8, 10.0), (40, 0, 15.0), + (45, 1, 20.0)] + return [estimator.update(c, r, now=t) for c, r, t in samples] + + # 10 + 30 + 5 completions over the three non-idle 5s intervals. + self.assertEqual(run('window_mean'), [None, 2.0, 4.0, None, 3.0]) + self.assertEqual(run('window_max'), [None, 2.0, 6.0, None, 6.0]) + self.assertEqual(run('lifetime_max'), [None, 2.0, 6.0, None, 6.0]) + self.assertEqual( + run('lifetime_max', multiplier=0.5), [None, 1.0, 3.0, None, 3.0]) + + def test_capacity_estimator_max_modes_differ_once_peak_leaves_window(self): + def run(aggregation): + estimator = autoscaling_metrics.CompletionRateCapacityEstimator( + min_value=0, + min_requests=0, + window_secs=10.0, + aggregation=aggregation) + # Rates: 6 req/s, then 1 req/s for three intervals. + samples = [(0, 1, 0.0), (30, 1, 5.0), (35, 1, 10.0), (40, 1, 15.0), + (45, 1, 20.0)] + return [estimator.update(c, r, now=t) for c, r, t in samples][-1] + + # The 6 req/s interval ended at t=5, outside the 10s window at t=20. + self.assertEqual(run('window_max'), 1.0) + self.assertEqual(run('lifetime_max'), 6.0) + + def test_capacity_estimator_completion_rate_works_without_in_flight(self): + """A server whose gauge always reads 0 still gets a completion rate.""" + concurrency = autoscaling_metrics.CompletionRateCapacityEstimator( + min_requests=0) + rate = autoscaling_metrics.CompletionRateCapacityEstimator( + min_value=0, min_requests=0, aggregation='window_mean') + for estimator in (concurrency, rate): + estimator.update(0, 0, now=0.0) + self.assertIsNone(concurrency.update(500, 0, now=5.0)) + self.assertEqual(rate.update(500, 0, now=5.0), 100.0) + + def test_capacity_estimator_rejects_unknown_aggregation(self): + with self.assertRaises(ValueError): + autoscaling_metrics.CompletionRateCapacityEstimator(aggregation='p90') + def test_publisher_publishes_signal_values_as_is(self): """Each capacity is published rounded; None skips the tick.""" cell = ( @@ -256,6 +308,11 @@ def test_publisher_config_and_defaults(self): config=autoscaling_metrics.AutoscalingMetricConfig( poll_interval_secs=7.0)) self.assertEqual(publisher._poll_interval_secs, 7.0) + self.assertEqual( + _publisher( + config=autoscaling_metrics.AutoscalingMetricConfig( + aggregation='lifetime_max'))._estimator._aggregation, + 'lifetime_max') # No config -> defaults. self.assertEqual( _publisher()._poll_interval_secs, From a84f728430535372ce1e2d51893719e1a035d293 Mon Sep 17 00:00:00 2001 From: AMOOOMA Date: Tue, 3 Nov 2026 01:27:35 +0000 Subject: [PATCH 10/14] Tidy the estimator docstring and fill test gaps Group the estimator docstring by mode, and make start() a no-op once a publisher is stopped: previously a superseded publisher that was started again would take its metric back from the newer one and then publish nothing. Add tests for scrape failures, non-numeric samples, negative gauge readings, running_percentile, window_secs=0, errors in the polling loop, and repeated start() calls; the module now has full line and branch coverage. --- .../ml/inference/autoscaling_metrics.py | 58 ++++++++-------- .../ml/inference/autoscaling_metrics_test.py | 68 ++++++++++++++++--- 2 files changed, 85 insertions(+), 41 deletions(-) diff --git a/sdks/python/apache_beam/ml/inference/autoscaling_metrics.py b/sdks/python/apache_beam/ml/inference/autoscaling_metrics.py index 3eb81d22b75c..660ec11ce55c 100644 --- a/sdks/python/apache_beam/ml/inference/autoscaling_metrics.py +++ b/sdks/python/apache_beam/ml/inference/autoscaling_metrics.py @@ -214,25 +214,36 @@ class CompletionRateCapacityEstimator: """Estimates how many requests a worker can handle. Feed it periodic samples of a model server's completed-requests counter and - in-flight gauge through :meth:`update`. Each sample closes an interval; + in-flight gauge through :meth:`update`. Each sample closes an interval, and ``aggregation`` picks how the intervals become an estimate. Every mode - returns ``max(min_value, value * multiplier)``. + returns ``max(min_value, value * multiplier)``. This class does no I/O. ``'concurrency'`` (the default) computes, over a trailing ``window_secs`` window:: mu = sum(completed) / sum(running * secs) - value = mu * percentile(running, running_percentile) + running_ref = percentile(running, running_percentile) + value = mu * running_ref ``mu`` is the throughput of one in-flight slot (roughly ``1 / latency``), computed as a ratio of windowed sums so discrete completions don't make it - jumpy. The percentile (``running_ref``) is the concurrency the worker - typically sustains; the p90 rather than the max avoids latching onto rare - backlog spikes. Their product is the worker's sustained throughput. + jumpy. ``running_ref`` is the concurrency the worker typically sustains; the + p90 rather than the max avoids latching onto rare backlog spikes. Their + product is the worker's sustained throughput. - The other modes use the raw completion rate and ignore concurrency: + Normalizing by concurrency keeps the estimate steady when load changes: right + after a scaling event, when the runner rebalances work and a worker is + briefly starved, its raw completion rate collapses, but ``mu`` and + ``running_ref`` hold, so the estimate does not crater and trigger a downscale + spiral. This mode suits servers whose requests take long enough to be seen + in flight by periodic sampling, such as LLM servers with latencies of a + second or more; for millisecond latencies the gauge almost always reads 0. + + The other modes use the raw completion rate and ignore concurrency, so + they also work for millisecond-latency servers: - * ``'window_mean'``: completions per second over the trailing window. + * ``'window_mean'``: completions per second over the trailing window. Drops + when a worker is starved. * ``'window_max'``: the highest single-interval completion rate in the trailing window. * ``'lifetime_max'``: the highest single-interval completion rate seen since @@ -241,28 +252,13 @@ class CompletionRateCapacityEstimator: Single-interval rates are noisy when requests take seconds, since a few completions can land in one interval; the max modes latch onto such bursts. - Units: every mode is in requests per second. The runner reads the published value as how many requests one - worker can handle, so each worker is sized for about one second of its - sustained throughput; ``multiplier`` scales that, e.g. 2.0 sizes each worker - for two seconds' worth of requests. - - Normalizing by concurrency keeps the estimate steady when load changes: right - after a scaling event, when the runner rebalances work and a worker is - briefly starved, its raw completion rate collapses, but ``mu`` and - ``running_ref`` hold, so the estimate does not crater and trigger a downscale - spiral. ``'window_mean'`` has no such protection; the max modes avoid it by - holding the peak. - - ``'concurrency'`` suits servers whose requests take long enough to be seen - in flight by periodic sampling, such as LLM servers with latencies of a - second or more. For servers with millisecond latencies the in-flight gauge - reads almost always 0; use a completion-rate mode, or supply a different - signal through ``AutoscalingMetricConfig.load_signal_fn``. - In every mode, intervals with no completions and nothing in flight are skipped, so idle periods don't drag the estimate down. - The estimate is already smoothed; publish it as is. This class does no I/O. + Units: every mode is in requests per second. The runner reads the published + value as how many requests one worker can handle, so each worker is sized + for about one second of its sustained throughput; ``multiplier`` scales + that, e.g. 2.0 sizes each worker for two seconds' worth of requests. Args: multiplier: Dimensionless factor on the estimate; larger means fewer @@ -341,7 +337,7 @@ def update( cutoff = now - self._window_secs self._samples = [s for s in self._samples if s.end_ts >= cutoff] if completed == 0 and running == 0: - # Skip fully idle intervals so troughs don't drag running_ref to 0. + # Skip fully idle intervals so idle periods don't drag the estimate down. return None self._samples.append( _Interval(now, float(completed), float(running), float(secs))) @@ -490,9 +486,11 @@ def __init__( def start(self) -> None: """Starts the background polling thread. - The metric stays absent until the first sample is published. + The metric stays absent until the first sample is published. Does nothing + if already started or once stopped, so a superseded publisher can't take + its metric back from the newer one; create a new publisher instead. """ - if self._thread is not None: + if self._thread is not None or self._stop_event.is_set(): return with NumLoadedModelsPublisher._active_lock: previous = NumLoadedModelsPublisher._active.get(self._metric_name) diff --git a/sdks/python/apache_beam/ml/inference/autoscaling_metrics_test.py b/sdks/python/apache_beam/ml/inference/autoscaling_metrics_test.py index 3be4ba41e9b2..7b1d85ae40ee 100644 --- a/sdks/python/apache_beam/ml/inference/autoscaling_metrics_test.py +++ b/sdks/python/apache_beam/ml/inference/autoscaling_metrics_test.py @@ -65,6 +65,8 @@ def test_parse_prometheus_metric_exposition_details(self): 'vllm:request_success_total{finished_reason="length"} 5.0\n' # A metric whose name merely starts with the requested one is ignored. 'vllm:request_success_total_created 1.7e9\n' + # Values that aren't numbers are skipped. + 'vllm:request_success_total{finished_reason="abort"} n/a\n' # Optional trailing timestamps are ignored. 'nv_inference_request_success{model="a"} 3 1700000000000\n' # Quoted label values may contain commas, braces and escapes. @@ -83,6 +85,19 @@ def test_parse_prometheus_metric_exposition_details(self): self.assertIsNone( parse(prom_text, 'labeled', label_filter={'path': 'other'})) + def test_scrape_prometheus_metrics(self): + def fake_urlopen(url, timeout): + self.assertEqual(url, 'http://localhost:8000/metrics') + self.assertEqual(timeout, 5.0) + return io.BytesIO(b'running 4\n') + + with mock.patch('urllib.request.urlopen', side_effect=fake_urlopen): + self.assertEqual( + autoscaling_metrics.scrape_prometheus_metrics(8000), 'running 4\n') + # Any failure (e.g. the server is restarting) -> None. + with mock.patch('urllib.request.urlopen', side_effect=OSError('refused')): + self.assertIsNone(autoscaling_metrics.scrape_prometheus_metrics(8000)) + def test_capacity_estimator(self): estimator = autoscaling_metrics.CompletionRateCapacityEstimator( multiplier=2.0, min_value=1, window_secs=60.0, running_percentile=90.0) @@ -100,8 +115,25 @@ def test_capacity_estimator_skips_counter_reset_and_clock_skew(self): self.assertIsNone(estimator.update(5.0, 2.0, now=105.0)) # Non-increasing timestamp -> skipped. self.assertIsNone(estimator.update(9.0, 2.0, now=105.0)) - # delta=4 over 5s with running=2 -> mu=0.4, p90=2 -> 0.8 -> min_value 1. - self.assertAlmostEqual(estimator.update(13.0, 2.0, now=110.0), 1.0) + # Negative gauge reading -> skipped. + self.assertIsNone(estimator.update(9.0, -1.0, now=107.0)) + # delta=4 over 3s with running=2 -> mu=0.67, p90=2 -> 1.33. + self.assertAlmostEqual(estimator.update(13.0, 2.0, now=110.0), 4 / 3) + + def test_capacity_estimator_running_percentile(self): + def run(running_percentile): + estimator = autoscaling_metrics.CompletionRateCapacityEstimator( + min_value=0, min_requests=0, running_percentile=running_percentile) + estimator.update(0.0, 1.0, now=0.0) + # Ten 1s intervals with 1..10 in flight, each completing as many + # requests: sum(completed) == sum(running * secs), so mu = 1. + for i in range(1, 11): + value = estimator.update(i * (i + 1) / 2, float(i), now=float(i)) + return value + + self.assertEqual(run(90.0), 9.0) # Nearest rank: the 9th of 10. + self.assertEqual(run(50.0), 5.0) + self.assertEqual(run(100.0), 10.0) def test_capacity_estimator_warmup_min_requests_and_timeout(self): # Requires >= 3 completed requests before the 10s warmup timeout. @@ -171,11 +203,11 @@ def run(aggregation, **kwargs): run('lifetime_max', multiplier=0.5), [None, 1.0, 3.0, None, 3.0]) def test_capacity_estimator_max_modes_differ_once_peak_leaves_window(self): - def run(aggregation): + def run(aggregation, window_secs=10.0): estimator = autoscaling_metrics.CompletionRateCapacityEstimator( min_value=0, min_requests=0, - window_secs=10.0, + window_secs=window_secs, aggregation=aggregation) # Rates: 6 req/s, then 1 req/s for three intervals. samples = [(0, 1, 0.0), (30, 1, 5.0), (35, 1, 10.0), (40, 1, 15.0), @@ -185,6 +217,8 @@ def run(aggregation): # The 6 req/s interval ended at t=5, outside the 10s window at t=20. self.assertEqual(run('window_max'), 1.0) self.assertEqual(run('lifetime_max'), 6.0) + # window_secs=0 keeps every interval. + self.assertEqual(run('window_max', window_secs=0), 6.0) def test_capacity_estimator_completion_rate_works_without_in_flight(self): """A server whose gauge always reads 0 still gets a completion rate.""" @@ -202,7 +236,7 @@ def test_capacity_estimator_rejects_unknown_aggregation(self): autoscaling_metrics.CompletionRateCapacityEstimator(aggregation='p90') def test_publisher_publishes_signal_values_as_is(self): - """Each capacity is published rounded; None skips the tick.""" + """Each capacity is published rounded; None or an error skips the tick.""" cell = ( autoscaling_metrics.MetricsEnvironment.process_wide_container(). get_distribution( @@ -212,7 +246,7 @@ def test_publisher_publishes_signal_values_as_is(self): cell.reset() self.addCleanup(cell.reset) published = [] - signal_values = [3.4, None, 7.6] + signal_values = [3.4, None, RuntimeError('scrape broke'), 7.6] def load_signal_fn(port): del port @@ -223,7 +257,10 @@ def load_signal_fn(port): pub._stop_event.set() # pylint: disable=protected-access return None published.append(None) - return signal_values.pop(0) + value = signal_values.pop(0) + if isinstance(value, Exception): + raise value + return value pub = _publisher( 'sink_test', @@ -231,9 +268,11 @@ def load_signal_fn(port): config=autoscaling_metrics.AutoscalingMetricConfig( poll_interval_secs=0, load_signal_fn=load_signal_fn)) # Run the polling loop on this thread so the test is deterministic. - pub._run() # pylint: disable=protected-access - # 3.4 -> 3; None leaves 3 in place; 7.6 -> 8. - self.assertEqual(published, [3, 3, 8]) + with self.assertLogs(level='ERROR') as logs: + pub._run() # pylint: disable=protected-access + self.assertIn('num_loaded_models_sink_test', logs.output[0]) + # 3.4 -> 3; None and the error leave 3 in place; 7.6 -> 8. + self.assertEqual(published, [3, 3, 3, 8]) self.assertEqual(cell.get_cumulative().count, 1) def test_publish_resets_process_wide_distribution_cell(self): @@ -273,9 +312,16 @@ def test_publisher_for_same_tag_supersedes_previous(self): self.assertTrue(second._thread.is_alive()) self.assertTrue(other._thread.is_alive()) self.assertIs(active['num_loaded_models_reload'], second) - # Stopping the stale one again must not unregister the new owner. + # Starting again is a no-op. + thread = second._thread + second.start() + self.assertIs(second._thread, thread) + # The stale one can neither take the tag back nor unregister the owner. + first.start() + self.assertIsNone(first._thread) first.stop() self.assertIs(active['num_loaded_models_reload'], second) + self.assertTrue(second._thread.is_alive()) finally: first.stop() second.stop() From 2cd0ec0a0c9ada00038b6c7a040bd9c476c3a530 Mon Sep 17 00:00:00 2001 From: AMOOOMA Date: Tue, 3 Nov 2026 01:47:34 +0000 Subject: [PATCH 11/14] Publish the autoscaling metric atomically and time samples monotonically The SDK harness reads process-wide metrics from another thread, and reset() followed by update() briefly left the distribution empty (count=0), which a concurrent report could pick up. Swap the cell's data in one step under its lock instead. Time estimator samples with time.monotonic() so wall-clock adjustments can't distort intervals. --- .../ml/inference/autoscaling_metrics.py | 17 ++++--- .../ml/inference/autoscaling_metrics_test.py | 46 +++++++++++++++---- 2 files changed, 48 insertions(+), 15 deletions(-) diff --git a/sdks/python/apache_beam/ml/inference/autoscaling_metrics.py b/sdks/python/apache_beam/ml/inference/autoscaling_metrics.py index 660ec11ce55c..be88366a5c88 100644 --- a/sdks/python/apache_beam/ml/inference/autoscaling_metrics.py +++ b/sdks/python/apache_beam/ml/inference/autoscaling_metrics.py @@ -63,6 +63,7 @@ from typing import NamedTuple from typing import Optional +from apache_beam.metrics.cells import DistributionData from apache_beam.metrics.execution import MetricsEnvironment from apache_beam.metrics.metricbase import MetricName @@ -315,7 +316,7 @@ def update( Args: completed_total: Current value of the completed-requests counter. running: Current number of in-flight requests. - now: Sample timestamp in seconds; defaults to ``time.time()``. + now: Sample timestamp in seconds; defaults to ``time.monotonic()``. Returns: The capacity estimate, or ``None`` when there is not enough data (first @@ -323,7 +324,7 @@ def update( keep its previous value. """ if now is None: - now = time.time() + now = time.monotonic() prev_count, prev_ts = self._last_count, self._last_ts self._last_count, self._last_ts = completed_total, now if prev_count is None or prev_ts is None: @@ -554,12 +555,14 @@ def _read_capacity(self) -> Optional[float]: def _publish(self, value: float) -> None: published = int(round(value)) # The SDK harness reports process-wide distributions without resetting - # them, and Dataflow reads their cumulative mean (sum / count). Resetting - # before each update keeps the cell at (sum=published, count=1), so the - # runner sees the latest value rather than an average since worker start. + # them, and Dataflow reads their cumulative mean (sum / count). Keeping the + # cell at (sum=published, count=1) makes the runner see the latest value + # rather than an average since worker start. The harness reads the cell + # from another thread, so swap the data in one step: reset() then + # update() would briefly expose an empty cell (count=0). cell = MetricsEnvironment.process_wide_container().get_distribution( MetricName(MODEL_MANAGER_METRICS_NAMESPACE, self._metric_name)) current = cell.get_cumulative() if current.count != 1 or current.sum != published: - cell.reset() - cell.update(published) + with cell._lock: # pylint: disable=protected-access + cell.data = DistributionData(published, 1, published, published) diff --git a/sdks/python/apache_beam/ml/inference/autoscaling_metrics_test.py b/sdks/python/apache_beam/ml/inference/autoscaling_metrics_test.py index 7b1d85ae40ee..7260e94ee496 100644 --- a/sdks/python/apache_beam/ml/inference/autoscaling_metrics_test.py +++ b/sdks/python/apache_beam/ml/inference/autoscaling_metrics_test.py @@ -15,6 +15,7 @@ # limitations under the License. # import io +import threading import unittest from unittest import mock @@ -172,11 +173,11 @@ def test_capacity_estimator_skips_idle_ticks_without_diluting_p90(self): # running=0): total_completed=30, total_slot_secs=25 -> mu=1.2, p90=5 -> 6. self.assertAlmostEqual(estimator.update(30.0, 0.0, now=125.0), 6.0) - def test_capacity_estimator_defaults_to_wall_clock(self): + def test_capacity_estimator_defaults_to_monotonic_clock(self): estimator = autoscaling_metrics.CompletionRateCapacityEstimator( min_requests=0) with mock.patch.object(autoscaling_metrics.time, - 'time', + 'monotonic', side_effect=[100.0, 105.0]): self.assertIsNone(estimator.update(0.0, 1.0)) self.assertAlmostEqual(estimator.update(10.0, 1.0), 2.0) @@ -285,13 +286,42 @@ def test_publish_resets_process_wide_distribution_cell(self): 'num_loaded_models_reset_test'))) cell.reset() self.addCleanup(cell.reset) + cell.update(7) # A stale value, e.g. from before a model reload. for _ in range(10): pub._publish(2.0) - self.assertEqual(cell.get_cumulative().count, 1) - self.assertEqual(cell.get_cumulative().sum, 2) + self.assertEqual( + cell.get_cumulative(), autoscaling_metrics.DistributionData(2, 1, 2, 2)) pub._publish(5.0) - self.assertEqual(cell.get_cumulative().count, 1) - self.assertEqual(cell.get_cumulative().sum, 5) + self.assertEqual( + cell.get_cumulative(), autoscaling_metrics.DistributionData(5, 1, 5, 5)) + + def test_publish_never_exposes_an_empty_cell(self): + """A concurrent reader must never see count=0 between updates.""" + pub = _publisher('atomic_test') + cell = ( + autoscaling_metrics.MetricsEnvironment.process_wide_container(). + get_distribution( + autoscaling_metrics.MetricName( + autoscaling_metrics.MODEL_MANAGER_METRICS_NAMESPACE, + 'num_loaded_models_atomic_test'))) + self.addCleanup(cell.reset) + pub._publish(1.0) + stop = threading.Event() + counts = set() + + def read(): + while not stop.is_set(): + counts.add(cell.get_cumulative().count) + + reader = threading.Thread(target=read) + reader.start() + try: + for i in range(20000): + pub._publish(float(i % 7 + 1)) + finally: + stop.set() + reader.join() + self.assertEqual(counts, {1}) def test_publisher_for_same_tag_supersedes_previous(self): """A reloaded model's publisher takes over its tag from the stale one.""" @@ -393,7 +423,7 @@ def fake_urlopen(url, timeout=5.0): with mock.patch('urllib.request.urlopen', side_effect=fake_urlopen): with mock.patch.object(autoscaling_metrics.time, - 'time', + 'monotonic', side_effect=[100.0, 105.0]): self.assertIsNone(signal()) # Baseline sample. # delta=20 over 5s with running=4 -> mu=1, p90=4, multiplier 2 -> 8. @@ -418,7 +448,7 @@ def test_publisher_skips_unreadable_scrapes(self): 'scrape_prometheus_metrics', side_effect=lambda port: scrapes.pop(0)): with mock.patch.object(autoscaling_metrics.time, - 'time', + 'monotonic', side_effect=[100.0, 110.0]): self.assertIsNone(signal()) self.assertIsNone(signal()) From 0227bbada6e0a4c7ea7322fbbb56d0caa2ffd61c Mon Sep 17 00:00:00 2001 From: AMOOOMA Date: Thu, 5 Nov 2026 18:46:47 +0000 Subject: [PATCH 12/14] Publish the autoscaling metric without touching cell internals The previous atomic swap took DistributionCell._lock and assigned cell.data, neither of which is reachable from Python when apache_beam.metrics is compiled with Cython (as in CI and on workers), so _publish raised AttributeError. Install a fresh DistributionCell holding (sum=value, count=1) in the process-wide container with a single dict store instead; that is still atomic under the GIL, so a concurrent reader never sees an empty cell. Tests re-fetch the cell for every read since it is now replaced rather than mutated. --- .../ml/inference/autoscaling_metrics.py | 21 ++++--- .../ml/inference/autoscaling_metrics_test.py | 61 +++++++++---------- 2 files changed, 44 insertions(+), 38 deletions(-) diff --git a/sdks/python/apache_beam/ml/inference/autoscaling_metrics.py b/sdks/python/apache_beam/ml/inference/autoscaling_metrics.py index be88366a5c88..b12ae7a8cff1 100644 --- a/sdks/python/apache_beam/ml/inference/autoscaling_metrics.py +++ b/sdks/python/apache_beam/ml/inference/autoscaling_metrics.py @@ -63,8 +63,9 @@ from typing import NamedTuple from typing import Optional -from apache_beam.metrics.cells import DistributionData +from apache_beam.metrics.cells import DistributionCell from apache_beam.metrics.execution import MetricsEnvironment +from apache_beam.metrics.execution import _TypedMetricName from apache_beam.metrics.metricbase import MetricName __all__ = [ @@ -558,11 +559,17 @@ def _publish(self, value: float) -> None: # them, and Dataflow reads their cumulative mean (sum / count). Keeping the # cell at (sum=published, count=1) makes the runner see the latest value # rather than an average since worker start. The harness reads the cell - # from another thread, so swap the data in one step: reset() then - # update() would briefly expose an empty cell (count=0). - cell = MetricsEnvironment.process_wide_container().get_distribution( + # from another thread, so the update must be atomic: reset() then update() + # would briefly expose an empty cell (count=0). The cell's lock and data + # are not reachable from Python when apache_beam.metrics is compiled with + # Cython (as on workers), so install a fresh cell holding the value with a + # single dict store instead, which is atomic under the GIL. + container = MetricsEnvironment.process_wide_container() + key = _TypedMetricName( + DistributionCell, MetricName(MODEL_MANAGER_METRICS_NAMESPACE, self._metric_name)) - current = cell.get_cumulative() + current = container.get_metric_cell(key).get_cumulative() if current.count != 1 or current.sum != published: - with cell._lock: # pylint: disable=protected-access - cell.data = DistributionData(published, 1, published, published) + cell = DistributionCell() + cell.update(published) + container.metrics[key] = cell diff --git a/sdks/python/apache_beam/ml/inference/autoscaling_metrics_test.py b/sdks/python/apache_beam/ml/inference/autoscaling_metrics_test.py index 7260e94ee496..9a42d11f77b7 100644 --- a/sdks/python/apache_beam/ml/inference/autoscaling_metrics_test.py +++ b/sdks/python/apache_beam/ml/inference/autoscaling_metrics_test.py @@ -19,6 +19,7 @@ import unittest from unittest import mock +from apache_beam.metrics.cells import DistributionData from apache_beam.ml.inference import autoscaling_metrics @@ -32,6 +33,23 @@ def _publisher(model_identifier='step-id', port=None, **kwargs): **kwargs) +def _cell(model_identifier): + """The process-wide ``num_loaded_models_`` cell. + + The publisher replaces the cell rather than mutating it, so fetch it again + for every read instead of holding on to it. + """ + return autoscaling_metrics.MetricsEnvironment.process_wide_container( + ).get_distribution( + autoscaling_metrics.MetricName( + autoscaling_metrics.MODEL_MANAGER_METRICS_NAMESPACE, + f'num_loaded_models_{model_identifier}')) + + +def _published(model_identifier): + return _cell(model_identifier).get_cumulative() + + class AutoscalingMetricsTest(unittest.TestCase): def test_sanitize_metric_tag(self): self.assertEqual( @@ -238,14 +256,8 @@ def test_capacity_estimator_rejects_unknown_aggregation(self): def test_publisher_publishes_signal_values_as_is(self): """Each capacity is published rounded; None or an error skips the tick.""" - cell = ( - autoscaling_metrics.MetricsEnvironment.process_wide_container(). - get_distribution( - autoscaling_metrics.MetricName( - autoscaling_metrics.MODEL_MANAGER_METRICS_NAMESPACE, - 'num_loaded_models_sink_test'))) - cell.reset() - self.addCleanup(cell.reset) + _cell('sink_test').reset() + self.addCleanup(lambda: _cell('sink_test').reset()) published = [] signal_values = [3.4, None, RuntimeError('scrape broke'), 7.6] @@ -253,7 +265,7 @@ def load_signal_fn(port): del port if published: # Record what the previous tick left in the cell. - published[-1] = cell.get_cumulative().sum + published[-1] = _published('sink_test').sum if not signal_values: pub._stop_event.set() # pylint: disable=protected-access return None @@ -274,44 +286,31 @@ def load_signal_fn(port): self.assertIn('num_loaded_models_sink_test', logs.output[0]) # 3.4 -> 3; None and the error leave 3 in place; 7.6 -> 8. self.assertEqual(published, [3, 3, 3, 8]) - self.assertEqual(cell.get_cumulative().count, 1) + self.assertEqual(_published('sink_test').count, 1) def test_publish_resets_process_wide_distribution_cell(self): pub = _publisher('reset_test') - cell = ( - autoscaling_metrics.MetricsEnvironment.process_wide_container(). - get_distribution( - autoscaling_metrics.MetricName( - autoscaling_metrics.MODEL_MANAGER_METRICS_NAMESPACE, - 'num_loaded_models_reset_test'))) - cell.reset() - self.addCleanup(cell.reset) - cell.update(7) # A stale value, e.g. from before a model reload. + _cell('reset_test').reset() + self.addCleanup(lambda: _cell('reset_test').reset()) + # A stale value, e.g. from before a model reload. + _cell('reset_test').update(7) for _ in range(10): pub._publish(2.0) - self.assertEqual( - cell.get_cumulative(), autoscaling_metrics.DistributionData(2, 1, 2, 2)) + self.assertEqual(_published('reset_test'), DistributionData(2, 1, 2, 2)) pub._publish(5.0) - self.assertEqual( - cell.get_cumulative(), autoscaling_metrics.DistributionData(5, 1, 5, 5)) + self.assertEqual(_published('reset_test'), DistributionData(5, 1, 5, 5)) def test_publish_never_exposes_an_empty_cell(self): """A concurrent reader must never see count=0 between updates.""" pub = _publisher('atomic_test') - cell = ( - autoscaling_metrics.MetricsEnvironment.process_wide_container(). - get_distribution( - autoscaling_metrics.MetricName( - autoscaling_metrics.MODEL_MANAGER_METRICS_NAMESPACE, - 'num_loaded_models_atomic_test'))) - self.addCleanup(cell.reset) + self.addCleanup(lambda: _cell('atomic_test').reset()) pub._publish(1.0) stop = threading.Event() counts = set() def read(): while not stop.is_set(): - counts.add(cell.get_cumulative().count) + counts.add(_published('atomic_test').count) reader = threading.Thread(target=read) reader.start() From e2ff8e32e6f2314a55455d3837cdf5f07e34f09f Mon Sep 17 00:00:00 2001 From: AMOOOMA Date: Thu, 5 Nov 2026 19:18:35 +0000 Subject: [PATCH 13/14] Add DistributionCell.reset_to and skip estimates with no completions - Add DistributionCell.reset_to(value), which replaces the distribution with a single sample under the cell's lock. It works when apache_beam.metrics is compiled with Cython, where _lock and data are not reachable from outside the class. NumLoadedModelsPublisher uses it instead of swapping the cell in the container. - CompletionRateCapacityEstimator returns None when requests are in flight but none completed within the window (slow first requests, first poll after an idle period), keeping the previously published capacity instead of reporting the floor. lifetime_max keeps reporting its peak. --- sdks/python/apache_beam/metrics/cells.py | 12 ++++++ sdks/python/apache_beam/metrics/cells_test.py | 31 ++++++++++++++ .../ml/inference/autoscaling_metrics.py | 30 ++++++------- .../ml/inference/autoscaling_metrics_test.py | 42 ++++++++++++++++--- 4 files changed, 95 insertions(+), 20 deletions(-) diff --git a/sdks/python/apache_beam/metrics/cells.py b/sdks/python/apache_beam/metrics/cells.py index 8300d7f063e9..ca91410cb98c 100644 --- a/sdks/python/apache_beam/metrics/cells.py +++ b/sdks/python/apache_beam/metrics/cells.py @@ -181,6 +181,18 @@ def reset(self): # type: () -> None self.data = DistributionData.identity_element() + def reset_to(self, value): + """Atomically replaces the distribution with a single sample of value. + + Like ``reset()`` followed by ``update(value)``, except that a concurrent + reader never sees the empty distribution in between. Useful for publishing + the latest value of a gauge-like quantity as a distribution whose + cumulative mean (sum / count) is that value. + """ + ivalue = int(value) + with self._lock: + self.data = DistributionData(ivalue, 1, ivalue, ivalue) + def combine(self, other): # type: (DistributionCell) -> DistributionCell result = DistributionCell() diff --git a/sdks/python/apache_beam/metrics/cells_test.py b/sdks/python/apache_beam/metrics/cells_test.py index 11ea20ed6f6d..0ec509d52b91 100644 --- a/sdks/python/apache_beam/metrics/cells_test.py +++ b/sdks/python/apache_beam/metrics/cells_test.py @@ -131,6 +131,37 @@ def test_basic_operations(self): d.update(900) self.assertEqual(d.get_cumulative(), DistributionData(912, 3, 2, 900)) + def test_reset_to(self): + d = DistributionCell() + d.update(10) + d.update(2) + d.reset_to(7) + self.assertEqual(d.get_cumulative(), DistributionData(7, 1, 7, 7)) + d.reset_to(3.6) + self.assertEqual(d.get_cumulative(), DistributionData(3, 1, 3, 3)) + d.update(5) + self.assertEqual(d.get_cumulative(), DistributionData(8, 2, 3, 5)) + + def test_reset_to_never_exposes_an_empty_distribution(self): + d = DistributionCell() + d.reset_to(1) + stop = threading.Event() + counts = set() + + def read(): + while not stop.is_set(): + counts.add(d.get_cumulative().count) + + reader = threading.Thread(target=read) + reader.start() + try: + for i in range(20000): + d.reset_to(i % 7 + 1) + finally: + stop.set() + reader.join() + self.assertEqual(counts, {1}) + def test_integer_only(self): d = DistributionCell() d.update(3.1) diff --git a/sdks/python/apache_beam/ml/inference/autoscaling_metrics.py b/sdks/python/apache_beam/ml/inference/autoscaling_metrics.py index b12ae7a8cff1..cd40ca22422b 100644 --- a/sdks/python/apache_beam/ml/inference/autoscaling_metrics.py +++ b/sdks/python/apache_beam/ml/inference/autoscaling_metrics.py @@ -63,9 +63,7 @@ from typing import NamedTuple from typing import Optional -from apache_beam.metrics.cells import DistributionCell from apache_beam.metrics.execution import MetricsEnvironment -from apache_beam.metrics.execution import _TypedMetricName from apache_beam.metrics.metricbase import MetricName __all__ = [ @@ -255,7 +253,9 @@ class CompletionRateCapacityEstimator: completions can land in one interval; the max modes latch onto such bursts. In every mode, intervals with no completions and nothing in flight are - skipped, so idle periods don't drag the estimate down. + skipped, so idle periods don't drag the estimate down. While requests are in + flight but none has completed within the window, there is no estimate either + (except in ``'lifetime_max'`` mode once a peak has been seen). Units: every mode is in requests per second. The runner reads the published value as how many requests one worker can handle, so each worker is sized @@ -351,6 +351,13 @@ def update( return None # Nothing was in flight during the window. if self._first_active_ts is None: self._first_active_ts = prev_ts + if total_completed <= 0 and not (self._aggregation == LIFETIME_MAX and + self._peak_rate > 0): + # Requests are in flight but none has completed in the window yet, e.g. + # slow first requests or the first poll after an idle period. That says + # nothing about capacity, so keep the previous value instead of + # reporting the floor. + return None if (total_completed < self._min_requests and now - self._first_active_ts < self._warmup_timeout_secs): # Wait until a new worker is calibrated, so it doesn't briefly report @@ -559,17 +566,10 @@ def _publish(self, value: float) -> None: # them, and Dataflow reads their cumulative mean (sum / count). Keeping the # cell at (sum=published, count=1) makes the runner see the latest value # rather than an average since worker start. The harness reads the cell - # from another thread, so the update must be atomic: reset() then update() - # would briefly expose an empty cell (count=0). The cell's lock and data - # are not reachable from Python when apache_beam.metrics is compiled with - # Cython (as on workers), so install a fresh cell holding the value with a - # single dict store instead, which is atomic under the GIL. - container = MetricsEnvironment.process_wide_container() - key = _TypedMetricName( - DistributionCell, + # from another thread, so replace the data in one step with reset_to(): + # reset() then update() would briefly expose an empty cell (count=0). + cell = MetricsEnvironment.process_wide_container().get_distribution( MetricName(MODEL_MANAGER_METRICS_NAMESPACE, self._metric_name)) - current = container.get_metric_cell(key).get_cumulative() + current = cell.get_cumulative() if current.count != 1 or current.sum != published: - cell = DistributionCell() - cell.update(published) - container.metrics[key] = cell + cell.reset_to(published) diff --git a/sdks/python/apache_beam/ml/inference/autoscaling_metrics_test.py b/sdks/python/apache_beam/ml/inference/autoscaling_metrics_test.py index 9a42d11f77b7..19fb4db24c40 100644 --- a/sdks/python/apache_beam/ml/inference/autoscaling_metrics_test.py +++ b/sdks/python/apache_beam/ml/inference/autoscaling_metrics_test.py @@ -34,11 +34,7 @@ def _publisher(model_identifier='step-id', port=None, **kwargs): def _cell(model_identifier): - """The process-wide ``num_loaded_models_`` cell. - - The publisher replaces the cell rather than mutating it, so fetch it again - for every read instead of holding on to it. - """ + """The process-wide ``num_loaded_models_`` cell.""" return autoscaling_metrics.MetricsEnvironment.process_wide_container( ).get_distribution( autoscaling_metrics.MetricName( @@ -191,6 +187,42 @@ def test_capacity_estimator_skips_idle_ticks_without_diluting_p90(self): # running=0): total_completed=30, total_slot_secs=25 -> mu=1.2, p90=5 -> 6. self.assertAlmostEqual(estimator.update(30.0, 0.0, now=125.0), 6.0) + def test_capacity_estimator_no_estimate_until_a_request_completes(self): + """In-flight requests with no completions never report the floor.""" + # Slow first requests: still nothing completed after the warmup timeout. + slow = autoscaling_metrics.CompletionRateCapacityEstimator( + min_requests=3, warmup_timeout_secs=10.0) + self.assertIsNone(slow.update(0.0, 4.0, now=100.0)) # Baseline. + for t in (105.0, 110.0, 115.0): + self.assertIsNone(slow.update(0.0, 4.0, now=t)) + # The first completions produce an estimate. + self.assertIsNotNone(slow.update(8.0, 4.0, now=120.0)) + + # First poll after an idle period longer than the window: the earlier + # samples have left the window and the new requests haven't finished yet. + for aggregation in ('concurrency', 'window_mean', 'window_max'): + resumed = autoscaling_metrics.CompletionRateCapacityEstimator( + min_value=1, + window_secs=10.0, + min_requests=0, + aggregation=aggregation) + resumed.update(0.0, 2.0, now=0.0) + self.assertGreater(resumed.update(20.0, 2.0, now=5.0), 1.0) + for t in (10.0, 15.0, 20.0): + self.assertIsNone(resumed.update(20.0, 0.0, now=t)) # Idle. + self.assertIsNone(resumed.update(20.0, 3.0, now=25.0), aggregation) + + # lifetime_max keeps reporting its peak. + peak = autoscaling_metrics.CompletionRateCapacityEstimator( + min_value=1, + window_secs=10.0, + min_requests=0, + aggregation='lifetime_max') + peak.update(0.0, 2.0, now=0.0) + self.assertEqual(peak.update(20.0, 2.0, now=5.0), 4.0) + self.assertIsNone(peak.update(20.0, 0.0, now=20.0)) # Idle. + self.assertEqual(peak.update(20.0, 3.0, now=25.0), 4.0) + def test_capacity_estimator_defaults_to_monotonic_clock(self): estimator = autoscaling_metrics.CompletionRateCapacityEstimator( min_requests=0) From 465a1d7c83177c8afef5b387825a1fd50aca3946 Mon Sep 17 00:00:00 2001 From: AMOOOMA Date: Thu, 5 Nov 2026 21:34:42 +0000 Subject: [PATCH 14/14] Narrow the autoscaling_metrics public API and state its requirements Address review feedback: - Rename NumLoadedModelsPublisher to PrometheusNumLoadedModelsPublisher and document what the model server must provide: a local Prometheus text-format /metrics endpoint with a completed-requests counter and an in-flight gauge (or AutoscalingMetricConfig.load_signal_fn otherwise). - Export only AutoscalingMetricConfig and PrometheusNumLoadedModelsPublisher. The capacity estimator and the Prometheus scrape/parse helpers are now private; the aggregation modes are documented on AutoscalingMetricConfig. --- .../ml/inference/autoscaling_metrics.py | 106 ++++++++++-------- .../ml/inference/autoscaling_metrics_test.py | 55 ++++----- 2 files changed, 89 insertions(+), 72 deletions(-) diff --git a/sdks/python/apache_beam/ml/inference/autoscaling_metrics.py b/sdks/python/apache_beam/ml/inference/autoscaling_metrics.py index cd40ca22422b..17945d427c57 100644 --- a/sdks/python/apache_beam/ml/inference/autoscaling_metrics.py +++ b/sdks/python/apache_beam/ml/inference/autoscaling_metrics.py @@ -25,12 +25,20 @@ model server (vLLM, Triton, ...) always has exactly one copy per worker, so it publishes an estimate of the worker's capacity under that metric instead. -Nothing here depends on a particular model server. A handler for a model server -that exposes a completed-requests counter and an in-flight gauge creates a -:class:`NumLoadedModelsPublisher`, supplying only what it knows: the two metric -names, how to find the server's port, and the RunInference model identifier:: +:class:`PrometheusNumLoadedModelsPublisher` reads the server's metrics over +HTTP, so the server must: - publisher = NumLoadedModelsPublisher( +* run on the same worker as the handler and serve the Prometheus text format + at ``http://localhost:/metrics``; +* export a counter of completed requests (monotonically increasing, reset only + when the server restarts), e.g. vLLM's ``vllm:request_success_total``; and +* export a gauge of the requests currently in flight, e.g. vLLM's + ``vllm:num_requests_running``. + +A handler for such a server (vLLM, for example) supplies only the two metric +names, a way to find the port, and the RunInference model identifier:: + + publisher = PrometheusNumLoadedModelsPublisher( metrics_port_fn=lambda: server.port if server.is_up else None, completed_metric='vllm:request_success_total', running_metric='vllm:num_requests_running', @@ -39,15 +47,12 @@ config=autoscaling_config) # An AutoscalingMetricConfig, or None. publisher.start() # Once the server is reachable; publisher.stop() at exit. -The pieces: +If the server exports different metrics, +``AutoscalingMetricConfig.load_signal_fn`` can replace the scrape and the +estimate: it is given the server's port and returns the capacity directly. -* :class:`NumLoadedModelsPublisher` reads the server, estimates its capacity - and publishes it from a background thread. -* :class:`CompletionRateCapacityEstimator` is the estimate itself: pure - computation on the counter and the gauge, no I/O. -* :class:`AutoscalingMetricConfig` holds the user-facing knobs. Plain data. -* :func:`scrape_prometheus_metrics` and :func:`parse_prometheus_metric` read a - server's Prometheus ``/metrics`` endpoint. +The public API is :class:`PrometheusNumLoadedModelsPublisher` and +:class:`AutoscalingMetricConfig`; everything else is an implementation detail. """ # pytype: skip-file @@ -68,11 +73,7 @@ __all__ = [ 'AutoscalingMetricConfig', - 'CompletionRateCapacityEstimator', - 'NumLoadedModelsPublisher', - 'parse_prometheus_metric', - 'sanitize_metric_tag', - 'scrape_prometheus_metrics', + 'PrometheusNumLoadedModelsPublisher', ] # The metric the runner reads (shared with model_manager.py). @@ -90,7 +91,7 @@ DEFAULT_MIN_REQUESTS = 3 DEFAULT_WARMUP_TIMEOUT_SECS = 10.0 -# How CompletionRateCapacityEstimator turns its samples into an estimate. +# How _CompletionRateCapacityEstimator turns its samples into an estimate. CONCURRENCY = 'concurrency' WINDOW_MEAN = 'window_mean' WINDOW_MAX = 'window_max' @@ -109,7 +110,7 @@ _PROMETHEUS_ESCAPE_RE = re.compile(r'\\(.)') -def sanitize_metric_tag(name: str) -> str: +def _sanitize_metric_tag(name: str) -> str: """Turns a model name into a metric-tag-safe token. Only used for the fallback tag of handlers running outside RunInference; the @@ -124,7 +125,7 @@ def sanitize_metric_tag(name: str) -> str: return ''.join(c if (c.isalnum() or c == '_') else '_' for c in name) -def scrape_prometheus_metrics( +def _scrape_prometheus_metrics( port: int, timeout_secs: float = 5.0, path: str = '/metrics') -> Optional[str]: @@ -158,7 +159,7 @@ def unescape(match: re.Match) -> str: } -def parse_prometheus_metric( +def _parse_prometheus_metric( metrics_text: str, metric_name: str, label_filter: Optional[Mapping[str, str]] = None, @@ -210,7 +211,7 @@ class _Interval(NamedTuple): secs: float # Length of the interval. -class CompletionRateCapacityEstimator: +class _CompletionRateCapacityEstimator: """Estimates how many requests a worker can handle. Feed it periodic samples of a model server's completed-requests counter and @@ -394,9 +395,16 @@ class AutoscalingMetricConfig: min_requests: Completions a newly started worker needs before it publishes, unless ``warmup_timeout_secs`` passes first. warmup_timeout_secs: Longest time a new worker waits for ``min_requests``. - aggregation: How the estimate is computed: ``'concurrency'`` (default), - ``'window_mean'``, ``'window_max'`` or ``'lifetime_max'``. See - :class:`CompletionRateCapacityEstimator`. + aggregation: How the estimate, in requests per second, is computed: + + * ``'concurrency'`` (default): per-slot throughput over the window times + the p90 in-flight concurrency. Holds steady when a worker is briefly + starved; suits servers with request latencies of a second or more. + * ``'window_mean'``: completions per second over the window. + * ``'window_max'``: the highest single-interval completion rate in the + window. + * ``'lifetime_max'``: the highest single-interval completion rate seen + so far; never decays. load_signal_fn: Optional replacement for the built-in estimate: given the server's metrics port, returns the capacity directly. tag: Optional metric tag override. Leave unset so the RunInference model @@ -414,12 +422,18 @@ class AutoscalingMetricConfig: tag: Optional[str] = None -class NumLoadedModelsPublisher: - """Publishes a model server's capacity as ``num_loaded_models_``. +class PrometheusNumLoadedModelsPublisher: + """Publishes a Prometheus model server's capacity for autoscaling. + + The capacity is published as ``num_loaded_models_``. - Every ``poll_interval_secs`` a daemon thread reads the server's Prometheus - ``/metrics`` endpoint, feeds the two metrics to a - :class:`CompletionRateCapacityEstimator`, and writes the estimate, rounded + The server must meet the requirements in the module docstring: a Prometheus + text-format ``/metrics`` endpoint on ``localhost`` exporting a + completed-requests counter and an in-flight gauge. + + Every ``poll_interval_secs`` a daemon thread reads that endpoint, feeds the + two metrics to a + capacity estimator, and writes the estimate, rounded to an int, to the process-wide distribution the runner reads. If ``config.load_signal_fn`` is set, it replaces the scrape and the estimate. @@ -434,8 +448,8 @@ class NumLoadedModelsPublisher: The metric tag is ``config.tag`` if set; else ``model_identifier``, used verbatim because it must equal the identifier RunInference annotates the step with (see ``ModelHandler.set_model_identifier``); else ``model_name`` - passed through :func:`sanitize_metric_tag`, for handlers used outside - RunInference. + with characters outside ``[A-Za-z0-9_]`` replaced by ``_``, for handlers + used outside RunInference. Args: metrics_port_fn: Returns the port serving ``/metrics`` while the server is @@ -456,7 +470,7 @@ class NumLoadedModelsPublisher: """ # The publisher owning each metric name in this process, so a reloaded # model's publisher can supersede the stale one. - _active: dict[str, 'NumLoadedModelsPublisher'] = {} + _active: dict[str, 'PrometheusNumLoadedModelsPublisher'] = {} _active_lock = threading.Lock() def __init__( @@ -471,7 +485,7 @@ def __init__( label_filter: Optional[Mapping[str, str]] = None, ): config = config or AutoscalingMetricConfig() - tag = config.tag or model_identifier or sanitize_metric_tag(model_name) + tag = config.tag or model_identifier or _sanitize_metric_tag(model_name) if not tag: raise ValueError( 'Cannot derive a num_loaded_models tag: pass a model_identifier, a ' @@ -483,7 +497,7 @@ def __init__( self._label_filter = dict(label_filter) if label_filter else None self._load_signal_fn = config.load_signal_fn self._poll_interval_secs = config.poll_interval_secs - self._estimator = CompletionRateCapacityEstimator( + self._estimator = _CompletionRateCapacityEstimator( multiplier=config.multiplier, window_secs=config.signal_window_secs, min_requests=config.min_requests, @@ -501,9 +515,10 @@ def start(self) -> None: """ if self._thread is not None or self._stop_event.is_set(): return - with NumLoadedModelsPublisher._active_lock: - previous = NumLoadedModelsPublisher._active.get(self._metric_name) - NumLoadedModelsPublisher._active[self._metric_name] = self + with PrometheusNumLoadedModelsPublisher._active_lock: + previous = PrometheusNumLoadedModelsPublisher._active.get( + self._metric_name) + PrometheusNumLoadedModelsPublisher._active[self._metric_name] = self if previous is not None and previous is not self: logging.info( 'Replacing the active publisher for %s (model reload or server ' @@ -522,9 +537,10 @@ def stop(self) -> None: further downscales. The drained worker's value ages out instead. """ self._stop_event.set() - with NumLoadedModelsPublisher._active_lock: - if NumLoadedModelsPublisher._active.get(self._metric_name) is self: - del NumLoadedModelsPublisher._active[self._metric_name] + with PrometheusNumLoadedModelsPublisher._active_lock: + if PrometheusNumLoadedModelsPublisher._active.get( + self._metric_name) is self: + del PrometheusNumLoadedModelsPublisher._active[self._metric_name] thread = self._thread if thread is not None: thread.join(timeout=self._poll_interval_secs + 5) @@ -549,12 +565,12 @@ def _read_capacity(self) -> Optional[float]: return None # The server is not up. if self._load_signal_fn is not None: return self._load_signal_fn(port) - text = scrape_prometheus_metrics(port) + text = _scrape_prometheus_metrics(port) if text is None: return None - completed = parse_prometheus_metric( + completed = _parse_prometheus_metric( text, self._completed_metric, self._label_filter) - running = parse_prometheus_metric( + running = _parse_prometheus_metric( text, self._running_metric, self._label_filter) if completed is None or running is None: return None diff --git a/sdks/python/apache_beam/ml/inference/autoscaling_metrics_test.py b/sdks/python/apache_beam/ml/inference/autoscaling_metrics_test.py index 19fb4db24c40..aec5dae9afda 100644 --- a/sdks/python/apache_beam/ml/inference/autoscaling_metrics_test.py +++ b/sdks/python/apache_beam/ml/inference/autoscaling_metrics_test.py @@ -25,7 +25,7 @@ def _publisher(model_identifier='step-id', port=None, **kwargs): """A publisher for a fake server exposing ``done_total`` and ``running``.""" - return autoscaling_metrics.NumLoadedModelsPublisher( + return autoscaling_metrics.PrometheusNumLoadedModelsPublisher( metrics_port_fn=lambda: port, completed_metric='done_total', running_metric='running', @@ -49,10 +49,11 @@ def _published(model_identifier): class AutoscalingMetricsTest(unittest.TestCase): def test_sanitize_metric_tag(self): self.assertEqual( - autoscaling_metrics.sanitize_metric_tag('Qwen/Qwen2.5-7B:v1'), + autoscaling_metrics._sanitize_metric_tag('Qwen/Qwen2.5-7B:v1'), 'Qwen_Qwen2_5_7B_v1') self.assertEqual( - autoscaling_metrics.sanitize_metric_tag('normal_tag_1'), 'normal_tag_1') + autoscaling_metrics._sanitize_metric_tag('normal_tag_1'), + 'normal_tag_1') def test_parse_prometheus_metric(self): prom_text = ( @@ -61,17 +62,17 @@ def test_parse_prometheus_metric(self): 'vllm:num_requests_running{model_name="qwen"} 4.0\n' 'vllm:num_requests_running{model_name="other"} 6.0\n') self.assertEqual( - autoscaling_metrics.parse_prometheus_metric( + autoscaling_metrics._parse_prometheus_metric( prom_text, 'vllm:num_requests_running', label_filter={'model_name': 'qwen'}), 4.0) self.assertEqual( - autoscaling_metrics.parse_prometheus_metric( + autoscaling_metrics._parse_prometheus_metric( prom_text, 'vllm:num_requests_running'), 10.0) self.assertIsNone( - autoscaling_metrics.parse_prometheus_metric(prom_text, 'vllm:missing')) + autoscaling_metrics._parse_prometheus_metric(prom_text, 'vllm:missing')) def test_parse_prometheus_metric_exposition_details(self): prom_text = ( @@ -88,7 +89,7 @@ def test_parse_prometheus_metric_exposition_details(self): 'labeled{path="a,b}c",quote="say \\"hi\\"",n="x\\ny"} 7\n' ' \n' 'not a sample line\n') - parse = autoscaling_metrics.parse_prometheus_metric + parse = autoscaling_metrics._parse_prometheus_metric self.assertEqual(parse(prom_text, 'vllm:request_success_total'), 15.0) self.assertEqual(parse(prom_text, 'nv_inference_request_success'), 3.0) self.assertEqual( @@ -108,13 +109,13 @@ def fake_urlopen(url, timeout): with mock.patch('urllib.request.urlopen', side_effect=fake_urlopen): self.assertEqual( - autoscaling_metrics.scrape_prometheus_metrics(8000), 'running 4\n') + autoscaling_metrics._scrape_prometheus_metrics(8000), 'running 4\n') # Any failure (e.g. the server is restarting) -> None. with mock.patch('urllib.request.urlopen', side_effect=OSError('refused')): - self.assertIsNone(autoscaling_metrics.scrape_prometheus_metrics(8000)) + self.assertIsNone(autoscaling_metrics._scrape_prometheus_metrics(8000)) def test_capacity_estimator(self): - estimator = autoscaling_metrics.CompletionRateCapacityEstimator( + estimator = autoscaling_metrics._CompletionRateCapacityEstimator( multiplier=2.0, min_value=1, window_secs=60.0, running_percentile=90.0) # First sample establishes baseline -> None. self.assertIsNone(estimator.update(10.0, 4.0, now=100.0)) @@ -123,7 +124,7 @@ def test_capacity_estimator(self): self.assertAlmostEqual(estimator.update(30.0, 4.0, now=105.0), 8.0) def test_capacity_estimator_skips_counter_reset_and_clock_skew(self): - estimator = autoscaling_metrics.CompletionRateCapacityEstimator( + estimator = autoscaling_metrics._CompletionRateCapacityEstimator( window_secs=60.0, min_requests=0) self.assertIsNone(estimator.update(50.0, 2.0, now=100.0)) # Counter went down (server restart) -> skipped, new baseline. @@ -137,7 +138,7 @@ def test_capacity_estimator_skips_counter_reset_and_clock_skew(self): def test_capacity_estimator_running_percentile(self): def run(running_percentile): - estimator = autoscaling_metrics.CompletionRateCapacityEstimator( + estimator = autoscaling_metrics._CompletionRateCapacityEstimator( min_value=0, min_requests=0, running_percentile=running_percentile) estimator.update(0.0, 1.0, now=0.0) # Ten 1s intervals with 1..10 in flight, each completing as many @@ -152,7 +153,7 @@ def run(running_percentile): def test_capacity_estimator_warmup_min_requests_and_timeout(self): # Requires >= 3 completed requests before the 10s warmup timeout. - estimator = autoscaling_metrics.CompletionRateCapacityEstimator( + estimator = autoscaling_metrics._CompletionRateCapacityEstimator( min_requests=3, warmup_timeout_secs=10.0) self.assertIsNone(estimator.update(0.0, 2.0, now=100.0)) # Baseline. # 1 completion (< 3) at elapsed=5s (< 10s) -> None. @@ -161,7 +162,7 @@ def test_capacity_estimator_warmup_min_requests_and_timeout(self): # (4 / (2*5 + 2*3)) * 2 = 0.5 -> clamped to min_value=1.0. self.assertAlmostEqual(estimator.update(4.0, 2.0, now=108.0), 1.0) - slow = autoscaling_metrics.CompletionRateCapacityEstimator( + slow = autoscaling_metrics._CompletionRateCapacityEstimator( min_requests=3, warmup_timeout_secs=10.0) self.assertIsNone(slow.update(0.0, 1.0, now=100.0)) # Baseline. # 1 completion (< 3) at elapsed=5s -> None. @@ -170,7 +171,7 @@ def test_capacity_estimator_warmup_min_requests_and_timeout(self): self.assertAlmostEqual(slow.update(2.0, 1.0, now=110.0), 1.0) def test_capacity_estimator_skips_idle_ticks_without_diluting_p90(self): - estimator = autoscaling_metrics.CompletionRateCapacityEstimator( + estimator = autoscaling_metrics._CompletionRateCapacityEstimator( multiplier=1.0, min_value=1, window_secs=60.0, @@ -190,7 +191,7 @@ def test_capacity_estimator_skips_idle_ticks_without_diluting_p90(self): def test_capacity_estimator_no_estimate_until_a_request_completes(self): """In-flight requests with no completions never report the floor.""" # Slow first requests: still nothing completed after the warmup timeout. - slow = autoscaling_metrics.CompletionRateCapacityEstimator( + slow = autoscaling_metrics._CompletionRateCapacityEstimator( min_requests=3, warmup_timeout_secs=10.0) self.assertIsNone(slow.update(0.0, 4.0, now=100.0)) # Baseline. for t in (105.0, 110.0, 115.0): @@ -201,7 +202,7 @@ def test_capacity_estimator_no_estimate_until_a_request_completes(self): # First poll after an idle period longer than the window: the earlier # samples have left the window and the new requests haven't finished yet. for aggregation in ('concurrency', 'window_mean', 'window_max'): - resumed = autoscaling_metrics.CompletionRateCapacityEstimator( + resumed = autoscaling_metrics._CompletionRateCapacityEstimator( min_value=1, window_secs=10.0, min_requests=0, @@ -213,7 +214,7 @@ def test_capacity_estimator_no_estimate_until_a_request_completes(self): self.assertIsNone(resumed.update(20.0, 3.0, now=25.0), aggregation) # lifetime_max keeps reporting its peak. - peak = autoscaling_metrics.CompletionRateCapacityEstimator( + peak = autoscaling_metrics._CompletionRateCapacityEstimator( min_value=1, window_secs=10.0, min_requests=0, @@ -224,7 +225,7 @@ def test_capacity_estimator_no_estimate_until_a_request_completes(self): self.assertEqual(peak.update(20.0, 3.0, now=25.0), 4.0) def test_capacity_estimator_defaults_to_monotonic_clock(self): - estimator = autoscaling_metrics.CompletionRateCapacityEstimator( + estimator = autoscaling_metrics._CompletionRateCapacityEstimator( min_requests=0) with mock.patch.object(autoscaling_metrics.time, 'monotonic', @@ -235,7 +236,7 @@ def test_capacity_estimator_defaults_to_monotonic_clock(self): def test_capacity_estimator_completion_rate_aggregations(self): """The completion-rate modes ignore concurrency.""" def run(aggregation, **kwargs): - estimator = autoscaling_metrics.CompletionRateCapacityEstimator( + estimator = autoscaling_metrics._CompletionRateCapacityEstimator( min_value=0, min_requests=0, window_secs=60.0, @@ -255,7 +256,7 @@ def run(aggregation, **kwargs): def test_capacity_estimator_max_modes_differ_once_peak_leaves_window(self): def run(aggregation, window_secs=10.0): - estimator = autoscaling_metrics.CompletionRateCapacityEstimator( + estimator = autoscaling_metrics._CompletionRateCapacityEstimator( min_value=0, min_requests=0, window_secs=window_secs, @@ -273,9 +274,9 @@ def run(aggregation, window_secs=10.0): def test_capacity_estimator_completion_rate_works_without_in_flight(self): """A server whose gauge always reads 0 still gets a completion rate.""" - concurrency = autoscaling_metrics.CompletionRateCapacityEstimator( + concurrency = autoscaling_metrics._CompletionRateCapacityEstimator( min_requests=0) - rate = autoscaling_metrics.CompletionRateCapacityEstimator( + rate = autoscaling_metrics._CompletionRateCapacityEstimator( min_value=0, min_requests=0, aggregation='window_mean') for estimator in (concurrency, rate): estimator.update(0, 0, now=0.0) @@ -284,7 +285,7 @@ def test_capacity_estimator_completion_rate_works_without_in_flight(self): def test_capacity_estimator_rejects_unknown_aggregation(self): with self.assertRaises(ValueError): - autoscaling_metrics.CompletionRateCapacityEstimator(aggregation='p90') + autoscaling_metrics._CompletionRateCapacityEstimator(aggregation='p90') def test_publisher_publishes_signal_values_as_is(self): """Each capacity is published rounded; None or an error skips the tick.""" @@ -356,7 +357,7 @@ def read(): def test_publisher_for_same_tag_supersedes_previous(self): """A reloaded model's publisher takes over its tag from the stale one.""" - active = autoscaling_metrics.NumLoadedModelsPublisher._active + active = autoscaling_metrics.PrometheusNumLoadedModelsPublisher._active fast = autoscaling_metrics.AutoscalingMetricConfig(poll_interval_secs=0.01) first = _publisher('reload', config=fast) second = _publisher('reload', config=fast) @@ -428,7 +429,7 @@ def test_publisher_config_and_defaults(self): def test_publisher_reads_server_capacity(self): """The handler supplies metric names and a port; the rest is built in.""" port = [None] - publisher = autoscaling_metrics.NumLoadedModelsPublisher( + publisher = autoscaling_metrics.PrometheusNumLoadedModelsPublisher( metrics_port_fn=lambda: port[0], completed_metric='done_total', running_metric='running', @@ -476,7 +477,7 @@ def test_publisher_skips_unreadable_scrapes(self): 'running{model="a"} 2\nrunning{model="b"} 50\n', ] with mock.patch.object(autoscaling_metrics, - 'scrape_prometheus_metrics', + '_scrape_prometheus_metrics', side_effect=lambda port: scrapes.pop(0)): with mock.patch.object(autoscaling_metrics.time, 'monotonic',