nexaroa 0.0.111__py3-none-any.whl

neuroshard/core/swarm/router.py

@@ -0,0 +1,658 @@
+"""
+Swarm Router - Fault-Tolerant Multipath Routing for NeuroShard
+
+Implements the "Swarm" network layer with:
+- K-candidate routing (multiple peers per layer range)
+- Automatic failover within 200ms
+- Weighted scoring (latency + queue_depth)
+- Probabilistic send with retry
+
+Key Directive: "If Node A hangs, the packet must automatically flow to Node B 
+without crashing the run."
+
+Reference: SWARM Parallelism papers for randomized pipeline construction.
+"""
+
+import asyncio
+import grpc
+import logging
+import random
+import time
+import torch
+from dataclasses import dataclass, field
+from typing import Dict, List, Optional, Tuple, Any
+from collections import defaultdict
+
+# Try importing proto definitions
+try:
+    from protos import neuroshard_pb2
+    from protos import neuroshard_pb2_grpc
+    GRPC_AVAILABLE = True
+except ImportError:
+    GRPC_AVAILABLE = False
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class PeerCandidate:
+    """
+    A candidate peer for routing.
+    
+    Candidates are scored by a weighted combination of latency and queue depth.
+    Lower score = better candidate.
+    """
+    node_id: str
+    grpc_addr: str
+    layer_range: Tuple[int, int]  # (start_layer, end_layer) exclusive
+    
+    # Performance metrics (rolling averages)
+    latency_ms: float = 100.0             # Default 100ms
+    queue_depth: int = 0                   # Current work queue size
+    
+    # Health tracking
+    last_heartbeat: float = field(default_factory=time.time)
+    consecutive_failures: int = 0
+    total_requests: int = 0
+    successful_requests: int = 0
+    
+    # Capacity info from heartbeat
+    available_memory_mb: int = 0
+    gpu_utilization: float = 0.0
+    is_accepting_activations: bool = True
+    
+    def score(self, weights: Optional[Dict[str, float]] = None) -> float:
+        """
+        Compute weighted score for peer selection.
+        
+        Lower score = better candidate.
+        
+        Default weights prioritize queue depth (60%) over latency (40%)
+        because queue depth is a better indicator of immediate availability.
+        """
+        w = weights or {"latency": 0.4, "queue": 0.6}
+        
+        # Normalize: latency 0-500ms maps to 0-1, queue 0-100 maps to 0-1
+        latency_norm = min(1.0, self.latency_ms / 500.0)
+        queue_norm = min(1.0, self.queue_depth / 100.0)
+        
+        base_score = w["latency"] * latency_norm + w["queue"] * queue_norm
+        
+        # Penalty for recent failures (exponential backoff)
+        if self.consecutive_failures > 0:
+            failure_penalty = min(1.0, 0.1 * (2 ** self.consecutive_failures))
+            base_score += failure_penalty
+            
+        # Penalty for stale heartbeat (>15s)
+        staleness = time.time() - self.last_heartbeat
+        if staleness > 15.0:
+            stale_penalty = min(1.0, staleness / 60.0)
+            base_score += stale_penalty
+            
+        return base_score
+    
+    @property
+    def success_rate(self) -> float:
+        """Success rate of requests to this peer."""
+        if self.total_requests == 0:
+            return 1.0  # Assume good until proven otherwise
+        return self.successful_requests / self.total_requests
+    
+    def covers_layer(self, layer: int) -> bool:
+        """Check if this peer handles the given layer."""
+        return self.layer_range[0] <= layer < self.layer_range[1]
+    
+    def update_latency(self, latency_ms: float, alpha: float = 0.3):
+        """Update latency with exponential moving average."""
+        self.latency_ms = alpha * latency_ms + (1 - alpha) * self.latency_ms
+        
+    def record_success(self, latency_ms: float):
+        """Record a successful request."""
+        self.total_requests += 1
+        self.successful_requests += 1
+        self.consecutive_failures = 0
+        self.update_latency(latency_ms)
+        
+    def record_failure(self):
+        """Record a failed request."""
+        self.total_requests += 1
+        self.consecutive_failures += 1
+
+
+@dataclass
+class RoutingResult:
+    """Result of a routing operation."""
+    success: bool
+    peer_used: Optional[str] = None
+    latency_ms: float = 0.0
+    attempts: int = 0
+    error: Optional[str] = None
+    response_data: Any = None
+
+
+class SwarmRouter:
+    """
+    Fault-tolerant multipath router for distributed inference/training.
+    
+    Key Features:
+    - Returns K candidates per layer range (not just one)
+    - Automatic failover: if primary doesn't ACK in 200ms, try secondary
+    - Weighted scoring: 0.4 * latency + 0.6 * queue_depth
+    - Probabilistic load balancing for equal-scored candidates
+    
+    Integration:
+    - Replaces single-hop routing in p2p.py
+    - Works with SwarmHeartbeatService for capacity awareness
+    - Uses DHT for peer discovery fallback
+    """
+    
+    ACK_TIMEOUT_MS = 200       # Failover if no ACK in 200ms
+    K_CANDIDATES = 3           # Return top-K candidates per layer range
+    CACHE_TTL_SECONDS = 30     # How long to cache peer info
+    
+    def __init__(
+        self, 
+        dht_protocol: Optional[Any] = None,
+        layer_pool: Optional[Any] = None,
+        tracker_url: Optional[str] = None,
+    ):
+        """
+        Initialize SwarmRouter.
+        
+        Args:
+            dht_protocol: DHT protocol for peer discovery
+            layer_pool: DynamicLayerPool for local layer info
+            tracker_url: Fallback tracker URL for peer discovery
+        """
+        self.dht = dht_protocol
+        self.layer_pool = layer_pool
+        self.tracker_url = tracker_url
+        
+        # Peer state cache
+        self.peer_stats: Dict[str, PeerCandidate] = {}
+        self._peer_lock = asyncio.Lock()
+        
+        # gRPC connection pool
+        self._channels: Dict[str, grpc.aio.Channel] = {}
+        self._stubs: Dict[str, Any] = {}
+        
+        # Routing metrics
+        self.total_routes = 0
+        self.successful_routes = 0
+        self.failover_count = 0
+        
+        # Background tasks
+        self._cleanup_task: Optional[asyncio.Task] = None
+    
+    # Property aliases for compatibility with SwarmServiceMixin
+    @property
+    def total_sends(self) -> int:
+        """Alias for total_routes."""
+        return self.total_routes
+    
+    @property
+    def successful_sends(self) -> int:
+        """Alias for successful_routes."""
+        return self.successful_routes
+        
+    async def start(self):
+        """Start background maintenance tasks."""
+        self._cleanup_task = asyncio.create_task(self._cleanup_loop())
+        logger.info("SwarmRouter started")
+        
+    async def stop(self):
+        """Stop router and cleanup."""
+        if self._cleanup_task:
+            self._cleanup_task.cancel()
+            try:
+                await self._cleanup_task
+            except asyncio.CancelledError:
+                pass
+                
+        # Close all gRPC channels
+        for channel in self._channels.values():
+            await channel.close()
+        self._channels.clear()
+        self._stubs.clear()
+        
+        logger.info("SwarmRouter stopped")
+        
+    async def _cleanup_loop(self):
+        """Periodically cleanup stale peers and connections."""
+        while True:
+            try:
+                await asyncio.sleep(60)  # Every minute
+                await self._cleanup_stale_peers()
+            except asyncio.CancelledError:
+                break
+            except Exception as e:
+                logger.error(f"Cleanup error: {e}")
+                
+    async def _cleanup_stale_peers(self):
+        """Remove peers that haven't sent heartbeat in >60s."""
+        now = time.time()
+        stale_threshold = 60.0
+        
+        async with self._peer_lock:
+            stale_peers = [
+                node_id for node_id, peer in self.peer_stats.items()
+                if (now - peer.last_heartbeat) > stale_threshold
+            ]
+            
+            for node_id in stale_peers:
+                del self.peer_stats[node_id]
+                if node_id in self._channels:
+                    await self._channels[node_id].close()
+                    del self._channels[node_id]
+                if node_id in self._stubs:
+                    del self._stubs[node_id]
+                    
+            if stale_peers:
+                logger.info(f"Cleaned up {len(stale_peers)} stale peers")
+    
+    def register_peer(self, peer: PeerCandidate):
+        """
+        Register or update a peer from heartbeat.
+        
+        Called by SwarmHeartbeatService when receiving capacity broadcasts.
+        """
+        self.peer_stats[peer.node_id] = peer
+        
+    def update_peer_from_heartbeat(
+        self,
+        node_id: str,
+        grpc_addr: str,
+        layer_range: Tuple[int, int],
+        queue_depth: int,
+        available_memory_mb: int,
+        gpu_utilization: float,
+        is_accepting: bool,
+    ):
+        """Update peer info from heartbeat message."""
+        if node_id in self.peer_stats:
+            peer = self.peer_stats[node_id]
+            peer.queue_depth = queue_depth
+            peer.available_memory_mb = available_memory_mb
+            peer.gpu_utilization = gpu_utilization
+            peer.is_accepting_activations = is_accepting
+            peer.last_heartbeat = time.time()
+        else:
+            # New peer
+            self.peer_stats[node_id] = PeerCandidate(
+                node_id=node_id,
+                grpc_addr=grpc_addr,
+                layer_range=layer_range,
+                queue_depth=queue_depth,
+                available_memory_mb=available_memory_mb,
+                gpu_utilization=gpu_utilization,
+                is_accepting_activations=is_accepting,
+                last_heartbeat=time.time(),
+            )
+    
+    def get_candidates(self, target_layer: int) -> List[PeerCandidate]:
+        """
+        Get K candidates for a target layer, sorted by score.
+        
+        Returns candidates from:
+        1. Local cache (fastest) - peers we know from heartbeats
+        2. DHT lookup (if cache miss)
+        3. Tracker fallback (if DHT fails)
+        
+        Args:
+            target_layer: The layer index we need to route to
+            
+        Returns:
+            List of up to K candidates, sorted by score (best first)
+        """
+        candidates: List[PeerCandidate] = []
+        
+        # Strategy 1: Local cache (from heartbeats)
+        for node_id, info in self.peer_stats.items():
+            if info.covers_layer(target_layer) and info.is_accepting_activations:
+                candidates.append(info)
+        
+        # Strategy 2: DHT lookup (if not enough candidates)
+        if len(candidates) < self.K_CANDIDATES and self.dht:
+            dht_peers = self._dht_lookup_layer(target_layer)
+            for peer in dht_peers:
+                if peer.node_id not in {c.node_id for c in candidates}:
+                    candidates.append(peer)
+        
+        # Strategy 3: Tracker fallback (if still not enough)
+        if len(candidates) < self.K_CANDIDATES and self.tracker_url:
+            tracker_peers = self._tracker_lookup_layer(target_layer)
+            for peer in tracker_peers:
+                if peer.node_id not in {c.node_id for c in candidates}:
+                    candidates.append(peer)
+        
+        # Sort by score, return top K
+        candidates.sort(key=lambda c: c.score())
+        
+        # Add small random shuffle among similarly-scored candidates
+        # to distribute load
+        if len(candidates) > 1:
+            candidates = self._probabilistic_shuffle(candidates)
+            
+        return candidates[:self.K_CANDIDATES]
+    
+    def _probabilistic_shuffle(
+        self, 
+        candidates: List[PeerCandidate], 
+        epsilon: float = 0.1
+    ) -> List[PeerCandidate]:
+        """
+        Add small random shuffle among similarly-scored candidates.
+        
+        Candidates with scores within epsilon of each other may be swapped.
+        This provides load balancing without completely ignoring scores.
+        """
+        if len(candidates) < 2:
+            return candidates
+            
+        result = candidates.copy()
+        
+        for i in range(len(result) - 1):
+            # Check if next candidate has similar score
+            score_diff = abs(result[i].score() - result[i + 1].score())
+            if score_diff < epsilon:
+                # 30% chance to swap similar candidates
+                if random.random() < 0.3:
+                    result[i], result[i + 1] = result[i + 1], result[i]
+                    
+        return result
+    
+    def _dht_lookup_layer(self, target_layer: int) -> List[PeerCandidate]:
+        """Look up peers for a layer via DHT."""
+        if not self.dht:
+            return []
+            
+        try:
+            # DHT key format: "layer:{layer_num}"
+            key = f"layer:{target_layer}"
+            values = self.dht.lookup_values(key, k=self.K_CANDIDATES * 2)
+            
+            candidates = []
+            for value in values:
+                # Parse peer info from DHT value
+                try:
+                    import json
+                    info = json.loads(value)
+                    peer = PeerCandidate(
+                        node_id=info.get("node_id", ""),
+                        grpc_addr=info.get("grpc_addr", ""),
+                        layer_range=(info.get("start_layer", 0), info.get("end_layer", 0)),
+                    )
+                    candidates.append(peer)
+                except:
+                    continue
+                    
+            return candidates
+        except Exception as e:
+            logger.debug(f"DHT lookup failed: {e}")
+            return []
+            
+    def _tracker_lookup_layer(self, target_layer: int) -> List[PeerCandidate]:
+        """Look up peers for a layer via centralized tracker."""
+        if not self.tracker_url:
+            return []
+            
+        try:
+            import requests
+            resp = requests.get(
+                f"{self.tracker_url}/peers",
+                params={"layer": target_layer},
+                timeout=2.0
+            )
+            
+            if resp.status_code != 200:
+                return []
+                
+            candidates = []
+            for peer_info in resp.json().get("peers", []):
+                try:
+                    # Parse shard range (e.g., "0-12")
+                    shard_str = peer_info.get("shard_range", "0-0")
+                    parts = shard_str.split("-")
+                    layer_range = (int(parts[0]), int(parts[1]) if len(parts) > 1 else int(parts[0]) + 1)
+                    
+                    peer = PeerCandidate(
+                        node_id=peer_info.get("node_id", peer_info.get("url", "")),
+                        grpc_addr=peer_info.get("url", ""),
+                        layer_range=layer_range,
+                    )
+                    candidates.append(peer)
+                except:
+                    continue
+                    
+            return candidates
+        except Exception as e:
+            logger.debug(f"Tracker lookup failed: {e}")
+            return []
+    
+    async def _get_channel(self, grpc_addr: str) -> grpc.aio.Channel:
+        """Get or create gRPC channel for address."""
+        if grpc_addr not in self._channels:
+            # Create new channel for P2P network
+            # Fast keepalive to detect dead nodes quickly in decentralized network
+            options = [
+                ('grpc.keepalive_time_ms', 30000),  # Ping every 30 seconds
+                ('grpc.keepalive_timeout_ms', 10000),  # 10 second timeout
+                ('grpc.keepalive_permit_without_calls', True),  # Ping even when idle
+                ('grpc.http2.max_pings_without_data', 0),  # Unlimited pings
+            ]
+            
+            # Handle different address formats
+            if grpc_addr.startswith("http://"):
+                addr = grpc_addr.replace("http://", "")
+            elif grpc_addr.startswith("https://"):
+                addr = grpc_addr.replace("https://", "")
+            else:
+                addr = grpc_addr
+                
+            self._channels[grpc_addr] = grpc.aio.insecure_channel(addr, options=options)
+            
+        return self._channels[grpc_addr]
+    
+    async def _get_stub(self, grpc_addr: str) -> Any:
+        """Get or create gRPC stub for address."""
+        if not GRPC_AVAILABLE:
+            raise RuntimeError("gRPC not available")
+            
+        if grpc_addr not in self._stubs:
+            channel = await self._get_channel(grpc_addr)
+            self._stubs[grpc_addr] = neuroshard_pb2_grpc.NeuroShardServiceStub(channel)
+            
+        return self._stubs[grpc_addr]
+    
+    async def send_with_failover(
+        self,
+        tensor: torch.Tensor,
+        target_layer: int,
+        session_id: str,
+        metadata: Optional[Dict[str, Any]] = None,
+    ) -> RoutingResult:
+        """
+        Send tensor to target layer with automatic failover.
+        
+        Algorithm:
+        1. Get K candidates sorted by score
+        2. Try primary candidate
+        3. If no ACK in 200ms, try secondary
+        4. Continue until success or all candidates exhausted
+        
+        Args:
+            tensor: Activation tensor to send
+            target_layer: Target layer index
+            session_id: Session identifier for routing
+            metadata: Optional metadata to include
+            
+        Returns:
+            RoutingResult with success status and response data
+        """
+        self.total_routes += 1
+        candidates = self.get_candidates(target_layer)
+        
+        if not candidates:
+            logger.error(f"No candidates available for layer {target_layer}")
+            return RoutingResult(
+                success=False,
+                error=f"No candidates for layer {target_layer}",
+                attempts=0,
+            )
+        
+        last_error: Optional[Exception] = None
+        
+        for i, candidate in enumerate(candidates):
+            start_time = time.time()
+            
+            try:
+                # Attempt send with timeout
+                result = await asyncio.wait_for(
+                    self._send_to_peer(candidate, tensor, session_id, metadata),
+                    timeout=self.ACK_TIMEOUT_MS / 1000.0
+                )
+                
+                latency_ms = (time.time() - start_time) * 1000
+                
+                # Record success
+                candidate.record_success(latency_ms)
+                self.successful_routes += 1
+                
+                if i > 0:
+                    self.failover_count += 1
+                    logger.info(f"Failover to candidate {i}: {candidate.node_id[:8]}...")
+                
+                return RoutingResult(
+                    success=True,
+                    peer_used=candidate.node_id,
+                    latency_ms=latency_ms,
+                    attempts=i + 1,
+                    response_data=result,
+                )
+                
+            except asyncio.TimeoutError:
+                candidate.record_failure()
+                logger.warning(
+                    f"Peer {candidate.node_id[:8]}... timed out after {self.ACK_TIMEOUT_MS}ms, "
+                    f"trying next ({i+1}/{len(candidates)})"
+                )
+                last_error = TimeoutError(f"Peer {candidate.node_id[:8]} timed out")
+                continue
+                
+            except grpc.aio.AioRpcError as e:
+                candidate.record_failure()
+                logger.warning(
+                    f"Peer {candidate.node_id[:8]}... gRPC error: {e.code()}, "
+                    f"trying next ({i+1}/{len(candidates)})"
+                )
+                last_error = e
+                continue
+                
+            except Exception as e:
+                candidate.record_failure()
+                logger.warning(
+                    f"Peer {candidate.node_id[:8]}... failed: {e}, "
+                    f"trying next ({i+1}/{len(candidates)})"
+                )
+                last_error = e
+                continue
+        
+        # All candidates failed
+        error_msg = f"All {len(candidates)} candidates failed for layer {target_layer}"
+        if last_error:
+            error_msg += f": {last_error}"
+            
+        logger.error(error_msg)
+        
+        return RoutingResult(
+            success=False,
+            error=error_msg,
+            attempts=len(candidates),
+        )
+    
+    async def _send_to_peer(
+        self,
+        candidate: PeerCandidate,
+        tensor: torch.Tensor,
+        session_id: str,
+        metadata: Optional[Dict[str, Any]] = None,
+    ) -> Any:
+        """
+        Send activation tensor to a specific peer.
+        
+        Uses gRPC SwarmForward RPC (or falls back to existing Forward).
+        """
+        if not GRPC_AVAILABLE:
+            raise RuntimeError("gRPC not available")
+            
+        stub = await self._get_stub(candidate.grpc_addr)
+        
+        # Serialize tensor
+        tensor_bytes = tensor.cpu().numpy().tobytes()
+        shape = list(tensor.shape)
+        dtype_str = str(tensor.dtype).replace("torch.", "")
+        
+        # Build request - use SwarmForward if available
+        try:
+            # Try SwarmForward (new swarm-aware RPC)
+            request = neuroshard_pb2.SwarmForwardRequest(
+                session_id=session_id,
+                request_id=f"swarm_{time.time()}_{session_id[:8]}",
+                hidden_states=tensor_bytes,
+                hidden_shape=shape,
+                # tensor_dtype inferred from shape/context
+                target_layer=candidate.layer_range[0],
+                sender_url=metadata.get("source_node", "") if metadata else "",
+                micro_batch_id=metadata.get("micro_batch_id", 0) if metadata else 0,
+                priority=metadata.get("priority", 10) if metadata else 10,
+            )
+            
+            response = await stub.SwarmForward(request)
+            return response
+            
+        except grpc.aio.AioRpcError as e:
+            if e.code() == grpc.StatusCode.UNIMPLEMENTED:
+                logger.warning(f"SwarmForward not implemented on {candidate.node_id}, falling back")
+                # Fall back to PipelineForward
+                request = neuroshard_pb2.PipelineForwardRequest(
+                    session_id=session_id,
+                    request_id=f"pipe_{time.time()}",
+                    hidden_states=tensor_bytes,
+                    hidden_shape=shape,
+                    target_shard=candidate.layer_range[0],
+                    sender_url=metadata.get("source_node", "") if metadata else "",
+                )
+                response = await stub.PipelineForward(request)
+                return response
+            raise
+    
+    def get_stats(self) -> Dict[str, Any]:
+        """Get router statistics."""
+        success_rate = (
+            self.successful_routes / self.total_routes 
+            if self.total_routes > 0 else 0.0
+        )
+        failover_rate = (
+            self.failover_count / self.total_routes 
+            if self.total_routes > 0 else 0.0
+        )
+        
+        return {
+            "total_routes": self.total_routes,
+            "successful_routes": self.successful_routes,
+            "success_rate": success_rate,
+            "failover_count": self.failover_count,
+            "failover_rate": failover_rate,
+            "known_peers": len(self.peer_stats),
+            "active_channels": len(self._channels),
+        }
+
+
+# Convenience function for backwards compatibility
+async def get_swarm_candidates(
+    router: SwarmRouter,
+    target_layer: int,
+) -> List[PeerCandidate]:
+    """Get swarm routing candidates for a layer."""
+    return router.get_candidates(target_layer)
+