main.rs 5.04 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
// 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.

use std::env;

use clap::Parser;

20
21
use dynamo_llm::entrypoint::input::Input;
use dynamo_run::Output;
Neelay Shah's avatar
Neelay Shah committed
22
use dynamo_runtime::logging;
23
24

const HELP: &str = r#"
Neelay Shah's avatar
Neelay Shah committed
25
dynamo-run is a single binary that wires together the various inputs (http, text, network) and workers (network, engine), that runs the services. It is the simplest way to use dynamo locally.
26

27
28
29
30
31
Verbosity:
- -v enables debug logs
- -vv enables full trace logs
- Default is info level logging

32
Example:
33
34
- cargo build --features cuda -p dynamo-run
- cd target/debug
35
- ./dynamo-run Qwen/Qwen3-0.6B
36
- OR: ./dynamo-run /data/models/Llama-3.2-1B-Instruct-Q4_K_M.gguf
37
38
"#;

39
const USAGE: &str = "USAGE: dynamo-run in=[http|text|dyn://<path>|batch:<folder>] out=ENGINE_LIST|dyn [--http-port 8080] [--model-path <path>] [--model-name <served-model-name>] [--model-config <hf-repo>] [--tensor-parallel-size=1] [--context-length=N] [--kv-cache-block-size=16] [--num-nodes=1] [--node-rank=0] [--leader-addr=127.0.0.1:9876] [--base-gpu-id=0] [--extra-engine-args=args.json] [--router-mode random|round-robin|kv] [--kv-overlap-score-weight=2.0] [--kv-gpu-cache-usage-weight=1.0] [--kv-waiting-requests-weight=1.0] [--verbosity (-v|-vv)]";
40
41

fn main() -> anyhow::Result<()> {
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
    // Set log level based on verbosity flag
    let log_level = match dynamo_run::Flags::try_parse() {
        Ok(flags) => match flags.verbosity {
            0 => "info",
            1 => "debug",
            2 => "trace",
            _ => {
                return Err(anyhow::anyhow!(
                    "Invalid verbosity level. Valid values are v (debug) or vv (trace)"
                ))
            }
        },
        Err(_) => "info",
    };

    if log_level != "info" {
        std::env::set_var("DYN_LOG", log_level);
    }

61
62
63
    logging::init();

    // max_worker_threads and max_blocking_threads from env vars or config file.
Neelay Shah's avatar
Neelay Shah committed
64
    let rt_config = dynamo_runtime::RuntimeConfig::from_settings()?;
65
66

    // One per process. Wraps a Runtime with holds two tokio runtimes.
Neelay Shah's avatar
Neelay Shah committed
67
    let worker = dynamo_runtime::Worker::from_config(rt_config)?;
68

69
    worker.execute(wrapper)
70
71
}

Neelay Shah's avatar
Neelay Shah committed
72
async fn wrapper(runtime: dynamo_runtime::Runtime) -> anyhow::Result<()> {
73
74
75
    let mut in_opt = None;
    let mut out_opt = None;
    let args: Vec<String> = env::args().skip(1).collect();
76
77
78
79
80
    if args.is_empty()
        || args[0] == "-h"
        || args[0] == "--help"
        || (args.iter().all(|arg| arg == "-v" || arg == "-vv"))
    {
81
82
83
        let engine_list = Output::available_engines().join("|");
        let usage = USAGE.replace("ENGINE_LIST", &engine_list);
        println!("{usage}");
84
85
        println!("{HELP}");
        return Ok(());
86
87
88
89
90
91
92
    } else if args[0] == "--version" {
        if let Some(describe) = option_env!("VERGEN_GIT_DESCRIBE") {
            println!("dynamo-run {}", describe);
        } else {
            println!("Version not available (git describe not available)");
        }
        return Ok(());
93
94
95
    }
    for arg in env::args().skip(1).take(2) {
        let Some((in_out, val)) = arg.split_once('=') else {
96
97
            // Probably we're defaulting in and/or out, and this is a flag
            continue;
98
99
100
101
102
103
104
105
106
107
108
109
110
        };
        match in_out {
            "in" => {
                in_opt = Some(val.try_into()?);
            }
            "out" => {
                out_opt = Some(val.try_into()?);
            }
            _ => {
                anyhow::bail!("Invalid argument, must start with 'in' or 'out. {USAGE}");
            }
        }
    }
111
112
113
114
115
116
    let mut non_flag_params = 1; // binary name
    let in_opt = match in_opt {
        Some(x) => {
            non_flag_params += 1;
            x
        }
117
        None => Input::default(),
118
    };
119
120
121
    if out_opt.is_some() {
        non_flag_params += 1;
    }
122
123

    // Clap skips the first argument expecting it to be the binary name, so add it back
124
    // Note `--model-path` has index=1 (in lib.rs) so that doesn't need a flag.
Neelay Shah's avatar
Neelay Shah committed
125
126
    let flags = dynamo_run::Flags::try_parse_from(
        ["dynamo-run".to_string()]
127
128
129
            .into_iter()
            .chain(env::args().skip(non_flag_params)),
    )?;
130

131
132
133
134
    if is_in_dynamic(&in_opt) && is_out_dynamic(&out_opt) {
        anyhow::bail!("Cannot use endpoint for both in and out");
    }

135
    dynamo_run::run(runtime, in_opt, out_opt, flags).await
136
}
137
138
139
140
141
142
143
144

fn is_in_dynamic(in_opt: &Input) -> bool {
    matches!(in_opt, Input::Endpoint(_))
}

fn is_out_dynamic(out_opt: &Option<Output>) -> bool {
    matches!(out_opt, Some(Output::Dynamic))
}