lib.rs 6.44 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;
9
use dynamo_runtime::{DistributedRuntime, Runtime};
10

11
mod flags;
12
use either::Either;
13
pub use flags::Flags;
14
mod opt;
15
pub use dynamo_llm::request_template::RequestTemplate;
16
pub use opt::Output;
17

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

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

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

36
37
38
39
40
41
42
43
    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.
44
45
46
            // 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?)
47
48
49
        }
    };

50
51
52
53
54
55
    //
    // Configure
    //

    let mut builder = LocalModelBuilder::default();
    builder
56
        .model_name(flags.model_name.clone().or(original_model_identifier))
57
58
59
        .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
60
61
62
        .http_port(flags.http_port)
        .tls_cert_path(flags.tls_cert_path.take())
        .tls_key_path(flags.tls_key_path.take())
63
        .router_config(Some(flags.router_config()))
64
        .request_template(flags.request_template.clone())
65
66
        .migration_limit(flags.migration_limit)
        .is_mocker(matches!(out_opt, Some(Output::Mocker)));
67

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

73
    // TODO: old, address this later:
74
75
    // 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.
76
    let mut rt = Either::Left(runtime.clone());
77
    if let Input::Endpoint(path) = &in_opt {
78
79
        builder.endpoint_id(Some(path.parse().with_context(|| path.clone())?));

80
81
        let dst_config = DistributedConfig::from_settings(flags.static_worker);
        let distributed_runtime = DistributedRuntime::new(runtime.clone(), dst_config).await?;
82
        rt = Either::Right(distributed_runtime);
83
    };
84
85
86
    if let Some(Output::Static(path)) = &out_opt {
        builder.endpoint_id(Some(path.parse().with_context(|| path.clone())?));
    }
87

88
    let local_model = builder.build().await?;
89

90
91
92
    //
    // Create an engine
    //
93

94
    let out_opt = out_opt.unwrap_or_else(|| default_engine_for(&local_model));
95
96
    print_cuda(&out_opt);

97
    // Now that we know the output we're targeting, check if we expect it to work
98
    flags.validate(&in_opt, &out_opt)?;
99

100
    // Make an engine from the local_model, flags and output.
101
102
103
    let engine_config = engine_for(out_opt, flags.clone(), local_model, rt.clone()).await?;

    // Run it from an input
104
    dynamo_llm::entrypoint::input::run_input(rt, in_opt, engine_config).await?;
105

106
107
    Ok(())
}
108

109
110
111
112
113
114
/// 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,
115
    rt: Either<Runtime, DistributedRuntime>,
116
) -> anyhow::Result<EngineConfig> {
117
    match out_opt {
118
119
120
121
122
123
124
125
        Output::Auto => {
            // Auto-discover backends
            Ok(EngineConfig::Dynamic(Box::new(local_model)))
        }
        Output::Static(_) => {
            // A single static backend, no etcd
            Ok(EngineConfig::StaticRemote(Box::new(local_model)))
        }
126
        Output::Echo => Ok(EngineConfig::StaticFull {
127
            model: Box::new(local_model),
128
            engine: dynamo_llm::engines::make_echo_engine(),
129
            is_static: flags.static_worker,
130
        }),
131
        #[cfg(feature = "mistralrs")]
132
133
134
        Output::MistralRs => Ok(EngineConfig::StaticFull {
            engine: dynamo_engine_mistralrs::make_engine(&local_model).await?,
            model: Box::new(local_model),
135
            is_static: flags.static_worker,
136
        }),
137
138
139
140
141
142
143
144
145
146
147
        Output::Mocker => {
            let Either::Right(drt) = rt else {
                panic!("Mocker requires a distributed runtime to run.");
            };

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

148
149
150
            Ok(EngineConfig::StaticCore {
                engine,
                model: Box::new(local_model),
151
                is_static: flags.static_worker,
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
    }
}