watcher.rs 15.9 KB
Newer Older
Ryan Olson's avatar
Ryan Olson committed
1
2
3
// SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0

Ryan Olson's avatar
Ryan Olson committed
4
use std::sync::Arc;
5
use tokio::sync::mpsc::Sender;
6

7
use anyhow::Context as _;
8
use tokio::sync::{Notify, mpsc::Receiver};
Ryan Olson's avatar
Ryan Olson committed
9

Neelay Shah's avatar
Neelay Shah committed
10
use dynamo_runtime::{
11
    DistributedRuntime,
12
    pipeline::{
13
14
        ManyOut, Operator, RouterMode, SegmentSource, ServiceBackend, SingleIn, Source,
        network::egress::push_router::PushRouter,
15
    },
16
17
    protocols::annotated::Annotated,
    transports::etcd::{KeyValue, WatchEvent},
18
};
Ryan Olson's avatar
Ryan Olson committed
19

20
21
use crate::{
    backend::Backend,
22
23
    entrypoint,
    kv_router::KvRouterConfig,
24
    model_type::ModelType,
25
26
27
28
29
30
31
32
33
34
    preprocessor::{OpenAIPreprocessor, PreprocessedEmbeddingRequest},
    protocols::{
        common::llm_backend::EmbeddingsEngineOutput,
        openai::{
            chat_completions::{
                NvCreateChatCompletionRequest, NvCreateChatCompletionStreamResponse,
            },
            completions::{NvCreateCompletionRequest, NvCreateCompletionResponse},
            embeddings::{NvCreateEmbeddingRequest, NvCreateEmbeddingResponse},
        },
35
36
    },
};
37

38
use super::{MODEL_ROOT_PATH, ModelEntry, ModelManager};
39

40
41
42
43
44
45
#[derive(Debug, Clone, Copy, PartialEq)]
pub enum ModelUpdate {
    Added(ModelType),
    Removed(ModelType),
}

46
pub struct ModelWatcher {
47
    manager: Arc<ModelManager>,
48
    drt: DistributedRuntime,
49
    router_mode: RouterMode,
50
    notify_on_model: Notify,
51
    model_update_tx: Option<Sender<ModelUpdate>>,
52
    kv_router_config: Option<KvRouterConfig>,
53
    busy_threshold: Option<f64>,
Ryan Olson's avatar
Ryan Olson committed
54
55
}

56
57
58
const ALL_MODEL_TYPES: &[ModelType] =
    &[ModelType::Chat, ModelType::Completion, ModelType::Embedding];

59
impl ModelWatcher {
60
    pub fn new(
61
        runtime: DistributedRuntime,
62
        model_manager: Arc<ModelManager>,
63
        router_mode: RouterMode,
64
        kv_router_config: Option<KvRouterConfig>,
65
        busy_threshold: Option<f64>,
66
67
    ) -> ModelWatcher {
        Self {
68
            manager: model_manager,
69
            drt: runtime,
70
            router_mode,
71
            notify_on_model: Notify::new(),
72
            model_update_tx: None,
73
            kv_router_config,
74
            busy_threshold,
75
        }
76
    }
Ryan Olson's avatar
Ryan Olson committed
77

78
79
80
81
    pub fn set_notify_on_model_update(&mut self, tx: Sender<ModelUpdate>) {
        self.model_update_tx = Some(tx);
    }

82
83
84
85
86
87
88
89
90
91
92
    /// Wait until we have at least one chat completions model and return it's name.
    pub async fn wait_for_chat_model(&self) -> String {
        // Loop in case it gets added and immediately deleted
        loop {
            if let Some(model_name) = self.manager.list_chat_completions_models().first() {
                return model_name.to_owned();
            }
            self.notify_on_model.notified().await
        }
    }

93
    pub async fn watch(&self, mut events_rx: Receiver<WatchEvent>) {
94
95
96
97
98
99
100
101
        tracing::debug!("model watcher started");

        while let Some(event) = events_rx.recv().await {
            match event {
                WatchEvent::Put(kv) => {
                    let model_entry = match serde_json::from_slice::<ModelEntry>(kv.value()) {
                        Ok(model_entry) => model_entry,
                        Err(err) => {
102
103
104
105
106
107
108
109
                            match kv.value_str() {
                                Ok(value) => {
                                    tracing::error!(%err, value, "Invalid JSON in model entry")
                                }
                                Err(value_str_err) => {
                                    tracing::error!(original_error = %err, %value_str_err, "Invalid UTF-8 string in model entry, expected JSON")
                                }
                            }
110
111
112
                            continue;
                        }
                    };
113
114
115
116
117
118
119
120
121
                    let key = match kv.key_str() {
                        Ok(k) => k,
                        Err(err) => {
                            tracing::error!(%err, ?kv, "Invalid UTF-8 string in model entry key, skipping");
                            continue;
                        }
                    };
                    self.manager.save_model_entry(key, model_entry.clone());

122
123
124
125
126
127
                    if let Some(tx) = &self.model_update_tx {
                        tx.send(ModelUpdate::Added(model_entry.model_type))
                            .await
                            .ok();
                    }

128
                    if self.manager.has_model_any(&model_entry.name) {
129
                        tracing::trace!(name = model_entry.name, "New endpoint for existing model");
130
                        self.notify_on_model.notify_waiters();
131
132
133
                        continue;
                    }

134
                    match self.handle_put(&model_entry).await {
135
136
                        Ok(()) => {
                            tracing::info!(model_name = model_entry.name, "added model");
137
                            self.notify_on_model.notify_waiters();
138
                        }
139
140
141
142
143
144
                        Err(err) => {
                            tracing::error!(
                                error = format!("{err:#}"),
                                "error adding model {}",
                                model_entry.name
                            );
145
146
147
                        }
                    }
                }
148
149
                WatchEvent::Delete(kv) => match self.handle_delete(&kv).await {
                    Ok(Some(model_name)) => {
150
                        tracing::info!("removed model {}", model_name);
151
                    }
152
153
154
                    Ok(None) => {
                        // There are other instances running this model, nothing to do
                    }
155
                    Err(e) => {
156
                        tracing::error!("error removing model: {}", e);
157
                    }
158
                },
159
            }
Ryan Olson's avatar
Ryan Olson committed
160
161
162
        }
    }

163
164
165
    /// If the last instance running this model has gone delete it.
    /// Returns the name of the model we just deleted, if any.
    async fn handle_delete(&self, kv: &KeyValue) -> anyhow::Result<Option<String>> {
166
        let key = kv.key_str()?;
167
        let model_entry = match self.manager.remove_model_entry(key) {
168
169
170
171
172
            Some(entry) => entry,
            None => {
                anyhow::bail!("Missing ModelEntry for {key}");
            }
        };
173
174
175
176
177
178
179
180
        let model_name = model_entry.name;
        let active_instances = self
            .entries_for_model(&model_name)
            .await
            .with_context(|| model_name.clone())?;
        if !active_instances.is_empty() {
            return Ok(None);
        }
181

182
        // Ignore the errors because model could be either type
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
        let chat_model_remove_err = self.manager.remove_chat_completions_model(&model_name);
        let completions_model_remove_err = self.manager.remove_completions_model(&model_name);
        let embeddings_model_remove_err = self.manager.remove_embeddings_model(&model_name);

        let mut chat_model_removed = false;
        let mut completions_model_removed = false;
        let mut embeddings_model_removed = false;

        if chat_model_remove_err.is_ok() && self.manager.list_chat_completions_models().is_empty() {
            chat_model_removed = true;
        }
        if completions_model_remove_err.is_ok() && self.manager.list_completions_models().is_empty()
        {
            completions_model_removed = true;
        }
        if embeddings_model_remove_err.is_ok() && self.manager.list_embeddings_models().is_empty() {
            embeddings_model_removed = true;
        }

        if !chat_model_removed && !completions_model_removed && !embeddings_model_removed {
            tracing::debug!(
                "No updates to send for model {}: chat_model_removed: {}, completions_model_removed: {}, embeddings_model_removed: {}",
                model_name,
                chat_model_removed,
                completions_model_removed,
                embeddings_model_removed
            );
        } else {
            for model_type in ALL_MODEL_TYPES {
212
                if ((chat_model_removed && *model_type == ModelType::Chat)
213
                    || (completions_model_removed && *model_type == ModelType::Completion)
214
215
                    || (embeddings_model_removed && *model_type == ModelType::Embedding))
                    && let Some(tx) = &self.model_update_tx
216
                {
217
                    tx.send(ModelUpdate::Removed(*model_type)).await.ok();
218
219
220
                }
            }
        }
221

222
        Ok(Some(model_name))
223
    }
Ryan Olson's avatar
Ryan Olson committed
224

225
226
    // Handles a PUT event from etcd, this usually means adding a new model to the list of served
    // models.
227
    async fn handle_put(&self, model_entry: &ModelEntry) -> anyhow::Result<()> {
228
        let endpoint_id = &model_entry.endpoint_id;
229
        let component = self
230
231
            .drt
            .namespace(&endpoint_id.namespace)?
232
            .component(&endpoint_id.component)?;
233
        let client = component.endpoint(&endpoint_id.name).client().await?;
234

235
236
237
238
        let Some(etcd_client) = self.drt.etcd_client() else {
            // Should be impossible because we only get here on an etcd event
            anyhow::bail!("Missing etcd_client");
        };
239
        let card = match model_entry.load_mdc(&etcd_client).await {
240
241
242
243
244
245
246
247
248
            Ok(card) => {
                tracing::debug!(card.display_name, "adding model");
                Some(card)
            }
            Err(err) => {
                tracing::info!(%err, "load_mdc did not complete");
                None
            }
        };
249

250
251
252
253
254
255
256
257
258
259
260
261
262
263
        match model_entry.model_type {
            ModelType::Backend => {
                // A Backend model expects pre-processed requests meaning it's up to us whether we
                // handle Chat or Completions requests, so handle both.

                let Some(mut card) = card else {
                    anyhow::bail!("Missing model deployment card");
                };
                // Download tokenizer.json etc to local disk
                // This cache_dir is a tempfile::TempDir will be deleted on drop. I _think_
                // OpenAIPreprocessor::new loads the files, so we can delete them after this
                // function. Needs checking carefully, possibly we need to store it in state.
                let _cache_dir = Some(card.move_from_nats(self.drt.nats_client()).await?);

264
265
266
                let kv_chooser = if self.router_mode == RouterMode::KV {
                    Some(
                        self.manager
267
268
269
270
                            .kv_chooser_for(
                                &model_entry.name,
                                &component,
                                card.kv_cache_block_size,
271
                                self.kv_router_config,
272
                            )
273
274
275
276
                            .await?,
                    )
                } else {
                    None
277
                };
278

279
280
281
282
283
284
285
286
287
288
289
                let chat_engine = entrypoint::build_routed_pipeline::<
                    NvCreateChatCompletionRequest,
                    NvCreateChatCompletionStreamResponse,
                >(
                    &card,
                    &client,
                    self.router_mode,
                    self.busy_threshold,
                    kv_chooser.clone(),
                )
                .await?;
290
291
292
                self.manager
                    .add_chat_completions_model(&model_entry.name, chat_engine)?;

293
294
295
296
297
298
299
300
301
302
303
                let completions_engine = entrypoint::build_routed_pipeline::<
                    NvCreateCompletionRequest,
                    NvCreateCompletionResponse,
                >(
                    &card,
                    &client,
                    self.router_mode,
                    self.busy_threshold,
                    kv_chooser,
                )
                .await?;
304
305
306
307
308
309
310
                self.manager
                    .add_completions_model(&model_entry.name, completions_engine)?;
            }
            ModelType::Chat => {
                let push_router = PushRouter::<
                    NvCreateChatCompletionRequest,
                    Annotated<NvCreateChatCompletionStreamResponse>,
311
312
313
                >::from_client_with_threshold(
                    client, Default::default(), self.busy_threshold
                )
314
315
316
317
318
319
                .await?;
                let engine = Arc::new(push_router);
                self.manager
                    .add_chat_completions_model(&model_entry.name, engine)?;
            }
            ModelType::Completion => {
320
321
                let push_router = PushRouter::<
                    NvCreateCompletionRequest,
322
                    Annotated<NvCreateCompletionResponse>,
323
324
325
                >::from_client_with_threshold(
                    client, Default::default(), self.busy_threshold
                )
326
                .await?;
327
328
329
330
                let engine = Arc::new(push_router);
                self.manager
                    .add_completions_model(&model_entry.name, engine)?;
            }
331
            ModelType::Embedding => {
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
                let Some(mut card) = card else {
                    anyhow::bail!("Missing model deployment card for embedding model");
                };

                // Download tokenizer files to local disk
                let _cache_dir = Some(card.move_from_nats(self.drt.nats_client()).await?);

                // Create preprocessing pipeline similar to Backend
                let frontend = SegmentSource::<
                    SingleIn<NvCreateEmbeddingRequest>,
                    ManyOut<Annotated<NvCreateEmbeddingResponse>>,
                >::new();

                let preprocessor = OpenAIPreprocessor::new(card.clone()).await?.into_operator();
                let backend = Backend::from_mdc(card.clone()).await?.into_operator();

                let router = PushRouter::<
                    PreprocessedEmbeddingRequest,
                    Annotated<EmbeddingsEngineOutput>,
351
352
353
                >::from_client_with_threshold(
                    client, self.router_mode, self.busy_threshold
                )
354
                .await?;
355
356
357
358
359
360
361
362
363
364
365
366
367

                // Note: Embeddings don't need KV routing complexity
                let service_backend = ServiceBackend::from_engine(Arc::new(router));

                // Link the pipeline: frontend -> preprocessor -> backend -> service_backend -> backend -> preprocessor -> frontend
                let embedding_engine = frontend
                    .link(preprocessor.forward_edge())?
                    .link(backend.forward_edge())?
                    .link(service_backend)?
                    .link(backend.backward_edge())?
                    .link(preprocessor.backward_edge())?
                    .link(frontend)?;

368
                self.manager
369
                    .add_embeddings_model(&model_entry.name, embedding_engine)?;
370
            }
371
        }
Ryan Olson's avatar
Ryan Olson committed
372

373
374
        Ok(())
    }
375

376
    /// All the registered ModelEntry, one per instance
377
    pub async fn all_entries(&self) -> anyhow::Result<Vec<ModelEntry>> {
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
        let Some(etcd_client) = self.drt.etcd_client() else {
            anyhow::bail!("all_entries: Missing etcd client");
        };
        let kvs = etcd_client.kv_get_prefix(MODEL_ROOT_PATH).await?;
        let mut entries = Vec::with_capacity(kvs.len());
        for kv in kvs {
            let model_entry = match serde_json::from_slice::<ModelEntry>(kv.value()) {
                Ok(model_entry) => model_entry,
                Err(err) => {
                    match kv.value_str() {
                        Ok(value) => {
                            tracing::error!(%err, value, "Invalid JSON in model entry")
                        }
                        Err(value_str_err) => {
                            tracing::error!(original_error = %err, %value_str_err, "Invalid UTF-8 string in model entry, expected JSON")
                        }
                    }
                    continue;
                }
            };
            entries.push(model_entry);
399
        }
400
401
402
        Ok(entries)
    }

403
    pub async fn entries_for_model(&self, model_name: &str) -> anyhow::Result<Vec<ModelEntry>> {
404
405
406
        let mut all = self.all_entries().await?;
        all.retain(|entry| entry.name == model_name);
        Ok(all)
407
    }
Ryan Olson's avatar
Ryan Olson committed
408
}