"""工作流调度器。 轮询 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