common.rs 11.8 KB
Newer Older
1
// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2
3
// SPDX-License-Identifier: Apache-2.0

4
5
use std::pin::Pin;

6
use crate::{
7
    backend::{Backend, ExecutionContext},
8
    discovery::{KvWorkerMonitor, ModelManager, ModelWatcher},
9
    engines::StreamingEngineAdapter,
10
    entrypoint::{EngineConfig, RouterConfig},
11
    http::service::metrics::Metrics,
12
13
14
    kv_router::{
        DirectRoutingRouter, KvPushRouter, KvRouter, PrefillRouter, metrics::RouterRequestMetrics,
    },
15
    migration::Migration,
16
    model_card::ModelDeploymentCard,
17
    namespace::NamespaceFilter,
18
    preprocessor::{OpenAIPreprocessor, prompt::PromptFormatter},
19
    protocols::common::llm_backend::{BackendOutput, LLMEngineOutput, PreprocessedRequest},
20
    request_template::RequestTemplate,
21
    types::{
22
        Annotated,
23
24
25
26
27
28
        openai::chat_completions::{
            NvCreateChatCompletionRequest, NvCreateChatCompletionStreamResponse,
            OpenAIChatCompletionsStreamingEngine,
        },
    },
};
29

30
use anyhow::Context as _;
31
use dynamo_runtime::{
32
    DistributedRuntime,
33
    component::Client,
34
    engine::{AsyncEngineStream, Data},
35
36
37
38
    pipeline::{
        Context, ManyOut, Operator, PushRouter, RouterMode, SegmentSource, ServiceBackend,
        ServiceEngine, ServiceFrontend, SingleIn, Source,
    },
39
40
41
};
use std::sync::Arc;

42
43
44
45
pub struct PreparedEngine {
    pub service_name: String,
    pub engine: OpenAIChatCompletionsStreamingEngine,
    pub inspect_template: bool,
46
47
48
49
50
51
52
53
54
55
56
57
    pub card: Option<ModelDeploymentCard>,
    pub request_template: Option<RequestTemplate>,
}

impl PreparedEngine {
    pub fn has_tokenizer(&self) -> bool {
        if let Some(card) = self.card.as_ref() {
            card.has_tokenizer()
        } else {
            false
        }
    }
58
59
}

60
/// Turns an EngineConfig into an OpenAI chat-completions and completions supported StreamingEngine.
61
pub async fn prepare_engine(
62
    distributed_runtime: DistributedRuntime,
63
    engine_config: EngineConfig,
64
) -> anyhow::Result<PreparedEngine> {
65
    match engine_config {
66
67
68
        EngineConfig::Dynamic {
            model: local_model, ..
        } => {
69
            let model_manager = Arc::new(ModelManager::new());
70
71
            // Create metrics for migration tracking (not exposed via /metrics in Dynamic engine mode)
            let metrics = Arc::new(Metrics::new());
72
            let watch_obj = Arc::new(ModelWatcher::new(
73
                distributed_runtime.clone(),
74
                model_manager.clone(),
75
                RouterConfig::default(),
76
                local_model.migration_limit(),
77
                None,
78
                metrics,
79
            ));
80
81
82
83
84
85
86
            let discovery = distributed_runtime.discovery();
            let discovery_stream = discovery
                .list_and_watch(
                    dynamo_runtime::discovery::DiscoveryQuery::AllModels,
                    Some(distributed_runtime.primary_token().clone()),
                )
                .await?;
87
            let inner_watch_obj = watch_obj.clone();
88
89
90
91
            let namespace_filter = NamespaceFilter::from_namespace_and_prefix(
                local_model.namespace(),
                local_model.namespace_prefix(),
            );
92
            let _watcher_task = tokio::spawn(async move {
93
94
95
                inner_watch_obj
                    .watch(discovery_stream, namespace_filter)
                    .await;
96
            });
97
            tracing::info!("Waiting for remote model..");
98

99
100
101
102
103
            // TODO: We use the first model to appear, usually we have only one
            // We should add slash commands to text input `/model <name>` to choose,
            // '/models` to list, and notifications when models are added / removed.

            let model_service_name = watch_obj.wait_for_chat_model().await;
104
            tracing::info!("Connected to {model_service_name}");
105
            let engine = model_manager.get_chat_completions_engine(&model_service_name)?;
106
            Ok(PreparedEngine {
107
                service_name: model_service_name,
108
109
                engine,
                inspect_template: false,
110
111
                card: None,
                request_template: local_model.request_template(),
112
            })
113
        }
114
        EngineConfig::InProcessText { engine, model, .. } => {
115
            let service_name = model.service_name().to_string();
116
            tracing::debug!("Model: {service_name} with engine pre-processing");
117
            let engine = Arc::new(StreamingEngineAdapter::new(engine));
118
119
120
121
            Ok(PreparedEngine {
                service_name,
                engine,
                inspect_template: false,
122
123
                request_template: model.request_template(),
                card: Some(model.into_card()),
124
            })
125
        }
126
        EngineConfig::InProcessTokens {
127
            engine: inner_engine,
128
            model,
129
            ..
130
        } => {
131
132
133
            let pipeline = build_pipeline::<
                NvCreateChatCompletionRequest,
                NvCreateChatCompletionStreamResponse,
Nikita's avatar
Nikita committed
134
            >(model.card(), inner_engine, model.card().tokenizer()?)
135
            .await?;
136

137
            let service_name = model.service_name().to_string();
138
139
140
141
142
            tracing::debug!("Model: {service_name} with Dynamo pre-processing");
            Ok(PreparedEngine {
                service_name,
                engine: pipeline,
                inspect_template: true,
143
144
                request_template: model.request_template(),
                card: Some(model.into_card()),
145
            })
146
147
148
        }
    }
}
149
150
151
152

pub async fn build_pipeline<Req, Resp>(
    card: &ModelDeploymentCard,
    engine: ExecutionContext,
Nikita's avatar
Nikita committed
153
    tokenizer: crate::tokenizers::Tokenizer,
154
155
156
157
158
) -> anyhow::Result<Arc<ServiceFrontend<SingleIn<Req>, ManyOut<Annotated<Resp>>>>>
where
    Req: Data,
    Resp: Data,
    OpenAIPreprocessor: Operator<
159
160
161
162
163
            Context<Req>,
            Pin<Box<dyn AsyncEngineStream<Annotated<Resp>>>>,
            Context<PreprocessedRequest>,
            Pin<Box<dyn AsyncEngineStream<Annotated<BackendOutput>>>>,
        >,
164
165
{
    let frontend = ServiceFrontend::<SingleIn<Req>, ManyOut<Annotated<Resp>>>::new();
166
167
    let PromptFormatter::OAI(formatter) = PromptFormatter::from_mdc(card)?;
    let preprocessor =
Nikita's avatar
Nikita committed
168
        OpenAIPreprocessor::new_with_parts(card.clone(), formatter, tokenizer.clone())?
169
            .into_operator();
Nikita's avatar
Nikita committed
170
    let backend = Backend::from_tokenizer(tokenizer).into_operator();
171
172
173
174
175
176
177
178
179
180
181
    let engine = ServiceBackend::from_engine(engine);

    Ok(frontend
        .link(preprocessor.forward_edge())?
        .link(backend.forward_edge())?
        .link(engine)?
        .link(backend.backward_edge())?
        .link(preprocessor.backward_edge())?
        .link(frontend)?)
}

182
#[allow(clippy::too_many_arguments)]
183
184
185
pub async fn build_routed_pipeline<Req, Resp>(
    card: &ModelDeploymentCard,
    client: &Client,
186
    model_manager: Arc<crate::discovery::ModelManager>,
187
    router_mode: RouterMode,
188
    worker_monitor: Option<KvWorkerMonitor>,
189
    chooser: Option<Arc<KvRouter>>,
Nikita's avatar
Nikita committed
190
    tokenizer: crate::tokenizers::Tokenizer,
191
    prefill_chooser: Option<Arc<PrefillRouter>>,
192
    decode_fallback: bool,
193
    migration_limit: u32,
194
    metrics: Arc<Metrics>,
195
) -> anyhow::Result<ServiceEngine<SingleIn<Req>, ManyOut<Annotated<Resp>>>>
196
197
198
199
200
201
202
203
204
205
where
    Req: Data,
    Resp: Data,
    OpenAIPreprocessor: Operator<
            Context<Req>,
            Pin<Box<dyn AsyncEngineStream<Annotated<Resp>>>>,
            Context<PreprocessedRequest>,
            Pin<Box<dyn AsyncEngineStream<Annotated<BackendOutput>>>>,
        >,
{
206
207
    let PromptFormatter::OAI(formatter) =
        PromptFormatter::from_mdc(card).context("PromptFormatter.from_mdc")?;
208
    let preprocessor =
Nikita's avatar
Nikita committed
209
        OpenAIPreprocessor::new_with_parts(card.clone(), formatter, tokenizer.clone())
210
            .context("OpenAIPreprocessor.new_with_parts")?;
211
212
213
    build_routed_pipeline_with_preprocessor(
        card,
        client,
214
        model_manager,
215
        router_mode,
216
        worker_monitor,
217
218
        chooser,
        preprocessor,
Nikita's avatar
Nikita committed
219
        tokenizer,
220
        prefill_chooser,
221
        decode_fallback,
222
        migration_limit,
223
        metrics,
224
225
226
227
    )
    .await
}

228
#[allow(clippy::too_many_arguments)]
229
230
231
pub async fn build_routed_pipeline_with_preprocessor<Req, Resp>(
    card: &ModelDeploymentCard,
    client: &Client,
232
    model_manager: Arc<crate::discovery::ModelManager>,
233
    router_mode: RouterMode,
234
    worker_monitor: Option<KvWorkerMonitor>,
235
236
    chooser: Option<Arc<KvRouter>>,
    preprocessor: Arc<OpenAIPreprocessor>,
Nikita's avatar
Nikita committed
237
    tokenizer: crate::tokenizers::Tokenizer,
238
    prefill_chooser: Option<Arc<PrefillRouter>>,
239
    decode_fallback: bool,
240
    migration_limit: u32,
241
    metrics: Arc<Metrics>,
242
) -> anyhow::Result<ServiceEngine<SingleIn<Req>, ManyOut<Annotated<Resp>>>>
243
244
245
246
where
    Req: Data,
    Resp: Data,
    OpenAIPreprocessor: Operator<
247
248
249
250
251
            Context<Req>,
            Pin<Box<dyn AsyncEngineStream<Annotated<Resp>>>>,
            Context<PreprocessedRequest>,
            Pin<Box<dyn AsyncEngineStream<Annotated<BackendOutput>>>>,
        >,
252
253
{
    let frontend = SegmentSource::<SingleIn<Req>, ManyOut<Annotated<Resp>>>::new();
254
    let preprocessor_op = preprocessor.into_operator();
Nikita's avatar
Nikita committed
255
    let backend = Backend::from_tokenizer(tokenizer).into_operator();
256
    let migration = Migration::from_mdc(card, migration_limit, metrics).into_operator();
257

258
259
260
261
262
263
264
265
266
267
    // For KV routing, use the client from the chooser to ensure shared state
    let router_client = if router_mode == RouterMode::KV {
        let Some(ref chooser) = chooser else {
            anyhow::bail!("RouterMode::KV requires KVRouter to not be null");
        };
        chooser.client().clone()
    } else {
        client.clone()
    };

268
    // Get threshold value and wrap monitor for PushRouter
269
270
271
272
    // Note: PushRouter uses active_decode_blocks_threshold for its internal logic
    let threshold_value = worker_monitor
        .as_ref()
        .map(|m| m.active_decode_blocks_threshold());
273
274
    let monitor_arc =
        worker_monitor.map(|m| Arc::new(m) as Arc<dyn dynamo_runtime::pipeline::WorkerLoadMonitor>);
275

276
277
    let router =
        PushRouter::<PreprocessedRequest, Annotated<LLMEngineOutput>>::from_client_with_threshold(
278
            router_client,
279
            router_mode,
280
281
            threshold_value,
            monitor_arc,
282
283
        )
        .await?;
284

285
286
287
288
289
290
    // Eagerly register router request metrics so they appear as zeros even in
    // non-KV modes (Direct, Random, RoundRobin) where KvPushRouter is never created.
    // In KV mode, KvPushRouter::new() also calls from_component() (idempotent via
    // OnceLock), which covers the standalone router path as well.
    RouterRequestMetrics::from_component(client.endpoint.component());

291
    let service_backend = match router_mode {
292
293
294
295
        RouterMode::Direct => {
            ServiceBackend::from_engine(Arc::new(DirectRoutingRouter::new(router)))
        }
        RouterMode::Random | RouterMode::RoundRobin => {
296
297
298
299
300
301
            ServiceBackend::from_engine(Arc::new(router))
        }
        RouterMode::KV => {
            let Some(chooser) = chooser else {
                anyhow::bail!("RouterMode::KV requires KVRouter to not be null");
            };
302
            ServiceBackend::from_engine(Arc::new(KvPushRouter::new(router, chooser)))
303
304
305
        }
    };

306
    // Use the provided prefill chooser, or create a disabled one if not provided
307
    let prefill_chooser = prefill_chooser
308
        .unwrap_or_else(|| PrefillRouter::disabled(model_manager, router_mode, decode_fallback));
309
310
311
    let prefill_op = prefill_chooser.into_operator();

    // Link with prefill chooser including backward edge for response flow
312
    let engine = frontend
313
        .link(preprocessor_op.forward_edge())?
314
        .link(migration.forward_edge())?
315
        .link(backend.forward_edge())?
316
        .link(prefill_op.forward_edge())?
317
        .link(service_backend)?
318
        .link(prefill_op.backward_edge())?
319
        .link(backend.backward_edge())?
320
        .link(migration.backward_edge())?
321
        .link(preprocessor_op.backward_edge())?
322
        .link(frontend)?;
323

324
325
    Ok(engine)
}