Upload core/aggregation_protocol/scheduler.py with huggingface_hub

core/aggregation_protocol/scheduler.py

@@ -0,0 +1,328 @@
+#!/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),
+            }