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>
147 lines
4.8 KiB
Python
147 lines
4.8 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)
|
||
|
||
|
||
def test_wait_until_loaded_blocks_until_ready(stub_llamafactory) -> None:
|
||
session = InferenceSession()
|
||
result = session.load("/models/qwen")
|
||
assert result["status"] == "loading"
|
||
# 同步等待后台加载线程完成
|
||
outcome = session.wait_until_loaded(timeout=3.0)
|
||
assert outcome["loaded"] is True
|
||
assert outcome["status"] == "ready"
|
||
|
||
|
||
def test_wait_until_loaded_reports_load_error(stub_failing_llamafactory) -> None:
|
||
session = InferenceSession()
|
||
session.load("/models/bad")
|
||
outcome = session.wait_until_loaded(timeout=3.0)
|
||
assert outcome["loaded"] is False
|
||
assert outcome["status"] == "error"
|
||
assert "boom" in outcome["error"]
|