This is an automated email from the ASF dual-hosted git repository. damccorm pushed a commit to branch feature/adk-local-model in repository https://gitbox.apache.org/repos/asf/beam.git
commit ded44c97aa262a44f651f438be5f874f7a460360 Author: Danny McCormick <[email protected]> AuthorDate: Tue Jun 23 16:04:04 2026 +0000 Extend ADKAgentModelHandler to support deploying local models This change extends ADKAgentModelHandler to allow deploying a local model serving process instead of relying on a remote model endpoint. It introduces a generic SubProcessModel wrapper that runs any ModelHandler in a separate FastAPI-based subprocess, exposing OpenAI-compatible endpoints as well as a transparent pickle-based endpoint for Beam integration. ADKAgentModelHandler now accepts this SubProcessModel (or any SubprocessModelHandler) and automatically manages its lifecycle, including port updates and recovery on failures, and propagates the local model configuration to subagents and tools. TAG=agy CONV=b2f450cf-8b63-49a9-8ee8-430ed676e116 --- .../ml/inference/agent_development_kit.py | 140 +++++++++++++- .../ml/inference/agent_development_kit_test.py | 140 +++++++++++++- sdks/python/apache_beam/ml/inference/base.py | 215 +++++++++++++++++++++ sdks/python/apache_beam/ml/inference/base_test.py | 59 ++++++ .../apache_beam/ml/inference/subprocess_server.py | 145 ++++++++++++++ .../apache_beam/ml/inference/vllm_inference.py | 60 ++++-- 6 files changed, 737 insertions(+), 22 deletions(-) diff --git a/sdks/python/apache_beam/ml/inference/agent_development_kit.py b/sdks/python/apache_beam/ml/inference/agent_development_kit.py index 386955b0dfa..9190754fa4b 100644 --- a/sdks/python/apache_beam/ml/inference/agent_development_kit.py +++ b/sdks/python/apache_beam/ml/inference/agent_development_kit.py @@ -61,6 +61,7 @@ from typing import Union from apache_beam.ml.inference.base import ModelHandler from apache_beam.ml.inference.base import PredictionResult +from apache_beam.ml.inference.base import SubprocessModelHandler try: from google.adk import sessions @@ -73,6 +74,30 @@ try: ADK_AVAILABLE = True except ImportError: ADK_AVAILABLE = False + +if ADK_AVAILABLE: + try: + from google.adk.models.base_llm import BaseLlm + except ImportError: + try: + from google.adk.models import BaseLlm + except ImportError: + BaseLlm = object + + class BeamPlaceholderModel(BaseLlm): + """Placeholder model to be used when the model will be injected by ADKAgentModelHandler.""" + def __init__(self): + pass + async def generate_content_async(self, *args, **kwargs): + raise NotImplementedError("Placeholder model cannot be used for inference.") +else: + class BeamPlaceholderModel(str): + """Placeholder model to be used when the model will be injected by ADKAgentModelHandler. + + Fallback when ADK is not available. + """ + def __new__(cls): + return super().__new__(cls, "beam-placeholder-model") genai_Content = Any # type: ignore[assignment, misc] genai_Part = Any # type: ignore[assignment, misc] @@ -128,6 +153,7 @@ class ADKAgentModelHandler(ModelHandler[str | genai_Content, app_name: str = "beam_inference", session_service_factory: Optional[Callable[[], "BaseSessionService"]] = None, + underlying_model_handler: Optional[SubprocessModelHandler] = None, *, min_batch_size: Optional[int] = None, max_batch_size: Optional[int] = None, @@ -146,6 +172,8 @@ class ADKAgentModelHandler(ModelHandler[str | genai_Content, self._agent_or_factory = agent self._app_name = app_name self._session_service_factory = session_service_factory + self._underlying_model_handler = underlying_model_handler + self._current_port = None super().__init__( min_batch_size=min_batch_size, @@ -165,11 +193,54 @@ class ADKAgentModelHandler(ModelHandler[str | genai_Content, Returns: A fully initialised :class:`~google.adk.runners.Runner`. """ + local_model = None + underlying_model = None + + if self._underlying_model_handler is not None: + underlying_model = self._underlying_model_handler.load_model() + self._current_port = self._underlying_model_handler.get_port(underlying_model) + model_name = self._underlying_model_handler.get_model_name() + + from google.adk.models.lite_llm import LiteLlm + local_model = LiteLlm( + model=model_name, + api_base=f"http://localhost:{self._current_port}/v1"; + ) + + # Resolve agent and inject model if callable(self._agent_or_factory) and not isinstance( self._agent_or_factory, Agent): - agent = self._agent_or_factory() + import inspect + sig = inspect.signature(self._agent_or_factory) + params = list(sig.parameters.values()) + + if len(params) == 1: + if local_model is None: + raise ValueError("Agent factory expects 1 argument but no local model was configured.") + agent = self._agent_or_factory(local_model) + elif len(params) == 0: + agent = self._agent_or_factory() + if local_model is not None: + if not isinstance(agent.model, BeamPlaceholderModel) and agent.model is not None: + raise ValueError( + f"Agent model must be BeamPlaceholderModel or None when using local model. " + f"Found: {agent.model}") + self._set_agent_model(agent, local_model, is_root=True) + else: + raise ValueError("Agent factory must take 0 or 1 argument.") else: agent = self._agent_or_factory + if local_model is not None: + if not isinstance(agent.model, BeamPlaceholderModel) and agent.model is not None: + raise ValueError( + f"Agent model must be BeamPlaceholderModel or None when using local model. " + f"Found: {agent.model}") + self._set_agent_model(agent, local_model, is_root=True) + + # Validation when local model is NOT used + if local_model is None: + if isinstance(agent.model, BeamPlaceholderModel): + raise ValueError("Agent model cannot be BeamPlaceholderModel when no local model is configured.") if self._session_service_factory is not None: session_service = self._session_service_factory() @@ -181,6 +252,10 @@ class ADKAgentModelHandler(ModelHandler[str | genai_Content, app_name=self._app_name, session_service=session_service, ) + + if underlying_model is not None: + runner._underlying_model = underlying_model + LOGGER.info( "Loaded ADK Runner for agent '%s' (app_name='%s')", agent.name, @@ -188,6 +263,22 @@ class ADKAgentModelHandler(ModelHandler[str | genai_Content, ) return runner + def _set_agent_model(self, agent: "Agent", model: Any, is_root: bool = False): + if is_root: + if isinstance(agent.model, BeamPlaceholderModel) or agent.model is None: + agent.model = model + else: + if isinstance(agent.model, BeamPlaceholderModel): + agent.model = model + + # Speculative propagation to subagents/tools + if hasattr(agent, 'tools'): + for tool in agent.tools: + if hasattr(tool, 'agent'): + self._set_agent_model(tool.agent, model, is_root=False) + elif isinstance(tool, Agent): + self._set_agent_model(tool, model, is_root=False) + def run_inference( self, batch: Sequence[str | genai_Content], @@ -219,6 +310,33 @@ class ADKAgentModelHandler(ModelHandler[str | genai_Content, An iterable of :class:`~apache_beam.ml.inference.base.PredictionResult`, one per input element. """ + underlying_model = None + if self._underlying_model_handler is not None: + underlying_model = getattr(model, '_underlying_model', None) + if underlying_model is not None: + port = self._underlying_model_handler.get_port(underlying_model) + if port != self._current_port: + LOGGER.info("Local model server port changed to %d, updating agent.", port) + self._update_agent_port(model.agent, port) + self._current_port = port + + try: + return self._run_inference_internal(batch, model, inference_args) + except Exception as e: + if self._underlying_model_handler is not None and underlying_model is not None: + LOGGER.warning("Inference failed, triggering local server connectivity check.") + try: + self._underlying_model_handler.check_connectivity(underlying_model) + except Exception as recovery_err: + LOGGER.error("Failed during connectivity check: %s", recovery_err) + raise e + + def _run_inference_internal( + self, + batch: Sequence[str | genai_Content], + model: "Runner", + inference_args: Optional[dict[str, Any]] = None, + ) -> Iterable[PredictionResult]: if inference_args is None: inference_args = {} @@ -259,6 +377,26 @@ class ADKAgentModelHandler(ModelHandler[str | genai_Content, return results + def _update_agent_port(self, agent: "Agent", port: int): + if ADK_AVAILABLE: + from google.adk.models.lite_llm import LiteLlm + if hasattr(agent, 'model') and isinstance(agent.model, LiteLlm): + agent.model = LiteLlm( + model=agent.model.model, + api_base=f"http://localhost:{port}/v1"; + ) + if hasattr(agent, 'tools'): + for tool in agent.tools: + if hasattr(tool, 'agent'): + self._update_agent_port(tool.agent, port) + elif isinstance(tool, Agent): + self._update_agent_port(tool, port) + + def share_model_across_processes(self) -> bool: + if self._underlying_model_handler is not None: + return self._underlying_model_handler.share_model_across_processes() + return super().share_model_across_processes() + @staticmethod async def _invoke_agent( runner: "Runner", diff --git a/sdks/python/apache_beam/ml/inference/agent_development_kit_test.py b/sdks/python/apache_beam/ml/inference/agent_development_kit_test.py index 6c8b5c5b351..214a0379f93 100644 --- a/sdks/python/apache_beam/ml/inference/agent_development_kit_test.py +++ b/sdks/python/apache_beam/ml/inference/agent_development_kit_test.py @@ -24,7 +24,9 @@ try: from google.adk.agents import Agent from apache_beam.ml.inference.agent_development_kit import ADKAgentModelHandler + from apache_beam.ml.inference.agent_development_kit import BeamPlaceholderModel from apache_beam.ml.inference.base import PredictionResult + from apache_beam.ml.inference.base import SubprocessModelHandler except ImportError: raise unittest.SkipTest('google-adk dependencies are not installed') @@ -129,12 +131,46 @@ class TestLoadModel(unittest.TestCase): @mock.patch(f"{_MODULE}.InMemorySessionService") def test_load_model_calls_factory(self, mock_session_cls, mock_runner_cls): agent = _make_mock_agent() - factory = mock.MagicMock(return_value=agent) + mock_factory = mock.MagicMock(return_value=agent) + def factory(): + return mock_factory() handler = ADKAgentModelHandler(agent=factory) handler.load_model() - factory.assert_called_once() + mock_factory.assert_called_once() + mock_runner_cls.assert_called_once_with( + agent=agent, + app_name="beam_inference", + session_service=mock_session_cls.return_value, + ) + + @mock.patch(f"{_MODULE}.Runner") + @mock.patch(f"{_MODULE}.InMemorySessionService") + def test_load_model_calls_factory_with_model(self, mock_session_cls, mock_runner_cls): + agent = _make_mock_agent() + mock_factory = mock.MagicMock(return_value=agent) + def factory(model): + return mock_factory(model) + + mock_underlying_handler = mock.MagicMock(spec=SubprocessModelHandler) + mock_underlying_handler.get_model_name.return_value = "mock_model" + mock_underlying_handler.get_port.return_value = 12345 + mock_underlying_handler.share_model_across_processes.return_value = False + + handler = ADKAgentModelHandler( + agent=factory, + underlying_model_handler=mock_underlying_handler + ) + handler.load_model() + + from google.adk.models.lite_llm import LiteLlm + mock_factory.assert_called_once() + called_model = mock_factory.call_args[0][0] + self.assertIsInstance(called_model, LiteLlm) + self.assertEqual(called_model.model, "mock_model") + self.assertEqual(called_model._additional_args.get("api_base"), "http://localhost:12345/v1";) + mock_runner_cls.assert_called_once_with( agent=agent, app_name="beam_inference", @@ -340,5 +376,105 @@ class TestResponseExtraction(unittest.TestCase): self.assertEqual(result, "direct result") +class TestLocalModelIntegration(unittest.TestCase): + def setUp(self): + self.mock_handler = mock.MagicMock(spec=SubprocessModelHandler) + self.mock_handler.get_model_name.return_value = "mock_model" + self.mock_handler.share_model_across_processes.return_value = False + + self.mock_model = mock.MagicMock() + self.mock_handler.load_model.return_value = self.mock_model + self.mock_handler.get_port.return_value = 12345 + + def test_placeholder_validation_fails_no_local_model(self): + agent = Agent(model=BeamPlaceholderModel(), name="test_agent") + handler = ADKAgentModelHandler(agent=agent) + with self.assertRaises(ValueError) as ctx: + handler.load_model() + self.assertIn("cannot be BeamPlaceholderModel when no local model is configured", str(ctx.exception)) + + def test_validation_fails_with_remote_model_and_local_configured(self): + agent = Agent(model="gemini-1.5-pro", name="test_agent") + handler = ADKAgentModelHandler(agent=agent, underlying_model_handler=self.mock_handler) + with self.assertRaises(ValueError) as ctx: + handler.load_model() + self.assertIn("Agent model must be BeamPlaceholderModel or None", str(ctx.exception)) + + def test_local_model_injection_and_propagation(self): + subagent = Agent(model=BeamPlaceholderModel(), name="subagent") + + def mock_tool(): + pass + mock_tool.agent = subagent + + root_agent = Agent( + model=BeamPlaceholderModel(), + name="root_agent", + tools=[mock_tool] + ) + + handler = ADKAgentModelHandler(agent=root_agent, underlying_model_handler=self.mock_handler) + runner = handler.load_model() + + from google.adk.models.lite_llm import LiteLlm + self.assertIsInstance(runner.agent.model, LiteLlm) + self.assertEqual(runner.agent.model.model, "mock_model") + self.assertEqual(runner.agent.model._additional_args.get("api_base"), "http://localhost:12345/v1";) + + self.assertIsInstance(subagent.model, LiteLlm) + self.assertEqual(subagent.model._additional_args.get("api_base"), "http://localhost:12345/v1";) + + def test_port_update_propagation(self): + subagent = Agent(model=BeamPlaceholderModel(), name="subagent") + + def mock_tool(): + pass + mock_tool.agent = subagent + + root_agent = Agent(model=BeamPlaceholderModel(), name="root_agent", tools=[mock_tool]) + + handler = ADKAgentModelHandler(agent=root_agent, underlying_model_handler=self.mock_handler) + runner = handler.load_model() + + self.assertEqual(root_agent.model._additional_args.get("api_base"), "http://localhost:12345/v1";) + self.assertEqual(subagent.model._additional_args.get("api_base"), "http://localhost:12345/v1";) + + self.mock_handler.get_port.return_value = 54321 + + handler._run_inference_internal = mock.MagicMock(return_value=[]) + + handler.run_inference(batch=["test"], model=runner) + + self.assertEqual(root_agent.model._additional_args.get("api_base"), "http://localhost:54321/v1";) + self.assertEqual(subagent.model._additional_args.get("api_base"), "http://localhost:54321/v1";) + self.assertEqual(handler._current_port, 54321) + + def test_recovery_on_failure(self): + root_agent = Agent(model=BeamPlaceholderModel(), name="root_agent") + handler = ADKAgentModelHandler(agent=root_agent, underlying_model_handler=self.mock_handler) + runner = handler.load_model() + + handler._run_inference_internal = mock.MagicMock(side_effect=Exception("Connection lost")) + + with self.assertRaises(Exception): + handler.run_inference(batch=["test"], model=runner) + + self.mock_handler.check_connectivity.assert_called_once_with(self.mock_model) + + def test_recovery_fails_does_not_mask_original_error(self): + root_agent = Agent(model=BeamPlaceholderModel(), name="root_agent") + handler = ADKAgentModelHandler(agent=root_agent, underlying_model_handler=self.mock_handler) + runner = handler.load_model() + + handler._run_inference_internal = mock.MagicMock(side_effect=ValueError("Original error")) + self.mock_handler.check_connectivity.side_effect = Exception("Recovery failed") + + with self.assertRaises(ValueError) as ctx: + handler.run_inference(batch=["test"], model=runner) + + self.assertEqual(str(ctx.exception), "Original error") + self.mock_handler.check_connectivity.assert_called_once_with(self.mock_model) + + 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 f81382bbeec..30f296042c0 100644 --- a/sdks/python/apache_beam/ml/inference/base.py +++ b/sdks/python/apache_beam/ml/inference/base.py @@ -31,7 +31,9 @@ import functools import logging import os import pickle +import subprocess import sys +import tempfile import threading import time import uuid @@ -406,6 +408,219 @@ class ModelHandler(Generic[ExampleT, PredictionT, ModelT]): return self.share_model_across_processes() +class SubprocessModelHandler(ModelHandler[ExampleT, PredictionT, ModelT], ABC): + """Base class for model handlers that spin up a subprocess server.""" + @abstractmethod + def get_port(self, model: ModelT) -> int: + """Returns the port the subprocess server is listening on.""" + pass + + @abstractmethod + def get_model_name(self) -> str: + """Returns the model name.""" + pass + + @abstractmethod + def check_connectivity(self, model: ModelT) -> None: + """Checks connectivity to the server and attempts to recover/mark for restart.""" + pass + + +class SubProcessModelServer: + """Manages the lifecycle of a generic subprocess model server.""" + def __init__(self, handler_path: str, model_name: str, port: int = None, temp_dir: tempfile.TemporaryDirectory = None): + self._handler_path = handler_path + self._model_name = model_name + self._port = port + self._temp_dir = temp_dir + self._process = None + self._server_started = False + self._server_process_lock = threading.RLock() + self.start_server() + + def start_server(self, retries=3): + with self._server_process_lock: + if not self._server_started: + if self._process: + logging.info("Terminating existing generic subprocess model server before restart") + try: + self._process.terminate() + self._process.wait(timeout=5) + except Exception: + try: + self._process.kill() + except Exception: + pass + self._process = None + self._port = None + + from apache_beam.utils import subprocess_server + if self._port is None: + self._port, = subprocess_server.pick_port(None) + + cmd = [ + sys.executable, + '-m', + 'apache_beam.ml.inference.subprocess_server', + '--handler_path', + self._handler_path, + '--port', + str(self._port), + ] + logging.info("Starting generic model server with %s", cmd) + self._process = subprocess.Popen( + cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT) + + # Emit the output of this command as info level logging. + def log_stdout(): + line = self._process.stdout.readline() + while line: + logging.info(line.decode(errors='backslashreplace').rstrip()) + line = self._process.stdout.readline() + + t = threading.Thread(target=log_stdout) + t.daemon = True + t.start() + + self.check_connectivity(retries) + + def get_server_port(self) -> int: + if not self._server_started: + self.start_server() + return self._port + + def check_connectivity(self, retries=3): + import urllib.request + import urllib.error + + url = f"http://localhost:{self._port}/v1/models"; + attempts = 0 + max_attempts = 12 # 12 * 5s = 60s timeout + while self._process.poll() is None and attempts < max_attempts: + try: + # Use standard library to check connectivity to avoid extra dependencies + req = urllib.request.Request(url, method="GET") + with urllib.request.urlopen(req, timeout=5) as response: + if response.status == 200: + self._server_started = True + return + except urllib.error.URLError: + pass + except Exception as e: + logging.warning("Error checking connectivity: %s", e) + attempts += 1 + time.sleep(5) + + if retries == 0: + self._server_started = False + raise Exception( + "Failed to start generic subprocess server, polling process exited with code " + + f"{self._process.poll()}. Next time a request is tried, the server will be restarted" + ) + else: + self.start_server(retries - 1) + + def __del__(self): + if self._process: + logging.info("Terminating generic subprocess model server") + try: + self._process.terminate() + self._process.wait(timeout=5) + except Exception: + try: + self._process.kill() + except Exception: + pass + if self._temp_dir: + try: + self._temp_dir.cleanup() + except Exception: + pass + + +class SubProcessModel(SubprocessModelHandler[ExampleT, PredictionT, Any]): + """Wrapper to adapt any ModelHandler to SubprocessModelHandler.""" + def __init__( + self, + handler: ModelHandler[ExampleT, PredictionT, ModelT], + model_name: str): + super().__init__() + self._handler = handler + self._model_name = model_name + self._handler_path = None + + def load_model(self) -> Any: + if isinstance(self._handler, SubprocessModelHandler): + return self._handler.load_model() + + import tempfile + temp_dir = tempfile.TemporaryDirectory(prefix="beam-subprocess-") + self._handler_path = os.path.join(temp_dir.name, "handler.pickle") + with open(self._handler_path, "wb") as f: + pickle.dump(self._handler, f) + + return SubProcessModelServer(self._handler_path, self._model_name, temp_dir=temp_dir) + + def run_inference(self, batch, model, inference_args=None): + if isinstance(model, SubProcessModelServer): + return self._run_inference_via_http(batch, model.get_server_port(), inference_args) + else: + return self._handler.run_inference(batch, model, inference_args) + + def _run_inference_via_http(self, batch, port, inference_args): + import urllib.request + import urllib.error + + url = f"http://localhost:{port}/v1/beam/inference"; + payload = { + "batch": batch, + "inference_args": inference_args + } + data = pickle.dumps(payload) + + req = urllib.request.Request(url, data=data, method="POST") + req.add_header('Content-Type', 'application/octet-stream') + + try: + with urllib.request.urlopen(req, timeout=30) as response: + resp_data = response.read() + results = pickle.loads(resp_data) + return results + except urllib.error.HTTPError as e: + try: + err_details = e.read().decode('utf-8') + logging.error("Subprocess server returned error: %s", err_details) + except Exception: + pass + logging.exception("Failed to run inference via HTTP raw endpoint (HTTPError)") + raise e + except Exception as e: + logging.exception("Failed to run inference via HTTP raw endpoint") + raise e + + def get_port(self, model: Any) -> int: + if hasattr(model, 'get_server_port'): + return model.get_server_port() + elif hasattr(model, 'port'): + return model.port + raise ValueError(f"Could not determine port from model of type {type(model)}") + + def get_model_name(self) -> str: + return self._model_name + + def check_connectivity(self, model: Any) -> None: + if hasattr(model, 'check_connectivity'): + model.check_connectivity() + elif hasattr(self._handler, 'check_connectivity'): + self._handler.check_connectivity(model) + + def share_model_across_processes(self) -> bool: + return self._handler.share_model_across_processes() + + def __getattr__(self, name): + return getattr(self._handler, name) + + class RemoteModelHandler(ABC, ModelHandler[ExampleT, PredictionT, ModelT]): """Has the ability to call a model at a remote endpoint.""" def __init__( diff --git a/sdks/python/apache_beam/ml/inference/base_test.py b/sdks/python/apache_beam/ml/inference/base_test.py index 3f0bea7fbe1..1fb0955623f 100644 --- a/sdks/python/apache_beam/ml/inference/base_test.py +++ b/sdks/python/apache_beam/ml/inference/base_test.py @@ -2433,5 +2433,64 @@ class ModelManagerTest(unittest.TestCase): assert_that(actual, equal_to([2, 6, 4, 11]), label='assert:inferences') +class StringFakeModel: + def predict(self, example: str) -> str: + return example + "_processed" + + +class StringFakeModelHandler(base.ModelHandler[str, base.PredictionResult, StringFakeModel]): + def load_model(self): + return StringFakeModel() + def run_inference(self, batch, model, inference_args=None): + return [base.PredictionResult(x, model.predict(x)) for x in batch] + + +class RawStringFakeModelHandler(base.ModelHandler[str, str, StringFakeModel]): + def load_model(self): + return StringFakeModel() + def run_inference(self, batch, model, inference_args=None): + return [model.predict(x) for x in batch] + + +class SubProcessModelTest(unittest.TestCase): + def test_subprocess_wrapper(self): + handler = StringFakeModelHandler() + wrapper = base.SubProcessModel(handler, model_name="string_fake_model") + + model = wrapper.load_model() + self.assertIsInstance(model, base.SubProcessModelServer) + + results = wrapper.run_inference(["hello", "world"], model) + results = list(results) + + self.assertEqual(len(results), 2) + self.assertEqual(results[0].example, "hello") + self.assertEqual(results[0].inference, "hello_processed") + self.assertEqual(results[1].example, "world") + self.assertEqual(results[1].inference, "world_processed") + + # Process is cleaned up when wrapper/model is garbage collected, + # but we can also trigger it manually by deleting them. + del model + del wrapper + + def test_subprocess_wrapper_raw_string(self): + handler = RawStringFakeModelHandler() + wrapper = base.SubProcessModel(handler, model_name="raw_string_fake_model") + + model = wrapper.load_model() + self.assertIsInstance(model, base.SubProcessModelServer) + + results = wrapper.run_inference(["hello", "world"], model) + results = list(results) + + self.assertEqual(len(results), 2) + self.assertEqual(results[0], "hello_processed") + self.assertEqual(results[1], "world_processed") + + del model + del wrapper + + if __name__ == '__main__': unittest.main() diff --git a/sdks/python/apache_beam/ml/inference/subprocess_server.py b/sdks/python/apache_beam/ml/inference/subprocess_server.py new file mode 100644 index 00000000000..f24dfe416cd --- /dev/null +++ b/sdks/python/apache_beam/ml/inference/subprocess_server.py @@ -0,0 +1,145 @@ +# +# 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 argparse +import logging +import pickle +import sys +import uvicorn +from fastapi import FastAPI, Response, Request, HTTPException +from pydantic import BaseModel +from typing import List, Any + +# Set up logging +logging.basicConfig(level=logging.INFO) +logger = logging.getLogger("subprocess_server") + +app = FastAPI() +handler = None +model = None + +def _extract_inference(result) -> str: + if hasattr(result, "inference"): + val = result.inference + elif hasattr(result, "text"): + val = result.text + else: + val = result + return str(val) + +class ChatMessage(BaseModel): + role: str + content: str + +class ChatCompletionRequest(BaseModel): + model: str + messages: List[ChatMessage] + temperature: float = 1.0 + max_tokens: int = 1024 + +class CompletionRequest(BaseModel): + model: str + prompt: str + [email protected]("/v1/chat/completions") +async def chat_completions(request: ChatCompletionRequest): + logger.info("Received chat completion request") + # Map to handler input. + # For compatibility with standard text handlers, we take the last user message content. + # If the handler expects the full history, we might need a more complex mapping. + prompt = request.messages[-1].content + + # Run inference + # handler.run_inference expects a sequence of inputs + try: + results = handler.run_inference([prompt], model) + result = list(results)[0] + inference_text = _extract_inference(result) + except Exception as e: + logger.exception("Error during inference") + raise HTTPException(status_code=500, detail=str(e)) + + return { + "choices": [{ + "message": { + "role": "assistant", + "content": inference_text + } + }] + } + [email protected]("/v1/completions") +async def completions(request: CompletionRequest): + logger.info("Received completion request") + try: + results = handler.run_inference([request.prompt], model) + result = list(results)[0] + inference_text = _extract_inference(result) + except Exception as e: + logger.exception("Error during inference") + raise HTTPException(status_code=500, detail=str(e)) + + return { + "choices": [{ + "text": inference_text + }] + } + [email protected]("/v1/beam/inference") +async def beam_inference(request: Request): + logger.info("Received Beam raw inference request") + try: + body = await request.body() + payload = pickle.loads(body) + batch = payload["batch"] + inference_args = payload.get("inference_args") + + results = handler.run_inference(batch, model, inference_args) + results_list = list(results) + pickled_results = pickle.dumps(results_list) + return Response(content=pickled_results, media_type="application/octet-stream") + except Exception as e: + logger.exception("Error during raw inference") + raise HTTPException(status_code=500, detail=str(e)) + [email protected]("/v1/models") +async def list_models(): + # Endpoint to check connectivity and list models + model_name = getattr(handler, "_model_name", "unknown") + return { + "data": [ + { + "id": model_name, + "object": "model", + } + ] + } + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--handler_path", required=True) + parser.add_argument("--port", type=int, required=True) + args = parser.parse_args() + + logger.info("Loading handler from %s", args.handler_path) + with open(args.handler_path, "rb") as f: + handler = pickle.load(f) + + logger.info("Loading model...") + model = handler.load_model() + logger.info("Starting uvicorn server on port %d", args.port) + uvicorn.run(app, host="127.0.0.1", port=args.port) diff --git a/sdks/python/apache_beam/ml/inference/vllm_inference.py b/sdks/python/apache_beam/ml/inference/vllm_inference.py index 38283f1efd4..0158e0d2dec 100644 --- a/sdks/python/apache_beam/ml/inference/vllm_inference.py +++ b/sdks/python/apache_beam/ml/inference/vllm_inference.py @@ -31,6 +31,10 @@ from collections.abc import Sequence from dataclasses import dataclass from typing import Any from typing import Optional +from typing import TypeVar + +ExampleT = TypeVar('ExampleT') +PredictionT = TypeVar('PredictionT') from openai import AsyncOpenAI from openai import OpenAI @@ -38,6 +42,7 @@ 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.ml.inference.base import SubprocessModelHandler from apache_beam.utils import subprocess_server try: @@ -170,9 +175,36 @@ class _VLLMModelServer(): self.start_server(retries - 1) -class VLLMCompletionsModelHandler(ModelHandler[str, - PredictionResult, - _VLLMModelServer]): +class _VLLMBaseModelHandler(SubprocessModelHandler[ExampleT, + PredictionT, + _VLLMModelServer]): + def __init__( + self, + model_name: str, + vllm_server_kwargs: Optional[dict[str, str]] = None, + **kwargs): + super().__init__(**kwargs) + self._model_name = model_name + self._vllm_server_kwargs = vllm_server_kwargs or {} + + def share_model_across_processes(self) -> bool: + return True + + def get_port(self, model: _VLLMModelServer) -> int: + return model.get_server_port() + + def get_model_name(self) -> str: + return self._model_name + + def check_connectivity(self, model: _VLLMModelServer) -> None: + model.check_connectivity() + + def load_model(self) -> _VLLMModelServer: + return _VLLMModelServer(self._model_name, self._vllm_server_kwargs) + + +class VLLMCompletionsModelHandler(_VLLMBaseModelHandler[str, + PredictionResult]): def __init__( self, model_name: str, @@ -221,6 +253,8 @@ class VLLMCompletionsModelHandler(ModelHandler[str, values for length-aware batching buckets. """ super().__init__( + model_name=model_name, + vllm_server_kwargs=vllm_server_kwargs, min_batch_size=min_batch_size, max_batch_size=max_batch_size, max_batch_duration_secs=max_batch_duration_secs, @@ -228,11 +262,6 @@ class VLLMCompletionsModelHandler(ModelHandler[str, 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, str] = vllm_server_kwargs or {} - - def load_model(self) -> _VLLMModelServer: - return _VLLMModelServer(self._model_name, self._vllm_server_kwargs) async def _async_run_inference( self, @@ -274,13 +303,9 @@ class VLLMCompletionsModelHandler(ModelHandler[str, """ return asyncio.run(self._async_run_inference(batch, model, inference_args)) - def share_model_across_processes(self) -> bool: - return True - -class VLLMChatModelHandler(ModelHandler[Sequence[OpenAIChatMessage], - PredictionResult, - _VLLMModelServer]): +class VLLMChatModelHandler(_VLLMBaseModelHandler[Sequence[OpenAIChatMessage], + PredictionResult]): def __init__( self, model_name: str, @@ -332,6 +357,8 @@ class VLLMChatModelHandler(ModelHandler[Sequence[OpenAIChatMessage], values for length-aware batching buckets. """ super().__init__( + model_name=model_name, + vllm_server_kwargs=vllm_server_kwargs, min_batch_size=min_batch_size, max_batch_size=max_batch_size, max_batch_duration_secs=max_batch_duration_secs, @@ -339,8 +366,6 @@ class VLLMChatModelHandler(ModelHandler[Sequence[OpenAIChatMessage], 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, str] = vllm_server_kwargs or {} self._chat_template_path = chat_template_path self._chat_file = f'template-{uuid.uuid4().hex}.jinja' @@ -399,6 +424,3 @@ class VLLMChatModelHandler(ModelHandler[Sequence[OpenAIChatMessage], An Iterable of type PredictionResult. """ return asyncio.run(self._async_run_inference(batch, model, inference_args)) - - def share_model_across_processes(self) -> bool: - return True
