875 lines
33 KiB
Python
875 lines
33 KiB
Python
"""Managed-worker failure, affinity, admission, and 70-client regression tests."""
|
|
|
|
import asyncio
|
|
from types import SimpleNamespace
|
|
|
|
import httpx
|
|
import pytest
|
|
from fastapi import HTTPException
|
|
|
|
from app.gateway.worker_pool.admission import Admission, bounded_worker_run
|
|
from app.gateway.worker_pool.controller import Controller, PoolUnavailable
|
|
from app.gateway.worker_pool.docker import DockerError, worker_spec
|
|
from app.gateway.worker_pool.settings import PoolSettings, SettingsStore
|
|
|
|
|
|
def source():
|
|
return {
|
|
"Image": "sha256:fixed",
|
|
"Config": {"Env": ["AUTH_JWT_SECRET=fixed", "WEB_CONCURRENCY=8"], "WorkingDir": "/app"},
|
|
"NetworkSettings": {"Networks": {"deerflow": {}}},
|
|
"Mounts": [{"Type": "volume", "Name": "runtime", "Destination": "/data", "RW": True}, {"Type": "bind", "Source": "/var/run/docker.sock", "Destination": "/var/run/docker.sock", "RW": True}],
|
|
}
|
|
|
|
|
|
def test_secrets_encrypted_and_preserved(tmp_path):
|
|
store = SettingsStore(tmp_path)
|
|
store.save(PoolSettings(password="secret-sentinel"))
|
|
assert "secret-sentinel" not in store.path.read_text()
|
|
assert "password" not in store.public()
|
|
assert store.public()["password_set"]
|
|
store.save(PoolSettings(desired=2))
|
|
assert SettingsStore(tmp_path).load().password == "secret-sentinel"
|
|
assert store.load().desired == 2
|
|
|
|
|
|
def test_worker_template_has_limits_no_scheduler_no_socket():
|
|
spec = worker_spec(source(), PoolSettings(network="deerflow"), "instance", "token", "/private")
|
|
assert spec["Image"] == "sha256:fixed"
|
|
assert spec["HostConfig"]["Memory"] == 2048 * 1024**2
|
|
assert len(spec["HostConfig"]["Mounts"]) == 1
|
|
assert "SCHEDULED_TASKS_ENABLED=false" in spec["Env"]
|
|
assert not any(e.startswith("WEB_CONCURRENCY=") for e in spec["Env"])
|
|
assert "PortBindings" not in spec["HostConfig"]
|
|
assert spec["Cmd"][-2:] == ["--workers", "1"]
|
|
|
|
|
|
def test_worker_template_uses_snapshot_of_main_runtime():
|
|
spec = worker_spec(
|
|
source(), PoolSettings(network="deerflow"), "instance", "token", "/private",
|
|
runtime_image="sha256:runtime-snapshot",
|
|
)
|
|
assert spec["Image"] == "sha256:runtime-snapshot"
|
|
|
|
|
|
def test_postgres_worker_requires_shared_persistent_chat_archive_directory():
|
|
template = source()
|
|
template["Config"]["Env"].append("DEER_FLOW_HOME=/data")
|
|
spec = worker_spec(
|
|
template, PoolSettings(network="deerflow"), "instance", "token", "/private",
|
|
require_shared_archive=True,
|
|
)
|
|
assert any(mount["Target"] == "/data" and mount["ReadOnly"] is False for mount in spec["HostConfig"]["Mounts"])
|
|
|
|
template["Mounts"] = [
|
|
{"Type": "volume", "Name": "other", "Destination": "/other", "RW": True},
|
|
]
|
|
with pytest.raises(DockerError, match="共享的可写持久卷"):
|
|
worker_spec(
|
|
template, PoolSettings(network="deerflow"), "instance", "token", "/private",
|
|
require_shared_archive=True,
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_live_storage_backend_change_reprobes_and_snapshots_main_runtime(tmp_path, monkeypatch):
|
|
from app.gateway.worker_pool import controller as controller_module
|
|
|
|
store = SettingsStore(tmp_path)
|
|
store.save(PoolSettings(
|
|
enabled=True, source_container="main", network="deerflow", desired=1,
|
|
))
|
|
pool = Controller(store)
|
|
storage_backends = iter(("sqlite", "postgres"))
|
|
discovered = []
|
|
snapshots = []
|
|
template = source()
|
|
template.update({
|
|
"Id": "main-container-id",
|
|
"Image": "sha256:base-image",
|
|
"Config": {"Env": ["AUTH_JWT_SECRET=fixed", "DEER_FLOW_HOME=/data"], "Labels": {}, "User": ""},
|
|
"NetworkSettings": {"Networks": {"deerflow": {}}},
|
|
})
|
|
|
|
class Engine:
|
|
def request(self, method, path, body=None):
|
|
assert method == "GET" and path == "/containers/main/json?size=true"
|
|
return template
|
|
|
|
def configured_storage_backend(self, source_detail, preferred, working_dir):
|
|
return next(storage_backends)
|
|
|
|
def discover_worker_python(self, source_detail, preferred, working_dir):
|
|
discovered.append(working_dir)
|
|
return "/opt/venv/bin/python"
|
|
|
|
def snapshot_worker_runtime(self, source_detail, instance):
|
|
image = f"sha256:runtime-{len(snapshots) + 1}"
|
|
snapshots.append((source_detail["Id"], instance))
|
|
return image
|
|
|
|
engine = Engine()
|
|
monkeypatch.setattr(controller_module, "DockerEngine", lambda settings: engine)
|
|
try:
|
|
_, sqlite_spec = await pool.inspect_template(store.load())
|
|
_, postgres_spec = await pool.inspect_template(store.load())
|
|
assert sqlite_spec["Image"] == "sha256:runtime-1"
|
|
assert postgres_spec["Image"] == "sha256:runtime-2"
|
|
assert len(discovered) == 2
|
|
assert len(snapshots) == 2
|
|
finally:
|
|
await pool.close()
|
|
|
|
|
|
def test_worker_keeps_service_networks_without_main_aliases_and_uses_private_health():
|
|
template = source()
|
|
template["NetworkSettings"]["Networks"]["knowledge"] = {"Aliases": ["main-only"], "IPAddress": "172.20.0.3"}
|
|
spec = worker_spec(template, PoolSettings(network="deerflow"), "instance", "token", "/private")
|
|
assert spec["NetworkingConfig"]["EndpointsConfig"] == {"deerflow": {}, "knowledge": {}}
|
|
health = spec["Healthcheck"]["Test"]
|
|
assert health[:3] == ["CMD", "/opt/venv/bin/python", "-c"]
|
|
assert "/internal/worker-status" in health[3]
|
|
assert "os.environ['DEER_FLOW_WORKER_TOKEN']" in health[3]
|
|
|
|
|
|
def test_dedicated_private_volume_is_excluded():
|
|
template = source()
|
|
template["Mounts"].append({"Type": "volume", "Name": "secrets", "Destination": "/private", "RW": True})
|
|
spec = worker_spec(template, PoolSettings(network="deerflow"), "instance", "token", "/private")
|
|
assert all(mount["Target"] != "/private" for mount in spec["HostConfig"]["Mounts"])
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_orphaned_run_becomes_terminal_without_retry():
|
|
from app.gateway.worker_pool.manager import WorkerRunManager
|
|
from deerflow.runtime.runs.store.memory import MemoryRunStore
|
|
|
|
store = MemoryRunStore()
|
|
await store.put("old", thread_id="thread", status="running")
|
|
manager = WorkerRunManager(store=store)
|
|
assert await manager.get_persisted("old", authorized_thread_id="different-thread") is None
|
|
assert (await store.get("old"))["status"] == "running"
|
|
record = await manager.get_persisted("old", authorized_thread_id="thread")
|
|
assert record.status == "error"
|
|
assert (await store.get("old"))["status"] == "error"
|
|
assert record.task is None
|
|
|
|
|
|
def test_reject_shared_controller_secrets():
|
|
with pytest.raises(DockerError, match="密钥目录"):
|
|
worker_spec(source(), PoolSettings(network="deerflow"), "i", "t", "/data/control")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_70_simultaneous_runs_are_bounded_and_complete():
|
|
gate = Admission(4, 128, 10)
|
|
maximum = 0
|
|
completed = set()
|
|
|
|
async def run(i):
|
|
nonlocal maximum
|
|
await gate.acquire()
|
|
try:
|
|
maximum = max(maximum, gate.active)
|
|
await asyncio.sleep(0.003)
|
|
assert i not in completed
|
|
completed.add(i)
|
|
finally:
|
|
gate.release()
|
|
|
|
await asyncio.gather(*(run(i) for i in range(70)))
|
|
assert len(completed) == 70
|
|
assert maximum == 4
|
|
assert gate.active == gate.waiting == 0
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cancelled_waiter_and_timeout_do_not_leak_capacity():
|
|
gate = Admission(1, 1, 0.03)
|
|
await gate.acquire()
|
|
waiter = asyncio.create_task(gate.acquire())
|
|
await asyncio.sleep(0)
|
|
with pytest.raises(HTTPException) as error:
|
|
await gate.acquire()
|
|
assert error.value.status_code == 429
|
|
waiter.cancel()
|
|
await asyncio.gather(waiter, return_exceptions=True)
|
|
with pytest.raises(HTTPException):
|
|
await gate.acquire()
|
|
gate.release()
|
|
await gate.acquire()
|
|
gate.release()
|
|
assert gate.active == gate.waiting == 0
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_background_run_holds_slot_after_http_returns():
|
|
gate = Admission(1)
|
|
event = asyncio.Event()
|
|
request = SimpleNamespace(app=SimpleNamespace(state=SimpleNamespace(worker_admission=gate)))
|
|
|
|
@bounded_worker_run
|
|
async def start(body, thread, request):
|
|
return SimpleNamespace(task=asyncio.create_task(event.wait()))
|
|
|
|
record = await start(None, "thread", request)
|
|
assert gate.active == 1
|
|
event.set()
|
|
await record.task
|
|
await asyncio.sleep(0)
|
|
assert gate.active == 0
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_affinity_stable_on_scaleup_fail_closed_and_persisted(tmp_path):
|
|
store = SettingsStore(tmp_path)
|
|
store.save(PoolSettings(enabled=True))
|
|
pool = Controller(store)
|
|
pool.nodes["a"] = {"name": "a", "ready": True, "active": 0, "draining": False}
|
|
try:
|
|
first = await pool.reserve("thread", create=True)
|
|
pool.release("a")
|
|
pool.nodes["b"] = {"name": "b", "ready": True, "active": 0, "draining": False}
|
|
assert (await pool.reserve("thread", create=True))["name"] == first["name"]
|
|
pool.release("a")
|
|
pool.nodes["a"]["ready"] = False
|
|
with pytest.raises(PoolUnavailable):
|
|
await pool.reserve("thread", create=True)
|
|
assert pool._ensure_db().execute("SELECT node FROM affinity WHERE thread='thread'").fetchone()[0] == "a"
|
|
finally:
|
|
await pool.close()
|
|
restarted = Controller(store)
|
|
try:
|
|
with pytest.raises(PoolUnavailable):
|
|
await restarted.reserve("thread", create=True)
|
|
finally:
|
|
await restarted.close()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_scale_down_waits_for_active_runs_and_requests(tmp_path, monkeypatch):
|
|
store = SettingsStore(tmp_path)
|
|
store.save(PoolSettings(enabled=True, source_container="main", network="deerflow", desired=0))
|
|
pool = Controller(store)
|
|
name = f"dfw-{store.raw()['instance']}-0"
|
|
calls = []
|
|
|
|
class Engine:
|
|
def request(self, method, path, body=None):
|
|
calls.append((method, path))
|
|
if path.startswith("/containers/json"):
|
|
return [{"Names": ["/" + name], "Id": "id", "State": "running"}]
|
|
return {}
|
|
|
|
async def template(settings):
|
|
return Engine(), {}
|
|
|
|
monkeypatch.setattr(pool, "inspect_template", template)
|
|
active = 1
|
|
|
|
async def status(request):
|
|
return httpx.Response(200, json={"ready": True, "active": active, "capacity": 4})
|
|
|
|
await pool.http.aclose()
|
|
pool.http = httpx.AsyncClient(transport=httpx.MockTransport(status))
|
|
try:
|
|
await pool.reconcile()
|
|
assert not any(method == "POST" for method, _ in calls)
|
|
active = 0
|
|
pool.inflight[name] = 1
|
|
await pool.reconcile()
|
|
assert not any(method == "POST" for method, _ in calls)
|
|
pool.release(name)
|
|
await pool.reconcile()
|
|
assert ("POST", "/containers/id/stop?t=30") in calls
|
|
assert ("DELETE", "/containers/id") in calls
|
|
assert not pool.nodes
|
|
finally:
|
|
await pool.close()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_changed_runtime_image_drains_and_recreates_old_workers(tmp_path, monkeypatch):
|
|
store = SettingsStore(tmp_path)
|
|
store.save(PoolSettings(
|
|
enabled=True, source_container="main", network="deerflow", desired=1,
|
|
))
|
|
pool = Controller(store)
|
|
name = f"dfw-{store.raw()['instance']}-0"
|
|
calls = []
|
|
|
|
class Engine:
|
|
def request(self, method, path, body=None):
|
|
calls.append((method, path, body))
|
|
if path.startswith("/containers/json"):
|
|
return [{
|
|
"Names": ["/" + name], "Id": "old-worker", "State": "running",
|
|
"ImageID": "sha256:old-runtime",
|
|
}]
|
|
if path.startswith("/containers/create?"):
|
|
assert body["Image"] == "sha256:new-runtime"
|
|
return {"Id": "new-worker"}
|
|
return {}
|
|
|
|
async def template(_settings):
|
|
return Engine(), {"Image": "sha256:new-runtime", "Entrypoint": ["python"], "WorkingDir": "/app"}
|
|
|
|
async def status(_request):
|
|
return httpx.Response(200, json={"ready": True, "active": 0, "capacity": 4})
|
|
|
|
monkeypatch.setattr(pool, "inspect_template", template)
|
|
await pool.http.aclose()
|
|
pool.http = httpx.AsyncClient(transport=httpx.MockTransport(status))
|
|
try:
|
|
await pool.reconcile()
|
|
paths = [(method, path) for method, path, _ in calls]
|
|
assert ("POST", "/containers/old-worker/stop?t=30") in paths
|
|
assert ("DELETE", "/containers/old-worker") in paths
|
|
create_index = next(i for i, (_, path) in enumerate(paths) if path.startswith("/containers/create?"))
|
|
assert paths.index(("DELETE", "/containers/old-worker")) < create_index
|
|
assert ("POST", "/containers/new-worker/start") in paths
|
|
assert pool.nodes[name]["id"] == "new-worker"
|
|
finally:
|
|
await pool.close()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_70_http_clients_distribute_to_four_workers_with_sse(tmp_path, monkeypatch):
|
|
from fastapi import FastAPI
|
|
|
|
from app.gateway import deps
|
|
from app.gateway.worker_pool.proxy import WorkerProxy
|
|
|
|
store = SettingsStore(tmp_path)
|
|
store.save(PoolSettings(enabled=True, concurrency=32))
|
|
pool = Controller(store)
|
|
for i in range(4):
|
|
name = f"worker-{i}"
|
|
pool.nodes[name] = {"name": name, "ready": True, "active": 0, "capacity": 32, "draining": False, "url": f"http://{name}:8001"}
|
|
|
|
async def user(request):
|
|
return SimpleNamespace(id="user")
|
|
|
|
async def access(*args, **kwargs):
|
|
return True
|
|
|
|
async def metadata(*args, **kwargs):
|
|
return {"metadata": {}}
|
|
|
|
monkeypatch.setattr(deps, "get_current_user_from_request", user)
|
|
monkeypatch.setattr(deps, "get_thread_store", lambda request: SimpleNamespace(check_access=access, get=metadata))
|
|
gates = {name: Admission(32) for name in pool.nodes}
|
|
counts = {name: 0 for name in pool.nodes}
|
|
|
|
class EventStream(httpx.AsyncByteStream):
|
|
async def __aiter__(self):
|
|
yield b"event: end\ndata: {}\n\n"
|
|
|
|
async def worker(request):
|
|
assert request.headers["cookie"] == "session=preserved"
|
|
assert request.headers["host"] == "gateway"
|
|
assert request.headers["x-deerflow-worker-token"] == store.raw()["token"]
|
|
name = request.url.host
|
|
await gates[name].acquire()
|
|
try:
|
|
await asyncio.sleep(0.005)
|
|
counts[name] += 1
|
|
return httpx.Response(200, stream=EventStream(), headers={"content-type": "text/event-stream", "location": "/threads/t/runs/r/stream"})
|
|
finally:
|
|
gates[name].release()
|
|
|
|
await pool.http.aclose()
|
|
pool.http = httpx.AsyncClient(transport=httpx.MockTransport(worker))
|
|
app = FastAPI()
|
|
app.state.worker_pool = pool
|
|
app.add_middleware(WorkerProxy)
|
|
try:
|
|
async with httpx.AsyncClient(transport=httpx.ASGITransport(app), base_url="http://gateway") as client:
|
|
responses = await asyncio.gather(*(client.post(f"/api/threads/t-{i}/runs/stream", headers={"cookie": "session=preserved"}, json={"input": {}}) for i in range(70)))
|
|
assert all(response.status_code == 200 and "event: end" in response.text for response in responses)
|
|
assert all(response.headers["location"] == "/threads/t/runs/r/stream" for response in responses)
|
|
assert sorted(counts.values()) == [17, 17, 18, 18]
|
|
assert all(value == 0 for value in pool.inflight.values())
|
|
finally:
|
|
await pool.close()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_70_concurrent_runs_fall_back_to_main_without_worker_queue(tmp_path, monkeypatch):
|
|
from fastapi import FastAPI, Request
|
|
from starlette.responses import StreamingResponse
|
|
|
|
from app.gateway import deps
|
|
from app.gateway.worker_pool.proxy import WorkerProxy
|
|
|
|
store = SettingsStore(tmp_path)
|
|
store.save(PoolSettings(enabled=True, concurrency=4))
|
|
pool = Controller(store)
|
|
pool.nodes["worker-0"] = {
|
|
"name": "worker-0", "ready": True, "active": 0, "capacity": 4,
|
|
"draining": False, "url": "http://worker-0:8001",
|
|
}
|
|
|
|
async def user(request):
|
|
return SimpleNamespace(id="user")
|
|
|
|
async def access(*args, **kwargs):
|
|
return True
|
|
|
|
async def metadata(*args, **kwargs):
|
|
return {"metadata": {}}
|
|
|
|
monkeypatch.setattr(deps, "get_current_user_from_request", user)
|
|
monkeypatch.setattr(deps, "get_thread_store", lambda request: SimpleNamespace(check_access=access, get=metadata))
|
|
worker_count = 0
|
|
main_count = 0
|
|
|
|
class EventStream(httpx.AsyncByteStream):
|
|
async def __aiter__(self):
|
|
await asyncio.sleep(0)
|
|
yield b"event: end\ndata: {}\n\n"
|
|
|
|
async def worker(request):
|
|
nonlocal worker_count
|
|
assert request.headers["x-deerflow-worker-token"] == store.raw()["token"]
|
|
worker_count += 1
|
|
await asyncio.sleep(0.01)
|
|
return httpx.Response(200, stream=EventStream(), headers={"content-type": "text/event-stream"})
|
|
|
|
await pool.http.aclose()
|
|
pool.http = httpx.AsyncClient(transport=httpx.MockTransport(worker))
|
|
app = FastAPI()
|
|
app.state.worker_pool = pool
|
|
app.state.checkpointer = object()
|
|
app.state.store = object()
|
|
app.add_middleware(WorkerProxy)
|
|
|
|
@app.post("/api/threads/{thread_id}/runs/stream")
|
|
async def main_run(thread_id: str, request: Request):
|
|
nonlocal main_count
|
|
assert request.state.worker_pool_fallback is True
|
|
main_count += 1
|
|
return StreamingResponse(iter([b"event: end\ndata: {}\n\n"]), media_type="text/event-stream")
|
|
|
|
try:
|
|
async with httpx.AsyncClient(transport=httpx.ASGITransport(app), base_url="http://gateway") as client:
|
|
responses = await asyncio.gather(*(
|
|
client.post(f"/api/threads/t-{i}/runs/stream", json={"input": {}})
|
|
for i in range(70)
|
|
))
|
|
assert all(response.status_code == 200 and "event: end" in response.text for response in responses)
|
|
assert worker_count == 4
|
|
assert main_count == 66
|
|
assert pool._ensure_db().execute("SELECT count(*) FROM affinity").fetchone()[0] == 4
|
|
assert all(value == 0 for value in pool.inflight.values())
|
|
finally:
|
|
await pool.close()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("worker_thread_active,status_code", [(False, 200), (True, 409)])
|
|
async def test_affined_thread_uses_main_fallback_only_when_worker_thread_is_idle(tmp_path, monkeypatch, worker_thread_active, status_code):
|
|
from fastapi import FastAPI, Request
|
|
|
|
from app.gateway import deps
|
|
from app.gateway.worker_pool.proxy import WorkerProxy
|
|
|
|
store = SettingsStore(tmp_path)
|
|
store.save(PoolSettings(enabled=True, concurrency=1))
|
|
pool = Controller(store)
|
|
pool.nodes["worker-0"] = {
|
|
"name": "worker-0", "ready": True, "active": 1, "capacity": 1,
|
|
"draining": False, "url": "http://worker-0:8001",
|
|
}
|
|
pool._ensure_db().execute("INSERT INTO affinity VALUES (?, ?)", ("thread-1", "worker-0"))
|
|
pool._ensure_db().commit()
|
|
|
|
async def user(request):
|
|
return SimpleNamespace(id="user")
|
|
|
|
async def access(*args, **kwargs):
|
|
return True
|
|
|
|
async def metadata(*args, **kwargs):
|
|
return {"metadata": {}}
|
|
|
|
monkeypatch.setattr(deps, "get_current_user_from_request", user)
|
|
monkeypatch.setattr(deps, "get_thread_store", lambda request: SimpleNamespace(check_access=access, get=metadata))
|
|
|
|
checked = []
|
|
|
|
async def worker(request):
|
|
checked.append(request.url.path)
|
|
return httpx.Response(200, json={"active": worker_thread_active})
|
|
|
|
await pool.http.aclose()
|
|
pool.http = httpx.AsyncClient(transport=httpx.MockTransport(worker))
|
|
app = FastAPI()
|
|
app.state.worker_pool = pool
|
|
app.state.checkpointer = object()
|
|
app.state.store = object()
|
|
app.add_middleware(WorkerProxy)
|
|
|
|
@app.post("/api/threads/thread-1/runs/stream")
|
|
async def main_run(request: Request):
|
|
assert request.state.worker_pool_fallback is True
|
|
return {"handled_by": "main"}
|
|
|
|
try:
|
|
async with httpx.AsyncClient(transport=httpx.ASGITransport(app), base_url="http://gateway") as client:
|
|
result = await client.post("/api/threads/thread-1/runs/stream", json={"input": {}})
|
|
assert result.status_code == status_code
|
|
assert checked == ["/internal/worker-thread-active/thread-1"]
|
|
if status_code == 200:
|
|
assert result.json() == {"handled_by": "main"}
|
|
finally:
|
|
await pool.close()
|
|
|
|
|
|
def test_unified_pg_and_legacy_override(monkeypatch):
|
|
from app.gateway.worker_pool.preflight import shared_postgres_configured
|
|
|
|
cfg = SimpleNamespace(database=SimpleNamespace(backend="postgres"), checkpointer=None)
|
|
assert shared_postgres_configured(cfg)
|
|
cfg.checkpointer = SimpleNamespace(type="sqlite", use_postgres=True)
|
|
assert not shared_postgres_configured(cfg)
|
|
cfg.checkpointer.type = "postgres"
|
|
assert shared_postgres_configured(cfg)
|
|
monkeypatch.setenv("DEER_FLOW_USE_PG_CHECKPOINTER", "false")
|
|
assert not shared_postgres_configured(cfg)
|
|
|
|
|
|
def test_mysql_business_supports_both_chat_backends(monkeypatch, tmp_path):
|
|
from app.gateway.worker_pool.preflight import shared_storage_configured, validate_sqlite_mount
|
|
|
|
cfg = SimpleNamespace(database=SimpleNamespace(backend="mysql"), checkpointer=SimpleNamespace(type="postgres", use_postgres=True))
|
|
monkeypatch.delenv("DEER_FLOW_USE_PG_CHECKPOINTER", raising=False)
|
|
assert shared_storage_configured(cfg)
|
|
cfg.checkpointer = SimpleNamespace(type="sqlite", connection_string=str(tmp_path / "chat.db"))
|
|
assert shared_storage_configured(cfg)
|
|
mountinfo = f"31 1 8:1 / {tmp_path.as_posix()} rw - ext4 /dev/sda rw"
|
|
validate_sqlite_mount(cfg, mountinfo)
|
|
with pytest.raises(RuntimeError, match="network filesystem"):
|
|
validate_sqlite_mount(cfg, mountinfo.replace("ext4", "nfs4"))
|
|
with pytest.raises(RuntimeError, match="persistently mounted"):
|
|
validate_sqlite_mount(cfg, "31 1 8:1 / / rw - ext4 /dev/sda rw")
|
|
cfg.checkpointer.connection_string = ":memory:"
|
|
assert not shared_storage_configured(cfg)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_delayed_health_probe_does_not_hide_existing_run(tmp_path):
|
|
from app.gateway.worker_pool.controller import Controller, PoolUnavailable
|
|
from app.gateway.worker_pool.settings import SettingsStore
|
|
|
|
pool = Controller(SettingsStore(tmp_path))
|
|
pool.nodes['worker'] = dict(name='worker', ready=False, ever_ready=True, draining=False, url='http://worker:8001')
|
|
pool._ensure_db().execute("INSERT INTO affinity VALUES (?, ?)", ('thread', 'worker'))
|
|
pool._ensure_db().commit()
|
|
try:
|
|
assert (await pool.reserve('thread', create=False))['name'] == 'worker'
|
|
pool.release('worker')
|
|
with pytest.raises(PoolUnavailable):
|
|
await pool.reserve('thread', create=True)
|
|
pool.nodes['worker']['ever_ready'] = False
|
|
with pytest.raises(PoolUnavailable):
|
|
await pool.reserve('thread', create=False)
|
|
finally:
|
|
await pool.close()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_reconcile_retains_previous_readiness_during_probe_timeout(tmp_path, monkeypatch):
|
|
import httpx
|
|
|
|
store = SettingsStore(tmp_path)
|
|
store.save(PoolSettings(enabled=True, desired=1, source_container='main', network='network'))
|
|
pool = Controller(store)
|
|
name = f"dfw-{store.raw()['instance']}-0"
|
|
pool.nodes[name] = dict(name=name, ready=True, ever_ready=True, active=1, draining=False)
|
|
pool._ensure_db().execute('INSERT INTO affinity VALUES (?, ?)', ('thread', name))
|
|
pool._ensure_db().commit()
|
|
|
|
class Engine:
|
|
def request(self, method, path, body=None):
|
|
assert path.startswith('/containers/json')
|
|
return [{'Names': ['/' + name], 'Id': 'id', 'State': 'running'}]
|
|
|
|
async def inspect(_settings):
|
|
return Engine(), {}
|
|
|
|
async def timeout(*args, **kwargs):
|
|
raise httpx.ReadTimeout('busy')
|
|
|
|
monkeypatch.setattr(pool, 'inspect_template', inspect)
|
|
monkeypatch.setattr(pool.http, 'get', timeout)
|
|
try:
|
|
await pool.reconcile()
|
|
await pool.reconcile()
|
|
assert (await pool.reserve('thread', create=False))['name'] == name
|
|
with pytest.raises(PoolUnavailable):
|
|
await pool.reserve('thread', create=True)
|
|
finally:
|
|
await pool.close()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_read_transport_retries_but_never_replays_model_submission():
|
|
import httpx
|
|
|
|
from app.gateway.worker_pool.proxy import send_upstream
|
|
|
|
attempts = []
|
|
|
|
async def transport(request):
|
|
attempts.append(request.method)
|
|
if request.method == 'POST' or len(attempts) == 1:
|
|
raise httpx.RemoteProtocolError('server closed keep-alive')
|
|
return httpx.Response(200, json={'status': 'success'})
|
|
|
|
async with httpx.AsyncClient(transport=httpx.MockTransport(transport)) as client:
|
|
response = await send_upstream(client, 'GET', 'http://worker/run', [], None)
|
|
assert response.status_code == 200
|
|
await response.aclose()
|
|
with pytest.raises(httpx.RemoteProtocolError):
|
|
await send_upstream(client, 'POST', 'http://worker/run', [], b'{}')
|
|
assert attempts == ['GET', 'GET', 'POST']
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_second_gateway_cannot_start_schedulers(tmp_path, monkeypatch):
|
|
from app.gateway.worker_pool.lifecycle import managed_controller
|
|
from app.gateway.worker_pool.settings import PoolSettings, SettingsStore
|
|
|
|
monkeypatch.setenv("DEER_FLOW_WORKER_CONTROL_HOME", str(tmp_path))
|
|
async with managed_controller(SimpleNamespace(state=SimpleNamespace())):
|
|
async with managed_controller(SimpleNamespace(state=SimpleNamespace())):
|
|
pass
|
|
assert not (tmp_path / "affinity.sqlite3").exists()
|
|
assert not (tmp_path / "diagnostics.sqlite3").exists()
|
|
assert not (tmp_path / "controller.lock").exists()
|
|
assert not (tmp_path / "settings.json").exists()
|
|
assert not (tmp_path / "secret.key").exists()
|
|
|
|
store = SettingsStore(tmp_path)
|
|
store.save(PoolSettings(enabled=True, source_container="main"))
|
|
async with managed_controller(SimpleNamespace(state=SimpleNamespace())) as first:
|
|
assert first is not None
|
|
first.start()
|
|
with pytest.raises(RuntimeError, match="只能由一个主服务进程"):
|
|
async with managed_controller(SimpleNamespace(state=SimpleNamespace())) as second:
|
|
second.start()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_disabled_pool_ignores_saved_docker_settings(tmp_path, monkeypatch):
|
|
store = SettingsStore(tmp_path)
|
|
store.save(PoolSettings(enabled=False, transport="socket", source_container="main", network="deerflow"))
|
|
pool = Controller(store)
|
|
|
|
def unexpected_request(*args, **kwargs):
|
|
pytest.fail("disabled Worker pool must not contact Docker")
|
|
|
|
monkeypatch.setattr("app.gateway.worker_pool.controller.DockerEngine.request", unexpected_request)
|
|
pool.start()
|
|
await asyncio.sleep(0)
|
|
assert pool.task is None
|
|
assert pool.db is None
|
|
assert not (tmp_path / "controller.lock").exists()
|
|
assert not (tmp_path / "affinity.sqlite3").exists()
|
|
await pool.close()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_non_admin_cannot_read_host_credentials(monkeypatch):
|
|
from app.gateway.worker_pool import router
|
|
|
|
async def user(request):
|
|
return SimpleNamespace(system_role="user")
|
|
|
|
monkeypatch.setattr(router, "get_current_user_from_request", user)
|
|
with pytest.raises(HTTPException) as error:
|
|
await router.controller(SimpleNamespace())
|
|
assert error.value.status_code == 403
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_disabled_proxy_and_reserved_search_are_unchanged(tmp_path):
|
|
from fastapi import FastAPI
|
|
|
|
from app.gateway.worker_pool.proxy import WorkerProxy
|
|
|
|
pool = Controller(SettingsStore(tmp_path))
|
|
app = FastAPI()
|
|
app.state.worker_pool = pool
|
|
app.add_middleware(WorkerProxy)
|
|
|
|
@app.post("/api/threads/search")
|
|
@app.post("/api/threads/t/runs/stream")
|
|
async def endpoint():
|
|
return {"local": True}
|
|
|
|
try:
|
|
async with httpx.AsyncClient(transport=httpx.ASGITransport(app), base_url="http://test") as client:
|
|
assert (await client.post("/api/threads/search")).json() == {"local": True}
|
|
assert (await client.post("/api/threads/t/runs/stream")).json() == {"local": True}
|
|
finally:
|
|
await pool.close()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("allowed,system,status_code", [(False, False, 404), (True, True, 200)])
|
|
async def test_private_and_scheduler_threads_are_not_assigned(tmp_path, monkeypatch, allowed, system, status_code):
|
|
from fastapi import FastAPI
|
|
|
|
from app.gateway import deps
|
|
from app.gateway.worker_pool.proxy import WorkerProxy
|
|
|
|
store = SettingsStore(tmp_path)
|
|
store.save(PoolSettings(enabled=True))
|
|
pool = Controller(store)
|
|
|
|
async def user(request):
|
|
return SimpleNamespace(id="user")
|
|
|
|
async def access(*args, **kwargs):
|
|
return allowed
|
|
|
|
async def metadata(*args, **kwargs):
|
|
return {"metadata": {"thread_type": "scheduler"} if system else {}}
|
|
|
|
monkeypatch.setattr(deps, "get_current_user_from_request", user)
|
|
monkeypatch.setattr(deps, "get_thread_store", lambda request: SimpleNamespace(check_access=access, get=metadata))
|
|
app = FastAPI()
|
|
app.state.worker_pool = pool
|
|
app.add_middleware(WorkerProxy)
|
|
|
|
@app.post("/api/threads/t/runs/stream")
|
|
async def endpoint():
|
|
return {"scheduler_stays_local": True}
|
|
|
|
try:
|
|
async with httpx.AsyncClient(transport=httpx.ASGITransport(app), base_url="http://gateway") as client:
|
|
result = await client.post("/api/threads/t/runs/stream")
|
|
assert result.status_code == status_code
|
|
assert pool._ensure_db().execute("SELECT count(*) FROM affinity").fetchone()[0] == 0
|
|
finally:
|
|
await pool.close()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_worker_boot_disables_background_and_refuses_memory_fallback(monkeypatch):
|
|
from contextlib import asynccontextmanager
|
|
|
|
from langgraph.checkpoint.memory import InMemorySaver
|
|
|
|
from app.gateway.worker_pool import worker
|
|
from deerflow.config.app_config import AppConfig
|
|
from deerflow.config.database_config import DatabaseConfig
|
|
|
|
cfg = AppConfig.model_construct(
|
|
database=DatabaseConfig.model_construct(backend="postgres"),
|
|
checkpointer=None,
|
|
sandbox=SimpleNamespace(use="LocalSandboxProvider"),
|
|
)
|
|
monkeypatch.setattr(worker, "get_app_config", lambda: cfg)
|
|
|
|
async def configure(config):
|
|
pass
|
|
|
|
monkeypatch.setattr(worker, "configure_runtime_executor", configure)
|
|
|
|
@asynccontextmanager
|
|
async def runtime(app, *, background):
|
|
assert background is False
|
|
app.state.checkpointer = InMemorySaver()
|
|
app.state.store = object()
|
|
yield
|
|
|
|
monkeypatch.setattr(worker, "langgraph_runtime", runtime)
|
|
app = SimpleNamespace(state=SimpleNamespace())
|
|
with pytest.raises(RuntimeError, match="fallback"):
|
|
async with worker.worker_lifespan(app):
|
|
pytest.fail("memory worker became ready")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_skill_disable_invalidates_other_worker_cache_immediately(monkeypatch):
|
|
from app.gateway.worker_pool.refresh import SharedSkillRefresh
|
|
from deerflow.agents.lead_agent import prompt
|
|
from deerflow.config import extensions_config
|
|
|
|
refresh = SharedSkillRefresh()
|
|
stamp = 1
|
|
calls = []
|
|
monkeypatch.setattr(refresh, "extension_stamp", lambda: stamp)
|
|
monkeypatch.setattr(refresh, "skill_stamp", lambda: ())
|
|
monkeypatch.setattr(extensions_config, "reload_extensions_config", lambda: calls.append("reload"))
|
|
|
|
async def invalidate():
|
|
calls.append("refresh")
|
|
|
|
monkeypatch.setattr(prompt, "refresh_skills_system_prompt_cache_async", invalidate)
|
|
await refresh.check()
|
|
await refresh.check()
|
|
assert calls == ["reload", "refresh"]
|
|
stamp = 2
|
|
await refresh.check()
|
|
assert calls == ["reload", "refresh", "reload", "refresh"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_scaleup_partial_failure_remains_owned(tmp_path, monkeypatch):
|
|
store = SettingsStore(tmp_path)
|
|
store.save(PoolSettings(enabled=True, source_container="main", network="deerflow", desired=1))
|
|
pool = Controller(store)
|
|
|
|
class Engine:
|
|
def request(self, method, path, body=None):
|
|
if path.startswith("/containers/json"):
|
|
return []
|
|
if path.startswith("/containers/create?"):
|
|
return {"Id": "created-before-failure"}
|
|
raise DockerError("start failed")
|
|
|
|
async def template(settings):
|
|
return Engine(), {}
|
|
|
|
monkeypatch.setattr(pool, "inspect_template", template)
|
|
try:
|
|
with pytest.raises(DockerError):
|
|
await pool.reconcile()
|
|
assert len(pool.nodes) == 1
|
|
assert next(iter(pool.nodes.values()))["id"] == "created-before-failure"
|
|
assert not next(iter(pool.nodes.values()))["ready"]
|
|
finally:
|
|
await pool.close()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_local_preparation_cannot_race_remote_assignment(tmp_path):
|
|
store = SettingsStore(tmp_path)
|
|
store.save(PoolSettings(enabled=True))
|
|
pool = Controller(store)
|
|
request = SimpleNamespace(app=SimpleNamespace(state=SimpleNamespace(worker_pool=pool)))
|
|
entered, finish = asyncio.Event(), asyncio.Event()
|
|
|
|
@bounded_worker_run
|
|
async def start(body, thread_id, request):
|
|
entered.set()
|
|
await finish.wait()
|
|
return SimpleNamespace(task=None)
|
|
|
|
task = asyncio.create_task(start(None, "thread", request))
|
|
try:
|
|
await entered.wait()
|
|
with pytest.raises(PoolUnavailable, match="主服务执行"):
|
|
await pool.reserve("thread", create=True)
|
|
finish.set()
|
|
await task
|
|
assert not pool.local_threads
|
|
finally:
|
|
task.cancel()
|
|
await asyncio.gather(task, return_exceptions=True)
|
|
await pool.close()
|