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]] = {} 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("/health") async def health_check() -> dict[str, str]: return { "status": "ok", "compute_host_id": os.getenv("COMPUTE_HOST_ID", "unknown"), } @app.get("/api/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("/api/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("/compute/resources/gpus") async def list_gpus() -> dict[str, Any]: return {"items": gpu_resources(), "compute_host_id": host_id()} @app.post("/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("/compute/jobs") async def list_jobs() -> dict[str, Any]: return {"items": [job_status(job) for job in jobs.values()]} @app.get("/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("/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("/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("/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("/compute/files/{file_id}/download") async def download_file(file_id: str) -> dict[str, Any]: return {"id": file_id, "status": "ready", "download_url": f"/compute/files/{file_id}/download"} return app app = create_app()