feat: Failure Detection while Responses are returning (#1671)

Author: kthui

lib/bindings/python/rust/lib.rs

@@ -214,7 +214,7 @@ struct Endpoint {
 #[pyclass]
 #[derive(Clone)]
 struct Client {
-    router: rs::pipeline::PushRouter<serde_json::Value, serde_json::Value>,
+    router: rs::pipeline::PushRouter<serde_json::Value, RsAnnotated<serde_json::Value>>,
 }
 
 #[pyclass(eq, eq_int)]
@@ -485,13 +485,12 @@ impl Endpoint {
         let inner = self.inner.clone();
         pyo3_async_runtimes::tokio::future_into_py(py, async move {
             let client = inner.client().await.map_err(to_pyerr)?;
-            let push_router =
-                rs::pipeline::PushRouter::<serde_json::Value, serde_json::Value>::from_client(
-                    client,
-                    Default::default(),
-                )
-                .await
-                .map_err(to_pyerr)?;
+            let push_router = rs::pipeline::PushRouter::<
+                serde_json::Value,
+                RsAnnotated<serde_json::Value>,
+            >::from_client(client, Default::default())
+            .await
+            .map_err(to_pyerr)?;
             Ok(Client {
                 router: push_router,
             })
@@ -757,23 +756,13 @@ impl Client {
 }
 
 async fn process_stream(
-    stream: EngineStream<serde_json::Value>,
+    stream: EngineStream<RsAnnotated<serde_json::Value>>,
     tx: tokio::sync::mpsc::Sender<RsAnnotated<PyObject>>,
 ) {
     let mut stream = stream;
     while let Some(response) = stream.next().await {
         // Convert the response to a PyObject using Python's GIL
-        // TODO: Remove the clone, but still log the full JSON string on error. But how?
-        let annotated: RsAnnotated<serde_json::Value> = match serde_json::from_value(
-            response.clone(),
-        ) {
-            Ok(a) => a,
-            Err(err) => {
-                tracing::error!(%err, %response, "process_stream: Failed de-serializing JSON into RsAnnotated");
-                break;
-            }
-        };
-
+        let annotated: RsAnnotated<serde_json::Value> = response;
         let annotated: RsAnnotated<PyObject> = annotated.map_data(|data| {
             let result = Python::with_gil(|py| match pythonize::pythonize(py, &data) {
                 Ok(pyobj) => Ok(pyobj.into()),

lib/llm/src/protocols/common/llm_backend.rs

@@ -18,6 +18,7 @@ use serde::{Deserialize, Serialize};
 pub use super::preprocessor::PreprocessedRequest;
 pub use super::FinishReason;
 use crate::protocols::TokenIdType;
+use dynamo_runtime::protocols::maybe_error::MaybeError;
 
 pub type TokenType = Option<String>;
 pub type LogProbs = Vec<f64>;
@@ -134,6 +135,20 @@ impl LLMEngineOutput {
     }
 }
 
+impl MaybeError for LLMEngineOutput {
+    fn from_err(err: Box<dyn std::error::Error>) -> Self {
+        LLMEngineOutput::error(format!("{:?}", err))
+    }
+
+    fn err(&self) -> Option<Box<dyn std::error::Error>> {
+        if let Some(FinishReason::Error(err_msg)) = &self.finish_reason {
+            Some(anyhow::Error::msg(err_msg.clone()).into())
+        } else {
+            None
+        }
+    }
+}
+
 /// Raw output from embedding engines containing embedding vectors
 #[derive(Serialize, Deserialize, Debug, Clone, PartialEq)]
 pub struct EmbeddingsEngineOutput {
@@ -144,3 +159,26 @@ pub struct EmbeddingsEngineOutput {
     pub prompt_tokens: u32,
     pub total_tokens: u32,
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_maybe_error() {
+        let output = LLMEngineOutput::stop();
+        assert!(output.err().is_none());
+        assert!(output.is_ok());
+        assert!(!output.is_err());
+
+        let output = LLMEngineOutput::error("Test error".to_string());
+        assert_eq!(format!("{}", output.err().unwrap()), "Test error");
+        assert!(!output.is_ok());
+        assert!(output.is_err());
+
+        let output = LLMEngineOutput::from_err(anyhow::Error::msg("Test error 2").into());
+        assert_eq!(format!("{}", output.err().unwrap()), "Test error 2");
+        assert!(!output.is_ok());
+        assert!(output.is_err());
+    }
+}

lib/runtime/src/pipeline/network.rs

@@ -323,3 +323,54 @@ impl<Req: PipelineIO, Resp: PipelineIO> Ingress<Req, Resp> {
 pub trait PushWorkHandler: Send + Sync {
     async fn handle_payload(&self, payload: Bytes) -> Result<(), PipelineError>;
 }
+
+/*
+/// `NetworkStreamWrapper` is a simple wrapper used to detect proper stream termination
+/// in network communication between ingress and egress components.
+///
+/// **Purpose**: This wrapper solves the problem of detecting whether a stream ended
+/// gracefully or was cut off prematurely (e.g., due to network issues).
+///
+/// **Design Rationale**:
+/// - Cannot use `Annotated` directly because the generic type `U` varies:
+///   - Sometimes `U = Annotated<...>`
+///   - Sometimes `U = LLMEngineOutput<...>`
+/// - Using `Annotated` would require double-wrapping like `Annotated<Annotated<...>>`
+/// - A simple wrapper is cleaner and more straightforward
+///
+/// **Stream Flow**:
+/// ```
+/// At AsyncEngine:
+///   response 1 -> response 2 -> response 3 -> <end>
+///
+/// Between ingress/egress:
+///   response 1 <end=false> -> response 2 <end=false> -> response 3 <end=false> -> (null) <end=true>
+///
+/// At client:
+///   response 1 -> response 2 -> response 3 -> <end>
+/// ```
+///
+/// **Error Handling**:
+/// If the stream is cut off before proper termination, the egress is responsible for
+/// injecting an error response to communicate the incomplete stream to the client:
+/// ```
+/// At AsyncEngine:
+///   response 1 -> ... <without end flag>
+///
+/// At egress:
+///   response 1 <end=false> -> <stream ended without end flag -> convert to error>
+///
+/// At client:
+///   response 1 -> error response
+/// ```
+///
+/// The detection must be done at egress level because premature stream termination
+/// can be due to network issues that only the egress component can detect.
+*/
+/// TODO: Detect end-of-stream using Server-Sent Events (SSE). This will be removed.
+#[derive(Serialize, Deserialize, Debug)]
+pub struct NetworkStreamWrapper<U> {
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub data: Option<U>,
+    pub complete_final: bool,
+}

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

@@ -17,7 +17,8 @@ use async_nats::client::Client;
 use tracing as log;
 
 use super::*;
-use crate::Result;
+use crate::{protocols::maybe_error::MaybeError, Result};
+use tokio_stream::{wrappers::ReceiverStream, StreamExt, StreamNotifyClose};
 
 #[derive(Debug, Clone, Serialize, Deserialize)]
 #[serde(rename_all = "snake_case")]
@@ -80,7 +81,7 @@ impl AddressedPushRouter {
 impl<T, U> AsyncEngine<SingleIn<AddressedRequest<T>>, ManyOut<U>, Error> for AddressedPushRouter
 where
     T: Data + Serialize,
-    U: Data + for<'de> Deserialize<'de>,
+    U: Data + for<'de> Deserialize<'de> + MaybeError,
 {
     async fn generate(&self, request: SingleIn<AddressedRequest<T>>) -> Result<ManyOut<U>, Error> {
         let request_id = request.context().id().to_string();
@@ -160,16 +161,49 @@ where
             .map_err(|_| PipelineError::DetatchedStreamReceiver)?
             .map_err(PipelineError::ConnectionFailed)?;
 
-        let stream = tokio_stream::wrappers::ReceiverStream::new(response_stream.rx);
-
-        let stream = stream.filter_map(|msg| async move {
-            match serde_json::from_slice::<U>(&msg) {
-                Ok(r) => Some(r),
-                Err(err) => {
-                    let json_str = String::from_utf8_lossy(&msg);
-                    log::warn!(%err, %json_str, "Failed deserializing JSON to response");
-                    None
+        // TODO: Detect end-of-stream using Server-Sent Events (SSE)
+        let mut is_complete_final = false;
+        let stream = tokio_stream::StreamNotifyClose::new(
+            tokio_stream::wrappers::ReceiverStream::new(response_stream.rx),
+        )
+        .filter_map(move |res| {
+            if let Some(res_bytes) = res {
+                if is_complete_final {
+                    return Some(U::from_err(
+                        Error::msg(
+                            "Response received after generation ended - this should never happen",
+                        )
+                        .into(),
+                    ));
+                }
+                match serde_json::from_slice::<NetworkStreamWrapper<U>>(&res_bytes) {
+                    Ok(item) => {
+                        is_complete_final = item.complete_final;
+                        if let Some(data) = item.data {
+                            Some(data)
+                        } else if is_complete_final {
+                            None
+                        } else {
+                            Some(U::from_err(
+                                Error::msg("Empty response received - this should never happen")
+                                    .into(),
+                            ))
+                        }
+                    }
+                    Err(err) => {
+                        // legacy log print
+                        let json_str = String::from_utf8_lossy(&res_bytes);
+                        log::warn!(%err, %json_str, "Failed deserializing JSON to response");
+
+                        Some(U::from_err(Error::new(err).into()))
+                    }
                 }
+            } else if is_complete_final {
+                None
+            } else {
+                Some(U::from_err(
+                    Error::msg("Stream ended before generation completed").into(),
+                ))
             }
         });