dynamo_runtime/
pipeline.rs

1// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4/// In a Pipeline, the [`AsyncEngine`] is constrained to take a [`Context`] as input and return
5/// a [`super::engine::ResponseStream`] as output.
6use serde::{Deserialize, Serialize};
7
8mod nodes;
9pub use nodes::{
10    Operator, PipelineNode, PipelineOperator, SegmentSink, SegmentSource, Service, ServiceBackend,
11    ServiceFrontend, Sink, Source,
12};
13
14pub mod context;
15pub mod error;
16pub mod network;
17pub use crate::routing_policy::{
18    BuiltinRoutePicker, OccupancyReservation, OccupancySelection, RouteTarget,
19    RoutingOccupancyState,
20};
21pub use network::egress::addressed_router::{
22    AddressedPushRouter, AddressedRequest, StreamingDispatch, attach_first_response_guard,
23    propagate_first_response_guard,
24};
25pub use network::egress::push_router::{
26    MultimodalCacheIndex, MultimodalCacheKeyExtractor, PushRouter, RouterMode, WorkerLoadMonitor,
27};
28pub mod registry;
29
30pub use crate::engine::{
31    self as engine, AsyncEngine, AsyncEngineContext, AsyncEngineContextProvider, AsyncEngineStream,
32    Data, DataStream, Engine, EngineStream, EngineUnary, ResponseStream, async_trait,
33};
34pub use anyhow::Error;
35pub use context::Context;
36pub use error::{PipelineError, PipelineErrorExt, TwoPartCodecError};
37
38/// Pipeline inputs carry a [`Context`] which can be used to carry metadata or additional information
39/// about the request. This information propagates through the stages, both local and distributed.
40pub type SingleIn<T> = Context<T>;
41
42/// `Sync` ownership cell for a non-`Sync` [`DataStream<T>`]. [`Self::take`]
43/// moves the inner stream out; the mutex serialises concurrent attempts so
44/// the first caller observes `Some(stream)` and all later callers see
45/// `None`. Iteration happens on the returned `DataStream<T>`.
46pub struct RequestStream<T: Data> {
47    inner: std::sync::Mutex<Option<DataStream<T>>>,
48}
49
50impl<T: Data> RequestStream<T> {
51    /// Wrap a [`DataStream<T>`] in a `Sync` ownership cell.
52    pub fn new(stream: DataStream<T>) -> Self {
53        Self {
54            inner: std::sync::Mutex::new(Some(stream)),
55        }
56    }
57
58    /// Atomically move the inner stream out. Returns `Some(stream)` exactly
59    /// once across all threads racing on the same `RequestStream`; every
60    /// subsequent call (on any thread) returns `None`. The returned stream
61    /// is the unique owner; the wrapper retains nothing.
62    pub fn take(&self) -> Option<DataStream<T>> {
63        self.inner.lock().unwrap().take()
64    }
65}
66
67impl<T: Data> std::fmt::Debug for RequestStream<T> {
68    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
69        let taken = self.inner.lock().map(|g| g.is_none()).unwrap_or(true);
70        f.debug_struct("RequestStream")
71            .field("taken", &taken)
72            .finish()
73    }
74}
75
76/// Pipeline input for streaming-request engines: a [`RequestStream<T>`]
77/// payload wrapped in a [`Context`], symmetric to [`SingleIn<T>`] for unary
78/// inputs.
79pub type ManyIn<T> = Context<RequestStream<T>>;
80
81/// Type alias for the output of pipeline that returns a single value
82pub type SingleOut<T> = EngineUnary<T>;
83
84/// Type alias for the output of pipeline that returns multiple values
85pub type ManyOut<T> = EngineStream<T>;
86
87pub type ServiceEngine<T, U> = Engine<T, U, Error>;
88
89/// Unary Engine is a pipeline that takes a single input and returns a single output
90pub type UnaryEngine<T, U> = ServiceEngine<SingleIn<T>, SingleOut<U>>;
91
92/// `ClientStreaming` Engine is a pipeline that takes multiple inputs and returns a single output
93/// Typically the engine will consume the entire input stream; however, it can also decided to exit
94/// early and emit a response without consuming the entire input stream.
95pub type ClientStreamingEngine<T, U> = ServiceEngine<ManyIn<T>, SingleOut<U>>;
96
97/// `ServerStreaming` takes a single input and returns multiple outputs.
98pub type ServerStreamingEngine<T, U> = ServiceEngine<SingleIn<T>, ManyOut<U>>;
99
100/// `BidirectionalStreaming` takes multiple inputs and returns multiple outputs. Input and output values
101/// are considered independent of each other; however, they could be constrained to be related.
102pub type BidirectionalStreamingEngine<T, U> = ServiceEngine<ManyIn<T>, ManyOut<U>>;
103
104pub trait AsyncTransportEngine<T: Data + PipelineIO, U: Data + PipelineIO>:
105    AsyncEngine<T, U, Error> + Send + Sync + 'static
106{
107}
108
109// pub type TransportEngine<T, U> = Arc<dyn AsyncTransportEngine<T, U>>;
110
111mod sealed {
112    use super::*;
113
114    #[allow(dead_code)]
115    pub struct Token;
116
117    pub trait Connectable {
118        type DataType: Data;
119    }
120
121    impl<T: Data> Connectable for Context<T> {
122        type DataType = T;
123    }
124    impl<T: Data> Connectable for EngineUnary<T> {
125        type DataType = T;
126    }
127    impl<T: Data> Connectable for EngineStream<T> {
128        type DataType = T;
129    }
130}
131
132pub trait PipelineIO: sealed::Connectable + AsyncEngineContextProvider + 'static {
133    fn id(&self) -> String;
134}
135
136impl<T: Data> PipelineIO for Context<T> {
137    fn id(&self) -> String {
138        self.id().to_string()
139    }
140}
141impl<T: Data> PipelineIO for EngineUnary<T> {
142    fn id(&self) -> String {
143        self.context().id().to_string()
144    }
145}
146impl<T: Data> PipelineIO for EngineStream<T> {
147    fn id(&self) -> String {
148        self.context().id().to_string()
149    }
150}
151
152#[derive(Serialize, Deserialize, Debug, Clone)]
153pub struct Event {
154    pub id: String,
155}
dynamo_runtime/pipeline.rs

dynamo_runtime/
pipeline.rs
