YG_FT/compute/engines/llama_factory/inference.py

from __future__ import annotations

import os
import threading
import time
import uuid
from typing import Any, Iterator


def _ensure_peft_transformers_compat() -> None:
    """Bridge a removed PEFT helper used by the bundled Transformers build.

    The offline Compute image currently contains Transformers 5.8.0 and PEFT
    0.18.1. Transformers imports this private helper when a model directory
    contains PEFT metadata, but PEFT 0.18.1 does not expose it. Evaluation
    runs in single-process HuggingFace mode, so tensor-parallel sharding is
    not applicable and a no-op compatibility hook is the correct behavior.
    """
    try:
        from peft.utils import save_and_load
    except Exception:
        return
    if hasattr(save_and_load, "_maybe_shard_state_dict_for_tp"):
        return

    def _maybe_shard_state_dict_for_tp(_model: Any, _state_dict: dict[str, Any], _adapter_name: str) -> None:
        return None

    save_and_load._maybe_shard_state_dict_for_tp = _maybe_shard_state_dict_for_tp


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._gpu_indices: list[int] = []
        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,
                "gpu_indices": list(self._gpu_indices),
            }

    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",
        gpu_indices=None,
        **kwargs,
    ) -> dict[str, Any]:
        requested_gpus = sorted({int(item) for item in (gpu_indices or [])})
        if any(item < 0 for item in requested_gpus):
            return {"loaded": False, "status": "error", "error": "GPU index must be non-negative"}
        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._gpu_indices = requested_gpus
            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."""
        with self._state_lock:
            requested_gpus = list(self._gpu_indices)
        model = None
        tokenizer = None
        generating_args: dict[str, Any] = {}
        error = ""
        previous_visible_devices = os.environ.get("CUDA_VISIBLE_DEVICES")
        try:
            # Set visibility before LLaMA-Factory/PyTorch initializes CUDA.
            if requested_gpus:
                os.environ["CUDA_VISIBLE_DEVICES"] = ",".join(str(item) for item in requested_gpus)
            if self._teardown_old:
                self._release_model()
                with self._state_lock:
                    self._gpu_indices = requested_gpus
            from llamafactory.chat import ChatModel
            from llamafactory.hparams import get_infer_args

            args = dict(self._load_args)
            infer_result = get_infer_args(args)
            _ensure_peft_transformers_compat()
            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)
        finally:
            if requested_gpus:
                if previous_visible_devices is None:
                    os.environ.pop("CUDA_VISIBLE_DEVICES", None)
                else:
                    os.environ["CUDA_VISIBLE_DEVICES"] = previous_visible_devices
        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"
                self._gpu_indices = []
                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 = ""
                self._gpu_indices = []

    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 = ""
                self._gpu_indices = []
        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