optional threads w.r.t model

erhant · erhant · commit 0365ded86858 · 2024-11-29T13:58:59.000+03:00
diff --git a/compute/src/handlers/pingpong.rs b/compute/src/handlers/pingpong.rs
@@ -24,7 +24,7 @@ struct PingpongResponse {
     /// Models available in the node.
     pub(crate) models: Vec<(ModelProvider, Model)>,
     /// Number of tasks in the channel currently, `single` and `batch`.
-    pub(crate) active_task_count: [usize; 2],
+    pub(crate) pending_tasks: [usize; 2],
 }
 
 impl PingpongHandler {
@@ -66,7 +66,7 @@ impl PingpongHandler {
         let response_body = PingpongResponse {
             uuid: pingpong.uuid.clone(),
             models: node.config.workflows.models.clone(),
-            active_task_count: node.get_active_task_count(),
+            pending_tasks: node.get_pending_task_count(),
         };
 
         // publish message
diff --git a/compute/src/handlers/workflow.rs b/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, bool)>> {
+    ) -> Result<Either<MessageAcceptance, WorkflowsWorkerInput>> {
         let stats = TaskStats::new().record_received_at();
         let task = compute_message
             .parse_payload::<TaskRequestPayload<WorkflowPayload>>(true)
@@ -100,18 +100,16 @@ 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,
-        )))
+        }))
     }
 
     /// Handles the result of a workflow task.
diff --git a/compute/src/main.rs b/compute/src/main.rs
@@ -1,11 +1,12 @@
 use dkn_compute::*;
 use eyre::Result;
 use std::env;
-use tokio_util::sync::CancellationToken;
+use tokio_util::{sync::CancellationToken, task::TaskTracker};
 
 #[tokio::main]
 async fn main() -> Result<()> {
     let dotenv_result = dotenvy::dotenv();
+
     // TODO: remove me later when the launcher is fixed
     amend_log_levels();
 
@@ -28,88 +29,80 @@ async fn main() -> Result<()> {
 "#
     );
 
-    let token = CancellationToken::new();
-    let cancellation_token = token.clone();
+    // task tracker for multiple threads
+    let task_tracker = TaskTracker::new();
+    let cancellation = CancellationToken::new();
+
+    // spawn the background task to wait for termination signals
+    let task_tracker_to_close = task_tracker.clone();
+    let cancellation_token = cancellation.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>())
         {
+            // the timeout is done for profiling only, and should not be used in production
             log::warn!("Waiting for {} seconds before exiting.", duration_secs);
             tokio::time::sleep(tokio::time::Duration::from_secs(duration_secs)).await;
 
             log::warn!("Exiting due to DKN_EXIT_TIMEOUT.");
+
             cancellation_token.cancel();
         } else if let Err(err) = wait_for_termination(cancellation_token.clone()).await {
+            // if there is no timeout, we wait for termination signals here
             log::error!("Error waiting for termination: {:?}", err);
             log::error!("Cancelling due to unexpected error.");
             cancellation_token.cancel();
         };
+
+        // close tracker in any case
+        task_tracker_to_close.close();
     });
 
     // create configurations & check required services & address in use
     let mut config = DriaComputeNodeConfig::new();
     config.assert_address_not_in_use()?;
-    let service_check_token = token.clone();
-    let config = tokio::spawn(async move {
-        tokio::select! {
-            result = config.workflows.check_services() => {
-                if let Err(err) = result {
-                    log::error!("Error checking services: {:?}", err);
-                    panic!("Service check failed.")
-                }
-                log::warn!("Using models: {:#?}", config.workflows.models);
-                config
-            }
-            _ = service_check_token.cancelled() => {
-                log::info!("Service check cancelled.");
-                config
-            }
+    // check services & models, will exit if there is an error
+    // since service check can take time, we allow early-exit here as well
+    tokio::select! {
+        result = config.workflows.check_services() => result,
+        _ = cancellation.cancelled() => {
+            log::info!("Service check cancelled, exiting.");
+            return Ok(());
         }
-    })
-    .await?;
-
-    // check early exit due to failed service check
-    if token.is_cancelled() {
-        log::warn!("Not launching node due to early exit, bye!");
-        return Ok(());
-    }
+    }?;
+    log::warn!("Using models: {:#?}", config.workflows.models);
 
     // create the node
-    let (mut node, p2p, mut worker_batch, mut worker_single) = DriaComputeNode::new(config).await?;
+    let (mut node, p2p, worker_batch, worker_single) = DriaComputeNode::new(config).await?;
 
+    // spawn threads
     log::info!("Spawning peer-to-peer client thread.");
-    let p2p_handle = tokio::spawn(async move { p2p.run().await });
+    task_tracker.spawn(async move { p2p.run().await });
 
-    log::info!("Spawning workflows batch worker thread.");
-    let worker_batch_handle = tokio::spawn(async move { worker_batch.run_batch().await });
+    if let Some(mut worker_batch) = worker_batch {
+        log::info!("Spawning workflows batch worker thread.");
+        task_tracker.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 });
+    if let Some(mut worker_single) = worker_single {
+        log::info!("Spawning workflows single worker thread.");
+        task_tracker.spawn(async move { worker_single.run().await });
+    }
 
     // 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 {
+    let node_token = cancellation.clone();
+    task_tracker.spawn(async move {
         if let Err(err) = node.run(node_token).await {
             log::error!("Node launch error: {}", err);
             panic!("Node failed.")
         };
+        log::info!("Closing node.")
     });
 
-    // wait for tasks to complete
-    if let Err(err) = node_handle.await {
-        log::error!("Node 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);
-    };
+    // wait for all tasks to finish
+    task_tracker.wait().await;
+    log::info!("All tasks have exited succesfully.");
 
     log::info!("Bye!");
     Ok(())
@@ -168,7 +161,7 @@ async fn wait_for_termination(cancellation: CancellationToken) -> Result<()> {
         cancellation.cancel();
     }
 
-    log::info!("Terminating the node...");
+    log::info!("Terminating the application...");
 
     Ok(())
 }
diff --git a/compute/src/node.rs b/compute/src/node.rs
@@ -6,6 +6,7 @@ use dkn_p2p::{
     DriaP2PClient, DriaP2PCommander, DriaP2PProtocol,
 };
 use eyre::Result;
+use std::collections::HashSet;
 use tokio::{sync::mpsc, time::Duration};
 use tokio_util::{either::Either, sync::CancellationToken};
 
@@ -32,9 +33,15 @@ pub struct DriaComputeNode {
     /// 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>,
+    workflow_batch_tx: Option<mpsc::Sender<WorkflowsWorkerInput>>,
     /// Workflow transmitter to send single tasks.
-    workflow_single_tx: mpsc::Sender<WorkflowsWorkerInput>,
+    workflow_single_tx: Option<mpsc::Sender<WorkflowsWorkerInput>>,
+    // TODO: instead of piggybacking task metadata within channels, we can store them here
+    // in a hashmap alone, and then use the task_id to get the metadata when needed
+    // Single tasks hash-map
+    pending_tasks_single: HashSet<String>,
+    // Batch tasks hash-map
+    pending_tasks_batch: HashSet<String>,
 }
 
 impl DriaComputeNode {
@@ -46,8 +53,8 @@ impl DriaComputeNode {
     ) -> Result<(
         DriaComputeNode,
         DriaP2PClient,
-        WorkflowsWorker,
-        WorkflowsWorker,
+        Option<WorkflowsWorker>,
+        Option<WorkflowsWorker>,
     )> {
         // create the keypair from secret key
         let keypair = secret_to_keypair(&config.secret_key);
@@ -77,8 +84,24 @@ impl DriaComputeNode {
 
         // 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);
+
+        // check if we should create a worker for batchable workflows
+        let (workflows_batch_worker, workflow_batch_tx) = if config.workflows.has_batchable_models()
+        {
+            let worker = WorkflowsWorker::new(publish_tx.clone());
+            (Some(worker.0), Some(worker.1))
+        } else {
+            (None, None)
+        };
+
+        // check if we should create a worker for single workflows
+        let (workflows_single_worker, workflow_single_tx) =
+            if config.workflows.has_non_batchable_models() {
+                let worker = WorkflowsWorker::new(publish_tx);
+                (Some(worker.0), Some(worker.1))
+            } else {
+                (None, None)
+            };
 
         Ok((
             DriaComputeNode {
@@ -89,6 +112,8 @@ impl DriaComputeNode {
                 publish_rx,
                 workflow_batch_tx,
                 workflow_single_tx,
+                pending_tasks_single: HashSet::new(),
+                pending_tasks_batch: HashSet::new(),
             },
             p2p_client,
             workflows_batch_worker,
@@ -119,10 +144,10 @@ impl DriaComputeNode {
     }
 
     /// Returns the task count within the channels, `single` and `batch`.
-    pub fn get_active_task_count(&self) -> [usize; 2] {
+    pub fn get_pending_task_count(&self) -> [usize; 2] {
         [
-            self.workflow_single_tx.max_capacity() - self.workflow_single_tx.capacity(),
-            self.workflow_batch_tx.max_capacity() - self.workflow_batch_tx.capacity(),
+            self.pending_tasks_single.len(),
+            self.pending_tasks_batch.len(),
         ]
     }
 
@@ -202,10 +227,32 @@ impl DriaComputeNode {
                             // we got acceptance, so something was not right about the workflow and we can ignore it
                             Ok(Either::Left(acceptance)) => Ok(acceptance),
                             // 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,
+                            Ok(Either::Right(workflow_message)) => {
+                                if let Err(e) = match workflow_message.batchable {
+                                    // this is a batchable task, send it to batch worker
+                                    // and keep track of the task id in pending tasks
+                                    true => match self.workflow_batch_tx {
+                                        Some(ref mut tx) => {
+                                            self.pending_tasks_batch
+                                                .insert(workflow_message.task_id.clone());
+                                            tx.send(workflow_message).await
+                                        }
+                                        None => unreachable!(
+                                            "Batchable workflow received but no worker available."
+                                        ),
+                                    },
+                                    // this is a single task, send it to single worker
+                                    // and keep track of the task id in pending tasks
+                                    false => match self.workflow_single_tx {
+                                        Some(ref mut tx) => {
+                                            self.pending_tasks_single
+                                                .insert(workflow_message.task_id.clone());
+                                            tx.send(workflow_message).await
+                                        }
+                                        None => unreachable!(
+                                            "Single workflow received but no worker available."
+                                        ),
+                                    },
                                 } {
                                     log::error!("Error sending workflow message: {:?}", e);
                                 };
@@ -266,18 +313,25 @@ 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_rx.recv() => {
-                    if let Some(result) = publish_msg {
-                        WorkflowHandler::handle_publish(self, result).await?;
+                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
+                        match publish_msg.batchable {
+                            true => self.pending_tasks_batch.remove(&publish_msg.task_id),
+                            false => self.pending_tasks_single.remove(&publish_msg.task_id),
+                        };
+
+                        // publish the message
+                        WorkflowHandler::handle_publish(self, publish_msg).await?;
                     } else {
                         log::error!("Publish channel closed unexpectedly.");
                         break;
                     };
                 },
                 // a GossipSub message is received from the channel
                 // this is expected to be sent by the p2p client
-                gossipsub_msg = self.message_rx.recv() => {
-                    if let Some((peer_id, message_id, message)) = gossipsub_msg {
+                gossipsub_msg_opt = self.message_rx.recv() => {
+                    if let Some((peer_id, message_id, message)) = gossipsub_msg_opt {
                         // handle the message, returning a message acceptance for the received one
                         let acceptance = self.handle_message((peer_id, &message_id, message)).await;
 
@@ -332,8 +386,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_pending_task_count();
+        log::info!("Pending Task Count (single/batch): {} / {}", single, batch);
     }
 
     /// Updates the local list of available nodes by refreshing it.
diff --git a/compute/src/workers/workflow.rs b/compute/src/workers/workflow.rs
@@ -13,6 +13,7 @@ pub struct WorkflowsWorkerInput {
     pub task_id: String,
     pub model_name: String,
     pub stats: TaskStats,
+    pub batchable: bool,
 }
 
 pub struct WorkflowsWorkerOutput {
@@ -22,6 +23,7 @@ pub struct WorkflowsWorkerOutput {
     pub task_id: String,
     pub model_name: String,
     pub stats: TaskStats,
+    pub batchable: bool,
 }
 
 pub struct WorkflowsWorker {
@@ -217,6 +219,7 @@ impl WorkflowsWorker {
             public_key: input.public_key,
             task_id: input.task_id,
             model_name: input.model_name,
+            batchable: input.batchable,
             stats: input.stats.record_execution_time(started_at),
         }
     }
diff --git a/workflows/src/config.rs b/workflows/src/config.rs

Original file line number	Diff line number	Diff line change
`@@ -13,6 +13,7 @@ pub struct WorkflowsWorkerInput {`
`13`	`13`	`pub task_id: String,`
`14`	`14`	`pub model_name: String,`
`15`	`15`	`pub stats: TaskStats,`
	`16`	`+ pub batchable: bool,`
`16`	`17`	`}`
`17`	`18`
`18`	`19`	`pub struct WorkflowsWorkerOutput {`
`@@ -22,6 +23,7 @@ pub struct WorkflowsWorkerOutput {`
`22`	`23`	`pub task_id: String,`
`23`	`24`	`pub model_name: String,`
`24`	`25`	`pub stats: TaskStats,`
	`26`	`+ pub batchable: bool,`
`25`	`27`	`}`
`26`	`28`
`27`	`29`	`pub struct WorkflowsWorker {`
`@@ -217,6 +219,7 @@ impl WorkflowsWorker {`
`217`	`219`	`public_key: input.public_key,`
`218`	`220`	`task_id: input.task_id,`
`219`	`221`	`model_name: input.model_name,`
	`222`	`+ batchable: input.batchable,`
`220`	`223`	`stats: input.stats.record_execution_time(started_at),`
`221`	`224`	`}`
`222`	`225`	`}`