from __future__ import annotations import os import math import hashlib import shutil import subprocess import time from pathlib import Path from typing import Any from fastapi import FastAPI, File, Form, HTTPException, Query, Request, UploadFile from fastapi.responses import FileResponse, JSONResponse from compute.agent.process_manager import ProcessManager 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" process_manager = ProcessManager(os.getenv("TRAINING_LOG_ROOT", "/opt/yg-ft/logs/training")) @app.middleware("http") async def compute_token_auth(request: Request, call_next): token = os.getenv("COMPUTE_SERVICE_TOKEN", "") auth_enabled = os.getenv("COMPUTE_AUTH_ENABLED", "true").lower() == "true" public_paths = {f"{route_prefix}/health", "/health"} if auth_enabled and token and request.url.path not in public_paths: header_token = request.headers.get("x-compute-token", "") auth_header = request.headers.get("authorization", "") bearer_token = auth_header.removeprefix("Bearer ").strip() if auth_header.startswith("Bearer ") else "" if header_token != token and bearer_token != token: return JSONResponse({"detail": "invalid compute service token"}, status_code=401) return await call_next(request) def now() -> float: return time.time() def host_id() -> str: return os.getenv("COMPUTE_HOST_ID", "gpu-node-01") def execution_mode() -> str: return os.getenv("COMPUTE_EXECUTION_MODE", os.getenv("COMPUTE_MODE", "real")).lower() def _int_env(name: str, default: int) -> int: raw = os.getenv(name) if raw is None or raw == "": return default return int(raw) def _float_env(name: str, default: float) -> float: raw = os.getenv(name) if raw is None or raw == "": return default return float(raw) def _path_inside(root: Path, candidate: Path) -> bool: try: candidate.resolve().relative_to(root.resolve()) return True except ValueError: return False def _llama_factory_version() -> str: for command in (["llamafactory-cli", "version"], ["llamafactory-cli", "--version"]): try: result = subprocess.run(command, capture_output=True, text=True, timeout=5) except Exception: continue output = (result.stdout or result.stderr).strip() if result.returncode == 0 and output: return output.splitlines()[0][:120] return "" def _slice_log_content( content: str, tail_lines: int | None = None, offset: int | None = None, limit: int | None = None, ) -> dict[str, Any]: lines = content.splitlines() total = len(lines) if offset is not None or limit is not None: start = max(0, offset or 0) end = start + limit if limit else total selected = lines[start:end] else: tail = tail_lines or 200 start = max(0, total - tail) selected = lines[start:] next_offset = start + len(selected) return { "content": "\n".join(selected), "total_lines": total, "offset": start, "limit": len(selected), "has_more": next_offset < total, "next_offset": next_offset if next_offset < total else None, } def _safe_float(value: Any, default: float = 0) -> float: try: return float(str(value).replace("[N/A]", "").strip() or default) except (TypeError, ValueError): return default def job_status(job: dict[str, Any]) -> dict[str, Any]: if execution_mode() != "simulator": return job 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", 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 real_gpu_resources() -> list[dict[str, Any]]: query = ( "index,uuid,name,memory.total,memory.used,utilization.gpu," "temperature.gpu,power.draw,power.limit" ) try: result = subprocess.run( ["nvidia-smi", f"--query-gpu={query}", "--format=csv,noheader,nounits"], check=True, capture_output=True, text=True, timeout=5, ) except Exception: return fallback_gpu_resources() items: list[dict[str, Any]] = [] for line in result.stdout.splitlines(): parts = [part.strip() for part in line.split(",")] if len(parts) < 9: continue idx, uuid, name, mem_total, mem_used, util, temp, power, power_limit = parts[:9] total_gb = round(_safe_float(mem_total) / 1024, 2) used_gb = round(_safe_float(mem_used) / 1024, 2) memory_percent = round(used_gb / total_gb * 100, 1) if total_gb else 0 gpu_percent = int(_safe_float(util)) items.append( { "id": int(idx), "gpu_index": int(idx), "uuid": uuid, "name": name, "status": "busy" if gpu_percent >= 5 or used_gb > 1 else "idle", "gpu_percent": gpu_percent, "memory_used_gb": used_gb, "memory_total_gb": total_gb, "memory_percent": memory_percent, "temperature": int(_safe_float(temp)), "power_w": round(_safe_float(power), 1), "power_limit_w": round(_safe_float(power_limit), 1), "processes": [], } ) return items def fallback_gpu_resources() -> list[dict[str, Any]]: count = _int_env("COMPUTE_GPU_COUNT", 0) if count <= 0: return [] name = os.getenv("COMPUTE_GPU_NAME", "Configured GPU") memory_total = _float_env("COMPUTE_GPU_MEMORY_GB", 80.0) power_limit = _float_env("COMPUTE_GPU_POWER_LIMIT_W", 300.0) return [ { "id": idx, "gpu_index": idx, "uuid": f"GPU-{host_id().upper()}-{idx}", "name": name, "status": "idle", "gpu_percent": 0, "memory_used_gb": 0, "memory_total_gb": memory_total, "memory_percent": 0, "temperature": _int_env("COMPUTE_GPU_BASE_TEMPERATURE", 35), "power_w": 0, "power_limit_w": power_limit, "processes": [], } for idx in range(count) ] def gpu_resources() -> list[dict[str, Any]]: if execution_mode() != "simulator": return real_gpu_resources() 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 def _check_path_item(item: dict[str, Any]) -> dict[str, Any]: path = Path(str(item.get("path") or "")) exists = path.exists() expected_type = str(item.get("type") or "any") ok = exists if exists and expected_type == "dir": ok = path.is_dir() if exists and expected_type == "file": ok = path.is_file() return { "name": item.get("name") or "", "path": str(path), "type": expected_type, "required": bool(item.get("required", True)), "exists": exists, "is_dir": path.is_dir() if exists else False, "is_file": path.is_file() if exists else False, "ok": ok or not item.get("required", True), } def _job_preview(payload: dict[str, Any], check_paths: bool) -> dict[str, Any]: warnings: list[str] = [] try: command = build_command(payload, os.getenv("LLAMA_FACTORY_HOME", "/app/LLaMA-Factory")) except ValueError as exc: return { "valid": False, "errors": [part.strip() for part in str(exc).split(";") if part.strip()], "warnings": warnings, "engine": str(payload.get("engine") or payload.get("training_engine") or "llama_factory"), "command": [], "command_text": "", "work_dir": os.getenv("LLAMA_FACTORY_HOME", "/app/LLaMA-Factory"), "env": {}, "path_checks": [], } errors: list[str] = [] engine = str(payload.get("engine") or payload.get("training_engine") or "llama_factory") path_checks: list[dict[str, Any]] = [] if check_paths and engine != "smoke": path_checks = [ _check_path_item( { "name": "model_name_or_path", "path": payload.get("model_name_or_path") or payload.get("base_model") or "", "type": "any", "required": True, } ) ] if payload.get("dataset_dir"): path_checks.append( _check_path_item( { "name": "dataset_dir", "path": payload.get("dataset_dir"), "type": "dir", "required": True, } ) ) output_dir = Path(str(payload.get("output_dir") or "/data/yg-ft/outputs/training-job")) path_checks.append( _check_path_item( { "name": "output_parent", "path": str(output_dir.parent), "type": "dir", "required": False, } ) ) errors.extend( [f"{item['name']} path not available: {item['path']}" for item in path_checks if not item["ok"] and item["required"]] ) if shutil.which(command.command[0]) is None: errors.append(f"training command not found: {command.command[0]}") if not Path(command.work_dir).exists(): errors.append(f"llama_factory_home not found: {command.work_dir}") elif engine == "smoke": warnings.append("smoke engine skips model and dataset path checks") return { "valid": not errors, "errors": errors, "warnings": warnings, "engine": engine, "command": command.command, "command_text": " ".join(command.command), "work_dir": command.work_dir, "env": command.env, "path_checks": path_checks, } @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("/health") async def health_check_root() -> dict[str, str]: return await health_check() @app.get(f"{route_prefix}/v1/compute/health") async def compute_health_check() -> dict[str, Any]: data_root = Path(os.getenv("YG_FT_DATA_ROOT", "/data/yg-ft")) dataset_root = Path(os.getenv("YG_FT_DATASET_ROOT", str(data_root / "datasets"))) output_root = Path(os.getenv("YG_FT_OUTPUT_ROOT", str(data_root / "outputs"))) llama_factory_home = Path(os.getenv("LLAMA_FACTORY_HOME", "/app/LLaMA-Factory")) return { "status": "ok", "api_version": "v1", "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(), "model_root": os.getenv("YG_FT_MODEL_ROOT", str(data_root / "models")), "dataset_root": str(dataset_root), "dataset_root_exists": dataset_root.exists(), "output_root": str(output_root), "output_root_exists": output_root.exists(), "log_root": os.getenv("TRAINING_LOG_ROOT", "/opt/yg-ft/logs/training"), "llama_factory_home": str(llama_factory_home), "llama_factory_home_exists": llama_factory_home.exists(), "llama_factory_version": os.getenv("LLAMA_FACTORY_VERSION", ""), "execution_mode": execution_mode(), "gpu_count": _int_env("COMPUTE_GPU_COUNT", 0), "gpu_discovery_endpoint": f"{route_prefix}/compute/resources/gpus", "capabilities": ["gpu_discovery", "llama_factory", "file_gateway", "job_polling"], } @app.get(f"{route_prefix}/v1/compute/jobs") async def list_jobs_alias() -> dict[str, list[dict[str, Any]]]: items = process_manager.list_jobs() if execution_mode() != "simulator" else [job_status(job) for job in jobs.values()] return {"items": items} @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.get(f"{route_prefix}/v1/compute/resources/gpus") async def list_gpus_v1() -> dict[str, Any]: return {"items": gpu_resources(), "compute_host_id": host_id()} @app.post(f"{route_prefix}/compute/jobs/preview") async def preview_job(payload: dict[str, Any]) -> dict[str, Any]: return _job_preview(payload, check_paths=False) @app.post(f"{route_prefix}/compute/jobs/validate") async def validate_job(payload: dict[str, Any]) -> dict[str, Any]: return _job_preview(payload, check_paths=True) @app.post(f"{route_prefix}/v1/compute/jobs/preview") async def preview_job_v1(payload: dict[str, Any]) -> dict[str, Any]: return await preview_job(payload) @app.post(f"{route_prefix}/v1/compute/jobs/validate") async def validate_job_v1(payload: dict[str, Any]) -> dict[str, Any]: return await validate_job(payload) @app.post(f"{route_prefix}/compute/files/check-paths") async def check_paths(payload: dict[str, Any]) -> dict[str, Any]: items = [_check_path_item(item) for item in payload.get("paths", []) if isinstance(item, dict)] return {"valid": all(item["ok"] for item in items), "items": items} @app.get(f"{route_prefix}/compute/files/list") async def list_files( root: str = Query(default="data"), relative_path: str = Query(default=""), directories_only: bool = Query(default=False), ) -> dict[str, Any]: roots = { "data": Path(os.getenv("YG_FT_DATA_ROOT", "/data/yg-ft")), "models": Path(os.getenv("YG_FT_MODEL_ROOT", os.getenv("YG_FT_DATA_ROOT", "/data/yg-ft") + "/models")), "datasets": Path(os.getenv("YG_FT_DATASET_ROOT", os.getenv("YG_FT_DATA_ROOT", "/data/yg-ft") + "/datasets")), "outputs": Path(os.getenv("YG_FT_OUTPUT_ROOT", os.getenv("YG_FT_DATA_ROOT", "/data/yg-ft") + "/outputs")), } base = roots.get(root) if base is None: raise HTTPException(status_code=400, detail="invalid root") target = (base / relative_path.lstrip("/\\")).resolve() if not _path_inside(base, target): raise HTTPException(status_code=400, detail="path must stay inside selected root") if not target.exists(): return {"root": root, "base_path": str(base), "relative_path": relative_path, "items": []} items = [] for child in sorted(target.iterdir(), key=lambda path: (not path.is_dir(), path.name.lower())): if directories_only and not child.is_dir(): continue items.append( { "name": child.name, "path": str(child), "relative_path": str(child.relative_to(base)).replace("\\", "/"), "type": "directory" if child.is_dir() else "file", "byte_size": child.stat().st_size if child.is_file() else 0, } ) return {"root": root, "base_path": str(base), "relative_path": relative_path, "items": items} @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)}") if execution_mode() != "simulator": try: return process_manager.create_job({**payload, "id": job_id}, command.command, command.work_dir) except FileNotFoundError as exc: raise HTTPException(status_code=500, detail=f"training command not found: {exc.filename}") except ValueError as exc: raise HTTPException(status_code=409, detail=str(exc)) 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]: items = process_manager.list_jobs() if execution_mode() != "simulator" else [job_status(job) for job in jobs.values()] return {"items": items} @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 execution_mode() != "simulator": job = process_manager.get_job(job_id) if not job: raise HTTPException(status_code=404, detail="job not found") return job 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]: if execution_mode() != "simulator": job = process_manager.stop_job(job_id) if not job: raise HTTPException(status_code=404, detail="job not found") return job 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, tail_lines: int | None = Query(default=200, ge=1, le=5000), offset: int | None = Query(default=None, ge=0), limit: int | None = Query(default=None, ge=1, le=5000), ) -> dict[str, Any]: if execution_mode() != "simulator": job = process_manager.get_job(job_id) if not job: raise HTTPException(status_code=404, detail="job not found") content = process_manager.logs(job_id) else: job = jobs.get(job_id) if not job: raise HTTPException(status_code=404, detail="job not found") job = job_status(job) content = job["logs"] window = _slice_log_content(content, tail_lines, offset, limit) metrics = [parse_log_line(line) for line in window["content"].splitlines()] return {"job_id": job_id, **window, "metrics": [m for m in metrics if m]} @app.post(f"{route_prefix}/compute/files/upload") async def upload_file( file: UploadFile | None = File(default=None), target_relative_path: str | None = Form(default=None), resource_type: str | None = Form(default=None), resource_id: str | None = Form(default=None), ) -> dict[str, Any]: file_id = f"file_{int(now() * 1000)}" data_root = Path(os.getenv("YG_FT_DATA_ROOT", "/data/yg-ft")) data_root.mkdir(parents=True, exist_ok=True) filename = Path(file.filename if file else file_id).name if target_relative_path: target = (data_root / target_relative_path.lstrip("/\\")).resolve() if not _path_inside(data_root, target): raise HTTPException(status_code=400, detail="target path must stay inside YG_FT_DATA_ROOT") else: target = data_root / "uploads" / f"{file_id}_{filename}" if file: target.parent.mkdir(parents=True, exist_ok=True) with target.open("wb") as output: while chunk := await file.read(1024 * 1024): output.write(chunk) else: target.parent.mkdir(parents=True, exist_ok=True) target.write_text("", encoding="utf-8") return { "id": file_id, "resource_type": resource_type, "resource_id": resource_id, "status": "available", "local_path": str(target), "byte_size": target.stat().st_size, "checksum_sha256": hashlib.sha256(target.read_bytes()).hexdigest() if target.is_file() else "", } @app.post(f"{route_prefix}/compute/files/import-local") async def import_local_file(payload: dict[str, Any]) -> dict[str, Any]: source = Path(str(payload.get("source_path") or "")) if not source.exists(): raise HTTPException(status_code=404, detail="source path not found") data_root = Path(os.getenv("YG_FT_DATA_ROOT", "/data/yg-ft")) data_root.mkdir(parents=True, exist_ok=True) relative = str(payload.get("target_relative_path") or f"imports/{source.name}").lstrip("/\\") target = (data_root / relative).resolve() if not _path_inside(data_root, target): raise HTTPException(status_code=400, detail="target path must stay inside YG_FT_DATA_ROOT") target.parent.mkdir(parents=True, exist_ok=True) if source.is_dir(): if target.exists(): shutil.rmtree(target) shutil.copytree(source, target) byte_size = sum(path.stat().st_size for path in target.rglob("*") if path.is_file()) checksum = "" else: shutil.copy2(source, target) byte_size = target.stat().st_size checksum = hashlib.sha256(target.read_bytes()).hexdigest() return { "id": str(payload.get("id") or f"file_{int(now() * 1000)}"), "resource_type": payload.get("resource_type"), "resource_id": payload.get("resource_id"), "status": "available", "local_path": str(target), "byte_size": byte_size, "checksum_sha256": checksum, } @app.get(f"{route_prefix}/compute/files/{{file_id}}/download") async def download_file(file_id: str) -> FileResponse: upload_root = Path(os.getenv("YG_FT_DATA_ROOT", "/data/yg-ft")) / "uploads" matches = list(upload_root.glob(f"{file_id}_*")) if not matches: raise HTTPException(status_code=404, detail="file not found") return FileResponse(matches[0]) return app app = create_app()