|
|
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 |
