feat: 完成工作流调度与上传下载 API
This commit is contained in:
cat-shark
@@ -0,0 +1,215 @@
|
||||
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:
|
||||
return datetime.now(timezone.utc).isoformat()
|
||||
|
||||
|
||||
def topological_sort(definition: WorkflowDefinition) -> list[str]:
|
||||
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)
|
||||
|
||||
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:
|
||||
if ref.startswith("input."):
|
||||
return run_input_uri
|
||||
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:
|
||||
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
|
||||
|
||||
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),
|
||||
),
|
||||
)
|
||||
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():
|
||||
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)
|
||||
|
||||
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:
|
||||
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:
|
||||
try:
|
||||
return Path(uri).stat().st_size
|
||||
except OSError:
|
||||
return 0
|
||||
Reference in New Issue
Block a user