more eyre'ing, slight folder restructures

Author: erhant

src/config/mod.rs

@@ -3,6 +3,7 @@ mod ollama;
 mod openai;
 
 use crate::utils::crypto::to_address;
+use eyre::{eyre, Result};
 use libsecp256k1::{PublicKey, SecretKey};
 use models::ModelConfig;
 use ollama::OllamaConfig;
@@ -129,7 +130,7 @@ impl DriaComputeNodeConfig {
     /// 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> {
+    pub async fn check_services(&mut self) -> Result<()> {
         log::info!("Checking configured services.");
 
         // TODO: can refactor (provider, model) logic here
@@ -171,7 +172,7 @@ impl DriaComputeNodeConfig {
 
         // update good models
         if good_models.is_empty() {
-            Err("No good models found, please check logs for errors.".into())
+            Err(eyre!("No good models found, please check logs for errors."))
         } else {
             self.model_config.models = good_models;
             Ok(())

src/config/ollama.rs

@@ -1,5 +1,6 @@
 use std::time::Duration;
 
+use eyre::{eyre, Result};
 use ollama_workflows::{
     ollama_rs::{
         generation::{
@@ -81,7 +82,7 @@ impl OllamaConfig {
         external_models: Vec<Model>,
         timeout: Duration,
         min_tps: f64,
-    ) -> Result<Vec<Model>, String> {
+    ) -> Result<Vec<Model>> {
         log::info!(
             "Checking Ollama requirements (auto-pull {}, workflow timeout: {}s)",
             if self.auto_pull { "on" } else { "off" },
@@ -96,7 +97,7 @@ impl OllamaConfig {
             Err(e) => {
                 return {
                     log::error!("Could not fetch local models from Ollama, is it online?");
-                    Err(e.to_string())
+                    Err(e.into())
                 }
             }
         };
@@ -137,25 +138,22 @@ impl OllamaConfig {
     }
 
     /// Pulls a model if `auto_pull` exists, otherwise returns an error.
-    async fn try_pull(&self, ollama: &Ollama, model: String) -> Result<(), String> {
+    async fn try_pull(&self, ollama: &Ollama, model: String) -> Result<()> {
         log::warn!("Model {} not found in Ollama", model);
         if self.auto_pull {
             // if auto-pull is enabled, pull the model
             log::info!(
                 "Downloading missing model {} (this may take a while)",
                 model
             );
-            let status = ollama
-                .pull_model(model, false)
-                .await
-                .map_err(|e| format!("Error pulling model with Ollama: {}", e))?;
+            let status = ollama.pull_model(model, false).await?;
             log::debug!("Pulled model with Ollama, final status: {:#?}", status);
             Ok(())
         } else {
             // otherwise, give error
             log::error!("Please download missing model with: ollama pull {}", model);
             log::error!("Or, set OLLAMA_AUTO_PULL=true to pull automatically.");
-            Err("Required model not pulled in Ollama.".into())
+            Err(eyre!("Required model not pulled in Ollama."))
         }
     }
 

src/config/openai.rs

@@ -1,5 +1,6 @@
 #![allow(unused)]
 
+use eyre::{eyre, Context, Result};
 use ollama_workflows::Model;
 use serde::Deserialize;
 
@@ -39,15 +40,13 @@ impl OpenAIConfig {
         Self { api_key }
     }
 
-    /// Check if requested models exist &
-    ///
-    ///
-    pub async fn check(&self, models: Vec<Model>) -> Result<Vec<Model>, String> {
+    /// Check if requested models exist & are available in the OpenAI account.
+    pub async fn check(&self, models: Vec<Model>) -> Result<Vec<Model>> {
         log::info!("Checking OpenAI requirements");
 
         // check API key
         let Some(api_key) = &self.api_key else {
-            return Err("OpenAI API key not found".into());
+            return Err(eyre!("OpenAI API key not found"));
         };
 
         // fetch models
@@ -56,24 +55,21 @@ impl OpenAIConfig {
             .get(OPENAI_MODELS_API)
             .header("Authorization", format!("Bearer {}", api_key))
             .build()
-            .map_err(|e| format!("Failed to build request: {}", e))?;
+            .wrap_err("Failed to build request")?;
 
         let response = client
             .execute(request)
             .await
-            .map_err(|e| format!("Failed to send request: {}", e))?;
+            .wrap_err("Failed to send request")?;
 
         // parse response
         if response.status().is_client_error() {
-            return Err(format!(
+            return Err(eyre!(
                 "Failed to fetch OpenAI models:\n{}",
                 response.text().await.unwrap_or_default()
             ));
         }
-        let openai_models = response
-            .json::<OpenAIModelsResponse>()
-            .await
-            .map_err(|e| e.to_string())?;
+        let openai_models = response.json::<OpenAIModelsResponse>().await?;
 
         // check if models exist and select those that are available
         let mut available_models = Vec::new();

src/handlers/mod.rs

@@ -1,4 +1,5 @@
 use async_trait::async_trait;
+use eyre::Result;
 use libp2p::gossipsub::MessageAcceptance;
 
 mod pingpong;
@@ -7,13 +8,13 @@ pub use pingpong::PingpongHandler;
 mod workflow;
 pub use workflow::WorkflowHandler;
 
-use crate::{p2p::P2PMessage, DriaComputeNode};
+use crate::{utils::P2PMessage, DriaComputeNode};
 
 #[async_trait]
 pub trait ComputeHandler {
     async fn handle_compute(
         node: &mut DriaComputeNode,
         message: P2PMessage,
         result_topic: &str,
-    ) -> eyre::Result<MessageAcceptance>;
+    ) -> Result<MessageAcceptance>;
 }

src/handlers/pingpong.rs

@@ -1,4 +1,7 @@
-use crate::{node::DriaComputeNode, p2p::P2PMessage, utils::get_current_time_nanos};
+use crate::{
+    node::DriaComputeNode,
+    utils::{get_current_time_nanos, P2PMessage},
+};
 use async_trait::async_trait;
 use eyre::Result;
 use libp2p::gossipsub::MessageAcceptance;
@@ -66,10 +69,10 @@ impl ComputeHandler for PingpongHandler {
 #[cfg(test)]
 mod tests {
     use crate::{
-        p2p::P2PMessage,
         utils::{
             crypto::{sha256hash, to_address},
             filter::FilterPayload,
+            P2PMessage,
         },
         DriaComputeNodeConfig,
     };

src/handlers/workflow.rs

@@ -5,9 +5,8 @@ use ollama_workflows::{Entry, Executor, ModelProvider, ProgramMemory, Workflow};
 use serde::Deserialize;
 
 use crate::node::DriaComputeNode;
-use crate::p2p::P2PMessage;
-use crate::utils::get_current_time_nanos;
 use crate::utils::payload::{TaskRequest, TaskRequestPayload};
+use crate::utils::{get_current_time_nanos, P2PMessage};
 
 use super::ComputeHandler;
 

src/main.rs

@@ -1,8 +1,9 @@
 use dkn_compute::{DriaComputeNode, DriaComputeNodeConfig};
+use eyre::Result;
 use tokio_util::sync::CancellationToken;
 
 #[tokio::main]
-async fn main() -> Result<(), Box<dyn std::error::Error>> {
+async fn main() -> Result<()> {
     if let Err(e) = dotenvy::dotenv() {
         log::warn!("Could not load .env file: {}", e);
     }

src/node.rs

@@ -6,8 +6,8 @@ use tokio_util::sync::CancellationToken;
 use crate::{
     config::DriaComputeNodeConfig,
     handlers::{ComputeHandler, PingpongHandler, WorkflowHandler},
-    p2p::{P2PClient, P2PMessage},
-    utils::{crypto::secret_to_keypair, AvailableNodes},
+    p2p::P2PClient,
+    utils::{crypto::secret_to_keypair, AvailableNodes, P2PMessage},
 };
 
 /// Number of seconds between refreshing the Admin RPC PeerIDs from Dria server.
@@ -38,10 +38,9 @@ impl DriaComputeNode {
     pub async fn new(
         config: DriaComputeNodeConfig,
         cancellation: CancellationToken,
-    ) -> Result<Self, String> {
+    ) -> Result<Self> {
         let keypair = secret_to_keypair(&config.secret_key);
-        let listen_addr =
-            Multiaddr::from_str(config.p2p_listen_addr.as_str()).map_err(|e| e.to_string())?;
+        let listen_addr = Multiaddr::from_str(config.p2p_listen_addr.as_str())?;
 
         // get available nodes (bootstrap, relay, rpc) for p2p
         let available_nodes = AvailableNodes::default()
@@ -266,9 +265,8 @@ impl DriaComputeNode {
 
 #[cfg(test)]
 mod tests {
-    use crate::{p2p::P2PMessage, DriaComputeNode, DriaComputeNodeConfig};
+    use super::*;
     use std::env;
-    use tokio_util::sync::CancellationToken;
 
     #[tokio::test]
     #[ignore = "run this manually"]

src/p2p/client.rs

@@ -1,3 +1,4 @@
+use eyre::Result;
 use libp2p::futures::StreamExt;
 use libp2p::gossipsub::{
     Message, MessageAcceptance, MessageId, PublishError, SubscriptionError, TopicHash,
@@ -17,6 +18,7 @@ use super::{DriaBehaviour, DriaBehaviourEvent, P2P_KADEMLIA_PROTOCOL, P2P_PROTOC
 
 /// Underlying libp2p client.
 pub struct P2PClient {
+    /// `Swarm` instance, everything is accesses through this one.
     swarm: Swarm<DriaBehaviour>,
     /// Peer count for (All, Mesh).
     peer_count: (usize, usize),
@@ -36,7 +38,7 @@ impl P2PClient {
         keypair: Keypair,
         listen_addr: Multiaddr,
         available_nodes: &AvailableNodes,
-    ) -> Result<Self, String> {
+    ) -> Result<Self> {
         // this is our peerId
         let node_peerid = keypair.public().to_peer_id();
         log::info!("Compute node peer address: {}", node_peerid);
@@ -47,13 +49,10 @@ impl P2PClient {
                 tcp::Config::default(),
                 noise::Config::new,
                 yamux::Config::default,
-            )
-            .map_err(|e| e.to_string())?
+            )?
             .with_quic()
-            .with_relay_client(noise::Config::new, yamux::Config::default)
-            .map_err(|e| e.to_string())?
-            .with_behaviour(|key, relay_behavior| Ok(DriaBehaviour::new(key, relay_behavior)))
-            .map_err(|e| e.to_string())?
+            .with_relay_client(noise::Config::new, yamux::Config::default)?
+            .with_behaviour(|key, relay_behavior| Ok(DriaBehaviour::new(key, relay_behavior)))?
             .with_swarm_config(|c| {
                 c.with_idle_connection_timeout(Duration::from_secs(IDLE_CONNECTION_TIMEOUT_SECS))
             })
@@ -76,7 +75,7 @@ impl P2PClient {
                 _ => None,
             }) {
                 log::info!("Dialling peer: {}", addr);
-                swarm.dial(addr.clone()).map_err(|e| e.to_string())?;
+                swarm.dial(addr.clone())?;
                 log::info!("Adding {} to Kademlia routing table", addr);
                 swarm
                     .behaviour_mut()
@@ -94,24 +93,18 @@ impl P2PClient {
             .behaviour_mut()
             .kademlia
             .get_closest_peers(random_peer);
-        swarm
-            .behaviour_mut()
-            .kademlia
-            .bootstrap()
-            .map_err(|e| e.to_string())?;
+        swarm.behaviour_mut().kademlia.bootstrap()?;
 
         // listen on all interfaces for incoming connections
         log::info!("Listening p2p network on: {}", listen_addr);
-        swarm.listen_on(listen_addr).map_err(|e| e.to_string())?;
+        swarm.listen_on(listen_addr)?;
 
         log::info!(
             "Listening to relay nodes: {:#?}",
             available_nodes.relay_nodes
         );
         for addr in &available_nodes.relay_nodes {
-            swarm
-                .listen_on(addr.clone().with(Protocol::P2pCircuit))
-                .map_err(|e| e.to_string())?;
+            swarm.listen_on(addr.clone().with(Protocol::P2pCircuit))?;
         }
 
         Ok(Self {
@@ -138,6 +131,8 @@ impl P2PClient {
     }
 
     /// Publish a message to a topic.
+    ///
+    /// Returns the message ID.
     pub fn publish(
         &mut self,
         topic_name: &str,
@@ -168,7 +163,7 @@ impl P2PClient {
         msg_id: &MessageId,
         propagation_source: &PeerId,
         acceptance: MessageAcceptance,
-    ) -> Result<(), PublishError> {
+    ) -> Result<()> {
         log::trace!("Validating message ({}): {:?}", msg_id, acceptance);
 
         let msg_was_in_cache = self
@@ -240,6 +235,7 @@ 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`
+        // TODO: this may be wrong
         let addr = info.observed_addr;
 
         // check protocol string

src/p2p/mod.rs

@@ -37,8 +37,4 @@ pub use behaviour::{DriaBehaviour, DriaBehaviourEvent};
 mod client;
 pub use client::P2PClient;
 
-mod message;
-
-pub use message::P2PMessage;
-
 mod data_transform;