YG-Soft/YG_FT
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: 模型推理异步加载与对话链路修复，同步基线

Browse Source 
模型推理全异步化改造：
- 计算节点 InferenceSession 改为后台线程异步加载模型，load 立即返回，
  加载期间事件循环保持响应（/inference/status 与 /health 不阻塞）
- 后端模型加载改为异步派发 + 轮询对账器（reconcile_inference_loads），
  任务状态由 starting 自动推进到 ready/error，解决多节点启动超时
  （timeout of 120000ms exceeded）
- 推理删除/卸载改为任务感知 + 短超时，删除先删记录再 best-effort 卸载，
  不再被不可达节点阻塞；同节点新模型替换旧任务标记失效
- 流式对话透传 task_id/node_id 路由到真正加载模型的算力节点，
  useStreamChat 解析 SSE 错误帧以干净文案展示
- 对话历史按任务 id 本地持久化，退出重进可恢复；移除页脚提示文本
- 新增后端推理异步加载与计算节点异步状态机单元测试

Co-Authored-By: Claude <noreply@anthropic.com>
This commit is contained in:
wuyongtao
2026-08-04 16:59:34 +08:00
parent 250e060271
commit 0271942ba5
21 changed files with 1272 additions and 245 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
compute/api/main.py
View File

@@ -1,5 +1,6 @@
from __future__ import annotations
import asyncio
import json
import os
import math
@@ -717,7 +718,9 @@ def create_app() -> FastAPI:
@app.post(f"{route_prefix}/inference/unload")
async def inference_unload() -> dict[str, Any]:
"""Unload the currently loaded model and free GPU memory."""
return get_inference_session().unload()
# Teardown (gc.collect + cuda.empty_cache) can take a while; run it off
# the event loop so /health and /inference/status stay responsive.
return await asyncio.to_thread(get_inference_session().unload)
@app.get(f"{route_prefix}/inference/status")
async def inference_status() -> dict[str, Any]:
@@ -737,7 +740,10 @@ def create_app() -> FastAPI:
messages = payload.get("messages") or []
if not messages:
raise HTTPException(status_code=400, detail="messages is required")
result = get_inference_session().chat(
# Generation is long-running; run it in a thread so the event loop keeps
# serving /inference/status and /health during inference.
result = await asyncio.to_thread(
get_inference_session().chat,
messages=messages,
temperature=float(payload.get("temperature", 0.95)),
top_p=float(payload.get("top_p", 0.7)),

264
compute/engines/llama_factory/inference.py
View File

@@ -2,114 +2,208 @@ from __future__ import annotations
import threading
import time
from typing import Any
import uuid
from typing import Any, Iterator
class InferenceSession:
"""Manages a loaded model for inference with LLaMA-Factory ChatModel."""
"""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._lock = threading.Lock()
self._loaded_at: float = 0.0
self._status: str = "idle"
@property
def status(self) -> str:
return self._status
@property
def model_name(self) -> str:
return self._model_name
@property
def adapter_path(self) -> str:
return self._adapter_path
@property
def loaded_at(self) -> float:
return self._loaded_at
with self._state_lock:
return self._status
def info(self) -> dict[str, Any]:
return {
"loaded": self._status == "ready",
"status": self._status,
"model_name": self._model_name,
"adapter_path": self._adapter_path,
"loaded_at": self._loaded_at,
}
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 load(self, model_name_or_path, adapter_name_or_path="", template="qwen", infer_backend="huggingface", infer_dtype="auto", **kwargs):
with self._lock:
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":
return {"loaded": False, "error": "model is already loading"}
if self._status == "ready":
self.unload()
# 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
try:
from llamafactory.chat import ChatModel
from llamafactory.hparams import get_infer_args
args = {"model_name_or_path": model_name_or_path, "template": template, "infer_backend": infer_backend, "infer_dtype": infer_dtype}
if adapter_name_or_path:
args["adapter_name_or_path"] = adapter_name_or_path
args.update(kwargs)
infer_result = get_infer_args(args)
# ChatModel internally re-parses the args dict via get_infer_args,
# so pass the original args (not the parsed dataclass objects).
self._model = ChatModel(args)
self._tokenizer = getattr(self._model, 'tokenizer', None) or self._model.engine.tokenizer
# Extract generating_args (last element) for later use in chat()
generating_args = infer_result[-1]
if hasattr(generating_args, '__dataclass_fields__'):
self._generating_args = {k: v for k, v in vars(generating_args).items()
if not k.startswith('_')}
else:
self._generating_args = dict(generating_args)
self._loaded_at = time.time()
self._status = "ready"
return {"loaded": True, "status": "ready"}
except Exception as exc:
self._status = "error"
self._model = None
return {"loaded": False, "status": "error", "error": str(exc)}
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 unload(self):
with self._lock:
if self._model is not None:
try:
del self._model
except Exception:
pass
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:
except Exception: # noqa: BLE001 - teardown must not raise
pass
self._status = "idle"
self._model_name = ""
self._adapter_path = ""
self._loaded_at = 0.0
return {"unloaded": True}
with self._state_lock:
self._status = "idle"
self._model_name = ""
self._adapter_path = ""
self._loaded_at = 0.0
self._error = ""
def chat(self, messages, temperature=0.95, top_p=0.7, max_new_tokens=1024, do_sample=True, **kwargs):
with self._lock:
if self._status != "ready" or self._model is None:
return {"error": "model not loaded", "response": ""}
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 = {
"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"]
@@ -118,27 +212,35 @@ class InferenceSession:
responses.append(response)
full_response = "".join(str(r) for r in responses)
return {"response": full_response}
except Exception as exc:
except Exception as exc: # noqa: BLE001 - return generation error to caller
return {"error": str(exc), "response": ""}
def chat_stream(self, messages, **kwargs):
with self._lock:
if self._status != "ready" or self._model is None:
yield 'data: {"error": "model not loaded"}\n\n'
return
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:
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():
def get_inference_session() -> InferenceSession:
global _inference_session
if _inference_session is None:
_inference_session = InferenceSession()

127
compute/tests/test_inference_session.py Normal file
View File

@@ -0,0 +1,127 @@
from __future__ import annotations
import sys
import time
import types
from typing import Any
import pytest
from compute.engines.llama_factory.inference import InferenceSession
# 模拟模型加载耗时，用于验证 load() 立即返回、info() 不阻塞
LOAD_DELAY = 0.2
class FakeChatModel:
def __init__(self, args: dict[str, Any]) -> None:
time.sleep(LOAD_DELAY)
self.tokenizer = object()
self.engine = types.SimpleNamespace(tokenizer=object())
self._output = "hello from model"
def stream_chat(self, *args, **kwargs):
for _ in range(1):
yield self._output
class FailingChatModel:
def __init__(self, args: dict[str, Any]) -> None:
time.sleep(LOAD_DELAY)
raise RuntimeError("boom: fake load failure")
def _get_infer_args(args: dict[str, Any]) -> list[Any]:
# 最后一个元素为 generating_args，worker 会转成 dict
return [None, None, {"temperature": 0.7}]
def _install_llamafactory(monkeypatch, chat_model: type) -> None:
llmf = types.ModuleType("llamafactory")
chat_mod = types.ModuleType("llamafactory.chat")
hparams_mod = types.ModuleType("llamafactory.hparams")
chat_mod.ChatModel = chat_model
hparams_mod.get_infer_args = _get_infer_args
llmf.chat = chat_mod
llmf.hparams = hparams_mod
monkeypatch.setitem(sys.modules, "llamafactory", llmf)
monkeypatch.setitem(sys.modules, "llamafactory.chat", chat_mod)
monkeypatch.setitem(sys.modules, "llamafactory.hparams", hparams_mod)
@pytest.fixture
def stub_llamafactory(monkeypatch) -> None:
_install_llamafactory(monkeypatch, FakeChatModel)
@pytest.fixture
def stub_failing_llamafactory(monkeypatch) -> None:
_install_llamafactory(monkeypatch, FailingChatModel)
def _wait_for_status(session: InferenceSession, status: str, timeout: float = 3.0) -> bool:
deadline = time.time() + timeout
while time.time() < deadline:
if session.info()["status"] == status:
return True
time.sleep(0.02)
return False
def test_load_returns_immediately_then_ready(stub_llamafactory) -> None:
session = InferenceSession()
started = time.time()
result = session.load("/models/qwen")
assert result["status"] == "loading"
assert result["loaded"] is False
assert result["request_id"]
# 在慢加载完成前就返回，且 info() 加载期间可响应
assert time.time() - started < LOAD_DELAY
assert session.info()["status"] == "loading"
assert _wait_for_status(session, "ready")
info = session.info()
assert info["loaded"] is True
assert info["status"] == "ready"
assert info["model_name"] == "/models/qwen"
def test_second_load_while_loading_deduped(stub_llamafactory) -> None:
session = InferenceSession()
r1 = session.load("/models/a")
r2 = session.load("/models/b")
assert r2["status"] == "loading"
assert r2["request_id"] == r1["request_id"]
assert _wait_for_status(session, "ready")
assert session.info()["status"] == "ready"
def test_load_error_surfaces_in_status(stub_failing_llamafactory) -> None:
session = InferenceSession()
session.load("/models/bad")
assert _wait_for_status(session, "error")
assert "boom" in session.info()["error"]
def test_unload_while_loading_cancels(stub_llamafactory) -> None:
session = InferenceSession()
session.load("/models/qwen")
result = session.unload()
assert result["status"] == "cancelling"
assert _wait_for_status(session, "idle")
def test_chat_while_loading_returns_loading_error(stub_llamafactory) -> None:
session = InferenceSession()
session.load("/models/qwen")
out = session.chat([{"role": "user", "content": "hi"}])
assert "still loading" in (out.get("error") or "")
assert _wait_for_status(session, "ready")
out = session.chat([{"role": "user", "content": "hi"}])
assert out.get("response") == "hello from model"
def test_chat_stream_while_loading_yields_error(stub_llamafactory) -> None:
session = InferenceSession()
session.load("/models/qwen")
chunks = list(session.chat_stream([{"role": "user", "content": "hi"}]))
assert any("still loading" in c for c in chunks)