added better RPC discovery logic

Author: erhant

compute/src/config.rs

@@ -1,15 +1,12 @@
-use dkn_p2p::{
-    libp2p::{Multiaddr, PeerId},
-    DriaNetworkType,
-};
+use dkn_p2p::libp2p::{Multiaddr, PeerId};
 use dkn_workflows::DriaWorkflowsConfig;
 use eyre::{eyre, Result};
 use libsecp256k1::{PublicKey, SecretKey};
 use std::{env, str::FromStr};
 
 use dkn_utils::{
     crypto::{public_key_to_address, secret_to_keypair},
-    SemanticVersion,
+    DriaNetwork, SemanticVersion,
 };
 
 const DEFAULT_TASK_BATCH_SIZE: usize = 5;
@@ -32,7 +29,7 @@ pub struct DriaComputeNodeConfig {
     /// Workflow configurations, e.g. models and providers.
     pub workflows: DriaWorkflowsConfig,
     /// Network type of the node.
-    pub network_type: DriaNetworkType,
+    pub network_type: DriaNetwork,
     /// Batch size for batchable tasks (e.g. API-based ones).
     ///
     /// A higher value will help execute more tasks concurrently,
@@ -95,7 +92,7 @@ impl DriaComputeNodeConfig {
 
         // parse network type
         let network_type = env::var("DKN_NETWORK")
-            .map(|s| DriaNetworkType::from(s.as_str()))
+            .map(|s| DriaNetwork::from(s.as_str()))
             .unwrap_or_default();
 
         // parse batch size

compute/src/node/diagnostic.rs

@@ -113,7 +113,7 @@ impl DriaComputeNode {
                 "Connection to RPC {} is lost, geting a new one!",
                 self.dria_rpc.addr,
             );
-            match DriaRPC::new_for_network(self.dria_rpc.network).await {
+            match DriaRPC::new_for_network(self.dria_rpc.network, &self.config.version).await {
                 Ok(new_rpc) => {
                     self.dria_rpc = new_rpc;
 

compute/src/node/mod.rs

@@ -82,7 +82,7 @@ impl DriaComputeNode {
             log::info!("Using initial RPC address: {}", addr);
             DriaRPC::new(addr, config.network_type).expect("could not get RPC to connect to")
         } else {
-            DriaRPC::new_for_network(config.network_type)
+            DriaRPC::new_for_network(config.network_type, &config.version)
                 .await
                 .expect("could not get RPC to connect to")
         };

compute/src/node/rpc.rs

@@ -1,5 +1,5 @@
 use dkn_p2p::libp2p::{multiaddr::Protocol, Multiaddr, PeerId};
-use dkn_p2p::DriaNetworkType;
+use dkn_utils::{DriaNetwork, SemanticVersion};
 use eyre::{Context, OptionExt, Result};
 use std::fmt::Debug;
 
@@ -8,12 +8,12 @@ use std::fmt::Debug;
 pub struct DriaRPC {
     pub addr: Multiaddr,
     pub peer_id: PeerId,
-    pub network: DriaNetworkType,
+    pub network: DriaNetwork,
 }
 
 impl DriaRPC {
     /// Creates a new RPC target at the given type, along with a network type for refreshing the RPC address.
-    pub fn new(addr: Multiaddr, network: DriaNetworkType) -> Result<Self> {
+    pub fn new(addr: Multiaddr, network: DriaNetwork) -> Result<Self> {
         let peer_id = addr
             .iter()
             .find_map(|p| match p {
@@ -29,28 +29,31 @@ impl DriaRPC {
         })
     }
 
-    /// Creates a new RPC target for the given network type.
-    pub async fn new_for_network(network: DriaNetworkType) -> Result<Self> {
-        let addr = get_rpc_for_network(&network).await?;
+    /// Creates a new RPC target for the given network type and version.
+    pub async fn new_for_network(network: DriaNetwork, version: &SemanticVersion) -> Result<Self> {
+        let addr = get_rpc_for_network(&network, version).await?;
         Self::new(addr, network)
     }
 }
 
 /// Calls the DKN API to get an RPC address for the given network type.
 ///
 /// The peer id is expected to be within the multi-address.
-async fn get_rpc_for_network(network: &DriaNetworkType) -> Result<Multiaddr> {
+async fn get_rpc_for_network(
+    network: &DriaNetwork,
+    version: &SemanticVersion,
+) -> Result<Multiaddr> {
     #[derive(serde::Deserialize, Debug)]
     struct DriaNodesApiResponse {
         pub rpc: Multiaddr,
     }
 
     // url to be used is determined by the network type
-    let url = match network {
-        DriaNetworkType::Community => "https://dkn.dria.co/v5/available-nodes",
-        DriaNetworkType::Pro => "https://dkn.dria.co/v5/sdk/available-nodes",
-        DriaNetworkType::Test => "https://dkn.dria.co/v5/test/available-nodes",
+    let base_url = match network {
+        DriaNetwork::Mainnet => "https://dkn.dria.co/available-nodes",
+        DriaNetwork::Testnet => "https://dkn.dria.co/available-nodes",
     };
+    let url = format!("{}/{}", base_url, version.as_major_minor());
 
     // make the request
     let response = reqwest::get(url).await?;
@@ -69,7 +72,9 @@ mod tests {
 
     #[tokio::test]
     async fn test_dria_nodes() {
-        let node = DriaRPC::new_for_network(DriaNetworkType::Community).await;
+        let node =
+            DriaRPC::new_for_network(DriaNetwork::Mainnet, &SemanticVersion::from_crate_version())
+                .await;
         assert!(node.is_ok());
     }
 }

p2p/src/lib.rs

@@ -9,9 +9,6 @@ pub use commands::{DriaP2PCommand, DriaP2PCommander};
 mod protocol;
 pub use protocol::DriaP2PProtocol;
 
-mod network;
-pub use network::DriaNetworkType;
-
 // re-exports
 pub use libp2p;
 pub use libp2p_identity;

p2p/src/network.rs

@@ -12,6 +12,9 @@ pub use csv::split_csv_line;
 mod env;
 pub use env::safe_read_env;
 
+mod network;
+pub use network::DriaNetwork;
+
 mod version;
 pub use version::SemanticVersion;
 

utils/src/lib.rs

@@ -0,0 +1,77 @@
+use crate::SemanticVersion;
+
+/// Network type, either mainnet or testnet.
+#[derive(Default, Debug, Clone, Copy, PartialEq)]
+pub enum DriaNetwork {
+    Mainnet,
+    #[default]
+    Testnet,
+}
+
+impl From<&str> for DriaNetwork {
+    fn from(s: &str) -> Self {
+        match s {
+            "mainnet" => DriaNetwork::Mainnet,
+            "testnet" => DriaNetwork::Testnet,
+            _ => Default::default(),
+        }
+    }
+}
+
+impl std::fmt::Display for DriaNetwork {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            DriaNetwork::Mainnet => write!(f, "mainnet"),
+            DriaNetwork::Testnet => write!(f, "testnet"),
+        }
+    }
+}
+
+impl DriaNetwork {
+    pub fn protocol_name(&self) -> &str {
+        match self {
+            DriaNetwork::Mainnet => "dria",
+            DriaNetwork::Testnet => "dria-test",
+        }
+    }
+
+    pub fn discovery_url(&self, version: &SemanticVersion) -> String {
+        let base_url = match self {
+            DriaNetwork::Mainnet => "https://mainnet.dkn.dria.co/discovery/available-nodes",
+            DriaNetwork::Testnet => "https://testnet.dkn.dria.co/discovery/available-nodes",
+        };
+
+        format!("{}/{}", base_url, version.as_major_minor())
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_dria_network() {
+        let mainnet = DriaNetwork::Mainnet;
+        let testnet = DriaNetwork::Testnet;
+        let version = SemanticVersion {
+            major: 1,
+            minor: 0,
+            patch: 42,
+        };
+
+        assert_eq!(mainnet.to_string(), "mainnet");
+        assert_eq!(testnet.to_string(), "testnet");
+
+        assert_eq!(mainnet.protocol_name(), "dria");
+        assert_eq!(testnet.protocol_name(), "dria-test");
+
+        assert_eq!(
+            mainnet.discovery_url(&version),
+            "https://mainnet.dkn.dria.co/discovery/available-nodes/1.0"
+        );
+        assert_eq!(
+            testnet.discovery_url(&version),
+            "https://testnet.dkn.dria.co/discovery/available-nodes/1.0"
+        );
+    }
+}