main.rs 53.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::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
    /// 4 bit quantization. Requires a specific GTPQ quantized model: <https://hf.co/models?search=gptq>.
59
    /// text-generation-inference will use exllama (faster) kernels wherever possible, and use
60
61
62
63
64
65
66
67
68
    /// triton kernel (wider support) when it's not.
    /// AWQ has faster kernels.
    Gptq,
    /// Bitsandbytes 8bit. Can be applied on any model, will cut the memory requirement in half,
    /// but it is known that the model will be much slower to run than the native f16.
    #[deprecated(
        since = "1.1.0",
        note = "Use `eetq` instead, which provides better latencies overall and is drop-in in most cases"
    )]
69
    Bitsandbytes,
70
71
    /// 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
72
    BitsandbytesNF4,
73
74
    /// 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
75
    BitsandbytesFP4,
Nicolas Patry's avatar
Nicolas Patry committed
76
77
78
79
80
    /// [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,
81
82
83
84
85
86
}

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 {
87
88
            #[allow(deprecated)]
            // Use `eetq` instead, which provides better latencies overall and is drop-in in most cases
89
90
91
            Quantization::Bitsandbytes => {
                write!(f, "bitsandbytes")
            }
Nicolas Patry's avatar
Nicolas Patry committed
92
93
94
95
96
97
            Quantization::BitsandbytesNF4 => {
                write!(f, "bitsandbytes-nf4")
            }
            Quantization::BitsandbytesFP4 => {
                write!(f, "bitsandbytes-fp4")
            }
98
99
100
            Quantization::Gptq => {
                write!(f, "gptq")
            }
101
102
103
            Quantization::Awq => {
                write!(f, "awq")
            }
104
105
106
            Quantization::Eetq => {
                write!(f, "eetq")
            }
Nicolas Patry's avatar
Nicolas Patry committed
107
108
109
            Quantization::Fp8 => {
                write!(f, "fp8")
            }
110
111
112
113
        }
    }
}

114
115
116
#[derive(Clone, Copy, Debug, ValueEnum)]
enum Dtype {
    Float16,
117
    #[clap(name = "bfloat16")]
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
    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
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
#[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
155
156
157
158
/// App Configuration
#[derive(Parser, Debug)]
#[clap(author, version, about, long_about = None)]
struct Args {
159
160
161
162
163
    /// 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
164
    #[clap(default_value = "bigscience/bloom-560m", long, env)]
165
    model_id: String,
166
167
168

    /// 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
169
    #[clap(long, env)]
170
    revision: Option<String>,
171

172
173
174
175
176
    /// The number of tokenizer workers used for payload validation and truncation inside the
    /// router.
    #[clap(default_value = "2", long, env)]
    validation_workers: usize,

177
    /// Whether to shard the model across multiple GPUs
178
179
    /// By default text-generation-inference will use all available GPUs to run
    /// the model. Setting it to `false` deactivates `num_shard`.
180
181
    #[clap(long, env)]
    sharded: Option<bool>,
182
183

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

190
    /// Whether you want the model to be quantized.
191
192
    #[clap(long, env, value_enum)]
    quantize: Option<Quantization>,
193

Nicolas Patry's avatar
Nicolas Patry committed
194
195
196
197
198
199
200
    /// 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>,

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

205
206
207
208
209
210
    /// 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,

211
212
213
    /// 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
214
215
    #[clap(default_value = "128", long, env)]
    max_concurrent_requests: usize,
216
217
218
219

    /// 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
220
221
    #[clap(default_value = "2", long, env)]
    max_best_of: usize,
222
223
224
225
226
227

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

Nicolas Patry's avatar
Nicolas Patry committed
231
232
233
234
235
236
237
238
    /// This is the maximum allowed value for clients to set `top_n_tokens`.
    /// `top_n_tokens is used to return information about the the `n` most likely
    /// tokens at each generation step, instead of just the sampled token. This
    /// information can be used for downstream tasks like for classification or
    /// ranking.
    #[clap(default_value = "5", long, env)]
    max_top_n_tokens: u32,

239
240
241
242
    /// 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.
243
244
245
246
247
248
249
    /// 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>,
250
251
252
253
254
255
256
257
258

    /// 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.
259
260
261
    /// Default to min(max_position_embeddings, 4096)
    #[clap(long, env)]
    max_total_tokens: Option<usize>,
262
263
264
265
266
267
268
269
270
271
272

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

276
277
278
    /// 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.
279
280
281
    /// Default to `max_input_tokens + 50` to give a bit of room.
    #[clap(long, env)]
    max_batch_prefill_tokens: Option<u32>,
282

283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
    /// **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.
300
301
    #[clap(long, env)]
    max_batch_total_tokens: Option<u32>,
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319

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

323
324
325
326
327
    /// 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>,

328
329
    /// Specify the batch sizes to compute cuda graphs for.
    /// Use "0" to disable.
330
331
332
    /// Default = "1,2,4,8,16,32"
    #[clap(long, env, value_delimiter = ',')]
    cuda_graphs: Option<Vec<usize>>,
333

334
335
336
337
    /// The IP address to listen on
    #[clap(default_value = "0.0.0.0", long, env)]
    hostname: String,

338
    /// The port to listen on.
339
    #[clap(default_value = "3000", long, short, env)]
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
340
    port: u16,
341
342
343

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

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

    /// The address the master port will listen on. (setting used by torch distributed)
352
    #[clap(default_value = "29500", long, env)]
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
353
    master_port: usize,
354
355
356

    /// The location of the huggingface hub cache.
    /// Used to override the location if you want to provide a mounted disk for instance
357
    #[clap(long, env)]
358
    huggingface_hub_cache: Option<String>,
359
360
361

    /// The location of the huggingface hub cache.
    /// Used to override the location if you want to provide a mounted disk for instance
362
363
    #[clap(long, env)]
    weights_cache_override: Option<String>,
364
365
366
367
368

    /// 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.
369
    #[clap(long, env)]
370
    disable_custom_kernels: bool,
371

372
373
374
375
376
    /// 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
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
    /// 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>,

397
    /// Outputs the logs in JSON format (useful for telemetry)
398
    #[clap(long, env)]
399
    json_output: bool,
400

401
402
    #[clap(long, env)]
    otlp_endpoint: Option<String>,
403

404
405
    #[clap(long, env)]
    cors_allow_origin: Vec<String>,
406
407
408
409
    #[clap(long, env)]
    watermark_gamma: Option<f32>,
    #[clap(long, env)]
    watermark_delta: Option<f32>,
410

411
412
413
414
415
416
417
418
    /// Enable ngrok tunneling
    #[clap(long, env)]
    ngrok: bool,

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

419
    /// ngrok edge
420
    #[clap(long, env)]
421
    ngrok_edge: Option<String>,
422

423
424
425
426
427
    /// 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
428
429
430
431
432
    /// 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,

433
434
435
    /// Display a lot of information about your runtime environment
    #[clap(long, short, action)]
    env: bool,
436
437
438
439

    /// 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
440
441
}

442
443
444
#[derive(Debug)]
enum ShardStatus {
    Ready,
445
    Failed(usize),
446
}
447

448
449
450
451
#[allow(clippy::too_many_arguments)]
fn shard_manager(
    model_id: String,
    revision: Option<String>,
452
    quantize: Option<Quantization>,
Nicolas Patry's avatar
Nicolas Patry committed
453
    speculate: Option<usize>,
454
    dtype: Option<Dtype>,
455
    trust_remote_code: bool,
456
457
458
459
460
461
462
463
464
465
    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>,
466
    cuda_graphs: Vec<usize>,
467
    cuda_memory_fraction: f32,
Nicolas Patry's avatar
Nicolas Patry committed
468
469
    rope_scaling: Option<RopeScaling>,
    rope_factor: Option<f32>,
470
471
    max_total_tokens: usize,
    max_batch_size: Option<usize>,
472
473
    otlp_endpoint: Option<String>,
    status_sender: mpsc::Sender<ShardStatus>,
474
    shutdown: Arc<AtomicBool>,
475
476
    _shutdown_sender: mpsc::Sender<()>,
) {
477
478
479
    // Enter shard-manager tracing span
    let _span = tracing::span!(tracing::Level::INFO, "shard-manager", rank = rank).entered();

480
481
482
483
    // Get UDS path
    let uds_string = format!("{uds_path}-{rank}");
    let uds = Path::new(&uds_string);
    // Clean previous runs
484
485
486
    if uds.exists() {
        fs::remove_file(uds).unwrap();
    }
487
488

    // Process args
OlivierDehaene's avatar
OlivierDehaene committed
489
    let mut shard_args = vec![
490
491
492
493
494
495
496
497
498
        "serve".to_string(),
        model_id,
        "--uds-path".to_string(),
        uds_path,
        "--logger-level".to_string(),
        "INFO".to_string(),
        "--json-output".to_string(),
    ];

499
500
    // Activate trust remote code
    if trust_remote_code {
OlivierDehaene's avatar
OlivierDehaene committed
501
        shard_args.push("--trust-remote-code".to_string());
502
503
    }

504
505
    // Activate tensor parallelism
    if world_size > 1 {
OlivierDehaene's avatar
OlivierDehaene committed
506
        shard_args.push("--sharded".to_string());
507
508
    }

509
    if let Some(quantize) = quantize {
OlivierDehaene's avatar
OlivierDehaene committed
510
511
        shard_args.push("--quantize".to_string());
        shard_args.push(quantize.to_string())
512
    }
513

Nicolas Patry's avatar
Nicolas Patry committed
514
515
516
517
518
    if let Some(speculate) = speculate {
        shard_args.push("--speculate".to_string());
        shard_args.push(speculate.to_string())
    }

519
    if let Some(dtype) = dtype {
OlivierDehaene's avatar
OlivierDehaene committed
520
521
        shard_args.push("--dtype".to_string());
        shard_args.push(dtype.to_string())
522
523
    }

524
525
    // Model optional revision
    if let Some(revision) = revision {
OlivierDehaene's avatar
OlivierDehaene committed
526
527
        shard_args.push("--revision".to_string());
        shard_args.push(revision)
528
    }
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
529

Nicolas Patry's avatar
Nicolas Patry committed
530
531
532
533
534
535
    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)),
    };
536

537
538
    // OpenTelemetry
    if let Some(otlp_endpoint) = otlp_endpoint {
OlivierDehaene's avatar
OlivierDehaene committed
539
540
        shard_args.push("--otlp-endpoint".to_string());
        shard_args.push(otlp_endpoint);
541
542
543
    }

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

546
547
548
    // Remove LOG_LEVEL if present
    envs.retain(|(name, _)| name != "LOG_LEVEL");

549
    // Torch Distributed Env vars
OlivierDehaene's avatar
OlivierDehaene committed
550
551
552
553
    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()));
554
    envs.push(("TORCH_NCCL_AVOID_RECORD_STREAMS".into(), "1".into()));
555

556
557
558
559
560
561
    // CUDA memory fraction
    envs.push((
        "CUDA_MEMORY_FRACTION".into(),
        cuda_memory_fraction.to_string().into(),
    ));

562
    // Safetensors load fast
OlivierDehaene's avatar
OlivierDehaene committed
563
    envs.push(("SAFETENSORS_FAST_GPU".into(), "1".into()));
564

565
566
567
    // Disable progress bar
    envs.push(("HF_HUB_DISABLE_PROGRESS_BARS".into(), "1".into()));

568
569
    // 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
570
    envs.push((
571
572
573
574
575
576
        "HF_HUB_ENABLE_HF_TRANSFER".into(),
        enable_hf_transfer.into(),
    ));

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

Nicolas Patry's avatar
Nicolas Patry committed
580
581
582
583
584
585
586
587
588
    // 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()));
    }

589
590
591
592
593
594
595
596
    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()));
    }

597
598
599
    // 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
600
        envs.push(("HUGGINGFACE_HUB_CACHE".into(), huggingface_hub_cache.into()));
601
602
603
604
605
    };

    // 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
606
        envs.push((
607
608
609
610
611
            "WEIGHTS_CACHE_OVERRIDE".into(),
            weights_cache_override.into(),
        ));
    };

612
    // Enable experimental support for cuda graphs
613
614
615
616
617
618
619
620
621
622
    if !cuda_graphs.is_empty() {
        envs.push((
            "CUDA_GRAPHS".into(),
            cuda_graphs
                .into_iter()
                .map(|c| c.to_string())
                .collect::<Vec<_>>()
                .join(",")
                .into(),
        ));
623
624
    }

625
626
    // If disable_custom_kernels is true, pass it to the shard as an env var
    if disable_custom_kernels {
OlivierDehaene's avatar
OlivierDehaene committed
627
        envs.push(("DISABLE_CUSTOM_KERNELS".into(), "True".into()))
628
629
630
631
    }

    // Watermark Gamma
    if let Some(watermark_gamma) = watermark_gamma {
OlivierDehaene's avatar
OlivierDehaene committed
632
        envs.push(("WATERMARK_GAMMA".into(), watermark_gamma.to_string().into()))
633
634
635
636
    }

    // Watermark Delta
    if let Some(watermark_delta) = watermark_delta {
OlivierDehaene's avatar
OlivierDehaene committed
637
        envs.push(("WATERMARK_DELTA".into(), watermark_delta.to_string().into()))
638
639
640
    }

    // Start process
641
    tracing::info!("Starting shard");
642
    let mut p = match Command::new("text-generation-server")
OlivierDehaene's avatar
OlivierDehaene committed
643
        .args(shard_args)
644
        .env_clear()
OlivierDehaene's avatar
OlivierDehaene committed
645
        .envs(envs)
646
647
648
649
650
        .stdout(Stdio::piped())
        .stderr(Stdio::piped())
        .process_group(0)
        .spawn()
    {
651
652
        Ok(p) => p,
        Err(err) => {
653
654
655
            if err.kind() == io::ErrorKind::NotFound {
                tracing::error!("text-generation-server not found in PATH");
                tracing::error!("Please install it with `make install-server`")
656
657
            }
            {
658
                tracing::error!("{}", err);
659
            }
660

661
            status_sender.send(ShardStatus::Failed(rank)).unwrap();
662
663
664
665
666
            return;
        }
    };

    // Redirect STDOUT to the console
667
    let shard_stdout_reader = BufReader::new(p.stdout.take().unwrap());
668
    let shard_stderr_reader = BufReader::new(p.stderr.take().unwrap());
669

670
    //stdout tracing thread
671
    thread::spawn(move || {
672
        log_lines(shard_stdout_reader.lines());
673
    });
674
675
676
    // 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
677
        for line in shard_stderr_reader.lines().map_while(Result::ok) {
678
679
680
            err_sender.send(line).unwrap_or(());
        }
    });
681
682
683
684
685
686

    let mut ready = false;
    let start_time = Instant::now();
    let mut wait_time = Instant::now();
    loop {
        // Process exited
687
        if let Some(exit_status) = p.try_wait().unwrap() {
688
689
690
691
            let mut err = String::new();
            while let Ok(line) = err_receiver.recv_timeout(Duration::from_millis(10)) {
                err = err + "\n" + &line;
            }
692

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

695
            if let Some(signal) = exit_status.signal() {
696
697
698
                tracing::error!("Shard process was signaled to shutdown with signal {signal}");
            }

699
            status_sender.send(ShardStatus::Failed(rank)).unwrap();
700
701
702
703
            return;
        }

        // We received a shutdown signal
704
        if shutdown.load(Ordering::SeqCst) {
705
            terminate("shard", p, Duration::from_secs(90)).unwrap();
706
707
708
709
710
            return;
        }

        // Shard is ready
        if uds.exists() && !ready {
711
            tracing::info!("Shard ready in {:?}", start_time.elapsed());
712
713
714
            status_sender.send(ShardStatus::Ready).unwrap();
            ready = true;
        } else if !ready && wait_time.elapsed() > Duration::from_secs(10) {
715
            tracing::info!("Waiting for shard to be ready...");
716
717
718
719
720
721
            wait_time = Instant::now();
        }
        sleep(Duration::from_millis(100));
    }
}

722
fn shutdown_shards(shutdown: Arc<AtomicBool>, shutdown_receiver: &mpsc::Receiver<()>) {
723
724
725
    tracing::info!("Shutting down shards");
    // Update shutdown value to true
    // This will be picked up by the shard manager
726
    shutdown.store(true, Ordering::SeqCst);
727
728
729
730
731
732
733

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

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

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

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

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

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

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

784
785
786
787
788
789
790
791
792
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
793
    for line in lines.map_while(Result::ok) {
794
795
796
797
798
799
800
        match PythonLogMessage::try_from(&line) {
            Ok(log) => log.trace(),
            Err(_) => tracing::debug!("{line}"),
        }
    }
}

801
802
803
804
fn find_num_shards(
    sharded: Option<bool>,
    num_shard: Option<usize>,
) -> Result<usize, LauncherError> {
805
806
807
808
    // 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
809
810
811
            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");
812
            if n_devices <= 1 {
813
814
815
                return Err(LauncherError::NotEnoughCUDADevices(format!(
                    "`sharded` is true but only found {n_devices} CUDA devices"
                )));
816
            }
817
            n_devices
818
        }
819
820
821
        (Some(true), Some(num_shard)) => {
            // we can't have only one shard while sharded
            if num_shard <= 1 {
822
823
824
                return Err(LauncherError::ArgumentValidation(
                    "`sharded` is true but `num_shard` <= 1".to_string(),
                ));
825
826
            }
            num_shard
827
        }
828
829
830
831
        (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,
832
    };
833
    if num_shard < 1 {
834
835
836
        return Err(LauncherError::ArgumentValidation(
            "`num_shard` cannot be < 1".to_string(),
        ));
837
    }
838
    Ok(num_shard)
839
}
840

841
#[derive(Debug, Error)]
842
enum LauncherError {
843
    #[error("Invalid argument: {0}")]
844
    ArgumentValidation(String),
845
    #[error("not enough cuda devices: {0}")]
846
    NotEnoughCUDADevices(String),
847
    #[error("Download error")]
848
    DownloadError,
849
    #[error("Shard cannot start")]
850
    ShardCannotStart,
851
    #[error("Shard disconnected")]
852
    ShardDisconnected,
853
    #[error("Shard failed")]
854
    ShardFailed,
855
    #[error("Webserver failed")]
856
    WebserverFailed,
857
    #[error("Webserver cannot start")]
858
859
    WebserverCannotStart,
}
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
860

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

OlivierDehaene's avatar
OlivierDehaene committed
865
    let mut download_args = vec![
866
867
868
869
870
871
872
873
        "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(),
    ];
874

875
876
    // Model optional revision
    if let Some(revision) = &args.revision {
OlivierDehaene's avatar
OlivierDehaene committed
877
878
        download_args.push("--revision".to_string());
        download_args.push(revision.to_string())
879
    }
880

881
882
883
884
885
    // Trust remote code for automatic peft fusion
    if args.trust_remote_code {
        download_args.push("--trust-remote-code".to_string());
    }

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

889
890
891
    // Remove LOG_LEVEL if present
    envs.retain(|(name, _)| name != "LOG_LEVEL");

892
893
894
    // Disable progress bar
    envs.push(("HF_HUB_DISABLE_PROGRESS_BARS".into(), "1".into()));

895
    // If huggingface_hub_cache is set, pass it to the download process
896
897
    // Useful when running inside a docker container
    if let Some(ref huggingface_hub_cache) = args.huggingface_hub_cache {
OlivierDehaene's avatar
OlivierDehaene committed
898
        envs.push(("HUGGINGFACE_HUB_CACHE".into(), huggingface_hub_cache.into()));
899
    };
900

901
902
    // 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
903
    envs.push((
904
905
906
        "HF_HUB_ENABLE_HF_TRANSFER".into(),
        enable_hf_transfer.into(),
    ));
907

908
909
    // Parse Inference API token
    if let Ok(api_token) = env::var("HF_API_TOKEN") {
OlivierDehaene's avatar
OlivierDehaene committed
910
        envs.push(("HUGGING_FACE_HUB_TOKEN".into(), api_token.into()))
911
    };
912

913
914
915
    // 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
916
        envs.push((
917
918
919
920
921
            "WEIGHTS_CACHE_OVERRIDE".into(),
            weights_cache_override.into(),
        ));
    };

922
923
    // Start process
    tracing::info!("Starting download process.");
924
    let mut download_process = match Command::new("text-generation-server")
OlivierDehaene's avatar
OlivierDehaene committed
925
        .args(download_args)
926
        .env_clear()
OlivierDehaene's avatar
OlivierDehaene committed
927
        .envs(envs)
928
929
930
931
932
        .stdout(Stdio::piped())
        .stderr(Stdio::piped())
        .process_group(0)
        .spawn()
    {
933
934
        Ok(p) => p,
        Err(err) => {
935
936
937
            if err.kind() == io::ErrorKind::NotFound {
                tracing::error!("text-generation-server not found in PATH");
                tracing::error!("Please install it with `make install-server`")
938
939
            } else {
                tracing::error!("{}", err);
940
            }
941

942
943
944
            return Err(LauncherError::DownloadError);
        }
    };
945

946
    let download_stdout = BufReader::new(download_process.stdout.take().unwrap());
947

948
    thread::spawn(move || {
949
950
951
952
953
954
955
956
        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
957
        for line in download_stderr.lines().map_while(Result::ok) {
958
959
            err_sender.send(line).unwrap_or(());
        }
960
    });
961

962
    loop {
963
964
965
966
        if let Some(status) = download_process.try_wait().unwrap() {
            if status.success() {
                tracing::info!("Successfully downloaded weights.");
                break;
967
            }
968
969

            let mut err = String::new();
970
971
972
973
            while let Ok(line) = err_receiver.recv_timeout(Duration::from_millis(10)) {
                err = err + "\n" + &line;
            }

974
975
976
977
978
979
980
981
982
            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);
983
        }
984
        if !running.load(Ordering::SeqCst) {
OlivierDehaene's avatar
OlivierDehaene committed
985
            terminate("download", download_process, Duration::from_secs(10)).unwrap();
986
987
988
            return Ok(());
        }
        sleep(Duration::from_millis(100));
989
    }
990
991
    Ok(())
}
992

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

1092
1093
1094
1095
1096
1097
1098
1099
1100
1101
1102
1103
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)
}

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

drbh's avatar
drbh committed
1149
1150
1151
1152
1153
    // Grammar support
    if args.disable_grammar_support {
        router_args.push("--disable-grammar-support".to_string());
    }

1154
1155
1156
1157
1158
1159
    // 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());
    }

1160
1161
1162
1163
1164
1165
    // 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());
    }

1166
1167
1168
1169
1170
1171
    // 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());
    }

1172
1173
    // Model optional revision
    if let Some(ref revision) = args.revision {
OlivierDehaene's avatar
OlivierDehaene committed
1174
1175
        router_args.push("--revision".to_string());
        router_args.push(revision.to_string())
1176
1177
    }

1178
    if args.json_output {
OlivierDehaene's avatar
OlivierDehaene committed
1179
        router_args.push("--json-output".to_string());
1180
1181
    }

1182
    // OpenTelemetry
1183
    if let Some(otlp_endpoint) = args.otlp_endpoint {
OlivierDehaene's avatar
OlivierDehaene committed
1184
1185
        router_args.push("--otlp-endpoint".to_string());
        router_args.push(otlp_endpoint);
1186
1187
1188
1189
    }

    // CORS origins
    for origin in args.cors_allow_origin.into_iter() {
OlivierDehaene's avatar
OlivierDehaene committed
1190
1191
        router_args.push("--cors-allow-origin".to_string());
        router_args.push(origin);
1192
1193
    }

1194
1195
    // Ngrok
    if args.ngrok {
OlivierDehaene's avatar
OlivierDehaene committed
1196
1197
        router_args.push("--ngrok".to_string());
        router_args.push("--ngrok-authtoken".to_string());
1198
1199
1200
        router_args.push(args.ngrok_authtoken.unwrap());
        router_args.push("--ngrok-edge".to_string());
        router_args.push(args.ngrok_edge.unwrap());
1201
1202
    }

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

1206
1207
    // Parse Inference API token
    if let Ok(api_token) = env::var("HF_API_TOKEN") {
OlivierDehaene's avatar
OlivierDehaene committed
1208
        envs.push(("HUGGING_FACE_HUB_TOKEN".into(), api_token.into()))
1209
    };
1210

1211
1212
1213
1214
1215
1216
1217
    // 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()))
    }

1218
    let mut webserver = match Command::new("text-generation-router")
OlivierDehaene's avatar
OlivierDehaene committed
1219
1220
        .args(router_args)
        .envs(envs)
1221
1222
1223
1224
1225
        .stdout(Stdio::piped())
        .stderr(Stdio::piped())
        .process_group(0)
        .spawn()
    {
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1226
1227
        Ok(p) => p,
        Err(err) => {
1228
            tracing::error!("Failed to start webserver: {}", err);
1229
1230
1231
            if err.kind() == io::ErrorKind::NotFound {
                tracing::error!("text-generation-router not found in PATH");
                tracing::error!("Please install it with `make install-router`")
1232
1233
            } else {
                tracing::error!("{}", err);
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1234
            }
1235

1236
            shutdown_shards(shutdown, shutdown_receiver);
1237
            return Err(LauncherError::WebserverCannotStart);
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1238
1239
1240
        }
    };

1241
1242
1243
    // Redirect STDOUT and STDERR to the console
    let webserver_stdout = webserver.stdout.take().unwrap();
    let webserver_stderr = webserver.stderr.take().unwrap();
1244
1245

    thread::spawn(move || {
1246
1247
        let stdout = BufReader::new(webserver_stdout);
        let stderr = BufReader::new(webserver_stderr);
1248
        for line in stdout.lines() {
1249
            println!("{}", line.unwrap());
1250
        }
1251
1252
        for line in stderr.lines() {
            println!("{}", line.unwrap());
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1253
        }
1254
1255
1256
    });
    Ok(webserver)
}
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1257

OlivierDehaene's avatar
OlivierDehaene committed
1258
1259
1260
1261
1262
1263
1264
1265
1266
1267
1268
1269
1270
1271
1272
1273
1274
1275
1276
1277
1278
1279
1280
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)
}

1281
1282
fn main() -> Result<(), LauncherError> {
    // Pattern match configuration
1283
    let args: Args = Args::parse();
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1284

1285
1286
1287
1288
    // Filter events with LOG_LEVEL
    let env_filter =
        EnvFilter::try_from_env("LOG_LEVEL").unwrap_or_else(|_| EnvFilter::new("info"));

1289
    if args.json_output {
1290
1291
1292
1293
        tracing_subscriber::fmt()
            .with_env_filter(env_filter)
            .json()
            .init();
1294
    } else {
1295
1296
1297
1298
        tracing_subscriber::fmt()
            .with_env_filter(env_filter)
            .compact()
            .init();
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1299
1300
    }

1301
1302
1303
1304
1305
    if args.env {
        let env_runtime = env_runtime::Env::new();
        tracing::info!("{}", env_runtime);
    }

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

1308
1309
1310
1311
1312
1313
1314
1315
1316
1317
1318
1319
1320
1321
1322
1323
1324
1325
1326
1327
1328
1329
    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)?;
1330
1331
        let config: RawConfig = serde_json::from_str(&content)?;
        let config: Config = config.into();
1332
1333
1334
1335

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

1336
1337
1338
1339
1340
1341
1342
1343
        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);
1344
                }
1345
1346
1347
                Ok(max_default)
            } else {
                Ok(max_position_embeddings)
1348
            }
1349
1350
1351
1352
1353
        } else {
            Err(Box::new(LauncherError::ArgumentValidation(
                "no max defined".to_string(),
            )))
        }
1354
1355
1356
1357
1358
1359
1360
1361
1362
1363
1364
1365
1366
1367
1368
1369
1370
1371
1372
1373
1374
1375
1376
1377
1378
1379
1380
1381
1382
1383
1384
1385
1386
1387
1388
1389
1390
1391
1392
1393
1394
1395
1396
1397
1398
    };
    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
            }
        }
    };

1399
    // Validate args
1400
    if max_input_tokens >= max_total_tokens {
1401
        return Err(LauncherError::ArgumentValidation(
1402
            "`max_input_tokens must be < `max_total_tokens`".to_string(),
1403
1404
        ));
    }
1405
    if max_input_tokens as u32 > max_batch_prefill_tokens {
1406
        return Err(LauncherError::ArgumentValidation(format!(
1407
1408
            "`max_batch_prefill_tokens` must be >= `max_input_tokens`. Given: {} and {}",
            max_batch_prefill_tokens, max_input_tokens
1409
1410
        )));
    }
1411

1412
    let cuda_graphs = match (&args.cuda_graphs, &args.quantize) {
Nicolas Patry's avatar
Nicolas Patry committed
1413
        (Some(cuda_graphs), _) => cuda_graphs.iter().cloned().filter(|&c| c > 0).collect(),
1414
1415
1416
1417
1418
1419
1420
1421
1422
1423
1424
1425
1426
1427
1428
1429
1430
1431
1432
        #[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
        }
    };

1433
1434
1435
1436
1437
    if args.validation_workers == 0 {
        return Err(LauncherError::ArgumentValidation(
            "`validation_workers` must be > 0".to_string(),
        ));
    }
1438
1439
1440
1441
1442
1443
    if args.trust_remote_code {
        tracing::warn!(
            "`trust_remote_code` is set. Trusting that model `{}` do not contain malicious code.",
            args.model_id
        );
    }
1444
1445

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

1450
    if let Some(ref max_batch_total_tokens) = args.max_batch_total_tokens {
1451
        if max_batch_prefill_tokens > *max_batch_total_tokens {
1452
1453
            return Err(LauncherError::ArgumentValidation(format!(
                "`max_batch_prefill_tokens` must be <= `max_batch_total_tokens`. Given: {} and {}",
1454
                max_batch_prefill_tokens, max_batch_total_tokens
1455
1456
            )));
        }
1457
        if max_total_tokens as u32 > *max_batch_total_tokens {
1458
1459
            return Err(LauncherError::ArgumentValidation(format!(
                "`max_total_tokens` must be <= `max_batch_total_tokens`. Given: {} and {}",
1460
                max_total_tokens, max_batch_total_tokens
1461
1462
1463
1464
            )));
        }
    }

1465
1466
1467
1468
1469
1470
1471
1472
1473
1474
1475
1476
1477
1478
    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(),
            ));
        }
    }

1479
1480
1481
1482
1483
1484
1485
    // 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");
1486

1487
    // Download and convert model weights
1488
    download_convert_model(&args, running.clone())?;
1489

OlivierDehaene's avatar
OlivierDehaene committed
1490
1491
1492
1493
1494
    if !running.load(Ordering::SeqCst) {
        // Launcher was asked to stop
        return Ok(());
    }

1495
    // Shared shutdown bool
1496
    let shutdown = Arc::new(AtomicBool::new(false));
1497
1498
1499
    // Shared shutdown channel
    // When shutting down, the main thread will wait for all senders to be dropped
    let (shutdown_sender, shutdown_receiver) = mpsc::channel();
1500

1501
1502
    // Shared channel to track shard status
    let (status_sender, status_receiver) = mpsc::channel();
1503

1504
1505
1506
    spawn_shards(
        num_shard,
        &args,
1507
        cuda_graphs,
1508
        max_total_tokens,
1509
1510
1511
1512
1513
1514
1515
        shutdown.clone(),
        &shutdown_receiver,
        shutdown_sender,
        &status_receiver,
        status_sender,
        running.clone(),
    )?;
1516

1517
1518
1519
1520
1521
    // We might have received a termination signal
    if !running.load(Ordering::SeqCst) {
        shutdown_shards(shutdown, &shutdown_receiver);
        return Ok(());
    }
1522

1523
1524
1525
1526
1527
1528
1529
1530
1531
1532
1533
1534
1535
    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
    })?;
1536
1537
1538
1539
1540

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

    while running.load(Ordering::SeqCst) {
1541
        if let Ok(ShardStatus::Failed(rank)) = status_receiver.try_recv() {
OlivierDehaene's avatar
OlivierDehaene committed
1542
            tracing::error!("Shard {rank} crashed");
1543
1544
1545
1546
            exit_code = Err(LauncherError::ShardFailed);
            break;
        };

1547
        match webserver.try_wait().unwrap() {
1548
1549
1550
1551
1552
1553
1554
1555
1556
            Some(_) => {
                tracing::error!("Webserver Crashed");
                shutdown_shards(shutdown, &shutdown_receiver);
                return Err(LauncherError::WebserverFailed);
            }
            None => {
                sleep(Duration::from_millis(100));
            }
        };
1557
    }
1558
1559

    // Graceful termination
OlivierDehaene's avatar
OlivierDehaene committed
1560
    terminate("webserver", webserver, Duration::from_secs(90)).unwrap();
1561
1562
1563
    shutdown_shards(shutdown, &shutdown_receiver);

    exit_code
1564
}