"""Bound executing runs, including detached/background runs, not HTTP streams.""" from __future__ import annotations import asyncio from contextvars import ContextVar from functools import wraps from fastapi import HTTPException _worker_ticket = ContextVar("worker_admission_ticket", default=None) class AdmissionTicket: def __init__(self, gate): if gate.active + gate.waiting >= gate.capacity + gate.queue_limit: raise HTTPException(429, "执行节点繁忙，请稍后重试", headers={"Retry-After": "3"}) self.gate = gate self.deadline = asyncio.get_running_loop().time() + gate.timeout self.waiting = True self.acquired = False self.released = False self.started = False gate.waiting += 1 async def acquire(self): try: remaining = max(0.001, self.deadline - asyncio.get_running_loop().time()) await asyncio.wait_for(self.gate.semaphore.acquire(), remaining) self.acquired = True self.gate.active += 1 except TimeoutError: raise HTTPException(429, "等待执行超时，请稍后重试") from None finally: if self.waiting: self.gate.waiting -= 1 self.waiting = False def release(self, _task=None): if self.released: return self.released = True if self.waiting: self.gate.waiting -= 1 self.waiting = False if self.acquired: self.gate.active -= 1 self.gate.semaphore.release() async def _finish_queued_run(kwargs, *, cancelled=False, message="等待执行超时，请稍后重试"): from deerflow.runtime.runs.schemas import RunStatus manager, record, bridge = kwargs["run_manager"], kwargs["record"], kwargs["bridge"] try: async with asyncio.timeout(5): await manager.set_status(record.run_id, RunStatus.interrupted if cancelled else RunStatus.error, error=None if cancelled else message) ctx = kwargs.get("ctx") store = getattr(ctx, "thread_store", None) if store is not None and not getattr(record, "superseded", False): await store.update_status(record.thread_id, "interrupted" if cancelled else "error") finally: try: if not cancelled: await bridge.publish(record.run_id, "error", {"name": "WorkerQueueTimeout", "message": message}) finally: await bridge.publish_end(record.run_id) asyncio.create_task(bridge.cleanup(record.run_id, delay=60)) async def run_with_worker_admission(runner, **kwargs): ticket = _worker_ticket.get() if ticket is None: return await runner(**kwargs) ticket.started = True record, bridge = kwargs["record"], kwargs["bridge"] # Publish a real, persisted run id before waiting for an execution slot. # The router keeps its canonical Location/reconnect headers unchanged. try: await bridge.publish(record.run_id, "metadata", {"run_id": record.run_id, "thread_id": record.thread_id, "worker_queue": {"waiting": ticket.gate.waiting, "capacity": ticket.gate.capacity}}) await ticket.acquire() except asyncio.CancelledError: await _finish_queued_run(kwargs, cancelled=True) raise except HTTPException as exc: await _finish_queued_run(kwargs, message=str(exc.detail)) return except Exception: await _finish_queued_run(kwargs, message="执行节点排队初始化失败，请查看节点日志") return return await runner(**kwargs) class Admission: def __init__(self, capacity: int, queue_limit: int = 128, timeout: float = 60): self.capacity = capacity self.active = 0 self.waiting = 0 self.queue_limit = queue_limit self.timeout = timeout self.semaphore = asyncio.Semaphore(capacity) self.lock = asyncio.Lock() async def acquire(self): async with self.lock: if self.waiting >= self.queue_limit: raise HTTPException(429, "执行节点繁忙，请稍后重试", headers={"Retry-After": "3"}) self.waiting += 1 try: await asyncio.wait_for(self.semaphore.acquire(), self.timeout) except TimeoutError: raise HTTPException(429, "等待执行超时，请稍后重试", headers={"Retry-After": "3"}) from None finally: self.waiting -= 1 self.active += 1 def release(self, _task=None): self.active -= 1 self.semaphore.release() def bounded_worker_run(function): @wraps(function) async def wrapped(body, thread_id, request, *args, **kwargs): pool = getattr(request.app.state, "worker_pool", None) fallback_to_main = bool(getattr(getattr(request, "state", None), "worker_pool_fallback", False)) if pool is not None and pool.db is not None and pool.db.execute("SELECT 1 FROM affinity WHERE thread=?", (thread_id,)).fetchone() and not fallback_to_main: raise HTTPException(409, "该会话由独立 Worker 管理，请通过会话接口提交，不能在主服务重复执行") gate = getattr(request.app.state, "worker_admission", None) if gate is None: if pool is None: return await function(body, thread_id, request, *args, **kwargs) if pool.local_threads.get(thread_id, 0): raise HTTPException(409, "该会话正在主服务执行，请等待当前回答完成") # Reserve before the first await, including start_run's preparation # phase before RunRecord.task exists. Otherwise a concurrent routed # request could assign this same thread to another process. pool.local_threads[thread_id] = pool.local_threads.get(thread_id, 0) + 1 def release_local(_task=None): remaining = pool.local_threads.get(thread_id, 1) - 1 if remaining: pool.local_threads[thread_id] = remaining else: pool.local_threads.pop(thread_id, None) try: record = await function(body, thread_id, request, *args, **kwargs) except BaseException: release_local() raise if record.task is None: release_local() else: record.task.add_done_callback(release_local) return record if getattr(request.app.state, "worker_deferred_admission", False) and getattr(request.app.state, "concurrency_gate", None) is None: ticket = AdmissionTicket(gate) token = _worker_ticket.set(ticket) try: from deerflow.runtime.thread_paths import aresolve_path_user_id await aresolve_path_user_id(thread_id) refresh = getattr(request.app.state, "worker_skill_refresh", None) if refresh is not None: await refresh.check() record = await function(body, thread_id, request, *args, **kwargs) except BaseException: ticket.release() raise finally: _worker_ticket.reset(token) if record.task is None: ticket.release() else: def finish_ticket(task): ticket.release() if task.cancelled() and not ticket.started: manager = getattr(request.app.state, "run_manager", None) bridge = getattr(request.app.state, "stream_bridge", None) if manager is not None and bridge is not None: asyncio.create_task(_finish_queued_run({"run_manager": manager, "record": record, "bridge": bridge}, cancelled=True)) record.task.add_done_callback(finish_ticket) return record await gate.acquire() try: # File/path middleware is synchronous. Resolve the persisted owner # on the database's home loop before it touches those helpers. from deerflow.runtime.thread_paths import aresolve_path_user_id await aresolve_path_user_id(thread_id) refresh = getattr(request.app.state, "worker_skill_refresh", None) if refresh is not None: await refresh.check() record = await function(body, thread_id, request, *args, **kwargs) except BaseException: gate.release() raise if record.task is None: gate.release() else: record.task.add_done_callback(gate.release) return record return wrapped