added error reporting

Author: erhant

src/handlers/mod.rs

@@ -2,6 +2,9 @@ use async_trait::async_trait;
 use eyre::Result;
 use libp2p::gossipsub::MessageAcceptance;
 
+mod topics;
+pub use topics::*;
+
 mod pingpong;
 pub use pingpong::PingpongHandler;
 
@@ -10,8 +13,10 @@ pub use workflow::WorkflowHandler;
 
 use crate::{utils::DKNMessage, DriaComputeNode};
 
+/// A DKN task is to be handled by the compute node, respecting this trait.
 #[async_trait]
 pub trait ComputeHandler {
+    /// A generic handler for DKN tasks.
     async fn handle_compute(
         node: &mut DriaComputeNode,
         message: DKNMessage,

src/handlers/pingpong.rs

@@ -1,6 +1,6 @@
 use crate::{
-    node::DriaComputeNode,
     utils::{get_current_time_nanos, DKNMessage},
+    DriaComputeNode,
 };
 use async_trait::async_trait;
 use eyre::Result;
@@ -55,14 +55,14 @@ impl ComputeHandler for PingpongHandler {
             timestamp: get_current_time_nanos(),
         };
 
+        // publish message
         let message = DKNMessage::new_signed(
             serde_json::json!(response_body).to_string(),
             result_topic,
             &node.config.secret_key,
         );
         node.publish(message)?;
 
-        // accept message, someone else may be included in the filter
         Ok(MessageAcceptance::Accept)
     }
 }

src/handlers/topics.rs

@@ -0,0 +1,4 @@
+pub const PINGPONG_LISTEN_TOPIC: &str = "ping";
+pub const PINGPONG_RESPONSE_TOPIC: &str = "pong";
+pub const WORKFLOW_LISTEN_TOPIC: &str = "task";
+pub const WORKFLOW_RESPONSE_TOPIC: &str = "results";

src/handlers/workflow.rs

@@ -4,9 +4,9 @@ use libp2p::gossipsub::MessageAcceptance;
 use ollama_workflows::{Entry, Executor, ModelProvider, ProgramMemory, Workflow};
 use serde::Deserialize;
 
-use crate::node::DriaComputeNode;
-use crate::utils::payload::{TaskRequestPayload, TaskResponsePayload};
+use crate::payloads::{TaskErrorPayload, TaskRequestPayload, TaskResponsePayload};
 use crate::utils::{get_current_time_nanos, DKNMessage};
+use crate::DriaComputeNode;
 
 use super::ComputeHandler;
 
@@ -32,6 +32,7 @@ impl ComputeHandler for WorkflowHandler {
         message: DKNMessage,
         result_topic: &str,
     ) -> Result<MessageAcceptance> {
+        let config = &node.config;
         let task = message.parse_payload::<TaskRequestPayload<WorkflowPayload>>(true)?;
 
         // check if deadline is past or not
@@ -49,7 +50,7 @@ impl ComputeHandler for WorkflowHandler {
         }
 
         // check task inclusion via the bloom filter
-        if !task.filter.contains(&node.config.address)? {
+        if !task.filter.contains(&config.address)? {
             log::info!(
                 "Task {} does not include this node within the filter.",
                 task.task_id
@@ -59,23 +60,15 @@ impl ComputeHandler for WorkflowHandler {
             return Ok(MessageAcceptance::Accept);
         }
 
-        // obtain public key from the payload
-        let task_public_key = hex::decode(&task.public_key)?;
-
         // read model / provider from the task
-        let (model_provider, model) = node
-            .config
+        let (model_provider, model) = config
             .model_config
             .get_any_matching_model(task.input.model)?;
         log::info!("Using model {} for task {}", model, task.task_id);
 
         // prepare workflow executor
         let executor = if model_provider == ModelProvider::Ollama {
-            Executor::new_at(
-                model,
-                &node.config.ollama_config.host,
-                node.config.ollama_config.port,
-            )
+            Executor::new_at(model, &config.ollama_config.host, config.ollama_config.port)
         } else {
             Executor::new(model)
         };
@@ -86,38 +79,52 @@ impl ComputeHandler for WorkflowHandler {
             .map(|prompt| Entry::try_value_or_str(&prompt));
 
         // execute workflow with cancellation
-        // TODO: is there a better way to handle this?
-        let result: String;
+        let exec_result: Result<String>;
         tokio::select! {
             _ = node.cancellation.cancelled() => {
                 log::info!("Received cancellation, quitting all tasks.");
-                return Ok(MessageAcceptance::Accept)
+                return Ok(MessageAcceptance::Accept);
             },
-            exec_result = executor.execute(entry.as_ref(), task.input.workflow, &mut memory) => {
-                match exec_result {
-                    Ok(exec_result) => {
-                        result = exec_result;
-                    }
-                    Err(e) => {
-                        return Err(eyre!("Workflow failed with error {}", e));
-                    }
-                }
+            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()));
             }
         }
 
-        // prepare signed and encrypted payload
-        let payload = TaskResponsePayload::new(
-            result,
-            &task.task_id,
-            &task_public_key,
-            &node.config.secret_key,
-        )?;
-        let payload_str = payload.to_string()?;
+        match exec_result {
+            Ok(result) => {
+                // obtain public key from the payload
+                let task_public_key = hex::decode(&task.public_key)?;
+
+                // prepare signed and encrypted payload
+                let payload = TaskResponsePayload::new(
+                    result,
+                    &task.task_id,
+                    &task_public_key,
+                    &config.secret_key,
+                )?;
+                let payload_str = serde_json::to_string(&payload)?;
+
+                // publish the result
+                let message = DKNMessage::new(payload_str, result_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);
+
+                // prepare error payload
+                let error_payload = TaskErrorPayload::new(task.task_id, err.to_string());
+                let error_payload_str = serde_json::to_string(&error_payload)?;
 
-        // publish the result
-        let message = DKNMessage::new(payload_str, result_topic);
-        node.publish(message)?;
+                // publish the error result for diagnostics
+                let message = DKNMessage::new(error_payload_str, result_topic);
+                node.publish(message)?;
 
-        Ok(MessageAcceptance::Accept)
+                // ignore just in case, workflow may be bugged
+                Ok(MessageAcceptance::Ignore)
+            }
+        }
     }
 }

src/lib.rs

@@ -5,6 +5,7 @@ pub(crate) mod config;
 pub(crate) mod handlers;
 pub(crate) mod node;
 pub(crate) mod p2p;
+pub(crate) mod payloads;
 pub(crate) mod utils;
 
 /// Crate version of the compute node.

src/node.rs

@@ -5,7 +5,7 @@ use tokio_util::sync::CancellationToken;
 
 use crate::{
     config::DriaComputeNodeConfig,
-    handlers::{ComputeHandler, PingpongHandler, WorkflowHandler},
+    handlers::*,
     p2p::P2PClient,
     utils::{crypto::secret_to_keypair, AvailableNodes, DKNMessage},
 };
@@ -105,11 +105,6 @@ impl DriaComputeNode {
     /// Launches the main loop of the compute node.
     /// This method is not expected to return until cancellation occurs.
     pub async fn launch(&mut self) -> Result<()> {
-        const PINGPONG_LISTEN_TOPIC: &str = "ping";
-        const PINGPONG_RESPONSE_TOPIC: &str = "pong";
-        const WORKFLOW_LISTEN_TOPIC: &str = "task";
-        const WORKFLOW_RESPONSE_TOPIC: &str = "results";
-
         // subscribe to topics
         self.subscribe(PINGPONG_LISTEN_TOPIC)?;
         self.subscribe(PINGPONG_RESPONSE_TOPIC)?;
@@ -172,7 +167,7 @@ impl DriaComputeNode {
                         };
 
                         // then handle the prepared message
-                        let handle_result = match topic_str {
+                        let handler_result = match topic_str {
                             WORKFLOW_LISTEN_TOPIC => {
                                 WorkflowHandler::handle_compute(self, message, WORKFLOW_RESPONSE_TOPIC).await
                             }
@@ -185,26 +180,23 @@ impl DriaComputeNode {
                         };
 
                         // validate the message based on the result
-                        match handle_result {
+                        match handler_result {
                             Ok(acceptance) => {
                                 self.p2p.validate_message(&message_id, &peer_id, acceptance)?;
                             },
                             Err(err) => {
                                 log::error!("Error handling {} message: {}", topic_str, err);
-                                self.p2p.validate_message(&message_id, &peer_id, gossipsub::MessageAcceptance::Reject)?;
+                                self.p2p.validate_message(&message_id, &peer_id, gossipsub::MessageAcceptance::Ignore)?;
                             }
                         }
                     } else if std::matches!(topic_str, PINGPONG_RESPONSE_TOPIC | WORKFLOW_RESPONSE_TOPIC) {
                         // since we are responding to these topics, we might receive messages from other compute nodes
-                        // we can gracefully ignore them
+                        // we can gracefully ignore them and propagate it to to others
                         log::debug!("Ignoring message for topic: {}", topic_str);
-
-                        // accept this message for propagation
                         self.p2p.validate_message(&message_id, &peer_id, gossipsub::MessageAcceptance::Accept)?;
                     } else {
-                        log::warn!("Received message from unexpected topic: {}", topic_str);
-
                         // reject this message as its from a foreign topic
+                        log::warn!("Received message from unexpected topic: {}", topic_str);
                         self.p2p.validate_message(&message_id, &peer_id, gossipsub::MessageAcceptance::Reject)?;
                     }
                 },

src/payloads/error.rs

@@ -0,0 +1,17 @@
+use serde::{Deserialize, Serialize};
+
+/// A generic task request, given by Dria.
+#[derive(Debug, Clone, Serialize, Deserialize)]
+#[serde(rename_all = "camelCase")]
+pub struct TaskErrorPayload {
+    /// The unique identifier of the task.
+    pub task_id: String,
+    /// The stringified error object
+    pub(crate) error: String,
+}
+
+impl TaskErrorPayload {
+    pub fn new(task_id: String, error: String) -> Self {
+        Self { task_id, error }
+    }
+}

src/payloads/mod.rs

@@ -0,0 +1,8 @@
+mod error;
+pub use error::TaskErrorPayload;
+
+mod request;
+pub use request::TaskRequestPayload;
+
+mod response;
+pub use response::TaskResponsePayload;

src/payloads/request.rs

@@ -0,0 +1,33 @@
+use crate::utils::{filter::FilterPayload, get_current_time_nanos};
+use fastbloom_rs::BloomFilter;
+use serde::{Deserialize, Serialize};
+use uuid::Uuid;
+
+/// A generic task request, given by Dria.
+#[derive(Debug, Clone, Serialize, Deserialize)]
+#[serde(rename_all = "camelCase")]
+pub struct TaskRequestPayload<T> {
+    /// The unique identifier of the task.
+    pub task_id: String,
+    /// The deadline of the task in nanoseconds.
+    pub(crate) deadline: u128,
+    /// The input to the compute function.
+    pub(crate) input: T,
+    /// The Bloom filter of the task.
+    pub(crate) filter: FilterPayload,
+    /// The public key of the requester, in hexadecimals.
+    pub(crate) public_key: String,
+}
+
+impl<T> TaskRequestPayload<T> {
+    #[allow(unused)]
+    pub fn new(input: T, filter: BloomFilter, time_ns: u128, public_key: Option<String>) -> Self {
+        Self {
+            task_id: Uuid::new_v4().into(),
+            deadline: get_current_time_nanos() + time_ns,
+            input,
+            filter: filter.into(),
+            public_key: public_key.unwrap_or_default(),
+        }
+    }
+}

src/utils/payload.rs renamed to src/payloads/response.rs

@@ -1,10 +1,7 @@
-use super::crypto::sha256hash;
-use crate::utils::{filter::FilterPayload, get_current_time_nanos};
+use crate::utils::crypto::sha256hash;
 use eyre::Result;
-use fastbloom_rs::BloomFilter;
 use libsecp256k1::SecretKey;
 use serde::{Deserialize, Serialize};
-use uuid::Uuid;
 
 /// A computation task is the task of computing a result from a given input. The result is encrypted with the public key of the requester.
 /// Plain result is signed by the compute node's private key, and a commitment is computed from the signature and plain result.
@@ -23,24 +20,20 @@ pub struct TaskResponsePayload {
 }
 
 impl TaskResponsePayload {
-    pub fn to_string(&self) -> Result<String> {
-        serde_json::to_string(&serde_json::json!(self)).map_err(Into::into)
-    }
-
     /// Creates the payload of a computation result.
     ///
     /// - Sign `task_id || payload` with node `self.secret_key`
     /// - Encrypt `result` with `task_public_key`
     pub fn new(
-        payload: impl AsRef<[u8]>,
+        result: impl AsRef<[u8]>,
         task_id: &str,
         encrypting_public_key: &[u8],
         signing_secret_key: &SecretKey,
     ) -> Result<Self> {
         // create the message `task_id || payload`
         let mut preimage = Vec::new();
         preimage.extend_from_slice(task_id.as_ref());
-        preimage.extend_from_slice(payload.as_ref());
+        preimage.extend_from_slice(result.as_ref());
 
         // sign the message
         // TODO: use `sign_recoverable` here instead?
@@ -50,7 +43,7 @@ impl TaskResponsePayload {
         let recid: [u8; 1] = [recid.serialize()];
 
         // encrypt payload itself
-        let ciphertext = ecies::encrypt(encrypting_public_key, payload.as_ref())?;
+        let ciphertext = ecies::encrypt(encrypting_public_key, result.as_ref())?;
 
         Ok(TaskResponsePayload {
             ciphertext: hex::encode(ciphertext),
@@ -59,32 +52,3 @@ impl TaskResponsePayload {
         })
     }
 }
-
-/// A generic task request, given by Dria.
-#[derive(Debug, Clone, Serialize, Deserialize)]
-#[serde(rename_all = "camelCase")]
-pub struct TaskRequestPayload<T> {
-    /// The unique identifier of the task.
-    pub task_id: String,
-    /// The deadline of the task in nanoseconds.
-    pub(crate) deadline: u128,
-    /// The input to the compute function.
-    pub(crate) input: T,
-    /// The Bloom filter of the task.
-    pub(crate) filter: FilterPayload,
-    /// The public key of the requester.
-    pub(crate) public_key: String,
-}
-
-impl<T> TaskRequestPayload<T> {
-    #[allow(unused)]
-    pub fn new(input: T, filter: BloomFilter, time_ns: u128, public_key: Option<String>) -> Self {
-        Self {
-            task_id: Uuid::new_v4().into(),
-            deadline: get_current_time_nanos() + time_ns,
-            input,
-            filter: filter.into(),
-            public_key: public_key.unwrap_or_default(),
-        }
-    }
-}