| # |
| # 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. |
| # |
| |
| # pytype: skip-file |
| |
| import asyncio |
| import logging |
| import os |
| import shutil |
| import subprocess |
| import sys |
| import threading |
| import time |
| import urllib.request |
| import uuid |
| from collections.abc import Callable |
| from collections.abc import Iterable |
| from collections.abc import Sequence |
| from dataclasses import dataclass |
| from typing import Any |
| from typing import Optional |
| |
| from openai import AsyncOpenAI |
| from openai import OpenAI |
| |
| from apache_beam.io.filesystems import FileSystems |
| from apache_beam.ml.inference.base import ModelHandler |
| from apache_beam.ml.inference.base import PredictionResult |
| from apache_beam.utils import subprocess_server |
| |
| try: |
| # VLLM logging config breaks beam logging. |
| os.environ["VLLM_CONFIGURE_LOGGING"] = "0" |
| import vllm # pylint: disable=unused-import |
| logging.info('vllm module successfully imported.') |
| os.environ["VLLM_CONFIGURE_LOGGING"] = "1" |
| except ModuleNotFoundError: |
| msg = 'vllm module was not found. This is ok as long as the specified ' \ |
| 'runner has vllm dependencies installed.' |
| logging.warning(msg) |
| |
| __all__ = [ |
| 'OpenAIChatMessage', |
| 'VLLMCompletionsModelHandler', |
| 'VLLMChatModelHandler', |
| ] |
| |
| |
| @dataclass(frozen=True) |
| class OpenAIChatMessage(): |
| """" |
| Dataclass containing previous chat messages in conversation. |
| Role is the entity that sent the message (either 'user' or 'system'). |
| Content is the contents of the message. |
| """ |
| role: str |
| content: str |
| |
| |
| def start_process(cmd) -> tuple[subprocess.Popen, int]: |
| 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) |
| |
| # Emit the output of this command as info level logging. |
| def log_stdout(): |
| line = process.stdout.readline() |
| while line: |
| # The log obtained from stdout is bytes, decode it into string. |
| # Remove newline via rstrip() to not print an empty line. |
| logging.info(line.decode(errors='backslashreplace').rstrip()) |
| line = process.stdout.readline() |
| |
| t = threading.Thread(target=log_stdout) |
| t.daemon = True |
| t.start() |
| return process, port |
| |
| |
| def getVLLMClient(port) -> OpenAI: |
| openai_api_key = "EMPTY" |
| openai_api_base = f"http://localhost:{port}/v1" |
| return OpenAI( |
| api_key=openai_api_key, |
| base_url=openai_api_base, |
| ) |
| |
| |
| def getAsyncVLLMClient(port) -> AsyncOpenAI: |
| openai_api_key = "EMPTY" |
| openai_api_base = f"http://localhost:{port}/v1" |
| return AsyncOpenAI( |
| api_key=openai_api_key, |
| base_url=openai_api_base, |
| ) |
| |
| |
| # Embedded Dynamo runtime defaults proven on the smoke test: etcd discovery, |
| # TCP request plane, ZMQ event plane, KV events disabled. KV-aware routing, |
| # disaggregated prefill/decode, and the Planner are not active in this mode. |
| _DYNAMO_FRONTEND_DEFAULT_KWARGS: dict[str, Optional[str]] = { |
| 'discovery-backend': 'etcd', |
| 'request-plane': 'tcp', |
| 'event-plane': 'zmq', |
| 'router-mode': 'round-robin', |
| 'no-router-kv-events': None, |
| } |
| |
| _DYNAMO_ENGINE_DEFAULT_KWARGS: dict[str, Optional[str]] = { |
| 'discovery-backend': 'etcd', |
| 'request-plane': 'tcp', |
| 'event-plane': 'zmq', |
| 'kv-events-config': '{"enable_kv_cache_events": false}', |
| } |
| |
| |
| def _append_kwargs(cmd: list[str], kwargs: dict[str, Optional[str]]) -> None: |
| for k, v in kwargs.items(): |
| cmd.append(f'--{k}') |
| # Only add values for commands with value part. |
| if v is not None: |
| cmd.append(v) |
| |
| |
| def _uses_etcd_discovery(kwargs: dict[str, Optional[str]]) -> bool: |
| return kwargs.get('discovery-backend') == 'etcd' |
| |
| |
| 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): |
| self._model_name = model_name |
| self._vllm_server_kwargs = vllm_server_kwargs |
| self._dynamo_frontend_kwargs = dynamo_frontend_kwargs or {} |
| self._server_started = False |
| self._server_process = None |
| self._dynamo_process = None |
| self._etcd_process = None |
| self._etcd_data_dir: Optional[str] = None |
| self._managed_etcd_endpoint = None |
| self._server_port: int = -1 |
| self._server_process_lock = threading.RLock() |
| self._use_dynamo = use_dynamo |
| |
| self.start_server() |
| |
| @staticmethod |
| def _stop_process(process: Optional[subprocess.Popen]) -> None: |
| if process is None or process.poll() is not None: |
| return |
| # A process may exit between poll() and terminate() / kill(), in which |
| # case the OS raises ProcessLookupError (or another OSError). Treat that |
| # as already-stopped so we don't bail out of the broader cleanup. |
| try: |
| process.terminate() |
| try: |
| process.wait(timeout=10) |
| except subprocess.TimeoutExpired: |
| process.kill() |
| process.wait() |
| except OSError: |
| pass |
| |
| def _stop_processes(self) -> None: |
| self._stop_process(self._dynamo_process) |
| self._stop_process(self._server_process) |
| self._stop_process(self._etcd_process) |
| if (self._managed_etcd_endpoint is not None and |
| os.environ.get('ETCD_ENDPOINTS') == self._managed_etcd_endpoint): |
| os.environ.pop('ETCD_ENDPOINTS', None) |
| if self._etcd_data_dir is not None: |
| shutil.rmtree(self._etcd_data_dir, ignore_errors=True) |
| self._etcd_data_dir = None |
| self._dynamo_process = None |
| self._server_process = None |
| self._etcd_process = None |
| self._managed_etcd_endpoint = None |
| self._server_started = False |
| self._server_port = -1 |
| |
| def _process_status(self) -> str: |
| process_status = [] |
| if self._server_process is not None: |
| process_status.append( |
| 'frontend/server exit code: %s' % self._server_process.poll()) |
| if self._dynamo_process is not None: |
| process_status.append( |
| 'dynamo worker exit code: %s' % self._dynamo_process.poll()) |
| if self._etcd_process is not None: |
| process_status.append('etcd exit code: %s' % self._etcd_process.poll()) |
| return ', '.join(process_status) or 'no process status available' |
| |
| 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: |
| self._stop_processes() |
| except Exception: # pylint: disable=broad-except |
| pass |
| |
| def _uses_embedded_etcd(self) -> bool: |
| return ( |
| self._use_dynamo and |
| _uses_etcd_discovery(self._dynamo_frontend_kwargs) and |
| _uses_etcd_discovery(self._vllm_server_kwargs) and |
| 'ETCD_ENDPOINTS' not in os.environ) |
| |
| def _wait_for_etcd(self, endpoint: str, timeout_secs=30) -> None: |
| deadline = time.time() + timeout_secs |
| health_url = endpoint.rstrip('/') + '/health' |
| while time.time() < deadline and self._etcd_process.poll() is None: |
| try: |
| with urllib.request.urlopen(health_url, timeout=2) as response: |
| if response.status < 500: |
| return |
| except Exception: # pylint: disable=broad-except |
| time.sleep(1) |
| |
| process_status = self._process_status() |
| self._stop_processes() |
| raise RuntimeError( |
| "Failed to start embedded etcd for Dynamo. Process status: " |
| f"{process_status}. Install etcd in the worker container or set " |
| "ETCD_ENDPOINTS to an external etcd service.") |
| |
| def _ensure_etcd(self) -> None: |
| if not self._uses_embedded_etcd(): |
| return |
| if shutil.which('etcd') is None: |
| raise RuntimeError( |
| "Embedded Dynamo mode requires etcd when ETCD_ENDPOINTS is not " |
| "set. Install etcd in the worker container or set ETCD_ENDPOINTS " |
| "to an external etcd service.") |
| |
| etcd_name = f'beam-dynamo-etcd-{uuid.uuid4().hex}' |
| self._etcd_data_dir = f'/tmp/{etcd_name}' |
| peer_port, = subprocess_server.pick_port(None) |
| etcd_cmd = [ |
| 'etcd', |
| '--name', |
| etcd_name, |
| '--listen-client-urls', |
| 'http://127.0.0.1:{{PORT}}', |
| '--advertise-client-urls', |
| 'http://127.0.0.1:{{PORT}}', |
| '--listen-peer-urls', |
| f'http://127.0.0.1:{peer_port}', |
| '--initial-advertise-peer-urls', |
| f'http://127.0.0.1:{peer_port}', |
| '--initial-cluster', |
| f'{etcd_name}=http://127.0.0.1:{peer_port}', |
| '--data-dir', |
| self._etcd_data_dir, |
| '--log-level', |
| 'warn', |
| ] |
| self._etcd_process, etcd_port = start_process(etcd_cmd) |
| endpoint = f'http://127.0.0.1:{etcd_port}' |
| os.environ['ETCD_ENDPOINTS'] = endpoint |
| self._managed_etcd_endpoint = endpoint |
| self._wait_for_etcd(endpoint) |
| |
| def start_server(self, retries=3): |
| with self._server_process_lock: |
| if not self._server_started: |
| self._stop_processes() |
| self._ensure_etcd() |
| if self._use_dynamo: |
| # Dynamo embedded mode uses the frontend as its OpenAI-compatible |
| # local endpoint and a separate vLLM worker process. |
| server_cmd = [ |
| sys.executable, |
| '-m', |
| 'dynamo.frontend', |
| '--http-port', |
| '{{PORT}}', |
| ] |
| _append_kwargs(server_cmd, self._dynamo_frontend_kwargs) |
| else: |
| server_cmd = [ |
| sys.executable, |
| '-m', |
| 'vllm.entrypoints.openai.api_server', |
| '--model', |
| self._model_name, |
| '--port', |
| '{{PORT}}', |
| ] |
| _append_kwargs(server_cmd, self._vllm_server_kwargs) |
| self._server_process, self._server_port = start_process(server_cmd) |
| |
| if self._use_dynamo: |
| server_cmd = [ |
| sys.executable, |
| '-m', |
| 'dynamo.vllm', |
| '--model', |
| self._model_name, |
| ] |
| _append_kwargs(server_cmd, self._vllm_server_kwargs) |
| self._dynamo_process, _ = start_process(server_cmd) |
| |
| self.check_connectivity(retries) |
| |
| def get_server_port(self) -> int: |
| if not self._server_started: |
| self.start_server() |
| return self._server_port |
| |
| def check_connectivity(self, retries=3, timeout_secs=600): |
| start_time = time.time() |
| with getVLLMClient(self._server_port) as client: |
| while (time.time() - start_time < timeout_secs and |
| self._server_process.poll() is None and |
| (self._dynamo_process is None or |
| self._dynamo_process.poll() is None) and |
| (self._etcd_process is None or self._etcd_process.poll() is None)): |
| try: |
| models = client.models.list().data |
| logging.info('models: %s' % models) |
| if len(models) > 0: |
| self._server_started = True |
| return |
| except: # pylint: disable=bare-except |
| pass |
| # Sleep while bringing up the process |
| time.sleep(5) |
| |
| process_status = self._process_status() |
| self._stop_processes() |
| if retries == 0: |
| raise Exception( |
| "Failed to start vLLM server. Process status: " |
| f"{process_status}. Next time a request is tried, the server " |
| "will be restarted") |
| else: |
| self.start_server(retries - 1) |
| |
| |
| class VLLMCompletionsModelHandler(ModelHandler[str, |
| PredictionResult, |
| _VLLMModelServer]): |
| def __init__( |
| self, |
| model_name: str, |
| vllm_server_kwargs: Optional[dict[str, Optional[str]]] = None, |
| *, |
| use_dynamo: bool = False, |
| dynamo_frontend_kwargs: Optional[dict[str, Optional[str]]] = None, |
| min_batch_size: Optional[int] = None, |
| max_batch_size: Optional[int] = None, |
| max_batch_duration_secs: Optional[int] = None, |
| 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): |
| """Implementation of the ModelHandler interface for vLLM using text as |
| input. |
| |
| Example Usage:: |
| |
| pcoll | RunInference(VLLMModelHandler(model_name='facebook/opt-125m')) |
| |
| Args: |
| model_name: The vLLM model. See |
| https://docs.vllm.ai/en/latest/models/supported_models.html for |
| supported models. |
| vllm_server_kwargs: Any additional kwargs to be passed into your vllm |
| server when it is being created. When ``use_dynamo`` is disabled, |
| this is invoked using ``python -m vllm.entrypoints.openai.api_server |
| <beam provided args> <vllm_server_kwargs>``. When ``use_dynamo`` is |
| enabled, these kwargs are passed to the ``dynamo.vllm`` worker |
| process. For example, you could pass ``{'echo': 'true'}`` to prepend |
| new messages with the previous message. On ~16GB GPUs, pass lower |
| ``max-num-seqs`` and ``gpu-memory-utilization`` values (see |
| ``apache_beam.examples.inference.vllm_text_completion``). For a list |
| of possible kwargs, see |
| https://docs.vllm.ai/en/latest/serving/openai_compatible_server.html#extra-parameters-for-completions-api |
| use_dynamo: Whether to use NVIDIA Dynamo as the underlying vLLM engine. |
| Requires installing Dynamo in your runtime environment |
| (``pip install ai-dynamo[vllm]``). This is an opt-in single-worker |
| embedded mode; KV-aware routing, disaggregated prefill/decode, KVBM |
| offload across nodes, the Planner, and Grove are not active in |
| embedded mode. Dynamo also requires an etcd-style discovery service: |
| when ``ETCD_ENDPOINTS`` is unset, Beam starts a local etcd, which |
| requires the ``etcd`` binary in the worker environment. |
| dynamo_frontend_kwargs: Additional kwargs to be passed to the |
| ``dynamo.frontend`` process when ``use_dynamo`` is enabled. By |
| default, embedded Dynamo uses etcd discovery, TCP request plane, ZMQ |
| event plane, round-robin routing, and disables router KV events. |
| min_batch_size: optional. the minimum batch size to use when batching |
| inputs. |
| max_batch_size: optional. the maximum batch size to use when batching |
| inputs. |
| max_batch_duration_secs: optional. the maximum amount of time to buffer |
| a batch before emitting; used in streaming contexts. |
| max_batch_weight: optional. the maximum total weight of a batch. |
| element_size_fn: optional. a function that returns the size (weight) of |
| an element. |
| batch_length_fn: optional. a callable that returns the length of an |
| element for length-aware batching. |
| batch_bucket_boundaries: optional. a sorted list of positive boundary |
| values for length-aware batching buckets. |
| """ |
| 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) |
| self._model_name = model_name |
| self._vllm_server_kwargs: dict[str, Optional[str]] = ({ |
| **_DYNAMO_ENGINE_DEFAULT_KWARGS, **(vllm_server_kwargs or {}) |
| } if use_dynamo else vllm_server_kwargs or {}) |
| self._dynamo_frontend_kwargs: dict[str, Optional[str]] = { |
| **_DYNAMO_FRONTEND_DEFAULT_KWARGS, **(dynamo_frontend_kwargs or {}) |
| } |
| self._use_dynamo = use_dynamo |
| |
| def load_model(self) -> _VLLMModelServer: |
| return _VLLMModelServer( |
| self._model_name, |
| self._vllm_server_kwargs, |
| self._dynamo_frontend_kwargs, |
| self._use_dynamo) |
| |
| async def _async_run_inference( |
| self, |
| batch: Sequence[str], |
| model: _VLLMModelServer, |
| inference_args: Optional[dict[str, Any]] = None |
| ) -> Iterable[PredictionResult]: |
| inference_args = inference_args or {} |
| |
| async with getAsyncVLLMClient(model.get_server_port()) as client: |
| try: |
| async_predictions = [ |
| client.completions.create( |
| model=self._model_name, prompt=prompt, **inference_args) |
| for prompt in batch |
| ] |
| responses = await asyncio.gather(*async_predictions) |
| except Exception as e: |
| model.check_connectivity() |
| raise e |
| |
| return [PredictionResult(x, y) for x, y in zip(batch, responses)] |
| |
| def run_inference( |
| self, |
| batch: Sequence[str], |
| model: _VLLMModelServer, |
| inference_args: Optional[dict[str, Any]] = None |
| ) -> Iterable[PredictionResult]: |
| """Runs inferences on a batch of text strings. |
| |
| Args: |
| batch: A sequence of examples as text strings. |
| model: A _VLLMModelServer containing info for connecting to the server. |
| inference_args: Any additional arguments for an inference. |
| |
| Returns: |
| An Iterable of type PredictionResult. |
| """ |
| return asyncio.run(self._async_run_inference(batch, model, inference_args)) |
| |
| def validate_inference_args(self, inference_args: Optional[dict[str, Any]]): |
| # Override the base validator so OpenAI-compatible request kwargs such as |
| # ``max_tokens`` can be passed through ``RunInference`` to the vLLM / |
| # Dynamo server. |
| pass |
| |
| def share_model_across_processes(self) -> bool: |
| return True |
| |
| |
| class VLLMChatModelHandler(ModelHandler[Sequence[OpenAIChatMessage], |
| PredictionResult, |
| _VLLMModelServer]): |
| def __init__( |
| self, |
| model_name: str, |
| chat_template_path: Optional[str] = None, |
| vllm_server_kwargs: Optional[dict[str, Optional[str]]] = None, |
| *, |
| use_dynamo: bool = False, |
| dynamo_frontend_kwargs: Optional[dict[str, Optional[str]]] = None, |
| min_batch_size: Optional[int] = None, |
| max_batch_size: Optional[int] = None, |
| max_batch_duration_secs: Optional[int] = None, |
| 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): |
| """ Implementation of the ModelHandler interface for vLLM using previous |
| messages as input. |
| |
| Example Usage:: |
| |
| pcoll | RunInference(VLLMModelHandler(model_name='facebook/opt-125m')) |
| |
| Args: |
| model_name: The vLLM model. See |
| https://docs.vllm.ai/en/latest/models/supported_models.html for |
| supported models. |
| chat_template_path: Path to a chat template. This file must be accessible |
| from your runner's execution environment, so it is recommended to use |
| a cloud based file storage system (e.g. Google Cloud Storage). |
| For info on chat templates, see: |
| https://docs.vllm.ai/en/latest/serving/openai_compatible_server.html#chat-template |
| vllm_server_kwargs: Any additional kwargs to be passed into your vllm |
| server when it is being created. When ``use_dynamo`` is disabled, |
| this is invoked using ``python -m vllm.entrypoints.openai.api_server |
| <beam provided args> <vllm_server_kwargs>``. When ``use_dynamo`` is |
| enabled, these kwargs are passed to the ``dynamo.vllm`` worker |
| process. For example, you could pass ``{'echo': 'true'}`` to prepend |
| new messages with the previous message. For a list of possible |
| kwargs, see |
| https://docs.vllm.ai/en/latest/serving/openai_compatible_server.html#extra-parameters-for-chat-api |
| use_dynamo: Whether to use NVIDIA Dynamo as the underlying vLLM engine. |
| Requires installing Dynamo in your runtime environment |
| (``pip install ai-dynamo[vllm]``). This is an opt-in single-worker |
| embedded mode; KV-aware routing, disaggregated prefill/decode, KVBM |
| offload across nodes, the Planner, and Grove are not active in |
| embedded mode. Dynamo also requires an etcd-style discovery service: |
| when ``ETCD_ENDPOINTS`` is unset, Beam starts a local etcd, which |
| requires the ``etcd`` binary in the worker environment. |
| dynamo_frontend_kwargs: Additional kwargs to be passed to the |
| ``dynamo.frontend`` process when ``use_dynamo`` is enabled. By |
| default, embedded Dynamo uses etcd discovery, TCP request plane, ZMQ |
| event plane, round-robin routing, and disables router KV events. |
| min_batch_size: optional. the minimum batch size to use when batching |
| inputs. |
| max_batch_size: optional. the maximum batch size to use when batching |
| inputs. |
| max_batch_duration_secs: optional. the maximum amount of time to buffer |
| a batch before emitting; used in streaming contexts. |
| max_batch_weight: optional. the maximum total weight of a batch. |
| element_size_fn: optional. a function that returns the size (weight) of |
| an element. |
| batch_length_fn: optional. a callable that returns the length of an |
| element for length-aware batching. |
| batch_bucket_boundaries: optional. a sorted list of positive boundary |
| values for length-aware batching buckets. |
| """ |
| 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) |
| self._model_name = model_name |
| self._vllm_server_kwargs: dict[str, Optional[str]] = ({ |
| **_DYNAMO_ENGINE_DEFAULT_KWARGS, **(vllm_server_kwargs or {}) |
| } if use_dynamo else vllm_server_kwargs or {}) |
| self._dynamo_frontend_kwargs: dict[str, Optional[str]] = { |
| **_DYNAMO_FRONTEND_DEFAULT_KWARGS, **(dynamo_frontend_kwargs or {}) |
| } |
| self._chat_template_path = chat_template_path |
| self._chat_file = f'template-{uuid.uuid4().hex}.jinja' |
| self._use_dynamo = use_dynamo |
| |
| def load_model(self) -> _VLLMModelServer: |
| chat_template_contents = '' |
| if self._chat_template_path is not None: |
| local_chat_template_path = os.path.join(os.getcwd(), self._chat_file) |
| if not os.path.exists(local_chat_template_path): |
| with FileSystems.open(self._chat_template_path) as fin: |
| chat_template_contents = fin.read().decode() |
| with open(local_chat_template_path, 'a') as f: |
| f.write(chat_template_contents) |
| self._vllm_server_kwargs['chat_template'] = local_chat_template_path |
| |
| return _VLLMModelServer( |
| self._model_name, |
| self._vllm_server_kwargs, |
| self._dynamo_frontend_kwargs, |
| self._use_dynamo) |
| |
| async def _async_run_inference( |
| self, |
| batch: Sequence[Sequence[OpenAIChatMessage]], |
| model: _VLLMModelServer, |
| inference_args: Optional[dict[str, Any]] = None |
| ) -> Iterable[PredictionResult]: |
| inference_args = inference_args or {} |
| |
| async with getAsyncVLLMClient(model.get_server_port()) as client: |
| try: |
| async_predictions = [ |
| client.chat.completions.create( |
| model=self._model_name, |
| messages=[{ |
| "role": message.role, "content": message.content |
| } for message in messages], |
| **inference_args) for messages in batch |
| ] |
| predictions = await asyncio.gather(*async_predictions) |
| except Exception as e: |
| model.check_connectivity() |
| raise e |
| |
| return [PredictionResult(x, y) for x, y in zip(batch, predictions)] |
| |
| def run_inference( |
| self, |
| batch: Sequence[Sequence[OpenAIChatMessage]], |
| model: _VLLMModelServer, |
| inference_args: Optional[dict[str, Any]] = None |
| ) -> Iterable[PredictionResult]: |
| """Runs inferences on a batch of text strings. |
| |
| Args: |
| batch: A sequence of examples as OpenAI messages. |
| model: A _VLLMModelServer for connecting to the spun up server. |
| inference_args: Any additional arguments for an inference. |
| |
| Returns: |
| An Iterable of type PredictionResult. |
| """ |
| return asyncio.run(self._async_run_inference(batch, model, inference_args)) |
| |
| def validate_inference_args(self, inference_args: Optional[dict[str, Any]]): |
| # Override the base validator so OpenAI-compatible request kwargs such as |
| # ``max_tokens`` can be passed through ``RunInference`` to the vLLM / |
| # Dynamo server. |
| pass |
| |
| def share_model_across_processes(self) -> bool: |
| return True |
