better error reporting

Author: erhant

compute/src/reqres/task.rs

@@ -1,7 +1,9 @@
 use colored::Colorize;
-use dkn_executor::TaskBody;
+use dkn_executor::{CompletionError, ModelProvider, PromptError, TaskBody};
 use dkn_p2p::libp2p::request_response::ResponseChannel;
-use dkn_utils::payloads::{TaskRequestPayload, TaskResponsePayload, TaskStats, TASK_RESULT_TOPIC};
+use dkn_utils::payloads::{
+    TaskError, TaskRequestPayload, TaskResponsePayload, TaskStats, TASK_RESULT_TOPIC,
+};
 use dkn_utils::DriaMessage;
 use eyre::{Context, Result};
 
@@ -25,27 +27,23 @@ impl TaskResponder {
         let task = compute_message
             .parse_payload::<TaskRequestPayload<serde_json::Value>>()
             .wrap_err("could not parse task request payload")?;
-        let task_body = match serde_json::from_value::<TaskBody>(task.input)
-            .wrap_err("could not parse task body")
-        {
+        let task_body = match serde_json::from_value::<TaskBody>(task.input) {
             Ok(task_body) => task_body,
             Err(err) => {
-                let err_string = format!("{:#}", err);
                 log::error!(
-                    "Task {}/{} failed due to parsing error: {}",
+                    "Task {}/{} failed due to parsing error: {err}",
                     task.file_id,
                     task.row_id,
-                    err_string
                 );
 
                 // prepare error payload
                 let error_payload = TaskResponsePayload {
                     result: None,
-                    error: Some(err_string),
+                    error: Some(TaskError::ParseError(err.to_string())),
                     row_id: task.row_id,
                     file_id: task.file_id,
                     task_id: task.task_id,
-                    model: Default::default(),
+                    model: "<n/a>".to_string(), // no model available due to parsing error
                     stats: TaskStats::new(),
                 };
 
@@ -56,7 +54,8 @@ impl TaskResponder {
                 let response = node.new_message(error_payload_str, TASK_RESULT_TOPIC);
                 node.p2p.respond(response.into(), channel).await?;
 
-                return Err(err);
+                // return with error
+                eyre::bail!("could not parse task body: {err}")
             }
         };
 
@@ -75,7 +74,7 @@ impl TaskResponder {
         let task_metadata = TaskWorkerMetadata {
             task_id: task.task_id,
             file_id: task.file_id,
-            model_name: task_body.model.to_string(),
+            model: task_body.model,
             channel,
         };
         let task_input = TaskWorkerInput {
@@ -112,7 +111,7 @@ impl TaskResponder {
                     file_id: task_metadata.file_id,
                     task_id: task_metadata.task_id,
                     row_id: task_output.row_id,
-                    model: task_metadata.model_name,
+                    model: task_metadata.model.to_string(),
                     stats: task_output
                         .stats
                         .record_published_at()
@@ -125,22 +124,21 @@ impl TaskResponder {
             }
             Err(err) => {
                 // use pretty display string for error logging with causes
-                let err_string = format!("{:#}", err);
                 log::error!(
-                    "Task {}/{} failed: {}",
+                    "Task {}/{} failed: {:#}",
                     task_metadata.file_id,
                     task_output.row_id,
-                    err_string
+                    err
                 );
 
                 // prepare error payload
                 let error_payload = TaskResponsePayload {
                     result: None,
-                    error: Some(err_string),
+                    error: Some(map_prompt_error(task_metadata.model.provider(), err)),
                     row_id: task_output.row_id,
                     file_id: task_metadata.file_id,
                     task_id: task_metadata.task_id,
-                    model: task_metadata.model_name,
+                    model: task_metadata.model.to_string(),
                     stats: task_output
                         .stats
                         .record_published_at()
@@ -161,3 +159,104 @@ impl TaskResponder {
         Ok(())
     }
 }
+
+/// Maps a [`PromptError`] to a [`DriaExecutorError`] with respect to the given provider.
+fn map_prompt_error(provider: ModelProvider, err: PromptError) -> TaskError {
+    if let PromptError::CompletionError(CompletionError::ProviderError(err_inner)) = &err {
+        /// A wrapper for `{ error: T }` to match the provider error format.
+        #[derive(Clone, serde::Deserialize)]
+        struct ErrorObject<T> {
+            error: T,
+        }
+
+        match provider {
+            ModelProvider::Gemini => {
+                /// Gemini API [error object](https://github.com/googleapis/go-genai/blob/main/api_client.go#L273).
+                #[derive(Clone, serde::Deserialize)]
+                pub struct GeminiError {
+                    code: u32,
+                    message: String,
+                    status: String,
+                }
+
+                serde_json::from_str::<ErrorObject<GeminiError>>(err_inner).map(
+                    |ErrorObject {
+                         error: gemini_error,
+                     }| TaskError::ProviderError {
+                        code: format!("{} ({})", gemini_error.code, gemini_error.status),
+                        message: gemini_error.message,
+                        provider: provider.to_string(),
+                    },
+                )
+            }
+            ModelProvider::OpenAI => {
+                /// OpenAI API [error object](https://github.com/openai/openai-go/blob/main/internal/apierror/apierror.go#L17).
+                #[derive(Clone, serde::Deserialize)]
+                pub struct OpenAIError {
+                    code: String,
+                    message: String,
+                }
+
+                serde_json::from_str::<ErrorObject<OpenAIError>>(err_inner).map(
+                    |ErrorObject {
+                         error: openai_error,
+                     }| TaskError::ProviderError {
+                        code: openai_error.code,
+                        message: openai_error.message,
+                        provider: provider.to_string(),
+                    },
+                )
+            }
+            ModelProvider::OpenRouter => {
+                /// OpenRouter API [error object](https://openrouter.ai/docs/api-reference/errors).
+                #[derive(Clone, serde::Deserialize)]
+                pub struct OpenRouterError {
+                    code: u32,
+                    message: String,
+                }
+
+                serde_json::from_str::<ErrorObject<OpenRouterError>>(err_inner).map(
+                    |ErrorObject {
+                         error: openrouter_error,
+                     }| {
+                        TaskError::ProviderError {
+                            code: openrouter_error.code.to_string(),
+                            message: openrouter_error.message,
+                            provider: provider.to_string(),
+                        }
+                    },
+                )
+            }
+            ModelProvider::Ollama => serde_json::from_str::<ErrorObject<String>>(err_inner).map(
+                // Ollama just returns a string error message
+                |ErrorObject {
+                     error: ollama_error,
+                 }| {
+                    // based on the error message, we can come up with out own "dummy" codes
+                    let code = if ollama_error.contains("server busy, please try again.") {
+                        "server_busy"
+                    } else if ollama_error.contains("model requires more system memory") {
+                        "model_requires_more_memory"
+                    } else if ollama_error.contains("cudaMalloc failed: out of memory") {
+                        "cuda_malloc_failed"
+                    } else if ollama_error.contains("CUDA error: out of memory") {
+                        "cuda_oom"
+                    } else {
+                        "unknown"
+                    };
+
+                    TaskError::ProviderError {
+                        code: code.to_string(),
+                        message: ollama_error,
+                        provider: provider.to_string(),
+                    }
+                },
+            ),
+        }
+        // if we couldn't parse it, just return a generic prompt error
+        .unwrap_or(TaskError::Other(err.to_string()))
+    } else {
+        // not a provider error, fallback to generic prompt error
+        TaskError::Other(err.to_string())
+    }
+}

compute/src/workers/task.rs

@@ -1,5 +1,5 @@
 use colored::Colorize;
-use dkn_executor::{DriaExecutor, TaskBody};
+use dkn_executor::{DriaExecutor, Model, TaskBody};
 use dkn_p2p::libp2p::request_response::ResponseChannel;
 use dkn_utils::payloads::TaskStats;
 use tokio::sync::mpsc;
@@ -9,7 +9,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 model: Model,
     pub task_id: String,
     pub file_id: Uuid,
     /// If for any reason this object is dropped before `channel` is responded to,

utils/src/payloads/mod.rs

@@ -1,5 +1,5 @@
 mod tasks;
-pub use tasks::{TaskRequestPayload, TaskResponsePayload, TaskStats};
+pub use tasks::{TaskError, TaskRequestPayload, TaskResponsePayload, TaskStats};
 pub use tasks::{TASK_REQUEST_TOPIC, TASK_RESULT_TOPIC};
 
 mod heartbeat;

utils/src/payloads/tasks.rs

@@ -33,11 +33,11 @@ pub struct TaskResponsePayload {
     /// If this is `None`, the task failed, and you should check the `error` field.
     #[serde(skip_serializing_if = "Option::is_none")]
     pub result: Option<String>,
-    /// An error message, if any.
+    /// An error, if any.
     ///
     /// If this is `Some`, you can ignore the `result` field.
     #[serde(skip_serializing_if = "Option::is_none")]
-    pub error: Option<String>,
+    pub error: Option<TaskError>,
 }
 
 /// A generic task request, given by Dria.
@@ -57,6 +57,35 @@ pub struct TaskRequestPayload<T> {
     pub input: T,
 }
 
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub enum TaskError {
+    /// A parse error occurred while parsing the task request or response.
+    ParseError(String),
+    /// An error returned from the model provider.
+    ProviderError {
+        /// Not necessarily an HTTP status code, but a code that the provider uses to identify the error.
+        ///
+        /// For example, OpenAI uses a string code like "invalid_request_error".
+        code: String,
+        /// The error message returned by the provider.
+        ///
+        /// May contain additional information about the error.
+        message: String,
+        /// The source of the error.
+        ///
+        /// Can be a provider name, or RPC etc.
+        provider: String,
+    },
+    /// The task request had failed for some network reason.
+    OutboundRequestError {
+        code: String,
+        /// The error message returned by the network.
+        message: String,
+    },
+    /// An error that returned by executor.
+    Other(String),
+}
+
 /// Task stats for diagnostics.
 ///
 /// Returning this as the payload helps to debug the errors received at client side, and latencies.