from __future__ import annotations import threading import time import uuid from typing import Any, Iterator class InferenceSession: """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._loaded_at: float = 0.0 @property def status(self) -> str: with self._state_lock: return self._status def info(self) -> dict[str, Any]: 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 wait_until_loaded(self, timeout: float | None = None) -> dict[str, Any]: """Wait for an in-flight async load to finish and return its outcome. 供同步消费方（如 eval_runner 子进程）使用：``load()`` 立即返回 loading 后， 调用本方法等待后台加载线程完成，拿到最终的 loaded/error 结果。 若在 timeout 秒内仍未加载完成，返回 ``status == "loading"`` 并附上超时提示。 """ with self._state_lock: thread = self._load_thread if thread is not None and thread.is_alive(): thread.join(timeout=timeout) with self._state_lock: loaded = self._status == "ready" status = self._status error = self._error if not loaded and status == "loading": error = error or f"model load timed out after {timeout or 'N/A'}s" return { "loaded": loaded, "status": status, "model_name": self._model_name, "adapter_path": self._adapter_path, "error": error, } 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": # 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 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 _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: # noqa: BLE001 - teardown must not raise pass with self._state_lock: self._status = "idle" self._model_name = "" self._adapter_path = "" self._loaded_at = 0.0 self._error = "" 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.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"] responses = [] for response in self._model.stream_chat(user_messages, system=system, **generate_kwargs): responses.append(response) full_response = "".join(str(r) for r in responses) return {"response": full_response} except Exception as exc: # noqa: BLE001 - return generation error to caller return {"error": str(exc), "response": ""} 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: # noqa: BLE001 - stream error as SSE event yield 'data: {"error": "' + str(exc) + '"}\n\n' _inference_session = None def get_inference_session() -> InferenceSession: global _inference_session if _inference_session is None: _inference_session = InferenceSession() return _inference_session