import time from typing import Dict, List, Optional, Set from dataclasses import dataclass, field from enum import Enum from app.core.models import JobType, WorkerRuntimeType, NodeType from app.phone.workers import get_worker, list_workers_for_session from app.phone.jobs import create_job, get_job, complete_job from app.gbridge.leases import list_active_leases class MeshNodeStatus(str, Enum): ONLINE = "online" BUSY = "busy" OFFLINE = "offline" @dataclass class MeshNode: node_id: str node_type: NodeType runtime_type: WorkerRuntimeType capabilities: List[str] = field(default_factory=list) status: MeshNodeStatus = MeshNodeStatus.ONLINE last_heartbeat: float = field(default_factory=time.time) tokens_per_second: float = 0.0 ram_used_mb: float = 0.0 ram_total_mb: float = 0.0 cpu_usage: float = 0.0 thermal_state: str = "nominal" battery_level: float = 100.0 current_job_id: Optional[str] = None supported_models: List[str] = field(default_factory=list) max_tokens_per_job: int = 512 class MeshRegistry: """Global mesh registry for all compute nodes across all sessions.""" _nodes: Dict[str, MeshNode] = {} _session_nodes: Dict[str, Set[str]] = {} @classmethod def register_node(cls, node: MeshNode) -> None: cls._nodes[node.node_id] = node @classmethod def update_node_heartbeat(cls, node_id: str, metrics: dict) -> None: node = cls._nodes.get(node_id) if node: node.last_heartbeat = time.time() node.tokens_per_second = metrics.get("tokens_per_second", 0.0) node.ram_used_mb = metrics.get("ram_used_mb", 0.0) node.ram_total_mb = metrics.get("ram_total_mb", 0.0) node.cpu_usage = metrics.get("cpu_usage", 0.0) node.thermal_state = metrics.get("thermal_state", "nominal") node.battery_level = metrics.get("battery_level", 100.0) node.status = MeshNodeStatus.BUSY if metrics.get("is_generating") else MeshNodeStatus.ONLINE if node.current_job_id and not metrics.get("is_generating"): node.current_job_id = None @classmethod def mark_node_busy(cls, node_id: str, job_id: str) -> None: node = cls._nodes.get(node_id) if node: node.status = MeshNodeStatus.BUSY node.current_job_id = job_id @classmethod def mark_node_free(cls, node_id: str) -> None: node = cls._nodes.get(node_id) if node: node.status = MeshNodeStatus.ONLINE node.current_job_id = None @classmethod def get_node(cls, node_id: str) -> Optional[MeshNode]: return cls._nodes.get(node_id) @classmethod def list_nodes(cls) -> List[MeshNode]: cls._prune_stale_nodes() return list(cls._nodes.values()) @classmethod def list_available_nodes(cls, capability: Optional[str] = None) -> List[MeshNode]: cls._prune_stale_nodes() nodes = [n for n in cls._nodes.values() if n.status != MeshNodeStatus.OFFLINE] if capability: nodes = [n for n in nodes if capability in n.capabilities] return nodes @classmethod def get_mesh_for_session(cls, session_id: str) -> List[MeshNode]: """Return all mesh nodes visible to a session (global mesh).""" cls._prune_stale_nodes() return list(cls._nodes.values()) @classmethod def schedule_job( cls, job_type: JobType, session_id: str, payload: dict, strategy: str = "best_node" ) -> Optional[str]: """Schedule a job across the mesh and return assigned node_id.""" cls._prune_stale_nodes() available = cls.list_available_nodes() if not available: return None if strategy == "best_node": # Pick node with highest tokens_per_second that supports this job type candidates = [n for n in available if cls._node_supports(n, job_type)] if not candidates: return None best = max(candidates, key=lambda n: n.tokens_per_second or 1.0) return best.node_id elif strategy == "round_robin": candidates = [n for n in available if cls._node_supports(n, job_type)] if not candidates: return None return candidates[0].node_id elif strategy == "parallel_split": # For split strategy, return the primary node (Mac/Ollama) # The secondary task (embedding/classification) goes to iPhone mac_nodes = [n for n in available if n.node_type == NodeType.MAC_AGENT and cls._node_supports(n, job_type)] if mac_nodes: return mac_nodes[0].node_id return available[0].node_id if available else None return None @classmethod def get_combined_throughput(cls) -> dict: """Return aggregate mesh stats.""" nodes = cls.list_nodes() total_tps = sum(n.tokens_per_second for n in nodes) total_ram_used = sum(n.ram_used_mb for n in nodes) total_ram = sum(n.ram_total_mb for n in nodes) active_jobs = sum(1 for n in nodes if n.status == MeshNodeStatus.BUSY) return { "nodes_online": len([n for n in nodes if n.status != MeshNodeStatus.OFFLINE]), "nodes_busy": active_jobs, "total_nodes": len(nodes), "combined_tokens_per_second": round(total_tps, 1), "combined_ram_used_mb": round(total_ram_used, 0), "combined_ram_total_mb": round(total_ram, 0), "active_jobs": active_jobs, } @classmethod def _node_supports(cls, node: MeshNode, job_type: JobType) -> bool: cap_map = { JobType.TEXT_EMBEDDING: ["iphone.text.echo.private", "mac.embed.text.local"], JobType.SMALL_LLM_GENERATE: ["iphone.text.echo.private", "mac.llm.generate.local"], JobType.PRIVACY_REDACTION: ["iphone.privacy.redact.local"], JobType.IMAGE_CLASSIFICATION: ["iphone.image.classify.local"], } required_caps = cap_map.get(job_type, []) return any(cap in node.capabilities for cap in required_caps) or not required_caps @classmethod def _prune_stale_nodes(cls, timeout_seconds: float = 60.0) -> None: now = time.time() stale = [nid for nid, n in cls._nodes.items() if now - n.last_heartbeat > timeout_seconds] for nid in stale: cls._nodes[nid].status = MeshNodeStatus.OFFLINE # Initialize mesh from existing workers def sync_mesh_from_workers() -> None: """Sync mesh registry with existing phone worker state.""" from app.phone.workers import _workers_by_session for session_id, workers in _workers_by_session.items(): for worker in workers.values(): node = MeshNode( node_id=worker.worker_id, node_type=NodeType.IPHONE_WORKER, runtime_type=worker.runtime_type or WorkerRuntimeType.SAFARI_WASM, capabilities=[c.capability_name for c in (worker.capabilities or [])], status=MeshNodeStatus.ONLINE if worker.is_connected else MeshNodeStatus.OFFLINE, tokens_per_second=worker.tokens_per_second or 0.0, ) MeshRegistry.register_node(node) def register_mac_agent_node( node_id: str, capabilities: List[str], supported_models: List[str] = None, ) -> None: node = MeshNode( node_id=node_id, node_type=NodeType.MAC_AGENT, runtime_type=WorkerRuntimeType.OLLAMA, capabilities=capabilities, supported_models=supported_models or ["llama3.2", "mistral"], ) MeshRegistry.register_node(node)