feat(image): parallelize image_generate batches
This commit is contained in:
parent
e6f1d613b6
commit
c0b0cc3925
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
Loading…
Reference in New Issue