mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-10-26 13:16:24 +00:00
a684ae2fe3
* fix(sandbox): leave SIGTERM/SIGINT teardown to the Gateway lifespan AIO and E2B providers built on the event-loop thread installed SIGTERM/SIGINT handlers that force-destroyed every sandbox synchronously, before uvicorn started the run drain. The Gateway lifespan now holds sandbox shutdown ownership from its first step, so providers skip those handlers and the existing hook closes the provider after run drain and MCP teardown. SIGHUP keeps the provider fallback. AIO acquire now refuses once shutdown began. * docs(sandbox): document Gateway-owned sandbox shutdown and AIO admission fence * test(gateway): inspect the lifespan body now that lifespan wraps it * docs(changelog): reference #7560 for the sandbox SIGTERM drain-order fix * ci: retrigger sandbox image smoke after a registry pull timeout
313 lines
13 KiB
Python
313 lines
13 KiB
Python
"""Gateway owns sandbox cleanup even when lazy creation cannot bind signals."""
|
|
|
|
import asyncio
|
|
import importlib
|
|
import signal
|
|
import threading
|
|
import time
|
|
from contextlib import asynccontextmanager
|
|
from types import SimpleNamespace
|
|
from unittest.mock import AsyncMock, MagicMock
|
|
|
|
import pytest
|
|
from fastapi import FastAPI
|
|
from langchain_core.messages import ToolMessage
|
|
from langgraph.prebuilt.tool_node import ToolCallRequest
|
|
|
|
|
|
def _runtime(monkeypatch, events):
|
|
@asynccontextmanager
|
|
async def runtime(_app, _config):
|
|
try:
|
|
yield
|
|
finally:
|
|
events.append("runtime_drained")
|
|
|
|
monkeypatch.setattr("app.gateway.app.langgraph_runtime", runtime)
|
|
pool = SimpleNamespace(close_all=AsyncMock(side_effect=lambda: events.append("mcp_closed")))
|
|
monkeypatch.setattr("deerflow.mcp.session_pool.get_session_pool", lambda: pool)
|
|
return pool
|
|
|
|
|
|
def test_shutdown_closes_provider_created_by_lazy_first_write_after_run_drain(monkeypatch):
|
|
from app.gateway.app import _runtime_with_mcp_pool_shutdown
|
|
from deerflow.agents.middlewares.read_before_write_middleware import ReadBeforeWriteMiddleware
|
|
from deerflow.runtime import RunManager, RunStatus
|
|
from deerflow.sandbox import sandbox_provider as registry
|
|
|
|
async def scenario():
|
|
events = []
|
|
loop_thread = threading.get_ident()
|
|
|
|
class Provider:
|
|
def __init__(self):
|
|
self.created_thread = threading.get_ident()
|
|
self.closed = False
|
|
|
|
def shutdown(self):
|
|
assert threading.get_ident() != loop_thread
|
|
events.append("sandbox_closed")
|
|
self.closed = True
|
|
|
|
monkeypatch.setattr(registry, "_default_sandbox_provider", None)
|
|
monkeypatch.setattr(registry, "get_app_config", lambda: SimpleNamespace(sandbox=SimpleNamespace(use="fixture:Provider")))
|
|
monkeypatch.setattr(registry, "resolve_class", lambda *_args: Provider)
|
|
pool = _runtime(monkeypatch, events)
|
|
manager = RunManager()
|
|
record = await manager.create("lazy-first-write")
|
|
await manager.set_status(record.run_id, RunStatus.running)
|
|
resume = asyncio.Event()
|
|
first_write_finished = asyncio.Event()
|
|
created = []
|
|
|
|
def read_new_file(_runtime, _path):
|
|
created.append(registry.get_sandbox_provider())
|
|
raise FileNotFoundError("new output")
|
|
|
|
middleware = ReadBeforeWriteMiddleware(content_reader=read_new_file)
|
|
state = {"messages": []}
|
|
runtime = SimpleNamespace(state=state, context={"thread_id": "lazy-first-write"})
|
|
call = ToolCallRequest(tool_call={"name": "write_file", "args": {"path": "/mnt/user-data/outputs/new.txt", "content": "new"}, "id": "write"}, tool=None, state=state, runtime=runtime)
|
|
|
|
async def write(_request):
|
|
return ToolMessage(content="OK", tool_call_id="write", name="write_file")
|
|
|
|
async def active_run():
|
|
await resume.wait()
|
|
await middleware.awrap_tool_call(call, write)
|
|
first_write_finished.set()
|
|
await asyncio.Event().wait()
|
|
|
|
@asynccontextmanager
|
|
async def draining_runtime(_app, _config):
|
|
record.task = asyncio.create_task(active_run())
|
|
try:
|
|
yield
|
|
finally:
|
|
resume.set()
|
|
await asyncio.wait_for(first_write_finished.wait(), timeout=1)
|
|
await manager.shutdown(timeout=1.0)
|
|
events.append("runtime_drained")
|
|
|
|
monkeypatch.setattr("app.gateway.app.langgraph_runtime", draining_runtime)
|
|
try:
|
|
async with _runtime_with_mcp_pool_shutdown(FastAPI(), SimpleNamespace()):
|
|
assert registry.get_initialized_sandbox_provider() is None
|
|
assert record.task.done()
|
|
assert len(created) == 1 and created[0].created_thread != loop_thread
|
|
assert created[0].closed
|
|
assert registry.get_initialized_sandbox_provider() is None
|
|
assert events == ["runtime_drained", "mcp_closed", "sandbox_closed"]
|
|
pool.close_all.assert_awaited_once()
|
|
finally:
|
|
resume.set()
|
|
if record.task is not None and not record.task.done():
|
|
record.task.cancel()
|
|
await asyncio.gather(record.task, return_exceptions=True)
|
|
|
|
asyncio.run(scenario())
|
|
|
|
|
|
def test_shutdown_without_provider_does_not_allocate_one(monkeypatch):
|
|
from app.gateway.app import _runtime_with_mcp_pool_shutdown
|
|
from deerflow.sandbox import sandbox_provider as registry
|
|
|
|
async def scenario():
|
|
events = []
|
|
monkeypatch.setattr(registry, "_default_sandbox_provider", None)
|
|
monkeypatch.setattr(registry, "resolve_class", lambda *_args: pytest.fail("shutdown allocated a sandbox provider"))
|
|
_runtime(monkeypatch, events)
|
|
async with _runtime_with_mcp_pool_shutdown(FastAPI(), SimpleNamespace()):
|
|
assert registry.get_initialized_sandbox_provider() is None
|
|
assert registry.get_initialized_sandbox_provider() is None
|
|
assert events == ["runtime_drained", "mcp_closed"]
|
|
|
|
asyncio.run(scenario())
|
|
|
|
|
|
def test_sandbox_shutdown_failure_is_logged_and_remains_retryable(monkeypatch, caplog):
|
|
from app.gateway.app import _runtime_with_mcp_pool_shutdown
|
|
from deerflow.sandbox import sandbox_provider as registry
|
|
|
|
async def scenario():
|
|
events = []
|
|
pool = _runtime(monkeypatch, events)
|
|
|
|
def fail_close():
|
|
raise RuntimeError("sandbox close failed")
|
|
|
|
provider = SimpleNamespace(shutdown=fail_close)
|
|
monkeypatch.setattr(registry, "_default_sandbox_provider", provider)
|
|
async with _runtime_with_mcp_pool_shutdown(FastAPI(), SimpleNamespace()):
|
|
pass
|
|
assert registry.get_initialized_sandbox_provider() is provider
|
|
assert "Failed to shut down sandbox provider" in caplog.text
|
|
pool.close_all.assert_awaited_once()
|
|
|
|
asyncio.run(scenario())
|
|
|
|
|
|
@pytest.mark.parametrize("cancel_during", ["mcp", "sandbox"])
|
|
def test_shutdown_drains_sandbox_worker_across_cancellation(monkeypatch, cancel_during):
|
|
from app.gateway.app import _runtime_with_mcp_pool_shutdown
|
|
from deerflow.sandbox import sandbox_provider as registry
|
|
|
|
async def scenario():
|
|
events = []
|
|
_runtime(monkeypatch, events)
|
|
pool_started = asyncio.Event()
|
|
worker_started = threading.Event()
|
|
release_worker = threading.Event()
|
|
worker_finished = threading.Event()
|
|
loop_thread = threading.get_ident()
|
|
|
|
async def close_pool():
|
|
events.append("mcp_close_started")
|
|
pool_started.set()
|
|
if cancel_during == "mcp":
|
|
await asyncio.Event().wait()
|
|
|
|
def close_sandbox():
|
|
assert threading.get_ident() != loop_thread
|
|
worker_started.set()
|
|
assert release_worker.wait(5.0)
|
|
worker_finished.set()
|
|
|
|
monkeypatch.setattr("deerflow.mcp.session_pool.get_session_pool", lambda: SimpleNamespace(close_all=close_pool))
|
|
monkeypatch.setattr(registry, "_default_sandbox_provider", SimpleNamespace(shutdown=close_sandbox))
|
|
context = _runtime_with_mcp_pool_shutdown(FastAPI(), SimpleNamespace())
|
|
await context.__aenter__()
|
|
task = asyncio.create_task(context.__aexit__(None, None, None))
|
|
try:
|
|
await asyncio.wait_for(pool_started.wait(), timeout=1)
|
|
if cancel_during == "mcp":
|
|
task.cancel()
|
|
assert await asyncio.to_thread(worker_started.wait, 1.0)
|
|
task.cancel()
|
|
await asyncio.sleep(0)
|
|
task.cancel()
|
|
await asyncio.sleep(0)
|
|
assert not task.done(), "shutdown abandoned the sandbox worker"
|
|
release_worker.set()
|
|
with pytest.raises(asyncio.CancelledError):
|
|
await task
|
|
assert worker_finished.is_set()
|
|
assert registry.get_initialized_sandbox_provider() is None
|
|
finally:
|
|
release_worker.set()
|
|
await asyncio.gather(task, return_exceptions=True)
|
|
|
|
asyncio.run(scenario())
|
|
|
|
|
|
@pytest.fixture
|
|
def restore_signal_handlers():
|
|
signums = [signal.SIGTERM, signal.SIGINT] + ([signal.SIGHUP] if hasattr(signal, "SIGHUP") else [])
|
|
saved = {signum: signal.getsignal(signum) for signum in signums}
|
|
try:
|
|
yield
|
|
finally:
|
|
for signum, handler in saved.items():
|
|
signal.signal(signum, handler)
|
|
|
|
|
|
def _build_aio_provider(monkeypatch, backend, sandbox_ids):
|
|
"""Run the real AIO constructor (signal registration included) over a fake backend."""
|
|
from deerflow.community.aio_sandbox.sandbox_info import SandboxInfo
|
|
|
|
aio_mod = importlib.import_module("deerflow.community.aio_sandbox.aio_sandbox_provider")
|
|
config = {"replicas": 10, "idle_timeout": 0, "ownership": None, "stream_bridge": None, "command_timeout": 600.0}
|
|
monkeypatch.setattr(aio_mod.AioSandboxProvider, "_load_config", lambda _self: config)
|
|
monkeypatch.setattr(aio_mod.AioSandboxProvider, "_create_backend", lambda _self: backend)
|
|
monkeypatch.setattr(aio_mod.AioSandboxProvider, "_reconcile_orphans", lambda _self: None)
|
|
provider = aio_mod.AioSandboxProvider()
|
|
for sandbox_id in sandbox_ids:
|
|
sandbox = MagicMock(id=sandbox_id, requires_container_recycle=False, has_pending_session_creates=False)
|
|
with provider._lock:
|
|
provider._sandbox_infos[sandbox_id] = SandboxInfo(sandbox_id=sandbox_id, sandbox_url=f"http://fake/{sandbox_id}")
|
|
provider._sandboxes[sandbox_id] = sandbox
|
|
provider._last_activity[sandbox_id] = time.monotonic()
|
|
return provider
|
|
|
|
|
|
@pytest.mark.parametrize("signum", [signal.SIGTERM, signal.SIGINT], ids=["SIGTERM", "SIGINT"])
|
|
def test_stop_signal_with_main_thread_provider_destroys_sandboxes_only_after_run_drain(monkeypatch, restore_signal_handlers, signum):
|
|
"""G6/S0: a provider built on the event-loop thread must not tear down from its signal handler."""
|
|
from app.gateway.app import _runtime_with_mcp_pool_shutdown
|
|
from deerflow.sandbox import sandbox_provider as registry
|
|
|
|
sandbox_ids = ["sbx-0", "sbx-1", "sbx-2"]
|
|
events = []
|
|
|
|
class Backend:
|
|
def destroy(self, info):
|
|
events.append(f"destroy:{info.sandbox_id}")
|
|
|
|
def list_running(self):
|
|
return []
|
|
|
|
def uvicorn_handle_exit(received, _frame):
|
|
events.append(f"handle_exit:{signal.Signals(received).name}")
|
|
|
|
# uvicorn installs its handlers before lifespan startup builds anything.
|
|
signal.signal(signal.SIGTERM, uvicorn_handle_exit)
|
|
signal.signal(signal.SIGINT, uvicorn_handle_exit)
|
|
monkeypatch.setattr(registry, "_default_sandbox_provider", None)
|
|
pool = _runtime(monkeypatch, events)
|
|
|
|
async def scenario():
|
|
provider = None
|
|
with registry.hold_sandbox_shutdown_ownership():
|
|
try:
|
|
provider = _build_aio_provider(monkeypatch, Backend(), sandbox_ids)
|
|
assert threading.current_thread() is threading.main_thread()
|
|
registry.set_sandbox_provider(provider)
|
|
|
|
@asynccontextmanager
|
|
async def draining_runtime(_app, _config):
|
|
try:
|
|
yield
|
|
finally:
|
|
# The in-flight run still reaches its sandbox while it drains.
|
|
events.append("runtime_drained" if provider.get("sbx-0") is not None else "runtime_drained_without_sandbox")
|
|
|
|
monkeypatch.setattr("app.gateway.app.langgraph_runtime", draining_runtime)
|
|
async with _runtime_with_mcp_pool_shutdown(FastAPI(), SimpleNamespace()):
|
|
signal.raise_signal(signum)
|
|
await asyncio.sleep(0)
|
|
assert events == [f"handle_exit:{signal.Signals(signum).name}"]
|
|
finally:
|
|
if provider is not None:
|
|
provider.shutdown()
|
|
|
|
assert events[:3] == [f"handle_exit:{signal.Signals(signum).name}", "runtime_drained", "mcp_closed"]
|
|
assert sorted(events[3:]) == [f"destroy:{sandbox_id}" for sandbox_id in sandbox_ids]
|
|
assert registry.get_initialized_sandbox_provider() is None
|
|
pool.close_all.assert_awaited_once()
|
|
|
|
asyncio.run(scenario())
|
|
|
|
|
|
def test_gateway_lifespan_owns_sandbox_shutdown_from_its_first_step(monkeypatch):
|
|
"""The declaration must precede every startup step that can build a provider."""
|
|
from app.gateway import app as gateway_app
|
|
from deerflow.sandbox import sandbox_provider as registry
|
|
|
|
observed = []
|
|
|
|
def failing_config():
|
|
observed.append(registry.host_owns_sandbox_shutdown())
|
|
raise RuntimeError("stop startup here")
|
|
|
|
monkeypatch.setattr(gateway_app, "get_app_config", failing_config)
|
|
|
|
async def scenario():
|
|
assert registry.host_owns_sandbox_shutdown() is False
|
|
with pytest.raises(RuntimeError, match="stop startup here"):
|
|
async with gateway_app.lifespan(FastAPI()):
|
|
pytest.fail("startup should have failed")
|
|
assert observed == [True]
|
|
assert registry.host_owns_sandbox_shutdown() is False
|
|
|
|
asyncio.run(scenario())
|