207 lines
8.7 KiB
Python
207 lines
8.7 KiB
Python
"""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
|