Agent强化学习的最佳实践：并行任务处理与性能优化

┌─────────────┐
│ Task Queue  │  ← 用户请求进入队列
└──────┬──────┘
       │
       ▼
┌─────────────────────────────┐
│   Task Scheduler            │  ← 智能调度器
│   - Priority Management     │
│   - Load Balancing          │
│   - Dependency Resolution   │
└──────┬──────────────────────┘
       │
       ▼
┌──────────────────────────────────────┐
│     Agent Worker Pool                │
│  ┌─────┐ ┌─────┐ ┌─────┐ ┌─────┐   │
│  │ W1  │ │ W2  │ │ W3  │ │ W4  │   │ ← 并发Worker
│  └─────┘ └─────┘ └─────┘ └─────┘   │
└──────┬───────────────────────────────┘
       │
       ▼
┌─────────────────────────────┐
│   Shared Context Store      │  ← 共享上下文和学习经验
│   - Vector Database         │
│   - Experience Replay Buffer│
└─────────────────────────────┘

import asyncio
import time
from typing import List, Dict, Any
from dataclasses import dataclass
from enum import Enum

class TaskStatus(Enum):
    PENDING = "pending"
    RUNNING = "running"
    COMPLETED = "completed"
    FAILED = "failed"

@dataclass
class AgentTask:
    """Agent任务定义"""
    task_id: str
    prompt: str
    priority: int = 0
    context: Dict[str, Any] = None
    status: TaskStatus = TaskStatus.PENDING
    result: Any = None

class ParallelAgent:
    """支持并行任务处理的Agent实现"""

    def __init__(self, max_concurrent_tasks: int = 5):
        self.max_concurrent_tasks = max_concurrent_tasks
        self.task_queue = asyncio.Queue()
        self.active_tasks = set()
        self.results = {}

        # 模拟的经验回放缓冲区（用于强化学习）
        self.experience_buffer = []

    async def process_task(self, task: AgentTask) -> Any:
        """
        处理单个任务的核心逻辑

        在真实场景中，这里会调用LLM API、执行工具、搜索知识库等
        """
        task.status = TaskStatus.RUNNING

        try:
            # 模拟LLM推理延迟（实际应该是await llm_client.generate()）
            await asyncio.sleep(1.0)

            # 模拟任务处理
            result = f"Processed: {task.prompt}"

            # 记录经验用于后续学习
            self._record_experience(task, result, reward=1.0)

            task.status = TaskStatus.COMPLETED
            task.result = result
            return result

        except Exception as e:
            task.status = TaskStatus.FAILED
            task.result = str(e)
            self._record_experience(task, None, reward=-1.0)
            raise

    def _record_experience(self, task: AgentTask, result: Any, reward: float):
        """记录任务执行经验用于强化学习"""
        experience = {
            'state': task.context,
            'action': task.prompt,
            'reward': reward,
            'next_state': result,
            'timestamp': time.time()
        }
        self.experience_buffer.append(experience)

        # 保持缓冲区大小
        if len(self.experience_buffer) > 10000:
            self.experience_buffer.pop(0)

    async def worker(self):
        """工作协程，持续从队列获取任务并处理"""
        while True:
            task = await self.task_queue.get()

            if task is None:  # 停止信号
                break

            try:
                await self.process_task(task)
                self.results[task.task_id] = task.result
            except Exception as e:
                print(f"Task {task.task_id} failed: {e}")
            finally:
                self.active_tasks.discard(task.task_id)
                self.task_queue.task_done()

    async def submit_task(self, task: AgentTask):
        """提交任务到队列"""
        self.active_tasks.add(task.task_id)
        await self.task_queue.put(task)

    async def run(self, tasks: List[AgentTask]) -> Dict[str, Any]:
        """
        运行Agent处理所有任务

        Args:
            tasks: 待处理的任务列表

        Returns:
            任务ID到结果的映射
        """
        # 创建worker池
        workers = [
            asyncio.create_task(self.worker())
            for _ in range(self.max_concurrent_tasks)
        ]

        # 提交所有任务
        for task in tasks:
            await self.submit_task(task)

        # 等待所有任务完成
        await self.task_queue.join()

        # 停止workers
        for _ in workers:
            await self.task_queue.put(None)

        await asyncio.gather(*workers)

        return self.results

# 使用示例
async def main():
    # 创建并行Agent
    agent = ParallelAgent(max_concurrent_tasks=5)

    # 准备10个任务
    tasks = [
        AgentTask(
            task_id=f"task_{i}",
            prompt=f"Analyze code file {i}",
            context={"file_id": i}
        )
        for i in range(10)
    ]

    # 执行并计时
    start_time = time.time()
    results = await agent.run(tasks)
    elapsed_time = time.time() - start_time

    print(f"Processed {len(results)} tasks in {elapsed_time:.2f} seconds")
    print(f"Average time per task: {elapsed_time/len(results):.2f} seconds")
    print(f"Experience buffer size: {len(agent.experience_buffer)}")

# 运行
# asyncio.run(main())

# generated by AI

import heapq
from typing import Optional, Set
from collections import defaultdict

class TaskScheduler:
    """
    智能任务调度器
    - 支持优先级调度
    - 处理任务依赖关系
    - 动态负载均衡
    """

    def __init__(self):
        self.priority_queue = []  # 优先级队列（最小堆）
        self.task_counter = 0  # 用于打破优先级相同时的tie

        # 任务依赖图
        self.dependencies = defaultdict(set)  # task_id -> set of dependencies
        self.dependents = defaultdict(set)    # task_id -> set of dependents

        # 已完成任务集合
        self.completed_tasks = set()

    def add_task(self, task: AgentTask, depends_on: Optional[List[str]] = None):
        """
        添加任务到调度器

        Args:
            task: Agent任务
            depends_on: 依赖的任务ID列表
        """
        if depends_on:
            self.dependencies[task.task_id] = set(depends_on)
            for dep_id in depends_on:
                self.dependents[dep_id].add(task.task_id)

        # 如果没有未完成的依赖，立即加入优先级队列
        if self._can_schedule(task.task_id):
            self._enqueue(task)

    def _can_schedule(self, task_id: str) -> bool:
        """检查任务是否可以调度（所有依赖都已完成）"""
        deps = self.dependencies.get(task_id, set())
        return deps.issubset(self.completed_tasks)

    def _enqueue(self, task: AgentTask):
        """将任务加入优先级队列"""
        # 使用负优先级实现最大堆（Python heapq是最小堆）
        # task_counter确保相同优先级的任务按FIFO顺序
        heapq.heappush(
            self.priority_queue,
            (-task.priority, self.task_counter, task)
        )
        self.task_counter += 1

    def get_next_task(self) -> Optional[AgentTask]:
        """获取下一个应该执行的任务"""
        if not self.priority_queue:
            return None

        _, _, task = heapq.heappop(self.priority_queue)
        return task

    def mark_completed(self, task_id: str):
        """
        标记任务完成，并检查是否可以解锁依赖任务

        这是强化学习中的关键步骤：完成一个任务后，
        评估哪些后续任务可以被触发
        """
        self.completed_tasks.add(task_id)

        # 查找所有依赖此任务的任务
        newly_ready = []
        for dependent_id in self.dependents.get(task_id, set()):
            if self._can_schedule(dependent_id):
                newly_ready.append(dependent_id)

        return newly_ready

    def get_metrics(self) -> Dict[str, Any]:
        """获取调度器性能指标"""
        return {
            'pending_tasks': len(self.priority_queue),
            'completed_tasks': len(self.completed_tasks),
            'dependency_chains': len(self.dependencies)
        }

# 使用示例：处理有依赖关系的任务
async def example_with_dependencies():
    scheduler = TaskScheduler()
    agent = ParallelAgent(max_concurrent_tasks=3)

    # 任务DAG：
    # task_1 → task_3 → task_5
    # task_2 → task_4 → task_5

    tasks = {
        'task_1': AgentTask('task_1', 'Read file A', priority=10),
        'task_2': AgentTask('task_2', 'Read file B', priority=10),
        'task_3': AgentTask('task_3', 'Analyze A', priority=5),
        'task_4': AgentTask('task_4', 'Analyze B', priority=5),
        'task_5': AgentTask('task_5', 'Merge results', priority=1),
    }

    # 添加任务及其依赖关系
    scheduler.add_task(tasks['task_1'])
    scheduler.add_task(tasks['task_2'])
    scheduler.add_task(tasks['task_3'], depends_on=['task_1'])
    scheduler.add_task(tasks['task_4'], depends_on=['task_2'])
    scheduler.add_task(tasks['task_5'], depends_on=['task_3', 'task_4'])

    # 执行任务（带依赖解析）
    while True:
        task = scheduler.get_next_task()
        if task is None:
            break

        await agent.process_task(task)

        # 标记完成并解锁依赖任务
        newly_ready_ids = scheduler.mark_completed(task.task_id)
        for task_id in newly_ready_ids:
            scheduler._enqueue(tasks[task_id])

    print(scheduler.get_metrics())

# generated by AI

import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
from collections import deque
import random

class SchedulerStateEncoder:
    """将调度器状态编码为向量"""

    @staticmethod
    def encode(scheduler: TaskScheduler, system_metrics: Dict) -> np.ndarray:
        """
        编码当前调度器状态

        返回特征向量：
        - 队列长度
        - 平均任务优先级
        - 系统CPU/内存使用率
        - 最近10个任务的平均完成时间
        - 待处理依赖关系数量
        """
        metrics = scheduler.get_metrics()

        features = [
            metrics['pending_tasks'] / 100.0,  # 归一化
            metrics['completed_tasks'] / 1000.0,
            metrics['dependency_chains'] / 50.0,
            system_metrics.get('cpu_usage', 0) / 100.0,
            system_metrics.get('memory_usage', 0) / 100.0,
            system_metrics.get('avg_completion_time', 1.0) / 10.0,
        ]

        return np.array(features, dtype=np.float32)

class DQNSchedulerPolicy(nn.Module):
    """
    使用DQN学习任务调度策略

    输入：调度器状态向量
    输出：每个候选任务的Q值
    """

    def __init__(self, state_dim: int, action_dim: int, hidden_dim: int = 128):
        super().__init__()

        self.network = nn.Sequential(
            nn.Linear(state_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, action_dim)
        )

    def forward(self, state: torch.Tensor) -> torch.Tensor:
        """
        前向传播：给定状态，输出所有动作的Q值

        Args:
            state: 状态向量 [batch_size, state_dim]

        Returns:
            Q值向量 [batch_size, action_dim]
        """
        return self.network(state)

class RLScheduler:
    """
    基于强化学习的自适应任务调度器

    通过在线学习不断优化调度策略，降低平均任务完成时间
    """

    def __init__(
        self,
        state_dim: int = 6,
        max_tasks: int = 10,
        learning_rate: float = 0.001,
        gamma: float = 0.99
    ):
        self.state_dim = state_dim
        self.action_dim = max_tasks

        # Q网络（主网络和目标网络）
        self.q_network = DQNSchedulerPolicy(state_dim, max_tasks)
        self.target_network = DQNSchedulerPolicy(state_dim, max_tasks)
        self.target_network.load_state_dict(self.q_network.state_dict())

        self.optimizer = optim.Adam(self.q_network.parameters(), lr=learning_rate)
        self.gamma = gamma

        # 经验回放缓冲区
        self.replay_buffer = deque(maxlen=10000)
        self.batch_size = 64

        # ε-greedy探索策略
        self.epsilon = 1.0
        self.epsilon_decay = 0.995
        self.epsilon_min = 0.01

    def select_action(
        self,
        state: np.ndarray,
        available_tasks: List[AgentTask]
    ) -> int:
        """
        根据当前状态选择任务（动作）

        使用ε-greedy策略平衡探索与利用
        """
        if random.random() < self.epsilon:
            # 探索：随机选择
            return random.randint(0, len(available_tasks) - 1)
        else:
            # 利用：选择Q值最高的动作
            with torch.no_grad():
                state_tensor = torch.FloatTensor(state).unsqueeze(0)
                q_values = self.q_network(state_tensor)[0]

                # 只考虑可用任务的Q值
                valid_q_values = q_values[:len(available_tasks)]
                return torch.argmax(valid_q_values).item()

    def store_experience(
        self,
        state: np.ndarray,
        action: int,
        reward: float,
        next_state: np.ndarray,
        done: bool
    ):
        """存储经验到回放缓冲区"""
        self.replay_buffer.append((state, action, reward, next_state, done))

    def train_step(self):
        """执行一次训练步骤"""
        if len(self.replay_buffer) < self.batch_size:
            return

        # 从经验回放缓冲区采样
        batch = random.sample(self.replay_buffer, self.batch_size)
        states, actions, rewards, next_states, dones = zip(*batch)

        states = torch.FloatTensor(np.array(states))
        actions = torch.LongTensor(actions)
        rewards = torch.FloatTensor(rewards)
        next_states = torch.FloatTensor(np.array(next_states))
        dones = torch.FloatTensor(dones)

        # 计算当前Q值
        current_q_values = self.q_network(states).gather(1, actions.unsqueeze(1))

        # 计算目标Q值（使用目标网络）
        with torch.no_grad():
            next_q_values = self.target_network(next_states).max(1)[0]
            target_q_values = rewards + (1 - dones) * self.gamma * next_q_values

        # 计算损失并反向传播
        loss = nn.MSELoss()(current_q_values.squeeze(), target_q_values)

        self.optimizer.zero_grad()
        loss.backward()
        self.optimizer.step()

        # 衰减探索率
        self.epsilon = max(self.epsilon_min, self.epsilon * self.epsilon_decay)

        return loss.item()

    def update_target_network(self):
        """定期更新目标网络"""
        self.target_network.load_state_dict(self.q_network.state_dict())

# 集成到Agent系统
class RLEnhancedAgent(ParallelAgent):
    """集成强化学习调度策略的Agent"""

    def __init__(self, max_concurrent_tasks: int = 5):
        super().__init__(max_concurrent_tasks)
        self.rl_scheduler = RLScheduler()
        self.state_encoder = SchedulerStateEncoder()

        # 性能指标追踪
        self.episode_rewards = []
        self.episode_completion_times = []

    async def run_with_rl(self, tasks: List[AgentTask]) -> Dict[str, Any]:
        """
        使用强化学习调度策略运行Agent

        在执行过程中在线学习最优调度策略
        """
        pending_tasks = tasks.copy()
        completed_count = 0
        episode_start_time = time.time()

        while pending_tasks:
            # 编码当前状态
            system_metrics = {
                'cpu_usage': 50.0,  # 实际应从系统获取
                'memory_usage': 60.0,
                'avg_completion_time': 2.0
            }
            state = self.state_encoder.encode(self.scheduler, system_metrics)

            # 使用RL策略选择下一个任务
            action = self.rl_scheduler.select_action(state, pending_tasks)
            selected_task = pending_tasks[action]

            # 执行任务
            task_start_time = time.time()
            await self.process_task(selected_task)
            task_completion_time = time.time() - task_start_time

            # 计算奖励（负完成时间，鼓励快速完成）
            reward = -task_completion_time

            # 移除已完成任务
            pending_tasks.pop(action)
            completed_count += 1

            # 编码下一个状态
            next_system_metrics = system_metrics.copy()
            next_system_metrics['avg_completion_time'] = task_completion_time
            next_state = self.state_encoder.encode(self.scheduler, next_system_metrics)

            # 存储经验
            done = len(pending_tasks) == 0
            self.rl_scheduler.store_experience(state, action, reward, next_state, done)

            # 训练RL模型
            if completed_count % 10 == 0:
                loss = self.rl_scheduler.train_step()
                if loss is not None:
                    print(f"Training loss: {loss:.4f}, epsilon: {self.rl_scheduler.epsilon:.3f}")

            # 定期更新目标网络
            if completed_count % 100 == 0:
                self.rl_scheduler.update_target_network()

        episode_time = time.time() - episode_start_time
        self.episode_completion_times.append(episode_time)

        return {
            'completed_tasks': completed_count,
            'total_time': episode_time,
            'avg_time_per_task': episode_time / completed_count
        }

# generated by AI

import ray
from ray import serve
from typing import List

# 初始化Ray集群
ray.init(address='auto')  # 连接到已有集群，或者启动本地集群

@ray.remote
class DistributedAgentWorker:
    """
    分布式Agent Worker

    每个Worker运行在独立的进程/机器上，
    可以处理CPU密集型或I/O密集型任务
    """

    def __init__(self, worker_id: int):
        self.worker_id = worker_id
        # 每个worker有独立的模型实例（避免序列化开销）
        self.local_model = self._load_model()

    def _load_model(self):
        """加载本地模型（可以是LLM、工具等）"""
        # 实际应用中，这里会加载真实的模型
        return {"model": "gpt-4", "worker_id": self.worker_id}

    def process_task(self, task: AgentTask) -> Dict[str, Any]:
        """
        处理单个任务

        注意：这是同步方法，Ray会自动处理并发
        """
        import time
        time.sleep(1)  # 模拟处理时间

        return {
            'task_id': task.task_id,
            'result': f"Processed by worker {self.worker_id}",
            'worker_id': self.worker_id
        }

class DistributedAgentCluster:
    """分布式Agent集群管理器"""

    def __init__(self, num_workers: int = 10):
        # 创建分布式worker池
        self.workers = [
            DistributedAgentWorker.remote(worker_id=i)
            for i in range(num_workers)
        ]

    def process_tasks_batch(self, tasks: List[AgentTask]) -> List[Dict[str, Any]]:
        """
        批量处理任务，自动分配到不同workers

        Ray会自动：
        1. 负载均衡到不同机器
        2. 容错处理（worker崩溃时重试）
        3. 结果收集和聚合
        """
        # 将任务分配给workers（轮询策略）
        futures = []
        for i, task in enumerate(tasks):
            worker = self.workers[i % len(self.workers)]
            future = worker.process_task.remote(task)
            futures.append(future)

        # 等待所有任务完成（并行执行）
        results = ray.get(futures)
        return results

    def process_tasks_dynamic(self, tasks: List[AgentTask]) -> List[Dict[str, Any]]:
        """
        动态任务分配：worker完成任务后立即分配新任务

        这种方式在任务耗时差异大时效率更高
        """
        results = []
        remaining_tasks = tasks.copy()

        # 初始分配：每个worker一个任务
        active_futures = {}
        for worker in self.workers:
            if not remaining_tasks:
                break
            task = remaining_tasks.pop(0)
            future = worker.process_task.remote(task)
            active_futures[future] = worker

        # 动态调度：完成一个，分配一个
        while active_futures:
            # 等待任何一个任务完成
            done_futures, _ = ray.wait(list(active_futures.keys()), num_returns=1)

            for future in done_futures:
                result = ray.get(future)
                results.append(result)

                worker = active_futures.pop(future)

                # 如果还有待处理任务，分配给刚完成的worker
                if remaining_tasks:
                    next_task = remaining_tasks.pop(0)
                    new_future = worker.process_task.remote(next_task)
                    active_futures[new_future] = worker

        return results

# 使用示例
def run_distributed_example():
    # 创建包含10个worker的集群
    cluster = DistributedAgentCluster(num_workers=10)

    # 准备100个任务
    tasks = [
        AgentTask(task_id=f"task_{i}", prompt=f"Task {i}")
        for i in range(100)
    ]

    # 动态调度执行
    import time
    start_time = time.time()
    results = cluster.process_tasks_dynamic(tasks)
    elapsed_time = time.time() - start_time

    print(f"Processed {len(results)} tasks in {elapsed_time:.2f}s")
    print(f"Throughput: {len(results)/elapsed_time:.1f} tasks/sec")

# ray.shutdown()

# generated by AI

# 不佳：串行调用10次
for task in tasks:
    result = await llm_client.generate(task.prompt)

# 优化：批量调用1次
prompts = [task.prompt for task in tasks]
results = await llm_client.generate_batch(prompts, batch_size=10)

from functools import lru_cache
import hashlib

class CachedAgent:
    def __init__(self):
        self.cache = {}

    def _compute_cache_key(self, prompt: str, context: Dict) -> str:
        """计算缓存键（考虑prompt和context）"""
        cache_input = f"{prompt}:{sorted(context.items())}"
        return hashlib.md5(cache_input.encode()).hexdigest()

    async def process_with_cache(self, task: AgentTask) -> Any:
        cache_key = self._compute_cache_key(task.prompt, task.context or {})

        if cache_key in self.cache:
            return self.cache[cache_key]  # 缓存命中

        result = await self.process_task(task)
        self.cache[cache_key] = result
        return result

# generated by AI

async def speculative_execution(task: AgentTask):
    # 主任务：LLM推理
    llm_future = asyncio.create_task(llm_client.generate(task.prompt))

    # 推测任务：预加载可能需要的文件
    if 'file_path' in task.context:
        file_future = asyncio.create_task(load_file(task.context['file_path']))

    # 等待主任务
    result = await llm_future

    # 如果推测正确，文件已经加载好了
    if 'load_file' in result.actions:
        file_content = await file_future  # 立即可用
        return process_with_file(result, file_content)

    return result

# generated by AI

class AdaptiveModelSelector:
    """根据任务特征自动选择最合适的模型"""

    def select_model(self, task: AgentTask) -> str:
        # 基于任务复杂度、延迟要求、成本预算选择模型
        if task.context.get('urgent', False):
            return 'gpt-3.5-turbo'  # 低延迟
        elif task.context.get('complexity') == 'high':
            return 'gpt-4'  # 高质量
        else:
            return 'local-model'  # 低成本

# generated by AI

import prometheus_client
from prometheus_client import Counter, Histogram, Gauge

# 定义Prometheus指标
task_counter = Counter('agent_tasks_total', 'Total tasks processed', ['status'])
task_duration = Histogram('agent_task_duration_seconds', 'Task processing duration')
queue_depth = Gauge('agent_queue_depth', 'Number of tasks in queue')
concurrency = Gauge('agent_active_workers', 'Number of active workers')

class MonitoredAgent(ParallelAgent):
    """带监控的Agent实现"""

    async def process_task(self, task: AgentTask) -> Any:
        queue_depth.dec()  # 从队列移除
        concurrency.inc()  # 开始处理

        start_time = time.time()
        try:
            result = await super().process_task(task)
            task_counter.labels(status='success').inc()
            return result
        except Exception as e:
            task_counter.labels(status='failed').inc()
            raise
        finally:
            duration = time.time() - start_time
            task_duration.observe(duration)
            concurrency.dec()

    async def submit_task(self, task: AgentTask):
        queue_depth.inc()
        await super().submit_task(task)

# 启动Prometheus HTTP服务器
prometheus_client.start_http_server(8000)

# generated by AI