main.rs 55.5 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 thiserror::Error;
20
use tracing_subscriber::{filter::LevelFilter, EnvFilter};
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
21

22
23
mod env_runtime;

24
#[derive(Deserialize)]
25
struct RawConfig {
26
    max_position_embeddings: Option<usize>,
27
    n_positions: Option<usize>,
28
29
30
    max_seq_len: Option<usize>,
}

31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
#[derive(Deserialize)]
struct Config {
    max_position_embeddings: Option<usize>,
}

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);
        Config {
            max_position_embeddings,
        }
    }
}

48
49
#[derive(Clone, Copy, Debug, ValueEnum)]
enum Quantization {
50
    /// 4 bit quantization. Requires a specific AWQ quantized model:
51
    ///   <https://hf.co/models?search=awq>.
52
    /// Should replace GPTQ models wherever possible because of the better latency
53
54
55
    Awq,
    /// 8 bit quantization, doesn't require specific model.
    /// Should be a drop-in replacement to bitsandbytes with much better performance.
56
    /// Kernels are from <https://github.com/NetEase-FuXi/EETQ.git>
57
    Eetq,
58
59
60
61
    /// 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,
62
    /// 4 bit quantization. Requires a specific GTPQ quantized model: <https://hf.co/models?search=gptq>.
63
    /// text-generation-inference will use exllama (faster) kernels wherever possible, and use
64
65
66
    /// triton kernel (wider support) when it's not.
    /// AWQ has faster kernels.
    Gptq,
67
68
    /// 4 bit quantization. Requires a specific Marlin quantized model: <https://hf.co/models?search=marlin>.
    Marlin,
69
70
71
72
73
74
    /// 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"
    )]
75
    Bitsandbytes,
76
77
    /// 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
78
    BitsandbytesNF4,
79
80
    /// 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
81
    BitsandbytesFP4,
Nicolas Patry's avatar
Nicolas Patry committed
82
83
84
85
86
    /// [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,
87
88
89
90
91
92
}

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 {
93
94
            #[allow(deprecated)]
            // Use `eetq` instead, which provides better latencies overall and is drop-in in most cases
95
96
97
            Quantization::Bitsandbytes => {
                write!(f, "bitsandbytes")
            }
Nicolas Patry's avatar
Nicolas Patry committed
98
99
100
101
102
103
            Quantization::BitsandbytesNF4 => {
                write!(f, "bitsandbytes-nf4")
            }
            Quantization::BitsandbytesFP4 => {
                write!(f, "bitsandbytes-fp4")
            }
104
105
106
            Quantization::Exl2 => {
                write!(f, "exl2")
            }
107
108
109
            Quantization::Gptq => {
                write!(f, "gptq")
            }
110
111
112
            Quantization::Marlin => {
                write!(f, "marlin")
            }
113
114
115
            Quantization::Awq => {
                write!(f, "awq")
            }
116
117
118
            Quantization::Eetq => {
                write!(f, "eetq")
            }
Nicolas Patry's avatar
Nicolas Patry committed
119
120
121
            Quantization::Fp8 => {
                write!(f, "fp8")
            }
122
123
124
125
        }
    }
}

126
127
128
#[derive(Clone, Copy, Debug, ValueEnum)]
enum Dtype {
    Float16,
129
    #[clap(name = "bfloat16")]
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
    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
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
#[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
167
168
169
170
/// App Configuration
#[derive(Parser, Debug)]
#[clap(author, version, about, long_about = None)]
struct Args {
171
172
173
174
175
    /// 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
176
    #[clap(default_value = "bigscience/bloom-560m", long, env)]
177
    model_id: String,
178
179
180

    /// 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
181
    #[clap(long, env)]
182
    revision: Option<String>,
183

184
185
186
187
188
    /// The number of tokenizer workers used for payload validation and truncation inside the
    /// router.
    #[clap(default_value = "2", long, env)]
    validation_workers: usize,

189
    /// Whether to shard the model across multiple GPUs
190
191
    /// By default text-generation-inference will use all available GPUs to run
    /// the model. Setting it to `false` deactivates `num_shard`.
192
193
    #[clap(long, env)]
    sharded: Option<bool>,
194
195

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

202
    /// Whether you want the model to be quantized.
203
204
    #[clap(long, env, value_enum)]
    quantize: Option<Quantization>,
205

Nicolas Patry's avatar
Nicolas Patry committed
206
207
208
209
210
211
212
    /// 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>,

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

217
218
219
220
221
222
    /// 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,

223
224
225
    /// 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
226
227
    #[clap(default_value = "128", long, env)]
    max_concurrent_requests: usize,
228
229
230
231

    /// 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
232
233
    #[clap(default_value = "2", long, env)]
    max_best_of: usize,
234
235
236
237
238
239

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

Nicolas Patry's avatar
Nicolas Patry committed
243
    /// This is the maximum allowed value for clients to set `top_n_tokens`.
244
    /// `top_n_tokens` is used to return information about the the `n` most likely
Nicolas Patry's avatar
Nicolas Patry committed
245
246
247
248
249
250
    /// 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,

251
252
253
254
    /// 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.
255
256
257
258
259
260
261
    /// 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>,
262
263
264
265
266
267
268
269
270

    /// 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.
271
272
273
    /// Default to min(max_position_embeddings, 4096)
    #[clap(long, env)]
    max_total_tokens: Option<usize>,
274
275
276
277
278
279
280
281
282
283
284

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

288
289
290
    /// 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.
291
292
293
    /// Default to `max_input_tokens + 50` to give a bit of room.
    #[clap(long, env)]
    max_batch_prefill_tokens: Option<u32>,
294

295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
    /// **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.
312
313
    #[clap(long, env)]
    max_batch_total_tokens: Option<u32>,
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331

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

335
336
337
338
339
    /// 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>,

340
341
    /// Specify the batch sizes to compute cuda graphs for.
    /// Use "0" to disable.
342
343
344
    /// Default = "1,2,4,8,16,32"
    #[clap(long, env, value_delimiter = ',')]
    cuda_graphs: Option<Vec<usize>>,
345

346
347
348
349
    /// The IP address to listen on
    #[clap(default_value = "0.0.0.0", long, env)]
    hostname: String,

350
    /// The port to listen on.
351
    #[clap(default_value = "3000", long, short, env)]
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
352
    port: u16,
353
354
355

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

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

    /// The address the master port will listen on. (setting used by torch distributed)
364
    #[clap(default_value = "29500", long, env)]
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
365
    master_port: usize,
366
367
368

    /// The location of the huggingface hub cache.
    /// Used to override the location if you want to provide a mounted disk for instance
369
    #[clap(long, env)]
370
    huggingface_hub_cache: Option<String>,
371
372
373

    /// The location of the huggingface hub cache.
    /// Used to override the location if you want to provide a mounted disk for instance
374
375
    #[clap(long, env)]
    weights_cache_override: Option<String>,
376
377
378
379
380

    /// 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.
381
    #[clap(long, env)]
382
    disable_custom_kernels: bool,
383

384
385
386
387
388
    /// 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
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
    /// 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>,

409
    /// Outputs the logs in JSON format (useful for telemetry)
410
    #[clap(long, env)]
411
    json_output: bool,
412

413
414
    #[clap(long, env)]
    otlp_endpoint: Option<String>,
415

416
417
    #[clap(long, env)]
    cors_allow_origin: Vec<String>,
418
419
420
421
    #[clap(long, env)]
    watermark_gamma: Option<f32>,
    #[clap(long, env)]
    watermark_delta: Option<f32>,
422

423
424
425
426
427
428
429
430
    /// Enable ngrok tunneling
    #[clap(long, env)]
    ngrok: bool,

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

431
    /// ngrok edge
432
    #[clap(long, env)]
433
    ngrok_edge: Option<String>,
434

435
436
437
438
439
    /// 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
440
441
442
443
444
    /// 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,

445
446
447
    /// Display a lot of information about your runtime environment
    #[clap(long, short, action)]
    env: bool,
448
449
450
451

    /// 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,
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
452
453
}

454
455
456
#[derive(Debug)]
enum ShardStatus {
    Ready,
457
    Failed(usize),
458
}
459

460
461
462
463
#[allow(clippy::too_many_arguments)]
fn shard_manager(
    model_id: String,
    revision: Option<String>,
464
    quantize: Option<Quantization>,
Nicolas Patry's avatar
Nicolas Patry committed
465
    speculate: Option<usize>,
466
    dtype: Option<Dtype>,
467
    trust_remote_code: bool,
468
469
470
471
472
473
474
475
476
477
    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>,
478
    cuda_graphs: Vec<usize>,
479
    cuda_memory_fraction: f32,
Nicolas Patry's avatar
Nicolas Patry committed
480
481
    rope_scaling: Option<RopeScaling>,
    rope_factor: Option<f32>,
482
483
    max_total_tokens: usize,
    max_batch_size: Option<usize>,
484
    otlp_endpoint: Option<String>,
485
    log_level: LevelFilter,
486
    status_sender: mpsc::Sender<ShardStatus>,
487
    shutdown: Arc<AtomicBool>,
488
489
    _shutdown_sender: mpsc::Sender<()>,
) {
490
491
492
    // Enter shard-manager tracing span
    let _span = tracing::span!(tracing::Level::INFO, "shard-manager", rank = rank).entered();

493
494
495
496
    // Get UDS path
    let uds_string = format!("{uds_path}-{rank}");
    let uds = Path::new(&uds_string);
    // Clean previous runs
497
498
499
    if uds.exists() {
        fs::remove_file(uds).unwrap();
    }
500
501

    // Process args
OlivierDehaene's avatar
OlivierDehaene committed
502
    let mut shard_args = vec![
503
504
505
506
507
        "serve".to_string(),
        model_id,
        "--uds-path".to_string(),
        uds_path,
        "--logger-level".to_string(),
508
        log_level.to_string().to_uppercase(),
509
510
511
        "--json-output".to_string(),
    ];

512
513
    // Activate trust remote code
    if trust_remote_code {
OlivierDehaene's avatar
OlivierDehaene committed
514
        shard_args.push("--trust-remote-code".to_string());
515
516
    }

517
518
    // Activate tensor parallelism
    if world_size > 1 {
OlivierDehaene's avatar
OlivierDehaene committed
519
        shard_args.push("--sharded".to_string());
520
521
    }

522
    if let Some(quantize) = quantize {
OlivierDehaene's avatar
OlivierDehaene committed
523
524
        shard_args.push("--quantize".to_string());
        shard_args.push(quantize.to_string())
525
    }
526

Nicolas Patry's avatar
Nicolas Patry committed
527
528
529
530
531
    if let Some(speculate) = speculate {
        shard_args.push("--speculate".to_string());
        shard_args.push(speculate.to_string())
    }

532
    if let Some(dtype) = dtype {
OlivierDehaene's avatar
OlivierDehaene committed
533
534
        shard_args.push("--dtype".to_string());
        shard_args.push(dtype.to_string())
535
536
    }

537
538
    // Model optional revision
    if let Some(revision) = revision {
OlivierDehaene's avatar
OlivierDehaene committed
539
540
        shard_args.push("--revision".to_string());
        shard_args.push(revision)
541
    }
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
542

Nicolas Patry's avatar
Nicolas Patry committed
543
544
545
546
547
548
    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)),
    };
549

550
551
    // OpenTelemetry
    if let Some(otlp_endpoint) = otlp_endpoint {
OlivierDehaene's avatar
OlivierDehaene committed
552
553
        shard_args.push("--otlp-endpoint".to_string());
        shard_args.push(otlp_endpoint);
554
555
556
    }

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

559
560
561
    // Remove LOG_LEVEL if present
    envs.retain(|(name, _)| name != "LOG_LEVEL");

562
    // Torch Distributed Env vars
OlivierDehaene's avatar
OlivierDehaene committed
563
564
565
566
    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()));
567
    envs.push(("TORCH_NCCL_AVOID_RECORD_STREAMS".into(), "1".into()));
568

569
570
571
572
573
574
    // CUDA memory fraction
    envs.push((
        "CUDA_MEMORY_FRACTION".into(),
        cuda_memory_fraction.to_string().into(),
    ));

575
    // Safetensors load fast
OlivierDehaene's avatar
OlivierDehaene committed
576
    envs.push(("SAFETENSORS_FAST_GPU".into(), "1".into()));
577

578
579
580
    // Disable progress bar
    envs.push(("HF_HUB_DISABLE_PROGRESS_BARS".into(), "1".into()));

581
582
    // 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
583
    envs.push((
584
585
586
587
588
589
        "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
590
        envs.push(("HUGGING_FACE_HUB_TOKEN".into(), api_token.into()))
591
592
    };

Nicolas Patry's avatar
Nicolas Patry committed
593
594
595
596
597
598
599
600
601
    // 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()));
    }

602
603
604
605
606
607
608
609
    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()));
    }

610
611
612
    // 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
613
        envs.push(("HUGGINGFACE_HUB_CACHE".into(), huggingface_hub_cache.into()));
614
615
616
617
618
    };

    // 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
619
        envs.push((
620
621
622
623
624
            "WEIGHTS_CACHE_OVERRIDE".into(),
            weights_cache_override.into(),
        ));
    };

625
    // Enable experimental support for cuda graphs
626
627
628
629
630
631
632
633
634
635
    if !cuda_graphs.is_empty() {
        envs.push((
            "CUDA_GRAPHS".into(),
            cuda_graphs
                .into_iter()
                .map(|c| c.to_string())
                .collect::<Vec<_>>()
                .join(",")
                .into(),
        ));
636
637
    }

638
639
    // If disable_custom_kernels is true, pass it to the shard as an env var
    if disable_custom_kernels {
OlivierDehaene's avatar
OlivierDehaene committed
640
        envs.push(("DISABLE_CUSTOM_KERNELS".into(), "True".into()))
641
642
643
644
    }

    // Watermark Gamma
    if let Some(watermark_gamma) = watermark_gamma {
OlivierDehaene's avatar
OlivierDehaene committed
645
        envs.push(("WATERMARK_GAMMA".into(), watermark_gamma.to_string().into()))
646
647
648
649
    }

    // Watermark Delta
    if let Some(watermark_delta) = watermark_delta {
OlivierDehaene's avatar
OlivierDehaene committed
650
        envs.push(("WATERMARK_DELTA".into(), watermark_delta.to_string().into()))
651
652
653
    }

    // Start process
654
    tracing::info!("Starting shard");
655
    let mut p = match Command::new("text-generation-server")
OlivierDehaene's avatar
OlivierDehaene committed
656
        .args(shard_args)
657
        .env_clear()
OlivierDehaene's avatar
OlivierDehaene committed
658
        .envs(envs)
659
660
661
662
663
        .stdout(Stdio::piped())
        .stderr(Stdio::piped())
        .process_group(0)
        .spawn()
    {
664
665
        Ok(p) => p,
        Err(err) => {
666
667
668
            if err.kind() == io::ErrorKind::NotFound {
                tracing::error!("text-generation-server not found in PATH");
                tracing::error!("Please install it with `make install-server`")
669
670
            }
            {
671
                tracing::error!("{}", err);
672
            }
673

674
            status_sender.send(ShardStatus::Failed(rank)).unwrap();
675
676
677
678
679
            return;
        }
    };

    // Redirect STDOUT to the console
680
    let shard_stdout_reader = BufReader::new(p.stdout.take().unwrap());
681
    let shard_stderr_reader = BufReader::new(p.stderr.take().unwrap());
682

683
    //stdout tracing thread
684
    thread::spawn(move || {
685
        log_lines(shard_stdout_reader.lines());
686
    });
687
688
689
    // 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
690
        for line in shard_stderr_reader.lines().map_while(Result::ok) {
691
692
693
            err_sender.send(line).unwrap_or(());
        }
    });
694
695
696
697
698
699

    let mut ready = false;
    let start_time = Instant::now();
    let mut wait_time = Instant::now();
    loop {
        // Process exited
700
        if let Some(exit_status) = p.try_wait().unwrap() {
701
702
703
704
            let mut err = String::new();
            while let Ok(line) = err_receiver.recv_timeout(Duration::from_millis(10)) {
                err = err + "\n" + &line;
            }
705

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

708
            if let Some(signal) = exit_status.signal() {
709
710
711
                tracing::error!("Shard process was signaled to shutdown with signal {signal}");
            }

712
            status_sender.send(ShardStatus::Failed(rank)).unwrap();
713
714
715
716
            return;
        }

        // We received a shutdown signal
717
        if shutdown.load(Ordering::SeqCst) {
718
            terminate("shard", p, Duration::from_secs(90)).unwrap();
719
720
721
722
723
            return;
        }

        // Shard is ready
        if uds.exists() && !ready {
724
            tracing::info!("Shard ready in {:?}", start_time.elapsed());
725
726
727
            status_sender.send(ShardStatus::Ready).unwrap();
            ready = true;
        } else if !ready && wait_time.elapsed() > Duration::from_secs(10) {
728
            tracing::info!("Waiting for shard to be ready...");
729
730
731
732
733
734
            wait_time = Instant::now();
        }
        sleep(Duration::from_millis(100));
    }
}

735
fn shutdown_shards(shutdown: Arc<AtomicBool>, shutdown_receiver: &mpsc::Receiver<()>) {
736
737
738
    tracing::info!("Shutting down shards");
    // Update shutdown value to true
    // This will be picked up by the shard manager
739
    shutdown.store(true, Ordering::SeqCst);
740
741
742
743
744
745
746

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

fn num_cuda_devices() -> Option<usize> {
747
748
749
750
    let devices = match env::var("CUDA_VISIBLE_DEVICES") {
        Ok(devices) => devices,
        Err(_) => env::var("NVIDIA_VISIBLE_DEVICES").ok()?,
    };
751
752
    let n_devices = devices.split(',').count();
    Some(n_devices)
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
}

#[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 {
786
787
788
789
790
791
792
            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()),
793
794
795
796
        }
    }
}

797
798
799
800
801
802
803
804
805
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
806
    for line in lines.map_while(Result::ok) {
807
808
809
810
811
812
813
        match PythonLogMessage::try_from(&line) {
            Ok(log) => log.trace(),
            Err(_) => tracing::debug!("{line}"),
        }
    }
}

814
815
816
817
fn find_num_shards(
    sharded: Option<bool>,
    num_shard: Option<usize>,
) -> Result<usize, LauncherError> {
818
819
820
821
    // 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
822
823
824
            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");
825
            if n_devices <= 1 {
826
827
828
                return Err(LauncherError::NotEnoughCUDADevices(format!(
                    "`sharded` is true but only found {n_devices} CUDA devices"
                )));
829
            }
830
            n_devices
831
        }
832
833
834
        (Some(true), Some(num_shard)) => {
            // we can't have only one shard while sharded
            if num_shard <= 1 {
835
836
837
                return Err(LauncherError::ArgumentValidation(
                    "`sharded` is true but `num_shard` <= 1".to_string(),
                ));
838
839
            }
            num_shard
840
        }
841
842
843
844
        (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,
845
    };
846
    if num_shard < 1 {
847
848
849
        return Err(LauncherError::ArgumentValidation(
            "`num_shard` cannot be < 1".to_string(),
        ));
850
    }
851
    Ok(num_shard)
852
}
853

854
#[derive(Debug, Error)]
855
enum LauncherError {
856
    #[error("Invalid argument: {0}")]
857
    ArgumentValidation(String),
858
    #[error("not enough cuda devices: {0}")]
859
    NotEnoughCUDADevices(String),
860
    #[error("Download error")]
861
    DownloadError,
862
    #[error("Shard cannot start")]
863
    ShardCannotStart,
864
    #[error("Shard disconnected")]
865
    ShardDisconnected,
866
    #[error("Shard failed")]
867
    ShardFailed,
868
    #[error("Webserver failed")]
869
    WebserverFailed,
870
    #[error("Webserver cannot start")]
871
872
    WebserverCannotStart,
}
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
873

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

OlivierDehaene's avatar
OlivierDehaene committed
878
    let mut download_args = vec![
879
880
881
882
883
884
885
886
        "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(),
    ];
887

888
889
    // Model optional revision
    if let Some(revision) = &args.revision {
OlivierDehaene's avatar
OlivierDehaene committed
890
891
        download_args.push("--revision".to_string());
        download_args.push(revision.to_string())
892
    }
893

894
895
896
897
898
    // Trust remote code for automatic peft fusion
    if args.trust_remote_code {
        download_args.push("--trust-remote-code".to_string());
    }

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

902
903
904
    // Remove LOG_LEVEL if present
    envs.retain(|(name, _)| name != "LOG_LEVEL");

905
906
907
    // Disable progress bar
    envs.push(("HF_HUB_DISABLE_PROGRESS_BARS".into(), "1".into()));

908
    // If huggingface_hub_cache is set, pass it to the download process
909
910
    // Useful when running inside a docker container
    if let Some(ref huggingface_hub_cache) = args.huggingface_hub_cache {
OlivierDehaene's avatar
OlivierDehaene committed
911
        envs.push(("HUGGINGFACE_HUB_CACHE".into(), huggingface_hub_cache.into()));
912
    };
913

914
915
    // 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
916
    envs.push((
917
918
919
        "HF_HUB_ENABLE_HF_TRANSFER".into(),
        enable_hf_transfer.into(),
    ));
920

921
922
    // Parse Inference API token
    if let Ok(api_token) = env::var("HF_API_TOKEN") {
OlivierDehaene's avatar
OlivierDehaene committed
923
        envs.push(("HUGGING_FACE_HUB_TOKEN".into(), api_token.into()))
924
    };
925

926
927
928
    // 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
929
        envs.push((
930
931
932
933
934
            "WEIGHTS_CACHE_OVERRIDE".into(),
            weights_cache_override.into(),
        ));
    };

935
936
    // Start process
    tracing::info!("Starting download process.");
937
    let mut download_process = match Command::new("text-generation-server")
OlivierDehaene's avatar
OlivierDehaene committed
938
        .args(download_args)
939
        .env_clear()
OlivierDehaene's avatar
OlivierDehaene committed
940
        .envs(envs)
941
942
943
944
945
        .stdout(Stdio::piped())
        .stderr(Stdio::piped())
        .process_group(0)
        .spawn()
    {
946
947
        Ok(p) => p,
        Err(err) => {
948
949
950
            if err.kind() == io::ErrorKind::NotFound {
                tracing::error!("text-generation-server not found in PATH");
                tracing::error!("Please install it with `make install-server`")
951
952
            } else {
                tracing::error!("{}", err);
953
            }
954

955
956
957
            return Err(LauncherError::DownloadError);
        }
    };
958

959
    let download_stdout = BufReader::new(download_process.stdout.take().unwrap());
960

961
    thread::spawn(move || {
962
963
964
965
966
967
968
969
        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
970
        for line in download_stderr.lines().map_while(Result::ok) {
971
972
            err_sender.send(line).unwrap_or(());
        }
973
    });
974

975
    loop {
976
977
978
979
        if let Some(status) = download_process.try_wait().unwrap() {
            if status.success() {
                tracing::info!("Successfully downloaded weights.");
                break;
980
            }
981
982

            let mut err = String::new();
983
984
985
986
            while let Ok(line) = err_receiver.recv_timeout(Duration::from_millis(10)) {
                err = err + "\n" + &line;
            }

987
988
989
990
991
992
993
994
995
            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);
996
        }
997
        if !running.load(Ordering::SeqCst) {
OlivierDehaene's avatar
OlivierDehaene committed
998
            terminate("download", download_process, Duration::from_secs(10)).unwrap();
999
1000
1001
            return Ok(());
        }
        sleep(Duration::from_millis(100));
1002
    }
1003
1004
    Ok(())
}
1005

1006
#[allow(clippy::too_many_arguments)]
1007
1008
1009
fn spawn_shards(
    num_shard: usize,
    args: &Args,
1010
    cuda_graphs: Vec<usize>,
1011
    max_total_tokens: usize,
1012
    max_log_level: LevelFilter,
1013
    shutdown: Arc<AtomicBool>,
1014
1015
1016
1017
1018
1019
    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
1020
1021
    // Start shard processes
    for rank in 0..num_shard {
1022
1023
1024
1025
1026
1027
        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
1028
1029
1030
        let status_sender = status_sender.clone();
        let shutdown = shutdown.clone();
        let shutdown_sender = shutdown_sender.clone();
1031
        let otlp_endpoint = args.otlp_endpoint.clone();
1032
        let quantize = args.quantize;
Nicolas Patry's avatar
Nicolas Patry committed
1033
        let speculate = args.speculate;
1034
        let dtype = args.dtype;
1035
        let trust_remote_code = args.trust_remote_code;
1036
1037
1038
1039
        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;
1040
        let cuda_graphs_clone = cuda_graphs.clone();
1041
        let cuda_memory_fraction = args.cuda_memory_fraction;
Nicolas Patry's avatar
Nicolas Patry committed
1042
1043
        let rope_scaling = args.rope_scaling;
        let rope_factor = args.rope_factor;
1044
        let max_batch_size = args.max_batch_size;
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1045
1046
        thread::spawn(move || {
            shard_manager(
1047
                model_id,
1048
                revision,
1049
                quantize,
Nicolas Patry's avatar
Nicolas Patry committed
1050
                speculate,
1051
                dtype,
1052
                trust_remote_code,
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1053
1054
1055
1056
1057
                uds_path,
                rank,
                num_shard,
                master_addr,
                master_port,
1058
1059
                huggingface_hub_cache,
                weights_cache_override,
1060
                disable_custom_kernels,
1061
1062
                watermark_gamma,
                watermark_delta,
1063
                cuda_graphs_clone,
1064
                cuda_memory_fraction,
Nicolas Patry's avatar
Nicolas Patry committed
1065
1066
                rope_scaling,
                rope_factor,
1067
1068
                max_total_tokens,
                max_batch_size,
1069
                otlp_endpoint,
1070
                max_log_level,
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1071
1072
1073
1074
1075
1076
1077
1078
1079
1080
1081
1082
1083
1084
1085
1086
1087
1088
1089
1090
1091
                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));
            }
1092
            Ok(ShardStatus::Failed(rank)) => {
1093
                tracing::error!("Shard {rank} failed to start");
1094
                shutdown_shards(shutdown, shutdown_receiver);
1095
                return Err(LauncherError::ShardCannotStart);
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1096
1097
1098
            }
            Err(TryRecvError::Disconnected) => {
                tracing::error!("Shard status channel disconnected");
1099
                shutdown_shards(shutdown, shutdown_receiver);
1100
                return Err(LauncherError::ShardDisconnected);
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1101
1102
1103
            }
        }
    }
1104
1105
    Ok(())
}
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1106

1107
1108
1109
1110
1111
1112
1113
1114
1115
1116
1117
1118
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)
}

1119
fn spawn_webserver(
1120
    num_shard: usize,
1121
    args: Args,
1122
1123
1124
    max_input_tokens: usize,
    max_total_tokens: usize,
    max_batch_prefill_tokens: u32,
1125
    shutdown: Arc<AtomicBool>,
1126
    shutdown_receiver: &mpsc::Receiver<()>,
1127
) -> Result<Child, LauncherError> {
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1128
1129
1130
    // All shard started
    // Start webserver
    tracing::info!("Starting Webserver");
OlivierDehaene's avatar
OlivierDehaene committed
1131
    let mut router_args = vec![
1132
1133
        "--max-client-batch-size".to_string(),
        args.max_client_batch_size.to_string(),
1134
        "--max-concurrent-requests".to_string(),
1135
        args.max_concurrent_requests.to_string(),
1136
        "--max-best-of".to_string(),
1137
        args.max_best_of.to_string(),
1138
        "--max-stop-sequences".to_string(),
1139
        args.max_stop_sequences.to_string(),
Nicolas Patry's avatar
Nicolas Patry committed
1140
1141
        "--max-top-n-tokens".to_string(),
        args.max_top_n_tokens.to_string(),
1142
1143
        "--max-input-tokens".to_string(),
        max_input_tokens.to_string(),
1144
        "--max-total-tokens".to_string(),
1145
        max_total_tokens.to_string(),
1146
        "--max-batch-prefill-tokens".to_string(),
1147
        max_batch_prefill_tokens.to_string(),
1148
        "--waiting-served-ratio".to_string(),
1149
        args.waiting_served_ratio.to_string(),
1150
        "--max-waiting-tokens".to_string(),
1151
        args.max_waiting_tokens.to_string(),
1152
1153
        "--validation-workers".to_string(),
        args.validation_workers.to_string(),
1154
1155
        "--hostname".to_string(),
        args.hostname.to_string(),
1156
        "--port".to_string(),
1157
        args.port.to_string(),
1158
        "--master-shard-uds-path".to_string(),
1159
        format!("{}-0", args.shard_uds_path),
1160
        "--tokenizer-name".to_string(),
1161
        args.model_id,
1162
1163
    ];

drbh's avatar
drbh committed
1164
1165
1166
1167
1168
    // Grammar support
    if args.disable_grammar_support {
        router_args.push("--disable-grammar-support".to_string());
    }

1169
1170
1171
1172
1173
1174
    // 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());
    }

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

1181
1182
1183
1184
1185
1186
    // 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());
    }

1187
1188
    // Model optional revision
    if let Some(ref revision) = args.revision {
OlivierDehaene's avatar
OlivierDehaene committed
1189
1190
        router_args.push("--revision".to_string());
        router_args.push(revision.to_string())
1191
1192
    }

1193
    if args.json_output {
OlivierDehaene's avatar
OlivierDehaene committed
1194
        router_args.push("--json-output".to_string());
1195
1196
    }

1197
    // OpenTelemetry
1198
    if let Some(otlp_endpoint) = args.otlp_endpoint {
OlivierDehaene's avatar
OlivierDehaene committed
1199
1200
        router_args.push("--otlp-endpoint".to_string());
        router_args.push(otlp_endpoint);
1201
1202
1203
1204
    }

    // CORS origins
    for origin in args.cors_allow_origin.into_iter() {
OlivierDehaene's avatar
OlivierDehaene committed
1205
1206
        router_args.push("--cors-allow-origin".to_string());
        router_args.push(origin);
1207
1208
    }

1209
1210
    // Ngrok
    if args.ngrok {
OlivierDehaene's avatar
OlivierDehaene committed
1211
1212
        router_args.push("--ngrok".to_string());
        router_args.push("--ngrok-authtoken".to_string());
1213
1214
1215
        router_args.push(args.ngrok_authtoken.unwrap());
        router_args.push("--ngrok-edge".to_string());
        router_args.push(args.ngrok_edge.unwrap());
1216
1217
    }

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

1221
1222
    // Parse Inference API token
    if let Ok(api_token) = env::var("HF_API_TOKEN") {
OlivierDehaene's avatar
OlivierDehaene committed
1223
        envs.push(("HUGGING_FACE_HUB_TOKEN".into(), api_token.into()))
1224
    };
1225

1226
1227
1228
1229
1230
1231
1232
    // 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()))
    }

1233
    let mut webserver = match Command::new("text-generation-router")
OlivierDehaene's avatar
OlivierDehaene committed
1234
1235
        .args(router_args)
        .envs(envs)
1236
1237
1238
1239
1240
        .stdout(Stdio::piped())
        .stderr(Stdio::piped())
        .process_group(0)
        .spawn()
    {
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1241
1242
        Ok(p) => p,
        Err(err) => {
1243
            tracing::error!("Failed to start webserver: {}", err);
1244
1245
1246
            if err.kind() == io::ErrorKind::NotFound {
                tracing::error!("text-generation-router not found in PATH");
                tracing::error!("Please install it with `make install-router`")
1247
1248
            } else {
                tracing::error!("{}", err);
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1249
            }
1250

1251
            shutdown_shards(shutdown, shutdown_receiver);
1252
            return Err(LauncherError::WebserverCannotStart);
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1253
1254
1255
        }
    };

1256
1257
1258
    // Redirect STDOUT and STDERR to the console
    let webserver_stdout = webserver.stdout.take().unwrap();
    let webserver_stderr = webserver.stderr.take().unwrap();
1259
1260

    thread::spawn(move || {
1261
1262
        let stdout = BufReader::new(webserver_stdout);
        let stderr = BufReader::new(webserver_stderr);
1263
        for line in stdout.lines() {
1264
            println!("{}", line.unwrap());
1265
        }
1266
1267
        for line in stderr.lines() {
            println!("{}", line.unwrap());
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1268
        }
1269
1270
1271
    });
    Ok(webserver)
}
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1272

OlivierDehaene's avatar
OlivierDehaene committed
1273
1274
1275
1276
1277
1278
1279
1280
1281
1282
1283
1284
1285
1286
1287
1288
1289
1290
1291
1292
1293
1294
1295
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)
}

1296
1297
fn main() -> Result<(), LauncherError> {
    // Pattern match configuration
1298
    let args: Args = Args::parse();
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1299

1300
    // Filter events with LOG_LEVEL
1301
1302
1303
1304
1305
1306
1307
1308
1309
1310
1311
1312
1313
1314
1315
1316
    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);
1317

1318
    if args.json_output {
1319
1320
1321
1322
        tracing_subscriber::fmt()
            .with_env_filter(env_filter)
            .json()
            .init();
1323
    } else {
1324
1325
1326
1327
        tracing_subscriber::fmt()
            .with_env_filter(env_filter)
            .compact()
            .init();
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1328
1329
    }

1330
1331
1332
1333
1334
    if args.env {
        let env_runtime = env_runtime::Env::new();
        tracing::info!("{}", env_runtime);
    }

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

1337
1338
1339
1340
1341
1342
1343
1344
1345
1346
1347
1348
1349
1350
1351
1352
1353
1354
1355
1356
1357
1358
    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)?;
1359
1360
        let config: RawConfig = serde_json::from_str(&content)?;
        let config: Config = config.into();
1361
1362
1363
1364

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

1365
1366
1367
1368
1369
1370
1371
1372
        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);
1373
                }
1374
1375
1376
                Ok(max_default)
            } else {
                Ok(max_position_embeddings)
1377
            }
1378
1379
1380
1381
1382
        } else {
            Err(Box::new(LauncherError::ArgumentValidation(
                "no max defined".to_string(),
            )))
        }
1383
1384
1385
1386
1387
1388
1389
1390
1391
1392
1393
1394
1395
1396
1397
1398
1399
1400
1401
1402
1403
1404
1405
1406
1407
1408
1409
1410
1411
1412
1413
1414
1415
1416
1417
1418
1419
1420
1421
1422
1423
1424
1425
1426
1427
    };
    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
            }
        }
    };

1428
    // Validate args
1429
    if max_input_tokens >= max_total_tokens {
1430
        return Err(LauncherError::ArgumentValidation(
1431
            "`max_input_tokens must be < `max_total_tokens`".to_string(),
1432
1433
        ));
    }
1434
    if max_input_tokens as u32 > max_batch_prefill_tokens {
1435
        return Err(LauncherError::ArgumentValidation(format!(
1436
1437
            "`max_batch_prefill_tokens` must be >= `max_input_tokens`. Given: {} and {}",
            max_batch_prefill_tokens, max_input_tokens
1438
1439
        )));
    }
1440

1441
    let cuda_graphs = match (&args.cuda_graphs, &args.quantize) {
Nicolas Patry's avatar
Nicolas Patry committed
1442
        (Some(cuda_graphs), _) => cuda_graphs.iter().cloned().filter(|&c| c > 0).collect(),
1443
1444
1445
1446
1447
1448
1449
1450
1451
1452
1453
1454
1455
1456
1457
1458
1459
1460
1461
        #[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
        }
    };

1462
1463
1464
1465
1466
    if args.validation_workers == 0 {
        return Err(LauncherError::ArgumentValidation(
            "`validation_workers` must be > 0".to_string(),
        ));
    }
1467
1468
1469
1470
1471
1472
    if args.trust_remote_code {
        tracing::warn!(
            "`trust_remote_code` is set. Trusting that model `{}` do not contain malicious code.",
            args.model_id
        );
    }
1473
1474

    let num_shard = find_num_shards(args.sharded, args.num_shard)?;
1475
    if num_shard > 1 {
1476
1477
1478
1479
1480
        if matches!(args.quantize, Some(Quantization::Exl2)) {
            return Err(LauncherError::ArgumentValidation(
                "Sharding is currently not supported with `exl2` quantization".into(),
            ));
        }
1481
        tracing::info!("Sharding model on {num_shard} processes");
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1482
1483
    }

1484
    if let Some(ref max_batch_total_tokens) = args.max_batch_total_tokens {
1485
        if max_batch_prefill_tokens > *max_batch_total_tokens {
1486
1487
            return Err(LauncherError::ArgumentValidation(format!(
                "`max_batch_prefill_tokens` must be <= `max_batch_total_tokens`. Given: {} and {}",
1488
                max_batch_prefill_tokens, max_batch_total_tokens
1489
1490
            )));
        }
1491
        if max_total_tokens as u32 > *max_batch_total_tokens {
1492
1493
            return Err(LauncherError::ArgumentValidation(format!(
                "`max_total_tokens` must be <= `max_batch_total_tokens`. Given: {} and {}",
1494
                max_total_tokens, max_batch_total_tokens
1495
1496
1497
1498
            )));
        }
    }

1499
1500
1501
1502
1503
1504
1505
1506
1507
1508
1509
1510
1511
1512
    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(),
            ));
        }
    }

1513
1514
1515
1516
1517
1518
1519
    // 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");
1520

1521
    // Download and convert model weights
1522
    download_convert_model(&args, running.clone())?;
1523

OlivierDehaene's avatar
OlivierDehaene committed
1524
1525
1526
1527
1528
    if !running.load(Ordering::SeqCst) {
        // Launcher was asked to stop
        return Ok(());
    }

1529
    // Shared shutdown bool
1530
    let shutdown = Arc::new(AtomicBool::new(false));
1531
1532
1533
    // Shared shutdown channel
    // When shutting down, the main thread will wait for all senders to be dropped
    let (shutdown_sender, shutdown_receiver) = mpsc::channel();
1534

1535
1536
    // Shared channel to track shard status
    let (status_sender, status_receiver) = mpsc::channel();
1537

1538
1539
1540
    spawn_shards(
        num_shard,
        &args,
1541
        cuda_graphs,
1542
        max_total_tokens,
1543
        max_log_level,
1544
1545
1546
1547
1548
1549
1550
        shutdown.clone(),
        &shutdown_receiver,
        shutdown_sender,
        &status_receiver,
        status_sender,
        running.clone(),
    )?;
1551

1552
1553
1554
1555
1556
    // We might have received a termination signal
    if !running.load(Ordering::SeqCst) {
        shutdown_shards(shutdown, &shutdown_receiver);
        return Ok(());
    }
1557

1558
1559
1560
1561
1562
1563
1564
1565
1566
1567
1568
1569
1570
    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
    })?;
1571
1572
1573
1574
1575

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

    while running.load(Ordering::SeqCst) {
1576
        if let Ok(ShardStatus::Failed(rank)) = status_receiver.try_recv() {
OlivierDehaene's avatar
OlivierDehaene committed
1577
            tracing::error!("Shard {rank} crashed");
1578
1579
1580
1581
            exit_code = Err(LauncherError::ShardFailed);
            break;
        };

1582
        match webserver.try_wait().unwrap() {
1583
1584
1585
1586
1587
1588
1589
1590
1591
            Some(_) => {
                tracing::error!("Webserver Crashed");
                shutdown_shards(shutdown, &shutdown_receiver);
                return Err(LauncherError::WebserverFailed);
            }
            None => {
                sleep(Duration::from_millis(100));
            }
        };
1592
    }
1593
1594

    // Graceful termination
OlivierDehaene's avatar
OlivierDehaene committed
1595
    terminate("webserver", webserver, Duration::from_secs(90)).unwrap();
1596
1597
1598
    shutdown_shards(shutdown, &shutdown_receiver);

    exit_code
1599
}