added versioning to protocols

Author: erhant

Cargo.lock

@@ -1,6 +1,6 @@
 [package]
 name = "dkn-compute"
-version = "0.1.5"
+version = "0.1.6"
 edition = "2021"
 license = "Apache-2.0"
 readme = "README.md"
@@ -31,6 +31,7 @@ url = "2.5.0"
 urlencoding = "2.1.3"
 uuid = { version = "1.8.0", features = ["v4"] }
 rand = "0.8.5"
+semver = "1.0.23"
 
 # logging
 env_logger = "0.11.3"
@@ -64,11 +65,9 @@ libp2p = { git = "https://github.com/anilaltuner/rust-libp2p.git", rev = "be2ed5
     "quic",
     "kad",
 ] }
-
 libp2p-identity = { version = "0.2.9", features = ["secp256k1"] }
 tracing = { version = "0.1.40" }
 tracing-subscriber = { version = "0.3.18", features = ["env-filter"] }
-semver = "1.0.23"
 
 
 [dev-dependencies]

Cargo.toml

@@ -105,11 +105,20 @@ impl DriaComputeNodeConfig {
         }
     }
 
-    /// Check if the required compute services are running, e.g. if Ollama
-    /// is detected as a provider for the chosen models, it will check that
-    /// Ollama is running.
+    /// Check if the required compute services are running.
+    /// This has several steps:
+    ///
+    /// - If Ollama models are used, hardcoded models are checked locally, and for
+    ///   external models, the workflow is tested with a simple task with timeout.
+    /// - If OpenAI models are used, the API key is checked and the models are tested
+    ///
+    /// If both type of models are used, both services are checked.
+    /// In the end, bad models are filtered out and we simply check if we are left if any valid models at all.
+    /// If not, an error is returned.
     pub async fn check_services(&mut self) -> Result<(), String> {
         log::info!("Checking configured services.");
+
+        // TODO: can refactor (provider, model) logic here
         let unique_providers = self.model_config.get_providers();
 
         let mut good_models = Vec::new();
@@ -121,8 +130,10 @@ impl DriaComputeNodeConfig {
                 .get_models_for_provider(ModelProvider::Ollama);
 
             // ensure that the models are pulled / pull them if not
-            let timeout = Duration::from_secs(30);
-            let good_ollama_models = self.ollama_config.check(ollama_models, timeout).await?;
+            let good_ollama_models = self
+                .ollama_config
+                .check(ollama_models, Duration::from_secs(30))
+                .await?;
             good_models.extend(
                 good_ollama_models
                     .into_iter()

src/config/mod.rs

@@ -91,12 +91,12 @@ impl OllamaConfig {
         // check hardcoded models & pull them if available
         // these are not used directly by the user, but are needed for the workflows
         log::debug!("Checking hardcoded models: {:#?}", self.hardcoded_models);
+        // only check if model is contained in local_models
+        // we dont check workflows for hardcoded models
         for model in &self.hardcoded_models {
             if !local_models.contains(model) {
                 self.try_pull(&ollama, model.to_owned()).await?;
             }
-
-            // we dont check workflows for hardcoded models
         }
 
         // check external models & pull them if available
@@ -107,10 +107,10 @@ impl OllamaConfig {
                 self.try_pull(&ollama, model.to_string()).await?;
             }
 
-            let ok = self
+            if self
                 .test_workflow(model.clone(), test_workflow_timeout)
-                .await;
-            if ok {
+                .await
+            {
                 good_models.push(model);
             }
         }
@@ -207,11 +207,11 @@ impl OllamaConfig {
         let mut memory = ProgramMemory::new();
         tokio::select! {
             _ = tokio::time::sleep(timeout) => {
-                log::warn!("Ignoring model {}: Timeout", model);
+                log::warn!("Ignoring model {}: Workflow timed out", model);
             },
             result = executor.execute(None, workflow, &mut memory) => {
                 if result.is_empty() {
-                    log::warn!("Ignoring model {}: Empty Result", model);
+                    log::warn!("Ignoring model {}: Workflow returned empty result", model);
                 } else {
                     log::info!("Accepting model {}", model);
                     return true;

src/config/ollama.rs

@@ -115,6 +115,8 @@ fn create_gossipsub_behavior(author: PeerId) -> gossipsub::Behaviour {
         MessageId::from(hasher.finish().to_string())
     };
 
+    // TODO: add data transform here later
+
     Behaviour::new(
         MessageAuthenticity::Author(author),
         ConfigBuilder::default()

src/p2p/behaviour.rs

@@ -1,4 +1,4 @@
-use crate::p2p::{AvailableNodes, P2P_PROTOCOL_STRING};
+use crate::p2p::AvailableNodes;
 use crate::DRIA_COMPUTE_NODE_VERSION;
 use libp2p::futures::StreamExt;
 use libp2p::gossipsub::{
@@ -14,7 +14,11 @@ use semver::Version;
 use tokio::time::Duration;
 use tokio::time::Instant;
 
-use super::{DriaBehaviour, DriaBehaviourEvent, P2P_KADEMLIA_PROTOCOL};
+use super::{DriaBehaviour, DriaBehaviourEvent};
+
+/// Used as default for unparsable versions.
+/// Will not match with any valid version.
+const ZERO_VERSION: Version = Version::new(0, 0, 0);
 
 /// Underlying libp2p client.
 pub struct P2PClient {
@@ -247,67 +251,55 @@ impl P2PClient {
     ///
     /// - For Kademlia, we check the kademlia protocol and then add the address to the Kademlia routing table.
     fn handle_identify_event(&mut self, peer_id: PeerId, info: identify::Info) {
+        // we only care about the observed address, although there may be other addresses at `info.listen_addrs`
         let addr = info.observed_addr;
 
         // check protocol string
-        let (left, right) = info.protocol_version.split_at(5); // equals "dria/".len()
-        if left != "dria/" {
+        let protocol_ok = self.check_version_with_prefix(&info.protocol_version, "/dria/");
+        if !protocol_ok {
             log::warn!(
-                "Identify: Peer {} is from different protocol: (have {}, want {})",
+                "Identify: Peer {} has different Identify protocol: (have {}, want {})",
                 peer_id,
                 info.protocol_version,
-                P2P_PROTOCOL_STRING
+                self.version
             );
             return;
         }
 
-        // check version
-        match semver::Version::parse(right) {
-            Ok(peer_version) => {
-                if peer_version.minor != self.version.minor {
-                    log::warn!(
-                        "Identify: Peer {} has different version: (have {}, want {})",
-                        peer_id,
-                        peer_version,
-                        self.version
-                    );
-                    return;
-                }
-            }
-            Err(err) => {
-                log::error!(
-                    "Identify: Peer {} version could not be parsed: {}",
-                    peer_id,
-                    err,
-                );
-                return;
-            }
-        }
+        // check kademlia protocol
+        if let Some(kad_protocol) = info
+            .protocols
+            .iter()
+            .find(|p| p.to_string().starts_with("/dria/kad/"))
+        {
+            let protocol_ok =
+                self.check_version_with_prefix(&kad_protocol.to_string(), "/dria/kad/");
 
-        // we only care about the observed address, although there may be other addresses at `info.listen_addrs`
-        if info.protocols.iter().any(|p| *p == P2P_KADEMLIA_PROTOCOL) {
             // if it matches our protocol, add it to the Kademlia routing table
-            log::info!(
-                "Identify: {} peer {} identified at {}",
-                info.protocol_version,
-                peer_id,
-                addr
-            );
+            if protocol_ok {
+                log::info!(
+                    "Identify: {} peer {} identified at {}",
+                    kad_protocol,
+                    peer_id,
+                    addr
+                );
 
-            self.swarm
-                .behaviour_mut()
-                .kademlia
-                .add_address(&peer_id, addr);
-        } else {
-            log::trace!(
-                "Identify: Incoming from different protocol, address {}. PeerID is {}",
-                addr,
-                peer_id
-            );
+                self.swarm
+                    .behaviour_mut()
+                    .kademlia
+                    .add_address(&peer_id, addr);
+            } else {
+                log::warn!(
+                    "Identify: Peer {} has different Kademlia version: (have {}, want {})",
+                    peer_id,
+                    kad_protocol,
+                    self.version
+                );
+            }
         }
     }
 
-    /// Handles the results of a Kademlia closest peers search, either adding peers to Gossipsub or logging timeout errors.
+    /// Handles the results of a Kademlia closest peers search, simply logs it.
     fn handle_closest_peers_result(
         &mut self,
         result: Result<GetClosestPeersOk, GetClosestPeersError>,
@@ -373,4 +365,22 @@ impl P2PClient {
             }
         }
     }
+
+    /// Generic function to split a string such as `prefix || version` and check that the major & minor versions are the same.
+    ///    
+    /// Some examples:
+    /// - `self.check_version_with_prefix("dria/")` for identity
+    /// - `self.check_version_with_prefix("dria/kad/")` for Kademlia
+    ///
+    /// Returns whether the version is ok.
+    fn check_version_with_prefix(&self, p: &str, prefix: &str) -> bool {
+        let parsed_version = p
+            .strip_prefix(prefix)
+            .and_then(|v| Version::parse(v).ok())
+            .unwrap_or(ZERO_VERSION);
+
+        p.starts_with(prefix)
+            && parsed_version.major == self.version.major
+            && parsed_version.minor == self.version.minor
+    }
 }

src/p2p/client.rs

@@ -1,6 +1,11 @@
 use libp2p::StreamProtocol;
 
-pub(crate) const P2P_KADEMLIA_PROTOCOL: StreamProtocol = StreamProtocol::new("/dria/kad/1.0.0");
+/// Kademlia protocol version, in the form of `/dria/kad/<version>`.
+/// Notice the `/` at the start.
+pub(crate) const P2P_KADEMLIA_PROTOCOL: StreamProtocol =
+    StreamProtocol::new(concat!("/dria/kad/", env!("CARGO_PKG_VERSION")));
+
+/// Kademlia protocol version, in the form of `dria/<version>`.
 pub(crate) const P2P_PROTOCOL_STRING: &str = concat!("dria/", env!("CARGO_PKG_VERSION"));
 
 mod behaviour;