main.rs 67.1 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")
            }
        }
    }
}

Nicolas Patry's avatar
Nicolas Patry committed
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
#[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")
            }
        }
    }
}

324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
#[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
351
352
353
354
/// App Configuration
#[derive(Parser, Debug)]
#[clap(author, version, about, long_about = None)]
struct Args {
355
356
357
358
359
    /// 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
360
    #[clap(default_value = "bigscience/bloom-560m", long, env)]
361
    model_id: String,
362
363
364

    /// 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
365
    #[clap(long, env)]
366
    revision: Option<String>,
367

368
369
370
371
372
    /// The number of tokenizer workers used for payload validation and truncation inside the
    /// router.
    #[clap(default_value = "2", long, env)]
    validation_workers: usize,

373
    /// Whether to shard the model across multiple GPUs
374
375
    /// By default text-generation-inference will use all available GPUs to run
    /// the model. Setting it to `false` deactivates `num_shard`.
376
377
    #[clap(long, env)]
    sharded: Option<bool>,
378
379

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

386
387
388
389
390
    /// 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.
391
392
    #[clap(long, env, value_enum)]
    quantize: Option<Quantization>,
393

Nicolas Patry's avatar
Nicolas Patry committed
394
395
396
397
398
399
400
    /// 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>,

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

405
406
407
408
409
410
    /// 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,

411
412
413
    /// 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
414
415
    #[clap(default_value = "128", long, env)]
    max_concurrent_requests: usize,
416
417
418
419

    /// 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
420
421
    #[clap(default_value = "2", long, env)]
    max_best_of: usize,
422
423
424
425
426
427

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

Nicolas Patry's avatar
Nicolas Patry committed
431
    /// This is the maximum allowed value for clients to set `top_n_tokens`.
432
    /// `top_n_tokens` is used to return information about the the `n` most likely
Nicolas Patry's avatar
Nicolas Patry committed
433
434
435
436
437
438
    /// 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,

439
440
441
442
    /// 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.
443
444
445
446
447
448
449
    /// 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>,
450
451
452
453
454
455
456
457
458

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

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

476
477
478
    /// 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.
479
480
481
    /// Default to `max_input_tokens + 50` to give a bit of room.
    #[clap(long, env)]
    max_batch_prefill_tokens: Option<u32>,
482

483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
    /// **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.
500
501
    #[clap(long, env)]
    max_batch_total_tokens: Option<u32>,
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519

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

523
524
525
526
527
    /// 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>,

528
529
    /// Specify the batch sizes to compute cuda graphs for.
    /// Use "0" to disable.
530
531
532
    /// Default = "1,2,4,8,16,32"
    #[clap(long, env, value_delimiter = ',')]
    cuda_graphs: Option<Vec<usize>>,
533

534
535
536
537
    /// The IP address to listen on
    #[clap(default_value = "0.0.0.0", long, env)]
    hostname: String,

538
    /// The port to listen on.
539
    #[clap(default_value = "3000", long, short, env)]
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
540
    port: u16,
541
542
543

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

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

    /// The address the master port will listen on. (setting used by torch distributed)
552
    #[clap(default_value = "29500", long, env)]
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
553
    master_port: usize,
554
555
556

    /// The location of the huggingface hub cache.
    /// Used to override the location if you want to provide a mounted disk for instance
557
    #[clap(long, env)]
558
    huggingface_hub_cache: Option<String>,
559
560
561

    /// The location of the huggingface hub cache.
    /// Used to override the location if you want to provide a mounted disk for instance
562
563
    #[clap(long, env)]
    weights_cache_override: Option<String>,
564
565
566
567
568

    /// 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.
569
    #[clap(long, env)]
570
    disable_custom_kernels: bool,
571

572
573
574
575
576
    /// 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
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
    /// 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>,

597
    /// Outputs the logs in JSON format (useful for telemetry)
598
    #[clap(long, env)]
599
    json_output: bool,
600

601
602
    #[clap(long, env)]
    otlp_endpoint: Option<String>,
603

604
605
606
    #[clap(default_value = "text-generation-inference.router", long, env)]
    otlp_service_name: String,

607
608
    #[clap(long, env)]
    cors_allow_origin: Vec<String>,
Erik Kaunismäki's avatar
Erik Kaunismäki committed
609
610
611
612

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

613
614
615
616
    #[clap(long, env)]
    watermark_gamma: Option<f32>,
    #[clap(long, env)]
    watermark_delta: Option<f32>,
617

618
619
620
621
622
623
624
625
    /// Enable ngrok tunneling
    #[clap(long, env)]
    ngrok: bool,

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

626
    /// ngrok edge
627
    #[clap(long, env)]
628
    ngrok_edge: Option<String>,
629

630
631
632
633
634
    /// 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
635
636
637
638
639
    /// 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,

640
641
642
    /// Display a lot of information about your runtime environment
    #[clap(long, short, action)]
    env: bool,
643
644
645
646

    /// 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
647
648
649
650
651

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

653
654
655
656
657
    /// 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
658
659
}

660
661
662
#[derive(Debug)]
enum ShardStatus {
    Ready,
663
    Failed(usize),
664
}
665

666
667
668
669
#[allow(clippy::too_many_arguments)]
fn shard_manager(
    model_id: String,
    revision: Option<String>,
670
    quantize: Option<Quantization>,
Nicolas Patry's avatar
Nicolas Patry committed
671
    speculate: Option<usize>,
672
    dtype: Option<Dtype>,
673
    trust_remote_code: bool,
674
675
676
677
678
679
680
681
682
683
    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>,
684
    cuda_graphs: Vec<usize>,
685
    cuda_memory_fraction: f32,
Nicolas Patry's avatar
Nicolas Patry committed
686
687
    rope_scaling: Option<RopeScaling>,
    rope_factor: Option<f32>,
688
689
    max_total_tokens: usize,
    max_batch_size: Option<usize>,
690
    max_input_tokens: usize,
drbh's avatar
drbh committed
691
    lora_adapters: Option<String>,
692
    otlp_endpoint: Option<String>,
693
    otlp_service_name: String,
694
    log_level: LevelFilter,
695
    status_sender: mpsc::Sender<ShardStatus>,
696
    shutdown: Arc<AtomicBool>,
697
698
    _shutdown_sender: mpsc::Sender<()>,
) {
699
700
701
    // Enter shard-manager tracing span
    let _span = tracing::span!(tracing::Level::INFO, "shard-manager", rank = rank).entered();

702
703
704
705
    // Get UDS path
    let uds_string = format!("{uds_path}-{rank}");
    let uds = Path::new(&uds_string);
    // Clean previous runs
706
707
708
    if uds.exists() {
        fs::remove_file(uds).unwrap();
    }
709
710

    // Process args
OlivierDehaene's avatar
OlivierDehaene committed
711
    let mut shard_args = vec![
712
713
714
715
716
        "serve".to_string(),
        model_id,
        "--uds-path".to_string(),
        uds_path,
        "--logger-level".to_string(),
717
        log_level.to_string().to_uppercase(),
718
719
720
        "--json-output".to_string(),
    ];

721
722
    // Activate trust remote code
    if trust_remote_code {
OlivierDehaene's avatar
OlivierDehaene committed
723
        shard_args.push("--trust-remote-code".to_string());
724
725
    }

726
727
    // Activate tensor parallelism
    if world_size > 1 {
OlivierDehaene's avatar
OlivierDehaene committed
728
        shard_args.push("--sharded".to_string());
729
730
    }

731
    if let Some(quantize) = quantize {
OlivierDehaene's avatar
OlivierDehaene committed
732
733
        shard_args.push("--quantize".to_string());
        shard_args.push(quantize.to_string())
734
    }
735

Nicolas Patry's avatar
Nicolas Patry committed
736
737
738
739
740
    if let Some(speculate) = speculate {
        shard_args.push("--speculate".to_string());
        shard_args.push(speculate.to_string())
    }

741
    if let Some(dtype) = dtype {
OlivierDehaene's avatar
OlivierDehaene committed
742
743
        shard_args.push("--dtype".to_string());
        shard_args.push(dtype.to_string())
744
745
    }

746
747
    // Model optional revision
    if let Some(revision) = revision {
OlivierDehaene's avatar
OlivierDehaene committed
748
749
        shard_args.push("--revision".to_string());
        shard_args.push(revision)
750
    }
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
751

Nicolas Patry's avatar
Nicolas Patry committed
752
753
754
755
756
757
    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)),
    };
758

759
    // OpenTelemetry Endpoint
760
    if let Some(otlp_endpoint) = otlp_endpoint {
OlivierDehaene's avatar
OlivierDehaene committed
761
762
        shard_args.push("--otlp-endpoint".to_string());
        shard_args.push(otlp_endpoint);
763
764
    }

765
766
767
768
    // OpenTelemetry Service Name
    shard_args.push("--otlp-service-name".to_string());
    shard_args.push(otlp_service_name);

769
770
771
772
    // 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());

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

776
777
778
    // Remove LOG_LEVEL if present
    envs.retain(|(name, _)| name != "LOG_LEVEL");

779
    // Torch Distributed Env vars
OlivierDehaene's avatar
OlivierDehaene committed
780
781
782
783
    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()));
784
    envs.push(("TORCH_NCCL_AVOID_RECORD_STREAMS".into(), "1".into()));
785

786
787
788
789
790
791
    // CUDA memory fraction
    envs.push((
        "CUDA_MEMORY_FRACTION".into(),
        cuda_memory_fraction.to_string().into(),
    ));

792
    // Safetensors load fast
OlivierDehaene's avatar
OlivierDehaene committed
793
    envs.push(("SAFETENSORS_FAST_GPU".into(), "1".into()));
794

795
796
797
    // Disable progress bar
    envs.push(("HF_HUB_DISABLE_PROGRESS_BARS".into(), "1".into()));

798
799
    // 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
800
    envs.push((
801
802
803
804
805
806
        "HF_HUB_ENABLE_HF_TRANSFER".into(),
        enable_hf_transfer.into(),
    ));

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

Nicolas Patry's avatar
Nicolas Patry committed
810
811
812
813
814
815
816
817
818
    // 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()));
    }

819
820
821
822
823
824
825
826
    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
827
828
829
830
831
    // Lora Adapters
    if let Some(lora_adapters) = lora_adapters {
        envs.push(("LORA_ADAPTERS".into(), lora_adapters.into()));
    }

832
833
834
    // 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
835
        envs.push(("HUGGINGFACE_HUB_CACHE".into(), huggingface_hub_cache.into()));
836
837
838
839
840
    };

    // 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
841
        envs.push((
842
843
844
845
846
            "WEIGHTS_CACHE_OVERRIDE".into(),
            weights_cache_override.into(),
        ));
    };

847
    // Enable experimental support for cuda graphs
848
849
850
851
852
853
854
855
856
857
    if !cuda_graphs.is_empty() {
        envs.push((
            "CUDA_GRAPHS".into(),
            cuda_graphs
                .into_iter()
                .map(|c| c.to_string())
                .collect::<Vec<_>>()
                .join(",")
                .into(),
        ));
858
859
    }

860
861
    // If disable_custom_kernels is true, pass it to the shard as an env var
    if disable_custom_kernels {
OlivierDehaene's avatar
OlivierDehaene committed
862
        envs.push(("DISABLE_CUSTOM_KERNELS".into(), "True".into()))
863
864
865
866
    }

    // Watermark Gamma
    if let Some(watermark_gamma) = watermark_gamma {
OlivierDehaene's avatar
OlivierDehaene committed
867
        envs.push(("WATERMARK_GAMMA".into(), watermark_gamma.to_string().into()))
868
869
870
871
    }

    // Watermark Delta
    if let Some(watermark_delta) = watermark_delta {
OlivierDehaene's avatar
OlivierDehaene committed
872
        envs.push(("WATERMARK_DELTA".into(), watermark_delta.to_string().into()))
873
874
875
    }

    // Start process
876
    tracing::info!("Starting shard");
877
    let mut p = match Command::new("text-generation-server")
OlivierDehaene's avatar
OlivierDehaene committed
878
        .args(shard_args)
879
        .env_clear()
OlivierDehaene's avatar
OlivierDehaene committed
880
        .envs(envs)
881
        .stdin(Stdio::piped())
882
883
884
885
886
        .stdout(Stdio::piped())
        .stderr(Stdio::piped())
        .process_group(0)
        .spawn()
    {
887
888
        Ok(p) => p,
        Err(err) => {
889
890
891
            if err.kind() == io::ErrorKind::NotFound {
                tracing::error!("text-generation-server not found in PATH");
                tracing::error!("Please install it with `make install-server`")
892
893
            }
            {
894
                tracing::error!("{}", err);
895
            }
896

897
            status_sender.send(ShardStatus::Failed(rank)).unwrap();
898
899
900
901
902
            return;
        }
    };

    // Redirect STDOUT to the console
903
    let mut pstdin = p.stdin.take().unwrap();
904
    let shard_stdout_reader = BufReader::new(p.stdout.take().unwrap());
905
    let shard_stderr_reader = BufReader::new(p.stderr.take().unwrap());
906

907
    //stdout tracing thread
908
    thread::spawn(move || {
909
        log_lines(shard_stdout_reader);
910
    });
911
912
913
    // 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
914
        for line in shard_stderr_reader.lines().map_while(Result::ok) {
915
916
917
            err_sender.send(line).unwrap_or(());
        }
    });
918
919
920
921
922
923
924
925
926
927
928
929
    // We read stdin in another thread as it seems that lines() can block in some cases
    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]);
                }
            }
        }
    });
930
931
932
933
934
935

    let mut ready = false;
    let start_time = Instant::now();
    let mut wait_time = Instant::now();
    loop {
        // Process exited
936
        if let Some(exit_status) = p.try_wait().unwrap() {
937
938
939
940
            let mut err = String::new();
            while let Ok(line) = err_receiver.recv_timeout(Duration::from_millis(10)) {
                err = err + "\n" + &line;
            }
941

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

944
            if let Some(signal) = exit_status.signal() {
945
946
947
                tracing::error!("Shard process was signaled to shutdown with signal {signal}");
            }

948
            status_sender.send(ShardStatus::Failed(rank)).unwrap();
949
950
951
952
            return;
        }

        // We received a shutdown signal
953
        if shutdown.load(Ordering::SeqCst) {
954
            terminate("shard", p, Duration::from_secs(90)).unwrap();
955
956
957
958
959
            return;
        }

        // Shard is ready
        if uds.exists() && !ready {
960
            tracing::info!("Shard ready in {:?}", start_time.elapsed());
961
962
963
            status_sender.send(ShardStatus::Ready).unwrap();
            ready = true;
        } else if !ready && wait_time.elapsed() > Duration::from_secs(10) {
964
            tracing::info!("Waiting for shard to be ready...");
965
966
967
968
969
970
            wait_time = Instant::now();
        }
        sleep(Duration::from_millis(100));
    }
}

971
fn shutdown_shards(shutdown: Arc<AtomicBool>, shutdown_receiver: &mpsc::Receiver<()>) {
972
973
974
    tracing::info!("Shutting down shards");
    // Update shutdown value to true
    // This will be picked up by the shard manager
975
    shutdown.store(true, Ordering::SeqCst);
976
977
978
979
980
981
982

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

fn num_cuda_devices() -> Option<usize> {
983
984
    let devices = match env::var("CUDA_VISIBLE_DEVICES") {
        Ok(devices) => devices,
985
986
987
        Err(_) => match env::var("NVIDIA_VISIBLE_DEVICES") {
            Ok(devices) => devices,
            Err(_) => env::var("ZE_AFFINITY_MASK").ok()?,
Nicolas Patry's avatar
Nicolas Patry committed
988
        },
989
    };
990
991
    let n_devices = devices.split(',').count();
    Some(n_devices)
992
993
994
995
996
997
998
999
1000
1001
1002
1003
1004
1005
1006
1007
1008
1009
1010
1011
1012
1013
1014
1015
1016
1017
1018
1019
1020
1021
1022
1023
1024
}

#[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 {
1025
1026
1027
1028
1029
1030
1031
            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()),
1032
1033
1034
1035
        }
    }
}

1036
impl TryFrom<&[u8]> for PythonLogMessage {
1037
1038
    type Error = serde_json::Error;

1039
1040
    fn try_from(value: &[u8]) -> Result<Self, Self::Error> {
        serde_json::from_slice::<Self>(value)
1041
1042
1043
    }
}

1044
1045
1046
1047
1048
1049
1050
1051
1052
1053
1054
1055
1056
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(_) => {
1057
1058
1059
1060
1061
1062
                            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();
1063
1064
1065
1066
1067
                            }
                        }
                    }
                }
            }
1068
1069
1070
1071
        }
    }
}

1072
1073
1074
1075
fn find_num_shards(
    sharded: Option<bool>,
    num_shard: Option<usize>,
) -> Result<usize, LauncherError> {
1076
1077
1078
1079
    // 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
1080
            tracing::info!("Parsing num_shard from CUDA_VISIBLE_DEVICES/NVIDIA_VISIBLE_DEVICES/ZE_AFFINITY_MASK");
1081
            let n_devices = num_cuda_devices()
1082
                .expect("--num-shard and CUDA_VISIBLE_DEVICES/NVIDIA_VISIBLE_DEVICES/ZE_AFFINITY_MASK are not set");
1083
            if n_devices <= 1 {
1084
1085
1086
                return Err(LauncherError::NotEnoughCUDADevices(format!(
                    "`sharded` is true but only found {n_devices} CUDA devices"
                )));
1087
            }
1088
            n_devices
1089
        }
1090
1091
1092
        (Some(true), Some(num_shard)) => {
            // we can't have only one shard while sharded
            if num_shard <= 1 {
1093
1094
1095
                return Err(LauncherError::ArgumentValidation(
                    "`sharded` is true but `num_shard` <= 1".to_string(),
                ));
1096
1097
            }
            num_shard
1098
        }
1099
1100
1101
1102
        (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,
1103
    };
1104
    if num_shard < 1 {
1105
1106
1107
        return Err(LauncherError::ArgumentValidation(
            "`num_shard` cannot be < 1".to_string(),
        ));
1108
    }
1109
    Ok(num_shard)
1110
}
1111

1112
#[derive(Debug, Error)]
1113
enum LauncherError {
1114
    #[error("Invalid argument: {0}")]
1115
    ArgumentValidation(String),
1116
    #[error("not enough cuda devices: {0}")]
1117
    NotEnoughCUDADevices(String),
1118
    #[error("Download error")]
1119
    DownloadError,
1120
    #[error("Shard cannot start")]
1121
    ShardCannotStart,
1122
    #[error("Shard disconnected")]
1123
    ShardDisconnected,
1124
    #[error("Shard failed")]
1125
    ShardFailed,
1126
    #[error("Webserver failed")]
1127
    WebserverFailed,
1128
    #[error("Webserver cannot start")]
1129
1130
    WebserverCannotStart,
}
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1131

1132
1133
1134
1135
1136
1137
1138
1139
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> {
1140
1141
1142
    // Enter download tracing span
    let _span = tracing::span!(tracing::Level::INFO, "download").entered();

OlivierDehaene's avatar
OlivierDehaene committed
1143
    let mut download_args = vec![
1144
        "download-weights".to_string(),
1145
        model_id.to_string(),
1146
1147
1148
1149
1150
1151
        "--extension".to_string(),
        ".safetensors".to_string(),
        "--logger-level".to_string(),
        "INFO".to_string(),
        "--json-output".to_string(),
    ];
1152

1153
    // Model optional revision
1154
    if let Some(revision) = &revision {
OlivierDehaene's avatar
OlivierDehaene committed
1155
1156
        download_args.push("--revision".to_string());
        download_args.push(revision.to_string())
1157
    }
1158

1159
    // Trust remote code for automatic peft fusion
1160
    if trust_remote_code {
1161
1162
1163
        download_args.push("--trust-remote-code".to_string());
    }

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

1167
1168
1169
    // Remove LOG_LEVEL if present
    envs.retain(|(name, _)| name != "LOG_LEVEL");

1170
1171
1172
    // Disable progress bar
    envs.push(("HF_HUB_DISABLE_PROGRESS_BARS".into(), "1".into()));

1173
    // If huggingface_hub_cache is set, pass it to the download process
1174
    // Useful when running inside a docker container
1175
    if let Some(ref huggingface_hub_cache) = huggingface_hub_cache {
OlivierDehaene's avatar
OlivierDehaene committed
1176
        envs.push(("HUGGINGFACE_HUB_CACHE".into(), huggingface_hub_cache.into()));
1177
    };
1178

1179
1180
    // 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
1181
    envs.push((
1182
1183
1184
        "HF_HUB_ENABLE_HF_TRANSFER".into(),
        enable_hf_transfer.into(),
    ));
1185

1186
1187
    // Parse Inference API token
    if let Ok(api_token) = env::var("HF_API_TOKEN") {
1188
        envs.push(("HF_TOKEN".into(), api_token.into()))
1189
    };
1190

1191
1192
    // If args.weights_cache_override is some, pass it to the download process
    // Useful when running inside a HuggingFace Inference Endpoint
1193
    if let Some(weights_cache_override) = &weights_cache_override {
OlivierDehaene's avatar
OlivierDehaene committed
1194
        envs.push((
1195
1196
1197
1198
1199
            "WEIGHTS_CACHE_OVERRIDE".into(),
            weights_cache_override.into(),
        ));
    };

1200
    // Start process
1201
    tracing::info!("Starting check and download process for {model_id}");
1202
    let mut download_process = match Command::new("text-generation-server")
OlivierDehaene's avatar
OlivierDehaene committed
1203
        .args(download_args)
1204
        .env_clear()
OlivierDehaene's avatar
OlivierDehaene committed
1205
        .envs(envs)
1206
1207
1208
1209
1210
        .stdout(Stdio::piped())
        .stderr(Stdio::piped())
        .process_group(0)
        .spawn()
    {
1211
1212
        Ok(p) => p,
        Err(err) => {
1213
1214
1215
            if err.kind() == io::ErrorKind::NotFound {
                tracing::error!("text-generation-server not found in PATH");
                tracing::error!("Please install it with `make install-server`")
1216
1217
            } else {
                tracing::error!("{}", err);
1218
            }
1219

1220
1221
1222
            return Err(LauncherError::DownloadError);
        }
    };
1223

1224
    let download_stdout = BufReader::new(download_process.stdout.take().unwrap());
1225

1226
    thread::spawn(move || {
1227
        log_lines(download_stdout);
1228
1229
1230
1231
1232
1233
1234
    });

    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
1235
        for line in download_stderr.lines().map_while(Result::ok) {
1236
1237
            err_sender.send(line).unwrap_or(());
        }
1238
    });
1239

1240
    loop {
1241
1242
        if let Some(status) = download_process.try_wait().unwrap() {
            if status.success() {
1243
                tracing::info!("Successfully downloaded weights for {model_id}");
1244
                break;
1245
            }
1246
1247

            let mut err = String::new();
1248
1249
1250
1251
            while let Ok(line) = err_receiver.recv_timeout(Duration::from_millis(10)) {
                err = err + "\n" + &line;
            }

1252
1253
1254
1255
1256
1257
1258
1259
1260
            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);
1261
        }
1262
        if !running.load(Ordering::SeqCst) {
OlivierDehaene's avatar
OlivierDehaene committed
1263
            terminate("download", download_process, Duration::from_secs(10)).unwrap();
1264
1265
1266
            return Ok(());
        }
        sleep(Duration::from_millis(100));
1267
    }
1268
1269
    Ok(())
}
1270

1271
#[allow(clippy::too_many_arguments)]
1272
1273
1274
fn spawn_shards(
    num_shard: usize,
    args: &Args,
1275
    cuda_graphs: Vec<usize>,
1276
    max_total_tokens: usize,
1277
    max_input_tokens: usize,
1278
    quantize: Option<Quantization>,
1279
    max_log_level: LevelFilter,
1280
    shutdown: Arc<AtomicBool>,
1281
1282
1283
1284
1285
1286
    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
1287
1288
    // Start shard processes
    for rank in 0..num_shard {
1289
1290
1291
1292
1293
1294
        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
1295
1296
1297
        let status_sender = status_sender.clone();
        let shutdown = shutdown.clone();
        let shutdown_sender = shutdown_sender.clone();
1298
        let otlp_endpoint = args.otlp_endpoint.clone();
1299
        let otlp_service_name = args.otlp_service_name.clone();
Nicolas Patry's avatar
Nicolas Patry committed
1300
        let speculate = args.speculate;
1301
        let dtype = args.dtype;
1302
        let trust_remote_code = args.trust_remote_code;
1303
1304
1305
1306
        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;
1307
        let cuda_graphs_clone = cuda_graphs.clone();
1308
        let cuda_memory_fraction = args.cuda_memory_fraction;
Nicolas Patry's avatar
Nicolas Patry committed
1309
1310
        let rope_scaling = args.rope_scaling;
        let rope_factor = args.rope_factor;
1311
        let max_batch_size = args.max_batch_size;
drbh's avatar
drbh committed
1312
        let lora_adapters = args.lora_adapters.clone();
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1313
1314
        thread::spawn(move || {
            shard_manager(
1315
                model_id,
1316
                revision,
1317
                quantize,
Nicolas Patry's avatar
Nicolas Patry committed
1318
                speculate,
1319
                dtype,
1320
                trust_remote_code,
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1321
1322
1323
1324
1325
                uds_path,
                rank,
                num_shard,
                master_addr,
                master_port,
1326
1327
                huggingface_hub_cache,
                weights_cache_override,
1328
                disable_custom_kernels,
1329
1330
                watermark_gamma,
                watermark_delta,
1331
                cuda_graphs_clone,
1332
                cuda_memory_fraction,
Nicolas Patry's avatar
Nicolas Patry committed
1333
1334
                rope_scaling,
                rope_factor,
1335
1336
                max_total_tokens,
                max_batch_size,
1337
                max_input_tokens,
drbh's avatar
drbh committed
1338
                lora_adapters,
1339
                otlp_endpoint,
1340
                otlp_service_name,
1341
                max_log_level,
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1342
1343
1344
1345
1346
1347
1348
1349
1350
1351
1352
1353
1354
1355
1356
1357
1358
1359
1360
1361
1362
                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));
            }
1363
            Ok(ShardStatus::Failed(rank)) => {
1364
                tracing::error!("Shard {rank} failed to start");
1365
                shutdown_shards(shutdown, shutdown_receiver);
1366
                return Err(LauncherError::ShardCannotStart);
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1367
1368
1369
            }
            Err(TryRecvError::Disconnected) => {
                tracing::error!("Shard status channel disconnected");
1370
                shutdown_shards(shutdown, shutdown_receiver);
1371
                return Err(LauncherError::ShardDisconnected);
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1372
1373
1374
            }
        }
    }
1375
1376
    Ok(())
}
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1377

1378
1379
1380
1381
1382
1383
1384
1385
1386
1387
1388
1389
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)
}

1390
fn spawn_webserver(
1391
    num_shard: usize,
1392
    args: Args,
1393
1394
1395
    max_input_tokens: usize,
    max_total_tokens: usize,
    max_batch_prefill_tokens: u32,
1396
    shutdown: Arc<AtomicBool>,
1397
    shutdown_receiver: &mpsc::Receiver<()>,
1398
) -> Result<Child, LauncherError> {
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1399
1400
1401
    // All shard started
    // Start webserver
    tracing::info!("Starting Webserver");
OlivierDehaene's avatar
OlivierDehaene committed
1402
    let mut router_args = vec![
1403
1404
        "--max-client-batch-size".to_string(),
        args.max_client_batch_size.to_string(),
1405
        "--max-concurrent-requests".to_string(),
1406
        args.max_concurrent_requests.to_string(),
1407
        "--max-best-of".to_string(),
1408
        args.max_best_of.to_string(),
1409
        "--max-stop-sequences".to_string(),
1410
        args.max_stop_sequences.to_string(),
Nicolas Patry's avatar
Nicolas Patry committed
1411
1412
        "--max-top-n-tokens".to_string(),
        args.max_top_n_tokens.to_string(),
1413
1414
        "--max-input-tokens".to_string(),
        max_input_tokens.to_string(),
1415
        "--max-total-tokens".to_string(),
1416
        max_total_tokens.to_string(),
1417
        "--max-batch-prefill-tokens".to_string(),
1418
        max_batch_prefill_tokens.to_string(),
1419
        "--waiting-served-ratio".to_string(),
1420
        args.waiting_served_ratio.to_string(),
1421
        "--max-waiting-tokens".to_string(),
1422
        args.max_waiting_tokens.to_string(),
1423
1424
        "--validation-workers".to_string(),
        args.validation_workers.to_string(),
1425
1426
        "--hostname".to_string(),
        args.hostname.to_string(),
1427
        "--port".to_string(),
1428
        args.port.to_string(),
1429
        "--master-shard-uds-path".to_string(),
1430
        format!("{}-0", args.shard_uds_path),
1431
        "--tokenizer-name".to_string(),
1432
        args.model_id,
1433
1434
    ];

1435
    // Pass usage stats flags to router
1436
1437
    router_args.push("--usage-stats".to_string());
    router_args.push(args.usage_stats.to_string());
1438

drbh's avatar
drbh committed
1439
1440
1441
1442
1443
    // Grammar support
    if args.disable_grammar_support {
        router_args.push("--disable-grammar-support".to_string());
    }

1444
1445
1446
1447
1448
1449
    // 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());
    }

1450
1451
1452
1453
1454
1455
    // 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());
    }

1456
1457
1458
1459
1460
1461
    // 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());
    }

1462
1463
    // Model optional revision
    if let Some(ref revision) = args.revision {
OlivierDehaene's avatar
OlivierDehaene committed
1464
1465
        router_args.push("--revision".to_string());
        router_args.push(revision.to_string())
1466
1467
    }

1468
    if args.json_output {
OlivierDehaene's avatar
OlivierDehaene committed
1469
        router_args.push("--json-output".to_string());
1470
1471
    }

1472
    // OpenTelemetry
1473
    if let Some(otlp_endpoint) = args.otlp_endpoint {
OlivierDehaene's avatar
OlivierDehaene committed
1474
1475
        router_args.push("--otlp-endpoint".to_string());
        router_args.push(otlp_endpoint);
1476
1477
    }

1478
1479
1480
1481
1482
    // OpenTelemetry
    let otlp_service_name = args.otlp_service_name;
    router_args.push("--otlp-service-name".to_string());
    router_args.push(otlp_service_name);

1483
1484
    // CORS origins
    for origin in args.cors_allow_origin.into_iter() {
OlivierDehaene's avatar
OlivierDehaene committed
1485
1486
        router_args.push("--cors-allow-origin".to_string());
        router_args.push(origin);
1487
1488
    }

Erik Kaunismäki's avatar
Erik Kaunismäki committed
1489
1490
1491
1492
1493
    // API Key
    if let Some(api_key) = args.api_key {
        router_args.push("--api-key".to_string());
        router_args.push(api_key);
    }
1494
1495
    // Ngrok
    if args.ngrok {
OlivierDehaene's avatar
OlivierDehaene committed
1496
1497
        router_args.push("--ngrok".to_string());
        router_args.push("--ngrok-authtoken".to_string());
1498
1499
1500
        router_args.push(args.ngrok_authtoken.unwrap());
        router_args.push("--ngrok-edge".to_string());
        router_args.push(args.ngrok_edge.unwrap());
1501
1502
    }

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

1506
1507
    // Parse Inference API token
    if let Ok(api_token) = env::var("HF_API_TOKEN") {
1508
        envs.push(("HF_TOKEN".into(), api_token.into()))
1509
    };
1510

1511
1512
1513
1514
1515
1516
1517
    // 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()))
    }

1518
    let mut webserver = match Command::new("text-generation-router")
OlivierDehaene's avatar
OlivierDehaene committed
1519
1520
        .args(router_args)
        .envs(envs)
1521
1522
1523
1524
1525
        .stdout(Stdio::piped())
        .stderr(Stdio::piped())
        .process_group(0)
        .spawn()
    {
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1526
1527
        Ok(p) => p,
        Err(err) => {
1528
            tracing::error!("Failed to start webserver: {}", err);
1529
1530
1531
            if err.kind() == io::ErrorKind::NotFound {
                tracing::error!("text-generation-router not found in PATH");
                tracing::error!("Please install it with `make install-router`")
1532
1533
            } else {
                tracing::error!("{}", err);
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1534
            }
1535

1536
            shutdown_shards(shutdown, shutdown_receiver);
1537
            return Err(LauncherError::WebserverCannotStart);
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1538
1539
1540
        }
    };

1541
1542
1543
    // Redirect STDOUT and STDERR to the console
    let webserver_stdout = webserver.stdout.take().unwrap();
    let webserver_stderr = webserver.stderr.take().unwrap();
1544
1545

    thread::spawn(move || {
1546
1547
        let stdout = BufReader::new(webserver_stdout);
        let stderr = BufReader::new(webserver_stderr);
1548
        for line in stdout.lines() {
1549
            println!("{}", line.unwrap());
1550
        }
1551
1552
        for line in stderr.lines() {
            println!("{}", line.unwrap());
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1553
        }
1554
1555
1556
    });
    Ok(webserver)
}
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1557

OlivierDehaene's avatar
OlivierDehaene committed
1558
1559
1560
1561
1562
1563
1564
1565
1566
1567
1568
1569
1570
1571
1572
1573
1574
1575
1576
1577
1578
1579
1580
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)
}

1581
1582
fn main() -> Result<(), LauncherError> {
    // Pattern match configuration
1583
    let args: Args = Args::parse();
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1584

1585
    // Filter events with LOG_LEVEL
1586
1587
1588
1589
1590
1591
1592
1593
1594
1595
1596
1597
1598
1599
1600
1601
    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);
1602

1603
    if args.json_output {
1604
1605
1606
1607
        tracing_subscriber::fmt()
            .with_env_filter(env_filter)
            .json()
            .init();
1608
    } else {
1609
1610
1611
1612
        tracing_subscriber::fmt()
            .with_env_filter(env_filter)
            .compact()
            .init();
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1613
1614
    }

1615
1616
1617
1618
1619
    if args.env {
        let env_runtime = env_runtime::Env::new();
        tracing::info!("{}", env_runtime);
    }

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

1622
1623
1624
1625
1626
1627
1628
1629
1630
1631
1632
1633
1634
1635
    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);
1636
                }
1637
                max_default
1638
            } else {
1639
                max_position_embeddings
1640
            }
1641
1642
1643
1644
1645
1646
1647
1648
1649
1650
        } 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);
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
1676
1677
1678
1679
1680
1681
1682
1683
1684
1685
1686
1687
1688
1689
1690
1691
1692
1693

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

1694
    // Validate args
1695
    if max_input_tokens >= max_total_tokens {
1696
        return Err(LauncherError::ArgumentValidation(
1697
            "`max_input_tokens must be < `max_total_tokens`".to_string(),
1698
1699
        ));
    }
1700
    if max_input_tokens as u32 > max_batch_prefill_tokens {
1701
        return Err(LauncherError::ArgumentValidation(format!(
1702
1703
            "`max_batch_prefill_tokens` must be >= `max_input_tokens`. Given: {} and {}",
            max_batch_prefill_tokens, max_input_tokens
1704
1705
        )));
    }
1706

1707
1708
1709
1710
1711
    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
1712
        (Some(cuda_graphs), _) => cuda_graphs.iter().cloned().filter(|&c| c > 0).collect(),
1713
1714
1715
1716
1717
        #[allow(deprecated)]
        (
            None,
            Some(
                Quantization::Bitsandbytes
1718
1719
                | Quantization::BitsandbytesNf4
                | Quantization::BitsandbytesFp4,
1720
1721
            ),
        ) => {
1722
1723
1724
1725
1726
            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");
1727
1728
1729
1730
1731
1732
1733
1734
1735
            vec![]
        }
        _ => {
            let cuda_graphs = vec![1, 2, 4, 8, 16, 32];
            tracing::info!("Using default cuda graphs {cuda_graphs:?}");
            cuda_graphs
        }
    };

1736
1737
1738
1739
1740
    if args.validation_workers == 0 {
        return Err(LauncherError::ArgumentValidation(
            "`validation_workers` must be > 0".to_string(),
        ));
    }
1741
1742
1743
1744
1745
1746
    if args.trust_remote_code {
        tracing::warn!(
            "`trust_remote_code` is set. Trusting that model `{}` do not contain malicious code.",
            args.model_id
        );
    }
1747
1748

    let num_shard = find_num_shards(args.sharded, args.num_shard)?;
1749
    if num_shard > 1 {
1750
1751
1752
1753
1754
        if matches!(args.quantize, Some(Quantization::Exl2)) {
            return Err(LauncherError::ArgumentValidation(
                "Sharding is currently not supported with `exl2` quantization".into(),
            ));
        }
1755
        tracing::info!("Sharding model on {num_shard} processes");
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1756
1757
    }

1758
    if let Some(ref max_batch_total_tokens) = args.max_batch_total_tokens {
1759
        if max_batch_prefill_tokens > *max_batch_total_tokens {
1760
1761
            return Err(LauncherError::ArgumentValidation(format!(
                "`max_batch_prefill_tokens` must be <= `max_batch_total_tokens`. Given: {} and {}",
1762
                max_batch_prefill_tokens, max_batch_total_tokens
1763
1764
            )));
        }
1765
        if max_total_tokens as u32 > *max_batch_total_tokens {
1766
1767
            return Err(LauncherError::ArgumentValidation(format!(
                "`max_total_tokens` must be <= `max_batch_total_tokens`. Given: {} and {}",
1768
                max_total_tokens, max_batch_total_tokens
1769
1770
1771
1772
            )));
        }
    }

1773
1774
1775
1776
1777
1778
1779
1780
1781
1782
1783
1784
1785
1786
    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(),
            ));
        }
    }

1787
1788
1789
1790
1791
1792
1793
    // 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");
1794

1795
    // Download and convert model weights
1796
1797
1798
1799
1800
1801
1802
1803
1804
1805
1806
1807
    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(',') {
1808
1809
1810
1811
            // skip download if a path is provided
            if adapter.contains('=') {
                continue;
            }
1812
1813
1814
1815
1816
1817
1818
1819
1820
1821
1822
1823
1824
1825
1826
1827
1828
1829
1830
1831
1832
1833
1834
1835
1836
1837
1838
1839
1840
1841
1842

            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
                )));
            }
1843
1844
        }
    }
1845

OlivierDehaene's avatar
OlivierDehaene committed
1846
1847
1848
1849
1850
    if !running.load(Ordering::SeqCst) {
        // Launcher was asked to stop
        return Ok(());
    }

1851
    // Shared shutdown bool
1852
    let shutdown = Arc::new(AtomicBool::new(false));
1853
1854
1855
    // Shared shutdown channel
    // When shutting down, the main thread will wait for all senders to be dropped
    let (shutdown_sender, shutdown_receiver) = mpsc::channel();
1856

1857
1858
    // Shared channel to track shard status
    let (status_sender, status_receiver) = mpsc::channel();
1859

1860
1861
1862
    spawn_shards(
        num_shard,
        &args,
1863
        cuda_graphs,
1864
        max_total_tokens,
1865
        max_input_tokens,
1866
        quantize,
1867
        max_log_level,
1868
1869
1870
1871
1872
1873
1874
        shutdown.clone(),
        &shutdown_receiver,
        shutdown_sender,
        &status_receiver,
        status_sender,
        running.clone(),
    )?;
1875

1876
1877
1878
1879
1880
    // We might have received a termination signal
    if !running.load(Ordering::SeqCst) {
        shutdown_shards(shutdown, &shutdown_receiver);
        return Ok(());
    }
1881

1882
1883
1884
1885
1886
1887
1888
1889
1890
    let mut webserver = spawn_webserver(
        num_shard,
        args,
        max_input_tokens,
        max_total_tokens,
        max_batch_prefill_tokens,
        shutdown.clone(),
        &shutdown_receiver,
    )
1891
    .inspect_err(|_| {
1892
1893
        shutdown_shards(shutdown.clone(), &shutdown_receiver);
    })?;
1894
1895
1896
1897
1898

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

    while running.load(Ordering::SeqCst) {
1899
        if let Ok(ShardStatus::Failed(rank)) = status_receiver.try_recv() {
OlivierDehaene's avatar
OlivierDehaene committed
1900
            tracing::error!("Shard {rank} crashed");
1901
1902
1903
1904
            exit_code = Err(LauncherError::ShardFailed);
            break;
        };

1905
        match webserver.try_wait().unwrap() {
1906
1907
1908
1909
1910
1911
1912
1913
1914
            Some(_) => {
                tracing::error!("Webserver Crashed");
                shutdown_shards(shutdown, &shutdown_receiver);
                return Err(LauncherError::WebserverFailed);
            }
            None => {
                sleep(Duration::from_millis(100));
            }
        };
1915
    }
1916
1917

    // Graceful termination
OlivierDehaene's avatar
OlivierDehaene committed
1918
    terminate("webserver", webserver, Duration::from_secs(90)).unwrap();
1919
1920
1921
    shutdown_shards(shutdown, &shutdown_receiver);

    exit_code
1922
}