You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@beam.apache.org by "damccorm (via GitHub)" <gi...@apache.org> on 2026/06/23 16:04:28 UTC
[PR] Extend ADKAgentModelHandler to support deploying local models [beam]
damccorm opened a new pull request, #39074:
URL: https://github.com/apache/beam/pull/39074
## Problem
Currently, `ADKAgentModelHandler` only supports deploying and calling remote model endpoints. This limits its usability in scenarios where the user wants to deploy a local model (e.g., using a local PyTorch or HuggingFace model) on the Beam workers.
## Solution
This PR extends `ADKAgentModelHandler` to allow deploying a local model serving process instead of relying on a remote model endpoint.
Key changes:
1. **Generic SubProcess Wrapper**: Introduced `SubProcessModel` (inherits from `SubprocessModelHandler`) in `apache_beam/ml/inference/base.py`. This wrapper can wrap *any* standard `ModelHandler` and serve it via a FastAPI server running in a subprocess.
2. **FastAPI Subprocess Server**: Added `subprocess_server.py` which runs in the subprocess. It exposes:
* OpenAI-compatible `/v1/chat/completions` and `/v1/completions` endpoints (for compatibility with ADK's `LiteLlm` client).
* A transparent pickle-based binary `/v1/beam/inference` endpoint to preserve exact input/output types when used directly in Beam pipelines.
3. **ADK Integration**: Updated `ADKAgentModelHandler` to:
* Accept an `underlying_model_handler` (which must be a `SubprocessModelHandler`, such as `SubProcessModel` or `VLLMCompletionsModelHandler`).
* Manage the lifecycle of the local server process.
* Propagate the local model configuration (endpoint base URL with dynamically selected port) to the root agent, subagents, and tools.
* Handle dynamic port updates if the server restarts.
* Implement recovery/connectivity checks on inference failures.
4. **Refactoring**: Refactored `vllm_inference.py` to use a new generic `_VLLMBaseModelHandler` to eliminate code duplication between completions and chat handlers.
5. **Robustness & Safety**:
* Improved temporary directory lifecycle management to prevent premature cleanup by GC.
* Bound the subprocess server to `127.0.0.1` instead of `0.0.0.0` for security.
* Improved error propagation and logging for both client and server sides.
## Verification
* Added comprehensive unit tests in `base_test.py` (`SubProcessModelTest`) and `agent_development_kit_test.py` (`TestLocalModelIntegration`).
* Verified all tests pass.
TAG=agy
CONV=b2f450cf-8b63-49a9-8ee8-430ed676e116
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Extend ADKAgentModelHandler to support deploying local models [beam]
Posted by "damccorm (via GitHub)" <gi...@apache.org>.
damccorm commented on code in PR #39074:
URL: https://github.com/apache/beam/pull/39074#discussion_r3461189889
##########
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
+
+@app.post("/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
+ }
+ }]
+ }
+
+@app.post("/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
+ }]
+ }
+
+@app.post("/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)
Review Comment:
Not a concern - this traffic is coming from a separate Beam process (with the same permissions)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Extend ADKAgentModelHandler to support deploying local models [beam]
Posted by "damccorm (via GitHub)" <gi...@apache.org>.
damccorm commented on PR #39074:
URL: https://github.com/apache/beam/pull/39074#issuecomment-4781102576
/gemini review
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Extend ADKAgentModelHandler to support deploying local models [beam]
Posted by "github-advanced-security[bot] (via GitHub)" <gi...@apache.org>.
github-advanced-security[bot] commented on code in PR #39074:
URL: https://github.com/apache/beam/pull/39074#discussion_r3461182592
##########
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
+
+@app.post("/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
+ }
+ }]
+ }
+
+@app.post("/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
+ }]
+ }
+
+@app.post("/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)
Review Comment:
## CodeQL / Deserialization of user-controlled data
Unsafe deserialization depends on a [user-provided value](1).
[Show more details](https://github.com/apache/beam/security/code-scanning/93)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Extend ADKAgentModelHandler to support deploying local models [beam]
Posted by "gemini-code-assist[bot] (via GitHub)" <gi...@apache.org>.
gemini-code-assist[bot] commented on code in PR #39074:
URL: https://github.com/apache/beam/pull/39074#discussion_r3461219111
##########
sdks/python/apache_beam/ml/inference/agent_development_kit.py:
##########
@@ -259,6 +389,26 @@ async def _run_concurrently():
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)
Review Comment:
Similarly, if `agent.tools` is `None`, iterating over it here will raise a `TypeError`. We should safely check if `agent.tools` is not `None` before iterating.
```suggestion
if getattr(agent, 'tools', None) is not None:
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)
```
##########
sdks/python/apache_beam/ml/inference/base.py:
##########
@@ -406,6 +408,225 @@ def should_garbage_collect_on_timeout(self) -> bool:
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:
Review Comment:
If the subprocess server crashes, `self._server_started` remains `True`. If `get_server_port()` is called subsequently, it will return the port without restarting the server because `self._server_started` is still `True`. We should check if the process has exited and reset `self._server_started = False` at the beginning of `start_server` to ensure recovery.
```python
def start_server(self, retries=3):
with self._server_process_lock:
if self._process and self._process.poll() is not None:
self._server_started = False
if not self._server_started:
```
##########
sdks/python/apache_beam/ml/inference/agent_development_kit.py:
##########
@@ -181,13 +264,33 @@ def load_model(self) -> "Runner":
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,
self._app_name,
)
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)
Review Comment:
If `agent.tools` is `None` (which is the default in ADK when no tools are provided), calling `hasattr(agent, 'tools')` will return `True`, but attempting to iterate over it will raise a `TypeError: 'NoneType' object is not iterable`. We should use `getattr(agent, 'tools', None) is not None` to safely guard against this.
```suggestion
if getattr(agent, 'tools', None) is not None:
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)
```
##########
sdks/python/apache_beam/ml/inference/base.py:
##########
@@ -406,6 +408,225 @@ def should_garbage_collect_on_timeout(self) -> bool:
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)
Review Comment:
There are two critical issues in `check_connectivity` and `start_server`:
1. **Race Condition / Redundant Restarts**: If multiple threads concurrently call `check_connectivity` due to an inference failure, they will all call `start_server`. While `start_server` is protected by `self._server_process_lock`, once the first thread finishes restarting the server and releases the lock, the waiting threads will acquire the lock and restart the server again, terminating the newly started healthy server. Since `self._server_process_lock` is an `RLock`, we should wrap the entire `check_connectivity` method with this lock so that subsequent threads immediately see the healthy server and return without redundant restarts.
2. **No-op Restart Loop**: If the server crashes, `self._server_started` remains `True`. When `check_connectivity` is called, the polling loop is skipped because `self._process.poll()` is not `None`. It then calls `start_server(retries - 1)`. However, because `self._server_started` is still `True`, `start_server` skips the spawning block entirely and just calls `check_connectivity` again. This results in a loop of no-op retries that never actually restarts the process. We must reset `self._server_started = False` before calling `start_server` to ensure the server is actually restarted.
```python
def check_connectivity(self, retries=3):
with self._server_process_lock:
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._server_started = False
self.start_server(retries - 1)
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Extend ADKAgentModelHandler to support deploying local models [beam]
Posted by "damccorm (via GitHub)" <gi...@apache.org>.
damccorm commented on PR #39074:
URL: https://github.com/apache/beam/pull/39074#issuecomment-4781146101
I have addressed all the feedback from the code review agent. Pushed updates to the branch.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Extend ADKAgentModelHandler to support deploying local models [beam]
Posted by "github-actions[bot] (via GitHub)" <gi...@apache.org>.
github-actions[bot] commented on PR #39074:
URL: https://github.com/apache/beam/pull/39074#issuecomment-5570843096
This pull request has been marked as stale due to 60 days of inactivity. It will be closed in 1 week if no further activity occurs. If you think that’s incorrect or this pull request requires a review, please simply write any comment. If closed, you can revive the PR at any time and @mention a reviewer or discuss it on the dev@beam.apache.org list. Thank you for your contributions.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Extend ADKAgentModelHandler to support deploying local models [beam]
Posted by "github-actions[bot] (via GitHub)" <gi...@apache.org>.
github-actions[bot] commented on PR #39074:
URL: https://github.com/apache/beam/pull/39074#issuecomment-5664181676
This pull request has been closed due to lack of activity. If you think that is incorrect, or the pull request requires review, you can revive the PR at any time.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Extend ADKAgentModelHandler to support deploying local models [beam]
Posted by "github-actions[bot] (via GitHub)" <gi...@apache.org>.
github-actions[bot] closed pull request #39074: Extend ADKAgentModelHandler to support deploying local models
URL: https://github.com/apache/beam/pull/39074
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Extend ADKAgentModelHandler to support deploying local models [beam]
Posted by "gemini-code-assist[bot] (via GitHub)" <gi...@apache.org>.
gemini-code-assist[bot] commented on code in PR #39074:
URL: https://github.com/apache/beam/pull/39074#discussion_r3461188359
##########
sdks/python/apache_beam/ml/inference/agent_development_kit.py:
##########
@@ -165,11 +193,54 @@ def load_model(self) -> "Runner":
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.")
Review Comment:
The current factory inspection logic strictly expects either exactly 0 or 1 parameter. This raises a `ValueError` if the factory has optional parameters (e.g., `def factory(model=None)` or `def factory(model, config=None)`), even though they are perfectly valid to call. Inspecting only the required parameters (those without default values) makes the factory invocation much more robust.
```python
import inspect
sig = inspect.signature(self._agent_or_factory)
params = list(sig.parameters.values())
required_params = [
p for p in params
if p.default is inspect.Parameter.empty and p.kind not in (
inspect.Parameter.VAR_POSITIONAL, inspect.Parameter.VAR_KEYWORD)
]
if len(required_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(required_params) == 0:
if local_model is not None and len(params) > 0:
agent = self._agent_or_factory(local_model)
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)
else:
raise ValueError("Agent factory must take 0 or 1 required argument.")
```
##########
sdks/python/apache_beam/ml/inference/base.py:
##########
@@ -406,6 +408,219 @@ def should_garbage_collect_on_timeout(self) -> bool:
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)
Review Comment:
During unpickling or deepcopying, `__getattr__` can be called before `_handler` is fully restored in `self.__dict__`. This leads to infinite recursion and a `RecursionError` because `getattr(self, '_handler')` will recursively trigger `__getattr__`. Checking if `_handler` is present in `self.__dict__` avoids this issue.
```suggestion
def __getattr__(self, name):
if name == '_handler':
raise AttributeError()
if '_handler' not in self.__dict__:
raise AttributeError(f"'{type(self).__name__}' object has no attribute '{name}'")
return getattr(self._handler, name)
```
##########
sdks/python/apache_beam/ml/inference/agent_development_kit.py:
##########
@@ -73,6 +74,30 @@
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]
Review Comment:
When `ADK_AVAILABLE` is `False`, the `Agent` and `Runner` classes are not imported or defined. This will cause a `NameError` at runtime when `isinstance(tool, Agent)` is evaluated in `_set_agent_model` and `_update_agent_port`. Defining dummy fallback classes in the `else` block safely avoids this issue.
```suggestion
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")
class Agent:
pass
class Runner:
pass
genai_Content = Any # type: ignore[assignment, misc]
genai_Part = Any # type: ignore[assignment, misc]
```
##########
sdks/python/apache_beam/ml/inference/base.py:
##########
@@ -406,6 +408,219 @@ def should_garbage_collect_on_timeout(self) -> bool:
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
Review Comment:
The HTTP inference timeout is currently hardcoded to 30 seconds. For heavy ML models or large batches, this can easily lead to premature timeouts. Adding a configurable `timeout` parameter (defaulting to a safer 300 seconds) provides better flexibility and production reliability.
```suggestion
def __init__(
self,
handler: ModelHandler[ExampleT, PredictionT, ModelT],
model_name: str,
timeout: int = 300):
super().__init__()
self._handler = handler
self._model_name = model_name
self._timeout = timeout
self._handler_path = None
```
##########
sdks/python/apache_beam/ml/inference/base.py:
##########
@@ -406,6 +408,219 @@ def should_garbage_collect_on_timeout(self) -> bool:
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
Review Comment:
Use the configurable `self._timeout` instead of the hardcoded 30-second timeout to prevent timeout failures on larger batches or slower workers.
```suggestion
try:
with urllib.request.urlopen(req, timeout=self._timeout) as response:
resp_data = response.read()
results = pickle.loads(resp_data)
return results
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Extend ADKAgentModelHandler to support deploying local models [beam]
Posted by "damccorm (via GitHub)" <gi...@apache.org>.
damccorm commented on PR #39074:
URL: https://github.com/apache/beam/pull/39074#issuecomment-4781150189
/gemini review
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Extend ADKAgentModelHandler to support deploying local models [beam]
Posted by "damccorm (via GitHub)" <gi...@apache.org>.
damccorm commented on code in PR #39074:
URL: https://github.com/apache/beam/pull/39074#discussion_r3499997220
##########
sdks/python/apache_beam/ml/inference/agent_development_kit.py:
##########
@@ -259,6 +389,26 @@ async def _run_concurrently():
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)
Review Comment:
Fixed. Safely checking if tools is not None before iterating. Added test_local_model_injection_with_none_tools to verify.
##########
sdks/python/apache_beam/ml/inference/base.py:
##########
@@ -406,6 +408,225 @@ def should_garbage_collect_on_timeout(self) -> bool:
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:
Review Comment:
Fixed. Resetting _server_started to False if the process has exited, which triggers a restart on the next start_server call. Added test_subprocess_server_recovery to verify.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Extend ADKAgentModelHandler to support deploying local models [beam]
Posted by "damccorm (via GitHub)" <gi...@apache.org>.
damccorm commented on code in PR #39074:
URL: https://github.com/apache/beam/pull/39074#discussion_r4048131931
##########
sdks/python/apache_beam/ml/inference/base.py:
##########
@@ -406,6 +408,225 @@ def should_garbage_collect_on_timeout(self) -> bool:
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)
Review Comment:
Fixed. Wrapped check_connectivity with self._server_process_lock and reset self._server_started = False before calling start_server.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org