diff --git a/lib/bindings/python/Cargo.lock b/lib/bindings/python/Cargo.lock index e7c59c8a1584..d76e5ba28043 100644 --- a/lib/bindings/python/Cargo.lock +++ b/lib/bindings/python/Cargo.lock @@ -2319,6 +2319,7 @@ dependencies = [ "anyhow", "async-stream", "async-trait", + "bytes", "clap", "dashmap", "dynamo-backend-common", @@ -2337,6 +2338,7 @@ dependencies = [ "pythonize", "rmp", "serde", + "serde-transcode", "serde_json", "thiserror 2.0.18", "tokio", @@ -7108,6 +7110,15 @@ dependencies = [ "serde_derive", ] +[[package]] +name = "serde-transcode" +version = "1.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "590c0e25c2a5bb6e85bf5c1bce768ceb86b316e7a01bdf07d2cb4ec2271990e2" +dependencies = [ + "serde", +] + [[package]] name = "serde-untagged" version = "0.1.9" diff --git a/lib/bindings/python/Cargo.toml b/lib/bindings/python/Cargo.toml index 6a76e8ee536d..5a6973b9545e 100644 --- a/lib/bindings/python/Cargo.toml +++ b/lib/bindings/python/Cargo.toml @@ -50,6 +50,7 @@ dynamo-backend-common = { path = "../../backend-common" } anyhow = { version = "1" } async-stream = { version = "0.3" } async-trait = { version = "0.1" } +bytes = { version = "1" } dashmap = { version = "6.1" } futures = { version = "0.3" } rmp = { version = "0.8" } @@ -57,6 +58,7 @@ once_cell = { version = "1.20.3" } parking_lot = { version = "0.12.4" } serde = { version = "1" } serde_json = { version = "1.0.138" } +serde-transcode = "1.1" thiserror = { version = "2.0" } tokio = { version = "1.46.0", features = ["full"] } tokio-stream = { version = "0" } diff --git a/lib/bindings/python/rust/engine.rs b/lib/bindings/python/rust/engine.rs index d932484331ef..114824fa7cbd 100644 --- a/lib/bindings/python/rust/engine.rs +++ b/lib/bindings/python/rust/engine.rs @@ -3,6 +3,7 @@ use std::pin::Pin; use std::sync::Arc; +use std::task::{Context as TaskContext, Poll}; use anyhow::{Error, Result}; use pyo3::prelude::*; @@ -19,8 +20,8 @@ use dynamo_runtime::error::{BackendError, DynamoError, ErrorType}; use dynamo_runtime::logging::get_distributed_tracing_context; pub use dynamo_runtime::{ pipeline::{ - AsyncEngine, AsyncEngineContext, AsyncEngineContextProvider, Data, ManyOut, ResponseStream, - SingleIn, + AsyncEngine, AsyncEngineContext, AsyncEngineContextProvider, Data, DataStream, ManyOut, + ResponseStream, SingleIn, }, protocols::{annotated::Annotated, maybe_error::MaybeError}, }; @@ -30,6 +31,7 @@ use dynamo_runtime::pipeline::ManyIn; use super::context::{Context, callable_accepts_kwarg}; use super::errors::{extract_http_like_error, py_exception_to_backend_error}; +use crate::python_payload::{PythonPayload, PythonResponseItem}; /// Add bindings from this crate to the provided module pub fn add_to_module(m: &Bound<'_, PyModule>) -> PyResult<()> { @@ -67,6 +69,8 @@ type PyItemStream = Pin>> + Send>>; /// /// The GIL is acquired on a blocking task rather than inline: under contention /// it can block for an unbounded time, which would park the tokio reactor. +/// The returned stream polls `__anext__` only when its consumer requests an +/// item, preventing Python from mutating a reused object before it is consumed. async fn invoke_generator( generator: Arc, event_loop: Arc, @@ -93,15 +97,38 @@ where }?; let locals = TaskLocals::new(event_loop.bind(py).clone()); - pyo3_async_runtimes::tokio::into_stream_with_locals_v1( - locals, - gen_result.into_bound(py), - ) + demand_driven_python_stream(locals, gen_result.into_bound(py)) }) }) .await .map_err(|e| anyhow::anyhow!("failed to offload python call to blocking task: {e}"))??; + Ok(stream) +} + +fn demand_driven_python_stream( + locals: TaskLocals, + generator: Bound<'_, PyAny>, +) -> PyResult { + let anext = generator.getattr("__anext__")?.unbind(); + let stream = futures::stream::unfold((anext, locals), |(anext, locals)| async move { + let next = Python::with_gil(|py| { + pyo3_async_runtimes::into_future_with_locals(&locals, anext.bind(py).call0()?) + }); + let item = match next { + Ok(next) => next.await, + Err(error) => Err(error), + }; + if item.as_ref().is_err_and(|error| { + Python::with_gil(|py| { + error.is_instance_of::(py) + }) + }) { + None + } else { + Some((item, (anext, locals))) + } + }); Ok(Box::pin(stream)) } @@ -159,6 +186,36 @@ impl PythonAsyncEngine { } } +impl PythonAsyncEngine { + pub(crate) fn network_engine(&self) -> PythonNetworkEngine { + PythonNetworkEngine(self.0.clone()) + } +} + +/// Network-only Python engine. Requests and responses stay as Python-owned +/// objects so the ingress adapter can transcode them directly to and from the +/// configured request-plane codec. +#[derive(Clone)] +pub(crate) struct PythonNetworkEngine(PythonServerStreamingEngine); + +#[async_trait::async_trait] +impl AsyncEngine, ManyOut, Error> + for PythonNetworkEngine +{ + async fn generate( + &self, + request: SingleIn, + ) -> Result, Error> { + generate_python_stream( + &self.0, + request, + |_py, request| Ok(request.into_inner()), + unbuffered_python_response_stream, + ) + .await + } +} + #[async_trait::async_trait] impl AsyncEngine, ManyOut>, Error> for PythonAsyncEngine where @@ -215,40 +272,52 @@ where Resp: Data + for<'de> Deserialize<'de>, { async fn generate(&self, request: SingleIn) -> Result>, Error> { - // Create a context - let (request, context) = request.transfer(()); - let ctx = context.context(); - - let id = context.id().to_string(); - tracing::trace!("processing request: {}", id); - - // Capture current trace context - let current_trace_context = get_distributed_tracing_context(); - let metadata = context.metadata().clone(); - - let stream = invoke_generator( - self.generator.clone(), - self.event_loop.clone(), - move |py| Ok(pythonize(py, &request)?.unbind()), - self.has_context.then_some({ - let ctx = ctx.clone(); - move |py: Python<'_>| { - Py::new(py, Context::new(ctx, current_trace_context, None, metadata)) - .map(|c| c.into_any()) - } - }), + generate_python_stream( + self, + request, + |py, request| Ok(pythonize(py, &request)?.unbind()), + buffered_typed_response_stream::, ) - .await?; + .await + } +} - // Drain the Python response stream on a dedicated task, mapping any - // generator error to a typed annotated error frame. - let rx = spawn_response_forwarder::(stream, ctx, id); +async fn generate_python_stream( + engine: &PythonServerStreamingEngine, + request: SingleIn, + to_python_input: ToPythonInput, + forward_responses: ForwardResponses, +) -> Result, Error> +where + Req: Data, + Resp: Data, + ToPythonInput: FnOnce(Python, Req) -> PyResult> + Send + 'static, + ForwardResponses: + FnOnce(PyItemStream, Arc, String) -> DataStream + Send, +{ + let (request, context) = request.transfer(()); + let ctx = context.context(); + let id = context.id().to_string(); + tracing::trace!("processing request: {}", id); + + let current_trace_context = get_distributed_tracing_context(); + let metadata = context.metadata().clone(); + let stream = invoke_generator( + engine.generator.clone(), + engine.event_loop.clone(), + move |py| to_python_input(py, request), + engine.has_context.then_some({ + let ctx = ctx.clone(); + move |py: Python<'_>| { + Py::new(py, Context::new(ctx, current_trace_context, None, metadata)) + .map(|context| context.into_any()) + } + }), + ) + .await?; - Ok(ResponseStream::new( - Box::pin(ReceiverStream::new(rx)), - context.context(), - )) - } + let response_stream = forward_responses(stream, ctx, id); + Ok(ResponseStream::new(response_stream, context.context())) } async fn process_item( @@ -257,84 +326,7 @@ async fn process_item( where Resp: Data + for<'de> Deserialize<'de>, { - let item = item.map_err(|e| { - Python::with_gil(|py| { - e.display(py); - - // Check if the Python exception is a Dynamo error type. - // Wrap as Backend* since this is the backend engine context. - if let Some((backend_err, message)) = py_exception_to_backend_error(py, &e) { - return ResponseProcessingError::Dynamo( - DynamoError::builder() - .error_type(ErrorType::Backend(backend_err)) - .message(message) - .build(), - ); - } - - // openai.rs::extract_backend_error_if_present parses the DynamoError - // message as JSON {message, code}; emit that shape so the HTTP status - // survives instead of defaulting to 500. - if let Some((code, message)) = extract_http_like_error(py, &e) { - let backend_err = if (400..500).contains(&code) { - BackendError::InvalidArgument - } else { - BackendError::Unknown - }; - let json_msg = serde_json::json!({ - "message": message, - "code": code, - }) - .to_string(); - return ResponseProcessingError::Dynamo( - DynamoError::builder() - .error_type(ErrorType::Backend(backend_err)) - .message(json_msg) - .build(), - ); - } - - // GeneratorExit from Python's generator protocol (e.g., GC closing - // a generator) is treated as an engine shutdown. - if e.is_instance_of::(py) { - return ResponseProcessingError::Dynamo( - DynamoError::builder() - .error_type(ErrorType::Backend(BackendError::EngineShutdown)) - .message("engine shutting down") - .build(), - ); - } - - // Map well-known Python exceptions to specific Backend error types. - // Order matters: check subclasses before their parents - // (e.g., ConnectionRefusedError before ConnectionError). - let backend_err = if e.is_instance_of::(py) - || e.is_instance_of::(py) - { - BackendError::InvalidArgument - } else if e.is_instance_of::(py) { - BackendError::ConnectionTimeout - } else if e.is_instance_of::(py) { - BackendError::CannotConnect - } else if e.is_instance_of::(py) - || e.is_instance_of::(py) - || e.is_instance_of::(py) - { - BackendError::Disconnected - } else if e.is_instance_of::(py) { - BackendError::Cancelled - } else { - BackendError::Unknown - }; - - ResponseProcessingError::Dynamo( - DynamoError::builder() - .error_type(ErrorType::Backend(backend_err)) - .message(e.to_string()) - .build(), - ) - }) - })?; + let item = item.map_err(|e| ResponseProcessingError::Dynamo(map_python_exception(e)))?; let response = tokio::task::spawn_blocking(move || { Python::with_gil(|py| { let bound = item.into_bound(py); @@ -360,6 +352,67 @@ where Ok(response) } +pub(crate) fn map_python_exception(error: PyErr) -> DynamoError { + Python::with_gil(|py| { + error.display(py); + + if let Some((backend_err, message)) = py_exception_to_backend_error(py, &error) { + return DynamoError::builder() + .error_type(ErrorType::Backend(backend_err)) + .message(message) + .build(); + } + + if let Some((code, message)) = extract_http_like_error(py, &error) { + let backend_err = if (400..500).contains(&code) { + BackendError::InvalidArgument + } else { + BackendError::Unknown + }; + let json_msg = serde_json::json!({ + "message": message, + "code": code, + }) + .to_string(); + return DynamoError::builder() + .error_type(ErrorType::Backend(backend_err)) + .message(json_msg) + .build(); + } + + if error.is_instance_of::(py) { + return DynamoError::builder() + .error_type(ErrorType::Backend(BackendError::EngineShutdown)) + .message("engine shutting down") + .build(); + } + + let backend_err = if error.is_instance_of::(py) + || error.is_instance_of::(py) + { + BackendError::InvalidArgument + } else if error.is_instance_of::(py) { + BackendError::ConnectionTimeout + } else if error.is_instance_of::(py) { + BackendError::CannotConnect + } else if error.is_instance_of::(py) + || error.is_instance_of::(py) + || error.is_instance_of::(py) + { + BackendError::Disconnected + } else if error.is_instance_of::(py) { + BackendError::Cancelled + } else { + BackendError::Unknown + }; + + DynamoError::builder() + .error_type(ErrorType::Backend(backend_err)) + .message(error.to_string()) + .build() + }) +} + /// Channel depth between the response-forwarding task and the consumer of /// the engine's output stream. const RESPONSE_CHANNEL_DEPTH: usize = 128; @@ -373,6 +426,9 @@ const RESPONSE_CHANNEL_DEPTH: usize = 128; /// and emitted as annotated error frames, so a failing generator returns as /// an error to the client rather than a silently truncated stream. /// On a deserialize mismatch the request context is told to stop generating. +/// The generator is not polled again until [`process_item`] has converted the +/// current `PyObject` into an owned Rust value, so this channel never buffers +/// mutable Python objects that a generator could reuse for a later yield. fn spawn_response_forwarder( stream: PyItemStream, ctx: Arc, @@ -458,6 +514,103 @@ where rx } +fn buffered_typed_response_stream( + stream: PyItemStream, + ctx: Arc, + request_id: String, +) -> DataStream> +where + Resp: Data + for<'de> Deserialize<'de>, +{ + Box::pin(ReceiverStream::new(spawn_response_forwarder::( + stream, ctx, request_id, + ))) +} + +fn unbuffered_python_response_stream( + stream: PyItemStream, + ctx: Arc, + request_id: String, +) -> DataStream { + // Do not poll the generator again until ingress has encoded the current + // Python object. A generator may reuse and mutate the same dict/list for + // later yields; buffering raw PyObject handles would make earlier frames + // observe those later mutations. + Box::pin(DirectPythonResponseStream { + stream: Some(stream), + ctx, + request_id, + exhausted: false, + }) +} + +/// Demand-driven network response stream with cooperative Python cancellation. +/// +/// Ingress stops the request context when the client response connection +/// closes, then drops this stream. Give the Python generator one final poll so +/// it can observe `context.is_stopped()` and run its cancellation path. Normal +/// response delivery remains unbuffered: the next generator item is not polled +/// until ingress has encoded the current one. +struct DirectPythonResponseStream { + stream: Option, + ctx: Arc, + request_id: String, + exhausted: bool, +} + +impl Stream for DirectPythonResponseStream { + type Item = PythonResponseItem; + + fn poll_next(mut self: Pin<&mut Self>, cx: &mut TaskContext<'_>) -> Poll> { + let poll = self + .stream + .as_mut() + .expect("direct Python response stream missing before exhaustion") + .as_mut() + .poll_next(cx); + match poll { + Poll::Ready(Some(item)) => { + if item.is_err() { + self.exhausted = true; + self.stream.take(); + } + Poll::Ready(Some(PythonResponseItem::new(item))) + } + Poll::Ready(None) => { + self.exhausted = true; + self.stream.take(); + Poll::Ready(None) + } + Poll::Pending => Poll::Pending, + } + } +} + +impl Drop for DirectPythonResponseStream { + fn drop(&mut self) { + if self.exhausted || !self.ctx.is_stopped() { + return; + } + + let Some(mut stream) = self.stream.take() else { + return; + }; + let request_id = self.request_id.clone(); + let Ok(runtime) = tokio::runtime::Handle::try_current() else { + return; + }; + runtime.spawn(async move { + // Cooperative Python generators inspect their Context when polled. + // Discard the result: the client has already dropped its stream. + let _ = stream.next().await; + tracing::trace!( + request_id, + "polled direct Python response generator after cancellation" + ); + }); + } +} + /// Channel depth between the inbound forwarder and the Python iterator. /// Mirrors the depth used by the wire-side bidirectional ingress forwarder /// in `lib/runtime/src/pipeline/network/ingress/push_handler.rs`. @@ -465,23 +618,18 @@ const BIDIRECTIONAL_INPUT_CHANNEL_DEPTH: usize = 8; /// Rust-side adapter that bridges a Python `async def generate(request_stream, context)` /// callable into an [`AsyncEngine`] of the ManyIn / ManyOut -/// shape (Req=serde_json::Value, Resp=Annotated). +/// shape (`PythonPayload` request frames and raw Python response items). /// /// The adapter: /// -/// 1. Transforms the inbound `RequestStream` into a +/// 1. Transforms the inbound `RequestStream` into a /// `PyAsyncRequestStream` and `context`. Similar to unary engine, /// cancellation observation is the Python engine's responsibility via the /// `context` argument. Input stream can end early if no more inputs are expected. /// 2. Invokes the Python generator with `(request_stream, context)`, then /// wraps the returned async generator into a Rust `Stream>`. -/// 3. Depythonizes each item and wraps it as `Annotated`, -/// and forwards it on the response stream. -/// -/// Wire types are fixed to `serde_json::Value` on the request side and -/// `Annotated` on the response side. The Python user -/// works with dicts on both sides and any schema enforcement is handled in -/// Python. +/// 3. Forwards each raw Python response to the ingress payload adapter, which +/// performs annotation handling and wire serialization in one blocking step. pub struct PythonBidirectionalEngine { generator: Arc, event_loop: Arc, @@ -504,13 +652,13 @@ impl PythonBidirectionalEngine { } #[async_trait::async_trait] -impl AsyncEngine, ManyOut>, Error> +impl AsyncEngine, ManyOut, Error> for PythonBidirectionalEngine { async fn generate( &self, - input: ManyIn, - ) -> Result>, Error> { + input: ManyIn, + ) -> Result, Error> { let (request_stream, ctx_unit) = input.into_parts(); let ctx = ctx_unit.context(); let request_id = ctx_unit.id().to_string(); @@ -523,8 +671,7 @@ impl AsyncEngine, ManyOut // dispatching task; needed when constructing the Python `Context`. let current_trace_context = get_distributed_tracing_context(); - // Forwarder: pull `serde_json::Value` frames off the inbound stream, - // pythonize each, and hand the resulting `PyObject` to the Python + // Forwarder: move Python-owned frames directly into the Python // iterator. The `frame_tx.closed()` arm cancels the forwarder as soon // as the Python iterator drops the receiver, so it shuts down promptly // instead of blocking on the next inbound frame. @@ -539,21 +686,7 @@ impl AsyncEngine, ManyOut let Some(value) = value else { break; }; - let pyobj = match Python::with_gil(|py| { - pythonize(py, &value).map(|bound| bound.unbind()) - }) { - Ok(pyobj) => pyobj, - Err(e) => { - tracing::error!( - request_id = %forwarder_request_id, - error = %e, - "failed to pythonize bidirectional request frame; \ - closing input forwarder" - ); - break; - } - }; - if frame_tx.send(pyobj).await.is_err() { + if frame_tx.send(value.into_inner()).await.is_err() { tracing::debug!( request_id = %forwarder_request_id, "python engine dropped request stream; input forwarder exiting" @@ -581,13 +714,7 @@ impl AsyncEngine, ManyOut ) .await?; - // Drain the Python response stream on a dedicated task. Sharing - // `spawn_response_forwarder` gives the bidirectional engine the same - // typed error mapping as the unary engine: a generator that raises now - // yields a structured annotated error frame instead of a silently - // truncated stream. - let rx = spawn_response_forwarder::(stream, ctx.clone(), request_id); - - Ok(ResponseStream::new(Box::pin(ReceiverStream::new(rx)), ctx)) + let response_stream = unbuffered_python_response_stream(stream, ctx.clone(), request_id); + Ok(ResponseStream::new(response_stream, ctx)) } } diff --git a/lib/bindings/python/rust/lib.rs b/lib/bindings/python/rust/lib.rs index 162faee7d816..7b94590671bd 100644 --- a/lib/bindings/python/rust/lib.rs +++ b/lib/bindings/python/rust/lib.rs @@ -81,12 +81,19 @@ mod llm; mod parsers; mod planner; mod prometheus_metrics; +mod python_payload; -type JsonServerStreamingIngress = - Ingress, ManyOut>>; +type PythonServerStreamingIngress = Ingress< + SingleIn, + ManyOut, + python_payload::PythonIngressPayloadAdapter, +>; -type JsonBidirectionalIngress = - Ingress, ManyOut>>; +type PythonBidirectionalIngress = Ingress< + rs::pipeline::ManyIn, + ManyOut, + python_payload::PythonIngressPayloadAdapter, +>; static INIT: OnceCell<()> = OnceCell::new(); @@ -1124,7 +1131,12 @@ impl Endpoint { generator, self.event_loop.clone(), )?); - let ingress = JsonServerStreamingIngress::for_engine(engine.clone()).map_err(to_pyerr)?; + let network_engine = Arc::new(engine.network_engine()); + let ingress = PythonServerStreamingIngress::for_engine_with_adapter( + network_engine, + python_payload::PythonIngressPayloadAdapter, + ) + .map_err(to_pyerr)?; // Convert Python dict to serde_json::Value if provided and validate it's an object let health_payload_json = health_check_payload @@ -1176,10 +1188,9 @@ impl Endpoint { /// `async def generate(request_stream, context)` coroutine that /// returns an async generator. `request_stream` is a /// [`PyAsyncRequestStream`] yielding inbound frames as plain Python - /// objects (dicts/lists/etc., the depythonization of - /// `serde_json::Value`). The generator yields response frames as - /// plain Python objects that are then pythonized back to JSON values - /// on the wire. + /// objects (dicts/lists/etc.) decoded directly from the configured + /// request-plane payload codec. The generator yields plain Python + /// response objects that are serialized directly to that codec. /// /// Request-stream end (when `__anext__` raises `StopAsyncIteration`) /// is *not* a cancellation signal: the caller has merely stopped @@ -1197,8 +1208,9 @@ impl Endpoint { generator, self.event_loop.clone(), )?); - let ingress: Arc = - Ingress::for_engine(engine).map_err(to_pyerr)?; + let ingress: Arc = + Ingress::for_engine_with_adapter(engine, python_payload::PythonIngressPayloadAdapter) + .map_err(to_pyerr)?; let builder = self .inner @@ -1664,11 +1676,10 @@ impl AsyncResponseStream { } /// Python-visible inbound iterator for bidirectional engines. Wraps an -/// mpsc receiver of pre-pythonized request frames; `__anext__` is a thin +/// mpsc receiver of Python-owned request frames; `__anext__` is a thin /// `.recv()` that returns the next `PyObject` directly, with no per-frame -/// GIL acquisition on the consumer side. The producer (the forwarder -/// spawned by `PythonBidirectionalEngine::generate`) acquires the GIL -/// once per frame and pushes the converted `PyObject` onto the channel. +/// GIL acquisition or value conversion on the consumer side. The producer +/// moves the object decoded by the ingress adapter onto the channel. /// /// Termination follows the same shape as `AsyncResponseStream`: when the /// channel returns `None`, `__anext__` raises `PyStopAsyncIteration` and @@ -1697,7 +1708,7 @@ impl PyAsyncRequestStream { } /// Required by the `AsyncIterator` protocol. Returns an awaitable - /// resolving to the next pre-pythonized frame, or raises + /// resolving to the next Python-owned frame, or raises /// `StopAsyncIteration` when the inbound channel is closed. #[pyo3(name = "__anext__")] fn next<'p>(&self, py: Python<'p>) -> PyResult> { diff --git a/lib/bindings/python/rust/python_payload.rs b/lib/bindings/python/rust/python_payload.rs new file mode 100644 index 000000000000..0d4920cbd3db --- /dev/null +++ b/lib/bindings/python/rust/python_payload.rs @@ -0,0 +1,276 @@ +// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use bytes::Bytes; +use dynamo_runtime::pipeline::PipelineError; +use dynamo_runtime::pipeline::network::{ + EncodedResponseFrame, IngressRequestDecoder, IngressResponseEncoder, NetworkStreamWrapper, + RequestPlanePayloadCodec, +}; +use dynamo_runtime::protocols::annotated::Annotated; +use dynamo_runtime::protocols::maybe_error::MaybeError; +use pyo3::prelude::*; +use pyo3::types::PyDict; +use pythonize::{Depythonizer, Pythonizer, depythonize}; +use serde::de::Error as _; +use serde::ser::Error as _; +use serde::{Deserialize, Deserializer, Serialize, Serializer}; + +use crate::engine::map_python_exception; + +/// Python-owned request value used only by the network ingress fast path. +/// Serde events are transcoded directly to or from Python objects without an +/// intermediate Rust value tree. +#[derive(Clone)] +pub(crate) struct PythonPayload(Py); + +impl PythonPayload { + pub(crate) fn into_inner(self) -> Py { + self.0 + } +} + +impl std::fmt::Debug for PythonPayload { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_str("PythonPayload()") + } +} + +impl<'de> Deserialize<'de> for PythonPayload { + fn deserialize(deserializer: D) -> Result + where + D: Deserializer<'de>, + { + Python::with_gil(|py| { + serde_transcode::transcode(deserializer, Pythonizer::new(py)) + .map(|value| Self(value.unbind())) + .map_err(D::Error::custom) + }) + } +} + +impl Serialize for PythonPayload { + fn serialize(&self, serializer: S) -> Result + where + S: Serializer, + { + Python::with_gil(|py| { + let mut depythonizer = Depythonizer::from_object(self.0.bind(py)); + serde_transcode::transcode(&mut depythonizer, serializer).map_err(S::Error::custom) + }) + } +} + +/// One raw item yielded by a Python async generator. +pub(crate) struct PythonResponseItem(PyResult>); + +impl PythonResponseItem { + pub(crate) fn new(item: PyResult>) -> Self { + Self(item) + } + + fn into_result(self) -> PyResult> { + self.0 + } +} + +impl std::fmt::Debug for PythonResponseItem { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match &self.0 { + Ok(_) => f.write_str("PythonResponseItem::Data()"), + Err(_) => f.write_str("PythonResponseItem::Error()"), + } + } +} + +#[derive(Debug, Default)] +pub(crate) struct PythonIngressPayloadAdapter; + +impl IngressRequestDecoder for PythonIngressPayloadAdapter { + async fn decode_request( + &self, + payload_codec: RequestPlanePayloadCodec, + bytes: Bytes, + ) -> Result { + tokio::task::spawn_blocking(move || payload_codec.decode::(&bytes)) + .await + .map_err(|error| { + PipelineError::DeserializationError(format!( + "failed to offload {} Python request decode: {error}", + payload_codec.name() + )) + })? + .map_err(|error| { + PipelineError::DeserializationError(format!( + "Failed deserializing {} Python request payload: {error}", + payload_codec.name() + )) + }) + } +} + +impl IngressResponseEncoder for PythonIngressPayloadAdapter { + async fn encode_response( + &self, + payload_codec: RequestPlanePayloadCodec, + response: Option, + complete_final: bool, + ) -> Result { + if complete_final { + let wrapper = NetworkStreamWrapper::> { + data: None, + complete_final: true, + }; + let bytes = payload_codec.encode(&wrapper).map_err(|error| { + PipelineError::SerializationError(format!( + "Failed serializing {} request-plane final response: {error}", + payload_codec.name() + )) + })?; + return Ok(EncodedResponseFrame { + bytes: bytes.into(), + is_error: false, + stop_stream: false, + }); + } + + let response = response.ok_or_else(|| { + PipelineError::SerializationError( + "request-plane response item missing before final frame".to_string(), + ) + })?; + tokio::task::spawn_blocking(move || encode_python_response(payload_codec, response)) + .await + .map_err(|error| { + PipelineError::SerializationError(format!( + "failed to offload {} Python response encode: {error}", + payload_codec.name() + )) + })? + } +} + +fn encode_python_response( + payload_codec: RequestPlanePayloadCodec, + response: PythonResponseItem, +) -> Result { + let (annotated, stop_stream) = match response.into_result() { + Ok(item) => match Python::with_gil(|py| parse_python_response(item, py)) { + Ok(annotated) => (annotated, false), + Err(error) => ( + Annotated::from_error(format!( + "critical error: invalid response object from Python async generator; \ + application-logic-mismatch: {error}" + )), + true, + ), + }, + Err(error) => (Annotated::from_err(map_python_exception(error)), true), + }; + let is_error = annotated.is_error(); + let wrapper = NetworkStreamWrapper { + data: Some(annotated), + complete_final: false, + }; + + match payload_codec.encode(&wrapper) { + Ok(bytes) => Ok(EncodedResponseFrame { + bytes: bytes.into(), + is_error, + stop_stream, + }), + Err(error) => { + let fallback = NetworkStreamWrapper { + data: Some(Annotated::<()>::from_error(format!( + "critical error: failed serializing Python response as {}: {error}", + payload_codec.name() + ))), + complete_final: false, + }; + let bytes = payload_codec.encode(&fallback).map_err(|fallback_error| { + PipelineError::SerializationError(format!( + "failed to serialize Python response and fallback error as {}: {fallback_error}", + payload_codec.name() + )) + })?; + Ok(EncodedResponseFrame { + bytes: bytes.into(), + is_error: true, + stop_stream: true, + }) + } + } +} + +fn parse_python_response( + item: Py, + py: Python<'_>, +) -> Result, String> { + let bound = item.bind(py); + let Some(dict) = bound.downcast::().ok() else { + return Ok(Annotated::from_data(PythonPayload(item))); + }; + let is_envelope = dict + .get_item("_dynamo_annotated") + .map_err(|error| error.to_string())? + .and_then(|value| value.is_truthy().ok()) + .unwrap_or(false); + if !is_envelope { + return Ok(Annotated::from_data(PythonPayload(item))); + } + + // Keep the payload itself as the original Python object. Fully + // depythonizing `Annotated` would rebuild the nested data + // subtree and defeat the direct request-plane path's ownership reuse. + let data = optional_item(dict, "data")?.map(|value| PythonPayload(value.unbind())); + let id = extract_optional(dict, "id")?; + let event = extract_optional(dict, "event")?; + let comment = extract_optional(dict, "comment")?; + let error = optional_item(dict, "error")? + .map(|value| depythonize(&value).map_err(|error| error.to_string())) + .transpose()?; + + Ok(Annotated { + data, + id, + event, + comment, + error, + }) +} + +fn optional_item<'py>( + dict: &Bound<'py, PyDict>, + name: &str, +) -> Result>, String> { + dict.get_item(name) + .map_err(|error| error.to_string()) + .map(|value| value.filter(|value| !value.is_none())) +} + +fn extract_optional<'py, T>(dict: &Bound<'py, PyDict>, name: &str) -> Result, String> +where + T: FromPyObject<'py>, +{ + optional_item(dict, name)? + .map(|value| value.extract().map_err(|error| error.to_string())) + .transpose() +} + +#[cfg(test)] +mod tests { + // Keep Rust unit tests here free of Python C API calls. This crate uses + // PyO3's `extension-module` feature, so standalone `cargo test` binaries + // intentionally do not link libpython. Python behavior is covered by + // tests/test_request_plane_python_payload.py against the built extension. + #[test] + fn network_ingress_types_do_not_contain_serde_json_value() { + let unary = std::any::type_name::(); + let bidirectional = std::any::type_name::(); + assert!(!unary.contains("serde_json::value::Value"), "{unary}"); + assert!( + !bidirectional.contains("serde_json::value::Value"), + "{bidirectional}" + ); + } +} diff --git a/lib/bindings/python/tests/test_kserve_grpc.py b/lib/bindings/python/tests/test_kserve_grpc.py index 9bbf76868945..ba414112b3ed 100644 --- a/lib/bindings/python/tests/test_kserve_grpc.py +++ b/lib/bindings/python/tests/test_kserve_grpc.py @@ -3,6 +3,7 @@ import asyncio import contextlib +import queue from contextlib import asynccontextmanager from typing import Any, AsyncIterator, Optional, Tuple @@ -49,11 +50,18 @@ def __init__(self, model_name: str): def generate(self, request, context=None): async def _generator(): - yield { + response = { "model": self._model_name, "tensors": request.get("tensors", []), "parameters": request.get("parameters", {}), } + if request.get("parameters", {}).get("reused_mutable"): + for sequence in range(64): + response["model"] = f"{self._model_name}-{sequence}" + yield response + return + + yield response return _generator() @@ -156,3 +164,44 @@ async def test_model_config_missing_tensor_config_errors(tensor_service): client.close() assert "not found" in str(excinfo.value).lower() + + +@pytest.mark.asyncio +@pytest.mark.forked +@pytest.mark.timeout(30) +async def test_python_async_engine_snapshots_reused_mutable_responses(tensor_service): + """Snapshot each typed response before polling a reused Python object again.""" + import numpy as np + import tritonclient.grpc as grpcclient + + model_name = "tensor-reused-mutable" + async with tensor_service(model_name) as (host, port): + client = grpcclient.InferenceServerClient(url=f"{host}:{port}") + completed: queue.Queue = queue.Queue() + + def callback(result, error): + completed.put(error if error is not None else result) + + input_data = np.array([1], dtype=np.int32) + infer_input = grpcclient.InferInput("INPUT0", input_data.shape, "INT32") + infer_input.set_data_from_numpy(input_data) + + client.start_stream(callback=callback) + try: + client.async_stream_infer( + model_name=model_name, + inputs=[infer_input], + parameters={"reused_mutable": True}, + ) + responses = [ + await asyncio.to_thread(completed.get, True, 5) for _ in range(64) + ] + finally: + client.stop_stream() + client.close() + + errors = [response for response in responses if isinstance(response, Exception)] + assert not errors + assert [response.get_response().model_name for response in responses] == [ + f"{model_name}-{sequence}" for sequence in range(64) + ] diff --git a/lib/bindings/python/tests/test_request_plane_python_payload.py b/lib/bindings/python/tests/test_request_plane_python_payload.py new file mode 100644 index 000000000000..b8c8f31095fb --- /dev/null +++ b/lib/bindings/python/tests/test_request_plane_python_payload.py @@ -0,0 +1,130 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import asyncio +import contextlib + +import pytest + +pytestmark = [ + pytest.mark.gpu_0, + pytest.mark.pre_merge, + pytest.mark.integration, +] + + +async def _generate(request, context): + assert dict(context.metadata.items()) == request.get("expected_metadata", {}) + if request["kind"] == "error": + raise ValueError("direct adapter test error") + if request["kind"] == "malformed": + yield object() + return + if request["kind"] == "explicit-none": + yield { + "_dynamo_annotated": True, + "data": {"nested": {"value": 42}}, + "id": None, + "event": None, + "comment": None, + "error": None, + } + return + if request["kind"] == "reused-mutable": + shared = {"sequence": 0} + for sequence in range(64): + shared["sequence"] = sequence + yield shared + return + + yield request["payload"] + yield { + "_dynamo_annotated": True, + "data": {"annotated": request["payload"]}, + "id": "chunk-2", + "event": "delta", + "comment": ["direct", "python"], + } + + +@pytest.fixture +async def request_plane_client(runtime): + endpoint = runtime.endpoint("direct-python-msgpack.backend.generate") + health_payload = { + "kind": "normal", + "payload": {"health": True}, + "expected_metadata": {}, + } + server_task = asyncio.ensure_future( + endpoint.serve_endpoint(_generate, health_check_payload=health_payload) + ) + client = await endpoint.client() + await client.wait_for_instances() + + yield client + + server_task.cancel() + with contextlib.suppress(asyncio.CancelledError): + await server_task + + +@pytest.mark.asyncio +@pytest.mark.timeout(30) +@pytest.mark.parametrize("request_plane", ["tcp"], indirect=True) +async def test_python_request_plane_plain_annotated_error_and_malformed_frames( + request_plane_client, +): + payload = { + "text": "hello δΈ­", + "tokens": [1, 2, 65535], + "nested": {"stream": True, "temperature": 0.25}, + "nullable": None, + } + request = { + "kind": "normal", + "payload": payload, + "expected_metadata": {"trace": "adapter-test"}, + } + + from dynamo.runtime import Context + + context = Context() + context.metadata["trace"] = "adapter-test" + stream = await request_plane_client.generate(request, context=context) + responses = [response async for response in stream] + + assert len(responses) == 2 + assert responses[0].data() == payload + assert responses[1].data() == {"annotated": payload} + assert responses[1].id() == "chunk-2" + assert responses[1].event() == "delta" + assert responses[1].comments() == ["direct", "python"] + + stream = await request_plane_client.generate( + {"kind": "explicit-none", "expected_metadata": {}} + ) + explicit_none_responses = [response async for response in stream] + assert len(explicit_none_responses) == 1 + assert explicit_none_responses[0].data() == {"nested": {"value": 42}} + assert explicit_none_responses[0].id() is None + assert explicit_none_responses[0].event() is None + assert explicit_none_responses[0].comments() is None + + stream = await request_plane_client.generate( + {"kind": "reused-mutable", "expected_metadata": {}} + ) + reused_mutable_responses = [response.data() async for response in stream] + assert reused_mutable_responses == [ + {"sequence": sequence} for sequence in range(64) + ] + + for kind, message in [ + ("error", "direct adapter test error"), + ("malformed", "failed serializing Python response"), + ]: + stream = await request_plane_client.generate( + {"kind": kind, "expected_metadata": {}} + ) + with pytest.raises(ValueError, match=message): + async for _ in stream: + pass diff --git a/lib/runtime/src/metrics/prometheus_names.rs b/lib/runtime/src/metrics/prometheus_names.rs index 63a988ee67cf..bb0d55243ff4 100644 --- a/lib/runtime/src/metrics/prometheus_names.rs +++ b/lib/runtime/src/metrics/prometheus_names.rs @@ -446,6 +446,9 @@ pub mod work_handler { /// Generation error pub const GENERATE: &str = "generate"; + /// Response serialization error + pub const SERIALIZATION: &str = "serialization"; + /// Response publishing error pub const PUBLISH_RESPONSE: &str = "publish_response"; diff --git a/lib/runtime/src/pipeline/network.rs b/lib/runtime/src/pipeline/network.rs index 914121aa0562..be106489901d 100644 --- a/lib/runtime/src/pipeline/network.rs +++ b/lib/runtime/src/pipeline/network.rs @@ -56,53 +56,57 @@ pub(crate) fn get_tcp_max_message_size() -> usize { #[derive(Debug, Clone, Copy, Default, Serialize, Deserialize, PartialEq, Eq)] #[serde(rename_all = "snake_case")] -pub(crate) enum RequestPlanePayloadCodec { - // TODO(jthomson04): Migrate the default to Msgpack after the 1.3 release. +pub enum RequestPlanePayloadCodec { + /// The serde default deliberately remains JSON for wire compatibility with + /// control messages produced before the payload codec field existed. #[default] Json, Msgpack, } impl RequestPlanePayloadCodec { - pub(crate) fn configured() -> Self { + pub fn configured() -> Self { *REQUEST_PLANE_PAYLOAD_CODEC.get_or_init(Self::from_env) } fn from_env() -> Self { - match std::env::var( - crate::config::environment_names::request_plane::DYN_REQUEST_PLANE_CODEC, - ) - .as_deref() - { - Err(_) | Ok("") | Ok("json") => Self::Json, - Ok("msgpack") => Self::Msgpack, - Ok(other) => { + let value = + std::env::var(crate::config::environment_names::request_plane::DYN_REQUEST_PLANE_CODEC) + .ok(); + Self::from_config_value(value.as_deref()) + } + + fn from_config_value(value: Option<&str>) -> Self { + match value { + None | Some("") | Some("msgpack") => Self::Msgpack, + Some("json") => Self::Json, + Some(other) => { tracing::warn!( env_var = crate::config::environment_names::request_plane::DYN_REQUEST_PLANE_CODEC, value = other, - "invalid request plane payload codec, defaulting to json" + "invalid request plane payload codec, defaulting to msgpack" ); - Self::Json + Self::Msgpack } } } - pub(crate) fn name(&self) -> &'static str { + pub fn name(&self) -> &'static str { match self { Self::Json => "json", Self::Msgpack => "msgpack", } } - pub(crate) fn encode(&self, value: &T) -> Result> { + pub fn encode(&self, value: &T) -> Result> { match self { Self::Json => Ok(serde_json::to_vec(value)?), Self::Msgpack => Ok(rmp_serde::to_vec_named(value)?), } } - pub(crate) fn decode(&self, bytes: &[u8]) -> Result { + pub fn decode(&self, bytes: &[u8]) -> Result { match self { Self::Json => Ok(serde_json::from_slice(bytes)?), Self::Msgpack => Ok(rmp_serde::from_slice(bytes)?), @@ -515,6 +519,34 @@ mod tests { assert_eq!(message.payload_codec, RequestPlanePayloadCodec::Msgpack); } + #[test] + fn request_plane_payload_codec_configuration_defaults_to_msgpack() { + assert_eq!( + RequestPlanePayloadCodec::from_config_value(None), + RequestPlanePayloadCodec::Msgpack + ); + assert_eq!( + RequestPlanePayloadCodec::from_config_value(Some("")), + RequestPlanePayloadCodec::Msgpack + ); + assert_eq!( + RequestPlanePayloadCodec::from_config_value(Some("invalid")), + RequestPlanePayloadCodec::Msgpack + ); + } + + #[test] + fn request_plane_payload_codec_configuration_honors_explicit_overrides() { + assert_eq!( + RequestPlanePayloadCodec::from_config_value(Some("json")), + RequestPlanePayloadCodec::Json + ); + assert_eq!( + RequestPlanePayloadCodec::from_config_value(Some("msgpack")), + RequestPlanePayloadCodec::Msgpack + ); + } + #[test] fn request_plane_payload_codec_round_trips_response_wrapper_json_and_msgpack() { let wrapper = NetworkStreamWrapper { @@ -550,19 +582,158 @@ where } } -pub struct Ingress { +/// Result of encoding one response item for the request plane. +pub struct EncodedResponseFrame { + pub bytes: Bytes, + pub is_error: bool, + /// Stop consuming the engine stream after publishing this frame. The + /// normal complete-final frame is still sent. + pub stop_stream: bool, +} + +/// Converts request-plane bytes into the item consumed by an ingress engine. +pub trait IngressRequestDecoder: Send + Sync + 'static +where + T: Data, +{ + fn decode_request( + &self, + payload_codec: RequestPlanePayloadCodec, + bytes: Bytes, + ) -> impl std::future::Future> + Send; +} + +/// Converts an ingress engine response into its complete on-wire frame. +pub trait IngressResponseEncoder: Send + Sync + 'static +where + U: Data, +{ + fn encode_response( + &self, + payload_codec: RequestPlanePayloadCodec, + response: Option, + complete_final: bool, + ) -> impl std::future::Future> + Send; +} + +/// Complete request/response payload adapter for an ingress engine. +pub trait IngressPayloadAdapter: + IngressRequestDecoder + IngressResponseEncoder +where + T: Data, + U: Data, +{ +} + +impl IngressPayloadAdapter for Adapter +where + T: Data, + U: Data, + Adapter: IngressRequestDecoder + IngressResponseEncoder, +{ +} + +/// Default adapter for ordinary Rust request and response types. +#[derive(Debug, Default)] +pub struct SerdeIngressPayloadAdapter; + +impl IngressRequestDecoder for SerdeIngressPayloadAdapter +where + T: Data + DeserializeOwned, +{ + #[inline] + fn decode_request( + &self, + payload_codec: RequestPlanePayloadCodec, + bytes: Bytes, + ) -> impl std::future::Future> + Send { + let decoded = payload_codec.decode(&bytes).map_err(|err| { + PipelineError::DeserializationError(format!( + "Failed deserializing {} request payload: {}", + payload_codec.name(), + err + )) + }); + std::future::ready(decoded) + } +} + +impl IngressResponseEncoder for SerdeIngressPayloadAdapter +where + U: Data + Serialize + MaybeError, +{ + #[inline] + fn encode_response( + &self, + payload_codec: RequestPlanePayloadCodec, + response: Option, + complete_final: bool, + ) -> impl std::future::Future> + Send + { + let is_error = response + .as_ref() + .is_some_and(|response| response.err().is_some()); + let wrapper = NetworkStreamWrapper { + data: response, + complete_final, + }; + let encoded = payload_codec.encode(&wrapper).map_err(|err| { + PipelineError::SerializationError(format!( + "Failed serializing {} request-plane response: {}", + payload_codec.name(), + err + )) + }); + std::future::ready(encoded.map(|bytes| EncodedResponseFrame { + bytes: bytes.into(), + is_error, + stop_stream: false, + })) + } +} + +pub struct Ingress { segment: OnceLock>>, metrics: OnceLock>, /// Endpoint-specific notifier for health check timer resets endpoint_health_check_notifier: OnceLock>, + payload_adapter: Arc, } impl Ingress { pub fn new() -> Arc { + Ingress::new_with_adapter(SerdeIngressPayloadAdapter) + } + + pub fn link(segment: Arc>) -> Result> { + let ingress = Ingress::new(); + ingress.attach(segment)?; + Ok(ingress) + } + + pub fn for_pipeline(segment: Arc>) -> Result> { + let ingress = Ingress::new(); + ingress.attach(segment)?; + Ok(ingress) + } + + pub fn for_engine(engine: ServiceEngine) -> Result> { + Self::for_engine_with_adapter(engine, SerdeIngressPayloadAdapter) + } +} + +impl Ingress +where + Req: PipelineIO + Sync, + Resp: PipelineIO, + Adapter: Send + Sync + 'static, +{ + pub fn new_with_adapter(payload_adapter: Adapter) -> Arc { Arc::new(Self { segment: OnceLock::new(), metrics: OnceLock::new(), endpoint_health_check_notifier: OnceLock::new(), + payload_adapter: Arc::new(payload_adapter), }) } @@ -597,26 +768,17 @@ impl Ingress { .map_err(|_| anyhow::anyhow!("Metrics already set")) } - pub fn link(segment: Arc>) -> Result> { - let ingress = Ingress::new(); - ingress.attach(segment)?; - Ok(ingress) - } - - pub fn for_pipeline(segment: Arc>) -> Result> { - let ingress = Ingress::new(); - ingress.attach(segment)?; - Ok(ingress) - } - - pub fn for_engine(engine: ServiceEngine) -> Result> { + pub fn for_engine_with_adapter( + engine: ServiceEngine, + payload_adapter: Adapter, + ) -> Result> { let frontend = SegmentSource::::new(); let backend = ServiceBackend::from_engine(engine); // create the pipeline let pipeline = frontend.link(backend)?.link(frontend)?; - let ingress = Ingress::new(); + let ingress = Ingress::new_with_adapter(payload_adapter); ingress.attach(pipeline)?; Ok(ingress) diff --git a/lib/runtime/src/pipeline/network/ingress/push_handler.rs b/lib/runtime/src/pipeline/network/ingress/push_handler.rs index d796d93eb894..9e6c351d3252 100644 --- a/lib/runtime/src/pipeline/network/ingress/push_handler.rs +++ b/lib/runtime/src/pipeline/network/ingress/push_handler.rs @@ -9,10 +9,9 @@ use crate::metrics::work_handler_perf::{ WORK_HANDLER_NETWORK_TRANSIT_SECONDS, WORK_HANDLER_TIME_TO_FIRST_RESPONSE_SECONDS, }; use crate::pipeline::{ManyIn, RequestStream}; -use crate::protocols::maybe_error::MaybeError; use futures::StreamExt; use prometheus::{Histogram, IntCounter, IntCounterVec, IntGauge}; -use serde::{Deserialize, Serialize}; +use serde::Deserialize; use std::sync::Arc; use std::time::Instant; use tracing::Instrument; @@ -141,7 +140,12 @@ impl Drop for RequestMetricsGuard { } } -impl Ingress { +impl Ingress +where + Req: PipelineIO + Sync, + Resp: PipelineIO, + Adapter: Send + Sync + 'static, +{ /// Pump every chunk from the engine's response stream out to the /// upstream-side `StreamSender`, plus the terminal complete-final /// frame. Captures the per-frame metrics, the publish-failure error @@ -154,7 +158,8 @@ impl Ingress { publisher: &StreamSender, payload_codec: RequestPlanePayloadCodec, ) where - U: Data + Serialize + MaybeError + std::fmt::Debug, + U: Data + std::fmt::Debug, + Adapter: IngressResponseEncoder, { let context = stream.context(); @@ -163,21 +168,31 @@ impl Ingress { let mut saw_error_response = false; while let Some(resp) = stream.next().await { tracing::trace!("Sending response: {:?}", resp); - let is_error = resp.err().is_some(); - if is_error { - saw_error_response = true; - } - let resp_wrapper = NetworkStreamWrapper { - data: Some(resp), - complete_final: false, + let encoded = match self + .payload_adapter + .encode_response(payload_codec, Some(resp), false) + .await + { + Ok(encoded) => encoded, + Err(err) => { + tracing::error!(%err, "failed to encode request-plane response"); + saw_error_response = true; + send_complete_final = false; + if let Some(m) = self.metrics() { + m.error_counter + .with_label_values(&[work_handler::error_types::SERIALIZATION]) + .inc(); + } + break; + } }; - let resp_bytes = payload_codec - .encode(&resp_wrapper) - .expect("fatal error: invalid request-plane response object"); + let is_error = encoded.is_error; + saw_error_response |= is_error; + let resp_bytes = encoded.bytes; if let Some(m) = self.metrics() { m.response_bytes.inc_by(resp_bytes.len() as u64); } - if (publisher.send(resp_bytes.into()).await).is_err() { + if (publisher.send(resp_bytes).await).is_err() { send_complete_final = false; if context.is_stopped() { // Say there are 2 threads accessing `context`, the sequence can be either: @@ -209,19 +224,36 @@ impl Ingress { notifier.notify_one(); } } + if encoded.stop_stream { + // Dropping the engine stream after the terminal frame is sent + // propagates cancellation to a producer that is still running. + // Stopping the context here can close the response transport + // before the queued error and clean terminal frames are read. + break; + } } if send_complete_final { - let resp_wrapper = NetworkStreamWrapper:: { - data: None, - complete_final: true, + let encoded = match self + .payload_adapter + .encode_response(payload_codec, None, true) + .await + { + Ok(encoded) => encoded, + Err(err) => { + tracing::error!(%err, "failed to encode request-plane final response"); + if let Some(m) = self.metrics() { + m.error_counter + .with_label_values(&[work_handler::error_types::PUBLISH_FINAL]) + .inc(); + } + return; + } }; - let resp_bytes = payload_codec - .encode(&resp_wrapper) - .expect("fatal error: invalid request-plane response final object"); + let resp_bytes = encoded.bytes; if let Some(m) = self.metrics() { m.response_bytes.inc_by(resp_bytes.len() as u64); } - if (publisher.send(resp_bytes.into()).await).is_err() { + if (publisher.send(resp_bytes).await).is_err() { tracing::error!( "Failed to publish complete final for stream {}", context.id() @@ -320,10 +352,11 @@ trait IngressDispatch: Send + Sync { } #[async_trait] -impl IngressDispatch for Ingress, ManyOut> +impl IngressDispatch for Ingress, ManyOut, Adapter> where T: Data + for<'de> Deserialize<'de> + std::fmt::Debug, - U: Data + Serialize + MaybeError + std::fmt::Debug, + U: Data + std::fmt::Debug, + Adapter: IngressRequestDecoder + Send + Sync + 'static, { type Request = SingleIn; @@ -346,18 +379,17 @@ where )) })?; let payload_codec = control_msg.payload_codec; - let request_t: T = payload_codec.decode(&data).map_err(|err| { - if let Some(m) = self.metrics() { - m.error_counter - .with_label_values(&[work_handler::error_types::DESERIALIZATION]) - .inc(); - } - PipelineError::DeserializationError(format!( - "Failed deserializing {} request payload: {}", - payload_codec.name(), - err - )) - })?; + let request_t: T = self + .payload_adapter + .decode_request(payload_codec, data) + .await + .inspect_err(|_| { + if let Some(m) = self.metrics() { + m.error_counter + .with_label_values(&[work_handler::error_types::DESERIALIZATION]) + .inc(); + } + })?; tracing::trace!( request_id = %control_msg.id, @@ -379,10 +411,11 @@ where } #[async_trait] -impl IngressDispatch for Ingress, ManyOut> +impl IngressDispatch for Ingress, ManyOut, Adapter> where T: Data + for<'de> Deserialize<'de> + std::fmt::Debug, - U: Data + Serialize + MaybeError + std::fmt::Debug, + U: Data + std::fmt::Debug, + Adapter: IngressRequestDecoder + Send + Sync + 'static, { type Request = ManyIn; @@ -461,6 +494,7 @@ where // header-only. let (frame_tx, frame_rx) = tokio::sync::mpsc::channel::(8); let forwarder_ctx = context_arc.clone(); + let payload_adapter = self.payload_adapter.clone(); tokio::spawn(async move { let mut rx = request_stream_recv.rx; while let Some(bytes) = rx.recv().await { @@ -471,7 +505,7 @@ where if forwarder_ctx.is_killed() || forwarder_ctx.is_stopped() { break; } - match payload_codec.decode::(&bytes) { + match payload_adapter.decode_request(payload_codec, bytes).await { Ok(item) => { if frame_tx.send(item).await.is_err() { tracing::debug!( @@ -506,9 +540,11 @@ where } } -impl Ingress> +impl Ingress, Adapter> where - U: Data + Serialize + MaybeError + std::fmt::Debug, + Req: PipelineIO + Sync, + U: Data + std::fmt::Debug, + Adapter: IngressResponseEncoder + Send + Sync + 'static, { /// Shared body of `PushWorkHandler::handle_payload` for every /// `Ingress>` shape that has an [`IngressDispatch`] @@ -636,10 +672,11 @@ where } #[async_trait] -impl PushWorkHandler for Ingress, ManyOut> +impl PushWorkHandler for Ingress, ManyOut, Adapter> where T: Data + for<'de> Deserialize<'de> + std::fmt::Debug, - U: Data + Serialize + MaybeError + std::fmt::Debug, + U: Data + std::fmt::Debug, + Adapter: IngressPayloadAdapter + Send + Sync + 'static, { fn add_metrics( &self, @@ -667,10 +704,11 @@ where } #[async_trait] -impl PushWorkHandler for Ingress, ManyOut> +impl PushWorkHandler for Ingress, ManyOut, Adapter> where T: Data + for<'de> Deserialize<'de> + std::fmt::Debug, - U: Data + Serialize + MaybeError + std::fmt::Debug, + U: Data + std::fmt::Debug, + Adapter: IngressPayloadAdapter + Send + Sync + 'static, { fn add_metrics( &self, diff --git a/lib/runtime/tests/bidirectional_e2e.rs b/lib/runtime/tests/bidirectional_e2e.rs index 38fe2a986ac1..9c77d2599d83 100644 --- a/lib/runtime/tests/bidirectional_e2e.rs +++ b/lib/runtime/tests/bidirectional_e2e.rs @@ -25,7 +25,10 @@ use dynamo_runtime::{ protocols::maybe_error::MaybeError, }; -use dynamo_runtime::pipeline::network::egress::push_router::{PushRouter, RouterMode}; +use dynamo_runtime::pipeline::network::{ + RequestPlanePayloadCodec, + egress::push_router::{PushRouter, RouterMode}, +}; #[derive(Clone, Debug, Deserialize, Serialize)] struct EchoResponse { @@ -70,6 +73,16 @@ impl AsyncEngine, ManyOut, Error> for EchoEngine { #[tokio::test] async fn bidirectional_end_to_end_echo() { + // This integration test is a standalone process, so clear any inherited + // override before the request-plane codec cache is initialized. + unsafe { + std::env::remove_var("DYN_REQUEST_PLANE_CODEC"); + } + assert_eq!( + RequestPlanePayloadCodec::configured(), + RequestPlanePayloadCodec::Msgpack + ); + let rt = Runtime::from_current().unwrap(); let drt = DistributedRuntime::new(rt.clone(), DistributedConfig::process_local()) .await diff --git a/lib/runtime/tests/bidirectional_e2e_json.rs b/lib/runtime/tests/bidirectional_e2e_json.rs new file mode 100644 index 000000000000..063dc745563e --- /dev/null +++ b/lib/runtime/tests/bidirectional_e2e_json.rs @@ -0,0 +1,241 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Process-isolated coverage for the explicit JSON request-plane fallback. +//! The codec is globally cached, so this must remain in its own integration +//! test binary rather than sharing a process with the default-codec test. + +use std::sync::Arc; + +use anyhow::Error; +use async_trait::async_trait; +use bytes::Bytes; +use serde::{Deserialize, Serialize}; + +use dynamo_runtime::pipeline::network::{ + EncodedResponseFrame, IngressRequestDecoder, IngressResponseEncoder, NetworkStreamWrapper, + RequestPlanePayloadCodec, + egress::push_router::{PushRouter, RouterMode}, +}; +use dynamo_runtime::{ + DistributedRuntime, Runtime, + distributed::DistributedConfig, + engine::{AsyncEngine, AsyncEngineContextProvider, DataStream}, + error::DynamoError, + metrics::MetricsHierarchy, + pipeline::{ + ManyIn, ManyOut, PipelineError, RequestStream, ResponseStream, context::Context, + network::Ingress, + }, + protocols::maybe_error::MaybeError, +}; + +#[derive(Clone, Debug, Deserialize, Serialize)] +struct EchoResponse { + value: Option, + #[serde(default)] + error: Option, +} + +impl MaybeError for EchoResponse { + fn from_err(err: impl std::error::Error + 'static) -> Self { + Self { + value: None, + error: Some(DynamoError::from( + Box::new(err) as Box + )), + } + } + + fn err(&self) -> Option { + self.error.clone() + } +} + +struct EchoEngine; + +#[async_trait] +impl AsyncEngine, ManyOut, Error> for EchoEngine { + async fn generate(&self, input: ManyIn) -> Result, Error> { + let ctx = input.context(); + let (request_stream, _ctx_unit) = input.into_parts(); + let inner = request_stream + .take() + .expect("RequestStream::take called twice on EchoEngine input"); + let mapped = futures::StreamExt::map(inner, |value| EchoResponse { + value: Some(value), + error: None, + }); + let stream: DataStream = Box::pin(mapped); + Ok(ResponseStream::new(stream, ctx)) + } +} + +#[derive(Debug)] +struct NonSerdeResponse; + +struct FailingResponseEngine; + +#[async_trait] +impl AsyncEngine, ManyOut, Error> for FailingResponseEngine { + async fn generate(&self, input: ManyIn) -> Result, Error> { + let ctx = input.context(); + let (request_stream, _ctx_unit) = input.into_parts(); + let inner = request_stream + .take() + .expect("RequestStream::take called twice on FailingResponseEngine input"); + let stream: DataStream = + Box::pin(futures::StreamExt::map(inner, |_| NonSerdeResponse)); + Ok(ResponseStream::new(stream, ctx)) + } +} + +#[derive(Debug)] +struct FailingResponseAdapter; + +impl IngressRequestDecoder for FailingResponseAdapter { + async fn decode_request( + &self, + payload_codec: RequestPlanePayloadCodec, + bytes: Bytes, + ) -> Result { + payload_codec.decode(&bytes).map_err(|error| { + PipelineError::DeserializationError(format!( + "failed decoding test request as {}: {error}", + payload_codec.name() + )) + }) + } +} + +impl IngressResponseEncoder for FailingResponseAdapter { + async fn encode_response( + &self, + payload_codec: RequestPlanePayloadCodec, + response: Option, + complete_final: bool, + ) -> Result { + if response.is_some() { + return Err(PipelineError::SerializationError( + "intentional response encoding failure".to_string(), + )); + } + assert!( + complete_final, + "only a clean terminal frame has no response" + ); + let bytes = payload_codec + .encode(&NetworkStreamWrapper:: { + data: None, + complete_final: true, + }) + .map_err(|error| PipelineError::SerializationError(error.to_string()))?; + Ok(EncodedResponseFrame { + bytes: bytes.into(), + is_error: false, + stop_stream: false, + }) + } +} + +#[tokio::test] +async fn bidirectional_end_to_end_echo_with_explicit_json_codec() { + // This test binary owns the process and sets the value before the first + // request initializes the request-plane codec cache. + unsafe { + std::env::set_var("DYN_REQUEST_PLANE_CODEC", "json"); + } + assert_eq!( + RequestPlanePayloadCodec::configured(), + RequestPlanePayloadCodec::Json + ); + + let rt = Runtime::from_current().unwrap(); + let drt = DistributedRuntime::new(rt.clone(), DistributedConfig::process_local()) + .await + .unwrap(); + let ns = drt.namespace("test_bidi_e2e_json".to_string()).unwrap(); + let component = ns.component("echo_component".to_string()).unwrap(); + let endpoint = component.endpoint("echo_endpoint".to_string()); + + let ingress = Ingress::for_engine(Arc::new(EchoEngine)).unwrap(); + let endpoint_for_server = endpoint.clone(); + tokio::spawn(async move { + let _ = endpoint_for_server + .endpoint_builder() + .handler(ingress) + .start() + .await; + }); + + let client = endpoint.client().await.unwrap(); + client.wait_for_instances().await.unwrap(); + + let router = PushRouter::::from_client(client, RouterMode::RoundRobin) + .await + .unwrap(); + let input: ManyIn = Context::new(RequestStream::new(Box::pin(tokio_stream::iter(vec![ + 10u64, 20, 30, + ])))); + let response_stream = router.generate(input).await.unwrap(); + let responses: Vec = futures::StreamExt::collect(response_stream).await; + + assert_eq!( + responses + .iter() + .filter_map(|response| response.value) + .collect::>(), + vec![10u64, 20, 30] + ); + + let failing_endpoint = component.endpoint("failing_endpoint".to_string()); + let failing_ingress = + Ingress::for_engine_with_adapter(Arc::new(FailingResponseEngine), FailingResponseAdapter) + .unwrap(); + let failing_endpoint_for_server = failing_endpoint.clone(); + tokio::spawn(async move { + let _ = failing_endpoint_for_server + .endpoint_builder() + .handler(failing_ingress) + .start() + .await; + }); + + let failing_client = failing_endpoint.client().await.unwrap(); + failing_client.wait_for_instances().await.unwrap(); + let failing_router = + PushRouter::::from_client(failing_client, RouterMode::RoundRobin) + .await + .unwrap(); + let failing_input: ManyIn = + Context::new(RequestStream::new(Box::pin(tokio_stream::iter(vec![ + 99u64, + ])))); + let failure_stream = failing_router.generate(failing_input).await.unwrap(); + let failure_responses: Vec = futures::StreamExt::collect(failure_stream).await; + let error = failure_responses + .into_iter() + .find_map(|response| response.error) + .expect("encoding failure must not be reported as a clean terminal frame"); + assert!( + error + .to_string() + .contains("Stream ended before generation completed"), + "unexpected client error: {error}" + ); + + let metrics = failing_endpoint.metrics().prometheus_expfmt().unwrap(); + let serialization_error = metrics + .lines() + .find(|line| { + !line.starts_with('#') + && line.contains("errors_total") + && line.contains("error_type=\"serialization\"") + }) + .expect("serialization error metric must be present"); + assert!( + serialization_error.ends_with(" 1"), + "unexpected serialization error metric: {serialization_error}" + ); + rt.shutdown(); +} diff --git a/tests/frontend/realtime_echo_worker.py b/tests/frontend/realtime_echo_worker.py index 1ff8cd7e78e6..605db0554386 100644 --- a/tests/frontend/realtime_echo_worker.py +++ b/tests/frontend/realtime_echo_worker.py @@ -67,34 +67,50 @@ async def _python_realtime_echo(request_stream, context): audio = client_event.get("audio", "") response_id = f"resp_{uuid.uuid4().hex}" item_id = f"item_{uuid.uuid4().hex}" - - yield { + # Intentionally reuse this mutable dict across yields. The Rust + # network engine must serialize each frame before polling us again, + # otherwise queued PyObject handles would all observe later updates. + response_event = { "type": "response.created", "event_id": _event_id(), "response": _response_payload(response_id, "in_progress"), } - yield { - "type": "response.output_audio.delta", - "event_id": _event_id(), - "response_id": response_id, - "item_id": item_id, - "output_index": 0, - "content_index": 0, - "delta": audio, - } - yield { - "type": "response.output_audio.done", - "event_id": _event_id(), - "response_id": response_id, - "item_id": item_id, - "output_index": 0, - "content_index": 0, - } - yield { - "type": "response.done", - "event_id": _event_id(), - "response": _response_payload(response_id, "completed"), - } + + yield response_event + response_event.clear() + response_event.update( + { + "type": "response.output_audio.delta", + "event_id": _event_id(), + "response_id": response_id, + "item_id": item_id, + "output_index": 0, + "content_index": 0, + "delta": audio, + } + ) + yield response_event + response_event.clear() + response_event.update( + { + "type": "response.output_audio.done", + "event_id": _event_id(), + "response_id": response_id, + "item_id": item_id, + "output_index": 0, + "content_index": 0, + } + ) + yield response_event + response_event.clear() + response_event.update( + { + "type": "response.done", + "event_id": _event_id(), + "response": _response_payload(response_id, "completed"), + } + ) + yield response_event else: yield { "type": "error",