gpu-worker 1.0.0

package/engines/base.py

@@ -0,0 +1,57 @@
+"""推理引擎基类"""
+from abc import ABC, abstractmethod
+from typing import Any, Dict, Optional
+import torch
+import logging
+
+logger = logging.getLogger(__name__)
+
+
+class BaseEngine(ABC):
+    """推理引擎基类 - 所有引擎必须继承此类"""
+
+    def __init__(self, config: Dict[str, Any]):
+        self.config = config
+        self.model = None
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.loaded = False
+
+    @abstractmethod
+    def load_model(self) -> None:
+        """加载模型到内存"""
+        pass
+
+    @abstractmethod
+    def inference(self, params: Dict[str, Any]) -> Dict[str, Any]:
+        """执行推理"""
+        pass
+
+    @abstractmethod
+    def unload_model(self) -> None:
+        """卸载模型，释放内存"""
+        pass
+
+    def get_status(self) -> Dict[str, Any]:
+        """获取引擎状态"""
+        status = {
+            "loaded": self.loaded,
+            "device": self.device,
+        }
+
+        if torch.cuda.is_available():
+            status["gpu"] = {
+                "name": torch.cuda.get_device_name(0),
+                "memory_used_gb": torch.cuda.memory_allocated() / 1024**3,
+                "memory_total_gb": torch.cuda.get_device_properties(0).total_memory / 1024**3
+            }
+
+        return status
+
+    def _get_gpu_memory(self) -> Optional[Dict[str, float]]:
+        """获取GPU内存信息"""
+        if torch.cuda.is_available():
+            return {
+                "used_gb": torch.cuda.memory_allocated() / 1024**3,
+                "total_gb": torch.cuda.get_device_properties(0).total_memory / 1024**3
+            }
+        return None

package/engines/image_gen.py

@@ -0,0 +1,83 @@
+"""图像生成推理引擎"""
+from typing import Dict, Any
+import torch
+import io
+import base64
+import logging
+
+from .base import BaseEngine
+
+logger = logging.getLogger(__name__)
+
+
+class ImageGenEngine(BaseEngine):
+    """图像生成推理引擎"""
+
+    def load_model(self) -> None:
+        """加载图像生成模型"""
+        from diffusers import DiffusionPipeline
+
+        model_id = self.config.get("model_id", "Zhihu-ai/Z-Image-Turbo")
+        logger.info(f"Loading image generation model: {model_id}")
+
+        # 加载Pipeline
+        self.pipe = DiffusionPipeline.from_pretrained(
+            model_id,
+            torch_dtype=torch.bfloat16
+        )
+
+        # 内存优化
+        if self.config.get("enable_cpu_offload", True):
+            self.pipe.enable_sequential_cpu_offload()
+        else:
+            self.pipe = self.pipe.to(self.device)
+
+        self.loaded = True
+        logger.info("Image generation model loaded successfully")
+
+    def inference(self, params: Dict[str, Any]) -> Dict[str, Any]:
+        """执行图像生成"""
+        prompt = params.get("prompt", "")
+        negative_prompt = params.get("negative_prompt", "")
+        width = params.get("width", 1024)
+        height = params.get("height", 1024)
+        steps = params.get("steps", 4)
+        seed = params.get("seed", None)
+
+        # 设置随机种子
+        generator = None
+        if seed is not None:
+            generator = torch.Generator(device="cpu").manual_seed(seed)
+
+        # 生成图像
+        result = self.pipe(
+            prompt=prompt,
+            negative_prompt=negative_prompt if negative_prompt else None,
+            width=width,
+            height=height,
+            num_inference_steps=steps,
+            generator=generator
+        )
+
+        image = result.images[0]
+
+        # 转换为base64
+        buffered = io.BytesIO()
+        image.save(buffered, format="PNG")
+        image_base64 = base64.b64encode(buffered.getvalue()).decode()
+
+        return {
+            "image_base64": image_base64,
+            "width": width,
+            "height": height,
+            "seed": seed,
+            "format": "png"
+        }
+
+    def unload_model(self) -> None:
+        """卸载模型"""
+        if hasattr(self, "pipe"):
+            del self.pipe
+        torch.cuda.empty_cache()
+        self.loaded = False
+        logger.info("Image generation model unloaded")

package/engines/llm.py

@@ -0,0 +1,97 @@
+"""LLM推理引擎"""
+from typing import Dict, Any, List
+import torch
+import logging
+
+from .base import BaseEngine
+
+logger = logging.getLogger(__name__)
+
+
+class LLMEngine(BaseEngine):
+    """大语言模型推理引擎"""
+
+    def load_model(self) -> None:
+        """加载LLM模型"""
+        from transformers import AutoModelForCausalLM, AutoTokenizer
+
+        model_id = self.config.get("model_id", "Qwen/Qwen2.5-7B-Instruct")
+        logger.info(f"Loading LLM model: {model_id}")
+
+        # 加载tokenizer
+        self.tokenizer = AutoTokenizer.from_pretrained(
+            model_id,
+            trust_remote_code=True
+        )
+
+        # 加载模型
+        load_kwargs = {
+            "torch_dtype": torch.bfloat16,
+            "trust_remote_code": True,
+        }
+
+        if self.config.get("enable_cpu_offload", True):
+            load_kwargs["device_map"] = "auto"
+        else:
+            load_kwargs["device_map"] = {"": self.device}
+
+        self.model = AutoModelForCausalLM.from_pretrained(model_id, **load_kwargs)
+
+        self.loaded = True
+        logger.info("LLM model loaded successfully")
+
+    def inference(self, params: Dict[str, Any]) -> Dict[str, Any]:
+        """执行LLM推理"""
+        messages = params.get("messages", [])
+        max_tokens = params.get("max_tokens", 2048)
+        temperature = params.get("temperature", 0.7)
+        top_p = params.get("top_p", 0.9)
+
+        # 格式化输入
+        text = self.tokenizer.apply_chat_template(
+            messages,
+            tokenize=False,
+            add_generation_prompt=True
+        )
+
+        inputs = self.tokenizer(text, return_tensors="pt").to(self.model.device)
+        input_length = inputs.input_ids.shape[1]
+
+        # 生成
+        with torch.no_grad():
+            outputs = self.model.generate(
+                **inputs,
+                max_new_tokens=max_tokens,
+                temperature=temperature,
+                top_p=top_p,
+                do_sample=temperature > 0,
+                pad_token_id=self.tokenizer.eos_token_id
+            )
+
+        # 解码响应
+        response = self.tokenizer.decode(
+            outputs[0][input_length:],
+            skip_special_tokens=True
+        )
+
+        output_length = outputs.shape[1] - input_length
+
+        return {
+            "response": response,
+            "usage": {
+                "prompt_tokens": input_length,
+                "completion_tokens": output_length,
+                "total_tokens": input_length + output_length
+            }
+        }
+
+    def unload_model(self) -> None:
+        """卸载模型"""
+        if self.model:
+            del self.model
+            self.model = None
+        if hasattr(self, "tokenizer"):
+            del self.tokenizer
+        torch.cuda.empty_cache()
+        self.loaded = False
+        logger.info("LLM model unloaded")

package/engines/llm_base.py

@@ -0,0 +1,216 @@
+"""LLM引擎扩展基类 - 支持高性能推理后端"""
+from abc import abstractmethod
+from typing import Dict, Any, List, Optional, AsyncIterator
+from dataclasses import dataclass
+from enum import Enum
+import asyncio
+import threading
+from concurrent.futures import Future
+import logging
+
+from .base import BaseEngine
+
+logger = logging.getLogger(__name__)
+
+
+class LLMBackend(Enum):
+    """LLM 推理后端类型"""
+    NATIVE = "native"       # 原生 Transformers
+    SGLANG = "sglang"       # SGLang (推荐)
+    VLLM = "vllm"           # vLLM
+
+
+@dataclass
+class GenerationConfig:
+    """生成配置"""
+    max_tokens: int = 2048
+    temperature: float = 0.7
+    top_p: float = 0.9
+    top_k: int = 50
+    stop_sequences: Optional[List[str]] = None
+    stream: bool = False
+
+
+@dataclass
+class GenerationResult:
+    """生成结果"""
+    text: str
+    prompt_tokens: int
+    completion_tokens: int
+    total_tokens: int
+    finish_reason: str = "stop"
+    cached_tokens: int = 0  # 前缀缓存命中的 token 数
+
+
+class LLMBaseEngine(BaseEngine):
+    """
+    LLM 引擎扩展基类
+
+    提供高性能推理后端的统一接口，支持：
+    - 异步生成
+    - 批量处理
+    - 流式输出
+    - 前缀缓存
+    """
+
+    def __init__(self, config: Dict[str, Any]):
+        super().__init__(config)
+        self.backend_type: LLMBackend = LLMBackend.NATIVE
+        self.tokenizer = None
+        self._batch_processor = None
+
+    @abstractmethod
+    async def generate_async(
+        self,
+        messages: List[Dict[str, str]],
+        config: Optional[GenerationConfig] = None
+    ) -> GenerationResult:
+        """
+        异步生成接口
+
+        Args:
+            messages: 对话消息列表 [{"role": "user", "content": "..."}]
+            config: 生成配置
+
+        Returns:
+            生成结果
+        """
+        pass
+
+    @abstractmethod
+    async def batch_generate(
+        self,
+        batch_messages: List[List[Dict[str, str]]],
+        config: Optional[GenerationConfig] = None
+    ) -> List[GenerationResult]:
+        """
+        批量生成接口
+
+        Args:
+            batch_messages: 批量对话消息
+            config: 生成配置
+
+        Returns:
+            批量生成结果
+        """
+        pass
+
+    async def stream_generate(
+        self,
+        messages: List[Dict[str, str]],
+        config: Optional[GenerationConfig] = None
+    ) -> AsyncIterator[str]:
+        """
+        流式生成接口（默认实现：单次返回完整结果）
+
+        Args:
+            messages: 对话消息列表
+            config: 生成配置
+
+        Yields:
+            生成的文本片段
+        """
+        result = await self.generate_async(messages, config)
+        yield result.text
+
+    @staticmethod
+    def _run_coroutine_in_new_thread(coro):
+        future: Future = Future()
+
+        def runner() -> None:
+            try:
+                future.set_result(asyncio.run(coro))
+            except BaseException as exc:
+                future.set_exception(exc)
+
+        threading.Thread(target=runner, daemon=True).start()
+        return future.result()
+
+    def inference(self, params: Dict[str, Any]) -> Dict[str, Any]:
+        """
+        同步推理接口（兼容 BaseEngine）
+
+        包装异步接口，供同步调用
+        """
+        messages = params.get("messages", [])
+        config = GenerationConfig(
+            max_tokens=params.get("max_tokens", 2048),
+            temperature=params.get("temperature", 0.7),
+            top_p=params.get("top_p", 0.9),
+            top_k=params.get("top_k", 50),
+            stop_sequences=params.get("stop", None),
+            stream=params.get("stream", False)
+        )
+
+        try:
+            asyncio.get_running_loop()
+        except RuntimeError:
+            result = asyncio.run(self.generate_async(messages, config))
+        else:
+            result = self._run_coroutine_in_new_thread(self.generate_async(messages, config))
+
+        return {
+            "response": result.text,
+            "usage": {
+                "prompt_tokens": result.prompt_tokens,
+                "completion_tokens": result.completion_tokens,
+                "total_tokens": result.total_tokens,
+                "cached_tokens": result.cached_tokens
+            },
+            "finish_reason": result.finish_reason
+        }
+
+    def supports_streaming(self) -> bool:
+        """是否支持流式输出"""
+        return False
+
+    def supports_prefix_caching(self) -> bool:
+        """是否支持前缀缓存"""
+        return False
+
+    def supports_batch_inference(self) -> bool:
+        """是否支持批量推理"""
+        return False
+
+    def get_backend_info(self) -> Dict[str, Any]:
+        """获取后端信息"""
+        return {
+            "backend": self.backend_type.value,
+            "supports_streaming": self.supports_streaming(),
+            "supports_prefix_caching": self.supports_prefix_caching(),
+            "supports_batch_inference": self.supports_batch_inference()
+        }
+
+    def get_status(self) -> Dict[str, Any]:
+        """获取引擎状态（扩展父类）"""
+        status = super().get_status()
+        status["backend_info"] = self.get_backend_info()
+        return status
+
+
+def create_llm_engine(config: Dict[str, Any]) -> LLMBaseEngine:
+    """
+    LLM 引擎工厂函数
+
+    根据配置创建对应的 LLM 引擎实例
+
+    Args:
+        config: 引擎配置
+
+    Returns:
+        LLM 引擎实例
+    """
+    backend = config.get("backend", "native").lower()
+
+    if backend == "sglang":
+        from .llm_sglang import SGLangEngine
+        return SGLangEngine(config)
+
+    elif backend == "vllm":
+        from .llm_vllm import VLLMEngine
+        return VLLMEngine(config)
+
+    else:
+        # 默认使用原生 Transformers
+        from .llm import LLMEngine
+        return LLMEngine(config)