lib.rs 6.34 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 anyhow::Context as _;
use dynamo_llm::entrypoint::EngineConfig;
6
use dynamo_llm::entrypoint::input::Input;
7
use dynamo_llm::local_model::{LocalModel, LocalModelBuilder};
8
use dynamo_runtime::distributed::{DistributedConfig, RequestPlaneMode};
9
10
use dynamo_runtime::storage::key_value_store::KeyValueStoreSelect;
use dynamo_runtime::transports::nats;
11
use dynamo_runtime::{DistributedRuntime, Runtime};
12

13
14
mod flags;
pub use flags::Flags;
15
mod opt;
16
pub use dynamo_llm::request_template::RequestTemplate;
17
pub use opt::Output;
18

19
pub async fn run(
20
    runtime: Runtime,
21
    in_opt: Input,
22
    out_opt: Option<Output>,
Graham King's avatar
Graham King committed
23
    mut flags: Flags,
24
) -> anyhow::Result<()> {
25
26
27
28
29
30
31
32
    //
    // Download
    //

    let maybe_remote_repo = flags
        .model_path_pos
        .clone()
        .or_else(|| flags.model_path_flag.clone());
33
34
35
36

    // Preserve the original model identifier before downloading (for default model name)
    let original_model_identifier = maybe_remote_repo.as_ref().map(|p| p.display().to_string());

37
38
39
40
41
42
43
44
    let model_path = match maybe_remote_repo {
        None => None,
        Some(p) if p.exists() => {
            // Already a local path
            Some(p)
        }
        Some(p) => {
            // model_path might be an HF repo, not a local path. Resolve it by downloading.
45
46
47
            // Mocker only needs tokenizer, not weights
            let ignore_weights = matches!(out_opt, Some(Output::Mocker));
            Some(LocalModel::fetch(&p.display().to_string(), ignore_weights).await?)
48
49
50
        }
    };

51
52
53
54
55
56
    //
    // Configure
    //

    let mut builder = LocalModelBuilder::default();
    builder
57
        .model_name(flags.model_name.clone().or(original_model_identifier))
58
59
60
        .kv_cache_block_size(flags.kv_cache_block_size)
        // Only set if user provides. Usually loaded from tokenizer_config.json
        .context_length(flags.context_length)
Graham King's avatar
Graham King committed
61
62
63
        .http_port(flags.http_port)
        .tls_cert_path(flags.tls_cert_path.take())
        .tls_key_path(flags.tls_key_path.take())
64
        .router_config(Some(flags.router_config()))
65
        .request_template(flags.request_template.clone())
66
67
        .migration_limit(flags.migration_limit)
        .is_mocker(matches!(out_opt, Some(Output::Mocker)));
68

69
70
71
72
73
    // Only the worker has a model path
    if let Some(model_path) = model_path {
        builder.model_path(model_path);
    }

74
    // TODO: old, address this later:
75
76
77
    // If `in=dyn` we want the trtllm/sglang/vllm subprocess to listen on that endpoint.
    // If not, then the endpoint isn't exposed so we let LocalModel invent one.
    if let Input::Endpoint(path) = &in_opt {
78
        builder.endpoint_id(Some(path.parse().with_context(|| path.clone())?));
79
80
    }
    let selected_store: KeyValueStoreSelect = flags.store_kv.parse()?;
81
    let request_plane: RequestPlaneMode = flags.request_plane.parse()?;
82
83
    let dst_config = DistributedConfig {
        store_backend: selected_store,
84
85
86
87
88
89
        // We only need NATS here to monitor it's metrics, so only if it's our request plane.
        nats_config: if request_plane.is_nats() {
            Some(nats::ClientOptions::default())
        } else {
            None
        },
90
        request_plane,
91
    };
92
    let distributed_runtime = DistributedRuntime::new(runtime.clone(), dst_config).await?;
93
    let local_model = builder.build().await?;
94

95
96
97
    //
    // Create an engine
    //
98

99
    let out_opt = out_opt.unwrap_or_else(|| default_engine_for(&local_model));
100
101
    print_cuda(&out_opt);

102
    // Now that we know the output we're targeting, check if we expect it to work
103
    flags.validate(&out_opt)?;
104

105
    // Make an engine from the local_model, flags and output.
106
107
108
109
110
111
112
    let engine_config = engine_for(
        out_opt,
        flags.clone(),
        local_model,
        distributed_runtime.clone(),
    )
    .await?;
113
114

    // Run it from an input
115
    dynamo_llm::entrypoint::input::run_input(distributed_runtime, in_opt, engine_config).await?;
116

117
118
    Ok(())
}
119

120
121
122
123
124
125
/// Create the engine matching `out_opt`
/// Note validation happens in Flags::validate. In here assume everything is going to work.
async fn engine_for(
    out_opt: Output,
    flags: Flags,
    local_model: LocalModel,
126
    drt: DistributedRuntime,
127
) -> anyhow::Result<EngineConfig> {
128
    match out_opt {
129
130
131
132
        Output::Auto => {
            // Auto-discover backends
            Ok(EngineConfig::Dynamic(Box::new(local_model)))
        }
133
        Output::Echo => Ok(EngineConfig::StaticFull {
134
            model: Box::new(local_model),
135
            engine: dynamo_llm::engines::make_echo_engine(),
136
        }),
137
        #[cfg(feature = "mistralrs")]
138
139
140
141
        Output::MistralRs => Ok(EngineConfig::StaticFull {
            engine: dynamo_engine_mistralrs::make_engine(&local_model).await?,
            model: Box::new(local_model),
        }),
142
143
144
145
146
147
148
        Output::Mocker => {
            let args = flags.mocker_config();
            let endpoint = local_model.endpoint_id().clone();

            let engine =
                dynamo_llm::mocker::engine::make_mocker_engine(drt, endpoint, args).await?;

149
150
151
            Ok(EngineConfig::StaticCore {
                engine,
                model: Box::new(local_model),
152
                is_prefill: false,
153
            })
154
155
156
        }
    }
}
157

158
/// If the user will benefit from CUDA or Metal, remind them to build with it.
159
/// If they have it, celebrate!
160
// Only mistralrs needs to be built with CUDA.
161
// The Python engines only need it at runtime.
162
#[cfg(feature = "mistralrs")]
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
fn print_cuda(output: &Output) {
    // These engines maybe be compiled in, but are they the chosen one?
    match output {
        #[cfg(feature = "mistralrs")]
        Output::MistralRs => {}
        _ => {
            return;
        }
    }

    #[cfg(feature = "cuda")]
    {
        tracing::info!("CUDA on");
    }
    #[cfg(feature = "metal")]
    {
        tracing::info!("Metal on");
    }
181
182
    #[cfg(not(any(feature = "cuda", feature = "metal")))]
    tracing::info!("CPU mode. Rebuild with `--features cuda|metal` for better performance");
183
184
}

185
#[cfg(not(feature = "mistralrs"))]
186
fn print_cuda(_output: &Output) {}
187

188
189
fn default_engine_for(_local_model: &LocalModel) -> Output {
    safetensors_default()
190
191
192
193
}

fn safetensors_default() -> Output {
    #[cfg(feature = "mistralrs")]
194
195
196
    {
        Output::MistralRs
    }
197
198

    #[cfg(not(feature = "mistralrs"))]
199
    {
200
        Output::Echo
201
202
    }
}