slight log and naming rfks, update executor impls

Author: erhant

Cargo.lock

@@ -7,7 +7,7 @@ default-members = ["compute"]
 
 [workspace.package]
 edition = "2021"
-version = "0.5.5"
+version = "0.5.6"
 license = "Apache-2.0"
 readme = "README.md"
 

Cargo.toml

@@ -69,11 +69,8 @@ async fn main() -> Result<()> {
     });
 
     // create configurations
-    let model_names = dkn_utils::split_csv_line(&env::var("DKN_MODELS").unwrap_or_default())
-        .into_iter()
-        .filter_map(|s| Model::try_from(s.as_str()).ok())
-        .collect::<Vec<_>>();
-    let executors_config = DriaExecutorsManager::new_from_env_for_models(model_names)?;
+    let models = Model::from_csv(env::var("DKN_MODELS").unwrap_or_default());
+    let executors_config = DriaExecutorsManager::new_from_env_for_models(models.into_iter())?;
     if executors_config.models.is_empty() {
         return Err(eyre::eyre!("No models were provided, make sure to restart with at least one model provided within DKN_MODELS."));
     }

compute/src/main.rs

@@ -100,7 +100,7 @@ impl DriaComputeNode {
             protocol,
         )?;
 
-        // create workflow workers, all workers use the same publish channel
+        // create channel for task executors, all workers use the same publish channel
         let (publish_tx, publish_rx) = mpsc::channel(PUBLISH_CHANNEL_BUFSIZE);
 
         // check if we should create a worker for batch executor

compute/src/node/mod.rs

@@ -4,8 +4,11 @@ use dkn_p2p::libp2p::{
     PeerId,
 };
 use dkn_p2p::DriaReqResMessage;
-use dkn_utils::payloads::{HEARTBEAT_TOPIC, SPECS_TOPIC, TASK_REQUEST_TOPIC};
-use eyre::{eyre, Result};
+use dkn_utils::{
+    payloads::{HEARTBEAT_TOPIC, SPECS_TOPIC, TASK_REQUEST_TOPIC},
+    DriaMessage,
+};
+use eyre::Result;
 
 use crate::{reqres::*, workers::task::TaskWorkerOutput};
 
@@ -33,7 +36,7 @@ impl DriaComputeNode {
                 if self.dria_rpc.peer_id != peer_id {
                     log::warn!("Received request from unauthorized source: {}", peer_id);
                     log::debug!("Allowed source: {}", self.dria_rpc.peer_id);
-                } else if let Err(e) = self.handle_request(peer_id, request, channel).await {
+                } else if let Err(e) = self.handle_request(peer_id, &request, channel).await {
                     log::error!("Error handling request: {:?}", e);
                 }
             }
@@ -61,6 +64,11 @@ impl DriaComputeNode {
         request_id: OutboundRequestId,
         data: Vec<u8>,
     ) -> Result<()> {
+        if peer_id != self.dria_rpc.peer_id {
+            log::warn!("Received response from unauthorized source: {}", peer_id);
+            log::debug!("Allowed source: {}", self.dria_rpc.peer_id);
+        }
+
         if let Ok(heartbeat_response) = HeartbeatRequester::try_parse_response(&data) {
             log::info!(
                 "Received a {} response ({request_id}) from {peer_id}",
@@ -85,14 +93,18 @@ impl DriaComputeNode {
     async fn handle_request(
         &mut self,
         peer_id: PeerId,
-        data: Vec<u8>,
+        message_data: &[u8],
         channel: ResponseChannel<Vec<u8>>,
     ) -> Result<()> {
-        if let Ok(task_request) = TaskResponder::try_parse_request(&data) {
-            self.handle_task_request(peer_id, task_request, channel)
-                .await
-        } else {
-            Err(eyre::eyre!("Received unhandled request from {peer_id}"))
+        let message = DriaMessage::from_slice_checked(
+            message_data,
+            self.p2p.protocol().name.clone(),
+            self.config.version,
+        )?;
+
+        match message.topic.as_str() {
+            TASK_REQUEST_TOPIC => self.handle_task_request(peer_id, message, channel).await,
+            _ => Err(eyre::eyre!("Received unhandled request from {peer_id}")),
         }
     }
 
@@ -113,7 +125,7 @@ impl DriaComputeNode {
         );
 
         let (task_input, task_metadata) =
-            TaskResponder::prepare_worker_input(self, &task_request, channel).await?;
+            TaskResponder::parse_task_request(self, &task_request, channel).await?;
         if let Err(e) = match task_input.task.is_batchable() {
             // this is a batchable task, send it to batch worker
             // and keep track of the task id in pending tasks
@@ -123,11 +135,7 @@ impl DriaComputeNode {
                         .insert(task_input.row_id, task_metadata);
                     tx.send(task_input).await
                 }
-                None => {
-                    return Err(eyre!(
-                        "Batchable workflow received but no worker available."
-                    ));
-                }
+                None => eyre::bail!("Batchable task received but no worker available."),
             },
 
             // this is a single task, send it to single worker
@@ -138,12 +146,10 @@ impl DriaComputeNode {
                         .insert(task_input.row_id, task_metadata);
                     tx.send(task_input).await
                 }
-                None => {
-                    return Err(eyre!("Single workflow received but no worker available."));
-                }
+                None => eyre::bail!("Single task received but no worker available."),
             },
         } {
-            log::error!("Error sending workflow message: {:?}", e);
+            log::error!("Could not send task to worker: {:?}", e);
         };
 
         Ok(())
@@ -165,7 +171,7 @@ impl DriaComputeNode {
         // respond to the response channel with the result
         match task_metadata {
             Some(task_metadata) => {
-                TaskResponder::send_output(self, task_response, task_metadata).await?;
+                TaskResponder::send_task_output(self, task_response, task_metadata).await?;
             }
             None => {
                 // totally unexpected case, wont happen at all

compute/src/node/reqres.rs

@@ -11,60 +11,54 @@ use crate::DriaComputeNode;
 pub struct TaskResponder;
 
 impl super::IsResponder for TaskResponder {
-    type Request = DriaMessage; // TODO: TaskRequestPayload<TaskWorkflow>;
-    type Response = DriaMessage; // TODO: TaskResponsePayload;
+    type Request = DriaMessage; // TODO: can we do this typed?
+    type Response = DriaMessage; // TODO: can we do this typed?
 }
 
 impl TaskResponder {
-    pub(crate) async fn prepare_worker_input(
+    pub(crate) async fn parse_task_request(
         node: &mut DriaComputeNode,
         compute_message: &DriaMessage,
         channel: ResponseChannel<Vec<u8>>,
     ) -> Result<(TaskWorkerInput, TaskWorkerMetadata)> {
-        // parse payload
         let task = compute_message
             .parse_payload::<TaskRequestPayload<TaskBody>>()
-            .wrap_err("could not parse workflow task")?;
-        log::info!("Handling task {}", task.row_id);
-
-        // record received time
+            .wrap_err("could not parse task payload")?;
         let stats = TaskStats::new().record_received_at();
-
         log::info!(
-            "Using model {} for {} {}",
-            task.input.model.to_string().yellow(),
+            "Handling {} {} with model {}",
             "task".yellow(),
-            task.row_id
+            task.row_id,
+            task.input.model.to_string().yellow()
         );
-        let task_body = task.input;
 
         // check if the model is available in this node, if so
         // it will return an executor that can run this model
         let executor = node
             .config
             .executors
-            .get_executor(&task_body.model)
+            .get_executor(&task.input.model)
             .await
             .wrap_err("could not get an executor")?;
 
         let task_metadata = TaskWorkerMetadata {
             task_id: task.task_id,
             file_id: task.file_id,
-            model_name: task_body.model.to_string(),
+            model_name: task.input.model.to_string(),
             channel,
         };
         let task_input = TaskWorkerInput {
             executor,
-            task: task_body,
+            task: task.input,
             row_id: task.row_id,
             stats,
         };
 
         Ok((task_input, task_metadata))
     }
 
-    /// Handles the result of a workflow task.
-    pub(crate) async fn send_output(
+    /// Handles the result of a task.
+    pub(crate) async fn send_task_output(
         node: &mut DriaComputeNode,
         task_output: TaskWorkerOutput,
         task_metadata: TaskWorkerMetadata,

compute/src/reqres/task.rs

@@ -40,14 +40,14 @@ pub struct TaskWorkerOutput {
 
 /// It is expected to be spawned in another thread, with [`Self::run_batch`] for batch processing and [`Self::run_series`] for single processing.
 pub struct TaskWorker {
-    /// Workflow message channel receiver, the sender is most likely the compute node itself.
+    /// Task channel receiver, the sender is most likely the compute node itself.
     task_rx: mpsc::Receiver<TaskWorkerInput>,
     /// Publish message channel sender, the receiver is most likely the compute node itself.
     publish_tx: mpsc::Sender<TaskWorkerOutput>,
     // TODO: batch size must be defined here
 }
 
-/// Buffer size for workflow tasks (per worker).
+/// Buffer size for task channels (per worker).
 const TASK_RX_CHANNEL_BUFSIZE: usize = 1024;
 
 impl TaskWorker {
@@ -225,11 +225,7 @@ impl TaskWorker {
     ) {
         let batchable = input.task.is_batchable();
         input.stats = input.stats.record_execution_started_at();
-        let result = input
-            .executor
-            // takes no explicit prompt input, everything is in the workflow
-            .execute(input.task)
-            .await;
+        let result = input.executor.execute(input.task).await;
         input.stats = input.stats.record_execution_ended_at();
 
         let output = TaskWorkerOutput {
@@ -240,7 +236,7 @@ impl TaskWorker {
         };
 
         if let Err(e) = publish_tx.send(output).await {
-            log::error!("Error sending workflow result: {}", e);
+            log::error!("Error sending task result: {}", e);
         }
     }
 }
@@ -255,7 +251,7 @@ mod tests {
     /// ## Run command
     ///
     /// ```sh
-    /// cargo test --package dkn-compute --lib --all-features -- workers::workflow::tests::test_executor_worker --exact --show-output --nocapture --ignored
+    /// cargo test --package dkn-compute --lib --all-features -- workers::task::tests::test_executor_worker --exact --show-output --nocapture --ignored
     /// ```
     #[tokio::test]
     #[ignore = "run manually"]
@@ -269,7 +265,7 @@ mod tests {
         let (publish_tx, mut publish_rx) = mpsc::channel(1024);
         let (mut worker, task_tx) = TaskWorker::new(publish_tx);
 
-        // create batch workflow worker
+        // create batch worker
         let worker_handle = tokio::spawn(async move {
             worker.run_batch(4).await;
         });
@@ -290,7 +286,7 @@ mod tests {
                 stats: TaskStats::default(),
             };
 
-            // send workflow to worker
+            // send task to worker
             task_tx.send(task_input).await.unwrap();
         }
 

compute/src/workers/task.rs

@@ -6,7 +6,7 @@ async fn main() -> eyre::Result<()> {
 
     let model = Model::Gemini2_0Flash;
     let models = vec![model];
-    let mut config = DriaExecutorsManager::new_from_env_for_models(models)?;
+    let mut config = DriaExecutorsManager::new_from_env_for_models(models.into_iter())?;
     config.check_services().await?;
 
     assert!(config.models.contains(&model));

executor/examples/gemini.rs

@@ -6,7 +6,7 @@ async fn main() -> eyre::Result<()> {
 
     let model = Model::Llama3_2_1bInstructQ4Km;
     let models = vec![model];
-    let mut config = DriaExecutorsManager::new_from_env_for_models(models)?;
+    let mut config = DriaExecutorsManager::new_from_env_for_models(models.into_iter())?;
     config.check_services().await?;
 
     assert!(config.models.contains(&model));

executor/examples/ollama.rs

@@ -6,7 +6,7 @@ async fn main() -> eyre::Result<()> {
 
     let model = Model::GPT4o;
     let models = vec![model];
-    let mut config = DriaExecutorsManager::new_from_env_for_models(models)?;
+    let mut config = DriaExecutorsManager::new_from_env_for_models(models.into_iter())?;
     config.check_services().await?;
 
     assert!(config.models.contains(&model));