main.rs 37.2 KB
Newer Older
1
use clap::{Parser, ValueEnum};
2
use serde::Deserialize;
Nicolas Patry's avatar
Nicolas Patry committed
3
use std::env;
4
use std::ffi::OsString;
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
5
6
7
8
use std::io::{BufRead, BufReader, Read};
use std::path::Path;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::mpsc::TryRecvError;
9
use std::sync::{mpsc, Arc};
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
10
11
12
13
use std::thread;
use std::thread::sleep;
use std::time::{Duration, Instant};
use std::{fs, io};
14
use subprocess::{ExitStatus, Popen, PopenConfig, PopenError, Redirection};
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
15

16
17
mod env_runtime;

18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
#[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")
            }
        }
    }
}

38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
#[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
58
59
60
61
/// App Configuration
#[derive(Parser, Debug)]
#[clap(author, version, about, long_about = None)]
struct Args {
62
63
64
65
66
    /// 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
67
    #[clap(default_value = "bigscience/bloom-560m", long, env)]
68
    model_id: String,
69
70
71

    /// 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
72
    #[clap(long, env)]
73
    revision: Option<String>,
74

75
    /// Whether to shard the model across multiple GPUs
76
77
    /// By default text-generation-inference will use all available GPUs to run
    /// the model. Setting it to `false` deactivates `num_shard`.
78
79
    #[clap(long, env)]
    sharded: Option<bool>,
80
81

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

88
    /// Whether you want the model to be quantized. This will use `bitsandbytes` for
89
90
91
    /// quantization on the fly, or `gptq`.
    #[clap(long, env, value_enum)]
    quantize: Option<Quantization>,
92

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

97
98
99
100
101
102
    /// 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,

103
104
105
    /// 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
106
107
    #[clap(default_value = "128", long, env)]
    max_concurrent_requests: usize,
108
109
110
111

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

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

    /// 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.
127
    #[clap(default_value = "1024", long, env)]
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
128
    max_input_length: usize,
129
130
131
132
133
134
135
136
137

    /// 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.
138
    #[clap(default_value = "2048", long, env)]
139
    max_total_tokens: usize,
140
141
142
143
144
145
146
147
148
149
150

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

154
155
156
157
158
159
    /// 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,

160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
    /// **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.
177
    #[clap(default_value = "16000", long, env)]
178
    max_batch_total_tokens: u32,
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196

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

200
201
202
203
    /// The IP address to listen on
    #[clap(default_value = "0.0.0.0", long, env)]
    hostname: String,

204
    /// The port to listen on.
205
    #[clap(default_value = "3000", long, short, env)]
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
206
    port: u16,
207
208
209

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

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

    /// The address the master port will listen on. (setting used by torch distributed)
218
    #[clap(default_value = "29500", long, env)]
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
219
    master_port: usize,
220
221
222

    /// The location of the huggingface hub cache.
    /// Used to override the location if you want to provide a mounted disk for instance
223
    #[clap(long, env)]
224
    huggingface_hub_cache: Option<String>,
225
226
227

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

    /// 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.
235
    #[clap(long, env)]
236
    disable_custom_kernels: bool,
237
238

    /// Outputs the logs in JSON format (useful for telemetry)
239
    #[clap(long, env)]
240
    json_output: bool,
241

242
243
    #[clap(long, env)]
    otlp_endpoint: Option<String>,
244

245
246
    #[clap(long, env)]
    cors_allow_origin: Vec<String>,
247
248
249
250
    #[clap(long, env)]
    watermark_gamma: Option<f32>,
    #[clap(long, env)]
    watermark_delta: Option<f32>,
251

252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
    /// 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>,

272
273
274
    /// 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
275
276
}

277
278
279
#[derive(Debug)]
enum ShardStatus {
    Ready,
280
    Failed((usize, Option<String>)),
281
}
282

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

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

323
324
325
326
327
    // Activate trust remote code
    if trust_remote_code {
        shard_argv.push("--trust-remote-code".to_string());
    }

328
329
330
    // Activate tensor parallelism
    if world_size > 1 {
        shard_argv.push("--sharded".to_string());
331
332
    }

333
334
335
    if let Some(quantize) = quantize {
        shard_argv.push("--quantize".to_string());
        shard_argv.push(quantize.to_string())
336
    }
337

338
339
340
341
342
    if let Some(dtype) = dtype {
        shard_argv.push("--dtype".to_string());
        shard_argv.push(dtype.to_string())
    }

343
344
345
346
347
    // 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
348

349
350
351
352
353
354
355
356
357
    // 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();

358
359
360
361
362
363
    // Use cuda allocator. It leads to less memory fragmentation
    env.push((
        "PYTORCH_CUDA_ALLOC_CONF".into(),
        "backend:cudaMallocAsync".into(),
    ));

364
365
366
367
368
369
370
371
372
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
427
428
429
430
431
432
433
434
435
    // 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}");
    let mut p = match Popen::create(
        &shard_argv,
        PopenConfig {
            stdout: Redirection::Pipe,
            stderr: Redirection::Pipe,
            // Needed for the shutdown procedure
            setpgid: true,
            // NCCL env vars
            env: Some(env),
            ..Default::default()
        },
    ) {
        Ok(p) => p,
        Err(err) => {
            if let PopenError::IoError(ref err) = err {
                if err.kind() == io::ErrorKind::NotFound {
                    tracing::error!("text-generation-server not found in PATH");
                    tracing::error!("Please install it with `make install-server`")
436
437
                }
            }
438
            status_sender
439
                .send(ShardStatus::Failed((rank, Some(err.to_string()))))
440
441
442
443
444
445
                .unwrap();
            return;
        }
    };

    // Redirect STDOUT to the console
446
447
    let shard_stdout_reader = BufReader::new(p.stdout.take().unwrap());
    let mut shard_stderr_reader = BufReader::new(p.stderr.take().unwrap());
448
449
450
451

    thread::spawn(move || {
        // Enter shard-manager tracing span
        let _span = tracing::span!(tracing::Level::INFO, "shard-manager", rank = rank).entered();
452
        for line in shard_stdout_reader.lines() {
453
454
455
456
457
458
459
460
461
462
463
464
            // 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
465
        if let Some(exit_status) = p.poll() {
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
            // 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();
482
483
484
485
486

            if let ExitStatus::Signaled(signal) = exit_status {
                tracing::error!("Shard process was signaled to shutdown with signal {signal}");
            }

487
488
489
490
491
492
493
            status_sender
                .send(ShardStatus::Failed((rank, err)))
                .unwrap();
            return;
        }

        // We received a shutdown signal
494
        if shutdown.load(Ordering::SeqCst) {
495
            p.kill().unwrap();
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
            let _ = p.wait_timeout(Duration::from_secs(90));
            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));
    }
}

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

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

fn num_cuda_devices() -> Option<usize> {
526
527
528
529
    let devices = match env::var("CUDA_VISIBLE_DEVICES") {
        Ok(devices) => devices,
        Err(_) => env::var("NVIDIA_VISIBLE_DEVICES").ok()?,
    };
530
531
    let n_devices = devices.split(',').count();
    Some(n_devices)
532
533
534
535
536
537
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
}

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

fn find_num_shards(sharded: Option<bool>, num_shard: Option<usize>) -> usize {
    // 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
581
582
583
            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");
584
585
            if n_devices <= 1 {
                panic!("`sharded` is true but only found {n_devices} CUDA devices");
586
            }
587
            n_devices
588
        }
589
590
591
592
593
594
        (Some(true), Some(num_shard)) => {
            // we can't have only one shard while sharded
            if num_shard <= 1 {
                panic!("`sharded` is true but `num_shard` <= 1");
            }
            num_shard
595
        }
596
597
598
599
        (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,
600
    };
601
602
603
    if num_shard < 1 {
        panic!("`num_shard` cannot be < 1");
    }
604
605
    num_shard
}
606

607
608
609
610
611
612
613
614
615
#[derive(Debug)]
enum LauncherError {
    DownloadError,
    ShardCannotStart,
    ShardDisconnected,
    ShardFailed,
    WebserverFailed,
    WebserverCannotStart,
}
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
616

617
fn download_convert_model(args: &Args, running: Arc<AtomicBool>) -> Result<(), LauncherError> {
618
619
620
621
622
623
624
625
626
627
    let mut download_argv = vec![
        "text-generation-server".to_string(),
        "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(),
    ];
628

629
630
631
632
633
    // Model optional revision
    if let Some(revision) = &args.revision {
        download_argv.push("--revision".to_string());
        download_argv.push(revision.to_string())
    }
634

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

638
    // If huggingface_hub_cache is set, pass it to the download process
639
640
641
642
    // 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()));
    };
643

644
645
646
647
648
649
    // 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(),
    ));
650

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

656
657
658
659
660
661
662
663
664
    // 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(),
        ));
    };

665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
    // Start process
    tracing::info!("Starting download process.");
    let mut download_process = match Popen::create(
        &download_argv,
        PopenConfig {
            stdout: Redirection::Pipe,
            stderr: Redirection::Pipe,
            // Needed for the shutdown procedure
            setpgid: true,
            env: Some(env),
            ..Default::default()
        },
    ) {
        Ok(p) => p,
        Err(err) => {
            if let PopenError::IoError(ref err) = err {
                if err.kind() == io::ErrorKind::NotFound {
                    tracing::error!("text-generation-server not found in PATH");
                    tracing::error!("Please install it with `make install-server`")
684
685
                }
            }
686
687
688
            return Err(LauncherError::DownloadError);
        }
    };
689

690
691
692
693
694
695
696
697
698
699
    // 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();
700
            }
701
702
        }
    });
703

704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
    loop {
        if let Some(status) = download_process.poll() {
            match status {
                ExitStatus::Exited(exit_code) => {
                    if exit_code == 0 {
                        tracing::info!("Successfully downloaded weights.");
                        break;
                    } else {
                        let mut err = String::new();
                        download_process
                            .stderr
                            .take()
                            .unwrap()
                            .read_to_string(&mut err)
                            .unwrap();
                        tracing::error!("Download encountered an error: {err}");
                        return Err(LauncherError::DownloadError);
721
722
                    }
                }
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
                ExitStatus::Signaled(signal) => {
                    let mut err = String::new();
                    download_process
                        .stderr
                        .take()
                        .unwrap()
                        .read_to_string(&mut err)
                        .unwrap();
                    tracing::error!(
                        "Download process was signaled to shutdown with signal {signal}: {err}"
                    );
                    return Err(LauncherError::DownloadError);
                }
                e => {
                    tracing::error!("Download process exited with an unknown status.: {e:?}");
738
739
                    return Err(LauncherError::DownloadError);
                }
740
741
            }
        }
742
743
744
745
746
747
748
749
750
751
        if !running.load(Ordering::SeqCst) {
            download_process.terminate().unwrap();
            tracing::info!("Waiting for download process to gracefully shutdown");
            download_process
                .wait_timeout(Duration::from_secs(90))
                .unwrap();
            tracing::info!("Download process terminated");
            return Ok(());
        }
        sleep(Duration::from_millis(100));
752
    }
753
754
    Ok(())
}
755

756
#[allow(clippy::too_many_arguments)]
757
758
759
fn spawn_shards(
    num_shard: usize,
    args: &Args,
760
    shutdown: Arc<AtomicBool>,
761
762
763
764
765
766
    shutdown_receiver: &mpsc::Receiver<()>,
    shutdown_sender: mpsc::Sender<()>,
    status_receiver: &mpsc::Receiver<ShardStatus>,
    status_sender: mpsc::Sender<ShardStatus>,
    running: Arc<AtomicBool>,
) -> Result<(), LauncherError> {
767
768
769
770
771
772
773
774
775
776
    if args.trust_remote_code {
        tracing::warn!(
            "`trust_remote_code` is set. Trusting that model `{}` do not contain malicious code.",
            args.model_id
        );
        if args.revision.is_none() {
            tracing::warn!("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.");
        }
    }

Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
777
778
    // Start shard processes
    for rank in 0..num_shard {
779
780
781
782
783
784
        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
785
786
787
        let status_sender = status_sender.clone();
        let shutdown = shutdown.clone();
        let shutdown_sender = shutdown_sender.clone();
788
        let otlp_endpoint = args.otlp_endpoint.clone();
789
        let quantize = args.quantize;
790
        let dtype = args.dtype;
791
        let trust_remote_code = args.trust_remote_code;
792
793
794
795
        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
796
797
        thread::spawn(move || {
            shard_manager(
798
                model_id,
799
                revision,
800
                quantize,
801
                dtype,
802
                trust_remote_code,
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
803
804
805
806
807
                uds_path,
                rank,
                num_shard,
                master_addr,
                master_port,
808
809
                huggingface_hub_cache,
                weights_cache_override,
810
                disable_custom_kernels,
811
812
                watermark_gamma,
                watermark_delta,
813
                otlp_endpoint,
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
                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));
            }
835
            Ok(ShardStatus::Failed((rank, err))) => {
836
837
838
839
                tracing::error!("Shard {rank} failed to start");
                if let Some(err) = err {
                    tracing::error!("{err}");
                }
840
                shutdown_shards(shutdown, shutdown_receiver);
841
                return Err(LauncherError::ShardCannotStart);
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
842
843
844
            }
            Err(TryRecvError::Disconnected) => {
                tracing::error!("Shard status channel disconnected");
845
                shutdown_shards(shutdown, shutdown_receiver);
846
                return Err(LauncherError::ShardDisconnected);
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
847
848
849
            }
        }
    }
850
851
    Ok(())
}
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
852

853
854
fn spawn_webserver(
    args: Args,
855
    shutdown: Arc<AtomicBool>,
856
857
    shutdown_receiver: &mpsc::Receiver<()>,
) -> Result<Popen, LauncherError> {
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
858
859
860
    // All shard started
    // Start webserver
    tracing::info!("Starting Webserver");
861
862
863
    let mut argv = vec![
        "text-generation-router".to_string(),
        "--max-concurrent-requests".to_string(),
864
        args.max_concurrent_requests.to_string(),
865
        "--max-best-of".to_string(),
866
        args.max_best_of.to_string(),
867
        "--max-stop-sequences".to_string(),
868
        args.max_stop_sequences.to_string(),
869
        "--max-input-length".to_string(),
870
        args.max_input_length.to_string(),
871
        "--max-total-tokens".to_string(),
872
        args.max_total_tokens.to_string(),
873
874
875
876
        "--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(),
877
        "--waiting-served-ratio".to_string(),
878
        args.waiting_served_ratio.to_string(),
879
        "--max-waiting-tokens".to_string(),
880
        args.max_waiting_tokens.to_string(),
881
882
        "--hostname".to_string(),
        args.hostname.to_string(),
883
        "--port".to_string(),
884
        args.port.to_string(),
885
        "--master-shard-uds-path".to_string(),
886
        format!("{}-0", args.shard_uds_path),
887
        "--tokenizer-name".to_string(),
888
        args.model_id,
889
890
    ];

891
892
893
894
    // Model optional revision
    if let Some(ref revision) = args.revision {
        argv.push("--revision".to_string());
        argv.push(revision.to_string())
895
896
    }

897
898
    if args.json_output {
        argv.push("--json-output".to_string());
899
900
    }

901
    // OpenTelemetry
902
903
904
905
906
907
908
909
910
    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);
911
912
    }

913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
    // 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);
        }
    }

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

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

945
946
    let mut webserver = match Popen::create(
        &argv,
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
947
948
949
950
951
        PopenConfig {
            stdout: Redirection::Pipe,
            stderr: Redirection::Pipe,
            // Needed for the shutdown procedure
            setpgid: true,
Nicolas Patry's avatar
Nicolas Patry committed
952
            env: Some(env),
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
953
954
955
956
957
            ..Default::default()
        },
    ) {
        Ok(p) => p,
        Err(err) => {
958
959
            tracing::error!("Failed to start webserver: {}", err);
            if let PopenError::IoError(err) = err {
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
960
                if err.kind() == io::ErrorKind::NotFound {
961
962
                    tracing::error!("text-generation-router not found in PATH");
                    tracing::error!("Please install it with `make install-router`")
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
963
                }
964
965
            } else {
                tracing::error!("{}", err);
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
966
            }
967

968
            shutdown_shards(shutdown, shutdown_receiver);
969
            return Err(LauncherError::WebserverCannotStart);
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
970
971
972
        }
    };

973
974
975
    // Redirect STDOUT and STDERR to the console
    let webserver_stdout = webserver.stdout.take().unwrap();
    let webserver_stderr = webserver.stderr.take().unwrap();
976
977

    thread::spawn(move || {
978
979
        let stdout = BufReader::new(webserver_stdout);
        let stderr = BufReader::new(webserver_stderr);
980
        for line in stdout.lines() {
981
            println!("{}", line.unwrap());
982
        }
983
984
        for line in stderr.lines() {
            println!("{}", line.unwrap());
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
985
        }
986
987
988
    });
    Ok(webserver)
}
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
989

990
991
992
fn main() -> Result<(), LauncherError> {
    // Pattern match configuration
    let args = Args::parse();
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
993

994
995
996
997
    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
998
999
    }

1000
1001
1002
1003
1004
    if args.env {
        let env_runtime = env_runtime::Env::new();
        tracing::info!("{}", env_runtime);
    }

1005
1006
1007
1008
1009
    tracing::info!("{:?}", args);

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

1012
1013
1014
1015
1016
1017
1018
    // 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");
1019

1020
    // Download and convert model weights
1021
    download_convert_model(&args, running.clone())?;
1022

1023
    // Shared shutdown bool
1024
    let shutdown = Arc::new(AtomicBool::new(false));
1025
1026
1027
    // Shared shutdown channel
    // When shutting down, the main thread will wait for all senders to be dropped
    let (shutdown_sender, shutdown_receiver) = mpsc::channel();
1028

1029
1030
    // Shared channel to track shard status
    let (status_sender, status_receiver) = mpsc::channel();
1031

1032
1033
1034
1035
1036
1037
1038
1039
1040
1041
    spawn_shards(
        num_shard,
        &args,
        shutdown.clone(),
        &shutdown_receiver,
        shutdown_sender,
        &status_receiver,
        status_sender,
        running.clone(),
    )?;
1042

1043
1044
1045
1046
1047
    // We might have received a termination signal
    if !running.load(Ordering::SeqCst) {
        shutdown_shards(shutdown, &shutdown_receiver);
        return Ok(());
    }
1048

OlivierDehaene's avatar
OlivierDehaene committed
1049
1050
1051
1052
1053
    let mut webserver =
        spawn_webserver(args, shutdown.clone(), &shutdown_receiver).map_err(|err| {
            shutdown_shards(shutdown.clone(), &shutdown_receiver);
            err
        })?;
1054
1055
1056
1057
1058
1059

    // 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
1060
            tracing::error!("Shard {rank} crashed");
1061
1062
1063
            if let Some(err) = err {
                tracing::error!("{err}");
            }
1064
1065
1066
1067
1068
1069
1070
1071
1072
1073
1074
1075
1076
1077
            exit_code = Err(LauncherError::ShardFailed);
            break;
        };

        match webserver.poll() {
            Some(_) => {
                tracing::error!("Webserver Crashed");
                shutdown_shards(shutdown, &shutdown_receiver);
                return Err(LauncherError::WebserverFailed);
            }
            None => {
                sleep(Duration::from_millis(100));
            }
        };
1078
    }
1079
1080
1081
1082
1083
1084
1085
1086
1087

    // Graceful termination
    webserver.terminate().unwrap();
    tracing::info!("Waiting for webserver to gracefully shutdown");
    webserver.wait_timeout(Duration::from_secs(90)).unwrap();
    tracing::info!("Webserver terminated");
    shutdown_shards(shutdown, &shutdown_receiver);

    exit_code
1088
}