diff --git a/lib/llm/src/engines.rs b/lib/llm/src/engines.rs index dc497e0944df..0488252ccd36 100644 --- a/lib/llm/src/engines.rs +++ b/lib/llm/src/engines.rs @@ -10,13 +10,12 @@ use async_stream::stream; use async_trait::async_trait; use dynamo_runtime::engine::{AsyncEngine, AsyncEngineContextProvider, ResponseStream}; +#[cfg(test)] +use dynamo_runtime::pipeline::RequestStream; use dynamo_runtime::pipeline::{Error, ManyIn, ManyOut, SingleIn}; use dynamo_runtime::protocols::annotated::Annotated; use futures::StreamExt; -#[cfg(test)] -use dynamo_runtime::engine::RequestStream; - use crate::protocols::openai::{ chat_completions::{NvCreateChatCompletionRequest, NvCreateChatCompletionStreamResponse}, completions::{NvCreateCompletionRequest, NvCreateCompletionResponse, prompt_to_string}, @@ -152,11 +151,15 @@ impl AsyncEngine, ManyOut, + incoming: ManyIn, ) -> Result>, Error> { let ctx = incoming.context(); let session_id = ctx.id().to_string(); let ctx_for_stream = ctx.clone(); + let (request_stream, _ctx_unit) = incoming.into_parts(); + let mut incoming = request_stream.take().ok_or_else(|| { + anyhow::anyhow!("RequestStream::take called twice on bidirectional input") + })?; let output = stream! { let ctx = ctx_for_stream; @@ -561,7 +564,7 @@ mod tests { use futures::stream; fn make_input(events: Vec) -> ManyIn { - RequestStream::new(Box::pin(stream::iter(events)), Context::new(()).context()) + Context::new(RequestStream::new(Box::pin(stream::iter(events)))) } fn parse_client_event(json: serde_json::Value) -> RealtimeClientEvent { diff --git a/lib/llm/src/http/service/realtime.rs b/lib/llm/src/http/service/realtime.rs index ce6f1b018881..87bcf339eca9 100644 --- a/lib/llm/src/http/service/realtime.rs +++ b/lib/llm/src/http/service/realtime.rs @@ -43,8 +43,8 @@ use axum::{ response::Response, routing::get, }; -use dynamo_runtime::engine::{AsyncEngineContextProvider, RequestStream}; -use dynamo_runtime::pipeline::Context; +use dynamo_runtime::engine::AsyncEngineContextProvider; +use dynamo_runtime::pipeline::{Context, RequestStream}; use futures::{SinkExt, StreamExt, stream::SplitSink}; use tokio::sync::mpsc; use tokio_stream::wrappers::ReceiverStream; @@ -145,7 +145,7 @@ async fn handle_socket(socket: WebSocket, state: Arc) { } let request_stream = Box::pin(ReceiverStream::new(req_rx)); - let input = RequestStream::new(request_stream, Context::new(()).context()); + let input = Context::new(RequestStream::new(request_stream)); let mut response_stream = match engine.generate(input).await { Ok(s) => s, diff --git a/lib/runtime/src/engine.rs b/lib/runtime/src/engine.rs index 32e30824dd94..f6702457cf71 100644 --- a/lib/runtime/src/engine.rs +++ b/lib/runtime/src/engine.rs @@ -180,14 +180,9 @@ pub trait AsyncEngineUnary: /// /// This trait combines `Stream` semantics with context provider capabilities, /// representing a continuous async operation that produces multiple messages over time. -/// -/// - **Output side:** wrapped as [`EngineStream`] = `crate::pipeline::ManyOut` -/// — the stream of response chunks an engine emits. -/// - **Input side:** same `EngineStream` shape, exposed as -/// `crate::pipeline::ManyIn` for documentary clarity at the call site. -/// -/// [`ResponseStream`] is the canonical concrete implementor; [`RequestStream`] -/// is a type alias of it for the input side. +/// The canonical boxed form is [`EngineStream`] (= `crate::pipeline::ManyOut`), +/// the stream of response chunks an engine emits; [`ResponseStream`] is the +/// canonical concrete implementor. pub trait AsyncEngineStream: Stream + AsyncEngineContextProvider + Send {} /// Engine is a trait that defines the interface for a streaming engine. @@ -267,16 +262,6 @@ impl Debug for ResponseStream { } } -/// Input-side type alias of [`ResponseStream`] — same struct, different name to -/// signal role at the call site. -/// -/// The shape is identical: a `(stream, ctx)` pair that implements [`Stream`], -/// [`AsyncEngineContextProvider`], and [`AsyncEngineStream`]. Use `RequestStream` -/// when you're constructing a value to feed into the `Req` slot of an engine, -/// and [`ResponseStream`] when constructing a value to emit from the `Resp` slot. -/// Functionally interchangeable. -pub type RequestStream = ResponseStream; - impl AsyncEngineContextProvider for Pin>> { fn context(&self) -> Arc { AsyncEngineContextProvider::context(&**self) diff --git a/lib/runtime/src/pipeline.rs b/lib/runtime/src/pipeline.rs index 8f81fb0719e7..5a6f119efb4f 100644 --- a/lib/runtime/src/pipeline.rs +++ b/lib/runtime/src/pipeline.rs @@ -19,8 +19,8 @@ pub use network::egress::push_router::{PushRouter, RouterMode, WorkerLoadMonitor pub mod registry; pub use crate::engine::{ - self as engine, AsyncEngine, AsyncEngineContext, AsyncEngineContextProvider, Data, DataStream, - Engine, EngineStream, EngineUnary, RequestStream, ResponseStream, async_trait, + self as engine, AsyncEngine, AsyncEngineContext, AsyncEngineContextProvider, AsyncEngineStream, + Data, DataStream, Engine, EngineStream, EngineUnary, ResponseStream, async_trait, }; pub use anyhow::Error; pub use context::Context; @@ -30,18 +30,44 @@ pub use error::{PipelineError, PipelineErrorExt, TwoPartCodecError}; /// about the request. This information propagates through the stages, both local and distributed. pub type SingleIn = Context; -/// Pipeline input for streaming-request engines: a trait-object alias around a -/// stream that carries its own cancellation context (via the -/// [`AsyncEngineContextProvider`] half of [`AsyncEngineStream`]). -/// -/// This is the input-side mirror of [`ManyOut`]. Both aliases resolve to the -/// same underlying [`EngineStream`] type; the directional names are -/// documentary. Earlier definitions wrapped the stream in a `Context<…>` -/// (`Context>`); that shape was uninstantiable because -/// `DataStream` is `!Sync` while `Context` requires `Sync`. The -/// trait-object form solves that cleanly: the cancellation surface is part of -/// the trait contract rather than an outer wrapper. -pub type ManyIn = EngineStream; +/// `Sync` ownership cell for a non-`Sync` [`DataStream`]. [`Self::take`] +/// moves the inner stream out; the mutex serialises concurrent attempts so +/// the first caller observes `Some(stream)` and all later callers see +/// `None`. Iteration happens on the returned `DataStream`. +pub struct RequestStream { + inner: std::sync::Mutex>>, +} + +impl RequestStream { + /// Wrap a [`DataStream`] in a `Sync` ownership cell. + pub fn new(stream: DataStream) -> Self { + Self { + inner: std::sync::Mutex::new(Some(stream)), + } + } + + /// Atomically move the inner stream out. Returns `Some(stream)` exactly + /// once across all threads racing on the same `RequestStream`; every + /// subsequent call (on any thread) returns `None`. The returned stream + /// is the unique owner; the wrapper retains nothing. + pub fn take(&self) -> Option> { + self.inner.lock().unwrap().take() + } +} + +impl std::fmt::Debug for RequestStream { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + let taken = self.inner.lock().map(|g| g.is_none()).unwrap_or(true); + f.debug_struct("RequestStream") + .field("taken", &taken) + .finish() + } +} + +/// Pipeline input for streaming-request engines: a [`RequestStream`] +/// payload wrapped in a [`Context`], symmetric to [`SingleIn`] for unary +/// inputs. +pub type ManyIn = Context>; /// Type alias for the output of pipeline that returns a single value pub type SingleOut = EngineUnary; diff --git a/lib/runtime/src/pipeline/network.rs b/lib/runtime/src/pipeline/network.rs index df211f524b8a..c0a3c3ede6c2 100644 --- a/lib/runtime/src/pipeline/network.rs +++ b/lib/runtime/src/pipeline/network.rs @@ -28,10 +28,12 @@ use super::{AsyncEngine, AsyncEngineContext, AsyncEngineContextProvider, Respons use serde::{Deserialize, Serialize}; use super::{ - AsyncTransportEngine, Context, Data, Error, ManyOut, PipelineError, PipelineIO, SegmentSource, - ServiceBackend, ServiceEngine, SingleIn, Source, context, + AsyncTransportEngine, Context, Data, Error, ManyIn, ManyOut, PipelineError, PipelineIO, + SegmentSource, ServiceBackend, ServiceEngine, SingleIn, Source, context, }; use crate::metrics::MetricsHierarchy; +use crate::metrics::prometheus_names::work_handler; +use crate::protocols::maybe_error::MaybeError; use ingress::push_handler::WorkHandlerMetrics; use prometheus::{CounterVec, Histogram, IntCounter, IntCounterVec, IntGauge}; @@ -94,6 +96,11 @@ pub(crate) struct RequestControlMessage { /// Reliable for single-machine profiling; treat cross-host values as approximate. #[serde(default, skip_serializing_if = "Option::is_none")] pub(crate) frontend_send_ts_ns: Option, + /// For bidirectional dispatch (`request_type == ManyIn`): connection info the + /// worker dials back to in order to receive subsequent request frames. `None` + /// for the unary path, which is the wire-compatible default. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub(crate) request_stream_connection_info: Option, } #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] diff --git a/lib/runtime/src/pipeline/network/egress/addressed_router.rs b/lib/runtime/src/pipeline/network/egress/addressed_router.rs index 13e52da18af8..f82518f51b1c 100644 --- a/lib/runtime/src/pipeline/network/egress/addressed_router.rs +++ b/lib/runtime/src/pipeline/network/egress/addressed_router.rs @@ -21,23 +21,100 @@ use crate::metrics::request_plane::{ use crate::pipeline::network::ConnectionInfo; use crate::pipeline::network::NetworkStreamWrapper; use crate::pipeline::network::PendingConnections; +use crate::pipeline::network::RegisteredStream; +use crate::pipeline::network::RequestControlMessage; +use crate::pipeline::network::RequestType; +use crate::pipeline::network::ResponseType; use crate::pipeline::network::StreamOptions; +use crate::pipeline::network::StreamProvider; +use crate::pipeline::network::StreamReceiver; +use crate::pipeline::network::StreamSender; use crate::pipeline::network::TwoPartCodec; use crate::pipeline::network::codec::TwoPartMessage; use crate::pipeline::network::tcp; -use crate::pipeline::{ManyOut, PipelineError, ResponseStream, SingleIn}; +use crate::pipeline::{ManyIn, ManyOut, PipelineError, ResponseStream, SingleIn}; use crate::protocols::maybe_error::MaybeError; use crate::traits::DistributedRuntimeProvider; use anyhow::{Error, Result}; use futures::stream::Stream; -use serde::Deserialize; -use serde::Serialize; use std::pin::Pin; use std::task::{Context, Poll}; use tokio_stream::{StreamExt, StreamNotifyClose, wrappers::ReceiverStream}; use tracing::Instrument; +/// Stream transformation helper that: +/// - decodes a response byte stream from network into the fully-shaped `ManyOut` +/// - emits TTFT and transport-roundtrip metrics on first response +/// - hands off the `InflightGuard` to a stream-lifetime `InflightDecStream` so +/// the inflight gauge stays accurate for the whole response lifetime. +fn decode_response_stream( + response_rx: tokio::sync::mpsc::Receiver, + engine_ctx: Arc, + queue_start: Instant, + tx_start: Instant, + inflight_guard: InflightGuard, +) -> ManyOut +where + U: Data + for<'de> Deserialize<'de> + MaybeError, +{ + let engine_ctx_for_stream = engine_ctx.clone(); + let mut is_complete_final = false; + let mut first_response = true; + let stream = StreamNotifyClose::new(ReceiverStream::new(response_rx)).filter_map(move |res| { + if let Some(res_bytes) = res { + if first_response { + first_response = false; + REQUEST_PLANE_ROUNDTRIP_TTFT_SECONDS.observe(tx_start.elapsed().as_secs_f64()); + STAGE_DURATION_SECONDS + .with_label_values(&["transport_roundtrip"]) + .observe(queue_start.elapsed().as_secs_f64()); + } + if is_complete_final { + let err = DynamoError::msg( + "Response received after generation ended - this should never happen", + ); + return Some(U::from_err(err)); + } + match serde_json::from_slice::>(&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 { + let err = + DynamoError::msg("Empty response received - this should never happen"); + Some(U::from_err(err)) + } + } + Err(err) => { + let json_str = String::from_utf8_lossy(&res_bytes); + tracing::warn!(%err, %json_str, "Failed deserializing JSON to response"); + Some(U::from_err(DynamoError::msg(err.to_string()))) + } + } + } else if is_complete_final { + None + } else if engine_ctx_for_stream.is_stopped() { + tracing::debug!("Request cancelled and then trying to read a response"); + None + } else { + let err = DynamoError::builder() + .error_type(ErrorType::Disconnected) + .message("Stream ended before generation completed") + .build(); + tracing::debug!("{err}"); + Some(U::from_err(err)) + } + }); + + inflight_guard.disarm(); + let stream = InflightDecStream { inner: stream }; + ResponseStream::new(Box::pin(stream), engine_ctx) +} + const CONTROL_MESSAGE_MAX_BYTES: usize = 128 * 1024; fn serialize_control_message(control_message: &RequestControlMessage) -> Result, Error> { @@ -53,6 +130,145 @@ fn serialize_control_message(control_message: &RequestControlMessage) -> Result< Ok(ctrl) } +/// Build the request control message, and serialize for transfer. +/// +/// `request` provides the optional unary request payload. Should set for +/// SingleIn generation. +/// `send_conn_info` provides the connection info for the request stream. +/// Should set for ManyIn generation. +fn build_request_envelope( + context: &context::Context<()>, + recv_conn_info: ConnectionInfo, + send_conn_info: Option, + request: Option<&T>, +) -> Result +where + T: serde::Serialize + ?Sized, +{ + let request_id = context.id(); + let request_type = if send_conn_info.is_some() { + RequestType::ManyIn + } else { + RequestType::SingleIn + }; + let control_message = RequestControlMessage { + id: request_id.to_string(), + request_type, + response_type: ResponseType::ManyOut, + connection_info: recv_conn_info, + metadata: context.metadata().clone(), + frontend_send_ts_ns: None, + request_stream_connection_info: send_conn_info, + }; + + let ctrl = serialize_control_message(&control_message)?; + let data: Option> = match request { + Some(req) => Some(serde_json::to_vec(req)?), + None => None, + }; + + let msg = match data { + Some(d) => { + tracing::trace!( + request_id, + "packaging two-part message; ctrl: {} bytes, data: {} bytes", + ctrl.len(), + d.len(), + ); + TwoPartMessage::from_parts(ctrl.into(), d.into()) + } + None => { + tracing::trace!( + request_id, + "packaging bidirectional header-only envelope; ctrl: {} bytes", + ctrl.len(), + ); + TwoPartMessage::from_header(ctrl.into()) + } + }; + + let codec = TwoPartCodec::default(); + let buffer = codec.encode_message(msg)?; + Ok(buffer) +} + +/// Await the network request-stream dial-in (if `request_stream_provider` is `Some`) +/// and spawn a detached task that forwards every item from `input_stream` onto +/// the request stream. Returns once the forwarder is spawned; `Err` if request-stream +/// dial-in fails. +async fn spawn_request_stream_forwarder( + request_stream_provider: Option>, + mut input_stream: crate::engine::DataStream, + engine_ctx: Arc, +) -> Result<(), Error> +where + T: serde::Serialize + Send + 'static, +{ + let Some(provider) = request_stream_provider else { + return Ok(()); + }; + + let request_sender = match provider.await { + Ok(Ok(sender)) => sender, + Ok(Err(e)) => { + return Err(anyhow::anyhow!( + DynamoError::builder() + .error_type(ErrorType::CannotConnect) + .message(format!("Worker dial-in failed for request stream: {e}")) + .build() + )); + } + Err(_) => { + return Err(anyhow::anyhow!( + DynamoError::builder() + .error_type(ErrorType::Disconnected) + .message("Worker disconnected before request stream was established") + .build() + )); + } + }; + + // The task exits on stream end, context kill/stop, send error (worker + // dropped its receiver), or local serialize failure. On any exit + // `request_sender` drops and triggers transport shutdown (see server.rs for details) + // which closes the upstream mpsc, triggering the server-side handler to emit + // `Sentinel`, which signals the worker's reader to end cleanly. + tokio::spawn(async move { + loop { + let item = tokio::select! { + biased; + _ = engine_ctx.killed() => break, + _ = engine_ctx.stopped() => break, + item = input_stream.next() => match item { + Some(item) => item, + None => break, + }, + }; + let bytes = match serde_json::to_vec(&item) { + Ok(b) => b, + Err(e) => { + // Stream-side framing failure: the engine sees a + // partial input, so kill the context to abort both + // directions consistently rather than silently + // dropping frames. + tracing::error!( + error = %e, + "failed to serialize bidirectional request frame; killing context" + ); + engine_ctx.kill(); + break; + } + }; + if request_sender.send(bytes.into()).await.is_err() { + tracing::debug!("worker request-stream receiver dropped; forwarder exiting"); + break; + } + } + }); + + Ok(()) +} + /// RAII guard that decrements REQUEST_PLANE_INFLIGHT on drop unless disarmed. /// Protects against gauge leaks when `?` operators cause early returns between /// the increment and `InflightDecStream` construction. @@ -102,6 +318,15 @@ impl Drop for InflightDecStream { } } +/// Extract the TCP stream subject from a [`ConnectionInfo`], if it carries a +/// well-formed [`tcp::TcpStreamConnectionInfo`]. Used for the pre-dispatch +/// tombstone check. +fn subject_of(conn_info: &ConnectionInfo) -> Option { + serde_json::from_str::(&conn_info.info) + .ok() + .map(|ci| ci.subject) +} + pub struct AddressedRequest { request: T, address: String, @@ -188,173 +413,132 @@ impl AddressedPushRouter { .clear_instance_tombstone(instance_id) .await } -} -#[async_trait::async_trait] -impl AsyncEngine>, ManyOut, Error> for AddressedPushRouter -where - T: Data + Serialize, - U: Data + for<'de> Deserialize<'de> + MaybeError, -{ - async fn generate(&self, request: SingleIn>) -> Result, Error> { - let queue_start = Instant::now(); - REQUEST_PLANE_INFLIGHT.inc(); - let inflight_guard = InflightGuard::new(); + /// Bidirectional generation. Note that it doesn't implement the AsyncEngine trait directly + /// because there is no trivial way to wrap (instance and address) into ManyIn style. + /// May wrap as SingleIn> and unwrap here but really just syntax + /// sugar, so we just do it inline here. Will consider only if we do want to call this from + /// typed erased AsyncEngine impls. + pub async fn generate_bidirectional( + &self, + instance: Instance, + address: String, + input: ManyIn, + ) -> Result, Error> + where + T: Data + Serialize, + U: Data + for<'de> Deserialize<'de> + MaybeError, + { + let (request_stream, context) = input.into_parts(); + let input_stream = request_stream.take().ok_or_else(|| { + anyhow::anyhow!("RequestStream::take called twice on bidirectional dispatch input") + })?; + + self.dispatch_and_finalize::( + &context, + address, + Some(&instance), + None, + Some(input_stream), + ) + .await + } - let request_id = request.context().id().to_string(); - let (addressed_request, context) = request.transfer(()); - let (request, address, instance_info) = addressed_request.into_parts(); + /// Shared dispatch core for both unary and bidirectional requests. Wire + /// shape is inferred from the inputs: + /// - `input_stream = Some(_)` + `request = None` → bidirectional, + /// header-only envelope. The worker dials back for both halves and + /// pulls request frames off the spawned forwarder. + /// - `input_stream = None` + `request = Some(_)` → unary, two-part + /// `[ctrl, data]` envelope. The payload travels in the data part. + async fn dispatch_and_finalize( + &self, + context: &context::Context<()>, + address: String, + instance: Option<&Instance>, + request: Option<&T>, + input_stream: Option>, + ) -> Result, Error> + where + T: Data + Serialize, + U: Data + for<'de> Deserialize<'de> + MaybeError, + { let engine_ctx = context.context(); - let engine_ctx_ = engine_ctx.clone(); - // registration options for the data plane in a singe in / many out configuration - let options = StreamOptions::builder() - .context(engine_ctx.clone()) - .enable_request_stream(false) - .enable_response_stream(true) - .build() - .unwrap(); - - // register our needs with the data plane - // todo - generalize this with a generic data plane object which hides the specific transports - let pending_connections: PendingConnections = self.resp_transport.register(options).await; - - // validate and unwrap the RegisteredStream object - let pending_response_stream = match pending_connections.into_parts() { - (None, Some(recv_stream)) => recv_stream, - _ => { - panic!("Invalid data plane registration for a SingleIn/ManyOut transport"); - } - }; - - // separate out the connection info and the stream provider from the registered stream - let (connection_info, response_stream_provider) = pending_response_stream.into_parts(); - - // Snapshot subject before connection_info is moved; used for cleanup. - let recv_subject: Option = - serde_json::from_str::(&connection_info.info) - .ok() - .map(|ci| ci.subject); + let queue_start = Instant::now(); + REQUEST_PLANE_INFLIGHT.inc(); + let inflight_guard = InflightGuard::new(); - // If the instance is already tombstoned, fail fast with a migratable - // error instead of writing to the request plane. - if let (Some(subject), Some(inst)) = (&recv_subject, &instance_info) { - let endpoint_instance_id = inst.endpoint_instance_id(); - if !self + let enable_request_stream = input_stream.is_some(); + + // Hold the `RegisteredStream` as their RAII cleanup stays armed while held, + // which simplifies the cancellation of registration on error. Each side is + // disarmed by `into_parts()` on awaiting stream provider: past that point the + // subject is reaped by the worker's dial-in (instance healthy) or the discovery + // watcher (instance dropped), so no cleanup is owed. + let (send_registered, recv_registered) = self + .register_streams(engine_ctx.clone(), enable_request_stream, true) + .await?; + let recv_registered = recv_registered.ok_or_else(|| { + anyhow::anyhow!("response stream registration missing despite enable_response_stream") + })?; + + // Tombstone check: if discovery already removed the worker, fail fast + // with a migratable error rather than writing to the request plane. + // Dropping the held registrations on this return runs their cleanup. + let recv_subject = subject_of(&recv_registered.connection_info); + let send_subject = send_registered + .as_ref() + .and_then(|r| subject_of(&r.connection_info)); + if let (Some(subject), Some(inst)) = (&recv_subject, instance) + && !self .resp_transport - .associate_instance(subject, None, &endpoint_instance_id) + .associate_instance( + subject, + send_subject.as_deref(), + &inst.endpoint_instance_id(), + ) .await - { - return Err(anyhow::anyhow!( - DynamoError::builder() - .error_type(ErrorType::Disconnected) - .message( - "Worker removed before request could be sent (tombstoned instance)" - ) - .build() - )); - } + { + return Err(anyhow::anyhow!( + DynamoError::builder() + .error_type(ErrorType::Disconnected) + .message("Worker removed before request could be sent (tombstoned instance)") + .build() + )); } - // package up the connection info as part of the "header" component of the two part message - // used to issue the request on the - // todo -- this object should be automatically created by the register call, and achieved by to the two into_parts() - // calls. all the information here is provided by the [`StreamOptions`] object and/or the dataplane object - let control_message = RequestControlMessage { - id: engine_ctx.id().to_string(), - request_type: RequestType::SingleIn, - response_type: ResponseType::ManyOut, - connection_info, - metadata: context.metadata().clone(), - frontend_send_ts_ns: None, - }; - - // next build the two part message where we package the connection info and the request into - // a single Vec that can be sent over the wire. - // --- package this up in the WorkQueuePublisher --- - let ctrl = match serialize_control_message(&control_message) { - Ok(v) => v, - Err(e) => { - if let Some(subject) = &recv_subject { - self.resp_transport.cancel_recv_stream(subject).await; - } - return Err(e); - } - }; - let data = match serde_json::to_vec(&request) { - Ok(v) => v, - Err(e) => { - if let Some(subject) = &recv_subject { - self.resp_transport.cancel_recv_stream(subject).await; - } - return Err(e.into()); - } - }; - - tracing::trace!( - request_id, - "packaging two-part message; ctrl: {} bytes, data: {} bytes", - ctrl.len(), - data.len() - ); - - let msg = TwoPartMessage::from_parts(ctrl.into(), data.into()); - - // the request plane / work queue should provide a two part message codec that can be used - // or it should take a two part message directly - // todo - update this - let codec = TwoPartCodec::default(); - let buffer = match codec.encode_message(msg) { - Ok(v) => v, - Err(e) => { - if let Some(subject) = &recv_subject { - self.resp_transport.cancel_recv_stream(subject).await; - } - return Err(e.into()); - } - }; - + let buffer = build_request_envelope( + context, + recv_registered.connection_info.clone(), + send_registered.as_ref().map(|r| r.connection_info.clone()), + request, + )?; REQUEST_PLANE_QUEUE_SECONDS.observe(queue_start.elapsed().as_secs_f64()); - let tx_start = Instant::now(); - - // TRANSPORT ABSTRACT REQUIRED - END HERE - - // Send request using unified client interface - tracing::trace!( - request_id, - transport = self.req_client.transport_name(), - address = %address, - "Sending request via request plane client" - ); - // Prepare trace headers using shared helper - let mut headers = std::collections::HashMap::new(); - inject_trace_headers_into_map(&mut headers); - headers.insert("request-id".to_string(), request_id.clone()); - - // Stamp send time right before the transport write so the network - // transit metric excludes serialization/encoding overhead. - let send_ts_ns = std::time::SystemTime::now() - .duration_since(std::time::UNIX_EPOCH) - .unwrap_or_default() - .as_nanos() as u64; - headers.insert("x-frontend-send-ts-ns".to_string(), send_ts_ns.to_string()); - - // Phase A: Frontend → Backend (network + queue + ack) - let _nvtx_send = dynamo_nvtx_range!("transport.tcp.send"); - let send_result = self.req_client.send_request(address, buffer, headers).await; - drop(_nvtx_send); + let tx_start = Instant::now(); + self.dispatch_buffer(address, buffer, context.id()).await?; + REQUEST_PLANE_SEND_SECONDS.observe(tx_start.elapsed().as_secs_f64()); - if let Err(e) = send_result { - if let Some(subject) = &recv_subject { - self.resp_transport.cancel_recv_stream(subject).await; - } - return Err(e); + // Spawn the forwarder before awaiting the response prologue so request + // frames pre-load into the worker's input buffer while the engine + // initialises in parallel. The response provider only resolves after + // `engine.generate()` returns; awaiting it second avoids stalling the + // request-side handshake on engine setup latency. + if let Some(stream) = input_stream { + let request_stream_provider = send_registered.map(|r| { + let (_conn_info, provider) = r.into_parts(); + provider + }); + spawn_request_stream_forwarder(request_stream_provider, stream, engine_ctx.clone()) + .await?; } - REQUEST_PLANE_SEND_SECONDS.observe(tx_start.elapsed().as_secs_f64()); let _nvtx_wait = dynamo_nvtx_range!("transport.tcp.wait_backend"); - tracing::trace!(request_id, "awaiting transport handshake"); + tracing::trace!(request_id = context.id(), "awaiting transport handshake"); + + // Disarms the recv-side cleanup; see the holding rationale above. + let (_recv_conn_info, response_stream_provider) = recv_registered.into_parts(); // RecvError → migratable Disconnected (watcher cancelled the subject // or the worker died before establishing the response stream). @@ -367,9 +551,6 @@ where // opaque string today, so app-level rejections also retry // -- safe because no side effects are visible yet. Follow-up: // structured prologue error type for finer routing. - if let Some(subject) = &recv_subject { - self.resp_transport.cancel_recv_stream(subject).await; - } return Err(anyhow::anyhow!( DynamoError::builder() .error_type(ErrorType::CannotConnect) @@ -382,9 +563,6 @@ where Err(_recv_err) => { // oneshot dropped: either the discovery watcher cancelled // this subject or the worker died mid-handshake. - if let Some(subject) = &recv_subject { - self.resp_transport.cancel_recv_stream(subject).await; - } return Err(anyhow::anyhow!( DynamoError::builder() .error_type(ErrorType::Disconnected) @@ -395,73 +573,104 @@ where }; drop(_nvtx_wait); - // TODO: Detect end-of-stream using Server-Sent Events (SSE) - let mut is_complete_final = false; - let mut first_response = true; - 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 first_response { - first_response = false; - let roundtrip_ttft = tx_start.elapsed().as_secs_f64(); - REQUEST_PLANE_ROUNDTRIP_TTFT_SECONDS.observe(roundtrip_ttft); - STAGE_DURATION_SECONDS - .with_label_values(&["transport_roundtrip"]) - .observe(queue_start.elapsed().as_secs_f64()); - } - if is_complete_final { - let err = DynamoError::msg( - "Response received after generation ended - this should never happen", - ); - return Some(U::from_err(err)); - } - match serde_json::from_slice::>(&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 { - let err = DynamoError::msg( - "Empty response received - this should never happen", - ); - Some(U::from_err(err)) - } - } - Err(err) => { - // legacy log print - let json_str = String::from_utf8_lossy(&res_bytes); - tracing::warn!(%err, %json_str, "Failed deserializing JSON to response"); + Ok(decode_response_stream( + response_stream.rx, + engine_ctx, + queue_start, + tx_start, + inflight_guard, + )) + } - Some(U::from_err(DynamoError::msg(err.to_string()))) - } - } - } else if is_complete_final { - // end of stream - None - } else if engine_ctx_.is_stopped() { - // Gracefully end the stream if 'stop_generating()' was called. Do NOT check for - // 'is_killed()' here because it implies the stream ended abnormally which should be - // handled by the error branch below. - tracing::debug!("Request cancelled and then trying to read a response"); - None - } else { - // stream ended unexpectedly - let err = DynamoError::builder() - .error_type(ErrorType::Disconnected) - .message("Stream ended before generation completed") - .build(); - tracing::debug!("{err}"); - Some(U::from_err(err)) - } - }); + /// Register the requested halves of a data-plane stream with the response + /// transport. Returns `(send_stream, recv_stream)` mirroring the + /// `PendingConnections::into_parts` shape — either side is `None` when not + /// requested. Asserts post-registration that the transport produced + /// exactly the requested shape; a mismatch is a transport-layer bug, not + /// a runtime error path. + async fn register_streams( + &self, + engine_ctx: Arc, + enable_request_stream: bool, + enable_response_stream: bool, + ) -> Result< + ( + Option>, + Option>, + ), + Error, + > { + let options = StreamOptions::builder() + .context(engine_ctx) + .enable_request_stream(enable_request_stream) + .enable_response_stream(enable_response_stream) + .build()?; + + let pending: PendingConnections = self.resp_transport.register(options).await; + let (send_stream, recv_stream) = pending.into_parts(); + + // Transport-layer invariant: the data plane produces exactly the halves + // we requested. A mismatch is a bug in the transport, not a runtime + // error path, so assert only in debug builds rather than panicking prod. + debug_assert_eq!( + send_stream.is_some(), + enable_request_stream, + "data-plane registration: request-stream presence does not match request" + ); + debug_assert_eq!( + recv_stream.is_some(), + enable_response_stream, + "data-plane registration: response-stream presence does not match request" + ); - inflight_guard.disarm(); - let stream = InflightDecStream { inner: stream }; - Ok(ResponseStream::new(Box::pin(stream), engine_ctx)) + Ok((send_stream, recv_stream)) + } + + /// Build standard request-plane headers (trace propagation, request-id, + /// frontend send-timestamp) and write the encoded buffer through the + /// request-plane client. + async fn dispatch_buffer( + &self, + address: String, + buffer: bytes::Bytes, + request_id: &str, + ) -> Result<(), Error> { + let mut headers = std::collections::HashMap::new(); + inject_trace_headers_into_map(&mut headers); + headers.insert("request-id".to_string(), request_id.to_string()); + let send_ts_ns = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap_or_default() + .as_nanos() as u64; + headers.insert("x-frontend-send-ts-ns".to_string(), send_ts_ns.to_string()); + + let _nvtx_send = dynamo_nvtx_range!("transport.tcp.send"); + self.req_client + .send_request(address, buffer, headers) + .await?; + drop(_nvtx_send); + Ok(()) + } +} + +#[async_trait::async_trait] +impl AsyncEngine>, ManyOut, Error> for AddressedPushRouter +where + T: Data + Serialize, + U: Data + for<'de> Deserialize<'de> + MaybeError, +{ + async fn generate(&self, request: SingleIn>) -> Result, Error> { + let (addressed_request, context) = request.transfer(()); + let (request, address, instance_info) = addressed_request.into_parts(); + + self.dispatch_and_finalize::( + &context, + address, + instance_info.as_ref(), + Some(&request), + None, + ) + .await } } @@ -484,6 +693,7 @@ mod tests { }, metadata, frontend_send_ts_ns: None, + request_stream_connection_info: None, } } diff --git a/lib/runtime/src/pipeline/network/egress/push_router.rs b/lib/runtime/src/pipeline/network/egress/push_router.rs index cbe2e8ee5723..44b37853aaf3 100644 --- a/lib/runtime/src/pipeline/network/egress/push_router.rs +++ b/lib/runtime/src/pipeline/network/egress/push_router.rs @@ -768,7 +768,7 @@ where async fn generate_with_fault_detection( &self, - mut instance_id: u64, + instance_id: u64, request: SingleIn, ) -> anyhow::Result> { let route_start = Instant::now(); @@ -784,113 +784,10 @@ where ) }; - // Check if the selected worker is overloaded (when fault detection is enabled). - if self.fault_detection_enabled { - let routing_instances = self.client.routing_instances(); - let selected_worker_overloaded = routing_instances.is_overloaded(instance_id); - let counts = routing_instances.counts(); - if tracing::enabled!(tracing::Level::DEBUG) { - tracing::debug!( - request_id = %request_id, - instance_id, - router_mode = ?self.router_mode, - free_workers = counts.free, - overloaded_workers = counts.overloaded, - total_workers = counts.discovered, - selected_worker_overloaded, - "checked worker overload state" - ); - } - if selected_worker_overloaded { - tracing::warn!( - instance_id, - overloaded_workers = counts.overloaded, - total_workers = counts.discovered, - "Rejecting request: selected worker is overloaded" - ); - let cause = PipelineError::ServiceOverloaded( - "Selected worker is overloaded, please retry later".to_string(), - ); - return Err(DynamoError::builder() - .error_type(ErrorType::ResourceExhausted) - .message("Selected worker is overloaded, please retry later") - .cause(cause) - .build() - .into()); - } - } - - // Resolve transport address; if the selected instance disappeared - // between selection and dispatch, fall back to another available one. - let (address, _transport_kind, instance) = { - use crate::component::TransportType; - - let resolve_transport = |id: u64| { - let instances = self.client.instances(); - instances - .iter() - .find(|i| i.instance_id == id) - .map(|instance| { - let (addr, kind) = match &instance.transport { - TransportType::Tcp(tcp_endpoint) => { - tracing::debug!( - instance_id = id, - tcp_endpoint = %tcp_endpoint, - "Using TCP transport for instance" - ); - (tcp_endpoint.clone(), "transport.tcp.request") - } - TransportType::Nats(subject) => { - tracing::debug!( - instance_id = id, - subject = %subject, - "Using NATS transport for instance" - ); - (subject.clone(), "transport.nats.request") - } - }; - (addr, kind, instance.clone()) - }) - }; + self.check_workers_available(instance_id, &request_id)?; - if let Some(result) = resolve_transport(instance_id) { - result - } else { - // Instance vanished — pick another from free_ids (same filter - // as pre-selection) and retry the lookup once. - let routing_instances = self.client.routing_instances(); - let fallback_id = routing_instances - .free_ids() - .iter() - .copied() - .find(|&id| id != instance_id); - match fallback_id { - Some(id) => { - tracing::warn!( - original_instance = instance_id, - fallback_instance = id, - "Instance disappeared during routing, reselecting" - ); - instance_id = id; - resolve_transport(id).ok_or_else(|| { - anyhow::anyhow!( - "Fallback instance {} also not found for endpoint {}", - id, - self.client.endpoint.id() - ) - })? - } - None => { - return Err(anyhow::anyhow!( - "Instance {} not found and no other instances available \ - for endpoint {}", - instance_id, - self.client.endpoint.id() - )); - } - } - } - }; + let (instance_id, address, transport_kind, instance) = + self.resolve_transport(instance_id)?; let request = request.map(|req| AddressedRequest::with_instance(req, address, instance)); @@ -898,82 +795,194 @@ where .with_label_values(&[STAGE_ROUTE]) .observe(route_start.elapsed().as_secs_f64()); - let _nvtx_transport = dynamo_nvtx_range!(_transport_kind); + let _nvtx_transport = dynamo_nvtx_range!(transport_kind); let stream: anyhow::Result> = self .addressed .generate(request) .instrument(route_span) .await; - match stream { - Ok(stream) => { - if !self.fault_detection_enabled { - return Ok(stream); - } - let engine_ctx = stream.context(); - let client = self.client.clone(); - let client_for_timeout = self.client.clone(); - let stream = stream.map(move |res| { - // Check if the error is migratable (indicates worker/connection failure) - if let Some(err) = res.err() - && is_inhibited(&err) - { - tracing::debug!( - "Reporting instance {instance_id} down due to migratable error: {err}" - ); - client.report_instance_down(instance_id); - } - res - }); - - // Request-plane inactivity timeout: emit a ResponseTimeout error item - // when the backend stops producing output. This triggers is_inhibited() - // → report_instance_down() to quarantine the worker. - let stream: Pin + Send>> = if let Some(timeout) = - self.response_timeout - { - Box::pin(async_stream::stream! { - let mut inner = Box::pin(stream); - loop { - tokio::select! { - biased; - item = inner.next() => { - match item { - Some(item) => yield item, - None => break, - } - } - _ = tokio::time::sleep(timeout) => { - tracing::warn!( - instance_id, - timeout_secs = timeout.as_secs(), - "backend response inactivity timeout — quarantining worker" - ); - client_for_timeout.report_instance_down(instance_id); - yield U::from_err( - crate::error::DynamoError::builder() - .error_type(crate::error::ErrorType::ResponseTimeout) - .message("backend response inactivity timeout") - .build() - ); - break; - } - } + self.wrap_with_fault_detection(stream, instance_id) + } + + /// Reject early if the selected worker is overloaded and fault detection + /// is enabled. The request_id is only used for the debug-level "checked + /// worker overload state" trace; pass an empty string from callers that + /// don't have one handy. + fn check_workers_available(&self, instance_id: u64, request_id: &str) -> anyhow::Result<()> { + if !self.fault_detection_enabled { + return Ok(()); + } + let routing_instances = self.client.routing_instances(); + let selected_worker_overloaded = routing_instances.is_overloaded(instance_id); + let counts = routing_instances.counts(); + if tracing::enabled!(tracing::Level::DEBUG) { + tracing::debug!( + request_id, + instance_id, + router_mode = ?self.router_mode, + free_workers = counts.free, + overloaded_workers = counts.overloaded, + total_workers = counts.discovered, + selected_worker_overloaded, + "checked worker overload state" + ); + } + if !selected_worker_overloaded { + return Ok(()); + } + tracing::warn!( + instance_id, + overloaded_workers = counts.overloaded, + total_workers = counts.discovered, + "Rejecting request: selected worker is overloaded" + ); + let cause = PipelineError::ServiceOverloaded( + "Selected worker is overloaded, please retry later".into(), + ); + Err(DynamoError::builder() + .error_type(ErrorType::ResourceExhausted) + .message("Selected worker is overloaded, please retry later") + .cause(cause) + .build() + .into()) + } + + /// Resolve `(instance_id, address, transport_kind_label, Instance)` for + /// the selected worker. If the instance has disappeared between selection + /// and dispatch, fall back to one other instance from `free_ids` (same + /// filter as pre-selection) and return the updated id so the caller can + /// `report_instance_down` the right worker on later failures. + fn resolve_transport( + &self, + instance_id: u64, + ) -> anyhow::Result<(u64, String, &'static str, Instance)> { + use crate::component::TransportType; + + let lookup = |id: u64| { + self.client + .instances() + .iter() + .find(|i| i.instance_id == id) + .map(|instance| { + let (addr, kind) = match &instance.transport { + TransportType::Tcp(tcp_endpoint) => { + (tcp_endpoint.clone(), "transport.tcp.request") } - }) - } else { - Box::pin(stream) - }; + TransportType::Nats(subject) => (subject.clone(), "transport.nats.request"), + }; + (addr, kind, instance.clone()) + }) + }; + + if let Some((addr, kind, inst)) = lookup(instance_id) { + return Ok((instance_id, addr, kind, inst)); + } - Ok(ResponseStream::new(stream, engine_ctx)) + let routing_instances = self.client.routing_instances(); + let fallback_id = routing_instances + .free_ids() + .iter() + .copied() + .find(|&id| id != instance_id); + match fallback_id { + Some(id) => { + tracing::warn!( + original_instance = instance_id, + fallback_instance = id, + "Instance disappeared during routing, reselecting" + ); + let (addr, kind, inst) = lookup(id).ok_or_else(|| { + anyhow::anyhow!( + "Fallback instance {} also not found for endpoint {}", + id, + self.client.endpoint.id() + ) + })?; + Ok((id, addr, kind, inst)) } + None => Err(anyhow::anyhow!( + "Instance {} not found and no other instances available for endpoint {}", + instance_id, + self.client.endpoint.id() + )), + } + } + + /// Wrap a dispatched stream with fault detection + inactivity timeout. + /// `is_inhibited` errors trigger `report_instance_down`; the timeout + /// (driven by `DYN_HTTP_BACKEND_STREAM_TIMEOUT_SECS`) yields a synthetic + /// `ResponseTimeout` and quarantines the worker. + fn wrap_with_fault_detection( + &self, + stream: anyhow::Result>, + instance_id: u64, + ) -> anyhow::Result> { + let stream = match stream { + Ok(stream) => stream, Err(err) => { if self.fault_detection_enabled && is_inhibited(err.as_ref()) { tracing::debug!("Reporting instance {instance_id} down due to error: {err}"); self.client.report_instance_down(instance_id); } - Err(err) + return Err(err); } + }; + + if !self.fault_detection_enabled { + return Ok(stream); } + + let engine_ctx = stream.context(); + let client = self.client.clone(); + let client_for_timeout = self.client.clone(); + let stream = stream.map(move |res| { + if let Some(err) = res.err() + && is_inhibited(&err) + { + tracing::debug!( + "Reporting instance {instance_id} down due to migratable error: {err}" + ); + client.report_instance_down(instance_id); + } + res + }); + + let stream: Pin + Send>> = + if let Some(timeout) = self.response_timeout { + Box::pin(async_stream::stream! { + let mut inner = Box::pin(stream); + loop { + tokio::select! { + biased; + item = inner.next() => { + match item { + Some(item) => yield item, + None => break, + } + } + _ = tokio::time::sleep(timeout) => { + tracing::warn!( + instance_id, + timeout_secs = timeout.as_secs(), + "backend response inactivity timeout — quarantining worker" + ); + client_for_timeout.report_instance_down(instance_id); + yield U::from_err( + crate::error::DynamoError::builder() + .error_type(crate::error::ErrorType::ResponseTimeout) + .message("backend response inactivity timeout") + .build() + ); + break; + } + } + } + }) + } else { + Box::pin(stream) + }; + + Ok(ResponseStream::new(stream, engine_ctx)) } } @@ -1002,23 +1011,64 @@ where } } +impl PushRouter +where + T: Data + Serialize, + U: Data + for<'de> Deserialize<'de> + MaybeError, +{ + /// Bidirectional sibling of [`Self::generate_with_fault_detection`]. + async fn bidirectional_dispatch( + &self, + instance_id: u64, + input: ManyIn, + ) -> anyhow::Result> { + let route_start = Instant::now(); + let request_id = input.context().id().to_string(); + let route_span = tracing::info_span!( + "router.route_request_bidirectional", + request_id = %request_id, + worker_id = instance_id, + router_mode = ?self.router_mode, + ); + + self.check_workers_available(instance_id, &request_id)?; + let (instance_id, address, transport_kind, instance) = + self.resolve_transport(instance_id)?; + + STAGE_DURATION_SECONDS + .with_label_values(&[STAGE_ROUTE]) + .observe(route_start.elapsed().as_secs_f64()); + + let _nvtx_transport = dynamo_nvtx_range!(transport_kind); + let stream: anyhow::Result> = self + .addressed + .generate_bidirectional(instance, address, input) + .instrument(route_span) + .await; + self.wrap_with_fault_detection(stream, instance_id) + } +} + /// Bidirectional `AsyncEngine` impl for streaming-input workloads (e.g. the -/// OpenAI Realtime API). Selects a sticky instance on the first inbound frame -/// and binds the whole input stream to that worker. Required so engines of -/// shape `BidirectionalStreamingEngine` can be stored as a `PushRouter` -/// in `WorkerSet`. -/// -/// Remote per-frame dispatch over `AddressedPushRouter` / `PushWorkHandler` -/// is not yet implemented; this impl currently bails after selecting the +/// OpenAI Realtime API). Reserves a sticky worker up front — before any +/// inbound frame is observed — and binds the whole input stream to that /// worker. KV and Direct modes inherit the same `bail!` invariants as the /// unary impl. +/// +/// **Reserve-before-observe rationale.** The router-mode strategies +/// (`RoundRobin`, `Random`, `PowerOfTwoChoices`, `LeastLoaded`, +/// `DeviceAwareWeighted`) don't depend on frame contents, so selection +/// runs immediately and connection setup proceeds in parallel with the +/// client producing its first frame. A client that connects but never +/// sends one still releases the slot via the response-stream-drop path; +/// the dispatch-side `cancel_both` cleanup covers the early-bail case. #[async_trait] impl AsyncEngine, ManyOut, Error> for PushRouter where T: Data + Serialize, U: Data + for<'de> Deserialize<'de> + MaybeError, { - async fn generate(&self, mut input: ManyIn) -> Result, Error> { + async fn generate(&self, input: ManyIn) -> Result, Error> { match self.router_mode { RouterMode::KV => { anyhow::bail!("KV routing should not call generate on PushRouter"); @@ -1028,24 +1078,28 @@ where "Direct routing should not call generate on PushRouter directly; use DirectRoutingRouter wrapper" ); } - _ => {} + // These modes drive `select_next_worker()` to `None` — they rely on + // the occupancy/load-aware selection the bidirectional path does not + // wire yet, which would otherwise surface as a misleading "no + // instances available" error below. Reject them explicitly until + // bidirectional support lands; tracked in + // https://github.com/ai-dynamo/dynamo/issues/10320. + RouterMode::PowerOfTwoChoices + | RouterMode::LeastLoaded + | RouterMode::DeviceAwareWeighted => { + anyhow::bail!( + "{:?} routing is not yet supported for bidirectional dispatch", + self.router_mode + ); + } + RouterMode::RoundRobin | RouterMode::Random => {} } - // Wait for the first frame so the sticky-instance pick reflects the - // session's actual start, not router construction time. - if input.next().await.is_none() { - anyhow::bail!("bidirectional input stream closed before first frame"); - } let instance_id = self .select_next_worker() .ok_or_else(|| anyhow::anyhow!("no instances available for bidirectional routing"))?; - // Per-frame remote dispatch over AddressedPushRouter / PushWorkHandler - // is tracked in #9361. Until that lands, callers must register engines - // in-process via ModelManager rather than rely on discovered workers. - anyhow::bail!( - "bidirectional remote dispatch is not yet implemented (selected instance {instance_id})" - ) + self.bidirectional_dispatch(instance_id, input).await } } @@ -1095,7 +1149,10 @@ mod tests { DistributedRuntime, Runtime, distributed::DistributedConfig, error::DynamoError, - pipeline::{ResponseStream, context::Controller}, + pipeline::{ + RequestStream, ResponseStream, + context::{Context, Controller}, + }, }; use serde::{Deserialize, Serialize}; @@ -1263,9 +1320,10 @@ mod tests { .await .unwrap(); - let ctx: Arc = Arc::new(Controller::default()); let input: ManyIn = - ResponseStream::new(Box::pin(tokio_stream::iter(vec![1u64, 2u64])), ctx); + Context::new(RequestStream::new(Box::pin(tokio_stream::iter(vec![ + 1u64, 2u64, + ])))); let result = router.generate(input).await; assert!( result.is_err(), @@ -1290,8 +1348,8 @@ mod tests { .await .unwrap(); - let ctx: Arc = Arc::new(Controller::default()); - let input: ManyIn = ResponseStream::new(Box::pin(tokio_stream::iter(vec![1u64])), ctx); + let input: ManyIn = + Context::new(RequestStream::new(Box::pin(tokio_stream::iter(vec![1u64])))); let result = router.generate(input).await; assert!( result.is_err(), @@ -1321,8 +1379,8 @@ mod tests { .await .unwrap(); - let ctx: Arc = Arc::new(Controller::default()); - let input: ManyIn = ResponseStream::new(Box::pin(tokio_stream::iter(vec![1u64])), ctx); + let input: ManyIn = + Context::new(RequestStream::new(Box::pin(tokio_stream::iter(vec![1u64])))); let result = router.generate(input).await; assert!( result.is_err(), @@ -1337,6 +1395,43 @@ mod tests { rt.shutdown(); } + #[tokio::test] + async fn bidirectional_generate_rejects_unsupported_load_aware_modes() { + let rt = Runtime::from_current().unwrap(); + let drt = DistributedRuntime::new(rt.clone(), DistributedConfig::process_local()) + .await + .unwrap(); + let ns = drt.namespace("test_bidi_load_aware".to_string()).unwrap(); + let component = ns.component("test_component".to_string()).unwrap(); + + for mode in [ + RouterMode::PowerOfTwoChoices, + RouterMode::LeastLoaded, + RouterMode::DeviceAwareWeighted, + ] { + let endpoint = component.endpoint("test_endpoint".to_string()); + let client = endpoint.client().await.unwrap(); + let router = PushRouter::::from_client(client, mode) + .await + .unwrap(); + + let input: ManyIn = + Context::new(RequestStream::new(Box::pin(tokio_stream::iter(vec![1u64])))); + let result = router.generate(input).await; + assert!( + result.is_err(), + "bidirectional generate must reject {mode:?} (not yet supported)" + ); + let err_msg = format!("{:?}", result.unwrap_err()); + assert!( + err_msg.contains("not yet supported for bidirectional dispatch"), + "error should explain the mode is unsupported, not 'no instances': got {err_msg}" + ); + } + + rt.shutdown(); + } + #[tokio::test] async fn least_loaded_select_and_peek_return_none_with_available_worker() { let rt = Runtime::from_current().unwrap(); diff --git a/lib/runtime/src/pipeline/network/ingress/push_handler.rs b/lib/runtime/src/pipeline/network/ingress/push_handler.rs index 9e93ff2c26b0..77703120f919 100644 --- a/lib/runtime/src/pipeline/network/ingress/push_handler.rs +++ b/lib/runtime/src/pipeline/network/ingress/push_handler.rs @@ -3,11 +3,14 @@ use super::*; +use crate::engine::AsyncEngineContext; use crate::metrics::prometheus_names::work_handler; 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 std::sync::Arc; @@ -232,6 +235,56 @@ impl Ingress { } } } + + /// Decode the wire envelope into its [`RequestControlMessage`] and the + /// optional data payload, shared by every [`IngressDispatch`] shape: + /// - `HeaderAndData` → `(control, Some(data))` — the unary wire shape, + /// where the request body travels in the data half. + /// - `HeaderOnly` → `(control, None)` — the bidirectional wire shape, + /// where request frames flow on the request-stream socket instead. + /// + /// The caller decides whether its path expects the data payload. The + /// deserialization and invalid-message error counters are incremented + /// here so every shape reports them consistently. + fn decode_control_message( + &self, + payload: Bytes, + ) -> Result<(RequestControlMessage, Option), PipelineError> { + let msg = TwoPartCodec::default() + .decode_message(payload)? + .into_message_type(); + + let (header, data) = match msg { + TwoPartMessageType::HeaderAndData(header, data) => (header, Some(data)), + TwoPartMessageType::HeaderOnly(header) => (header, None), + _ => { + if let Some(m) = self.metrics() { + m.error_counter + .with_label_values(&[work_handler::error_types::INVALID_MESSAGE]) + .inc(); + } + return Err(PipelineError::Generic(String::from( + "Unexpected message from work queue; expected a header-only or header-and-data TwoPartMessage", + ))); + } + }; + + let control_msg: RequestControlMessage = + serde_json::from_slice(&header).map_err(|err| { + if let Some(m) = self.metrics() { + m.error_counter + .with_label_values(&[work_handler::error_types::DESERIALIZATION]) + .inc(); + } + let json_str = String::from_utf8_lossy(&header); + PipelineError::DeserializationError(format!( + "Failed deserializing to RequestControlMessage. err={err}, json_str={json_str}, header_len={}", + header.len(), + )) + })?; + + Ok((control_msg, data)) + } } /// The output of [`IngressDispatch::parse_and_build_request`]: the typed /// request the engine consumes, plus the bits of the on-wire control @@ -244,14 +297,11 @@ struct ParsedRequest { } /// Per-shape strategy for turning a raw payload into a typed engine -/// request. Captures the wire-shape-specific parsing of the request -/// envelope; everything else — metrics-guard, response stream open, -/// `segment.generate`, prologue, pump — lives in +/// request. Captures the wire-shape divergence between the unary +/// (`HeaderAndData`) and bidirectional (`HeaderOnly` + dial-in for the +/// request stream) paths; everything else — metrics-guard, response stream +/// open, `segment.generate`, prologue, pump — lives in /// [`Ingress::handle_payload_shared`] below. -/// -/// Currently has a single impl (the unary `HeaderAndData` shape). The -/// abstraction exists to keep an additional impl for the bidirectional -/// `HeaderOnly` + dial-in shape addition cheap when that lands. #[async_trait] trait IngressDispatch: Send + Sync { type Request: PipelineIO; @@ -274,50 +324,22 @@ where &self, payload: Bytes, ) -> Result>, PipelineError> { - // decode the control message and the request - let msg = TwoPartCodec::default() - .decode_message(payload)? - .into_message_type(); + let (control_msg, data) = self.decode_control_message(payload)?; - // we must have a header and a body - let (control_msg, request_t) = match msg { - TwoPartMessageType::HeaderAndData(header, data) => { - tracing::trace!( - "received two part message with ctrl: {} bytes, data: {} bytes", - header.len(), - data.len() - ); - let control_msg: RequestControlMessage = match serde_json::from_slice(&header) { - Ok(cm) => cm, - Err(err) => { - let json_str = String::from_utf8_lossy(&header); - if let Some(m) = self.metrics() { - m.error_counter - .with_label_values(&[work_handler::error_types::DESERIALIZATION]) - .inc(); - } - return Err(PipelineError::DeserializationError(format!( - "Failed deserializing to RequestControlMessage. err={err}, json_str={json_str}, header_len={}", - header.len(), - ))); - } - }; - let request_t: T = serde_json::from_slice(&data)?; - (control_msg, request_t) - } - _ => { - if let Some(m) = self.metrics() { - m.error_counter - .with_label_values(&[work_handler::error_types::INVALID_MESSAGE]) - .inc(); - } - return Err(PipelineError::Generic(String::from( - "Unexpected message from work queue; unable extract a TwoPartMessage with a header and data", - ))); + // The unary path carries the request body in the data half; a + // header-only envelope means the sender used the bidirectional shape. + let data = data.ok_or_else(|| { + if let Some(m) = self.metrics() { + m.error_counter + .with_label_values(&[work_handler::error_types::INVALID_MESSAGE]) + .inc(); } - }; + PipelineError::Generic(String::from( + "unary engine received a header-only envelope; expected a request payload", + )) + })?; + let request_t: T = serde_json::from_slice(&data)?; - // extend request with context tracing::trace!( request_id = %control_msg.id, metadata_entries = control_msg.metadata.len(), @@ -336,6 +358,131 @@ where } } +#[async_trait] +impl IngressDispatch for Ingress, ManyOut> +where + T: Data + for<'de> Deserialize<'de> + std::fmt::Debug, + U: Data + Serialize + MaybeError + std::fmt::Debug, +{ + type Request = ManyIn; + + async fn parse_and_build_request( + &self, + payload: Bytes, + ) -> Result>, PipelineError> { + let (control_msg, data) = self.decode_control_message(payload)?; + + // Bidirectional envelopes are header-only — all request frames + // (including the first) flow on the request-stream socket once it's + // dialed in. A data payload means the sender used the unary wire + // shape; reject it. + if data.is_some() { + if let Some(m) = self.metrics() { + m.error_counter + .with_label_values(&[work_handler::error_types::INVALID_MESSAGE]) + .inc(); + } + return Err(PipelineError::Generic(String::from( + "bidirectional engine received a non-header-only envelope", + ))); + } + + if !matches!(control_msg.request_type, RequestType::ManyIn) { + if let Some(m) = self.metrics() { + m.error_counter + .with_label_values(&[work_handler::error_types::INVALID_MESSAGE]) + .inc(); + } + return Err(PipelineError::Generic(String::from( + "bidirectional engine received a non-ManyIn request envelope", + ))); + } + + let req_stream_conn_info = control_msg + .request_stream_connection_info + .clone() + .ok_or_else(|| { + PipelineError::Generic(String::from( + "bidirectional control message missing request_stream_connection_info", + )) + })?; + + let request_context: context::Context<()> = context::Context::with_id_and_metadata( + (), + control_msg.id.clone(), + control_msg.metadata.clone(), + ); + let context_arc: Arc = request_context.context(); + + // Open the request stream (upstream → worker) up front. The shared + // handler opens the response stream uniformly after we return. If + // response-stream open subsequently fails, the forwarder task + // spawned below exits cleanly when `frame_tx.send` observes the + // dropped `frame_rx`. + let request_stream_recv = tcp::client::TcpClient::create_request_stream( + context_arc.clone(), + req_stream_conn_info, + None, + ) + .await + .map_err(|e| { + if let Some(m) = self.metrics() { + m.error_counter + .with_label_values(&[work_handler::error_types::RESPONSE_STREAM]) + .inc(); + } + PipelineError::Generic(format!("Failed to create request stream: {e}")) + })?; + + // Forwarder: deserialize raw bytes off the request socket into `T` + // and feed the engine's `ManyIn` input. Every request frame + // (including the first) flows over this socket — the envelope is + // header-only. + let (frame_tx, frame_rx) = tokio::sync::mpsc::channel::(8); + let forwarder_ctx = context_arc.clone(); + tokio::spawn(async move { + let mut rx = request_stream_recv.rx; + while let Some(bytes) = rx.recv().await { + // Stop forwarding on either kill or soft-stop, matching the + // send-side `spawn_request_stream_forwarder`. Without the + // `stopped()` check, a `stop_generating()` would leave this + // task pumping frames into a channel the engine has abandoned. + if forwarder_ctx.is_killed() || forwarder_ctx.is_stopped() { + break; + } + match serde_json::from_slice::(&bytes) { + Ok(item) => { + if frame_tx.send(item).await.is_err() { + tracing::debug!( + "engine consumer dropped; bidirectional input forwarder exiting" + ); + break; + } + } + Err(e) => { + tracing::error!( + error = %e, + "failed to deserialize bidirectional request frame; killing context" + ); + forwarder_ctx.kill(); + break; + } + } + } + }); + + let input_stream: crate::engine::DataStream = + Box::pin(tokio_stream::wrappers::ReceiverStream::new(frame_rx)); + let request: ManyIn = request_context.map(|_| RequestStream::new(input_stream)); + + Ok(ParsedRequest { + request, + response_connection_info: control_msg.connection_info, + frontend_send_ts_ns: control_msg.frontend_send_ts_ns, + }) + } +} + impl Ingress> where U: Data + Serialize + MaybeError + std::fmt::Debug, @@ -404,7 +551,7 @@ where .with_label_values(&[work_handler::error_types::RESPONSE_STREAM]) .inc(); } - PipelineError::Generic(format!("Failed to create response stream: {:?}", e,)) + PipelineError::Generic(format!("Failed to create response stream: {e}")) })?; tracing::trace!("calling generate"); @@ -474,13 +621,42 @@ where endpoint: &crate::component::Endpoint, metrics_labels: Option<&[(&str, &str)]>, ) -> Result<()> { - // Call the Ingress-specific add_metrics implementation - use crate::pipeline::network::Ingress; + // Call the inherent `Ingress::add_metrics`, not this trait method. + Ingress::add_metrics(self, endpoint, metrics_labels) + } + + fn set_endpoint_health_check_notifier(&self, notifier: Arc) -> Result<()> { + self.endpoint_health_check_notifier + .set(notifier) + .map_err(|_| anyhow::anyhow!("Endpoint health check notifier already set"))?; + Ok(()) + } + + async fn handle_payload( + &self, + payload: Bytes, + request_id: Option, + ) -> Result<(), PipelineError> { + self.handle_payload_shared(payload, request_id).await + } +} + +#[async_trait] +impl PushWorkHandler for Ingress, ManyOut> +where + T: Data + for<'de> Deserialize<'de> + std::fmt::Debug, + U: Data + Serialize + MaybeError + std::fmt::Debug, +{ + fn add_metrics( + &self, + endpoint: &crate::component::Endpoint, + metrics_labels: Option<&[(&str, &str)]>, + ) -> Result<()> { + // Call the inherent `Ingress::add_metrics`, not this trait method. Ingress::add_metrics(self, endpoint, metrics_labels) } fn set_endpoint_health_check_notifier(&self, notifier: Arc) -> Result<()> { - use crate::pipeline::network::Ingress; self.endpoint_health_check_notifier .set(notifier) .map_err(|_| anyhow::anyhow!("Endpoint health check notifier already set"))?; diff --git a/lib/runtime/tests/bidirectional_e2e.rs b/lib/runtime/tests/bidirectional_e2e.rs new file mode 100644 index 000000000000..38fe2a986ac1 --- /dev/null +++ b/lib/runtime/tests/bidirectional_e2e.rs @@ -0,0 +1,112 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! End-to-end bidirectional round trip over the full Rust stack: a real +//! `Ingress, ManyOut>` worker, a `PushRouter` +//! client, and a three-frame `ManyIn` producing three response frames. +//! +//! Lives in its own `tests/` binary so the process-global TCP server's +//! accept loop runs for the whole test process. + +use std::sync::Arc; + +use anyhow::Error; +use async_trait::async_trait; +use serde::{Deserialize, Serialize}; + +use dynamo_runtime::{ + DistributedRuntime, Runtime, + distributed::DistributedConfig, + engine::{AsyncEngine, AsyncEngineContextProvider, DataStream}, + error::DynamoError, + pipeline::{ + ManyIn, ManyOut, RequestStream, ResponseStream, context::Context, network::Ingress, + }, + protocols::maybe_error::MaybeError, +}; + +use dynamo_runtime::pipeline::network::egress::push_router::{PushRouter, RouterMode}; + +#[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, |v| EchoResponse { + value: Some(v), + error: None, + }); + let stream: DataStream = Box::pin(mapped); + Ok(ResponseStream::new(stream, ctx)) + } +} + +#[tokio::test] +async fn bidirectional_end_to_end_echo() { + let rt = Runtime::from_current().unwrap(); + let drt = DistributedRuntime::new(rt.clone(), DistributedConfig::process_local()) + .await + .unwrap(); + let ns = drt.namespace("test_bidi_e2e".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![ + 1u64, 2, 3, + ])))); + let response_stream = router.generate(input).await.unwrap(); + let responses: Vec = futures::StreamExt::collect(response_stream).await; + + let values: Vec = responses.iter().filter_map(|r| r.value).collect(); + assert_eq!( + values, + vec![1u64, 2, 3], + "echo engine should reflect each input frame back; got {responses:?}" + ); + + rt.shutdown(); +}