some kademlia fixes

Author: erhant

README.md

@@ -34,7 +34,7 @@ Compute nodes can technically do any arbitrary task, from computing the square r
 
 - **Workflows**: Each task is given in the form of a [workflow](https://github.com/andthattoo/ollama-workflows). Every workflow defines an agentic behavior for the chosen LLM, all captured in a single JSON file, and can represent things ranging from simple LLM generations to iterative web searching & reasoning.
 
-## Node Running
+### Running a Node
 
 Refer to [node guide](./docs/NODE_GUIDE.md) to quickly get started and run your own node!
 

compute/src/node.rs

@@ -26,18 +26,17 @@ const PUBLISH_CHANNEL_BUFSIZE: usize = 1024;
 
 pub struct DriaComputeNode {
     pub config: DriaComputeNodeConfig,
-    pub p2p: DriaP2PCommander,
     pub available_nodes: AvailableNodes,
-    /// Gossipsub message receiver.
+    /// Peer-to-peer client commander to interact with the network.
+    pub p2p: DriaP2PCommander,
+    /// Gossipsub message receiver, used by peer-to-peer client in a separate thread.
     message_rx: mpsc::Receiver<(PeerId, MessageId, Message)>,
-    /// Publish receiver to receive messages to be published.
+    /// Publish receiver to receive messages to be published,
     publish_rx: mpsc::Receiver<WorkflowsWorkerOutput>,
     /// Workflow transmitter to send batchable tasks.
     workflow_batch_tx: Option<mpsc::Sender<WorkflowsWorkerInput>>,
     /// Workflow transmitter to send single tasks.
     workflow_single_tx: Option<mpsc::Sender<WorkflowsWorkerInput>>,
-    // TODO: instead of piggybacking task metadata within channels, we can store them here
-    // in a hashmap alone, and then use the task_id to get the metadata when needed
     // Single tasks hash-map
     pending_tasks_single: HashSet<String>,
     // Batch tasks hash-map
@@ -402,7 +401,7 @@ impl DriaComputeNode {
 
         // print tasks count
         let [single, batch] = self.get_pending_task_count();
-        log::info!("Pending Tasks (single/batch):   {} / {}", single, batch);
+        log::info!("Pending Tasks (single/batch): {} / {}", single, batch);
 
         // completed tasks count
         log::debug!(

compute/src/workers/workflow.rs

@@ -4,6 +4,8 @@ use tokio::sync::mpsc;
 
 use crate::payloads::TaskStats;
 
+// TODO: instead of piggybacking stuff here, maybe node can hold it in a hashmap w.r.t taskId
+
 pub struct WorkflowsWorkerInput {
     pub entry: Option<Entry>,
     pub executor: Executor,
@@ -59,7 +61,7 @@ impl WorkflowsWorker {
 
     /// Closes the workflow receiver channel.
     fn shutdown(&mut self) {
-        log::warn!("Closing workflows worker.");
+        log::info!("Closing workflows worker.");
         self.workflow_rx.close();
     }
 

p2p/README.md

@@ -1,4 +1,4 @@
-# DKN Peer-to-Peer Client
+# Dria Peer-to-Peer Client
 
 Dria Knowledge Network is a peer-to-peer network, built over libp2p. This crate is a wrapper client to easily interact with DKN.
 

p2p/src/behaviour.rs

@@ -65,12 +65,14 @@ fn create_kademlia_behaviour(
 ) -> kad::Behaviour<MemoryStore> {
     use kad::{Behaviour, Config};
 
-    const QUERY_TIMEOUT_SECS: u64 = 5 * 60;
-    const RECORD_TTL_SECS: u64 = 30;
+    const KADEMLIA_BOOTSTRAP_INTERVAL_SECS: u64 = 5 * 60; // default is 5 minutes
+    const QUERY_TIMEOUT_SECS: u64 = 3 * 60; // default is 1 minute
 
     let mut cfg = Config::new(protocol_name);
     cfg.set_query_timeout(Duration::from_secs(QUERY_TIMEOUT_SECS))
-        .set_record_ttl(Some(Duration::from_secs(RECORD_TTL_SECS)));
+        .set_periodic_bootstrap_interval(Some(Duration::from_secs(
+            KADEMLIA_BOOTSTRAP_INTERVAL_SECS,
+        )));
 
     Behaviour::with_config(local_peer_id, MemoryStore::new(local_peer_id), cfg)
 }
@@ -157,7 +159,6 @@ fn create_gossipsub_behaviour(author: PeerId) -> Result<gossipsub::Behaviour> {
         MessageId::from(digest.to_be_bytes())
     };
 
-    // TODO: add data transform here later
     Behaviour::new(
         MessageAuthenticity::Author(author),
         ConfigBuilder::default()

p2p/src/client.rs

@@ -154,7 +154,7 @@ impl DriaP2PClient {
                     Some(c) => self.handle_command(c).await,
                     // channel closed, thus shutting down the network event loop
                     None=>  {
-                        log::warn!("Closing P2P client.");
+                        log::info!("Closing peer-to-peer client.");
                         return
                     },
                 },
@@ -253,6 +253,7 @@ impl DriaP2PClient {
 
                 // remove own peerId from Autonat server list
                 self.swarm.behaviour_mut().autonat.remove_server(&peer_id);
+
                 let _ = sender.send(());
             }
         }
@@ -272,32 +273,78 @@ impl DriaP2PClient {
                 }
             }
 
+            // kademlia events
             SwarmEvent::Behaviour(DriaBehaviourEvent::Kademlia(
                 kad::Event::OutboundQueryProgressed {
                     result: QueryResult::GetClosestPeers(result),
                     ..
                 },
             )) => self.handle_closest_peers_result(result),
+
+            // identify events
             SwarmEvent::Behaviour(DriaBehaviourEvent::Identify(identify::Event::Received {
                 peer_id,
                 info,
                 ..
             })) => self.handle_identify_event(peer_id, info),
 
+            // autonat events
             SwarmEvent::Behaviour(DriaBehaviourEvent::Autonat(autonat::Event::StatusChanged {
                 old,
                 new,
             })) => {
                 log::warn!("AutoNAT status changed from {:?} to {:?}", old, new);
             }
 
+            // log listen addreses
             SwarmEvent::NewListenAddr { address, .. } => {
                 log::warn!("Local node is listening on {}", address);
             }
+
+            // add external address of peers to Kademlia routing table
+            SwarmEvent::NewExternalAddrOfPeer { peer_id, address } => {
+                self.swarm
+                    .behaviour_mut()
+                    .kademlia
+                    .add_address(&peer_id, address);
+            }
+            // add your own peer_id to kademlia as well
             SwarmEvent::ExternalAddrConfirmed { address } => {
                 // this is usually the external address via relay
                 log::info!("External address confirmed: {}", address);
+                let peer_id = *self.swarm.local_peer_id();
+                self.swarm
+                    .behaviour_mut()
+                    .kademlia
+                    .add_address(&peer_id, address);
             }
+
+            // SwarmEvent::IncomingConnectionError {
+            //     local_addr,
+            //     send_back_addr,
+            //     error,
+            //     connection_id,
+            // } => {
+            //     log::debug!(
+            //         "Incoming connection {} error: from {} to {} - {:?}",
+            //         connection_id,
+            //         local_addr,
+            //         send_back_addr,
+            //         error
+            //     );
+            // }
+            // SwarmEvent::IncomingConnection {
+            //     connection_id,
+            //     local_addr,
+            //     send_back_addr,
+            // } => {
+            //     log::debug!(
+            //         "Incoming connection {} attepmt: from {} to {}",
+            //         connection_id,
+            //         local_addr,
+            //         send_back_addr
+            //     );
+            // }
             // SwarmEvent::OutgoingConnectionError { peer_id, error, .. } => {
             //     if let Some(peer_id) = peer_id {
             //         log::warn!("Could not connect to peer {}: {:?}", peer_id, error);

workflows/README.md

@@ -1,7 +1,6 @@
-# DKN Workflows
+# Dria Workflows
 
-We make use of Ollama Workflows in DKN; however, we also want to make sure that the chosen models are valid and is performant enough (i.e. have enough TPS).
-This crate handles the configurations of models to be used, and implements various service checks.
+We make use of [Ollama Workflows](https://github.com/andthattoo/ollama-workflows) in Dria Knowledge Network; however, we also want to make sure that the chosen models are valid and is performant enough (i.e. have enough TPS). This crate handles the configurations of models to be used, and implements various service checks.
 
 There are two types of services: