parallel workflows first version works

erhant · erhant · commit 601e0e87713c · 2024-11-27T23:25:52.000+03:00
diff --git a/compute/src/handlers/workflow.rs b/compute/src/handlers/workflow.rs
@@ -1,12 +1,13 @@
 use dkn_p2p::libp2p::gossipsub::MessageAcceptance;
-use dkn_workflows::{Entry, Executor, ModelProvider, ProgramMemory, Workflow};
-use eyre::{eyre, Context, Result};
+use dkn_workflows::{Entry, Executor, ModelProvider, Workflow};
+use eyre::{Context, Result};
 use libsecp256k1::PublicKey;
 use serde::Deserialize;
-use std::time::Instant;
+use tokio_util::either::Either;
 
 use crate::payloads::{TaskErrorPayload, TaskRequestPayload, TaskResponsePayload, TaskStats};
 use crate::utils::{get_current_time_nanos, DKNMessage};
+use crate::workers::workflow::*;
 use crate::DriaComputeNode;
 
 pub struct WorkflowHandler;
@@ -31,11 +32,13 @@ impl WorkflowHandler {
     pub(crate) async fn handle_compute(
         node: &mut DriaComputeNode,
         message: DKNMessage,
-    ) -> Result<MessageAcceptance> {
+    ) -> Result<Either<MessageAcceptance, WorkflowsWorkerInput>> {
         let task = message
             .parse_payload::<TaskRequestPayload<WorkflowPayload>>(true)
             .wrap_err("Could not parse workflow task")?;
-        let mut task_stats = TaskStats::default().record_received_at();
+
+        // TODO: !!!
+        let task_stats = TaskStats::default().record_received_at();
 
         // check if deadline is past or not
         let current_time = get_current_time_nanos();
@@ -48,7 +51,7 @@ impl WorkflowHandler {
             );
 
             // ignore the message
-            return Ok(MessageAcceptance::Ignore);
+            return Ok(Either::Left(MessageAcceptance::Ignore));
         }
 
         // check task inclusion via the bloom filter
@@ -59,9 +62,15 @@ impl WorkflowHandler {
             );
 
             // accept the message, someone else may be included in filter
-            return Ok(MessageAcceptance::Accept);
+            return Ok(Either::Left(MessageAcceptance::Accept));
         }
 
+        // obtain public key from the payload
+        // do this early to avoid unnecessary processing
+        let task_public_key_bytes =
+            hex::decode(&task.public_key).wrap_err("could not decode public key")?;
+        let task_public_key = PublicKey::parse_slice(&task_public_key_bytes, None)?;
+
         // read model / provider from the task
         let (model_provider, model) = node
             .config
@@ -80,65 +89,66 @@ impl WorkflowHandler {
         } else {
             Executor::new(model)
         };
+
+        // prepare entry from prompt
         let entry: Option<Entry> = task
             .input
             .prompt
             .map(|prompt| Entry::try_value_or_str(&prompt));
 
-        // execute workflow with cancellation
-        let mut memory = ProgramMemory::new();
-
-        let exec_started_at = Instant::now();
-        let exec_result = executor
-            .execute(entry.as_ref(), &task.input.workflow, &mut memory)
-            .await
-            .map_err(|e| eyre!("Execution error: {}", e.to_string()));
-        task_stats = task_stats.record_execution_time(exec_started_at);
-
-        Ok(MessageAcceptance::Accept)
+        // get workflow as well
+        let workflow = task.input.workflow;
+
+        Ok(Either::Right(WorkflowsWorkerInput {
+            entry,
+            executor,
+            workflow,
+            model_name,
+            task_id: task.task_id,
+            public_key: task_public_key,
+            stats: task_stats,
+        }))
     }
 
-    async fn handle_publish(
+    pub(crate) async fn handle_publish(
         node: &mut DriaComputeNode,
-        result: String,
-        task_id: String,
-    ) -> Result<()> {
-        let (message, acceptance) = match exec_result {
+        task: WorkflowsWorkerOutput,
+    ) -> Result<MessageAcceptance> {
+        let (message, acceptance) = match task.result {
             Ok(result) => {
-                // obtain public key from the payload
-                let task_public_key_bytes =
-                    hex::decode(&task.public_key).wrap_err("Could not decode public key")?;
-                let task_public_key = PublicKey::parse_slice(&task_public_key_bytes, None)?;
-
                 // prepare signed and encrypted payload
                 let payload = TaskResponsePayload::new(
                     result,
-                    &task_id,
-                    &task_public_key,
+                    &task.task_id,
+                    &task.public_key,
                     &node.config.secret_key,
-                    model_name,
-                    task_stats.record_published_at(),
+                    task.model_name,
+                    task.stats.record_published_at(),
                 )?;
                 let payload_str = serde_json::to_string(&payload)
                     .wrap_err("Could not serialize response payload")?;
 
                 // prepare signed message
-                log::debug!("Publishing result for task {}\n{}", task_id, payload_str);
+                log::debug!(
+                    "Publishing result for task {}\n{}",
+                    task.task_id,
+                    payload_str
+                );
                 let message = DKNMessage::new(payload_str, Self::RESPONSE_TOPIC);
                 // accept so that if there are others included in filter they can do the task
                 (message, MessageAcceptance::Accept)
             }
             Err(err) => {
                 // use pretty display string for error logging with causes
                 let err_string = format!("{:#}", err);
-                log::error!("Task {} failed: {}", task_id, err_string);
+                log::error!("Task {} failed: {}", task.task_id, err_string);
 
                 // prepare error payload
                 let error_payload = TaskErrorPayload {
-                    task_id,
+                    task_id: task.task_id.clone(),
                     error: err_string,
-                    model: model_name,
-                    stats: task_stats.record_published_at(),
+                    model: task.model_name,
+                    stats: task.stats.record_published_at(),
                 };
                 let error_payload_str = serde_json::to_string(&error_payload)
                     .wrap_err("Could not serialize error payload")?;
@@ -160,7 +170,7 @@ impl WorkflowHandler {
             log::error!("{}", err_msg);
 
             let payload = serde_json::json!({
-                "taskId": task_id,
+                "taskId": task.task_id,
                 "error": err_msg,
             });
             let message = DKNMessage::new_signed(
@@ -171,6 +181,6 @@ impl WorkflowHandler {
             node.publish(message).await?;
         };
 
-        Ok(())
+        Ok(acceptance)
     }
 }
diff --git a/compute/src/lib.rs b/compute/src/lib.rs
@@ -6,6 +6,7 @@ pub(crate) mod handlers;
 pub(crate) mod node;
 pub(crate) mod payloads;
 pub(crate) mod utils;
+pub(crate) mod workers;
 
 /// Crate version of the compute node.
 /// This value is attached within the published messages.
diff --git a/compute/src/main.rs b/compute/src/main.rs
@@ -75,12 +75,16 @@ async fn main() -> Result<()> {
     }
 
     let node_token = token.clone();
-    let (mut node, p2p) = DriaComputeNode::new(config, node_token).await?;
+    let (mut node, p2p, mut workflows) = DriaComputeNode::new(config, node_token).await?;
 
     // launch the p2p in a separate thread
     log::info!("Spawning peer-to-peer client thread.");
     let p2p_handle = tokio::spawn(async move { p2p.run().await });
 
+    // launch the workflows in a separate thread
+    log::info!("Spawning workflows worker thread.");
+    let workflows_handle = tokio::spawn(async move { workflows.run().await });
+
     // launch the node in a separate thread
     log::info!("Spawning compute node thread.");
     let node_handle = tokio::spawn(async move {
@@ -94,6 +98,9 @@ async fn main() -> Result<()> {
     if let Err(err) = node_handle.await {
         log::error!("Node handle error: {}", err);
     };
+    if let Err(err) = workflows_handle.await {
+        log::error!("Workflows handle error: {}", err);
+    };
     if let Err(err) = p2p_handle.await {
         log::error!("P2P handle error: {}", err);
     };
diff --git a/compute/src/node.rs b/compute/src/node.rs
@@ -10,12 +10,13 @@ use tokio::{
     sync::mpsc,
     time::{Duration, Instant},
 };
-use tokio_util::sync::CancellationToken;
+use tokio_util::{either::Either, sync::CancellationToken};
 
 use crate::{
     config::*,
     handlers::*,
     utils::{crypto::secret_to_keypair, AvailableNodes, DKNMessage},
+    workers::workflow::{WorkflowsWorker, WorkflowsWorkerInput, WorkflowsWorkerOutput},
 };
 
 /// Number of seconds between refreshing the Kademlia DHT.
@@ -27,7 +28,10 @@ pub struct DriaComputeNode {
     pub available_nodes: AvailableNodes,
     pub cancellation: CancellationToken,
     peer_last_refreshed: Instant,
-    msg_rx: mpsc::Receiver<(PeerId, MessageId, Message)>,
+    // channels
+    message_rx: mpsc::Receiver<(PeerId, MessageId, Message)>,
+    worklow_tx: mpsc::Sender<WorkflowsWorkerInput>,
+    publish_rx: mpsc::Receiver<WorkflowsWorkerOutput>,
 }
 
 impl DriaComputeNode {
@@ -37,7 +41,7 @@ impl DriaComputeNode {
     pub async fn new(
         config: DriaComputeNodeConfig,
         cancellation: CancellationToken,
-    ) -> Result<(DriaComputeNode, DriaP2PClient)> {
+    ) -> Result<(DriaComputeNode, DriaP2PClient, WorkflowsWorker)> {
         // create the keypair from secret key
         let keypair = secret_to_keypair(&config.secret_key);
 
@@ -55,7 +59,7 @@ impl DriaComputeNode {
         log::info!("Using identity: {}", protocol);
 
         // create p2p client
-        let (p2p_client, p2p_commander, msg_rx) = DriaP2PClient::new(
+        let (p2p_client, p2p_commander, message_rx) = DriaP2PClient::new(
             keypair,
             config.p2p_listen_addr.clone(),
             available_nodes.bootstrap_nodes.clone().into_iter(),
@@ -64,16 +68,24 @@ impl DriaComputeNode {
             protocol,
         )?;
 
+        // create workflow worker
+        let (worklow_tx, workflow_rx) = mpsc::channel(256);
+        let (publish_tx, publish_rx) = mpsc::channel(256);
+        let workflows_worker = WorkflowsWorker::new(workflow_rx, publish_tx);
+
         Ok((
             DriaComputeNode {
                 config,
                 p2p: p2p_commander,
                 cancellation,
                 available_nodes,
-                msg_rx,
+                message_rx,
+                worklow_tx,
+                publish_rx,
                 peer_last_refreshed: Instant::now(),
             },
             p2p_client,
+            workflows_worker,
         ))
     }
 
@@ -207,7 +219,18 @@ impl DriaComputeNode {
             // then handle the prepared message
             let handler_result = match topic_str {
                 WorkflowHandler::LISTEN_TOPIC => {
-                    WorkflowHandler::handle_compute(self, message).await
+                    let compute_result = WorkflowHandler::handle_compute(self, message).await;
+                    match compute_result {
+                        Ok(Either::Left(acceptance)) => Ok(acceptance),
+                        Ok(Either::Right(workflow_message)) => {
+                            if let Err(e) = self.worklow_tx.send(workflow_message).await {
+                                log::error!("Error sending workflow message: {:?}", e);
+                            };
+
+                            Ok(MessageAcceptance::Accept)
+                        }
+                        Err(err) => Err(err),
+                    }
                 }
                 PingpongHandler::LISTEN_TOPIC => PingpongHandler::handle_ping(self, message).await,
                 _ => unreachable!(), // unreachable because of the if condition
@@ -251,7 +274,12 @@ impl DriaComputeNode {
         // the underlying p2p client is expected to handle the rest within its own loop
         loop {
             tokio::select! {
-                gossipsub_msg = self.msg_rx.recv() => {
+                publish_msg = self.publish_rx.recv() => {
+                    if let Some(result) = publish_msg {
+                        WorkflowHandler::handle_publish(self, result).await?;
+                    }
+                },
+                gossipsub_msg = self.message_rx.recv() => {
                     if let Some((peer_id, message_id, message)) = gossipsub_msg {
                         // handle the message, returning a message acceptance for the received one
                         let acceptance = self.handle_message((peer_id, &message_id, message)).await;
@@ -282,15 +310,16 @@ impl DriaComputeNode {
         Ok(())
     }
 
-    /// Shutdown channels between p2p and yourself.
+    /// Shutdown channels between p2p, worker and yourself.
     pub async fn shutdown(&mut self) -> Result<()> {
-        // send shutdown signal
         log::debug!("Sending shutdown command to p2p client.");
         self.p2p.shutdown().await?;
 
-        // close message channel
         log::debug!("Closing message channel.");
-        self.msg_rx.close();
+        self.message_rx.close();
+
+        log::debug!("Closing publish channel.");
+        self.publish_rx.close();
 
         Ok(())
     }
@@ -329,7 +358,7 @@ mod tests {
 
         // create node
         let cancellation = CancellationToken::new();
-        let (mut node, p2p) =
+        let (mut node, p2p, _) =
             DriaComputeNode::new(DriaComputeNodeConfig::default(), cancellation.clone())
                 .await
                 .expect("should create node");
diff --git a/compute/src/workers/mod.rs b/compute/src/workers/mod.rs
@@ -0,0 +1 @@
+pub mod workflow;
diff --git a/compute/src/workers/workflow.rs b/compute/src/workers/workflow.rs
