inference_actor.py

"""
InferenceActor: 部署在 GPU 1，运行 vLLM + CLIP 双模块。
- vLLM: 高速多模态生成 (AWQ 量化)，支持 GRPO 采样 n=8。
- CLIP: 提取视觉特征 Tensor (B, 576, 1024)，通过 Ray Object Store 传给 GPU 0 做训练。
显存控制：vLLM gpu_memory_utilization=0.6，为 CLIP 预留空间。
"""

import os
from typing import List, Tuple, Optional, Dict, Any

import ray
import torch
from PIL import Image
from transformers import CLIPVisionModel, CLIPImageProcessor
from vllm import LLM, SamplingParams


def _get_device():
    return torch.device("cuda:0" if torch.cuda.is_available() else "cpu")


@ray.remote(num_gpus=1)
class InferenceActor:
    """
    GPU 1：vLLM 高速生成 + CLIP 视觉特征提取。
    - extract_features(): 返回 last_hidden_state (B, 576, 1024) 存入 Object Store。
    - generate_rollouts(): vLLM 多模态采样，返回文本结果。
    """

    def __init__(
        self,
        llm_model_name: str = "llava-hf/llava-1.5-7b-hf",
        clip_model_name: str = "openai/clip-vit-large-patch14-336",
        gpu_memory_utilization: float = 0.6,
        max_model_len: int = 4096,
    ):
        self._device = _get_device()

        # 1. vLLM：多模态生成（AWQ 量化，控制显存）
        self.llm = LLM(
            model=llm_model_name,
            trust_remote_code=True,
            gpu_memory_utilization=gpu_memory_utilization,
            max_model_len=max_model_len,
            limit_mm_per_prompt={"image": 1},
            # quantization="awq",  # 若使用 AWQ 版本模型则开启
        )
        self.tokenizer = self.llm.get_tokenizer()

        # 2. CLIP：单独加载用于提取视觉特征给 GPU 0
        # 选择 336 版本与 LLaVA 1.5 一致，patch=14 -> 336/14=24 -> 24*24=576 tokens
        self.clip_processor = CLIPImageProcessor.from_pretrained(clip_model_name)
        self.clip_vision = CLIPVisionModel.from_pretrained(
            clip_model_name, torch_dtype=torch.float16
        )
        self.clip_vision.to(self._device)
        self.clip_vision.eval()
        self.clip_hidden_size = self.clip_vision.config.hidden_size  # 1024

        # 采样参数
        self.default_sampling_params = SamplingParams(
            temperature=0.7,
            top_p=0.95,
            max_tokens=512,
        )

    def _load_image(self, image_path: str) -> Image.Image:
        """加载单张图片。"""
        return Image.open(image_path).convert("RGB")

    # -------------------------------------------------------------------------
    # 方法 1: extract_features (提取 CLIP 特征给 GPU 0)
    # -------------------------------------------------------------------------
    def extract_features(self, image_paths: List[str]) -> torch.Tensor:
        """
        提取 CLIP 视觉特征，返回 last_hidden_state (B, 576, 1024)。
        必须 .detach().cpu() 放入 Ray Object Store，供 GPU 0 拉取。

        Args:
            image_paths: 图片路径列表

        Returns:
            visual_features: (B, num_patches, hidden_size) = (B, 576, 1024)
        """
        if not image_paths:
            raise ValueError("image_paths 不能为空")

        all_features = []
        with torch.no_grad():
            for path in image_paths:
                if path and os.path.isfile(path):
                    image = self._load_image(path)
                    inputs = self.clip_processor(
                        images=image, return_tensors="pt"
                    )
                    pixel_values = inputs["pixel_values"].to(
                        self._device, dtype=torch.float16
                    )
                    outputs = self.clip_vision(pixel_values, output_hidden_states=True)
                    # last_hidden_state: (1, 577, 1024)，去掉 CLS token -> (1, 576, 1024)
                    last_hidden = outputs.last_hidden_state[:, 1:, :]
                    all_features.append(last_hidden)
                else:
                    # 联调用零向量占位：(1, 576, 1024)
                    all_features.append(
                        torch.zeros(
                            1, 576, self.clip_hidden_size,
                            device=self._device, dtype=torch.float16
                        )
                    )

        # (B, 576, 1024) -> .detach().cpu() 放入 Object Store
        features = torch.cat(all_features, dim=0)
        return features.detach().cpu()

    # -------------------------------------------------------------------------
    # 方法 2: generate_rollouts (vLLM 多模态生成)
    # -------------------------------------------------------------------------
    def generate_rollouts(
        self,
        prompts: List[str],
        image_paths: List[str],
        n: int = 8,
        max_tokens: int = 512,
        temperature: float = 0.7,
        top_p: float = 0.95,
    ) -> List[List[str]]:
        """
        使用 vLLM 进行多模态生成，每个 prompt 采样 n 条候选。

        Args:
            prompts: 文本 prompt 列表（应包含 <image> 占位符）
            image_paths: 对应的图片路径
            n: 每个 prompt 采样数量（GRPO Group Size）

        Returns:
            responses_per_prompt: List[List[str]]，每个 prompt 对应 n 条生成结果
        """
        if len(prompts) != len(image_paths):
            raise ValueError("prompts 和 image_paths 长度必须一致")

        sampling_params = SamplingParams(
            temperature=temperature,
            top_p=top_p,
            max_tokens=max_tokens,
            n=n,
        )

        # 构造 vLLM 多模态输入
        llm_inputs = []
        for prompt, img_path in zip(prompts, image_paths):
            if img_path and os.path.isfile(img_path):
                image = self._load_image(img_path)
                llm_inputs.append({
                    "prompt": prompt,
                    "multi_modal_data": {"image": image},
                })
            else:
                # 无图片时仅文本
                llm_inputs.append({"prompt": prompt})

        outputs = self.llm.generate(llm_inputs, sampling_params)

        # 整理为 List[List[str]]
        responses_per_prompt: List[List[str]] = []
        for output in outputs:
            responses_per_prompt.append([o.text for o in output.outputs])

        return responses_per_prompt

    # -------------------------------------------------------------------------
    # 工具方法
    # -------------------------------------------------------------------------
    def get_clip_hidden_size(self) -> int:
        return self.clip_hidden_size

    def get_num_patches(self) -> int:
        """返回 CLIP 输出的 patch 数量（不含 CLS）。"""
        return 576  # 336/14=24, 24*24=576

    def get_tokenizer(self):
        """返回 vLLM tokenizer，供 main 使用。"""
        return self.tokenizer