| | from typing import List, Dict, Any, Optional
|
| | import asyncio
|
| | import json
|
| | from fastapi import FastAPI, WebSocket
|
| | from pydantic import BaseModel
|
| | import torch
|
| |
|
| | class PeerNetwork:
|
| | def __init__(self, host: str = "localhost", port: int = 8000):
|
| | self.app = FastAPI()
|
| | self.active_peers: Dict[str, WebSocket] = {}
|
| | self.host = host
|
| | self.port = port
|
| |
|
| |
|
| | @self.app.websocket("/ws/{peer_id}")
|
| | async def websocket_endpoint(websocket: WebSocket, peer_id: str):
|
| | await self.connect_peer(websocket, peer_id)
|
| | try:
|
| | while True:
|
| | data = await websocket.receive_text()
|
| | await self.broadcast(data, peer_id)
|
| | except Exception:
|
| | await self.disconnect_peer(peer_id)
|
| |
|
| | async def connect_peer(self, websocket: WebSocket, peer_id: str):
|
| | """Connect a new peer to the network"""
|
| | await websocket.accept()
|
| | self.active_peers[peer_id] = websocket
|
| |
|
| | async def disconnect_peer(self, peer_id: str):
|
| | """Remove a peer from the network"""
|
| | if peer_id in self.active_peers:
|
| | await self.active_peers[peer_id].close()
|
| | del self.active_peers[peer_id]
|
| |
|
| | async def broadcast(self, message: str, sender_id: str):
|
| | """Broadcast a message to all peers except the sender"""
|
| | for peer_id, websocket in self.active_peers.items():
|
| | if peer_id != sender_id:
|
| | await websocket.send_text(message)
|
| |
|
| | class OpenPeerClient:
|
| | def __init__(self, network_url: str):
|
| | self.network_url = network_url
|
| | self.websocket: Optional[WebSocket] = None
|
| | self.peer_id: Optional[str] = None
|
| |
|
| | async def connect(self, peer_id: str):
|
| | """Connect to the peer network"""
|
| | self.peer_id = peer_id
|
| | self.websocket = await WebSocket.connect(f"{self.network_url}/ws/{peer_id}")
|
| |
|
| | async def send_model_update(self, model_state: Dict[str, torch.Tensor]):
|
| | """Send model state updates to the network"""
|
| | if not self.websocket:
|
| | raise RuntimeError("Not connected to network")
|
| |
|
| | serialized_state = {
|
| | "type": "model_update",
|
| | "peer_id": self.peer_id,
|
| | "state": {k: v.cpu().numpy().tolist() for k, v in model_state.items()}
|
| | }
|
| | await self.websocket.send_text(json.dumps(serialized_state))
|
| |
|
| | async def receive_updates(self):
|
| | """Receive updates from the network"""
|
| | if not self.websocket:
|
| | raise RuntimeError("Not connected to network")
|
| |
|
| | while True:
|
| | data = await self.websocket.receive_text()
|
| | yield json.loads(data)
|
| |
|
| | def create_peer_network(host: str = "localhost", port: int = 8000) -> PeerNetwork:
|
| | """Create and start a peer network server"""
|
| | network = PeerNetwork(host, port)
|
| | import uvicorn
|
| | uvicorn.run(network.app, host=host, port=port)
|
| | return network |
