Merge pull request #197 from firstbatchxyz/erhant/task-id-is-string

feat: task id is now string

Author: erhant

Cargo.lock

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

Cargo.toml

@@ -48,9 +48,9 @@ pub struct DriaComputeNode {
     task_request_batch_tx: Option<mpsc::Sender<TaskWorkerInput>>,
     /// Task worker transmitter to send single tasks.
     task_request_single_tx: Option<mpsc::Sender<TaskWorkerInput>>,
-    // Single tasks, key is `row_id`
+    /// Single tasks, key is `row_id`, which has negligible probability of collision.
     pub pending_tasks_single: HashMap<Uuid, TaskWorkerMetadata>,
-    // Batchable tasks, key is `row_id`
+    // Batchable tasks, key is `row_id`, which has negligible probability of collision.
     pub pending_tasks_batch: HashMap<Uuid, TaskWorkerMetadata>,
     /// Completed single tasks count
     completed_tasks_single: usize,

compute/src/node/mod.rs

@@ -34,7 +34,7 @@ impl TaskResponder {
                 log::error!(
                     "Task {}/{} failed due to parsing error: {}",
                     task.file_id,
-                    task.task_id,
+                    task.row_id,
                     err_string
                 );
 
@@ -106,7 +106,7 @@ impl TaskResponder {
                     "Publishing {} result for {}/{}",
                     "task".yellow(),
                     task_metadata.file_id,
-                    task_metadata.task_id
+                    task_output.row_id
                 );
 
                 // TODO: will get better token count from `TaskWorkerOutput`
@@ -134,7 +134,7 @@ impl TaskResponder {
                 log::error!(
                     "Task {}/{} failed: {}",
                     task_metadata.file_id,
-                    task_metadata.task_id,
+                    task_output.row_id,
                     err_string
                 );
 

compute/src/reqres/task.rs

@@ -10,7 +10,7 @@ use uuid::Uuid;
 /// This is put into a map before execution, and then removed after the task is done.
 pub struct TaskWorkerMetadata {
     pub model_name: String,
-    pub task_id: Uuid,
+    pub task_id: String,
     pub file_id: Uuid,
     /// If for any reason this object is dropped before `channel` is responded to,
     /// the task will be lost and the channel will be abruptly closed, causing an error on

compute/src/workers/task.rs

@@ -12,15 +12,18 @@ pub const TASK_RESULT_TOPIC: &str = "results";
 /// `result` and `error` are mutually-exclusive, only one of them can be `Some`:
 /// - if `result` is `Some`, then it contains the result.
 /// - if `error` is `Some`, then it contains the error message.
+///
+/// Each task belongs to a file (uniquely identified by `file_id`), and has a unique identifier (`row_id`).
+/// THe `task_id` is a custom identifier given by a user.
 #[derive(Debug, Clone, Serialize, Deserialize)]
 #[serde(rename_all = "camelCase")]
 pub struct TaskResponsePayload {
-    /// The unique identifier of the task.
-    pub row_id: Uuid,
-    /// The custom identifier of the task.
-    pub task_id: Uuid,
     /// The file that this task is associated with.
     pub file_id: Uuid,
+    /// The unique identifier of the task.
+    pub row_id: Uuid,
+    /// The custom identifier of the task, not necessarily unique.
+    pub task_id: String,
     /// Name of the model used for this task.
     pub model: String,
     /// Stats about the task execution.
@@ -38,15 +41,18 @@ pub struct TaskResponsePayload {
 }
 
 /// A generic task request, given by Dria.
+///
+/// Each task belongs to a file (uniquely identified by `file_id`), and has a unique identifier (`row_id`).
+/// THe `task_id` is a custom identifier given by a user.
 #[derive(Debug, Clone, Serialize, Deserialize)]
 #[serde(rename_all = "camelCase")]
 pub struct TaskRequestPayload<T> {
-    /// The unique identifier of the task.
-    pub row_id: Uuid,
-    /// The unique identifier of the task.
-    pub task_id: Uuid,
     /// The file that this task is associated with.
     pub file_id: Uuid,
+    /// The unique identifier of the task.
+    pub row_id: Uuid,
+    /// The custom identifier of the task, not necessarily unique.
+    pub task_id: String,
     /// The input to the compute function.
     pub input: T,
 }