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

25
26
mod env_runtime;

27
#[derive(Deserialize)]
28
struct RawConfig {
29
    max_position_embeddings: Option<usize>,
30
    n_positions: Option<usize>,
31
    model_type: Option<String>,
32
33
34
    max_seq_len: Option<usize>,
}

35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
#[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,
        }
    }
}

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

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

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

    /// 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
185
    #[clap(long, env)]
186
    revision: Option<String>,
187

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

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

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

206
    /// Whether you want the model to be quantized.
207
208
    #[clap(long, env, value_enum)]
    quantize: Option<Quantization>,
209

Nicolas Patry's avatar
Nicolas Patry committed
210
211
212
213
214
215
216
    /// 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>,

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

221
222
223
224
225
226
    /// 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,

227
228
229
    /// 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
230
231
    #[clap(default_value = "128", long, env)]
    max_concurrent_requests: usize,
232
233
234
235

    /// 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
236
237
    #[clap(default_value = "2", long, env)]
    max_best_of: usize,
238
239
240
241
242
243

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

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

255
256
257
258
    /// 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.
259
260
261
262
263
264
265
    /// 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>,
266
267
268
269
270
271
272
273
274

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

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

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

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

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

339
340
341
342
343
    /// 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>,

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

350
351
352
353
    /// The IP address to listen on
    #[clap(default_value = "0.0.0.0", long, env)]
    hostname: String,

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

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

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

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

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

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

    /// 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.
385
    #[clap(long, env)]
386
    disable_custom_kernels: bool,
387

388
389
390
391
392
    /// 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
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
    /// 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>,

413
    /// Outputs the logs in JSON format (useful for telemetry)
414
    #[clap(long, env)]
415
    json_output: bool,
416

417
418
    #[clap(long, env)]
    otlp_endpoint: Option<String>,
419

420
421
422
    #[clap(default_value = "text-generation-inference.router", long, env)]
    otlp_service_name: String,

423
424
    #[clap(long, env)]
    cors_allow_origin: Vec<String>,
Erik Kaunismäki's avatar
Erik Kaunismäki committed
425
426
427
428

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

429
430
431
432
    #[clap(long, env)]
    watermark_gamma: Option<f32>,
    #[clap(long, env)]
    watermark_delta: Option<f32>,
433

434
435
436
437
438
439
440
441
    /// Enable ngrok tunneling
    #[clap(long, env)]
    ngrok: bool,

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

442
    /// ngrok edge
443
    #[clap(long, env)]
444
    ngrok_edge: Option<String>,
445

446
447
448
449
450
    /// 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
451
452
453
454
455
    /// 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,

456
457
458
    /// Display a lot of information about your runtime environment
    #[clap(long, short, action)]
    env: bool,
459
460
461
462

    /// Control the maximum number of inputs that a client can send in a single request
    #[clap(default_value = "4", long, env)]
    max_client_batch_size: usize,
drbh's avatar
drbh committed
463
464
465
466
467

    /// Lora Adapters a list of adapter ids i.e. `repo/adapter1,repo/adapter2` to load during
    /// startup that will be available to callers via the `adapter_id` field in a request.
    #[clap(long, env)]
    lora_adapters: Option<String>,
468
469
470
471
472
473
474
475

    /// Disable sending of all usage statistics
    #[clap(default_value = "false", long, env)]
    disable_usage_stats: bool,

    /// Disable sending of crash reports, but allow anonymous usage statistics
    #[clap(default_value = "false", long, env)]
    disable_crash_reports: bool,
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
476
477
}

478
479
480
#[derive(Debug)]
enum ShardStatus {
    Ready,
481
    Failed(usize),
482
}
483

484
485
486
487
#[allow(clippy::too_many_arguments)]
fn shard_manager(
    model_id: String,
    revision: Option<String>,
488
    quantize: Option<Quantization>,
Nicolas Patry's avatar
Nicolas Patry committed
489
    speculate: Option<usize>,
490
    dtype: Option<Dtype>,
491
    trust_remote_code: bool,
492
493
494
495
496
497
498
499
500
501
    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>,
502
    cuda_graphs: Vec<usize>,
503
    cuda_memory_fraction: f32,
Nicolas Patry's avatar
Nicolas Patry committed
504
505
    rope_scaling: Option<RopeScaling>,
    rope_factor: Option<f32>,
506
507
    max_total_tokens: usize,
    max_batch_size: Option<usize>,
508
    max_input_tokens: usize,
drbh's avatar
drbh committed
509
    lora_adapters: Option<String>,
510
    otlp_endpoint: Option<String>,
511
    otlp_service_name: String,
512
    log_level: LevelFilter,
513
    status_sender: mpsc::Sender<ShardStatus>,
514
    shutdown: Arc<AtomicBool>,
515
516
    _shutdown_sender: mpsc::Sender<()>,
) {
517
518
519
    // Enter shard-manager tracing span
    let _span = tracing::span!(tracing::Level::INFO, "shard-manager", rank = rank).entered();

520
521
522
523
    // Get UDS path
    let uds_string = format!("{uds_path}-{rank}");
    let uds = Path::new(&uds_string);
    // Clean previous runs
524
525
526
    if uds.exists() {
        fs::remove_file(uds).unwrap();
    }
527
528

    // Process args
OlivierDehaene's avatar
OlivierDehaene committed
529
    let mut shard_args = vec![
530
531
532
533
534
        "serve".to_string(),
        model_id,
        "--uds-path".to_string(),
        uds_path,
        "--logger-level".to_string(),
535
        log_level.to_string().to_uppercase(),
536
537
538
        "--json-output".to_string(),
    ];

539
540
    // Activate trust remote code
    if trust_remote_code {
OlivierDehaene's avatar
OlivierDehaene committed
541
        shard_args.push("--trust-remote-code".to_string());
542
543
    }

544
545
    // Activate tensor parallelism
    if world_size > 1 {
OlivierDehaene's avatar
OlivierDehaene committed
546
        shard_args.push("--sharded".to_string());
547
548
    }

549
    if let Some(quantize) = quantize {
OlivierDehaene's avatar
OlivierDehaene committed
550
551
        shard_args.push("--quantize".to_string());
        shard_args.push(quantize.to_string())
552
    }
553

Nicolas Patry's avatar
Nicolas Patry committed
554
555
556
557
558
    if let Some(speculate) = speculate {
        shard_args.push("--speculate".to_string());
        shard_args.push(speculate.to_string())
    }

559
    if let Some(dtype) = dtype {
OlivierDehaene's avatar
OlivierDehaene committed
560
561
        shard_args.push("--dtype".to_string());
        shard_args.push(dtype.to_string())
562
563
    }

564
565
    // Model optional revision
    if let Some(revision) = revision {
OlivierDehaene's avatar
OlivierDehaene committed
566
567
        shard_args.push("--revision".to_string());
        shard_args.push(revision)
568
    }
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
569

Nicolas Patry's avatar
Nicolas Patry committed
570
571
572
573
574
575
    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)),
    };
576

577
    // OpenTelemetry Endpoint
578
    if let Some(otlp_endpoint) = otlp_endpoint {
OlivierDehaene's avatar
OlivierDehaene committed
579
580
        shard_args.push("--otlp-endpoint".to_string());
        shard_args.push(otlp_endpoint);
581
582
    }

583
584
585
586
    // OpenTelemetry Service Name
    shard_args.push("--otlp-service-name".to_string());
    shard_args.push(otlp_service_name);

587
588
589
590
    // 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());

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

594
595
596
    // Remove LOG_LEVEL if present
    envs.retain(|(name, _)| name != "LOG_LEVEL");

597
    // Torch Distributed Env vars
OlivierDehaene's avatar
OlivierDehaene committed
598
599
600
601
    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()));
602
    envs.push(("TORCH_NCCL_AVOID_RECORD_STREAMS".into(), "1".into()));
603

604
605
606
607
608
609
    // CUDA memory fraction
    envs.push((
        "CUDA_MEMORY_FRACTION".into(),
        cuda_memory_fraction.to_string().into(),
    ));

610
    // Safetensors load fast
OlivierDehaene's avatar
OlivierDehaene committed
611
    envs.push(("SAFETENSORS_FAST_GPU".into(), "1".into()));
612

613
614
615
    // Disable progress bar
    envs.push(("HF_HUB_DISABLE_PROGRESS_BARS".into(), "1".into()));

616
617
    // 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
618
    envs.push((
619
620
621
622
623
624
        "HF_HUB_ENABLE_HF_TRANSFER".into(),
        enable_hf_transfer.into(),
    ));

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

Nicolas Patry's avatar
Nicolas Patry committed
628
629
630
631
632
633
634
635
636
    // 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()));
    }

637
638
639
640
641
642
643
644
    envs.push((
        "MAX_TOTAL_TOKENS".into(),
        max_total_tokens.to_string().into(),
    ));
    if let Some(max_batch_size) = max_batch_size {
        envs.push(("MAX_BATCH_SIZE".into(), max_batch_size.to_string().into()));
    }

drbh's avatar
drbh committed
645
646
647
648
649
    // Lora Adapters
    if let Some(lora_adapters) = lora_adapters {
        envs.push(("LORA_ADAPTERS".into(), lora_adapters.into()));
    }

650
651
652
    // 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
653
        envs.push(("HUGGINGFACE_HUB_CACHE".into(), huggingface_hub_cache.into()));
654
655
656
657
658
    };

    // 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
659
        envs.push((
660
661
662
663
664
            "WEIGHTS_CACHE_OVERRIDE".into(),
            weights_cache_override.into(),
        ));
    };

665
    // Enable experimental support for cuda graphs
666
667
668
669
670
671
672
673
674
675
    if !cuda_graphs.is_empty() {
        envs.push((
            "CUDA_GRAPHS".into(),
            cuda_graphs
                .into_iter()
                .map(|c| c.to_string())
                .collect::<Vec<_>>()
                .join(",")
                .into(),
        ));
676
677
    }

678
679
    // If disable_custom_kernels is true, pass it to the shard as an env var
    if disable_custom_kernels {
OlivierDehaene's avatar
OlivierDehaene committed
680
        envs.push(("DISABLE_CUSTOM_KERNELS".into(), "True".into()))
681
682
683
684
    }

    // Watermark Gamma
    if let Some(watermark_gamma) = watermark_gamma {
OlivierDehaene's avatar
OlivierDehaene committed
685
        envs.push(("WATERMARK_GAMMA".into(), watermark_gamma.to_string().into()))
686
687
688
689
    }

    // Watermark Delta
    if let Some(watermark_delta) = watermark_delta {
OlivierDehaene's avatar
OlivierDehaene committed
690
        envs.push(("WATERMARK_DELTA".into(), watermark_delta.to_string().into()))
691
692
693
    }

    // Start process
694
    tracing::info!("Starting shard");
695
    let mut p = match Command::new("text-generation-server")
OlivierDehaene's avatar
OlivierDehaene committed
696
        .args(shard_args)
697
        .env_clear()
OlivierDehaene's avatar
OlivierDehaene committed
698
        .envs(envs)
699
700
701
702
703
        .stdout(Stdio::piped())
        .stderr(Stdio::piped())
        .process_group(0)
        .spawn()
    {
704
705
        Ok(p) => p,
        Err(err) => {
706
707
708
            if err.kind() == io::ErrorKind::NotFound {
                tracing::error!("text-generation-server not found in PATH");
                tracing::error!("Please install it with `make install-server`")
709
710
            }
            {
711
                tracing::error!("{}", err);
712
            }
713

714
            status_sender.send(ShardStatus::Failed(rank)).unwrap();
715
716
717
718
719
            return;
        }
    };

    // Redirect STDOUT to the console
720
    let shard_stdout_reader = BufReader::new(p.stdout.take().unwrap());
721
    let shard_stderr_reader = BufReader::new(p.stderr.take().unwrap());
722

723
    //stdout tracing thread
724
    thread::spawn(move || {
725
        log_lines(shard_stdout_reader.lines());
726
    });
727
728
729
    // 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
730
        for line in shard_stderr_reader.lines().map_while(Result::ok) {
731
732
733
            err_sender.send(line).unwrap_or(());
        }
    });
734
735
736
737
738
739

    let mut ready = false;
    let start_time = Instant::now();
    let mut wait_time = Instant::now();
    loop {
        // Process exited
740
        if let Some(exit_status) = p.try_wait().unwrap() {
741
742
743
744
            let mut err = String::new();
            while let Ok(line) = err_receiver.recv_timeout(Duration::from_millis(10)) {
                err = err + "\n" + &line;
            }
745

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

748
            if let Some(signal) = exit_status.signal() {
749
750
751
                tracing::error!("Shard process was signaled to shutdown with signal {signal}");
            }

752
            status_sender.send(ShardStatus::Failed(rank)).unwrap();
753
754
755
756
            return;
        }

        // We received a shutdown signal
757
        if shutdown.load(Ordering::SeqCst) {
758
            terminate("shard", p, Duration::from_secs(90)).unwrap();
759
760
761
762
763
            return;
        }

        // Shard is ready
        if uds.exists() && !ready {
764
            tracing::info!("Shard ready in {:?}", start_time.elapsed());
765
766
767
            status_sender.send(ShardStatus::Ready).unwrap();
            ready = true;
        } else if !ready && wait_time.elapsed() > Duration::from_secs(10) {
768
            tracing::info!("Waiting for shard to be ready...");
769
770
771
772
773
774
            wait_time = Instant::now();
        }
        sleep(Duration::from_millis(100));
    }
}

775
fn shutdown_shards(shutdown: Arc<AtomicBool>, shutdown_receiver: &mpsc::Receiver<()>) {
776
777
778
    tracing::info!("Shutting down shards");
    // Update shutdown value to true
    // This will be picked up by the shard manager
779
    shutdown.store(true, Ordering::SeqCst);
780
781
782
783
784
785
786

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

fn num_cuda_devices() -> Option<usize> {
787
788
    let devices = match env::var("CUDA_VISIBLE_DEVICES") {
        Ok(devices) => devices,
789
790
791
        Err(_) => match env::var("NVIDIA_VISIBLE_DEVICES") {
            Ok(devices) => devices,
            Err(_) => env::var("ZE_AFFINITY_MASK").ok()?,
Nicolas Patry's avatar
Nicolas Patry committed
792
        },
793
    };
794
795
    let n_devices = devices.split(',').count();
    Some(n_devices)
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
}

#[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 {
829
830
831
832
833
834
835
            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()),
836
837
838
839
        }
    }
}

840
841
842
843
844
845
846
847
848
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
849
    for line in lines.map_while(Result::ok) {
850
851
852
853
854
855
856
        match PythonLogMessage::try_from(&line) {
            Ok(log) => log.trace(),
            Err(_) => tracing::debug!("{line}"),
        }
    }
}

857
858
859
860
fn find_num_shards(
    sharded: Option<bool>,
    num_shard: Option<usize>,
) -> Result<usize, LauncherError> {
861
862
863
864
    // 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
865
            tracing::info!("Parsing num_shard from CUDA_VISIBLE_DEVICES/NVIDIA_VISIBLE_DEVICES/ZE_AFFINITY_MASK");
866
            let n_devices = num_cuda_devices()
867
                .expect("--num-shard and CUDA_VISIBLE_DEVICES/NVIDIA_VISIBLE_DEVICES/ZE_AFFINITY_MASK are not set");
868
            if n_devices <= 1 {
869
870
871
                return Err(LauncherError::NotEnoughCUDADevices(format!(
                    "`sharded` is true but only found {n_devices} CUDA devices"
                )));
872
            }
873
            n_devices
874
        }
875
876
877
        (Some(true), Some(num_shard)) => {
            // we can't have only one shard while sharded
            if num_shard <= 1 {
878
879
880
                return Err(LauncherError::ArgumentValidation(
                    "`sharded` is true but `num_shard` <= 1".to_string(),
                ));
881
882
            }
            num_shard
883
        }
884
885
886
887
        (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,
888
    };
889
    if num_shard < 1 {
890
891
892
        return Err(LauncherError::ArgumentValidation(
            "`num_shard` cannot be < 1".to_string(),
        ));
893
    }
894
    Ok(num_shard)
895
}
896

897
#[derive(Debug, Error)]
898
enum LauncherError {
899
    #[error("Invalid argument: {0}")]
900
    ArgumentValidation(String),
901
    #[error("not enough cuda devices: {0}")]
902
    NotEnoughCUDADevices(String),
903
    #[error("Download error")]
904
    DownloadError,
905
    #[error("Shard cannot start")]
906
    ShardCannotStart,
907
    #[error("Shard disconnected")]
908
    ShardDisconnected,
909
    #[error("Shard failed")]
910
    ShardFailed,
911
    #[error("Webserver failed")]
912
    WebserverFailed,
913
    #[error("Webserver cannot start")]
914
915
    WebserverCannotStart,
}
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
916

917
918
919
920
921
922
923
924
fn download_convert_model(
    model_id: &str,
    revision: Option<&str>,
    trust_remote_code: bool,
    huggingface_hub_cache: Option<&str>,
    weights_cache_override: Option<&str>,
    running: Arc<AtomicBool>,
) -> Result<(), LauncherError> {
925
926
927
    // Enter download tracing span
    let _span = tracing::span!(tracing::Level::INFO, "download").entered();

OlivierDehaene's avatar
OlivierDehaene committed
928
    let mut download_args = vec![
929
        "download-weights".to_string(),
930
        model_id.to_string(),
931
932
933
934
935
936
        "--extension".to_string(),
        ".safetensors".to_string(),
        "--logger-level".to_string(),
        "INFO".to_string(),
        "--json-output".to_string(),
    ];
937

938
    // Model optional revision
939
    if let Some(revision) = &revision {
OlivierDehaene's avatar
OlivierDehaene committed
940
941
        download_args.push("--revision".to_string());
        download_args.push(revision.to_string())
942
    }
943

944
    // Trust remote code for automatic peft fusion
945
    if trust_remote_code {
946
947
948
        download_args.push("--trust-remote-code".to_string());
    }

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

952
953
954
    // Remove LOG_LEVEL if present
    envs.retain(|(name, _)| name != "LOG_LEVEL");

955
956
957
    // Disable progress bar
    envs.push(("HF_HUB_DISABLE_PROGRESS_BARS".into(), "1".into()));

958
    // If huggingface_hub_cache is set, pass it to the download process
959
    // Useful when running inside a docker container
960
    if let Some(ref huggingface_hub_cache) = huggingface_hub_cache {
OlivierDehaene's avatar
OlivierDehaene committed
961
        envs.push(("HUGGINGFACE_HUB_CACHE".into(), huggingface_hub_cache.into()));
962
    };
963

964
965
    // 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
966
    envs.push((
967
968
969
        "HF_HUB_ENABLE_HF_TRANSFER".into(),
        enable_hf_transfer.into(),
    ));
970

971
972
    // Parse Inference API token
    if let Ok(api_token) = env::var("HF_API_TOKEN") {
973
        envs.push(("HF_TOKEN".into(), api_token.into()))
974
    };
975

976
977
    // If args.weights_cache_override is some, pass it to the download process
    // Useful when running inside a HuggingFace Inference Endpoint
978
    if let Some(weights_cache_override) = &weights_cache_override {
OlivierDehaene's avatar
OlivierDehaene committed
979
        envs.push((
980
981
982
983
984
            "WEIGHTS_CACHE_OVERRIDE".into(),
            weights_cache_override.into(),
        ));
    };

985
    // Start process
986
    tracing::info!("Starting check and download process for {model_id}");
987
    let mut download_process = match Command::new("text-generation-server")
OlivierDehaene's avatar
OlivierDehaene committed
988
        .args(download_args)
989
        .env_clear()
OlivierDehaene's avatar
OlivierDehaene committed
990
        .envs(envs)
991
992
993
994
995
        .stdout(Stdio::piped())
        .stderr(Stdio::piped())
        .process_group(0)
        .spawn()
    {
996
997
        Ok(p) => p,
        Err(err) => {
998
999
1000
            if err.kind() == io::ErrorKind::NotFound {
                tracing::error!("text-generation-server not found in PATH");
                tracing::error!("Please install it with `make install-server`")
1001
1002
            } else {
                tracing::error!("{}", err);
1003
            }
1004

1005
1006
1007
            return Err(LauncherError::DownloadError);
        }
    };
1008

1009
    let download_stdout = BufReader::new(download_process.stdout.take().unwrap());
1010

1011
    thread::spawn(move || {
1012
1013
1014
1015
1016
1017
1018
1019
        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
1020
        for line in download_stderr.lines().map_while(Result::ok) {
1021
1022
            err_sender.send(line).unwrap_or(());
        }
1023
    });
1024

1025
    loop {
1026
1027
        if let Some(status) = download_process.try_wait().unwrap() {
            if status.success() {
1028
                tracing::info!("Successfully downloaded weights for {model_id}");
1029
                break;
1030
            }
1031
1032

            let mut err = String::new();
1033
1034
1035
1036
            while let Ok(line) = err_receiver.recv_timeout(Duration::from_millis(10)) {
                err = err + "\n" + &line;
            }

1037
1038
1039
1040
1041
1042
1043
1044
1045
            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);
1046
        }
1047
        if !running.load(Ordering::SeqCst) {
OlivierDehaene's avatar
OlivierDehaene committed
1048
            terminate("download", download_process, Duration::from_secs(10)).unwrap();
1049
1050
1051
            return Ok(());
        }
        sleep(Duration::from_millis(100));
1052
    }
1053
1054
    Ok(())
}
1055

1056
#[allow(clippy::too_many_arguments)]
1057
1058
1059
fn spawn_shards(
    num_shard: usize,
    args: &Args,
1060
    cuda_graphs: Vec<usize>,
1061
    max_total_tokens: usize,
1062
    max_input_tokens: usize,
1063
    max_log_level: LevelFilter,
1064
    shutdown: Arc<AtomicBool>,
1065
1066
1067
1068
1069
1070
    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
1071
1072
    // Start shard processes
    for rank in 0..num_shard {
1073
1074
1075
1076
1077
1078
        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
1079
1080
1081
        let status_sender = status_sender.clone();
        let shutdown = shutdown.clone();
        let shutdown_sender = shutdown_sender.clone();
1082
        let otlp_endpoint = args.otlp_endpoint.clone();
1083
        let otlp_service_name = args.otlp_service_name.clone();
1084
        let quantize = args.quantize;
Nicolas Patry's avatar
Nicolas Patry committed
1085
        let speculate = args.speculate;
1086
        let dtype = args.dtype;
1087
        let trust_remote_code = args.trust_remote_code;
1088
1089
1090
1091
        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;
1092
        let cuda_graphs_clone = cuda_graphs.clone();
1093
        let cuda_memory_fraction = args.cuda_memory_fraction;
Nicolas Patry's avatar
Nicolas Patry committed
1094
1095
        let rope_scaling = args.rope_scaling;
        let rope_factor = args.rope_factor;
1096
        let max_batch_size = args.max_batch_size;
drbh's avatar
drbh committed
1097
        let lora_adapters = args.lora_adapters.clone();
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1098
1099
        thread::spawn(move || {
            shard_manager(
1100
                model_id,
1101
                revision,
1102
                quantize,
Nicolas Patry's avatar
Nicolas Patry committed
1103
                speculate,
1104
                dtype,
1105
                trust_remote_code,
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1106
1107
1108
1109
1110
                uds_path,
                rank,
                num_shard,
                master_addr,
                master_port,
1111
1112
                huggingface_hub_cache,
                weights_cache_override,
1113
                disable_custom_kernels,
1114
1115
                watermark_gamma,
                watermark_delta,
1116
                cuda_graphs_clone,
1117
                cuda_memory_fraction,
Nicolas Patry's avatar
Nicolas Patry committed
1118
1119
                rope_scaling,
                rope_factor,
1120
1121
                max_total_tokens,
                max_batch_size,
1122
                max_input_tokens,
drbh's avatar
drbh committed
1123
                lora_adapters,
1124
                otlp_endpoint,
1125
                otlp_service_name,
1126
                max_log_level,
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1127
1128
1129
1130
1131
1132
1133
1134
1135
1136
1137
1138
1139
1140
1141
1142
1143
1144
1145
1146
1147
                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));
            }
1148
            Ok(ShardStatus::Failed(rank)) => {
1149
                tracing::error!("Shard {rank} failed to start");
1150
                shutdown_shards(shutdown, shutdown_receiver);
1151
                return Err(LauncherError::ShardCannotStart);
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1152
1153
1154
            }
            Err(TryRecvError::Disconnected) => {
                tracing::error!("Shard status channel disconnected");
1155
                shutdown_shards(shutdown, shutdown_receiver);
1156
                return Err(LauncherError::ShardDisconnected);
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1157
1158
1159
            }
        }
    }
1160
1161
    Ok(())
}
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1162

1163
1164
1165
1166
1167
1168
1169
1170
1171
1172
1173
1174
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)
}

1175
fn spawn_webserver(
1176
    num_shard: usize,
1177
    args: Args,
1178
1179
1180
    max_input_tokens: usize,
    max_total_tokens: usize,
    max_batch_prefill_tokens: u32,
1181
    shutdown: Arc<AtomicBool>,
1182
    shutdown_receiver: &mpsc::Receiver<()>,
1183
) -> Result<Child, LauncherError> {
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1184
1185
1186
    // All shard started
    // Start webserver
    tracing::info!("Starting Webserver");
OlivierDehaene's avatar
OlivierDehaene committed
1187
    let mut router_args = vec![
1188
1189
        "--max-client-batch-size".to_string(),
        args.max_client_batch_size.to_string(),
1190
        "--max-concurrent-requests".to_string(),
1191
        args.max_concurrent_requests.to_string(),
1192
        "--max-best-of".to_string(),
1193
        args.max_best_of.to_string(),
1194
        "--max-stop-sequences".to_string(),
1195
        args.max_stop_sequences.to_string(),
Nicolas Patry's avatar
Nicolas Patry committed
1196
1197
        "--max-top-n-tokens".to_string(),
        args.max_top_n_tokens.to_string(),
1198
1199
        "--max-input-tokens".to_string(),
        max_input_tokens.to_string(),
1200
        "--max-total-tokens".to_string(),
1201
        max_total_tokens.to_string(),
1202
        "--max-batch-prefill-tokens".to_string(),
1203
        max_batch_prefill_tokens.to_string(),
1204
        "--waiting-served-ratio".to_string(),
1205
        args.waiting_served_ratio.to_string(),
1206
        "--max-waiting-tokens".to_string(),
1207
        args.max_waiting_tokens.to_string(),
1208
1209
        "--validation-workers".to_string(),
        args.validation_workers.to_string(),
1210
1211
        "--hostname".to_string(),
        args.hostname.to_string(),
1212
        "--port".to_string(),
1213
        args.port.to_string(),
1214
        "--master-shard-uds-path".to_string(),
1215
        format!("{}-0", args.shard_uds_path),
1216
        "--tokenizer-name".to_string(),
1217
        args.model_id,
1218
1219
    ];

1220
1221
1222
1223
1224
1225
1226
1227
    // Pass usage stats flags to router
    if args.disable_usage_stats {
        router_args.push("--disable-usage-stats".to_string());
    }
    if args.disable_crash_reports {
        router_args.push("--disable-crash-reports".to_string());
    }

drbh's avatar
drbh committed
1228
1229
1230
1231
1232
    // Grammar support
    if args.disable_grammar_support {
        router_args.push("--disable-grammar-support".to_string());
    }

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

1239
1240
1241
1242
1243
1244
    // 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());
    }

1245
1246
1247
1248
1249
1250
    // 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());
    }

1251
1252
    // Model optional revision
    if let Some(ref revision) = args.revision {
OlivierDehaene's avatar
OlivierDehaene committed
1253
1254
        router_args.push("--revision".to_string());
        router_args.push(revision.to_string())
1255
1256
    }

1257
    if args.json_output {
OlivierDehaene's avatar
OlivierDehaene committed
1258
        router_args.push("--json-output".to_string());
1259
1260
    }

1261
    // OpenTelemetry
1262
    if let Some(otlp_endpoint) = args.otlp_endpoint {
OlivierDehaene's avatar
OlivierDehaene committed
1263
1264
        router_args.push("--otlp-endpoint".to_string());
        router_args.push(otlp_endpoint);
1265
1266
    }

1267
1268
1269
1270
1271
    // OpenTelemetry
    let otlp_service_name = args.otlp_service_name;
    router_args.push("--otlp-service-name".to_string());
    router_args.push(otlp_service_name);

1272
1273
    // CORS origins
    for origin in args.cors_allow_origin.into_iter() {
OlivierDehaene's avatar
OlivierDehaene committed
1274
1275
        router_args.push("--cors-allow-origin".to_string());
        router_args.push(origin);
1276
1277
    }

Erik Kaunismäki's avatar
Erik Kaunismäki committed
1278
1279
1280
1281
1282
    // API Key
    if let Some(api_key) = args.api_key {
        router_args.push("--api-key".to_string());
        router_args.push(api_key);
    }
1283
1284
    // Ngrok
    if args.ngrok {
OlivierDehaene's avatar
OlivierDehaene committed
1285
1286
        router_args.push("--ngrok".to_string());
        router_args.push("--ngrok-authtoken".to_string());
1287
1288
1289
        router_args.push(args.ngrok_authtoken.unwrap());
        router_args.push("--ngrok-edge".to_string());
        router_args.push(args.ngrok_edge.unwrap());
1290
1291
    }

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

1295
1296
    // Parse Inference API token
    if let Ok(api_token) = env::var("HF_API_TOKEN") {
1297
        envs.push(("HF_TOKEN".into(), api_token.into()))
1298
    };
1299

1300
1301
1302
1303
1304
1305
1306
    // 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()))
    }

1307
    let mut webserver = match Command::new("text-generation-router")
OlivierDehaene's avatar
OlivierDehaene committed
1308
1309
        .args(router_args)
        .envs(envs)
1310
1311
1312
1313
1314
        .stdout(Stdio::piped())
        .stderr(Stdio::piped())
        .process_group(0)
        .spawn()
    {
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1315
1316
        Ok(p) => p,
        Err(err) => {
1317
            tracing::error!("Failed to start webserver: {}", err);
1318
1319
1320
            if err.kind() == io::ErrorKind::NotFound {
                tracing::error!("text-generation-router not found in PATH");
                tracing::error!("Please install it with `make install-router`")
1321
1322
            } else {
                tracing::error!("{}", err);
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1323
            }
1324

1325
            shutdown_shards(shutdown, shutdown_receiver);
1326
            return Err(LauncherError::WebserverCannotStart);
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1327
1328
1329
        }
    };

1330
1331
1332
    // Redirect STDOUT and STDERR to the console
    let webserver_stdout = webserver.stdout.take().unwrap();
    let webserver_stderr = webserver.stderr.take().unwrap();
1333
1334

    thread::spawn(move || {
1335
1336
        let stdout = BufReader::new(webserver_stdout);
        let stderr = BufReader::new(webserver_stderr);
1337
        for line in stdout.lines() {
1338
            println!("{}", line.unwrap());
1339
        }
1340
1341
        for line in stderr.lines() {
            println!("{}", line.unwrap());
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1342
        }
1343
1344
1345
    });
    Ok(webserver)
}
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1346

OlivierDehaene's avatar
OlivierDehaene committed
1347
1348
1349
1350
1351
1352
1353
1354
1355
1356
1357
1358
1359
1360
1361
1362
1363
1364
1365
1366
1367
1368
1369
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)
}

1370
1371
fn main() -> Result<(), LauncherError> {
    // Pattern match configuration
1372
    let args: Args = Args::parse();
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1373

1374
    // Filter events with LOG_LEVEL
1375
1376
1377
1378
1379
1380
1381
1382
1383
1384
1385
1386
1387
1388
1389
1390
    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);
1391

1392
    if args.json_output {
1393
1394
1395
1396
        tracing_subscriber::fmt()
            .with_env_filter(env_filter)
            .json()
            .init();
1397
    } else {
1398
1399
1400
1401
        tracing_subscriber::fmt()
            .with_env_filter(env_filter)
            .compact()
            .init();
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1402
1403
    }

1404
1405
1406
1407
1408
    if args.env {
        let env_runtime = env_runtime::Env::new();
        tracing::info!("{}", env_runtime);
    }

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

1411
1412
1413
1414
1415
    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
Nicolas Patry's avatar
Nicolas Patry committed
1416
1417
1418
1419
1420
1421
1422

            let api = if let Ok(token) = std::env::var("HF_TOKEN") {
                // env variable has precedence over on file token.
                ApiBuilder::new().with_token(Some(token)).build()?
            } else {
                Api::new()?
            };
1423
1424
1425
1426
1427
1428
1429
1430
1431
1432
1433
1434
1435
1436
1437
1438
            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)?;
1439
        let config: RawConfig = serde_json::from_str(&content)?;
1440
1441
1442
1443
1444

        if config.model_type == Some("gemma2".to_string()) {
            tracing::info!("Forcing flash decoding because of softcap usage");
            std::env::set_var("FLASH_DECODING", "1");
        }
1445
        let config: Config = config.into();
1446
1447
1448
1449

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

1450
1451
1452
1453
1454
1455
1456
1457
        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);
1458
                }
1459
1460
1461
                Ok(max_default)
            } else {
                Ok(max_position_embeddings)
1462
            }
1463
1464
1465
1466
1467
        } else {
            Err(Box::new(LauncherError::ArgumentValidation(
                "no max defined".to_string(),
            )))
        }
1468
1469
1470
1471
1472
1473
1474
1475
1476
1477
1478
1479
1480
1481
1482
1483
1484
1485
1486
1487
1488
1489
1490
1491
1492
1493
1494
1495
1496
1497
1498
1499
1500
1501
1502
1503
1504
1505
1506
1507
1508
1509
1510
1511
1512
    };
    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
            }
        }
    };

1513
    // Validate args
1514
    if max_input_tokens >= max_total_tokens {
1515
        return Err(LauncherError::ArgumentValidation(
1516
            "`max_input_tokens must be < `max_total_tokens`".to_string(),
1517
1518
        ));
    }
1519
    if max_input_tokens as u32 > max_batch_prefill_tokens {
1520
        return Err(LauncherError::ArgumentValidation(format!(
1521
1522
            "`max_batch_prefill_tokens` must be >= `max_input_tokens`. Given: {} and {}",
            max_batch_prefill_tokens, max_input_tokens
1523
1524
        )));
    }
1525

1526
    let cuda_graphs = match (&args.cuda_graphs, &args.quantize) {
Nicolas Patry's avatar
Nicolas Patry committed
1527
        (Some(cuda_graphs), _) => cuda_graphs.iter().cloned().filter(|&c| c > 0).collect(),
1528
1529
1530
1531
1532
1533
1534
1535
1536
1537
1538
1539
1540
1541
1542
1543
1544
1545
1546
        #[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
        }
    };

1547
1548
1549
1550
1551
    if args.validation_workers == 0 {
        return Err(LauncherError::ArgumentValidation(
            "`validation_workers` must be > 0".to_string(),
        ));
    }
1552
1553
1554
1555
1556
1557
    if args.trust_remote_code {
        tracing::warn!(
            "`trust_remote_code` is set. Trusting that model `{}` do not contain malicious code.",
            args.model_id
        );
    }
1558
1559

    let num_shard = find_num_shards(args.sharded, args.num_shard)?;
1560
    if num_shard > 1 {
1561
1562
1563
1564
1565
        if matches!(args.quantize, Some(Quantization::Exl2)) {
            return Err(LauncherError::ArgumentValidation(
                "Sharding is currently not supported with `exl2` quantization".into(),
            ));
        }
1566
        tracing::info!("Sharding model on {num_shard} processes");
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1567
1568
    }

1569
    if let Some(ref max_batch_total_tokens) = args.max_batch_total_tokens {
1570
        if max_batch_prefill_tokens > *max_batch_total_tokens {
1571
1572
            return Err(LauncherError::ArgumentValidation(format!(
                "`max_batch_prefill_tokens` must be <= `max_batch_total_tokens`. Given: {} and {}",
1573
                max_batch_prefill_tokens, max_batch_total_tokens
1574
1575
            )));
        }
1576
        if max_total_tokens as u32 > *max_batch_total_tokens {
1577
1578
            return Err(LauncherError::ArgumentValidation(format!(
                "`max_total_tokens` must be <= `max_batch_total_tokens`. Given: {} and {}",
1579
                max_total_tokens, max_batch_total_tokens
1580
1581
1582
1583
            )));
        }
    }

1584
1585
1586
1587
1588
1589
1590
1591
1592
1593
1594
1595
1596
1597
    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(),
            ));
        }
    }

1598
1599
1600
1601
1602
1603
1604
    // 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");
1605

1606
    // Download and convert model weights
1607
1608
1609
1610
1611
1612
1613
1614
1615
1616
1617
1618
    download_convert_model(
        &args.model_id,
        args.revision.as_deref(),
        args.trust_remote_code,
        args.huggingface_hub_cache.as_deref(),
        args.weights_cache_override.as_deref(),
        running.clone(),
    )?;

    // Download and convert lora adapters if any
    if let Some(lora_adapters) = &args.lora_adapters {
        for adapter in lora_adapters.split(',') {
1619
1620
1621
1622
            // skip download if a path is provided
            if adapter.contains('=') {
                continue;
            }
1623
1624
1625
1626
1627
1628
1629
1630
1631
1632
            download_convert_model(
                adapter,
                None,
                args.trust_remote_code,
                args.huggingface_hub_cache.as_deref(),
                args.weights_cache_override.as_deref(),
                running.clone(),
            )?;
        }
    }
1633

OlivierDehaene's avatar
OlivierDehaene committed
1634
1635
1636
1637
1638
    if !running.load(Ordering::SeqCst) {
        // Launcher was asked to stop
        return Ok(());
    }

1639
    // Shared shutdown bool
1640
    let shutdown = Arc::new(AtomicBool::new(false));
1641
1642
1643
    // Shared shutdown channel
    // When shutting down, the main thread will wait for all senders to be dropped
    let (shutdown_sender, shutdown_receiver) = mpsc::channel();
1644

1645
1646
    // Shared channel to track shard status
    let (status_sender, status_receiver) = mpsc::channel();
1647

1648
1649
1650
    spawn_shards(
        num_shard,
        &args,
1651
        cuda_graphs,
1652
        max_total_tokens,
1653
        max_input_tokens,
1654
        max_log_level,
1655
1656
1657
1658
1659
1660
1661
        shutdown.clone(),
        &shutdown_receiver,
        shutdown_sender,
        &status_receiver,
        status_sender,
        running.clone(),
    )?;
1662

1663
1664
1665
1666
1667
    // We might have received a termination signal
    if !running.load(Ordering::SeqCst) {
        shutdown_shards(shutdown, &shutdown_receiver);
        return Ok(());
    }
1668

1669
1670
1671
1672
1673
1674
1675
1676
1677
1678
1679
1680
1681
    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
    })?;
1682
1683
1684
1685
1686

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

    while running.load(Ordering::SeqCst) {
1687
        if let Ok(ShardStatus::Failed(rank)) = status_receiver.try_recv() {
OlivierDehaene's avatar
OlivierDehaene committed
1688
            tracing::error!("Shard {rank} crashed");
1689
1690
1691
1692
            exit_code = Err(LauncherError::ShardFailed);
            break;
        };

1693
        match webserver.try_wait().unwrap() {
1694
1695
1696
1697
1698
1699
1700
1701
1702
            Some(_) => {
                tracing::error!("Webserver Crashed");
                shutdown_shards(shutdown, &shutdown_receiver);
                return Err(LauncherError::WebserverFailed);
            }
            None => {
                sleep(Duration::from_millis(100));
            }
        };
1703
    }
1704
1705

    // Graceful termination
OlivierDehaene's avatar
OlivierDehaene committed
1706
    terminate("webserver", webserver, Duration::from_secs(90)).unwrap();
1707
1708
1709
    shutdown_shards(shutdown, &shutdown_receiver);

    exit_code
1710
}