YG_FT/compute/tests/test_inference_session.py

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"]