mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-09-09 21:49:37 +00:00
fix(subagents): harden background-task registry and capacity snapshot edge cases (#5086)
* fix(subagents): harden background-task registry and capacity snapshot edge cases - execute_async drops the just-registered background entry when submitting to the isolated loop fails. The caller sees the exception and never polls, and cleanup_background_task refuses non-terminal entries, so the entry would otherwise stay as a PENDING zombie forever. - SubagentExecutionCapacity.snapshot derives queued from len(_waiters) instead of iterating it. snapshot is read from non-loop threads (e.g. configure_subagent_execution_capacity) while the loop thread mutates the deque, so iteration can raise 'deque mutated during iteration'. The raw length may count a waiter that just timed out but has not removed itself yet, which only makes the busy-check more conservative. * fix(subagents): close failure-path gaps around background submit Address both review findings on the background submission lifecycle: - execute_async() copies the isolated-loop context before registering the _background_tasks entry, so a context-copy failure (callback- manager copy or loop-bound handler filtering) can no longer strand a permanent PENDING entry the caller will never poll. - _submit_to_isolated_loop_in_context() resolves the loop before calling the coroutine factory. As direct run_coroutine_threadsafe arguments the coroutine was created first, so a loop-startup failure stranded a never-awaited coroutine (RuntimeWarning + retained captures until collection). Both call sites share the fix. New tests verified red on the previous implementation, green after: - context-copy failure leaves no registry residue - the real submit helper (only the loop getter patched) never invokes the coroutine factory when loop startup fails * fix(subagents): close the coroutine when scheduling rejects it run_coroutine_threadsafe can itself raise once the coroutine exists (e.g. the loop closes between the lookup and the internal call_soon_threadsafe). Wrap the call, close the rejected coroutine, and re-raise; a focused test patches only run_coroutine_threadsafe and asserts the created coroutine reaches CORO_CLOSED.
This commit is contained in:
parent
8a830f6354
commit
2f8d1cfc21
@ -43,11 +43,18 @@ class SubagentExecutionCapacity:
|
||||
self._waiters: deque[asyncio.Future[None]] = deque()
|
||||
|
||||
def snapshot(self) -> SubagentCapacitySnapshot:
|
||||
# snapshot() is read from non-loop threads (notably
|
||||
# configure_subagent_execution_capacity), while the loop thread mutates
|
||||
# the waiters deque under its own asyncio lock — iterating cross-thread
|
||||
# can raise "deque mutated during iteration". ``len()`` reads the
|
||||
# deque's size atomically instead. The raw length may count a waiter
|
||||
# that just timed out but has not removed itself yet; that only makes
|
||||
# the "capacity busy" answer more conservative, never less.
|
||||
return SubagentCapacitySnapshot(
|
||||
max_running=self._config.max_running,
|
||||
running=self._running,
|
||||
max_queued=self._config.max_queued,
|
||||
queued=sum(not waiter.done() for waiter in self._waiters),
|
||||
queued=len(self._waiters),
|
||||
admission_policy=self._config.admission_policy,
|
||||
)
|
||||
|
||||
|
||||
@ -435,13 +435,34 @@ def _submit_to_isolated_loop_in_context(
|
||||
context: Context,
|
||||
coro_factory: Callable[[], Coroutine[Any, Any, SubagentResult]],
|
||||
) -> Future[SubagentResult]:
|
||||
"""Submit a coroutine to the isolated loop while preserving ContextVar state."""
|
||||
return context.run(
|
||||
lambda: asyncio.run_coroutine_threadsafe(
|
||||
coro_factory(),
|
||||
_get_isolated_subagent_loop(),
|
||||
)
|
||||
)
|
||||
"""Submit a coroutine to the isolated loop while preserving ContextVar state.
|
||||
|
||||
The loop must be resolved before the coroutine is created: as direct
|
||||
``run_coroutine_threadsafe(coro_factory(), ...)`` arguments, Python
|
||||
evaluates the coroutine first, so a loop-startup failure would strand a
|
||||
created-but-never-scheduled coroutine (``RuntimeWarning: coroutine ...
|
||||
was never awaited``) holding its captures until collection.
|
||||
|
||||
Scheduling itself can still reject an already-created coroutine — e.g.
|
||||
the loop closes between the lookup above and the ``call_soon_threadsafe``
|
||||
inside ``run_coroutine_threadsafe`` — so a rejected coroutine is closed
|
||||
before the error propagates.
|
||||
"""
|
||||
|
||||
def _submit() -> Future[SubagentResult]:
|
||||
loop = _get_isolated_subagent_loop()
|
||||
coroutine = coro_factory()
|
||||
try:
|
||||
return asyncio.run_coroutine_threadsafe(coroutine, loop)
|
||||
except BaseException:
|
||||
# run_coroutine_threadsafe has no cleanup path for this window.
|
||||
# The coroutine has not started (CORO_CREATED), so close() cannot
|
||||
# run any of its body — it only releases the object and its
|
||||
# captures instead of leaving them until collection.
|
||||
coroutine.close()
|
||||
raise
|
||||
|
||||
return context.run(_submit)
|
||||
|
||||
|
||||
def _copy_isolated_subagent_context() -> Context:
|
||||
@ -1441,11 +1462,16 @@ class SubagentExecutor:
|
||||
self.config.timeout_seconds,
|
||||
)
|
||||
|
||||
# Copy the parent context before registering: context copying can
|
||||
# itself fail (callback-manager copy or loop-bound handler filtering),
|
||||
# and a failure after registration would strand a PENDING entry —
|
||||
# the caller never receives an execution_id to poll, and
|
||||
# cleanup_background_task() refuses non-terminal entries.
|
||||
parent_context = _copy_isolated_subagent_context()
|
||||
|
||||
with _background_tasks_lock:
|
||||
_background_tasks[execution_id] = result
|
||||
|
||||
parent_context = _copy_isolated_subagent_context()
|
||||
|
||||
async def run_with_timeout() -> SubagentResult:
|
||||
try:
|
||||
return await asyncio.wait_for(
|
||||
@ -1473,7 +1499,18 @@ class SubagentExecutor:
|
||||
result.try_set_terminal(SubagentStatus.FAILED, error=str(exc))
|
||||
return result
|
||||
|
||||
execution_future = _submit_to_isolated_loop_in_context(parent_context, run_with_timeout)
|
||||
try:
|
||||
execution_future = _submit_to_isolated_loop_in_context(parent_context, run_with_timeout)
|
||||
except Exception:
|
||||
# Submitting can fail before any coroutine starts (e.g. the
|
||||
# persistent loop failed to spin up). The caller then sees the
|
||||
# exception and never polls this execution_id, and
|
||||
# cleanup_background_task() refuses non-terminal entries — so the
|
||||
# just-registered entry must be dropped here, not left as a
|
||||
# PENDING zombie nothing will ever remove.
|
||||
with _background_tasks_lock:
|
||||
_background_tasks.pop(execution_id, None)
|
||||
raise
|
||||
with _background_tasks_lock:
|
||||
_background_futures[execution_id] = execution_future
|
||||
|
||||
|
||||
59
backend/tests/test_subagent_capacity.py
Normal file
59
backend/tests/test_subagent_capacity.py
Normal file
@ -0,0 +1,59 @@
|
||||
"""Cross-thread safety of the shared subagent execution capacity."""
|
||||
|
||||
import asyncio
|
||||
from collections import deque
|
||||
|
||||
import pytest
|
||||
|
||||
from deerflow.config.subagent_runtime_config import SubagentRuntimeConfig
|
||||
from deerflow.subagents.capacity import SubagentExecutionCapacity
|
||||
|
||||
|
||||
class _NoIterationDeque(deque):
|
||||
"""A deque whose iteration always fails, as if mutated concurrently."""
|
||||
|
||||
def __iter__(self):
|
||||
raise RuntimeError("deque mutated during iteration")
|
||||
|
||||
|
||||
def test_snapshot_does_not_iterate_waiters():
|
||||
"""snapshot() must be readable from a non-loop thread.
|
||||
|
||||
``configure_subagent_execution_capacity`` reads a snapshot while the loop
|
||||
thread owns the waiters deque; iterating it cross-thread can raise
|
||||
``deque mutated during iteration``. The snapshot must therefore derive
|
||||
``queued`` without iterating.
|
||||
"""
|
||||
capacity = SubagentExecutionCapacity(SubagentRuntimeConfig(max_running=2, max_queued=4))
|
||||
capacity._waiters = _NoIterationDeque(range(2))
|
||||
|
||||
snapshot = capacity.snapshot()
|
||||
|
||||
assert snapshot.max_running == 2
|
||||
assert snapshot.max_queued == 4
|
||||
assert snapshot.queued == 2
|
||||
assert snapshot.running == 0
|
||||
assert snapshot.admission_policy == "queue"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_snapshot_reports_running_and_queued_waiters():
|
||||
"""The derived queued count still reflects pending waiters."""
|
||||
capacity = SubagentExecutionCapacity(SubagentRuntimeConfig(max_running=1, max_queued=4, queue_timeout_seconds=5))
|
||||
|
||||
async def hold_slot():
|
||||
async with capacity.slot():
|
||||
await asyncio.sleep(60)
|
||||
|
||||
holder = asyncio.create_task(hold_slot())
|
||||
queued = asyncio.create_task(capacity._acquire())
|
||||
try:
|
||||
await asyncio.sleep(0.05)
|
||||
|
||||
snapshot = capacity.snapshot()
|
||||
assert snapshot.running == 1
|
||||
assert snapshot.queued == 1
|
||||
finally:
|
||||
queued.cancel()
|
||||
holder.cancel()
|
||||
await asyncio.gather(queued, holder, return_exceptions=True)
|
||||
@ -16,6 +16,7 @@ the real implementation in isolation.
|
||||
|
||||
import asyncio
|
||||
import importlib
|
||||
import inspect
|
||||
import sys
|
||||
import threading
|
||||
import time
|
||||
@ -2082,6 +2083,122 @@ class TestCleanupBackgroundTask:
|
||||
|
||||
return _patch_default_get_app_config(importlib.reload(executor))
|
||||
|
||||
def test_execute_async_removes_entry_when_submit_fails(self, executor_module, classes, base_config):
|
||||
"""A failed submit must not leave a PENDING entry nothing will ever poll.
|
||||
|
||||
The registry entry is created before the coroutine is submitted to the
|
||||
isolated loop. When submission itself raises (e.g. the loop failed to
|
||||
start), the caller sees the exception and never polls — and
|
||||
``cleanup_background_task`` refuses non-terminal entries — so the fix
|
||||
must drop the entry on the submit-failure path.
|
||||
"""
|
||||
SubagentExecutor = classes["SubagentExecutor"]
|
||||
|
||||
executor = SubagentExecutor(
|
||||
config=base_config,
|
||||
tools=[],
|
||||
thread_id="test-thread",
|
||||
trace_id="submit-failure-trace",
|
||||
)
|
||||
|
||||
def failing_submit(_context, _coro_factory):
|
||||
raise RuntimeError("isolated subagent event loop failed to start")
|
||||
|
||||
with patch.object(executor_module, "_submit_to_isolated_loop_in_context", side_effect=failing_submit):
|
||||
with pytest.raises(RuntimeError, match="isolated subagent event loop"):
|
||||
executor.execute_async("Task")
|
||||
|
||||
leftovers = [r for r in executor_module.list_background_tasks() if r.trace_id == "submit-failure-trace"]
|
||||
assert leftovers == []
|
||||
|
||||
def test_execute_async_registers_nothing_when_context_copy_fails(self, executor_module, classes, base_config):
|
||||
"""A context-copy failure must not leave a PENDING entry either.
|
||||
|
||||
``_copy_isolated_subagent_context`` (callback-manager copy or
|
||||
loop-bound handler filtering) can raise before the coroutine is ever
|
||||
submitted. The registry entry must not exist at that point yet —
|
||||
the caller gets no execution_id to poll, and
|
||||
``cleanup_background_task`` refuses non-terminal entries, so a
|
||||
registration before the copy would strand the same permanent
|
||||
PENDING entry the submit-failure path already guards against.
|
||||
"""
|
||||
SubagentExecutor = classes["SubagentExecutor"]
|
||||
|
||||
executor = SubagentExecutor(
|
||||
config=base_config,
|
||||
tools=[],
|
||||
thread_id="test-thread",
|
||||
trace_id="context-copy-failure-trace",
|
||||
)
|
||||
|
||||
def failing_context_copy():
|
||||
raise RuntimeError("callback manager copy failed")
|
||||
|
||||
with patch.object(executor_module, "_copy_isolated_subagent_context", side_effect=failing_context_copy):
|
||||
with pytest.raises(RuntimeError, match="callback manager copy"):
|
||||
executor.execute_async("Task")
|
||||
|
||||
leftovers = [r for r in executor_module.list_background_tasks() if r.trace_id == "context-copy-failure-trace"]
|
||||
assert leftovers == []
|
||||
|
||||
def test_submit_helper_skips_coroutine_creation_when_loop_startup_fails(self, executor_module):
|
||||
"""Loop-startup failure must not strand an unscheduled coroutine.
|
||||
|
||||
Exercises the real ``_submit_to_isolated_loop_in_context`` (only the
|
||||
loop getter is patched) rather than mocking the whole helper: the
|
||||
loop must be resolved before the coroutine is created. If the
|
||||
coroutine factory ran first, the created coroutine would be neither
|
||||
scheduled nor closed — ``RuntimeWarning: coroutine ... was never
|
||||
awaited`` — retaining its captures until collection.
|
||||
"""
|
||||
factory_calls = []
|
||||
|
||||
def coro_factory():
|
||||
factory_calls.append("created")
|
||||
|
||||
async def never_scheduled(): # pragma: no cover - must not run
|
||||
return None
|
||||
|
||||
return never_scheduled()
|
||||
|
||||
def failing_loop():
|
||||
raise RuntimeError("Timed out starting isolated subagent event loop")
|
||||
|
||||
with patch.object(executor_module, "_get_isolated_subagent_loop", side_effect=failing_loop):
|
||||
with pytest.raises(RuntimeError, match="Timed out starting"):
|
||||
executor_module._submit_to_isolated_loop_in_context(executor_module.copy_context(), coro_factory)
|
||||
|
||||
assert factory_calls == []
|
||||
|
||||
def test_submit_helper_closes_coroutine_when_scheduling_rejects_it(self, executor_module):
|
||||
"""Scheduling rejection after creation must close the coroutine.
|
||||
|
||||
Resolving the loop before calling the factory covers loop-startup
|
||||
failure, but ``run_coroutine_threadsafe`` can itself raise once the
|
||||
coroutine exists (e.g. the loop closes between the lookup and the
|
||||
internal ``call_soon_threadsafe``). Only ``run_coroutine_threadsafe``
|
||||
is patched: the helper must close the rejected coroutine — otherwise
|
||||
it stays in ``CORO_CREATED`` and re-triggers the never-awaited
|
||||
warning and retained captures the startup fix already guards against.
|
||||
"""
|
||||
created = []
|
||||
|
||||
def coro_factory():
|
||||
async def pending():
|
||||
return None # pragma: no cover - must never run
|
||||
|
||||
coroutine = pending()
|
||||
created.append(coroutine)
|
||||
return coroutine
|
||||
|
||||
with patch.object(executor_module, "_get_isolated_subagent_loop", return_value=object()):
|
||||
with patch.object(executor_module.asyncio, "run_coroutine_threadsafe", side_effect=RuntimeError("Event loop is closed")):
|
||||
with pytest.raises(RuntimeError, match="Event loop is closed"):
|
||||
executor_module._submit_to_isolated_loop_in_context(executor_module.copy_context(), coro_factory)
|
||||
|
||||
assert len(created) == 1
|
||||
assert inspect.getcoroutinestate(created[0]) is inspect.CORO_CLOSED
|
||||
|
||||
def test_cleanup_removes_terminal_completed_task(self, executor_module, classes):
|
||||
"""Test that cleanup removes a COMPLETED task."""
|
||||
SubagentResult = classes["SubagentResult"]
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user