09 - AI网络专题¶

# 流式API客户端示例
import requests
import json

def stream_chat_completion():
    """OpenAI风格流式API调用"""
    url = "https://api.openai.com/v1/chat/completions"
    headers = {
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json"
    }
    data = {
        "model": "gpt-4",
        "messages": [{"role": "user", "content": "讲个故事"}],
        "stream": True  # 关键：启用流式传输
    }

    response = requests.post(url, headers=headers, json=data, stream=True)

    for line in response.iter_lines():
        if line:
            line = line.decode('utf-8')
            if line.startswith('data: '):
                json_str = line[6:]
                if json_str != '[DONE]':
                    chunk = json.loads(json_str)
                    content = chunk['choices'][0]['delta'].get('content', '')
                    print(content, end='', flush=True)

# 使用aiohttp的异步版本
import aiohttp
import asyncio

async def async_stream_completion():
    """异步流式调用，性能更好"""
    async with aiohttp.ClientSession() as session:
        async with session.post(url, headers=headers, json=data) as resp:
            async for line in resp.content:
                # 处理流式数据...
                pass

import requests
from requests.adapters import HTTPAdapter
from urllib3.util.retry import Retry

# 创建带连接池的session
session = requests.Session()

# 配置连接池
adapter = HTTPAdapter(
    pool_connections=10,    # 连接池大小
    pool_maxsize=20,        # 最大连接数
    max_retries=Retry(
        total=3,
        backoff_factor=1,
        status_forcelist=[429, 500, 502, 503, 504]
    )
)

session.mount("https://", adapter)

# 复用session进行多次调用
for prompt in prompts:
    response = session.post(api_url, json={"prompt": prompt})

# 使用WebSocket实现真正的双向实时通信
import websockets
import asyncio
import json

class LLMWebSocketClient:
    def __init__(self, uri):
        self.uri = uri
        self.websocket = None

    async def connect(self):
        self.websocket = await websockets.connect(self.uri)

    async def stream_generate(self, prompt):
        """通过WebSocket流式生成"""
        await self.websocket.send(json.dumps({
            "type": "generate",
            "prompt": prompt,
            "stream": True
        }))

        while True:
            message = await self.websocket.recv()
            data = json.loads(message)

            if data.get("type") == "token":
                yield data["content"]
            elif data.get("type") == "done":
                break

    async def close(self):
        await self.websocket.close()

# 使用示例
async def main():
    client = LLMWebSocketClient("wss://api.example.com/llm")
    await client.connect()

    async for token in client.stream_generate("你好"):
        print(token, end='')

    await client.close()

# 使用Redis实现分布式限流
import redis
import time
from functools import wraps

class APIRateLimiter:
    def __init__(self, redis_client):
        self.redis = redis_client

    def is_allowed(self, api_key, max_requests=60, window=60):
        """滑动窗口限流"""
        key = f"rate_limit:{api_key}"
        now = time.time()
        window_start = now - window

        # 移除窗口外的请求记录
        self.redis.zremrangebyscore(key, 0, window_start)

        # 获取当前窗口内的请求数
        current = self.redis.zcard(key)

        if current < max_requests:
            # 记录本次请求
            self.redis.zadd(key, {str(now): now})
            self.redis.expire(key, window)
            return True
        return False

# Flask限流装饰器
from flask import Flask, request, jsonify

app = Flask(__name__)
limiter = APIRateLimiter(redis.Redis())

@app.route('/v1/chat/completions', methods=['POST'])
def chat_completions():
    api_key = request.headers.get('Authorization', '').replace('Bearer ', '')

    if not limiter.is_allowed(api_key, max_requests=60):
        return jsonify({
            "error": "Rate limit exceeded",
            "retry_after": 60
        }), 429

    # 处理请求...

class ModelRouter:
    """根据请求特征智能选择模型"""

    def __init__(self):
        self.models = {
            'gpt-4': {'cost': 0.03, 'speed': 'slow', 'quality': 'high'},
            'gpt-3.5': {'cost': 0.002, 'speed': 'fast', 'quality': 'medium'},
        }

    def route(self, request_data):
        """根据请求特征路由到合适模型"""
        prompt_length = len(request_data['messages'][-1]['content'])
        complexity = self._estimate_complexity(request_data)

        # 简单请求使用轻量级模型
        if prompt_length < 100 and complexity < 0.3:
            return 'gpt-3.5'

        # 复杂任务使用强模型
        return 'gpt-4'

    def _estimate_complexity(self, request):
        """估算任务复杂度"""
        prompt = request['messages'][-1]['content'].lower()
        complex_keywords = ['分析', '比较', '推理', '证明', '代码']
        return sum(1 for kw in complex_keywords if kw in prompt) / len(complex_keywords)

import asyncio
from typing import AsyncGenerator, List
import time

class BatchStreamProcessor:
    """批量流式处理器 - 减少网络往返次数"""

    def __init__(self, batch_size: int = 5, max_latency_ms: float = 50):
        self.batch_size = batch_size
        self.max_latency_ms = max_latency_ms
        self.buffer = []
        self.last_flush_time = time.time()

    async def process_stream(
        self,
        token_generator: AsyncGenerator[str, None]
    ) -> AsyncGenerator[List[str], None]:
        """
        将单个token流转换为批量token流

        优势：
        - 减少网络包数量 5-10x
        - 降低客户端渲染开销
        - 更好的压缩率
        """
        async for token in token_generator:
            self.buffer.append(token)

            # 触发条件1：缓冲区满
            buffer_full = len(self.buffer) >= self.batch_size

            # 触发条件2：超时
            time_since_flush = (time.time() - self.last_flush_time) * 1000
            timeout_reached = time_since_flush >= self.max_latency_ms

            if buffer_full or timeout_reached:
                yield self.buffer.copy()
                self.buffer.clear()
                self.last_flush_time = time.time()

        # 刷新剩余token
        if self.buffer:
            yield self.buffer.copy()
            self.buffer.clear()

# 使用示例
async def optimized_stream():
    processor = BatchStreamProcessor(batch_size=5, max_latency_ms=30)

    async for token_batch in processor.process_stream(raw_token_stream()):
        # 一次发送5个token，减少网络往返
        yield {
            "tokens": token_batch,
            "count": len(token_batch)
        }

import time
from dataclasses import dataclass
from collections import deque

@dataclass
class FlowMetrics:
    """流量控制指标"""
    tokens_per_second: float
    network_latency_ms: float
    client_buffer_health: float  # 0-1, 1表示缓冲区充足
    timestamp: float

class AdaptiveFlowController:
    """自适应流控 - 根据网络状况动态调整发送速率"""

    def __init__(self):
        self.metrics_history = deque(maxlen=10)
        self.base_tokens_per_chunk = 1
        self.min_chunk_size = 1
        self.max_chunk_size = 20

    def update_metrics(self, metrics: FlowMetrics):
        """更新网络指标"""
        self.metrics_history.append(metrics)

    def get_optimal_chunk_size(self) -> int:
        """计算最优分块大小"""
        if len(self.metrics_history) < 3:
            return self.base_tokens_per_chunk

        # 计算平均指标
        avg_latency = sum(m.network_latency_ms for m in self.metrics_history) / len(self.metrics_history)
        avg_buffer_health = sum(m.client_buffer_health for m in self.metrics_history) / len(self.metrics_history)

        # 高延迟 + 缓冲区健康 → 增大chunk size
        if avg_latency > 100 and avg_buffer_health > 0.7:
            return min(self.max_chunk_size, int(self.base_tokens_per_chunk * 2))

        # 低延迟 → 减小chunk size，提高实时性
        if avg_latency < 30:
            return max(self.min_chunk_size, int(self.base_tokens_per_chunk * 0.5))

        # 缓冲区不足 → 减小chunk size
        if avg_buffer_health < 0.3:
            return self.min_chunk_size

        return self.base_tokens_per_chunk

    def should_backpressure(self) -> bool:
        """判断是否需要背压（减缓发送）"""
        if len(self.metrics_history) < 3:
            return False

        recent = list(self.metrics_history)[-3:]
        # 连续3次缓冲区不足
        return all(m.client_buffer_health < 0.2 for m in recent)

# 服务端集成
class StreamingLLMService:
    def __init__(self):
        self.flow_controller = AdaptiveFlowController()

    async def generate_with_flow_control(self, request):
        chunk_size = self.flow_controller.get_optimal_chunk_size()

        tokens = []
        async for token in self.model.generate(request):
            tokens.append(token)

            if len(tokens) >= chunk_size:
                # 检查背压
                if self.flow_controller.should_backpressure():
                    await asyncio.sleep(0.01)  # 短暂停顿

                yield tokens
                tokens = []

                # 动态调整chunk size
                chunk_size = self.flow_controller.get_optimal_chunk_size()

import hashlib
import numpy as np
from sentence_transformers import SentenceTransformer
import redis

class SemanticCache:
    """语义缓存 - 缓存相似请求的响应"""

    def __init__(self, redis_client, similarity_threshold=0.95):
        self.redis = redis_client
        self.encoder = SentenceTransformer('all-MiniLM-L6-v2')
        self.similarity_threshold = similarity_threshold
        self.cache_ttl = 3600  # 1小时

    def _get_embedding(self, text: str) -> np.ndarray:
        """获取文本的向量表示"""
        return self.encoder.encode(text)

    def _compute_similarity(self, emb1: np.ndarray, emb2: np.ndarray) -> float:
        """计算余弦相似度"""
        return np.dot(emb1, emb2) / (np.linalg.norm(emb1) * np.linalg.norm(emb2))

    async def get_cached_response(self, query: str) -> tuple[bool, str]:
        """
        查找语义相似的缓存

        返回: (是否命中, 缓存内容)
        """
        query_emb = self._get_embedding(query)

        # 获取所有缓存键（实际生产中使用向量数据库如Milvus/Pinecone）
        cache_keys = self.redis.keys("semantic_cache:*")

        for key in cache_keys:
            cached_data = self.redis.get(key)
            if not cached_data:
                continue

            data = json.loads(cached_data)  # json.loads将JSON字符串转为Python对象
            cached_emb = np.array(data['embedding'])

            similarity = self._compute_similarity(query_emb, cached_emb)

            if similarity >= self.similarity_threshold:
                return True, data['response']

        return False, None

    async def cache_response(self, query: str, response: str):
        """缓存响应"""
        query_emb = self._get_embedding(query)

        cache_key = f"semantic_cache:{hashlib.md5(query.encode()).hexdigest()}"

        data = {
            'query': query,
            'embedding': query_emb.tolist(),
            'response': response,
            'timestamp': time.time()
        }

        self.redis.setex(
            cache_key,
            self.cache_ttl,
            json.dumps(data)  # json.dumps将Python对象转为JSON字符串
        )

# 使用示例
class CachedLLMService:
    def __init__(self):
        self.cache = SemanticCache(redis.Redis())
        self.model = LLMModel()

    async def chat(self, request):
        query = request['messages'][-1]['content']  # 负索引：从末尾倒数访问元素

        # 1. 检查语义缓存
        cache_hit, cached_response = await self.cache.get_cached_response(query)

        if cache_hit:
            return {
                "response": cached_response,
                "cached": True,
                "source": "semantic_cache"
            }

        # 2. 调用模型
        response = await self.model.generate(request)

        # 3. 缓存结果
        await self.cache.cache_response(query, response)

        return {
            "response": response,
            "cached": False
        }

import asyncio
from typing import List, Dict, Any
from dataclasses import dataclass

@dataclass
class PendingRequest:
    """待处理的请求"""
    request_id: str
    prompt: str
    future: asyncio.Future
    timestamp: float

class RequestCoalescer:
    """
    请求合并器 - 合并相似的同时请求

    场景：多个用户同时询问相似问题，只调用一次模型
    """

    def __init__(self, merge_window_ms: float = 50):
        self.merge_window_ms = merge_window_ms
        self.pending_requests: List[PendingRequest] = []
        self.lock = asyncio.Lock()
        self.processing = False

    async def submit_request(self, request_id: str, prompt: str) -> str:
        """提交请求，可能与其他请求合并"""
        future = asyncio.get_event_loop().create_future()

        pending = PendingRequest(
            request_id=request_id,
            prompt=prompt,
            future=future,
            timestamp=time.time()
        )

        async with self.lock:
            self.pending_requests.append(pending)

            # 启动合并窗口
            if not self.processing:
                self.processing = True
                asyncio.create_task(self._process_batch())

        # 等待结果
        return await future

    async def _process_batch(self):
        """处理一批请求"""
        await asyncio.sleep(self.merge_window_ms / 1000)

        async with self.lock:
            batch = self.pending_requests.copy()
            self.pending_requests.clear()
            self.processing = False

        # 按相似度分组
        groups = self._group_similar_requests(batch)

        # 为每组调用一次模型
        for group in groups:
            response = await self._call_model_once(group)

            # 为组内所有请求设置结果
            for pending in group:
                pending.future.set_result(response)

    def _group_similar_requests(
        self,
        requests: List[PendingRequest]
    ) -> List[List[PendingRequest]]:
        """将请求按相似度分组"""
        # 简化实现：使用前缀匹配
        groups = []
        used = set()

        for i, req1 in enumerate(requests):
            if i in used:
                continue

            group = [req1]
            used.add(i)

            for j, req2 in enumerate(requests[i+1:], start=i+1):
                if j in used:
                    continue

                # 简单的相似度检查（实际使用embedding）
                if self._is_similar(req1.prompt, req2.prompt):
                    group.append(req2)
                    used.add(j)

            groups.append(group)

        return groups

    def _is_similar(self, prompt1: str, prompt2: str) -> bool:
        """判断两个prompt是否相似"""
        # 简化实现：检查共同子串长度
        # 实际生产中使用语义相似度
        words1 = set(prompt1.lower().split())
        words2 = set(prompt2.lower().split())

        if not words1 or not words2:
            return False

        intersection = words1 & words2
        similarity = len(intersection) / max(len(words1), len(words2))

        return similarity > 0.8

    async def _call_model_once(self, group: List[PendingRequest]) -> str:
        """为一组相似请求调用一次模型"""
        # 选择代表性prompt（最长的，通常信息最全）
        representative = max(group, key=lambda x: len(x.prompt))

        # 调用模型
        response = await call_llm_api(representative.prompt)

        return response

# 使用示例
coalescer = RequestCoalescer(merge_window_ms=100)

async def handle_user_request(user_id: str, prompt: str):
    """处理用户请求，自动合并相似请求"""
    response = await coalescer.submit_request(user_id, prompt)
    return response

# PyTorch分布式训练基础
import torch
import torch.distributed as dist
import torch.multiprocessing as mp
from torch.nn.parallel import DistributedDataParallel as DDP

def setup_distributed(rank, world_size):
    """初始化分布式环境"""
    # 使用NCCL后端（NVIDIA GPU最优）
    dist.init_process_group(
        backend='nccl',
        init_method='tcp://192.168.1.100:29500',
        world_size=world_size,
        rank=rank
    )

    # 设置当前设备
    torch.cuda.set_device(rank)

def train(rank, world_size):
    setup_distributed(rank, world_size)

    # 创建模型并包装为DDP
    model = MyModel().to(rank)
    ddp_model = DDP(model, device_ids=[rank])

    # 训练循环
    for data, target in dataloader:
        output = ddp_model(data)
        loss = criterion(output, target)
        loss.backward()
        optimizer.step()
        # DDP自动处理梯度同步

# 启动多进程训练
if __name__ == '__main__':
    world_size = 4  # 4个GPU
    mp.spawn(train, args=(world_size,), nprocs=world_size)

# DeepSpeed ZeRO 配置示例
from deepspeed import DeepSpeedEngine
import deepspeed

# ZeRO Stage 3 配置 - 极致内存优化
zero_stage3_config = {
    "fp16": {
        "enabled": True,
        "loss_scale": 0,
        "loss_scale_window": 1000,
        "initial_scale_power": 16,
        "hysteresis": 2,
        "min_loss_scale": 1
    },
    "zero_optimization": {
        "stage": 3,  # ZeRO Stage 3
        "offload_optimizer": {
            "device": "cpu",
            "pin_memory": True
        },
        "offload_param": {
            "device": "cpu",
            "pin_memory": True
        },
        "overlap_comm": True,  # 重叠通信和计算
        "contiguous_gradients": True,
        "sub_group_size": 1e9,
        "reduce_bucket_size": "auto",
        "stage3_prefetch_bucket_size": "auto",
        "stage3_param_persistence_threshold": "auto",
        "stage3_max_live_parameters": 1e9,
        "stage3_max_reuse_distance": 1e9,
        "stage3_gather_16bit_weights_on_model_save": True
    },
    "gradient_accumulation_steps": 4,
    "gradient_clipping": 1.0,
    "steps_per_print": 2000,
    "wall_clock_breakdown": False
}

# 初始化模型
model_engine, optimizer, _, _ = deepspeed.initialize(
    model=model,
    model_parameters=model.parameters(),
    config=zero_stage3_config
)

# 训练循环
for step, batch in enumerate(dataloader):
    loss = model_engine(batch)
    model_engine.backward(loss)
    model_engine.step()

# Megatron-LM 3D并行配置
from megatron import get_args
from megatron.model import GPTModel
from megatron.training import pretrain

# 3D并行配置
parallel_config = {
    # 数据并行 (Data Parallelism)
    "data_parallel_size": 8,      # 8路数据并行

    # 张量并行 (Tensor Parallelism)
    "tensor_model_parallel_size": 4,  # 4路张量并行

    # 流水线并行 (Pipeline Parallelism)
    "pipeline_model_parallel_size": 2,  # 2路流水线并行

    # 总GPU数 = 8 * 4 * 2 = 64
}

# 通信组划分
"""
总GPU数: 64

数据并行组 (8组，每组8个GPU):
- DP Group 0: GPU 0-7
- DP Group 1: GPU 8-15
- ...

张量并行组 (16组，每组4个GPU):
- TP Group 0: GPU 0, 16, 32, 48
- TP Group 1: GPU 1, 17, 33, 49
- ...

流水线并行组 (32组，每组2个GPU):
- PP Group 0: GPU 0, 8
- PP Group 1: GPU 1, 9
- ...
"""

# 通信优化参数
communication_config = {
    # 梯度累积
    "gradient_accumulation_fusion": True,

    # 异步通信
    "async_tensor_model_parallel_allreduce": True,

    # 批量通信
    "bucket_size": 50 * 1024 * 1024,  # 50MB

    # 重叠计算和通信
    "overlap_p2p_comm": True,
    "overlap_grad_reduce": True,
}

import torch.distributed as dist
from torch.distributed.distributed_c10d import _get_default_group

class TopologyAwareScheduler:
    """拓扑感知调度器 - 最小化跨节点通信"""

    def __init__(self, num_nodes, gpus_per_node):
        self.num_nodes = num_nodes
        self.gpus_per_node = gpus_per_node
        self.total_gpus = num_nodes * gpus_per_node

        # 构建拓扑矩阵
        self.topology = self._build_topology_matrix()

    def _build_topology_matrix(self):
        """
        构建网络拓扑矩阵

        带宽矩阵 (GB/s):
        - 同节点GPU: 900 (NVLink)
        - 同机架节点: 100 (InfiniBand)
        - 跨机架节点: 25 (以太网)
        """
        topology = {}

        for i in range(self.total_gpus):
            for j in range(self.total_gpus):
                node_i = i // self.gpus_per_node
                node_j = j // self.gpus_per_node

                if node_i == node_j:
                    # 同节点 - NVLink
                    topology[(i, j)] = {"bandwidth": 900, "latency": 1}
                elif abs(node_i - node_j) <= 2:
                    # 同机架 - InfiniBand
                    topology[(i, j)] = {"bandwidth": 100, "latency": 5}
                else:
                    # 跨机架 - 以太网
                    topology[(i, j)] = {"bandwidth": 25, "latency": 50}

        return topology

    def optimize_placement(self, model_layers):
        """
        优化模型层放置，最小化通信开销

        策略：
        1. 张量并行放在同节点（NVLink高速通信）
        2. 流水线并行放在同机架（InfiniBand）
        3. 数据并行跨机架（通信量小）
        """
        placement = {}

        layers_per_pipeline = len(model_layers) // self.num_nodes

        for layer_idx, layer in enumerate(model_layers):
            # 确定流水线阶段
            pipeline_stage = layer_idx // layers_per_pipeline

            # 该阶段使用的节点
            node_id = pipeline_stage % self.num_nodes

            # 该层使用的GPU（同节点内张量并行）
            gpus = list(range(
                node_id * self.gpus_per_node,
                (node_id + 1) * self.gpus_per_node
            ))

            placement[layer] = {
                "node": node_id,
                "gpus": gpus,
                "pipeline_stage": pipeline_stage
            }

        return placement

    def estimate_communication_cost(self, placement, batch_size):
        """估算通信开销"""
        total_cost = 0

        # 计算每种通信的代价
        for layer, info in placement.items():
            # 张量并行通信（同节点，低成本）
            tp_cost = self._tensor_parallel_cost(info["gpus"])

            # 流水线并行通信（跨层，中成本）
            pp_cost = self._pipeline_parallel_cost(info["pipeline_stage"])

            total_cost += tp_cost + pp_cost

        return total_cost

    def _tensor_parallel_cost(self, gpus):
        """张量并行通信成本"""
        # 同节点NVLink，成本极低
        return len(gpus) * 0.1  # 归一化成本

    def _pipeline_parallel_cost(self, stage):
        """流水线并行通信成本"""
        # 跨节点通信
        return 1.0 if stage > 0 else 0.0

# 使用示例
scheduler = TopologyAwareScheduler(num_nodes=8, gpus_per_node=8)
placement = scheduler.optimize_placement(model_layers=range(32))

import torch
from torch.profiler import profile, ProfilerActivity
import time

class NetworkMonitor:
    """分布式训练网络监控"""

    def __init__(self):
        self.metrics = {
            "allreduce_times": [],
            "bandwidth_utilization": [],
            "latency_ms": [],
        }

    def monitor_allreduce(self, tensor_size_bytes):
        """监控AllReduce操作"""
        start = time.time()

        # 执行AllReduce
        tensor = torch.randn(tensor_size_bytes // 4, device="cuda")
        dist.all_reduce(tensor)
        torch.cuda.synchronize()

        elapsed = time.time() - start

        # 计算带宽利用率
        # AllReduce通信量 = 2 * (n-1)/n * tensor_size (Ring算法)
        world_size = dist.get_world_size()
        comm_volume = 2 * (world_size - 1) / world_size * tensor_size_bytes
        bandwidth = comm_volume / elapsed / 1e9  # GB/s

        self.metrics["allreduce_times"].append(elapsed)
        self.metrics["bandwidth_utilization"].append(bandwidth)

        return {
            "time_ms": elapsed * 1000,
            "bandwidth_gbps": bandwidth,
            "efficiency": bandwidth / 100 * 100  # 假设理论带宽100GB/s
        }

    def generate_report(self):
        """生成网络性能报告"""
        if not self.metrics["allreduce_times"]:
            return "No data collected"

        avg_time = sum(self.metrics["allreduce_times"]) / len(self.metrics["allreduce_times"])
        avg_bw = sum(self.metrics["bandwidth_utilization"]) / len(self.metrics["bandwidth_utilization"])

        report = f"""
=== Network Performance Report ===
Average AllReduce Time: {avg_time * 1000:.2f} ms
Average Bandwidth: {avg_bw:.2f} GB/s
Bandwidth Efficiency: {avg_bw / 100 * 100:.1f}%

Recommendations:
"""

        if avg_bw < 50:
            report += "- WARNING: Low bandwidth utilization. Check network configuration.\n"
            report += "  - Verify NCCL settings: NCCL_IB_HCA, NCCL_SOCKET_IFNAME\n"
            report += "  - Check for network congestion\n"
        elif avg_bw < 80:
            report += "- Moderate bandwidth utilization. Consider optimization:\n"
            report += "  - Enable NCCL tree algorithm for large messages\n"
            report += "  - Adjust NCCL_BUFFSIZE\n"
        else:
            report += "- Good bandwidth utilization.\n"

        return report

# NCCL调试环境变量
"""
# 启用NCCL调试
export NCCL_DEBUG=INFO
export NCCL_DEBUG_SUBSYS=ALL

# 指定InfiniBand设备
export NCCL_IB_HCA=mlx5_0,mlx5_1

# 禁用TCP回退（强制使用RDMA）
export NCCL_IB_DISABLE=0
export NCCL_SOCKET_IFNAME=ib0

# 调整缓冲区大小
export NCCL_BUFFSIZE=2097152  # 2MB

# 启用树形算法（大消息优化）
export NCCL_ALGO=RING,TREE
"""

# 网络诊断脚本
def diagnose_network():
    """网络诊断"""
    print("=== Network Diagnostics ===\n")

    # 1. 检查NCCL版本
    print("1. NCCL Version:")
    !python -c "import torch; print(torch.cuda.nccl.version())"

    # 2. 检查GPU互联
    print("\n2. GPU Interconnect:")
    !nvidia-smi topo -m

    # 3. 检查InfiniBand设备
    print("\n3. InfiniBand Devices:")
    !ibstat

    # 4. 测试节点间带宽
    print("\n4. Testing Inter-node Bandwidth:")
    # 使用nccl-tests
    !mpirun -np 2 -H node1:1,node2:1 ./build/all_reduce_perf -b 8M -e 128M

    # 5. 检查网络配置
    print("\n5. Network Configuration:")
    print(f"NCCL_DEBUG: {os.environ.get('NCCL_DEBUG', 'Not set')}")
    print(f"NCCL_IB_HCA: {os.environ.get('NCCL_IB_HCA', 'Not set')}")
    print(f"NCCL_SOCKET_IFNAME: {os.environ.get('NCCL_SOCKET_IFNAME', 'Not set')}")

# 运行诊断
# diagnose_network()

import secrets
import hashlib
import hmac
from datetime import datetime, timedelta

class APIKeyManager:
    """安全的API密钥管理系统"""

    def __init__(self):
        self.master_key = secrets.token_hex(32)  # 256位主密钥

    def generate_key(self, user_id, tier='standard'):
        """生成带权限层级的API密钥"""
        # 生成随机密钥
        random_part = secrets.token_urlsafe(32)

        # 构造密钥：前缀 + 用户ID + 随机部分
        api_key = f"sk-{tier}-{user_id[:8]}-{random_part}"  # 切片操作：[start:end:step]提取子序列

        # 存储密钥哈希（不存明文）
        key_hash = hashlib.sha256(api_key.encode()).hexdigest()

        return {
            'api_key': api_key,  # 仅显示一次
            'key_hash': key_hash,
            'created_at': datetime.now(),
            'tier': tier,
            'rate_limit': self._get_tier_limit(tier)
        }

    def verify_key(self, api_key):
        """验证API密钥"""
        key_hash = hashlib.sha256(api_key.encode()).hexdigest()
        # 查询数据库验证hash...
        return self._lookup_key(key_hash)

    def _get_tier_limit(self, tier):
        limits = {
            'free': {'rpm': 20, 'tpd': 1000},
            'standard': {'rpm': 60, 'tpd': 10000},
            'premium': {'rpm': 1000, 'tpd': 100000}
        }
        return limits.get(tier, limits['free'])

class KeyRotation:
    """自动密钥轮换"""

    def rotate_key(self, user_id):
        """轮换用户API密钥"""
        # 1. 生成新密钥
        new_key = self.generate_key(user_id)

        # 2. 标记旧密钥（宽限期24小时）
        self.mark_key_deprecated(user_id, grace_period=24)

        # 3. 通知用户
        self.notify_user(user_id, new_key['api_key'])

        return new_key

    def cleanup_deprecated_keys(self):
        """清理过期的废弃密钥"""
        cutoff = datetime.now() - timedelta(hours=24)
        # 删除宽限期结束的旧密钥
        self.db.keys.delete_many({
            'deprecated': True,
            'deprecated_at': {'$lt': cutoff}
        })

import jwt
from datetime import datetime, timedelta

class JWTAuth:
    def __init__(self, secret_key):
        self.secret = secret_key

    def create_token(self, user_id, permissions):
        """创建JWT访问令牌"""
        payload = {
            'sub': user_id,
            'permissions': permissions,
            'iat': datetime.utcnow(),
            'exp': datetime.utcnow() + timedelta(hours=1),
            'type': 'access'
        }
        return jwt.encode(payload, self.secret, algorithm='HS256')

    def verify_token(self, token):
        """验证JWT令牌"""
        try:
            payload = jwt.decode(token, self.secret, algorithms=['HS256'])
            return payload
        except jwt.ExpiredSignatureError:
            raise Exception("Token已过期")
        except jwt.InvalidTokenError:
            raise Exception("无效的Token")

# FastAPI集成示例
from fastapi import FastAPI, Depends, HTTPException
from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials

app = FastAPI()
security = HTTPBearer()
auth = JWTAuth(secret_key="your-secret-key")

async def get_current_user(credentials: HTTPAuthorizationCredentials = Depends(security)):
    token = credentials.credentials
    try:
        payload = auth.verify_token(token)
        return payload
    except Exception as e:
        raise HTTPException(status_code=401, detail=str(e))

@app.post("/v1/chat/completions")
async def chat(
    request: ChatRequest,
    user: dict = Depends(get_current_user)
):
    # 检查权限
    if 'chat:write' not in user['permissions']:
        raise HTTPException(status_code=403, detail="权限不足")

    # 处理请求...

import geoip2.database
from ipaddress import ip_address, ip_network

class AccessControl:
    def __init__(self):
        self.geo_reader = geoip2.database.Reader('GeoLite2-Country.mmdb')

    def check_access(self, client_ip, user_config):
        """检查访问权限"""
        # IP白名单检查
        if user_config.get('ip_whitelist'):
            if not self._ip_in_whitelist(client_ip, user_config['ip_whitelist']):
                return False, "IP不在白名单中"

        # 地理限制检查
        if user_config.get('allowed_countries'):
            country = self._get_country(client_ip)
            if country not in user_config['allowed_countries']:
                return False, f"该服务在{country}不可用"

        return True, "允许访问"

    def _ip_in_whitelist(self, ip, whitelist):
        client_ip = ip_address(ip)
        for allowed in whitelist:
            if client_ip in ip_network(allowed, strict=False):
                return True
        return False

import hmac
import hashlib
import time

class RequestSigner:
    """请求签名验证，防止重放攻击"""

    def sign_request(self, api_key, secret, method, path, body):
        """生成请求签名"""
        timestamp = str(int(time.time()))
        nonce = secrets.token_hex(16)

        # 构造签名字符串
        string_to_sign = f"{method}\n{path}\n{timestamp}\n{nonce}\n{body}"

        # HMAC-SHA256签名
        signature = hmac.new(
            secret.encode(),
            string_to_sign.encode(),
            hashlib.sha256
        ).hexdigest()

        return {
            'X-Timestamp': timestamp,
            'X-Nonce': nonce,
            'X-Signature': signature
        }

    def verify_signature(self, headers, secret, method, path, body):
        """验证请求签名"""
        timestamp = int(headers.get('X-Timestamp'))

        # 检查时间戳（5分钟有效期）
        if abs(time.time() - timestamp) > 300:
            return False, "请求已过期"

        # 检查nonce是否已使用（防重放）
        nonce = headers.get('X-Nonce')
        if self._is_nonce_used(nonce):
            return False, "重复的请求"

        # 验证签名
        expected = self.sign_request(None, secret, method, path, body)
        if headers.get('X-Signature') != expected['X-Signature']:
            return False, "签名无效"

        # 记录nonce
        self._record_nonce(nonce)
        return True, "验证通过"

import re

class ContentFilter:
    """输入/输出内容安全过滤"""

    # 敏感模式（简化示例）
    SENSITIVE_PATTERNS = [
        r'\b\d{16,19}\b',  # 银行卡号
        r'\b\d{18}\b',     # 身份证号
        r'[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}',  # 邮箱
    ]

    def filter_input(self, text):
        """过滤用户输入"""
        # 检测提示词注入
        if self._detect_prompt_injection(text):
            raise ValueError("检测到提示词注入攻击")

        # 检测敏感信息
        for pattern in self.SENSITIVE_PATTERNS:
            if re.search(pattern, text):
                return self._mask_sensitive(text, pattern)

        return text

    def _detect_prompt_injection(self, text):
        """检测提示词注入"""
        injection_patterns = [
            r'ignore previous instructions',
            r'system prompt',
            r'you are now',
            r'forget (everything|all)',
        ]
        return any(re.search(p, text, re.I) for p in injection_patterns)

    def filter_output(self, text, user_tier):
        """过滤模型输出"""
        # 根据用户等级过滤内容
        if user_tier == 'free':
            # 免费用户添加水印
            text += "\n\n[Generated by AI Model]"

        return text

# 完整的高性能API网关实现
import asyncio
import time
from typing import Optional, Dict, Any
from dataclasses import dataclass
from enum import Enum
import redis.asyncio as redis
from fastapi import FastAPI, HTTPException, Request, Response
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import StreamingResponse
import prometheus_client
from prometheus_client import Counter, Histogram, Gauge  # Counter计数器：统计元素出现次数

# Prometheus指标
REQUEST_COUNT = Counter('llm_requests_total', 'Total requests', ['model', 'status'])
REQUEST_DURATION = Histogram('llm_request_duration_seconds', 'Request duration', ['model'])
TOKENS_GENERATED = Counter('llm_tokens_generated_total', 'Total tokens', ['model'])
ACTIVE_CONNECTIONS = Gauge('llm_active_connections', 'Active connections')
QUEUE_SIZE = Gauge('llm_queue_size', 'Queue size')

class ModelBackend(Enum):
    """模型后端类型"""
    VLLM = "vllm"
    TGI = "tgi"
    TENSORRT = "tensorrt"

@dataclass  # 自动生成__init__等方法
class ModelConfig:
    """模型配置"""
    name: str
    backend: ModelBackend
    endpoint: str
    max_concurrency: int
    timeout: float
    priority: int  # 优先级，数字越小优先级越高

class AdaptiveRateLimiter:
    """自适应限流器 - 根据后端负载动态调整"""

    def __init__(self, redis_client: redis.Redis, base_rate: int = 100):
        self.redis = redis_client
        self.base_rate = base_rate
        self.current_rate = base_rate
        self.load_history = []

    async def update_load_metrics(self, model: str, latency: float, queue_size: int):
        """更新负载指标并调整限流"""
        self.load_history.append({
            'latency': latency,
            'queue': queue_size,
            'timestamp': time.time()
        })

        # 只保留最近10个数据点
        self.load_history = self.load_history[-10:]

        # 计算平均负载
        avg_latency = sum(h['latency'] for h in self.load_history) / len(self.load_history)
        avg_queue = sum(h['queue'] for h in self.load_history) / len(self.load_history)

        # 动态调整限流
        if avg_latency > 5.0 or avg_queue > 10:  # 高负载
            self.current_rate = max(self.base_rate // 2, 10)
        elif avg_latency < 1.0 and avg_queue < 2:  # 低负载
            self.current_rate = min(self.base_rate * 2, 1000)

    async def is_allowed(self, api_key: str) -> bool:
        """检查是否允许请求"""
        key = f"rate_limit:{api_key}"
        current = await self.redis.get(key)

        if current and int(current) >= self.current_rate:
            return False

        pipe = self.redis.pipeline()
        pipe.incr(key)
        pipe.expire(key, 60)
        await pipe.execute()
        return True

class CircuitBreaker:
    """熔断器 - 防止级联故障"""

    def __init__(self, failure_threshold: int = 5, recovery_timeout: float = 30.0):
        self.failure_threshold = failure_threshold
        self.recovery_timeout = recovery_timeout
        self.failures = 0
        self.last_failure_time = None
        self.state = "closed"  # closed, open, half-open

    def call(self, func):
        """装饰器：熔断保护"""
        async def wrapper(*args, **kwargs):  # *args接收任意位置参数；**kwargs接收任意关键字参数
            if self.state == "open":
                if time.time() - self.last_failure_time > self.recovery_timeout:
                    self.state = "half-open"
                else:
                    raise HTTPException(503, "Service temporarily unavailable")

            try:  # try/except捕获异常
                result = await func(*args, **kwargs)
                if self.state == "half-open":
                    self.state = "closed"
                    self.failures = 0
                return result
            except Exception as e:
                self.failures += 1
                self.last_failure_time = time.time()

                if self.failures >= self.failure_threshold:
                    self.state = "open"

                raise e

        return wrapper

class LLMAPIGateway:
    """大模型API网关"""

    def __init__(self):
        self.app = FastAPI(title="LLM API Gateway", version="2.0.0")
        self.redis = redis.Redis(host='localhost', port=6379, db=0)
        self.rate_limiter = AdaptiveRateLimiter(self.redis)
        self.circuit_breakers: Dict[str, CircuitBreaker] = {}
        self.model_configs: Dict[str, ModelConfig] = {}
        self.request_queue = asyncio.Queue(maxsize=1000)

        self._setup_middleware()
        self._setup_routes()
        self._start_background_tasks()

    def _setup_middleware(self):
        """设置中间件"""
        # CORS
        self.app.add_middleware(
            CORSMiddleware,
            allow_origins=["*"],
            allow_methods=["*"],
            allow_headers=["*"],
        )

        # 请求日志和指标
        @self.app.middleware("http")
        async def metrics_middleware(request: Request, call_next):
            start_time = time.time()
            ACTIVE_CONNECTIONS.inc()

            response = await call_next(request)

            duration = time.time() - start_time
            REQUEST_DURATION.labels(model=request.path_params.get('model', 'unknown')).observe(duration)
            ACTIVE_CONNECTIONS.dec()

            return response

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

        @self.app.post("/v1/chat/completions")
        async def chat_completions(request: Request):
            """主聊天接口"""
            body = await request.json()
            model = body.get('model', 'default')
            stream = body.get('stream', False)

            # 限流检查
            api_key = request.headers.get('Authorization', '').replace('Bearer ', '')
            if not await self.rate_limiter.is_allowed(api_key):
                raise HTTPException(429, "Rate limit exceeded")

            # 路由到合适的后端
            backend = self._select_backend(model, body)

            # 熔断保护
            breaker = self.circuit_breakers.get(backend.endpoint, CircuitBreaker())
            self.circuit_breakers[backend.endpoint] = breaker

            if stream:
                return StreamingResponse(
                    self._stream_generate(body, backend, breaker),
                    media_type="text/event-stream"
                )
            else:
                return await breaker.call(self._generate)(body, backend)

        @self.app.get("/health")
        async def health_check():
            """健康检查"""
            health_status = {}
            for name, config in self.model_configs.items():
                try:
                    # 检查后端健康
                    healthy = await self._check_backend_health(config)
                    health_status[name] = "healthy" if healthy else "unhealthy"
                except:
                    health_status[name] = "unknown"

            return {
                "status": "healthy" if all(s == "healthy" for s in health_status.values()) else "degraded",
                "models": health_status,
                "rate_limit": self.rate_limiter.current_rate
            }

        @self.app.get("/metrics")
        async def metrics():
            """Prometheus指标"""
            return Response(
                content=prometheus_client.generate_latest(),
                media_type="text/plain"
            )

    def _select_backend(self, model: str, request_body: Dict) -> ModelConfig:
        """智能路由选择"""
        # 根据模型名称、负载、成本选择后端
        candidates = [c for c in self.model_configs.values() if c.name == model]

        if not candidates:
            raise HTTPException(404, f"Model {model} not found")

        # 选择负载最低的后端
        return min(candidates, key=lambda c: self._get_backend_load(c))  # lambda匿名函数：简洁的单行函数

    async def _stream_generate(self, body: Dict, backend: ModelConfig, breaker: CircuitBreaker):
        """流式生成"""
        try:
            async for chunk in self._call_backend_stream(body, backend):
                yield f"data: {chunk}\n\n"
            yield "data: [DONE]\n\n"
        except Exception as e:
            yield f"data: {{\"error\": \"{str(e)}\"}}\n\n"

    async def _generate(self, body: Dict, backend: ModelConfig) -> Dict:
        """非流式生成"""
        # 实际调用后端
        response = await self._call_backend(body, backend)

        # 更新指标
        REQUEST_COUNT.labels(model=backend.name, status="200").inc()
        if 'usage' in response:
            TOKENS_GENERATED.labels(model=backend.name).inc(response['usage'].get('total_tokens', 0))

        return response

    def _start_background_tasks(self):
        """启动后台任务"""
        @self.app.on_event("startup")
        async def startup():
            asyncio.create_task(self._metrics_updater())

        @self.app.on_event("shutdown")
        async def shutdown():
            await self.redis.close()

    async def _metrics_updater(self):
        """定期更新指标"""
        while True:
            await asyncio.sleep(10)
            # 更新负载指标
            for name, config in self.model_configs.items():
                load = await self._get_backend_load(config)
                await self.rate_limiter.update_load_metrics(name, load, 0)

# 启动网关
gateway = LLMAPIGateway()
app = gateway.app

# 运行: uvicorn gateway:app --host 0.0.0.0 --port 8000