Initial AI conversion with some work

2025-02-12 07:03:39 +00:00
parent fed48a2868
commit 70d58995ec
4 changed files with 401 additions and 5 deletions
@@ -0,0 +1,260 @@
+use crate::topology::DeviceCapabilities;
+use serde::{Deserialize, Serialize};
+use std::{
+    collections::HashMap,
+    net::{IpAddr, SocketAddr},
+    sync::Arc,
+    time::{Duration, Instant},
+};
+use thiserror::Error;
+use tokio::{net::UdpSocket, sync::RwLock, time};
+use tracing::{error};
+
+#[derive(Error, Debug)]
+pub enum DiscoveryError {
+    #[error("IO error: {0}")]
+    Io(#[from] std::io::Error),
+
+    #[error("JSON serialization error: {0}")]
+    Json(#[from] serde_json::Error),
+
+    #[error("UTF-8 conversion error: {0}")]
+    Utf8(#[from] std::string::FromUtf8Error),
+
+    #[error("Address parse error: {0}")]
+    AddrParse(#[from] std::net::AddrParseError),
+
+    #[error("Network interface error: {0}")]
+    NetworkInterface(String),
+
+    #[error("Peer error: {0}")]
+    Peer(String),
+}
+
+// Constants
+const DEBUG: i32 = 0;
+const BROADCAST_ADDR: &str = "255.255.255.255";
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+struct DiscoveryMessage {
+    message_type: String,
+    node_id: String,
+    grpc_port: u16,
+    device_capabilities: DeviceCapabilities,
+    priority: i32,
+    interface_name: String,
+    interface_type: String,
+}
+
+struct PeerInfo {
+    peer_handle: PeerHandle,
+    connected_at: Instant,
+    last_seen: Instant,
+    priority: i32,
+}
+
+struct UdpDiscovery {
+    node_id: String,
+    node_port: u16,
+    listen_port: u16,
+    broadcast_port: u16,
+    broadcast_interval: Duration,
+    discovery_timeout: Duration,
+    device_capabilities: DeviceCapabilities,
+    allowed_node_ids: Option<Vec<String>>,
+    allowed_interface_types: Option<Vec<String>>,
+    known_peers: Arc<RwLock<HashMap<String, PeerInfo>>>,
+}
+
+impl UdpDiscovery {
+    pub fn new(
+        node_id: String,
+        node_port: u16,
+        listen_port: u16,
+        broadcast_port: u16,
+        broadcast_interval: Duration,
+        discovery_timeout: Duration,
+        device_capabilities: DeviceCapabilities,
+        allowed_node_ids: Option<Vec<String>>,
+        allowed_interface_types: Option<Vec<String>>,
+    ) -> Self {
+        Self {
+            node_id,
+            node_port,
+            listen_port,
+            broadcast_port,
+            broadcast_interval,
+            discovery_timeout,
+            device_capabilities,
+            allowed_node_ids,
+            allowed_interface_types,
+            known_peers: Arc::new(RwLock::new(HashMap::new())),
+        }
+    }
+
+    pub async fn start(&self) -> Result<(), DiscoveryError> {
+        let broadcast_task = self.task_broadcast_presence();
+        let listen_task = self.task_listen_for_peers();
+        let cleanup_task = self.task_cleanup_peers();
+
+        tokio::try_join!(broadcast_task, listen_task, cleanup_task)?;
+        Ok(())
+    }
+}
+
+impl UdpDiscovery {
+    async fn task_listen_for_peers(&self) -> Result<(), DiscoveryError> {
+        let socket = UdpSocket::bind(format!("0.0.0.0:{}", self.listen_port)).await?;
+        let mut buf = vec![0u8; 65535];
+
+        loop {
+            let (len, addr) = socket.recv_from(&mut buf).await?;
+            if len == 0 {
+                continue;
+            }
+
+            if let Ok(message) = String::from_utf8(buf[..len].to_vec()) {
+                if let Ok(discovery_message) = serde_json::from_str::<DiscoveryMessage>(&message) {
+                    self.handle_discovery_message(discovery_message, addr)
+                        .await
+                        .map_err(|e| DiscoveryError::Peer(e.to_string()))?;
+                }
+            }
+        }
+    }
+}
+
+impl UdpDiscovery {
+    async fn task_broadcast_presence(&self) -> Result<(), DiscoveryError> {
+        let socket = UdpSocket::bind("0.0.0.0:0").await?;
+        socket.set_broadcast(true)?;
+
+        loop {
+            let interfaces = get_all_ip_addresses_and_interfaces()?;
+            for (addr, interface_name) in interfaces {
+                let (interface_priority, interface_type) =
+                    get_interface_priority_and_type(&interface_name)
+                        .await
+                        .map_err(|e| DiscoveryError::NetworkInterface(e.to_string()))?;
+
+                let message = DiscoveryMessage {
+                    message_type: "discovery".to_string(),
+                    node_id: self.node_id.clone(),
+                    grpc_port: self.node_port,
+                    device_capabilities: self.device_capabilities.clone(),
+                    priority: interface_priority,
+                    interface_name: interface_name.clone(),
+                    interface_type: interface_type.clone(),
+                };
+
+                let message_json = serde_json::to_string(&message)?;
+                let broadcast_addr =
+                    SocketAddr::new(get_broadcast_address(&addr).parse()?, self.broadcast_port);
+
+                if let Err(e) = socket
+                    .send_to(message_json.as_bytes(), &broadcast_addr)
+                    .await
+                {
+                    error!("Error broadcasting to {}: {}", broadcast_addr, e);
+                }
+            }
+
+            time::sleep(self.broadcast_interval).await;
+        }
+    }
+}
+
+impl UdpDiscovery {
+    async fn task_cleanup_peers(&self) -> Result<(), DiscoveryError> {
+        loop {
+            let now = Instant::now();
+
+            // TODO: Do we really want a read then write lock or should we just take a write lock
+            //  to begin with?
+            let peers_to_remove = {
+                let mut peers_to_remove = Vec::new();
+
+                let peers = self.known_peers.read().await;
+                for (peer_id, peer_info) in peers.iter() {
+                    if self.should_remove_peer(peer_info, now).await {
+                        peers_to_remove.push(peer_id.clone());
+                    }
+                }
+
+                peers_to_remove
+            };
+
+            {
+                let mut peers = self.known_peers.write().await;
+                for peer_id in peers_to_remove {
+                    peers.remove(&peer_id);
+                }
+            };
+
+            time::sleep(self.broadcast_interval).await;
+        }
+    }
+
+    async fn should_remove_peer(&self, peer_info: &PeerInfo, now: Instant) -> bool {
+        let is_connected = peer_info
+            .peer_handle
+            .is_connected()
+            .await
+            .ok()
+            .unwrap_or(false);
+
+        if !is_connected {
+            return true;
+        }
+
+        if now.duration_since(peer_info.connected_at) > self.discovery_timeout {
+            return true;
+        }
+
+        if now.duration_since(peer_info.last_seen) > self.discovery_timeout {
+            return true;
+        }
+
+        peer_info
+            .peer_handle
+            .health_check()
+            .await
+            .ok()
+            .unwrap_or(false)
+    }
+}
+
+fn get_broadcast_address(ip_addr: &IpAddr) -> String {
+    match ip_addr {
+        IpAddr::V4(addr) => {
+            let octets = addr.octets();
+            format!("{}.{}.{}.255", octets[0], octets[1], octets[2])
+        }
+        _ => BROADCAST_ADDR.to_string(),
+    }
+}
+
+// You'll need to implement these functions based on your system
+fn get_all_ip_addresses_and_interfaces() -> Result<Vec<(IpAddr, String)>, DiscoveryError> {
+    // Implementation needed
+    Ok(vec![])
+}
+
+async fn get_interface_priority_and_type(
+    interface_name: &str,
+) -> Result<(i32, String), DiscoveryError> {
+    // Implementation needed
+    Ok((0, "unknown".to_string()))
+}
+
+struct PeerHandle;
+
+impl PeerHandle {
+    async fn is_connected(&self) -> Result<bool, DiscoveryError> {
+        Ok(true)
+    }
+
+    async fn health_check(&self) -> Result<bool, DiscoveryError> {
+        Ok(true)
+    }
+}
@@ -1,4 +1,5 @@
 mod topology;
+mod discovery;

 use serde::{Deserialize, Serialize};
 use serde_json::Value;
@@ -9,7 +10,7 @@ use crate::node_service::{
    PromptRequest, SendOpaqueStatusRequest, SendResultRequest, Tensor, Topology as TopologyProto,
 };
 use node_service::node_service_server::{NodeService, NodeServiceServer};
-use node_service::{Shard, TensorRequest};
+use node_service::{TensorRequest};
 use topology::Topology;

 pub mod node_service {
@@ -36,6 +37,15 @@ enum OpaqueStatus {
    SupportedInferenceEngines(SupportedInferenceEngines),
 }

+#[derive(Debug, Deserialize, Serialize, Clone)]
+struct Shard {
+    pub model_id: String,
+    pub start_layer: i32,
+    pub end_layer: i32,
+    pub n_layers: i32,
+}
+
+
 #[derive(Debug, Deserialize, Serialize, Clone)]
 struct NodeStatus {
    node_id: String,
@@ -69,7 +79,7 @@ struct SupportedInferenceEngines {
 }

 impl Node {
-    fn on_opaque_status(&self, request_id: String, status: String) {
+    fn on_opaque_status(&self, _request_id: String, status: String) {
        let status = serde_json::from_str::<OpaqueStatus>(&status).unwrap();

        match status {
@@ -140,7 +150,7 @@ impl NodeService for Node {
    async fn collect_topology(
        &self,
        request: Request<CollectTopologyRequest>,
-    ) -> Result<Response<Topology>, Status> {
+    ) -> Result<Response<TopologyProto>, Status> {
        todo!()
    }

@@ -155,8 +165,10 @@ impl NodeService for Node {
        &self,
        request: Request<SendOpaqueStatusRequest>,
    ) -> Result<Response<Empty>, Status> {
-        let request_id = request.into_inner().request_id;
-        let status = request.into_inner().status;
+        let request = request.into_inner();
+        let request_id = request.request_id;
+        let status = request.status;
+
        println!(
            "Received SendOpaqueStatus request: {} {}",
            request_id, status
@@ -174,6 +186,9 @@ impl NodeService for Node {

 #[tokio::main]
 async fn main() -> Result<(), Box<dyn std::error::Error>> {
+    // install global collector configured based on RUST_LOG env var.
+    tracing_subscriber::fmt::init();
+
    let grpc_addr = "[::1]:50051".parse()?;
    let node = Node::default();