YG_FT/backend/tests/test_compare_inference_async.py

"""
模型推理异步加载改造的单元测试。

覆盖：
- model_compare_load：异步派发，立即返回 starting + 节点信息（不等待加载完成）
- model_compare_delete：先删记录，卸载失败也不阻塞删除
- reconcile_inference_loads：starting -> ready/error/idle/不可达的状态迁移与封顶
- _unload_from_compute_node：任务感知，只命中记录中的节点
"""

from __future__ import annotations

import asyncio
from types import SimpleNamespace
from typing import Any

from app.api.v1.endpoints.platform import model_compare_delete, model_compare_load
import app.api.v1.endpoints.platform as platform
from app.modules.compute_gateway.client import ComputeNodeClient
from app.modules.compute_gateway.sync import MAX_STARTING_ATTEMPTS, reconcile_inference_loads


class FakeInferenceStore:
    """内存 store，仅实现推理加载/对账用到的接口。"""

    def __init__(self, tasks: list[dict[str, Any]] | None = None, nodes: list[dict[str, Any]] | None = None) -> None:
        self._tasks: dict[str, dict[str, Any]] = {t["id"]: dict(t) for t in (tasks or [])}
        self._nodes = nodes or []
        self._inference_nodes: set[str] = set()

    def compare_task(self, task_id: str) -> dict[str, Any]:
        if task_id not in self._tasks:
            raise KeyError(task_id)
        return dict(self._tasks[task_id])

    def compare_tasks(self) -> list[dict[str, Any]]:
        return [dict(t) for t in self._tasks.values()]

    def update_compare_task(self, task_id: str, payload: dict[str, Any]) -> dict[str, Any]:
        current = self._tasks[task_id]
        merged = {**current, **payload, "id": task_id}
        self._tasks[task_id] = merged
        return dict(merged)

    def delete_compare_task(self, task_id: str) -> None:
        self._tasks.pop(task_id, None)

    def compute_nodes(self) -> list[dict[str, Any]]:
        return [dict(n) for n in self._nodes]

    def model(self, model_id: str) -> dict[str, Any]:
        raise KeyError(model_id)

    def trained_models(self) -> list[dict[str, Any]]:
        return []

    def mark_inference_loaded(self, node_id: str) -> None:
        self._inference_nodes.add(node_id)

    def mark_inference_unloaded(self, node_id: str) -> None:
        self._inference_nodes.discard(node_id)

    def is_inference_loaded(self, node_id: str) -> bool:
        return node_id in self._inference_nodes


def _node(node_id: str, code: str = "") -> dict[str, Any]:
    return {
        "id": node_id,
        "code": code or node_id,
        "name": code or node_id,
        "api_base_url": f"http://{code or node_id}:19100",
        "enabled": True,
        "scheduler_status": "online",
    }


def _task(task_id: str, *, node_id: str | None = None, load_status: dict[str, Any] | None = None) -> dict[str, Any]:
    return {
        "id": task_id,
        "name": f"task-{task_id}",
        "status": "pending",
        "models": [
            {"model_id": "m_1", "model_name": "qwen", "model_path": "/models/qwen", "node_id": node_id}
        ],
        "load_status": load_status or {"loaded_models": []},
    }


async def _fake_inference_load(self, payload: dict[str, Any]) -> dict[str, Any]:
    return {"loaded": False, "status": "loading", "request_id": "req-1"}


async def _fake_inference_unload(self) -> dict[str, Any]:
    return {"unloaded": True, "status": "idle"}


def _patch_store(monkeypatch, store: FakeInferenceStore) -> None:
    monkeypatch.setattr(platform, "get_platform_store", lambda: store)
    monkeypatch.setattr(platform, "get_settings", lambda: SimpleNamespace(compute_mode="real"))


def test_select_eval_node_prefers_model_node(monkeypatch) -> None:
    from app.api.v1.endpoints.platform import _select_eval_node

    store = FakeInferenceStore(nodes=[_node("n1"), _node("n2")])
    # 指定模型所在节点时优先返回该节点
    assert _select_eval_node(store, "n2")["id"] == "n2"
    # 无指定节点时回退到第一个在线节点
    assert _select_eval_node(store, None)["id"] == "n1"


def test_select_eval_node_returns_none_when_model_node_offline(monkeypatch) -> None:
    from app.api.v1.endpoints.platform import _select_eval_node

    nodes = [_node("n1"), _node("n2")]
    nodes[1]["enabled"] = False
    store = FakeInferenceStore(nodes=nodes)
    # 模型所在节点不可用 → 明确失败，不派发到其它节点
    assert _select_eval_node(store, "n2") is None
    # 无指定节点时仍回退第一个在线节点
    assert _select_eval_node(store, None)["id"] == "n1"


def test_model_compare_load_dispatches_and_returns_starting(monkeypatch) -> None:
    store = FakeInferenceStore(tasks=[_task("t1", node_id="n1")], nodes=[_node("n1")])
    _patch_store(monkeypatch, store)
    monkeypatch.setattr(ComputeNodeClient, "inference_load", _fake_inference_load)

    result = asyncio.run(model_compare_load("t1"))
    assert result["code"] == 0
    updated = result["data"]
    assert updated["status"] == "starting"
    items = updated["load_status"]["loaded_models"]
    assert items[0]["status"] == "starting"
    assert items[0]["node_id"] == "n1"
    assert "n1" in store._inference_nodes


def test_model_compare_load_marks_error_when_all_nodes_fail(monkeypatch) -> None:
    store = FakeInferenceStore(tasks=[_task("t1", node_id="n1")], nodes=[_node("n1")])
    _patch_store(monkeypatch, store)

    async def _raise(self, payload: dict[str, Any]) -> dict[str, Any]:
        raise RuntimeError("conn refused")

    monkeypatch.setattr(ComputeNodeClient, "inference_load", _raise)

    result = asyncio.run(model_compare_load("t1"))
    updated = result["data"]
    assert updated["status"] == "failed"
    assert updated["load_status"]["loaded_models"][0]["status"] == "error"
    assert "conn refused" in updated["load_status"]["loaded_models"][0]["error"]


def test_model_compare_delete_removes_record_even_if_unload_raises(monkeypatch) -> None:
    task = _task(
        "t1",
        node_id="n1",
        load_status={"loaded_models": [{"model_id": "m_1", "status": "ready", "node_id": "n1"}]},
    )
    store = FakeInferenceStore(tasks=[task], nodes=[_node("n1")])
    _patch_store(monkeypatch, store)

    async def _raise(self) -> dict[str, Any]:
        raise RuntimeError("unload boom")

    monkeypatch.setattr(ComputeNodeClient, "inference_unload", _raise)

    result = asyncio.run(model_compare_delete("t1"))
    assert result["data"] == {"deleted": "t1"}
    assert "t1" not in store._tasks
    # finally 中仍清掉了节点标记
    assert "n1" not in store._inference_nodes


def test_unload_from_compute_node_only_hits_recorded_node(monkeypatch) -> None:
    task = _task(
        "t1",
        load_status={"loaded_models": [{"model_id": "m_1", "status": "ready", "node_id": "n1"}]},
    )
    store = FakeInferenceStore(tasks=[task], nodes=[_node("n1"), _node("n2")])
    _patch_store(monkeypatch, store)
    monkeypatch.setattr(ComputeNodeClient, "inference_unload", _fake_inference_unload)

    from app.api.v1.endpoints.platform import _unload_from_compute_node

    result = asyncio.run(_unload_from_compute_node(store, task=task))
    assert result["unloaded"] is True
    # 只命中任务记录中的节点 n1，n2 未被卸载
    assert [r["node_id"] for r in result["nodes"]] == ["n1"]
    assert "n1" not in store._inference_nodes


async def _status_ready(self) -> dict[str, Any]:
    return {"loaded": True, "status": "ready", "model_name": "qwen"}


def test_reconcile_transitions_starting_to_ready(monkeypatch) -> None:
    task = _task(
        "t1",
        node_id="n1",
        load_status={"loaded_models": [{"model_id": "m_1", "status": "starting", "node_id": "n1"}]},
    )
    store = FakeInferenceStore(tasks=[task], nodes=[_node("n1")])
    monkeypatch.setattr(ComputeNodeClient, "inference_status", _status_ready)

    reconciled = asyncio.run(reconcile_inference_loads(store))
    assert reconciled == [{"task_id": "t1", "status": "loaded"}]
    updated = store._tasks["t1"]
    assert updated["status"] == "loaded"
    assert updated["load_status"]["loaded_models"][0]["status"] == "ready"
    assert "n1" in store._inference_nodes


def test_reconcile_transitions_to_error_and_failed(monkeypatch) -> None:
    async def _status_error(self) -> dict[str, Any]:
        return {"loaded": False, "status": "error", "error": "CUDA out of memory"}

    task = _task(
        "t1",
        node_id="n1",
        load_status={"loaded_models": [{"model_id": "m_1", "status": "starting", "node_id": "n1"}]},
    )
    store = FakeInferenceStore(tasks=[task], nodes=[_node("n1")])
    monkeypatch.setattr(ComputeNodeClient, "inference_status", _status_error)

    reconciled = asyncio.run(reconcile_inference_loads(store))
    assert reconciled == [{"task_id": "t1", "status": "failed"}]
    item = store._tasks["t1"]["load_status"]["loaded_models"][0]
    assert item["status"] == "error"
    assert "CUDA out of memory" in item["error"]
    assert "n1" not in store._inference_nodes


def test_reconcile_idle_marks_model_disappeared(monkeypatch) -> None:
    async def _status_idle(self) -> dict[str, Any]:
        return {"loaded": False, "status": "idle"}

    task = _task(
        "t1",
        node_id="n1",
        load_status={"loaded_models": [{"model_id": "m_1", "status": "starting", "node_id": "n1"}]},
    )
    store = FakeInferenceStore(tasks=[task], nodes=[_node("n1")])
    monkeypatch.setattr(ComputeNodeClient, "inference_status", _status_idle)

    asyncio.run(reconcile_inference_loads(store))
    item = store._tasks["t1"]["load_status"]["loaded_models"][0]
    assert item["status"] == "error"
    assert "disappeared" in item["error"]
    assert store._tasks["t1"]["status"] == "failed"


def test_reconcile_unreachable_node_flips_to_error_after_cap(monkeypatch) -> None:
    async def _raise(self) -> dict[str, Any]:
        raise RuntimeError("conn refused")

    task = _task(
        "t1",
        node_id="n1",
        load_status={"loaded_models": [{"model_id": "m_1", "status": "starting", "node_id": "n1"}]},
    )
    store = FakeInferenceStore(tasks=[task], nodes=[_node("n1")])
    monkeypatch.setattr(ComputeNodeClient, "inference_status", _raise)

    # 每次轮询前重置节流时间戳，逐次推进 load_attempts 到封顶
    for _ in range(MAX_STARTING_ATTEMPTS):
        item = store._tasks["t1"]["load_status"]["loaded_models"][0]
        item["last_polled_at"] = 0
        asyncio.run(reconcile_inference_loads(store))

    item = store._tasks["t1"]["load_status"]["loaded_models"][0]
    assert item["status"] == "error"
    assert "unreachable" in item["error"]
    assert store._tasks["t1"]["status"] == "failed"