sync update from LAN repo (source dhk-1 @ 3b1ec5b4)
This commit is contained in:
gozeros committed
1 parent
da6521b6cb
commit
264e2200d6
995 files changed
+148105
-13399
No files matched your search
@@ -0,0 +1,874 @@
|
||||
"""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()
|
||||
Reference in new issue
Block a user