main.rs 4.56 KB
Newer Older
1
2
3
4
5
6
7
// SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0

use std::env;

use clap::Parser;

8
9
use dynamo_llm::entrypoint::input::Input;
use dynamo_run::Output;
Neelay Shah's avatar
Neelay Shah committed
10
use dynamo_runtime::logging;
11
12

const HELP: &str = r#"
Neelay Shah's avatar
Neelay Shah committed
13
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.
14

15
16
17
18
19
Verbosity:
- -v enables debug logs
- -vv enables full trace logs
- Default is info level logging

20
Example:
21
22
- cargo build --features cuda -p dynamo-run
- cd target/debug
23
- ./dynamo-run Qwen/Qwen3-0.6B
24
- OR: ./dynamo-run /data/models/Llama-3.2-1B-Instruct-Q4_K_M.gguf
25
26
"#;

27
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)]";
28
29

fn main() -> anyhow::Result<()> {
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
    // 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);
    }

49
50
51
    logging::init();

    // max_worker_threads and max_blocking_threads from env vars or config file.
Neelay Shah's avatar
Neelay Shah committed
52
    let rt_config = dynamo_runtime::RuntimeConfig::from_settings()?;
53
    tracing::debug!("Runtime config: {rt_config}");
54

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

58
    worker.execute(wrapper)
59
60
}

Neelay Shah's avatar
Neelay Shah committed
61
async fn wrapper(runtime: dynamo_runtime::Runtime) -> anyhow::Result<()> {
62
63
64
    let mut in_opt = None;
    let mut out_opt = None;
    let args: Vec<String> = env::args().skip(1).collect();
65
66
67
68
69
    if args.is_empty()
        || args[0] == "-h"
        || args[0] == "--help"
        || (args.iter().all(|arg| arg == "-v" || arg == "-vv"))
    {
70
71
72
        let engine_list = Output::available_engines().join("|");
        let usage = USAGE.replace("ENGINE_LIST", &engine_list);
        println!("{usage}");
73
74
        println!("{HELP}");
        return Ok(());
75
76
77
78
79
80
81
    } 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(());
82
83
84
    }
    for arg in env::args().skip(1).take(2) {
        let Some((in_out, val)) = arg.split_once('=') else {
85
86
            // Probably we're defaulting in and/or out, and this is a flag
            continue;
87
88
89
90
91
92
93
94
95
96
97
98
99
        };
        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}");
            }
        }
    }
100
101
102
103
104
105
    let mut non_flag_params = 1; // binary name
    let in_opt = match in_opt {
        Some(x) => {
            non_flag_params += 1;
            x
        }
106
        None => Input::default(),
107
    };
108
109
110
    if out_opt.is_some() {
        non_flag_params += 1;
    }
111
112

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

120
121
122
123
    if is_in_dynamic(&in_opt) && is_out_dynamic(&out_opt) {
        anyhow::bail!("Cannot use endpoint for both in and out");
    }

124
    dynamo_run::run(runtime, in_opt, out_opt, flags).await
125
}
126
127
128
129
130
131
132
133

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))
}