"megatron/vscode:/vscode.git/clone" did not exist on "051f58f1a5a8a7450ffea5c3aadaa2ea4b3a8630"
main.rs 46 KB
Newer Older
1
use clap::{Parser, ValueEnum};
2
3
use nix::sys::signal::{self, Signal};
use nix::unistd::Pid;
4
use serde::Deserialize;
Nicolas Patry's avatar
Nicolas Patry committed
5
use std::env;
6
use std::ffi::OsString;
7
use std::io::{BufRead, BufReader, Lines};
8
use std::os::unix::process::{CommandExt, ExitStatusExt};
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
9
use std::path::Path;
OlivierDehaene's avatar
OlivierDehaene committed
10
use std::process::{Child, Command, ExitStatus, Stdio};
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
11
12
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::mpsc::TryRecvError;
13
use std::sync::{mpsc, Arc};
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
14
15
16
17
use std::thread;
use std::thread::sleep;
use std::time::{Duration, Instant};
use std::{fs, io};
18
use tracing_subscriber::EnvFilter;
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
19

20
21
mod env_runtime;

22
23
#[derive(Clone, Copy, Debug, ValueEnum)]
enum Quantization {
24
    /// 4 bit quantization. Requires a specific AWQ quantized model:
25
    ///   https://hf.co/models?search=awq.
26
    /// Should replace GPTQ models wherever possible because of the better latency
27
28
29
30
31
32
    Awq,
    /// 8 bit quantization, doesn't require specific model.
    /// Should be a drop-in replacement to bitsandbytes with much better performance.
    /// Kernels are from https://github.com/NetEase-FuXi/EETQ.git
    Eetq,
    /// 4 bit quantization. Requires a specific GTPQ quantized model: https://hf.co/models?search=gptq.
33
    /// text-generation-inference will use exllama (faster) kernels wherever possible, and use
34
35
36
37
38
39
40
41
42
    /// triton kernel (wider support) when it's not.
    /// AWQ has faster kernels.
    Gptq,
    /// Bitsandbytes 8bit. Can be applied on any model, will cut the memory requirement in half,
    /// but it is known that the model will be much slower to run than the native f16.
    #[deprecated(
        since = "1.1.0",
        note = "Use `eetq` instead, which provides better latencies overall and is drop-in in most cases"
    )]
43
    Bitsandbytes,
44
45
    /// Bitsandbytes 4bit. Can be applied on any model, will cut the memory requirement by 4x,
    /// but it is known that the model will be much slower to run than the native f16.
Nicolas Patry's avatar
Nicolas Patry committed
46
    BitsandbytesNF4,
47
48
    /// Bitsandbytes 4bit. nf4 should be preferred in most cases but maybe this one has better
    /// perplexity performance for you model
Nicolas Patry's avatar
Nicolas Patry committed
49
    BitsandbytesFP4,
50
51
52
53
54
55
}

impl std::fmt::Display for Quantization {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        // To keep in track with `server`.
        match self {
56
57
            #[allow(deprecated)]
            // Use `eetq` instead, which provides better latencies overall and is drop-in in most cases
58
59
60
            Quantization::Bitsandbytes => {
                write!(f, "bitsandbytes")
            }
Nicolas Patry's avatar
Nicolas Patry committed
61
62
63
64
65
66
            Quantization::BitsandbytesNF4 => {
                write!(f, "bitsandbytes-nf4")
            }
            Quantization::BitsandbytesFP4 => {
                write!(f, "bitsandbytes-fp4")
            }
67
68
69
            Quantization::Gptq => {
                write!(f, "gptq")
            }
70
71
72
            Quantization::Awq => {
                write!(f, "awq")
            }
73
74
75
            Quantization::Eetq => {
                write!(f, "eetq")
            }
76
77
78
79
        }
    }
}

80
81
82
#[derive(Clone, Copy, Debug, ValueEnum)]
enum Dtype {
    Float16,
83
    #[clap(name = "bfloat16")]
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
    BFloat16,
}

impl std::fmt::Display for Dtype {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        // To keep in track with `server`.
        match self {
            Dtype::Float16 => {
                write!(f, "float16")
            }
            Dtype::BFloat16 => {
                write!(f, "bfloat16")
            }
        }
    }
}

Nicolas Patry's avatar
Nicolas Patry committed
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
#[derive(Clone, Copy, Debug, ValueEnum)]
enum RopeScaling {
    Linear,
    Dynamic,
}

impl std::fmt::Display for RopeScaling {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        // To keep in track with `server`.
        match self {
            RopeScaling::Linear => {
                write!(f, "linear")
            }
            RopeScaling::Dynamic => {
                write!(f, "dynamic")
            }
        }
    }
}

Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
121
122
123
124
/// App Configuration
#[derive(Parser, Debug)]
#[clap(author, version, about, long_about = None)]
struct Args {
125
126
127
128
129
    /// The name of the model to load.
    /// Can be a MODEL_ID as listed on <https://hf.co/models> like
    /// `gpt2` or `OpenAssistant/oasst-sft-1-pythia-12b`.
    /// Or it can be a local directory containing the necessary files
    /// as saved by `save_pretrained(...)` methods of transformers
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
130
    #[clap(default_value = "bigscience/bloom-560m", long, env)]
131
    model_id: String,
132
133
134

    /// The actual revision of the model if you're referring to a model
    /// on the hub. You can use a specific commit id or a branch like `refs/pr/2`.
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
135
    #[clap(long, env)]
136
    revision: Option<String>,
137

138
139
140
141
142
    /// The number of tokenizer workers used for payload validation and truncation inside the
    /// router.
    #[clap(default_value = "2", long, env)]
    validation_workers: usize,

143
    /// Whether to shard the model across multiple GPUs
144
145
    /// By default text-generation-inference will use all available GPUs to run
    /// the model. Setting it to `false` deactivates `num_shard`.
146
147
    #[clap(long, env)]
    sharded: Option<bool>,
148
149

    /// The number of shards to use if you don't want to use all GPUs on a given machine.
150
151
    /// You can use `CUDA_VISIBLE_DEVICES=0,1 text-generation-launcher... --num_shard 2`
    /// and `CUDA_VISIBLE_DEVICES=2,3 text-generation-launcher... --num_shard 2` to
152
    /// launch 2 copies with 2 shard each on a given machine with 4 GPUs for instance.
153
154
    #[clap(long, env)]
    num_shard: Option<usize>,
155

156
    /// Whether you want the model to be quantized.
157
158
    #[clap(long, env, value_enum)]
    quantize: Option<Quantization>,
159

Nicolas Patry's avatar
Nicolas Patry committed
160
161
162
163
164
165
166
    /// The number of input_ids to speculate on
    /// If using a medusa model, the heads will be picked up automatically
    /// Other wise, it will use n-gram speculation which is relatively free
    /// in terms of compute, but the speedup heavily depends on the task.
    #[clap(long, env)]
    speculate: Option<usize>,

167
168
169
170
    /// The dtype to be forced upon the model. This option cannot be used with `--quantize`.
    #[clap(long, env, value_enum)]
    dtype: Option<Dtype>,

171
172
173
174
175
176
    /// Whether you want to execute hub modelling code. Explicitly passing a `revision` is
    /// encouraged when loading a model with custom code to ensure no malicious code has been
    /// contributed in a newer revision.
    #[clap(long, env, value_enum)]
    trust_remote_code: bool,

177
178
179
    /// The maximum amount of concurrent requests for this particular deployment.
    /// Having a low limit will refuse clients requests instead of having them
    /// wait for too long and is usually good to handle backpressure correctly.
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
180
181
    #[clap(default_value = "128", long, env)]
    max_concurrent_requests: usize,
182
183
184
185

    /// This is the maximum allowed value for clients to set `best_of`.
    /// Best of makes `n` generations at the same time, and return the best
    /// in terms of overall log probability over the entire generated sequence
186
187
    #[clap(default_value = "2", long, env)]
    max_best_of: usize,
188
189
190
191
192
193

    /// This is the maximum allowed value for clients to set `stop_sequences`.
    /// Stop sequences are used to allow the model to stop on more than just
    /// the EOS token, and enable more complex "prompting" where users can preprompt
    /// the model in a specific way and define their "own" stop token aligned with
    /// their prompt.
194
195
    #[clap(default_value = "4", long, env)]
    max_stop_sequences: usize,
196

Nicolas Patry's avatar
Nicolas Patry committed
197
198
199
200
201
202
203
204
    /// This is the maximum allowed value for clients to set `top_n_tokens`.
    /// `top_n_tokens is used to return information about the the `n` most likely
    /// tokens at each generation step, instead of just the sampled token. This
    /// information can be used for downstream tasks like for classification or
    /// ranking.
    #[clap(default_value = "5", long, env)]
    max_top_n_tokens: u32,

205
206
207
208
    /// This is the maximum allowed input length (expressed in number of tokens)
    /// for users. The larger this value, the longer prompt users can send which
    /// can impact the overall memory required to handle the load.
    /// Please note that some models have a finite range of sequence they can handle.
209
    #[clap(default_value = "1024", long, env)]
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
210
    max_input_length: usize,
211
212
213
214
215
216
217
218
219

    /// This is the most important value to set as it defines the "memory budget"
    /// of running clients requests.
    /// Clients will send input sequences and ask to generate `max_new_tokens`
    /// on top. with a value of `1512` users can send either a prompt of
    /// `1000` and ask for `512` new tokens, or send a prompt of `1` and ask for
    /// `1511` max_new_tokens.
    /// The larger this value, the larger amount each request will be in your RAM
    /// and the less effective batching can be.
220
    #[clap(default_value = "2048", long, env)]
221
    max_total_tokens: usize,
222
223
224
225
226
227
228
229
230
231
232

    /// This represents the ratio of waiting queries vs running queries where
    /// you want to start considering pausing the running queries to include the waiting
    /// ones into the same batch.
    /// `waiting_served_ratio=1.2` Means when 12 queries are waiting and there's
    /// only 10 queries left in the current batch we check if we can fit those 12
    /// waiting queries into the batching strategy, and if yes, then batching happens
    /// delaying the 10 running queries by a `prefill` run.
    ///
    /// This setting is only applied if there is room in the batch
    /// as defined by `max_batch_total_tokens`.
233
234
    #[clap(default_value = "1.2", long, env)]
    waiting_served_ratio: f32,
235

236
237
238
239
240
241
    /// Limits the number of tokens for the prefill operation.
    /// Since this operation take the most memory and is compute bound, it is interesting
    /// to limit the number of requests that can be sent.
    #[clap(default_value = "4096", long, env)]
    max_batch_prefill_tokens: u32,

242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
    /// **IMPORTANT** This is one critical control to allow maximum usage
    /// of the available hardware.
    ///
    /// This represents the total amount of potential tokens within a batch.
    /// When using padding (not recommended) this would be equivalent of
    /// `batch_size` * `max_total_tokens`.
    ///
    /// However in the non-padded (flash attention) version this can be much finer.
    ///
    /// For `max_batch_total_tokens=1000`, you could fit `10` queries of `total_tokens=100`
    /// or a single query of `1000` tokens.
    ///
    /// Overall this number should be the largest possible amount that fits the
    /// remaining memory (after the model is loaded). Since the actual memory overhead
    /// depends on other parameters like if you're using quantization, flash attention
    /// or the model implementation, text-generation-inference cannot infer this number
    /// automatically.
259
260
    #[clap(long, env)]
    max_batch_total_tokens: Option<u32>,
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278

    /// This setting defines how many tokens can be passed before forcing the waiting
    /// queries to be put on the batch (if the size of the batch allows for it).
    /// New queries require 1 `prefill` forward, which is different from `decode`
    /// and therefore you need to pause the running batch in order to run `prefill`
    /// to create the correct values for the waiting queries to be able to join the batch.
    ///
    /// With a value too small, queries will always "steal" the compute to run `prefill`
    /// and running queries will be delayed by a lot.
    ///
    /// With a value too big, waiting queries could wait for a very long time
    /// before being allowed a slot in the running batch. If your server is busy
    /// that means that requests that could run in ~2s on an empty server could
    /// end up running in ~20s because the query had to wait for 18s.
    ///
    /// This number is expressed in number of tokens to make it a bit more
    /// "model" agnostic, but what should really matter is the overall latency
    /// for end users.
279
280
    #[clap(default_value = "20", long, env)]
    max_waiting_tokens: usize,
281

282
283
284
285
286
    /// Enforce a maximum number of requests per batch
    /// Specific flag for hardware targets that do not support unpadded inference
    #[clap(long, env)]
    max_batch_size: Option<usize>,

287
288
289
290
    /// Enable experimental support for cuda graphs
    #[clap(long, env)]
    enable_cuda_graphs: bool,

291
292
293
294
    /// The IP address to listen on
    #[clap(default_value = "0.0.0.0", long, env)]
    hostname: String,

295
    /// The port to listen on.
296
    #[clap(default_value = "3000", long, short, env)]
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
297
    port: u16,
298
299
300

    /// The name of the socket for gRPC communication between the webserver
    /// and the shards.
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
301
302
    #[clap(default_value = "/tmp/text-generation-server", long, env)]
    shard_uds_path: String,
303
304

    /// The address the master shard will listen on. (setting used by torch distributed)
305
    #[clap(default_value = "localhost", long, env)]
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
306
    master_addr: String,
307
308

    /// The address the master port will listen on. (setting used by torch distributed)
309
    #[clap(default_value = "29500", long, env)]
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
310
    master_port: usize,
311
312
313

    /// The location of the huggingface hub cache.
    /// Used to override the location if you want to provide a mounted disk for instance
314
    #[clap(long, env)]
315
    huggingface_hub_cache: Option<String>,
316
317
318

    /// The location of the huggingface hub cache.
    /// Used to override the location if you want to provide a mounted disk for instance
319
320
    #[clap(long, env)]
    weights_cache_override: Option<String>,
321
322
323
324
325

    /// For some models (like bloom), text-generation-inference implemented custom
    /// cuda kernels to speed up inference. Those kernels were only tested on A100.
    /// Use this flag to disable them if you're running on different hardware and
    /// encounter issues.
326
    #[clap(long, env)]
327
    disable_custom_kernels: bool,
328

329
330
331
332
333
    /// Limit the CUDA available memory.
    /// The allowed value equals the total visible memory multiplied by cuda-memory-fraction.
    #[clap(default_value = "1.0", long, env)]
    cuda_memory_fraction: f32,

Nicolas Patry's avatar
Nicolas Patry committed
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
    /// Rope scaling will only be used for RoPE models
    /// and allow rescaling the position rotary to accomodate for
    /// larger prompts.
    ///
    /// Goes together with `rope_factor`.
    ///
    /// `--rope-factor 2.0` gives linear scaling with a factor of 2.0
    /// `--rope-scaling dynamic` gives dynamic scaling with a factor of 1.0
    /// `--rope-scaling linear` gives linear scaling with a factor of 1.0 (Nothing will be changed
    /// basically)
    ///
    /// `--rope-scaling linear --rope-factor` fully describes the scaling you want
    #[clap(long, env)]
    rope_scaling: Option<RopeScaling>,

    /// Rope scaling will only be used for RoPE models
    /// See `rope_scaling`
    #[clap(long, env)]
    rope_factor: Option<f32>,

354
    /// Outputs the logs in JSON format (useful for telemetry)
355
    #[clap(long, env)]
356
    json_output: bool,
357

358
359
    #[clap(long, env)]
    otlp_endpoint: Option<String>,
360

361
362
    #[clap(long, env)]
    cors_allow_origin: Vec<String>,
363
364
365
366
    #[clap(long, env)]
    watermark_gamma: Option<f32>,
    #[clap(long, env)]
    watermark_delta: Option<f32>,
367

368
369
370
371
372
373
374
375
    /// Enable ngrok tunneling
    #[clap(long, env)]
    ngrok: bool,

    /// ngrok authentication token
    #[clap(long, env)]
    ngrok_authtoken: Option<String>,

376
    /// ngrok edge
377
    #[clap(long, env)]
378
    ngrok_edge: Option<String>,
379

380
381
382
383
384
    /// The path to the tokenizer config file. This path is used to load the tokenizer configuration which may
    /// include a `chat_template`. If not provided, the default config will be used from the model hub.
    #[clap(long, env)]
    tokenizer_config_path: Option<String>,

385
386
387
    /// Display a lot of information about your runtime environment
    #[clap(long, short, action)]
    env: bool,
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
388
389
}

390
391
392
#[derive(Debug)]
enum ShardStatus {
    Ready,
393
    Failed(usize),
394
}
395

396
397
398
399
#[allow(clippy::too_many_arguments)]
fn shard_manager(
    model_id: String,
    revision: Option<String>,
400
    quantize: Option<Quantization>,
Nicolas Patry's avatar
Nicolas Patry committed
401
    speculate: Option<usize>,
402
    dtype: Option<Dtype>,
403
    trust_remote_code: bool,
404
405
406
407
408
409
410
411
412
413
    uds_path: String,
    rank: usize,
    world_size: usize,
    master_addr: String,
    master_port: usize,
    huggingface_hub_cache: Option<String>,
    weights_cache_override: Option<String>,
    disable_custom_kernels: bool,
    watermark_gamma: Option<f32>,
    watermark_delta: Option<f32>,
414
    enable_cuda_graphs: bool,
415
    cuda_memory_fraction: f32,
Nicolas Patry's avatar
Nicolas Patry committed
416
417
    rope_scaling: Option<RopeScaling>,
    rope_factor: Option<f32>,
418
419
    otlp_endpoint: Option<String>,
    status_sender: mpsc::Sender<ShardStatus>,
420
    shutdown: Arc<AtomicBool>,
421
422
    _shutdown_sender: mpsc::Sender<()>,
) {
423
424
425
    // Enter shard-manager tracing span
    let _span = tracing::span!(tracing::Level::INFO, "shard-manager", rank = rank).entered();

426
427
428
429
    // Get UDS path
    let uds_string = format!("{uds_path}-{rank}");
    let uds = Path::new(&uds_string);
    // Clean previous runs
430
431
432
    if uds.exists() {
        fs::remove_file(uds).unwrap();
    }
433
434

    // Process args
OlivierDehaene's avatar
OlivierDehaene committed
435
    let mut shard_args = vec![
436
437
438
439
440
441
442
443
444
        "serve".to_string(),
        model_id,
        "--uds-path".to_string(),
        uds_path,
        "--logger-level".to_string(),
        "INFO".to_string(),
        "--json-output".to_string(),
    ];

445
446
    // Activate trust remote code
    if trust_remote_code {
OlivierDehaene's avatar
OlivierDehaene committed
447
        shard_args.push("--trust-remote-code".to_string());
448
449
    }

450
451
    // Activate tensor parallelism
    if world_size > 1 {
OlivierDehaene's avatar
OlivierDehaene committed
452
        shard_args.push("--sharded".to_string());
453
454
    }

455
    if let Some(quantize) = quantize {
OlivierDehaene's avatar
OlivierDehaene committed
456
457
        shard_args.push("--quantize".to_string());
        shard_args.push(quantize.to_string())
458
    }
459

Nicolas Patry's avatar
Nicolas Patry committed
460
461
462
463
464
    if let Some(speculate) = speculate {
        shard_args.push("--speculate".to_string());
        shard_args.push(speculate.to_string())
    }

465
    if let Some(dtype) = dtype {
OlivierDehaene's avatar
OlivierDehaene committed
466
467
        shard_args.push("--dtype".to_string());
        shard_args.push(dtype.to_string())
468
469
    }

470
471
    // Model optional revision
    if let Some(revision) = revision {
OlivierDehaene's avatar
OlivierDehaene committed
472
473
        shard_args.push("--revision".to_string());
        shard_args.push(revision)
474
    }
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
475

Nicolas Patry's avatar
Nicolas Patry committed
476
477
478
479
480
481
    let rope = match (rope_scaling, rope_factor) {
        (None, None) => None,
        (Some(scaling), None) => Some((scaling, 1.0)),
        (Some(scaling), Some(factor)) => Some((scaling, factor)),
        (None, Some(factor)) => Some((RopeScaling::Linear, factor)),
    };
482
483
    // OpenTelemetry
    if let Some(otlp_endpoint) = otlp_endpoint {
OlivierDehaene's avatar
OlivierDehaene committed
484
485
        shard_args.push("--otlp-endpoint".to_string());
        shard_args.push(otlp_endpoint);
486
487
488
    }

    // Copy current process env
OlivierDehaene's avatar
OlivierDehaene committed
489
    let mut envs: Vec<(OsString, OsString)> = env::vars_os().collect();
490
491

    // Torch Distributed Env vars
OlivierDehaene's avatar
OlivierDehaene committed
492
493
494
495
    envs.push(("RANK".into(), rank.to_string().into()));
    envs.push(("WORLD_SIZE".into(), world_size.to_string().into()));
    envs.push(("MASTER_ADDR".into(), master_addr.into()));
    envs.push(("MASTER_PORT".into(), master_port.to_string().into()));
496
    envs.push(("TORCH_NCCL_AVOID_RECORD_STREAMS".into(), "1".into()));
497

498
499
500
501
502
503
    // CUDA memory fraction
    envs.push((
        "CUDA_MEMORY_FRACTION".into(),
        cuda_memory_fraction.to_string().into(),
    ));

504
    // Safetensors load fast
OlivierDehaene's avatar
OlivierDehaene committed
505
    envs.push(("SAFETENSORS_FAST_GPU".into(), "1".into()));
506

507
508
509
    // Disable progress bar
    envs.push(("HF_HUB_DISABLE_PROGRESS_BARS".into(), "1".into()));

510
511
    // Enable hf transfer for insane download speeds
    let enable_hf_transfer = env::var("HF_HUB_ENABLE_HF_TRANSFER").unwrap_or("1".to_string());
OlivierDehaene's avatar
OlivierDehaene committed
512
    envs.push((
513
514
515
516
517
518
        "HF_HUB_ENABLE_HF_TRANSFER".into(),
        enable_hf_transfer.into(),
    ));

    // Parse Inference API token
    if let Ok(api_token) = env::var("HF_API_TOKEN") {
OlivierDehaene's avatar
OlivierDehaene committed
519
        envs.push(("HUGGING_FACE_HUB_TOKEN".into(), api_token.into()))
520
521
    };

Nicolas Patry's avatar
Nicolas Patry committed
522
523
524
525
526
527
528
529
530
    // Detect rope scaling
    // Sending as env instead of CLI args to not bloat everything
    // those only can be used by RoPE models, so passing information around
    // for all models will complexify code unnecessarily
    if let Some((scaling, factor)) = rope {
        envs.push(("ROPE_SCALING".into(), scaling.to_string().into()));
        envs.push(("ROPE_FACTOR".into(), factor.to_string().into()));
    }

531
532
533
    // If huggingface_hub_cache is some, pass it to the shard
    // Useful when running inside a docker container
    if let Some(huggingface_hub_cache) = huggingface_hub_cache {
OlivierDehaene's avatar
OlivierDehaene committed
534
        envs.push(("HUGGINGFACE_HUB_CACHE".into(), huggingface_hub_cache.into()));
535
536
537
538
539
    };

    // If weights_cache_override is some, pass it to the shard
    // Useful when running inside a HuggingFace Inference Endpoint
    if let Some(weights_cache_override) = weights_cache_override {
OlivierDehaene's avatar
OlivierDehaene committed
540
        envs.push((
541
542
543
544
545
            "WEIGHTS_CACHE_OVERRIDE".into(),
            weights_cache_override.into(),
        ));
    };

546
547
548
549
550
    // Enable experimental support for cuda graphs
    if enable_cuda_graphs {
        envs.push(("ENABLE_CUDA_GRAPHS".into(), "True".into()))
    }

551
552
    // If disable_custom_kernels is true, pass it to the shard as an env var
    if disable_custom_kernels {
OlivierDehaene's avatar
OlivierDehaene committed
553
        envs.push(("DISABLE_CUSTOM_KERNELS".into(), "True".into()))
554
555
556
557
    }

    // Watermark Gamma
    if let Some(watermark_gamma) = watermark_gamma {
OlivierDehaene's avatar
OlivierDehaene committed
558
        envs.push(("WATERMARK_GAMMA".into(), watermark_gamma.to_string().into()))
559
560
561
562
    }

    // Watermark Delta
    if let Some(watermark_delta) = watermark_delta {
OlivierDehaene's avatar
OlivierDehaene committed
563
        envs.push(("WATERMARK_DELTA".into(), watermark_delta.to_string().into()))
564
565
566
    }

    // Start process
567
    tracing::info!("Starting shard");
568
    let mut p = match Command::new("text-generation-server")
OlivierDehaene's avatar
OlivierDehaene committed
569
570
        .args(shard_args)
        .envs(envs)
571
572
573
574
575
        .stdout(Stdio::piped())
        .stderr(Stdio::piped())
        .process_group(0)
        .spawn()
    {
576
577
        Ok(p) => p,
        Err(err) => {
578
579
580
            if err.kind() == io::ErrorKind::NotFound {
                tracing::error!("text-generation-server not found in PATH");
                tracing::error!("Please install it with `make install-server`")
581
582
            }
            {
583
                tracing::error!("{}", err);
584
            }
585

586
            status_sender.send(ShardStatus::Failed(rank)).unwrap();
587
588
589
590
591
            return;
        }
    };

    // Redirect STDOUT to the console
592
    let shard_stdout_reader = BufReader::new(p.stdout.take().unwrap());
593
    let shard_stderr_reader = BufReader::new(p.stderr.take().unwrap());
594

595
    //stdout tracing thread
596
    thread::spawn(move || {
597
        log_lines(shard_stdout_reader.lines());
598
    });
599
600
601
602
603
604
605
    // We read stderr in another thread as it seems that lines() can block in some cases
    let (err_sender, err_receiver) = mpsc::channel();
    thread::spawn(move || {
        for line in shard_stderr_reader.lines().flatten() {
            err_sender.send(line).unwrap_or(());
        }
    });
606
607
608
609
610
611

    let mut ready = false;
    let start_time = Instant::now();
    let mut wait_time = Instant::now();
    loop {
        // Process exited
612
        if let Some(exit_status) = p.try_wait().unwrap() {
613
614
615
616
            let mut err = String::new();
            while let Ok(line) = err_receiver.recv_timeout(Duration::from_millis(10)) {
                err = err + "\n" + &line;
            }
617

618
            tracing::error!("Shard complete standard error output:\n{err}");
619

620
            if let Some(signal) = exit_status.signal() {
621
622
623
                tracing::error!("Shard process was signaled to shutdown with signal {signal}");
            }

624
            status_sender.send(ShardStatus::Failed(rank)).unwrap();
625
626
627
628
            return;
        }

        // We received a shutdown signal
629
        if shutdown.load(Ordering::SeqCst) {
630
            p.kill().unwrap();
631
            let _ = p.wait();
632
            tracing::info!("Shard terminated");
633
634
635
636
637
            return;
        }

        // Shard is ready
        if uds.exists() && !ready {
638
            tracing::info!("Shard ready in {:?}", start_time.elapsed());
639
640
641
            status_sender.send(ShardStatus::Ready).unwrap();
            ready = true;
        } else if !ready && wait_time.elapsed() > Duration::from_secs(10) {
642
            tracing::info!("Waiting for shard to be ready...");
643
644
645
646
647
648
            wait_time = Instant::now();
        }
        sleep(Duration::from_millis(100));
    }
}

649
fn shutdown_shards(shutdown: Arc<AtomicBool>, shutdown_receiver: &mpsc::Receiver<()>) {
650
651
652
    tracing::info!("Shutting down shards");
    // Update shutdown value to true
    // This will be picked up by the shard manager
653
    shutdown.store(true, Ordering::SeqCst);
654
655
656
657
658
659
660

    // Wait for shards to shutdown
    // This will block till all shutdown_sender are dropped
    let _ = shutdown_receiver.recv();
}

fn num_cuda_devices() -> Option<usize> {
661
662
663
664
    let devices = match env::var("CUDA_VISIBLE_DEVICES") {
        Ok(devices) => devices,
        Err(_) => env::var("NVIDIA_VISIBLE_DEVICES").ok()?,
    };
665
666
    let n_devices = devices.split(',').count();
    Some(n_devices)
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
}

#[derive(Deserialize)]
#[serde(rename_all = "UPPERCASE")]
enum PythonLogLevelEnum {
    Trace,
    Debug,
    Info,
    Success,
    Warning,
    Error,
    Critical,
}

#[derive(Deserialize)]
struct PythonLogLevel {
    name: PythonLogLevelEnum,
}

#[derive(Deserialize)]
struct PythonLogRecord {
    level: PythonLogLevel,
}

#[derive(Deserialize)]
struct PythonLogMessage {
    text: String,
    record: PythonLogRecord,
}

impl PythonLogMessage {
    fn trace(&self) {
        match self.record.level.name {
            PythonLogLevelEnum::Trace => tracing::trace!("{}", self.text),
            PythonLogLevelEnum::Debug => tracing::debug!("{}", self.text),
            PythonLogLevelEnum::Info => tracing::info!("{}", self.text),
            PythonLogLevelEnum::Success => tracing::info!("{}", self.text),
            PythonLogLevelEnum::Warning => tracing::warn!("{}", self.text),
            PythonLogLevelEnum::Error => tracing::error!("{}", self.text),
            PythonLogLevelEnum::Critical => tracing::error!("{}", self.text),
        }
    }
}

711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
impl TryFrom<&String> for PythonLogMessage {
    type Error = serde_json::Error;

    fn try_from(value: &String) -> Result<Self, Self::Error> {
        serde_json::from_str::<Self>(value)
    }
}

fn log_lines<S: Sized + BufRead>(lines: Lines<S>) {
    for line in lines.flatten() {
        match PythonLogMessage::try_from(&line) {
            Ok(log) => log.trace(),
            Err(_) => tracing::debug!("{line}"),
        }
    }
}

728
729
730
731
fn find_num_shards(
    sharded: Option<bool>,
    num_shard: Option<usize>,
) -> Result<usize, LauncherError> {
732
733
734
735
    // get the number of shards given `sharded` and `num_shard`
    let num_shard = match (sharded, num_shard) {
        (Some(true), None) => {
            // try to default to the number of available GPUs
736
737
738
            tracing::info!("Parsing num_shard from CUDA_VISIBLE_DEVICES/NVIDIA_VISIBLE_DEVICES");
            let n_devices = num_cuda_devices()
                .expect("--num-shard and CUDA_VISIBLE_DEVICES/NVIDIA_VISIBLE_DEVICES are not set");
739
            if n_devices <= 1 {
740
741
742
                return Err(LauncherError::NotEnoughCUDADevices(format!(
                    "`sharded` is true but only found {n_devices} CUDA devices"
                )));
743
            }
744
            n_devices
745
        }
746
747
748
        (Some(true), Some(num_shard)) => {
            // we can't have only one shard while sharded
            if num_shard <= 1 {
749
750
751
                return Err(LauncherError::ArgumentValidation(
                    "`sharded` is true but `num_shard` <= 1".to_string(),
                ));
752
753
            }
            num_shard
754
        }
755
756
757
758
        (Some(false), Some(num_shard)) => num_shard,
        (Some(false), None) => 1,
        (None, None) => num_cuda_devices().unwrap_or(1),
        (None, Some(num_shard)) => num_shard,
759
    };
760
    if num_shard < 1 {
761
762
763
        return Err(LauncherError::ArgumentValidation(
            "`num_shard` cannot be < 1".to_string(),
        ));
764
    }
765
    Ok(num_shard)
766
}
767

768
769
#[derive(Debug)]
enum LauncherError {
770
771
    ArgumentValidation(String),
    NotEnoughCUDADevices(String),
772
773
774
775
776
777
778
    DownloadError,
    ShardCannotStart,
    ShardDisconnected,
    ShardFailed,
    WebserverFailed,
    WebserverCannotStart,
}
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
779

780
fn download_convert_model(args: &Args, running: Arc<AtomicBool>) -> Result<(), LauncherError> {
781
782
783
    // Enter download tracing span
    let _span = tracing::span!(tracing::Level::INFO, "download").entered();

OlivierDehaene's avatar
OlivierDehaene committed
784
    let mut download_args = vec![
785
786
787
788
789
790
791
792
        "download-weights".to_string(),
        args.model_id.to_string(),
        "--extension".to_string(),
        ".safetensors".to_string(),
        "--logger-level".to_string(),
        "INFO".to_string(),
        "--json-output".to_string(),
    ];
793

794
795
    // Model optional revision
    if let Some(revision) = &args.revision {
OlivierDehaene's avatar
OlivierDehaene committed
796
797
        download_args.push("--revision".to_string());
        download_args.push(revision.to_string())
798
    }
799

800
801
802
803
804
    // Trust remote code for automatic peft fusion
    if args.trust_remote_code {
        download_args.push("--trust-remote-code".to_string());
    }

805
    // Copy current process env
OlivierDehaene's avatar
OlivierDehaene committed
806
    let mut envs: Vec<(OsString, OsString)> = env::vars_os().collect();
807

808
809
810
    // Disable progress bar
    envs.push(("HF_HUB_DISABLE_PROGRESS_BARS".into(), "1".into()));

811
    // If huggingface_hub_cache is set, pass it to the download process
812
813
    // Useful when running inside a docker container
    if let Some(ref huggingface_hub_cache) = args.huggingface_hub_cache {
OlivierDehaene's avatar
OlivierDehaene committed
814
        envs.push(("HUGGINGFACE_HUB_CACHE".into(), huggingface_hub_cache.into()));
815
    };
816

817
818
    // Enable hf transfer for insane download speeds
    let enable_hf_transfer = env::var("HF_HUB_ENABLE_HF_TRANSFER").unwrap_or("1".to_string());
OlivierDehaene's avatar
OlivierDehaene committed
819
    envs.push((
820
821
822
        "HF_HUB_ENABLE_HF_TRANSFER".into(),
        enable_hf_transfer.into(),
    ));
823

824
825
    // Parse Inference API token
    if let Ok(api_token) = env::var("HF_API_TOKEN") {
OlivierDehaene's avatar
OlivierDehaene committed
826
        envs.push(("HUGGING_FACE_HUB_TOKEN".into(), api_token.into()))
827
    };
828

829
830
831
    // If args.weights_cache_override is some, pass it to the download process
    // Useful when running inside a HuggingFace Inference Endpoint
    if let Some(weights_cache_override) = &args.weights_cache_override {
OlivierDehaene's avatar
OlivierDehaene committed
832
        envs.push((
833
834
835
836
837
            "WEIGHTS_CACHE_OVERRIDE".into(),
            weights_cache_override.into(),
        ));
    };

838
839
    // Start process
    tracing::info!("Starting download process.");
840
    let mut download_process = match Command::new("text-generation-server")
OlivierDehaene's avatar
OlivierDehaene committed
841
842
        .args(download_args)
        .envs(envs)
843
844
845
846
847
        .stdout(Stdio::piped())
        .stderr(Stdio::piped())
        .process_group(0)
        .spawn()
    {
848
849
        Ok(p) => p,
        Err(err) => {
850
851
852
            if err.kind() == io::ErrorKind::NotFound {
                tracing::error!("text-generation-server not found in PATH");
                tracing::error!("Please install it with `make install-server`")
853
854
            } else {
                tracing::error!("{}", err);
855
            }
856

857
858
859
            return Err(LauncherError::DownloadError);
        }
    };
860

861
    let download_stdout = BufReader::new(download_process.stdout.take().unwrap());
862

863
    thread::spawn(move || {
864
865
866
867
868
869
870
871
872
873
874
        log_lines(download_stdout.lines());
    });

    let download_stderr = BufReader::new(download_process.stderr.take().unwrap());

    // We read stderr in another thread as it seems that lines() can block in some cases
    let (err_sender, err_receiver) = mpsc::channel();
    thread::spawn(move || {
        for line in download_stderr.lines().flatten() {
            err_sender.send(line).unwrap_or(());
        }
875
    });
876

877
    loop {
878
879
880
881
        if let Some(status) = download_process.try_wait().unwrap() {
            if status.success() {
                tracing::info!("Successfully downloaded weights.");
                break;
882
            }
883
884

            let mut err = String::new();
885
886
887
888
            while let Ok(line) = err_receiver.recv_timeout(Duration::from_millis(10)) {
                err = err + "\n" + &line;
            }

889
890
891
892
893
894
895
896
897
            if let Some(signal) = status.signal() {
                tracing::error!(
                    "Download process was signaled to shutdown with signal {signal}: {err}"
                );
            } else {
                tracing::error!("Download encountered an error: {err}");
            }

            return Err(LauncherError::DownloadError);
898
        }
899
        if !running.load(Ordering::SeqCst) {
OlivierDehaene's avatar
OlivierDehaene committed
900
            terminate("download", download_process, Duration::from_secs(10)).unwrap();
901
902
903
            return Ok(());
        }
        sleep(Duration::from_millis(100));
904
    }
905
906
    Ok(())
}
907

908
#[allow(clippy::too_many_arguments)]
909
910
911
fn spawn_shards(
    num_shard: usize,
    args: &Args,
912
    shutdown: Arc<AtomicBool>,
913
914
915
916
917
918
    shutdown_receiver: &mpsc::Receiver<()>,
    shutdown_sender: mpsc::Sender<()>,
    status_receiver: &mpsc::Receiver<ShardStatus>,
    status_sender: mpsc::Sender<ShardStatus>,
    running: Arc<AtomicBool>,
) -> Result<(), LauncherError> {
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
919
920
    // Start shard processes
    for rank in 0..num_shard {
921
922
923
924
925
926
        let model_id = args.model_id.clone();
        let revision = args.revision.clone();
        let uds_path = args.shard_uds_path.clone();
        let master_addr = args.master_addr.clone();
        let huggingface_hub_cache = args.huggingface_hub_cache.clone();
        let weights_cache_override = args.weights_cache_override.clone();
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
927
928
929
        let status_sender = status_sender.clone();
        let shutdown = shutdown.clone();
        let shutdown_sender = shutdown_sender.clone();
930
        let otlp_endpoint = args.otlp_endpoint.clone();
931
        let quantize = args.quantize;
Nicolas Patry's avatar
Nicolas Patry committed
932
        let speculate = args.speculate;
933
        let dtype = args.dtype;
934
        let trust_remote_code = args.trust_remote_code;
935
936
937
938
        let master_port = args.master_port;
        let disable_custom_kernels = args.disable_custom_kernels;
        let watermark_gamma = args.watermark_gamma;
        let watermark_delta = args.watermark_delta;
939
        let enable_cuda_graphs = args.enable_cuda_graphs;
940
        let cuda_memory_fraction = args.cuda_memory_fraction;
Nicolas Patry's avatar
Nicolas Patry committed
941
942
        let rope_scaling = args.rope_scaling;
        let rope_factor = args.rope_factor;
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
943
944
        thread::spawn(move || {
            shard_manager(
945
                model_id,
946
                revision,
947
                quantize,
Nicolas Patry's avatar
Nicolas Patry committed
948
                speculate,
949
                dtype,
950
                trust_remote_code,
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
951
952
953
954
955
                uds_path,
                rank,
                num_shard,
                master_addr,
                master_port,
956
957
                huggingface_hub_cache,
                weights_cache_override,
958
                disable_custom_kernels,
959
960
                watermark_gamma,
                watermark_delta,
961
                enable_cuda_graphs,
962
                cuda_memory_fraction,
Nicolas Patry's avatar
Nicolas Patry committed
963
964
                rope_scaling,
                rope_factor,
965
                otlp_endpoint,
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
966
967
968
969
970
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
                status_sender,
                shutdown,
                shutdown_sender,
            )
        });
    }
    drop(shutdown_sender);

    // Wait for shard to start
    let mut shard_ready = 0;
    while running.load(Ordering::SeqCst) {
        match status_receiver.try_recv() {
            Ok(ShardStatus::Ready) => {
                shard_ready += 1;
                if shard_ready == num_shard {
                    break;
                }
            }
            Err(TryRecvError::Empty) => {
                sleep(Duration::from_millis(100));
            }
987
            Ok(ShardStatus::Failed(rank)) => {
988
                tracing::error!("Shard {rank} failed to start");
989
                shutdown_shards(shutdown, shutdown_receiver);
990
                return Err(LauncherError::ShardCannotStart);
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
991
992
993
            }
            Err(TryRecvError::Disconnected) => {
                tracing::error!("Shard status channel disconnected");
994
                shutdown_shards(shutdown, shutdown_receiver);
995
                return Err(LauncherError::ShardDisconnected);
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
996
997
998
            }
        }
    }
999
1000
    Ok(())
}
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1001

1002
1003
1004
1005
1006
1007
1008
1009
1010
1011
1012
1013
fn compute_type(num_shard: usize) -> Option<String> {
    let output = Command::new("nvidia-smi")
        .args(["--query-gpu=gpu_name", "--format=csv"])
        .output()
        .ok()?;
    let output = String::from_utf8(output.stdout).ok()?;
    let fullname = output.split('\n').nth(1)?;
    let cardname = fullname.replace(' ', "-").to_lowercase();
    let compute_type = format!("{num_shard}-{cardname}");
    Some(compute_type)
}

1014
fn spawn_webserver(
1015
    num_shard: usize,
1016
    args: Args,
1017
    shutdown: Arc<AtomicBool>,
1018
    shutdown_receiver: &mpsc::Receiver<()>,
1019
) -> Result<Child, LauncherError> {
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1020
1021
1022
    // All shard started
    // Start webserver
    tracing::info!("Starting Webserver");
OlivierDehaene's avatar
OlivierDehaene committed
1023
    let mut router_args = vec![
1024
        "--max-concurrent-requests".to_string(),
1025
        args.max_concurrent_requests.to_string(),
1026
        "--max-best-of".to_string(),
1027
        args.max_best_of.to_string(),
1028
        "--max-stop-sequences".to_string(),
1029
        args.max_stop_sequences.to_string(),
Nicolas Patry's avatar
Nicolas Patry committed
1030
1031
        "--max-top-n-tokens".to_string(),
        args.max_top_n_tokens.to_string(),
1032
        "--max-input-length".to_string(),
1033
        args.max_input_length.to_string(),
1034
        "--max-total-tokens".to_string(),
1035
        args.max_total_tokens.to_string(),
1036
1037
        "--max-batch-prefill-tokens".to_string(),
        args.max_batch_prefill_tokens.to_string(),
1038
        "--waiting-served-ratio".to_string(),
1039
        args.waiting_served_ratio.to_string(),
1040
        "--max-waiting-tokens".to_string(),
1041
        args.max_waiting_tokens.to_string(),
1042
1043
        "--validation-workers".to_string(),
        args.validation_workers.to_string(),
1044
1045
        "--hostname".to_string(),
        args.hostname.to_string(),
1046
        "--port".to_string(),
1047
        args.port.to_string(),
1048
        "--master-shard-uds-path".to_string(),
1049
        format!("{}-0", args.shard_uds_path),
1050
        "--tokenizer-name".to_string(),
1051
        args.model_id,
1052
1053
    ];

1054
1055
1056
1057
1058
1059
    // Tokenizer config path
    if let Some(ref tokenizer_config_path) = args.tokenizer_config_path {
        router_args.push("--tokenizer-config-path".to_string());
        router_args.push(tokenizer_config_path.to_string());
    }

1060
1061
1062
1063
1064
1065
    // Model optional max batch total tokens
    if let Some(max_batch_total_tokens) = args.max_batch_total_tokens {
        router_args.push("--max-batch-total-tokens".to_string());
        router_args.push(max_batch_total_tokens.to_string());
    }

1066
1067
1068
1069
1070
1071
    // Router optional max batch size
    if let Some(max_batch_size) = args.max_batch_size {
        router_args.push("--max-batch-size".to_string());
        router_args.push(max_batch_size.to_string());
    }

1072
1073
    // Model optional revision
    if let Some(ref revision) = args.revision {
OlivierDehaene's avatar
OlivierDehaene committed
1074
1075
        router_args.push("--revision".to_string());
        router_args.push(revision.to_string())
1076
1077
    }

1078
    if args.json_output {
OlivierDehaene's avatar
OlivierDehaene committed
1079
        router_args.push("--json-output".to_string());
1080
1081
    }

1082
    // OpenTelemetry
1083
    if let Some(otlp_endpoint) = args.otlp_endpoint {
OlivierDehaene's avatar
OlivierDehaene committed
1084
1085
        router_args.push("--otlp-endpoint".to_string());
        router_args.push(otlp_endpoint);
1086
1087
1088
1089
    }

    // CORS origins
    for origin in args.cors_allow_origin.into_iter() {
OlivierDehaene's avatar
OlivierDehaene committed
1090
1091
        router_args.push("--cors-allow-origin".to_string());
        router_args.push(origin);
1092
1093
    }

1094
1095
    // Ngrok
    if args.ngrok {
OlivierDehaene's avatar
OlivierDehaene committed
1096
1097
        router_args.push("--ngrok".to_string());
        router_args.push("--ngrok-authtoken".to_string());
1098
1099
1100
        router_args.push(args.ngrok_authtoken.unwrap());
        router_args.push("--ngrok-edge".to_string());
        router_args.push(args.ngrok_edge.unwrap());
1101
1102
    }

1103
    // Copy current process env
OlivierDehaene's avatar
OlivierDehaene committed
1104
    let mut envs: Vec<(OsString, OsString)> = env::vars_os().collect();
1105

1106
1107
    // Parse Inference API token
    if let Ok(api_token) = env::var("HF_API_TOKEN") {
OlivierDehaene's avatar
OlivierDehaene committed
1108
        envs.push(("HUGGING_FACE_HUB_TOKEN".into(), api_token.into()))
1109
    };
1110

1111
1112
1113
1114
1115
1116
1117
    // Parse Compute type
    if let Ok(compute_type) = env::var("COMPUTE_TYPE") {
        envs.push(("COMPUTE_TYPE".into(), compute_type.into()))
    } else if let Some(compute_type) = compute_type(num_shard) {
        envs.push(("COMPUTE_TYPE".into(), compute_type.into()))
    }

1118
    let mut webserver = match Command::new("text-generation-router")
OlivierDehaene's avatar
OlivierDehaene committed
1119
1120
        .args(router_args)
        .envs(envs)
1121
1122
1123
1124
1125
        .stdout(Stdio::piped())
        .stderr(Stdio::piped())
        .process_group(0)
        .spawn()
    {
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1126
1127
        Ok(p) => p,
        Err(err) => {
1128
            tracing::error!("Failed to start webserver: {}", err);
1129
1130
1131
            if err.kind() == io::ErrorKind::NotFound {
                tracing::error!("text-generation-router not found in PATH");
                tracing::error!("Please install it with `make install-router`")
1132
1133
            } else {
                tracing::error!("{}", err);
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1134
            }
1135

1136
            shutdown_shards(shutdown, shutdown_receiver);
1137
            return Err(LauncherError::WebserverCannotStart);
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1138
1139
1140
        }
    };

1141
1142
1143
    // Redirect STDOUT and STDERR to the console
    let webserver_stdout = webserver.stdout.take().unwrap();
    let webserver_stderr = webserver.stderr.take().unwrap();
1144
1145

    thread::spawn(move || {
1146
1147
        let stdout = BufReader::new(webserver_stdout);
        let stderr = BufReader::new(webserver_stderr);
1148
        for line in stdout.lines() {
1149
            println!("{}", line.unwrap());
1150
        }
1151
1152
        for line in stderr.lines() {
            println!("{}", line.unwrap());
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1153
        }
1154
1155
1156
    });
    Ok(webserver)
}
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1157

OlivierDehaene's avatar
OlivierDehaene committed
1158
1159
1160
1161
1162
1163
1164
1165
1166
1167
1168
1169
1170
1171
1172
1173
1174
1175
1176
1177
1178
1179
1180
1181
1182
fn terminate(process_name: &str, mut process: Child, timeout: Duration) -> io::Result<ExitStatus> {
    tracing::info!("Terminating {process_name}");

    let terminate_time = Instant::now();
    signal::kill(Pid::from_raw(process.id() as i32), Signal::SIGTERM).unwrap();

    tracing::info!("Waiting for {process_name} to gracefully shutdown");

    while terminate_time.elapsed() < timeout {
        if let Some(status) = process.try_wait()? {
            tracing::info!("{process_name} terminated");
            return Ok(status);
        }
        sleep(Duration::from_millis(100));
    }

    tracing::info!("Killing {process_name}");

    process.kill()?;
    let exit_status = process.wait()?;

    tracing::info!("{process_name} killed");
    Ok(exit_status)
}

1183
1184
fn main() -> Result<(), LauncherError> {
    // Pattern match configuration
1185
    let args: Args = Args::parse();
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1186

1187
1188
1189
1190
    // Filter events with LOG_LEVEL
    let env_filter =
        EnvFilter::try_from_env("LOG_LEVEL").unwrap_or_else(|_| EnvFilter::new("info"));

1191
    if args.json_output {
1192
1193
1194
1195
        tracing_subscriber::fmt()
            .with_env_filter(env_filter)
            .json()
            .init();
1196
    } else {
1197
1198
1199
1200
        tracing_subscriber::fmt()
            .with_env_filter(env_filter)
            .compact()
            .init();
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1201
1202
    }

1203
1204
1205
1206
1207
    if args.env {
        let env_runtime = env_runtime::Env::new();
        tracing::info!("{}", env_runtime);
    }

1208
1209
    tracing::info!("{:?}", args);

1210
1211
1212
1213
1214
1215
1216
1217
1218
1219
1220
1221
    // Validate args
    if args.max_input_length >= args.max_total_tokens {
        return Err(LauncherError::ArgumentValidation(
            "`max_input_length` must be < `max_total_tokens`".to_string(),
        ));
    }
    if args.max_input_length as u32 > args.max_batch_prefill_tokens {
        return Err(LauncherError::ArgumentValidation(format!(
            "`max_batch_prefill_tokens` must be >= `max_input_length`. Given: {} and {}",
            args.max_batch_prefill_tokens, args.max_input_length
        )));
    }
1222

1223
1224
1225
1226
1227
    if args.validation_workers == 0 {
        return Err(LauncherError::ArgumentValidation(
            "`validation_workers` must be > 0".to_string(),
        ));
    }
1228
1229
1230
1231
1232
1233
    if args.trust_remote_code {
        tracing::warn!(
            "`trust_remote_code` is set. Trusting that model `{}` do not contain malicious code.",
            args.model_id
        );
    }
1234
1235

    let num_shard = find_num_shards(args.sharded, args.num_shard)?;
1236
1237
    if num_shard > 1 {
        tracing::info!("Sharding model on {num_shard} processes");
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1238
1239
    }

1240
1241
1242
1243
1244
1245
1246
1247
1248
1249
1250
1251
1252
1253
1254
    if let Some(ref max_batch_total_tokens) = args.max_batch_total_tokens {
        if args.max_batch_prefill_tokens > *max_batch_total_tokens {
            return Err(LauncherError::ArgumentValidation(format!(
                "`max_batch_prefill_tokens` must be <= `max_batch_total_tokens`. Given: {} and {}",
                args.max_batch_prefill_tokens, max_batch_total_tokens
            )));
        }
        if args.max_total_tokens as u32 > *max_batch_total_tokens {
            return Err(LauncherError::ArgumentValidation(format!(
                "`max_total_tokens` must be <= `max_batch_total_tokens`. Given: {} and {}",
                args.max_total_tokens, max_batch_total_tokens
            )));
        }
    }

1255
1256
1257
1258
1259
1260
1261
1262
1263
1264
1265
1266
1267
1268
    if args.ngrok {
        if args.ngrok_authtoken.is_none() {
            return Err(LauncherError::ArgumentValidation(
                "`ngrok-authtoken` must be set when using ngrok tunneling".to_string(),
            ));
        }

        if args.ngrok_edge.is_none() {
            return Err(LauncherError::ArgumentValidation(
                "`ngrok-edge` must be set when using ngrok tunneling".to_string(),
            ));
        }
    }

1269
1270
1271
1272
1273
1274
1275
    // Signal handler
    let running = Arc::new(AtomicBool::new(true));
    let r = running.clone();
    ctrlc::set_handler(move || {
        r.store(false, Ordering::SeqCst);
    })
    .expect("Error setting Ctrl-C handler");
1276

1277
    // Download and convert model weights
1278
    download_convert_model(&args, running.clone())?;
1279

OlivierDehaene's avatar
OlivierDehaene committed
1280
1281
1282
1283
1284
    if !running.load(Ordering::SeqCst) {
        // Launcher was asked to stop
        return Ok(());
    }

1285
    // Shared shutdown bool
1286
    let shutdown = Arc::new(AtomicBool::new(false));
1287
1288
1289
    // Shared shutdown channel
    // When shutting down, the main thread will wait for all senders to be dropped
    let (shutdown_sender, shutdown_receiver) = mpsc::channel();
1290

1291
1292
    // Shared channel to track shard status
    let (status_sender, status_receiver) = mpsc::channel();
1293

1294
1295
1296
1297
1298
1299
1300
1301
1302
1303
    spawn_shards(
        num_shard,
        &args,
        shutdown.clone(),
        &shutdown_receiver,
        shutdown_sender,
        &status_receiver,
        status_sender,
        running.clone(),
    )?;
1304

1305
1306
1307
1308
1309
    // We might have received a termination signal
    if !running.load(Ordering::SeqCst) {
        shutdown_shards(shutdown, &shutdown_receiver);
        return Ok(());
    }
1310

1311
1312
    let mut webserver = spawn_webserver(num_shard, args, shutdown.clone(), &shutdown_receiver)
        .map_err(|err| {
OlivierDehaene's avatar
OlivierDehaene committed
1313
1314
1315
            shutdown_shards(shutdown.clone(), &shutdown_receiver);
            err
        })?;
1316
1317
1318
1319
1320

    // Default exit code
    let mut exit_code = Ok(());

    while running.load(Ordering::SeqCst) {
1321
        if let Ok(ShardStatus::Failed(rank)) = status_receiver.try_recv() {
OlivierDehaene's avatar
OlivierDehaene committed
1322
            tracing::error!("Shard {rank} crashed");
1323
1324
1325
1326
            exit_code = Err(LauncherError::ShardFailed);
            break;
        };

1327
        match webserver.try_wait().unwrap() {
1328
1329
1330
1331
1332
1333
1334
1335
1336
            Some(_) => {
                tracing::error!("Webserver Crashed");
                shutdown_shards(shutdown, &shutdown_receiver);
                return Err(LauncherError::WebserverFailed);
            }
            None => {
                sleep(Duration::from_millis(100));
            }
        };
1337
    }
1338
1339

    // Graceful termination
OlivierDehaene's avatar
OlivierDehaene committed
1340
    terminate("webserver", webserver, Duration::from_secs(90)).unwrap();
1341
1342
1343
    shutdown_shards(shutdown, &shutdown_receiver);

    exit_code
1344
}