feat: 模型推理异步加载与对话链路修复，同步基线

模型推理全异步化改造： - 计算节点 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>
2026-08-04 16:59:34 +08:00
parent 250e060271
commit 0271942ba5
21 changed files with 1272 additions and 245 deletions
--- a/compute/engines/llama_factory/inference.py
+++ b/compute/engines/llama_factory/inference.py
@@ -2,114 +2,208 @@ from __future__ import annotations

 import threading
 import time
-from typing import Any
+import uuid
+from typing import Any, Iterator


 class InferenceSession:
-    """Manages a loaded model for inference with LLaMA-Factory ChatModel."""
+    """Manages a loaded model for inference with LLaMA-Factory ChatModel.
+
+    Model loading is asynchronous: ``load()`` spawns a background daemon thread
+    and returns immediately with ``status == "loading"``. ``info()`` (served by
+    ``/inference/status``) is always responsive, so the platform backend can
+    poll loading progress without being blocked by a minutes-long model load —
+    which previously froze the whole compute node event loop.
+
+    State machine: idle -> loading -> ready | error, ready -> idle (unload),
+    loading -> idle (cancelled). Long operations (ChatModel build, teardown,
+    generation) never run while holding ``_state_lock``; they either run in the
+    worker thread or under ``_chat_lock`` only.
+    """

    def __init__(self) -> None:
+        self._state_lock = threading.Lock()  # brief state transitions only
+        self._chat_lock = threading.Lock()  # serialize chat/teardown
+        self._status: str = "idle"
+        self._error: str = ""
+        self._request_id: str = ""
+        self._load_args: dict[str, Any] = {}
+        self._teardown_old = False  # load-while-ready: unload old before loading new
+        self._cancel_requested = False  # unload-while-loading: tear down after load finishes
+        self._load_thread: threading.Thread | None = None
        self._model: Any = None
        self._tokenizer: Any = None
        self._generating_args: dict[str, Any] = {}
        self._model_name: str = ""
        self._adapter_path: str = ""
-        self._lock = threading.Lock()
        self._loaded_at: float = 0.0
-        self._status: str = "idle"

    @property
    def status(self) -> str:
-        return self._status
-
-    @property
-    def model_name(self) -> str:
-        return self._model_name
-
-    @property
-    def adapter_path(self) -> str:
-        return self._adapter_path
-
-    @property
-    def loaded_at(self) -> float:
-        return self._loaded_at
+        with self._state_lock:
+            return self._status

    def info(self) -> dict[str, Any]:
-        return {
-            "loaded": self._status == "ready",
-            "status": self._status,
-            "model_name": self._model_name,
-            "adapter_path": self._adapter_path,
-            "loaded_at": self._loaded_at,
-        }
+        with self._state_lock:
+            return {
+                "loaded": self._status == "ready",
+                "status": self._status,
+                "model_name": self._model_name,
+                "adapter_path": self._adapter_path,
+                "loaded_at": self._loaded_at,
+                "request_id": self._request_id,
+                "error": self._error,
+            }

-    def load(self, model_name_or_path, adapter_name_or_path="", template="qwen", infer_backend="huggingface", infer_dtype="auto", **kwargs):
-        with self._lock:
+    def load(
+        self,
+        model_name_or_path,
+        adapter_name_or_path="",
+        template="qwen",
+        infer_backend="huggingface",
+        infer_dtype="auto",
+        **kwargs,
+    ) -> dict[str, Any]:
+        with self._state_lock:
            if self._status == "loading":
-                return {"loaded": False, "error": "model is already loading"}
-            if self._status == "ready":
-                self.unload()
+                # A model is already loading — dedupe, reuse the same request id.
+                return {"loaded": False, "status": "loading", "request_id": self._request_id}
+            self._teardown_old = self._status == "ready"
            self._status = "loading"
+            self._error = ""
+            self._request_id = uuid.uuid4().hex[:12]
+            self._cancel_requested = False
+            self._load_args = {
+                "model_name_or_path": model_name_or_path,
+                "template": template,
+                "infer_backend": infer_backend,
+                "infer_dtype": infer_dtype,
+            }
+            if adapter_name_or_path:
+                self._load_args["adapter_name_or_path"] = adapter_name_or_path
+            self._load_args.update(kwargs)
            self._model_name = model_name_or_path
            self._adapter_path = adapter_name_or_path
-            try:
-                from llamafactory.chat import ChatModel
-                from llamafactory.hparams import get_infer_args
-                args = {"model_name_or_path": model_name_or_path, "template": template, "infer_backend": infer_backend, "infer_dtype": infer_dtype}
-                if adapter_name_or_path:
-                    args["adapter_name_or_path"] = adapter_name_or_path
-                args.update(kwargs)
-                infer_result = get_infer_args(args)
-                # ChatModel internally re-parses the args dict via get_infer_args,
-                # so pass the original args (not the parsed dataclass objects).
-                self._model = ChatModel(args)
-                self._tokenizer = getattr(self._model, 'tokenizer', None) or self._model.engine.tokenizer
-                # Extract generating_args (last element) for later use in chat()
-                generating_args = infer_result[-1]
-                if hasattr(generating_args, '__dataclass_fields__'):
-                    self._generating_args = {k: v for k, v in vars(generating_args).items()
-                                              if not k.startswith('_')}
-                else:
-                    self._generating_args = dict(generating_args)
-                self._loaded_at = time.time()
-                self._status = "ready"
-                return {"loaded": True, "status": "ready"}
-            except Exception as exc:
-                self._status = "error"
-                self._model = None
-                return {"loaded": False, "status": "error", "error": str(exc)}
+            self._load_thread = threading.Thread(target=self._load_worker, daemon=True)
+            self._load_thread.start()
+            return {"loaded": False, "status": "loading", "request_id": self._request_id}

-    def unload(self):
-        with self._lock:
-            if self._model is not None:
-                try:
-                    del self._model
-                except Exception:
-                    pass
+    def _load_worker(self) -> None:
+        """Build the ChatModel off the state lock so info() never blocks."""
+        model = None
+        tokenizer = None
+        generating_args: dict[str, Any] = {}
+        error = ""
+        try:
+            if self._teardown_old:
+                self._release_model()
+            from llamafactory.chat import ChatModel
+            from llamafactory.hparams import get_infer_args
+
+            args = dict(self._load_args)
+            infer_result = get_infer_args(args)
+            model = ChatModel(args)
+            tokenizer = getattr(model, "tokenizer", None) or model.engine.tokenizer
+            generating_args = infer_result[-1]
+            if hasattr(generating_args, "__dataclass_fields__"):
+                generating_args = {
+                    k: v for k, v in vars(generating_args).items() if not k.startswith("_")
+                }
+            else:
+                generating_args = dict(generating_args)
+        except Exception as exc:  # noqa: BLE001 - surface load failure via status
+            error = str(exc)
+        with self._state_lock:
+            if error:
                self._model = None
                self._tokenizer = None
+                self._status = "error"
+                self._error = error
+                return
+            if self._cancel_requested:
+                # Unload was requested while loading — drop the fresh model.
+                model = None
+                tokenizer = None
+                self._model = None
+                self._tokenizer = None
+                self._status = "idle"
+                return
+            self._model = model
+            self._tokenizer = tokenizer
+            self._generating_args = generating_args
+            self._loaded_at = time.time()
+            self._status = "ready"
+
+    def _release_model(self) -> None:
+        with self._chat_lock:
+            with self._state_lock:
+                self._status = "unloading"
+                model = self._model
+                self._model = None
+                self._tokenizer = None
+            if model is not None:
+                try:
+                    del model
+                except Exception:  # noqa: BLE001 - best-effort teardown
+                    pass
                # 强制释放 PyTorch CUDA 缓存，真正归还 GPU 显存
                try:
                    import gc
+
                    gc.collect()
                    import torch
+
                    if torch.cuda.is_available():
                        torch.cuda.empty_cache()
                        torch.cuda.synchronize()
-                except Exception:
+                except Exception:  # noqa: BLE001 - teardown must not raise
                    pass
-            self._status = "idle"
-            self._model_name = ""
-            self._adapter_path = ""
-            self._loaded_at = 0.0
-            return {"unloaded": True}
+            with self._state_lock:
+                self._status = "idle"
+                self._model_name = ""
+                self._adapter_path = ""
+                self._loaded_at = 0.0
+                self._error = ""

-    def chat(self, messages, temperature=0.95, top_p=0.7, max_new_tokens=1024, do_sample=True, **kwargs):
-        with self._lock:
-            if self._status != "ready" or self._model is None:
-                return {"error": "model not loaded", "response": ""}
+    def unload(self) -> dict[str, Any]:
+        with self._state_lock:
+            if self._status == "loading":
+                # Ask the worker to tear down right after the load finishes.
+                self._cancel_requested = True
+                return {"unloaded": False, "status": "cancelling", "request_id": self._request_id}
+            was_ready = self._status == "ready"
+        if was_ready:
+            self._release_model()
+        else:
+            with self._state_lock:
+                self._model = None
+                self._tokenizer = None
+                self._status = "idle"
+                self._model_name = ""
+                self._adapter_path = ""
+                self._loaded_at = 0.0
+                self._error = ""
+        return {"unloaded": True, "status": "idle"}
+
+    def chat(self, messages, temperature=0.95, top_p=0.7, max_new_tokens=1024, do_sample=True, **kwargs) -> dict[str, Any]:
+        with self._chat_lock:
+            with self._state_lock:
+                if self._status == "loading":
+                    return {
+                        "error": f"model is still loading (request_id={self._request_id}); please retry",
+                        "response": "",
+                    }
+                if self._status == "error":
+                    return {"error": f"model load failed: {self._error}", "response": ""}
+                if self._status != "ready" or self._model is None:
+                    return {"error": "model not loaded", "response": ""}
            try:
-                generate_kwargs = {"temperature": temperature, "top_p": top_p, "max_new_tokens": max_new_tokens, "do_sample": do_sample}
+                generate_kwargs = {
+                    "temperature": temperature,
+                    "top_p": top_p,
+                    "max_new_tokens": max_new_tokens,
+                    "do_sample": do_sample,
+                }
                generate_kwargs.update(kwargs)
                system = next((m["content"] for m in messages if m["role"] == "system"), None)
                user_messages = [m for m in messages if m["role"] != "system"]
@@ -118,27 +212,35 @@ class InferenceSession:
                    responses.append(response)
                full_response = "".join(str(r) for r in responses)
                return {"response": full_response}
-            except Exception as exc:
+            except Exception as exc:  # noqa: BLE001 - return generation error to caller
                return {"error": str(exc), "response": ""}

-    def chat_stream(self, messages, **kwargs):
-        with self._lock:
-            if self._status != "ready" or self._model is None:
-                yield 'data: {"error": "model not loaded"}\n\n'
-                return
+    def chat_stream(self, messages, **kwargs) -> Iterator[str]:
+        with self._chat_lock:
+            with self._state_lock:
+                if self._status == "loading":
+                    yield 'data: {"error": "model is still loading; please retry"}\n\n'
+                    return
+                if self._status == "error":
+                    yield 'data: {"error": "model load failed: ' + str(self._error) + '"}\n\n'
+                    return
+                if self._status != "ready" or self._model is None:
+                    yield 'data: {"error": "model not loaded"}\n\n'
+                    return
            try:
                generate_kwargs = {**kwargs}
                system = next((m["content"] for m in messages if m["role"] == "system"), None)
                user_messages = [m for m in messages if m["role"] != "system"]
                for new_text in self._model.stream_chat(user_messages, system=system, **generate_kwargs):
                    yield new_text
-            except Exception as exc:
+            except Exception as exc:  # noqa: BLE001 - stream error as SSE event
                yield 'data: {"error": "' + str(exc) + '"}\n\n'


 _inference_session = None

-def get_inference_session():
+
+def get_inference_session() -> InferenceSession:
    global _inference_session
    if _inference_session is None:
        _inference_session = InferenceSession()