much better eyre error printing

erhant · erhant · commit 2030b337b68d · 2024-10-02T18:21:31.000+03:00
diff --git a/Cargo.toml b/Cargo.toml
@@ -4,6 +4,7 @@ version = "0.2.6"
 edition = "2021"
 license = "Apache-2.0"
 readme = "README.md"
+authors = ["Erhan Tezcan <erhan@firstbatch.xyz>"]
 
 # profiling build for flamegraphs
 [profile.profiling]
diff --git a/src/handlers/mod.rs b/src/handlers/mod.rs
@@ -1,25 +1,30 @@
+use crate::{utils::DKNMessage, DriaComputeNode};
 use async_trait::async_trait;
 use eyre::Result;
 use libp2p::gossipsub::MessageAcceptance;
 
-mod topics;
-pub use topics::*;
-
 mod pingpong;
 pub use pingpong::PingpongHandler;
 
 mod workflow;
 pub use workflow::WorkflowHandler;
 
-use crate::{utils::DKNMessage, DriaComputeNode};
-
 /// A DKN task is to be handled by the compute node, respecting this trait.
+///
+/// It is expected for the implemented handler to handle messages coming from `LISTEN_TOPIC`,
+/// and then respond back to the `RESPONSE_TOPIC`.
 #[async_trait]
 pub trait ComputeHandler {
+    /// Gossipsub topic name to listen for incoming messages from the network.
+    const LISTEN_TOPIC: &'static str;
+    /// Gossipsub topic name to respond with messages to the network.
+    const RESPONSE_TOPIC: &'static str;
+
     /// A generic handler for DKN tasks.
+    ///
+    /// Returns a `MessageAcceptance` value that tells the P2P client to accept the incoming message.
     async fn handle_compute(
         node: &mut DriaComputeNode,
         message: DKNMessage,
-        result_topic: &str,
     ) -> Result<MessageAcceptance>;
 }
diff --git a/src/handlers/pingpong.rs b/src/handlers/pingpong.rs
@@ -27,10 +27,12 @@ struct PingpongResponse {
 
 #[async_trait]
 impl ComputeHandler for PingpongHandler {
+    const LISTEN_TOPIC: &'static str = "ping";
+    const RESPONSE_TOPIC: &'static str = "pong";
+
     async fn handle_compute(
         node: &mut DriaComputeNode,
         message: DKNMessage,
-        result_topic: &str,
     ) -> Result<MessageAcceptance> {
         let pingpong = message.parse_payload::<PingpongPayload>(true)?;
 
@@ -58,7 +60,7 @@ impl ComputeHandler for PingpongHandler {
         // publish message
         let message = DKNMessage::new_signed(
             serde_json::json!(response_body).to_string(),
-            result_topic,
+            Self::RESPONSE_TOPIC,
             &node.config.secret_key,
         );
         node.publish(message)?;
diff --git a/src/handlers/topics.rs b/src/handlers/topics.rs
diff --git a/src/handlers/workflow.rs b/src/handlers/workflow.rs
@@ -1,5 +1,5 @@
 use async_trait::async_trait;
-use eyre::{eyre, Result};
+use eyre::{eyre, Context, Result};
 use libp2p::gossipsub::MessageAcceptance;
 use ollama_workflows::{Entry, Executor, ModelProvider, ProgramMemory, Workflow};
 use serde::Deserialize;
@@ -27,13 +27,17 @@ struct WorkflowPayload {
 
 #[async_trait]
 impl ComputeHandler for WorkflowHandler {
+    const LISTEN_TOPIC: &'static str = "task";
+    const RESPONSE_TOPIC: &'static str = "results";
+
     async fn handle_compute(
         node: &mut DriaComputeNode,
         message: DKNMessage,
-        result_topic: &str,
     ) -> Result<MessageAcceptance> {
         let config = &node.config;
-        let task = message.parse_payload::<TaskRequestPayload<WorkflowPayload>>(true)?;
+        let task = message
+            .parse_payload::<TaskRequestPayload<WorkflowPayload>>(true)
+            .wrap_err("Could not parse error")?;
 
         // check if deadline is past or not
         let current_time = get_current_time_nanos();
@@ -56,7 +60,7 @@ impl ComputeHandler for WorkflowHandler {
                 task.task_id
             );
 
-            // accept the message, someonelse may be included in filter
+            // accept the message, someone else may be included in filter
             return Ok(MessageAcceptance::Accept);
         }
 
@@ -86,14 +90,15 @@ impl ComputeHandler for WorkflowHandler {
                 return Ok(MessageAcceptance::Accept);
             },
             exec_result_inner = executor.execute(entry.as_ref(), task.input.workflow, &mut memory) => {
-                exec_result = exec_result_inner.map_err(|e| eyre!("{}", e.to_string()));
+                exec_result = exec_result_inner.map_err(|e| eyre!("Execution error: {}", e.to_string()));
             }
         }
 
         match exec_result {
             Ok(result) => {
                 // obtain public key from the payload
-                let task_public_key = hex::decode(&task.public_key)?;
+                let task_public_key =
+                    hex::decode(&task.public_key).wrap_err("Could not decode public key")?;
 
                 // prepare signed and encrypted payload
                 let payload = TaskResponsePayload::new(
@@ -102,24 +107,27 @@ impl ComputeHandler for WorkflowHandler {
                     &task_public_key,
                     &config.secret_key,
                 )?;
-                let payload_str = serde_json::to_string(&payload)?;
+                let payload_str =
+                    serde_json::to_string(&payload).wrap_err("Could not serialize payload")?;
 
                 // publish the result
-                let message = DKNMessage::new(payload_str, result_topic);
+                let message = DKNMessage::new(payload_str, Self::RESPONSE_TOPIC);
                 node.publish(message)?;
 
                 // accept so that if there are others included in filter they can do the task
                 Ok(MessageAcceptance::Accept)
             }
             Err(err) => {
-                log::error!("Task {} failed: {}", task.task_id, err);
+                // use pretty display string for error logging with causes
+                let err_string = format!("{:#}", err);
+                log::error!("Task {} failed: {}", task.task_id, err_string);
 
                 // prepare error payload
-                let error_payload = TaskErrorPayload::new(task.task_id, err.to_string());
+                let error_payload = TaskErrorPayload::new(task.task_id, err_string);
                 let error_payload_str = serde_json::to_string(&error_payload)?;
 
                 // publish the error result for diagnostics
-                let message = DKNMessage::new(error_payload_str, result_topic);
+                let message = DKNMessage::new(error_payload_str, Self::RESPONSE_TOPIC);
                 node.publish(message)?;
 
                 // ignore just in case, workflow may be bugged
diff --git a/src/main.rs b/src/main.rs
@@ -38,7 +38,7 @@ async fn main() -> Result<()> {
 
         #[cfg(not(feature = "profiling"))]
         if let Err(err) = wait_for_termination(cancellation_token.clone()).await {
-            log::error!("Error waiting for termination: {}", err);
+            log::error!("Error waiting for termination: {:?}", err);
             log::error!("Cancelling due to unexpected error.");
             cancellation_token.cancel();
         };
@@ -55,7 +55,7 @@ async fn main() -> Result<()> {
             }
             result = config_clone.check_services() => {
                 if let Err(err) = result {
-                    log::error!("Error checking services: {}", err);
+                    log::error!("Error checking services: {:?}", err);
                     panic!("Service check failed.")
                 }
             }
@@ -101,7 +101,7 @@ async fn main() -> Result<()> {
 ///
 /// Handles Unix and Windows [target families](https://doc.rust-lang.org/reference/conditional-compilation.html#target_family).
 #[allow(unused)]
-async fn wait_for_termination(cancellation: CancellationToken) -> std::io::Result<()> {
+async fn wait_for_termination(cancellation: CancellationToken) -> Result<()> {
     #[cfg(unix)]
     {
         use tokio::signal::unix::{signal, SignalKind};
diff --git a/src/node.rs b/src/node.rs
@@ -4,7 +4,7 @@ use std::{str::FromStr, time::Duration};
 use tokio_util::sync::CancellationToken;
 
 use crate::{
-    config::DriaComputeNodeConfig,
+    config::*,
     handlers::*,
     p2p::P2PClient,
     utils::{crypto::secret_to_keypair, AvailableNodes, DKNMessage},
@@ -13,6 +13,19 @@ use crate::{
 /// Number of seconds between refreshing the Admin RPC PeerIDs from Dria server.
 const RPC_PEER_ID_REFRESH_INTERVAL_SECS: u64 = 30;
 
+/// **Dria Compute Node**
+///
+/// Internally, the node will create a new P2P client with the given secret key.
+/// This P2P client, although created synchronously, requires a tokio runtime.
+///
+/// ### Example
+///
+/// ```rs
+/// let config = DriaComputeNodeConfig::new();
+/// let mut node = DriaComputeNode::new(config, CancellationToken::new())?;
+/// node.check_services().await?;
+/// node.launch().await?;
+/// ```
 pub struct DriaComputeNode {
     pub config: DriaComputeNodeConfig,
     pub p2p: P2PClient,
@@ -22,19 +35,6 @@ pub struct DriaComputeNode {
 }
 
 impl DriaComputeNode {
-    /// Create a new compute node with the given configuration and cancellation token.
-    ///
-    /// Internally, the node will create a new P2P client with the given secret key.
-    /// This P2P client, although created synchronously, requires a tokio runtime.
-    ///
-    /// ### Example
-    ///
-    /// ```rs
-    /// let config = DriaComputeNodeConfig::new();
-    /// let mut node = DriaComputeNode::new(config, CancellationToken::new())?;
-    /// node.check_services().await?;
-    /// node.launch().await?;
-    /// ```
     pub async fn new(
         config: DriaComputeNodeConfig,
         cancellation: CancellationToken,
@@ -106,10 +106,10 @@ impl DriaComputeNode {
     /// This method is not expected to return until cancellation occurs.
     pub async fn launch(&mut self) -> Result<()> {
         // subscribe to topics
-        self.subscribe(PINGPONG_LISTEN_TOPIC)?;
-        self.subscribe(PINGPONG_RESPONSE_TOPIC)?;
-        self.subscribe(WORKFLOW_LISTEN_TOPIC)?;
-        self.subscribe(WORKFLOW_RESPONSE_TOPIC)?;
+        self.subscribe(PingpongHandler::LISTEN_TOPIC)?;
+        self.subscribe(PingpongHandler::RESPONSE_TOPIC)?;
+        self.subscribe(WorkflowHandler::LISTEN_TOPIC)?;
+        self.subscribe(WorkflowHandler::RESPONSE_TOPIC)?;
 
         // main loop, listens for message events in particular
         // the underlying p2p client is expected to handle the rest within its own loop
@@ -127,7 +127,7 @@ impl DriaComputeNode {
                     let topic_str = topic.as_str();
 
                     // handle message w.r.t topic
-                    if std::matches!(topic_str, PINGPONG_LISTEN_TOPIC | WORKFLOW_LISTEN_TOPIC) {
+                    if std::matches!(topic_str, PingpongHandler::LISTEN_TOPIC | WorkflowHandler::LISTEN_TOPIC) {
                         // ensure that the message is from a valid source (origin)
                         let source_peer_id = match message.source {
                             Some(peer) => peer,
@@ -159,7 +159,7 @@ impl DriaComputeNode {
                         let message = match self.parse_message_to_prepared_message(message.clone()) {
                             Ok(message) => message,
                             Err(e) => {
-                                log::error!("Error parsing message: {}", e);
+                                log::error!("Error parsing message: {:?}", e);
                                 log::debug!("Message: {}", String::from_utf8_lossy(&message.data));
                                 self.p2p.validate_message(&message_id, &peer_id, gossipsub::MessageAcceptance::Ignore)?;
                                 continue;
@@ -168,11 +168,11 @@ impl DriaComputeNode {
 
                         // then handle the prepared message
                         let handler_result = match topic_str {
-                            WORKFLOW_LISTEN_TOPIC => {
-                                WorkflowHandler::handle_compute(self, message, WORKFLOW_RESPONSE_TOPIC).await
+                            WorkflowHandler::LISTEN_TOPIC => {
+                                WorkflowHandler::handle_compute(self, message).await
                             }
-                            PINGPONG_LISTEN_TOPIC => {
-                                PingpongHandler::handle_compute(self, message, PINGPONG_RESPONSE_TOPIC).await
+                            PingpongHandler::LISTEN_TOPIC => {
+                                PingpongHandler::handle_compute(self, message).await
                             }
                             // TODO: can we do this in a nicer way?
                             // TODO: yes, cast to enum above and let type-casting do the work
@@ -185,11 +185,11 @@ impl DriaComputeNode {
                                 self.p2p.validate_message(&message_id, &peer_id, acceptance)?;
                             },
                             Err(err) => {
-                                log::error!("Error handling {} message: {}", topic_str, err);
+                                log::error!("Error handling {} message: {:?}", topic_str, err);
                                 self.p2p.validate_message(&message_id, &peer_id, gossipsub::MessageAcceptance::Ignore)?;
                             }
                         }
-                    } else if std::matches!(topic_str, PINGPONG_RESPONSE_TOPIC | WORKFLOW_RESPONSE_TOPIC) {
+                    } else if std::matches!(topic_str, PingpongHandler::RESPONSE_TOPIC | WorkflowHandler::RESPONSE_TOPIC) {
                         // since we are responding to these topics, we might receive messages from other compute nodes
                         // we can gracefully ignore them and propagate it to to others
                         log::debug!("Ignoring message for topic: {}", topic_str);
@@ -205,10 +205,10 @@ impl DriaComputeNode {
         }
 
         // unsubscribe from topics
-        self.unsubscribe(PINGPONG_LISTEN_TOPIC)?;
-        self.unsubscribe(PINGPONG_RESPONSE_TOPIC)?;
-        self.unsubscribe(WORKFLOW_LISTEN_TOPIC)?;
-        self.unsubscribe(WORKFLOW_RESPONSE_TOPIC)?;
+        self.unsubscribe(PingpongHandler::LISTEN_TOPIC)?;
+        self.unsubscribe(PingpongHandler::RESPONSE_TOPIC)?;
+        self.unsubscribe(WorkflowHandler::LISTEN_TOPIC)?;
+        self.unsubscribe(WorkflowHandler::RESPONSE_TOPIC)?;
 
         Ok(())
     }
diff --git a/src/payloads/error.rs b/src/payloads/error.rs
@@ -1,6 +1,7 @@
 use serde::{Deserialize, Serialize};
 
-/// A generic task request, given by Dria.
+/// A task error response.
+/// Returning this as the payload helps to debug the errors received at client side.
 #[derive(Debug, Clone, Serialize, Deserialize)]
 #[serde(rename_all = "camelCase")]
 pub struct TaskErrorPayload {

Original file line number	Diff line number	Diff line change
`@@ -38,7 +38,7 @@ async fn main() -> Result<()> {`
`38`	`38`
`39`	`39`	`#[cfg(not(feature = "profiling"))]`
`40`	`40`	`if let Err(err) = wait_for_termination(cancellation_token.clone()).await {`
`41`		`- log::error!("Error waiting for termination: {}", err);`
	`41`	`+ log::error!("Error waiting for termination: {:?}", err);`
`42`	`42`	`log::error!("Cancelling due to unexpected error.");`
`43`	`43`	`cancellation_token.cancel();`
`44`	`44`	`};`
`@@ -55,7 +55,7 @@ async fn main() -> Result<()> {`
`55`	`55`	`}`
`56`	`56`	`result = config_clone.check_services() => {`
`57`	`57`	`if let Err(err) = result {`
`58`		`- log::error!("Error checking services: {}", err);`
	`58`	`+ log::error!("Error checking services: {:?}", err);`
`59`	`59`	`panic!("Service check failed.")`
`60`	`60`	`}`
`61`	`61`	`}`
`@@ -101,7 +101,7 @@ async fn main() -> Result<()> {`
`101`	`101`	`///`
`102`	`102`	`/// Handles Unix and Windows [target families](https://doc.rust-lang.org/reference/conditional-compilation.html#target_family).`
`103`	`103`	`#[allow(unused)]`
`104`		`-async fn wait_for_termination(cancellation: CancellationToken) -> std::io::Result<()> {`
	`104`	`+async fn wait_for_termination(cancellation: CancellationToken) -> Result<()> {`
`105`	`105`	`#[cfg(unix)]`
`106`	`106`	`{`
`107`	`107`	`use tokio::signal::unix::{signal, SignalKind};`