0271942ba5
模型推理全异步化改造： - 计算节点 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>
128 lines
4.1 KiB
Python
128 lines
4.1 KiB
Python
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)
|