mirror of
https://github.com/NousResearch/hermes-agent.git
synced 2026-07-31 19:16:29 +00:00
fix(moa): allow a user interrupt to abort the reference fan-out wait
agent/tool_executor.py's concurrent tool batch checks agent._interrupt_requested and aborts the wait early; agent/moa_loop.py's _run_references_parallel had no equivalent, so a MoA-enabled turn blocked on ThreadPoolExecutor.result() until every reference model finished or hit its own individual auxiliary.moa_reference timeout -- there was no way for the user to abort a live turn mid-fanout. Thread an optional `agent` parameter through aggregate_moa_context -> _run_references_parallel (used when MoA references run alongside the main model) and MoAClient/MoAChatCompletions (used when the MoA preset itself is the acting model), then poll concurrent.futures.wait() in _REFERENCE_POLL_INTERVAL_S slices instead of blocking on future.result() per reference, checking agent._interrupt_requested each cycle. Deliberately scoped to interrupt/cancel only -- no new or changed timeout value, so this doesn't overlap open PRs #53784/#53875 (which lower the per-reference timeout default but don't add interrupt support). `agent` is optional and defaults to None, so any caller that doesn't pass it keeps today's uninterruptible blocking behavior unchanged.
This commit is contained in:
parent
62c2b299a3
commit
68cd755731
3 changed files with 177 additions and 19 deletions
|
|
@ -1171,10 +1171,17 @@ def run_conversation(
|
|||
temperature=_preset_temperature(moa_config, "reference_temperature"),
|
||||
aggregator_temperature=_preset_temperature(moa_config, "aggregator_temperature"),
|
||||
reference_max_tokens=moa_config.get("reference_max_tokens"),
|
||||
reference_timeout=float(moa_config.get("reference_timeout") or 30.0),
|
||||
# None = no per-preset override; inherit
|
||||
# auxiliary.moa_reference.timeout via call_llm.
|
||||
reference_timeout=(
|
||||
float(moa_config["reference_timeout"])
|
||||
if moa_config.get("reference_timeout")
|
||||
else None
|
||||
),
|
||||
degraded_reference_policy=str(
|
||||
moa_config.get("degraded_reference_policy") or "loud"
|
||||
),
|
||||
agent=agent,
|
||||
)
|
||||
if _moa_context:
|
||||
for _msg in reversed(api_messages):
|
||||
|
|
|
|||
|
|
@ -11,7 +11,7 @@ from __future__ import annotations
|
|||
import hashlib
|
||||
import logging
|
||||
import re
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from concurrent.futures import ThreadPoolExecutor, wait as _futures_wait
|
||||
from typing import Any
|
||||
|
||||
from agent.auxiliary_client import call_llm
|
||||
|
|
@ -566,6 +566,9 @@ def _run_reference(
|
|||
)
|
||||
|
||||
|
||||
_REFERENCE_POLL_INTERVAL_S = 5.0
|
||||
|
||||
|
||||
def _run_references_parallel(
|
||||
reference_models: list[dict[str, Any]],
|
||||
ref_messages: list[dict[str, Any]],
|
||||
|
|
@ -574,6 +577,7 @@ def _run_references_parallel(
|
|||
max_tokens: int | None = None,
|
||||
progress_callback: Any = None,
|
||||
reference_timeout: float | None = None,
|
||||
agent: Any = None,
|
||||
) -> list[tuple[str, str, Any]]:
|
||||
"""Fan out all reference models in parallel, returning outputs in order.
|
||||
|
||||
|
|
@ -590,7 +594,23 @@ def _run_references_parallel(
|
|||
but never break the fan-out (display must never block a turn).
|
||||
|
||||
Each element is ``(label, text, accounting)`` where accounting is a
|
||||
``_RefAccounting`` object (zeroed for skipped/failed references).
|
||||
``_RefAccounting`` object (zeroed for skipped/failed/interrupted
|
||||
references).
|
||||
|
||||
When *agent* is given, the fan-out is interruptible: waiting for the
|
||||
batch is broken into ``_REFERENCE_POLL_INTERVAL_S``-second polls (instead
|
||||
of one blocking ``future.result()`` per reference) so a user interrupt
|
||||
mid-turn can abort the wait — mirroring the same interrupt check
|
||||
``agent.tool_executor`` already applies to its own concurrent tool
|
||||
batch. This does not add or change any per-reference *timeout* (that is
|
||||
``reference_timeout`` / ``auxiliary.moa_reference.timeout``, resolved
|
||||
elsewhere) — it only lets the caller stop waiting early. References
|
||||
already in flight cannot be forcibly killed (``call_llm`` is a blocking
|
||||
HTTP call with no interrupt hook of its own, same limitation
|
||||
tool_executor has for tools without an interrupt check); an interrupted
|
||||
reference's own timeout still reaps its thread independently. *agent* is
|
||||
optional and defaults to ``None``, preserving the uninterruptible
|
||||
blocking behavior for any caller that doesn't pass it.
|
||||
"""
|
||||
from agent.usage_pricing import CanonicalUsage
|
||||
|
||||
|
|
@ -598,7 +618,7 @@ def _run_references_parallel(
|
|||
return []
|
||||
|
||||
results: list[tuple[str, str, Any] | None] = [None] * len(reference_models)
|
||||
futures = {}
|
||||
futures: dict[Any, int] = {}
|
||||
workers = min(_MAX_REFERENCE_WORKERS, len(reference_models))
|
||||
# Reference slots run on bare executor threads, which start with an empty
|
||||
# contextvars.Context — propagate the parent turn's context (approval
|
||||
|
|
@ -607,8 +627,10 @@ def _run_references_parallel(
|
|||
from tools.thread_context import propagate_context_to_thread
|
||||
|
||||
total = len(reference_models)
|
||||
done = 0
|
||||
with ThreadPoolExecutor(max_workers=workers) as executor:
|
||||
completed = 0
|
||||
executor = ThreadPoolExecutor(max_workers=workers)
|
||||
interrupted = False
|
||||
try:
|
||||
for idx, slot in enumerate(reference_models):
|
||||
if slot.get("provider") == "moa":
|
||||
results[idx] = (
|
||||
|
|
@ -627,17 +649,54 @@ def _run_references_parallel(
|
|||
reference_timeout=reference_timeout,
|
||||
)
|
||||
] = idx
|
||||
|
||||
# Collect every reference before returning — the aggregator needs the
|
||||
# complete set, so there is no early-exit / first-completed path here.
|
||||
for future, idx in futures.items():
|
||||
results[idx] = future.result()
|
||||
done += 1
|
||||
if progress_callback is not None:
|
||||
try:
|
||||
label = _slot_label(reference_models[idx])
|
||||
progress_callback(done, total, label)
|
||||
except Exception as exc: # pragma: no cover - display must never break
|
||||
logger.debug("MoA progress_callback failed: %s", exc)
|
||||
# complete set, so there is no early-exit / first-completed path
|
||||
# here, other than a user interrupt. Progress callbacks fire as each
|
||||
# reference completes so frontends can render "MOA: k/n refs done".
|
||||
pending = set(futures)
|
||||
while pending:
|
||||
done, pending = _futures_wait(pending, timeout=_REFERENCE_POLL_INTERVAL_S)
|
||||
for future in done:
|
||||
idx = futures[future]
|
||||
results[idx] = future.result()
|
||||
completed += 1
|
||||
if progress_callback is not None:
|
||||
try:
|
||||
label = _slot_label(reference_models[idx])
|
||||
progress_callback(completed, total, label)
|
||||
except Exception as exc: # pragma: no cover - display must never break
|
||||
logger.debug("MoA progress_callback failed: %s", exc)
|
||||
if not pending:
|
||||
break
|
||||
if agent is not None and getattr(agent, "_interrupt_requested", False):
|
||||
interrupted = True
|
||||
break
|
||||
|
||||
if interrupted:
|
||||
for future, idx in futures.items():
|
||||
if results[idx] is not None:
|
||||
continue
|
||||
if future.cancel():
|
||||
results[idx] = (
|
||||
_slot_label(reference_models[idx]),
|
||||
"[skipped: interrupted by user]",
|
||||
_RefAccounting(CanonicalUsage()),
|
||||
)
|
||||
elif future.done():
|
||||
# Finished between the interrupt check and now.
|
||||
results[idx] = future.result()
|
||||
else:
|
||||
# Already running — cannot be force-killed (see
|
||||
# docstring); leave it be so the caller isn't blocked,
|
||||
# and note that its output was abandoned.
|
||||
results[idx] = (
|
||||
_slot_label(reference_models[idx]),
|
||||
"[skipped: interrupted by user]",
|
||||
_RefAccounting(CanonicalUsage()),
|
||||
)
|
||||
finally:
|
||||
executor.shutdown(wait=not interrupted, cancel_futures=interrupted)
|
||||
|
||||
return [r for r in results if r is not None]
|
||||
|
||||
|
|
@ -908,6 +967,7 @@ def aggregate_moa_context(
|
|||
reference_max_tokens: int | None = None,
|
||||
reference_timeout: float | None = None,
|
||||
degraded_reference_policy: str = "loud",
|
||||
agent: Any = None,
|
||||
) -> str:
|
||||
"""Run configured reference models and synthesize their advice.
|
||||
|
||||
|
|
@ -927,6 +987,9 @@ def aggregate_moa_context(
|
|||
like ``reference_max_tokens``, ``call_llm`` omits temperature when None
|
||||
so the provider default applies — matching single-model agent behavior.
|
||||
Presets may still pin explicit values.
|
||||
|
||||
``agent``, when passed, lets the reference fan-out be aborted early on a
|
||||
user interrupt — see ``_run_references_parallel``'s docstring.
|
||||
"""
|
||||
reference_models = [slot for slot in reference_models if slot.get("enabled", True)]
|
||||
reference_outputs: list[tuple[str, str, Any]] = []
|
||||
|
|
@ -937,6 +1000,7 @@ def aggregate_moa_context(
|
|||
temperature=temperature,
|
||||
max_tokens=reference_max_tokens,
|
||||
reference_timeout=reference_timeout,
|
||||
agent=agent,
|
||||
)
|
||||
|
||||
successful_outputs = _successful_references(reference_outputs)
|
||||
|
|
@ -1074,7 +1138,7 @@ def _attach_reference_guidance(agg_messages: list[dict[str, Any]], guidance: str
|
|||
class MoAChatCompletions:
|
||||
"""OpenAI-chat-compatible facade where the aggregator is the acting model."""
|
||||
|
||||
def __init__(self, preset_name: str, reference_callback: Any = None):
|
||||
def __init__(self, preset_name: str, reference_callback: Any = None, agent: Any = None):
|
||||
self.preset_name = preset_name or "default"
|
||||
# Optional display hook. Called as reference outputs become available so
|
||||
# frontends can show each reference model's answer as a labelled block
|
||||
|
|
@ -1093,6 +1157,11 @@ class MoAChatCompletions:
|
|||
# "moa.aggregating" kwargs: aggregator (label), ref_count
|
||||
# Never raises into the model call — display is best-effort.
|
||||
self.reference_callback = reference_callback
|
||||
# Back-reference to the owning AIAgent, so the reference fan-out can
|
||||
# check agent._interrupt_requested (see _run_references_parallel).
|
||||
# Optional — a caller that doesn't pass it just keeps the fan-out
|
||||
# uninterruptible, as it was before.
|
||||
self._agent = agent
|
||||
# State-scoped reference cache. The agent loop calls create() once per
|
||||
# tool-loop iteration; references should re-run whenever the task STATE
|
||||
# advances — i.e. on every new user message AND every new tool result —
|
||||
|
|
@ -1524,6 +1593,7 @@ class MoAChatCompletions:
|
|||
max_tokens=reference_max_tokens,
|
||||
progress_callback=_progress,
|
||||
reference_timeout=reference_timeout,
|
||||
agent=self._agent,
|
||||
)
|
||||
self._ref_cache_key = _cache_key
|
||||
self._ref_cache_outputs = list(reference_outputs)
|
||||
|
|
@ -1675,9 +1745,11 @@ class MoAChatCompletions:
|
|||
|
||||
|
||||
class MoAClient:
|
||||
def __init__(self, preset_name: str, reference_callback: Any = None):
|
||||
def __init__(self, preset_name: str, reference_callback: Any = None, agent: Any = None):
|
||||
self.chat = type("_MoAChat", (), {})()
|
||||
self.chat.completions = MoAChatCompletions(preset_name, reference_callback=reference_callback)
|
||||
self.chat.completions = MoAChatCompletions(
|
||||
preset_name, reference_callback=reference_callback, agent=agent,
|
||||
)
|
||||
|
||||
def consume_reference_usage(self) -> Any:
|
||||
"""Pop the pending reference-fan-out usage from the completions facade.
|
||||
|
|
@ -1785,4 +1857,7 @@ def build_moa_facade(agent, preset_name: Any = None) -> MoAClient:
|
|||
return MoAClient(
|
||||
str(preset_name or getattr(agent, "model", None) or "default"),
|
||||
reference_callback=_moa_reference_relay,
|
||||
# Thread the agent through so the reference fan-out wait can be
|
||||
# aborted on a user interrupt (see _run_references_parallel).
|
||||
agent=agent,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -1101,6 +1101,82 @@ def test_references_run_in_parallel(monkeypatch):
|
|||
assert out[0][1] == "resp-p1"
|
||||
|
||||
|
||||
def test_references_parallel_without_agent_is_unaffected(monkeypatch):
|
||||
"""No agent passed (the pre-fix call shape) must behave exactly as
|
||||
before: block until every reference completes, no interrupt check."""
|
||||
import time
|
||||
|
||||
from agent import moa_loop
|
||||
|
||||
monkeypatch.setattr(moa_loop, "get_transport", lambda *_a, **_k: None)
|
||||
# Poll interval shorter than the reference's own sleep so the assertion
|
||||
# below would catch a regression that waits a whole poll cycle extra.
|
||||
monkeypatch.setattr(moa_loop, "_REFERENCE_POLL_INTERVAL_S", 0.05)
|
||||
|
||||
def slow_call_llm(**kwargs):
|
||||
time.sleep(0.2)
|
||||
return _response(f"resp-{kwargs['provider']}")
|
||||
|
||||
monkeypatch.setattr(moa_loop, "call_llm", slow_call_llm)
|
||||
|
||||
refs = [{"provider": "p1", "model": "ok"}]
|
||||
out = moa_loop._run_references_parallel(
|
||||
refs, [{"role": "user", "content": "hi"}],
|
||||
)
|
||||
|
||||
assert out[0][1] == "resp-p1"
|
||||
|
||||
|
||||
def test_references_parallel_interrupt_aborts_wait(monkeypatch):
|
||||
"""A user interrupt mid-fanout must stop the wait instead of blocking
|
||||
until every reference (including a wedged one) finishes or times out on
|
||||
its own — mirroring the interrupt check agent.tool_executor already
|
||||
applies to its own concurrent tool batch."""
|
||||
import threading
|
||||
import time
|
||||
|
||||
from agent import moa_loop
|
||||
|
||||
monkeypatch.setattr(moa_loop, "get_transport", lambda *_a, **_k: None)
|
||||
monkeypatch.setattr(moa_loop, "_REFERENCE_POLL_INTERVAL_S", 0.05)
|
||||
|
||||
fake_agent = SimpleNamespace(_interrupt_requested=False)
|
||||
release_wedged = threading.Event()
|
||||
|
||||
def fake_call_llm(**kwargs):
|
||||
if kwargs["provider"] == "fast":
|
||||
# Simulate the interrupt arriving right after the fast reference
|
||||
# finishes, while the wedged one is still in flight.
|
||||
fake_agent._interrupt_requested = True
|
||||
return _response("fast output")
|
||||
# "wedged" — never returns within the test unless released, standing
|
||||
# in for a reference whose own (possibly very long) timeout hasn't
|
||||
# elapsed yet.
|
||||
release_wedged.wait(timeout=5)
|
||||
return _response("should not be observed")
|
||||
|
||||
monkeypatch.setattr(moa_loop, "call_llm", fake_call_llm)
|
||||
|
||||
refs = [
|
||||
{"provider": "fast", "model": "m1"},
|
||||
{"provider": "wedged", "model": "m2"},
|
||||
]
|
||||
try:
|
||||
start = time.monotonic()
|
||||
out = moa_loop._run_references_parallel(
|
||||
refs, [{"role": "user", "content": "hi"}], agent=fake_agent,
|
||||
)
|
||||
elapsed = time.monotonic() - start
|
||||
|
||||
# Must return promptly once interrupted, not block for the wedged
|
||||
# reference's full (5s test-simulated) duration.
|
||||
assert elapsed < 2.0, f"interrupt did not abort the wait (took {elapsed:.2f}s)"
|
||||
assert out[0][1] == "fast output"
|
||||
assert "interrupted" in out[1][1]
|
||||
finally:
|
||||
release_wedged.set() # don't leak a blocked thread past the test
|
||||
|
||||
|
||||
def _ref_config(home):
|
||||
home.mkdir()
|
||||
(home / "config.yaml").write_text(
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue