lib.rs 17.4 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
// SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

16
use std::io::Read;
17
18
#[cfg(any(feature = "vllm", feature = "sglang"))]
use std::{future::Future, pin::Pin};
19

Neelay Shah's avatar
Neelay Shah committed
20
use dynamo_llm::{
21
22
    backend::ExecutionContext, model_card::model::ModelDeploymentCard,
    types::openai::chat_completions::OpenAIChatCompletionsStreamingEngine,
23
};
Neelay Shah's avatar
Neelay Shah committed
24
use dynamo_runtime::protocols::Endpoint;
25

26
27
mod flags;
pub use flags::Flags;
28
mod hub;
29
mod input;
30
#[cfg(any(feature = "vllm", feature = "sglang"))]
31
mod net;
32
33
34
35
mod opt;
mod output;
pub use opt::{Input, Output};

36
37
38
/// How we identify a namespace/component/endpoint URL.
/// Technically the '://' is not part of the scheme but it eliminates several string
/// concatenations.
39
const ENDPOINT_SCHEME: &str = "dyn://";
40

41
42
43
44
/// When `in=text` the user doesn't need to know the model name, and doesn't need to provide it on
/// the command line. Hence it's optional, and defaults to this.
const INVISIBLE_MODEL_NAME: &str = "dynamo-run";

45
46
47
48
/// How we identify a python string endpoint
#[cfg(feature = "python")]
const PYTHON_STR_SCHEME: &str = "pystr:";

49
50
51
52
/// How we identify a python token endpoint
#[cfg(feature = "python")]
const PYTHON_TOK_SCHEME: &str = "pytok:";

53
pub enum EngineConfig {
54
    /// An remote networked engine we don't know about yet
55
    Dynamic(Endpoint),
56

57
58
59
60
61
    /// A Full service engine does it's own tokenization and prompt formatting.
    StaticFull {
        service_name: String,
        engine: OpenAIChatCompletionsStreamingEngine,
    },
62
63
64
65
66
67
68

    /// A core engine expects to be wrapped with pre/post processors that handle tokenization.
    StaticCore {
        service_name: String,
        engine: ExecutionContext,
        card: Box<ModelDeploymentCard>,
    },
69

70
71
    /// vllm multi-node doesn't run an engine on nodes other than 0. 'ray' does all the work.
    None,
72
73
}

74
#[allow(unused_mut)]
75
pub async fn run(
Neelay Shah's avatar
Neelay Shah committed
76
    runtime: dynamo_runtime::Runtime,
77
    mut in_opt: Input, // mut because vllm and sglang multi-node can change it
78
79
    out_opt: Output,
    flags: Flags,
80
    #[allow(unused_variables)] zmq_socket_prefix: Option<String>,
81
) -> anyhow::Result<()> {
82
83
    let cancel_token = runtime.primary_token();

84
    // Turn relative paths into absolute paths
85
    let mut model_path = flags
86
        .model_path_pos
87
88
        .clone()
        .or(flags.model_path_flag.clone())
89
90
91
92
93
94
95
        .and_then(|p| {
            if p.exists() {
                p.canonicalize().ok()
            } else {
                Some(p)
            }
        });
96

Graham King's avatar
Graham King committed
97
    // Serve the model under the name provided, or the name of the GGUF file or HF repo.
98
    let mut model_name = flags
99
        .model_name
100
        .clone()
101
102
103
104
105
106
107
108
109
110
111
112
113
        .or_else(|| {
            model_path
                .as_ref()
                .and_then(|p| p.iter().last())
                .map(|n| n.to_string_lossy().into_owned())
        })
        .or_else(|| {
            if in_opt == Input::Text {
                Some(INVISIBLE_MODEL_NAME.to_string())
            } else {
                None
            }
        });
114
115
116
117
118
119
120
121
122
123
124
125

    // If it's an HF repo download it
    if let Some(inner_model_path) = model_path.as_ref() {
        if !inner_model_path.exists() {
            model_name = inner_model_path
                .iter()
                .last()
                .map(|s| s.to_string_lossy().to_string());
            model_path = Some(hub::from_hf(inner_model_path).await?);
        }
    }

Graham King's avatar
Graham King committed
126
127
128
129
130
131
    // Load the model deployment card, if any
    // Only used by some engines, so without those feature flags it's unused.
    #[allow(unused_variables)]
    let (maybe_card_path, maybe_card) = match (&model_path, &flags.model_config) {
        // --model-config takes precedence
        (_, Some(model_config)) => {
132
133
134
135
136
137
138
139
140
141
142
143
144
145
            let card =
                match ModelDeploymentCard::from_local_path(model_config, model_name.as_deref())
                    .await
                {
                    Ok(card) => Some(card),
                    Err(e) => {
                        tracing::error!(
                            "Failed to load model card from config path {}: {}",
                            model_config.display(),
                            e
                        );
                        None
                    }
                };
Graham King's avatar
Graham King committed
146
            (Some(model_config.clone()), card)
147
        }
Graham King's avatar
Graham King committed
148
149
        // If --model-path is an HF repo use that
        (Some(model_path), _) if model_path.is_dir() => {
150
            let card = match ModelDeploymentCard::from_local_path(model_path, model_name.as_deref())
Graham King's avatar
Graham King committed
151
                .await
152
153
154
155
156
157
158
159
160
161
162
            {
                Ok(card) => Some(card),
                Err(e) => {
                    tracing::error!(
                        "Failed to load model card from model path {}: {}",
                        model_path.display(),
                        e
                    );
                    None
                }
            };
Graham King's avatar
Graham King committed
163
164
165
            (Some(model_path.clone()), card)
        }
        // Otherwise we don't have one, but we only need it if we're tokenizing
166
167
168
169
        _ => {
            tracing::debug!("No model card path provided (neither --model-config nor a directory in --model-path)");
            (None, None)
        }
170
    };
171

Graham King's avatar
Graham King committed
172
    #[cfg(any(feature = "vllm", feature = "sglang"))]
173
    let mut extra: Option<Pin<Box<dyn Future<Output = ()> + Send>>> = None; // vllm and sglang sub-process
174

175
176
    // Create the engine matching `out`
    let engine_config = match out_opt {
177
178
179
180
181
182
183
184
185
186
187
        Output::EchoFull => {
            let Some(model_name) = model_name else {
                anyhow::bail!(
                    "Pass --model-name or --model-path so we know which model to imitate"
                );
            };
            EngineConfig::StaticFull {
                service_name: model_name,
                engine: output::echo_full::make_engine_full(),
            }
        }
188
189
190
191
192
193
194
195
196
197
198
199
200
        Output::EchoCore => {
            let Some(mut card) = maybe_card.clone() else {
                anyhow::bail!(
                    "out=echo_core need to find the tokenizer. Pass flag --model-path <path>"
                );
            };
            card.requires_preprocessing = true;
            EngineConfig::StaticCore {
                service_name: card.service_name.clone(),
                engine: output::echo_core::make_engine_core(),
                card: Box::new(card),
            }
        }
201
        Output::Endpoint(path) => {
202
            let endpoint: Endpoint = path.parse()?;
203
            EngineConfig::Dynamic(endpoint)
204
        }
205
206
207
208
209
210
211
212
213
214
        #[cfg(feature = "mistralrs")]
        Output::MistralRs => {
            let Some(model_path) = model_path else {
                anyhow::bail!("out=mistralrs requires flag --model-path=<full-path-to-model-gguf>");
            };
            let Some(model_name) = model_name else {
                unreachable!("We checked model_path earlier, and set model_name from model_path");
            };
            EngineConfig::StaticFull {
                service_name: model_name,
Neelay Shah's avatar
Neelay Shah committed
215
                engine: dynamo_llm::engines::mistralrs::make_engine(&model_path).await?,
216
217
            }
        }
218
219
        #[cfg(feature = "sglang")]
        Output::SgLang => {
Neelay Shah's avatar
Neelay Shah committed
220
            use dynamo_llm::engines::sglang;
221
222
223
224
225
226
227
228
229
230
231
            let Some(model_path) = model_path else {
                anyhow::bail!("out=sglang requires flag --model-path=<full-path-to-model-dir>");
            };
            if !model_path.is_dir() {
                anyhow::bail!("`--model-path should point at a HuggingFace repo checkout");
            }
            // Safety: Earlier we build maybe_card from model_path, which we checked right above
            let card = maybe_card.clone().unwrap();
            let Some(sock_prefix) = zmq_socket_prefix else {
                anyhow::bail!("sglang requires zmq_socket_prefix");
            };
Neelay Shah's avatar
Neelay Shah committed
232
            let node_conf = dynamo_llm::engines::MultiNodeConfig {
233
234
                num_nodes: flags.num_nodes,
                node_rank: flags.node_rank,
235
                leader_addr: flags.leader_addr.unwrap_or_default(),
236
237
238
239
240
241
242
            };
            if node_conf.num_nodes > 1 {
                if let Ok(Some(if_name)) = net::get_primary_interface().await {
                    tracing::info!("If you see 'gloo' errors from sglang try setting these environment variables:");
                    tracing::info!("export GLOO_SOCKET_IFNAME={if_name}");
                    tracing::info!("export NCCL_SOCKET_IFNAME={if_name}");
                }
243
244
245
246
247
                if node_conf.node_rank != 0 {
                    // Follower nodes take input from leader node over pytorch distributed, not
                    // from user.
                    in_opt = Input::None;
                }
248
249
250
251
252
253
254
255
256
257
258
            }

            let (engine, sglang_process) = sglang::make_engine(
                cancel_token.clone(),
                &model_path,
                &sock_prefix,
                node_conf,
                flags.tensor_parallel_size,
                flags.base_gpu_id,
            )
            .await?;
259
260
261
            extra = Some(Box::pin(async move {
                let _ = sglang_process.await;
            }));
262
263
            EngineConfig::StaticCore {
                service_name: card.service_name.clone(),
264
265
266
267
                engine,
                card: Box::new(card),
            }
        }
Graham King's avatar
Graham King committed
268
269
        #[cfg(feature = "vllm")]
        Output::Vllm => {
Neelay Shah's avatar
Neelay Shah committed
270
            use dynamo_llm::engines::vllm;
271
272
273
            if flags.base_gpu_id != 0 {
                anyhow::bail!("vllm does not support base_gpu_id. Set environment variable CUDA_VISIBLE_DEVICES instead.");
            }
Graham King's avatar
Graham King committed
274
275
276
277
278
279
280
281
282
283
284
285
286
287
            let Some(model_path) = model_path else {
                anyhow::bail!(
                    "out=vllm requires flag --model-path=<full-path-to-hf-repo-or-model-gguf>"
                );
            };
            let Some(card_path) = maybe_card_path else {
                // If we have a gguf we also need a model card because we don't currently parse
                // tokenizer et al out of gguf.
                anyhow::bail!(
                    "Running GGUF files also requires a `--model-config` for the tokenizer et al."
                );
            };
            let Some(card) = maybe_card.clone() else {
                anyhow::bail!(
288
289
                    "Failed to load model card: either unsupported HuggingFace repo format \
                    or for GGUF files --model-config is missing."
Graham King's avatar
Graham King committed
290
291
292
293
294
                );
            };
            let Some(sock_prefix) = zmq_socket_prefix else {
                anyhow::bail!("vllm requires zmq_socket_prefix");
            };
Neelay Shah's avatar
Neelay Shah committed
295
            let node_conf = dynamo_llm::engines::MultiNodeConfig {
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
                num_nodes: flags.num_nodes,
                node_rank: flags.node_rank,
                leader_addr: flags.leader_addr.unwrap_or_default(),
            };
            if node_conf.num_nodes > 1 {
                if let Ok(Some(if_name)) = net::get_primary_interface().await {
                    tracing::info!("If you see network errors from vllm try setting this environment variable:");
                    tracing::info!("export NCCL_SOCKET_IFNAME={if_name}");
                }
                if node_conf.node_rank != 0 {
                    // Only node 0 runs vllm, the others communicate over ray
                    in_opt = Input::None;
                }
            }
            if node_conf.node_rank == 0 {
                // vllm multi-node only the leader runs vllm
                let (engine, vllm_future) = vllm::make_leader_engine(
                    cancel_token.clone(),
                    &card_path,
                    &model_path,
                    &sock_prefix,
                    node_conf,
                    flags.tensor_parallel_size,
                )
                .await?;
                extra = Some(Box::pin(async move {
                    let _ = vllm_future.await;
                }));
                EngineConfig::StaticCore {
                    service_name: card.service_name.clone(),
                    engine,
                    card: Box::new(card),
                }
            } else {
                // Nodes rank > 0 only run 'ray'
                let stop_future = vllm::start_follower(cancel_token.clone(), node_conf).await?;
                extra = Some(Box::pin(stop_future));
                EngineConfig::None
Graham King's avatar
Graham King committed
334
335
            }
        }
336
337
        #[cfg(feature = "llamacpp")]
        Output::LlamaCpp => {
Neelay Shah's avatar
Neelay Shah committed
338
            use dynamo_llm::engines::llamacpp;
339
340
341
342
343
344
            let Some(model_path) = model_path else {
                anyhow::bail!("out=llamacpp requires flag --model-path=<full-path-to-model-gguf>");
            };
            if !model_path.is_file() {
                anyhow::bail!("--model-path should refer to a GGUF file. llama_cpp does not support safetensors.");
            }
345
            let Some(card) = maybe_card.clone() else {
Graham King's avatar
Graham King committed
346
347
348
                anyhow::bail!(
                    "Pass --model-config so we can find the tokenizer, should be an HF checkout."
                );
349
350
351
352
            };
            let engine = llamacpp::make_engine(cancel_token.clone(), &model_path).await?;
            EngineConfig::StaticCore {
                service_name: card.service_name.clone(),
353
                engine,
Graham King's avatar
Graham King committed
354
355
356
357
358
                card: Box::new(card),
            }
        }
        #[cfg(feature = "trtllm")]
        Output::TrtLLM => {
Neelay Shah's avatar
Neelay Shah committed
359
            use dynamo_llm::engines::trtllm;
Graham King's avatar
Graham King committed
360
361
362
363
364
365
366
367
368
369
370
371
372
373
            let Some(model_path) = model_path else {
                anyhow::bail!("out=trtllm requires flag --model-path=<full-path-to-model-dir>");
            };
            if !model_path.is_dir() {
                anyhow::bail!(
                    "--model-path should point at a directory containing `.engine` files."
                );
            }
            // Safety: Earlier we build maybe_card from model_path, which we checked right above
            let card = maybe_card.clone().unwrap();
            let engine = trtllm::make_engine(model_path.display(), flags.tensor_parallel_size)?;
            EngineConfig::StaticCore {
                service_name: card.service_name.clone(),
                engine,
374
375
376
                card: Box::new(card),
            }
        }
377
378
        #[cfg(feature = "python")]
        Output::PythonStr(path_str) => {
Neelay Shah's avatar
Neelay Shah committed
379
            use dynamo_llm::engines::python;
380
381
382
            let Some(model_name) = model_name else {
                anyhow::bail!("Provide model service name as `--model-name <this>`");
            };
383
            let py_args = flags.as_vec(&path_str, &model_name);
384
            let p = std::path::PathBuf::from(path_str);
385
            let engine = python::make_string_engine(cancel_token.clone(), &p, py_args).await?;
386
387
388
389
390
            EngineConfig::StaticFull {
                service_name: model_name,
                engine,
            }
        }
391
392
        #[cfg(feature = "python")]
        Output::PythonTok(path_str) => {
Neelay Shah's avatar
Neelay Shah committed
393
            use dynamo_llm::engines::python;
394
395
396
397
398
399
            let Some(card) = maybe_card.clone() else {
                anyhow::bail!("Could not find tokenizer. Pass flag --model-path <path>");
            };
            let Some(model_name) = model_name else {
                unreachable!("If we have a card we must have a model name");
            };
400
            let py_args = flags.as_vec(&path_str, &model_name);
401
            let p = std::path::PathBuf::from(path_str);
402
            let engine = python::make_token_engine(cancel_token.clone(), &p, py_args).await?;
403
404
405
406
407
408
            EngineConfig::StaticCore {
                service_name: model_name.clone(),
                engine,
                card: Box::new(card),
            }
        }
409
410
411
412
    };

    match in_opt {
        Input::Http => {
413
            crate::input::http::run(runtime.clone(), flags.http_port, engine_config).await?;
414
415
        }
        Input::Text => {
416
417
418
419
420
421
422
423
424
425
426
427
428
            crate::input::text::run(runtime.clone(), cancel_token.clone(), None, engine_config)
                .await?;
        }
        Input::Stdin => {
            let mut prompt = String::new();
            std::io::stdin().read_to_string(&mut prompt).unwrap();
            crate::input::text::run(
                runtime.clone(),
                cancel_token.clone(),
                Some(prompt),
                engine_config,
            )
            .await?;
429
        }
430
431
432
433
434
435
436
437
438
439
        Input::Batch(path) => {
            crate::input::batch::run(
                runtime.clone(),
                cancel_token.clone(),
                maybe_card,
                path,
                engine_config,
            )
            .await?;
        }
440
441
442
        Input::Endpoint(path) => {
            crate::input::endpoint::run(runtime.clone(), path, engine_config).await?;
        }
443
444
445
446
447
448
449
450
        Input::None => {
            // Multi-node setup. The engine sub-process has been started and is talking
            // to it's node_rank 0 controller. We do nothing.
            // TODO: Acquire an etcd lease, we are running
            cancel_token.cancelled().await;
        }
    }

Graham King's avatar
Graham King committed
451
    #[cfg(any(feature = "vllm", feature = "sglang"))]
452
453
    // Allow engines to ask main thread to wait on an extra future.
    if let Some(extra) = extra {
454
        extra.await;
455
456
457
458
    }

    Ok(())
}