main.rs 38.2 KB
Newer Older
1
use clap::{Parser, ValueEnum};
2
3
use nix::sys::signal::{self, Signal};
use nix::unistd::Pid;
4
use serde::Deserialize;
Nicolas Patry's avatar
Nicolas Patry committed
5
use std::env;
6
use std::ffi::OsString;
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
7
use std::io::{BufRead, BufReader, Read};
8
use std::os::unix::process::{CommandExt, ExitStatusExt};
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
9
use std::path::Path;
OlivierDehaene's avatar
OlivierDehaene committed
10
use std::process::{Child, Command, ExitStatus, Stdio};
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
11
12
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::mpsc::TryRecvError;
13
use std::sync::{mpsc, Arc};
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
14
15
16
17
18
use std::thread;
use std::thread::sleep;
use std::time::{Duration, Instant};
use std::{fs, io};

19
20
mod env_runtime;

21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
#[derive(Clone, Copy, Debug, ValueEnum)]
enum Quantization {
    Bitsandbytes,
    Gptq,
}

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 {
            Quantization::Bitsandbytes => {
                write!(f, "bitsandbytes")
            }
            Quantization::Gptq => {
                write!(f, "gptq")
            }
        }
    }
}

41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
#[derive(Clone, Copy, Debug, ValueEnum)]
enum Dtype {
    Float16,
    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")
            }
        }
    }
}

Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
61
62
63
64
/// App Configuration
#[derive(Parser, Debug)]
#[clap(author, version, about, long_about = None)]
struct Args {
65
66
67
68
69
    /// 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
70
    #[clap(default_value = "bigscience/bloom-560m", long, env)]
71
    model_id: String,
72
73
74

    /// 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
75
    #[clap(long, env)]
76
    revision: Option<String>,
77

78
79
80
81
82
    /// The number of tokenizer workers used for payload validation and truncation inside the
    /// router.
    #[clap(default_value = "2", long, env)]
    validation_workers: usize,

83
    /// Whether to shard the model across multiple GPUs
84
85
    /// By default text-generation-inference will use all available GPUs to run
    /// the model. Setting it to `false` deactivates `num_shard`.
86
87
    #[clap(long, env)]
    sharded: Option<bool>,
88
89

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

96
    /// Whether you want the model to be quantized. This will use `bitsandbytes` for
97
98
99
    /// quantization on the fly, or `gptq`.
    #[clap(long, env, value_enum)]
    quantize: Option<Quantization>,
100

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

105
106
107
108
109
110
    /// 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,

111
112
113
    /// 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
114
115
    #[clap(default_value = "128", long, env)]
    max_concurrent_requests: usize,
116
117
118
119

    /// 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
120
121
    #[clap(default_value = "2", long, env)]
    max_best_of: usize,
122
123
124
125
126
127

    /// 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.
128
129
    #[clap(default_value = "4", long, env)]
    max_stop_sequences: usize,
130
131
132
133
134

    /// 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.
135
    #[clap(default_value = "1024", long, env)]
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
136
    max_input_length: usize,
137
138
139
140
141
142
143
144
145

    /// 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.
146
    #[clap(default_value = "2048", long, env)]
147
    max_total_tokens: usize,
148
149
150
151
152
153
154
155
156
157
158

    /// 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`.
159
160
    #[clap(default_value = "1.2", long, env)]
    waiting_served_ratio: f32,
161

162
163
164
165
166
167
    /// 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.
    #[clap(default_value = "4096", long, env)]
    max_batch_prefill_tokens: u32,

168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
    /// **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.
185
    #[clap(default_value = "16000", long, env)]
186
    max_batch_total_tokens: u32,
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204

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

208
209
210
211
    /// The IP address to listen on
    #[clap(default_value = "0.0.0.0", long, env)]
    hostname: String,

212
    /// The port to listen on.
213
    #[clap(default_value = "3000", long, short, env)]
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
214
    port: u16,
215
216
217

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

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

    /// The address the master port will listen on. (setting used by torch distributed)
226
    #[clap(default_value = "29500", long, env)]
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
227
    master_port: usize,
228
229
230

    /// The location of the huggingface hub cache.
    /// Used to override the location if you want to provide a mounted disk for instance
231
    #[clap(long, env)]
232
    huggingface_hub_cache: Option<String>,
233
234
235

    /// The location of the huggingface hub cache.
    /// Used to override the location if you want to provide a mounted disk for instance
236
237
    #[clap(long, env)]
    weights_cache_override: Option<String>,
238
239
240
241
242

    /// 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.
243
    #[clap(long, env)]
244
    disable_custom_kernels: bool,
245
246

    /// Outputs the logs in JSON format (useful for telemetry)
247
    #[clap(long, env)]
248
    json_output: bool,
249

250
251
    #[clap(long, env)]
    otlp_endpoint: Option<String>,
252

253
254
    #[clap(long, env)]
    cors_allow_origin: Vec<String>,
255
256
257
258
    #[clap(long, env)]
    watermark_gamma: Option<f32>,
    #[clap(long, env)]
    watermark_delta: Option<f32>,
259

260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
    /// Enable ngrok tunneling
    #[clap(long, env)]
    ngrok: bool,

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

    /// ngrok domain name where the axum webserver will be available at
    #[clap(long, env)]
    ngrok_domain: Option<String>,

    /// ngrok basic auth username
    #[clap(long, env)]
    ngrok_username: Option<String>,

    /// ngrok basic auth password
    #[clap(long, env)]
    ngrok_password: Option<String>,

280
281
282
    /// Display a lot of information about your runtime environment
    #[clap(long, short, action)]
    env: bool,
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
283
284
}

285
286
287
#[derive(Debug)]
enum ShardStatus {
    Ready,
288
    Failed((usize, Option<String>)),
289
}
290

291
292
293
294
#[allow(clippy::too_many_arguments)]
fn shard_manager(
    model_id: String,
    revision: Option<String>,
295
    quantize: Option<Quantization>,
296
    dtype: Option<Dtype>,
297
    trust_remote_code: bool,
298
299
300
301
302
303
304
305
306
307
308
309
    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>,
    otlp_endpoint: Option<String>,
    status_sender: mpsc::Sender<ShardStatus>,
310
    shutdown: Arc<AtomicBool>,
311
312
313
314
315
316
    _shutdown_sender: mpsc::Sender<()>,
) {
    // Get UDS path
    let uds_string = format!("{uds_path}-{rank}");
    let uds = Path::new(&uds_string);
    // Clean previous runs
317
318
319
    if uds.exists() {
        fs::remove_file(uds).unwrap();
    }
320
321

    // Process args
OlivierDehaene's avatar
OlivierDehaene committed
322
    let mut shard_args = vec![
323
324
325
326
327
328
329
330
331
        "serve".to_string(),
        model_id,
        "--uds-path".to_string(),
        uds_path,
        "--logger-level".to_string(),
        "INFO".to_string(),
        "--json-output".to_string(),
    ];

332
333
    // Activate trust remote code
    if trust_remote_code {
OlivierDehaene's avatar
OlivierDehaene committed
334
        shard_args.push("--trust-remote-code".to_string());
335
336
    }

337
338
    // Activate tensor parallelism
    if world_size > 1 {
OlivierDehaene's avatar
OlivierDehaene committed
339
        shard_args.push("--sharded".to_string());
340
341
    }

342
    if let Some(quantize) = quantize {
OlivierDehaene's avatar
OlivierDehaene committed
343
344
        shard_args.push("--quantize".to_string());
        shard_args.push(quantize.to_string())
345
    }
346

347
    if let Some(dtype) = dtype {
OlivierDehaene's avatar
OlivierDehaene committed
348
349
        shard_args.push("--dtype".to_string());
        shard_args.push(dtype.to_string())
350
351
    }

352
353
    // Model optional revision
    if let Some(revision) = revision {
OlivierDehaene's avatar
OlivierDehaene committed
354
355
        shard_args.push("--revision".to_string());
        shard_args.push(revision)
356
    }
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
357

358
359
    // OpenTelemetry
    if let Some(otlp_endpoint) = otlp_endpoint {
OlivierDehaene's avatar
OlivierDehaene committed
360
361
        shard_args.push("--otlp-endpoint".to_string());
        shard_args.push(otlp_endpoint);
362
363
364
    }

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

367
    // Use cuda allocator. It leads to less memory fragmentation
OlivierDehaene's avatar
OlivierDehaene committed
368
    envs.push((
369
370
371
372
        "PYTORCH_CUDA_ALLOC_CONF".into(),
        "backend:cudaMallocAsync".into(),
    ));

373
    // Torch Distributed Env vars
OlivierDehaene's avatar
OlivierDehaene committed
374
375
376
377
378
    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()));
    envs.push(("NCCL_ASYNC_ERROR_HANDLING".into(), "1".into()));
379
380

    // Safetensors load fast
OlivierDehaene's avatar
OlivierDehaene committed
381
    envs.push(("SAFETENSORS_FAST_GPU".into(), "1".into()));
382
383
384

    // 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
385
    envs.push((
386
387
388
389
390
391
        "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
392
        envs.push(("HUGGING_FACE_HUB_TOKEN".into(), api_token.into()))
393
394
395
396
397
    };

    // 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
398
        envs.push(("HUGGINGFACE_HUB_CACHE".into(), huggingface_hub_cache.into()));
399
400
401
402
403
    };

    // 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
404
        envs.push((
405
406
407
408
409
410
411
            "WEIGHTS_CACHE_OVERRIDE".into(),
            weights_cache_override.into(),
        ));
    };

    // If disable_custom_kernels is true, pass it to the shard as an env var
    if disable_custom_kernels {
OlivierDehaene's avatar
OlivierDehaene committed
412
        envs.push(("DISABLE_CUSTOM_KERNELS".into(), "True".into()))
413
414
415
416
    }

    // Watermark Gamma
    if let Some(watermark_gamma) = watermark_gamma {
OlivierDehaene's avatar
OlivierDehaene committed
417
        envs.push(("WATERMARK_GAMMA".into(), watermark_gamma.to_string().into()))
418
419
420
421
    }

    // Watermark Delta
    if let Some(watermark_delta) = watermark_delta {
OlivierDehaene's avatar
OlivierDehaene committed
422
        envs.push(("WATERMARK_DELTA".into(), watermark_delta.to_string().into()))
423
424
425
426
    }

    // Start process
    tracing::info!("Starting shard {rank}");
427
    let mut p = match Command::new("text-generation-server")
OlivierDehaene's avatar
OlivierDehaene committed
428
429
        .args(shard_args)
        .envs(envs)
430
431
432
433
434
        .stdout(Stdio::piped())
        .stderr(Stdio::piped())
        .process_group(0)
        .spawn()
    {
435
436
        Ok(p) => p,
        Err(err) => {
437
438
439
440
441
            if err.kind() == io::ErrorKind::NotFound {
                tracing::error!("text-generation-server not found in PATH");
                tracing::error!("Please install it with `make install-server`")
            } else {
                tracing::error!("{}", err);
442
            }
443

444
            status_sender
445
                .send(ShardStatus::Failed((rank, Some(err.to_string()))))
446
447
448
449
450
451
                .unwrap();
            return;
        }
    };

    // Redirect STDOUT to the console
452
453
    let shard_stdout_reader = BufReader::new(p.stdout.take().unwrap());
    let mut shard_stderr_reader = BufReader::new(p.stderr.take().unwrap());
454
455
456
457

    thread::spawn(move || {
        // Enter shard-manager tracing span
        let _span = tracing::span!(tracing::Level::INFO, "shard-manager", rank = rank).entered();
458
        for line in shard_stdout_reader.lines() {
459
460
461
462
463
464
465
466
467
468
469
470
            // Parse loguru logs
            if let Ok(log) = serde_json::from_str::<PythonLogMessage>(&line.unwrap()) {
                log.trace();
            }
        }
    });

    let mut ready = false;
    let start_time = Instant::now();
    let mut wait_time = Instant::now();
    loop {
        // Process exited
471
        if let Some(exit_status) = p.try_wait().unwrap() {
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
            // We read stderr in another thread as it seems that `read_to_string` can block
            // indefinitely in some cases
            let (err_sender, err_receiver) = mpsc::channel();
            thread::spawn(move || {
                let mut err = String::new();
                shard_stderr_reader.read_to_string(&mut err).unwrap();
                err_sender.send(err).unwrap_or(());
            });

            let err = err_receiver
                .recv_timeout(Duration::from_millis(100))
                .map_err(|err| {
                    tracing::error!("Unable to read shard {rank} error from stderr");
                    err
                })
                .ok();
488

489
            if let Some(signal) = exit_status.signal() {
490
491
492
                tracing::error!("Shard process was signaled to shutdown with signal {signal}");
            }

493
494
495
496
497
498
499
            status_sender
                .send(ShardStatus::Failed((rank, err)))
                .unwrap();
            return;
        }

        // We received a shutdown signal
500
        if shutdown.load(Ordering::SeqCst) {
501
            p.kill().unwrap();
502
            let _ = p.wait();
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
            tracing::info!("Shard {rank} terminated");
            return;
        }

        // Shard is ready
        if uds.exists() && !ready {
            tracing::info!("Shard {rank} ready in {:?}", start_time.elapsed());
            status_sender.send(ShardStatus::Ready).unwrap();
            ready = true;
        } else if !ready && wait_time.elapsed() > Duration::from_secs(10) {
            tracing::info!("Waiting for shard {rank} to be ready...");
            wait_time = Instant::now();
        }
        sleep(Duration::from_millis(100));
    }
}

520
fn shutdown_shards(shutdown: Arc<AtomicBool>, shutdown_receiver: &mpsc::Receiver<()>) {
521
522
523
    tracing::info!("Shutting down shards");
    // Update shutdown value to true
    // This will be picked up by the shard manager
524
    shutdown.store(true, Ordering::SeqCst);
525
526
527
528
529
530
531

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

fn num_cuda_devices() -> Option<usize> {
532
533
534
535
    let devices = match env::var("CUDA_VISIBLE_DEVICES") {
        Ok(devices) => devices,
        Err(_) => env::var("NVIDIA_VISIBLE_DEVICES").ok()?,
    };
536
537
    let n_devices = devices.split(',').count();
    Some(n_devices)
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
}

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

582
583
584
585
fn find_num_shards(
    sharded: Option<bool>,
    num_shard: Option<usize>,
) -> Result<usize, LauncherError> {
586
587
588
589
    // 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
590
591
592
            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");
593
            if n_devices <= 1 {
594
595
596
                return Err(LauncherError::NotEnoughCUDADevices(format!(
                    "`sharded` is true but only found {n_devices} CUDA devices"
                )));
597
            }
598
            n_devices
599
        }
600
601
602
        (Some(true), Some(num_shard)) => {
            // we can't have only one shard while sharded
            if num_shard <= 1 {
603
604
605
                return Err(LauncherError::ArgumentValidation(
                    "`sharded` is true but `num_shard` <= 1".to_string(),
                ));
606
607
            }
            num_shard
608
        }
609
610
611
612
        (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,
613
    };
614
    if num_shard < 1 {
615
616
617
        return Err(LauncherError::ArgumentValidation(
            "`num_shard` cannot be < 1".to_string(),
        ));
618
    }
619
    Ok(num_shard)
620
}
621

622
623
#[derive(Debug)]
enum LauncherError {
624
625
    ArgumentValidation(String),
    NotEnoughCUDADevices(String),
626
627
628
629
630
631
632
    DownloadError,
    ShardCannotStart,
    ShardDisconnected,
    ShardFailed,
    WebserverFailed,
    WebserverCannotStart,
}
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
633

634
fn download_convert_model(args: &Args, running: Arc<AtomicBool>) -> Result<(), LauncherError> {
OlivierDehaene's avatar
OlivierDehaene committed
635
    let mut download_args = vec![
636
637
638
639
640
641
642
643
        "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(),
    ];
644

645
646
    // Model optional revision
    if let Some(revision) = &args.revision {
OlivierDehaene's avatar
OlivierDehaene committed
647
648
        download_args.push("--revision".to_string());
        download_args.push(revision.to_string())
649
    }
650

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

654
    // If huggingface_hub_cache is set, pass it to the download process
655
656
    // Useful when running inside a docker container
    if let Some(ref huggingface_hub_cache) = args.huggingface_hub_cache {
OlivierDehaene's avatar
OlivierDehaene committed
657
        envs.push(("HUGGINGFACE_HUB_CACHE".into(), huggingface_hub_cache.into()));
658
    };
659

660
661
    // 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
662
    envs.push((
663
664
665
        "HF_HUB_ENABLE_HF_TRANSFER".into(),
        enable_hf_transfer.into(),
    ));
666

667
668
    // Parse Inference API token
    if let Ok(api_token) = env::var("HF_API_TOKEN") {
OlivierDehaene's avatar
OlivierDehaene committed
669
        envs.push(("HUGGING_FACE_HUB_TOKEN".into(), api_token.into()))
670
    };
671

672
673
674
    // 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
675
        envs.push((
676
677
678
679
680
            "WEIGHTS_CACHE_OVERRIDE".into(),
            weights_cache_override.into(),
        ));
    };

681
682
    // Start process
    tracing::info!("Starting download process.");
683
    let mut download_process = match Command::new("text-generation-server")
OlivierDehaene's avatar
OlivierDehaene committed
684
685
        .args(download_args)
        .envs(envs)
686
687
688
689
690
        .stdout(Stdio::piped())
        .stderr(Stdio::piped())
        .process_group(0)
        .spawn()
    {
691
692
        Ok(p) => p,
        Err(err) => {
693
694
695
            if err.kind() == io::ErrorKind::NotFound {
                tracing::error!("text-generation-server not found in PATH");
                tracing::error!("Please install it with `make install-server`")
696
            }
697

698
699
700
            return Err(LauncherError::DownloadError);
        }
    };
701

702
703
704
705
706
707
708
709
710
711
    // Redirect STDOUT to the console
    let download_stdout = download_process.stdout.take().unwrap();
    thread::spawn(move || {
        // Enter download tracing span
        let stdout = BufReader::new(download_stdout);
        let _span = tracing::span!(tracing::Level::INFO, "download").entered();
        for line in stdout.lines() {
            // Parse loguru logs
            if let Ok(log) = serde_json::from_str::<PythonLogMessage>(&line.unwrap()) {
                log.trace();
712
            }
713
714
        }
    });
715

716
    loop {
717
718
719
720
        if let Some(status) = download_process.try_wait().unwrap() {
            if status.success() {
                tracing::info!("Successfully downloaded weights.");
                break;
721
            }
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738

            let mut err = String::new();
            download_process
                .stderr
                .take()
                .unwrap()
                .read_to_string(&mut err)
                .unwrap();
            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);
739
        }
740
        if !running.load(Ordering::SeqCst) {
OlivierDehaene's avatar
OlivierDehaene committed
741
            terminate("download", download_process, Duration::from_secs(10)).unwrap();
742
743
744
            return Ok(());
        }
        sleep(Duration::from_millis(100));
745
    }
746
747
    Ok(())
}
748

749
#[allow(clippy::too_many_arguments)]
750
751
752
fn spawn_shards(
    num_shard: usize,
    args: &Args,
753
    shutdown: Arc<AtomicBool>,
754
755
756
757
758
759
    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
760
761
    // Start shard processes
    for rank in 0..num_shard {
762
763
764
765
766
767
        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
768
769
770
        let status_sender = status_sender.clone();
        let shutdown = shutdown.clone();
        let shutdown_sender = shutdown_sender.clone();
771
        let otlp_endpoint = args.otlp_endpoint.clone();
772
        let quantize = args.quantize;
773
        let dtype = args.dtype;
774
        let trust_remote_code = args.trust_remote_code;
775
776
777
778
        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;
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
779
780
        thread::spawn(move || {
            shard_manager(
781
                model_id,
782
                revision,
783
                quantize,
784
                dtype,
785
                trust_remote_code,
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
786
787
788
789
790
                uds_path,
                rank,
                num_shard,
                master_addr,
                master_port,
791
792
                huggingface_hub_cache,
                weights_cache_override,
793
                disable_custom_kernels,
794
795
                watermark_gamma,
                watermark_delta,
796
                otlp_endpoint,
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
                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));
            }
818
            Ok(ShardStatus::Failed((rank, err))) => {
819
820
821
822
                tracing::error!("Shard {rank} failed to start");
                if let Some(err) = err {
                    tracing::error!("{err}");
                }
823
                shutdown_shards(shutdown, shutdown_receiver);
824
                return Err(LauncherError::ShardCannotStart);
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
825
826
827
            }
            Err(TryRecvError::Disconnected) => {
                tracing::error!("Shard status channel disconnected");
828
                shutdown_shards(shutdown, shutdown_receiver);
829
                return Err(LauncherError::ShardDisconnected);
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
830
831
832
            }
        }
    }
833
834
    Ok(())
}
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
835

836
837
fn spawn_webserver(
    args: Args,
838
    shutdown: Arc<AtomicBool>,
839
    shutdown_receiver: &mpsc::Receiver<()>,
840
) -> Result<Child, LauncherError> {
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
841
842
843
    // All shard started
    // Start webserver
    tracing::info!("Starting Webserver");
OlivierDehaene's avatar
OlivierDehaene committed
844
    let mut router_args = vec![
845
        "--max-concurrent-requests".to_string(),
846
        args.max_concurrent_requests.to_string(),
847
        "--max-best-of".to_string(),
848
        args.max_best_of.to_string(),
849
        "--max-stop-sequences".to_string(),
850
        args.max_stop_sequences.to_string(),
851
        "--max-input-length".to_string(),
852
        args.max_input_length.to_string(),
853
        "--max-total-tokens".to_string(),
854
        args.max_total_tokens.to_string(),
855
856
857
858
        "--max-batch-prefill-tokens".to_string(),
        args.max_batch_prefill_tokens.to_string(),
        "--max-batch-total-tokens".to_string(),
        args.max_batch_total_tokens.to_string(),
859
        "--waiting-served-ratio".to_string(),
860
        args.waiting_served_ratio.to_string(),
861
        "--max-waiting-tokens".to_string(),
862
        args.max_waiting_tokens.to_string(),
863
864
        "--validation-workers".to_string(),
        args.validation_workers.to_string(),
865
866
        "--hostname".to_string(),
        args.hostname.to_string(),
867
        "--port".to_string(),
868
        args.port.to_string(),
869
        "--master-shard-uds-path".to_string(),
870
        format!("{}-0", args.shard_uds_path),
871
        "--tokenizer-name".to_string(),
872
        args.model_id,
873
874
    ];

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

881
    if args.json_output {
OlivierDehaene's avatar
OlivierDehaene committed
882
        router_args.push("--json-output".to_string());
883
884
    }

885
    // OpenTelemetry
886
    if let Some(otlp_endpoint) = args.otlp_endpoint {
OlivierDehaene's avatar
OlivierDehaene committed
887
888
        router_args.push("--otlp-endpoint".to_string());
        router_args.push(otlp_endpoint);
889
890
891
892
    }

    // CORS origins
    for origin in args.cors_allow_origin.into_iter() {
OlivierDehaene's avatar
OlivierDehaene committed
893
894
        router_args.push("--cors-allow-origin".to_string());
        router_args.push(origin);
895
896
    }

897
898
899
900
901
902
903
    // Ngrok
    if args.ngrok {
        let authtoken = args.ngrok_authtoken.ok_or_else(|| {
            tracing::error!("`ngrok-authtoken` must be set when using ngrok tunneling");
            LauncherError::WebserverCannotStart
        })?;

OlivierDehaene's avatar
OlivierDehaene committed
904
905
906
        router_args.push("--ngrok".to_string());
        router_args.push("--ngrok-authtoken".to_string());
        router_args.push(authtoken);
907
908

        if let Some(domain) = args.ngrok_domain {
OlivierDehaene's avatar
OlivierDehaene committed
909
910
            router_args.push("--ngrok-domain".to_string());
            router_args.push(domain);
911
912
913
        }

        if let (Some(username), Some(password)) = (args.ngrok_username, args.ngrok_password) {
OlivierDehaene's avatar
OlivierDehaene committed
914
915
916
917
            router_args.push("--ngrok-username".to_string());
            router_args.push(username);
            router_args.push("--ngrok-password".to_string());
            router_args.push(password);
918
919
920
        }
    }

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

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

929
    let mut webserver = match Command::new("text-generation-router")
OlivierDehaene's avatar
OlivierDehaene committed
930
931
        .args(router_args)
        .envs(envs)
932
933
934
935
936
        .stdout(Stdio::piped())
        .stderr(Stdio::piped())
        .process_group(0)
        .spawn()
    {
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
937
938
        Ok(p) => p,
        Err(err) => {
939
            tracing::error!("Failed to start webserver: {}", err);
940
941
942
            if err.kind() == io::ErrorKind::NotFound {
                tracing::error!("text-generation-router not found in PATH");
                tracing::error!("Please install it with `make install-router`")
943
944
            } else {
                tracing::error!("{}", err);
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
945
            }
946

947
            shutdown_shards(shutdown, shutdown_receiver);
948
            return Err(LauncherError::WebserverCannotStart);
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
949
950
951
        }
    };

952
953
954
    // Redirect STDOUT and STDERR to the console
    let webserver_stdout = webserver.stdout.take().unwrap();
    let webserver_stderr = webserver.stderr.take().unwrap();
955
956

    thread::spawn(move || {
957
958
        let stdout = BufReader::new(webserver_stdout);
        let stderr = BufReader::new(webserver_stderr);
959
        for line in stdout.lines() {
960
            println!("{}", line.unwrap());
961
        }
962
963
        for line in stderr.lines() {
            println!("{}", line.unwrap());
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
964
        }
965
966
967
    });
    Ok(webserver)
}
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
968

OlivierDehaene's avatar
OlivierDehaene committed
969
970
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
991
992
993
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)
}

994
995
996
fn main() -> Result<(), LauncherError> {
    // Pattern match configuration
    let args = Args::parse();
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
997

998
999
1000
1001
    if args.json_output {
        tracing_subscriber::fmt().json().init();
    } else {
        tracing_subscriber::fmt().compact().init();
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1002
1003
    }

1004
1005
1006
1007
1008
    if args.env {
        let env_runtime = env_runtime::Env::new();
        tracing::info!("{}", env_runtime);
    }

1009
1010
    tracing::info!("{:?}", args);

1011
1012
1013
1014
1015
1016
1017
1018
1019
1020
1021
1022
1023
1024
1025
1026
1027
1028
1029
1030
1031
1032
1033
1034
1035
1036
1037
1038
1039
    // Validate args
    if args.max_input_length >= args.max_total_tokens {
        return Err(LauncherError::ArgumentValidation(
            "`max_input_length` must be < `max_total_tokens`".to_string(),
        ));
    }
    if args.max_input_length as u32 > args.max_batch_prefill_tokens {
        return Err(LauncherError::ArgumentValidation(format!(
            "`max_batch_prefill_tokens` must be >= `max_input_length`. Given: {} and {}",
            args.max_batch_prefill_tokens, args.max_input_length
        )));
    }
    if args.max_batch_prefill_tokens > args.max_batch_total_tokens {
        return Err(LauncherError::ArgumentValidation(format!(
            "`max_batch_prefill_tokens` must be <= `max_batch_total_tokens`. Given: {} and {}",
            args.max_batch_prefill_tokens, args.max_batch_total_tokens
        )));
    }
    if args.max_total_tokens as u32 > args.max_batch_total_tokens {
        return Err(LauncherError::ArgumentValidation(format!(
            "`max_total_tokens` must be <= `max_batch_total_tokens`. Given: {} and {}",
            args.max_total_tokens, args.max_batch_total_tokens
        )));
    }
    if args.validation_workers == 0 {
        return Err(LauncherError::ArgumentValidation(
            "`validation_workers` must be > 0".to_string(),
        ));
    }
1040
1041
1042
1043
1044
1045
    if args.trust_remote_code {
        tracing::warn!(
            "`trust_remote_code` is set. Trusting that model `{}` do not contain malicious code.",
            args.model_id
        );
    }
1046
1047

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

1052
1053
1054
1055
1056
1057
1058
    // 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");
1059

1060
    // Download and convert model weights
1061
    download_convert_model(&args, running.clone())?;
1062

OlivierDehaene's avatar
OlivierDehaene committed
1063
1064
1065
1066
1067
    if !running.load(Ordering::SeqCst) {
        // Launcher was asked to stop
        return Ok(());
    }

1068
    // Shared shutdown bool
1069
    let shutdown = Arc::new(AtomicBool::new(false));
1070
1071
1072
    // Shared shutdown channel
    // When shutting down, the main thread will wait for all senders to be dropped
    let (shutdown_sender, shutdown_receiver) = mpsc::channel();
1073

1074
1075
    // Shared channel to track shard status
    let (status_sender, status_receiver) = mpsc::channel();
1076

1077
1078
1079
1080
1081
1082
1083
1084
1085
1086
    spawn_shards(
        num_shard,
        &args,
        shutdown.clone(),
        &shutdown_receiver,
        shutdown_sender,
        &status_receiver,
        status_sender,
        running.clone(),
    )?;
1087

1088
1089
1090
1091
1092
    // We might have received a termination signal
    if !running.load(Ordering::SeqCst) {
        shutdown_shards(shutdown, &shutdown_receiver);
        return Ok(());
    }
1093

OlivierDehaene's avatar
OlivierDehaene committed
1094
1095
1096
1097
1098
    let mut webserver =
        spawn_webserver(args, shutdown.clone(), &shutdown_receiver).map_err(|err| {
            shutdown_shards(shutdown.clone(), &shutdown_receiver);
            err
        })?;
1099
1100
1101
1102
1103
1104

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

    while running.load(Ordering::SeqCst) {
        if let Ok(ShardStatus::Failed((rank, err))) = status_receiver.try_recv() {
OlivierDehaene's avatar
OlivierDehaene committed
1105
            tracing::error!("Shard {rank} crashed");
1106
1107
1108
            if let Some(err) = err {
                tracing::error!("{err}");
            }
1109
1110
1111
1112
            exit_code = Err(LauncherError::ShardFailed);
            break;
        };

1113
        match webserver.try_wait().unwrap() {
1114
1115
1116
1117
1118
1119
1120
1121
1122
            Some(_) => {
                tracing::error!("Webserver Crashed");
                shutdown_shards(shutdown, &shutdown_receiver);
                return Err(LauncherError::WebserverFailed);
            }
            None => {
                sleep(Duration::from_millis(100));
            }
        };
1123
    }
1124
1125

    // Graceful termination
OlivierDehaene's avatar
OlivierDehaene committed
1126
    terminate("webserver", webserver, Duration::from_secs(90)).unwrap();
1127
1128
1129
    shutdown_shards(shutdown, &shutdown_receiver);

    exit_code
1130
}