360 lines
18 KiB
Python
360 lines
18 KiB
Python
"""One gateway owns reconciliation and durable conversation affinity.
|
||
|
||
Never move conversations from a live worker. A draining worker is stopped only
|
||
after its run registry AND outstanding HTTP requests are empty. Failure does
|
||
not replay a request, because a model/tool may already have produced effects.
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import asyncio
|
||
import json
|
||
import os
|
||
import sqlite3
|
||
import time
|
||
from urllib.parse import quote
|
||
|
||
import httpx
|
||
|
||
from .docker import DockerEngine, DockerError, worker_spec
|
||
from .settings import SettingsStore
|
||
|
||
|
||
class PoolUnavailable(RuntimeError):
|
||
pass
|
||
|
||
|
||
class Controller:
|
||
def __init__(self, store: SettingsStore):
|
||
self.store = store
|
||
self.lock = asyncio.Lock()
|
||
self.route_lock = asyncio.Lock()
|
||
self.nodes: dict[str, dict] = {}
|
||
self.error: str | None = None
|
||
self.task: asyncio.Task | None = None
|
||
self._lock_handle = None
|
||
self.http = httpx.AsyncClient(timeout=httpx.Timeout(15, read=None), trust_env=False, limits=httpx.Limits(max_connections=256, max_keepalive_connections=64))
|
||
self.db: sqlite3.Connection | None = None
|
||
self.pending = 0
|
||
self.inflight: dict[str, int] = {}
|
||
self.local_threads: dict[str, int] = {}
|
||
self._python_probe_key: tuple | None = None
|
||
self._python_probe_result: str | None = None
|
||
self._worker_image_key: tuple | None = None
|
||
self._worker_image_result: str | None = None
|
||
self.discovery_complete = not bool(store.load().source_container)
|
||
|
||
def _ensure_db(self):
|
||
if self.db is None:
|
||
self.db = sqlite3.connect(self.store.directory / "affinity.sqlite3")
|
||
self.db.execute("PRAGMA journal_mode=WAL")
|
||
self.db.execute("CREATE TABLE IF NOT EXISTS affinity (thread TEXT PRIMARY KEY, node TEXT NOT NULL)")
|
||
self.db.commit()
|
||
return self.db
|
||
|
||
def _acquire_controller_lock(self):
|
||
if self._lock_handle is not None:
|
||
return
|
||
self.store.directory.mkdir(mode=0o700, parents=True, exist_ok=True)
|
||
handle = (self.store.directory / "controller.lock").open("a+b")
|
||
try:
|
||
if os.name == "nt":
|
||
import msvcrt
|
||
|
||
handle.seek(0)
|
||
if not handle.read(1):
|
||
handle.write(b"0")
|
||
handle.flush()
|
||
handle.seek(0)
|
||
msvcrt.locking(handle.fileno(), msvcrt.LK_NBLCK, 1)
|
||
else:
|
||
import fcntl
|
||
|
||
fcntl.flock(handle, fcntl.LOCK_EX | fcntl.LOCK_NB)
|
||
except OSError:
|
||
handle.close()
|
||
raise RuntimeError("Worker 池只能由一个主服务进程管理；请关闭重复的 Worker 池控制器") from None
|
||
self._lock_handle = handle
|
||
|
||
def start(self):
|
||
settings = self.store.load()
|
||
if settings.enabled:
|
||
self._acquire_controller_lock()
|
||
self._ensure_db()
|
||
if self.task is None or self.task.done():
|
||
self.task = asyncio.create_task(self.loop())
|
||
if hasattr(self, "diagnostics") and self.diagnostics.enabled:
|
||
self.diagnostics.start()
|
||
|
||
async def close(self):
|
||
if self.task:
|
||
self.task.cancel()
|
||
await asyncio.gather(self.task, return_exceptions=True)
|
||
if hasattr(self, "diagnostics"):
|
||
await self.diagnostics.close()
|
||
await self.http.aclose()
|
||
if self.db is not None:
|
||
self.db.close()
|
||
self.db = None
|
||
self._release_controller_lock()
|
||
|
||
async def loop(self):
|
||
try:
|
||
while True:
|
||
try:
|
||
await self.reconcile()
|
||
except Exception as exc:
|
||
self.error = str(exc) if isinstance(exc, (DockerError, PoolUnavailable)) else f"节点检查失败（{type(exc).__name__}）"
|
||
settings = self.store.load()
|
||
if not settings.enabled and not self.nodes:
|
||
return
|
||
await asyncio.sleep(5)
|
||
finally:
|
||
if not self.store.load().enabled:
|
||
self._release_controller_lock()
|
||
|
||
def _release_controller_lock(self):
|
||
if self._lock_handle is None:
|
||
return
|
||
if os.name == "nt":
|
||
import msvcrt
|
||
|
||
self._lock_handle.seek(0)
|
||
msvcrt.locking(self._lock_handle.fileno(), msvcrt.LK_UNLCK, 1)
|
||
else:
|
||
import fcntl
|
||
|
||
fcntl.flock(self._lock_handle, fcntl.LOCK_UN)
|
||
self._lock_handle.close()
|
||
self._lock_handle = None
|
||
|
||
async def inspect_template(self, settings):
|
||
config = settings
|
||
if config.desired * config.cpu > config.total_cpu or config.desired * config.memory_mb > config.total_memory_mb:
|
||
raise PoolUnavailable("Worker 总资源超过设置的预算，请减少数量或调整预算")
|
||
if config.desired * (config.database_pool_size + 2 * config.postgres_pool_size) > config.database_connection_budget:
|
||
raise PoolUnavailable("Worker 数据库连接总额度超过预算，请减少数量或单节点连接数")
|
||
if not config.source_container or not config.network:
|
||
raise PoolUnavailable("请填写源容器名称和 Docker 网络")
|
||
engine = DockerEngine(config)
|
||
source = await asyncio.to_thread(engine.request, "GET", f"/containers/{quote(config.source_container, safe='')}/json?size=true")
|
||
if source.get("Config", {}).get("Labels", {}).get("deerflow.worker-pool"):
|
||
raise PoolUnavailable("源容器必须是主服务，不能选择已有 Worker")
|
||
backend = await asyncio.to_thread(
|
||
engine.configured_storage_backend, source, config.python_executable, config.working_dir,
|
||
)
|
||
probe_key = (
|
||
config.transport, config.host, config.port, config.username, config.source_container,
|
||
source.get("Id"), source.get("Image"), config.python_executable,
|
||
config.working_dir, source.get("Config", {}).get("User", ""), backend,
|
||
)
|
||
if probe_key == self._python_probe_key and self._python_probe_result:
|
||
python_executable = self._python_probe_result
|
||
else:
|
||
python_executable = await asyncio.to_thread(
|
||
engine.discover_worker_python, source, config.python_executable, config.working_dir,
|
||
)
|
||
self._python_probe_key = probe_key
|
||
self._python_probe_result = python_executable
|
||
if python_executable != config.python_executable:
|
||
config = config.model_copy(update={"python_executable": python_executable})
|
||
self.store.save(config)
|
||
self._python_probe_key = (
|
||
config.transport, config.host, config.port, config.username, config.source_container,
|
||
source.get("Id"), source.get("Image"), config.python_executable,
|
||
config.working_dir, source.get("Config", {}).get("User", ""), backend,
|
||
)
|
||
runtime_image = None
|
||
if config.enabled and config.desired > 0:
|
||
image_key = (source.get("Id"), source.get("Image"), python_executable, backend)
|
||
if image_key != self._worker_image_key or not self._worker_image_result:
|
||
self._worker_image_result = await asyncio.to_thread(
|
||
engine.snapshot_worker_runtime, source, self.store.raw()["instance"],
|
||
)
|
||
self._worker_image_key = image_key
|
||
runtime_image = self._worker_image_result
|
||
spec = worker_spec(
|
||
source, config, self.store.raw()["instance"], self.store.raw()["token"],
|
||
str(self.store.directory), runtime_image=runtime_image,
|
||
require_shared_archive=backend == "postgres" and config.enabled and config.desired > 0,
|
||
)
|
||
return engine, spec
|
||
|
||
async def reconcile(self):
|
||
async with self.lock:
|
||
settings = self.store.load()
|
||
if not settings.enabled and not settings.source_container:
|
||
return
|
||
db = self._ensure_db()
|
||
engine, spec = await self.inspect_template(settings)
|
||
instance = self.store.raw()["instance"]
|
||
filters = quote(json.dumps({"label": [f"deerflow.worker-pool={instance}"]}))
|
||
containers = await asyncio.to_thread(engine.request, "GET", f"/containers/json?all=1&filters={filters}")
|
||
self.discovery_complete = True
|
||
desired = settings.desired if settings.enabled else 0
|
||
wanted = {f"dfw-{instance}-{i}" for i in range(desired)}
|
||
existing = {row["Names"][0].lstrip("/"): row for row in containers}
|
||
stale_templates = set()
|
||
for name, row in list(existing.items()):
|
||
self.nodes.setdefault(name, {"name": name, "id": row["Id"], "ready": False, "active": 0, "draining": name not in wanted, "url": f"http://{name}:8001"})
|
||
if name in wanted and row.get("State") == "running" and row.get("ImageID") and row["ImageID"] != spec["Image"]:
|
||
stale_templates.add(name)
|
||
self.nodes[name]["draining"] = True
|
||
|
||
# A stopped container cannot be executing a request. If its saved
|
||
# startup template has an old Python path, replace it with the
|
||
# resolved template instead of restarting the same broken command.
|
||
for name, row in list(existing.items()):
|
||
if name not in wanted or row["State"] == "running":
|
||
continue
|
||
detail = await asyncio.to_thread(engine.request, "GET", f"/containers/{quote(row['Id'], safe='')}/json")
|
||
worker_config = detail.get("Config", {})
|
||
stale_template = (
|
||
detail.get("Image") != spec["Image"]
|
||
or worker_config.get("Entrypoint") != spec["Entrypoint"]
|
||
or worker_config.get("WorkingDir") != spec["WorkingDir"]
|
||
)
|
||
if not stale_template:
|
||
continue
|
||
if hasattr(self, "diagnostics"):
|
||
await self.diagnostics.archive_logs(engine, name)
|
||
await asyncio.to_thread(engine.request, "DELETE", f"/containers/{quote(row['Id'], safe='')}")
|
||
db.execute("DELETE FROM affinity WHERE node=?", (name,))
|
||
db.commit()
|
||
existing.pop(name, None)
|
||
self.nodes.pop(name, None)
|
||
for name, row in list(existing.items()):
|
||
previous = self.nodes.get(name, {})
|
||
node = {
|
||
"name": name,
|
||
"id": row["Id"],
|
||
"ready": previous.get("ready", False),
|
||
"ever_ready": previous.get("ever_ready", False),
|
||
"active": previous.get("active", 0),
|
||
"capacity": previous.get("capacity"),
|
||
"draining": name not in wanted or name in stale_templates,
|
||
"url": f"http://{name}:8001",
|
||
"checked_at": time.time(),
|
||
}
|
||
for key in ("running", "waiting", "process_started_at", "runtime_revision", "event_loop_lag_ms", "http"):
|
||
if key in previous:
|
||
node[key] = previous[key]
|
||
self.nodes[name] = node
|
||
if row["State"] != "running":
|
||
node["ready"] = False
|
||
node["ever_ready"] = False
|
||
# A confirmed stopped process cannot keep executing. Old
|
||
# conversations can move only after this confirmation.
|
||
db.execute("DELETE FROM affinity WHERE node=?", (name,))
|
||
db.commit()
|
||
if name in wanted:
|
||
await asyncio.to_thread(engine.request, "POST", f"/containers/{row['Id']}/start")
|
||
else:
|
||
if hasattr(self, "diagnostics"):
|
||
await self.diagnostics.archive_logs(engine, name)
|
||
await asyncio.to_thread(engine.request, "DELETE", f"/containers/{row['Id']}")
|
||
self.nodes.pop(name, None)
|
||
continue
|
||
try:
|
||
diagnostics_enabled = bool(getattr(getattr(self, "diagnostics", None), "enabled", False))
|
||
response = await self.http.get(
|
||
node["url"] + "/internal/worker-status",
|
||
headers={
|
||
"x-deerflow-worker-token": self.store.raw()["token"],
|
||
"x-deerflow-diagnostics-enabled": "true" if diagnostics_enabled else "false",
|
||
},
|
||
timeout=3,
|
||
)
|
||
response.raise_for_status()
|
||
status = response.json()
|
||
node.update(ready=status["ready"], active=status["active"], capacity=status["capacity"], ever_ready=True)
|
||
for key in ("running", "waiting", "process_started_at", "runtime_revision", "event_loop_lag_ms", "http"):
|
||
if key in status:
|
||
node[key] = status[key]
|
||
except (httpx.HTTPError, ValueError, KeyError):
|
||
node["ready"] = False
|
||
node["error"] = "等待就绪或连接失败；请检查主容器是否位于同一网络、镜像代码和聊天存储配置"
|
||
continue
|
||
if node["draining"] and not node["active"] and not self.inflight.get(name, 0):
|
||
# Refuse all new reservations while stop is in flight.
|
||
node["ready"] = False
|
||
node["ever_ready"] = False
|
||
await asyncio.to_thread(engine.request, "POST", f"/containers/{row['Id']}/stop?t=30")
|
||
if hasattr(self, "diagnostics"):
|
||
await self.diagnostics.archive_logs(engine, name)
|
||
await asyncio.to_thread(engine.request, "DELETE", f"/containers/{row['Id']}")
|
||
db.execute("DELETE FROM affinity WHERE node=?", (name,))
|
||
db.commit()
|
||
self.nodes.pop(name, None)
|
||
existing.pop(name, None)
|
||
# Replacements are created only after the previous container has
|
||
# drained and been removed. This also refreshes Workers after the
|
||
# main service's runtime image or configured storage backend changes.
|
||
for name in sorted(wanted - existing.keys()):
|
||
row = await asyncio.to_thread(engine.request, "POST", f"/containers/create?name={name}", spec)
|
||
self.nodes[name] = {
|
||
"name": name,
|
||
"id": row["Id"],
|
||
"ready": False,
|
||
"active": 0,
|
||
"draining": False,
|
||
"url": f"http://{name}:8001",
|
||
}
|
||
await asyncio.to_thread(engine.request, "POST", f"/containers/{row['Id']}/start")
|
||
existing[name] = {"Id": row["Id"], "State": "running", "ImageID": spec["Image"]}
|
||
# Missing nodes are not silently reassigned: an operator must
|
||
# reconcile the actual container state before changing affinity.
|
||
self.error = None
|
||
|
||
async def reserve(self, thread: str, *, create: bool) -> dict | None:
|
||
settings = self.store.load()
|
||
if not settings.enabled and not self.nodes:
|
||
return None
|
||
db = self._ensure_db()
|
||
# Container creation/SSH can be slow; routing must never wait on that lock.
|
||
async with self.route_lock:
|
||
row = db.execute("SELECT node FROM affinity WHERE thread=?", (thread,)).fetchone()
|
||
if row is None and create and self.local_threads.get(thread, 0):
|
||
raise PoolUnavailable("该会话正在主服务执行，请等待当前轮完成")
|
||
if row:
|
||
node = self.nodes.get(row[0])
|
||
if node and node.get("draining") and create:
|
||
return {"_fallback_to_main": True, "_fallback_reason": "worker_draining", "_worker": dict(node)}
|
||
# A delayed health probe stops new work, but must not hide an
|
||
# existing run's result. Keep reads/cancellation on its owner;
|
||
# stopped containers clear ever_ready before affinity changes.
|
||
if not node or (not node["ready"] and (create or not node.get("ever_ready"))):
|
||
raise PoolUnavailable("该会话的执行节点暂不可用，请稍后重试；系统不会重复提交任务")
|
||
if node["draining"] and create:
|
||
raise PoolUnavailable("该会话节点正在停止接单，请等待当前任务完成后重试")
|
||
if create:
|
||
capacity = max(1, int(node.get("capacity") or settings.concurrency))
|
||
load = max(int(node.get("active") or 0), self.inflight.get(node["name"], 0))
|
||
if load >= capacity:
|
||
return {"_fallback_to_main": True, "_fallback_reason": "worker_saturated", "_worker": dict(node)}
|
||
elif not settings.enabled or not create:
|
||
return None
|
||
else:
|
||
choices = [n for n in self.nodes.values() if n["ready"] and not n["draining"]]
|
||
if not choices:
|
||
return {"_fallback_to_main": True, "_fallback_reason": "workers_unavailable"}
|
||
available = []
|
||
for candidate in choices:
|
||
capacity = max(1, int(candidate.get("capacity") or settings.concurrency))
|
||
load = max(int(candidate.get("active") or 0), self.inflight.get(candidate["name"], 0))
|
||
if load < capacity:
|
||
available.append((load, candidate))
|
||
if not available:
|
||
return {"_fallback_to_main": True, "_fallback_reason": "worker_pool_saturated"}
|
||
_, node = min(available, key=lambda item: (item[0], item[1]["name"]))
|
||
db.execute("INSERT INTO affinity VALUES (?, ?)", (thread, node["name"]))
|
||
db.commit()
|
||
self.inflight[node["name"]] = self.inflight.get(node["name"], 0) + 1
|
||
return dict(node)
|
||
|
||
def release(self, name):
|
||
self.inflight[name] = max(0, self.inflight.get(name, 0) - 1)
|
||
|
||
def status(self):
|
||
return {"settings": self.store.public(), "nodes": [{k: v for k, v in n.items() if k != "url"} for n in self.nodes.values()], "error": self.error}
|