Merge pull request #148 from firstbatchxyz/erhant/node-refresh-rfks-p…

…eer-fixes

feat: multi-network & peer disconnections

Cargo.toml

@@ -7,7 +7,7 @@ default-members = ["compute"]
 
 [workspace.package]
 edition = "2021"
-version = "0.2.22"
+version = "0.2.23"
 license = "Apache-2.0"
 readme = "README.md"
 

compute/src/config.rs

@@ -9,6 +9,33 @@ use libsecp256k1::{PublicKey, SecretKey};
 
 use std::{env, str::FromStr};
 
+/// Network type.
+#[derive(Default, Debug, Clone, Copy)]
+pub enum DriaNetworkType {
+    #[default]
+    Community,
+    Pro,
+}
+
+impl From<&str> for DriaNetworkType {
+    fn from(s: &str) -> Self {
+        match s {
+            "community" => DriaNetworkType::Community,
+            "pro" => DriaNetworkType::Pro,
+            _ => Default::default(),
+        }
+    }
+}
+
+impl DriaNetworkType {
+    pub fn protocol_name(&self) -> &str {
+        match self {
+            DriaNetworkType::Community => "dria",
+            DriaNetworkType::Pro => "dria-sdk",
+        }
+    }
+}
+
 #[derive(Debug, Clone)]
 pub struct DriaComputeNodeConfig {
     /// Wallet secret/private key.
@@ -23,6 +50,8 @@ pub struct DriaComputeNodeConfig {
     pub p2p_listen_addr: Multiaddr,
     /// Workflow configurations, e.g. models and providers.
     pub workflows: DriaWorkflowsConfig,
+    /// Network type of the node.
+    pub network_type: DriaNetworkType,
 }
 
 /// The default P2P network listen address.
@@ -104,13 +133,19 @@ impl DriaComputeNodeConfig {
         let p2p_listen_addr = Multiaddr::from_str(&p2p_listen_addr_str)
             .expect("Could not parse the given P2P listen address.");
 
+        // parse network type
+        let network_type = env::var("DKN_NETWORK")
+            .map(|s| DriaNetworkType::from(s.as_str()))
+            .unwrap_or_default();
+
         Self {
             admin_public_key,
             secret_key,
             public_key,
             address,
             workflows,
             p2p_listen_addr,
+            network_type,
         }
     }
 

compute/src/handlers/pingpong.rs

@@ -64,7 +64,7 @@ impl ComputeHandler for PingpongHandler {
             Self::RESPONSE_TOPIC,
             &node.config.secret_key,
         );
-        node.publish(message)?;
+        node.publish(message).await?;
 
         Ok(MessageAcceptance::Accept)
     }

compute/src/handlers/workflow.rs

@@ -161,8 +161,7 @@ impl ComputeHandler for WorkflowHandler {
         };
 
         // try publishing the result
-
-        if let Err(publish_err) = node.publish(message) {
+        if let Err(publish_err) = node.publish(message).await {
             let err_msg = format!("Could not publish result: {:?}", publish_err);
             log::error!("{}", err_msg);
 
@@ -175,7 +174,7 @@ impl ComputeHandler for WorkflowHandler {
                 Self::RESPONSE_TOPIC,
                 &node.config.secret_key,
             );
-            node.publish(message)?;
+            node.publish(message).await?;
         }
 
         Ok(acceptance)

compute/src/lib.rs

@@ -11,5 +11,5 @@ pub(crate) mod utils;
 /// This value is attached within the published messages.
 pub const DRIA_COMPUTE_NODE_VERSION: &str = env!("CARGO_PKG_VERSION");
 
-pub use config::DriaComputeNodeConfig;
+pub use config::{DriaComputeNodeConfig, DriaNetworkType};
 pub use node::DriaComputeNode;

compute/src/main.rs

@@ -1,5 +1,5 @@
 use dkn_compute::*;
-use eyre::{Context, Result};
+use eyre::Result;
 use std::env;
 use tokio_util::sync::CancellationToken;
 
@@ -25,7 +25,7 @@ async fn main() -> Result<()> {
 ██║  ██║██╔══██╗██║██╔══██║   https://dria.co
 ██████╔╝██║  ██║██║██║  ██║
 ╚═════╝ ╚═╝  ╚═╝╚═╝╚═╝  ╚═╝
-"#,
+"#
     );
 
     let token = CancellationToken::new();
@@ -52,10 +52,6 @@ async fn main() -> Result<()> {
     let service_check_token = token.clone();
     let config = tokio::spawn(async move {
         tokio::select! {
-            _ = service_check_token.cancelled() => {
-                log::info!("Service check cancelled.");
-                config
-            }
             result = config.workflows.check_services() => {
                 if let Err(err) = result {
                     log::error!("Error checking services: {:?}", err);
@@ -64,38 +60,44 @@ async fn main() -> Result<()> {
                 log::warn!("Using models: {:#?}", config.workflows.models);
                 config
             }
+            _ = service_check_token.cancelled() => {
+                log::info!("Service check cancelled.");
+                config
+            }
         }
     })
-    .await
-    .wrap_err("error during service checks")?;
-
-    if !token.is_cancelled() {
-        // launch the node in a separate thread
-        let node_token = token.clone();
-        let node_handle = tokio::spawn(async move {
-            match DriaComputeNode::new(config, node_token).await {
-                Ok(mut node) => {
-                    if let Err(err) = node.launch().await {
-                        log::error!("Node launch error: {}", err);
-                        panic!("Node failed.")
-                    };
-                }
-                Err(err) => {
-                    log::error!("Node setup error: {}", err);
-                    panic!("Could not setup node.")
-                }
-            }
-        });
+    .await?;
 
-        // wait for tasks to complete
-        if let Err(err) = node_handle.await {
-            log::error!("Node handle error: {}", err);
-            panic!("Could not exit Node thread handle.");
-        };
-    } else {
-        log::warn!("Not launching node due to early exit.");
+    // check early exit due to failed service check
+    if token.is_cancelled() {
+        log::warn!("Not launching node due to early exit, bye!");
+        return Ok(());
     }
 
+    let node_token = token.clone();
+    let (mut node, p2p) = DriaComputeNode::new(config, node_token).await?;
+
+    // launch the p2p in a separate thread
+    log::info!("Spawning peer-to-peer client thread.");
+    let p2p_handle = tokio::spawn(async move { p2p.run().await });
+
+    // launch the node in a separate thread
+    log::info!("Spawning compute node thread.");
+    let node_handle = tokio::spawn(async move {
+        if let Err(err) = node.launch().await {
+            log::error!("Node launch error: {}", err);
+            panic!("Node failed.")
+        };
+    });
+
+    // wait for tasks to complete
+    if let Err(err) = node_handle.await {
+        log::error!("Node handle error: {}", err);
+    };
+    if let Err(err) = p2p_handle.await {
+        log::error!("P2P handle error: {}", err);
+    };
+
     log::info!("Bye!");
     Ok(())
 }
@@ -158,6 +160,7 @@ async fn wait_for_termination(cancellation: CancellationToken) -> Result<()> {
     Ok(())
 }
 
+// #[deprecated]
 /// Very CRUDE fix due to launcher log level bug
 ///
 /// TODO: remove me later when the launcher is fixed