Merge pull request #198 from firstbatchxyz/erhant/better-error-reporting

feat: better error reporting

Author: erhant

Cargo.lock

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

Cargo.toml

@@ -86,21 +86,27 @@ async fn main() -> Result<()> {
 
     // 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! {
+    let model_perf = tokio::select! {
         result = config.executors.check_services() => result,
         _ = cancellation.cancelled() => {
             log::info!("Service check cancelled, exiting.");
             return Ok(());
         }
     }?;
-    log::warn!(
-        "Using models: {}",
-        config.executors.get_model_names().join(", ")
+    log::info!(
+        "Using models: {}\n{}",
+        config.executors.get_model_names().join(", "),
+        model_perf
+            .iter()
+            .map(|(model, perf)| format!("{}: {}", model, perf))
+            .collect::<Vec<_>>()
+            .join("\n")
     );
 
     // create the node
     let batch_size = config.batch_size;
-    let (mut node, p2p, worker_batch, worker_single) = DriaComputeNode::new(config).await?;
+    let (mut node, p2p, worker_batch, worker_single) =
+        DriaComputeNode::new(config, model_perf).await?;
 
     // spawn p2p client first
     log::info!("Spawning peer-to-peer client thread.");

compute/src/main.rs

@@ -76,7 +76,16 @@ impl DriaComputeNode {
                 _ = diagnostic_refresh_interval.tick() => self.handle_diagnostic_refresh().await,
 
                 // check RPC, and get a new one if we are disconnected
-                _ = rpc_liveness_refresh_interval.tick() => self.handle_rpc_liveness_check().await,
+                _ = rpc_liveness_refresh_interval.tick() => {
+                    let is_connected = self.handle_rpc_liveness_check().await;
+                    if !is_connected {
+                        // make sure we reset the heartbeat and specs intervals so that
+                        // we dont wait the entire duration for this new connection
+                        log::info!("Connecting was re-attempted, resetting timers.");
+                        heartbeat_interval.reset_after(Duration::from_secs(5));
+                        specs_interval.reset_after(Duration::from_secs(5));
+                    }
+                },
 
                 // log points every now and then
                 _ = points_refresh_interval.tick() => self.handle_points_refresh().await,

compute/src/node/core.rs

@@ -85,9 +85,10 @@ impl DriaComputeNode {
 
     /// Dials the existing RPC node if we are not connected to it.
     ///
-    /// If there is an error while doing that,
-    /// it will try to get a new RPC node and dial it.
-    pub(crate) async fn handle_rpc_liveness_check(&mut self) {
+    /// If there is an error while doing that, it will try to get a new RPC node and dial it.
+    ///
+    /// Returns `true` if the RPC is connected, `false` otherwise.
+    pub(crate) async fn handle_rpc_liveness_check(&mut self) -> bool {
         log::debug!("Checking RPC connections for diagnostics.");
 
         // check if we are connected
@@ -124,6 +125,9 @@ impl DriaComputeNode {
         } else {
             log::debug!("Connection with {} is intact.", self.dria_rpc.peer_id);
         }
+
+        // return the connection status
+        is_connected
     }
 
     /// Updates the points for the given address.

compute/src/node/diagnostic.rs

@@ -1,7 +1,8 @@
+use dkn_executor::Model;
 use dkn_p2p::{
     libp2p::PeerId, DriaP2PClient, DriaP2PCommander, DriaP2PProtocol, DriaReqResMessage,
 };
-use dkn_utils::crypto::secret_to_keypair;
+use dkn_utils::{crypto::secret_to_keypair, payloads::SpecModelPerformance};
 use eyre::Result;
 use std::collections::{HashMap, HashSet};
 use tokio::sync::mpsc;
@@ -68,6 +69,7 @@ impl DriaComputeNode {
     /// Returns the node instance and p2p client together. P2p MUST be run in a separate task before this node is used at all.
     pub async fn new(
         mut config: DriaComputeNodeConfig,
+        model_perf: HashMap<Model, SpecModelPerformance>,
     ) -> Result<(
         DriaComputeNode,
         DriaP2PClient,
@@ -124,7 +126,7 @@ impl DriaComputeNode {
         let model_names = config.executors.get_model_names();
         let points_client = DriaPointsClient::new(&config.address, &config.network)?;
 
-        let spec_collector = SpecCollector::new(model_names.clone(), config.version);
+        let spec_collector = SpecCollector::new(model_names.clone(), model_perf, config.version);
         Ok((
             DriaComputeNode {
                 config,

compute/src/node/mod.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,24 @@ 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_to_task_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 +162,111 @@ impl TaskResponder {
         Ok(())
     }
 }
+
+/// Maps a [`PromptError`] to a [`TaskError`] with respect to the given provider.
+fn map_prompt_error_to_task_error(provider: ModelProvider, err: PromptError) -> TaskError {
+    match &err {
+        // if the error is a provider error, we can try to parse it
+        PromptError::CompletionError(CompletionError::ProviderError(err_inner)) => {
+            /// 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::ExecutorError(err_inner.clone()))
+        }
+        // if its a http error, we can try to parse it as well
+        PromptError::CompletionError(CompletionError::HttpError(err_inner)) => {
+            TaskError::HttpError(err_inner.to_string())
+        }
+        // if it's not a completion error, we just return the error as is
+        err => TaskError::Other(err.to_string()),
+    }
+}