Files
YG_FT/compute/tests/test_inference_session.py
wuyongtao 0292bf5138 fix: 模型评测异步加载等待与多节点路由修复
- 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>
2026-08-04 18:21:16 +08:00

147 lines
4.8 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
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_argsworker 会转成 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"]