main.rs 37.5 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;
10
use std::process::{Child, Command, 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
322
323
324
325
326
327
328
329
330
331

    // Process args
    let mut shard_argv = vec![
        "serve".to_string(),
        model_id,
        "--uds-path".to_string(),
        uds_path,
        "--logger-level".to_string(),
        "INFO".to_string(),
        "--json-output".to_string(),
    ];

332
333
334
335
336
    // Activate trust remote code
    if trust_remote_code {
        shard_argv.push("--trust-remote-code".to_string());
    }

337
338
339
    // Activate tensor parallelism
    if world_size > 1 {
        shard_argv.push("--sharded".to_string());
340
341
    }

342
343
344
    if let Some(quantize) = quantize {
        shard_argv.push("--quantize".to_string());
        shard_argv.push(quantize.to_string())
345
    }
346

347
348
349
350
351
    if let Some(dtype) = dtype {
        shard_argv.push("--dtype".to_string());
        shard_argv.push(dtype.to_string())
    }

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

358
359
360
361
362
363
364
365
366
    // OpenTelemetry
    if let Some(otlp_endpoint) = otlp_endpoint {
        shard_argv.push("--otlp-endpoint".to_string());
        shard_argv.push(otlp_endpoint);
    }

    // Copy current process env
    let mut env: Vec<(OsString, OsString)> = env::vars_os().collect();

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

373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
    // Torch Distributed Env vars
    env.push(("RANK".into(), rank.to_string().into()));
    env.push(("WORLD_SIZE".into(), world_size.to_string().into()));
    env.push(("MASTER_ADDR".into(), master_addr.into()));
    env.push(("MASTER_PORT".into(), master_port.to_string().into()));
    env.push(("NCCL_ASYNC_ERROR_HANDLING".into(), "1".into()));

    // Safetensors load fast
    env.push(("SAFETENSORS_FAST_GPU".into(), "1".into()));

    // Enable hf transfer for insane download speeds
    let enable_hf_transfer = env::var("HF_HUB_ENABLE_HF_TRANSFER").unwrap_or("1".to_string());
    env.push((
        "HF_HUB_ENABLE_HF_TRANSFER".into(),
        enable_hf_transfer.into(),
    ));

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

    // 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 {
        env.push(("HUGGINGFACE_HUB_CACHE".into(), huggingface_hub_cache.into()));
    };

    // 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 {
        env.push((
            "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 {
        env.push(("DISABLE_CUSTOM_KERNELS".into(), "True".into()))
    }

    // Watermark Gamma
    if let Some(watermark_gamma) = watermark_gamma {
        env.push(("WATERMARK_GAMMA".into(), watermark_gamma.to_string().into()))
    }

    // Watermark Delta
    if let Some(watermark_delta) = watermark_delta {
        env.push(("WATERMARK_DELTA".into(), watermark_delta.to_string().into()))
    }

    // Start process
    tracing::info!("Starting shard {rank}");
427
428
429
430
431
432
433
434
    let mut p = match Command::new("text-generation-server")
        .args(shard_argv)
        .envs(env)
        .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> {
635
636
637
638
639
640
641
642
643
    let mut download_argv = vec![
        "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
647
648
649
    // Model optional revision
    if let Some(revision) = &args.revision {
        download_argv.push("--revision".to_string());
        download_argv.push(revision.to_string())
    }
650

651
652
    // Copy current process env
    let mut env: Vec<(OsString, OsString)> = env::vars_os().collect();
653

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

660
661
662
663
664
665
    // Enable hf transfer for insane download speeds
    let enable_hf_transfer = env::var("HF_HUB_ENABLE_HF_TRANSFER").unwrap_or("1".to_string());
    env.push((
        "HF_HUB_ENABLE_HF_TRANSFER".into(),
        enable_hf_transfer.into(),
    ));
666

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

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

681
682
    // Start process
    tracing::info!("Starting download process.");
683
684
685
686
687
688
689
690
    let mut download_process = match Command::new("text-generation-server")
        .args(download_argv)
        .envs(env)
        .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) {
741
            signal::kill(Pid::from_raw(download_process.id() as i32), Signal::SIGTERM).unwrap();
742
            tracing::info!("Waiting for download process to gracefully shutdown");
743
            download_process.wait().unwrap();
744
745
746
747
            tracing::info!("Download process terminated");
            return Ok(());
        }
        sleep(Duration::from_millis(100));
748
    }
749
750
    Ok(())
}
751

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

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

878
879
880
881
    // Model optional revision
    if let Some(ref revision) = args.revision {
        argv.push("--revision".to_string());
        argv.push(revision.to_string())
882
883
    }

884
885
    if args.json_output {
        argv.push("--json-output".to_string());
886
887
    }

888
    // OpenTelemetry
889
890
891
892
893
894
895
896
897
    if let Some(otlp_endpoint) = args.otlp_endpoint {
        argv.push("--otlp-endpoint".to_string());
        argv.push(otlp_endpoint);
    }

    // CORS origins
    for origin in args.cors_allow_origin.into_iter() {
        argv.push("--cors-allow-origin".to_string());
        argv.push(origin);
898
899
    }

900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
    // 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
        })?;

        argv.push("--ngrok".to_string());
        argv.push("--ngrok-authtoken".to_string());
        argv.push(authtoken);

        if let Some(domain) = args.ngrok_domain {
            argv.push("--ngrok-domain".to_string());
            argv.push(domain);
        }

        if let (Some(username), Some(password)) = (args.ngrok_username, args.ngrok_password) {
            argv.push("--ngrok-username".to_string());
            argv.push(username);
            argv.push("--ngrok-password".to_string());
            argv.push(password);
        }
    }

924
925
926
    // Copy current process env
    let mut env: Vec<(OsString, OsString)> = env::vars_os().collect();

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

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

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

955
956
957
    // Redirect STDOUT and STDERR to the console
    let webserver_stdout = webserver.stdout.take().unwrap();
    let webserver_stderr = webserver.stderr.take().unwrap();
958
959

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

972
973
974
fn main() -> Result<(), LauncherError> {
    // Pattern match configuration
    let args = Args::parse();
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
975

976
977
978
979
    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
980
981
    }

982
983
984
985
986
    if args.env {
        let env_runtime = env_runtime::Env::new();
        tracing::info!("{}", env_runtime);
    }

987
988
    tracing::info!("{:?}", args);

989
990
991
992
993
994
995
996
997
998
999
1000
1001
1002
1003
1004
1005
1006
1007
1008
1009
1010
1011
1012
1013
1014
1015
1016
1017
    // 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(),
        ));
    }
1018
1019
1020
1021
1022
1023
    if args.trust_remote_code {
        tracing::warn!(
            "`trust_remote_code` is set. Trusting that model `{}` do not contain malicious code.",
            args.model_id
        );
    }
1024
1025

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

1030
1031
1032
1033
1034
1035
1036
    // 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");
1037

1038
    // Download and convert model weights
1039
    download_convert_model(&args, running.clone())?;
1040

1041
    // Shared shutdown bool
1042
    let shutdown = Arc::new(AtomicBool::new(false));
1043
1044
1045
    // Shared shutdown channel
    // When shutting down, the main thread will wait for all senders to be dropped
    let (shutdown_sender, shutdown_receiver) = mpsc::channel();
1046

1047
1048
    // Shared channel to track shard status
    let (status_sender, status_receiver) = mpsc::channel();
1049

1050
1051
1052
1053
1054
1055
1056
1057
1058
1059
    spawn_shards(
        num_shard,
        &args,
        shutdown.clone(),
        &shutdown_receiver,
        shutdown_sender,
        &status_receiver,
        status_sender,
        running.clone(),
    )?;
1060

1061
1062
1063
1064
1065
    // We might have received a termination signal
    if !running.load(Ordering::SeqCst) {
        shutdown_shards(shutdown, &shutdown_receiver);
        return Ok(());
    }
1066

OlivierDehaene's avatar
OlivierDehaene committed
1067
1068
1069
1070
1071
    let mut webserver =
        spawn_webserver(args, shutdown.clone(), &shutdown_receiver).map_err(|err| {
            shutdown_shards(shutdown.clone(), &shutdown_receiver);
            err
        })?;
1072
1073
1074
1075
1076
1077

    // 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
1078
            tracing::error!("Shard {rank} crashed");
1079
1080
1081
            if let Some(err) = err {
                tracing::error!("{err}");
            }
1082
1083
1084
1085
            exit_code = Err(LauncherError::ShardFailed);
            break;
        };

1086
        match webserver.try_wait().unwrap() {
1087
1088
1089
1090
1091
1092
1093
1094
1095
            Some(_) => {
                tracing::error!("Webserver Crashed");
                shutdown_shards(shutdown, &shutdown_receiver);
                return Err(LauncherError::WebserverFailed);
            }
            None => {
                sleep(Duration::from_millis(100));
            }
        };
1096
    }
1097
1098

    // Graceful termination
1099
    signal::kill(Pid::from_raw(webserver.id() as i32), Signal::SIGTERM).unwrap();
1100
    tracing::info!("Waiting for webserver to gracefully shutdown");
1101
    webserver.wait().unwrap();
1102
1103
1104
1105
    tracing::info!("Webserver terminated");
    shutdown_shards(shutdown, &shutdown_receiver);

    exit_code
1106
}