main.rs 55.8 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
    max_input_tokens: usize,
485
    otlp_endpoint: Option<String>,
486
    log_level: LevelFilter,
487
    status_sender: mpsc::Sender<ShardStatus>,
488
    shutdown: Arc<AtomicBool>,
489
490
    _shutdown_sender: mpsc::Sender<()>,
) {
491
492
493
    // Enter shard-manager tracing span
    let _span = tracing::span!(tracing::Level::INFO, "shard-manager", rank = rank).entered();

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

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

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

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

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

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

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

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

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

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

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

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

564
565
566
    // Remove LOG_LEVEL if present
    envs.retain(|(name, _)| name != "LOG_LEVEL");

567
    // Torch Distributed Env vars
OlivierDehaene's avatar
OlivierDehaene committed
568
569
570
571
    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()));
572
    envs.push(("TORCH_NCCL_AVOID_RECORD_STREAMS".into(), "1".into()));
573

574
575
576
577
578
579
    // CUDA memory fraction
    envs.push((
        "CUDA_MEMORY_FRACTION".into(),
        cuda_memory_fraction.to_string().into(),
    ));

580
    // Safetensors load fast
OlivierDehaene's avatar
OlivierDehaene committed
581
    envs.push(("SAFETENSORS_FAST_GPU".into(), "1".into()));
582

583
584
585
    // Disable progress bar
    envs.push(("HF_HUB_DISABLE_PROGRESS_BARS".into(), "1".into()));

586
587
    // 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
588
    envs.push((
589
590
591
592
593
594
        "HF_HUB_ENABLE_HF_TRANSFER".into(),
        enable_hf_transfer.into(),
    ));

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

Nicolas Patry's avatar
Nicolas Patry committed
598
599
600
601
602
603
604
605
606
    // 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()));
    }

607
608
609
610
611
612
613
614
    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()));
    }

615
616
617
    // 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
618
        envs.push(("HUGGINGFACE_HUB_CACHE".into(), huggingface_hub_cache.into()));
619
620
621
622
623
    };

    // 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
624
        envs.push((
625
626
627
628
629
            "WEIGHTS_CACHE_OVERRIDE".into(),
            weights_cache_override.into(),
        ));
    };

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

643
644
    // If disable_custom_kernels is true, pass it to the shard as an env var
    if disable_custom_kernels {
OlivierDehaene's avatar
OlivierDehaene committed
645
        envs.push(("DISABLE_CUSTOM_KERNELS".into(), "True".into()))
646
647
648
649
    }

    // Watermark Gamma
    if let Some(watermark_gamma) = watermark_gamma {
OlivierDehaene's avatar
OlivierDehaene committed
650
        envs.push(("WATERMARK_GAMMA".into(), watermark_gamma.to_string().into()))
651
652
653
654
    }

    // Watermark Delta
    if let Some(watermark_delta) = watermark_delta {
OlivierDehaene's avatar
OlivierDehaene committed
655
        envs.push(("WATERMARK_DELTA".into(), watermark_delta.to_string().into()))
656
657
658
    }

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

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

    // Redirect STDOUT to the console
685
    let shard_stdout_reader = BufReader::new(p.stdout.take().unwrap());
686
    let shard_stderr_reader = BufReader::new(p.stderr.take().unwrap());
687

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

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

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

713
            if let Some(signal) = exit_status.signal() {
714
715
716
                tracing::error!("Shard process was signaled to shutdown with signal {signal}");
            }

717
            status_sender.send(ShardStatus::Failed(rank)).unwrap();
718
719
720
721
            return;
        }

        // We received a shutdown signal
722
        if shutdown.load(Ordering::SeqCst) {
723
            terminate("shard", p, Duration::from_secs(90)).unwrap();
724
725
726
727
728
            return;
        }

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

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

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

fn num_cuda_devices() -> Option<usize> {
752
753
754
755
    let devices = match env::var("CUDA_VISIBLE_DEVICES") {
        Ok(devices) => devices,
        Err(_) => env::var("NVIDIA_VISIBLE_DEVICES").ok()?,
    };
756
757
    let n_devices = devices.split(',').count();
    Some(n_devices)
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
786
787
788
789
790
}

#[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 {
791
792
793
794
795
796
797
            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()),
798
799
800
801
        }
    }
}

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

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

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

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

OlivierDehaene's avatar
OlivierDehaene committed
883
    let mut download_args = vec![
884
885
886
887
888
889
890
891
        "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(),
    ];
892

893
894
    // Model optional revision
    if let Some(revision) = &args.revision {
OlivierDehaene's avatar
OlivierDehaene committed
895
896
        download_args.push("--revision".to_string());
        download_args.push(revision.to_string())
897
    }
898

899
900
901
902
903
    // Trust remote code for automatic peft fusion
    if args.trust_remote_code {
        download_args.push("--trust-remote-code".to_string());
    }

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

907
908
909
    // Remove LOG_LEVEL if present
    envs.retain(|(name, _)| name != "LOG_LEVEL");

910
911
912
    // Disable progress bar
    envs.push(("HF_HUB_DISABLE_PROGRESS_BARS".into(), "1".into()));

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

919
920
    // 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
921
    envs.push((
922
923
924
        "HF_HUB_ENABLE_HF_TRANSFER".into(),
        enable_hf_transfer.into(),
    ));
925

926
927
    // Parse Inference API token
    if let Ok(api_token) = env::var("HF_API_TOKEN") {
928
        envs.push(("HF_TOKEN".into(), api_token.into()))
929
    };
930

931
932
933
    // 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
934
        envs.push((
935
936
937
938
939
            "WEIGHTS_CACHE_OVERRIDE".into(),
            weights_cache_override.into(),
        ));
    };

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

960
961
962
            return Err(LauncherError::DownloadError);
        }
    };
963

964
    let download_stdout = BufReader::new(download_process.stdout.take().unwrap());
965

966
    thread::spawn(move || {
967
968
969
970
971
972
973
974
        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
975
        for line in download_stderr.lines().map_while(Result::ok) {
976
977
            err_sender.send(line).unwrap_or(());
        }
978
    });
979

980
    loop {
981
982
983
984
        if let Some(status) = download_process.try_wait().unwrap() {
            if status.success() {
                tracing::info!("Successfully downloaded weights.");
                break;
985
            }
986
987

            let mut err = String::new();
988
989
990
991
            while let Ok(line) = err_receiver.recv_timeout(Duration::from_millis(10)) {
                err = err + "\n" + &line;
            }

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

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

1114
1115
1116
1117
1118
1119
1120
1121
1122
1123
1124
1125
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)
}

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

drbh's avatar
drbh committed
1171
1172
1173
1174
1175
    // Grammar support
    if args.disable_grammar_support {
        router_args.push("--disable-grammar-support".to_string());
    }

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

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

1188
1189
1190
1191
1192
1193
    // 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());
    }

1194
1195
    // Model optional revision
    if let Some(ref revision) = args.revision {
OlivierDehaene's avatar
OlivierDehaene committed
1196
1197
        router_args.push("--revision".to_string());
        router_args.push(revision.to_string())
1198
1199
    }

1200
    if args.json_output {
OlivierDehaene's avatar
OlivierDehaene committed
1201
        router_args.push("--json-output".to_string());
1202
1203
    }

1204
    // OpenTelemetry
1205
    if let Some(otlp_endpoint) = args.otlp_endpoint {
OlivierDehaene's avatar
OlivierDehaene committed
1206
1207
        router_args.push("--otlp-endpoint".to_string());
        router_args.push(otlp_endpoint);
1208
1209
1210
1211
    }

    // CORS origins
    for origin in args.cors_allow_origin.into_iter() {
OlivierDehaene's avatar
OlivierDehaene committed
1212
1213
        router_args.push("--cors-allow-origin".to_string());
        router_args.push(origin);
1214
1215
    }

1216
1217
    // Ngrok
    if args.ngrok {
OlivierDehaene's avatar
OlivierDehaene committed
1218
1219
        router_args.push("--ngrok".to_string());
        router_args.push("--ngrok-authtoken".to_string());
1220
1221
1222
        router_args.push(args.ngrok_authtoken.unwrap());
        router_args.push("--ngrok-edge".to_string());
        router_args.push(args.ngrok_edge.unwrap());
1223
1224
    }

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

1228
1229
    // Parse Inference API token
    if let Ok(api_token) = env::var("HF_API_TOKEN") {
1230
        envs.push(("HF_TOKEN".into(), api_token.into()))
1231
    };
1232

1233
1234
1235
1236
1237
1238
1239
    // 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()))
    }

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

1258
            shutdown_shards(shutdown, shutdown_receiver);
1259
            return Err(LauncherError::WebserverCannotStart);
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1260
1261
1262
        }
    };

1263
1264
1265
    // Redirect STDOUT and STDERR to the console
    let webserver_stdout = webserver.stdout.take().unwrap();
    let webserver_stderr = webserver.stderr.take().unwrap();
1266
1267

    thread::spawn(move || {
1268
1269
        let stdout = BufReader::new(webserver_stdout);
        let stderr = BufReader::new(webserver_stderr);
1270
        for line in stdout.lines() {
1271
            println!("{}", line.unwrap());
1272
        }
1273
1274
        for line in stderr.lines() {
            println!("{}", line.unwrap());
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1275
        }
1276
1277
1278
    });
    Ok(webserver)
}
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1279

OlivierDehaene's avatar
OlivierDehaene committed
1280
1281
1282
1283
1284
1285
1286
1287
1288
1289
1290
1291
1292
1293
1294
1295
1296
1297
1298
1299
1300
1301
1302
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)
}

1303
1304
fn main() -> Result<(), LauncherError> {
    // Pattern match configuration
1305
    let args: Args = Args::parse();
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1306

1307
    // Filter events with LOG_LEVEL
1308
1309
1310
1311
1312
1313
1314
1315
1316
1317
1318
1319
1320
1321
1322
1323
    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);
1324

1325
    if args.json_output {
1326
1327
1328
1329
        tracing_subscriber::fmt()
            .with_env_filter(env_filter)
            .json()
            .init();
1330
    } else {
1331
1332
1333
1334
        tracing_subscriber::fmt()
            .with_env_filter(env_filter)
            .compact()
            .init();
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1335
1336
    }

1337
1338
1339
1340
1341
    if args.env {
        let env_runtime = env_runtime::Env::new();
        tracing::info!("{}", env_runtime);
    }

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

1344
1345
1346
1347
1348
1349
1350
1351
1352
1353
1354
1355
1356
1357
1358
1359
1360
1361
1362
1363
1364
1365
    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)?;
1366
1367
        let config: RawConfig = serde_json::from_str(&content)?;
        let config: Config = config.into();
1368
1369
1370
1371

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

1372
1373
1374
1375
1376
1377
1378
1379
        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);
1380
                }
1381
1382
1383
                Ok(max_default)
            } else {
                Ok(max_position_embeddings)
1384
            }
1385
1386
1387
1388
1389
        } else {
            Err(Box::new(LauncherError::ArgumentValidation(
                "no max defined".to_string(),
            )))
        }
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
1428
1429
1430
1431
1432
1433
1434
    };
    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
            }
        }
    };

1435
    // Validate args
1436
    if max_input_tokens >= max_total_tokens {
1437
        return Err(LauncherError::ArgumentValidation(
1438
            "`max_input_tokens must be < `max_total_tokens`".to_string(),
1439
1440
        ));
    }
1441
    if max_input_tokens as u32 > max_batch_prefill_tokens {
1442
        return Err(LauncherError::ArgumentValidation(format!(
1443
1444
            "`max_batch_prefill_tokens` must be >= `max_input_tokens`. Given: {} and {}",
            max_batch_prefill_tokens, max_input_tokens
1445
1446
        )));
    }
1447

1448
    let cuda_graphs = match (&args.cuda_graphs, &args.quantize) {
Nicolas Patry's avatar
Nicolas Patry committed
1449
        (Some(cuda_graphs), _) => cuda_graphs.iter().cloned().filter(|&c| c > 0).collect(),
1450
1451
1452
1453
1454
1455
1456
1457
1458
1459
1460
1461
1462
1463
1464
1465
1466
1467
1468
        #[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
        }
    };

1469
1470
1471
1472
1473
    if args.validation_workers == 0 {
        return Err(LauncherError::ArgumentValidation(
            "`validation_workers` must be > 0".to_string(),
        ));
    }
1474
1475
1476
1477
1478
1479
    if args.trust_remote_code {
        tracing::warn!(
            "`trust_remote_code` is set. Trusting that model `{}` do not contain malicious code.",
            args.model_id
        );
    }
1480
1481

    let num_shard = find_num_shards(args.sharded, args.num_shard)?;
1482
    if num_shard > 1 {
1483
1484
1485
1486
1487
        if matches!(args.quantize, Some(Quantization::Exl2)) {
            return Err(LauncherError::ArgumentValidation(
                "Sharding is currently not supported with `exl2` quantization".into(),
            ));
        }
1488
        tracing::info!("Sharding model on {num_shard} processes");
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1489
1490
    }

1491
    if let Some(ref max_batch_total_tokens) = args.max_batch_total_tokens {
1492
        if max_batch_prefill_tokens > *max_batch_total_tokens {
1493
1494
            return Err(LauncherError::ArgumentValidation(format!(
                "`max_batch_prefill_tokens` must be <= `max_batch_total_tokens`. Given: {} and {}",
1495
                max_batch_prefill_tokens, max_batch_total_tokens
1496
1497
            )));
        }
1498
        if max_total_tokens as u32 > *max_batch_total_tokens {
1499
1500
            return Err(LauncherError::ArgumentValidation(format!(
                "`max_total_tokens` must be <= `max_batch_total_tokens`. Given: {} and {}",
1501
                max_total_tokens, max_batch_total_tokens
1502
1503
1504
1505
            )));
        }
    }

1506
1507
1508
1509
1510
1511
1512
1513
1514
1515
1516
1517
1518
1519
    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(),
            ));
        }
    }

1520
1521
1522
1523
1524
1525
1526
    // 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");
1527

1528
    // Download and convert model weights
1529
    download_convert_model(&args, running.clone())?;
1530

OlivierDehaene's avatar
OlivierDehaene committed
1531
1532
1533
1534
1535
    if !running.load(Ordering::SeqCst) {
        // Launcher was asked to stop
        return Ok(());
    }

1536
    // Shared shutdown bool
1537
    let shutdown = Arc::new(AtomicBool::new(false));
1538
1539
1540
    // Shared shutdown channel
    // When shutting down, the main thread will wait for all senders to be dropped
    let (shutdown_sender, shutdown_receiver) = mpsc::channel();
1541

1542
1543
    // Shared channel to track shard status
    let (status_sender, status_receiver) = mpsc::channel();
1544

1545
1546
1547
    spawn_shards(
        num_shard,
        &args,
1548
        cuda_graphs,
1549
        max_total_tokens,
1550
        max_input_tokens,
1551
        max_log_level,
1552
1553
1554
1555
1556
1557
1558
        shutdown.clone(),
        &shutdown_receiver,
        shutdown_sender,
        &status_receiver,
        status_sender,
        running.clone(),
    )?;
1559

1560
1561
1562
1563
1564
    // We might have received a termination signal
    if !running.load(Ordering::SeqCst) {
        shutdown_shards(shutdown, &shutdown_receiver);
        return Ok(());
    }
1565

1566
1567
1568
1569
1570
1571
1572
1573
1574
1575
1576
1577
1578
    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
    })?;
1579
1580
1581
1582
1583

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

    while running.load(Ordering::SeqCst) {
1584
        if let Ok(ShardStatus::Failed(rank)) = status_receiver.try_recv() {
OlivierDehaene's avatar
OlivierDehaene committed
1585
            tracing::error!("Shard {rank} crashed");
1586
1587
1588
1589
            exit_code = Err(LauncherError::ShardFailed);
            break;
        };

1590
        match webserver.try_wait().unwrap() {
1591
1592
1593
1594
1595
1596
1597
1598
1599
            Some(_) => {
                tracing::error!("Webserver Crashed");
                shutdown_shards(shutdown, &shutdown_receiver);
                return Err(LauncherError::WebserverFailed);
            }
            None => {
                sleep(Duration::from_millis(100));
            }
        };
1600
    }
1601
1602

    // Graceful termination
OlivierDehaene's avatar
OlivierDehaene committed
1603
    terminate("webserver", webserver, Duration::from_secs(90)).unwrap();
1604
1605
1606
    shutdown_shards(shutdown, &shutdown_receiver);

    exit_code
1607
}