AMOOOMA commented on code in PR #40760: URL: https://github.com/apache/beam/pull/40760#discussion_r4425809183
########## sdks/python/apache_beam/ml/inference/autoscaling_metrics.py: ########## @@ -0,0 +1,586 @@ +# +# 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. +# + +"""Publishes the ``num_loaded_models_<tag>`` autoscaling signal for a model +server running alongside a RunInference worker. + +Dataflow's horizontal autoscaler scales GPU/model pools on the process-wide +metric ``num_loaded_models_<tag>`` in the ``BeamML_ModelManager`` namespace +(see ``model_manager.py`` and +``//dist_proc/windmill/autoscaling/horizontal_scaling_policy.cc``), reading it +as *the number of messages each worker can hold*. A model server such as vLLM +or Triton always has exactly one loaded copy per worker, so instead the handler +publishes that metric with a per-worker capacity estimate derived from the +server's own Prometheus ``/metrics`` endpoint. + +Building blocks, from the bottom up: + +* :func:`scrape_prometheus_metrics` / :func:`parse_prometheus_metric` -- fetch + and read Prometheus text exposition. +* :class:`CompletionRateSignal` -- estimates per-worker capacity as completion + rate x in-flight concurrency over a trailing window (already smoothed). +* :class:`NumLoadedModelsPublisher` -- polls a signal on a background thread and + publishes it, optionally applying a policy such as ``'lifetime_max'`` on top. +* :class:`AutoscalingMetricConfig` -- the user-facing knobs; a handler takes + one as ``autoscaling_config`` and hands it to whatever owns the server + process. + +Nothing here is specific to a particular model server; the server-specific +piece is just the pair of metric names passed to :class:`CompletionRateSignal`. +""" + +# 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', + 'CompletionRateSignal', + 'NumLoadedModelsPublisher', + '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. Windmill sizes the pool 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 CompletionRateSignal. +DEFAULT_SIGNAL_WINDOW_SECS = 1800.0 # 30 minutes. +DEFAULT_RUNNING_PERCENTILE = 90.0 +DEFAULT_MIN_REQUESTS = 3 +DEFAULT_WARMUP_TIMEOUT_SECS = 10.0 +# CompletionRateSignal 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:<port><path>``; 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_<tag>`` 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<NN>'`` -- 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 CompletionRateSignal: Review Comment: Yes I was trying to make it generic, but wasn't sure what would the other pieces look like yet so it might look couploed with vLLM, originally i was hoping to do this in the triton PR but yeah i can try to update this now. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
