"vscode:/vscode.git/clone" did not exist on "9b928293df3f729b10db595773261c94c2d40dcf"
main.rs 68.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

32
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
59
60
61
62
63
64
65
66
67
68
69
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) {
70
    let compute_capability = gpu::get_cuda_capability();
71
72
73
74
75
76
77
78
79
80
81
82
    let mut prefix_caching: Option<String> = std::env::var("USE_PREFIX_CACHING").ok();
    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());
            }
        }
83
84
85
86
87
88
89

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

90
91
92
93
94
95
96
97
98
99
100
101
        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() {
                    Some("gemma2") | Some("falcon") | Some("deepseek_v2") => {
                        // Required because gemma2 needs bfloat16 which is not supported by
                        // flashinfer ?
                        if attention.is_none() {
                            tracing::info!(
102
                                "Forcing attention to '{fallback_attention}' because model {} requires it",
103
104
                                config.model_type.as_ref().unwrap()
                            );
105
106
107
108
109
                            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());
110
111
112
113
114
115
116
117
                        }
                    }
                    Some("t5") => {}
                    _ => {}
                }
            }
            _ => {
                if attention.is_none() {
118
119
                    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());
120
121
122
123
124
125
126
                }
                if prefix_caching.is_none() {
                    prefix_caching = Some("0".to_string());
                }
            }
        }
    }
127

128
    let attention = attention.unwrap_or("flashinfer".to_string());
129
130
    let prefix_caching = prefix_caching.unwrap_or("true".to_string());

131
132
133
    (prefix_caching, attention)
}

134
#[derive(Deserialize)]
135
struct RawConfig {
136
    max_position_embeddings: Option<usize>,
137
    n_positions: Option<usize>,
138
    model_type: Option<String>,
139
    max_seq_len: Option<usize>,
140
    quantization_config: Option<QuantizationConfig>,
141
142
143
144
145
146
    n_embd: Option<usize>,
    hidden_size: Option<usize>,
    num_attention_heads: Option<usize>,
    head_dim: Option<usize>,
    vision_config: Option<VisionConfig>,
    is_encoder_decoder: Option<bool>,
147
148
149
150
151
}

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

154
155
156
#[derive(Deserialize)]
struct VisionConfig {}

157
158
159
#[derive(Deserialize)]
struct Config {
    max_position_embeddings: Option<usize>,
160
    quantize: Option<Quantization>,
161
162
163
164
    head_dim: Option<usize>,
    model_type: Option<String>,
    vision_config: Option<VisionConfig>,
    is_encoder_decoder: bool,
165
166
167
168
169
170
171
172
}

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);
173
        let quantize = other.quantization_config.and_then(|q| q.quant_method);
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
        let head_dim = other.head_dim.or_else(|| {
            match (other.hidden_size, other.n_embd, other.num_attention_heads) {
                (Some(hidden_size), _, Some(num_attention_heads))
                    if hidden_size % num_attention_heads == 0 =>
                {
                    Some(hidden_size / num_attention_heads)
                }
                // Legacy
                (_, Some(hidden_size), Some(num_attention_heads))
                    if hidden_size % num_attention_heads == 0 =>
                {
                    Some(hidden_size / num_attention_heads)
                }
                _ => None,
            }
        });
        let model_type = other.model_type;
        let vision_config = other.vision_config;
        let is_encoder_decoder = other.is_encoder_decoder.unwrap_or(false);
193
194
        Config {
            max_position_embeddings,
195
            quantize,
196
197
198
199
            head_dim,
            model_type,
            vision_config,
            is_encoder_decoder,
200
201
202
203
        }
    }
}

204
205
#[derive(Clone, Copy, Debug, ValueEnum, Deserialize)]
#[serde(rename_all = "kebab-case")]
206
enum Quantization {
207
    /// 4 bit quantization. Requires a specific AWQ quantized model:
208
    ///   <https://hf.co/models?search=awq>.
209
    /// Should replace GPTQ models wherever possible because of the better latency
210
211
212
    Awq,
    /// 8 bit quantization, doesn't require specific model.
    /// Should be a drop-in replacement to bitsandbytes with much better performance.
213
    /// Kernels are from <https://github.com/NetEase-FuXi/EETQ.git>
214
    Eetq,
215
216
217
218
    /// 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,
219
    /// 4 bit quantization. Requires a specific GTPQ quantized model: <https://hf.co/models?search=gptq>.
220
    /// text-generation-inference will use exllama (faster) kernels wherever possible, and use
221
222
223
    /// triton kernel (wider support) when it's not.
    /// AWQ has faster kernels.
    Gptq,
224
225
    /// 4 bit quantization. Requires a specific Marlin quantized model: <https://hf.co/models?search=marlin>.
    Marlin,
226
227
    /// 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.
228
229
230
231
    // #[deprecated(
    //     since = "1.1.0",
    //     note = "Use `eetq` instead, which provides better latencies overall and is drop-in in most cases"
    // )]
232
    Bitsandbytes,
233
234
    /// 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.
235
    BitsandbytesNf4,
236
237
    /// Bitsandbytes 4bit. nf4 should be preferred in most cases but maybe this one has better
    /// perplexity performance for you model
238
    BitsandbytesFp4,
Nicolas Patry's avatar
Nicolas Patry committed
239
240
241
242
243
    /// [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,
244
245
246
247
248
249
}

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 {
250
251
            #[allow(deprecated)]
            // Use `eetq` instead, which provides better latencies overall and is drop-in in most cases
252
253
254
            Quantization::Bitsandbytes => {
                write!(f, "bitsandbytes")
            }
255
            Quantization::BitsandbytesNf4 => {
Nicolas Patry's avatar
Nicolas Patry committed
256
257
                write!(f, "bitsandbytes-nf4")
            }
258
            Quantization::BitsandbytesFp4 => {
Nicolas Patry's avatar
Nicolas Patry committed
259
260
                write!(f, "bitsandbytes-fp4")
            }
261
262
263
            Quantization::Exl2 => {
                write!(f, "exl2")
            }
264
265
266
            Quantization::Gptq => {
                write!(f, "gptq")
            }
267
268
269
            Quantization::Marlin => {
                write!(f, "marlin")
            }
270
271
272
            Quantization::Awq => {
                write!(f, "awq")
            }
273
274
275
            Quantization::Eetq => {
                write!(f, "eetq")
            }
Nicolas Patry's avatar
Nicolas Patry committed
276
277
278
            Quantization::Fp8 => {
                write!(f, "fp8")
            }
279
280
281
282
        }
    }
}

283
284
285
#[derive(Clone, Copy, Debug, ValueEnum)]
enum Dtype {
    Float16,
286
    #[clap(name = "bfloat16")]
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
    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")
            }
        }
    }
}

304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
#[derive(Clone, Copy, Debug, ValueEnum)]
enum KVCacheDtype {
    #[clap(name = "fp8_e5m2")]
    Fp8e5m2,
}

impl std::fmt::Display for KVCacheDtype {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            KVCacheDtype::Fp8e5m2 => {
                write!(f, "fp8_e5m2")
            }
        }
    }
}

Nicolas Patry's avatar
Nicolas Patry committed
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
#[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")
            }
        }
    }
}

340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
#[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
367
368
369
370
/// App Configuration
#[derive(Parser, Debug)]
#[clap(author, version, about, long_about = None)]
struct Args {
371
372
373
374
375
    /// 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
376
    #[clap(default_value = "bigscience/bloom-560m", long, env)]
377
    model_id: String,
378
379
380

    /// 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
381
    #[clap(long, env)]
382
    revision: Option<String>,
383

384
385
386
387
388
    /// The number of tokenizer workers used for payload validation and truncation inside the
    /// router.
    #[clap(default_value = "2", long, env)]
    validation_workers: usize,

389
    /// Whether to shard the model across multiple GPUs
390
391
    /// By default text-generation-inference will use all available GPUs to run
    /// the model. Setting it to `false` deactivates `num_shard`.
392
393
    #[clap(long, env)]
    sharded: Option<bool>,
394
395

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

402
403
404
405
406
    /// 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.
407
408
    #[clap(long, env, value_enum)]
    quantize: Option<Quantization>,
409

Nicolas Patry's avatar
Nicolas Patry committed
410
411
412
413
414
415
416
    /// 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>,

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

421
422
423
424
425
426
    /// 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
    /// the only supported value is `fp8_e5m2` on CUDA.
    #[clap(long, env, value_enum)]
    kv_cache_dtype: Option<KVCacheDtype>,

427
428
429
430
431
432
    /// 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,

433
434
435
    /// 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
436
437
    #[clap(default_value = "128", long, env)]
    max_concurrent_requests: usize,
438
439
440
441

    /// 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
442
443
    #[clap(default_value = "2", long, env)]
    max_best_of: usize,
444
445
446
447
448
449

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

Nicolas Patry's avatar
Nicolas Patry committed
453
    /// This is the maximum allowed value for clients to set `top_n_tokens`.
454
    /// `top_n_tokens` is used to return information about the the `n` most likely
Nicolas Patry's avatar
Nicolas Patry committed
455
456
457
458
459
460
    /// 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,

461
462
463
464
    /// 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.
465
466
467
468
469
470
471
    /// Default to min(max_position_embeddings - 1, 4095)
    #[clap(long, env)]
    max_input_tokens: Option<usize>,

    /// Legacy version of [`Args::max_input_tokens`].
    #[clap(long, env)]
    max_input_length: Option<usize>,
472
473
474
475
476
477
478
479
480

    /// 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.
481
482
483
    /// Default to min(max_position_embeddings, 4096)
    #[clap(long, env)]
    max_total_tokens: Option<usize>,
484
485
486
487
488
489
490
491
492
493
494

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

498
499
500
    /// 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.
501
502
503
    /// Default to `max_input_tokens + 50` to give a bit of room.
    #[clap(long, env)]
    max_batch_prefill_tokens: Option<u32>,
504

505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
    /// **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.
522
523
    #[clap(long, env)]
    max_batch_total_tokens: Option<u32>,
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541

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

545
546
547
548
549
    /// 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>,

550
551
    /// Specify the batch sizes to compute cuda graphs for.
    /// Use "0" to disable.
552
553
554
    /// Default = "1,2,4,8,16,32"
    #[clap(long, env, value_delimiter = ',')]
    cuda_graphs: Option<Vec<usize>>,
555

556
557
558
559
    /// The IP address to listen on
    #[clap(default_value = "0.0.0.0", long, env)]
    hostname: String,

560
    /// The port to listen on.
561
    #[clap(default_value = "3000", long, short, env)]
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
562
    port: u16,
563
564
565

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

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

    /// The address the master port will listen on. (setting used by torch distributed)
574
    #[clap(default_value = "29500", long, env)]
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
575
    master_port: usize,
576
577
578

    /// The location of the huggingface hub cache.
    /// Used to override the location if you want to provide a mounted disk for instance
579
    #[clap(long, env)]
580
    huggingface_hub_cache: Option<String>,
581
582
583

    /// The location of the huggingface hub cache.
    /// Used to override the location if you want to provide a mounted disk for instance
584
585
    #[clap(long, env)]
    weights_cache_override: Option<String>,
586
587
588
589
590

    /// 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.
591
    #[clap(long, env)]
592
    disable_custom_kernels: bool,
593

594
595
596
597
598
    /// 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
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
    /// 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>,

619
    /// Outputs the logs in JSON format (useful for telemetry)
620
    #[clap(long, env)]
621
    json_output: bool,
622

623
624
    #[clap(long, env)]
    otlp_endpoint: Option<String>,
625

626
627
628
    #[clap(default_value = "text-generation-inference.router", long, env)]
    otlp_service_name: String,

629
630
    #[clap(long, env)]
    cors_allow_origin: Vec<String>,
Erik Kaunismäki's avatar
Erik Kaunismäki committed
631
632
633
634

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

635
636
637
638
    #[clap(long, env)]
    watermark_gamma: Option<f32>,
    #[clap(long, env)]
    watermark_delta: Option<f32>,
639

640
641
642
643
644
645
646
647
    /// Enable ngrok tunneling
    #[clap(long, env)]
    ngrok: bool,

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

648
    /// ngrok edge
649
    #[clap(long, env)]
650
    ngrok_edge: Option<String>,
651

652
653
654
655
656
    /// 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
657
658
659
660
661
    /// 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,

662
663
664
    /// Display a lot of information about your runtime environment
    #[clap(long, short, action)]
    env: bool,
665
666
667
668

    /// 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
669
670
671
672
673

    /// 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>,
674

675
676
677
678
679
    /// 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,
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
680
681
}

682
683
684
#[derive(Debug)]
enum ShardStatus {
    Ready,
685
    Failed(usize),
686
}
687

688
689
690
691
#[allow(clippy::too_many_arguments)]
fn shard_manager(
    model_id: String,
    revision: Option<String>,
692
    quantize: Option<Quantization>,
Nicolas Patry's avatar
Nicolas Patry committed
693
    speculate: Option<usize>,
694
    dtype: Option<Dtype>,
695
    kv_cache_dtype: Option<KVCacheDtype>,
696
    trust_remote_code: bool,
697
698
699
700
701
702
703
704
705
706
    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>,
707
    cuda_graphs: Vec<usize>,
708
    cuda_memory_fraction: f32,
Nicolas Patry's avatar
Nicolas Patry committed
709
710
    rope_scaling: Option<RopeScaling>,
    rope_factor: Option<f32>,
711
712
    max_total_tokens: usize,
    max_batch_size: Option<usize>,
713
    max_input_tokens: usize,
drbh's avatar
drbh committed
714
    lora_adapters: Option<String>,
715
    otlp_endpoint: Option<String>,
716
    otlp_service_name: String,
717
    log_level: LevelFilter,
718
    status_sender: mpsc::Sender<ShardStatus>,
719
    shutdown: Arc<AtomicBool>,
720
721
    _shutdown_sender: mpsc::Sender<()>,
) {
722
723
724
    // Enter shard-manager tracing span
    let _span = tracing::span!(tracing::Level::INFO, "shard-manager", rank = rank).entered();

725
726
727
728
    // Get UDS path
    let uds_string = format!("{uds_path}-{rank}");
    let uds = Path::new(&uds_string);
    // Clean previous runs
729
730
731
    if uds.exists() {
        fs::remove_file(uds).unwrap();
    }
732
733

    // Process args
OlivierDehaene's avatar
OlivierDehaene committed
734
    let mut shard_args = vec![
735
736
737
738
739
        "serve".to_string(),
        model_id,
        "--uds-path".to_string(),
        uds_path,
        "--logger-level".to_string(),
740
        log_level.to_string().to_uppercase(),
741
742
743
        "--json-output".to_string(),
    ];

744
745
    // Activate trust remote code
    if trust_remote_code {
OlivierDehaene's avatar
OlivierDehaene committed
746
        shard_args.push("--trust-remote-code".to_string());
747
748
    }

749
750
    // Activate tensor parallelism
    if world_size > 1 {
OlivierDehaene's avatar
OlivierDehaene committed
751
        shard_args.push("--sharded".to_string());
752
753
    }

754
    if let Some(quantize) = quantize {
OlivierDehaene's avatar
OlivierDehaene committed
755
756
        shard_args.push("--quantize".to_string());
        shard_args.push(quantize.to_string())
757
    }
758

Nicolas Patry's avatar
Nicolas Patry committed
759
760
761
762
763
    if let Some(speculate) = speculate {
        shard_args.push("--speculate".to_string());
        shard_args.push(speculate.to_string())
    }

764
    if let Some(dtype) = dtype {
OlivierDehaene's avatar
OlivierDehaene committed
765
766
        shard_args.push("--dtype".to_string());
        shard_args.push(dtype.to_string())
767
768
    }

769
770
771
772
773
    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())
    }

774
775
    // Model optional revision
    if let Some(revision) = revision {
OlivierDehaene's avatar
OlivierDehaene committed
776
777
        shard_args.push("--revision".to_string());
        shard_args.push(revision)
778
    }
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
779

Nicolas Patry's avatar
Nicolas Patry committed
780
781
782
783
784
785
    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)),
    };
786

787
    // OpenTelemetry Endpoint
788
    if let Some(otlp_endpoint) = otlp_endpoint {
OlivierDehaene's avatar
OlivierDehaene committed
789
790
        shard_args.push("--otlp-endpoint".to_string());
        shard_args.push(otlp_endpoint);
791
792
    }

793
794
795
796
    // OpenTelemetry Service Name
    shard_args.push("--otlp-service-name".to_string());
    shard_args.push(otlp_service_name);

797
798
799
800
    // In case we use sliding window, we may ignore the sliding in flash for some backends depending on the parameter.
    shard_args.push("--max-input-tokens".to_string());
    shard_args.push(max_input_tokens.to_string());

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

804
805
806
    // Remove LOG_LEVEL if present
    envs.retain(|(name, _)| name != "LOG_LEVEL");

807
    // Torch Distributed Env vars
OlivierDehaene's avatar
OlivierDehaene committed
808
809
810
811
    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()));
812
    envs.push(("TORCH_NCCL_AVOID_RECORD_STREAMS".into(), "1".into()));
813

814
815
816
817
818
819
    // CUDA memory fraction
    envs.push((
        "CUDA_MEMORY_FRACTION".into(),
        cuda_memory_fraction.to_string().into(),
    ));

820
    // Safetensors load fast
OlivierDehaene's avatar
OlivierDehaene committed
821
    envs.push(("SAFETENSORS_FAST_GPU".into(), "1".into()));
822

823
824
825
    // Disable progress bar
    envs.push(("HF_HUB_DISABLE_PROGRESS_BARS".into(), "1".into()));

826
827
    // 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
828
    envs.push((
829
830
831
832
833
834
        "HF_HUB_ENABLE_HF_TRANSFER".into(),
        enable_hf_transfer.into(),
    ));

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

Nicolas Patry's avatar
Nicolas Patry committed
838
839
840
841
842
843
844
845
846
    // 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()));
    }

847
848
849
850
851
852
853
854
    envs.push((
        "MAX_TOTAL_TOKENS".into(),
        max_total_tokens.to_string().into(),
    ));
    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
855
856
857
858
859
    // Lora Adapters
    if let Some(lora_adapters) = lora_adapters {
        envs.push(("LORA_ADAPTERS".into(), lora_adapters.into()));
    }

860
861
862
    // 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
863
        envs.push(("HUGGINGFACE_HUB_CACHE".into(), huggingface_hub_cache.into()));
864
865
866
867
868
    };

    // 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
869
        envs.push((
870
871
872
873
874
            "WEIGHTS_CACHE_OVERRIDE".into(),
            weights_cache_override.into(),
        ));
    };

875
    // Enable experimental support for cuda graphs
876
877
878
879
880
881
882
883
884
885
    if !cuda_graphs.is_empty() {
        envs.push((
            "CUDA_GRAPHS".into(),
            cuda_graphs
                .into_iter()
                .map(|c| c.to_string())
                .collect::<Vec<_>>()
                .join(",")
                .into(),
        ));
886
887
    }

888
889
    // If disable_custom_kernels is true, pass it to the shard as an env var
    if disable_custom_kernels {
OlivierDehaene's avatar
OlivierDehaene committed
890
        envs.push(("DISABLE_CUSTOM_KERNELS".into(), "True".into()))
891
892
893
894
    }

    // Watermark Gamma
    if let Some(watermark_gamma) = watermark_gamma {
OlivierDehaene's avatar
OlivierDehaene committed
895
        envs.push(("WATERMARK_GAMMA".into(), watermark_gamma.to_string().into()))
896
897
898
899
    }

    // Watermark Delta
    if let Some(watermark_delta) = watermark_delta {
OlivierDehaene's avatar
OlivierDehaene committed
900
        envs.push(("WATERMARK_DELTA".into(), watermark_delta.to_string().into()))
901
902
903
    }

    // Start process
904
    tracing::info!("Starting shard");
905
    let mut p = match Command::new("text-generation-server")
OlivierDehaene's avatar
OlivierDehaene committed
906
        .args(shard_args)
907
        .env_clear()
OlivierDehaene's avatar
OlivierDehaene committed
908
        .envs(envs)
909
        .stdin(Stdio::piped())
910
911
912
913
914
        .stdout(Stdio::piped())
        .stderr(Stdio::piped())
        .process_group(0)
        .spawn()
    {
915
916
        Ok(p) => p,
        Err(err) => {
917
918
919
            if err.kind() == io::ErrorKind::NotFound {
                tracing::error!("text-generation-server not found in PATH");
                tracing::error!("Please install it with `make install-server`")
920
921
            }
            {
922
                tracing::error!("{}", err);
923
            }
924

925
            status_sender.send(ShardStatus::Failed(rank)).unwrap();
926
927
928
929
930
            return;
        }
    };

    // Redirect STDOUT to the console
931
    let mut pstdin = p.stdin.take().unwrap();
932
    let shard_stdout_reader = BufReader::new(p.stdout.take().unwrap());
933
    let shard_stderr_reader = BufReader::new(p.stderr.take().unwrap());
934

935
    //stdout tracing thread
936
    thread::spawn(move || {
937
        log_lines(shard_stdout_reader);
938
    });
939
940
941
    // 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
942
        for line in shard_stderr_reader.lines().map_while(Result::ok) {
943
944
945
            err_sender.send(line).unwrap_or(());
        }
    });
946
    // We read stdin in another thread as it seems that lines() can block in some cases
Nicolas Patry's avatar
Nicolas Patry committed
947
948
949
950
951
952
953
954
955
    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]);
                    }
956
957
                }
            }
Nicolas Patry's avatar
Nicolas Patry committed
958
959
        });
    }
960
961
962
963
964
965

    let mut ready = false;
    let start_time = Instant::now();
    let mut wait_time = Instant::now();
    loop {
        // Process exited
966
        if let Some(exit_status) = p.try_wait().unwrap() {
967
968
969
970
            let mut err = String::new();
            while let Ok(line) = err_receiver.recv_timeout(Duration::from_millis(10)) {
                err = err + "\n" + &line;
            }
971

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

974
            if let Some(signal) = exit_status.signal() {
975
976
977
                tracing::error!("Shard process was signaled to shutdown with signal {signal}");
            }

978
            status_sender.send(ShardStatus::Failed(rank)).unwrap();
979
980
981
982
            return;
        }

        // We received a shutdown signal
983
        if shutdown.load(Ordering::SeqCst) {
984
            terminate("shard", p, Duration::from_secs(90)).unwrap();
985
986
987
988
989
            return;
        }

        // Shard is ready
        if uds.exists() && !ready {
990
            tracing::info!("Shard ready in {:?}", start_time.elapsed());
991
992
993
            status_sender.send(ShardStatus::Ready).unwrap();
            ready = true;
        } else if !ready && wait_time.elapsed() > Duration::from_secs(10) {
994
            tracing::info!("Waiting for shard to be ready...");
995
996
997
998
999
1000
            wait_time = Instant::now();
        }
        sleep(Duration::from_millis(100));
    }
}

1001
fn shutdown_shards(shutdown: Arc<AtomicBool>, shutdown_receiver: &mpsc::Receiver<()>) {
1002
1003
1004
    tracing::info!("Shutting down shards");
    // Update shutdown value to true
    // This will be picked up by the shard manager
1005
    shutdown.store(true, Ordering::SeqCst);
1006
1007
1008
1009
1010
1011
1012

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

fn num_cuda_devices() -> Option<usize> {
1013
1014
    let devices = match env::var("CUDA_VISIBLE_DEVICES") {
        Ok(devices) => devices,
1015
1016
1017
        Err(_) => match env::var("NVIDIA_VISIBLE_DEVICES") {
            Ok(devices) => devices,
            Err(_) => env::var("ZE_AFFINITY_MASK").ok()?,
Nicolas Patry's avatar
Nicolas Patry committed
1018
        },
1019
    };
1020
1021
    let n_devices = devices.split(',').count();
    Some(n_devices)
1022
1023
1024
1025
1026
1027
1028
1029
1030
1031
1032
1033
1034
1035
1036
1037
1038
1039
1040
1041
1042
1043
1044
1045
1046
1047
1048
1049
1050
1051
1052
1053
1054
}

#[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 {
1055
1056
1057
1058
1059
1060
1061
            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()),
1062
1063
1064
1065
        }
    }
}

1066
impl TryFrom<&[u8]> for PythonLogMessage {
1067
1068
    type Error = serde_json::Error;

1069
1070
    fn try_from(value: &[u8]) -> Result<Self, Self::Error> {
        serde_json::from_slice::<Self>(value)
1071
1072
1073
    }
}

1074
1075
1076
1077
1078
1079
1080
1081
1082
1083
1084
1085
1086
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(_) => {
1087
1088
1089
1090
1091
1092
                            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();
1093
1094
1095
1096
1097
                            }
                        }
                    }
                }
            }
1098
1099
1100
1101
        }
    }
}

1102
1103
1104
1105
fn find_num_shards(
    sharded: Option<bool>,
    num_shard: Option<usize>,
) -> Result<usize, LauncherError> {
1106
1107
1108
1109
    // 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
1110
            tracing::info!("Parsing num_shard from CUDA_VISIBLE_DEVICES/NVIDIA_VISIBLE_DEVICES/ZE_AFFINITY_MASK");
1111
            let n_devices = num_cuda_devices()
1112
                .expect("--num-shard and CUDA_VISIBLE_DEVICES/NVIDIA_VISIBLE_DEVICES/ZE_AFFINITY_MASK are not set");
1113
            if n_devices <= 1 {
1114
1115
1116
                return Err(LauncherError::NotEnoughCUDADevices(format!(
                    "`sharded` is true but only found {n_devices} CUDA devices"
                )));
1117
            }
1118
            n_devices
1119
        }
1120
1121
1122
        (Some(true), Some(num_shard)) => {
            // we can't have only one shard while sharded
            if num_shard <= 1 {
1123
1124
1125
                return Err(LauncherError::ArgumentValidation(
                    "`sharded` is true but `num_shard` <= 1".to_string(),
                ));
1126
1127
            }
            num_shard
1128
        }
1129
1130
1131
1132
        (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,
1133
    };
1134
    if num_shard < 1 {
1135
1136
1137
        return Err(LauncherError::ArgumentValidation(
            "`num_shard` cannot be < 1".to_string(),
        ));
1138
    }
1139
    Ok(num_shard)
1140
}
1141

1142
#[derive(Debug, Error)]
1143
enum LauncherError {
1144
    #[error("Invalid argument: {0}")]
1145
    ArgumentValidation(String),
1146
    #[error("not enough cuda devices: {0}")]
1147
    NotEnoughCUDADevices(String),
1148
    #[error("Download error")]
1149
    DownloadError,
1150
    #[error("Shard cannot start")]
1151
    ShardCannotStart,
1152
    #[error("Shard disconnected")]
1153
    ShardDisconnected,
1154
    #[error("Shard failed")]
1155
    ShardFailed,
1156
    #[error("Webserver failed")]
1157
    WebserverFailed,
1158
    #[error("Webserver cannot start")]
1159
1160
    WebserverCannotStart,
}
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1161

1162
1163
1164
1165
1166
1167
1168
1169
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>,
) -> Result<(), LauncherError> {
1170
1171
1172
    // Enter download tracing span
    let _span = tracing::span!(tracing::Level::INFO, "download").entered();

OlivierDehaene's avatar
OlivierDehaene committed
1173
    let mut download_args = vec![
1174
        "download-weights".to_string(),
1175
        model_id.to_string(),
1176
1177
1178
1179
1180
1181
        "--extension".to_string(),
        ".safetensors".to_string(),
        "--logger-level".to_string(),
        "INFO".to_string(),
        "--json-output".to_string(),
    ];
1182

1183
    // Model optional revision
1184
    if let Some(revision) = &revision {
OlivierDehaene's avatar
OlivierDehaene committed
1185
1186
        download_args.push("--revision".to_string());
        download_args.push(revision.to_string())
1187
    }
1188

1189
    // Trust remote code for automatic peft fusion
1190
    if trust_remote_code {
1191
1192
1193
        download_args.push("--trust-remote-code".to_string());
    }

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

1197
1198
1199
    // Remove LOG_LEVEL if present
    envs.retain(|(name, _)| name != "LOG_LEVEL");

1200
1201
1202
    // Disable progress bar
    envs.push(("HF_HUB_DISABLE_PROGRESS_BARS".into(), "1".into()));

1203
    // If huggingface_hub_cache is set, pass it to the download process
1204
    // Useful when running inside a docker container
1205
    if let Some(ref huggingface_hub_cache) = huggingface_hub_cache {
OlivierDehaene's avatar
OlivierDehaene committed
1206
        envs.push(("HUGGINGFACE_HUB_CACHE".into(), huggingface_hub_cache.into()));
1207
    };
1208

1209
1210
    // 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
1211
    envs.push((
1212
1213
1214
        "HF_HUB_ENABLE_HF_TRANSFER".into(),
        enable_hf_transfer.into(),
    ));
1215

1216
1217
    // Parse Inference API token
    if let Ok(api_token) = env::var("HF_API_TOKEN") {
1218
        envs.push(("HF_TOKEN".into(), api_token.into()))
1219
    };
1220

1221
1222
    // If args.weights_cache_override is some, pass it to the download process
    // Useful when running inside a HuggingFace Inference Endpoint
1223
    if let Some(weights_cache_override) = &weights_cache_override {
OlivierDehaene's avatar
OlivierDehaene committed
1224
        envs.push((
1225
1226
1227
1228
1229
            "WEIGHTS_CACHE_OVERRIDE".into(),
            weights_cache_override.into(),
        ));
    };

1230
    // Start process
1231
    tracing::info!("Starting check and download process for {model_id}");
1232
    let mut download_process = match Command::new("text-generation-server")
OlivierDehaene's avatar
OlivierDehaene committed
1233
        .args(download_args)
1234
        .env_clear()
OlivierDehaene's avatar
OlivierDehaene committed
1235
        .envs(envs)
1236
1237
1238
1239
1240
        .stdout(Stdio::piped())
        .stderr(Stdio::piped())
        .process_group(0)
        .spawn()
    {
1241
1242
        Ok(p) => p,
        Err(err) => {
1243
1244
1245
            if err.kind() == io::ErrorKind::NotFound {
                tracing::error!("text-generation-server not found in PATH");
                tracing::error!("Please install it with `make install-server`")
1246
1247
            } else {
                tracing::error!("{}", err);
1248
            }
1249

1250
1251
1252
            return Err(LauncherError::DownloadError);
        }
    };
1253

1254
    let download_stdout = BufReader::new(download_process.stdout.take().unwrap());
1255

1256
    thread::spawn(move || {
1257
        log_lines(download_stdout);
1258
1259
1260
1261
1262
1263
1264
    });

    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
1265
        for line in download_stderr.lines().map_while(Result::ok) {
1266
1267
            err_sender.send(line).unwrap_or(());
        }
1268
    });
1269

1270
    loop {
1271
1272
        if let Some(status) = download_process.try_wait().unwrap() {
            if status.success() {
1273
                tracing::info!("Successfully downloaded weights for {model_id}");
1274
                break;
1275
            }
1276
1277

            let mut err = String::new();
1278
1279
1280
1281
            while let Ok(line) = err_receiver.recv_timeout(Duration::from_millis(10)) {
                err = err + "\n" + &line;
            }

1282
1283
1284
1285
1286
1287
1288
1289
1290
            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);
1291
        }
1292
        if !running.load(Ordering::SeqCst) {
OlivierDehaene's avatar
OlivierDehaene committed
1293
            terminate("download", download_process, Duration::from_secs(10)).unwrap();
1294
1295
1296
            return Ok(());
        }
        sleep(Duration::from_millis(100));
1297
    }
1298
1299
    Ok(())
}
1300

1301
#[allow(clippy::too_many_arguments)]
1302
1303
1304
fn spawn_shards(
    num_shard: usize,
    args: &Args,
1305
    cuda_graphs: Vec<usize>,
1306
    max_total_tokens: usize,
1307
    max_input_tokens: usize,
1308
    quantize: Option<Quantization>,
1309
    max_log_level: LevelFilter,
1310
    shutdown: Arc<AtomicBool>,
1311
1312
1313
1314
1315
1316
    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
1317
1318
    // Start shard processes
    for rank in 0..num_shard {
1319
1320
1321
1322
1323
1324
        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
1325
1326
1327
        let status_sender = status_sender.clone();
        let shutdown = shutdown.clone();
        let shutdown_sender = shutdown_sender.clone();
1328
        let otlp_endpoint = args.otlp_endpoint.clone();
1329
        let otlp_service_name = args.otlp_service_name.clone();
Nicolas Patry's avatar
Nicolas Patry committed
1330
        let speculate = args.speculate;
1331
        let dtype = args.dtype;
1332
        let kv_cache_dtype = args.kv_cache_dtype;
1333
        let trust_remote_code = args.trust_remote_code;
1334
1335
1336
1337
        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;
1338
        let cuda_graphs_clone = cuda_graphs.clone();
1339
        let cuda_memory_fraction = args.cuda_memory_fraction;
Nicolas Patry's avatar
Nicolas Patry committed
1340
1341
        let rope_scaling = args.rope_scaling;
        let rope_factor = args.rope_factor;
1342
        let max_batch_size = args.max_batch_size;
drbh's avatar
drbh committed
1343
        let lora_adapters = args.lora_adapters.clone();
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1344
1345
        thread::spawn(move || {
            shard_manager(
1346
                model_id,
1347
                revision,
1348
                quantize,
Nicolas Patry's avatar
Nicolas Patry committed
1349
                speculate,
1350
                dtype,
1351
                kv_cache_dtype,
1352
                trust_remote_code,
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1353
1354
1355
1356
1357
                uds_path,
                rank,
                num_shard,
                master_addr,
                master_port,
1358
1359
                huggingface_hub_cache,
                weights_cache_override,
1360
                disable_custom_kernels,
1361
1362
                watermark_gamma,
                watermark_delta,
1363
                cuda_graphs_clone,
1364
                cuda_memory_fraction,
Nicolas Patry's avatar
Nicolas Patry committed
1365
1366
                rope_scaling,
                rope_factor,
1367
1368
                max_total_tokens,
                max_batch_size,
1369
                max_input_tokens,
drbh's avatar
drbh committed
1370
                lora_adapters,
1371
                otlp_endpoint,
1372
                otlp_service_name,
1373
                max_log_level,
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1374
1375
1376
1377
1378
1379
1380
1381
1382
1383
1384
1385
1386
1387
1388
1389
1390
1391
1392
1393
1394
                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));
            }
1395
            Ok(ShardStatus::Failed(rank)) => {
1396
                tracing::error!("Shard {rank} failed to start");
1397
                shutdown_shards(shutdown, shutdown_receiver);
1398
                return Err(LauncherError::ShardCannotStart);
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1399
1400
1401
            }
            Err(TryRecvError::Disconnected) => {
                tracing::error!("Shard status channel disconnected");
1402
                shutdown_shards(shutdown, shutdown_receiver);
1403
                return Err(LauncherError::ShardDisconnected);
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1404
1405
1406
            }
        }
    }
1407
1408
    Ok(())
}
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1409

1410
1411
1412
1413
1414
1415
1416
1417
1418
1419
1420
1421
fn compute_type(num_shard: usize) -> Option<String> {
    let output = Command::new("nvidia-smi")
        .args(["--query-gpu=gpu_name", "--format=csv"])
        .output()
        .ok()?;
    let output = String::from_utf8(output.stdout).ok()?;
    let fullname = output.split('\n').nth(1)?;
    let cardname = fullname.replace(' ', "-").to_lowercase();
    let compute_type = format!("{num_shard}-{cardname}");
    Some(compute_type)
}

1422
fn spawn_webserver(
1423
    num_shard: usize,
1424
    args: Args,
1425
1426
1427
    max_input_tokens: usize,
    max_total_tokens: usize,
    max_batch_prefill_tokens: u32,
1428
    shutdown: Arc<AtomicBool>,
1429
    shutdown_receiver: &mpsc::Receiver<()>,
1430
) -> Result<Child, LauncherError> {
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1431
1432
1433
    // All shard started
    // Start webserver
    tracing::info!("Starting Webserver");
OlivierDehaene's avatar
OlivierDehaene committed
1434
    let mut router_args = vec![
1435
1436
        "--max-client-batch-size".to_string(),
        args.max_client_batch_size.to_string(),
1437
        "--max-concurrent-requests".to_string(),
1438
        args.max_concurrent_requests.to_string(),
1439
        "--max-best-of".to_string(),
1440
        args.max_best_of.to_string(),
1441
        "--max-stop-sequences".to_string(),
1442
        args.max_stop_sequences.to_string(),
Nicolas Patry's avatar
Nicolas Patry committed
1443
1444
        "--max-top-n-tokens".to_string(),
        args.max_top_n_tokens.to_string(),
1445
1446
        "--max-input-tokens".to_string(),
        max_input_tokens.to_string(),
1447
        "--max-total-tokens".to_string(),
1448
        max_total_tokens.to_string(),
1449
        "--max-batch-prefill-tokens".to_string(),
1450
        max_batch_prefill_tokens.to_string(),
1451
        "--waiting-served-ratio".to_string(),
1452
        args.waiting_served_ratio.to_string(),
1453
        "--max-waiting-tokens".to_string(),
1454
        args.max_waiting_tokens.to_string(),
1455
1456
        "--validation-workers".to_string(),
        args.validation_workers.to_string(),
1457
1458
        "--hostname".to_string(),
        args.hostname.to_string(),
1459
        "--port".to_string(),
1460
        args.port.to_string(),
1461
        "--master-shard-uds-path".to_string(),
1462
        format!("{}-0", args.shard_uds_path),
1463
        "--tokenizer-name".to_string(),
1464
        args.model_id,
1465
1466
    ];

1467
    // Pass usage stats flags to router
1468
1469
    router_args.push("--usage-stats".to_string());
    router_args.push(args.usage_stats.to_string());
1470

drbh's avatar
drbh committed
1471
1472
1473
1474
1475
    // Grammar support
    if args.disable_grammar_support {
        router_args.push("--disable-grammar-support".to_string());
    }

1476
1477
1478
1479
1480
1481
    // 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());
    }

1482
1483
1484
1485
1486
1487
    // 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());
    }

1488
1489
1490
1491
1492
1493
    // 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());
    }

1494
1495
    // Model optional revision
    if let Some(ref revision) = args.revision {
OlivierDehaene's avatar
OlivierDehaene committed
1496
1497
        router_args.push("--revision".to_string());
        router_args.push(revision.to_string())
1498
1499
    }

1500
    if args.json_output {
OlivierDehaene's avatar
OlivierDehaene committed
1501
        router_args.push("--json-output".to_string());
1502
1503
    }

1504
    // OpenTelemetry
1505
    if let Some(otlp_endpoint) = args.otlp_endpoint {
OlivierDehaene's avatar
OlivierDehaene committed
1506
1507
        router_args.push("--otlp-endpoint".to_string());
        router_args.push(otlp_endpoint);
1508
1509
    }

1510
1511
1512
1513
1514
    // OpenTelemetry
    let otlp_service_name = args.otlp_service_name;
    router_args.push("--otlp-service-name".to_string());
    router_args.push(otlp_service_name);

1515
1516
    // CORS origins
    for origin in args.cors_allow_origin.into_iter() {
OlivierDehaene's avatar
OlivierDehaene committed
1517
1518
        router_args.push("--cors-allow-origin".to_string());
        router_args.push(origin);
1519
1520
    }

Erik Kaunismäki's avatar
Erik Kaunismäki committed
1521
1522
1523
1524
1525
    // API Key
    if let Some(api_key) = args.api_key {
        router_args.push("--api-key".to_string());
        router_args.push(api_key);
    }
1526
1527
    // Ngrok
    if args.ngrok {
OlivierDehaene's avatar
OlivierDehaene committed
1528
1529
        router_args.push("--ngrok".to_string());
        router_args.push("--ngrok-authtoken".to_string());
1530
1531
1532
        router_args.push(args.ngrok_authtoken.unwrap());
        router_args.push("--ngrok-edge".to_string());
        router_args.push(args.ngrok_edge.unwrap());
1533
1534
    }

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

1538
1539
    // Parse Inference API token
    if let Ok(api_token) = env::var("HF_API_TOKEN") {
1540
        envs.push(("HF_TOKEN".into(), api_token.into()))
1541
    };
1542

1543
1544
1545
1546
1547
1548
1549
    // 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()))
    }

1550
    let mut webserver = match Command::new("text-generation-router")
OlivierDehaene's avatar
OlivierDehaene committed
1551
1552
        .args(router_args)
        .envs(envs)
1553
1554
1555
1556
1557
        .stdout(Stdio::piped())
        .stderr(Stdio::piped())
        .process_group(0)
        .spawn()
    {
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1558
1559
        Ok(p) => p,
        Err(err) => {
1560
            tracing::error!("Failed to start webserver: {}", err);
1561
1562
1563
            if err.kind() == io::ErrorKind::NotFound {
                tracing::error!("text-generation-router not found in PATH");
                tracing::error!("Please install it with `make install-router`")
1564
1565
            } else {
                tracing::error!("{}", err);
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1566
            }
1567

1568
            shutdown_shards(shutdown, shutdown_receiver);
1569
            return Err(LauncherError::WebserverCannotStart);
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1570
1571
1572
        }
    };

1573
1574
1575
    // Redirect STDOUT and STDERR to the console
    let webserver_stdout = webserver.stdout.take().unwrap();
    let webserver_stderr = webserver.stderr.take().unwrap();
1576
1577

    thread::spawn(move || {
1578
1579
        let stdout = BufReader::new(webserver_stdout);
        let stderr = BufReader::new(webserver_stderr);
1580
        for line in stdout.lines() {
1581
            println!("{}", line.unwrap());
1582
        }
1583
1584
        for line in stderr.lines() {
            println!("{}", line.unwrap());
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1585
        }
1586
1587
1588
    });
    Ok(webserver)
}
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1589

OlivierDehaene's avatar
OlivierDehaene committed
1590
1591
1592
1593
1594
1595
1596
1597
1598
1599
1600
1601
1602
1603
1604
1605
1606
1607
1608
1609
1610
1611
1612
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)
}

1613
1614
fn main() -> Result<(), LauncherError> {
    // Pattern match configuration
1615
    let args: Args = Args::parse();
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1616

1617
    // Filter events with LOG_LEVEL
1618
1619
1620
1621
1622
1623
1624
1625
1626
1627
1628
1629
1630
1631
1632
1633
    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);
1634

1635
    if args.json_output {
1636
1637
1638
1639
        tracing_subscriber::fmt()
            .with_env_filter(env_filter)
            .json()
            .init();
1640
    } else {
1641
1642
1643
1644
        tracing_subscriber::fmt()
            .with_env_filter(env_filter)
            .compact()
            .init();
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1645
1646
    }

1647
1648
1649
1650
1651
    if args.env {
        let env_runtime = env_runtime::Env::new();
        tracing::info!("{}", env_runtime);
    }

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

1654
1655
1656
1657
1658
1659
1660
1661
1662
1663
1664
1665
1666
1667
    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 max_default = 4096;

    let max_position_embeddings = if let Some(config) = &config {
        if let Some(max_position_embeddings) = config.max_position_embeddings {
            if max_position_embeddings > max_default {
                let max = max_position_embeddings;
                if args.max_input_tokens.is_none()
                    && args.max_total_tokens.is_none()
                    && args.max_batch_prefill_tokens.is_none()
                {
                    tracing::info!("Model supports up to {max} but tgi will now set its default to {max_default} instead. This is to save VRAM by refusing large prompts in order to allow more users on the same hardware. You can increase that size using `--max-batch-prefill-tokens={} --max-total-tokens={max} --max-input-tokens={}`.", max + 50, max - 1);
1668
                }
1669
                max_default
1670
            } else {
1671
                max_position_embeddings
1672
            }
1673
1674
1675
1676
1677
1678
1679
1680
1681
1682
        } else {
            max_default
        }
    } else {
        max_default
    };
    let (prefix_caching, attention) = resolve_attention(&config, &args.lora_adapters);
    tracing::info!("Using attention {attention} - Prefix caching {prefix_caching}");
    std::env::set_var("USE_PREFIX_CACHING", prefix_caching);
    std::env::set_var("ATTENTION", attention);
1683
1684
1685
1686
1687
1688
1689
1690
1691
1692
1693
1694
1695
1696
1697
1698
1699
1700
1701
1702
1703
1704
1705
1706
1707
1708
1709
1710
1711
1712
1713
1714
1715
1716
1717
1718
1719
1720
1721
1722
1723
1724
1725

    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.",
                )));
            }
            (Some(max_input_tokens), None) | (None, Some(max_input_tokens)) => max_input_tokens,
            (None, None) => {
                let value = max_position_embeddings - 1;
                tracing::info!("Default `max_input_tokens` to {value}");
                value
            }
        }
    };
    let max_total_tokens = {
        match args.max_total_tokens {
            Some(max_total_tokens) => max_total_tokens,
            None => {
                let value = max_position_embeddings;
                tracing::info!("Default `max_total_tokens` to {value}");
                value
            }
        }
    };
    let max_batch_prefill_tokens = {
        match args.max_batch_prefill_tokens {
            Some(max_batch_prefill_tokens) => max_batch_prefill_tokens,
            None => {
                let value: u32 = if let Some(max_batch_size) = args.max_batch_size {
                    max_batch_size * max_input_tokens
                } else {
                    // Adding some edge in order to account for potential block_size alignement
                    // issue.
                    max_input_tokens + 50
                } as u32;
                tracing::info!("Default `max_batch_prefill_tokens` to {value}");
                value
            }
        }
    };

1726
    // Validate args
1727
    if max_input_tokens >= max_total_tokens {
1728
        return Err(LauncherError::ArgumentValidation(
1729
            "`max_input_tokens must be < `max_total_tokens`".to_string(),
1730
1731
        ));
    }
1732
    if max_input_tokens as u32 > max_batch_prefill_tokens {
1733
        return Err(LauncherError::ArgumentValidation(format!(
1734
1735
            "`max_batch_prefill_tokens` must be >= `max_input_tokens`. Given: {} and {}",
            max_batch_prefill_tokens, max_input_tokens
1736
1737
        )));
    }
1738

1739
1740
1741
1742
1743
    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
1744
        (Some(cuda_graphs), _) => cuda_graphs.iter().cloned().filter(|&c| c > 0).collect(),
1745
1746
1747
1748
1749
        #[allow(deprecated)]
        (
            None,
            Some(
                Quantization::Bitsandbytes
1750
1751
                | Quantization::BitsandbytesNf4
                | Quantization::BitsandbytesFp4,
1752
1753
            ),
        ) => {
1754
1755
1756
1757
1758
            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");
1759
1760
1761
1762
1763
1764
1765
1766
1767
            vec![]
        }
        _ => {
            let cuda_graphs = vec![1, 2, 4, 8, 16, 32];
            tracing::info!("Using default cuda graphs {cuda_graphs:?}");
            cuda_graphs
        }
    };

1768
1769
1770
1771
1772
    if args.validation_workers == 0 {
        return Err(LauncherError::ArgumentValidation(
            "`validation_workers` must be > 0".to_string(),
        ));
    }
1773
1774
1775
1776
1777
1778
    if args.trust_remote_code {
        tracing::warn!(
            "`trust_remote_code` is set. Trusting that model `{}` do not contain malicious code.",
            args.model_id
        );
    }
1779
1780

    let num_shard = find_num_shards(args.sharded, args.num_shard)?;
1781
    if num_shard > 1 {
1782
1783
1784
1785
1786
        if matches!(args.quantize, Some(Quantization::Exl2)) {
            return Err(LauncherError::ArgumentValidation(
                "Sharding is currently not supported with `exl2` quantization".into(),
            ));
        }
1787
        tracing::info!("Sharding model on {num_shard} processes");
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1788
1789
    }

1790
    if let Some(ref max_batch_total_tokens) = args.max_batch_total_tokens {
1791
        if max_batch_prefill_tokens > *max_batch_total_tokens {
1792
1793
            return Err(LauncherError::ArgumentValidation(format!(
                "`max_batch_prefill_tokens` must be <= `max_batch_total_tokens`. Given: {} and {}",
1794
                max_batch_prefill_tokens, max_batch_total_tokens
1795
1796
            )));
        }
1797
        if max_total_tokens as u32 > *max_batch_total_tokens {
1798
1799
            return Err(LauncherError::ArgumentValidation(format!(
                "`max_total_tokens` must be <= `max_batch_total_tokens`. Given: {} and {}",
1800
                max_total_tokens, max_batch_total_tokens
1801
1802
1803
1804
            )));
        }
    }

1805
1806
1807
1808
1809
1810
1811
1812
1813
1814
1815
1816
1817
1818
    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(),
            ));
        }
    }

1819
1820
1821
1822
1823
1824
1825
    // 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");
1826

1827
    // Download and convert model weights
1828
1829
1830
1831
1832
1833
1834
1835
1836
1837
1838
1839
    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(),
    )?;

    // Download and convert lora adapters if any
    if let Some(lora_adapters) = &args.lora_adapters {
        for adapter in lora_adapters.split(',') {
1840
1841
1842
1843
            // skip download if a path is provided
            if adapter.contains('=') {
                continue;
            }
1844
1845
1846
1847
1848
1849
1850
1851
1852
1853
1854
1855
1856
1857
1858
1859
1860
1861
1862
1863
1864
1865
1866
1867
1868
1869
1870
1871
1872
1873
1874

            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(),
                )?;
            } else {
                return Err(LauncherError::ArgumentValidation(format!(
                    "Invalid LoRA adapter format: {}",
                    adapter
                )));
            }
1875
1876
        }
    }
1877

OlivierDehaene's avatar
OlivierDehaene committed
1878
1879
1880
1881
1882
    if !running.load(Ordering::SeqCst) {
        // Launcher was asked to stop
        return Ok(());
    }

1883
    // Shared shutdown bool
1884
    let shutdown = Arc::new(AtomicBool::new(false));
1885
1886
1887
    // Shared shutdown channel
    // When shutting down, the main thread will wait for all senders to be dropped
    let (shutdown_sender, shutdown_receiver) = mpsc::channel();
1888

1889
1890
    // Shared channel to track shard status
    let (status_sender, status_receiver) = mpsc::channel();
1891

1892
1893
1894
    spawn_shards(
        num_shard,
        &args,
1895
        cuda_graphs,
1896
        max_total_tokens,
1897
        max_input_tokens,
1898
        quantize,
1899
        max_log_level,
1900
1901
1902
1903
1904
1905
1906
        shutdown.clone(),
        &shutdown_receiver,
        shutdown_sender,
        &status_receiver,
        status_sender,
        running.clone(),
    )?;
1907

1908
1909
1910
1911
1912
    // We might have received a termination signal
    if !running.load(Ordering::SeqCst) {
        shutdown_shards(shutdown, &shutdown_receiver);
        return Ok(());
    }
1913

1914
1915
1916
1917
1918
1919
1920
1921
1922
    let mut webserver = spawn_webserver(
        num_shard,
        args,
        max_input_tokens,
        max_total_tokens,
        max_batch_prefill_tokens,
        shutdown.clone(),
        &shutdown_receiver,
    )
1923
    .inspect_err(|_| {
1924
1925
        shutdown_shards(shutdown.clone(), &shutdown_receiver);
    })?;
1926
1927
1928
1929
1930

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

    while running.load(Ordering::SeqCst) {
1931
        if let Ok(ShardStatus::Failed(rank)) = status_receiver.try_recv() {
OlivierDehaene's avatar
OlivierDehaene committed
1932
            tracing::error!("Shard {rank} crashed");
1933
1934
1935
1936
            exit_code = Err(LauncherError::ShardFailed);
            break;
        };

1937
        match webserver.try_wait().unwrap() {
1938
1939
1940
1941
1942
1943
1944
1945
1946
            Some(_) => {
                tracing::error!("Webserver Crashed");
                shutdown_shards(shutdown, &shutdown_receiver);
                return Err(LauncherError::WebserverFailed);
            }
            None => {
                sleep(Duration::from_millis(100));
            }
        };
1947
    }
1948
1949

    // Graceful termination
OlivierDehaene's avatar
OlivierDehaene committed
1950
    terminate("webserver", webserver, Duration::from_secs(90)).unwrap();
1951
1952
1953
    shutdown_shards(shutdown, &shutdown_receiver);

    exit_code
1954
}