From fa3bb438c9f009c7818a370314bc89b71407a70c Mon Sep 17 00:00:00 2001 From: NanPan <111261006+poijygfdyy@users.noreply.github.com> Date: Thu, 8 Oct 2026 09:55:54 +0800 Subject: [PATCH] fix(gateway): bound subagent batch shutdown (#6434) * fix(gateway): bound subagent batch shutdown * test(gateway): cover bounded subagent batch shutdown * fix(subagents): cancel batch work before stop drain * test(subagents): pin cancellation before stop drain * test(gateway): account for subagent shutdown budget * docs(helm): refresh gateway shutdown budget * docs(compose): refresh gateway shutdown budget * docs(compose): correct gateway shutdown total --------- Co-authored-by: Willem Jiang --- backend/app/gateway/app.py | 32 ++++--- .../deerflow/subagents/batch_service.py | 14 ++- backend/tests/_gateway_shutdown_budget.py | 7 +- .../test_subagent_batch_shutdown_timeout.py | 86 +++++++++++++++++++ deploy/helm/deer-flow/README.md | 2 +- docker/docker-compose.yaml | 7 +- 6 files changed, 128 insertions(+), 20 deletions(-) create mode 100644 backend/tests/test_subagent_batch_shutdown_timeout.py diff --git a/backend/app/gateway/app.py b/backend/app/gateway/app.py index 07e5349853..fc98d741e3 100644 --- a/backend/app/gateway/app.py +++ b/backend/app/gateway/app.py @@ -416,6 +416,27 @@ async def _shutdown_scheduled_task_service(app: FastAPI) -> None: logger.exception("Failed to stop scheduled task service") +async def _shutdown_subagent_batch_service(app: FastAPI) -> None: + """Bound durable subagent batch stop so Gateway exit cannot hang forever.""" + service = getattr(app.state, "subagent_batch_service", None) + if service is None: + return + app.state.subagent_batches_available = False + try: + await asyncio.wait_for(service.stop(), timeout=_SHUTDOWN_HOOK_TIMEOUT_SECONDS) + except TimeoutError: + logger.warning( + "Subagent batch service shutdown exceeded %.1fs; proceeding with worker exit.", + _SHUTDOWN_HOOK_TIMEOUT_SECONDS, + ) + except Exception: + logger.exception("Failed to stop subagent batch service") + finally: + from deerflow.subagents.batch_runtime import set_subagent_batch_submitter + + set_subagent_batch_submitter(None) + + def _scheduled_task_notification_repos(startup_config: AppConfig): """Return ``(connection_repo, notification_repo)`` for the scheduled-run outbox (issue #4254). @@ -848,16 +869,7 @@ async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]: set_mcp_task_config_snapshot(None) - if getattr(app.state, "subagent_batch_service", None) is not None: - app.state.subagent_batches_available = False - try: - await app.state.subagent_batch_service.stop() - except Exception: - logger.exception("Failed to stop subagent batch service") - finally: - from deerflow.subagents.batch_runtime import set_subagent_batch_submitter - - set_subagent_batch_submitter(None) + await _shutdown_subagent_batch_service(app) # Browser sessions have their own bounded teardown. MCP sessions close # after the runtime drains runs, since those runs may still call tools. diff --git a/backend/packages/harness/deerflow/subagents/batch_service.py b/backend/packages/harness/deerflow/subagents/batch_service.py index c8550c43f5..d57c9112f5 100644 --- a/backend/packages/harness/deerflow/subagents/batch_service.py +++ b/backend/packages/harness/deerflow/subagents/batch_service.py @@ -77,15 +77,23 @@ class SubagentBatchService: self._stop.set() poller = self._poller self._poller = None - if poller is not None: - poller.cancel() - await asyncio.gather(poller, return_exceptions=True) + + # Issue every owned-work cancellation before the first await. The + # Gateway wraps this stop hook in a deadline; if poller teardown is + # slow, cancellation of this coroutine must not prevent native/item + # cancellation from being requested. execution_ids = list(self._execution_ids.values()) for execution_id in execution_ids: request_cancel_background_task(execution_id) tasks = list(self._executions.values()) + + if poller is not None: + poller.cancel() for task in tasks: task.cancel() + + if poller is not None: + await asyncio.gather(poller, return_exceptions=True) if tasks: await asyncio.gather(*tasks, return_exceptions=True) self._executions.clear() diff --git a/backend/tests/_gateway_shutdown_budget.py b/backend/tests/_gateway_shutdown_budget.py index 9ce3ee41f3..e8d100fa83 100644 --- a/backend/tests/_gateway_shutdown_budget.py +++ b/backend/tests/_gateway_shutdown_budget.py @@ -24,6 +24,7 @@ HOOKS_BOUNDED_BY_SHUTDOWN_HOOK_TIMEOUT: tuple[str, ...] = ( "notification delivery worker stop", "channel service stop", "scheduled task service stop", + "subagent batch service stop", "browser session close", "MCP session pool close", ) @@ -62,8 +63,8 @@ def memory_flush_seconds() -> float: def lifespan_shutdown_seconds() -> float: """Worst-case lifespan shutdown: every bounded hook + retrieval-warm wait + run drain + memory flush. - Unbounded steps (OIDC close, MCP/subagent task-service stops, backend close) - are expected to be quick and are not modeled; deployments keep slack on top - of this figure for them. + Unbounded steps (OIDC close, MCP task-service stop, backend close) are + expected to be quick and are not modeled; deployments keep slack on top of + this figure for them. """ return bounded_hooks_seconds() + retrieval_warm_wait_seconds() + run_drain_seconds() + memory_flush_seconds() diff --git a/backend/tests/test_subagent_batch_shutdown_timeout.py b/backend/tests/test_subagent_batch_shutdown_timeout.py new file mode 100644 index 0000000000..be417fca4c --- /dev/null +++ b/backend/tests/test_subagent_batch_shutdown_timeout.py @@ -0,0 +1,86 @@ +from __future__ import annotations + +import asyncio +import inspect +import logging +from unittest.mock import AsyncMock, MagicMock + +import pytest +from fastapi import FastAPI + +from app.gateway import app as gateway_app +from deerflow.subagents import batch_service as batch_service_module +from deerflow.subagents.batch_service import SubagentBatchService + + +@pytest.mark.asyncio +async def test_subagent_batch_service_shutdown_is_bounded( + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, +) -> None: + async def hang() -> None: + await asyncio.sleep(3600) + + app = FastAPI() + service = MagicMock() + service.stop = AsyncMock(side_effect=hang) + app.state.subagent_batch_service = service + app.state.subagent_batches_available = True + monkeypatch.setattr(gateway_app, "_SHUTDOWN_HOOK_TIMEOUT_SECONDS", 0.01) + + started = asyncio.get_running_loop().time() + with caplog.at_level(logging.WARNING, logger="app.gateway.app"): + await gateway_app._shutdown_subagent_batch_service(app) + elapsed = asyncio.get_running_loop().time() - started + + assert elapsed < 0.5 + assert app.state.subagent_batches_available is False + service.stop.assert_awaited_once() + assert "Subagent batch service shutdown exceeded" in caplog.text + + +def test_lifespan_uses_bounded_subagent_batch_shutdown() -> None: + source = inspect.getsource(gateway_app.lifespan) + assert "await _shutdown_subagent_batch_service(app)" in source + + +@pytest.mark.asyncio +async def test_batch_stop_requests_owned_cancellation_before_poller_drain( + monkeypatch: pytest.MonkeyPatch, +) -> None: + release_poller = asyncio.Event() + poller_cancelled = asyncio.Event() + requested: list[str] = [] + + async def slow_poller() -> None: + try: + await asyncio.sleep(3600) + except asyncio.CancelledError: + poller_cancelled.set() + await release_poller.wait() + + async def item_work() -> None: + await asyncio.sleep(3600) + + service = SubagentBatchService.__new__(SubagentBatchService) + service._stop = asyncio.Event() + service._poller = asyncio.create_task(slow_poller()) + item_task = asyncio.create_task(item_work()) + service._executions = {"item-1": item_task} + service._execution_ids = {"item-1": "execution-1"} + service._item_batches = {"item-1": "batch-1"} + + monkeypatch.setattr( + batch_service_module, + "request_cancel_background_task", + requested.append, + ) + + stop_task = asyncio.create_task(service.stop()) + await poller_cancelled.wait() + + assert requested == ["execution-1"] + assert item_task.cancelled() or item_task.cancelling() + + release_poller.set() + await stop_task diff --git a/deploy/helm/deer-flow/README.md b/deploy/helm/deer-flow/README.md index f8e97f10b3..1695599d48 100644 --- a/deploy/helm/deer-flow/README.md +++ b/deploy/helm/deer-flow/README.md @@ -277,7 +277,7 @@ kubectl -n deer-flow exec deploy/deer-flow-provisioner -- curl -s localhost:8002 `%40`). The chart uses an external `databaseUrl` verbatim and does not rewrite the DSN in a user-managed Secret. -- **Graceful shutdown & memory drain.** The gateway pod sets `terminationGracePeriodSeconds` (default 90s, overridable via `gateway.terminationGracePeriodSeconds`) plus an optional `preStop` sleep (`gateway.preStopSleepSeconds`, default 5s), and bounds uvicorn's `--timeout-graceful-shutdown` (`gateway.uvicornGracefulShutdownSeconds`, default 10s) so an idle SSE connection cannot hold up lifespan shutdown indefinitely. The grace period MUST exceed the Gateway's graceful-shutdown work — the preStop sleep, the uvicorn timeout, and the lifespan's worst case: six hooks bounded at 5s each (startup trash sweep, notification delivery worker, channel service, scheduled task service, browser sessions, MCP session pool), the 1s retrieval-warm wait, the in-flight run drain (5s) and the memory-queue drain (`memory.shutdown_flush_timeout_seconds`, default 30s), about 66s in total, plus a buffer — because the drain runs on a daemon thread and K8s SIGKILLs anything still running at the end of the grace window. K8s defaults to 30s, which SIGKILLs the drain mid-flight and silently re-introduces the memory loss the drain is fixing. `backend/tests/_gateway_shutdown_budget.py` reads these bounds from the Gateway and pins the chart and compose budgets against them. **When you raise `memory.shutdown_flush_timeout_seconds`, raise `gateway.terminationGracePeriodSeconds` to match** (preStop + uvicorn timeout + ~36s of bounded hooks and drains + memory drain + buffer). +- **Graceful shutdown & memory drain.** The gateway pod sets `terminationGracePeriodSeconds` (default 90s, overridable via `gateway.terminationGracePeriodSeconds`) plus an optional `preStop` sleep (`gateway.preStopSleepSeconds`, default 5s), and bounds uvicorn's `--timeout-graceful-shutdown` (`gateway.uvicornGracefulShutdownSeconds`, default 10s) so an idle SSE connection cannot hold up lifespan shutdown indefinitely. The grace period MUST exceed the Gateway's graceful-shutdown work — the preStop sleep, the uvicorn timeout, and the lifespan's worst case: seven hooks bounded at 5s each (startup trash sweep, notification delivery worker, channel service, scheduled task service, subagent batch service, browser sessions, MCP session pool), the 1s retrieval-warm wait, the in-flight run drain (5s) and the memory-queue drain (`memory.shutdown_flush_timeout_seconds`, default 30s), about 71s in total, plus a buffer — because the drain runs on a daemon thread and K8s SIGKILLs anything still running at the end of the grace window. K8s defaults to 30s, which SIGKILLs the drain mid-flight and silently re-introduces the memory loss the drain is fixing. `backend/tests/_gateway_shutdown_budget.py` reads these bounds from the Gateway and pins the chart and compose budgets against them. **When you raise `memory.shutdown_flush_timeout_seconds`, raise `gateway.terminationGracePeriodSeconds` to match** (preStop + uvicorn timeout + ~41s of bounded hooks and drains + memory drain + buffer). - **Gateway replicas.** Run control is cross-pod-safe since the work tracked by [issue #3948](https://github.com/bytedance/deer-flow/issues/3948) landed (#4003, #4064, #4500): admission is a durable one-active-run-per-thread diff --git a/docker/docker-compose.yaml b/docker/docker-compose.yaml index f4367d28af..6a0d4e723f 100644 --- a/docker/docker-compose.yaml +++ b/docker/docker-compose.yaml @@ -124,11 +124,12 @@ services: # QR login and browser tools stay worker-local, so keep one worker while # they are enabled. --timeout-graceful-shutdown bounds how long open SSE # connections may delay lifespan shutdown after SIGTERM; stop_grace_period - # below must cover that bound plus the lifespan's worst case: five hooks + # below must cover that bound plus the lifespan's worst case: seven hooks # bounded at 5s each (trash sweep, notification worker, channel service, - # browser sessions, MCP session pool), the 1s retrieval-warm wait, the 5s + # scheduled task service, subagent batch service, browser sessions, + # MCP session pool), the 1s retrieval-warm wait, the 5s # run drain and the 30s memory flush (memory.shutdown_flush_timeout_seconds) - # -- about 61s -- plus slack. backend/tests/_gateway_shutdown_budget.py + # -- about 71s -- plus slack. backend/tests/_gateway_shutdown_budget.py # pins the arithmetic. command: sh -c "cd backend && PYTHONPATH=. uv run --no-sync uvicorn app.gateway.app:app --host 0.0.0.0 --port 8001 --workers ${GATEWAY_WORKERS:-1} --timeout-graceful-shutdown 10" # Docker's 10s default SIGKILLed the 30s memory queue flush