propagate service overloaded error

Author: PeaBrane

lib/llm/src/http/service/openai.rs

@@ -108,6 +108,24 @@ impl ErrorMessage {
     /// If successful, it will return the [`HttpError`] as an [`ErrorMessage::internal_server_error`]
     /// with the details of the error.
     pub fn from_anyhow(err: anyhow::Error, alt_msg: &str) -> ErrorResponse {
+        // First check for PipelineError::ServiceOverloaded
+        if let Some(pipeline_err) =
+            err.downcast_ref::<dynamo_runtime::pipeline::error::PipelineError>()
+        {
+            if matches!(
+                pipeline_err,
+                dynamo_runtime::pipeline::error::PipelineError::ServiceOverloaded(_)
+            ) {
+                return (
+                    StatusCode::SERVICE_UNAVAILABLE,
+                    Json(ErrorMessage {
+                        error: pipeline_err.to_string(),
+                    }),
+                );
+            }
+        }
+
+        // Then check for HttpError
         match err.downcast::<HttpError>() {
             Ok(http_error) => ErrorMessage::from_http_error(http_error),
             Err(err) => ErrorMessage::internal_server_error(&format!("{alt_msg}: {err}")),
@@ -1150,6 +1168,22 @@ mod tests {
         );
     }
 
+    #[test]
+    fn test_service_overloaded_error_response_from_anyhow() {
+        use dynamo_runtime::pipeline::error::PipelineError;
+
+        let err: anyhow::Error = PipelineError::ServiceOverloaded(
+            "All workers are busy, please retry later".to_string(),
+        )
+        .into();
+        let (status, response) = ErrorMessage::from_anyhow(err, BACKUP_ERROR_MESSAGE);
+        assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
+        assert_eq!(
+            response.error,
+            "Service temporarily unavailable: All workers are busy, please retry later"
+        );
+    }
+
     #[test]
     fn test_validate_input_is_text_only_accepts_text() {
         let request = make_base_request();

lib/runtime/src/pipeline/error.rs

@@ -131,6 +131,10 @@ pub enum PipelineError {
 
     #[error("NATS KV Err: {0} for bucket '{1}")]
     KeyValueError(String, String),
+
+    /// All instances are busy and cannot handle new requests
+    #[error("Service temporarily unavailable: {0}")]
+    ServiceOverloaded(String),
 }
 
 #[derive(Debug, thiserror::Error)]

lib/runtime/src/pipeline/network/egress/push_router.rs

@@ -7,7 +7,8 @@ use crate::{
     component::{Client, Endpoint, InstanceSource},
     engine::{AsyncEngine, Data},
     pipeline::{
-        error::PipelineErrorExt, AddressedPushRouter, AddressedRequest, Error, ManyOut, SingleIn,
+        error::{PipelineError, PipelineErrorExt},
+        AddressedPushRouter, AddressedRequest, Error, ManyOut, SingleIn,
     },
     protocols::maybe_error::MaybeError,
     traits::DistributedRuntimeProvider,
@@ -210,7 +211,10 @@ where
             // Check if we actually have any instances at all
             let all_instances = self.client.instance_ids();
             if !all_instances.is_empty() {
-                return Err(anyhow::anyhow!("All workers are busy, please retry later"));
+                return Err(PipelineError::ServiceOverloaded(
+                    "All workers are busy, please retry later".to_string(),
+                )
+                .into());
             }
         }
 
@@ -356,7 +360,7 @@ where
                             let busy_instances: Vec<i64> = states
                                 .iter()
                                 .filter_map(|(&id, state)| {
-                                    state.is_busy(BUSY_THRESHOLD).then(|| id)
+                                    state.is_busy(BUSY_THRESHOLD).then_some(id)
                                 })
                                 .collect();
                             drop(states);