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

22
23
mod env_runtime;

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

30
31
#[derive(Clone, Copy, Debug, ValueEnum)]
enum Quantization {
32
    /// 4 bit quantization. Requires a specific AWQ quantized model:
33
    ///   <https://hf.co/models?search=awq>.
34
    /// Should replace GPTQ models wherever possible because of the better latency
35
36
37
    Awq,
    /// 8 bit quantization, doesn't require specific model.
    /// Should be a drop-in replacement to bitsandbytes with much better performance.
38
    /// Kernels are from <https://github.com/NetEase-FuXi/EETQ.git>
39
    Eetq,
40
    /// 4 bit quantization. Requires a specific GTPQ quantized model: <https://hf.co/models?search=gptq>.
41
    /// text-generation-inference will use exllama (faster) kernels wherever possible, and use
42
43
44
45
46
47
48
49
50
    /// 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"
    )]
51
    Bitsandbytes,
52
53
    /// 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
54
    BitsandbytesNF4,
55
56
    /// 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
57
    BitsandbytesFP4,
Nicolas Patry's avatar
Nicolas Patry committed
58
59
60
61
62
    /// [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,
63
64
65
66
67
68
}

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

96
97
98
#[derive(Clone, Copy, Debug, ValueEnum)]
enum Dtype {
    Float16,
99
    #[clap(name = "bfloat16")]
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
    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
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
#[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
137
138
139
140
/// App Configuration
#[derive(Parser, Debug)]
#[clap(author, version, about, long_about = None)]
struct Args {
141
142
143
144
145
    /// 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
146
    #[clap(default_value = "bigscience/bloom-560m", long, env)]
147
    model_id: String,
148
149
150

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

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

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

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

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

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

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

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

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

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

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

Nicolas Patry's avatar
Nicolas Patry committed
213
214
215
216
217
218
219
220
    /// 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,

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

415
416
417
    /// Display a lot of information about your runtime environment
    #[clap(long, short, action)]
    env: bool,
418
419
420
421

    /// 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
422
423
}

424
425
426
#[derive(Debug)]
enum ShardStatus {
    Ready,
427
    Failed(usize),
428
}
429

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

462
463
464
465
    // Get UDS path
    let uds_string = format!("{uds_path}-{rank}");
    let uds = Path::new(&uds_string);
    // Clean previous runs
466
467
468
    if uds.exists() {
        fs::remove_file(uds).unwrap();
    }
469
470

    // Process args
OlivierDehaene's avatar
OlivierDehaene committed
471
    let mut shard_args = vec![
472
473
474
475
476
477
478
479
480
        "serve".to_string(),
        model_id,
        "--uds-path".to_string(),
        uds_path,
        "--logger-level".to_string(),
        "INFO".to_string(),
        "--json-output".to_string(),
    ];

481
482
    // Activate trust remote code
    if trust_remote_code {
OlivierDehaene's avatar
OlivierDehaene committed
483
        shard_args.push("--trust-remote-code".to_string());
484
485
    }

486
487
    // Activate tensor parallelism
    if world_size > 1 {
OlivierDehaene's avatar
OlivierDehaene committed
488
        shard_args.push("--sharded".to_string());
489
490
    }

491
    if let Some(quantize) = quantize {
OlivierDehaene's avatar
OlivierDehaene committed
492
493
        shard_args.push("--quantize".to_string());
        shard_args.push(quantize.to_string())
494
    }
495

Nicolas Patry's avatar
Nicolas Patry committed
496
497
498
499
500
    if let Some(speculate) = speculate {
        shard_args.push("--speculate".to_string());
        shard_args.push(speculate.to_string())
    }

501
    if let Some(dtype) = dtype {
OlivierDehaene's avatar
OlivierDehaene committed
502
503
        shard_args.push("--dtype".to_string());
        shard_args.push(dtype.to_string())
504
505
    }

506
507
    // Model optional revision
    if let Some(revision) = revision {
OlivierDehaene's avatar
OlivierDehaene committed
508
509
        shard_args.push("--revision".to_string());
        shard_args.push(revision)
510
    }
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
511

Nicolas Patry's avatar
Nicolas Patry committed
512
513
514
515
516
517
    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)),
    };
518

519
520
    // OpenTelemetry
    if let Some(otlp_endpoint) = otlp_endpoint {
OlivierDehaene's avatar
OlivierDehaene committed
521
522
        shard_args.push("--otlp-endpoint".to_string());
        shard_args.push(otlp_endpoint);
523
524
525
    }

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

528
529
530
    // Remove LOG_LEVEL if present
    envs.retain(|(name, _)| name != "LOG_LEVEL");

531
    // Torch Distributed Env vars
OlivierDehaene's avatar
OlivierDehaene committed
532
533
534
535
    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()));
536
    envs.push(("TORCH_NCCL_AVOID_RECORD_STREAMS".into(), "1".into()));
537

538
539
540
541
542
543
    // CUDA memory fraction
    envs.push((
        "CUDA_MEMORY_FRACTION".into(),
        cuda_memory_fraction.to_string().into(),
    ));

544
    // Safetensors load fast
OlivierDehaene's avatar
OlivierDehaene committed
545
    envs.push(("SAFETENSORS_FAST_GPU".into(), "1".into()));
546

547
548
549
    // Disable progress bar
    envs.push(("HF_HUB_DISABLE_PROGRESS_BARS".into(), "1".into()));

550
551
    // 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
552
    envs.push((
553
554
555
556
557
558
        "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
559
        envs.push(("HUGGING_FACE_HUB_TOKEN".into(), api_token.into()))
560
561
    };

Nicolas Patry's avatar
Nicolas Patry committed
562
563
564
565
566
567
568
569
570
    // 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()));
    }

571
572
573
574
575
576
577
578
    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()));
    }

579
580
581
    // 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
582
        envs.push(("HUGGINGFACE_HUB_CACHE".into(), huggingface_hub_cache.into()));
583
584
585
586
587
    };

    // 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
588
        envs.push((
589
590
591
592
593
            "WEIGHTS_CACHE_OVERRIDE".into(),
            weights_cache_override.into(),
        ));
    };

594
    // Enable experimental support for cuda graphs
595
596
597
598
599
600
601
602
603
604
    if !cuda_graphs.is_empty() {
        envs.push((
            "CUDA_GRAPHS".into(),
            cuda_graphs
                .into_iter()
                .map(|c| c.to_string())
                .collect::<Vec<_>>()
                .join(",")
                .into(),
        ));
605
606
    }

607
608
    // If disable_custom_kernels is true, pass it to the shard as an env var
    if disable_custom_kernels {
OlivierDehaene's avatar
OlivierDehaene committed
609
        envs.push(("DISABLE_CUSTOM_KERNELS".into(), "True".into()))
610
611
612
613
    }

    // Watermark Gamma
    if let Some(watermark_gamma) = watermark_gamma {
OlivierDehaene's avatar
OlivierDehaene committed
614
        envs.push(("WATERMARK_GAMMA".into(), watermark_gamma.to_string().into()))
615
616
617
618
    }

    // Watermark Delta
    if let Some(watermark_delta) = watermark_delta {
OlivierDehaene's avatar
OlivierDehaene committed
619
        envs.push(("WATERMARK_DELTA".into(), watermark_delta.to_string().into()))
620
621
622
    }

    // Start process
623
    tracing::info!("Starting shard");
624
    let mut p = match Command::new("text-generation-server")
OlivierDehaene's avatar
OlivierDehaene committed
625
        .args(shard_args)
626
        .env_clear()
OlivierDehaene's avatar
OlivierDehaene committed
627
        .envs(envs)
628
629
630
631
632
        .stdout(Stdio::piped())
        .stderr(Stdio::piped())
        .process_group(0)
        .spawn()
    {
633
634
        Ok(p) => p,
        Err(err) => {
635
636
637
            if err.kind() == io::ErrorKind::NotFound {
                tracing::error!("text-generation-server not found in PATH");
                tracing::error!("Please install it with `make install-server`")
638
639
            }
            {
640
                tracing::error!("{}", err);
641
            }
642

643
            status_sender.send(ShardStatus::Failed(rank)).unwrap();
644
645
646
647
648
            return;
        }
    };

    // Redirect STDOUT to the console
649
    let shard_stdout_reader = BufReader::new(p.stdout.take().unwrap());
650
    let shard_stderr_reader = BufReader::new(p.stderr.take().unwrap());
651

652
    //stdout tracing thread
653
    thread::spawn(move || {
654
        log_lines(shard_stdout_reader.lines());
655
    });
656
657
658
    // 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
659
        for line in shard_stderr_reader.lines().map_while(Result::ok) {
660
661
662
            err_sender.send(line).unwrap_or(());
        }
    });
663
664
665
666
667
668

    let mut ready = false;
    let start_time = Instant::now();
    let mut wait_time = Instant::now();
    loop {
        // Process exited
669
        if let Some(exit_status) = p.try_wait().unwrap() {
670
671
672
673
            let mut err = String::new();
            while let Ok(line) = err_receiver.recv_timeout(Duration::from_millis(10)) {
                err = err + "\n" + &line;
            }
674

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

677
            if let Some(signal) = exit_status.signal() {
678
679
680
                tracing::error!("Shard process was signaled to shutdown with signal {signal}");
            }

681
            status_sender.send(ShardStatus::Failed(rank)).unwrap();
682
683
684
685
            return;
        }

        // We received a shutdown signal
686
        if shutdown.load(Ordering::SeqCst) {
687
            terminate("shard", p, Duration::from_secs(90)).unwrap();
688
689
690
691
692
            return;
        }

        // Shard is ready
        if uds.exists() && !ready {
693
            tracing::info!("Shard ready in {:?}", start_time.elapsed());
694
695
696
            status_sender.send(ShardStatus::Ready).unwrap();
            ready = true;
        } else if !ready && wait_time.elapsed() > Duration::from_secs(10) {
697
            tracing::info!("Waiting for shard to be ready...");
698
699
700
701
702
703
            wait_time = Instant::now();
        }
        sleep(Duration::from_millis(100));
    }
}

704
fn shutdown_shards(shutdown: Arc<AtomicBool>, shutdown_receiver: &mpsc::Receiver<()>) {
705
706
707
    tracing::info!("Shutting down shards");
    // Update shutdown value to true
    // This will be picked up by the shard manager
708
    shutdown.store(true, Ordering::SeqCst);
709
710
711
712
713
714
715

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

fn num_cuda_devices() -> Option<usize> {
716
717
718
719
    let devices = match env::var("CUDA_VISIBLE_DEVICES") {
        Ok(devices) => devices,
        Err(_) => env::var("NVIDIA_VISIBLE_DEVICES").ok()?,
    };
720
721
    let n_devices = devices.split(',').count();
    Some(n_devices)
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
752
753
754
755
756
757
758
759
760
761
762
763
764
765
}

#[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),
        }
    }
}

766
767
768
769
770
771
772
773
774
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
775
    for line in lines.map_while(Result::ok) {
776
777
778
779
780
781
782
        match PythonLogMessage::try_from(&line) {
            Ok(log) => log.trace(),
            Err(_) => tracing::debug!("{line}"),
        }
    }
}

783
784
785
786
fn find_num_shards(
    sharded: Option<bool>,
    num_shard: Option<usize>,
) -> Result<usize, LauncherError> {
787
788
789
790
    // 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
791
792
793
            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");
794
            if n_devices <= 1 {
795
796
797
                return Err(LauncherError::NotEnoughCUDADevices(format!(
                    "`sharded` is true but only found {n_devices} CUDA devices"
                )));
798
            }
799
            n_devices
800
        }
801
802
803
        (Some(true), Some(num_shard)) => {
            // we can't have only one shard while sharded
            if num_shard <= 1 {
804
805
806
                return Err(LauncherError::ArgumentValidation(
                    "`sharded` is true but `num_shard` <= 1".to_string(),
                ));
807
808
            }
            num_shard
809
        }
810
811
812
813
        (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,
814
    };
815
    if num_shard < 1 {
816
817
818
        return Err(LauncherError::ArgumentValidation(
            "`num_shard` cannot be < 1".to_string(),
        ));
819
    }
820
    Ok(num_shard)
821
}
822

823
#[derive(Debug, Error)]
824
enum LauncherError {
825
    #[error("Invalid argument: {0}")]
826
    ArgumentValidation(String),
827
    #[error("not enough cuda devices: {0}")]
828
    NotEnoughCUDADevices(String),
829
    #[error("Download error")]
830
    DownloadError,
831
    #[error("Shard cannot start")]
832
    ShardCannotStart,
833
    #[error("Shard disconnected")]
834
    ShardDisconnected,
835
    #[error("Shard failed")]
836
    ShardFailed,
837
    #[error("Webserver failed")]
838
    WebserverFailed,
839
    #[error("Webserver cannot start")]
840
841
    WebserverCannotStart,
}
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
842

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

OlivierDehaene's avatar
OlivierDehaene committed
847
    let mut download_args = vec![
848
849
850
851
852
853
854
855
        "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(),
    ];
856

857
858
    // Model optional revision
    if let Some(revision) = &args.revision {
OlivierDehaene's avatar
OlivierDehaene committed
859
860
        download_args.push("--revision".to_string());
        download_args.push(revision.to_string())
861
    }
862

863
864
865
866
867
    // Trust remote code for automatic peft fusion
    if args.trust_remote_code {
        download_args.push("--trust-remote-code".to_string());
    }

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

871
872
873
    // Remove LOG_LEVEL if present
    envs.retain(|(name, _)| name != "LOG_LEVEL");

874
875
876
    // Disable progress bar
    envs.push(("HF_HUB_DISABLE_PROGRESS_BARS".into(), "1".into()));

877
    // If huggingface_hub_cache is set, pass it to the download process
878
879
    // Useful when running inside a docker container
    if let Some(ref huggingface_hub_cache) = args.huggingface_hub_cache {
OlivierDehaene's avatar
OlivierDehaene committed
880
        envs.push(("HUGGINGFACE_HUB_CACHE".into(), huggingface_hub_cache.into()));
881
    };
882

883
884
    // 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
885
    envs.push((
886
887
888
        "HF_HUB_ENABLE_HF_TRANSFER".into(),
        enable_hf_transfer.into(),
    ));
889

890
891
    // Parse Inference API token
    if let Ok(api_token) = env::var("HF_API_TOKEN") {
OlivierDehaene's avatar
OlivierDehaene committed
892
        envs.push(("HUGGING_FACE_HUB_TOKEN".into(), api_token.into()))
893
    };
894

895
896
897
    // 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
898
        envs.push((
899
900
901
902
903
            "WEIGHTS_CACHE_OVERRIDE".into(),
            weights_cache_override.into(),
        ));
    };

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

924
925
926
            return Err(LauncherError::DownloadError);
        }
    };
927

928
    let download_stdout = BufReader::new(download_process.stdout.take().unwrap());
929

930
    thread::spawn(move || {
931
932
933
934
935
936
937
938
        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
939
        for line in download_stderr.lines().map_while(Result::ok) {
940
941
            err_sender.send(line).unwrap_or(());
        }
942
    });
943

944
    loop {
945
946
947
948
        if let Some(status) = download_process.try_wait().unwrap() {
            if status.success() {
                tracing::info!("Successfully downloaded weights.");
                break;
949
            }
950
951

            let mut err = String::new();
952
953
954
955
            while let Ok(line) = err_receiver.recv_timeout(Duration::from_millis(10)) {
                err = err + "\n" + &line;
            }

956
957
958
959
960
961
962
963
964
            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);
965
        }
966
        if !running.load(Ordering::SeqCst) {
OlivierDehaene's avatar
OlivierDehaene committed
967
            terminate("download", download_process, Duration::from_secs(10)).unwrap();
968
969
970
            return Ok(());
        }
        sleep(Duration::from_millis(100));
971
    }
972
973
    Ok(())
}
974

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

1074
1075
1076
1077
1078
1079
1080
1081
1082
1083
1084
1085
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)
}

1086
fn spawn_webserver(
1087
    num_shard: usize,
1088
    args: Args,
1089
1090
1091
    max_input_tokens: usize,
    max_total_tokens: usize,
    max_batch_prefill_tokens: u32,
1092
    shutdown: Arc<AtomicBool>,
1093
    shutdown_receiver: &mpsc::Receiver<()>,
1094
) -> Result<Child, LauncherError> {
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1095
1096
1097
    // All shard started
    // Start webserver
    tracing::info!("Starting Webserver");
OlivierDehaene's avatar
OlivierDehaene committed
1098
    let mut router_args = vec![
1099
1100
        "--max-client-batch-size".to_string(),
        args.max_client_batch_size.to_string(),
1101
        "--max-concurrent-requests".to_string(),
1102
        args.max_concurrent_requests.to_string(),
1103
        "--max-best-of".to_string(),
1104
        args.max_best_of.to_string(),
1105
        "--max-stop-sequences".to_string(),
1106
        args.max_stop_sequences.to_string(),
Nicolas Patry's avatar
Nicolas Patry committed
1107
1108
        "--max-top-n-tokens".to_string(),
        args.max_top_n_tokens.to_string(),
1109
1110
        "--max-input-tokens".to_string(),
        max_input_tokens.to_string(),
1111
        "--max-total-tokens".to_string(),
1112
        max_total_tokens.to_string(),
1113
        "--max-batch-prefill-tokens".to_string(),
1114
        max_batch_prefill_tokens.to_string(),
1115
        "--waiting-served-ratio".to_string(),
1116
        args.waiting_served_ratio.to_string(),
1117
        "--max-waiting-tokens".to_string(),
1118
        args.max_waiting_tokens.to_string(),
1119
1120
        "--validation-workers".to_string(),
        args.validation_workers.to_string(),
1121
1122
        "--hostname".to_string(),
        args.hostname.to_string(),
1123
        "--port".to_string(),
1124
        args.port.to_string(),
1125
        "--master-shard-uds-path".to_string(),
1126
        format!("{}-0", args.shard_uds_path),
1127
        "--tokenizer-name".to_string(),
1128
        args.model_id,
1129
1130
    ];

drbh's avatar
drbh committed
1131
1132
1133
1134
1135
    // Grammar support
    if args.disable_grammar_support {
        router_args.push("--disable-grammar-support".to_string());
    }

1136
1137
1138
1139
1140
1141
    // 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());
    }

1142
1143
1144
1145
1146
1147
    // 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());
    }

1148
1149
1150
1151
1152
1153
    // 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());
    }

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

1160
    if args.json_output {
OlivierDehaene's avatar
OlivierDehaene committed
1161
        router_args.push("--json-output".to_string());
1162
1163
    }

1164
    // OpenTelemetry
1165
    if let Some(otlp_endpoint) = args.otlp_endpoint {
OlivierDehaene's avatar
OlivierDehaene committed
1166
1167
        router_args.push("--otlp-endpoint".to_string());
        router_args.push(otlp_endpoint);
1168
1169
1170
1171
    }

    // CORS origins
    for origin in args.cors_allow_origin.into_iter() {
OlivierDehaene's avatar
OlivierDehaene committed
1172
1173
        router_args.push("--cors-allow-origin".to_string());
        router_args.push(origin);
1174
1175
    }

1176
1177
    // Ngrok
    if args.ngrok {
OlivierDehaene's avatar
OlivierDehaene committed
1178
1179
        router_args.push("--ngrok".to_string());
        router_args.push("--ngrok-authtoken".to_string());
1180
1181
1182
        router_args.push(args.ngrok_authtoken.unwrap());
        router_args.push("--ngrok-edge".to_string());
        router_args.push(args.ngrok_edge.unwrap());
1183
1184
    }

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

1188
1189
    // Parse Inference API token
    if let Ok(api_token) = env::var("HF_API_TOKEN") {
OlivierDehaene's avatar
OlivierDehaene committed
1190
        envs.push(("HUGGING_FACE_HUB_TOKEN".into(), api_token.into()))
1191
    };
1192

1193
1194
1195
1196
1197
1198
1199
    // 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()))
    }

1200
    let mut webserver = match Command::new("text-generation-router")
OlivierDehaene's avatar
OlivierDehaene committed
1201
1202
        .args(router_args)
        .envs(envs)
1203
1204
1205
1206
1207
        .stdout(Stdio::piped())
        .stderr(Stdio::piped())
        .process_group(0)
        .spawn()
    {
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1208
1209
        Ok(p) => p,
        Err(err) => {
1210
            tracing::error!("Failed to start webserver: {}", err);
1211
1212
1213
            if err.kind() == io::ErrorKind::NotFound {
                tracing::error!("text-generation-router not found in PATH");
                tracing::error!("Please install it with `make install-router`")
1214
1215
            } else {
                tracing::error!("{}", err);
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1216
            }
1217

1218
            shutdown_shards(shutdown, shutdown_receiver);
1219
            return Err(LauncherError::WebserverCannotStart);
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1220
1221
1222
        }
    };

1223
1224
1225
    // Redirect STDOUT and STDERR to the console
    let webserver_stdout = webserver.stdout.take().unwrap();
    let webserver_stderr = webserver.stderr.take().unwrap();
1226
1227

    thread::spawn(move || {
1228
1229
        let stdout = BufReader::new(webserver_stdout);
        let stderr = BufReader::new(webserver_stderr);
1230
        for line in stdout.lines() {
1231
            println!("{}", line.unwrap());
1232
        }
1233
1234
        for line in stderr.lines() {
            println!("{}", line.unwrap());
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1235
        }
1236
1237
1238
    });
    Ok(webserver)
}
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1239

OlivierDehaene's avatar
OlivierDehaene committed
1240
1241
1242
1243
1244
1245
1246
1247
1248
1249
1250
1251
1252
1253
1254
1255
1256
1257
1258
1259
1260
1261
1262
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)
}

1263
1264
fn main() -> Result<(), LauncherError> {
    // Pattern match configuration
1265
    let args: Args = Args::parse();
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1266

1267
1268
1269
1270
    // Filter events with LOG_LEVEL
    let env_filter =
        EnvFilter::try_from_env("LOG_LEVEL").unwrap_or_else(|_| EnvFilter::new("info"));

1271
    if args.json_output {
1272
1273
1274
1275
        tracing_subscriber::fmt()
            .with_env_filter(env_filter)
            .json()
            .init();
1276
    } else {
1277
1278
1279
1280
        tracing_subscriber::fmt()
            .with_env_filter(env_filter)
            .compact()
            .init();
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1281
1282
    }

1283
1284
1285
1286
1287
    if args.env {
        let env_runtime = env_runtime::Env::new();
        tracing::info!("{}", env_runtime);
    }

Nicolas Patry's avatar
Nicolas Patry committed
1288
    tracing::info!("{:#?}", args);
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
    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;
Nicolas Patry's avatar
Nicolas Patry committed
1321
1322
1323
1324
1325
1326
                    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);
                    }
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
1361
1362
1363
1364
1365
1366
1367
1368
1369
1370
1371
1372
1373
1374
1375
1376
1377
1378
1379
1380
1381
1382
1383
                    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
            }
        }
    };

1384
    // Validate args
1385
    if max_input_tokens >= max_total_tokens {
1386
        return Err(LauncherError::ArgumentValidation(
1387
            "`max_input_tokens must be < `max_total_tokens`".to_string(),
1388
1389
        ));
    }
1390
    if max_input_tokens as u32 > max_batch_prefill_tokens {
1391
        return Err(LauncherError::ArgumentValidation(format!(
1392
1393
            "`max_batch_prefill_tokens` must be >= `max_input_tokens`. Given: {} and {}",
            max_batch_prefill_tokens, max_input_tokens
1394
1395
        )));
    }
1396

1397
    let cuda_graphs = match (&args.cuda_graphs, &args.quantize) {
Nicolas Patry's avatar
Nicolas Patry committed
1398
        (Some(cuda_graphs), _) => cuda_graphs.iter().cloned().filter(|&c| c > 0).collect(),
1399
1400
1401
1402
1403
1404
1405
1406
1407
1408
1409
1410
1411
1412
1413
1414
1415
1416
1417
        #[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
        }
    };

1418
1419
1420
1421
1422
    if args.validation_workers == 0 {
        return Err(LauncherError::ArgumentValidation(
            "`validation_workers` must be > 0".to_string(),
        ));
    }
1423
1424
1425
1426
1427
1428
    if args.trust_remote_code {
        tracing::warn!(
            "`trust_remote_code` is set. Trusting that model `{}` do not contain malicious code.",
            args.model_id
        );
    }
1429
1430

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

1435
    if let Some(ref max_batch_total_tokens) = args.max_batch_total_tokens {
1436
        if max_batch_prefill_tokens > *max_batch_total_tokens {
1437
1438
            return Err(LauncherError::ArgumentValidation(format!(
                "`max_batch_prefill_tokens` must be <= `max_batch_total_tokens`. Given: {} and {}",
1439
                max_batch_prefill_tokens, max_batch_total_tokens
1440
1441
            )));
        }
1442
        if max_total_tokens as u32 > *max_batch_total_tokens {
1443
1444
            return Err(LauncherError::ArgumentValidation(format!(
                "`max_total_tokens` must be <= `max_batch_total_tokens`. Given: {} and {}",
1445
                max_total_tokens, max_batch_total_tokens
1446
1447
1448
1449
            )));
        }
    }

1450
1451
1452
1453
1454
1455
1456
1457
1458
1459
1460
1461
1462
1463
    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(),
            ));
        }
    }

1464
1465
1466
1467
1468
1469
1470
    // 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");
1471

1472
    // Download and convert model weights
1473
    download_convert_model(&args, running.clone())?;
1474

OlivierDehaene's avatar
OlivierDehaene committed
1475
1476
1477
1478
1479
    if !running.load(Ordering::SeqCst) {
        // Launcher was asked to stop
        return Ok(());
    }

1480
    // Shared shutdown bool
1481
    let shutdown = Arc::new(AtomicBool::new(false));
1482
1483
1484
    // Shared shutdown channel
    // When shutting down, the main thread will wait for all senders to be dropped
    let (shutdown_sender, shutdown_receiver) = mpsc::channel();
1485

1486
1487
    // Shared channel to track shard status
    let (status_sender, status_receiver) = mpsc::channel();
1488

1489
1490
1491
    spawn_shards(
        num_shard,
        &args,
1492
        cuda_graphs,
1493
        max_total_tokens,
1494
1495
1496
1497
1498
1499
1500
        shutdown.clone(),
        &shutdown_receiver,
        shutdown_sender,
        &status_receiver,
        status_sender,
        running.clone(),
    )?;
1501

1502
1503
1504
1505
1506
    // We might have received a termination signal
    if !running.load(Ordering::SeqCst) {
        shutdown_shards(shutdown, &shutdown_receiver);
        return Ok(());
    }
1507

1508
1509
1510
1511
1512
1513
1514
1515
1516
1517
1518
1519
1520
    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
    })?;
1521
1522
1523
1524
1525

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

    while running.load(Ordering::SeqCst) {
1526
        if let Ok(ShardStatus::Failed(rank)) = status_receiver.try_recv() {
OlivierDehaene's avatar
OlivierDehaene committed
1527
            tracing::error!("Shard {rank} crashed");
1528
1529
1530
1531
            exit_code = Err(LauncherError::ShardFailed);
            break;
        };

1532
        match webserver.try_wait().unwrap() {
1533
1534
1535
1536
1537
1538
1539
1540
1541
            Some(_) => {
                tracing::error!("Webserver Crashed");
                shutdown_shards(shutdown, &shutdown_receiver);
                return Err(LauncherError::WebserverFailed);
            }
            None => {
                sleep(Duration::from_millis(100));
            }
        };
1542
    }
1543
1544

    // Graceful termination
OlivierDehaene's avatar
OlivierDehaene committed
1545
    terminate("webserver", webserver, Duration::from_secs(90)).unwrap();
1546
1547
1548
    shutdown_shards(shutdown, &shutdown_receiver);

    exit_code
1549
}