Spaces:
Running
Running
| from __future__ import annotations | |
| import logging | |
| import shutil | |
| import tempfile | |
| import time | |
| import zipfile | |
| from pathlib import Path | |
| from typing import Callable | |
| from app.config import settings | |
| from app.processing import JobContext, ModuleSkip, ArtifactRef | |
| from app.processing.registry import get as get_module | |
| from app.workflow_types import migrate_graph_types | |
| from worker.client import WorkerClient | |
| logger = logging.getLogger("egodata.worker") | |
| def _topo_sort(nodes: list[dict], edges: list[dict]) -> list[str]: | |
| degrees = {node.get("id"): 0 for node in nodes if node.get("id")} | |
| children = {node_id: [] for node_id in degrees} | |
| for edge in edges: | |
| source, target = edge.get("source"), edge.get("target") | |
| if source in children and target in degrees: | |
| children[source].append(target) | |
| degrees[target] += 1 | |
| ready = [node_id for node_id, degree in degrees.items() if degree == 0] | |
| order: list[str] = [] | |
| while ready: | |
| current = ready.pop(0) | |
| order.append(current) | |
| for child in children[current]: | |
| degrees[child] -= 1 | |
| if degrees[child] == 0: | |
| ready.append(child) | |
| return order + [node_id for node_id in degrees if node_id not in order] | |
| def _safe_extract(archive_path: Path, destination: Path) -> None: | |
| destination.mkdir(parents=True, exist_ok=True) | |
| root = destination.resolve() | |
| with zipfile.ZipFile(archive_path) as archive: | |
| for info in archive.infolist(): | |
| relative = Path(info.filename) | |
| if relative.is_absolute() or ".." in relative.parts: | |
| raise ValueError("Unsafe input archive path") | |
| target = (destination / relative).resolve() | |
| if root != target and root not in target.parents: | |
| raise ValueError("Unsafe input archive path") | |
| if info.is_dir(): | |
| target.mkdir(parents=True, exist_ok=True) | |
| else: | |
| target.parent.mkdir(parents=True, exist_ok=True) | |
| with archive.open(info) as source, target.open("wb") as output: | |
| shutil.copyfileobj(source, output, length=1024 * 1024) | |
| def _zip_directory(source: Path, destination: Path) -> None: | |
| with zipfile.ZipFile(destination, "w", allowZip64=True) as archive: | |
| for path in source.rglob("*"): | |
| if path.is_file(): | |
| compression = zipfile.ZIP_STORED if path.suffix.lower() in {".mp4", ".parquet"} else zipfile.ZIP_DEFLATED | |
| archive.write(path, str(path.relative_to(source)), compress_type=compression) | |
| def _incoming_artifacts(node_id: str, edges: list[dict], node_outputs: dict[str, dict[str, ArtifactRef]]) -> dict[str, ArtifactRef]: | |
| """Aggregate upstream artifacts for a node. | |
| 三条规则(单边/单输出行为与旧实现逐位一致): | |
| 1. 同一 targetHandle 多条边 → 后缀去重(video → video#1、video#2…), | |
| 解决"双目左右目两条边指向同一识别节点,后边覆盖前边"的问题; | |
| 2. 命中 sourceHandle 时,同源 "#" 兄弟键一并并入 | |
| (mediapipe_hand 多路输出 hand_keypoints / hand_keypoints#1 的场景); | |
| 3. sourceHandle 未命中(透传链,如 annotation/reviewed 在透传输出中 | |
| 不存在)→ 平铺 source 的全部输出,而不是只取第一个。 | |
| """ | |
| inputs: dict[str, ArtifactRef] = {} | |
| for edge in edges: | |
| if edge.get("target") != node_id: | |
| continue | |
| source_outputs = node_outputs.get(edge.get("source"), {}) | |
| source_handle = edge.get("sourceHandle") | |
| target_handle = edge.get("targetHandle") or "data" | |
| if source_handle and source_handle in source_outputs: | |
| # 规则 2:主键 + "#" 兄弟键(如 hand_keypoints, hand_keypoints#1) | |
| refs = [source_outputs[source_handle]] + [ | |
| r for k, r in source_outputs.items() | |
| if k.startswith(source_handle + "#")] | |
| elif len(source_outputs) == 1: | |
| refs = [next(iter(source_outputs.values()))] # 现状行为 | |
| elif source_outputs: | |
| # 规则 3:透传链平铺全部输出 | |
| refs = list(source_outputs.values()) | |
| else: | |
| continue | |
| seen: set[tuple] = set() | |
| for ref in refs: | |
| dedup_key = (ref.kind, str(ref.path), ref.source_key) | |
| if dedup_key in seen: | |
| continue | |
| seen.add(dedup_key) | |
| # 规则 1:同 targetHandle 冲突 → 后缀去重,边序确定性 | |
| key, n = target_handle, 1 | |
| while key in inputs: | |
| key = f"{target_handle}#{n}" | |
| n += 1 | |
| inputs[key] = ref | |
| return inputs | |
| def execute_job( | |
| job: dict, | |
| input_root: Path, | |
| output_root: Path, | |
| heartbeat: Callable[[float, dict], None] | None = None, | |
| ) -> tuple[dict, dict]: | |
| graph, _ = migrate_graph_types(job.get("graph") or {}) | |
| nodes = graph.get("nodes") or [] | |
| edges = graph.get("edges") or [] | |
| node_map = {node.get("id"): node for node in nodes} | |
| node_outputs: dict[str, dict[str, ArtifactRef]] = {} | |
| node_states: dict[str, dict] = {} | |
| output_manifest: dict = {"artifacts": {}} | |
| order = _topo_sort(nodes, edges) | |
| for index, node_id in enumerate(order): | |
| node = node_map.get(node_id) | |
| if not node: | |
| continue | |
| data = node.get("data") or {} | |
| node_type = data.get("nodeType", "") | |
| config = dict(data.get("config") or {}) | |
| incoming = _incoming_artifacts(node_id, edges, node_outputs) | |
| node_states[node_id] = {"type": node_type, "status": "running", "progress": 0.0} | |
| if heartbeat: | |
| heartbeat(index / max(1, len(order)), node_states) | |
| # 统一分发:按 node_type 查模块注册表执行(可插拔,新增模块无需改这里) | |
| module = get_module(node_type) | |
| if module is None: | |
| raise ValueError(f"Unsupported worker module: {node_type}") | |
| def node_progress(p: float): | |
| if heartbeat: | |
| heartbeat((index + p) / max(1, len(order)), node_states) | |
| ctx = JobContext( | |
| node_id=node_id, node_type=node_type, config=config, job=job, | |
| input_root=input_root, output_root=output_root, | |
| incoming=incoming, progress=node_progress, node_states=node_states, | |
| ) | |
| try: | |
| outputs = module.run(ctx) | |
| except ModuleSkip as exc: | |
| # 模块主动跳过(无有效输入)—— 标记 skipped,不失败,继续后续节点 | |
| node_states[node_id] = {"type": node_type, "status": "skipped", | |
| "note": exc.reason} | |
| continue | |
| if outputs is None: | |
| node_states[node_id] = {"type": node_type, "status": "skipped"} | |
| continue | |
| node_outputs[node_id] = outputs | |
| output_manifest["artifacts"][node_id] = {key: ref.to_dict() for key, ref in outputs.items()} | |
| warnings = [] | |
| for ref in outputs.values(): | |
| for warning in (ref.metadata or {}).get("processing_warnings") or []: | |
| if warning and warning not in warnings: | |
| warnings.append(str(warning)) | |
| state = {"type": node_type, "status": "completed", "progress": 1.0} | |
| if warnings: | |
| state["warnings"] = warnings | |
| node_states[node_id] = state | |
| if heartbeat: | |
| heartbeat((index + 1) / max(1, len(order)), node_states) | |
| return node_states, output_manifest | |
| def process_one(client: WorkerClient, job: dict, worker_id: str, work_dir: Path) -> None: | |
| run_id = str(job["run_id"]) | |
| lease_token = str(job.get("lease_token") or "") | |
| # ignore_cleanup_errors:Windows 上 OpenCV 句柄可能短暂占用文件, | |
| # 清理失败不应让已完成的 job 报错。 | |
| with tempfile.TemporaryDirectory(prefix=f"egodata-{run_id}-", dir=str(work_dir), | |
| ignore_cleanup_errors=True) as temp: | |
| root = Path(temp) | |
| input_zip = root / "input.zip" | |
| input_root = root / "input" | |
| output_root = root / "outputs" | |
| client.download_input(run_id, input_zip) | |
| _safe_extract(input_zip, input_root) | |
| def heartbeat(progress: float, states: dict): | |
| # 心跳失败不中断处理(后端热加载/瞬时故障时继续跑), | |
| # 否则 MediaPipe 进度回调异常会中断整个 job。 | |
| try: | |
| client.heartbeat(run_id, worker_id, lease_token, | |
| min(0.99, progress), states) | |
| except Exception as exc: | |
| # A 409 means this worker lost the lease to a newer attempt. | |
| # Continuing would waste GPU time and could let a stale worker | |
| # submit results after the newer attempt has started. | |
| if getattr(getattr(exc, "response", None), "status_code", None) == 409: | |
| raise RuntimeError("Worker lease lost; aborting stale attempt") from exc | |
| logger.warning("Heartbeat failed (continuing): %s", exc) | |
| states, manifest = execute_job(job, input_root, output_root, heartbeat) | |
| result_zip = root / "outputs.zip" | |
| _zip_directory(output_root, result_zip) | |
| client.complete(run_id, worker_id, lease_token, result_zip, states, manifest) | |
| def run_forever(server_url: str, api_key: str, worker_id: str, device: str = "auto", poll_seconds: float = 2.0, work_dir: Path | None = None): | |
| work_dir = work_dir or settings.temp_root / "worker" | |
| work_dir.mkdir(parents=True, exist_ok=True) | |
| client = WorkerClient(server_url, api_key) | |
| try: | |
| while True: | |
| job = None | |
| try: | |
| job = client.claim(worker_id, ["mediapipe", "hand_keypoints", "video_overlay"], device) | |
| if job is None: | |
| time.sleep(poll_seconds) | |
| continue | |
| logger.info("Claimed workflow run %s", job["run_id"]) | |
| process_one(client, job, worker_id, work_dir) | |
| logger.info("Completed workflow run %s", job["run_id"]) | |
| except KeyboardInterrupt: | |
| raise | |
| except Exception as exc: | |
| logger.exception("Worker job failed: %s", exc) | |
| if "job" in locals() and job: | |
| try: | |
| client.fail(str(job["run_id"]), worker_id, | |
| str(job.get("lease_token") or ""), | |
| str(exc), retry=True) | |
| except Exception: | |
| logger.exception("Failed to report worker failure") | |
| time.sleep(poll_seconds) | |
| finally: | |
| client.close() | |