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

@@ -1,5 +1,6 @@
from __future__ import annotations
import os
import threading
import time
import uuid
@@ -36,6 +37,7 @@ class InferenceSession:
self._generating_args: dict[str, Any] = {}
self._model_name: str = ""
self._adapter_path: str = ""
self._gpu_indices: list[int] = []
self._loaded_at: float = 0.0
@property
@@ -53,6 +55,7 @@ class InferenceSession:
"loaded_at": self._loaded_at,
"request_id": self._request_id,
"error": self._error,
"gpu_indices": list(self._gpu_indices),
}
def wait_until_loaded(self, timeout: float | None = None) -> dict[str, Any]:
@@ -87,8 +90,12 @@ class InferenceSession:
template="qwen",
infer_backend="huggingface",
infer_dtype="auto",
gpu_indices=None,
**kwargs,
) -> dict[str, Any]:
requested_gpus = sorted({int(item) for item in (gpu_indices or [])})
if any(item < 0 for item in requested_gpus):
return {"loaded": False, "status": "error", "error": "GPU index must be non-negative"}
with self._state_lock:
if self._status == "loading":
# A model is already loading — dedupe, reuse the same request id.
@@ -97,6 +104,7 @@ class InferenceSession:
self._status = "loading"
self._error = ""
self._request_id = uuid.uuid4().hex[:12]
self._gpu_indices = requested_gpus
self._cancel_requested = False
self._load_args = {
"model_name_or_path": model_name_or_path,
@@ -115,13 +123,21 @@ class InferenceSession:
def _load_worker(self) -> None:
"""Build the ChatModel off the state lock so info() never blocks."""
with self._state_lock:
requested_gpus = list(self._gpu_indices)
model = None
tokenizer = None
generating_args: dict[str, Any] = {}
error = ""
previous_visible_devices = os.environ.get("CUDA_VISIBLE_DEVICES")
try:
# Set visibility before LLaMA-Factory/PyTorch initializes CUDA.
if requested_gpus:
os.environ["CUDA_VISIBLE_DEVICES"] = ",".join(str(item) for item in requested_gpus)
if self._teardown_old:
self._release_model()
with self._state_lock:
self._gpu_indices = requested_gpus
from llamafactory.chat import ChatModel
from llamafactory.hparams import get_infer_args
@@ -138,6 +154,12 @@ class InferenceSession:
generating_args = dict(generating_args)
except Exception as exc: # noqa: BLE001 - surface load failure via status
error = str(exc)
finally:
if requested_gpus:
if previous_visible_devices is None:
os.environ.pop("CUDA_VISIBLE_DEVICES", None)
else:
os.environ["CUDA_VISIBLE_DEVICES"] = previous_visible_devices
with self._state_lock:
if error:
self._model = None
@@ -152,6 +174,7 @@ class InferenceSession:
self._model = None
self._tokenizer = None
self._status = "idle"
self._gpu_indices = []
return
self._model = model
self._tokenizer = tokenizer
@@ -189,6 +212,7 @@ class InferenceSession:
self._adapter_path = ""
self._loaded_at = 0.0
self._error = ""
self._gpu_indices = []
def unload(self) -> dict[str, Any]:
with self._state_lock:
@@ -208,6 +232,7 @@ class InferenceSession:
self._adapter_path = ""
self._loaded_at = 0.0
self._error = ""
self._gpu_indices = []
return {"unloaded": True, "status": "idle"}
def chat(self, messages, temperature=0.95, top_p=0.7, max_new_tokens=1024, do_sample=True, **kwargs) -> dict[str, Any]: