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..965765eafef8 --- /dev/null +++ b/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_`` 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_`` 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:``; 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 CompletionRateSignal: + """Stateful ``num_loaded_models`` signal: per-slot throughput x concurrency. + + Works for any model server that exposes a Prometheus ``/metrics`` endpoint + with a monotonically increasing *completed requests* counter + (``completed_metric``) and an *in-flight requests* gauge + (``running_metric``), e.g. vLLM's ``vllm:request_success_total`` / + ``vllm:num_requests_running``. Each call scrapes both from + ``http://localhost:/metrics``. 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 reports a per-worker message capacity:: + + reported = mu * running_ref * multiplier + + i.e. the worker's sustained throughput scaled by a dimensionless tuning + ``multiplier``. Windmill's GPU-intensive autoscaler sizes the pool as + ``desired = ceil(backlog / reported)`` (see + //dist_proc/windmill/autoscaling/horizontal_scaling_policy.cc). Normalizing by + ``running`` makes the estimate load-independent: right after an autoscaling + event, when Windmill rebalances keys 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 ``reported`` stays + steady instead of cratering to the floor and triggering a downscale spiral. + + The output 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 capacity estimate). + """ + def __init__( + self, + completed_metric: str, + running_metric: str, + label_filter: Optional[Mapping[str, str]] = None, + 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._completed_metric = completed_metric + self._running_metric = running_metric + self._label_filter = dict(label_filter) if label_filter else None + 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 __call__(self, port: int, timeout_secs: float = 5.0) -> Optional[float]: + text = scrape_prometheus_metrics(port, timeout_secs=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 + now = time.time() + prev_count, prev_ts = self._last_count, self._last_ts + self._last_count, self._last_ts = completed, 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 - 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) + reported = mu * running_ref * self._multiplier + return max(float(self._min_value), reported) + + +@dataclass(frozen=True) +class AutoscalingMetricConfig: + """User-facing knobs for the ``num_loaded_models_`` autoscaling signal. + + Pass an instance to a model handler (e.g. + ``VLLMCompletionsModelHandler(..., autoscaling_config=...)``) 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 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:`CompletionRateSignal` 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 signal. + + Whatever owns the server process calls :meth:`make_completion_rate_signal` / + :meth:`make_publisher` on the resolved config. + """ + 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(self, default_tag: str) -> 'AutoscalingMetricConfig': + """Returns a copy whose ``tag`` is set, falling back to ``default_tag``. + + ``default_tag`` is used verbatim: when it is the identifier RunInference + reports for the step it must match that annotation exactly. + """ + if self.tag: + return self + if not default_tag: + raise ValueError('default_tag must be a non-empty string.') + return dataclasses.replace(self, tag=default_tag) + + def make_completion_rate_signal( + self, + completed_metric: str, + running_metric: str, + label_filter: Optional[Mapping[str, str]] = None, + ) -> CompletionRateSignal: + return CompletionRateSignal( + completed_metric=completed_metric, + running_metric=running_metric, + label_filter=label_filter, + 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(default_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..ac890f565ee4 --- /dev/null +++ b/sdks/python/apache_beam/ml/inference/autoscaling_metrics_test.py @@ -0,0 +1,320 @@ +# +# 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_completion_rate_signal(self): + signal = autoscaling_metrics.CompletionRateSignal( + completed_metric='vllm:request_success_total', + running_metric='vllm:num_requests_running', + multiplier=2.0, + min_value=1, + window_secs=60.0, + running_percentile=90.0) + responses = [ + ('vllm:request_success_total 10.0\n' + 'vllm:num_requests_running 4.0\n'), + ('vllm:request_success_total 30.0\n' + 'vllm:num_requests_running 4.0\n'), + ] + + def fake_urlopen(url, timeout=5.0): + del url, timeout + body = responses.pop(0).encode('utf-8') + return io.BytesIO(body) + + with mock.patch('urllib.request.urlopen', side_effect=fake_urlopen): + with mock.patch.object(autoscaling_metrics.time, + 'time', + side_effect=[100.0, 105.0]): + # First sample establishes baseline -> None. + self.assertIsNone(signal(8000)) + # Second sample: 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(signal(8000), 8.0) + + def test_completion_rate_signal_label_filter(self): + signal = autoscaling_metrics.CompletionRateSignal( + completed_metric='nv_inference_request_success', + running_metric='nv_inference_pending_request_count', + label_filter={'model': 'a'}, + window_secs=60.0) + responses = [ + ( + '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'), + ( + 'nv_inference_request_success{model="a"} 20\n' + 'nv_inference_request_success{model="b"} 999\n' + 'nv_inference_pending_request_count{model="a"} 2\n' + 'nv_inference_pending_request_count{model="b"} 50\n'), + ] + + def fake_urlopen(url, timeout=5.0): + del url, 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, 110.0]): + self.assertIsNone(signal(8002)) + # Only model "a" counts: delta=10 over 10s with running=2 -> + # mu = 0.5, running_ref = 2 -> 1.0. + self.assertAlmostEqual(signal(8002), 1.0) + + def test_completion_rate_signal_warmup_min_requests_and_timeout(self): + # Requires >= 3 completed requests before 10s warmup timeout. + signal = autoscaling_metrics.CompletionRateSignal( + completed_metric='vllm:request_success_total', + running_metric='vllm:num_requests_running', + min_requests=3, + warmup_timeout_secs=10.0) + responses = [ + # t=100: baseline. + 'vllm:request_success_total 0.0\nvllm:num_requests_running 2.0\n', + # t=105: only 1 request completed (< 3) and elapsed=5s (< 10s) -> None. + 'vllm:request_success_total 1.0\nvllm:num_requests_running 2.0\n', + # t=108: 4 total requests completed (>= 3) at elapsed=8s (< 10s) -> + # publishes (4 / (2*5 + 2*3) * 2 = 0.5 -> clamped to min_value=1.0). + 'vllm:request_success_total 4.0\nvllm:num_requests_running 2.0\n', + ] + + def fake_urlopen(url, timeout=5.0): + del url, 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, 108.0]): + self.assertIsNone(signal(8000)) + self.assertIsNone(signal(8000)) + self.assertAlmostEqual(signal(8000), 1.0) + + # Also verify warmup_timeout_secs: if < 3 requests complete, it publishes + # once elapsed >= 10s from the first active interval. + slow_signal = autoscaling_metrics.CompletionRateSignal( + completed_metric='vllm:request_success_total', + running_metric='vllm:num_requests_running', + min_requests=3, + warmup_timeout_secs=10.0) + slow_responses = [ + # t=100: baseline. + 'vllm:request_success_total 0.0\nvllm:num_requests_running 1.0\n', + # t=105: 1 request (< 3), elapsed=5s (< 10s) -> None. + 'vllm:request_success_total 1.0\nvllm:num_requests_running 1.0\n', + # t=110: 2 requests (< 3), elapsed=10s (>= 10s) -> publishes. + 'vllm:request_success_total 2.0\nvllm:num_requests_running 1.0\n', + ] + + def fake_slow_urlopen(url, timeout=5.0): + del url, timeout + return io.BytesIO(slow_responses.pop(0).encode('utf-8')) + + with mock.patch('urllib.request.urlopen', side_effect=fake_slow_urlopen): + with mock.patch.object(autoscaling_metrics.time, + 'time', + side_effect=[100.0, 105.0, 110.0]): + self.assertIsNone(slow_signal(8000)) + self.assertIsNone(slow_signal(8000)) + self.assertAlmostEqual(slow_signal(8000), 1.0) + + def test_completion_rate_signal_skips_idle_ticks_without_diluting_p90(self): + signal = autoscaling_metrics.CompletionRateSignal( + completed_metric='vllm:request_success_total', + running_metric='vllm:num_requests_running', + multiplier=1.0, + min_value=1, + window_secs=60.0, + running_percentile=90.0, + min_requests=1) + responses = [ + # t=100: baseline. + 'vllm:request_success_total 0.0\nvllm:num_requests_running 5.0\n', + # t=105: 25 completions over 5s with running=5 -> mu=1.0, p90=5 -> 5.0. + 'vllm:request_success_total 25.0\nvllm:num_requests_running 5.0\n', + # t=110..120: 3 completely idle ticks (delta=0, running=0) -> None, + # without adding r=0 samples that would dilute p90 concurrency. + 'vllm:request_success_total 25.0\nvllm:num_requests_running 0.0\n', + 'vllm:request_success_total 25.0\nvllm:num_requests_running 0.0\n', + 'vllm:request_success_total 25.0\nvllm:num_requests_running 0.0\n', + # t=125: 5 completions over 5s that already finished before scrape + # (delta=5, running=0) -> uses positive running samples for p90=5. + 'vllm:request_success_total 30.0\nvllm:num_requests_running 0.0\n', + ] + + def fake_urlopen(url, timeout=5.0): + del url, 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, + 110.0, + 115.0, + 120.0, + 125.0]): + self.assertIsNone(signal(8000)) + self.assertAlmostEqual(signal(8000), 5.0) + self.assertIsNone(signal(8000)) + self.assertIsNone(signal(8000)) + self.assertIsNone(signal(8000)) + # total_completed = 30, total_slot_secs = 25 -> mu = 1.2, p90 = 5 -> 6.0 + self.assertAlmostEqual(signal(8000), 6.0) + + def test_autoscaling_metric_config_resolve(self): + # An unset tag is derived from the model name; an explicit one is kept. + 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).resolve('step-id') + # The default tag is used verbatim (it has to match the annotation). + self.assertEqual(cfg.tag, 'step-id') + self.assertEqual( + autoscaling_metrics.AutoscalingMetricConfig( + tag='custom').resolve('step-id').tag, + 'custom') + with self.assertRaises(ValueError): + autoscaling_metrics.AutoscalingMetricConfig().resolve('') + with self.assertRaises(ValueError): + autoscaling_metrics.AutoscalingMetricConfig().make_publisher(lambda: None) + signal = cfg.make_completion_rate_signal('done_total', 'running') + self.assertEqual(signal._multiplier, 2.0) + self.assertEqual(signal._window_secs, 30.0) + self.assertEqual(signal._min_requests, 5) + self.assertEqual(signal._warmup_timeout_secs, 15.0) + publisher = cfg.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 signal, never a second time here. + self.assertEqual(publisher._multiplier, 1.0) + + +if __name__ == '__main__': + unittest.main() diff --git a/sdks/python/apache_beam/ml/inference/base.py b/sdks/python/apache_beam/ml/inference/base.py index c2e084d0e03f..5ad199e72e19 100644 --- a/sdks/python/apache_beam/ml/inference/base.py +++ b/sdks/python/apache_beam/ml/inference/base.py @@ -210,6 +210,10 @@ def __init__( element_size_fn: Optional[Callable[[Any], int]] = None, batch_length_fn: Optional[Callable[[Any], int]] = None, batch_bucket_boundaries: Optional[list[int]] = None, + batch_key_fn: Optional[Callable[[Any], Any]] = None, + batch_key_num_shards: Optional[Union[int, + Mapping[Any, int], + Callable[[Any], int]]] = None, large_model: bool = False, model_copies: Optional[int] = None, **kwargs): @@ -240,6 +244,19 @@ def __init__( semantics): bucket i covers lengths in [boundaries[i-1], boundaries[i]). Requires batch_length_fn. Defaults to [16, 32, 64, 128, 256, 512] when batch_length_fn is set. + batch_key_fn: a callable mapping an element to a grouping key. It + receives the element as it enters batching, i.e. *after* any + ``with_preprocess_fn`` functions have been applied. When set, + BatchElements keys elements via WithBatchKey (without a per-worker UUID) + and groups them across workers using GroupIntoBatches so that elements + sharing the same key (e.g., prompt prefix or category) are routed to + the same worker and batched together. + batch_key_num_shards: bounds hot-key fanout when batch_key_fn is set. + Either an int applied to every key, a mapping from batch key to shard + count for a known skewed distribution (keys not in the mapping get 1 + shard), or a callable from batch key to shard count. Each key is + spread across at most that many random sub-keys / workers. Defaults + to 1. large_model: set to true if your model is large enough to run into memory pressure if you load multiple copies. model_copies: The exact number of models that you would like loaded @@ -270,6 +287,10 @@ def __init__( self._batching_kwargs['length_fn'] = batch_length_fn if batch_bucket_boundaries is not None: self._batching_kwargs['bucket_boundaries'] = batch_bucket_boundaries + if batch_key_fn is not None: + self._batching_kwargs['key_fn'] = batch_key_fn + if batch_key_num_shards is not None: + self._batching_kwargs['num_shards'] = batch_key_num_shards self._large_model = large_model self._model_copies = model_copies self._share_across_processes = large_model or (model_copies is not None) @@ -357,6 +378,19 @@ def update_model_paths( """ pass + def set_model_identifier(self, model_identifier: str): + """Receives the identifier RunInference annotates this model with. + + Called by RunInference in the worker before the first ``load_model`` with + the ``model_identifier`` passed to ``RunInference`` (or the random id it + generated when none was passed). This is the value Dataflow's autoscaler + uses to find the ``num_loaded_models_`` metric for this step, + and it is stable across side-input model updates, so handlers that publish + that metric themselves should tag it with this identifier. The default is + a no-op. + """ + pass + def get_preprocess_fns(self) -> Iterable[Callable[[Any], Any]]: """Gets all preprocessing functions to be run before batching/inference. Functions are in order that they should be applied.""" @@ -1122,6 +1156,12 @@ def update_model_path(self, model_path: Optional[str] = None): 'KeyedModelHandlers with multiple different per-key ' + 'ModelHandlers.') + def set_model_identifier(self, model_identifier: str): + if self._single_model: + return self._unkeyed.set_model_identifier(model_identifier) + for mh in self._id_to_mh_map.values(): + mh.set_model_identifier(model_identifier) + def share_model_across_processes(self) -> bool: if self._single_model: return self._unkeyed.share_model_across_processes() @@ -1248,6 +1288,9 @@ def validate_inference_args(self, inference_args: Optional[dict[str, Any]]): def update_model_path(self, model_path: Optional[str] = None): return self._unkeyed.update_model_path(model_path=model_path) + def set_model_identifier(self, model_identifier: str): + return self._unkeyed.set_model_identifier(model_identifier) + def get_preprocess_fns(self) -> Iterable[Callable[[Any], Any]]: return self._unkeyed.get_preprocess_fns() @@ -1319,6 +1362,9 @@ def validate_inference_args(self, inference_args: Optional[dict[str, Any]]): def update_model_path(self, model_path: Optional[str] = None): return self._base.update_model_path(model_path=model_path) + def set_model_identifier(self, model_identifier: str): + return self._base.set_model_identifier(model_identifier) + def get_preprocess_fns(self) -> Iterable[Callable[[Any], Any]]: return self._base.get_preprocess_fns() @@ -1397,6 +1443,9 @@ def validate_inference_args(self, inference_args: Optional[dict[str, Any]]): def update_model_path(self, model_path: Optional[str] = None): return self._base.update_model_path(model_path=model_path) + def set_model_identifier(self, model_identifier: str): + return self._base.set_model_identifier(model_identifier) + def get_preprocess_fns(self) -> Iterable[Callable[[Any], Any]]: return [self._preprocess_fn] + self._base.get_preprocess_fns() @@ -1474,6 +1523,9 @@ def validate_inference_args(self, inference_args: Optional[dict[str, Any]]): def update_model_path(self, model_path: Optional[str] = None): return self._base.update_model_path(model_path=model_path) + def set_model_identifier(self, model_identifier: str): + return self._base.set_model_identifier(model_identifier) + def get_preprocess_fns(self) -> Iterable[Callable[[Any], Any]]: return self._base.get_preprocess_fns() @@ -2276,6 +2328,7 @@ def get_metrics_collector(self, prefix: str = ''): def setup(self): self._metrics_collector = self.get_metrics_collector() self._model_handler.set_environment_vars() + self._model_handler.set_model_identifier(self._model_tag) self._model_metadata = load_model_status( self._model_tag, self._model_handler.share_model_across_processes()) if not self._load_model_at_runtime: diff --git a/sdks/python/apache_beam/ml/inference/base_test.py b/sdks/python/apache_beam/ml/inference/base_test.py index f1bb5d1245e1..9abef15dfd66 100644 --- a/sdks/python/apache_beam/ml/inference/base_test.py +++ b/sdks/python/apache_beam/ml/inference/base_test.py @@ -2799,6 +2799,20 @@ def test_batching_kwargs_none_values_omitted(self): self.assertNotIn('bucket_boundaries', kwargs) self.assertEqual(kwargs['min_batch_size'], 5) + def test_batch_key_fn_and_num_shards_passed_to_kwargs(self): + key_fn = lambda s: s[:4] + handler = FakeModelHandlerForBatching( + batch_key_fn=key_fn, batch_key_num_shards=4) + kwargs = handler.batch_elements_kwargs() + self.assertIs(kwargs['key_fn'], key_fn) + self.assertEqual(kwargs['num_shards'], 4) + + def test_batch_key_num_shards_mapping_passed_to_kwargs(self): + shard_map = {'hot_prefix': 8} + handler = FakeModelHandlerForBatching( + batch_key_fn=lambda s: s, batch_key_num_shards=shard_map) + self.assertEqual(handler.batch_elements_kwargs()['num_shards'], shard_map) + class PaddingReportingStringModelHandler(base.ModelHandler[str, str, FakeModel]): @@ -2811,6 +2825,79 @@ def run_inference(self, batch, model, inference_args=None): return [f'{s}:{max_len}' for s in batch] +class IdentifierReportingModelHandler(base.ModelHandler[int, str, FakeModel]): + """Emits the identifier RunInference reported, as seen at load_model time. + + The DoFn works on a pickled copy of the handler, so the only way to observe + the hook is through the pipeline output. + """ + def __init__(self): + super().__init__() + self._identifier = None + self._identifier_at_load = None + + def set_model_identifier(self, model_identifier): + self._identifier = model_identifier + + def load_model(self): + # Captures whether set_model_identifier ran before load_model. + self._identifier_at_load = self._identifier + return FakeModel() + + def run_inference(self, batch, model, inference_args=None): + return [f'{x}:{self._identifier_at_load}' for x in batch] + + +class RunInferenceModelIdentifierTest(unittest.TestCase): + """RunInference hands handlers the identifier the autoscaler looks up.""" + def test_explicit_model_identifier_is_reported_before_load(self): + with TestPipeline('FnApiRunner') as p: + out = ( + p | beam.Create([1, 2], reshuffle=False) + | base.RunInference( + IdentifierReportingModelHandler(), + model_identifier='my-step-model')) + assert_that(out, equal_to(['1:my-step-model', '2:my-step-model'])) + + def test_generated_uuid_identifier_matches_annotation(self): + # No model_identifier given: RunInference generates a uuid hex, annotates + # the inference step with it, and reports that same value to the handler. + with TestPipeline('FnApiRunner') as p: + out = ( + p | beam.Create([1], reshuffle=False) + | base.RunInference(IdentifierReportingModelHandler())) + annotated = [ + t.annotations['model_identifier'].decode() + for t in p.to_runner_api().components.transforms.values() + if t.unique_name == 'RunInference/BeamML_RunInference' + ] + self.assertEqual(len(annotated), 1) + self.assertRegex(annotated[0], r'^[0-9a-f]{32}$') + assert_that(out, equal_to([f'1:{annotated[0]}'])) + + def test_identifier_forwarded_through_wrappers(self): + handler = base.KeyedModelHandler( + IdentifierReportingModelHandler()).with_preprocess_fn( + lambda kv: kv).with_postprocess_fn(lambda kv: kv) + with TestPipeline('FnApiRunner') as p: + out = ( + p | beam.Create([('k', 1)], reshuffle=False) + | base.RunInference(handler, model_identifier='wrapped-id')) + assert_that(out, equal_to([('k', '1:wrapped-id')])) + + +class PrefixPeerReportingStringModelHandler(base.ModelHandler[str, + str, + FakeModel]): + """Reports each element with the sorted prefix set of the batch it ran in.""" + def load_model(self): + return FakeModel() + + def run_inference(self, batch, model, inference_args=None): + prefixes = ','.join(sorted({s.split(':', 1)[0] for s in batch})) + return [f'{s}|batch_prefixes={prefixes}|size={len(batch)}' for s in batch] + + class RunInferenceLengthAwareBatchingTest(unittest.TestCase): """End-to-end tests for PR2 length-aware batching in RunInference.""" def test_run_inference_with_length_aware_batch_elements(self): @@ -2829,6 +2916,44 @@ def test_run_inference_with_length_aware_batch_elements(self): | base.RunInference(handler)) assert_that(results, equal_to(['a:2', 'bb:2', 'cccccc:7', 'ddddddd:7'])) + def test_run_inference_with_batch_key_fn_and_preprocess_fn(self): + # batch_key_fn receives the element *after* preprocessing (the prompt + # string here), not the raw input record. + handler = PrefixPeerReportingStringModelHandler( + min_batch_size=2, + max_batch_size=2, + max_batch_duration_secs=1.0, + batch_key_fn=lambda prompt: prompt.split(':', 1)[0]).with_preprocess_fn( + lambda rec: f"{rec['sector']}:{rec['id']}") + + records = [ + { + 'sector': 'fin', 'id': '1' + }, + { + 'sector': 'med', 'id': '1' + }, + { + 'sector': 'fin', 'id': '2' + }, + { + 'sector': 'med', 'id': '2' + }, + ] + with TestPipeline('FnApiRunner') as p: + results = ( + p + | beam.Create(records, reshuffle=False) + | base.RunInference(handler)) + assert_that( + results, + equal_to([ + 'fin:1|batch_prefixes=fin|size=2', + 'fin:2|batch_prefixes=fin|size=2', + 'med:1|batch_prefixes=med|size=2', + 'med:2|batch_prefixes=med|size=2', + ])) + class BatchLengthsModelHandler(base.ModelHandler[str, tuple, FakeModel]): """Predicts the lengths of the examples in the batch of each example.""" diff --git a/sdks/python/apache_beam/ml/inference/prefix_batching.py b/sdks/python/apache_beam/ml/inference/prefix_batching.py new file mode 100644 index 000000000000..5b8cc02d3df7 --- /dev/null +++ b/sdks/python/apache_beam/ml/inference/prefix_batching.py @@ -0,0 +1,199 @@ +# +# 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. +# + +"""Prefix-aware batching for RunInference model handlers that front an LLM +server with a prefix (KV) cache, such as vLLM. + +``ModelHandler(batch_key_fn=...)`` makes ``BatchElements`` route elements that +share a key to the same worker via ``GroupIntoBatches``, so each batch the +server sees is homogeneous and its prefix cache gets hit instead of thrashed. +This module packages the knobs for that into one :class:`PrefixBatchingConfig` +plus the default key extractors a handler installs when the user does not +supply their own: + +* :class:`PromptPrefixKeyFn` -- for plain-text prompts (completions API). +* :class:`ChatPrefixKeyFn` -- for chat message lists (chat API). + +Both see the element *after* any ``with_preprocess_fn`` has run, i.e. the +prompt / message list the handler will actually send. Nothing here is specific +to a particular server; a handler only decides which extractor is its default. +""" + +# pytype: skip-file + +import dataclasses +from collections.abc import Callable +from collections.abc import Mapping +from collections.abc import Sequence +from dataclasses import dataclass +from typing import Any +from typing import Optional +from typing import Union + +__all__ = [ + 'ChatPrefixKeyFn', + 'PrefixBatchingConfig', + 'PromptPrefixKeyFn', +] + +# Characters of prompt (or leading message content) the default extractors key +# on. Long enough to cover a typical system prompt, short enough that keys are +# cheap to hash and shuffle. +DEFAULT_PREFIX_CHAR_LEN = 128 +# GroupIntoBatches buffering applied when the handler's own +# ``max_batch_duration_secs`` is unset, so low-traffic keys still flush in +# streaming instead of waiting for a full batch. +DEFAULT_KEYED_BATCH_DURATION_SECS = 1.0 + +# Accepted forms of ``num_shards``: one count for every key, a per-key table, +# or a callable computed per key (see ``WithBatchKey`` in transforms/util.py). +NumShardsSpec = Union[int, Mapping[Any, int], Callable[[Any], int]] + + +class PromptPrefixKeyFn: + """Default prefix key extractor for plain-text prompts. + + Keys a prompt string by its leading ``prefix_char_len`` characters, or by + everything before the first ``\\n\\n`` when that comes sooner, so prompts + sharing the same system prompt or instruction header are routed to the same + worker and batched together. + """ + def __init__(self, prefix_char_len: int = DEFAULT_PREFIX_CHAR_LEN): + if prefix_char_len <= 0: + raise ValueError('prefix_char_len must be positive.') + self._prefix_char_len = prefix_char_len + + def __call__(self, element: str) -> str: + if not isinstance(element, str): + raise TypeError( + f'Expected str prompt for prefix key extraction, got ' + f'{type(element).__name__}.') + sep_idx = element.find('\n\n') + if 0 < sep_idx <= self._prefix_char_len: + return element[:sep_idx] + return element[:self._prefix_char_len] + + +class ChatPrefixKeyFn: + """Default prefix key extractor for chat message lists. + + Keys a conversation by the role and leading ``prefix_char_len`` characters of + its first (typically system) message, so conversations sharing the same + system prompt are routed to the same worker and batched together. Messages + may be objects with ``role``/``content`` attributes or dicts with those keys. + """ + def __init__(self, prefix_char_len: int = DEFAULT_PREFIX_CHAR_LEN): + if prefix_char_len <= 0: + raise ValueError('prefix_char_len must be positive.') + self._prefix_char_len = prefix_char_len + + def __call__(self, element: Sequence[Any]) -> str: + if not isinstance(element, (list, tuple)): + raise TypeError( + f'Expected a sequence of chat messages for prefix key extraction, ' + f'got {type(element).__name__}.') + if not element: + return '' + first_msg = element[0] + role = getattr(first_msg, 'role', None) + content = getattr(first_msg, 'content', None) + if role is None and isinstance(first_msg, dict): + role = first_msg.get('role', '') + content = first_msg.get('content', '') + content_str = str(content or '') + return f'{role or ""}:{content_str[:self._prefix_char_len]}' + + +@dataclass(frozen=True) +class PrefixBatchingConfig: + """User-facing knobs for prefix-aware batching. + + Pass an instance to a model handler (e.g. + ``VLLMCompletionsModelHandler(..., prefix_batching_config=...)``) to turn + keyed batching on; leave it unset to keep the default unkeyed batching. Each + field: + + * ``key_fn`` -- callable mapping an element (as it enters batching, i.e. + after any ``with_preprocess_fn``) to a grouping key. Leave unset (the + default) to use the handler's default prefix extractor + (:class:`PromptPrefixKeyFn` or :class:`ChatPrefixKeyFn`). + * ``prefix_char_len`` -- how many leading characters the default extractor + keys on (default 128). Ignored when ``key_fn`` is set. + * ``num_shards`` -- bounds hot-key fanout: an ``int`` applied to every key, + a ``Mapping[key, int]`` with per-key shard counts for a known skewed + prefix distribution (e.g. ``{'': 8}``; unlisted + keys get 1), or a ``Callable[[key], int]``. Each key is spread across at + most that many workers (default 1). + + The handler calls :meth:`resolve` with its default extractor and then + :meth:`batching_kwargs` to obtain the ``ModelHandler`` arguments. + """ + key_fn: Optional[Callable[[Any], Any]] = None + prefix_char_len: int = DEFAULT_PREFIX_CHAR_LEN + num_shards: Optional[NumShardsSpec] = None + + def resolve( + self, default_key_fn: Callable[[int], Callable[[Any], Any]] + ) -> 'PrefixBatchingConfig': + """Returns a copy whose ``key_fn`` is set. + + Args: + default_key_fn: Factory (e.g. :class:`PromptPrefixKeyFn`) invoked with + ``prefix_char_len`` when ``key_fn`` is unset. + """ + if self.key_fn is not None: + return self + return dataclasses.replace( + self, key_fn=default_key_fn(self.prefix_char_len)) + + def batching_kwargs(self, + max_batch_duration_secs: Optional[float] = None + ) -> dict[str, Any]: + """Returns the ``ModelHandler`` batching kwargs for a resolved config. + + Args: + max_batch_duration_secs: The handler's own buffering duration; when + unset, ``DEFAULT_KEYED_BATCH_DURATION_SECS`` is used so low-traffic + keys still flush in streaming. + """ + if self.key_fn is None: + raise ValueError( + 'PrefixBatchingConfig.key_fn is unset; call resolve(default_key_fn) ' + 'first.') + return { + 'batch_key_fn': self.key_fn, + 'batch_key_num_shards': self.num_shards, + 'max_batch_duration_secs': ( + DEFAULT_KEYED_BATCH_DURATION_SECS + if max_batch_duration_secs is None else max_batch_duration_secs), + } + + +def resolve_batching_kwargs( + config: Optional[PrefixBatchingConfig], + default_key_fn: Callable[[int], Callable[[Any], Any]], + max_batch_duration_secs: Optional[float] = None, +) -> dict[str, Any]: + """Batching kwargs for an optional config; unkeyed when ``config`` is None. + + Convenience for handlers: always yields ``max_batch_duration_secs`` and, when + prefix batching is on, ``batch_key_fn`` / ``batch_key_num_shards`` too, so + the result can be splatted straight into ``ModelHandler.__init__``. + """ + if config is None: + return {'max_batch_duration_secs': max_batch_duration_secs} + return config.resolve(default_key_fn).batching_kwargs(max_batch_duration_secs) diff --git a/sdks/python/apache_beam/ml/inference/prefix_batching_test.py b/sdks/python/apache_beam/ml/inference/prefix_batching_test.py new file mode 100644 index 000000000000..5e991ce443fa --- /dev/null +++ b/sdks/python/apache_beam/ml/inference/prefix_batching_test.py @@ -0,0 +1,99 @@ +# +# 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 unittest + +from apache_beam.ml.inference import prefix_batching +from apache_beam.ml.inference.prefix_batching import ChatPrefixKeyFn +from apache_beam.ml.inference.prefix_batching import PrefixBatchingConfig +from apache_beam.ml.inference.prefix_batching import PromptPrefixKeyFn + + +class _Msg: + def __init__(self, role, content): + self.role = role + self.content = content + + +class PrefixBatchingTest(unittest.TestCase): + def test_prompt_prefix_key_fn(self): + key_fn = PromptPrefixKeyFn(prefix_char_len=5) + self.assertEqual(key_fn('abcdefghij'), 'abcde') + # Stops at the first blank line when it comes before prefix_char_len. + self.assertEqual(PromptPrefixKeyFn(prefix_char_len=20)('ab\n\ncdef'), 'ab') + self.assertEqual(PromptPrefixKeyFn(prefix_char_len=1)('ab\n\ncdef'), 'a') + with self.assertRaises(TypeError): + key_fn({'prompt': 'abcdefghij'}) + with self.assertRaises(ValueError): + PromptPrefixKeyFn(prefix_char_len=0) + + def test_chat_prefix_key_fn(self): + key_fn = ChatPrefixKeyFn(prefix_char_len=3) + msgs = [_Msg('system', 'You are A'), _Msg('user', 'Hello')] + self.assertEqual(key_fn(msgs), 'system:You') + # Dict-style messages and empty conversations are handled too. + self.assertEqual( + key_fn([{ + 'role': 'system', 'content': 'You are B' + }]), 'system:You') + self.assertEqual(key_fn([]), '') + with self.assertRaises(TypeError): + key_fn({'messages': msgs}) + with self.assertRaises(ValueError): + ChatPrefixKeyFn(prefix_char_len=0) + + def test_config_resolve_and_batching_kwargs(self): + # Default: the handler's extractor is built with prefix_char_len. + cfg = PrefixBatchingConfig(prefix_char_len=32, num_shards=4) + resolved = cfg.resolve(PromptPrefixKeyFn) + self.assertIsInstance(resolved.key_fn, PromptPrefixKeyFn) + self.assertEqual(resolved.key_fn._prefix_char_len, 32) + self.assertEqual( + resolved.batching_kwargs(), + { + 'batch_key_fn': resolved.key_fn, + 'batch_key_num_shards': 4, + 'max_batch_duration_secs': prefix_batching. + DEFAULT_KEYED_BATCH_DURATION_SECS, + }) + # An explicit handler duration wins over the keyed default. + self.assertEqual( + resolved.batching_kwargs( + max_batch_duration_secs=2.5)['max_batch_duration_secs'], + 2.5) + # A custom key_fn is kept verbatim. + custom = PrefixBatchingConfig(key_fn=len).resolve(PromptPrefixKeyFn) + self.assertIs(custom.key_fn, len) + with self.assertRaises(ValueError): + PrefixBatchingConfig().batching_kwargs() + + def test_resolve_batching_kwargs_without_config(self): + # No config -> unkeyed batching with the handler's duration passed through. + self.assertEqual( + prefix_batching.resolve_batching_kwargs(None, PromptPrefixKeyFn, 3.0), + {'max_batch_duration_secs': 3.0}) + self.assertEqual( + prefix_batching.resolve_batching_kwargs(None, PromptPrefixKeyFn), + {'max_batch_duration_secs': None}) + kwargs = prefix_batching.resolve_batching_kwargs( + PrefixBatchingConfig(), ChatPrefixKeyFn) + self.assertIsInstance(kwargs['batch_key_fn'], ChatPrefixKeyFn) + self.assertIsNone(kwargs['batch_key_num_shards']) + self.assertEqual(kwargs['max_batch_duration_secs'], 1.0) + + +if __name__ == '__main__': + unittest.main() diff --git a/sdks/python/apache_beam/ml/inference/vllm_inference.py b/sdks/python/apache_beam/ml/inference/vllm_inference.py index e3bc203494c7..cd4ff60010cb 100644 --- a/sdks/python/apache_beam/ml/inference/vllm_inference.py +++ b/sdks/python/apache_beam/ml/inference/vllm_inference.py @@ -38,8 +38,12 @@ from openai import OpenAI from apache_beam.io.filesystems import FileSystems +from apache_beam.ml.inference import autoscaling_metrics +from apache_beam.ml.inference import prefix_batching +from apache_beam.ml.inference.autoscaling_metrics import AutoscalingMetricConfig from apache_beam.ml.inference.base import ModelHandler from apache_beam.ml.inference.base import PredictionResult +from apache_beam.ml.inference.prefix_batching import PrefixBatchingConfig from apache_beam.utils import subprocess_server try: @@ -59,6 +63,30 @@ 'VLLMChatModelHandler', ] +# vLLM Prometheus metrics used by the ``num_loaded_models`` autoscaling signal +# (see apache_beam.ml.inference.autoscaling_metrics). +_VLLM_COMPLETED_REQUESTS_METRIC = 'vllm:request_success_total' +_VLLM_RUNNING_REQUESTS_METRIC = 'vllm:num_requests_running' + + +def _resolve_autoscaling( + enabled: bool, + config: Optional[AutoscalingMetricConfig], + model_identifier: Optional[str], + model_name: str) -> Optional[AutoscalingMetricConfig]: + """Picks the ``num_loaded_models`` tag for a vLLM handler, or None if off. + + Precedence: an explicit ``config.tag``; else the identifier RunInference + reported via ``set_model_identifier`` (what the autoscaler looks up, stable + across model updates); else a sanitized ``model_name`` for handlers used + outside RunInference. + """ + if not enabled: + return None + config = config or AutoscalingMetricConfig() + return config.resolve( + model_identifier or autoscaling_metrics.sanitize_metric_tag(model_name)) + @dataclass(frozen=True) class OpenAIChatMessage(): @@ -148,7 +176,8 @@ def __init__( model_name: str, vllm_server_kwargs: dict[str, Optional[str]], dynamo_frontend_kwargs: Optional[dict[str, Optional[str]]] = None, - use_dynamo: bool = False): + use_dynamo: bool = False, + autoscaling: Optional[AutoscalingMetricConfig] = None): self._model_name = model_name self._vllm_server_kwargs = vllm_server_kwargs self._dynamo_frontend_kwargs = dynamo_frontend_kwargs or {} @@ -161,9 +190,36 @@ def __init__( self._server_port: int = -1 self._server_process_lock = threading.RLock() self._use_dynamo = use_dynamo + # Autoscaling: a capacity signal polled by a publisher thread that is + # started once the server is confirmed up (see start_server) and stopped + # in __del__. ``autoscaling=None`` disables the metric entirely. + self._autoscaling = autoscaling + self._completion_rate_signal = ( + autoscaling.make_completion_rate_signal( + _VLLM_COMPLETED_REQUESTS_METRIC, _VLLM_RUNNING_REQUESTS_METRIC) + if autoscaling else None) + self._num_loaded_models_publisher: Optional[ + autoscaling_metrics.NumLoadedModelsPublisher] = ( + autoscaling.make_publisher(self._get_load_signal) + if autoscaling else None) self.start_server() + def _get_load_signal(self) -> Optional[float]: + """Returns the current per-worker capacity estimate, or None. + + Returns None while the server is down/restarting so the publisher skips the + tick instead of reporting a misleading 0. + """ + if not self._server_started or self._server_port < 0: + return None + if self._autoscaling is not None and (self._autoscaling.load_signal_fn + is not None): + return self._autoscaling.load_signal_fn(self._server_port) + if self._completion_rate_signal is None: + return None + return self._completion_rate_signal(self._server_port) + @staticmethod def _stop_process(process: Optional[subprocess.Popen]) -> None: if process is None or process.poll() is not None: @@ -214,6 +270,11 @@ def __del__(self): # __del__ may run during interpreter shutdown when module globals can # already be torn down; swallow any cleanup failures so we don't print # a noisy traceback. + try: + if self._num_loaded_models_publisher is not None: + self._num_loaded_models_publisher.stop() + except Exception: # pylint: disable=broad-except + pass try: self._stop_processes() except Exception: # pylint: disable=broad-except @@ -322,6 +383,8 @@ def start_server(self, retries=3): self._dynamo_process, _ = start_process(server_cmd) self.check_connectivity(retries) + if self._num_loaded_models_publisher is not None: + self._num_loaded_models_publisher.start() def get_server_port(self) -> int: if not self._server_started: @@ -374,7 +437,10 @@ def __init__( max_batch_weight: Optional[int] = None, element_size_fn: Optional[Callable[[Any], int]] = None, batch_length_fn: Optional[Callable[[Any], int]] = None, - batch_bucket_boundaries: Optional[list[int]] = None): + batch_bucket_boundaries: Optional[list[int]] = None, + prefix_batching_config: Optional[PrefixBatchingConfig] = None, + enable_autoscaling_metric: bool = True, + autoscaling_config: Optional[AutoscalingMetricConfig] = None): """Implementation of the ModelHandler interface for vLLM using text as input. @@ -425,15 +491,37 @@ def __init__( batch_bucket_boundaries: optional. a sorted list of positive boundary values for length-aware batching buckets. Cannot be used together with max_batch_weight. + prefix_batching_config: optional ``PrefixBatchingConfig``. When set, + prompts sharing a prefix are shuffled via ``GroupIntoBatches`` to the + same worker and batched together to maximize vLLM prefix cache hits. + By default prompts are keyed on their leading ``prefix_char_len`` + characters (or up to the first ``\\n\\n`` delimiter); set ``key_fn`` + for a custom key and ``num_shards`` to bound hot-key fanout (see + ``apache_beam.ml.inference.prefix_batching``). + enable_autoscaling_metric: if True (default), publish the + ``num_loaded_models_`` autoscaling signal: the worker's estimated + message capacity derived from the vLLM server's request completion + rate and in-flight concurrency (see + ``apache_beam.ml.inference.autoscaling_metrics``). + autoscaling_config: optional ``AutoscalingMetricConfig`` overriding the + signal's defaults (poll interval, multiplier, windows, aggregation + policy, or a custom ``load_signal_fn``). The metric tag is taken from + the identifier RunInference reports for this step + (``model_identifier``, or its generated id), so it matches what the + autoscaler looks up and survives model updates; set ``tag`` only to + override that. """ super().__init__( min_batch_size=min_batch_size, max_batch_size=max_batch_size, - max_batch_duration_secs=max_batch_duration_secs, max_batch_weight=max_batch_weight, element_size_fn=element_size_fn, batch_length_fn=batch_length_fn, - batch_bucket_boundaries=batch_bucket_boundaries) + batch_bucket_boundaries=batch_bucket_boundaries, + **prefix_batching.resolve_batching_kwargs( + prefix_batching_config, + prefix_batching.PromptPrefixKeyFn, + max_batch_duration_secs)) self._model_name = model_name self._vllm_server_kwargs: dict[str, Optional[str]] = ({ **_DYNAMO_ENGINE_DEFAULT_KWARGS, **(vllm_server_kwargs or {}) @@ -442,13 +530,24 @@ def __init__( **_DYNAMO_FRONTEND_DEFAULT_KWARGS, **(dynamo_frontend_kwargs or {}) } self._use_dynamo = use_dynamo + self._enable_autoscaling_metric = enable_autoscaling_metric + self._autoscaling_config = autoscaling_config + self._model_identifier: Optional[str] = None + + def set_model_identifier(self, model_identifier: str): + self._model_identifier = model_identifier def load_model(self) -> _VLLMModelServer: return _VLLMModelServer( self._model_name, self._vllm_server_kwargs, self._dynamo_frontend_kwargs, - self._use_dynamo) + self._use_dynamo, + autoscaling=_resolve_autoscaling( + self._enable_autoscaling_metric, + self._autoscaling_config, + self._model_identifier, + self._model_name)) async def _async_run_inference( self, @@ -517,7 +616,10 @@ def __init__( max_batch_weight: Optional[int] = None, element_size_fn: Optional[Callable[[Any], int]] = None, batch_length_fn: Optional[Callable[[Any], int]] = None, - batch_bucket_boundaries: Optional[list[int]] = None): + batch_bucket_boundaries: Optional[list[int]] = None, + prefix_batching_config: Optional[PrefixBatchingConfig] = None, + enable_autoscaling_metric: bool = True, + autoscaling_config: Optional[AutoscalingMetricConfig] = None): """ Implementation of the ModelHandler interface for vLLM using previous messages as input. @@ -571,15 +673,37 @@ def __init__( batch_bucket_boundaries: optional. a sorted list of positive boundary values for length-aware batching buckets. Cannot be used together with max_batch_weight. + prefix_batching_config: optional ``PrefixBatchingConfig``. When set, + conversations sharing a prefix are shuffled via ``GroupIntoBatches`` + to the same worker and batched together to maximize vLLM prefix cache + hits. By default conversations are keyed on the role and leading + ``prefix_char_len`` characters of their first (system) message; set + ``key_fn`` for a custom key and ``num_shards`` to bound hot-key fanout + (see ``apache_beam.ml.inference.prefix_batching``). + enable_autoscaling_metric: if True (default), publish the + ``num_loaded_models_`` autoscaling signal: the worker's estimated + message capacity derived from the vLLM server's request completion + rate and in-flight concurrency (see + ``apache_beam.ml.inference.autoscaling_metrics``). + autoscaling_config: optional ``AutoscalingMetricConfig`` overriding the + signal's defaults (poll interval, multiplier, windows, aggregation + policy, or a custom ``load_signal_fn``). The metric tag is taken from + the identifier RunInference reports for this step + (``model_identifier``, or its generated id), so it matches what the + autoscaler looks up and survives model updates; set ``tag`` only to + override that. """ super().__init__( min_batch_size=min_batch_size, max_batch_size=max_batch_size, - max_batch_duration_secs=max_batch_duration_secs, max_batch_weight=max_batch_weight, element_size_fn=element_size_fn, batch_length_fn=batch_length_fn, - batch_bucket_boundaries=batch_bucket_boundaries) + batch_bucket_boundaries=batch_bucket_boundaries, + **prefix_batching.resolve_batching_kwargs( + prefix_batching_config, + prefix_batching.ChatPrefixKeyFn, + max_batch_duration_secs)) self._model_name = model_name self._vllm_server_kwargs: dict[str, Optional[str]] = ({ **_DYNAMO_ENGINE_DEFAULT_KWARGS, **(vllm_server_kwargs or {}) @@ -590,6 +714,12 @@ def __init__( self._chat_template_path = chat_template_path self._chat_file = f'template-{uuid.uuid4().hex}.jinja' self._use_dynamo = use_dynamo + self._enable_autoscaling_metric = enable_autoscaling_metric + self._autoscaling_config = autoscaling_config + self._model_identifier: Optional[str] = None + + def set_model_identifier(self, model_identifier: str): + self._model_identifier = model_identifier def load_model(self) -> _VLLMModelServer: chat_template_contents = '' @@ -606,7 +736,12 @@ def load_model(self) -> _VLLMModelServer: self._model_name, self._vllm_server_kwargs, self._dynamo_frontend_kwargs, - self._use_dynamo) + self._use_dynamo, + autoscaling=_resolve_autoscaling( + self._enable_autoscaling_metric, + self._autoscaling_config, + self._model_identifier, + self._model_name)) async def _async_run_inference( self, diff --git a/sdks/python/apache_beam/ml/inference/vllm_inference_test.py b/sdks/python/apache_beam/ml/inference/vllm_inference_test.py index 4ff5186178e3..3e45b7c763ae 100644 --- a/sdks/python/apache_beam/ml/inference/vllm_inference_test.py +++ b/sdks/python/apache_beam/ml/inference/vllm_inference_test.py @@ -20,6 +20,16 @@ import unittest from unittest import mock +import apache_beam as beam +from apache_beam.ml.inference import prefix_batching +from apache_beam.ml.inference.autoscaling_metrics import AutoscalingMetricConfig +from apache_beam.ml.inference.base import PredictionResult +from apache_beam.ml.inference.base import RunInference +from apache_beam.ml.inference.prefix_batching import PrefixBatchingConfig +from apache_beam.testing.test_pipeline import TestPipeline +from apache_beam.testing.util import assert_that +from apache_beam.testing.util import equal_to + # Protect against environments where the OpenAI python library is not # available. The command-construction tests below do not actually need a # real OpenAI client; stubbing the module is enough for vllm_inference to @@ -81,6 +91,24 @@ def start_process(cmd): return start_process +class _RecordingVLLMHandler(vllm_inference.VLLMCompletionsModelHandler): + """Stub VLLMCompletionsModelHandler that records batches without a server.""" + def __init__(self, **kwargs): + super().__init__(**kwargs) + self.recorded_batches = [] + + def load_model(self): + return self.recorded_batches + + def run_inference(self, batch, model, inference_args=None): + batch_list = list(batch) + model.append(batch_list) + return [ + PredictionResult(example=prompt, inference={'batch': batch_list}) + for prompt in batch_list + ] + + class VLLMInferenceTest(unittest.TestCase): def test_native_vllm_starts_single_server_process(self): commands = [] @@ -93,7 +121,8 @@ def test_native_vllm_starts_single_server_process(self): model_name='test-model', vllm_server_kwargs={ 'gpu-memory-utilization': '0.9' - }).load_model() + }, + enable_autoscaling_metric=False).load_model() self.assertEqual(1, len(commands)) self.assertIn('vllm.entrypoints.openai.api_server', commands[0]) self.assertIn('--model', commands[0]) @@ -120,7 +149,8 @@ def test_dynamo_starts_frontend_and_engine_with_separate_kwargs(self): use_dynamo=True, dynamo_frontend_kwargs={ 'router-mode': 'round-robin' - }).load_model() + }, + enable_autoscaling_metric=False).load_model() self.assertEqual(2, len(commands)) frontend_cmd = commands[0] engine_cmd = commands[1] @@ -152,6 +182,139 @@ def test_validate_inference_args_accepts_openai_request_kwargs(self): vllm_inference.VLLMChatModelHandler('test-model').validate_inference_args( {'max_tokens': 8}) + def test_autoscaling_defaults_do_not_double_smooth(self): + """The signal smooths; the publisher must forward it unchanged by default.""" + handler = vllm_inference.VLLMCompletionsModelHandler( + model_name='test-model', + autoscaling_config=AutoscalingMetricConfig( + multiplier=3.0, signal_window_secs=120.0)) + with mock.patch.object(vllm_inference._VLLMModelServer, 'start_server'): + server = handler.load_model() + publisher = server._num_loaded_models_publisher + self.assertIsNotNone(publisher) + self.assertEqual(publisher._aggregation, 'last') + # The multiplier lives in the signal; the publisher must not apply a + # second one. + self.assertEqual(publisher._multiplier, 1.0) + signal = server._completion_rate_signal + self.assertEqual(signal._multiplier, 3.0) + self.assertEqual(signal._window_secs, 120.0) + self.assertEqual(signal._completed_metric, 'vllm:request_success_total') + self.assertEqual(signal._running_metric, 'vllm:num_requests_running') + self.assertEqual(publisher._metric_name, 'num_loaded_models_test_model') + + def test_autoscaling_tag_prefers_run_inference_identifier(self): + """The metric tag is the identifier RunInference reports (uuid or given). + + That is what the autoscaler looks up and it is stable across model + updates; the sanitized model name is only a fallback for handlers used + outside RunInference, and an explicit config tag always wins. + """ + def tag_for(handler): + with mock.patch.object(vllm_inference._VLLMModelServer, 'start_server'): + return handler.load_model()._num_loaded_models_publisher._metric_name + + # Outside RunInference: sanitized model name. + handler = vllm_inference.VLLMCompletionsModelHandler( + model_name='org/model:v1') + self.assertEqual(tag_for(handler), 'num_loaded_models_org_model_v1') + + # Inside RunInference: the reported identifier, verbatim (uuid-like or + # user-provided with characters sanitize would have changed). + handler.set_model_identifier('3f2a9c0e1b4d4f6e8a7b6c5d4e3f2a1b') + self.assertEqual( + tag_for(handler), 'num_loaded_models_3f2a9c0e1b4d4f6e8a7b6c5d4e3f2a1b') + handler.set_model_identifier('my-step/model') + self.assertEqual(tag_for(handler), 'num_loaded_models_my-step/model') + + # Explicit config tag overrides both. + pinned = vllm_inference.VLLMCompletionsModelHandler( + model_name='org/model:v1', + autoscaling_config=AutoscalingMetricConfig(tag='pinned')) + pinned.set_model_identifier('ignored') + self.assertEqual(tag_for(pinned), 'num_loaded_models_pinned') + + # Chat handler follows the same rule. + chat = vllm_inference.VLLMChatModelHandler(model_name='org/model:v1') + chat.set_model_identifier('chat-id') + self.assertEqual(tag_for(chat), 'num_loaded_models_chat-id') + + def test_prefix_batching_configuration(self): + comp_handler = vllm_inference.VLLMCompletionsModelHandler( + model_name='test-model', + prefix_batching_config=PrefixBatchingConfig( + prefix_char_len=32, num_shards=4)) + comp_kwargs = comp_handler.batch_elements_kwargs() + self.assertIsInstance( + comp_kwargs['key_fn'], prefix_batching.PromptPrefixKeyFn) + self.assertEqual(comp_kwargs['key_fn']._prefix_char_len, 32) + self.assertEqual(comp_kwargs['num_shards'], 4) + self.assertEqual(comp_kwargs['max_batch_duration_secs'], 1.0) + + chat_handler = vllm_inference.VLLMChatModelHandler( + model_name='test-model', + prefix_batching_config=PrefixBatchingConfig(prefix_char_len=64), + max_batch_duration_secs=2.5) + chat_kwargs = chat_handler.batch_elements_kwargs() + self.assertIsInstance( + chat_kwargs['key_fn'], prefix_batching.ChatPrefixKeyFn) + self.assertEqual(chat_kwargs['key_fn']._prefix_char_len, 64) + self.assertEqual(chat_kwargs['max_batch_duration_secs'], 2.5) + + # A custom key_fn is used verbatim; unset config leaves batching unkeyed. + custom_handler = vllm_inference.VLLMCompletionsModelHandler( + model_name='test-model', + prefix_batching_config=PrefixBatchingConfig(key_fn=len)) + self.assertIs(custom_handler.batch_elements_kwargs()['key_fn'], len) + plain_handler = vllm_inference.VLLMCompletionsModelHandler( + model_name='test-model') + self.assertNotIn('key_fn', plain_handler.batch_elements_kwargs()) + + def test_run_inference_with_prefix_batching_and_preprocess_fn(self): + records = [ + { + 'id': 1, 'prompt': 'PREFIX_ALPHA: question 1' + }, + { + 'id': 2, 'prompt': 'PREFIX_BETA_: question 2' + }, + { + 'id': 3, 'prompt': 'PREFIX_ALPHA: question 3' + }, + { + 'id': 4, 'prompt': 'PREFIX_BETA_: question 4' + }, + ] + handler = _RecordingVLLMHandler( + model_name='test-model', + prefix_batching_config=PrefixBatchingConfig(prefix_char_len=12), + min_batch_size=2, + max_batch_size=2, + max_batch_duration_secs=1.0).with_preprocess_fn(lambda r: r['prompt']) + + with TestPipeline('FnApiRunner') as p: + preds = ( + p + | beam.Create(records, reshuffle=False) + | RunInference(handler)) + grouped_batches = preds | beam.Map( + lambda r: (r.example, sorted(r.inference['batch']))) + expected = [ + ( + 'PREFIX_ALPHA: question 1', + ['PREFIX_ALPHA: question 1', 'PREFIX_ALPHA: question 3']), + ( + 'PREFIX_ALPHA: question 3', + ['PREFIX_ALPHA: question 1', 'PREFIX_ALPHA: question 3']), + ( + 'PREFIX_BETA_: question 2', + ['PREFIX_BETA_: question 2', 'PREFIX_BETA_: question 4']), + ( + 'PREFIX_BETA_: question 4', + ['PREFIX_BETA_: question 2', 'PREFIX_BETA_: question 4']), + ] + assert_that(grouped_batches, equal_to(expected)) + if __name__ == '__main__': unittest.main() diff --git a/sdks/python/apache_beam/transforms/util.py b/sdks/python/apache_beam/transforms/util.py index df9607977950..568c626bec44 100644 --- a/sdks/python/apache_beam/transforms/util.py +++ b/sdks/python/apache_beam/transforms/util.py @@ -33,6 +33,7 @@ import uuid from collections.abc import Callable from collections.abc import Iterable +from collections.abc import Mapping from typing import TYPE_CHECKING from typing import Any from typing import Optional @@ -1206,6 +1207,91 @@ def process(self, element): yield ((self.key, bucket), element) +def _validate_num_shards(num_shards: Any) -> Any: + """Validates a ``num_shards`` spec for ``WithBatchKey`` / ``BatchElements``. + + Accepts an int (uniform shard count), a mapping from batch key to shard count + (keys absent from the mapping get 1 shard), or a callable from batch key to + shard count. Returns the normalized spec (``None`` becomes ``1``). + """ + if num_shards is None: + return 1 + if isinstance(num_shards, bool): + raise ValueError('num_shards must be an int, mapping, or callable.') + if isinstance(num_shards, int): + if num_shards < 1: + raise ValueError(f'num_shards ({num_shards}) must be >= 1.') + return num_shards + if isinstance(num_shards, Mapping): + bad = { + k: v + for k, v in num_shards.items() + if not (isinstance(v, int) and not isinstance(v, bool) and v >= 1) + } + if bad: + raise ValueError( + f'num_shards mapping values must be ints >= 1; got {bad}.') + return dict(num_shards) + if callable(num_shards): + return num_shards + raise ValueError( + 'num_shards must be an int, a mapping from key to int, or a callable ' + f'from key to int; got {type(num_shards).__name__}.') + + +class WithBatchKey(DoFn): + """Keys elements with a deterministic string key for cross-worker + GroupIntoBatches batching. + + Unlike ``WithSharedKey`` and ``WithLengthBucketKey`` (which include a + per-worker UUID so elements remain worker-local), ``WithBatchKey`` omits any + worker UUID so that elements sharing the same key across all upstream workers + shuffle to the same downstream worker and are grouped into homogeneous + batches (e.g., for LLM prefix caching). + + ``num_shards`` bounds hot-key fanout: when a key resolves to ``n > 1`` + shards, a random suffix ``:`` in ``[0, n)`` is appended so that key is + spread across at most ``n`` workers. It may be + + * an ``int`` -- the same shard count for every key; + * a ``Mapping[key, int]`` -- per-key shard counts for a known skewed + distribution (e.g. ``{'default_system_prompt': 8}``); keys not in the + mapping get 1 shard; + * a ``Callable[[key], int]`` -- computed per key (e.g. from a frequency + table loaded at ``setup``). + + The mapping / callable is looked up with the raw ``key_fn`` result first and + its ``str()`` form second, so either may be used as the mapping key. + """ + def __init__(self, key_fn: Callable[[Any], Any], num_shards: Any = 1): + self._key_fn = key_fn + self._num_shards = _validate_num_shards(num_shards) + + def _shards_for(self, raw_key: Any, key: str) -> int: + spec = self._num_shards + if isinstance(spec, int): + return spec + if isinstance(spec, Mapping): + n = spec.get(raw_key) + if n is None: + n = spec.get(key, 1) + return n + n = spec(raw_key) + if not isinstance(n, int) or n < 1: + raise ValueError( + f'num_shards callable returned {n!r} for key {key!r}; expected an ' + 'int >= 1.') + return n + + def process(self, element): + raw_key = self._key_fn(element) + key = str(raw_key) if raw_key is not None else '' + num_shards = self._shards_for(raw_key, key) + if num_shards > 1: + key = f'{key}:{random.randrange(num_shards)}' + yield (key, element) + + @typehints.with_input_types(T) @typehints.with_output_types(list[T]) class BatchElements(PTransform): @@ -1283,6 +1369,16 @@ class BatchElements(PTransform): semantics): e.g., for boundaries [10, 50], buckets are (-inf, 10), [10, 50), [50, inf). Defaults to [16, 32, 64, 128, 256, 512] when length_fn is set. Requires length_fn. + key_fn: (optional) a callable mapping an element to a grouping key. When + set, elements are keyed via ``WithBatchKey`` (without a per-worker UUID) + and grouped across workers using ``GroupIntoBatches`` up to + ``max_batch_size`` (and ``max_batch_duration_secs`` if provided), + producing homogeneous batches per key (e.g., for LLM prefix caching). + num_shards: (optional) bounds hot-key fanout when ``key_fn`` is set: an + ``int`` applied to every key, a ``Mapping[key, int]`` giving per-key + shard counts for a known skewed distribution (unlisted keys get 1), or + a ``Callable[[key], int]``. Each key is spread across at most that + many random sub-keys / workers. Defaults to 1. See ``WithBatchKey``. """ _DEFAULT_BUCKET_BOUNDARIES = [16, 32, 64, 128, 256, 512] @@ -1301,7 +1397,9 @@ def __init__( clock=time.time, record_metrics=True, length_fn=None, - bucket_boundaries=None): + bucket_boundaries=None, + key_fn=None, + num_shards=1): if max_batch_weight is not None: if element_size_fn is None: raise ValueError('max_batch_weight requires element_size_fn to be set.') @@ -1320,6 +1418,11 @@ def __init__( _LOGGER.warning( 'BatchElements ignores length_fn and bucket_boundaries unless ' 'max_batch_duration_secs is set.') + if key_fn is not None and not callable(key_fn): + raise ValueError('key_fn must be callable.') + num_shards = _validate_num_shards(num_shards) + if num_shards != 1 and key_fn is None: + raise ValueError('num_shards requires key_fn to be set.') self._batch_size_estimator = _BatchSizeEstimator( min_batch_size=min_batch_size, max_batch_size=max_batch_size, @@ -1339,9 +1442,26 @@ def __init__( self._bucket_boundaries = self._DEFAULT_BUCKET_BOUNDARIES else: self._bucket_boundaries = bucket_boundaries + self._key_fn = key_fn + self._num_shards = num_shards def expand(self, pcoll): - if getattr(pcoll.pipeline.runner, 'is_streaming', False): + if self._key_fn is not None: + element_type = pcoll.element_type + buffering_secs = ( + self._max_batch_dur if self._max_batch_dur is not None else 0) + return ( + pcoll + | 'WithBatchKey' >> ParDo( + WithBatchKey(self._key_fn, self._num_shards)).with_output_types( + tuple[str, element_type]) + | 'GroupIntoBatchesByKey' >> GroupIntoBatches( + self._batch_size_estimator._max_batch_size, + max_buffering_duration_secs=buffering_secs, + clock=self._clock) + | 'DropBatchKey' >> Map(lambda kv: list(kv[1])).with_output_types( + list[element_type])) + elif getattr(pcoll.pipeline.runner, 'is_streaming', False): raise NotImplementedError("Requires stateful processing (BEAM-2687)") elif self._max_batch_dur is not None: coder = coders.registry.get_coder(pcoll) diff --git a/sdks/python/apache_beam/transforms/util_test.py b/sdks/python/apache_beam/transforms/util_test.py index 4456fea18536..ff76271e629d 100644 --- a/sdks/python/apache_beam/transforms/util_test.py +++ b/sdks/python/apache_beam/transforms/util_test.py @@ -1281,6 +1281,90 @@ def check_batch_homogeneity(batch): checks = batches | beam.Map(check_batch_homogeneity) assert_that(checks, is_not_empty()) + def test_batch_elements_key_fn_groups_homogeneous_batches(self): + """BatchElements with key_fn groups elements by key via GroupIntoBatches.""" + elements = ['sysA:q1', 'sysB:q1', 'sysA:q2', 'sysB:q2', 'sysA:q3'] + with TestPipeline('FnApiRunner') as p: + batches = ( + p + | beam.Create(elements, reshuffle=False) + | util.BatchElements( + min_batch_size=1, + max_batch_size=2, + max_batch_duration_secs=1.0, + key_fn=lambda s: s.split(':', 1)[0])) + + def check_homogeneous(batch_list): + flattened = [] + for b in batch_list: + prefixes = {s.split(':', 1)[0] for s in b} + assert len(prefixes) == 1, f'Mixed prefixes in batch: {b}' + assert len(b) <= 2, f'Batch exceeded max_batch_size: {b}' + flattened.extend(b) + assert sorted(flattened) == sorted(elements) + + assert_that(batches, check_homogeneous) + + def test_with_batch_key_num_shards(self): + """WithBatchKey appends a bounded random shard suffix to the key.""" + dofn = util.WithBatchKey(key_fn=lambda d: d['cat'], num_shards=4) + element = {'cat': 'math', 'q': '1+1'} + out = list(dofn.process(element)) + self.assertEqual(len(out), 1) + key, val = out[0] + self.assertIs(val, element) + self.assertTrue(key.startswith('math:')) + shard = int(key.split(':', 1)[1]) + self.assertIn(shard, range(4)) + + def test_with_batch_key_num_shards_mapping_handles_skew(self): + """A per-key shard mapping fans out only the listed hot keys.""" + dofn = util.WithBatchKey( + key_fn=lambda d: d['cat'], num_shards={ + 'hot': 4, 'warm': 2 + }) + hot_keys = {list(dofn.process({'cat': 'hot'}))[0][0] for _ in range(200)} + warm_keys = {list(dofn.process({'cat': 'warm'}))[0][0] for _ in range(200)} + cold_keys = {list(dofn.process({'cat': 'cold'}))[0][0] for _ in range(200)} + self.assertEqual(hot_keys, {f'hot:{i}' for i in range(4)}) + self.assertEqual(warm_keys, {'warm:0', 'warm:1'}) + # Keys absent from the mapping are not sharded at all. + self.assertEqual(cold_keys, {'cold'}) + + def test_with_batch_key_num_shards_mapping_matches_str_key(self): + """Mapping lookup falls back to the str() form of the key.""" + dofn = util.WithBatchKey(key_fn=lambda d: d['tenant'], num_shards={'7': 3}) + keys = {list(dofn.process({'tenant': 7}))[0][0] for _ in range(100)} + self.assertEqual(keys, {'7:0', '7:1', '7:2'}) + + def test_with_batch_key_num_shards_callable(self): + """A callable computes the shard count per key.""" + dofn = util.WithBatchKey( + key_fn=lambda s: s, num_shards=lambda k: 3 if k == 'hot' else 1) + hot_keys = {list(dofn.process('hot'))[0][0] for _ in range(100)} + self.assertEqual(hot_keys, {'hot:0', 'hot:1', 'hot:2'}) + self.assertEqual(list(dofn.process('cold')), [('cold', 'cold')]) + bad = util.WithBatchKey(key_fn=lambda s: s, num_shards=lambda k: 0) + with self.assertRaises(ValueError): + list(bad.process('x')) + + def test_batch_elements_key_fn_validation(self): + with self.assertRaises(ValueError): + util.BatchElements(key_fn='not_callable') + with self.assertRaises(ValueError): + util.BatchElements(key_fn=len, num_shards=0) + with self.assertRaises(ValueError): + util.BatchElements(num_shards=2) + with self.assertRaises(ValueError): + util.BatchElements(num_shards={'a': 2}) + with self.assertRaises(ValueError): + util.BatchElements(key_fn=len, num_shards={'a': 0}) + with self.assertRaises(ValueError): + util.BatchElements(key_fn=len, num_shards='4') + # Valid specs. + util.BatchElements(key_fn=len, num_shards={'a': 2}) + util.BatchElements(key_fn=len, num_shards=lambda k: 2) + def test_global_batching_dofn_single_vs_multiple_bundles(self): # This test directly verifies how bundling affects the batch sizes produced by # the internal _GlobalWindowsBatchingDoFn of BatchElements.