discovery.rs 17.2 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

6
use anyhow::Context as _;
Ryan Olson's avatar
Ryan Olson committed
7
use serde::{Deserialize, Serialize};
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
    component::{self, Component, Instance},
12
    pipeline::{
13
14
        network::egress::push_router::PushRouter, ManyOut, Operator, RouterMode, SegmentSource,
        ServiceBackend, SingleIn, Source,
15
    },
Ryan Olson's avatar
Ryan Olson committed
16
    protocols::{self, annotated::Annotated},
17
18
    slug::Slug,
    transports::etcd::{self, KeyValue, WatchEvent},
19
    DistributedRuntime,
20
};
Ryan Olson's avatar
Ryan Olson committed
21
22

use super::ModelManager;
23
use crate::protocols::openai::completions::{CompletionRequest, CompletionResponse};
24
25
26
27
28
29
use crate::{
    backend::Backend,
    model_type::ModelType,
    preprocessor::{BackendInput, OpenAIPreprocessor},
    protocols::common::llm_backend::LLMEngineOutput,
};
30
31
32
33
use crate::{
    key_value_store::{EtcdStorage, KeyValueStore, KeyValueStoreManager},
    model_card::{self, ModelDeploymentCard},
};
34
35
36
37
38
use crate::{
    kv_router::{scheduler::DefaultWorkerSelector, KvPushRouter, KvRouter},
    protocols::openai::chat_completions::{
        NvCreateChatCompletionRequest, NvCreateChatCompletionStreamResponse,
    },
39
    protocols::openai::embeddings::{NvCreateEmbeddingRequest, NvCreateEmbeddingResponse},
40
};
41
use tracing;
42

Ryan Olson's avatar
Ryan Olson committed
43
44
45
46
47
48
/// [ModelEntry] is a struct that contains the information for the HTTP service to discover models
/// from the etcd cluster.
#[derive(Debug, Clone, Serialize, Deserialize, Eq, PartialEq)]
pub struct ModelEntry {
    /// Public name of the model
    /// This will be used to identify the model in the HTTP service and the value used in an
49
    /// an [OAI ChatRequest][crate::protocols::openai::chat_completions::NvCreateChatCompletionRequest].
Ryan Olson's avatar
Ryan Olson committed
50
51
52
53
    pub name: String,

    /// Component of the endpoint.
    pub endpoint: protocols::Endpoint,
54
55
56

    /// Specifies whether the model is a chat or completion model.s
    pub model_type: ModelType,
Ryan Olson's avatar
Ryan Olson committed
57
58
}

59
impl ModelEntry {
60
61
62
63
64
65
    pub fn requires_preprocessing(&self) -> bool {
        matches!(self.model_type, ModelType::Backend)
    }

    /// Fetch the ModelDeploymentCard from NATS.
    /// This does not touch it's fields so you may need to call move_from_nats on it.
66
67
    pub async fn load_mdc(
        &self,
68
        etcd_client: &etcd::Client,
69
    ) -> anyhow::Result<ModelDeploymentCard> {
70
        let kvstore: Box<dyn KeyValueStore> = Box::new(EtcdStorage::new(etcd_client.clone()));
71
72
73
        let card_store = Arc::new(KeyValueStoreManager::new(kvstore));
        let card_key = ModelDeploymentCard::service_name_slug(&self.name);
        match card_store
74
            .load::<ModelDeploymentCard>(model_card::ROOT_PATH, &card_key)
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
            .await
        {
            Ok(Some(mdc)) => Ok(mdc),
            Ok(None) => {
                anyhow::bail!("Missing ModelDeploymentCard in etcd under key {card_key}");
            }
            Err(err) => {
                anyhow::bail!(
                    "Error fetching ModelDeploymentCard from etcd under key {card_key}. {err}"
                );
            }
        }
    }
}

#[derive(Debug, Clone)]
pub struct ModelNetworkName(String);

impl ModelNetworkName {
    /// Key to store this model entry in networked key-value store (etcd).
    ///
    /// It looks like this:
    /// ns.cp.ep-694d967ca5efd804
    fn from_parts(namespace: &str, component: &str, endpoint: &str, lease_id: i64) -> Self {
99
100
101
        let model_root = component::MODEL_ROOT_PATH;
        let slug = Slug::slugify(&format!("{namespace}.{component}.{endpoint}-{lease_id:x}"));
        ModelNetworkName(format!("{model_root}/{slug}"))
102
103
104
105
106
107
108
109
110
111
112
113
    }

    // We can't do From<&component::Endpoint> here because we also need the lease_id
    pub fn from_local(endpoint: &component::Endpoint, lease_id: i64) -> Self {
        Self::from_parts(
            &endpoint.component().namespace().to_string(),
            &endpoint.component().name(),
            endpoint.name(),
            lease_id,
        )
    }

114
    /// Fetch the ModelEntry from etcd.
115
    pub async fn load_entry(&self, etcd_client: &etcd::Client) -> anyhow::Result<ModelEntry> {
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
        let mut model_entries = etcd_client.kv_get(self.to_string(), None).await?;
        if model_entries.is_empty() {
            anyhow::bail!("No ModelEntry in etcd for key {self}");
        }
        let model_entry = model_entries.remove(0);
        serde_json::from_slice(model_entry.value()).with_context(|| {
            format!(
                "Error deserializing JSON. Key={self}. JSON={}",
                model_entry.value_str().unwrap_or("INVALID UTF-8")
            )
        })
    }

    /// Fetch the ModelDeploymentCard from NATS.
    /// This does not touch it's fields so you may need to call move_from_nats on it.
    /// TODO We have potentially two for each endpoint, one Chat and one Completion.
132
133
    pub async fn load_mdc(
        &self,
134
        etcd_client: &etcd::Client,
135
    ) -> anyhow::Result<ModelDeploymentCard> {
136
        let entry = self.load_entry(etcd_client).await?;
137
        entry.load_mdc(etcd_client).await
138
139
140
    }
}

141
142
143
144
145
146
147
148
impl From<&Instance> for ModelNetworkName {
    fn from(cei: &Instance) -> Self {
        Self::from_parts(
            &cei.namespace,
            &cei.component,
            &cei.endpoint,
            cei.instance_id,
        )
149
150
151
152
153
154
155
156
157
    }
}

impl std::fmt::Display for ModelNetworkName {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{}", self.0)
    }
}

158
159
160
pub struct ModelWatcher {
    manager: ModelManager,
    drt: DistributedRuntime,
161
    router_mode: RouterMode,
162
    notify_on_model: Notify,
Ryan Olson's avatar
Ryan Olson committed
163
164
}

165
166
impl ModelWatcher {
    pub async fn new(
167
        runtime: DistributedRuntime,
168
        model_manager: ModelManager,
169
        router_mode: RouterMode,
170
171
172
    ) -> anyhow::Result<ModelWatcher> {
        Ok(Self {
            manager: model_manager,
173
            drt: runtime,
174
            router_mode,
175
            notify_on_model: Notify::new(),
176
177
        })
    }
Ryan Olson's avatar
Ryan Olson committed
178

179
180
181
182
183
184
185
186
187
188
189
    /// 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
        }
    }

190
191
192
193
194
195
196
197
198
    pub async fn watch(self: Arc<Self>, mut events_rx: Receiver<WatchEvent>) {
        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) => {
199
200
201
202
203
204
205
206
                            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")
                                }
                            }
207
208
209
210
211
212
213
214
                            continue;
                        }
                    };
                    if self.manager.has_model_any(&model_entry.name) {
                        tracing::trace!(
                            service_name = model_entry.name,
                            "New endpoint for existing model"
                        );
215
                        self.notify_on_model.notify_waiters();
216
217
218
                        continue;
                    }

219
                    match self.clone().handle_put(&kv, &model_entry).await {
220
221
                        Ok(()) => {
                            tracing::info!(model_name = model_entry.name, "added model");
222
                            self.notify_on_model.notify_waiters();
223
224
225
226
227
228
229
230
231
                        }
                        Err(e) => {
                            tracing::error!(%e, "error adding model {}", model_entry.name);
                        }
                    }
                }
                WatchEvent::Delete(kv) => match self.clone().handle_delete(&kv).await {
                    Ok(model_name) => {
                        tracing::info!("removed model {}", model_name);
232
233
                    }
                    Err(e) => {
234
                        tracing::error!("error removing model: {}", e);
235
                    }
236
                },
237
            }
Ryan Olson's avatar
Ryan Olson committed
238
239
240
        }
    }

241
242
    /// Returns the name of the model we just deleted
    async fn handle_delete(self: Arc<ModelWatcher>, kv: &KeyValue) -> anyhow::Result<String> {
243
        let key = kv.key_str()?;
244
245
246
247
248
249
250
251
        let model_entry = match self.manager.state.entries.lock().unwrap().remove(key) {
            Some(entry) => entry,
            None => {
                anyhow::bail!("Missing ModelEntry for {key}");
            }
        };
        let model_name = &model_entry.name;
        tracing::debug!(model_name, "removing model");
252

253
254
255
        // Ignore the errors because model could be either type
        let _ = self.manager.remove_chat_completions_model(model_name);
        let _ = self.manager.remove_completions_model(model_name);
256
        let _ = self.manager.remove_embeddings_model(model_name);
257

258
259
        // We own model_entry now so take ownership of the name
        Ok(model_entry.name)
260
    }
Ryan Olson's avatar
Ryan Olson committed
261

262
263
264
265
    // Handles a PUT event from etcd, this usually means adding a new model to the list of served
    // models.
    //
    // If this method errors, for the near term, we will delete the offending key.
266
267
268
269
270
271
    async fn handle_put(
        self: Arc<ModelWatcher>,
        kv: &KeyValue,
        model_entry: &ModelEntry,
    ) -> anyhow::Result<()> {
        let key = kv.key_str()?;
272
        let endpoint_id = model_entry.endpoint.clone();
273
        let component = self
274
275
            .drt
            .namespace(&endpoint_id.namespace)?
276
277
            .component(&endpoint_id.component)?;
        let client = component.endpoint(&endpoint_id.name).client().await?;
278

279
280
281
282
        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");
        };
283
        let card = match model_entry.load_mdc(&etcd_client).await {
284
285
286
287
288
289
290
291
292
293
            Ok(card) => {
                tracing::debug!(card.display_name, "adding model");
                Some(card)
            }
            Err(err) => {
                // `dynamo serve` isn't using MDC yet so can't be an error
                tracing::info!(%err, "load_mdc did not complete");
                None
            }
        };
294
295
296
297
298
299
300
301
        // We need to save the entry to know what the model is called when we delete it
        self.manager
            .state
            .entries
            .lock()
            .unwrap()
            .insert(key.to_string(), model_entry.clone());

302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
        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?);

                let frontend = SegmentSource::<
                    SingleIn<NvCreateChatCompletionRequest>,
                    ManyOut<Annotated<NvCreateChatCompletionStreamResponse>>,
                >::new();
                let preprocessor = OpenAIPreprocessor::new(card.clone()).await?.into_operator();
                let backend = Backend::from_mdc(card.clone()).await?.into_operator();
                let router = PushRouter::<BackendInput, Annotated<LLMEngineOutput>>::from_client(
                    client.clone(),
324
                    self.router_mode,
325
326
327
                )
                .await?;
                let service_backend = match self.router_mode {
328
                    RouterMode::Random | RouterMode::RoundRobin | RouterMode::Direct(_) => {
329
330
                        ServiceBackend::from_engine(Arc::new(router))
                    }
331
                    RouterMode::KV => {
332
333
                        let chooser = self.kv_chooser_for(&model_entry.name, &component).await?;
                        let kv_push_router = KvPushRouter::new(router, chooser);
334
335
336
                        ServiceBackend::from_engine(Arc::new(kv_push_router))
                    }
                };
337

338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
                let chat_engine = frontend
                    .link(preprocessor.forward_edge())?
                    .link(backend.forward_edge())?
                    .link(service_backend)?
                    .link(backend.backward_edge())?
                    .link(preprocessor.backward_edge())?
                    .link(frontend)?;
                self.manager
                    .add_chat_completions_model(&model_entry.name, chat_engine)?;

                let frontend = SegmentSource::<
                    SingleIn<CompletionRequest>,
                    ManyOut<Annotated<CompletionResponse>>,
                >::new();
                let preprocessor = OpenAIPreprocessor::new(card.clone()).await?.into_operator();
                let backend = Backend::from_mdc(card.clone()).await?.into_operator();
                let router = PushRouter::<BackendInput, Annotated<LLMEngineOutput>>::from_client(
355
                    client,
356
                    self.router_mode,
357
                )
358
                .await?;
359
                let service_backend = match self.router_mode {
360
                    RouterMode::Random | RouterMode::RoundRobin | RouterMode::Direct(_) => {
361
362
                        ServiceBackend::from_engine(Arc::new(router))
                    }
363
                    RouterMode::KV => {
364
365
                        let chooser = self.kv_chooser_for(&model_entry.name, &component).await?;
                        let kv_push_router = KvPushRouter::new(router, chooser);
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
                        ServiceBackend::from_engine(Arc::new(kv_push_router))
                    }
                };

                let completions_engine = frontend
                    .link(preprocessor.forward_edge())?
                    .link(backend.forward_edge())?
                    .link(service_backend)?
                    .link(backend.backward_edge())?
                    .link(preprocessor.backward_edge())?
                    .link(frontend)?;
                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 => {
                let push_router =
                    PushRouter::<CompletionRequest, Annotated<CompletionResponse>>::from_client(
                        client,
                        Default::default(),
                    )
                    .await?;
                let engine = Arc::new(push_router);
                self.manager
                    .add_completions_model(&model_entry.name, engine)?;
            }
401
402
403
404
405
406
407
408
409
410
            ModelType::Embedding => {
                let push_router = PushRouter::<
                    NvCreateEmbeddingRequest,
                    Annotated<NvCreateEmbeddingResponse>,
                >::from_client(client, Default::default())
                .await?;
                let engine = Arc::new(push_router);
                self.manager
                    .add_embeddings_model(&model_entry.name, engine)?;
            }
411
        }
Ryan Olson's avatar
Ryan Olson committed
412

413
414
        Ok(())
    }
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447

    async fn kv_chooser_for(
        &self,
        model_name: &str,
        component: &Component,
    ) -> anyhow::Result<Arc<KvRouter>> {
        if let Some(kv_chooser) = self
            .manager
            .state
            .kv_choosers
            .lock()
            .unwrap()
            .get(model_name)
        {
            // Return early to avoid holding the lock during the await later
            return Ok(Arc::clone(kv_chooser));
        }
        let selector = Box::new(DefaultWorkerSelector {});
        let chooser = KvRouter::new(
            component.clone(),
            crate::DEFAULT_KV_BLOCK_SIZE,
            Some(selector),
        )
        .await?;
        let new_kv_chooser = Arc::new(chooser);
        self.manager
            .state
            .kv_choosers
            .lock()
            .unwrap()
            .insert(model_name.to_string(), new_kv_chooser.clone());
        Ok(new_kv_chooser)
    }
Ryan Olson's avatar
Ryan Olson committed
448
}