From c0b0cc39256eac041790db305b29e4cca48290ec Mon Sep 17 00:00:00 2001 From: EndeavorYen Date: Fri, 29 May 2026 00:10:59 +0800 Subject: [PATCH] feat(image): parallelize image_generate batches --- agent/tool_dispatch_helpers.py | 1 + agent/tool_executor.py | 40 +++++++- .../run_agent/test_image_generate_parallel.py | 97 +++++++++++++++++++ 3 files changed, 137 insertions(+), 1 deletion(-) create mode 100644 tests/run_agent/test_image_generate_parallel.py diff --git a/agent/tool_dispatch_helpers.py b/agent/tool_dispatch_helpers.py index f7f003f24abb0..6e76e52081779 100644 --- a/agent/tool_dispatch_helpers.py +++ b/agent/tool_dispatch_helpers.py @@ -48,6 +48,7 @@ _PARALLEL_SAFE_TOOLS = frozenset({ "ha_get_state", "ha_list_entities", "ha_list_services", + "image_generate", "read_file", "search_files", "session_search", diff --git a/agent/tool_executor.py b/agent/tool_executor.py index c422c118e7275..a4ce0031a5e83 100644 --- a/agent/tool_executor.py +++ b/agent/tool_executor.py @@ -93,6 +93,7 @@ def _budget_for_agent(agent) -> BudgetConfig: # Maximum number of concurrent worker threads for parallel tool execution. # Mirrors the constant in ``run_agent`` for tests/imports that look here. _MAX_TOOL_WORKERS = 8 +_DEFAULT_IMAGE_PARALLEL_REQUESTS = 4 # Keep this above the stock auxiliary.web_extract timeout (360s) so the batch # guard does not preempt a slow-but-valid summarization attempt. _DEFAULT_CONCURRENT_TOOL_TIMEOUT_S = 420.0 @@ -159,6 +160,43 @@ def _flush_session_db_after_tool_progress( return False +def _image_generate_parallel_limit() -> int: + """Return the configured image-generation parallelism cap. + + Image-generation calls are slow enough that concurrent execution is useful, + but backend bursts can hit TTFB or rate-limit failures. Keep the default + intentionally conservative while allowing users to tune it per install. + """ + try: + from hermes_cli.config import load_config + + cfg = load_config() or {} + image_gen = cfg.get("image_gen") if isinstance(cfg, dict) else None + value = ( + image_gen.get("max_parallel_requests") + if isinstance(image_gen, dict) + else None + ) + except Exception: + value = None + + try: + limit = int(value) + except (TypeError, ValueError): + limit = _DEFAULT_IMAGE_PARALLEL_REQUESTS + return max(1, min(limit, _MAX_TOOL_WORKERS)) + + +def _max_workers_for_tool_batch(runnable_calls) -> int: + """Return the worker cap for a concurrent tool batch.""" + if not runnable_calls: + return 0 + max_workers = _MAX_TOOL_WORKERS + if any(name == "image_generate" for _, _, name, _ in runnable_calls): + max_workers = min(max_workers, _image_generate_parallel_limit()) + return min(len(runnable_calls), max_workers) + + def _ra(): """Lazy reference to ``run_agent`` so patches like ``run_agent._set_interrupt`` work.""" import run_agent @@ -922,7 +960,7 @@ def execute_tool_calls_concurrent(agent, assistant_message, messages: list, effe timeout_s = _resolve_concurrent_tool_timeout() deadline = time.monotonic() + timeout_s if timeout_s is not None else None if runnable_calls: - max_workers = min(len(runnable_calls), _MAX_TOOL_WORKERS) + max_workers = _max_workers_for_tool_batch(runnable_calls) # Daemon workers: an interrupted/timed-out batch is abandoned with # shutdown(wait=False), but stdlib ThreadPoolExecutor workers are # non-daemon and registered in concurrent.futures' atexit hook, diff --git a/tests/run_agent/test_image_generate_parallel.py b/tests/run_agent/test_image_generate_parallel.py new file mode 100644 index 0000000000000..3223be9d9d9a1 --- /dev/null +++ b/tests/run_agent/test_image_generate_parallel.py @@ -0,0 +1,97 @@ +"""Regression tests for parallel image-generation tool batches.""" + +import json +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +import run_agent +from agent import tool_executor + + +def _tool_call(name: str, args: dict, call_id: str) -> SimpleNamespace: + return SimpleNamespace( + id=call_id, + function=SimpleNamespace( + name=name, + arguments=json.dumps(args), + ), + ) + + +def test_image_generate_batch_routes_to_concurrent_executor(): + agent = SimpleNamespace() + agent._execute_tool_calls = run_agent.AIAgent._execute_tool_calls.__get__(agent) + agent._execute_tool_calls_concurrent = MagicMock() + agent._execute_tool_calls_sequential = MagicMock() + assistant_message = SimpleNamespace( + tool_calls=[ + _tool_call("image_generate", {"prompt": "variation one"}, "img_1"), + _tool_call("image_generate", {"prompt": "variation two"}, "img_2"), + ], + ) + + agent._execute_tool_calls(assistant_message, [], "task-image-batch") + + agent._execute_tool_calls_concurrent.assert_called_once() + agent._execute_tool_calls_sequential.assert_not_called() + + +def test_image_generate_parallel_worker_cap_defaults_to_four(): + runnable_calls = [ + ( + 0, + _tool_call("image_generate", {"prompt": "one"}, "img_1"), + "image_generate", + {}, + ), + ( + 1, + _tool_call("image_generate", {"prompt": "two"}, "img_2"), + "image_generate", + {}, + ), + ( + 2, + _tool_call("image_generate", {"prompt": "three"}, "img_3"), + "image_generate", + {}, + ), + ( + 3, + _tool_call("image_generate", {"prompt": "four"}, "img_4"), + "image_generate", + {}, + ), + ( + 4, + _tool_call("image_generate", {"prompt": "five"}, "img_5"), + "image_generate", + {}, + ), + ] + + with patch("hermes_cli.config.load_config", return_value={}): + assert tool_executor._max_workers_for_tool_batch(runnable_calls) == 4 + + +def test_image_generate_parallel_worker_cap_can_be_configured_lower(): + runnable_calls = [ + ( + 0, + _tool_call("image_generate", {"prompt": "one"}, "img_1"), + "image_generate", + {}, + ), + ( + 1, + _tool_call("image_generate", {"prompt": "two"}, "img_2"), + "image_generate", + {}, + ), + ] + + with patch( + "hermes_cli.config.load_config", + return_value={"image_gen": {"max_parallel_requests": 1}}, + ): + assert tool_executor._max_workers_for_tool_batch(runnable_calls) == 1