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>
255 lines
9.5 KiB
Python
255 lines
9.5 KiB
Python
"""
|
||
模型推理异步加载改造的单元测试。
|
||
|
||
覆盖：
|
||
- 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_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"
|