from __future__ import annotations import asyncio import json import os import math import hashlib import shutil import subprocess import time from pathlib import Path from typing import Any import httpx from fastapi import FastAPI, File, Form, HTTPException, Query, Request, UploadFile from fastapi.responses import FileResponse, JSONResponse, StreamingResponse from compute.agent.process_manager import ProcessManager from compute.api.security import docs_kwargs from compute.engines.llama_factory.adapter import build_command, parse_log_line, prepare_runtime_files from compute.engines.llama_factory.inference import get_inference_session def create_app() -> FastAPI: app = FastAPI(title="YG Zhilian Compute API", **docs_kwargs()) jobs: dict[str, dict[str, Any]] = {} cache_locks: dict[str, asyncio.Lock] = {} 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 torch_cuda_status() -> dict[str, Any]: try: import torch # type: ignore[import-not-found] except Exception as exc: # noqa: BLE001 - keep health endpoint resilient return { "available": False, "device_count": 0, "torch_version": "", "torch_cuda_version": "", "error": f"torch import failed: {exc}", } try: available = bool(torch.cuda.is_available()) device_count = int(torch.cuda.device_count()) devices = [] for index in range(device_count): props = torch.cuda.get_device_properties(index) devices.append( { "index": index, "name": props.name, "memory_total_gb": round(props.total_memory / 1024 / 1024 / 1024, 2), } ) return { "available": available, "device_count": device_count, "torch_version": str(torch.__version__), "torch_cuda_version": str(torch.version.cuda or ""), "devices": devices, "error": "" if available else "torch cuda is not available", } except Exception as exc: # noqa: BLE001 - expose CUDA initialization failures return { "available": False, "device_count": 0, "torch_version": str(getattr(torch, "__version__", "")), "torch_cuda_version": str(getattr(torch.version, "cuda", "") or ""), "devices": [], "error": str(exc), } 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]] = [] processes_by_uuid: dict[str, list[dict[str, Any]]] = {} try: process_result = subprocess.run( [ "nvidia-smi", "--query-compute-apps=gpu_uuid,pid,process_name,used_memory", "--format=csv,noheader,nounits", ], check=True, capture_output=True, text=True, timeout=5, ) for process_line in process_result.stdout.splitlines(): process_parts = [part.strip() for part in process_line.split(",")] if len(process_parts) < 4: continue process_uuid, pid, process_name, used_memory = process_parts[:4] processes_by_uuid.setdefault(process_uuid, []).append( { "pid": int(_safe_float(pid)), "name": process_name, "memory_used_gb": round(_safe_float(used_memory) / 1024, 2), } ) except Exception: # Some driver/runtime combinations do not expose compute-apps; # utilization and memory metrics remain useful without processes. pass 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": processes_by_uuid.get(uuid, []), } ) 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 _validate_training_accelerator(payload: dict[str, Any]) -> tuple[list[str], list[str], dict[str, Any]]: errors: list[str] = [] warnings: list[str] = [] if str(payload.get("engine") or payload.get("training_engine") or "llama_factory") == "smoke": return errors, warnings, {} requested_gpus = [int(item) for item in payload.get("gpus") or []] if not requested_gpus: warnings.append("no gpu selected; training will run on CPU") return errors, warnings, {} cuda = torch_cuda_status() if not cuda.get("available"): errors.append(f"torch cuda unavailable on compute node: {cuda.get('error') or 'unknown error'}") device_count = int(cuda.get("device_count") or 0) if device_count and max(requested_gpus) >= device_count: errors.append(f"requested gpu index out of torch device range: requested={requested_gpus}, device_count={device_count}") min_memory_gb = _float_env("MIN_TRAINING_GPU_MEMORY_GB", 4.0) gpus = {int(item["gpu_index"]): item for item in gpu_resources() if "gpu_index" in item} for gpu_index in requested_gpus: gpu = gpus.get(gpu_index) if not gpu: errors.append(f"requested gpu not found by nvidia-smi: {gpu_index}") continue memory_total = float(gpu.get("memory_total_gb") or 0) if memory_total and memory_total < min_memory_gb: errors.append( f"gpu {gpu_index} memory too small: {memory_total}GB < required {min_memory_gb}GB" ) return errors, warnings, cuda 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, "byte_size": sum(child.stat().st_size for child in path.rglob("*") if child.is_file()) if exists and path.is_dir() else path.stat().st_size if exists and path.is_file() else 0, "ok": ok or not item.get("required", True), } def _job_preview(payload: dict[str, Any], check_paths: bool) -> dict[str, Any]: warnings: list[str] = [] runtime_files: list[dict[str, str]] = [] command_payload = {**payload, "require_dataset_files": check_paths} if check_paths: try: runtime_files = prepare_runtime_files(command_payload) except OSError as exc: return { "valid": False, "errors": [f"prepare runtime files failed: {exc}"], "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": {}, "runtime_files": [], "path_checks": [], } try: command = build_command(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": {}, "runtime_files": runtime_files, "path_checks": [], } errors: list[str] = [] engine = str(payload.get("engine") or payload.get("training_engine") or "llama_factory") path_checks: list[dict[str, Any]] = [] accelerator: 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 payload.get("base_model_path") or "", "type": "any", "required": True, } ) ] if engine in {"merge", "export", "llama_factory_export"} and payload.get("adapter_name_or_path"): path_checks.append( _check_path_item( { "name": "adapter_name_or_path", "path": payload.get("adapter_name_or_path"), "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}") if engine not in {"merge", "export", "llama_factory_export"}: accelerator_errors, accelerator_warnings, accelerator = _validate_training_accelerator(payload) errors.extend(accelerator_errors) warnings.extend(accelerator_warnings) elif engine == "eval": # Eval engine: validate model path and dataset path if not payload.get("model_name_or_path"): errors.append("model_name_or_path is required for eval") else: path_checks.append(_check_path_item({ "name": "model_name_or_path", "path": payload.get("model_name_or_path", ""), "type": "any", "required": True, })) if payload.get("dataset_path"): path_checks.append(_check_path_item({ "name": "dataset_path", "path": payload.get("dataset_path", ""), "type": "file", "required": True, })) else: errors.append("dataset_path is required for eval") if shutil.which("python") is None: errors.append("python runtime not found") 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, "runtime_files": runtime_files, "accelerator": accelerator, "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")) gpu_items = gpu_resources() torch_cuda = torch_cuda_status() storage_available = False storage_error = "" try: data_root.mkdir(parents=True, exist_ok=True) probe = data_root / ".yg-ft-storage-healthcheck" probe.write_text(host_id(), encoding="utf-8") storage_available = probe.read_text(encoding="utf-8").strip() == host_id() probe.unlink(missing_ok=True) except Exception as exc: # noqa: BLE001 - health endpoint must remain available storage_error = str(exc) return { "status": "ok" if storage_available else "storage_unavailable", "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(), "storage_mode": os.getenv("STORAGE_MODE", "minio-cache"), "storage_available": storage_available, "storage_error": storage_error, "storage_root": str(data_root), "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), "nvidia_gpu_count": len(gpu_items), "torch_cuda": torch_cuda, "gpu_discovery_endpoint": f"{route_prefix}/compute/resources/gpus", "capabilities": ["gpu_discovery", "torch_cuda_diagnostics", "llama_factory", "file_gateway", "job_polling", "storage_health"], } @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/files/upload-to-url") async def upload_file_to_url(payload: dict[str, Any]) -> dict[str, Any]: """Upload one node-local artifact to a Backend-issued presigned URL.""" source = Path(str(payload.get("path") or "")).resolve() data_root = Path(os.getenv("YG_FT_DATA_ROOT", "/data/yg-ft")).resolve() if not _path_inside(data_root, source) or not source.is_file(): raise HTTPException(status_code=400, detail="artifact path must be an existing file inside YG_FT_DATA_ROOT") upload_url = str(payload.get("upload_url") or "") if not upload_url: raise HTTPException(status_code=400, detail="upload_url is required") digest = hashlib.sha256() byte_size = 0 try: async with httpx.AsyncClient(timeout=httpx.Timeout(900, connect=30)) as client: with source.open("rb") as handle: content = handle.read() digest.update(content) byte_size = len(content) response = await client.put(upload_url, content=content, headers={"Content-Type": str(payload.get("content_type") or "application/octet-stream")}) response.raise_for_status() except Exception as exc: raise HTTPException(status_code=502, detail=f"artifact upload failed: {exc}") from exc return {"status": "available", "path": str(source), "byte_size": byte_size, "checksum_sha256": digest.hexdigest(), "object_key": payload.get("object_key")} @app.post(f"{route_prefix}/compute/jobs") async def create_job(payload: dict[str, Any]) -> dict[str, Any]: payload = {**payload, "require_dataset_files": True} try: prepare_runtime_files(payload) except OSError as exc: raise HTTPException(status_code=400, detail=f"prepare runtime files failed: {exc}") 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 [], "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]} # ── Inference Endpoints ─────────────────────────────────────────── @app.post(f"{route_prefix}/inference/load") async def inference_load(payload: dict[str, Any]) -> dict[str, Any]: """Load a model for inference using LLaMA-Factory ChatModel. Expected payload: model_name_or_path: str (required) adapter_name_or_path: str (optional, for LoRA adapters) template: str (default: "qwen") infer_backend: str (default: "huggingface") infer_dtype: str (default: "auto") """ requested_gpus = payload.get("gpu_indices") if requested_gpus is None: requested_gpus = payload.get("gpus") or [] try: requested_gpus = sorted({int(item) for item in requested_gpus}) except (TypeError, ValueError) as exc: raise HTTPException(status_code=400, detail=f"invalid GPU selection: {exc}") from exc if any(item < 0 for item in requested_gpus): raise HTTPException(status_code=400, detail="GPU index must be non-negative") if requested_gpus: known_gpus = {int(item.get("gpu_index", item.get("id", -1))) for item in gpu_resources()} missing = sorted(set(requested_gpus) - known_gpus) if missing: raise HTTPException(status_code=409, detail=f"requested GPU not found: {missing}") conflict = sorted(set(requested_gpus).intersection(process_manager.locked_gpus())) if conflict: raise HTTPException(status_code=409, detail=f"GPU already used by another compute job: {conflict}") session = get_inference_session() result = session.load( model_name_or_path=payload.get("model_name_or_path", ""), adapter_name_or_path=payload.get("adapter_name_or_path", ""), template=payload.get("template", "qwen"), infer_backend=payload.get("infer_backend", "huggingface"), infer_dtype=payload.get("infer_dtype", "auto"), gpu_indices=requested_gpus, ) return result @app.post(f"{route_prefix}/inference/unload") async def inference_unload() -> dict[str, Any]: """Unload the currently loaded model and free GPU memory.""" # Teardown (gc.collect + cuda.empty_cache) can take a while; run it off # the event loop so /health and /inference/status stay responsive. return await asyncio.to_thread(get_inference_session().unload) @app.get(f"{route_prefix}/inference/status") async def inference_status() -> dict[str, Any]: """Get the current inference session status.""" return get_inference_session().info() @app.post(f"{route_prefix}/inference/chat") async def inference_chat(payload: dict[str, Any]) -> dict[str, Any]: """Chat with the loaded model (non-streaming). Expected payload: messages: list[dict] (OpenAI format) temperature: float (default 0.95) top_p: float (default 0.7) max_new_tokens: int (default 1024) """ messages = payload.get("messages") or [] if not messages: raise HTTPException(status_code=400, detail="messages is required") # Generation is long-running; run it in a thread so the event loop keeps # serving /inference/status and /health during inference. result = await asyncio.to_thread( get_inference_session().chat, messages=messages, temperature=float(payload.get("temperature", 0.95)), top_p=float(payload.get("top_p", 0.7)), max_new_tokens=int(payload.get("max_new_tokens", 1024)), do_sample=bool(payload.get("do_sample", True)), ) if result.get("error"): raise HTTPException(status_code=500, detail=result["error"]) return {"response": result["response"]} @app.post(f"{route_prefix}/inference/chat/stream") async def inference_chat_stream(payload: dict[str, Any]) -> StreamingResponse: """Chat with streaming response (Server-Sent Events).""" messages = payload.get("messages") or [] if not messages: raise HTTPException(status_code=400, detail="messages is required") def generate(): session = get_inference_session() for chunk in session.chat_stream( messages=messages, temperature=float(payload.get("temperature", 0.95)), top_p=float(payload.get("top_p", 0.7)), max_new_tokens=int(payload.get("max_new_tokens", 1024)), do_sample=bool(payload.get("do_sample", True)), ): yield chunk return StreamingResponse(generate(), media_type="text/event-stream") @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/cache/prepare") async def prepare_cache(payload: dict[str, Any]) -> dict[str, Any]: """Download one MinIO object into the node-local cache atomically.""" download_url = str(payload.get("download_url") or "") resource_id = str(payload.get("resource_id") or "") version_id = str(payload.get("version_id") or "latest") if not download_url or not resource_id: raise HTTPException(status_code=400, detail="download_url and resource_id are required") data_root = Path(os.getenv("YG_FT_DATA_ROOT", "/data/yg-ft")) cache_root = Path(os.getenv("YG_FT_CACHE_ROOT", str(data_root))) relative_path = str(payload.get("relative_path") or f"resources/{resource_id}/{version_id}/resource") target = (cache_root / relative_path.lstrip("/\\")).resolve() if not _path_inside(cache_root, target): raise HTTPException(status_code=400, detail="cache path must stay inside cache root") target.parent.mkdir(parents=True, exist_ok=True) temp_target = target.with_name(f".{target.name}.part") expected_checksum = str(payload.get("checksum_sha256") or "").lower() expected_size = int(payload.get("byte_size") or 0) lock = cache_locks.setdefault(str(target), asyncio.Lock()) async with lock: if target.is_file() and expected_checksum: existing_digest = hashlib.sha256() with target.open("rb") as existing: while chunk := existing.read(1024 * 1024): existing_digest.update(chunk) if existing_digest.hexdigest().lower() == expected_checksum and (not expected_size or target.stat().st_size == expected_size): return {"resource_id": resource_id, "version_id": version_id, "status": "ready", "local_path": str(target), "byte_size": target.stat().st_size, "checksum_sha256": existing_digest.hexdigest(), "reused": True} digest = hashlib.sha256() byte_size = 0 try: async with lock: for attempt in range(3): try: digest = hashlib.sha256() byte_size = 0 temp_target.unlink(missing_ok=True) async with httpx.AsyncClient(timeout=httpx.Timeout(900, connect=30), follow_redirects=True) as client: async with client.stream("GET", download_url) as response: response.raise_for_status() with temp_target.open("wb") as output: async for chunk in response.aiter_bytes(1024 * 1024): output.write(chunk) digest.update(chunk) byte_size += len(chunk) break except Exception: temp_target.unlink(missing_ok=True) if attempt == 2: raise await asyncio.sleep(2 ** attempt) checksum = digest.hexdigest() if expected_checksum and checksum != expected_checksum: temp_target.unlink(missing_ok=True) raise HTTPException(status_code=502, detail="cache checksum mismatch") if expected_size and byte_size != expected_size: temp_target.unlink(missing_ok=True) raise HTTPException(status_code=502, detail="cache byte size mismatch") temp_target.replace(target) except HTTPException: raise except Exception as exc: temp_target.unlink(missing_ok=True) raise HTTPException(status_code=502, detail=f"cache download failed: {exc}") from exc return { "resource_id": resource_id, "version_id": version_id, "status": "ready", "local_path": str(target), "byte_size": byte_size, "checksum_sha256": checksum, } @app.get(f"{route_prefix}/compute/cache/status") async def cache_status(resource_id: str = Query(...), version_id: str = Query(default="latest")) -> dict[str, Any]: data_root = Path(os.getenv("YG_FT_DATA_ROOT", "/data/yg-ft")) cache_root = Path(os.getenv("YG_FT_CACHE_ROOT", str(data_root))) target = cache_root / "resources" / resource_id / version_id / "resource" return { "resource_id": resource_id, "version_id": version_id, "status": "ready" if target.is_file() else "missing", "local_path": str(target), "byte_size": target.stat().st_size if target.is_file() else 0, } @app.delete(f"{route_prefix}/compute/cache") async def clear_cache(resource_id: str | None = Query(default=None), version_id: str | None = Query(default=None)) -> dict[str, Any]: data_root = Path(os.getenv("YG_FT_DATA_ROOT", "/data/yg-ft")) cache_root = Path(os.getenv("YG_FT_CACHE_ROOT", str(data_root))) target = cache_root / "resources" if resource_id: target = target / resource_id if version_id: target = target / version_id target = target.resolve() if not _path_inside(cache_root, target): raise HTTPException(status_code=400, detail="cache path must stay inside cache root") if target.exists(): shutil.rmtree(target) return {"status": "cleared", "resource_id": resource_id, "version_id": version_id} @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/read") async def read_file(path: str = Query(...)) -> JSONResponse: """Read a text file from within YG_FT_DATA_ROOT. Used by the backend to fetch eval results and other job outputs.""" data_root = Path(os.getenv("YG_FT_DATA_ROOT", "/data/yg-ft")) target = (data_root / path.lstrip("/\\")).resolve() if not _path_inside(data_root, target): raise HTTPException(status_code=400, detail="path must stay inside YG_FT_DATA_ROOT") if not target.is_file(): raise HTTPException(status_code=404, detail="file not found") try: content = target.read_text(encoding="utf-8") return JSONResponse(json.loads(content) if content.strip().startswith("{") else {"content": content}) except Exception as exc: raise HTTPException(status_code=500, detail=str(exc)) @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" # file_id 仅允许普通标识符，拒绝 ../、/、\ 等路径穿越字符。 if not file_id or not all(character.isalnum() or character in {"_", "-"} for character in file_id): raise HTTPException(status_code=400, detail="invalid file id") matches = list(upload_root.glob(f"{file_id}_*")) if not matches: raise HTTPException(status_code=404, detail="file not found") # 解析符号链接后仍必须位于 upload 根目录内，防止符号链接指向目录外文件。 resolved = matches[0].resolve() if not _path_inside(upload_root, resolved): raise HTTPException(status_code=404, detail="file not found") return FileResponse(resolved) return app app = create_app()