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
The following commit(s) were added to refs/heads/feature/adk-local-model by
this push:
new 692775a3ff4 Address PR review comments on ADK local model deployment
692775a3ff4 is described below
commit 692775a3ff44992b522092562ab45638c733b69f
Author: Danny McCormick <[email protected]>
AuthorDate: Tue Jun 30 15:37:57 2026 +0000
Address PR review comments on ADK local model deployment
- Safely check if agent.tools is not None before iterating in
ADKAgentModelHandler.
- Reset _server_started to False if the subprocess server crashes, allowing
recovery.
- Add unit tests for both fixes.
TAG=agy
CONV=b2f450cf-8b63-49a9-8ee8-430ed676e116
---
.../ml/inference/agent_development_kit.py | 4 +--
.../ml/inference/agent_development_kit_test.py | 31 ++++++++++++++++++++++
sdks/python/apache_beam/ml/inference/base.py | 2 ++
sdks/python/apache_beam/ml/inference/base_test.py | 28 +++++++++++++++++++
4 files changed, 63 insertions(+), 2 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 d134a92076b..1f5d5e5f745 100644
--- a/sdks/python/apache_beam/ml/inference/agent_development_kit.py
+++ b/sdks/python/apache_beam/ml/inference/agent_development_kit.py
@@ -284,7 +284,7 @@ class ADKAgentModelHandler(ModelHandler[str | genai_Content,
agent.model = model
# Speculative propagation to subagents/tools
- if hasattr(agent, 'tools'):
+ 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)
@@ -397,7 +397,7 @@ class ADKAgentModelHandler(ModelHandler[str | genai_Content,
model=agent.model.model,
api_base=f"http://localhost:{port}/v1";
)
- if hasattr(agent, 'tools'):
+ if getattr(agent, 'tools', None) is not None:
for tool in agent.tools:
if hasattr(tool, 'agent'):
self._update_agent_port(tool.agent, port)
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 214a0379f93..6fc7a79008c 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
@@ -475,6 +475,37 @@ class TestLocalModelIntegration(unittest.TestCase):
self.assertEqual(str(ctx.exception), "Original error")
self.mock_handler.check_connectivity.assert_called_once_with(self.mock_model)
+ def test_local_model_injection_with_none_tools(self):
+ root_agent = Agent(
+ model=BeamPlaceholderModel(),
+ name="root_agent"
+ )
+ object.__setattr__(root_agent, 'tools', None)
+
+ 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)
+
+ def test_port_update_propagation_with_none_tools(self):
+ root_agent = Agent(model=BeamPlaceholderModel(), name="root_agent")
+ object.__setattr__(root_agent, 'tools', None)
+
+ 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.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(handler._current_port, 54321)
+
+
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 4ae84ded69c..2a7e54680da 100644
--- a/sdks/python/apache_beam/ml/inference/base.py
+++ b/sdks/python/apache_beam/ml/inference/base.py
@@ -440,6 +440,8 @@ class SubProcessModelServer:
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:
if self._process:
logging.info("Terminating existing generic subprocess model server
before restart")
diff --git a/sdks/python/apache_beam/ml/inference/base_test.py
b/sdks/python/apache_beam/ml/inference/base_test.py
index 1fb0955623f..f904c1d8eea 100644
--- a/sdks/python/apache_beam/ml/inference/base_test.py
+++ b/sdks/python/apache_beam/ml/inference/base_test.py
@@ -2491,6 +2491,34 @@ class SubProcessModelTest(unittest.TestCase):
del model
del wrapper
+ def test_subprocess_server_recovery(self):
+ handler = StringFakeModelHandler()
+ wrapper = base.SubProcessModel(handler, model_name="string_fake_model")
+
+ model = wrapper.load_model()
+ self.assertIsInstance(model, base.SubProcessModelServer)
+
+ # Verify it works
+ results = wrapper.run_inference(["hello"], model)
+ results = list(results)
+ self.assertEqual(results[0].inference, "hello_processed")
+
+ # Forcefully kill the process
+ model._process.kill()
+ model._process.wait()
+
+ # Trigger connectivity check which should restart the server
+ model.check_connectivity()
+
+ # Verify it recovered and works again
+ results = wrapper.run_inference(["world"], model)
+ results = list(results)
+ self.assertEqual(results[0].inference, "world_processed")
+
+ del model
+ del wrapper
+
+
if __name__ == '__main__':
unittest.main()
