added separate ollama worker, small rfks

Author: erhant

.github/workflows/tests.yml

@@ -32,3 +32,6 @@ jobs:
 
       - name: Run tests
         run: cargo test --workspace
+
+      - name: Run linter
+        run: cargo clippy --workspace

compute/src/handlers/workflow.rs

@@ -32,7 +32,7 @@ impl WorkflowHandler {
     pub(crate) async fn handle_compute(
         node: &mut DriaComputeNode,
         compute_message: &DKNMessage,
-    ) -> Result<Either<MessageAcceptance, WorkflowsWorkerInput>> {
+    ) -> Result<Either<MessageAcceptance, (WorkflowsWorkerInput, bool)>> {
         let stats = TaskStats::new().record_received_at();
         let task = compute_message
             .parse_payload::<TaskRequestPayload<WorkflowPayload>>(true)
@@ -78,14 +78,17 @@ impl WorkflowHandler {
         log::info!("Using model {} for task {}", model_name, task.task_id);
 
         // prepare workflow executor
-        let executor = if model_provider == ModelProvider::Ollama {
-            Executor::new_at(
-                model,
-                &node.config.workflows.ollama.host,
-                node.config.workflows.ollama.port,
+        let (executor, batchable) = if model_provider == ModelProvider::Ollama {
+            (
+                Executor::new_at(
+                    model,
+                    &node.config.workflows.ollama.host,
+                    node.config.workflows.ollama.port,
+                ),
+                false,
             )
         } else {
-            Executor::new(model)
+            (Executor::new(model), true)
         };
 
         // prepare entry from prompt
@@ -97,15 +100,18 @@ impl WorkflowHandler {
         // 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,
-        }))
+        Ok(Either::Right((
+            WorkflowsWorkerInput {
+                entry,
+                executor,
+                workflow,
+                model_name,
+                task_id: task.task_id,
+                public_key: task_public_key,
+                stats,
+            },
+            batchable,
+        )))
     }
 
     pub(crate) async fn handle_publish(
@@ -123,16 +129,15 @@ impl WorkflowHandler {
                     task.model_name,
                     task.stats.record_published_at(),
                 )?;
+
+                // convert payload to message
                 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.task_id,
                     payload_str
                 );
-
                 DKNMessage::new(payload_str, Self::RESPONSE_TOPIC)
             }
             Err(err) => {

compute/src/main.rs

@@ -31,6 +31,7 @@ async fn main() -> Result<()> {
     let token = CancellationToken::new();
     let cancellation_token = token.clone();
     tokio::spawn(async move {
+        // the timeout is done for profiling only, and should not be used in production
         if let Ok(Ok(duration_secs)) =
             env::var("DKN_EXIT_TIMEOUT").map(|s| s.to_string().parse::<u64>())
         {
@@ -75,15 +76,17 @@ async fn main() -> Result<()> {
     }
 
     let node_token = token.clone();
-    let (mut node, p2p, mut workflows) = DriaComputeNode::new(config, node_token).await?;
+    let (mut node, p2p, mut worker_batch, mut worker_single) =
+        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 });
+    log::info!("Spawning workflows batch worker thread.");
+    let worker_batch_handle = tokio::spawn(async move { worker_batch.run_batch().await });
+
+    log::info!("Spawning workflows single worker thread.");
+    let worker_single_handle = tokio::spawn(async move { worker_single.run().await });
 
     // launch the node in a separate thread
     log::info!("Spawning compute node thread.");
@@ -98,8 +101,11 @@ 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) = worker_single_handle.await {
+        log::error!("Workflows single worker handle error: {}", err);
+    };
+    if let Err(err) = worker_batch_handle.await {
+        log::error!("Workflows batch worker handle error: {}", err);
     };
     if let Err(err) = p2p_handle.await {
         log::error!("P2P handle error: {}", err);

compute/src/node.rs

@@ -26,10 +26,16 @@ pub struct DriaComputeNode {
     pub p2p: DriaP2PCommander,
     pub available_nodes: AvailableNodes,
     pub cancellation: CancellationToken,
-    // channels
+    /// Gossipsub message receiver.
     message_rx: mpsc::Receiver<(PeerId, MessageId, Message)>,
-    worklow_tx: mpsc::Sender<WorkflowsWorkerInput>,
-    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 {
@@ -39,7 +45,12 @@ impl DriaComputeNode {
     pub async fn new(
         config: DriaComputeNodeConfig,
         cancellation: CancellationToken,
-    ) -> Result<(DriaComputeNode, DriaP2PClient, WorkflowsWorker)> {
+    ) -> Result<(
+        DriaComputeNode,
+        DriaP2PClient,
+        WorkflowsWorker,
+        WorkflowsWorker,
+    )> {
         // create the keypair from secret key
         let keypair = secret_to_keypair(&config.secret_key);
 
@@ -66,10 +77,10 @@ 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);
+        // 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();
 
         Ok((
             DriaComputeNode {
@@ -78,11 +89,14 @@ impl DriaComputeNode {
                 cancellation,
                 available_nodes,
                 message_rx,
-                worklow_tx,
-                publish_rx,
+                workflow_batch_tx,
+                publish_batch_rx,
+                workflow_single_tx,
+                publish_single_rx,
             },
             p2p_client,
-            workflows_worker,
+            workflows_batch_worker,
+            workflows_single_worker,
         ))
     }
 
@@ -164,7 +178,7 @@ impl DriaComputeNode {
                     return MessageAcceptance::Ignore;
                 }
 
-                // first, parse the raw gossipsub message to a prepared message
+                // parse the raw gossipsub message to a prepared DKN message
                 let message = match DKNMessage::try_from_gossipsub_message(
                     &message,
                     &self.config.admin_public_key,
@@ -177,19 +191,25 @@ impl DriaComputeNode {
                     }
                 };
 
-                // then handle the prepared message
+                // handle the DKN message with respect to the topic
                 let handler_result = match message.topic.as_str() {
                     WorkflowHandler::LISTEN_TOPIC => {
                         match WorkflowHandler::handle_compute(self, &message).await {
+                            // we got acceptance, so something was not right about the workflow and we can ignore it
                             Ok(Either::Left(acceptance)) => Ok(acceptance),
-                            Ok(Either::Right(workflow_message)) => {
-                                if let Err(e) = self.worklow_tx.send(workflow_message).await {
+                            // we got the parsed workflow itself, send to a worker thread w.r.t batchable
+                            Ok(Either::Right((workflow_message, batchable))) => {
+                                if let Err(e) = match batchable {
+                                    true => self.workflow_batch_tx.send(workflow_message).await,
+                                    false => self.workflow_single_tx.send(workflow_message).await,
+                                } {
                                     log::error!("Error sending workflow message: {:?}", e);
                                 };
 
                                 // accept the message in case others may be included in the filter as well
                                 Ok(MessageAcceptance::Accept)
                             }
+                            // something went wrong, handle this outside
                             Err(err) => Err(err),
                         }
                     }
@@ -241,7 +261,16 @@ impl DriaComputeNode {
                 _ = tokio::time::sleep(available_node_refresh_duration) => 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_rx.recv() => {
+                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() => {
                     if let Some(result) = publish_msg {
                         WorkflowHandler::handle_publish(self, result).await?;
                     } else {
@@ -293,7 +322,7 @@ impl DriaComputeNode {
         self.message_rx.close();
 
         log::debug!("Closing publish channel.");
-        self.publish_rx.close();
+        self.publish_batch_rx.close();
 
         Ok(())
     }
@@ -339,7 +368,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");

compute/src/utils/message.rs

@@ -133,7 +133,7 @@ impl DKNMessage {
 
         // check dria signature
         // NOTE: when we have many public keys, we should check the signature against all of them
-        if !message.is_signed(&public_key)? {
+        if !message.is_signed(public_key)? {
             return Err(eyre!("Invalid signature."));
         }