feat: 推理与评测支持多卡 GPU 选择,训练任务实时 GPU 监控

- 新增 GPU 选择归一化 helper,统一前端各形态的选择(gpu_indices/gpus/gpu_id)
- 评测与推理支持同一节点内多卡选择,校验所选 GPU 空闲后再派发
- 新增 /fine-tune/{id}/gpu-status 接口,训练日志页展示实时 GPU 指标
- 算力节点推理加载支持 CUDA_VISIBLE_DEVICES 多卡可见,nvidia-smi 进程级监控
- GPU 占用跟踪细化为按卡记录,覆盖评测任务、推理模型与对比任务

Co-Authored-By: Claude <noreply@anthropic.com>
This commit is contained in:
wuyongtao
2026-08-19 10:44:56 +08:00
parent c3e96ae61b
commit 9379f93633
13 changed files with 360 additions and 48 deletions

View File

@@ -57,6 +57,33 @@ def _select_first_online_node(store: Any) -> dict[str, Any] | None:
return None
def _normalize_gpu_indices(payload: dict[str, Any], *, allow_primary: bool = True) -> list[int]:
"""Normalize all frontend GPU selection shapes to sorted integer indexes."""
raw = payload.get("gpu_indices")
if raw is None:
raw = payload.get("gpus")
if raw is None and allow_primary and payload.get("gpu_id") is not None:
raw = [payload.get("gpu_id")]
if raw is None or raw == "":
return []
if isinstance(raw, str):
raw = [item.strip() for item in raw.split(",") if item.strip()]
if not isinstance(raw, (list, tuple, set)):
raw = [raw]
result: set[int] = set()
for item in raw:
if isinstance(item, str) and ":" in item:
item = item.rsplit(":", 1)[-1]
try:
index = int(item)
except (TypeError, ValueError) as exc:
raise ValueError(f"invalid GPU index: {item}") from exc
if index < 0:
raise ValueError("GPU index must be non-negative")
result.add(index)
return sorted(result)
async def _wait_for_object_storage() -> None:
"""Wait for MinIO before starting a resource task."""
settings = get_settings()
@@ -1527,10 +1554,10 @@ async def start_fine_tune(
if node_id and gpu_indices:
if not store.check_gpu_access(current_user["id"], node_id, gpu_indices):
raise fail(403, "无权使用所选 GPU请联系管理员分配")
# 记录创建者
if node_id and not gpu_indices:
payload["allowed_gpu_indices"] = store.assigned_gpu_indexes(current_user["id"], node_id)
payload["strict_node_selection"] = bool(node_id)
# 页面明确选择节点时,调度器必须保持节点约束;否则可能落到其它节点。
payload["strict_node_selection"] = bool(payload.get("compute_node_id") or payload.get("node_id"))
payload.setdefault("created_by", current_user.get("id"))
try:
return ok(await _submit_fine_tune_task(store, payload))
@@ -1661,6 +1688,42 @@ async def fine_tune_diagnostics(task_id: str) -> dict[str, Any]:
)
@router.get("/fine-tune/{task_id}/gpu-status")
async def fine_tune_gpu_status(task_id: str, current_user: dict[str, Any] = Depends(get_current_user)) -> dict[str, Any]:
"""Return live GPU metrics for the task's selected node and cards."""
store = get_platform_store()
try:
task = store.task(task_id)
except KeyError:
raise fail(404, "fine tune task not found")
if not has_resource_access("fine-tune", task_id, current_user, "read"):
raise fail(403, "no permission to access this task")
node = _node_for_task(task)
selected = set(_normalize_gpu_indices({"gpus": task.get("gpus") or []}, allow_primary=False))
if not node:
return ok({"source": "unavailable", "items": [], "selected_gpus": sorted(selected)})
try:
if get_settings().compute_mode == "simulator":
live_items = store.gpus()
else:
live_items = await ComputeNodeClient(node["api_base_url"]).gpu_resources()
items = []
for item in live_items:
index = int(item.get("gpu_index", item.get("id", -1)))
if selected and index not in selected:
continue
items.append({
**item,
"id": index,
"node_id": node["id"],
"node_code": node.get("code"),
"node_name": node.get("name"),
})
return ok({"source": "compute", "items": items, "selected_gpus": sorted(selected)})
except Exception as exc: # noqa: BLE001 - let UI retain last good snapshot
return ok({"source": "unavailable", "items": [], "selected_gpus": sorted(selected), "error": str(exc)})
@router.put("/fine-tune/{task_id}")
async def update_fine_tune(task_id: str, payload: dict[str, Any] = Body(...), current_user: dict = Depends(get_current_user)) -> dict[str, Any]:
if not has_resource_access("fine-tune", task_id, current_user, "write"):
@@ -1808,6 +1871,12 @@ async def model_eval_detail(task_id: str, current_user: dict = Depends(get_curre
async def model_eval_start(payload: dict[str, Any] = Body(...), current_user: dict = Depends(get_current_user)) -> dict[str, Any]:
"""Start an evaluation task: submit eval job to compute node."""
store = get_platform_store()
try:
gpu_indices = _normalize_gpu_indices(payload)
except ValueError as exc:
raise fail(400, str(exc))
if not gpu_indices:
raise fail(400, "请选择至少一张 GPU")
# 1. Create eval task record
payload.setdefault("created_by", current_user.get("id"))
task = store.create_eval_task({**payload, "status": "pending"})
@@ -1910,6 +1979,17 @@ async def model_eval_start(payload: dict[str, Any] = Body(...), current_user: di
store.update_eval_task(task["id"], {"status": "failed", "error": message})
return ok({"task_id": task["id"], "status": "failed", "error": message})
node_gpus = {
int(item.get("id", item.get("gpu_index", -1))): item
for item in store.gpus()
if item.get("node_id") == node["id"]
}
unavailable = [index for index in gpu_indices if node_gpus.get(index, {}).get("status") != "idle"]
if unavailable:
message = f"selected GPU is not idle on compute node {node.get('code')}: {unavailable}"
store.update_eval_task(task["id"], {"status": "failed", "error": message})
return ok({"task_id": task["id"], "status": "failed", "error": message})
# 6. Build eval job payload
output_dir = f"/data/yg-ft/outputs/{task['id']}"
job_payload = {
@@ -1923,7 +2003,9 @@ async def model_eval_start(payload: dict[str, Any] = Body(...), current_user: di
"output_dir": output_dir,
"basic_metrics": payload.get("basic_metrics", {}),
"dimension": dimension_cfg,
"gpus": [int(payload.get("gpu_id", 0))],
"gpu_id": gpu_indices[0],
"gpu_indices": gpu_indices,
"gpus": gpu_indices,
"temperature": payload.get("temperature", 0.1),
"max_new_tokens": payload.get("max_new_tokens", 512),
"compute_node_id": node["id"],
@@ -2222,6 +2304,15 @@ async def model_compare_load(task_id: str, current_user: dict = Depends(get_curr
"model_name_or_path": model_path,
"template": item.get("template", "qwen"),
}
try:
item_gpu_indices = _normalize_gpu_indices(item)
except ValueError as exc:
loaded_models.append({**item, "status": "error", "error": str(exc)})
continue
if not item_gpu_indices:
loaded_models.append({**item, "status": "error", "error": "no GPU selected"})
continue
load_payload["gpu_indices"] = item_gpu_indices
if item.get("adapter_path"):
load_payload["adapter_name_or_path"] = item["adapter_path"]
if get_settings().compute_mode == "simulator":
@@ -2230,14 +2321,28 @@ async def model_compare_load(task_id: str, current_user: dict = Depends(get_curr
# 只派发HTTP 响应成功即视为已接受节点会异步加载loaded 字段忽略
item_dispatched = False
errors = []
for node in _candidate_online_nodes(store, preferred_node_id):
candidate_nodes = _candidate_online_nodes(store, preferred_node_id)
if preferred_node_id:
candidate_nodes = candidate_nodes[:1]
for node in candidate_nodes:
try:
node_gpu_map = {
int(gpu.get("id", gpu.get("gpu_index", -1))): gpu
for gpu in store.gpus()
if gpu.get("node_id") == node["id"]
}
unavailable = [
index for index in item_gpu_indices
if node_gpu_map.get(index, {}).get("status") != "idle"
]
if unavailable:
raise RuntimeError(f"selected GPU is not idle on compute node {node.get('code')}: {unavailable}")
if get_settings().minio_enabled:
await _wait_for_object_storage()
client = ComputeNodeClient(node["api_base_url"])
await client.inference_load(load_payload)
store.mark_inference_loaded(node["id"])
loaded_models.append({**item, "status": "starting", "node_id": node["id"], "node_name": node.get("name")})
store.mark_inference_loaded(node["id"], item_gpu_indices)
loaded_models.append({**item, "gpu_indices": item_gpu_indices, "gpus": item_gpu_indices, "status": "starting", "node_id": node["id"], "node_name": node.get("name")})
item_dispatched = True
break
except Exception as exc: # noqa: BLE001 - try next candidate node
@@ -2340,7 +2445,7 @@ async def model_chat_local_preload(payload: dict[str, Any] = Body(...)) -> dict[
# 计算节点现在异步加载HTTP 接受loading/ready即视为派发成功
result = await client.inference_load(payload)
if result.get("loaded") or result.get("status") in {"loading", "ready"}:
store.mark_inference_loaded(node["id"])
store.mark_inference_loaded(node["id"], _normalize_gpu_indices(payload))
return ok(result)
except Exception as exc:
return ok({"loaded": False, "error": str(exc)})
@@ -2420,7 +2525,7 @@ async def model_chat_trained_preload(payload: dict[str, Any] = Body(...), curren
# 计算节点现在异步加载HTTP 接受loading/ready即视为派发成功
result = await client.inference_load({**payload, "compute_node_id": node["id"]})
if result.get("loaded") or result.get("status") in {"loading", "ready"}:
store.mark_inference_loaded(node["id"])
store.mark_inference_loaded(node["id"], _normalize_gpu_indices(payload))
return ok(result)
except Exception as exc:
return ok({"loaded": False, "error": str(exc)})

View File

@@ -454,15 +454,21 @@ class PlatformStore:
self.ensure_seed_data()
# Track which compute nodes have an active inference model loaded
self._inference_nodes: set[str] = set()
self._inference_gpu_indexes: dict[str, set[int]] = {}
self._last_runtime_refresh = 0.0
# ── inference node tracking ────────────────────────────────────
def mark_inference_loaded(self, node_id: str) -> None:
def mark_inference_loaded(self, node_id: str, gpu_indexes: list[int] | None = None) -> None:
self._inference_nodes.add(node_id)
if gpu_indexes is not None:
self._inference_gpu_indexes[node_id] = {int(item) for item in gpu_indexes}
else:
self._inference_gpu_indexes.pop(node_id, None)
def mark_inference_unloaded(self, node_id: str) -> None:
self._inference_nodes.discard(node_id)
self._inference_gpu_indexes.pop(node_id, None)
def is_inference_loaded(self, node_id: str) -> bool:
return node_id in self._inference_nodes
@@ -2766,7 +2772,35 @@ class PlatformStore:
"SELECT gpu_index FROM gpu_allocations WHERE node_id=? AND status IN ('allocated','running')",
(node_id,),
).fetchall()
return {int(row["gpu_index"]) for row in rows}
active = {int(row["gpu_index"]) for row in rows}
# Evaluation jobs use the same Compute ProcessManager GPU lock but do
# not have fine-tune allocation rows; derive their selected cards here
# so a training task cannot race onto an evaluation GPU.
for row in conn.execute(
"SELECT payload FROM eval_tasks WHERE status IN ('syncing','queued','running')"
).fetchall():
payload = json_loads(row["payload"], {})
if payload.get("compute_node_id") != node_id:
continue
selected = payload.get("gpu_indices") or payload.get("gpus")
if selected is None and payload.get("gpu_id") is not None:
selected = [payload.get("gpu_id")]
active.update(int(item) for item in selected or [])
# Loaded inference models also reserve only their selected cards.
active.update(self._inference_gpu_indexes.get(node_id, set()))
for row in conn.execute("SELECT payload FROM compare_tasks").fetchall():
payload = json_loads(row["payload"], {})
load_status = payload.get("load_status") or {}
if isinstance(load_status, str):
load_status = json_loads(load_status, {})
for item in load_status.get("loaded_models") or []:
if item.get("node_id") != node_id or item.get("status") not in {"starting", "ready", "running"}:
continue
selected = item.get("gpu_indices") or item.get("gpus")
if selected is None and item.get("gpu_id") is not None:
selected = [item.get("gpu_id")]
active.update(int(gpu) for gpu in selected or [])
return active
def _node_gpu_indexes(self, conn: PgConnection, node: dict[str, Any]) -> set[int]:
rows = conn.execute("SELECT gpu_index FROM gpus WHERE node_id=?", (node["id"],)).fetchall()
@@ -3195,6 +3229,9 @@ class PlatformStore:
# 推理模型占用算力节点同样计入:优先从 compare_tasks 持久化状态派生
# (重启后仍准确),并用内存标记兜底(直接 preload 的模型无 compare 记录)
inference_node_ids = set(self._inference_nodes)
inference_gpu_indexes: dict[str, set[int]] = {
node_id: set(indexes) for node_id, indexes in self._inference_gpu_indexes.items()
}
for ctr in conn.execute("SELECT payload FROM compare_tasks").fetchall():
ls = json_loads(ctr["payload"], {}).get("load_status") or {}
if isinstance(ls, str):
@@ -3204,8 +3241,13 @@ class PlatformStore:
ls = {}
for m in ls.get("loaded_models") or []:
if m.get("status") in {"ready", "running"} and m.get("node_id"):
inference_node_ids.add(m["node_id"])
for nid in inference_node_ids:
node_id = m["node_id"]
selected = m.get("gpu_indices") or m.get("gpus")
if selected:
inference_gpu_indexes.setdefault(node_id, set()).update(int(item) for item in selected)
else:
inference_node_ids.add(node_id)
for nid in set(inference_node_ids) | set(inference_gpu_indexes):
running_map[nid] = running_map.get(nid, 0) + 1
rows = conn.execute("SELECT * FROM compute_nodes ORDER BY scheduler_weight DESC, code").fetchall()
return [
@@ -3434,6 +3476,9 @@ class PlatformStore:
# 推理模型占用的节点:优先从 compare_tasks 持久化状态派生(重启后仍准确),
# 内存标记兜底(直接 preload 的模型无 compare 记录)
inference_node_ids = set(self._inference_nodes)
inference_gpu_indexes: dict[str, set[int]] = {
node_id: set(indexes) for node_id, indexes in self._inference_gpu_indexes.items()
}
for ctr in conn.execute("SELECT payload FROM compare_tasks").fetchall():
ls = json_loads(ctr["payload"], {}).get("load_status") or {}
if isinstance(ls, str):
@@ -3443,7 +3488,12 @@ class PlatformStore:
ls = {}
for m in ls.get("loaded_models") or []:
if m.get("status") in {"ready", "running"} and m.get("node_id"):
inference_node_ids.add(m["node_id"])
node_id = m["node_id"]
selected = m.get("gpu_indices") or m.get("gpus")
if selected:
inference_gpu_indexes.setdefault(node_id, set()).update(int(item) for item in selected)
else:
inference_node_ids.add(node_id)
items = []
for row in rows:
task = next(
@@ -3459,7 +3509,10 @@ class PlatformStore:
t
for t in eval_running
if t.get("compute_node_id") == row["node_id"]
and row["gpu_index"] == (int(t["gpu_id"]) if t.get("gpu_id") is not None else -1)
and row["gpu_index"] in {
int(item)
for item in (t.get("gpu_indices") or t.get("gpus") or ([t["gpu_id"]] if t.get("gpu_id") is not None else []))
}
),
None,
)
@@ -3468,7 +3521,8 @@ class PlatformStore:
eval_task is not None and eval_task.get("status") in {"syncing", "queued"}
)
# Also mark GPU as busy if an inference model is loaded on this node
if row["node_id"] in inference_node_ids and not busy:
inference_on_gpu = row["node_id"] in inference_node_ids or row["gpu_index"] in inference_gpu_indexes.get(row["node_id"], set())
if inference_on_gpu and not busy:
busy = True
reserved = False
memory_used = round(row["memory_total_gb"] * (0.72 if busy else 0.18 if reserved else 0.04), 1)

View File

@@ -240,6 +240,10 @@ class ComputeNodeClient:
async def inference_status(self) -> dict[str, Any]:
return await self._request("GET", "/inference/status", timeout=INFERENCE_STATUS_TIMEOUT)
async def gpu_resources(self) -> list[dict[str, Any]]:
"""Read live per-GPU metrics from this compute node."""
return await self.gpus()
async def inference_unload(self) -> dict[str, Any]:
return await self._request("POST", "/inference/unload", json_data={}, timeout=INFERENCE_UNLOAD_TIMEOUT)

View File

@@ -73,7 +73,8 @@ async def reconcile_inference_loads(store: Any) -> list[dict[str, Any]]:
if node_status == "ready":
item["status"] = "ready"
item.pop("error", None)
store.mark_inference_loaded(node["id"])
selected_gpus = item.get("gpu_indices") or item.get("gpus")
store.mark_inference_loaded(node["id"], selected_gpus)
elif node_status == "error":
item["status"] = "error"
item["error"] = status.get("error") or "model load failed on compute node"