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:
@@ -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]:
|
||||
|
||||
Reference in New Issue
Block a user