252 lines
9.7 KiB
Python
252 lines
9.7 KiB
Python
"""工作流调度器。
|
|
|
|
轮询 SQLite 中的排队任务，按工作流 DAG 的拓扑顺序依次调用节点，并把节点
|
|
产物登记为任务产物。MVP 使用进程内单线程顺序执行，后续可替换为分布式队列。
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import threading
|
|
import time
|
|
import uuid
|
|
from datetime import datetime, timezone
|
|
from pathlib import Path
|
|
|
|
from wov_sdk.models import InvokeRequest, WorkflowDefinition
|
|
|
|
from app.db import Database
|
|
from app.node_manager import NodeManager
|
|
|
|
|
|
def _now_iso() -> str:
|
|
"""返回当前 UTC 时间的 ISO 格式字符串。"""
|
|
return datetime.now(timezone.utc).isoformat()
|
|
|
|
|
|
def topological_sort(definition: WorkflowDefinition) -> list[str]:
|
|
"""对工作流 DAG 做拓扑排序，返回可执行的节点 ID 顺序。"""
|
|
nodes = {node.id: node for node in definition.nodes}
|
|
# 统计每个节点的入度，并记录依赖关系。
|
|
indegree = {node_id: 0 for node_id in nodes}
|
|
dependents: dict[str, list[str]] = {node_id: [] for node_id in nodes}
|
|
|
|
for edge in definition.edges:
|
|
# 边引用了不存在的节点时直接报错。
|
|
if edge.from_node not in nodes or edge.to_node not in nodes:
|
|
raise ValueError(f"unknown edge: {edge.from_node} -> {edge.to_node}")
|
|
indegree[edge.to_node] += 1
|
|
dependents[edge.from_node].append(edge.to_node)
|
|
|
|
# Kahn 算法：从入度为 0 的节点开始逐层取出。
|
|
queue = [node_id for node_id, degree in indegree.items() if degree == 0]
|
|
ordered: list[str] = []
|
|
while queue:
|
|
current = queue.pop(0)
|
|
ordered.append(current)
|
|
for dependent in dependents[current]:
|
|
indegree[dependent] -= 1
|
|
if indegree[dependent] == 0:
|
|
queue.append(dependent)
|
|
|
|
# 排序结果数量不足说明存在环，无法确定执行顺序。
|
|
if len(ordered) != len(nodes):
|
|
raise ValueError("workflow contains a cycle")
|
|
return ordered
|
|
|
|
|
|
class WorkflowScheduler:
|
|
"""后台任务调度器：单线程轮询并执行排队中的工作流运行。"""
|
|
|
|
def __init__(
|
|
self,
|
|
db: Database,
|
|
node_manager: NodeManager,
|
|
storage_dir: Path,
|
|
interval_seconds: float = 1.0,
|
|
) -> None:
|
|
"""保存依赖并初始化轮询线程控制字段。"""
|
|
self.db = db
|
|
self.node_manager = node_manager
|
|
self.storage_dir = storage_dir
|
|
self.interval_seconds = interval_seconds
|
|
self._thread: threading.Thread | None = None
|
|
self._stopping = False
|
|
|
|
def start(self) -> None:
|
|
"""启动调度线程；重复调用无副作用。"""
|
|
if self._thread is not None:
|
|
return
|
|
self._stopping = False
|
|
self._thread = threading.Thread(
|
|
target=self._loop,
|
|
name="wov-workflow-scheduler",
|
|
daemon=True,
|
|
)
|
|
self._thread.start()
|
|
|
|
def stop(self) -> None:
|
|
"""请求停止并等待轮询线程退出。"""
|
|
self._stopping = True
|
|
if self._thread is not None:
|
|
self._thread.join(timeout=5)
|
|
self._thread = None
|
|
|
|
def _loop(self) -> None:
|
|
"""轮询循环：有排队任务就立即执行，否则休眠一个间隔。"""
|
|
while not self._stopping:
|
|
run = self.db.next_queued_run()
|
|
if run is not None:
|
|
self.execute_run(run["id"])
|
|
else:
|
|
time.sleep(self.interval_seconds)
|
|
|
|
def _resolve_ref(
|
|
self,
|
|
ref: str,
|
|
run_input_uri: str | None,
|
|
outputs_by_node: dict[str, dict[str, str]],
|
|
) -> str | None:
|
|
"""解析输入引用：input.xxx 取任务入口，node.key 取前序节点产物。"""
|
|
# 入口引用以 input. 为前缀。
|
|
if ref.startswith("input."):
|
|
return run_input_uri
|
|
# 其余引用必须形如 "节点ID.输出名"。
|
|
node_id, separator, key = ref.partition(".")
|
|
if not separator:
|
|
return None
|
|
return outputs_by_node.get(node_id, {}).get(key)
|
|
|
|
def execute_run(self, run_id: str) -> None:
|
|
"""执行单个任务：加载 DAG、按拓扑顺序调用节点并登记产物。"""
|
|
run = self.db.get_run(run_id)
|
|
# 任务不存在或不在排队状态时直接返回，避免重复执行。
|
|
if run is None or run["status"] != "QUEUED":
|
|
return
|
|
|
|
# 工作流或版本记录丢失时把任务标记为失败。
|
|
workflow = self.db.get_workflow(run["workflow_id"])
|
|
if workflow is None:
|
|
self.db.update_run(run_id, status="FAILED", error="workflow not found", updated_at=_now_iso())
|
|
return
|
|
|
|
version = self.db.get_workflow_version(run["workflow_id"], run["workflow_version"])
|
|
if version is None:
|
|
self.db.update_run(run_id, status="FAILED", error="workflow version not found", updated_at=_now_iso())
|
|
return
|
|
|
|
# 解析并校验 DAG，随后计算拓扑执行顺序。
|
|
definition = WorkflowDefinition.from_dict(version["definition"])
|
|
definition.validate()
|
|
ordered = topological_sort(definition)
|
|
outputs_by_node: dict[str, dict[str, str]] = {}
|
|
self.db.update_run(run_id, status="RUNNING", progress=0, updated_at=_now_iso())
|
|
|
|
try:
|
|
for index, node_id in enumerate(ordered):
|
|
# 当前节点进度 = 已完成节点数 / 总节点数。
|
|
node_spec = next(item for item in definition.nodes if item.id == node_id)
|
|
self.db.update_run(
|
|
run_id,
|
|
current_node_id=node_id,
|
|
progress=index / len(ordered),
|
|
updated_at=_now_iso(),
|
|
)
|
|
# 解析节点声明的每个输入引用，缺任一输入即失败。
|
|
invoke_inputs: dict[str, str] = {}
|
|
for input_name, ref in node_spec.inputs.items():
|
|
value = self._resolve_ref(ref, run.get("input_uri"), outputs_by_node)
|
|
if value is None:
|
|
raise ValueError(f"missing input {input_name} for node {node_id}")
|
|
invoke_inputs[input_name] = value
|
|
|
|
# 每个任务的每个节点使用独立产物目录，避免并发冲突。
|
|
output_dir = (
|
|
self.storage_dir
|
|
/ "runs"
|
|
/ run_id
|
|
/ "steps"
|
|
/ node_id
|
|
)
|
|
response = self.node_manager.invoke(
|
|
node_spec.node_type,
|
|
InvokeRequest(
|
|
run_id=run_id,
|
|
node_instance_id="",
|
|
inputs=invoke_inputs,
|
|
params=node_spec.params,
|
|
output_dir=str(output_dir),
|
|
),
|
|
)
|
|
# 节点返回非 completed 即视为步骤失败。
|
|
if response.status != "completed":
|
|
raise RuntimeError(response.error or f"node {node_id} failed")
|
|
|
|
# 记录节点输出，供后续节点引用和最终产物映射使用。
|
|
outputs_by_node[node_id] = {
|
|
str(key): str(value) for key, value in response.outputs.items()
|
|
}
|
|
for key, uri in outputs_by_node[node_id].items():
|
|
# 产物名带节点前缀，例如 asr.srt_uri，避免跨节点重名。
|
|
artifact = {
|
|
"run_id": run_id,
|
|
"node_id": node_id,
|
|
"name": f"{node_id}.{key}",
|
|
"uri": uri,
|
|
"mime_type": self._mime_type(uri),
|
|
"size": self._file_size(uri),
|
|
}
|
|
self.db.create_artifact(artifact)
|
|
|
|
# 处理 final_outputs，为用户端提供简洁的下载别名。
|
|
for alias, ref in definition.final_outputs.items():
|
|
resolved = self._resolve_ref(ref, run.get("input_uri"), outputs_by_node)
|
|
if resolved is not None:
|
|
self.db.create_artifact(
|
|
{
|
|
"run_id": run_id,
|
|
"node_id": ref.partition(".")[0],
|
|
"name": alias,
|
|
"uri": resolved,
|
|
"mime_type": self._mime_type(resolved),
|
|
"size": self._file_size(resolved),
|
|
}
|
|
)
|
|
|
|
# 全部节点成功后任务标记为完成。
|
|
self.db.update_run(
|
|
run_id,
|
|
status="COMPLETED",
|
|
current_node_id=None,
|
|
progress=1.0,
|
|
updated_at=_now_iso(),
|
|
)
|
|
except Exception as exc: # noqa: BLE001
|
|
# 任一步骤异常都结束任务并记录错误，等待用户重试。
|
|
self.db.update_run(
|
|
run_id,
|
|
status="FAILED",
|
|
error=str(exc),
|
|
updated_at=_now_iso(),
|
|
)
|
|
|
|
@staticmethod
|
|
def _mime_type(uri: str) -> str:
|
|
"""按扩展名推断产物 MIME 类型，未知类型使用通用二进制类型。"""
|
|
path = Path(uri)
|
|
suffix = path.suffix.lower()
|
|
return {
|
|
".srt": "application/x-subrip",
|
|
".ass": "text/plain",
|
|
".wav": "audio/wav",
|
|
".mp4": "video/mp4",
|
|
".txt": "text/plain",
|
|
}.get(suffix, "application/octet-stream")
|
|
|
|
@staticmethod
|
|
def _file_size(uri: str) -> int:
|
|
"""读取产物文件大小；文件缺失时按 0 处理。"""
|
|
try:
|
|
return Path(uri).stat().st_size
|
|
except OSError:
|
|
return 0
|