core/aggregation_protocol/scheduler.py

#!/usr/bin/env python3
"""
虫群聚合协议 — 任务调度器

核心功能：
- 接收任务请求
- 组建临时服务器(TaskForce)
- 分配任务到各节点
- 聚合各节点结果
- 任务完成后解散临时服务器

类比GPU集群：
- 任务调度器 = SLURM/调度系统
- 临时服务器 = 临时分配的GPU组
- 节点 = 单块GPU
"""

import hashlib
import logging
import threading
import time
from collections import deque
from datetime import datetime
from typing import Callable, Dict, List, Optional

from .types import (
    AggregationStrategy, AggregationTask, TaskForce, TaskForceStatus,
    NodeInfo, NodeStatus, ProtocolMessage,
)
from .discovery import NodeRegistry
from .transport import MessageBus

logger = logging.getLogger(__name__)


class TaskForceManager:
    """
    临时服务器管理器
    
    生命周期：创建 → 组建 → 运行 → 完成 → 解散
    """

    def __init__(self, node_registry: NodeRegistry, message_bus: MessageBus):
        self.registry = node_registry
        self.bus = message_bus
        
        # 活跃的临时服务器 taskforce_id -> TaskForce
        self._taskforces: Dict[str, TaskForce] = {}
        self._lock = threading.RLock()
        
        # 任务队列
        self._task_queue = deque()
        
        # 统计
        self._stats = {
            "taskforces_created": 0,
            "taskforces_completed": 0,
            "taskforces_failed": 0,
            "tasks_processed": 0,
        }

    # ============================================================
    # 临时服务器生命周期
    # ============================================================

    def create_taskforce(self, task: AggregationTask) -> Optional[TaskForce]:
        """
        为任务创建临时服务器
        
        类似GPU集群分配资源：
        1. 分析任务需要的能力
        2. 从可用节点中选择
        3. 组建临时服务器
        4. 通知成员节点
        """
        # 1. 发现合适的节点
        candidates = self.registry.discover_for_task(
            required_caps=task.required_capabilities,
            min_nodes=task.min_nodes,
            max_nodes=task.max_nodes,
            exclude=[task.requester],
        )
        
        if len(candidates) < task.min_nodes:
            logger.warning(f"节点不足: 需要{task.min_nodes}，找到{len(candidates)}")
            return None
        
        # 2. 创建临时服务器
        tf_id = self._gen_id("tf")
        tf = TaskForce(
            taskforce_id=tf_id,
            name=f"TaskForce-{tf_id}",
            coordinator=task.requester,
            strategy=task.strategy,
            task_description=task.query,
        )
        
        # 3. 添加成员
        for node in candidates:
            tf.add_member(node.node_id)
            # 更新节点的当前任务列表
            node.current_taskforces.append(tf_id)
        
        # 4. 记录
        with self._lock:
            self._taskforces[tf_id] = tf
            task.taskforce_id = tf_id
            task.status = "assigned"
            self._stats["taskforces_created"] += 1
        
        # 5. 通知成员（通过消息总线）
        self.bus.broadcast("join_taskforce", {
            "taskforce_id": tf_id,
            "coordinator": task.requester,
            "task": task.query,
            "members": tf.members,
            "strategy": tf.strategy.value,
        })
        
        logger.info(f"临时服务器创建: {tf_id}, 成员: {tf.members}")
        return tf

    def complete_taskforce(self, tf_id: str, result: Dict):
        """完成任务，解散临时服务器"""
        with self._lock:
            tf = self._taskforces.get(tf_id)
            if not tf:
                return
            
            tf.status = TaskForceStatus.COMPLETED
            tf.completed_at = datetime.now()
            tf.results = result
            
            # 清理成员节点的任务列表
            for member_id in tf.members:
                node = self.registry.get_node(member_id)
                if node and tf_id in node.current_taskforces:
                    node.current_taskforces.remove(tf_id)
            
            self._stats["taskforces_completed"] += 1
        
        # 通知解散
        self.bus.broadcast("leave_taskforce", {
            "taskforce_id": tf_id,
            "status": "completed",
        })
        
        logger.info(f"临时服务器解散: {tf_id}")

    def fail_taskforce(self, tf_id: str, reason: str = ""):
        """临时服务器失败"""
        with self._lock:
            tf = self._taskforces.get(tf_id)
            if not tf:
                return
            
            tf.status = TaskForceStatus.FAILED
            tf.completed_at = datetime.now()
            
            for member_id in tf.members:
                node = self.registry.get_node(member_id)
                if node and tf_id in node.current_taskforces:
                    node.current_taskforces.remove(tf_id)
            
            self._stats["taskforces_failed"] += 1

    # ============================================================
    # 任务调度
    # ============================================================

    def submit_task(self, task: AggregationTask) -> Optional[str]:
        """提交任务"""
        tf = self.create_taskforce(task)
        if tf:
            return tf.taskforce_id
        return None

    def get_taskforce(self, tf_id: str) -> Optional[TaskForce]:
        with self._lock:
            return self._taskforces.get(tf_id)

    def get_active_taskforces(self) -> List[TaskForce]:
        with self._lock:
            return [tf for tf in self._taskforces.values() 
                    if tf.status == TaskForceStatus.ACTIVE]

    # ============================================================
    # 结果聚合
    # ============================================================

    def aggregate_results(self, tf_id: str, 
                          node_results: Dict[str, Dict]) -> Dict:
        """
        聚合各节点结果
        
        策略：
        - PARAMETER_AVERAGE: 参数平均（联邦学习风格）
        - ENSEMBLE_VOTE: 投票法（多数同意）
        - SEQUENTIAL_REFINE: 顺序精炼（每个节点改进上一个的结果）
        - ADAPTIVE_MIX: 自适应混合（按置信度加权）
        """
        tf = self.get_taskforce(tf_id)
        if not tf:
            return {"error": "临时服务器不存在"}
        
        strategy = tf.strategy
        results = {k: v for k, v in node_results.items() if v}
        
        if not results:
            return {"error": "无有效结果"}
        
        if strategy == AggregationStrategy.ENSEMBLE_VOTE:
            return self._vote_aggregate(results)
        elif strategy == AggregationStrategy.SEQUENTIAL_REFINE:
            return self._sequential_aggregate(results)
        elif strategy == AggregationStrategy.PARAMETER_AVERAGE:
            return self._parameter_average(results)
        else:  # ADAPTIVE_MIX
            return self._adaptive_mix(results)

    def _vote_aggregate(self, results: Dict[str, Dict]) -> Dict:
        """投票聚合 — 选择出现最多的回答"""
        from collections import Counter
        
        responses = []
        for node_id, result in results.items():
            resp = result.get("response", "")
            if resp:
                responses.append(resp)
        
        if not responses:
            return {"response": "", "confidence": 0.0}
        
        # 简单投票：选最长的回答（通常是信息最丰富的）
        counter = Counter(responses)
        if counter:
            best = counter.most_common(1)[0][0]
            confidence = counter.most_common(1)[0][1] / len(responses)
            return {"response": best, "confidence": confidence, "method": "vote"}
        
        return {"response": responses[0], "confidence": 0.5, "method": "vote"}

    def _sequential_aggregate(self, results: Dict[str, Dict]) -> Dict:
        """顺序精炼 — 每个节点改进上一个结果"""
        refined = ""
        confidence = 0.0
        
        for node_id, result in results.items():
            if refined:
                # 将前一个结果作为上下文传入
                refined = result.get("response", refined)
            else:
                refined = result.get("response", "")
            confidence = max(confidence, result.get("confidence", 0.0))
        
        return {"response": refined, "confidence": confidence, "method": "sequential"}

    def _parameter_average(self, results: Dict[str, Dict]) -> Dict:
        """参数平均 — 联邦学习风格"""
        # 对置信度做加权平均
        total_weight = 0.0
        weighted_confidence = 0.0
        best_response = ""
        best_conf = 0.0
        
        for node_id, result in results.items():
            conf = result.get("confidence", 0.5)
            total_weight += conf
            weighted_confidence += conf * conf
            
            if conf > best_conf:
                best_conf = conf
                best_response = result.get("response", "")
        
        avg_confidence = weighted_confidence / max(total_weight, 0.01)
        return {
            "response": best_response,
            "confidence": avg_confidence,
            "method": "parameter_average",
            "contributing_nodes": len(results),
        }

    def _adaptive_mix(self, results: Dict[str, Dict]) -> Dict:
        """自适应混合 — 按置信度和专长加权"""
        total_score = 0.0
        best_response = ""
        best_score = 0.0
        all_responses = []
        
        for node_id, result in results.items():
            conf = result.get("confidence", 0.5)
            # 考虑节点的专长匹配度
            node = self.registry.get_node(node_id)
            expertise_bonus = 0.0
            if node:
                expertise_bonus = node.capability.compute_score * 0.1
            
            score = conf + expertise_bonus
            total_score += score
            all_responses.append(result.get("response", ""))
            
            if score > best_score:
                best_score = score
                best_response = result.get("response", "")
        
        avg_confidence = total_score / max(len(results), 1)
        return {
            "response": best_response,
            "confidence": min(avg_confidence, 1.0),
            "method": "adaptive_mix",
            "contributing_nodes": len(results),
            "all_responses": all_responses[:3],  # 保留前3个备选
        }

    # ============================================================
    # 工具方法
    # ============================================================

    def _gen_id(self, prefix: str) -> str:
        return f"{prefix}_{hashlib.md5(f'{time.time()}{prefix}'.encode()).hexdigest()[:8]}"

    def get_stats(self) -> Dict:
        with self._lock:
            return {
                **self._stats,
                "active_taskforces": len(self.get_active_taskforces()),
                "total_taskforces": len(self._taskforces),
            }