0292bf5138
- eval_runner 等待异步模型加载完成（InferenceSession.wait_until_loaded）， 修复 "model load failed: unknown" - 评测算力节点选择：优先页面选择的节点 / 模型所在节点（_select_eval_node）， 多节点时不再派发到不可达节点导致连接超时 - 前端评测 GPU 选择改为节点感知（节点:GPU 复合值），透传 compute_node_id， 并检查 startEval 结果展示真实错误 - 大模型评价（judge）使用模型记录的真实 API 模型名（api_model）， 避免用平台内部名调用 LLM API 导致 HTTP 400 - 新增后端节点选择与 compute wait_until_loaded 单元测试 Co-Authored-By: Claude <noreply@anthropic.com>
273 lines
11 KiB
Python
273 lines
11 KiB
Python
from __future__ import annotations
|
|
|
|
import threading
|
|
import time
|
|
import uuid
|
|
from typing import Any, Iterator
|
|
|
|
|
|
class InferenceSession:
|
|
"""Manages a loaded model for inference with LLaMA-Factory ChatModel.
|
|
|
|
Model loading is asynchronous: ``load()`` spawns a background daemon thread
|
|
and returns immediately with ``status == "loading"``. ``info()`` (served by
|
|
``/inference/status``) is always responsive, so the platform backend can
|
|
poll loading progress without being blocked by a minutes-long model load —
|
|
which previously froze the whole compute node event loop.
|
|
|
|
State machine: idle -> loading -> ready | error, ready -> idle (unload),
|
|
loading -> idle (cancelled). Long operations (ChatModel build, teardown,
|
|
generation) never run while holding ``_state_lock``; they either run in the
|
|
worker thread or under ``_chat_lock`` only.
|
|
"""
|
|
|
|
def __init__(self) -> None:
|
|
self._state_lock = threading.Lock() # brief state transitions only
|
|
self._chat_lock = threading.Lock() # serialize chat/teardown
|
|
self._status: str = "idle"
|
|
self._error: str = ""
|
|
self._request_id: str = ""
|
|
self._load_args: dict[str, Any] = {}
|
|
self._teardown_old = False # load-while-ready: unload old before loading new
|
|
self._cancel_requested = False # unload-while-loading: tear down after load finishes
|
|
self._load_thread: threading.Thread | None = None
|
|
self._model: Any = None
|
|
self._tokenizer: Any = None
|
|
self._generating_args: dict[str, Any] = {}
|
|
self._model_name: str = ""
|
|
self._adapter_path: str = ""
|
|
self._loaded_at: float = 0.0
|
|
|
|
@property
|
|
def status(self) -> str:
|
|
with self._state_lock:
|
|
return self._status
|
|
|
|
def info(self) -> dict[str, Any]:
|
|
with self._state_lock:
|
|
return {
|
|
"loaded": self._status == "ready",
|
|
"status": self._status,
|
|
"model_name": self._model_name,
|
|
"adapter_path": self._adapter_path,
|
|
"loaded_at": self._loaded_at,
|
|
"request_id": self._request_id,
|
|
"error": self._error,
|
|
}
|
|
|
|
def wait_until_loaded(self, timeout: float | None = None) -> dict[str, Any]:
|
|
"""Wait for an in-flight async load to finish and return its outcome.
|
|
|
|
供同步消费方（如 eval_runner 子进程）使用：``load()`` 立即返回 loading 后，
|
|
调用本方法等待后台加载线程完成，拿到最终的 loaded/error 结果。
|
|
若在 timeout 秒内仍未加载完成，返回 ``status == "loading"`` 并附上超时提示。
|
|
"""
|
|
with self._state_lock:
|
|
thread = self._load_thread
|
|
if thread is not None and thread.is_alive():
|
|
thread.join(timeout=timeout)
|
|
with self._state_lock:
|
|
loaded = self._status == "ready"
|
|
status = self._status
|
|
error = self._error
|
|
if not loaded and status == "loading":
|
|
error = error or f"model load timed out after {timeout or 'N/A'}s"
|
|
return {
|
|
"loaded": loaded,
|
|
"status": status,
|
|
"model_name": self._model_name,
|
|
"adapter_path": self._adapter_path,
|
|
"error": error,
|
|
}
|
|
|
|
def load(
|
|
self,
|
|
model_name_or_path,
|
|
adapter_name_or_path="",
|
|
template="qwen",
|
|
infer_backend="huggingface",
|
|
infer_dtype="auto",
|
|
**kwargs,
|
|
) -> dict[str, Any]:
|
|
with self._state_lock:
|
|
if self._status == "loading":
|
|
# A model is already loading — dedupe, reuse the same request id.
|
|
return {"loaded": False, "status": "loading", "request_id": self._request_id}
|
|
self._teardown_old = self._status == "ready"
|
|
self._status = "loading"
|
|
self._error = ""
|
|
self._request_id = uuid.uuid4().hex[:12]
|
|
self._cancel_requested = False
|
|
self._load_args = {
|
|
"model_name_or_path": model_name_or_path,
|
|
"template": template,
|
|
"infer_backend": infer_backend,
|
|
"infer_dtype": infer_dtype,
|
|
}
|
|
if adapter_name_or_path:
|
|
self._load_args["adapter_name_or_path"] = adapter_name_or_path
|
|
self._load_args.update(kwargs)
|
|
self._model_name = model_name_or_path
|
|
self._adapter_path = adapter_name_or_path
|
|
self._load_thread = threading.Thread(target=self._load_worker, daemon=True)
|
|
self._load_thread.start()
|
|
return {"loaded": False, "status": "loading", "request_id": self._request_id}
|
|
|
|
def _load_worker(self) -> None:
|
|
"""Build the ChatModel off the state lock so info() never blocks."""
|
|
model = None
|
|
tokenizer = None
|
|
generating_args: dict[str, Any] = {}
|
|
error = ""
|
|
try:
|
|
if self._teardown_old:
|
|
self._release_model()
|
|
from llamafactory.chat import ChatModel
|
|
from llamafactory.hparams import get_infer_args
|
|
|
|
args = dict(self._load_args)
|
|
infer_result = get_infer_args(args)
|
|
model = ChatModel(args)
|
|
tokenizer = getattr(model, "tokenizer", None) or model.engine.tokenizer
|
|
generating_args = infer_result[-1]
|
|
if hasattr(generating_args, "__dataclass_fields__"):
|
|
generating_args = {
|
|
k: v for k, v in vars(generating_args).items() if not k.startswith("_")
|
|
}
|
|
else:
|
|
generating_args = dict(generating_args)
|
|
except Exception as exc: # noqa: BLE001 - surface load failure via status
|
|
error = str(exc)
|
|
with self._state_lock:
|
|
if error:
|
|
self._model = None
|
|
self._tokenizer = None
|
|
self._status = "error"
|
|
self._error = error
|
|
return
|
|
if self._cancel_requested:
|
|
# Unload was requested while loading — drop the fresh model.
|
|
model = None
|
|
tokenizer = None
|
|
self._model = None
|
|
self._tokenizer = None
|
|
self._status = "idle"
|
|
return
|
|
self._model = model
|
|
self._tokenizer = tokenizer
|
|
self._generating_args = generating_args
|
|
self._loaded_at = time.time()
|
|
self._status = "ready"
|
|
|
|
def _release_model(self) -> None:
|
|
with self._chat_lock:
|
|
with self._state_lock:
|
|
self._status = "unloading"
|
|
model = self._model
|
|
self._model = None
|
|
self._tokenizer = None
|
|
if model is not None:
|
|
try:
|
|
del model
|
|
except Exception: # noqa: BLE001 - best-effort teardown
|
|
pass
|
|
# 强制释放 PyTorch CUDA 缓存，真正归还 GPU 显存
|
|
try:
|
|
import gc
|
|
|
|
gc.collect()
|
|
import torch
|
|
|
|
if torch.cuda.is_available():
|
|
torch.cuda.empty_cache()
|
|
torch.cuda.synchronize()
|
|
except Exception: # noqa: BLE001 - teardown must not raise
|
|
pass
|
|
with self._state_lock:
|
|
self._status = "idle"
|
|
self._model_name = ""
|
|
self._adapter_path = ""
|
|
self._loaded_at = 0.0
|
|
self._error = ""
|
|
|
|
def unload(self) -> dict[str, Any]:
|
|
with self._state_lock:
|
|
if self._status == "loading":
|
|
# Ask the worker to tear down right after the load finishes.
|
|
self._cancel_requested = True
|
|
return {"unloaded": False, "status": "cancelling", "request_id": self._request_id}
|
|
was_ready = self._status == "ready"
|
|
if was_ready:
|
|
self._release_model()
|
|
else:
|
|
with self._state_lock:
|
|
self._model = None
|
|
self._tokenizer = None
|
|
self._status = "idle"
|
|
self._model_name = ""
|
|
self._adapter_path = ""
|
|
self._loaded_at = 0.0
|
|
self._error = ""
|
|
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]:
|
|
with self._chat_lock:
|
|
with self._state_lock:
|
|
if self._status == "loading":
|
|
return {
|
|
"error": f"model is still loading (request_id={self._request_id}); please retry",
|
|
"response": "",
|
|
}
|
|
if self._status == "error":
|
|
return {"error": f"model load failed: {self._error}", "response": ""}
|
|
if self._status != "ready" or self._model is None:
|
|
return {"error": "model not loaded", "response": ""}
|
|
try:
|
|
generate_kwargs = {
|
|
"temperature": temperature,
|
|
"top_p": top_p,
|
|
"max_new_tokens": max_new_tokens,
|
|
"do_sample": do_sample,
|
|
}
|
|
generate_kwargs.update(kwargs)
|
|
system = next((m["content"] for m in messages if m["role"] == "system"), None)
|
|
user_messages = [m for m in messages if m["role"] != "system"]
|
|
responses = []
|
|
for response in self._model.stream_chat(user_messages, system=system, **generate_kwargs):
|
|
responses.append(response)
|
|
full_response = "".join(str(r) for r in responses)
|
|
return {"response": full_response}
|
|
except Exception as exc: # noqa: BLE001 - return generation error to caller
|
|
return {"error": str(exc), "response": ""}
|
|
|
|
def chat_stream(self, messages, **kwargs) -> Iterator[str]:
|
|
with self._chat_lock:
|
|
with self._state_lock:
|
|
if self._status == "loading":
|
|
yield 'data: {"error": "model is still loading; please retry"}\n\n'
|
|
return
|
|
if self._status == "error":
|
|
yield 'data: {"error": "model load failed: ' + str(self._error) + '"}\n\n'
|
|
return
|
|
if self._status != "ready" or self._model is None:
|
|
yield 'data: {"error": "model not loaded"}\n\n'
|
|
return
|
|
try:
|
|
generate_kwargs = {**kwargs}
|
|
system = next((m["content"] for m in messages if m["role"] == "system"), None)
|
|
user_messages = [m for m in messages if m["role"] != "system"]
|
|
for new_text in self._model.stream_chat(user_messages, system=system, **generate_kwargs):
|
|
yield new_text
|
|
except Exception as exc: # noqa: BLE001 - stream error as SSE event
|
|
yield 'data: {"error": "' + str(exc) + '"}\n\n'
|
|
|
|
|
|
_inference_session = None
|
|
|
|
|
|
def get_inference_session() -> InferenceSession:
|
|
global _inference_session
|
|
if _inference_session is None:
|
|
_inference_session = InferenceSession()
|
|
return _inference_session
|