YG_FT/backend/app/modules/compute_gateway/sync.py

from __future__ import annotations

import json
import time
from typing import Any

from app.db.platform_store import get_platform_store
from app.modules.compute_gateway.client import ComputeNodeClient

# starting 状态允许的最大轮询次数（约 40 * 3s ≈ 2 分钟），超过即判定节点不可达
MAX_STARTING_ATTEMPTS = 40


def _node_for_task(task: dict[str, Any]) -> dict[str, Any] | None:
    return next((node for node in get_platform_store().compute_nodes() if node["id"] == task.get("compute_node_id")), None)


def _parse_inference_load_status(task: dict[str, Any]) -> tuple[list[dict[str, Any]], dict[str, Any]]:
    load_status = task.get("load_status") or {}
    if isinstance(load_status, str):
        try:
            load_status = json.loads(load_status)
        except (json.JSONDecodeError, TypeError):
            load_status = {}
    return load_status.get("loaded_models") or [], load_status


async def reconcile_inference_loads(store: Any) -> list[dict[str, Any]]:
    """推进处于 starting 状态的推理加载。

    模型加载已改为异步派发：/model-compare/{id}/load 立即返回，这里在每次
    轮询时查询对应计算节点的 /inference/status，把任务从 starting 推进到
    ready/error。使用短超时，单节点不可达不会阻塞整轮轮询。
    """
    reconciled: list[dict[str, Any]] = []
    now = time.time()
    for task in store.compare_tasks():
        items, _ = _parse_inference_load_status(task)
        if not any(item.get("status") == "starting" for item in items):
            continue
        # dirty 只要处理过任一 starting 项就置位：load_attempts / last_polled_at
        # 必须落库，否则节点不可达时计数不会累积，封顶逻辑永远触发不了
        dirty = False
        for item in items:
            if item.get("status") != "starting":
                continue
            # 节流：同一 item 每 3s 只查询一次
            if now - float(item.get("last_polled_at") or 0) < 3:
                continue
            item["last_polled_at"] = now
            item["load_attempts"] = int(item.get("load_attempts") or 0) + 1
            dirty = True
            node = next((n for n in store.compute_nodes() if n["id"] == item.get("node_id")), None)
            if not node:
                item["status"] = "error"
                item["error"] = "compute node deleted"
                store.mark_inference_unloaded(item.get("node_id") or "")
                continue
            if not node.get("enabled") or node.get("scheduler_status") != "online":
                item["status"] = "error"
                item["error"] = "compute node offline"
                store.mark_inference_unloaded(node["id"])
                continue
            try:
                status = await ComputeNodeClient(node["api_base_url"]).inference_status()
            except Exception as exc:  # noqa: BLE001 - node unreachable; keep retrying until cap
                if int(item.get("load_attempts") or 0) >= MAX_STARTING_ATTEMPTS:
                    item["status"] = "error"
                    item["error"] = f"compute node unreachable: {exc}"
                    store.mark_inference_unloaded(node["id"])
                continue
            node_status = status.get("status")
            if node_status == "ready":
                item["status"] = "ready"
                item.pop("error", None)
                store.mark_inference_loaded(node["id"])
            elif node_status == "error":
                item["status"] = "error"
                item["error"] = status.get("error") or "model load failed on compute node"
                store.mark_inference_unloaded(node["id"])
            elif node_status == "idle":
                # 节点重启导致已加载模型丢失
                item["status"] = "error"
                item["error"] = "model disappeared from compute node (node may have restarted)"
                store.mark_inference_unloaded(node["id"])
            # node_status == "loading" -> 保持 starting，下轮再查
        if dirty:
            if any(i.get("status") in {"ready", "running"} for i in items):
                new_status = "loaded"
            elif any(i.get("status") == "starting" for i in items):
                new_status = "starting"  # 仍在加载中，保持 starting
            else:
                new_status = "failed"
            store.update_compare_task(task["id"], {"status": new_status, "load_status": {"loaded_models": items}})
            reconciled.append({"task_id": task["id"], "status": new_status})
    return reconciled


async def fetch_eval_result_content(client: ComputeNodeClient, node: dict[str, Any], job: dict[str, Any]) -> dict[str, Any] | None:
    output_dir = job.get("output_dir")
    if not output_dir:
        return None
    full_path = f"{str(output_dir).rstrip('/')}/eval_results.json"
    data_root = "/data/yg-ft/"
    if full_path.startswith(data_root):
        full_path = full_path[len(data_root):]
    rel_path = full_path.lstrip("/")
    import httpx
    url = f"{node['api_base_url'].rstrip('/')}/modelTF/compute/files/read"
    async with httpx.AsyncClient(timeout=30, headers=client.headers()) as http:
        response = await http.get(url, params={"path": rel_path})
        response.raise_for_status()
        payload = response.json()
        return payload if isinstance(payload, dict) else None


async def poll_compute_jobs_once() -> dict[str, Any]:
    store = get_platform_store()
    synced: list[dict[str, Any]] = []
    failed: list[dict[str, str]] = []
    for task in store.running_compute_tasks():
        node = _node_for_task(task)
        if not node:
            failed.append({"task_id": task["id"], "error": "compute node not found"})
            continue
        try:
            client = ComputeNodeClient(node["api_base_url"])
            job = await client.get_job(task["compute_job_id"])
            try:
                logs = await client.job_logs(task["compute_job_id"], tail_lines=5000)
                store.record_training_log_metrics(task["id"], str(logs.get("content") or ""))
            except Exception:
                pass
            # P0-4: Force-fetch last log snippet when job reaches terminal state
            if job.get("status") in {"failed", "stopped"}:
                try:
                    last_logs = await client.job_logs(task["compute_job_id"], tail_lines=200)
                    job["log_snippet"] = str(last_logs.get("content") or "")[:8192]
                except Exception:
                    pass
            synced.append(store.apply_compute_job(task["id"], job))
        except Exception as exc:  # noqa: BLE001 - keep polling other jobs
            failed.append({"task_id": task["id"], "error": str(exc)})
    standalone_synced: list[dict[str, Any]] = []
    for record in store.active_standalone_compute_jobs():
        node = next((item for item in store.compute_nodes() if item["id"] == record.get("node_id")), None)
        if not node:
            failed.append({"job_id": record["id"], "error": "compute node not found"})
            continue
        try:
            job = await ComputeNodeClient(node["api_base_url"]).get_job(record["id"])
            standalone_synced.append(store.sync_model_merge_job(record["id"], job))
        except Exception as exc:  # noqa: BLE001 - keep polling other jobs
            failed.append({"job_id": record["id"], "error": str(exc)})

    # ── Eval job sync ────────────────────────────────────────────────
    eval_synced = 0
    for eval_task in store.running_eval_tasks():
        node = next(
            (item for item in store.compute_nodes() if item["id"] == eval_task.get("compute_node_id")),
            None,
        )
        if not node:
            failed.append({"eval_task_id": eval_task["id"], "error": "compute node not found"})
            continue
        try:
            client = ComputeNodeClient(node["api_base_url"])
            job = await client.get_job(eval_task["compute_job_id"])
            result_content = None
            # Try to read eval_results.json from the job output directory
            if job.get("status") == "completed" and job.get("output_dir"):
                try:
                    result_content = await fetch_eval_result_content(client, node, job)
                except Exception:
                    pass
            store.apply_eval_job_result(eval_task["id"], job, result_content)
            # If job completed, un-mark inference loaded
            if job.get("status") in {"completed", "failed", "stopped"}:
                store.mark_inference_unloaded(node["id"])
            eval_synced += 1
        except Exception as exc:  # noqa: BLE001
            failed.append({"eval_task_id": eval_task["id"], "error": str(exc)})

    # ── Inference load reconciliation ─────────────────────────────────────
    try:
        inference_reconciled = await reconcile_inference_loads(store)
    except Exception as exc:  # noqa: BLE001 - keep polling alive
        failed.append({"inference_reconcile": str(exc)})
        inference_reconciled = []

    return {"synced": len(synced) + len(standalone_synced) + eval_synced, "failed": failed,
            "items": synced, "standalone": standalone_synced, "eval_synced": eval_synced,
            "inference_reconciled": inference_reconciled}