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

4
5
use std::pin::Pin;

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

26
use dynamo_runtime::{
27
    DistributedRuntime,
28
    component::Client,
29
    engine::{AsyncEngineStream, Data},
30
31
32
33
    pipeline::{
        Context, ManyOut, Operator, PushRouter, RouterMode, SegmentSource, ServiceBackend,
        ServiceEngine, ServiceFrontend, SingleIn, Source,
    },
34
35
36
};
use std::sync::Arc;

37
38
39
40
pub struct PreparedEngine {
    pub service_name: String,
    pub engine: OpenAIChatCompletionsStreamingEngine,
    pub inspect_template: bool,
41
42
43
44
45
46
47
48
49
50
51
52
    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
        }
    }
53
54
}

55
/// Turns an EngineConfig into an OpenAI chat-completions and completions supported StreamingEngine.
56
pub async fn prepare_engine(
57
    distributed_runtime: DistributedRuntime,
58
    engine_config: EngineConfig,
59
) -> anyhow::Result<PreparedEngine> {
60
    match engine_config {
61
        EngineConfig::Dynamic(local_model) => {
62
            let store = Arc::new(distributed_runtime.store().clone());
63
            let model_manager = Arc::new(ModelManager::new());
64
            let watch_obj = Arc::new(ModelWatcher::new(
65
                distributed_runtime.clone(),
66
67
                model_manager.clone(),
                dynamo_runtime::pipeline::RouterMode::RoundRobin,
68
                None,
69
                None,
70
            ));
71
72
73
74
75
            let (_, receiver) = store.watch(
                model_card::ROOT_PATH,
                None,
                distributed_runtime.primary_token(),
            );
76
            let inner_watch_obj = watch_obj.clone();
77
            let _watcher_task = tokio::spawn(async move {
78
                inner_watch_obj.watch(receiver, None).await;
79
            });
80
            tracing::info!("Waiting for remote model..");
81

82
83
84
85
86
            // 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;
87
            tracing::info!("Connected to {model_service_name}");
88
            let engine = model_manager.get_chat_completions_engine(&model_service_name)?;
89
            Ok(PreparedEngine {
90
                service_name: model_service_name,
91
92
                engine,
                inspect_template: false,
93
94
                card: None,
                request_template: local_model.request_template(),
95
            })
96
        }
97
98
99
100
101
102
103
104
        EngineConfig::StaticRemote(local_model) => {
            // The card should have been loaded at 'build' phase earlier
            let card = local_model.card();
            let router_mode = local_model.router_config().router_mode;

            let endpoint_id = local_model.endpoint_id();
            let component = distributed_runtime
                .namespace(&endpoint_id.namespace)?
105
                .component(&endpoint_id.component)?;
106

107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
            let client = component.endpoint(&endpoint_id.name).client().await?;

            let kv_chooser = if router_mode == RouterMode::KV {
                let model_manager = Arc::new(ModelManager::new());
                Some(
                    model_manager
                        .kv_chooser_for(
                            &component,
                            card.kv_cache_block_size,
                            Some(local_model.router_config().kv_router_config),
                        )
                        .await?,
                )
            } else {
                None
            };

124
            let hf_tokenizer = card.tokenizer_hf()?;
125
126
127
            let chat_engine = entrypoint::build_routed_pipeline::<
                NvCreateChatCompletionRequest,
                NvCreateChatCompletionStreamResponse,
128
129
130
131
132
133
134
            >(
                card,
                &client,
                router_mode,
                None,
                kv_chooser.clone(),
                hf_tokenizer,
135
                None, // No prefill chooser in static mode
136
            )
137
138
139
140
141
142
143
144
145
146
147
148
149
            .await?;

            let service_name = local_model.service_name().to_string();
            tracing::info!("Static connecting to {service_name}");
            Ok(PreparedEngine {
                service_name,
                engine: chat_engine,
                inspect_template: false,
                request_template: local_model.request_template(),
                card: Some(local_model.into_card()),
            })
        }
        EngineConfig::StaticFull { engine, model, .. } => {
150
            let service_name = model.service_name().to_string();
151
            tracing::debug!("Model: {service_name} with engine pre-processing");
152
            let engine = Arc::new(StreamingEngineAdapter::new(engine));
153
154
155
156
            Ok(PreparedEngine {
                service_name,
                engine,
                inspect_template: false,
157
158
                request_template: model.request_template(),
                card: Some(model.into_card()),
159
            })
160
161
162
        }
        EngineConfig::StaticCore {
            engine: inner_engine,
163
            model,
164
            ..
165
        } => {
166
167
168
            let pipeline = build_pipeline::<
                NvCreateChatCompletionRequest,
                NvCreateChatCompletionStreamResponse,
169
            >(model.card(), inner_engine, model.card().tokenizer_hf()?)
170
            .await?;
171

172
            let service_name = model.service_name().to_string();
173
174
175
176
177
            tracing::debug!("Model: {service_name} with Dynamo pre-processing");
            Ok(PreparedEngine {
                service_name,
                engine: pipeline,
                inspect_template: true,
178
179
                request_template: model.request_template(),
                card: Some(model.into_card()),
180
            })
181
182
183
        }
    }
}
184
185
186
187

pub async fn build_pipeline<Req, Resp>(
    card: &ModelDeploymentCard,
    engine: ExecutionContext,
188
    hf_tokenizer: tokenizers::Tokenizer,
189
190
191
192
193
) -> anyhow::Result<Arc<ServiceFrontend<SingleIn<Req>, ManyOut<Annotated<Resp>>>>>
where
    Req: Data,
    Resp: Data,
    OpenAIPreprocessor: Operator<
194
195
196
197
198
            Context<Req>,
            Pin<Box<dyn AsyncEngineStream<Annotated<Resp>>>>,
            Context<PreprocessedRequest>,
            Pin<Box<dyn AsyncEngineStream<Annotated<BackendOutput>>>>,
        >,
199
200
{
    let frontend = ServiceFrontend::<SingleIn<Req>, ManyOut<Annotated<Resp>>>::new();
201
202
203
204
205
    let PromptFormatter::OAI(formatter) = PromptFormatter::from_mdc(card)?;
    let preprocessor =
        OpenAIPreprocessor::new_with_parts(card.clone(), formatter, hf_tokenizer.clone())?
            .into_operator();
    let backend = Backend::from_tokenizer(hf_tokenizer).into_operator();
206
207
208
209
210
211
212
213
214
215
216
    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)?)
}

217
218
219
220
pub async fn build_routed_pipeline<Req, Resp>(
    card: &ModelDeploymentCard,
    client: &Client,
    router_mode: RouterMode,
221
    busy_threshold: Option<f64>,
222
    chooser: Option<Arc<KvRouter>>,
223
    hf_tokenizer: tokenizers::Tokenizer,
224
    prefill_chooser: Option<Arc<PrefillRouter>>,
225
) -> anyhow::Result<ServiceEngine<SingleIn<Req>, ManyOut<Annotated<Resp>>>>
226
227
228
229
230
231
232
233
234
235
where
    Req: Data,
    Resp: Data,
    OpenAIPreprocessor: Operator<
            Context<Req>,
            Pin<Box<dyn AsyncEngineStream<Annotated<Resp>>>>,
            Context<PreprocessedRequest>,
            Pin<Box<dyn AsyncEngineStream<Annotated<BackendOutput>>>>,
        >,
{
236
237
238
    let PromptFormatter::OAI(formatter) = PromptFormatter::from_mdc(card)?;
    let preprocessor =
        OpenAIPreprocessor::new_with_parts(card.clone(), formatter, hf_tokenizer.clone())?;
239
240
241
242
243
244
245
    build_routed_pipeline_with_preprocessor(
        card,
        client,
        router_mode,
        busy_threshold,
        chooser,
        preprocessor,
246
        hf_tokenizer,
247
        prefill_chooser,
248
249
250
251
    )
    .await
}

252
#[allow(clippy::too_many_arguments)]
253
254
255
256
257
258
259
pub async fn build_routed_pipeline_with_preprocessor<Req, Resp>(
    card: &ModelDeploymentCard,
    client: &Client,
    router_mode: RouterMode,
    busy_threshold: Option<f64>,
    chooser: Option<Arc<KvRouter>>,
    preprocessor: Arc<OpenAIPreprocessor>,
260
    hf_tokenizer: tokenizers::Tokenizer,
261
    prefill_chooser: Option<Arc<PrefillRouter>>,
262
) -> anyhow::Result<ServiceEngine<SingleIn<Req>, ManyOut<Annotated<Resp>>>>
263
264
265
266
where
    Req: Data,
    Resp: Data,
    OpenAIPreprocessor: Operator<
267
268
269
270
271
            Context<Req>,
            Pin<Box<dyn AsyncEngineStream<Annotated<Resp>>>>,
            Context<PreprocessedRequest>,
            Pin<Box<dyn AsyncEngineStream<Annotated<BackendOutput>>>>,
        >,
272
273
{
    let frontend = SegmentSource::<SingleIn<Req>, ManyOut<Annotated<Resp>>>::new();
274
    let preprocessor_op = preprocessor.into_operator();
275
276
    let backend = Backend::from_tokenizer(hf_tokenizer).into_operator();
    let migration = Migration::from_mdc(card).into_operator();
277
278
279
280
281
282
283
284
285

    // Create worker monitor only if busy_threshold is set
    let worker_monitor = busy_threshold.map(|threshold| {
        Arc::new(crate::discovery::KvWorkerMonitor::new(
            Arc::new(client.clone()),
            threshold,
        )) as Arc<dyn dynamo_runtime::pipeline::WorkerLoadMonitor>
    });

286
287
288
289
290
    let router =
        PushRouter::<PreprocessedRequest, Annotated<LLMEngineOutput>>::from_client_with_threshold(
            client.clone(),
            router_mode,
            busy_threshold,
291
            worker_monitor,
292
293
        )
        .await?;
294

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

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

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

325
326
    Ok(engine)
}