164 lines
5.4 KiB
Python
164 lines
5.4 KiB
Python
"""派发队列：依赖拓扑排序、就绪筛选、预算判断与内存任务队列。
|
||
|
||
零第三方依赖，只用 Python 标准库；不依赖 host.py，纯函数 + 内存
|
||
状态实现。状态机与 host.py 对齐：
|
||
|
||
queued / running / completed / failed / interrupted / merged
|
||
/ merge-failed
|
||
|
||
注意：本模块只决定“按什么顺序派发”，不跑模型回合、不调 agent。
|
||
"""
|
||
|
||
VALID_STATUSES = frozenset({
|
||
"queued",
|
||
"running",
|
||
"completed",
|
||
"failed",
|
||
"interrupted",
|
||
"merged",
|
||
"merge-failed",
|
||
})
|
||
|
||
|
||
def _task_id(task):
|
||
task_id = task.get("id")
|
||
if task_id is None:
|
||
raise ValueError("任务缺少 id")
|
||
return task_id
|
||
|
||
|
||
def _depends_on(task):
|
||
deps = task.get("depends_on") or []
|
||
if isinstance(deps, str):
|
||
deps = [deps]
|
||
return list(deps)
|
||
|
||
|
||
def _priority(task):
|
||
return task.get("priority") or 0
|
||
|
||
|
||
def _topo_sort(tasks, initial_done=frozenset(), strict_missing=True):
|
||
"""按波次拓扑排序：依赖先排，同一波内 priority 大的先，
|
||
priority 相同时按 id 升序（保证确定性）。
|
||
|
||
initial_done：已视为完成的 id 集合（ready_tasks 用 completed_ids
|
||
填充它，这样就绪子集里的依赖不再被当成缺失依赖）。
|
||
strict_missing=True 时，依赖不在任务集且不在 initial_done 里
|
||
即 raise ValueError；成环同样 raise ValueError。
|
||
"""
|
||
by_id = {}
|
||
for task in tasks:
|
||
task_id = _task_id(task)
|
||
if task_id in by_id:
|
||
raise ValueError(f"任务 id 重复：{task_id}")
|
||
by_id[task_id] = task
|
||
for task in tasks:
|
||
for dep in _depends_on(task):
|
||
if dep == _task_id(task):
|
||
raise ValueError(f"任务 {_task_id(task)} 不能依赖自己")
|
||
if strict_missing and dep not in by_id and dep not in initial_done:
|
||
raise ValueError(
|
||
f"任务 {_task_id(task)} 依赖不存在的任务 {dep}")
|
||
|
||
ordered = []
|
||
done = set(initial_done)
|
||
remaining = list(tasks)
|
||
while remaining:
|
||
ready = [t for t in remaining
|
||
if all(d in done for d in _depends_on(t))]
|
||
if not ready:
|
||
names = ", ".join(str(_task_id(t)) for t in remaining)
|
||
raise ValueError(f"依赖成环，无法排序：{names}")
|
||
ready.sort(key=lambda t: (-_priority(t), str(_task_id(t))))
|
||
for task in ready:
|
||
ordered.append(task)
|
||
done.add(_task_id(task))
|
||
remaining = [t for t in remaining if _task_id(t) not in done]
|
||
return ordered
|
||
|
||
|
||
def order_tasks(tasks) -> list:
|
||
"""按依赖拓扑排序，同一波内 priority 大的先。
|
||
|
||
tasks 是 dict 列表，每项含 id、depends_on（list，可缺省）、
|
||
priority（int，可缺省 0）。缺失依赖或成环 raise ValueError。
|
||
不修改输入列表，返回新列表。
|
||
"""
|
||
return _topo_sort(list(tasks))
|
||
|
||
|
||
def ready_tasks(tasks, completed_ids) -> list:
|
||
"""返回依赖全部在 completed_ids 里、且状态为 queued 的任务，
|
||
再按 order_tasks 的规则（priority 降序、同波拓扑）排序。
|
||
|
||
未写 status 字段的任务按 queued 计；running/completed/failed 等
|
||
状态的任务一律不返回。
|
||
"""
|
||
completed = set(completed_ids or [])
|
||
ready = [t for t in tasks
|
||
if t.get("status", "queued") == "queued"
|
||
and all(d in completed for d in _depends_on(t))]
|
||
return _topo_sort(ready, initial_done=completed, strict_missing=False)
|
||
|
||
|
||
def budget_exceeded(usage_turns, max_turns, usage_cost_usd,
|
||
max_cost_usd) -> bool:
|
||
"""任一上限非 None 且用量 >= 上限即 True，否则 False。
|
||
|
||
上限为 None 表示该项不限；用量为 None 按 0 计。
|
||
"""
|
||
if max_turns is not None and (usage_turns or 0) >= max_turns:
|
||
return True
|
||
if max_cost_usd is not None and (usage_cost_usd or 0) >= max_cost_usd:
|
||
return True
|
||
return False
|
||
|
||
|
||
class TaskQueue:
|
||
"""内存队列：add / mark_running / mark_done / next_ready。
|
||
|
||
内部存任务 dict 的浅拷贝，add 时缺省状态为 queued；status
|
||
只能是 VALID_STATUSES 中的值，否则 raise ValueError。
|
||
"""
|
||
|
||
def __init__(self, tasks=None):
|
||
self._tasks = {}
|
||
for task in tasks or []:
|
||
self.add(task)
|
||
|
||
def add(self, task):
|
||
task_id = _task_id(task)
|
||
if task_id in self._tasks:
|
||
raise ValueError(f"任务 id 重复：{task_id}")
|
||
stored = dict(task)
|
||
status = stored.get("status", "queued")
|
||
if status not in VALID_STATUSES:
|
||
raise ValueError(f"非法任务状态：{status}")
|
||
stored["status"] = status
|
||
self._tasks[task_id] = stored
|
||
return stored
|
||
|
||
def mark_running(self, task_id):
|
||
task = self._get(task_id)
|
||
task["status"] = "running"
|
||
return task
|
||
|
||
def mark_done(self, task_id, status):
|
||
if status not in VALID_STATUSES:
|
||
raise ValueError(f"非法任务状态：{status}")
|
||
task = self._get(task_id)
|
||
task["status"] = status
|
||
return task
|
||
|
||
def next_ready(self, completed_ids):
|
||
"""按 ready_tasks 顺序的第一个可派任务，或 None。"""
|
||
ready = ready_tasks(list(self._tasks.values()), completed_ids)
|
||
return ready[0] if ready else None
|
||
|
||
def _get(self, task_id):
|
||
try:
|
||
return self._tasks[task_id]
|
||
except KeyError:
|
||
raise ValueError(f"未知任务：{task_id}") from None
|