YG_FT/compute/api/main.py

from __future__ import annotations

import os
import math
import time
from pathlib import Path
from typing import Any

from fastapi import FastAPI, HTTPException

from compute.engines.llama_factory.adapter import build_command, parse_log_line


def create_app() -> FastAPI:
    app = FastAPI(title="YG Fine-Tune Compute API")
    jobs: dict[str, dict[str, Any]] = {}
    route_prefix = os.getenv("MODELTF_ROUTE_PREFIX", "/modelTF").rstrip("/") or "/modelTF"

    def now() -> float:
        return time.time()

    def host_id() -> str:
        return os.getenv("COMPUTE_HOST_ID", "gpu-node-01")

    def job_status(job: dict[str, Any]) -> dict[str, Any]:
        elapsed = max(0, int(now() - job["created_at"]))
        if job["status"] not in {"stopped", "failed", "completed"}:
            if elapsed < 5:
                job["status"] = "queued"
                job["progress"] = 12 + elapsed * 3
            elif elapsed < 60:
                job["status"] = "running"
                job["progress"] = min(96, 25 + int((elapsed - 5) / 55 * 70))
            else:
                job["status"] = "completed"
                job["progress"] = 100
        job["logs"] = generate_logs(job)
        return job

    def generate_logs(job: dict[str, Any]) -> str:
        progress = int(job.get("progress", 0) or 0)
        points = max(1, min(80, progress))
        lines = [
            f"[INFO] compute_host_id={host_id()} job_id={job['id']} engine=llama_factory mode=simulator",
            f"[INFO] command={' '.join(job['command'])}",
        ]
        for step in range(1, points + 1):
            if step % 4 != 0 and step != points:
                continue
            loss = max(0.11, 2.5 * math.exp(-step / 40))
            grad_norm = 0.4 + (step % 5) * 0.04
            lr = 0.0002 * max(0.05, 1 - step / 100)
            epoch = round(step / points * 3, 4)
            lines.append(
                "{"
                f"'loss': {loss:.4f}, 'grad_norm': {grad_norm:.4f}, "
                f"'learning_rate': {lr:.8f}, 'epoch': {epoch:.4f}"
                "}"
            )
        if job.get("status") == "completed":
            lines.extend(
                [
                    "***** train metrics *****",
                    "epoch = 3",
                    "train_loss = 0.1181",
                    "train_runtime = 1m 0s",
                    "***** train metrics end *****",
                ]
            )
        return "\n".join(lines)

    def gpu_resources() -> list[dict[str, Any]]:
        active_jobs = [job_status(job) for job in jobs.values() if job["status"] in {"queued", "running"}]
        gpus: list[dict[str, Any]] = []
        for idx in range(4):
            task = next((job for job in active_jobs if idx in job.get("gpus", [])), None)
            busy = task is not None and task["status"] == "running"
            reserved = task is not None and task["status"] == "queued"
            gpus.append(
                {
                    "id": idx,
                    "uuid": f"GPU-{host_id().upper()}-{idx}",
                    "name": os.getenv("COMPUTE_GPU_NAME", "NVIDIA A800-SXM4-80GB"),
                    "status": "busy" if busy else "reserved" if reserved else "idle",
                    "gpu_percent": 88 if busy else 25 if reserved else 4,
                    "memory_used_gb": 58 if busy else 12 if reserved else 2,
                    "memory_total_gb": 80,
                    "temperature": 61 if busy else 45 if reserved else 36,
                    "power_w": 215 if busy else 80 if reserved else 25,
                    "power_limit_w": 300,
                    "processes": [
                        {
                            "pid": task["pid"],
                            "name": "llamafactory-cli",
                            "task_name": task["name"],
                            "memory_used_gb": 58 if busy else 12,
                        }
                    ]
                    if task
                    else [],
                }
            )
        return gpus

    @app.get(f"{route_prefix}/health")
    async def health_check() -> dict[str, str]:
        return {
            "status": "ok",
            "compute_host_id": os.getenv("COMPUTE_HOST_ID", "unknown"),
        }

    @app.get(f"{route_prefix}/v1/compute/health")
    async def compute_health_check() -> dict[str, str | bool]:
        data_root = Path(os.getenv("YG_FT_DATA_ROOT", "/data/yg-ft"))
        llama_factory_home = Path(os.getenv("LLAMA_FACTORY_HOME", "/app/LLaMA-Factory"))
        return {
            "status": "ok",
            "compute_host_id": os.getenv("COMPUTE_HOST_ID", "unknown"),
            "app_callback_enabled": os.getenv("ENABLE_APP_CALLBACK", "false").lower() == "true",
            "data_root": str(data_root),
            "data_root_exists": data_root.exists(),
            "llama_factory_home": str(llama_factory_home),
            "llama_factory_home_exists": llama_factory_home.exists(),
        }

    @app.get(f"{route_prefix}/v1/compute/jobs")
    async def list_jobs_alias() -> dict[str, list[dict[str, Any]]]:
        return {"items": [job_status(job) for job in jobs.values()]}

    @app.get(f"{route_prefix}/compute/resources/gpus")
    async def list_gpus() -> dict[str, Any]:
        return {"items": gpu_resources(), "compute_host_id": host_id()}

    @app.post(f"{route_prefix}/compute/jobs")
    async def create_job(payload: dict[str, Any]) -> dict[str, Any]:
        try:
            command = build_command(payload, os.getenv("LLAMA_FACTORY_HOME", "/app/LLaMA-Factory"))
        except ValueError as exc:
            raise HTTPException(status_code=400, detail=str(exc))
        job_id = str(payload.get("id") or f"job_{int(now() * 1000)}")
        job = {
            "id": job_id,
            "name": payload.get("name", job_id),
            "status": "queued",
            "progress": 10,
            "pid": int(52000 + now() % 10000),
            "gpus": payload.get("gpus") or [0],
            "created_at": now(),
            "command": command.command,
            "work_dir": command.work_dir,
            "artifacts": [],
            "logs": "",
        }
        jobs[job_id] = job
        return job_status(job)

    @app.get(f"{route_prefix}/compute/jobs")
    async def list_jobs() -> dict[str, Any]:
        return {"items": [job_status(job) for job in jobs.values()]}

    @app.get(f"{route_prefix}/compute/jobs/{{job_id}}")
    async def get_job(job_id: str) -> dict[str, Any]:
        job = jobs.get(job_id)
        if not job:
            raise HTTPException(status_code=404, detail="job not found")
        return job_status(job)

    @app.post(f"{route_prefix}/compute/jobs/{{job_id}}/stop")
    async def stop_job(job_id: str) -> dict[str, Any]:
        job = jobs.get(job_id)
        if not job:
            raise HTTPException(status_code=404, detail="job not found")
        job["status"] = "stopped"
        job["progress"] = min(job.get("progress", 0), 99)
        return job

    @app.get(f"{route_prefix}/compute/jobs/{{job_id}}/logs")
    async def job_logs(job_id: str) -> dict[str, Any]:
        job = jobs.get(job_id)
        if not job:
            raise HTTPException(status_code=404, detail="job not found")
        job = job_status(job)
        metrics = [parse_log_line(line) for line in job["logs"].splitlines()]
        return {"job_id": job_id, "content": job["logs"], "metrics": [m for m in metrics if m]}

    @app.post(f"{route_prefix}/compute/files/upload")
    async def upload_file(payload: dict[str, Any]) -> dict[str, Any]:
        file_id = str(payload.get("id") or f"file_{int(now() * 1000)}")
        return {"id": file_id, "status": "available", "local_path": f"/data/yg-ft/uploads/{file_id}"}

    @app.get(f"{route_prefix}/compute/files/{{file_id}}/download")
    async def download_file(file_id: str) -> dict[str, Any]:
        return {"id": file_id, "status": "ready", "download_url": f"{route_prefix}/compute/files/{file_id}/download"}

    return app


app = create_app()