模型服务框架

# vllm_server.py
from vllm import LLM, SamplingParams
from vllm.engine.arg_utils import AsyncEngineArgs
from vllm.engine.async_llm_engine import AsyncLLMEngine
from vllm.entrypoints.openai.api_server import (
    create_chat_completion,
    create_completion,
)
import asyncio
from fastapi import FastAPI, Request
from typing import Optional, List, Dict, Any
import uvicorn

# vLLM 配置类
class VLLMConfig:
    def __init__(
        self,
        model: str,
        # 硬件配置
        tensor_parallel_size: int = 1,
        pipeline_parallel_size: int = 1,
        gpu_memory_utilization: float = 0.9,
        # 批处理配置
        max_num_batched_tokens: int = 8192,
        max_num_seqs: int = 256,
        max_model_len: Optional[int] = None,
        # KV Cache 配置
        block_size: int = 16,
        swap_space: int = 4,  # GB
        # 量化配置
        quantization: Optional[str] = None,  # awq, gptq, squeezellm
        # 调度配置
        scheduler_delay_factor: float = 0.0,
        enable_chunked_prefill: bool = False,
        # 推测解码
        speculative_model: Optional[str] = None,
        num_speculative_tokens: int = 5,
    ):
        self.model = model
        self.tensor_parallel_size = tensor_parallel_size
        self.pipeline_parallel_size = pipeline_parallel_size
        self.gpu_memory_utilization = gpu_memory_utilization
        self.max_num_batched_tokens = max_num_batched_tokens
        self.max_num_seqs = max_num_seqs
        self.max_model_len = max_model_len
        self.block_size = block_size
        self.swap_space = swap_space
        self.quantization = quantization
        self.scheduler_delay_factor = scheduler_delay_factor
        self.enable_chunked_prefill = enable_chunked_prefill
        self.speculative_model = speculative_model
        self.num_speculative_tokens = num_speculative_tokens

    def to_engine_args(self) -> AsyncEngineArgs:
        return AsyncEngineArgs(
            model=self.model,
            tensor_parallel_size=self.tensor_parallel_size,
            pipeline_parallel_size=self.pipeline_parallel_size,
            gpu_memory_utilization=self.gpu_memory_utilization,
            max_num_batched_tokens=self.max_num_batched_tokens,
            max_num_seqs=self.max_num_seqs,
            max_model_len=self.max_model_len,
            block_size=self.block_size,
            swap_space=self.swap_space,
            quantization=self.quantization,
            scheduler_delay_factor=self.scheduler_delay_factor,
            enable_chunked_prefill=self.enable_chunked_prefill,
            speculative_model=self.speculative_model,
            num_speculative_tokens=self.num_speculative_tokens,
            trust_remote_code=True,
            dtype="auto",
        )


# 自定义 vLLM 服务
class VLLMServer:
    def __init__(self, config: VLLMConfig):
        self.config = config
        self.engine: Optional[AsyncLLMEngine] = None
        self.app = FastAPI(title="vLLM API Server")
        self._setup_routes()

    async def initialize(self):
        """初始化引擎"""
        engine_args = self.config.to_engine_args()
        self.engine = AsyncLLMEngine.from_engine_args(engine_args)

    def _setup_routes(self):
        """设置 API 路由"""

        @self.app.post("/v1/completions")
        async def completions(request: Request):
            """OpenAI Completions API"""
            body = await request.json()
            return await self._handle_completion(body)

        @self.app.post("/v1/chat/completions")
        async def chat_completions(request: Request):
            """OpenAI Chat Completions API"""
            body = await request.json()
            return await self._handle_chat_completion(body)

        @self.app.get("/health")
        async def health():
            """健康检查"""
            return {"status": "healthy"}

        @self.app.get("/v1/models")
        async def list_models():
            """列出模型"""
            return {
                "data": [
                    {
                        "id": self.config.model,
                        "object": "model",
                        "owned_by": "vllm",
                    }
                ]
            }

    async def _handle_completion(self, body: Dict[str, Any]):
        """处理 Completion 请求"""
        prompt = body.get("prompt", "")
        sampling_params = SamplingParams(
            temperature=body.get("temperature", 1.0),
            top_p=body.get("top_p", 1.0),
            top_k=body.get("top_k", -1),
            max_tokens=body.get("max_tokens", 16),
            stop=body.get("stop"),
            presence_penalty=body.get("presence_penalty", 0.0),
            frequency_penalty=body.get("frequency_penalty", 0.0),
        )

        request_id = f"cmpl-{self._generate_request_id()}"

        # 生成
        results = []
        async for output in self.engine.generate(prompt, sampling_params, request_id):
            results.append(output)

        final_output = results[-1]

        return {
            "id": request_id,
            "object": "text_completion",
            "model": self.config.model,
            "choices": [
                {
                    "text": output.text,
                    "index": i,
                    "finish_reason": output.finish_reason,
                }
                for i, output in enumerate(final_output.outputs)
            ],
            "usage": {
                "prompt_tokens": len(final_output.prompt_token_ids),
                "completion_tokens": sum(
                    len(o.token_ids) for o in final_output.outputs
                ),
                "total_tokens": len(final_output.prompt_token_ids)
                + sum(len(o.token_ids) for o in final_output.outputs),
            },
        }

    async def _handle_chat_completion(self, body: Dict[str, Any]):
        """处理 Chat Completion 请求"""
        messages = body.get("messages", [])

        # 转换消息为 prompt
        prompt = self._messages_to_prompt(messages)

        # 复用 completion 逻辑
        body["prompt"] = prompt
        return await self._handle_completion(body)

    def _messages_to_prompt(self, messages: List[Dict]) -> str:
        """将消息列表转换为 prompt"""
        # 根据模型类型使用不同的模板
        # 这里使用 ChatML 格式作为示例
        prompt_parts = []
        for msg in messages:
            role = msg["role"]
            content = msg["content"]
            if role == "system":
                prompt_parts.append(f"<|im_start|>system\n{content}<|im_end|>")
            elif role == "user":
                prompt_parts.append(f"<|im_start|>user\n{content}<|im_end|>")
            elif role == "assistant":
                prompt_parts.append(f"<|im_start|>assistant\n{content}<|im_end|>")

        prompt_parts.append("<|im_start|>assistant\n")
        return "\n".join(prompt_parts)

    def _generate_request_id(self) -> str:
        import uuid
        return str(uuid.uuid4())[:8]

    def run(self, host: str = "0.0.0.0", port: int = 8000):
        """运行服务"""
        asyncio.run(self.initialize())
        uvicorn.run(self.app, host=host, port=port)


# 启动脚本
if __name__ == "__main__":
    config = VLLMConfig(
        model="meta-llama/Llama-2-7b-chat-hf",
        tensor_parallel_size=1,
        gpu_memory_utilization=0.9,
        max_num_seqs=256,
        max_model_len=4096,
    )

    server = VLLMServer(config)
    server.run()

# tgi_client.py
import asyncio
from typing import AsyncIterator, Optional, List, Dict, Any
import aiohttp
import json

class TGIClient:
    """TGI 异步客户端"""

    def __init__(self, base_url: str, timeout: float = 60.0):
        self.base_url = base_url.rstrip("/")
        self.timeout = aiohttp.ClientTimeout(total=timeout)

    async def generate(
        self,
        prompt: str,
        max_new_tokens: int = 256,
        temperature: float = 1.0,
        top_p: float = 1.0,
        top_k: Optional[int] = None,
        repetition_penalty: float = 1.0,
        stop_sequences: Optional[List[str]] = None,
        stream: bool = False,
    ) -> Dict[str, Any]:
        """生成文本"""
        payload = {
            "inputs": prompt,
            "parameters": {
                "max_new_tokens": max_new_tokens,
                "temperature": temperature,
                "top_p": top_p,
                "repetition_penalty": repetition_penalty,
                "do_sample": temperature > 0,
            },
        }

        if top_k is not None:
            payload["parameters"]["top_k"] = top_k
        if stop_sequences:
            payload["parameters"]["stop"] = stop_sequences

        async with aiohttp.ClientSession(timeout=self.timeout) as session:
            if stream:
                return await self._stream_generate(session, payload)
            else:
                return await self._generate(session, payload)

    async def _generate(
        self,
        session: aiohttp.ClientSession,
        payload: Dict[str, Any],
    ) -> Dict[str, Any]:
        """非流式生成"""
        async with session.post(
            f"{self.base_url}/generate",
            json=payload,
        ) as response:
            response.raise_for_status()
            return await response.json()

    async def _stream_generate(
        self,
        session: aiohttp.ClientSession,
        payload: Dict[str, Any],
    ) -> AsyncIterator[Dict[str, Any]]:
        """流式生成"""
        payload["stream"] = True

        async with session.post(
            f"{self.base_url}/generate_stream",
            json=payload,
        ) as response:
            response.raise_for_status()

            async for line in response.content:
                line = line.decode("utf-8").strip()
                if line.startswith("data:"):
                    data = json.loads(line[5:])
                    yield data

    async def generate_stream(
        self,
        prompt: str,
        **kwargs,
    ) -> AsyncIterator[str]:
        """流式生成，返回文本流"""
        async for chunk in await self.generate(prompt, stream=True, **kwargs):
            if "token" in chunk:
                yield chunk["token"]["text"]

    async def health(self) -> Dict[str, Any]:
        """健康检查"""
        async with aiohttp.ClientSession(timeout=self.timeout) as session:
            async with session.get(f"{self.base_url}/health") as response:
                return await response.json()

    async def info(self) -> Dict[str, Any]:
        """获取模型信息"""
        async with aiohttp.ClientSession(timeout=self.timeout) as session:
            async with session.get(f"{self.base_url}/info") as response:
                return await response.json()


# 使用示例
async def main():
    client = TGIClient("http://tgi-server:80")

    # 检查健康状态
    health = await client.health()
    print(f"Health: {health}")

    # 获取模型信息
    info = await client.info()
    print(f"Model: {info['model_id']}")

    # 非流式生成
    response = await client.generate(
        prompt="What is the capital of France?",
        max_new_tokens=50,
        temperature=0.7,
    )
    print(f"Response: {response['generated_text']}")

    # 流式生成
    print("Streaming response: ", end="")
    async for token in client.generate_stream(
        prompt="Tell me a short story about a robot.",
        max_new_tokens=200,
        temperature=0.8,
    ):
        print(token, end="", flush=True)
    print()


if __name__ == "__main__":
    asyncio.run(main())

# build_trt_llm.py
import tensorrt_llm
from tensorrt_llm import BuildConfig, Mapping
from tensorrt_llm.models import LLaMAForCausalLM
from tensorrt_llm.quantization import QuantMode
import torch

def build_trt_llm_engine(
    model_dir: str,
    output_dir: str,
    dtype: str = "float16",
    tp_size: int = 1,
    pp_size: int = 1,
    max_batch_size: int = 256,
    max_input_len: int = 2048,
    max_output_len: int = 2048,
    quantization: str = None,  # int8, int4_awq, fp8
):
    """构建 TensorRT-LLM 引擎"""

    # 设置量化模式
    quant_mode = QuantMode(0)
    if quantization == "int8":
        quant_mode = QuantMode.use_weight_only(use_int4_weights=False)
    elif quantization == "int4_awq":
        quant_mode = QuantMode.use_weight_only(use_int4_weights=True)
    elif quantization == "fp8":
        quant_mode = QuantMode.use_fp8_qdq()

    # 创建映射配置（张量并行/流水线并行）
    mapping = Mapping(
        world_size=tp_size * pp_size,
        tp_size=tp_size,
        pp_size=pp_size,
    )

    # 构建配置
    build_config = BuildConfig(
        max_batch_size=max_batch_size,
        max_input_len=max_input_len,
        max_output_len=max_output_len,
        max_beam_width=1,
        max_num_tokens=max_batch_size * max_input_len,
        strongly_typed=True,
        builder_opt=5,  # 优化级别
    )

    # 加载模型
    if dtype == "float16":
        torch_dtype = torch.float16
    elif dtype == "bfloat16":
        torch_dtype = torch.bfloat16
    else:
        torch_dtype = torch.float32

    # 从 HuggingFace 格式转换
    model = LLaMAForCausalLM.from_hugging_face(
        model_dir,
        dtype=dtype,
        mapping=mapping,
        quant_mode=quant_mode,
    )

    # 构建引擎
    engine = tensorrt_llm.build(
        model,
        build_config,
    )

    # 保存引擎
    engine.save(output_dir)

    print(f"Engine saved to {output_dir}")
    return output_dir


# 多 GPU 并行构建脚本
def build_multi_gpu_engine():
    """构建多 GPU 引擎"""
    import subprocess
    import os

    # 使用 mpirun 启动多进程构建
    cmd = [
        "mpirun",
        "-n", "4",  # 4 GPU
        "--allow-run-as-root",
        "python", "build_trt_llm.py",
        "--model_dir", "/models/llama-70b",
        "--output_dir", "/engines/llama-70b-tp4",
        "--tp_size", "4",
        "--dtype", "float16",
        "--max_batch_size", "64",
    ]

    subprocess.run(cmd, check=True)


if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser()
    parser.add_argument("--model_dir", required=True)
    parser.add_argument("--output_dir", required=True)
    parser.add_argument("--tp_size", type=int, default=1)
    parser.add_argument("--dtype", default="float16")
    parser.add_argument("--max_batch_size", type=int, default=256)
    parser.add_argument("--quantization", default=None)

    args = parser.parse_args()

    build_trt_llm_engine(
        model_dir=args.model_dir,
        output_dir=args.output_dir,
        dtype=args.dtype,
        tp_size=args.tp_size,
        max_batch_size=args.max_batch_size,
        quantization=args.quantization,
    )

# triton_trt_llm_config.py
# Triton 模型配置生成

import json
import os

def generate_triton_config(
    model_name: str,
    engine_dir: str,
    output_dir: str,
    max_batch_size: int = 256,
    tp_size: int = 1,
    pp_size: int = 1,
):
    """生成 Triton 配置文件"""

    config = {
        "name": model_name,
        "backend": "tensorrtllm",
        "max_batch_size": max_batch_size,

        "model_transaction_policy": {
            "decoupled": True  # 流式输出
        },

        "input": [
            {
                "name": "input_ids",
                "data_type": "TYPE_INT32",
                "dims": [-1]
            },
            {
                "name": "input_lengths",
                "data_type": "TYPE_INT32",
                "dims": [1]
            },
            {
                "name": "request_output_len",
                "data_type": "TYPE_INT32",
                "dims": [1]
            },
            {
                "name": "end_id",
                "data_type": "TYPE_INT32",
                "dims": [1],
                "optional": True
            },
            {
                "name": "pad_id",
                "data_type": "TYPE_INT32",
                "dims": [1],
                "optional": True
            },
            {
                "name": "beam_width",
                "data_type": "TYPE_INT32",
                "dims": [1],
                "optional": True
            },
            {
                "name": "temperature",
                "data_type": "TYPE_FP32",
                "dims": [1],
                "optional": True
            },
            {
                "name": "top_k",
                "data_type": "TYPE_INT32",
                "dims": [1],
                "optional": True
            },
            {
                "name": "top_p",
                "data_type": "TYPE_FP32",
                "dims": [1],
                "optional": True
            },
            {
                "name": "streaming",
                "data_type": "TYPE_BOOL",
                "dims": [1],
                "optional": True
            }
        ],

        "output": [
            {
                "name": "output_ids",
                "data_type": "TYPE_INT32",
                "dims": [-1, -1]
            },
            {
                "name": "sequence_length",
                "data_type": "TYPE_INT32",
                "dims": [-1]
            }
        ],

        "instance_group": [
            {
                "count": 1,
                "kind": "KIND_GPU",
                "gpus": list(range(tp_size * pp_size))
            }
        ],

        "parameters": {
            "gpt_model_type": {"string_value": "llama"},
            "gpt_model_path": {"string_value": engine_dir},
            "batch_scheduler_policy": {"string_value": "inflight_fused_batching"},
            "decoupled_mode": {"string_value": "true"},
            "max_tokens_in_paged_kv_cache": {"string_value": ""},
            "kv_cache_free_gpu_mem_fraction": {"string_value": "0.9"},
            "enable_chunked_context": {"string_value": "true"},
        }
    }

    # 创建目录
    model_dir = os.path.join(output_dir, model_name, "1")
    os.makedirs(model_dir, exist_ok=True)

    # 写入配置
    config_path = os.path.join(output_dir, model_name, "config.pbtxt")
    with open(config_path, "w") as f:
        f.write(dict_to_pbtxt(config))

    print(f"Triton config saved to {config_path}")
    return config_path


def dict_to_pbtxt(d: dict, indent: int = 0) -> str:
    """将字典转换为 pbtxt 格式"""
    lines = []
    prefix = "  " * indent

    for key, value in d.items():
        if isinstance(value, dict):
            lines.append(f"{prefix}{key} {{")
            lines.append(dict_to_pbtxt(value, indent + 1))
            lines.append(f"{prefix}}}")
        elif isinstance(value, list):
            for item in value:
                if isinstance(item, dict):
                    lines.append(f"{prefix}{key} {{")
                    lines.append(dict_to_pbtxt(item, indent + 1))
                    lines.append(f"{prefix}}}")
                else:
                    lines.append(f'{prefix}{key}: {json.dumps(item)}')
        elif isinstance(value, bool):
            lines.append(f'{prefix}{key}: {str(value).lower()}')
        elif isinstance(value, str):
            lines.append(f'{prefix}{key}: "{value}"')
        else:
            lines.append(f'{prefix}{key}: {value}')

    return "\n".join(lines)