"""
大模型工程化監(jiān)控：日志分析工具選型與使用 - 基礎(chǔ)實現(xiàn)示例
大模型工程化核心代碼
"""

import torch
import torch.nn as nn
from transformers import AutoModelForCausalLM, AutoTokenizer
from typing import Optional, List, Dict, Any
import time
import logging

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


class LLMEngine:
    """
    大模型推理引擎基礎(chǔ)類
    
    提供模型加載、推理、優(yōu)化等核心功能
    """
    
    def __init__(self, 
                 model_name: str,
                 device: str = "cuda",
                 precision: str = "fp16"):
        """
        初始化推理引擎
        
        Args:
            model_name: 模型名稱或路徑
            device: 運行設(shè)備
            precision: 精度類型 (fp32/fp16/bf16)
        """
        self.model_name = model_name
        self.device = device
        self.precision = precision
        self.model = None
        self.tokenizer = None
        
        self._load_model()
        
    def _load_model(self):
        """加載模型和分詞器"""
        logger.info(f"正在加載模型: {self.model_name}")
        
        # 加載分詞器
        self.tokenizer = AutoTokenizer.from_pretrained(
            self.model_name,
            trust_remote_code=True
        )
        
        # 設(shè)置精度
        torch_dtype = {
            "fp32": torch.float32,
            "fp16": torch.float16,
            "bf16": torch.bfloat16
        }.get(self.precision, torch.float16)
        
        # 加載模型
        self.model = AutoModelForCausalLM.from_pretrained(
            self.model_name,
            torch_dtype=torch_dtype,
            device_map="auto",
            trust_remote_code=True
        )
        
        self.model.eval()
        logger.info("模型加載完成")
        
    def generate(self,
                 prompt: str,
                 max_new_tokens: int = 512,
                 temperature: float = 0.7,
                 top_p: float = 0.9,
                 **kwargs) -> str:
        """
        生成文本
        
        Args:
            prompt: 輸入提示詞
            max_new_tokens: 最大生成token數(shù)
            temperature: 溫度參數(shù)
            top_p: nucleus采樣參數(shù)
            
        Returns:
            生成的文本
        """
        # 編碼輸入
        inputs = self.tokenizer(prompt, return_tensors="pt")
        inputs = {k: v.to(self.device) for k, v in inputs.items()}
        
        # 生成
        with torch.no_grad():
            outputs = self.model.generate(
                **inputs,
                max_new_tokens=max_new_tokens,
                temperature=temperature,
                top_p=top_p,
                do_sample=True,
                pad_token_id=self.tokenizer.eos_token_id,
                **kwargs
            )
        
        # 解碼輸出
        generated_text = self.tokenizer.decode(
            outputs[0],
            skip_special_tokens=True
        )
        
        return generated_text
    
    def benchmark(self, 
                  prompt: str = "你好",
                  num_runs: int = 10) -> Dict[str, float]:
        """
        性能基準測試
        
        Args:
            prompt: 測試提示詞
            num_runs: 運行次數(shù)
            
        Returns:
            性能指標字典
        """
        latencies = []
        
        for i in range(num_runs):
            start_time = time.time()
            _ = self.generate(prompt, max_new_tokens=50)
            end_time = time.time()
            latencies.append(end_time - start_time)
            
        return {
            "avg_latency": sum(latencies) / len(latencies),
            "min_latency": min(latencies),
            "max_latency": max(latencies),
            "p50": sorted(latencies)[len(latencies) // 2],
            "p99": sorted(latencies)[int(len(latencies) * 0.99)]
        }
    
    def get_memory_usage(self) -> Dict[str, float]:
        """獲取顯存使用情況"""
        if torch.cuda.is_available():
            allocated = torch.cuda.memory_allocated() / 1024**3
            reserved = torch.cuda.memory_reserved() / 1024**3
            return {
                "allocated_gb": round(allocated, 2),
                "reserved_gb": round(reserved, 2)
            }
        return {}


class OptimizedLLMEngine(LLMEngine):
    """
    優(yōu)化后的大模型推理引擎
    
    包含量化、緩存等優(yōu)化技術(shù)
    """
    
    def __init__(self, 
                 model_name: str,
                 device: str = "cuda",
                 precision: str = "fp16",
                 enable_kv_cache: bool = True):
        """
        初始化優(yōu)化引擎
        
        Args:
            model_name: 模型名稱
            device: 運行設(shè)備
            precision: 精度類型
            enable_kv_cache: 是否啟用KV緩存
        """
        self.enable_kv_cache = enable_kv_cache
        self.kv_cache = {}
        super().__init__(model_name, device, precision)
        
    def generate_with_cache(self,
                            prompt: str,
                            session_id: Optional[str] = None,
                            **kwargs) -> str:
        """
        帶緩存的生成
        
        Args:
            prompt: 輸入提示詞
            session_id: 會話ID，用于緩存管理
            
        Returns:
            生成的文本
        """
        # 檢查緩存
        if session_id and session_id in self.kv_cache:
            # 使用緩存的KV
            past_key_values = self.kv_cache[session_id]
        else:
            past_key_values = None
            
        # 編碼輸入
        inputs = self.tokenizer(prompt, return_tensors="pt")
        inputs = {k: v.to(self.device) for k, v in inputs.items()}
        
        # 生成
        with torch.no_grad():
            outputs = self.model.generate(
                **inputs,
                past_key_values=past_key_values,
                use_cache=self.enable_kv_cache,
                **kwargs
            )
            
        # 更新緩存
        if session_id and self.enable_kv_cache:
            self.kv_cache[session_id] = outputs.past_key_values
            
        return self.tokenizer.decode(outputs[0], skip_special_tokens=True)
    
    def clear_cache(self, session_id: Optional[str] = None):
        """清除緩存"""
        if session_id:
            self.kv_cache.pop(session_id, None)
        else:
            self.kv_cache.clear()


# 使用示例
if __name__ == "__main__":
    # 創(chuàng)建引擎
    engine = LLMEngine(
        model_name="Qwen/Qwen2-1.5B",
        device="cuda",
        precision="fp16"
    )
    
    # 生成文本
    response = engine.generate("請介紹一下大模型工程化")
    print(f"生成結(jié)果: {response}")
    
    # 性能測試
    metrics = engine.benchmark()
    print(f"性能指標: {metrics}")
    
    # 顯存使用
    memory = engine.get_memory_usage()
    print(f"顯存使用: {memory}")

"""
大模型量化優(yōu)化實現(xiàn)
支持INT8/INT4量化
"""

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from typing import Optional
import bitsandbytes as bnb


class QuantizedLLMEngine:
    """
    量化大模型引擎
    
    支持INT8和INT4量化，大幅降低顯存占用
    """
    
    def __init__(self,
                 model_name: str,
                 quantization: str = "int8",
                 device_map: str = "auto"):
        """
        初始化量化引擎
        
        Args:
            model_name: 模型名稱
            quantization: 量化類型 (int8/int4/fp4/nf4)
            device_map: 設(shè)備映射策略
        """
        self.model_name = model_name
        self.quantization = quantization
        
        # 配置量化參數(shù)
        quantization_config = self._get_quantization_config()
        
        # 加載模型
        self.model = AutoModelForCausalLM.from_pretrained(
            model_name,
            quantization_config=quantization_config,
            device_map=device_map,
            trust_remote_code=True
        )
        
        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
        
    def _get_quantization_config(self):
        """獲取量化配置"""
        from transformers import BitsAndBytesConfig
        
        if self.quantization == "int8":
            return BitsAndBytesConfig(
                load_in_8bit=True,
                llm_int8_threshold=6.0
            )
        elif self.quantization == "int4":
            return BitsAndBytesConfig(
                load_in_4bit=True,
                bnb_4bit_compute_dtype=torch.float16,
                bnb_4bit_use_double_quant=True,
                bnb_4bit_quant_type="nf4"
            )
        else:
            return None
            
    def generate(self, prompt: str, **kwargs) -> str:
        """生成文本"""
        inputs = self.tokenizer(prompt, return_tensors="pt")
        inputs = {k: v.cuda() for k, v in inputs.items()}
        
        with torch.no_grad():
            outputs = self.model.generate(**inputs, **kwargs)
            
        return self.tokenizer.decode(outputs[0], skip_special_tokens=True)
    
    def get_model_size(self) -> Dict[str, float]:
        """獲取模型大小信息"""
        param_count = sum(p.numel() for p in self.model.parameters())
        
        # 估算顯存占用
        if self.quantization == "int4":
            size_gb = param_count * 0.5 / 1024**3  # INT4約0.5字節(jié)/參數(shù)
        elif self.quantization == "int8":
            size_gb = param_count * 1.0 / 1024**3  # INT8約1字節(jié)/參數(shù)
        else:
            size_gb = param_count * 2.0 / 1024**3  # FP16約2字節(jié)/參數(shù)
            
        return {
            "param_count_billion": round(param_count / 1e9, 2),
            "estimated_size_gb": round(size_gb, 2)
        }


# 使用示例
if __name__ == "__main__":
    # INT8量化
    engine_int8 = QuantizedLLMEngine(
        model_name="Qwen/Qwen2-7B",
        quantization="int8"
    )
    print(f"INT8模型大小: {engine_int8.get_model_size()}")
    
    # INT4量化
    engine_int4 = QuantizedLLMEngine(
        model_name="Qwen/Qwen2-7B",
        quantization="int4"
    )
    print(f"INT4模型大小: {engine_int4.get_model_size()}")

"""
大模型推理優(yōu)化技術(shù)
包含批處理、并行等優(yōu)化
"""

import torch
from typing import List, Optional
from dataclasses import dataclass
from queue import Queue
import threading
import time


@dataclass
class Request:
    """推理請求"""
    request_id: str
    prompt: str
    max_tokens: int = 100
    timestamp: float = time.time()


class DynamicBatcher:
    """
    動態(tài)批處理器
    
    自動將多個請求合并處理，提升吞吐量
    """
    
    def __init__(self,
                 model,
                 tokenizer,
                 max_batch_size: int = 32,
                 max_wait_time: float = 0.1):
        """
        初始化批處理器
        
        Args:
            model: 模型實例
            tokenizer: 分詞器
            max_batch_size: 最大批量大小
            max_wait_time: 最大等待時間
        """
        self.model = model
        self.tokenizer = tokenizer
        self.max_batch_size = max_batch_size
        self.max_wait_time = max_wait_time
        
        self.request_queue = Queue()
        self.results = {}
        self.running = True
        
        # 啟動處理線程
        self.process_thread = threading.Thread(target=self._process_loop)
        self.process_thread.start()
        
    def _process_loop(self):
        """批處理循環(huán)"""
        while self.running:
            batch = []
            start_time = time.time()
            
            # 收集請求
            while len(batch) < self.max_batch_size:
                if time.time() - start_time > self.max_wait_time:
                    break
                    
                try:
                    request = self.request_queue.get(timeout=0.01)
                    batch.append(request)
                except:
                    continue
                    
            if not batch:
                continue
                
            # 批量推理
            self._process_batch(batch)
            
    def _process_batch(self, batch: List[Request]):
        """處理批量請求"""
        prompts = [r.prompt for r in batch]
        
        # 批量編碼
        inputs = self.tokenizer(
            prompts,
            padding=True,
            return_tensors="pt"
        ).to(self.model.device)
        
        # 批量生成
        with torch.no_grad():
            outputs = self.model.generate(
                **inputs,
                max_new_tokens=max(r.max_tokens for r in batch)
            )
            
        # 解碼結(jié)果
        for i, request in enumerate(batch):
            result = self.tokenizer.decode(
                outputs[i],
                skip_special_tokens=True
            )
            self.results[request.request_id] = result
            
    def submit(self, request: Request):
        """提交請求"""
        self.request_queue.put(request)
        
    def get_result(self, request_id: str, timeout: float = 30) -> Optional[str]:
        """獲取結(jié)果"""
        start_time = time.time()
        while time.time() - start_time < timeout:
            if request_id in self.results:
                return self.results.pop(request_id)
            time.sleep(0.01)
        return None


class ModelParallel:
    """
    模型并行推理
    
    將大模型分割到多個GPU上運行
    """
    
    def __init__(self, model_name: str, num_gpus: int = 2):
        """
        初始化模型并行
        
        Args:
            model_name: 模型名稱
            num_gpus: GPU數(shù)量
        """
        self.num_gpus = num_gpus
        self.device_map = self._create_device_map()
        
        # 加載模型
        self.model = AutoModelForCausalLM.from_pretrained(
            model_name,
            device_map=self.device_map,
            trust_remote_code=True
        )
        
    def _create_device_map(self) -> Dict:
        """創(chuàng)建設(shè)備映射"""
        # 簡單的層分配策略
        return "auto"  # 使用自動分配
        
    def generate(self, prompt: str, **kwargs) -> str:
        """生成文本"""
        inputs = self.tokenizer(prompt, return_tensors="pt")
        # 自動路由到正確的設(shè)備
        inputs = {k: v.to("cuda:0") for k, v in inputs.items()}
        
        with torch.no_grad():
            outputs = self.model.generate(**inputs, **kwargs)
            
        return self.tokenizer.decode(outputs[0], skip_special_tokens=True)

"""
大模型性能監(jiān)控系統(tǒng)
"""

import time
import psutil
import torch
from dataclasses import dataclass, field
from typing import Dict, List
import json
from datetime import datetime


@dataclass
class PerformanceMetrics:
    """性能指標"""
    timestamp: str
    latency_ms: float
    throughput_qps: float
    gpu_memory_used_gb: float
    gpu_memory_total_gb: float
    gpu_utilization: float
    cpu_utilization: float
    request_count: int
    
    def to_dict(self) -> Dict:
        return {
            "timestamp": self.timestamp,
            "latency_ms": self.latency_ms,
            "throughput_qps": self.throughput_qps,
            "gpu_memory_used_gb": self.gpu_memory_used_gb,
            "gpu_memory_total_gb": self.gpu_memory_total_gb,
            "gpu_utilization": self.gpu_utilization,
            "cpu_utilization": self.cpu_utilization,
            "request_count": self.request_count
        }


class LLMPerformanceMonitor:
    """
    大模型性能監(jiān)控器
    
    實時監(jiān)控推理性能和資源使用
    """
    
    def __init__(self, collection_interval: float = 1.0):
        """
        初始化監(jiān)控器
        
        Args:
            collection_interval: 采集間隔（秒）
        """
        self.collection_interval = collection_interval
        self.metrics_history: List[PerformanceMetrics] = []
        self.request_times: List[float] = []
        self.request_count = 0
        self.running = False
        
    def start(self):
        """啟動監(jiān)控"""
        self.running = True
        
    def stop(self):
        """停止監(jiān)控"""
        self.running = False
        
    def record_request(self, latency: float):
        """記錄請求"""
        self.request_times.append(latency)
        self.request_count += 1
        
    def collect_metrics(self) -> PerformanceMetrics:
        """采集性能指標"""
        # GPU指標
        if torch.cuda.is_available():
            gpu_memory_used = torch.cuda.memory_allocated() / 1024**3
            gpu_memory_total = torch.cuda.get_device_properties(0).total_memory / 1024**3
            # GPU利用率需要nvidia-smi或其他工具
            gpu_utilization = 0.0  # 簡化處理
        else:
            gpu_memory_used = 0
            gpu_memory_total = 0
            gpu_utilization = 0
            
        # CPU指標
        cpu_utilization = psutil.cpu_percent()
        
        # 計算吞吐量
        if len(self.request_times) > 0:
            recent_requests = [t for t in self.request_times if time.time() - t < 60]
            throughput = len(recent_requests) / 60.0
        else:
            throughput = 0
            
        # 計算延遲
        if self.request_times:
            avg_latency = sum(self.request_times[-100:]) / len(self.request_times[-100:]) * 1000
        else:
            avg_latency = 0
            
        metrics = PerformanceMetrics(
            timestamp=datetime.now().isoformat(),
            latency_ms=avg_latency,
            throughput_qps=throughput,
            gpu_memory_used_gb=gpu_memory_used,
            gpu_memory_total_gb=gpu_memory_total,
            gpu_utilization=gpu_utilization,
            cpu_utilization=cpu_utilization,
            request_count=self.request_count
        )
        
        self.metrics_history.append(metrics)
        return metrics
    
    def get_summary(self) -> Dict:
        """獲取性能摘要"""
        if not self.metrics_history:
            return {}
            
        recent = self.metrics_history[-100:]
        
        return {
            "avg_latency_ms": sum(m.latency_ms for m in recent) / len(recent),
            "max_latency_ms": max(m.latency_ms for m in recent),
            "min_latency_ms": min(m.latency_ms for m in recent),
            "avg_throughput_qps": sum(m.throughput_qps for m in recent) / len(recent),
            "avg_gpu_memory_gb": sum(m.gpu_memory_used_gb for m in recent) / len(recent),
            "total_requests": self.request_count
        }
    
    def export_metrics(self, filepath: str):
        """導(dǎo)出指標"""
        data = [m.to_dict() for m in self.metrics_history]
        with open(filepath, 'w') as f:
            json.dump(data, f, indent=2)


# 使用示例
if __name__ == "__main__":
    monitor = LLMPerformanceMonitor()
    monitor.start()
    
    # 模擬請求
    for i in range(100):
        monitor.record_request(time.time())
        metrics = monitor.collect_metrics()
        print(f"指標: {metrics.to_dict()}")
        time.sleep(0.1)
        
    print(f"性能摘要: {monitor.get_summary()}")

"""
在線推理服務(wù)實現(xiàn)
FastAPI + 大模型
"""

from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from typing import Optional
import uvicorn
import asyncio
from concurrent.futures import ThreadPoolExecutor

app = FastAPI(title="LLM Inference API")

# 請求模型
class GenerateRequest(BaseModel):
    prompt: str
    max_tokens: int = 512
    temperature: float = 0.7
    top_p: float = 0.9
    
class GenerateResponse(BaseModel):
    text: str
    latency_ms: float
    tokens_generated: int

# 全局引擎
engine = None
executor = ThreadPoolExecutor(max_workers=4)

@app.on_event("startup")
async def startup():
    """啟動時加載模型"""
    global engine
    engine = LLMEngine(
        model_name="Qwen/Qwen2-1.5B",
        precision="fp16"
    )

@app.post("/generate", response_model=GenerateResponse)
async def generate(request: GenerateRequest):
    """生成接口"""
    import time
    start = time.time()
    
    # 異步執(zhí)行推理
    loop = asyncio.get_event_loop()
    result = await loop.run_in_executor(
        executor,
        engine.generate,
        request.prompt,
        request.max_tokens,
        request.temperature,
        request.top_p
    )
    
    latency = (time.time() - start) * 1000
    
    return GenerateResponse(
        text=result,
        latency_ms=latency,
        tokens_generated=len(result.split())
    )

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

@app.get("/metrics")
async def metrics():
    """性能指標"""
    return engine.get_memory_usage()

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8000)

# 安裝依賴
pip install torch transformers accelerate
pip install bitsandbytes  # 量化支持
pip install tensorrt      # TensorRT加速
pip install onnx onnxruntime  # ONNX支持

# 驗證GPU
python -c "import torch; print(torch.cuda.is_available())"

# 優(yōu)化方案實施
class OptimizedService:
    """優(yōu)化后的服務(wù)"""
    
    def __init__(self):
        # 1. 使用INT4量化
        self.model = QuantizedLLMEngine(
            model_name="model",
            quantization="int4"
        )
        
        # 2. 啟用動態(tài)批處理
        self.batcher = DynamicBatcher(
            self.model,
            max_batch_size=16
        )
        
        # 3. 配置監(jiān)控
        self.monitor = LLMPerformanceMonitor()
        
    def serve(self, request):
        """服務(wù)請求"""
        self.monitor.record_request(time.time())
        return self.batcher.submit(request)

# 顯存優(yōu)化策略
def optimize_memory():
    # 1. 清理緩存
    torch.cuda.empty_cache()
    
    # 2. 使用梯度檢查點
    model.gradient_checkpointing_enable()
    
    # 3. 降低精度
    model = model.half()
    
    # 4. 模型分片
    model = AutoModelForCausalLM.from_pretrained(
        model_name,
        device_map="auto"
    )