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

from __future__ import annotations

from typing import Any

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


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)


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

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