第一次提交
This commit is contained in:
wangjiming
1
compute/api/__init__.py
Normal file
1
compute/api/__init__.py
Normal file
@@ -0,0 +1 @@
|
||||
"""Compute API package."""
|
||||
211
compute/api/main.py
Normal file
211
compute/api/main.py
Normal file
@@ -0,0 +1,211 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import math
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from fastapi import FastAPI, HTTPException
|
||||
|
||||
from compute.engines.llama_factory.adapter import build_command, parse_log_line
|
||||
|
||||
|
||||
def create_app() -> FastAPI:
|
||||
app = FastAPI(title="YG Fine-Tune Compute API")
|
||||
jobs: dict[str, dict[str, Any]] = {}
|
||||
route_prefix = os.getenv("MODELTF_ROUTE_PREFIX", "/modelTF").rstrip("/") or "/modelTF"
|
||||
|
||||
def now() -> float:
|
||||
return time.time()
|
||||
|
||||
def host_id() -> str:
|
||||
return os.getenv("COMPUTE_HOST_ID", "gpu-node-01")
|
||||
|
||||
def execution_mode() -> str:
|
||||
return os.getenv("COMPUTE_EXECUTION_MODE", os.getenv("COMPUTE_MODE", "real")).lower()
|
||||
|
||||
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 gpu_resources() -> list[dict[str, Any]]:
|
||||
if execution_mode() != "simulator":
|
||||
return []
|
||||
active_jobs = [job_status(job) for job in jobs.values() if job["status"] in {"queued", "running"}]
|
||||
gpus: list[dict[str, Any]] = []
|
||||
for idx in range(4):
|
||||
task = next((job for job in active_jobs if idx in job.get("gpus", [])), None)
|
||||
busy = task is not None and task["status"] == "running"
|
||||
reserved = task is not None and task["status"] == "queued"
|
||||
gpus.append(
|
||||
{
|
||||
"id": idx,
|
||||
"uuid": f"GPU-{host_id().upper()}-{idx}",
|
||||
"name": os.getenv("COMPUTE_GPU_NAME", "NVIDIA A800-SXM4-80GB"),
|
||||
"status": "busy" if busy else "reserved" if reserved else "idle",
|
||||
"gpu_percent": 88 if busy else 25 if reserved else 4,
|
||||
"memory_used_gb": 58 if busy else 12 if reserved else 2,
|
||||
"memory_total_gb": 80,
|
||||
"temperature": 61 if busy else 45 if reserved else 36,
|
||||
"power_w": 215 if busy else 80 if reserved else 25,
|
||||
"power_limit_w": 300,
|
||||
"processes": [
|
||||
{
|
||||
"pid": task["pid"],
|
||||
"name": "llamafactory-cli",
|
||||
"task_name": task["name"],
|
||||
"memory_used_gb": 58 if busy else 12,
|
||||
}
|
||||
]
|
||||
if task
|
||||
else [],
|
||||
}
|
||||
)
|
||||
return gpus
|
||||
|
||||
@app.get(f"{route_prefix}/health")
|
||||
async def health_check() -> dict[str, str]:
|
||||
return {
|
||||
"status": "ok",
|
||||
"compute_host_id": os.getenv("COMPUTE_HOST_ID", "unknown"),
|
||||
}
|
||||
|
||||
@app.get(f"{route_prefix}/v1/compute/health")
|
||||
async def compute_health_check() -> dict[str, str | bool]:
|
||||
data_root = Path(os.getenv("YG_FT_DATA_ROOT", "/data/yg-ft"))
|
||||
llama_factory_home = Path(os.getenv("LLAMA_FACTORY_HOME", "/app/LLaMA-Factory"))
|
||||
return {
|
||||
"status": "ok",
|
||||
"compute_host_id": os.getenv("COMPUTE_HOST_ID", "unknown"),
|
||||
"app_callback_enabled": os.getenv("ENABLE_APP_CALLBACK", "false").lower() == "true",
|
||||
"data_root": str(data_root),
|
||||
"data_root_exists": data_root.exists(),
|
||||
"llama_factory_home": str(llama_factory_home),
|
||||
"llama_factory_home_exists": llama_factory_home.exists(),
|
||||
"execution_mode": execution_mode(),
|
||||
}
|
||||
|
||||
@app.get(f"{route_prefix}/v1/compute/jobs")
|
||||
async def list_jobs_alias() -> dict[str, list[dict[str, Any]]]:
|
||||
return {"items": [job_status(job) for job in jobs.values()]}
|
||||
|
||||
@app.get(f"{route_prefix}/compute/resources/gpus")
|
||||
async def list_gpus() -> dict[str, Any]:
|
||||
return {"items": gpu_resources(), "compute_host_id": host_id()}
|
||||
|
||||
@app.post(f"{route_prefix}/compute/jobs")
|
||||
async def create_job(payload: dict[str, Any]) -> dict[str, Any]:
|
||||
try:
|
||||
command = build_command(payload, os.getenv("LLAMA_FACTORY_HOME", "/app/LLaMA-Factory"))
|
||||
except ValueError as exc:
|
||||
raise HTTPException(status_code=400, detail=str(exc))
|
||||
if execution_mode() != "simulator":
|
||||
raise HTTPException(
|
||||
status_code=501,
|
||||
detail="real compute executor is not implemented yet; set COMPUTE_EXECUTION_MODE=simulator only for isolated development",
|
||||
)
|
||||
job_id = str(payload.get("id") or f"job_{int(now() * 1000)}")
|
||||
job = {
|
||||
"id": job_id,
|
||||
"name": payload.get("name", job_id),
|
||||
"status": "queued",
|
||||
"progress": 10,
|
||||
"pid": int(52000 + now() % 10000),
|
||||
"gpus": payload.get("gpus") or [0],
|
||||
"created_at": now(),
|
||||
"command": command.command,
|
||||
"work_dir": command.work_dir,
|
||||
"artifacts": [],
|
||||
"logs": "",
|
||||
}
|
||||
jobs[job_id] = job
|
||||
return job_status(job)
|
||||
|
||||
@app.get(f"{route_prefix}/compute/jobs")
|
||||
async def list_jobs() -> dict[str, Any]:
|
||||
return {"items": [job_status(job) for job in jobs.values()]}
|
||||
|
||||
@app.get(f"{route_prefix}/compute/jobs/{{job_id}}")
|
||||
async def get_job(job_id: str) -> dict[str, Any]:
|
||||
job = jobs.get(job_id)
|
||||
if not job:
|
||||
raise HTTPException(status_code=404, detail="job not found")
|
||||
return job_status(job)
|
||||
|
||||
@app.post(f"{route_prefix}/compute/jobs/{{job_id}}/stop")
|
||||
async def stop_job(job_id: str) -> dict[str, Any]:
|
||||
job = jobs.get(job_id)
|
||||
if not job:
|
||||
raise HTTPException(status_code=404, detail="job not found")
|
||||
job["status"] = "stopped"
|
||||
job["progress"] = min(job.get("progress", 0), 99)
|
||||
return job
|
||||
|
||||
@app.get(f"{route_prefix}/compute/jobs/{{job_id}}/logs")
|
||||
async def job_logs(job_id: str) -> dict[str, Any]:
|
||||
job = jobs.get(job_id)
|
||||
if not job:
|
||||
raise HTTPException(status_code=404, detail="job not found")
|
||||
job = job_status(job)
|
||||
metrics = [parse_log_line(line) for line in job["logs"].splitlines()]
|
||||
return {"job_id": job_id, "content": job["logs"], "metrics": [m for m in metrics if m]}
|
||||
|
||||
@app.post(f"{route_prefix}/compute/files/upload")
|
||||
async def upload_file(payload: dict[str, Any]) -> dict[str, Any]:
|
||||
file_id = str(payload.get("id") or f"file_{int(now() * 1000)}")
|
||||
return {"id": file_id, "status": "available", "local_path": f"/data/yg-ft/uploads/{file_id}"}
|
||||
|
||||
@app.get(f"{route_prefix}/compute/files/{{file_id}}/download")
|
||||
async def download_file(file_id: str) -> dict[str, Any]:
|
||||
return {"id": file_id, "status": "ready", "download_url": f"{route_prefix}/compute/files/{file_id}/download"}
|
||||
|
||||
return app
|
||||
|
||||
|
||||
app = create_app()
|
||||
Reference in New Issue
Block a user