watcher.rs 17.4 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::{mpsc::Receiver, Notify};
Ryan Olson's avatar
Ryan Olson committed
9

Neelay Shah's avatar
Neelay Shah committed
10
use dynamo_runtime::{
11
    pipeline::{
12
13
        network::egress::push_router::PushRouter, ManyOut, Operator, RouterMode, SegmentSource,
        ServiceBackend, SingleIn, Source,
14
    },
15
16
    protocols::annotated::Annotated,
    transports::etcd::{KeyValue, WatchEvent},
17
    DistributedRuntime,
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::{ModelEntry, ModelManager, MODEL_ROOT_PATH};
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>,
Ryan Olson's avatar
Ryan Olson committed
53
54
}

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

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

75
76
77
78
    pub fn set_notify_on_model_update(&mut self, tx: Sender<ModelUpdate>) {
        self.model_update_tx = Some(tx);
    }

79
80
81
82
83
84
85
86
87
88
89
    /// 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
        }
    }

90
    pub async fn watch(&self, mut events_rx: Receiver<WatchEvent>) {
91
92
93
94
95
96
97
98
        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) => {
99
100
101
102
103
104
105
106
                            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")
                                }
                            }
107
108
109
                            continue;
                        }
                    };
110
111
112
113
114
115
116
117
118
                    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());

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

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

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

160
161
162
    /// 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>> {
163
        let key = kv.key_str()?;
164
        let model_entry = match self.manager.remove_model_entry(key) {
165
166
167
168
169
            Some(entry) => entry,
            None => {
                anyhow::bail!("Missing ModelEntry for {key}");
            }
        };
170
171
172
173
174
175
        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() {
176
177
178
179
180
181
182
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
212
213
214
215
216
217
            let mut update_tx = true;
            let mut model_type: ModelType = model_entry.model_type;
            if model_entry.model_type == ModelType::Chat
                && self.manager.list_chat_completions_models().is_empty()
            {
                self.manager.remove_chat_completions_model(&model_name).ok();
                model_type = ModelType::Chat;
            } else if model_entry.model_type == ModelType::Completion
                && self.manager.list_completions_models().is_empty()
            {
                self.manager.remove_completions_model(&model_name).ok();
                model_type = ModelType::Completion;
            } else if model_entry.model_type == ModelType::Embedding
                && self.manager.list_embeddings_models().is_empty()
            {
                self.manager.remove_embeddings_model(&model_name).ok();
                model_type = ModelType::Embedding;
            } else if model_entry.model_type == ModelType::Backend {
                if self.manager.list_chat_completions_models().is_empty() {
                    self.manager.remove_chat_completions_model(&model_name).ok();
                    model_type = ModelType::Chat;
                }
                if self.manager.list_completions_models().is_empty() {
                    self.manager.remove_completions_model(&model_name).ok();
                    if model_type == ModelType::Chat {
                        model_type = ModelType::Backend;
                    } else {
                        model_type = ModelType::Completion;
                    }
                }
            } else {
                tracing::debug!(
                    "Model {} is still active in other instances, not removing",
                    model_name
                );
                update_tx = false;
            }
            if update_tx {
                if let Some(tx) = &self.model_update_tx {
                    tx.send(ModelUpdate::Removed(model_type)).await.ok();
                }
            }
218
219
            return Ok(None);
        }
220

221
        // Ignore the errors because model could be either type
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
        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 {
                if (chat_model_removed && *model_type == ModelType::Chat)
                    || (completions_model_removed && *model_type == ModelType::Completion)
                    || (embeddings_model_removed && *model_type == ModelType::Embedding)
                {
                    if let Some(tx) = &self.model_update_tx {
                        tx.send(ModelUpdate::Removed(*model_type)).await.ok();
                    }
                }
            }
        }
261

262
        Ok(Some(model_name))
263
    }
Ryan Olson's avatar
Ryan Olson committed
264

265
266
    // Handles a PUT event from etcd, this usually means adding a new model to the list of served
    // models.
267
    async fn handle_put(&self, model_entry: &ModelEntry) -> anyhow::Result<()> {
268
        let endpoint_id = model_entry.endpoint.clone();
269
        let component = self
270
271
            .drt
            .namespace(&endpoint_id.namespace)?
272
            .component(&endpoint_id.component)?;
273
        let client = component.endpoint(&endpoint_id.name).client().await?;
274

275
276
277
278
        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");
        };
279
        let card = match model_entry.load_mdc(&etcd_client).await {
280
281
282
283
284
285
286
287
288
            Ok(card) => {
                tracing::debug!(card.display_name, "adding model");
                Some(card)
            }
            Err(err) => {
                tracing::info!(%err, "load_mdc did not complete");
                None
            }
        };
289

290
291
292
293
294
295
296
297
298
299
300
301
302
303
        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?);

304
305
306
                let kv_chooser = if self.router_mode == RouterMode::KV {
                    Some(
                        self.manager
307
308
309
310
                            .kv_chooser_for(
                                &model_entry.name,
                                &component,
                                card.kv_cache_block_size,
311
                                self.kv_router_config,
312
                            )
313
314
315
316
                            .await?,
                    )
                } else {
                    None
317
                };
318

319
320
321
322
323
324
                let chat_engine =
                    entrypoint::build_routed_pipeline::<
                        NvCreateChatCompletionRequest,
                        NvCreateChatCompletionStreamResponse,
                    >(&card, &client, self.router_mode, kv_chooser.clone())
                    .await?;
325
326
327
                self.manager
                    .add_chat_completions_model(&model_entry.name, chat_engine)?;

328
329
330
331
332
                let completions_engine =
                    entrypoint::build_routed_pipeline::<
                        NvCreateCompletionRequest,
                        NvCreateCompletionResponse,
                    >(&card, &client, self.router_mode, kv_chooser)
333
                    .await?;
334
335
336
337
338
339
340
341
342
343
344
345
346
347
                self.manager
                    .add_completions_model(&model_entry.name, completions_engine)?;
            }
            ModelType::Chat => {
                let push_router = PushRouter::<
                    NvCreateChatCompletionRequest,
                    Annotated<NvCreateChatCompletionStreamResponse>,
                >::from_client(client, Default::default())
                .await?;
                let engine = Arc::new(push_router);
                self.manager
                    .add_chat_completions_model(&model_entry.name, engine)?;
            }
            ModelType::Completion => {
348
349
                let push_router = PushRouter::<
                    NvCreateCompletionRequest,
350
                    Annotated<NvCreateCompletionResponse>,
351
352
                >::from_client(client, Default::default())
                .await?;
353
354
355
356
                let engine = Arc::new(push_router);
                self.manager
                    .add_completions_model(&model_entry.name, engine)?;
            }
357
            ModelType::Embedding => {
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
                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>,
                >::from_client(client, self.router_mode)
378
                .await?;
379
380
381
382
383
384
385
386
387
388
389
390
391

                // 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)?;

392
                self.manager
393
                    .add_embeddings_model(&model_entry.name, embedding_engine)?;
394
            }
395
        }
Ryan Olson's avatar
Ryan Olson committed
396

397
398
        Ok(())
    }
399

400
    /// All the registered ModelEntry, one per instance
401
    pub async fn all_entries(&self) -> anyhow::Result<Vec<ModelEntry>> {
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
        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);
423
        }
424
425
426
        Ok(entries)
    }

427
    pub async fn entries_for_model(&self, model_name: &str) -> anyhow::Result<Vec<ModelEntry>> {
428
429
430
        let mut all = self.all_entries().await?;
        all.retain(|entry| entry.name == model_name);
        Ok(all)
431
    }
Ryan Olson's avatar
Ryan Olson committed
432
}