combine publish channels

Author: erhant

compute/src/handlers/pingpong.rs

@@ -19,11 +19,12 @@ struct PingpongPayload {
 
 #[derive(Serialize, Deserialize, Debug, Clone)]
 struct PingpongResponse {
+    /// UUID as given in the ping payload.
     pub(crate) uuid: String,
+    /// Models available in the node.
     pub(crate) models: Vec<(ModelProvider, Model)>,
-    pub(crate) timestamp: u128,
     /// Number of tasks in the channel currently, `single` and `batch`.
-    pub(crate) tasks: [usize; 2],
+    pub(crate) active_task_count: [usize; 2],
 }
 
 impl PingpongHandler {
@@ -65,8 +66,7 @@ impl PingpongHandler {
         let response_body = PingpongResponse {
             uuid: pingpong.uuid.clone(),
             models: node.config.workflows.models.clone(),
-            timestamp: get_current_time_nanos(),
-            tasks: node.get_active_task_count(),
+            active_task_count: node.get_active_task_count(),
         };
 
         // publish message

compute/src/handlers/workflow.rs

@@ -114,6 +114,7 @@ impl WorkflowHandler {
         )))
     }
 
+    /// Handles the result of a workflow task.
     pub(crate) async fn handle_publish(
         node: &mut DriaComputeNode,
         task: WorkflowsWorkerOutput,
@@ -131,8 +132,7 @@ impl WorkflowHandler {
                 )?;
 
                 // convert payload to message
-                let payload_str = serde_json::to_string(&payload)
-                    .wrap_err("could not serialize response payload")?;
+                let payload_str = serde_json::json!(payload).to_string();
                 log::debug!(
                     "Publishing result for task {}\n{}",
                     task.task_id,
@@ -152,8 +152,7 @@ impl WorkflowHandler {
                     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")?;
+                let error_payload_str = serde_json::json!(error_payload).to_string();
 
                 // prepare signed message
                 DKNMessage::new_signed(
@@ -178,6 +177,7 @@ impl WorkflowHandler {
                 Self::RESPONSE_TOPIC,
                 &node.config.secret_key,
             );
+
             node.publish(message).await?;
         };
 

compute/src/main.rs

@@ -75,9 +75,8 @@ async fn main() -> Result<()> {
         return Ok(());
     }
 
-    let node_token = token.clone();
-    let (mut node, p2p, mut worker_batch, mut worker_single) =
-        DriaComputeNode::new(config, node_token).await?;
+    // create the node
+    let (mut node, p2p, mut worker_batch, mut worker_single) = DriaComputeNode::new(config).await?;
 
     log::info!("Spawning peer-to-peer client thread.");
     let p2p_handle = tokio::spawn(async move { p2p.run().await });
@@ -90,8 +89,9 @@ async fn main() -> Result<()> {
 
     // launch the node in a separate thread
     log::info!("Spawning compute node thread.");
+    let node_token = token.clone();
     let node_handle = tokio::spawn(async move {
-        if let Err(err) = node.run().await {
+        if let Err(err) = node.run(node_token).await {
             log::error!("Node launch error: {}", err);
             panic!("Node failed.")
         };

compute/src/node.rs

@@ -20,22 +20,21 @@ use crate::{
 const DIAGNOSTIC_REFRESH_INTERVAL_SECS: u64 = 30;
 /// Number of seconds between refreshing the available nodes.
 const AVAILABLE_NODES_REFRESH_INTERVAL_SECS: u64 = 30 * 60; // 30 minutes
+/// Buffer size for message publishes.
+const PUBLISH_CHANNEL_BUFSIZE: usize = 1024;
 
 pub struct DriaComputeNode {
     pub config: DriaComputeNodeConfig,
     pub p2p: DriaP2PCommander,
     pub available_nodes: AvailableNodes,
-    pub cancellation: CancellationToken,
     /// Gossipsub message receiver.
     message_rx: mpsc::Receiver<(PeerId, MessageId, Message)>,
+    /// Publish receiver to receive messages to be published.
+    publish_rx: mpsc::Receiver<WorkflowsWorkerOutput>,
     /// Workflow transmitter to send batchable tasks.
     workflow_batch_tx: mpsc::Sender<WorkflowsWorkerInput>,
-    /// Publish receiver to receive messages to be published.
-    publish_batch_rx: mpsc::Receiver<WorkflowsWorkerOutput>,
     /// Workflow transmitter to send single tasks.
     workflow_single_tx: mpsc::Sender<WorkflowsWorkerInput>,
-    /// Publish receiver to receive messages to be published.
-    publish_single_rx: mpsc::Receiver<WorkflowsWorkerOutput>,
 }
 
 impl DriaComputeNode {
@@ -44,7 +43,6 @@ impl DriaComputeNode {
     /// Returns the node instance and p2p client together. P2p MUST be run in a separate task before this node is used at all.
     pub async fn new(
         config: DriaComputeNodeConfig,
-        cancellation: CancellationToken,
     ) -> Result<(
         DriaComputeNode,
         DriaP2PClient,
@@ -77,22 +75,20 @@ impl DriaComputeNode {
             protocol,
         )?;
 
-        // create workflow workers
-        let (workflows_batch_worker, workflow_batch_tx, publish_batch_rx) = WorkflowsWorker::new();
-        let (workflows_single_worker, workflow_single_tx, publish_single_rx) =
-            WorkflowsWorker::new();
+        // create workflow workers, all workers use the same publish channel
+        let (publish_tx, publish_rx) = mpsc::channel(PUBLISH_CHANNEL_BUFSIZE);
+        let (workflows_batch_worker, workflow_batch_tx) = WorkflowsWorker::new(publish_tx.clone());
+        let (workflows_single_worker, workflow_single_tx) = WorkflowsWorker::new(publish_tx);
 
         Ok((
             DriaComputeNode {
                 config,
                 p2p: p2p_commander,
-                cancellation,
                 available_nodes,
                 message_rx,
+                publish_rx,
                 workflow_batch_tx,
-                publish_batch_rx,
                 workflow_single_tx,
-                publish_single_rx,
             },
             p2p_client,
             workflows_batch_worker,
@@ -248,8 +244,8 @@ impl DriaComputeNode {
     }
 
     /// Runs the main loop of the compute node.
-    /// This method is not expected to return until cancellation occurs.
-    pub async fn run(&mut self) -> Result<()> {
+    /// This method is not expected to return until cancellation occurs for the given token.
+    pub async fn run(&mut self, cancellation: CancellationToken) -> Result<()> {
         // prepare durations for sleeps
         let mut peer_refresh_interval =
             tokio::time::interval(Duration::from_secs(DIAGNOSTIC_REFRESH_INTERVAL_SECS));
@@ -270,16 +266,7 @@ impl DriaComputeNode {
                 _ = available_node_refresh_interval.tick() => self.handle_available_nodes_refresh().await,
                 // a Workflow message to be published is received from the channel
                 // this is expected to be sent by the workflow worker
-                publish_msg = self.publish_batch_rx.recv() => {
-                    if let Some(result) = publish_msg {
-                        WorkflowHandler::handle_publish(self, result).await?;
-                    } else {
-                        log::error!("Publish channel closed unexpectedly.");
-                        break;
-                    };
-                },
-                // TODO: make the both receivers handled together somehow
-                publish_msg = self.publish_single_rx.recv() => {
+                publish_msg = self.publish_rx.recv() => {
                     if let Some(result) = publish_msg {
                         WorkflowHandler::handle_publish(self, result).await?;
                     } else {
@@ -306,7 +293,7 @@ impl DriaComputeNode {
                 },
                 // check if the cancellation token is cancelled
                 // this is expected to be cancelled by the main thread with signal handling
-                _ = self.cancellation.cancelled() => break,
+                _ = cancellation.cancelled() => break,
             }
         }
 
@@ -331,7 +318,7 @@ impl DriaComputeNode {
         self.message_rx.close();
 
         log::debug!("Closing publish channel.");
-        self.publish_batch_rx.close();
+        self.publish_rx.close();
 
         Ok(())
     }
@@ -345,8 +332,8 @@ impl DriaComputeNode {
         }
 
         // print task counts
-        let [single, batch] = self.get_active_task_count();
-        log::info!("Active Task Count (single/batch): {} / {}", single, batch);
+        // let [single, batch] = self.get_active_task_count();
+        // log::info!("Active Task Count (single/batch): {} / {}", single, batch);
     }
 
     /// Updates the local list of available nodes by refreshing it.
@@ -382,18 +369,18 @@ mod tests {
 
         // create node
         let cancellation = CancellationToken::new();
-        let (mut node, p2p, _, _) =
-            DriaComputeNode::new(DriaComputeNodeConfig::default(), cancellation.clone())
-                .await
-                .expect("should create node");
+        let (mut node, p2p, _, _) = DriaComputeNode::new(DriaComputeNodeConfig::default())
+            .await
+            .expect("should create node");
 
         // spawn p2p task
         let p2p_task = tokio::spawn(async move { p2p.run().await });
 
         // launch & wait for a while for connections
         log::info!("Waiting a bit for peer setup.");
+        let run_cancellation = cancellation.clone();
         tokio::select! {
-            _ = node.run() => (),
+            _ = node.run(run_cancellation) => (),
             _ = tokio::time::sleep(tokio::time::Duration::from_secs(20)) => cancellation.cancel(),
         }
         log::info!("Connected Peers:\n{:#?}", node.peers().await?);

compute/src/workers/workflow.rs

@@ -29,8 +29,8 @@ pub struct WorkflowsWorker {
     publish_tx: mpsc::Sender<WorkflowsWorkerOutput>,
 }
 
+/// Buffer size for workflow tasks (per worker).
 const WORKFLOW_CHANNEL_BUFSIZE: usize = 1024;
-const PUBLISH_CHANNEL_BUFSIZE: usize = 1024;
 
 impl WorkflowsWorker {
     /// Batch size that defines how many tasks can be executed in parallel at once.
@@ -39,24 +39,21 @@ impl WorkflowsWorker {
     const BATCH_SIZE: usize = 8;
 
     /// Creates a worker and returns the sender and receiver for the worker.
-    pub fn new() -> (
-        WorkflowsWorker,
-        mpsc::Sender<WorkflowsWorkerInput>,
-        mpsc::Receiver<WorkflowsWorkerOutput>,
-    ) {
+    pub fn new(
+        publish_tx: mpsc::Sender<WorkflowsWorkerOutput>,
+    ) -> (WorkflowsWorker, mpsc::Sender<WorkflowsWorkerInput>) {
         let (workflow_tx, workflow_rx) = mpsc::channel(WORKFLOW_CHANNEL_BUFSIZE);
-        let (publish_tx, publish_rx) = mpsc::channel(PUBLISH_CHANNEL_BUFSIZE);
 
         (
             Self {
                 workflow_rx,
                 publish_tx,
             },
             workflow_tx,
-            publish_rx,
         )
     }
 
+    /// Closes the workflow receiver channel.
     fn shutdown(&mut self) {
         log::warn!("Closing workflows worker.");
         self.workflow_rx.close();