main.rs 79.2 KB
Newer Older
1
use clap::{Parser, ValueEnum};
Nicolas Patry's avatar
Nicolas Patry committed
2
3
4
5
use hf_hub::{
    api::sync::{Api, ApiBuilder},
    Repo, RepoType,
};
6
7
use nix::sys::signal::{self, Signal};
use nix::unistd::Pid;
8
use regex::Regex;
9
use serde::Deserialize;
Nicolas Patry's avatar
Nicolas Patry committed
10
use std::env;
11
use std::ffi::OsString;
12
use std::io::{BufRead, BufReader};
13
use std::os::unix::process::{CommandExt, ExitStatusExt};
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
14
use std::path::Path;
OlivierDehaene's avatar
OlivierDehaene committed
15
use std::process::{Child, Command, ExitStatus, Stdio};
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
16
17
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::mpsc::TryRecvError;
18
use std::sync::{mpsc, Arc};
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
19
20
21
use std::thread;
use std::thread::sleep;
use std::time::{Duration, Instant};
22
23
24
25
use std::{
    fs, io,
    io::{Read, Write},
};
26
use thiserror::Error;
27
use tracing_subscriber::{filter::LevelFilter, EnvFilter};
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
28

29
mod env_runtime;
30
mod gpu;
31

Nicolas Patry's avatar
Nicolas Patry committed
32
fn compute_optimal(config: Option<&Config>, compute: Option<&ComputeType>) -> Option<usize> {
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
    let config = config?;
    let compute = compute?;
    let f16_max_compute = compute.f16_flop()?;
    let model_compute = config.flop()?;
    tracing::debug!(
        "Max compute {} model compute {}",
        human_size(f16_max_compute as usize, "flop"),
        human_size(model_compute as usize, "flop")
    );
    let optimal_size = (f16_max_compute / model_compute) as usize;
    if optimal_size > 100 {
        // Ignore calculations that's too low
        // Most likely an error
        Some(optimal_size)
    } else {
        None
    }
}

fn human_size(size: usize, suffix: &str) -> String {
    let mut size: f64 = size as f64;
    let mut p = "";
    for prefix in ["", "K", "M", "G", "T"] {
        p = prefix;
        if size > 1_000.0 {
            size /= 1_000.0;
Nicolas Patry's avatar
Nicolas Patry committed
59
        } else {
60
            break;
Nicolas Patry's avatar
Nicolas Patry committed
61
        }
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
    }
    format!("{size:.2}{p}{suffix}")
}

fn vram_maximum(
    config: Option<&Config>,
    compute: Option<&ComputeType>,
    memory_fraction: f32,
) -> Option<usize> {
    let config = config?;
    let compute = compute?;
    let available = compute.vram(memory_fraction)?;
    let model = config.model_vram()?;
    let token_vram = config.token_vram()?;
    if let Some(vram) = available.checked_sub(model) {
        let tokens_allowed = vram / token_vram;
        tracing::debug!(
        "Available vram {}: model needs {}, every tokens requires {}, maximum allocatable tokens {tokens_allowed}",
        human_size(available, "B"),
        human_size(model, "B"),
        human_size(token_vram, "B"),
    );
        Some(tokens_allowed)
Nicolas Patry's avatar
Nicolas Patry committed
85
    } else {
86
87
88
89
90
        tracing::warn!(
            "Not enough VRAM to run the model: Available: {} - Model {}.",
            human_size(available, "B"),
            human_size(model, "B")
        );
Nicolas Patry's avatar
Nicolas Patry committed
91
92
93
94
        None
    }
}

95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
fn get_config(
    model_id: &str,
    revision: &Option<String>,
) -> Result<Config, Box<dyn std::error::Error>> {
    let mut path = std::path::Path::new(model_id).to_path_buf();
    let model_id = model_id.to_string();
    let filename = if !path.exists() {
        // Assume it's a hub id

        let api = if let Ok(token) = std::env::var("HF_TOKEN") {
            // env variable has precedence over on file token.
            ApiBuilder::new().with_token(Some(token)).build()?
        } else {
            Api::new()?
        };
        let repo = if let Some(ref revision) = revision {
            api.repo(Repo::with_revision(
                model_id,
                RepoType::Model,
                revision.to_string(),
            ))
        } else {
            api.model(model_id)
        };
        repo.get("config.json")?
    } else {
        path.push("config.json");
        path
    };

    let content = std::fs::read_to_string(filename)?;
    let config: RawConfig = serde_json::from_str(&content)?;

    let config: Config = config.into();
    Ok(config)
}

fn resolve_attention(config: &Option<Config>, lora_adapters: &Option<String>) -> (String, String) {
133
    let compute_capability = gpu::get_cuda_capability();
134
    let mut prefix_caching: Option<String> = std::env::var("PREFIX_CACHING").ok();
135
136
137
138
139
140
141
142
143
144
145
    let mut attention: Option<String> = std::env::var("ATTENTION").ok();
    if let Some(config) = config {
        if prefix_caching.is_none() {
            if config.vision_config.is_some() {
                tracing::info!("Disabling prefix caching because of VLM model");
                prefix_caching = Some("0".to_string());
            } else if config.is_encoder_decoder {
                tracing::info!("Disabling prefix caching because of seq2seq model");
                prefix_caching = Some("0".to_string());
            }
        }
146
147
148
149
150
151
152

        let fallback_attention = if matches!(compute_capability, Some((major, _)) if major < 8) {
            "paged"
        } else {
            "flashdecoding"
        };

153
154
155
156
157
158
159
        match config.head_dim {
            Some(h) if h == 64 || h == 128 || h == 256 => {
                if lora_adapters.is_some() && prefix_caching.is_none() {
                    tracing::info!("Disabling prefix caching because of lora adapters");
                    prefix_caching = Some("0".to_string());
                }
                match config.model_type.as_deref() {
Daniël de Kok's avatar
Daniël de Kok committed
160
                    Some("falcon") | Some("deepseek_v2") => {
161
162
163
164
                        // Required because gemma2 needs bfloat16 which is not supported by
                        // flashinfer ?
                        if attention.is_none() {
                            tracing::info!(
165
                                "Forcing attention to '{fallback_attention}' because model {} requires it",
166
167
                                config.model_type.as_ref().unwrap()
                            );
168
169
170
171
172
                            attention = Some(fallback_attention.to_string());
                        }
                        if fallback_attention == "paged" && prefix_caching.is_none() {
                            tracing::info!("Disabling prefix caching because it is not supported with 'paged' attention");
                            prefix_caching = Some("0".to_string());
173
174
175
176
177
178
179
180
                        }
                    }
                    Some("t5") => {}
                    _ => {}
                }
            }
            _ => {
                if attention.is_none() {
181
182
                    tracing::info!("Forcing attention to '{fallback_attention}' because head dim is not supported by flashinfer, also disabling prefix caching");
                    attention = Some(fallback_attention.to_string());
183
184
185
186
187
188
189
                }
                if prefix_caching.is_none() {
                    prefix_caching = Some("0".to_string());
                }
            }
        }
    }
190
191
192
193
    if attention == Some("paged".to_string()) && prefix_caching.is_none() {
        tracing::info!("Disabling prefix caching on paged attention");
        prefix_caching = Some("0".to_string());
    }
194

195
    let attention = attention.unwrap_or("flashinfer".to_string());
196
197
    let prefix_caching = prefix_caching.unwrap_or("true".to_string());

198
199
200
    (prefix_caching, attention)
}

201
#[derive(Deserialize)]
202
struct RawConfig {
203
    max_position_embeddings: Option<usize>,
204
    n_positions: Option<usize>,
205
    model_type: Option<String>,
206
    max_seq_len: Option<usize>,
207
    quantization_config: Option<QuantizationConfig>,
208
209
    n_embd: Option<usize>,
    hidden_size: Option<usize>,
Nicolas Patry's avatar
Nicolas Patry committed
210
    intermediate_size: Option<usize>,
211
    num_attention_heads: Option<usize>,
Nicolas Patry's avatar
Nicolas Patry committed
212
213
    num_key_value_heads: Option<usize>,
    num_hidden_layers: Option<usize>,
214
215
216
    head_dim: Option<usize>,
    vision_config: Option<VisionConfig>,
    is_encoder_decoder: Option<bool>,
Nicolas Patry's avatar
Nicolas Patry committed
217
    #[serde(rename = "num_experts_per_tok")]
Nicolas Patry's avatar
Nicolas Patry committed
218
219
220
    num_experts_per_token: Option<usize>,
    #[serde(rename = "n_shared_experts")]
    num_shared_experts: Option<usize>,
221
222
223
    #[serde(rename = "num_local_experts")]
    num_experts: Option<usize>,
    vocab_size: Option<usize>,
224
225
226
227
228
}

#[derive(Deserialize)]
struct QuantizationConfig {
    quant_method: Option<Quantization>,
229
230
}

Nicolas Patry's avatar
Nicolas Patry committed
231
#[derive(Debug, Deserialize)]
232
233
struct VisionConfig {}

Nicolas Patry's avatar
Nicolas Patry committed
234
#[derive(Debug, Deserialize)]
235
236
struct Config {
    max_position_embeddings: Option<usize>,
237
    quantize: Option<Quantization>,
238
    head_dim: Option<usize>,
Nicolas Patry's avatar
Nicolas Patry committed
239
240
241
242
243
    num_heads: Option<usize>,
    num_kv_heads: Option<usize>,
    num_layers: Option<usize>,
    intermediate_size: Option<usize>,
    hidden_size: Option<usize>,
244
245
246
    model_type: Option<String>,
    vision_config: Option<VisionConfig>,
    is_encoder_decoder: bool,
Nicolas Patry's avatar
Nicolas Patry committed
247
248
    num_experts_per_token: usize,
    num_shared_experts: usize,
249
250
    num_experts: usize,
    vocab_size: Option<usize>,
Nicolas Patry's avatar
Nicolas Patry committed
251
252
253
254
255
256
257
258
259
260
261
262
263
}

impl Config {
    fn flop(&self) -> Option<u64> {
        if self.vision_config.is_some() {
            // VLM are much harder to predict and VRAM requirements
            // Are more complex.
            return None;
        }
        let num_heads = self.num_heads? as u64;
        let num_kv_heads = self.num_kv_heads? as u64;
        let head_dim = self.head_dim? as u64;
        let hidden_size = self.hidden_size? as u64;
Nicolas Patry's avatar
Nicolas Patry committed
264
265
266
        let intermediate_size = (self.intermediate_size?
            * (self.num_experts_per_token + self.num_shared_experts))
            as u64;
Nicolas Patry's avatar
Nicolas Patry committed
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
        let num_layers = self.num_layers? as u64;

        let q_flops = 2 * num_heads * head_dim * hidden_size;
        let k_flops = 2 * num_kv_heads * head_dim * hidden_size;
        let v_flops = 2 * num_kv_heads * head_dim * hidden_size;
        let attn_flops = 2 * num_heads * head_dim * hidden_size;
        let o_flops = 2 * num_heads * head_dim * hidden_size;
        let attn_layer_flops = q_flops + k_flops + v_flops + attn_flops + o_flops;

        let gate_up_down_flops = 2 * 3 * hidden_size * intermediate_size;

        let layer_flops = attn_layer_flops + gate_up_down_flops;
        let total = layer_flops * num_layers;
        Some(total)
    }
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324

    fn kv_vram_per_tok(&self) -> Option<usize> {
        if self.quantize.is_some() {
            // TODO handle quantization
            return None;
        }
        // 2 for key and values
        // 2 for f16 dtype?
        Some(self.num_kv_heads? * 2 * self.head_dim? * 2 * self.num_layers?)
    }

    fn mlp_vram_per_tok(&self) -> Option<usize> {
        // TODO handle quantization
        // TODO This calculation depends on the actual implementation
        let dtype_size = 2;
        let mlp_size = self.intermediate_size?;
        // calculation is overshooting here.
        // Coming from here: https://github.com/vllm-project/vllm/blob/d1c2e15eb31ef12e688ce0cb71895f88eaf4cd4f/vllm/model_executor/layers/fused_moe/fused_moe.py#L618-L624
        Some((mlp_size + mlp_size / 2) * self.num_experts * dtype_size * 3)
    }

    fn token_vram(&self) -> Option<usize> {
        let kv = self.kv_vram_per_tok()?;
        let mlp_intermediary = self.mlp_vram_per_tok()?;
        let per_tok = kv + mlp_intermediary;
        Some(per_tok)
    }

    fn model_vram(&self) -> Option<usize> {
        let attn_vram = (self.num_heads? + 2 * self.num_kv_heads?) * self.head_dim?;
        let o_vram = self.num_heads? * self.head_dim? * self.hidden_size?;
        // gate + up + down = 3
        let mlp_vram = 3 * self.intermediate_size? * self.num_experts * self.hidden_size?;
        let layer_vram = mlp_vram + attn_vram + o_vram;
        let vocab = self.hidden_size? * self.vocab_size?;
        let params = layer_vram * self.num_layers? + 2 * vocab;
        let dtype_size = 2;
        if self.quantize.is_some() {
            // TODO handle quantization
            return None;
        }
        Some(params * dtype_size)
    }
325
326
327
328
329
330
331
332
}

impl From<RawConfig> for Config {
    fn from(other: RawConfig) -> Self {
        let max_position_embeddings = other
            .max_position_embeddings
            .or(other.max_seq_len)
            .or(other.n_positions);
333
        let quantize = other.quantization_config.and_then(|q| q.quant_method);
Nicolas Patry's avatar
Nicolas Patry committed
334
335
336
337
338
        let hidden_size = other.hidden_size.or(other.n_embd);
        let head_dim = other
            .head_dim
            .or_else(|| match (hidden_size, other.num_attention_heads) {
                (Some(hidden_size), Some(num_attention_heads))
339
340
341
342
343
                    if hidden_size % num_attention_heads == 0 =>
                {
                    Some(hidden_size / num_attention_heads)
                }
                _ => None,
Nicolas Patry's avatar
Nicolas Patry committed
344
345
346
347
348
            });
        let num_heads = other.num_attention_heads;
        let num_layers = other.num_hidden_layers;
        let num_kv_heads = other.num_key_value_heads.or(other.num_attention_heads);
        let intermediate_size = other.intermediate_size;
349
350
351
        let model_type = other.model_type;
        let vision_config = other.vision_config;
        let is_encoder_decoder = other.is_encoder_decoder.unwrap_or(false);
Nicolas Patry's avatar
Nicolas Patry committed
352
353
        let num_experts_per_token = other.num_experts_per_token.unwrap_or(1);
        let num_shared_experts = other.num_shared_experts.unwrap_or(0);
354
355
        let num_experts = other.num_experts.unwrap_or(1);
        let vocab_size = other.vocab_size;
356
357
        Config {
            max_position_embeddings,
358
            quantize,
359
360
361
362
            head_dim,
            model_type,
            vision_config,
            is_encoder_decoder,
Nicolas Patry's avatar
Nicolas Patry committed
363
364
365
366
367
            hidden_size,
            num_heads,
            num_kv_heads,
            intermediate_size,
            num_layers,
Nicolas Patry's avatar
Nicolas Patry committed
368
369
            num_experts_per_token,
            num_shared_experts,
370
371
            num_experts,
            vocab_size,
372
373
374
375
        }
    }
}

376
377
#[derive(Clone, Copy, Debug, ValueEnum, Deserialize)]
#[serde(rename_all = "kebab-case")]
378
enum Quantization {
379
    /// 4 bit quantization. Requires a specific AWQ quantized model:
380
    ///   <https://hf.co/models?search=awq>.
381
    /// Should replace GPTQ models wherever possible because of the better latency
382
    Awq,
383
384
    /// Compressed tensors, which can be a mixture of different quantization methods.
    CompressedTensors,
385
386
    /// 8 bit quantization, doesn't require specific model.
    /// Should be a drop-in replacement to bitsandbytes with much better performance.
387
    /// Kernels are from <https://github.com/NetEase-FuXi/EETQ.git>
388
    Eetq,
389
390
391
392
    /// Variable bit quantization. Requires a specific EXL2 quantized model:
    /// <https://hf.co/models?search=exl2>. Requires exllama2 kernels and does
    /// not support tensor parallelism (num_shard > 1).
    Exl2,
393
    /// 4 bit quantization. Requires a specific GTPQ quantized model: <https://hf.co/models?search=gptq>.
394
    /// text-generation-inference will use exllama (faster) kernels wherever possible, and use
395
396
397
    /// triton kernel (wider support) when it's not.
    /// AWQ has faster kernels.
    Gptq,
398
399
    /// 4 bit quantization. Requires a specific Marlin quantized model: <https://hf.co/models?search=marlin>.
    Marlin,
400
401
    /// 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.
402
403
404
405
    // #[deprecated(
    //     since = "1.1.0",
    //     note = "Use `eetq` instead, which provides better latencies overall and is drop-in in most cases"
    // )]
406
    Bitsandbytes,
407
408
    /// 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.
409
    BitsandbytesNf4,
410
411
    /// Bitsandbytes 4bit. nf4 should be preferred in most cases but maybe this one has better
    /// perplexity performance for you model
412
    BitsandbytesFp4,
Nicolas Patry's avatar
Nicolas Patry committed
413
414
415
416
417
    /// [FP8](https://developer.nvidia.com/blog/nvidia-arm-and-intel-publish-fp8-specification-for-standardization-as-an-interchange-format-for-ai/) (e4m3) works on H100 and above
    /// This dtype has native ops should be the fastest if available.
    /// This is currently not the fastest because of local unpacking + padding to satisfy matrix
    /// multiplication limitations.
    Fp8,
418
419
420
421
422
423
}

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 {
424
425
            #[allow(deprecated)]
            // Use `eetq` instead, which provides better latencies overall and is drop-in in most cases
426
427
428
            Quantization::Bitsandbytes => {
                write!(f, "bitsandbytes")
            }
429
            Quantization::BitsandbytesNf4 => {
Nicolas Patry's avatar
Nicolas Patry committed
430
431
                write!(f, "bitsandbytes-nf4")
            }
432
            Quantization::BitsandbytesFp4 => {
Nicolas Patry's avatar
Nicolas Patry committed
433
434
                write!(f, "bitsandbytes-fp4")
            }
435
436
437
            Quantization::Exl2 => {
                write!(f, "exl2")
            }
438
439
440
            Quantization::Gptq => {
                write!(f, "gptq")
            }
441
442
443
            Quantization::Marlin => {
                write!(f, "marlin")
            }
444
445
446
            Quantization::Awq => {
                write!(f, "awq")
            }
447
448
449
            Quantization::CompressedTensors => {
                write!(f, "compressed-tensors")
            }
450
451
452
            Quantization::Eetq => {
                write!(f, "eetq")
            }
Nicolas Patry's avatar
Nicolas Patry committed
453
454
455
            Quantization::Fp8 => {
                write!(f, "fp8")
            }
456
457
458
459
        }
    }
}

460
461
462
#[derive(Clone, Copy, Debug, ValueEnum)]
enum Dtype {
    Float16,
463
    #[clap(name = "bfloat16")]
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
    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")
            }
        }
    }
}

481
482
#[derive(Clone, Copy, Debug, ValueEnum)]
enum KVCacheDtype {
483
484
485
    #[clap(name = "fp8_e4m3fn")]
    Fp8e4m3fn,

486
487
488
489
490
491
492
    #[clap(name = "fp8_e5m2")]
    Fp8e5m2,
}

impl std::fmt::Display for KVCacheDtype {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
493
494
495
            KVCacheDtype::Fp8e4m3fn => {
                write!(f, "fp8_e4m3fn")
            }
496
497
498
499
500
501
502
            KVCacheDtype::Fp8e5m2 => {
                write!(f, "fp8_e5m2")
            }
        }
    }
}

Nicolas Patry's avatar
Nicolas Patry committed
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
#[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")
            }
        }
    }
}

523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
#[derive(Clone, Copy, Debug, ValueEnum)]
pub enum UsageStatsLevel {
    /// Default option, usage statistics are collected anonymously
    On,
    /// Disables all collection of usage statistics
    Off,
    /// Doesn't send the error stack trace or error type, but allows sending a crash event
    NoStack,
}

impl std::fmt::Display for UsageStatsLevel {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        // To keep in track with `server`.
        match self {
            UsageStatsLevel::On => {
                write!(f, "on")
            }
            UsageStatsLevel::Off => {
                write!(f, "off")
            }
            UsageStatsLevel::NoStack => {
                write!(f, "no-stack")
            }
        }
    }
}

Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
550
551
552
553
/// App Configuration
#[derive(Parser, Debug)]
#[clap(author, version, about, long_about = None)]
struct Args {
554
555
556
557
558
    /// 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
559
    #[clap(default_value = "bigscience/bloom-560m", long, env)]
560
    model_id: String,
561
562
563

    /// 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
564
    #[clap(long, env)]
565
    revision: Option<String>,
566

567
568
569
570
571
    /// The number of tokenizer workers used for payload validation and truncation inside the
    /// router.
    #[clap(default_value = "2", long, env)]
    validation_workers: usize,

572
    /// Whether to shard the model across multiple GPUs
573
574
    /// By default text-generation-inference will use all available GPUs to run
    /// the model. Setting it to `false` deactivates `num_shard`.
575
576
    #[clap(long, env)]
    sharded: Option<bool>,
577
578

    /// The number of shards to use if you don't want to use all GPUs on a given machine.
579
580
    /// 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
581
    /// launch 2 copies with 2 shard each on a given machine with 4 GPUs for instance.
582
583
    #[clap(long, env)]
    num_shard: Option<usize>,
584

585
586
587
588
589
    /// Quantization method to use for the model. It is not necessary to specify this option
    /// for pre-quantized models, since the quantization method is read from the model
    /// configuration.
    ///
    /// Marlin kernels will be used automatically for GPTQ/AWQ models.
590
591
    #[clap(long, env, value_enum)]
    quantize: Option<Quantization>,
592

Nicolas Patry's avatar
Nicolas Patry committed
593
594
595
596
597
598
599
    /// 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>,

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

604
605
    /// Specify the dtype for the key-value cache. When this option is not provided,
    /// the dtype of the model is used (typically `float16` or `bfloat16`). Currently
606
    /// the only supported value are `fp8_e4m3fn` and `fp8_e5m2` on CUDA.
607
608
609
    #[clap(long, env, value_enum)]
    kv_cache_dtype: Option<KVCacheDtype>,

610
611
612
613
614
615
    /// 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,

616
617
618
    /// 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
619
620
    #[clap(default_value = "128", long, env)]
    max_concurrent_requests: usize,
621
622
623
624

    /// 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
625
626
    #[clap(default_value = "2", long, env)]
    max_best_of: usize,
627
628
629
630
631
632

    /// 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.
633
634
    #[clap(default_value = "4", long, env)]
    max_stop_sequences: usize,
635

Nicolas Patry's avatar
Nicolas Patry committed
636
    /// This is the maximum allowed value for clients to set `top_n_tokens`.
637
    /// `top_n_tokens` is used to return information about the the `n` most likely
Nicolas Patry's avatar
Nicolas Patry committed
638
639
640
641
642
643
    /// 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,

644
645
646
647
    /// 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.
648
    /// Default to min(max_allocatable, max_position_embeddings) - 1
649
650
651
652
653
654
    #[clap(long, env)]
    max_input_tokens: Option<usize>,

    /// Legacy version of [`Args::max_input_tokens`].
    #[clap(long, env)]
    max_input_length: Option<usize>,
655
656
657
658
659
660
661
662
663

    /// 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.
664
    /// Default to min(max_allocatable, max_position_embeddings)
665
666
    #[clap(long, env)]
    max_total_tokens: Option<usize>,
667
668
669
670
671
672
673
674
675
676
677

    /// 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`.
678
    #[clap(default_value = "0.3", long, env)]
679
    waiting_served_ratio: f32,
680

681
682
683
    /// 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.
684
685
686
    /// Default to `max_input_tokens + 50` to give a bit of room.
    #[clap(long, env)]
    max_batch_prefill_tokens: Option<u32>,
687

688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
    /// **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.
705
706
    #[clap(long, env)]
    max_batch_total_tokens: Option<u32>,
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724

    /// 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.
725
726
    #[clap(default_value = "20", long, env)]
    max_waiting_tokens: usize,
727

728
729
730
731
732
    /// 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>,

733
734
    /// Specify the batch sizes to compute cuda graphs for.
    /// Use "0" to disable.
735
736
737
    /// Default = "1,2,4,8,16,32"
    #[clap(long, env, value_delimiter = ',')]
    cuda_graphs: Option<Vec<usize>>,
738

739
740
741
742
    /// The IP address to listen on
    #[clap(default_value = "0.0.0.0", long, env)]
    hostname: String,

743
    /// The port to listen on.
744
    #[clap(default_value = "3000", long, short, env)]
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
745
    port: u16,
746
747
748

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

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

    /// The address the master port will listen on. (setting used by torch distributed)
757
    #[clap(default_value = "29500", long, env)]
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
758
    master_port: usize,
759
760
761

    /// The location of the huggingface hub cache.
    /// Used to override the location if you want to provide a mounted disk for instance
762
    #[clap(long, env)]
763
    huggingface_hub_cache: Option<String>,
764
765
766

    /// The location of the huggingface hub cache.
    /// Used to override the location if you want to provide a mounted disk for instance
767
768
    #[clap(long, env)]
    weights_cache_override: Option<String>,
769
770
771
772
773

    /// 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.
774
    #[clap(long, env)]
775
    disable_custom_kernels: bool,
776

777
778
779
780
781
    /// 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
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
    /// 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>,

802
    /// Outputs the logs in JSON format (useful for telemetry)
803
    #[clap(long, env)]
804
    json_output: bool,
805

806
807
    #[clap(long, env)]
    otlp_endpoint: Option<String>,
808

809
810
811
    #[clap(default_value = "text-generation-inference.router", long, env)]
    otlp_service_name: String,

812
813
    #[clap(long, env)]
    cors_allow_origin: Vec<String>,
Erik Kaunismäki's avatar
Erik Kaunismäki committed
814
815
816
817

    #[clap(long, env)]
    api_key: Option<String>,

818
819
820
821
    #[clap(long, env)]
    watermark_gamma: Option<f32>,
    #[clap(long, env)]
    watermark_delta: Option<f32>,
822

823
824
825
826
827
828
829
830
    /// Enable ngrok tunneling
    #[clap(long, env)]
    ngrok: bool,

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

831
    /// ngrok edge
832
    #[clap(long, env)]
833
    ngrok_edge: Option<String>,
834

835
836
837
838
839
    /// 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>,

drbh's avatar
drbh committed
840
841
842
843
844
    /// Disable outlines grammar constrained generation.
    /// This is a feature that allows you to generate text that follows a specific grammar.
    #[clap(long, env)]
    disable_grammar_support: bool,

845
846
847
    /// Display a lot of information about your runtime environment
    #[clap(long, short, action)]
    env: bool,
848
849
850
851

    /// Control the maximum number of inputs that a client can send in a single request
    #[clap(default_value = "4", long, env)]
    max_client_batch_size: usize,
drbh's avatar
drbh committed
852
853
854
855
856

    /// Lora Adapters a list of adapter ids i.e. `repo/adapter1,repo/adapter2` to load during
    /// startup that will be available to callers via the `adapter_id` field in a request.
    #[clap(long, env)]
    lora_adapters: Option<String>,
857

858
859
860
861
862
    /// Control if anonymous usage stats are collected.
    /// Options are "on", "off" and "no-stack"
    /// Defaul is on.
    #[clap(default_value = "on", long, env)]
    usage_stats: UsageStatsLevel,
863
864
865
866
867
868

    /// Payload size limit in bytes
    ///
    /// Default is 2MB
    #[clap(default_value = "2000000", long, env)]
    payload_limit: usize,
Nicolas Patry's avatar
Nicolas Patry committed
869
870
871
872
873
874
875
876

    /// Enables prefill logprobs
    ///
    /// Logprobs in the prompt are deactivated by default because they consume
    /// a large amount of VRAM (especially for long prompts).
    /// Using this flag reallows users to ask for them.
    #[clap(long, env)]
    enable_prefill_logprobs: bool,
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
877
878
}

879
880
881
#[derive(Debug)]
enum ShardStatus {
    Ready,
882
    Failed(usize),
883
}
884

885
886
887
888
#[allow(clippy::too_many_arguments)]
fn shard_manager(
    model_id: String,
    revision: Option<String>,
889
    quantize: Option<Quantization>,
Nicolas Patry's avatar
Nicolas Patry committed
890
    speculate: Option<usize>,
891
    dtype: Option<Dtype>,
892
    kv_cache_dtype: Option<KVCacheDtype>,
893
    trust_remote_code: bool,
894
895
896
897
898
899
900
901
902
903
    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>,
904
    cuda_graphs: Vec<usize>,
905
    cuda_memory_fraction: f32,
Nicolas Patry's avatar
Nicolas Patry committed
906
907
    rope_scaling: Option<RopeScaling>,
    rope_factor: Option<f32>,
908
    max_total_tokens: Option<usize>,
909
    max_batch_size: Option<usize>,
910
    max_input_tokens: Option<usize>,
drbh's avatar
drbh committed
911
    lora_adapters: Option<String>,
Nicolas Patry's avatar
Nicolas Patry committed
912
    enable_prefill_logprobs: bool,
913
    otlp_endpoint: Option<String>,
914
    otlp_service_name: String,
915
    log_level: LevelFilter,
916
    status_sender: mpsc::Sender<ShardStatus>,
917
    shutdown: Arc<AtomicBool>,
918
919
    _shutdown_sender: mpsc::Sender<()>,
) {
920
921
922
    // Enter shard-manager tracing span
    let _span = tracing::span!(tracing::Level::INFO, "shard-manager", rank = rank).entered();

923
924
925
926
    // Get UDS path
    let uds_string = format!("{uds_path}-{rank}");
    let uds = Path::new(&uds_string);
    // Clean previous runs
927
928
929
    if uds.exists() {
        fs::remove_file(uds).unwrap();
    }
930
931

    // Process args
OlivierDehaene's avatar
OlivierDehaene committed
932
    let mut shard_args = vec![
933
934
935
936
937
        "serve".to_string(),
        model_id,
        "--uds-path".to_string(),
        uds_path,
        "--logger-level".to_string(),
938
        log_level.to_string().to_uppercase(),
939
940
941
        "--json-output".to_string(),
    ];

942
943
    // Activate trust remote code
    if trust_remote_code {
OlivierDehaene's avatar
OlivierDehaene committed
944
        shard_args.push("--trust-remote-code".to_string());
945
946
    }

947
948
    // Activate tensor parallelism
    if world_size > 1 {
OlivierDehaene's avatar
OlivierDehaene committed
949
        shard_args.push("--sharded".to_string());
950
951
    }

952
    if let Some(quantize) = quantize {
OlivierDehaene's avatar
OlivierDehaene committed
953
954
        shard_args.push("--quantize".to_string());
        shard_args.push(quantize.to_string())
955
    }
956

Nicolas Patry's avatar
Nicolas Patry committed
957
958
959
960
961
    if let Some(speculate) = speculate {
        shard_args.push("--speculate".to_string());
        shard_args.push(speculate.to_string())
    }

962
    if let Some(dtype) = dtype {
OlivierDehaene's avatar
OlivierDehaene committed
963
964
        shard_args.push("--dtype".to_string());
        shard_args.push(dtype.to_string())
965
966
    }

967
968
969
970
971
    if let Some(kv_cache_dtype) = kv_cache_dtype {
        shard_args.push("--kv-cache-dtype".to_string());
        shard_args.push(kv_cache_dtype.to_string())
    }

972
973
    // Model optional revision
    if let Some(revision) = revision {
OlivierDehaene's avatar
OlivierDehaene committed
974
975
        shard_args.push("--revision".to_string());
        shard_args.push(revision)
976
    }
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
977

Nicolas Patry's avatar
Nicolas Patry committed
978
979
980
981
982
983
    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)),
    };
984

985
    // OpenTelemetry Endpoint
986
    if let Some(otlp_endpoint) = otlp_endpoint {
OlivierDehaene's avatar
OlivierDehaene committed
987
988
        shard_args.push("--otlp-endpoint".to_string());
        shard_args.push(otlp_endpoint);
989
990
    }

991
992
993
994
    // OpenTelemetry Service Name
    shard_args.push("--otlp-service-name".to_string());
    shard_args.push(otlp_service_name);

995
    // In case we use sliding window, we may ignore the sliding in flash for some backends depending on the parameter.
996
997
998
999
    if let Some(max_input_tokens) = max_input_tokens {
        shard_args.push("--max-input-tokens".to_string());
        shard_args.push(max_input_tokens.to_string());
    }
1000

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

1004
1005
1006
    // Remove LOG_LEVEL if present
    envs.retain(|(name, _)| name != "LOG_LEVEL");

1007
    // Torch Distributed Env vars
OlivierDehaene's avatar
OlivierDehaene committed
1008
1009
1010
1011
    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()));
1012
    envs.push(("TORCH_NCCL_AVOID_RECORD_STREAMS".into(), "1".into()));
1013

1014
1015
1016
1017
1018
1019
    // CUDA memory fraction
    envs.push((
        "CUDA_MEMORY_FRACTION".into(),
        cuda_memory_fraction.to_string().into(),
    ));

1020
    // Safetensors load fast
OlivierDehaene's avatar
OlivierDehaene committed
1021
    envs.push(("SAFETENSORS_FAST_GPU".into(), "1".into()));
1022

1023
1024
1025
    // Disable progress bar
    envs.push(("HF_HUB_DISABLE_PROGRESS_BARS".into(), "1".into()));

1026
1027
    // 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
1028
    envs.push((
1029
1030
1031
1032
1033
1034
        "HF_HUB_ENABLE_HF_TRANSFER".into(),
        enable_hf_transfer.into(),
    ));

    // Parse Inference API token
    if let Ok(api_token) = env::var("HF_API_TOKEN") {
1035
        envs.push(("HF_TOKEN".into(), api_token.into()))
1036
1037
    };

Nicolas Patry's avatar
Nicolas Patry committed
1038
1039
1040
1041
1042
1043
1044
1045
1046
    // 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()));
    }

1047
1048
1049
1050
1051
1052
    if let Some(max_total_tokens) = max_total_tokens {
        envs.push((
            "MAX_TOTAL_TOKENS".into(),
            max_total_tokens.to_string().into(),
        ));
    }
1053
1054
1055
1056
    if let Some(max_batch_size) = max_batch_size {
        envs.push(("MAX_BATCH_SIZE".into(), max_batch_size.to_string().into()));
    }

drbh's avatar
drbh committed
1057
1058
1059
1060
1061
    // Lora Adapters
    if let Some(lora_adapters) = lora_adapters {
        envs.push(("LORA_ADAPTERS".into(), lora_adapters.into()));
    }

Nicolas Patry's avatar
Nicolas Patry committed
1062
1063
1064
1065
1066
    // Logprobs
    if enable_prefill_logprobs {
        envs.push(("REQUEST_LOGPROBS".into(), "1".into()));
    }

1067
1068
1069
    // 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
1070
        envs.push(("HUGGINGFACE_HUB_CACHE".into(), huggingface_hub_cache.into()));
1071
1072
1073
1074
1075
    };

    // 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
1076
        envs.push((
1077
1078
1079
1080
1081
            "WEIGHTS_CACHE_OVERRIDE".into(),
            weights_cache_override.into(),
        ));
    };

1082
    // Enable experimental support for cuda graphs
1083
1084
1085
1086
1087
1088
1089
1090
1091
1092
    if !cuda_graphs.is_empty() {
        envs.push((
            "CUDA_GRAPHS".into(),
            cuda_graphs
                .into_iter()
                .map(|c| c.to_string())
                .collect::<Vec<_>>()
                .join(",")
                .into(),
        ));
1093
1094
    }

1095
1096
    // If disable_custom_kernels is true, pass it to the shard as an env var
    if disable_custom_kernels {
OlivierDehaene's avatar
OlivierDehaene committed
1097
        envs.push(("DISABLE_CUSTOM_KERNELS".into(), "True".into()))
1098
1099
1100
1101
    }

    // Watermark Gamma
    if let Some(watermark_gamma) = watermark_gamma {
OlivierDehaene's avatar
OlivierDehaene committed
1102
        envs.push(("WATERMARK_GAMMA".into(), watermark_gamma.to_string().into()))
1103
1104
1105
1106
    }

    // Watermark Delta
    if let Some(watermark_delta) = watermark_delta {
OlivierDehaene's avatar
OlivierDehaene committed
1107
        envs.push(("WATERMARK_DELTA".into(), watermark_delta.to_string().into()))
1108
1109
1110
    }

    // Start process
1111
    tracing::info!("Starting shard");
1112
    let mut p = match Command::new("text-generation-server")
OlivierDehaene's avatar
OlivierDehaene committed
1113
        .args(shard_args)
1114
        .env_clear()
OlivierDehaene's avatar
OlivierDehaene committed
1115
        .envs(envs)
1116
        .stdin(Stdio::piped())
1117
1118
1119
1120
1121
        .stdout(Stdio::piped())
        .stderr(Stdio::piped())
        .process_group(0)
        .spawn()
    {
1122
1123
        Ok(p) => p,
        Err(err) => {
1124
1125
1126
            if err.kind() == io::ErrorKind::NotFound {
                tracing::error!("text-generation-server not found in PATH");
                tracing::error!("Please install it with `make install-server`")
1127
1128
            }
            {
1129
                tracing::error!("{}", err);
1130
            }
1131

1132
            status_sender.send(ShardStatus::Failed(rank)).unwrap();
1133
1134
1135
1136
1137
            return;
        }
    };

    // Redirect STDOUT to the console
1138
    let mut pstdin = p.stdin.take().unwrap();
1139
    let shard_stdout_reader = BufReader::new(p.stdout.take().unwrap());
1140
    let shard_stderr_reader = BufReader::new(p.stderr.take().unwrap());
1141

1142
    //stdout tracing thread
1143
    thread::spawn(move || {
1144
        log_lines(shard_stdout_reader);
1145
    });
1146
1147
1148
    // 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 || {
OlivierDehaene's avatar
OlivierDehaene committed
1149
        for line in shard_stderr_reader.lines().map_while(Result::ok) {
1150
1151
1152
            err_sender.send(line).unwrap_or(());
        }
    });
1153
    // We read stdin in another thread as it seems that lines() can block in some cases
Nicolas Patry's avatar
Nicolas Patry committed
1154
1155
1156
1157
1158
1159
1160
1161
1162
    if LevelFilter::current() >= tracing::Level::DEBUG {
        thread::spawn(move || {
            let mut stdin = io::stdin(); // We get `Stdin` here.
            loop {
                let mut buffer = vec![0; 4096];
                if let Ok(n) = stdin.read(&mut buffer) {
                    if n > 0 {
                        let _ = pstdin.write_all(&buffer[..n]);
                    }
1163
1164
                }
            }
Nicolas Patry's avatar
Nicolas Patry committed
1165
1166
        });
    }
1167
1168
1169
1170
1171
1172

    let mut ready = false;
    let start_time = Instant::now();
    let mut wait_time = Instant::now();
    loop {
        // Process exited
1173
        if let Some(exit_status) = p.try_wait().unwrap() {
1174
1175
1176
1177
            let mut err = String::new();
            while let Ok(line) = err_receiver.recv_timeout(Duration::from_millis(10)) {
                err = err + "\n" + &line;
            }
1178

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

1181
            if let Some(signal) = exit_status.signal() {
1182
1183
1184
                tracing::error!("Shard process was signaled to shutdown with signal {signal}");
            }

1185
            status_sender.send(ShardStatus::Failed(rank)).unwrap();
1186
1187
1188
1189
            return;
        }

        // We received a shutdown signal
1190
        if shutdown.load(Ordering::SeqCst) {
1191
            terminate("shard", p, Duration::from_secs(90)).unwrap();
1192
1193
1194
1195
1196
            return;
        }

        // Shard is ready
        if uds.exists() && !ready {
1197
            tracing::info!("Shard ready in {:?}", start_time.elapsed());
1198
1199
1200
            status_sender.send(ShardStatus::Ready).unwrap();
            ready = true;
        } else if !ready && wait_time.elapsed() > Duration::from_secs(10) {
1201
            tracing::info!("Waiting for shard to be ready...");
1202
1203
1204
1205
1206
1207
            wait_time = Instant::now();
        }
        sleep(Duration::from_millis(100));
    }
}

1208
fn shutdown_shards(shutdown: Arc<AtomicBool>, shutdown_receiver: &mpsc::Receiver<()>) {
1209
1210
1211
    tracing::info!("Shutting down shards");
    // Update shutdown value to true
    // This will be picked up by the shard manager
1212
    shutdown.store(true, Ordering::SeqCst);
1213
1214
1215
1216
1217
1218
1219

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

fn num_cuda_devices() -> Option<usize> {
1220
1221
    let devices = match env::var("CUDA_VISIBLE_DEVICES") {
        Ok(devices) => devices,
1222
1223
1224
        Err(_) => match env::var("NVIDIA_VISIBLE_DEVICES") {
            Ok(devices) => devices,
            Err(_) => env::var("ZE_AFFINITY_MASK").ok()?,
Nicolas Patry's avatar
Nicolas Patry committed
1225
        },
1226
    };
1227
1228
    let n_devices = devices.split(',').count();
    Some(n_devices)
1229
1230
1231
1232
1233
1234
1235
1236
1237
1238
1239
1240
1241
1242
1243
1244
1245
1246
1247
1248
1249
1250
1251
1252
1253
1254
1255
1256
1257
1258
1259
1260
1261
}

#[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 {
1262
1263
1264
1265
1266
1267
1268
            PythonLogLevelEnum::Trace => tracing::trace!("{}", self.text.trim_end()),
            PythonLogLevelEnum::Debug => tracing::debug!("{}", self.text.trim_end()),
            PythonLogLevelEnum::Info => tracing::info!("{}", self.text.trim_end()),
            PythonLogLevelEnum::Success => tracing::info!("{}", self.text.trim_end()),
            PythonLogLevelEnum::Warning => tracing::warn!("{}", self.text.trim_end()),
            PythonLogLevelEnum::Error => tracing::error!("{}", self.text.trim_end()),
            PythonLogLevelEnum::Critical => tracing::error!("{}", self.text.trim_end()),
1269
1270
1271
1272
        }
    }
}

1273
impl TryFrom<&[u8]> for PythonLogMessage {
1274
1275
    type Error = serde_json::Error;

1276
1277
    fn try_from(value: &[u8]) -> Result<Self, Self::Error> {
        serde_json::from_slice::<Self>(value)
1278
1279
1280
    }
}

1281
1282
1283
1284
1285
1286
1287
1288
1289
1290
1291
1292
1293
fn log_lines<R: Sized + Read>(mut bufread: BufReader<R>) {
    let mut buffer = vec![0u8; 8 * 4096];
    let mut stdout = std::io::stdout();
    loop {
        let n = bufread.read(&mut buffer);
        if let Ok(n) = n {
            if n > 0 {
                let mut lines = buffer[..n].split(|i| *i == b'\n').peekable();
                while let Some(line) = lines.next() {
                    match PythonLogMessage::try_from(line) {
                        Ok(log) => log.trace(),
                        // For interactive debugging ?
                        Err(_) => {
1294
1295
1296
1297
1298
1299
                            if LevelFilter::current() >= tracing::Level::DEBUG {
                                stdout.write_all(line).unwrap();
                                if lines.peek().is_some() {
                                    stdout.write_all(b"\n").unwrap();
                                }
                                stdout.flush().unwrap();
1300
1301
1302
1303
                            }
                        }
                    }
                }
1304
1305
            } else {
                break;
1306
            }
1307
1308
1309
1310
        }
    }
}

1311
1312
1313
1314
fn find_num_shards(
    sharded: Option<bool>,
    num_shard: Option<usize>,
) -> Result<usize, LauncherError> {
1315
1316
1317
1318
    // 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
1319
            tracing::info!("Parsing num_shard from CUDA_VISIBLE_DEVICES/NVIDIA_VISIBLE_DEVICES/ZE_AFFINITY_MASK");
1320
            let n_devices = num_cuda_devices()
1321
                .expect("--num-shard and CUDA_VISIBLE_DEVICES/NVIDIA_VISIBLE_DEVICES/ZE_AFFINITY_MASK are not set");
1322
            if n_devices <= 1 {
1323
1324
1325
                return Err(LauncherError::NotEnoughCUDADevices(format!(
                    "`sharded` is true but only found {n_devices} CUDA devices"
                )));
1326
            }
1327
            n_devices
1328
        }
1329
1330
1331
        (Some(true), Some(num_shard)) => {
            // we can't have only one shard while sharded
            if num_shard <= 1 {
1332
1333
1334
                return Err(LauncherError::ArgumentValidation(
                    "`sharded` is true but `num_shard` <= 1".to_string(),
                ));
1335
1336
            }
            num_shard
1337
        }
1338
1339
1340
1341
        (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,
1342
    };
1343
    if num_shard < 1 {
1344
1345
1346
        return Err(LauncherError::ArgumentValidation(
            "`num_shard` cannot be < 1".to_string(),
        ));
1347
    }
1348
    Ok(num_shard)
1349
}
1350

1351
#[derive(Debug, Error)]
1352
enum LauncherError {
1353
    #[error("Invalid argument: {0}")]
1354
    ArgumentValidation(String),
1355
    #[error("not enough cuda devices: {0}")]
1356
    NotEnoughCUDADevices(String),
1357
    #[error("Download error")]
1358
    DownloadError,
1359
    #[error("Shard cannot start")]
1360
    ShardCannotStart,
1361
    #[error("Shard disconnected")]
1362
    ShardDisconnected,
1363
    #[error("Shard failed")]
1364
    ShardFailed,
1365
    #[error("Webserver failed")]
1366
    WebserverFailed,
1367
    #[error("Webserver cannot start")]
1368
1369
    WebserverCannotStart,
}
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1370

1371
1372
1373
1374
1375
1376
1377
fn download_convert_model(
    model_id: &str,
    revision: Option<&str>,
    trust_remote_code: bool,
    huggingface_hub_cache: Option<&str>,
    weights_cache_override: Option<&str>,
    running: Arc<AtomicBool>,
1378
    merge_lora: bool,
1379
) -> Result<(), LauncherError> {
1380
1381
1382
    // Enter download tracing span
    let _span = tracing::span!(tracing::Level::INFO, "download").entered();

OlivierDehaene's avatar
OlivierDehaene committed
1383
    let mut download_args = vec![
1384
        "download-weights".to_string(),
1385
        model_id.to_string(),
1386
1387
1388
1389
1390
1391
        "--extension".to_string(),
        ".safetensors".to_string(),
        "--logger-level".to_string(),
        "INFO".to_string(),
        "--json-output".to_string(),
    ];
1392

1393
1394
1395
1396
    if merge_lora {
        download_args.push("--merge-lora".to_string());
    }

1397
    // Model optional revision
1398
    if let Some(revision) = &revision {
OlivierDehaene's avatar
OlivierDehaene committed
1399
1400
        download_args.push("--revision".to_string());
        download_args.push(revision.to_string())
1401
    }
1402

1403
    // Trust remote code for automatic peft fusion
1404
    if trust_remote_code {
1405
1406
1407
        download_args.push("--trust-remote-code".to_string());
    }

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

1411
1412
1413
    // Remove LOG_LEVEL if present
    envs.retain(|(name, _)| name != "LOG_LEVEL");

1414
1415
1416
    // Disable progress bar
    envs.push(("HF_HUB_DISABLE_PROGRESS_BARS".into(), "1".into()));

1417
    // If huggingface_hub_cache is set, pass it to the download process
1418
    // Useful when running inside a docker container
1419
    if let Some(ref huggingface_hub_cache) = huggingface_hub_cache {
OlivierDehaene's avatar
OlivierDehaene committed
1420
        envs.push(("HUGGINGFACE_HUB_CACHE".into(), huggingface_hub_cache.into()));
1421
    };
1422

1423
1424
    // 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
1425
    envs.push((
1426
1427
1428
        "HF_HUB_ENABLE_HF_TRANSFER".into(),
        enable_hf_transfer.into(),
    ));
1429

1430
1431
    // Parse Inference API token
    if let Ok(api_token) = env::var("HF_API_TOKEN") {
1432
        envs.push(("HF_TOKEN".into(), api_token.into()))
1433
    };
1434

1435
1436
    // If args.weights_cache_override is some, pass it to the download process
    // Useful when running inside a HuggingFace Inference Endpoint
1437
    if let Some(weights_cache_override) = &weights_cache_override {
OlivierDehaene's avatar
OlivierDehaene committed
1438
        envs.push((
1439
1440
1441
1442
1443
            "WEIGHTS_CACHE_OVERRIDE".into(),
            weights_cache_override.into(),
        ));
    };

1444
    // Start process
1445
    tracing::info!("Starting check and download process for {model_id}");
1446
    let mut download_process = match Command::new("text-generation-server")
OlivierDehaene's avatar
OlivierDehaene committed
1447
        .args(download_args)
1448
        .env_clear()
OlivierDehaene's avatar
OlivierDehaene committed
1449
        .envs(envs)
1450
1451
1452
1453
1454
        .stdout(Stdio::piped())
        .stderr(Stdio::piped())
        .process_group(0)
        .spawn()
    {
1455
1456
        Ok(p) => p,
        Err(err) => {
1457
1458
1459
            if err.kind() == io::ErrorKind::NotFound {
                tracing::error!("text-generation-server not found in PATH");
                tracing::error!("Please install it with `make install-server`")
1460
1461
            } else {
                tracing::error!("{}", err);
1462
            }
1463

1464
1465
1466
            return Err(LauncherError::DownloadError);
        }
    };
1467

1468
    let download_stdout = BufReader::new(download_process.stdout.take().unwrap());
1469

1470
    thread::spawn(move || {
1471
        log_lines(download_stdout);
1472
1473
1474
1475
1476
1477
1478
    });

    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 || {
OlivierDehaene's avatar
OlivierDehaene committed
1479
        for line in download_stderr.lines().map_while(Result::ok) {
1480
1481
            err_sender.send(line).unwrap_or(());
        }
1482
    });
1483

1484
    loop {
1485
1486
        if let Some(status) = download_process.try_wait().unwrap() {
            if status.success() {
1487
                tracing::info!("Successfully downloaded weights for {model_id}");
1488
                break;
1489
            }
1490
1491

            let mut err = String::new();
1492
1493
1494
1495
            while let Ok(line) = err_receiver.recv_timeout(Duration::from_millis(10)) {
                err = err + "\n" + &line;
            }

1496
1497
1498
1499
1500
1501
1502
1503
1504
            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);
1505
        }
1506
        if !running.load(Ordering::SeqCst) {
OlivierDehaene's avatar
OlivierDehaene committed
1507
            terminate("download", download_process, Duration::from_secs(10)).unwrap();
1508
1509
1510
            return Ok(());
        }
        sleep(Duration::from_millis(100));
1511
    }
1512
1513
    Ok(())
}
1514

1515
#[allow(clippy::too_many_arguments)]
1516
1517
1518
fn spawn_shards(
    num_shard: usize,
    args: &Args,
1519
    cuda_graphs: Vec<usize>,
1520
1521
    max_total_tokens: Option<usize>,
    max_input_tokens: Option<usize>,
1522
    quantize: Option<Quantization>,
1523
    max_log_level: LevelFilter,
1524
    shutdown: Arc<AtomicBool>,
1525
1526
1527
1528
1529
1530
    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
1531
1532
    // Start shard processes
    for rank in 0..num_shard {
1533
1534
1535
1536
1537
1538
        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
1539
1540
1541
        let status_sender = status_sender.clone();
        let shutdown = shutdown.clone();
        let shutdown_sender = shutdown_sender.clone();
1542
        let otlp_endpoint = args.otlp_endpoint.clone();
1543
        let otlp_service_name = args.otlp_service_name.clone();
Nicolas Patry's avatar
Nicolas Patry committed
1544
        let speculate = args.speculate;
1545
        let dtype = args.dtype;
1546
        let kv_cache_dtype = args.kv_cache_dtype;
1547
        let trust_remote_code = args.trust_remote_code;
1548
1549
1550
1551
        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;
1552
        let cuda_graphs_clone = cuda_graphs.clone();
1553
        let cuda_memory_fraction = args.cuda_memory_fraction;
Nicolas Patry's avatar
Nicolas Patry committed
1554
1555
        let rope_scaling = args.rope_scaling;
        let rope_factor = args.rope_factor;
1556
        let max_batch_size = args.max_batch_size;
drbh's avatar
drbh committed
1557
        let lora_adapters = args.lora_adapters.clone();
Nicolas Patry's avatar
Nicolas Patry committed
1558
        let enable_prefill_logprobs = args.enable_prefill_logprobs;
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1559
1560
        thread::spawn(move || {
            shard_manager(
1561
                model_id,
1562
                revision,
1563
                quantize,
Nicolas Patry's avatar
Nicolas Patry committed
1564
                speculate,
1565
                dtype,
1566
                kv_cache_dtype,
1567
                trust_remote_code,
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1568
1569
1570
1571
1572
                uds_path,
                rank,
                num_shard,
                master_addr,
                master_port,
1573
1574
                huggingface_hub_cache,
                weights_cache_override,
1575
                disable_custom_kernels,
1576
1577
                watermark_gamma,
                watermark_delta,
1578
                cuda_graphs_clone,
1579
                cuda_memory_fraction,
Nicolas Patry's avatar
Nicolas Patry committed
1580
1581
                rope_scaling,
                rope_factor,
1582
1583
                max_total_tokens,
                max_batch_size,
1584
                max_input_tokens,
drbh's avatar
drbh committed
1585
                lora_adapters,
Nicolas Patry's avatar
Nicolas Patry committed
1586
                enable_prefill_logprobs,
1587
                otlp_endpoint,
1588
                otlp_service_name,
1589
                max_log_level,
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1590
1591
1592
1593
1594
1595
1596
1597
1598
1599
1600
1601
1602
1603
1604
1605
1606
1607
1608
1609
1610
                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));
            }
1611
            Ok(ShardStatus::Failed(rank)) => {
1612
                tracing::error!("Shard {rank} failed to start");
1613
                shutdown_shards(shutdown, shutdown_receiver);
1614
                return Err(LauncherError::ShardCannotStart);
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1615
1616
1617
            }
            Err(TryRecvError::Disconnected) => {
                tracing::error!("Shard status channel disconnected");
1618
                shutdown_shards(shutdown, shutdown_receiver);
1619
                return Err(LauncherError::ShardDisconnected);
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1620
1621
1622
            }
        }
    }
1623
1624
    Ok(())
}
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1625

1626
1627
1628
1629
1630
1631
1632
1633
1634
1635
1636
1637
1638
#[derive(Debug)]
enum Gpu {
    RTX4090,
    T4,
    L4,
    L40,
    L40S,
    A10G,
    H100,
    A100,
    Unknown(String),
}

Nicolas Patry's avatar
Nicolas Patry committed
1639
1640
1641
#[derive(Debug)]
struct ComputeType {
    count: usize,
1642
1643
1644
1645
1646
1647
1648
1649
1650
1651
1652
1653
1654
1655
1656
1657
1658
1659
1660
1661
1662
1663
1664
1665
1666
1667
1668
1669
1670
1671
1672
1673
1674
1675
    card: Gpu,
}

impl From<&str> for Gpu {
    fn from(value: &str) -> Self {
        match value {
            "nvidia-4090" => Gpu::RTX4090,
            "nvidia-t4" => Gpu::T4,
            "nvidia-l4" => Gpu::L4,
            "nvidia-l40" => Gpu::L40,
            "nvidia-l40s" => Gpu::L40S,
            "nvidia-a10g" => Gpu::A10G,
            "nvidia-h100-80gb-hbm3" => Gpu::H100,
            "nvidia-a100-sxm4-80gb" => Gpu::A100,
            "nvidia-a100" => Gpu::A100,
            card => Gpu::Unknown(card.to_string()),
        }
    }
}

impl std::fmt::Display for Gpu {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Gpu::RTX4090 => write!(f, "nvida-4090"),
            Gpu::T4 => write!(f, "nvida-t4"),
            Gpu::L4 => write!(f, "nvida-l4"),
            Gpu::L40 => write!(f, "nvida-l40"),
            Gpu::L40S => write!(f, "nvida-l40s"),
            Gpu::A10G => write!(f, "nvidia-a10g"),
            Gpu::H100 => write!(f, "nvidia-h100-80fb-hbm3"),
            Gpu::A100 => write!(f, "nvida-a100-sxm4-80gb"),
            Gpu::Unknown(card) => write!(f, "{}", card),
        }
    }
Nicolas Patry's avatar
Nicolas Patry committed
1676
1677
1678
1679
}

impl ComputeType {
    fn f16_flop(&self) -> Option<u64> {
1680
        let card_flop = match &self.card {
Nicolas Patry's avatar
Nicolas Patry committed
1681
1682
            // https://www.nvidia.com/en-us/geforce/graphics-cards/40-series/rtx-4090/
            // Specs are unclear https://www.itcreations.com/nvidia-gpu/nvidia-geforce-rtx-4090-gpu
1683
            Gpu::RTX4090 => Some(82 * 10u64.pow(12)),
Nicolas Patry's avatar
Nicolas Patry committed
1684
            // https://www.nvidia.com/en-us/data-center/tesla-t4/
1685
            Gpu::T4 => Some(65 * 10u64.pow(12)),
Nicolas Patry's avatar
Nicolas Patry committed
1686
            // https://www.nvidia.com/en-us/data-center/l4/
1687
1688
1689
1690
1691
            Gpu::L4 => Some(121 * 10u64.pow(12)),
            // https://www.nvidia.com/en-us/data-center/l40/
            Gpu::L40 => Some(181 * 10u64.pow(12)),
            // https://www.nvidia.com/en-us/data-center/l40s/
            Gpu::L40S => Some(363 * 10u64.pow(12)),
Nicolas Patry's avatar
Nicolas Patry committed
1692
            // https://www.nvidia.com/en-us/data-center/products/a10-gpu/
1693
            Gpu::A10G => Some(125 * 10u64.pow(12)),
Nicolas Patry's avatar
Nicolas Patry committed
1694
1695
            // https://www.nvidia.com/en-us/data-center/h100/
            // https://www.techpowerup.com/gpu-specs/docs/nvidia-gh100-architecture.pdf
1696
            Gpu::H100 => Some(900 * 10u64.pow(12)),
Nicolas Patry's avatar
Nicolas Patry committed
1697
            // https://www.nvidia.com/content/dam/en-zz/Solutions/Data-Center/a100/pdf/nvidia-a100-datasheet-us-nvidia-1758950-r4-web.pdf
1698
1699
            Gpu::A100 => Some(312 * 10u64.pow(12)),
            Gpu::Unknown(card) => {
Nicolas Patry's avatar
Nicolas Patry committed
1700
1701
1702
1703
1704
1705
                tracing::warn!("Unkown compute for card {card}");
                None
            }
        };
        card_flop.map(|f| f * self.count as u64)
    }
1706
1707
1708
1709
1710
1711
1712
1713
1714
1715
1716
1717
1718
1719
1720
1721
1722
1723
1724
1725
1726
1727
1728
1729
1730

    fn vram(&self, memory_fraction: f32) -> Option<usize> {
        let output = Command::new("nvidia-smi")
            .args(["--query-gpu=memory.total", "--format=csv"])
            .output()
            .ok()?;
        let output = String::from_utf8(output.stdout).ok()?;
        let fullname = output.split('\n').nth(1)?;
        let mut tokens = fullname.split(' ');
        let amount = tokens.next()?;
        let unit = tokens.next()?;
        if unit != "MiB" {
            tracing::warn!("Unexpected memory unit {unit}, expected MiB");
            return None;
        }
        let amount: usize = amount.parse().ok()?;
        let amount = amount * 2usize.pow(20);
        let wiggle_room: f32 = env::var("TGI_WIGGLE_ROOM")
            .ok()
            .and_then(|wiggle| wiggle.parse().ok())
            .unwrap_or(0.95);
        let total = amount * self.count;
        let adjusted = ((total as f32) * memory_fraction * wiggle_room) as usize;
        Some(adjusted)
    }
Nicolas Patry's avatar
Nicolas Patry committed
1731
1732
1733
1734
1735
1736
1737
1738
}

impl From<ComputeType> for OsString {
    fn from(value: ComputeType) -> Self {
        format!("{}-{}", value.count, value.card).into()
    }
}

1739
fn compute_type(count: usize) -> Option<ComputeType> {
1740
1741
1742
1743
1744
1745
1746
    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();
1747
1748
    let card = (&*cardname).into();
    Some(ComputeType { count, card })
1749
1750
}

1751
fn spawn_webserver(
1752
    num_shard: usize,
1753
    args: Args,
1754
1755
    max_input_tokens: Option<usize>,
    max_total_tokens: Option<usize>,
1756
    max_batch_prefill_tokens: u32,
1757
    shutdown: Arc<AtomicBool>,
1758
    shutdown_receiver: &mpsc::Receiver<()>,
1759
) -> Result<Child, LauncherError> {
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1760
1761
1762
    // All shard started
    // Start webserver
    tracing::info!("Starting Webserver");
OlivierDehaene's avatar
OlivierDehaene committed
1763
    let mut router_args = vec![
1764
1765
        "--max-client-batch-size".to_string(),
        args.max_client_batch_size.to_string(),
1766
        "--max-concurrent-requests".to_string(),
1767
        args.max_concurrent_requests.to_string(),
1768
        "--max-best-of".to_string(),
1769
        args.max_best_of.to_string(),
1770
        "--max-stop-sequences".to_string(),
1771
        args.max_stop_sequences.to_string(),
Nicolas Patry's avatar
Nicolas Patry committed
1772
1773
        "--max-top-n-tokens".to_string(),
        args.max_top_n_tokens.to_string(),
1774
        "--max-batch-prefill-tokens".to_string(),
1775
        max_batch_prefill_tokens.to_string(),
1776
        "--waiting-served-ratio".to_string(),
1777
        args.waiting_served_ratio.to_string(),
1778
        "--max-waiting-tokens".to_string(),
1779
        args.max_waiting_tokens.to_string(),
1780
1781
        "--validation-workers".to_string(),
        args.validation_workers.to_string(),
1782
1783
        "--hostname".to_string(),
        args.hostname.to_string(),
1784
        "--port".to_string(),
1785
        args.port.to_string(),
1786
        "--master-shard-uds-path".to_string(),
1787
        format!("{}-0", args.shard_uds_path),
1788
        "--tokenizer-name".to_string(),
1789
        args.model_id,
1790
1791
        "--payload-limit".to_string(),
        args.payload_limit.to_string(),
1792
    ];
1793
1794
1795
1796
1797
1798
1799
1800
1801
1802
1803
1804
    if let Some(max_input_tokens) = max_input_tokens {
        router_args.extend_from_slice(&[
            "--max-input-tokens".to_string(),
            max_input_tokens.to_string(),
        ]);
    }
    if let Some(max_total_tokens) = max_total_tokens {
        router_args.extend_from_slice(&[
            "--max-total-tokens".to_string(),
            max_total_tokens.to_string(),
        ]);
    }
1805

1806
    // Pass usage stats flags to router
1807
1808
    router_args.push("--usage-stats".to_string());
    router_args.push(args.usage_stats.to_string());
1809

drbh's avatar
drbh committed
1810
1811
1812
1813
1814
    // Grammar support
    if args.disable_grammar_support {
        router_args.push("--disable-grammar-support".to_string());
    }

1815
1816
1817
1818
1819
1820
    // 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());
    }

1821
1822
1823
1824
1825
1826
    // 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());
    }

1827
1828
1829
1830
1831
1832
    // 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());
    }

1833
1834
    // Model optional revision
    if let Some(ref revision) = args.revision {
OlivierDehaene's avatar
OlivierDehaene committed
1835
1836
        router_args.push("--revision".to_string());
        router_args.push(revision.to_string())
1837
1838
    }

1839
1840
1841
1842
    if args.trust_remote_code {
        router_args.push("--trust-remote-code".to_string());
    }

1843
    if args.json_output {
OlivierDehaene's avatar
OlivierDehaene committed
1844
        router_args.push("--json-output".to_string());
1845
1846
    }

1847
    // OpenTelemetry
1848
    if let Some(otlp_endpoint) = args.otlp_endpoint {
OlivierDehaene's avatar
OlivierDehaene committed
1849
1850
        router_args.push("--otlp-endpoint".to_string());
        router_args.push(otlp_endpoint);
1851
1852
    }

1853
1854
1855
1856
1857
    // OpenTelemetry
    let otlp_service_name = args.otlp_service_name;
    router_args.push("--otlp-service-name".to_string());
    router_args.push(otlp_service_name);

1858
1859
    // CORS origins
    for origin in args.cors_allow_origin.into_iter() {
OlivierDehaene's avatar
OlivierDehaene committed
1860
1861
        router_args.push("--cors-allow-origin".to_string());
        router_args.push(origin);
1862
1863
    }

Erik Kaunismäki's avatar
Erik Kaunismäki committed
1864
1865
1866
1867
1868
    // API Key
    if let Some(api_key) = args.api_key {
        router_args.push("--api-key".to_string());
        router_args.push(api_key);
    }
1869
1870
    // Ngrok
    if args.ngrok {
OlivierDehaene's avatar
OlivierDehaene committed
1871
1872
        router_args.push("--ngrok".to_string());
        router_args.push("--ngrok-authtoken".to_string());
1873
1874
1875
        router_args.push(args.ngrok_authtoken.unwrap());
        router_args.push("--ngrok-edge".to_string());
        router_args.push(args.ngrok_edge.unwrap());
1876
1877
    }

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

1881
1882
    // Parse Inference API token
    if let Ok(api_token) = env::var("HF_API_TOKEN") {
1883
        envs.push(("HF_TOKEN".into(), api_token.into()))
1884
    };
1885

1886
1887
1888
1889
1890
1891
1892
    // 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()))
    }

1893
    let mut webserver = match Command::new("text-generation-router")
OlivierDehaene's avatar
OlivierDehaene committed
1894
1895
        .args(router_args)
        .envs(envs)
1896
1897
1898
1899
1900
        .stdout(Stdio::piped())
        .stderr(Stdio::piped())
        .process_group(0)
        .spawn()
    {
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1901
1902
        Ok(p) => p,
        Err(err) => {
1903
            tracing::error!("Failed to start webserver: {}", err);
1904
1905
1906
            if err.kind() == io::ErrorKind::NotFound {
                tracing::error!("text-generation-router not found in PATH");
                tracing::error!("Please install it with `make install-router`")
1907
1908
            } else {
                tracing::error!("{}", err);
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1909
            }
1910

1911
            shutdown_shards(shutdown, shutdown_receiver);
1912
            return Err(LauncherError::WebserverCannotStart);
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1913
1914
1915
        }
    };

1916
1917
1918
    // Redirect STDOUT and STDERR to the console
    let webserver_stdout = webserver.stdout.take().unwrap();
    let webserver_stderr = webserver.stderr.take().unwrap();
1919
1920

    thread::spawn(move || {
1921
1922
        let stdout = BufReader::new(webserver_stdout);
        let stderr = BufReader::new(webserver_stderr);
1923
        for line in stdout.lines() {
1924
            println!("{}", line.unwrap());
1925
        }
1926
1927
        for line in stderr.lines() {
            println!("{}", line.unwrap());
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1928
        }
1929
1930
1931
    });
    Ok(webserver)
}
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1932

OlivierDehaene's avatar
OlivierDehaene committed
1933
1934
1935
1936
1937
1938
1939
1940
1941
1942
1943
1944
1945
1946
1947
1948
1949
1950
1951
1952
1953
1954
1955
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)
}

1956
1957
fn main() -> Result<(), LauncherError> {
    // Pattern match configuration
1958
    let args: Args = Args::parse();
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1959

1960
    // Filter events with LOG_LEVEL
1961
1962
1963
1964
1965
1966
1967
1968
1969
1970
1971
1972
1973
1974
1975
1976
    let varname = "LOG_LEVEL";
    let env_filter = if let Ok(log_level) = std::env::var(varname) {
        // Override to avoid simple logs to be spammed with tokio level informations
        let log_level = match &log_level[..] {
            "warn" => "text_generation_launcher=warn,text_generation_router=warn",
            "info" => "text_generation_launcher=info,text_generation_router=info",
            "debug" => "text_generation_launcher=debug,text_generation_router=debug",
            log_level => log_level,
        };
        EnvFilter::builder()
            .with_default_directive(LevelFilter::INFO.into())
            .parse_lossy(log_level)
    } else {
        EnvFilter::new("info")
    };
    let max_log_level = env_filter.max_level_hint().unwrap_or(LevelFilter::INFO);
1977

1978
    if args.json_output {
1979
1980
1981
1982
        tracing_subscriber::fmt()
            .with_env_filter(env_filter)
            .json()
            .init();
1983
    } else {
1984
1985
1986
1987
        tracing_subscriber::fmt()
            .with_env_filter(env_filter)
            .compact()
            .init();
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1988
1989
    }

1990
1991
1992
1993
1994
    if args.env {
        let env_runtime = env_runtime::Env::new();
        tracing::info!("{}", env_runtime);
    }

Nicolas Patry's avatar
Nicolas Patry committed
1995
    tracing::info!("{:#?}", args);
1996

1997
1998
1999
2000
2001
2002
    let config: Option<Config> = get_config(&args.model_id, &args.revision).ok();
    let quantize = config.as_ref().and_then(|c| c.quantize);
    // Quantization usually means you're even more RAM constrained.

    let (prefix_caching, attention) = resolve_attention(&config, &args.lora_adapters);
    tracing::info!("Using attention {attention} - Prefix caching {prefix_caching}");
2003
    std::env::set_var("PREFIX_CACHING", prefix_caching);
2004
    std::env::set_var("ATTENTION", attention);
2005

Nicolas Patry's avatar
Nicolas Patry committed
2006
2007
2008
2009
2010
2011
2012
2013
2014
2015
    let num_shard = find_num_shards(args.sharded, args.num_shard)?;
    if num_shard > 1 {
        if matches!(args.quantize, Some(Quantization::Exl2)) {
            return Err(LauncherError::ArgumentValidation(
                "Sharding is currently not supported with `exl2` quantization".into(),
            ));
        }
        tracing::info!("Sharding model on {num_shard} processes");
    }

2016
2017
2018
2019
2020
2021
2022
    let max_input_tokens = {
        match (args.max_input_tokens, args.max_input_length) {
            (Some(max_input_tokens), Some(max_input_length)) => {
                return Err(LauncherError::ArgumentValidation(
                    format!("Both `max_input_tokens` ({max_input_tokens}) and `max_input_length` ({max_input_length}) are set. Please define only `max_input_tokens` as `max_input_length is deprecated for naming consistency.",
                )));
            }
2023
2024
            (Some(max_input_tokens), None) | (None, Some(max_input_tokens)) => {
                Some(max_input_tokens)
2025
            }
2026
            (None, None) => None,
2027
2028
        }
    };
2029
    let max_total_tokens = args.max_total_tokens;
2030
2031
2032
2033
    let max_batch_prefill_tokens = {
        match args.max_batch_prefill_tokens {
            Some(max_batch_prefill_tokens) => max_batch_prefill_tokens,
            None => {
Nicolas Patry's avatar
Nicolas Patry committed
2034
2035
2036
                let compute_type = compute_type(num_shard);
                let compute_optimal = compute_optimal(config.as_ref(), compute_type.as_ref());
                let default = compute_optimal.unwrap_or(4096);
2037
2038
2039
2040
2041
                let vram_maximum = vram_maximum(
                    config.as_ref(),
                    compute_type.as_ref(),
                    args.cuda_memory_fraction,
                );
Nicolas Patry's avatar
Nicolas Patry committed
2042
2043
2044
2045
2046
2047
                let max_position_embeddings = config.and_then(|c| c.max_position_embeddings);
                let value = if let Some(max_position_embeddings) = max_position_embeddings {
                    default.min(max_position_embeddings)
                } else {
                    default
                };
2048
2049
2050
2051
2052
2053
2054
2055
                let value = if let Some(vram_maximum) = vram_maximum {
                    if vram_maximum < value {
                        tracing::warn!("Reducing the max batch prefill from {default} to {vram_maximum} because there is not enough VRAM to support it.");
                    }
                    value.min(vram_maximum)
                } else {
                    value
                };
2056
                tracing::info!("Default `max_batch_prefill_tokens` to {value}");
Nicolas Patry's avatar
Nicolas Patry committed
2057
                value as u32
2058
2059
2060
2061
            }
        }
    };

2062
    // Validate args
2063
2064
2065
2066
2067
2068
    if let (Some(max_input_tokens), Some(max_total_tokens)) = (max_input_tokens, max_total_tokens) {
        if max_input_tokens >= max_total_tokens {
            return Err(LauncherError::ArgumentValidation(
                    format!("`max_input_tokens`({max_input_tokens}) must be < `max_total_tokens`({max_total_tokens})"),
                ));
        }
2069
    }
2070

2071
2072
2073
2074
2075
    if matches!(args.quantize, Some(Quantization::Bitsandbytes)) {
        tracing::warn!("Bitsandbytes is deprecated, use `eetq` instead, which provides better latencies overall and is drop-in in most cases.");
    }
    let quantize = args.quantize.or(quantize);
    let cuda_graphs = match (&args.cuda_graphs, &quantize) {
Nicolas Patry's avatar
Nicolas Patry committed
2076
        (Some(cuda_graphs), _) => cuda_graphs.iter().cloned().filter(|&c| c > 0).collect(),
2077
2078
2079
2080
2081
        #[allow(deprecated)]
        (
            None,
            Some(
                Quantization::Bitsandbytes
2082
2083
                | Quantization::BitsandbytesNf4
                | Quantization::BitsandbytesFp4,
2084
2085
            ),
        ) => {
2086
2087
2088
2089
2090
            tracing::warn!("Bitsandbytes doesn't work with cuda graphs, deactivating them");
            vec![]
        }
        (None, Some(Quantization::Exl2)) => {
            tracing::warn!("Exl2 doesn't work with cuda graphs, deactivating them");
2091
2092
2093
2094
2095
2096
2097
2098
2099
            vec![]
        }
        _ => {
            let cuda_graphs = vec![1, 2, 4, 8, 16, 32];
            tracing::info!("Using default cuda graphs {cuda_graphs:?}");
            cuda_graphs
        }
    };

2100
2101
2102
2103
2104
    if args.validation_workers == 0 {
        return Err(LauncherError::ArgumentValidation(
            "`validation_workers` must be > 0".to_string(),
        ));
    }
2105
2106
2107
2108
2109
2110
    if args.trust_remote_code {
        tracing::warn!(
            "`trust_remote_code` is set. Trusting that model `{}` do not contain malicious code.",
            args.model_id
        );
    }
2111

2112
    if let Some(ref max_batch_total_tokens) = args.max_batch_total_tokens {
2113
2114
2115
2116
2117
2118
2119
        if let Some(max_total_tokens) = max_total_tokens {
            if 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 {}",
                    max_total_tokens, max_batch_total_tokens
                )));
            }
2120
2121
2122
        }
    }

2123
2124
2125
2126
2127
2128
2129
2130
2131
2132
2133
2134
2135
2136
    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(),
            ));
        }
    }

2137
2138
2139
2140
2141
2142
2143
    // 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");
2144

2145
    // Download and convert model weights
2146
2147
2148
2149
2150
2151
2152
    download_convert_model(
        &args.model_id,
        args.revision.as_deref(),
        args.trust_remote_code,
        args.huggingface_hub_cache.as_deref(),
        args.weights_cache_override.as_deref(),
        running.clone(),
2153
        true, // if its only a lora model - we should merge the lora adapters
2154
2155
2156
2157
2158
    )?;

    // Download and convert lora adapters if any
    if let Some(lora_adapters) = &args.lora_adapters {
        for adapter in lora_adapters.split(',') {
2159
2160
2161
2162
            // skip download if a path is provided
            if adapter.contains('=') {
                continue;
            }
2163
2164
2165
2166
2167
2168
2169
2170
2171
2172
2173
2174
2175
2176
2177
2178
2179
2180
2181
2182
2183
2184
2185
2186

            let adapter = adapter.trim();

            // check if adapter has more than 1 '@'
            if adapter.matches('@').count() > 1 {
                return Err(LauncherError::ArgumentValidation(format!(
                    "Invalid LoRA adapter format: {}",
                    adapter
                )));
            }

            // capture adapter_id, path, revision in format of adapter_id=path@revision
            let re = Regex::new(r"^([^=@]+)(?:=([^@]+))?(?:@(.+))?$").unwrap();
            if let Some(caps) = re.captures(adapter) {
                let adapter_id = caps.get(1).map_or("", |m| m.as_str());
                let revision = caps.get(3).map(|m| m.as_str());

                download_convert_model(
                    adapter_id,
                    revision,
                    args.trust_remote_code,
                    args.huggingface_hub_cache.as_deref(),
                    args.weights_cache_override.as_deref(),
                    running.clone(),
2187
                    false, // avoid merging lora adapters if using multi-lora
2188
2189
2190
2191
2192
2193
2194
                )?;
            } else {
                return Err(LauncherError::ArgumentValidation(format!(
                    "Invalid LoRA adapter format: {}",
                    adapter
                )));
            }
2195
2196
        }
    }
2197

OlivierDehaene's avatar
OlivierDehaene committed
2198
2199
2200
2201
2202
    if !running.load(Ordering::SeqCst) {
        // Launcher was asked to stop
        return Ok(());
    }

2203
    // Shared shutdown bool
2204
    let shutdown = Arc::new(AtomicBool::new(false));
2205
2206
2207
    // Shared shutdown channel
    // When shutting down, the main thread will wait for all senders to be dropped
    let (shutdown_sender, shutdown_receiver) = mpsc::channel();
2208

2209
2210
    // Shared channel to track shard status
    let (status_sender, status_receiver) = mpsc::channel();
2211

2212
2213
2214
    spawn_shards(
        num_shard,
        &args,
2215
        cuda_graphs,
2216
        max_total_tokens,
2217
        max_input_tokens,
2218
        quantize,
2219
        max_log_level,
2220
2221
2222
2223
2224
2225
2226
        shutdown.clone(),
        &shutdown_receiver,
        shutdown_sender,
        &status_receiver,
        status_sender,
        running.clone(),
    )?;
2227

2228
2229
2230
2231
2232
    // We might have received a termination signal
    if !running.load(Ordering::SeqCst) {
        shutdown_shards(shutdown, &shutdown_receiver);
        return Ok(());
    }
2233

2234
2235
2236
2237
2238
2239
2240
2241
2242
    let mut webserver = spawn_webserver(
        num_shard,
        args,
        max_input_tokens,
        max_total_tokens,
        max_batch_prefill_tokens,
        shutdown.clone(),
        &shutdown_receiver,
    )
2243
    .inspect_err(|_| {
2244
2245
        shutdown_shards(shutdown.clone(), &shutdown_receiver);
    })?;
2246
2247
2248
2249
2250

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

    while running.load(Ordering::SeqCst) {
2251
        if let Ok(ShardStatus::Failed(rank)) = status_receiver.try_recv() {
OlivierDehaene's avatar
OlivierDehaene committed
2252
            tracing::error!("Shard {rank} crashed");
2253
2254
2255
2256
            exit_code = Err(LauncherError::ShardFailed);
            break;
        };

2257
        match webserver.try_wait().unwrap() {
2258
2259
2260
2261
2262
2263
2264
2265
2266
            Some(_) => {
                tracing::error!("Webserver Crashed");
                shutdown_shards(shutdown, &shutdown_receiver);
                return Err(LauncherError::WebserverFailed);
            }
            None => {
                sleep(Duration::from_millis(100));
            }
        };
2267
    }
2268
2269

    // Graceful termination
OlivierDehaene's avatar
OlivierDehaene committed
2270
    terminate("webserver", webserver, Duration::from_secs(90)).unwrap();
2271
2272
2273
    shutdown_shards(shutdown, &shutdown_receiver);

    exit_code
2274
}