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

from __future__ import annotations

import json
import time
from pathlib import Path
from typing import Any

from app.db.platform_store import get_platform_store
from app.core.config import get_settings
from app.modules.compute_gateway.client import ComputeNodeClient
from app.modules.storage.minio_store import get_object_storage

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


async def _archive_node_directory(
    store: Any,
    client: ComputeNodeClient,
    node: dict[str, Any],
    source_path: str,
    resource_type: str,
    resource_id: str,
    version_id: str,
    object_prefix: str,
) -> list[dict[str, Any]]:
    """Archive a completed node directory to MinIO, preserving subdirectories."""
    data_root = Path(str(node.get("data_root") or "/data/yg-ft")).resolve()
    source = Path(source_path).resolve()
    try:
        relative_root = source.relative_to(data_root).as_posix()
    except ValueError as exc:
        raise RuntimeError(f"artifact path is outside compute data root: {source_path}") from exc
    queue = [relative_root]
    archived: list[dict[str, Any]] = []
    while queue:
        relative = queue.pop(0)
        listing = await client.list_files(root="data", relative_path=relative)
        for item in listing.get("items") or []:
            item_relative = str(item.get("relative_path") or "")
            if item.get("type") == "directory":
                queue.append(item_relative)
                continue
            path = str(item.get("path") or "")
            if not path:
                continue
            try:
                relative_file = Path(item_relative).relative_to(Path(relative_root)).as_posix()
            except ValueError:
                relative_file = Path(str(item.get("name") or Path(path).name)).name
            object_key = f"{object_prefix}/{version_id}/{relative_file}"
            upload_url = get_object_storage().presigned_put(object_key)
            result = await client.upload_file_to_url(path, upload_url, object_key)
            metadata = get_object_storage().stat(object_key)
            archived.append(
                store.create_storage_object(
                    {
                        "resource_type": resource_type,
                        "resource_id": resource_id,
                        "version_id": version_id,
                        "bucket": get_object_storage().bucket,
                        "object_key": object_key,
                        "file_name": relative_file,
                        "content_type": "application/octet-stream",
                        "byte_size": metadata.get("byte_size") or result.get("byte_size") or 0,
                        "checksum_sha256": result.get("checksum_sha256") or "",
                        "status": "available",
                    }
                )
            )
    return archived


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)
                selected_gpus = item.get("gpu_indices") or item.get("gpus")
                store.mark_inference_loaded(node["id"], selected_gpus)
            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
            updated_task = store.apply_compute_job(task["id"], job)
            if (
                get_settings().minio_enabled
                and job.get("status") == "completed"
                and job.get("output_dir")
            ):
                trained_model = next(
                    (
                        item
                        for item in store.trained_models()
                        if item.get("name")
                        == (task.get("output_model_name") or f"{task.get('name')}-lora")
                    ),
                    None,
                )
                if trained_model:
                    archived = await _archive_node_directory(
                        store,
                        client,
                        node,
                        str(job["output_dir"]),
                        "trained_model",
                        str(trained_model["id"]),
                        str(job.get("id") or task.get("compute_job_id") or task["id"]),
                        f"trained_models/{trained_model['id']}",
                    )
                    artifacts = store.model_artifacts(str(trained_model["id"]))
                    if archived and artifacts:
                        store.link_model_artifact_storage_object(
                            str(artifacts[0]["id"]), str(archived[0]["id"])
                        )
            synced.append(updated_task)
        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))
            if get_settings().minio_enabled and job.get("status") == "completed" and job.get("output_dir"):
                payload = (store.compute_job(record["id"]).get("payload") or {})
                trained_model_id = str(payload.get("trained_model_id") or payload.get("model_name") or "")
                if trained_model_id:
                    trained_model = next(
                        (item for item in store.trained_models() if item.get("id") == trained_model_id or item.get("name") == trained_model_id),
                        None,
                    )
                    if trained_model:
                        archived = await _archive_node_directory(
                            store,
                            ComputeNodeClient(node["api_base_url"], timeout=900),
                            node,
                            str(job["output_dir"]),
                            "trained_model",
                            str(trained_model["id"]),
                            str(job.get("id") or record["id"]),
                            f"trained_models/{trained_model['id']}",
                        )
                        artifacts = store.model_artifacts(str(trained_model["id"]))
                        if archived and artifacts:
                            store.link_model_artifact_storage_object(
                                str(artifacts[0]["id"]), str(archived[0]["id"])
                            )
        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 get_settings().minio_enabled and job.get("status") == "completed" and job.get("output_dir"):
                await _archive_node_directory(
                    store,
                    client,
                    node,
                    str(job["output_dir"]),
                    "eval",
                    str(eval_task["id"]),
                    str(job.get("id") or eval_task.get("compute_job_id") or eval_task["id"]),
                    f"evaluations/{eval_task['id']}",
                )
            # 评测 GPU 占用由 eval_tasks 状态派生，无需维护推理内存标记
            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}