Agent 系统可扩展性设计：从单机到分布式

from dataclasses import dataclass, field
from typing import Optional, Any
import json

# ❌ 错误做法：状态在内存中
class StatefulAgent:
    def __init__(self):
        self.conversations: dict[str, list] = {}  # 内存中的会话状态
    
    def chat(self, session_id: str, message: str) -> str:
        history = self.conversations[session_id]  # 只在本机有效
        history.append({"role": "user", "content": message})
        # ...处理...
        return "response"

# ✅ 正确做法：状态外部化
class StatelessAgent:
    """无状态 Agent，所有会话状态存入 Redis"""
    
    def __init__(self, state_store: 'SessionStore'):
        self.store = state_store
    
    def chat(self, session_id: str, message: str) -> str:
        # 每次从外部存储加载会话状态
        history = self.store.get_history(session_id)
        history.append({"role": "user", "content": message})
        
        # 处理逻辑
        response = self._generate(history)
        history.append({"role": "assistant", "content": response})
        
        # 持久化回外部存储
        self.store.save_history(session_id, history)
        return response


class SessionStore:
    """基于 Redis 的会话存储"""
    
    def __init__(self, redis_client):
        self.redis = redis_client
        self.ttl = 3600  # 会话过期时间 1 小时
    
    def get_history(self, session_id: str) -> list[dict]:
        data = self.redis.get(f"session:{session_id}")
        return json.loads(data) if data else []
    
    def save_history(self, session_id: str, history: list[dict]):
        self.redis.setex(
            f"session:{session_id}",
            self.ttl,
            json.dumps(history, ensure_ascii=False)
        )
    
    def extend_ttl(self, session_id: str):
        self.redis.expire(f"session:{session_id}", self.ttl)

"""
三层架构设计：
- 接入层：处理 HTTP/WebSocket 连接，会话路由
- 推理层：LLM 推理，可水平扩展
- 知识层：向量数据库、知识图谱
"""
from abc import ABC, abstractmethod

class GatewayLayer:
    """接入层：路由请求到推理节点"""
    
    def __init__(self, inference_pool: 'InferencePool'):
        self.pool = inference_pool
    
    async def handle_request(self, session_id: str, message: str):
        # 根据会话亲和性选择节点
        node = self.pool.select_node(session_id)
        return await node.process(session_id, message)

class InferenceNode:
    """推理层：单个推理节点"""
    
    def __init__(self, node_id: str, model_path: str):
        self.node_id = node_id
        self.model = self._load_model(model_path)
        self.active_sessions = 0
        self.max_sessions = 10
    
    async def process(self, session_id: str, message: str) -> str:
        self.active_sessions += 1
        try:
            # 推理逻辑
            result = await self._infer(message)
            return result
        finally:
            self.active_sessions -= 1
    
    @property
    def available(self) -> bool:
        return self.active_sessions < self.max_sessions

class InferencePool:
    """推理节点池"""
    
    def __init__(self):
        self.nodes: list[InferenceNode] = []
    
    def add_node(self, node: InferenceNode):
        self.nodes.append(node)
    
    def select_node(self, session_id: str = None) -> InferenceNode:
        # 优先选择可用且负载最低的节点
        available = [n for n in self.nodes if n.available]
        if not available:
            raise RuntimeError("No available inference nodes")
        return min(available, key=lambda n: n.active_sessions)

import asyncio
from typing import Any

class AgentOrchestrator:
    """编排器：协调推理服务和工具服务"""
    
    def __init__(self, llm_client: 'LLMClient', tool_service: 'ToolService'):
        self.llm = llm_client
        self.tools = tool_service
    
    async def run(self, query: str, session_id: str) -> str:
        messages = [{"role": "user", "content": query}]
        
        for _ in range(10):  # 最多 10 轮
            # 1. 调用 LLM（GPU 节点）
            response = await self.llm.chat(messages)
            
            if not response.tool_calls:
                return response.content
            
            # 2. 执行工具调用（CPU 节点，可并行）
            tool_results = await asyncio.gather(*[
                self.tools.execute(tc) for tc in response.tool_calls
            ])
            
            # 3. 合并结果，继续推理
            messages.append({"role": "assistant", "content": response.content})
            for tr in tool_results:
                messages.append({"role": "tool", "content": tr})
        
        return "达到最大轮次限制"

class LLMClient:
    """LLM 推理客户端"""
    
    def __init__(self, endpoint: str):
        self.endpoint = endpoint
    
    async def chat(self, messages: list[dict]) -> 'LLMResponse':
        # 调用推理服务
        pass

class ToolService:
    """工具执行服务"""
    
    def __init__(self, workers: int = 20):
        self.semaphore = asyncio.Semaphore(workers)
    
    async def execute(self, tool_call: dict) -> str:
        async with self.semaphore:
            # 执行工具调用
            pass

import hashlib
from collections import defaultdict

class LoadBalancer:
    """多种负载均衡策略"""
    
    def __init__(self, nodes: list[str]):
        self.nodes = nodes
        self._index = 0
        self._connections: dict[str, int] = defaultdict(int)
        self._session_map: dict[str, str] = {}  # session_id -> node
    
    def round_robin(self) -> str:
        """轮询：简单但不考虑负载"""
        node = self.nodes[self._index % len(self.nodes)]
        self._index += 1
        return node
    
    def least_connections(self) -> str:
        """最少连接：选择当前连接数最少的节点"""
        return min(self.nodes, key=lambda n: self._connections[n])
    
    def session_affinity(self, session_id: str) -> str:
        """会话亲和：同一会话路由到同一节点
        
        利用一致性哈希，节点增减时最小化会话迁移
        """
        if session_id in self._session_map:
            return self._session_map[session_id]
        
        # 一致性哈希
        ring = {}
        for node in self.nodes:
            for i in range(150):  # 虚拟节点
                key = f"{node}:{i}"
                ring[self._hash(key)] = node
        
        hash_val = self._hash(session_id)
        sorted_keys = sorted(ring.keys())
        # 找到第一个 >= hash_val 的节点
        for k in sorted_keys:
            if k >= hash_val:
                node = ring[k]
                self._session_map[session_id] = node
                return node
        
        node = ring[sorted_keys[0]]
        self._session_map[session_id] = node
        return node
    
    def _hash(self, key: str) -> int:
        return int(hashlib.md5(key.encode()).hexdigest(), 16)
    
    def on_connection_start(self, node: str):
        self._connections[node] += 1
    
    def on_connection_end(self, node: str):
        self._connections[node] -= 1

class StickySessionManager:
    """粘性会话管理器"""
    
    def __init__(self, ttl: int = 1800):
        self.ttl = ttl
        self._bindings: dict[str, tuple[str, float]] = {}  # session -> (node, timestamp)
    
    def get_node(self, session_id: str, available_nodes: list[str]) -> str:
        """获取会话绑定的节点，如过期或不可用则重新分配"""
        import time
        
        if session_id in self._bindings:
            node, bound_time = self._bindings[session_id]
            if time.time() - bound_time < self.ttl and node in available_nodes:
                return node
        
        # 重新分配
        node = min(available_nodes, key=lambda n: sum(
            1 for _, (nd, _) in self._bindings.values() if nd == n
        ))
        self._bindings[session_id] = (node, time.time())
        return node
    
    def remove_binding(self, session_id: str):
        self._bindings.pop(session_id, None)
    
    def cleanup_expired(self):
        import time
        now = time.time()
        expired = [sid for sid, (_, t) in self._bindings.items() if now - t > self.ttl]
        for sid in expired:
            del self._bindings[sid]

import uuid
from contextvars import ContextVar
from dataclasses import dataclass, field
from typing import Optional

# 上下文变量传播 trace
trace_id_var: ContextVar[str] = ContextVar("trace_id", default="")
span_id_var: ContextVar[str] = ContextVar("span_id", default="")

@dataclass
class Span:
    """追踪片段"""
    trace_id: str
    span_id: str
    parent_span_id: Optional[str]
    operation: str
    start_time: float
    end_time: Optional[float] = None
    tags: dict = field(default_factory=dict)
    status: str = "ok"

class Tracer:
    """分布式追踪器"""
    
    def __init__(self):
        self.spans: list[Span] = []
    
    def start_span(self, operation: str, **tags) -> str:
        import time
        
        trace_id = trace_id_var.get() or str(uuid.uuid4())
        parent_id = span_id_var.get() or None
        span_id = str(uuid.uuid4())
        
        span = Span(
            trace_id=trace_id,
            span_id=span_id,
            parent_span_id=parent_id,
            operation=operation,
            start_time=time.time(),
            tags=tags
        )
        self.spans.append(span)
        
        # 传播上下文
        trace_id_var.set(trace_id)
        span_id_var.set(span_id)
        return span_id
    
    def finish_span(self, span_id: str, status: str = "ok"):
        import time
        for span in self.spans:
            if span.span_id == span_id:
                span.end_time = time.time()
                span.status = status
                break

# 使用示例
tracer = Tracer()

async def handle_user_request(session_id: str, message: str):
    span_id = tracer.start_span("handle_request", session_id=session_id)
    try:
        # 子操作：LLM 推理
        llm_span = tracer.start_span("llm_inference")
        response = await call_llm(message)
        tracer.finish_span(llm_span)
        
        # 子操作：工具调用
        if needs_tool(response):
            tool_span = tracer.start_span("tool_execution")
            tool_result = await execute_tool(response)
            tracer.finish_span(tool_span)
        
        tracer.finish_span(span_id, "ok")
        return response
    except Exception as e:
        tracer.finish_span(span_id, "error")
        raise

async def call_llm(message: str) -> str:
    await asyncio.sleep(0.5)  # 模拟推理
    return f"Response to: {message}"

async def execute_tool(response: str) -> str:
    await asyncio.sleep(0.2)
    return "tool result"

def needs_tool(response: str) -> bool:
    return "tool" in response.lower()

请求链路示例:
trace_id: a1b2c3d4
├── handle_request (1250ms) 
│   ├── llm_inference (800ms) ← GPU 推理是瓶颈
│   ├── tool_execution (350ms)
│   │   ├── search_api (200ms)
│   │   └── parse_result (150ms)
│   └── format_response (100ms)

class AutoScaler:
    """基于负载的自动扩缩容"""
    
    def __init__(self, min_nodes: int = 2, max_nodes: int = 20,
                 scale_up_threshold: float = 0.7,
                 scale_down_threshold: float = 0.3):
        self.min_nodes = min_nodes
        self.max_nodes = max_nodes
        self.scale_up_threshold = scale_up_threshold
        self.scale_down_threshold = scale_down_threshold
        self.current_nodes = min_nodes
    
    def evaluate(self, metrics: dict) -> str:
        """根据指标决定扩缩容
        
        metrics: {
            "avg_gpu_utilization": 0.85,
            "avg_memory_usage": 0.72,
            "request_queue_depth": 15,
            "avg_response_time_ms": 3500
        }
        """
        load_score = self._calculate_load_score(metrics)
        
        if load_score > self.scale_up_threshold and self.current_nodes < self.max_nodes:
            return self._scale_up()
        elif load_score < self.scale_down_threshold and self.current_nodes > self.min_nodes:
            return self._scale_down()
        return "no_action"
    
    def _calculate_load_score(self, metrics: dict) -> float:
        """综合负载评分 [0, 1]"""
        weights = {
            "avg_gpu_utilization": 0.4,
            "avg_memory_usage": 0.2,
            "request_queue_depth_norm": 0.2,  # 归一化后的队列深度
            "avg_response_time_norm": 0.2
        }
        # 归一化
        queue_norm = min(metrics.get("request_queue_depth", 0) / 50, 1.0)
        latency_norm = min(metrics.get("avg_response_time_ms", 0) / 5000, 1.0)
        
        score = (
            metrics.get("avg_gpu_utilization", 0) * weights["avg_gpu_utilization"] +
            metrics.get("avg_memory_usage", 0) * weights["avg_memory_usage"] +
            queue_norm * weights["request_queue_depth_norm"] +
            latency_norm * weights["avg_response_time_norm"]
        )
        return score
    
    def _scale_up(self) -> str:
        # 扩容幅度基于负载严重程度
        self.current_nodes = min(self.current_nodes + 2, self.max_nodes)
        return f"scaled_up to {self.current_nodes}"
    
    def _scale_down(self) -> str:
        self.current_nodes = max(self.current_nodes - 1, self.min_nodes)
        return f"scaled_down to {self.current_nodes}"