813 lines
35 KiB
Python
813 lines
35 KiB
Python
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, StreamingResponse
|
|
|
|
from compute.agent.process_manager import ProcessManager
|
|
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 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 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]] = []
|
|
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 _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 == "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()
|
|
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),
|
|
"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", "inference"],
|
|
}
|
|
|
|
@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]:
|
|
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 [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]}
|
|
|
|
# ── 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."""
|
|
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"),
|
|
)
|
|
if not result.get("loaded"):
|
|
raise HTTPException(status_code=500, detail=result.get("error", "model load failed"))
|
|
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."""
|
|
return 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)."""
|
|
messages = payload.get("messages") or []
|
|
if not messages:
|
|
raise HTTPException(status_code=400, detail="messages is required")
|
|
result = 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/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()
|