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

21
22
mod env_runtime;

23
24
25
26
27
28
#[derive(Deserialize)]
struct Config {
    max_position_embeddings: Option<usize>,
    max_seq_len: Option<usize>,
}

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

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 {
68
69
            #[allow(deprecated)]
            // Use `eetq` instead, which provides better latencies overall and is drop-in in most cases
70
71
72
            Quantization::Bitsandbytes => {
                write!(f, "bitsandbytes")
            }
Nicolas Patry's avatar
Nicolas Patry committed
73
74
75
76
77
78
            Quantization::BitsandbytesNF4 => {
                write!(f, "bitsandbytes-nf4")
            }
            Quantization::BitsandbytesFP4 => {
                write!(f, "bitsandbytes-fp4")
            }
79
80
81
            Quantization::Gptq => {
                write!(f, "gptq")
            }
82
83
84
            Quantization::Awq => {
                write!(f, "awq")
            }
85
86
87
            Quantization::Eetq => {
                write!(f, "eetq")
            }
Nicolas Patry's avatar
Nicolas Patry committed
88
89
90
            Quantization::Fp8 => {
                write!(f, "fp8")
            }
91
92
93
94
        }
    }
}

95
96
97
#[derive(Clone, Copy, Debug, ValueEnum)]
enum Dtype {
    Float16,
98
    #[clap(name = "bfloat16")]
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
    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
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
#[derive(Clone, Copy, Debug, ValueEnum)]
enum RopeScaling {
    Linear,
    Dynamic,
}

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

Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
136
137
138
139
/// App Configuration
#[derive(Parser, Debug)]
#[clap(author, version, about, long_about = None)]
struct Args {
140
141
142
143
144
    /// 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
145
    #[clap(default_value = "bigscience/bloom-560m", long, env)]
146
    model_id: String,
147
148
149

    /// 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
150
    #[clap(long, env)]
151
    revision: Option<String>,
152

153
154
155
156
157
    /// The number of tokenizer workers used for payload validation and truncation inside the
    /// router.
    #[clap(default_value = "2", long, env)]
    validation_workers: usize,

158
    /// Whether to shard the model across multiple GPUs
159
160
    /// By default text-generation-inference will use all available GPUs to run
    /// the model. Setting it to `false` deactivates `num_shard`.
161
162
    #[clap(long, env)]
    sharded: Option<bool>,
163
164

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

171
    /// Whether you want the model to be quantized.
172
173
    #[clap(long, env, value_enum)]
    quantize: Option<Quantization>,
174

Nicolas Patry's avatar
Nicolas Patry committed
175
176
177
178
179
180
181
    /// 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>,

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

186
187
188
189
190
191
    /// 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,

192
193
194
    /// 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
195
196
    #[clap(default_value = "128", long, env)]
    max_concurrent_requests: usize,
197
198
199
200

    /// 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
201
202
    #[clap(default_value = "2", long, env)]
    max_best_of: usize,
203
204
205
206
207
208

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

Nicolas Patry's avatar
Nicolas Patry committed
212
213
214
215
216
217
218
219
    /// This is the maximum allowed value for clients to set `top_n_tokens`.
    /// `top_n_tokens is used to return information about the the `n` most likely
    /// tokens at each generation step, instead of just the sampled token. This
    /// information can be used for downstream tasks like for classification or
    /// ranking.
    #[clap(default_value = "5", long, env)]
    max_top_n_tokens: u32,

220
221
222
223
    /// 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.
224
225
226
227
228
229
230
    /// 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>,
231
232
233
234
235
236
237
238
239

    /// 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.
240
241
242
    /// Default to min(max_position_embeddings, 4096)
    #[clap(long, env)]
    max_total_tokens: Option<usize>,
243
244
245
246
247
248
249
250
251
252
253

    /// 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`.
254
255
    #[clap(default_value = "1.2", long, env)]
    waiting_served_ratio: f32,
256

257
258
259
    /// 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.
260
261
262
    /// Default to `max_input_tokens + 50` to give a bit of room.
    #[clap(long, env)]
    max_batch_prefill_tokens: Option<u32>,
263

264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
    /// **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.
281
282
    #[clap(long, env)]
    max_batch_total_tokens: Option<u32>,
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300

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

304
305
306
307
308
    /// 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>,

309
310
    /// Specify the batch sizes to compute cuda graphs for.
    /// Use "0" to disable.
311
312
313
    /// Default = "1,2,4,8,16,32"
    #[clap(long, env, value_delimiter = ',')]
    cuda_graphs: Option<Vec<usize>>,
314

315
316
317
318
    /// The IP address to listen on
    #[clap(default_value = "0.0.0.0", long, env)]
    hostname: String,

319
    /// The port to listen on.
320
    #[clap(default_value = "3000", long, short, env)]
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
321
    port: u16,
322
323
324

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

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

    /// The address the master port will listen on. (setting used by torch distributed)
333
    #[clap(default_value = "29500", long, env)]
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
334
    master_port: usize,
335
336
337

    /// The location of the huggingface hub cache.
    /// Used to override the location if you want to provide a mounted disk for instance
338
    #[clap(long, env)]
339
    huggingface_hub_cache: Option<String>,
340
341
342

    /// The location of the huggingface hub cache.
    /// Used to override the location if you want to provide a mounted disk for instance
343
344
    #[clap(long, env)]
    weights_cache_override: Option<String>,
345
346
347
348
349

    /// 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.
350
    #[clap(long, env)]
351
    disable_custom_kernels: bool,
352

353
354
355
356
357
    /// 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
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
    /// 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>,

378
    /// Outputs the logs in JSON format (useful for telemetry)
379
    #[clap(long, env)]
380
    json_output: bool,
381

382
383
    #[clap(long, env)]
    otlp_endpoint: Option<String>,
384

385
386
    #[clap(long, env)]
    cors_allow_origin: Vec<String>,
387
388
389
390
    #[clap(long, env)]
    watermark_gamma: Option<f32>,
    #[clap(long, env)]
    watermark_delta: Option<f32>,
391

392
393
394
395
396
397
398
399
    /// Enable ngrok tunneling
    #[clap(long, env)]
    ngrok: bool,

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

400
    /// ngrok edge
401
    #[clap(long, env)]
402
    ngrok_edge: Option<String>,
403

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

414
415
416
    /// Display a lot of information about your runtime environment
    #[clap(long, short, action)]
    env: bool,
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
417
418
}

419
420
421
#[derive(Debug)]
enum ShardStatus {
    Ready,
422
    Failed(usize),
423
}
424

425
426
427
428
#[allow(clippy::too_many_arguments)]
fn shard_manager(
    model_id: String,
    revision: Option<String>,
429
    quantize: Option<Quantization>,
Nicolas Patry's avatar
Nicolas Patry committed
430
    speculate: Option<usize>,
431
    dtype: Option<Dtype>,
432
    trust_remote_code: bool,
433
434
435
436
437
438
439
440
441
442
    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>,
443
    cuda_graphs: Vec<usize>,
444
    cuda_memory_fraction: f32,
Nicolas Patry's avatar
Nicolas Patry committed
445
446
    rope_scaling: Option<RopeScaling>,
    rope_factor: Option<f32>,
447
448
    otlp_endpoint: Option<String>,
    status_sender: mpsc::Sender<ShardStatus>,
449
    shutdown: Arc<AtomicBool>,
450
451
    _shutdown_sender: mpsc::Sender<()>,
) {
452
453
454
    // Enter shard-manager tracing span
    let _span = tracing::span!(tracing::Level::INFO, "shard-manager", rank = rank).entered();

455
456
457
458
    // Get UDS path
    let uds_string = format!("{uds_path}-{rank}");
    let uds = Path::new(&uds_string);
    // Clean previous runs
459
460
461
    if uds.exists() {
        fs::remove_file(uds).unwrap();
    }
462
463

    // Process args
OlivierDehaene's avatar
OlivierDehaene committed
464
    let mut shard_args = vec![
465
466
467
468
469
470
471
472
473
        "serve".to_string(),
        model_id,
        "--uds-path".to_string(),
        uds_path,
        "--logger-level".to_string(),
        "INFO".to_string(),
        "--json-output".to_string(),
    ];

474
475
    // Activate trust remote code
    if trust_remote_code {
OlivierDehaene's avatar
OlivierDehaene committed
476
        shard_args.push("--trust-remote-code".to_string());
477
478
    }

479
480
    // Activate tensor parallelism
    if world_size > 1 {
OlivierDehaene's avatar
OlivierDehaene committed
481
        shard_args.push("--sharded".to_string());
482
483
    }

484
    if let Some(quantize) = quantize {
OlivierDehaene's avatar
OlivierDehaene committed
485
486
        shard_args.push("--quantize".to_string());
        shard_args.push(quantize.to_string())
487
    }
488

Nicolas Patry's avatar
Nicolas Patry committed
489
490
491
492
493
    if let Some(speculate) = speculate {
        shard_args.push("--speculate".to_string());
        shard_args.push(speculate.to_string())
    }

494
    if let Some(dtype) = dtype {
OlivierDehaene's avatar
OlivierDehaene committed
495
496
        shard_args.push("--dtype".to_string());
        shard_args.push(dtype.to_string())
497
498
    }

499
500
    // Model optional revision
    if let Some(revision) = revision {
OlivierDehaene's avatar
OlivierDehaene committed
501
502
        shard_args.push("--revision".to_string());
        shard_args.push(revision)
503
    }
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
504

Nicolas Patry's avatar
Nicolas Patry committed
505
506
507
508
509
510
    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)),
    };
511
512
    // OpenTelemetry
    if let Some(otlp_endpoint) = otlp_endpoint {
OlivierDehaene's avatar
OlivierDehaene committed
513
514
        shard_args.push("--otlp-endpoint".to_string());
        shard_args.push(otlp_endpoint);
515
516
517
    }

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

520
521
522
    // Remove LOG_LEVEL if present
    envs.retain(|(name, _)| name != "LOG_LEVEL");

523
    // Torch Distributed Env vars
OlivierDehaene's avatar
OlivierDehaene committed
524
525
526
527
    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()));
528
    envs.push(("TORCH_NCCL_AVOID_RECORD_STREAMS".into(), "1".into()));
529

530
531
532
533
534
535
    // CUDA memory fraction
    envs.push((
        "CUDA_MEMORY_FRACTION".into(),
        cuda_memory_fraction.to_string().into(),
    ));

536
    // Safetensors load fast
OlivierDehaene's avatar
OlivierDehaene committed
537
    envs.push(("SAFETENSORS_FAST_GPU".into(), "1".into()));
538

539
540
541
    // Disable progress bar
    envs.push(("HF_HUB_DISABLE_PROGRESS_BARS".into(), "1".into()));

542
543
    // 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
544
    envs.push((
545
546
547
548
549
550
        "HF_HUB_ENABLE_HF_TRANSFER".into(),
        enable_hf_transfer.into(),
    ));

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

Nicolas Patry's avatar
Nicolas Patry committed
554
555
556
557
558
559
560
561
562
    // 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()));
    }

563
564
565
    // 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
566
        envs.push(("HUGGINGFACE_HUB_CACHE".into(), huggingface_hub_cache.into()));
567
568
569
570
571
    };

    // 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
572
        envs.push((
573
574
575
576
577
            "WEIGHTS_CACHE_OVERRIDE".into(),
            weights_cache_override.into(),
        ));
    };

578
    // Enable experimental support for cuda graphs
579
580
581
582
583
584
585
586
587
588
    if !cuda_graphs.is_empty() {
        envs.push((
            "CUDA_GRAPHS".into(),
            cuda_graphs
                .into_iter()
                .map(|c| c.to_string())
                .collect::<Vec<_>>()
                .join(",")
                .into(),
        ));
589
590
    }

591
592
    // If disable_custom_kernels is true, pass it to the shard as an env var
    if disable_custom_kernels {
OlivierDehaene's avatar
OlivierDehaene committed
593
        envs.push(("DISABLE_CUSTOM_KERNELS".into(), "True".into()))
594
595
596
597
    }

    // Watermark Gamma
    if let Some(watermark_gamma) = watermark_gamma {
OlivierDehaene's avatar
OlivierDehaene committed
598
        envs.push(("WATERMARK_GAMMA".into(), watermark_gamma.to_string().into()))
599
600
601
602
    }

    // Watermark Delta
    if let Some(watermark_delta) = watermark_delta {
OlivierDehaene's avatar
OlivierDehaene committed
603
        envs.push(("WATERMARK_DELTA".into(), watermark_delta.to_string().into()))
604
605
606
    }

    // Start process
607
    tracing::info!("Starting shard");
608
    let mut p = match Command::new("text-generation-server")
OlivierDehaene's avatar
OlivierDehaene committed
609
        .args(shard_args)
610
        .env_clear()
OlivierDehaene's avatar
OlivierDehaene committed
611
        .envs(envs)
612
613
614
615
616
        .stdout(Stdio::piped())
        .stderr(Stdio::piped())
        .process_group(0)
        .spawn()
    {
617
618
        Ok(p) => p,
        Err(err) => {
619
620
621
            if err.kind() == io::ErrorKind::NotFound {
                tracing::error!("text-generation-server not found in PATH");
                tracing::error!("Please install it with `make install-server`")
622
623
            }
            {
624
                tracing::error!("{}", err);
625
            }
626

627
            status_sender.send(ShardStatus::Failed(rank)).unwrap();
628
629
630
631
632
            return;
        }
    };

    // Redirect STDOUT to the console
633
    let shard_stdout_reader = BufReader::new(p.stdout.take().unwrap());
634
    let shard_stderr_reader = BufReader::new(p.stderr.take().unwrap());
635

636
    //stdout tracing thread
637
    thread::spawn(move || {
638
        log_lines(shard_stdout_reader.lines());
639
    });
640
641
642
    // 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
643
        for line in shard_stderr_reader.lines().map_while(Result::ok) {
644
645
646
            err_sender.send(line).unwrap_or(());
        }
    });
647
648
649
650
651
652

    let mut ready = false;
    let start_time = Instant::now();
    let mut wait_time = Instant::now();
    loop {
        // Process exited
653
        if let Some(exit_status) = p.try_wait().unwrap() {
654
655
656
657
            let mut err = String::new();
            while let Ok(line) = err_receiver.recv_timeout(Duration::from_millis(10)) {
                err = err + "\n" + &line;
            }
658

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

661
            if let Some(signal) = exit_status.signal() {
662
663
664
                tracing::error!("Shard process was signaled to shutdown with signal {signal}");
            }

665
            status_sender.send(ShardStatus::Failed(rank)).unwrap();
666
667
668
669
            return;
        }

        // We received a shutdown signal
670
        if shutdown.load(Ordering::SeqCst) {
671
            p.kill().unwrap();
672
            let _ = p.wait();
673
            tracing::info!("Shard terminated");
674
675
676
677
678
            return;
        }

        // Shard is ready
        if uds.exists() && !ready {
679
            tracing::info!("Shard ready in {:?}", start_time.elapsed());
680
681
682
            status_sender.send(ShardStatus::Ready).unwrap();
            ready = true;
        } else if !ready && wait_time.elapsed() > Duration::from_secs(10) {
683
            tracing::info!("Waiting for shard to be ready...");
684
685
686
687
688
689
            wait_time = Instant::now();
        }
        sleep(Duration::from_millis(100));
    }
}

690
fn shutdown_shards(shutdown: Arc<AtomicBool>, shutdown_receiver: &mpsc::Receiver<()>) {
691
692
693
    tracing::info!("Shutting down shards");
    // Update shutdown value to true
    // This will be picked up by the shard manager
694
    shutdown.store(true, Ordering::SeqCst);
695
696
697
698
699
700
701

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

fn num_cuda_devices() -> Option<usize> {
702
703
704
705
    let devices = match env::var("CUDA_VISIBLE_DEVICES") {
        Ok(devices) => devices,
        Err(_) => env::var("NVIDIA_VISIBLE_DEVICES").ok()?,
    };
706
707
    let n_devices = devices.split(',').count();
    Some(n_devices)
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
}

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

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

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

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

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

752
753
754
755
756
757
758
759
760
impl TryFrom<&String> for PythonLogMessage {
    type Error = serde_json::Error;

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

fn log_lines<S: Sized + BufRead>(lines: Lines<S>) {
OlivierDehaene's avatar
OlivierDehaene committed
761
    for line in lines.map_while(Result::ok) {
762
763
764
765
766
767
768
        match PythonLogMessage::try_from(&line) {
            Ok(log) => log.trace(),
            Err(_) => tracing::debug!("{line}"),
        }
    }
}

769
770
771
772
fn find_num_shards(
    sharded: Option<bool>,
    num_shard: Option<usize>,
) -> Result<usize, LauncherError> {
773
774
775
776
    // 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
777
778
779
            tracing::info!("Parsing num_shard from CUDA_VISIBLE_DEVICES/NVIDIA_VISIBLE_DEVICES");
            let n_devices = num_cuda_devices()
                .expect("--num-shard and CUDA_VISIBLE_DEVICES/NVIDIA_VISIBLE_DEVICES are not set");
780
            if n_devices <= 1 {
781
782
783
                return Err(LauncherError::NotEnoughCUDADevices(format!(
                    "`sharded` is true but only found {n_devices} CUDA devices"
                )));
784
            }
785
            n_devices
786
        }
787
788
789
        (Some(true), Some(num_shard)) => {
            // we can't have only one shard while sharded
            if num_shard <= 1 {
790
791
792
                return Err(LauncherError::ArgumentValidation(
                    "`sharded` is true but `num_shard` <= 1".to_string(),
                ));
793
794
            }
            num_shard
795
        }
796
797
798
799
        (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,
800
    };
801
    if num_shard < 1 {
802
803
804
        return Err(LauncherError::ArgumentValidation(
            "`num_shard` cannot be < 1".to_string(),
        ));
805
    }
806
    Ok(num_shard)
807
}
808

809
810
#[derive(Debug)]
enum LauncherError {
811
812
    ArgumentValidation(String),
    NotEnoughCUDADevices(String),
813
814
815
816
817
818
819
    DownloadError,
    ShardCannotStart,
    ShardDisconnected,
    ShardFailed,
    WebserverFailed,
    WebserverCannotStart,
}
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
820

821
822
823
824
825
826
827
828
impl core::fmt::Display for LauncherError {
    fn fmt(&self, f: &mut core::fmt::Formatter) -> core::fmt::Result {
        write!(f, "{self:?}")
    }
}

impl std::error::Error for LauncherError {}

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

OlivierDehaene's avatar
OlivierDehaene committed
833
    let mut download_args = vec![
834
835
836
837
838
839
840
841
        "download-weights".to_string(),
        args.model_id.to_string(),
        "--extension".to_string(),
        ".safetensors".to_string(),
        "--logger-level".to_string(),
        "INFO".to_string(),
        "--json-output".to_string(),
    ];
842

843
844
    // Model optional revision
    if let Some(revision) = &args.revision {
OlivierDehaene's avatar
OlivierDehaene committed
845
846
        download_args.push("--revision".to_string());
        download_args.push(revision.to_string())
847
    }
848

849
850
851
852
853
    // Trust remote code for automatic peft fusion
    if args.trust_remote_code {
        download_args.push("--trust-remote-code".to_string());
    }

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

857
858
859
    // Remove LOG_LEVEL if present
    envs.retain(|(name, _)| name != "LOG_LEVEL");

860
861
862
    // Disable progress bar
    envs.push(("HF_HUB_DISABLE_PROGRESS_BARS".into(), "1".into()));

863
    // If huggingface_hub_cache is set, pass it to the download process
864
865
    // Useful when running inside a docker container
    if let Some(ref huggingface_hub_cache) = args.huggingface_hub_cache {
OlivierDehaene's avatar
OlivierDehaene committed
866
        envs.push(("HUGGINGFACE_HUB_CACHE".into(), huggingface_hub_cache.into()));
867
    };
868

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

876
877
    // Parse Inference API token
    if let Ok(api_token) = env::var("HF_API_TOKEN") {
OlivierDehaene's avatar
OlivierDehaene committed
878
        envs.push(("HUGGING_FACE_HUB_TOKEN".into(), api_token.into()))
879
    };
880

881
882
883
    // If args.weights_cache_override is some, pass it to the download process
    // Useful when running inside a HuggingFace Inference Endpoint
    if let Some(weights_cache_override) = &args.weights_cache_override {
OlivierDehaene's avatar
OlivierDehaene committed
884
        envs.push((
885
886
887
888
889
            "WEIGHTS_CACHE_OVERRIDE".into(),
            weights_cache_override.into(),
        ));
    };

890
891
    // Start process
    tracing::info!("Starting download process.");
892
    let mut download_process = match Command::new("text-generation-server")
OlivierDehaene's avatar
OlivierDehaene committed
893
        .args(download_args)
894
        .env_clear()
OlivierDehaene's avatar
OlivierDehaene committed
895
        .envs(envs)
896
897
898
899
900
        .stdout(Stdio::piped())
        .stderr(Stdio::piped())
        .process_group(0)
        .spawn()
    {
901
902
        Ok(p) => p,
        Err(err) => {
903
904
905
            if err.kind() == io::ErrorKind::NotFound {
                tracing::error!("text-generation-server not found in PATH");
                tracing::error!("Please install it with `make install-server`")
906
907
            } else {
                tracing::error!("{}", err);
908
            }
909

910
911
912
            return Err(LauncherError::DownloadError);
        }
    };
913

914
    let download_stdout = BufReader::new(download_process.stdout.take().unwrap());
915

916
    thread::spawn(move || {
917
918
919
920
921
922
923
924
        log_lines(download_stdout.lines());
    });

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

    // We read stderr in another thread as it seems that lines() can block in some cases
    let (err_sender, err_receiver) = mpsc::channel();
    thread::spawn(move || {
OlivierDehaene's avatar
OlivierDehaene committed
925
        for line in download_stderr.lines().map_while(Result::ok) {
926
927
            err_sender.send(line).unwrap_or(());
        }
928
    });
929

930
    loop {
931
932
933
934
        if let Some(status) = download_process.try_wait().unwrap() {
            if status.success() {
                tracing::info!("Successfully downloaded weights.");
                break;
935
            }
936
937

            let mut err = String::new();
938
939
940
941
            while let Ok(line) = err_receiver.recv_timeout(Duration::from_millis(10)) {
                err = err + "\n" + &line;
            }

942
943
944
945
946
947
948
949
950
            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);
951
        }
952
        if !running.load(Ordering::SeqCst) {
OlivierDehaene's avatar
OlivierDehaene committed
953
            terminate("download", download_process, Duration::from_secs(10)).unwrap();
954
955
956
            return Ok(());
        }
        sleep(Duration::from_millis(100));
957
    }
958
959
    Ok(())
}
960

961
#[allow(clippy::too_many_arguments)]
962
963
964
fn spawn_shards(
    num_shard: usize,
    args: &Args,
965
    cuda_graphs: Vec<usize>,
966
    shutdown: Arc<AtomicBool>,
967
968
969
970
971
972
    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
973
974
    // Start shard processes
    for rank in 0..num_shard {
975
976
977
978
979
980
        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
981
982
983
        let status_sender = status_sender.clone();
        let shutdown = shutdown.clone();
        let shutdown_sender = shutdown_sender.clone();
984
        let otlp_endpoint = args.otlp_endpoint.clone();
985
        let quantize = args.quantize;
Nicolas Patry's avatar
Nicolas Patry committed
986
        let speculate = args.speculate;
987
        let dtype = args.dtype;
988
        let trust_remote_code = args.trust_remote_code;
989
990
991
992
        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;
993
        let cuda_graphs_clone = cuda_graphs.clone();
994
        let cuda_memory_fraction = args.cuda_memory_fraction;
Nicolas Patry's avatar
Nicolas Patry committed
995
996
        let rope_scaling = args.rope_scaling;
        let rope_factor = args.rope_factor;
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
997
998
        thread::spawn(move || {
            shard_manager(
999
                model_id,
1000
                revision,
1001
                quantize,
Nicolas Patry's avatar
Nicolas Patry committed
1002
                speculate,
1003
                dtype,
1004
                trust_remote_code,
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1005
1006
1007
1008
1009
                uds_path,
                rank,
                num_shard,
                master_addr,
                master_port,
1010
1011
                huggingface_hub_cache,
                weights_cache_override,
1012
                disable_custom_kernels,
1013
1014
                watermark_gamma,
                watermark_delta,
1015
                cuda_graphs_clone,
1016
                cuda_memory_fraction,
Nicolas Patry's avatar
Nicolas Patry committed
1017
1018
                rope_scaling,
                rope_factor,
1019
                otlp_endpoint,
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1020
1021
1022
1023
1024
1025
1026
1027
1028
1029
1030
1031
1032
1033
1034
1035
1036
1037
1038
1039
1040
                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));
            }
1041
            Ok(ShardStatus::Failed(rank)) => {
1042
                tracing::error!("Shard {rank} failed to start");
1043
                shutdown_shards(shutdown, shutdown_receiver);
1044
                return Err(LauncherError::ShardCannotStart);
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1045
1046
1047
            }
            Err(TryRecvError::Disconnected) => {
                tracing::error!("Shard status channel disconnected");
1048
                shutdown_shards(shutdown, shutdown_receiver);
1049
                return Err(LauncherError::ShardDisconnected);
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1050
1051
1052
            }
        }
    }
1053
1054
    Ok(())
}
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1055

1056
1057
1058
1059
1060
1061
1062
1063
1064
1065
1066
1067
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)
}

1068
fn spawn_webserver(
1069
    num_shard: usize,
1070
    args: Args,
1071
1072
1073
    max_input_tokens: usize,
    max_total_tokens: usize,
    max_batch_prefill_tokens: u32,
1074
    shutdown: Arc<AtomicBool>,
1075
    shutdown_receiver: &mpsc::Receiver<()>,
1076
) -> Result<Child, LauncherError> {
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1077
1078
1079
    // All shard started
    // Start webserver
    tracing::info!("Starting Webserver");
OlivierDehaene's avatar
OlivierDehaene committed
1080
    let mut router_args = vec![
1081
        "--max-concurrent-requests".to_string(),
1082
        args.max_concurrent_requests.to_string(),
1083
        "--max-best-of".to_string(),
1084
        args.max_best_of.to_string(),
1085
        "--max-stop-sequences".to_string(),
1086
        args.max_stop_sequences.to_string(),
Nicolas Patry's avatar
Nicolas Patry committed
1087
1088
        "--max-top-n-tokens".to_string(),
        args.max_top_n_tokens.to_string(),
1089
1090
        "--max-input-tokens".to_string(),
        max_input_tokens.to_string(),
1091
        "--max-total-tokens".to_string(),
1092
        max_total_tokens.to_string(),
1093
        "--max-batch-prefill-tokens".to_string(),
1094
        max_batch_prefill_tokens.to_string(),
1095
        "--waiting-served-ratio".to_string(),
1096
        args.waiting_served_ratio.to_string(),
1097
        "--max-waiting-tokens".to_string(),
1098
        args.max_waiting_tokens.to_string(),
1099
1100
        "--validation-workers".to_string(),
        args.validation_workers.to_string(),
1101
1102
        "--hostname".to_string(),
        args.hostname.to_string(),
1103
        "--port".to_string(),
1104
        args.port.to_string(),
1105
        "--master-shard-uds-path".to_string(),
1106
        format!("{}-0", args.shard_uds_path),
1107
        "--tokenizer-name".to_string(),
1108
        args.model_id,
1109
1110
    ];

drbh's avatar
drbh committed
1111
1112
1113
1114
1115
    // Grammar support
    if args.disable_grammar_support {
        router_args.push("--disable-grammar-support".to_string());
    }

1116
1117
1118
1119
1120
1121
    // 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());
    }

1122
1123
1124
1125
1126
1127
    // 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());
    }

1128
1129
1130
1131
1132
1133
    // 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());
    }

1134
1135
    // Model optional revision
    if let Some(ref revision) = args.revision {
OlivierDehaene's avatar
OlivierDehaene committed
1136
1137
        router_args.push("--revision".to_string());
        router_args.push(revision.to_string())
1138
1139
    }

1140
    if args.json_output {
OlivierDehaene's avatar
OlivierDehaene committed
1141
        router_args.push("--json-output".to_string());
1142
1143
    }

1144
    // OpenTelemetry
1145
    if let Some(otlp_endpoint) = args.otlp_endpoint {
OlivierDehaene's avatar
OlivierDehaene committed
1146
1147
        router_args.push("--otlp-endpoint".to_string());
        router_args.push(otlp_endpoint);
1148
1149
1150
1151
    }

    // CORS origins
    for origin in args.cors_allow_origin.into_iter() {
OlivierDehaene's avatar
OlivierDehaene committed
1152
1153
        router_args.push("--cors-allow-origin".to_string());
        router_args.push(origin);
1154
1155
    }

1156
1157
    // Ngrok
    if args.ngrok {
OlivierDehaene's avatar
OlivierDehaene committed
1158
1159
        router_args.push("--ngrok".to_string());
        router_args.push("--ngrok-authtoken".to_string());
1160
1161
1162
        router_args.push(args.ngrok_authtoken.unwrap());
        router_args.push("--ngrok-edge".to_string());
        router_args.push(args.ngrok_edge.unwrap());
1163
1164
    }

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

1168
1169
    // Parse Inference API token
    if let Ok(api_token) = env::var("HF_API_TOKEN") {
OlivierDehaene's avatar
OlivierDehaene committed
1170
        envs.push(("HUGGING_FACE_HUB_TOKEN".into(), api_token.into()))
1171
    };
1172

1173
1174
1175
1176
1177
1178
1179
    // 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()))
    }

1180
    let mut webserver = match Command::new("text-generation-router")
OlivierDehaene's avatar
OlivierDehaene committed
1181
1182
        .args(router_args)
        .envs(envs)
1183
1184
1185
1186
1187
        .stdout(Stdio::piped())
        .stderr(Stdio::piped())
        .process_group(0)
        .spawn()
    {
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1188
1189
        Ok(p) => p,
        Err(err) => {
1190
            tracing::error!("Failed to start webserver: {}", err);
1191
1192
1193
            if err.kind() == io::ErrorKind::NotFound {
                tracing::error!("text-generation-router not found in PATH");
                tracing::error!("Please install it with `make install-router`")
1194
1195
            } else {
                tracing::error!("{}", err);
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1196
            }
1197

1198
            shutdown_shards(shutdown, shutdown_receiver);
1199
            return Err(LauncherError::WebserverCannotStart);
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1200
1201
1202
        }
    };

1203
1204
1205
    // Redirect STDOUT and STDERR to the console
    let webserver_stdout = webserver.stdout.take().unwrap();
    let webserver_stderr = webserver.stderr.take().unwrap();
1206
1207

    thread::spawn(move || {
1208
1209
        let stdout = BufReader::new(webserver_stdout);
        let stderr = BufReader::new(webserver_stderr);
1210
        for line in stdout.lines() {
1211
            println!("{}", line.unwrap());
1212
        }
1213
1214
        for line in stderr.lines() {
            println!("{}", line.unwrap());
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1215
        }
1216
1217
1218
    });
    Ok(webserver)
}
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1219

OlivierDehaene's avatar
OlivierDehaene committed
1220
1221
1222
1223
1224
1225
1226
1227
1228
1229
1230
1231
1232
1233
1234
1235
1236
1237
1238
1239
1240
1241
1242
1243
1244
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)
}

1245
1246
fn main() -> Result<(), LauncherError> {
    // Pattern match configuration
1247
    let args: Args = Args::parse();
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1248

1249
1250
1251
1252
    // Filter events with LOG_LEVEL
    let env_filter =
        EnvFilter::try_from_env("LOG_LEVEL").unwrap_or_else(|_| EnvFilter::new("info"));

1253
    if args.json_output {
1254
1255
1256
1257
        tracing_subscriber::fmt()
            .with_env_filter(env_filter)
            .json()
            .init();
1258
    } else {
1259
1260
1261
1262
        tracing_subscriber::fmt()
            .with_env_filter(env_filter)
            .compact()
            .init();
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1263
1264
    }

1265
1266
1267
1268
1269
    if args.env {
        let env_runtime = env_runtime::Env::new();
        tracing::info!("{}", env_runtime);
    }

1270
1271
    tracing::info!("{:?}", args);

1272
1273
1274
1275
1276
1277
1278
1279
1280
1281
1282
1283
1284
1285
1286
1287
1288
1289
1290
1291
1292
1293
1294
1295
1296
1297
1298
1299
1300
1301
1302
1303
1304
1305
1306
1307
1308
1309
1310
1311
1312
1313
1314
1315
1316
1317
1318
1319
1320
1321
1322
1323
1324
1325
1326
1327
1328
1329
1330
1331
1332
1333
1334
1335
1336
1337
1338
1339
1340
1341
1342
1343
1344
1345
1346
1347
1348
1349
1350
1351
1352
1353
1354
1355
1356
1357
1358
1359
1360
    let get_max_position_embeddings = || -> Result<usize, Box<dyn std::error::Error>> {
        let model_id = args.model_id.clone();
        let mut path = std::path::Path::new(&args.model_id).to_path_buf();
        let filename = if !path.exists() {
            // Assume it's a hub id
            let api = Api::new()?;
            let repo = if let Some(ref revision) = args.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: Config = serde_json::from_str(&content)?;

        // Quantization usually means you're even more RAM constrained.
        let max_default = 4096;

        let max_position_embeddings = match (config.max_position_embeddings, config.max_seq_len) {
            (Some(max_position_embeddings), _) | (None, Some(max_position_embeddings)) => {
                if max_position_embeddings > max_default {
                    let max = max_position_embeddings;
                    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);
                    max_default
                } else {
                    max_position_embeddings
                }
            }
            _ => {
                return Err(Box::new(LauncherError::ArgumentValidation(
                    "no max defined".to_string(),
                )));
            }
        };
        Ok(max_position_embeddings)
    };
    let max_position_embeddings: usize = get_max_position_embeddings().unwrap_or(4096);

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

1361
    // Validate args
1362
    if max_input_tokens >= max_total_tokens {
1363
        return Err(LauncherError::ArgumentValidation(
1364
            "`max_input_tokens must be < `max_total_tokens`".to_string(),
1365
1366
        ));
    }
1367
    if max_input_tokens as u32 > max_batch_prefill_tokens {
1368
        return Err(LauncherError::ArgumentValidation(format!(
1369
1370
            "`max_batch_prefill_tokens` must be >= `max_input_tokens`. Given: {} and {}",
            max_batch_prefill_tokens, max_input_tokens
1371
1372
        )));
    }
1373

1374
1375
1376
1377
1378
1379
1380
1381
1382
1383
1384
1385
1386
1387
1388
1389
1390
1391
1392
1393
1394
    let cuda_graphs = match (&args.cuda_graphs, &args.quantize) {
        (Some(cuda_graphs), Some(_q)) => cuda_graphs.clone(),
        #[allow(deprecated)]
        (
            None,
            Some(
                Quantization::Bitsandbytes
                | Quantization::BitsandbytesNF4
                | Quantization::BitsandbytesFP4,
            ),
        ) => {
            tracing::info!("Bitsandbytes doesn't work with cuda graphs, deactivating them");
            vec![]
        }
        _ => {
            let cuda_graphs = vec![1, 2, 4, 8, 16, 32];
            tracing::info!("Using default cuda graphs {cuda_graphs:?}");
            cuda_graphs
        }
    };

1395
1396
1397
1398
1399
    if args.validation_workers == 0 {
        return Err(LauncherError::ArgumentValidation(
            "`validation_workers` must be > 0".to_string(),
        ));
    }
1400
1401
1402
1403
1404
1405
    if args.trust_remote_code {
        tracing::warn!(
            "`trust_remote_code` is set. Trusting that model `{}` do not contain malicious code.",
            args.model_id
        );
    }
1406
1407

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

1412
    if let Some(ref max_batch_total_tokens) = args.max_batch_total_tokens {
1413
        if max_batch_prefill_tokens > *max_batch_total_tokens {
1414
1415
            return Err(LauncherError::ArgumentValidation(format!(
                "`max_batch_prefill_tokens` must be <= `max_batch_total_tokens`. Given: {} and {}",
1416
                max_batch_prefill_tokens, max_batch_total_tokens
1417
1418
            )));
        }
1419
        if max_total_tokens as u32 > *max_batch_total_tokens {
1420
1421
            return Err(LauncherError::ArgumentValidation(format!(
                "`max_total_tokens` must be <= `max_batch_total_tokens`. Given: {} and {}",
1422
                max_total_tokens, max_batch_total_tokens
1423
1424
1425
1426
            )));
        }
    }

1427
1428
1429
1430
1431
1432
1433
1434
1435
1436
1437
1438
1439
1440
    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(),
            ));
        }
    }

1441
1442
1443
1444
1445
1446
1447
    // 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");
1448

1449
    // Download and convert model weights
1450
    download_convert_model(&args, running.clone())?;
1451

OlivierDehaene's avatar
OlivierDehaene committed
1452
1453
1454
1455
1456
    if !running.load(Ordering::SeqCst) {
        // Launcher was asked to stop
        return Ok(());
    }

1457
    // Shared shutdown bool
1458
    let shutdown = Arc::new(AtomicBool::new(false));
1459
1460
1461
    // Shared shutdown channel
    // When shutting down, the main thread will wait for all senders to be dropped
    let (shutdown_sender, shutdown_receiver) = mpsc::channel();
1462

1463
1464
    // Shared channel to track shard status
    let (status_sender, status_receiver) = mpsc::channel();
1465

1466
1467
1468
    spawn_shards(
        num_shard,
        &args,
1469
        cuda_graphs,
1470
1471
1472
1473
1474
1475
1476
        shutdown.clone(),
        &shutdown_receiver,
        shutdown_sender,
        &status_receiver,
        status_sender,
        running.clone(),
    )?;
1477

1478
1479
1480
1481
1482
    // We might have received a termination signal
    if !running.load(Ordering::SeqCst) {
        shutdown_shards(shutdown, &shutdown_receiver);
        return Ok(());
    }
1483

1484
1485
1486
1487
1488
1489
1490
1491
1492
1493
1494
1495
1496
    let mut webserver = spawn_webserver(
        num_shard,
        args,
        max_input_tokens,
        max_total_tokens,
        max_batch_prefill_tokens,
        shutdown.clone(),
        &shutdown_receiver,
    )
    .map_err(|err| {
        shutdown_shards(shutdown.clone(), &shutdown_receiver);
        err
    })?;
1497
1498
1499
1500
1501

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

    while running.load(Ordering::SeqCst) {
1502
        if let Ok(ShardStatus::Failed(rank)) = status_receiver.try_recv() {
OlivierDehaene's avatar
OlivierDehaene committed
1503
            tracing::error!("Shard {rank} crashed");
1504
1505
1506
1507
            exit_code = Err(LauncherError::ShardFailed);
            break;
        };

1508
        match webserver.try_wait().unwrap() {
1509
1510
1511
1512
1513
1514
1515
1516
1517
            Some(_) => {
                tracing::error!("Webserver Crashed");
                shutdown_shards(shutdown, &shutdown_receiver);
                return Err(LauncherError::WebserverFailed);
            }
            None => {
                sleep(Duration::from_millis(100));
            }
        };
1518
    }
1519
1520

    // Graceful termination
OlivierDehaene's avatar
OlivierDehaene committed
1521
    terminate("webserver", webserver, Duration::from_secs(90)).unwrap();
1522
1523
1524
    shutdown_shards(shutdown, &shutdown_receiver);

    exit_code
1525
}