fix piggybacking issues

Author: erhant

compute/src/node/core.rs

@@ -32,7 +32,7 @@ impl DriaComputeNode {
                 publish_msg_opt = self.publish_rx.recv() => {
                     if let Some(publish_msg) = publish_msg_opt {
                         // remove the task from pending tasks based on its batchability
-                        let channel = match publish_msg.batchable {
+                        let task_metadata = match publish_msg.batchable {
                             true => {
                                 self.completed_tasks_batch += 1;
                                 self.pending_tasks_batch.remove(&publish_msg.task_id)
@@ -44,7 +44,7 @@ impl DriaComputeNode {
                         };
 
                         // respond to the request
-                        match channel {
+                        match task_metadata {
                             Some(channel) => {
                               TaskResponder::handle_respond(self, publish_msg, channel).await?;
                             }

compute/src/node/mod.rs

@@ -14,7 +14,7 @@ use crate::{
     config::*,
     gossipsub::*,
     utils::{crypto::secret_to_keypair, refresh_dria_nodes, SpecCollector},
-    workers::task::{TaskWorker, TaskWorkerInput, TaskWorkerOutput},
+    workers::task::{TaskWorker, TaskWorkerInput, TaskWorkerMetadata, TaskWorkerOutput},
 };
 
 mod core;
@@ -45,9 +45,9 @@ pub struct DriaComputeNode {
     /// Task worker transmitter to send single tasks.
     task_single_tx: Option<mpsc::Sender<TaskWorkerInput>>,
     // Single tasks
-    pending_tasks_single: HashMap<String, ResponseChannel<Vec<u8>>>,
+    pending_tasks_single: HashMap<String, TaskWorkerMetadata>,
     // Batchable tasks
-    pending_tasks_batch: HashMap<String, ResponseChannel<Vec<u8>>>,
+    pending_tasks_batch: HashMap<String, TaskWorkerMetadata>,
     /// Completed single tasks count
     completed_tasks_single: usize,
     /// Completed batch tasks count

compute/src/node/reqres.rs

@@ -27,15 +27,16 @@ impl DriaComputeNode {
         } else if let Ok(task_request) = TaskResponder::try_parse_request(&data) {
             log::info!("Received a task request from {}", peer_id);
 
-            let task_tx_message = TaskResponder::handle_compute(self, &task_request).await?;
-            if let Err(e) = match task_tx_message.batchable {
+            let (task_input, task_metadata) =
+                TaskResponder::prepare_worker_input(self, &task_request, channel).await?;
+            if let Err(e) = match task_input.batchable {
                 // this is a batchable task, send it to batch worker
                 // and keep track of the task id in pending tasks
                 true => match self.task_batch_tx {
                     Some(ref mut tx) => {
                         self.pending_tasks_batch
-                            .insert(task_tx_message.task_id.clone(), channel);
-                        tx.send(task_tx_message).await
+                            .insert(task_input.task_id.clone(), task_metadata);
+                        tx.send(task_input).await
                     }
                     None => {
                         return Err(eyre!(
@@ -49,8 +50,8 @@ impl DriaComputeNode {
                 false => match self.task_single_tx {
                     Some(ref mut tx) => {
                         self.pending_tasks_single
-                            .insert(task_tx_message.task_id.clone(), channel);
-                        tx.send(task_tx_message).await
+                            .insert(task_input.task_id.clone(), task_metadata);
+                        tx.send(task_input).await
                     }
                     None => {
                         return Err(eyre!("Single workflow received but no worker available."));

compute/src/reqres/task.rs

@@ -36,13 +36,11 @@ pub struct TaskPayload {
 
 impl TaskResponder {
     /// Handles the compute message for workflows.
-    ///
-    /// - FIXME: DOES NOT CHECK FOR FILTER AS IT IS NO LONGER USED
-    /// - FIXME: GIVES ERROR ON DEADLINE PAST CASE, BUT WE DONT NEED DEADLINE AS WELL
-    pub(crate) async fn handle_compute(
+    pub(crate) async fn prepare_worker_input(
         node: &mut DriaComputeNode,
         compute_message: &DriaMessage,
-    ) -> Result<TaskWorkerInput> {
+        channel: ResponseChannel<Vec<u8>>,
+    ) -> Result<(TaskWorkerInput, TaskWorkerMetadata)> {
         // parse payload
         let task = compute_message
             .parse_payload::<TaskRequestPayload<TaskPayload>>()
@@ -52,7 +50,7 @@ impl TaskResponder {
         let stats = TaskStats::new().record_received_at();
 
         // check if deadline is past or not
-        // with request-response, we dont expect this to happen much
+        // FIXME: with request-response, we dont expect this to happen much
         if get_current_time_nanos() >= task.deadline {
             return Err(eyre!(
                 "Task {} is past the deadline, ignoring",
@@ -97,34 +95,40 @@ impl TaskResponder {
         // get workflow as well
         let workflow = task.input.workflow;
 
-        Ok(TaskWorkerInput {
+        let task_input = TaskWorkerInput {
             entry,
             executor,
             workflow,
-            model_name,
             task_id: task.task_id,
-            public_key: task_public_key,
             stats,
             batchable,
-        })
+        };
+
+        let task_metadata = TaskWorkerMetadata {
+            model_name,
+            public_key: task_public_key,
+            channel,
+        };
+
+        Ok((task_input, task_metadata))
     }
 
     /// Handles the result of a workflow task.
     pub(crate) async fn handle_respond(
         node: &mut DriaComputeNode,
-        task: TaskWorkerOutput,
-        channel: ResponseChannel<Vec<u8>>,
+        task_output: TaskWorkerOutput,
+        task_metadata: TaskWorkerMetadata,
     ) -> Result<()> {
-        let response = match task.result {
+        let response = match task_output.result {
             Ok(result) => {
                 // prepare signed and encrypted payload
-                log::info!("Publishing result for task {}", task.task_id);
+                log::info!("Publishing result for task {}", task_output.task_id);
                 let payload = TaskResponsePayload::new(
                     result,
-                    &task.task_id,
-                    &task.public_key,
-                    task.model_name,
-                    task.stats.record_published_at(),
+                    &task_output.task_id,
+                    &task_metadata.public_key,
+                    task_metadata.model_name,
+                    task_output.stats.record_published_at(),
                 )?;
 
                 // convert payload to message
@@ -135,14 +139,14 @@ impl TaskResponder {
             Err(err) => {
                 // use pretty display string for error logging with causes
                 let err_string = format!("{:#}", err);
-                log::error!("Task {} failed: {}", task.task_id, err_string);
+                log::error!("Task {} failed: {}", task_output.task_id, err_string);
 
                 // prepare error payload
                 let error_payload = TaskErrorPayload {
-                    task_id: task.task_id,
+                    task_id: task_output.task_id,
                     error: err_string,
-                    model: task.model_name,
-                    stats: task.stats.record_published_at(),
+                    model: task_metadata.model_name,
+                    stats: task_output.stats.record_published_at(),
                 };
                 let error_payload_str = serde_json::json!(error_payload).to_string();
 
@@ -152,7 +156,7 @@ impl TaskResponder {
 
         // respond through the channel
         let data = response.to_bytes()?;
-        node.p2p.respond(data, channel).await?;
+        node.p2p.respond(data, task_metadata.channel).await?;
 
         Ok(())
     }

compute/src/workers/task.rs

@@ -1,29 +1,28 @@
+use dkn_p2p::libp2p::request_response::ResponseChannel;
 use dkn_workflows::{Entry, ExecutionError, Executor, Workflow};
 use libsecp256k1::PublicKey;
 use tokio::sync::mpsc;
 
 use crate::payloads::TaskStats;
 
-// TODO: instead of piggybacking stuff here, maybe node can hold it in a hashmap w.r.t taskId
+pub struct TaskWorkerMetadata {
+    pub public_key: PublicKey,
+    pub model_name: String,
+    pub channel: ResponseChannel<Vec<u8>>,
+}
 
 pub struct TaskWorkerInput {
     pub entry: Option<Entry>,
     pub executor: Executor,
     pub workflow: Workflow,
     pub task_id: String,
-    // piggybacked
-    pub public_key: PublicKey,
-    pub model_name: String,
     pub stats: TaskStats,
     pub batchable: bool,
 }
 
 pub struct TaskWorkerOutput {
     pub result: Result<String, ExecutionError>,
     pub task_id: String,
-    // piggybacked
-    pub public_key: PublicKey,
-    pub model_name: String,
     pub stats: TaskStats,
     pub batchable: bool,
 }
@@ -129,6 +128,7 @@ impl TaskWorker {
                 "number of tasks cant be larger than batch size"
             );
             debug_assert!(num_tasks != 0, "number of tasks cant be zero");
+
             log::info!("Processing {} tasks in batch", num_tasks);
             let mut batch = tasks.into_iter().map(|b| (b, &self.publish_tx));
             match num_tasks {
@@ -226,9 +226,7 @@ impl TaskWorker {
 
         let output = TaskWorkerOutput {
             result,
-            public_key: input.public_key,
             task_id: input.task_id,
-            model_name: input.model_name,
             batchable: input.batchable,
             stats: input.stats,
         };
@@ -242,7 +240,6 @@ impl TaskWorker {
 #[cfg(test)]
 mod tests {
     use dkn_workflows::{Executor, Model};
-    use libsecp256k1::{PublicKey, SecretKey};
 
     use super::*;
     use crate::payloads::TaskStats;
@@ -311,9 +308,7 @@ mod tests {
                 entry: None,
                 executor,
                 workflow,
-                public_key: PublicKey::from_secret_key(&SecretKey::default()),
                 task_id: format!("task-{}", i + 1),
-                model_name: model.to_string(),
                 stats: TaskStats::default(),
                 batchable: true,
             };