big change: dropping Sync requirement from Resp in AsyncEngine<Req, Resp, Error>

Author: ryanolson

lib/llm/src/http/client.rs

@@ -27,7 +27,7 @@ use crate::protocols::openai::chat_completions::{
 };
 use crate::protocols::Annotated;
 use dynamo_runtime::engine::{
-    AsyncEngineContext, AsyncEngineContextProvider, AsyncEngineStream, Data,
+    AsyncEngineContext, AsyncEngineContextProvider, AsyncEngineStream, Data, DataStream,
 };
 
 /// Configuration for HTTP clients
@@ -226,43 +226,29 @@ impl BaseHttpClient {
 }
 
 /// Type alias for NV chat response stream
-pub type NvChatResponseStream = Pin<
-    Box<
-        dyn Stream<Item = Result<Annotated<NvCreateChatCompletionStreamResponse>, OpenAIError>>
-            + Send
-            + Sync,
-    >,
->;
+pub type NvChatResponseStream =
+    DataStream<Result<Annotated<NvCreateChatCompletionStreamResponse>, OpenAIError>>;
 
 /// Type alias for generic BYOT response stream
-pub type ByotResponseStream = Pin<Box<dyn Stream<Item = Result<Value, OpenAIError>> + Send + Sync>>;
+pub type ByotResponseStream = DataStream<Result<Value, OpenAIError>>;
 
 /// Type alias for pure OpenAI chat response stream
-pub type OpenAIChatResponseStream = Pin<
-    Box<
-        dyn Stream<
-                Item = Result<async_openai::types::CreateChatCompletionStreamResponse, OpenAIError>,
-            > + Send
-            + Sync,
-    >,
->;
+pub type OpenAIChatResponseStream =
+    DataStream<Result<async_openai::types::CreateChatCompletionStreamResponse, OpenAIError>>;
 
 /// A wrapped HTTP response stream that combines a stream with its context
 /// This provides a unified interface for HTTP client responses
 #[derive(Dissolve)]
 pub struct HttpResponseStream<T> {
     /// The underlying stream of responses
-    pub stream: Pin<Box<dyn Stream<Item = T> + Send>>,
+    pub stream: DataStream<T>,
     /// The context for this request
     pub context: Arc<dyn AsyncEngineContext>,
 }
 
 impl<T> HttpResponseStream<T> {
     /// Create a new HttpResponseStream
-    pub fn new(
-        stream: Pin<Box<dyn Stream<Item = T> + Send>>,
-        context: Arc<dyn AsyncEngineContext>,
-    ) -> Self {
+    pub fn new(stream: DataStream<T>, context: Arc<dyn AsyncEngineContext>) -> Self {
         Self { stream, context }
     }
 }
@@ -299,7 +285,7 @@ impl<T: Data> HttpResponseStream<T> {
 
 /// A wrapper that implements AsyncEngineStream for streams that are Send + Sync
 struct AsyncEngineStreamWrapper<T> {
-    stream: Pin<Box<dyn Stream<Item = T> + Send>>,
+    stream: DataStream<T>,
     context: Arc<dyn AsyncEngineContext>,
 }
 
@@ -317,9 +303,9 @@ impl<T: Data> AsyncEngineContextProvider for AsyncEngineStreamWrapper<T> {
     }
 }
 
-// This is unsafe because we're claiming the stream is Sync when it might not be
-// But this is needed for the AsyncEngineStream trait
-unsafe impl<T> Sync for AsyncEngineStreamWrapper<T> {}
+// // This is unsafe because we're claiming the stream is Sync when it might not be
+// // But this is needed for the AsyncEngineStream trait
+// unsafe impl<T> Sync for AsyncEngineStreamWrapper<T> {}
 
 impl<T: Data> AsyncEngineStream<T> for AsyncEngineStreamWrapper<T> {}
 

lib/llm/src/protocols.rs

@@ -19,9 +19,7 @@
 //! both publicly via the HTTP API and internally between Dynamo components.
 //!
 
-use std::pin::Pin;
-
-use futures::{Stream, StreamExt};
+use futures::StreamExt;
 use serde::{Deserialize, Serialize};
 
 pub mod codec;
@@ -30,7 +28,7 @@ pub mod openai;
 
 /// The token ID type
 pub type TokenIdType = u32;
-pub type DataStream<T> = Pin<Box<dyn Stream<Item = T> + Send + Sync>>;
+pub use dynamo_runtime::engine::DataStream;
 
 // TODO: This is an awkward dependency that we need to address
 // Originally, all the Annotated/SSE Codec bits where in the LLM protocol module; however, [Annotated]

lib/llm/src/protocols/openai/chat_completions/aggregator.rs

@@ -24,7 +24,7 @@ use crate::protocols::{
 };
 
 /// A type alias for a pinned, dynamically-dispatched stream that is `Send` and `Sync`.
-type DataStream<T> = Pin<Box<dyn Stream<Item = T> + Send + Sync>>;
+type DataStream<T> = Pin<Box<dyn Stream<Item = T> + Send>>;
 
 /// Aggregates a stream of [`NvCreateChatCompletionStreamResponse`]s into a single
 /// [`NvCreateChatCompletionResponse`]. This struct accumulates incremental responses

lib/llm/src/protocols/openai/embeddings/aggregator.rs

@@ -22,9 +22,7 @@ use crate::protocols::{
     codec::{Message, SseCodecError},
     convert_sse_stream, Annotated,
 };
-
-/// A type alias for a pinned, dynamically-dispatched stream that is `Send` and `Sync`.
-type DataStream<T> = Pin<Box<dyn Stream<Item = T> + Send + Sync>>;
+use dynamo_runtime::engine::DataStream;
 
 /// Aggregates a stream of [`NvCreateEmbeddingResponse`]s into a single
 /// [`NvCreateEmbeddingResponse`]. For embeddings, this is typically simpler

lib/llm/tests/data/replays/deepseek-r1-distill-llama-8b/chat-completions.stream.logprobs.1

@@ -210,20 +210,20 @@ fn create_multi_choice_stream() -> Arc<RecordedStream<NvCreateChatCompletionStre
     Arc::new(stream)
 }
 
-fn create_stream_from_recorded_sse_stream(
-    file: &str,
-) -> Arc<RecordedStream<NvCreateChatCompletionStreamResponse>> {
-    let data = std::fs::read_to_string(file).unwrap();
-    let sse_stream = create_message_stream(&data);
-    let response_stream =
-        convert_sse_stream::<NvCreateChatCompletionStreamResponse>(Box::pin(sse_stream));
-
-    let context = Arc::new(MockContext::new());
-    let response_stream = record_stream_with_context(response_stream, context, RecordingMode::Sink);
-    let filtered_stream = response_stream.filter_map(|annotated| async move { annotated.data });
-    let (recorded_stream, recording_rx) =
-        record_stream_with_context(Box::pin(filtered_stream), ctx, RecordingMode::Sink);
-}
+// fn create_stream_from_recorded_sse_stream(
+//     file: &str,
+// ) -> Arc<RecordedStream<NvCreateChatCompletionStreamResponse>> {
+//     let data = std::fs::read_to_string(file).unwrap();
+//     let sse_stream = create_message_stream(&data);
+//     let response_stream =
+//         convert_sse_stream::<NvCreateChatCompletionStreamResponse>(Box::pin(sse_stream));
+
+//     let context = Arc::new(MockContext::new());
+//     let response_stream = record_stream_with_context(response_stream, context, RecordingMode::Sink);
+//     let filtered_stream = response_stream.filter_map(|annotated| async move { annotated.data });
+//     let (recorded_stream, recording_rx) =
+//         record_stream_with_context(Box::pin(filtered_stream), ctx, RecordingMode::Sink);
+// }
 
 fn create_stream_with_multiple_close_tokens(
 ) -> Arc<RecordedStream<NvCreateChatCompletionStreamResponse>> {

lib/llm/tests/logprob_analysis_integration.rs

@@ -92,8 +92,8 @@ impl<T: Send + Sync + 'static> Data for T {}
 
 /// [`DataStream`] is a type alias for a stream of [`Data`] items. This can be adapted to a [`ResponseStream`]
 /// by associating it with a [`AsyncEngineContext`].
-pub type DataUnary<T> = Pin<Box<dyn Future<Output = T> + Send + Sync>>;
-pub type DataStream<T> = Pin<Box<dyn Stream<Item = T> + Send + Sync>>;
+pub type DataUnary<T> = Pin<Box<dyn Future<Output = T> + Send>>;
+pub type DataStream<T> = Pin<Box<dyn Stream<Item = T> + Send>>;
 
 pub type Engine<Req, Resp, E> = Arc<dyn AsyncEngine<Req, Resp, E>>;
 pub type EngineUnary<Resp> = Pin<Box<dyn AsyncEngineUnary<Resp>>>;
@@ -174,7 +174,7 @@ pub trait AsyncEngineContextProvider: Send + Debug {
 /// This trait combines `Future` semantics with context provider capabilities,
 /// representing a single async operation that produces one result.
 pub trait AsyncEngineUnary<Resp: Data>:
-    Future<Output = Resp> + AsyncEngineContextProvider + Send + Sync
+    Future<Output = Resp> + AsyncEngineContextProvider + Send
 {
 }
 
@@ -183,7 +183,7 @@ pub trait AsyncEngineUnary<Resp: Data>:
 /// This trait combines `Stream` semantics with context provider capabilities,
 /// representing a continuous async operation that produces multiple results over time.
 pub trait AsyncEngineStream<Resp: Data>:
-    Stream<Item = Resp> + AsyncEngineContextProvider + Send + Sync
+    Stream<Item = Resp> + AsyncEngineContextProvider + Send
 {
 }
 
@@ -204,7 +204,7 @@ pub trait AsyncEngineStream<Resp: Data>:
 /// Implementations should ensure proper error handling and resource management.
 /// The `generate` method should be cancellable via the response's context provider.
 #[async_trait]
-pub trait AsyncEngine<Req: Data, Resp: Data + AsyncEngineContextProvider, E: Data>:
+pub trait AsyncEngine<Req: Send + Sync + 'static, Resp: AsyncEngineContextProvider, E: Data>:
     Send + Sync
 {
     /// Generate a stream of completion responses.

lib/runtime/src/engine.rs

@@ -69,7 +69,7 @@ pub type ServerStreamingEngine<T, U> = ServiceEngine<SingleIn<T>, ManyOut<U>>;
 /// are considered independent of each other; however, they could be constrained to be related.
 pub type BidirectionalStreamingEngine<T, U> = ServiceEngine<ManyIn<T>, ManyOut<U>>;
 
-pub trait AsyncTransportEngine<T: PipelineIO, U: PipelineIO>:
+pub trait AsyncTransportEngine<T: Data + PipelineIO, U: Data + PipelineIO>:
     AsyncEngine<T, U, Error> + Send + Sync + 'static
 {
 }
@@ -97,7 +97,7 @@ mod sealed {
     }
 }
 
-pub trait PipelineIO: Data + sealed::Connectable + AsyncEngineContextProvider {
+pub trait PipelineIO: sealed::Connectable + AsyncEngineContextProvider + 'static {
     fn id(&self) -> String;
 }
 

lib/runtime/src/pipeline.rs

@@ -280,7 +280,7 @@ pub struct Ingress<Req: PipelineIO, Resp: PipelineIO> {
     segment: OnceLock<Arc<SegmentSource<Req, Resp>>>,
 }
 
-impl<Req: PipelineIO, Resp: PipelineIO> Ingress<Req, Resp> {
+impl<Req: PipelineIO + Sync, Resp: PipelineIO> Ingress<Req, Resp> {
     pub fn new() -> Arc<Self> {
         Arc::new(Self {
             segment: OnceLock::new(),

lib/runtime/src/pipeline/network.rs

@@ -221,8 +221,8 @@ where
 impl<UpIn, UpOut, DownIn, DownOut> AsyncEngine<UpIn, UpOut, Error>
     for PipelineOperator<UpIn, UpOut, DownIn, DownOut>
 where
-    UpIn: PipelineIO,
-    DownIn: PipelineIO,
+    UpIn: PipelineIO + Sync,
+    DownIn: PipelineIO + Sync,
     DownOut: PipelineIO,
     UpOut: PipelineIO,
 {
@@ -235,8 +235,8 @@ where
 impl<UpIn, UpOut, DownIn, DownOut> Sink<UpIn>
     for PipelineOperatorForwardEdge<UpIn, UpOut, DownIn, DownOut>
 where
-    UpIn: PipelineIO,
-    DownIn: PipelineIO,
+    UpIn: PipelineIO + Sync,
+    DownIn: PipelineIO + Sync,
     DownOut: PipelineIO,
     UpOut: PipelineIO,
 {