diff --git a/CHANGES.md b/CHANGES.md index bcf4911d5f9b..9300cbda4eb3 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -101,6 +101,7 @@ * [Beam SQL] Added `PRIMARY KEY (...) NOT ENFORCED` to `CREATE EXTERNAL TABLE`, and `ALTER TABLE ... ADD/DROP PRIMARY KEY` ([#40826](https://github.com/apache/beam/issues/40826)). * [Beam SQL] Added a `WITH (...)` clause to set per-statement connector options in `FROM` and `INSERT INTO` ([#40848](https://github.com/apache/beam/issues/40848)). * (Python) `BatchElements` can group elements by key across workers (`key_fn`, `num_shards`), and RunInference model handlers can use it through `batch_key_fn` / `batch_key_num_shards`. `apache_beam.ml.inference.prefix_batching` adds a `PrefixBatchingConfig` and prompt/chat prefix extractors so LLM handlers can batch requests that share a prompt prefix and hit the server's prefix cache ([#40905](https://github.com/apache/beam/pull/40905)). +* (Python) vLLM model handlers accept a `prefix_batching_config` to batch requests that share a prompt prefix, and publish a capacity-based autoscaling signal (`num_loaded_models_`, estimated from vLLM's completion rate and in-flight requests) so Dataflow's GPU autoscaling scales vLLM workers with load. The signal is on by default; tune or disable it with `autoscaling_config` / `enable_autoscaling_metric`. Other model-server handlers can publish the same signal with `apache_beam.ml.inference.autoscaling_metrics.PrometheusNumLoadedModelsPublisher` ([#40907](https://github.com/apache/beam/pull/40907), [#40908](https://github.com/apache/beam/pull/40908)). ### Performance diff --git a/sdks/python/apache_beam/ml/inference/vllm_inference.py b/sdks/python/apache_beam/ml/inference/vllm_inference.py index e3bc203494c7..50b5adce7bee 100644 --- a/sdks/python/apache_beam/ml/inference/vllm_inference.py +++ b/sdks/python/apache_beam/ml/inference/vllm_inference.py @@ -20,13 +20,16 @@ import asyncio import logging import os +import random import shutil +import socket import subprocess import sys import threading import time import urllib.request import uuid +import weakref from collections.abc import Callable from collections.abc import Iterable from collections.abc import Sequence @@ -38,8 +41,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 +66,19 @@ 'VLLMChatModelHandler', ] +# The vLLM-specific inputs to the generic capacity-based autoscaling signal +# (see apache_beam.ml.inference.autoscaling_metrics): vLLM's Prometheus counter +# of completed requests and gauge of in-flight requests. +_VLLM_COMPLETED_REQUESTS_METRIC = 'vllm:request_success_total' +_VLLM_RUNNING_REQUESTS_METRIC = 'vllm:num_requests_running' + + +def _autoscaling_config( + enabled: bool, config: Optional[AutoscalingMetricConfig] +) -> Optional[AutoscalingMetricConfig]: + """The config to publish the autoscaling signal with, or None if disabled.""" + return (config or AutoscalingMetricConfig()) if enabled else None + @dataclass(frozen=True) class OpenAIChatMessage(): @@ -71,12 +91,28 @@ class OpenAIChatMessage(): content: str -def start_process(cmd) -> tuple[subprocess.Popen, int]: +def start_process(cmd, + env: Optional[dict[str, str]] = None + ) -> tuple[subprocess.Popen, int]: + """Starts ``cmd`` with ``{{PORT}}`` replaced by a free port. + + Args: + cmd: The command; ``{{PORT}}`` in any argument is replaced by the port. + env: Extra environment variables, added to this process's environment. + + Returns: + The process and the port picked for ``{{PORT}}``. + """ port, = subprocess_server.pick_port(None) cmd = [arg.replace('{{PORT}}', str(port)) for arg in cmd] # pylint: disable=not-an-iterable logging.info("Starting service with %s", str(cmd).replace("',", "'")) process = subprocess.Popen( - cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT) + cmd, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + env={ + **os.environ, **env + } if env else None) # Emit the output of this command as info level logging. def log_stdout(): @@ -93,6 +129,25 @@ def log_stdout(): return process, port +# Dynamo parses DYN_SYSTEM_PORT as a signed 16-bit integer, so the worker's +# system port must be below 32768, under the usual ephemeral port range. +_DYNAMO_SYSTEM_PORT_RANGE = (20000, 32767) + + +def _pick_dynamo_system_port() -> int: + """Returns a free port that Dynamo accepts for DYN_SYSTEM_PORT.""" + low, high = _DYNAMO_SYSTEM_PORT_RANGE + for port in random.sample(range(low, high + 1), 100): + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + try: + s.bind(('localhost', port)) + except OSError: + continue + return port + raise RuntimeError( + f'No free port in {low}-{high} for the Dynamo worker system server.') + + def getVLLMClient(port) -> OpenAI: openai_api_key = "EMPTY" openai_api_base = f"http://localhost:{port}/v1" @@ -142,13 +197,33 @@ def _uses_etcd_discovery(kwargs: dict[str, Optional[str]]) -> bool: return kwargs.get('discovery-backend') == 'etcd' +def _weak_metrics_port_fn( + server: '_VLLMModelServer') -> Callable[[], Optional[int]]: + """Returns ``server._metrics_port`` without keeping ``server`` alive. + + The autoscaling publisher's thread holds this callable for as long as it + runs. A strong reference to the server would keep it from ever being + garbage-collected, so ``_VLLMModelServer.__del__``, which stops the publisher + and the vLLM processes, would never run. + """ + method = weakref.WeakMethod(server._metrics_port) + + def metrics_port() -> Optional[int]: + bound = method() + return None if bound is None else bound() + + return metrics_port + + class _VLLMModelServer(): def __init__( self, 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, + model_identifier: Optional[str] = None): self._model_name = model_name self._vllm_server_kwargs = vllm_server_kwargs self._dynamo_frontend_kwargs = dynamo_frontend_kwargs or {} @@ -159,11 +234,38 @@ def __init__( self._etcd_data_dir: Optional[str] = None self._managed_etcd_endpoint = None self._server_port: int = -1 + # In Dynamo mode, the port of the dynamo.vllm worker's system server, which + # serves vLLM's metrics; the frontend on _server_port does not. + self._dynamo_metrics_port: int = -1 self._server_process_lock = threading.RLock() self._use_dynamo = use_dynamo + # Autoscaling: a publisher thread, started once the server is confirmed up + # (see start_server) and stopped in __del__, reports this worker's capacity + # from vLLM's request metrics. ``autoscaling=None`` disables it. + self._num_loaded_models_publisher: Optional[ + autoscaling_metrics.PrometheusNumLoadedModelsPublisher] = ( + autoscaling_metrics.PrometheusNumLoadedModelsPublisher( + metrics_port_fn=_weak_metrics_port_fn(self), + completed_metric=_VLLM_COMPLETED_REQUESTS_METRIC, + running_metric=_VLLM_RUNNING_REQUESTS_METRIC, + model_identifier=model_identifier, + model_name=model_name, + config=autoscaling) if autoscaling else None) self.start_server() + def _metrics_port(self) -> Optional[int]: + """The port serving vLLM's ``/metrics`` while the server is up, or None. + + That is the vLLM API server itself, or in Dynamo mode the dynamo.vllm + worker's system server: the Dynamo frontend only serves its own + ``dynamo_frontend_*`` metrics. + """ + port = self._dynamo_metrics_port if self._use_dynamo else self._server_port + if not self._server_started or port < 0: + return None + return port + @staticmethod def _stop_process(process: Optional[subprocess.Popen]) -> None: if process is None or process.poll() is not None: @@ -197,6 +299,7 @@ def _stop_processes(self) -> None: self._managed_etcd_endpoint = None self._server_started = False self._server_port = -1 + self._dynamo_metrics_port = -1 def _process_status(self) -> str: process_status = [] @@ -214,6 +317,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 @@ -319,9 +427,16 @@ def start_server(self, retries=3): self._model_name, ] _append_kwargs(server_cmd, self._vllm_server_kwargs) - self._dynamo_process, _ = start_process(server_cmd) + # DYN_SYSTEM_PORT turns on the worker's system server, which also + # serves vLLM's engine metrics for the autoscaling signal. + self._dynamo_metrics_port = _pick_dynamo_system_port() + self._dynamo_process, _ = start_process( + server_cmd, + env={'DYN_SYSTEM_PORT': str(self._dynamo_metrics_port)}) 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 +489,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 +543,39 @@ 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 routed 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``). Each prefix is spread + over 8 shards by default; raise ``num_shards`` when one prefix carries + most of the traffic and the pool scales well past 8 workers. + enable_autoscaling_metric: if True (default), publish the + ``num_loaded_models_`` autoscaling signal: the worker's estimated + request 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, aggregation, estimator + window and warmup, 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 +584,18 @@ 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 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=_autoscaling_config( + self._enable_autoscaling_metric, self._autoscaling_config), + model_identifier=self._model_identifier) async def _async_run_inference( self, @@ -517,7 +664,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 +721,39 @@ 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 routed 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``). Each prefix is + spread over 8 shards by default; raise ``num_shards`` when one system + prompt carries most of the traffic and the pool scales well past 8 + workers. + enable_autoscaling_metric: if True (default), publish the + ``num_loaded_models_`` autoscaling signal: the worker's estimated + request 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, aggregation, estimator + window and warmup, 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 +764,8 @@ 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 def load_model(self) -> _VLLMModelServer: chat_template_contents = '' @@ -606,7 +782,10 @@ def load_model(self) -> _VLLMModelServer: self._model_name, self._vllm_server_kwargs, self._dynamo_frontend_kwargs, - self._use_dynamo) + self._use_dynamo, + autoscaling=_autoscaling_config( + self._enable_autoscaling_metric, self._autoscaling_config), + model_identifier=self._model_identifier) 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..19f774f67716 100644 --- a/sdks/python/apache_beam/ml/inference/vllm_inference_test.py +++ b/sdks/python/apache_beam/ml/inference/vllm_inference_test.py @@ -14,17 +14,29 @@ # See the License for the specific language governing permissions and # limitations under the License. # +import gc import os import sys import types import unittest from unittest import mock +import apache_beam as beam +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 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 # import cleanly. -# pylint: disable=wrong-import-order, wrong-import-position +# pylint: disable=wrong-import-order, wrong-import-position, ungrouped-imports try: import openai # pylint: disable=unused-import except ImportError: @@ -73,14 +85,34 @@ def __exit__(self, exc_type, exc_value, traceback): return False -def _record_start_process(commands): - def start_process(cmd): +def _record_start_process(commands, envs=None): + def start_process(cmd, env=None): commands.append(list(cmd)) + if envs is not None: + envs.append(env) return _FakeProcess(), 10000 + len(commands) 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 +125,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 +153,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] @@ -146,12 +180,267 @@ def test_dynamo_starts_frontend_and_engine_with_separate_kwargs(self): self.assertNotIn('--router-mode', engine_cmd) self.assertNotIn('--no-router-kv-events', engine_cmd) + def test_dynamo_autoscaling_reads_the_engine_system_port(self): + """In Dynamo mode the signal reads the dynamo.vllm worker's metrics. + + The frontend only serves dynamo_frontend_* metrics, so the worker is + started with DYN_SYSTEM_PORT and the publisher scrapes that port. + """ + commands, envs = [], [] + with mock.patch.dict(os.environ, + {'ETCD_ENDPOINTS': 'http://127.0.0.1:2379'}): + with mock.patch.object(vllm_inference, + 'start_process', + _record_start_process(commands, envs)): + with mock.patch.object(vllm_inference, + 'getVLLMClient', + return_value=_FakeClient()): + with mock.patch.object( + autoscaling_metrics.PrometheusNumLoadedModelsPublisher, 'start'): + server = vllm_inference.VLLMCompletionsModelHandler( + model_name='test-model', use_dynamo=True).load_model() + (frontend_cmd, frontend_env), (engine_cmd, engine_env) = zip(commands, envs) + self.assertIn('dynamo.frontend', frontend_cmd) + self.assertIsNone(frontend_env) + self.assertIn('dynamo.vllm', engine_cmd) + system_port = int(engine_env['DYN_SYSTEM_PORT']) + # Dynamo parses DYN_SYSTEM_PORT as a signed 16-bit integer. + self.assertLess(system_port, 2**15) + self.assertNotEqual(system_port, server._server_port) + self.assertEqual(server._metrics_port(), system_port) + # Down server -> no port, so the publisher skips the poll. + server._stop_processes() + self.assertIsNone(server._metrics_port()) + + def test_native_vllm_autoscaling_reads_the_api_server_port(self): + commands, envs = [], [] + with mock.patch.object(vllm_inference, + 'start_process', + _record_start_process(commands, envs)): + with mock.patch.object(vllm_inference, + 'getVLLMClient', + return_value=_FakeClient()): + with mock.patch.object( + autoscaling_metrics.PrometheusNumLoadedModelsPublisher, 'start'): + server = vllm_inference.VLLMCompletionsModelHandler( + model_name='test-model').load_model() + self.assertEqual(envs, [None]) + self.assertEqual(server._metrics_port(), server._server_port) + def test_validate_inference_args_accepts_openai_request_kwargs(self): vllm_inference.VLLMCompletionsModelHandler( 'test-model').validate_inference_args({'max_tokens': 8}) vllm_inference.VLLMChatModelHandler('test-model').validate_inference_args( {'max_tokens': 8}) + def test_autoscaling_publishes_capacity_from_vllm_metrics(self): + """The server publishes the capacity estimated from vLLM's metrics.""" + handler = vllm_inference.VLLMCompletionsModelHandler( + model_name='test-model', + autoscaling_config=AutoscalingMetricConfig( + multiplier=3.0, poll_interval_secs=7.0)) + with mock.patch.object(vllm_inference._VLLMModelServer, 'start_server'): + server = handler.load_model() + publisher = server._num_loaded_models_publisher + self.assertEqual(publisher._poll_interval_secs, 7.0) + signal = publisher._read_capacity # What the publisher's thread polls. + # The server is not up yet -> the tick is skipped. + self.assertIsNone(signal()) + + server._server_started = True + server._server_port = 8000 + scrapes = [ + 'vllm:request_success_total{finished_reason="stop"} 10\n' + 'vllm:num_requests_running 4\n', + 'vllm:request_success_total{finished_reason="stop"} 25\n' + 'vllm:request_success_total{finished_reason="length"} 5\n' + 'vllm:num_requests_running 4\n', + ] + with mock.patch.object(autoscaling_metrics, + '_scrape_prometheus_metrics', + side_effect=lambda port, **_: scrapes.pop(0)): + with mock.patch.object(autoscaling_metrics.time, + 'monotonic', + side_effect=[100.0, 105.0]): + self.assertIsNone(signal()) # Baseline sample. + # 20 completions over 5s with 4 in flight -> mu=1, p90=4, x3 -> 12. + self.assertAlmostEqual(signal(), 12.0) + + def test_autoscaling_publisher_follows_server_lifecycle(self): + """The publisher starts once the server is up and stops with it.""" + def load(**kwargs): + with mock.patch.object(vllm_inference, + 'start_process', + _record_start_process([])): + with mock.patch.object(vllm_inference, + 'getVLLMClient', + return_value=_FakeClient()): + return vllm_inference.VLLMCompletionsModelHandler( + model_name='test-model', **kwargs).load_model() + + self.assertIsNone( + load(enable_autoscaling_metric=False)._num_loaded_models_publisher) + + with mock.patch.object(autoscaling_metrics, + '_scrape_prometheus_metrics', + return_value=None): + server = load( + autoscaling_config=AutoscalingMetricConfig(poll_interval_secs=60)) + publisher = server._num_loaded_models_publisher + self.assertEqual(server._metrics_port(), 10001) + self.assertTrue(publisher._thread.is_alive()) + server.__del__() + self.assertIsNone(publisher._thread) + self.assertNotIn( + publisher._metric_name, + autoscaling_metrics.PrometheusNumLoadedModelsPublisher._active) + + def test_autoscaling_publisher_does_not_keep_server_alive(self): + """A dropped server is garbage-collected, which stops its publisher.""" + with mock.patch.object(vllm_inference, + 'start_process', + _record_start_process([])): + with mock.patch.object(vllm_inference, + 'getVLLMClient', + return_value=_FakeClient()): + server = vllm_inference.VLLMCompletionsModelHandler( + model_name='test-model', + autoscaling_config=AutoscalingMetricConfig( + poll_interval_secs=60)).load_model() + with mock.patch.object(autoscaling_metrics, + '_scrape_prometheus_metrics', + return_value=None): + publisher = server._num_loaded_models_publisher + thread = publisher._thread + self.assertTrue(thread.is_alive()) + with mock.patch.object(vllm_inference._VLLMModelServer, + '_stop_processes') as stop_processes: + del server + gc.collect() + stop_processes.assert_called_once() + thread.join(timeout=5) + self.assertFalse(thread.is_alive()) + self.assertTrue(publisher._stop_event.is_set()) + + 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) + # num_shards unset in the config -> the prefix-batching default reaches + # BatchElements. + self.assertEqual( + chat_kwargs['num_shards'], prefix_batching.DEFAULT_NUM_SHARDS) + + # 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', + # One shard per prefix so both prompts of a prefix deterministically + # land in the same batch (the default spreads a prefix over 8 shards). + prefix_batching_config=PrefixBatchingConfig( + prefix_char_len=12, num_shards=1), + 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()