Blame for processor/worker/runner.py

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()