"doc/en/xpu.md" did not exist on "333351c7c85760df23f5009b60c8b5b198f00d4c"
main.rs 40.6 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;
7
use std::io::{BufRead, BufReader, Lines, 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
use std::thread;
use std::thread::sleep;
use std::time::{Duration, Instant};
use std::{fs, io};
18
use tracing_subscriber::EnvFilter;
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
19

20
21
mod env_runtime;

22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
#[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")
            }
        }
    }
}

42
43
44
#[derive(Clone, Copy, Debug, ValueEnum)]
enum Dtype {
    Float16,
45
    #[clap(name = "bfloat16")]
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
    BFloat16,
}

impl std::fmt::Display for Dtype {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        // To keep in track with `server`.
        match self {
            Dtype::Float16 => {
                write!(f, "float16")
            }
            Dtype::BFloat16 => {
                write!(f, "bfloat16")
            }
        }
    }
}

Nicolas Patry's avatar
Nicolas Patry committed
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
#[derive(Clone, Copy, Debug, ValueEnum)]
enum RopeScaling {
    Linear,
    Dynamic,
}

impl std::fmt::Display for RopeScaling {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        // To keep in track with `server`.
        match self {
            RopeScaling::Linear => {
                write!(f, "linear")
            }
            RopeScaling::Dynamic => {
                write!(f, "dynamic")
            }
        }
    }
}

Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
83
84
85
86
/// App Configuration
#[derive(Parser, Debug)]
#[clap(author, version, about, long_about = None)]
struct Args {
87
88
89
90
91
    /// 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
92
    #[clap(default_value = "bigscience/bloom-560m", long, env)]
93
    model_id: String,
94
95
96

    /// 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
97
    #[clap(long, env)]
98
    revision: Option<String>,
99

100
101
102
103
104
    /// The number of tokenizer workers used for payload validation and truncation inside the
    /// router.
    #[clap(default_value = "2", long, env)]
    validation_workers: usize,

105
    /// Whether to shard the model across multiple GPUs
106
107
    /// By default text-generation-inference will use all available GPUs to run
    /// the model. Setting it to `false` deactivates `num_shard`.
108
109
    #[clap(long, env)]
    sharded: Option<bool>,
110
111

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

118
    /// Whether you want the model to be quantized. This will use `bitsandbytes` for
119
120
121
    /// quantization on the fly, or `gptq`.
    #[clap(long, env, value_enum)]
    quantize: Option<Quantization>,
122

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

127
128
129
130
131
132
    /// 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,

133
134
135
    /// 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
136
137
    #[clap(default_value = "128", long, env)]
    max_concurrent_requests: usize,
138
139
140
141

    /// 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
142
143
    #[clap(default_value = "2", long, env)]
    max_best_of: usize,
144
145
146
147
148
149

    /// 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.
150
151
    #[clap(default_value = "4", long, env)]
    max_stop_sequences: usize,
152
153
154
155
156

    /// 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.
157
    #[clap(default_value = "1024", long, env)]
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
158
    max_input_length: usize,
159
160
161
162
163
164
165
166
167

    /// 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.
168
    #[clap(default_value = "2048", long, env)]
169
    max_total_tokens: usize,
170
171
172
173
174
175
176
177
178
179
180

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

184
185
186
187
188
189
    /// 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,

190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
    /// **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.
207
208
    #[clap(long, env)]
    max_batch_total_tokens: Option<u32>,
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226

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

230
231
232
233
    /// The IP address to listen on
    #[clap(default_value = "0.0.0.0", long, env)]
    hostname: String,

234
    /// The port to listen on.
235
    #[clap(default_value = "3000", long, short, env)]
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
236
    port: u16,
237
238
239

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

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

    /// The address the master port will listen on. (setting used by torch distributed)
248
    #[clap(default_value = "29500", long, env)]
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
249
    master_port: usize,
250
251
252

    /// The location of the huggingface hub cache.
    /// Used to override the location if you want to provide a mounted disk for instance
253
    #[clap(long, env)]
254
    huggingface_hub_cache: Option<String>,
255
256
257

    /// The location of the huggingface hub cache.
    /// Used to override the location if you want to provide a mounted disk for instance
258
259
    #[clap(long, env)]
    weights_cache_override: Option<String>,
260
261
262
263
264

    /// 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.
265
    #[clap(long, env)]
266
    disable_custom_kernels: bool,
267

268
269
270
271
272
    /// Limit the CUDA available memory.
    /// The allowed value equals the total visible memory multiplied by cuda-memory-fraction.
    #[clap(default_value = "1.0", long, env)]
    cuda_memory_fraction: f32,

Nicolas Patry's avatar
Nicolas Patry committed
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
    /// Rope scaling will only be used for RoPE models
    /// and allow rescaling the position rotary to accomodate for
    /// larger prompts.
    ///
    /// Goes together with `rope_factor`.
    ///
    /// `--rope-factor 2.0` gives linear scaling with a factor of 2.0
    /// `--rope-scaling dynamic` gives dynamic scaling with a factor of 1.0
    /// `--rope-scaling linear` gives linear scaling with a factor of 1.0 (Nothing will be changed
    /// basically)
    ///
    /// `--rope-scaling linear --rope-factor` fully describes the scaling you want
    #[clap(long, env)]
    rope_scaling: Option<RopeScaling>,

    /// Rope scaling will only be used for RoPE models
    /// See `rope_scaling`
    #[clap(long, env)]
    rope_factor: Option<f32>,

293
    /// Outputs the logs in JSON format (useful for telemetry)
294
    #[clap(long, env)]
295
    json_output: bool,
296

297
298
    #[clap(long, env)]
    otlp_endpoint: Option<String>,
299

300
301
    #[clap(long, env)]
    cors_allow_origin: Vec<String>,
302
303
304
305
    #[clap(long, env)]
    watermark_gamma: Option<f32>,
    #[clap(long, env)]
    watermark_delta: Option<f32>,
306

307
308
309
310
311
312
313
314
    /// Enable ngrok tunneling
    #[clap(long, env)]
    ngrok: bool,

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

315
    /// ngrok edge
316
    #[clap(long, env)]
317
    ngrok_edge: Option<String>,
318

319
320
321
    /// 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
322
323
}

324
325
326
#[derive(Debug)]
enum ShardStatus {
    Ready,
327
    Failed(usize),
328
}
329

330
331
332
333
#[allow(clippy::too_many_arguments)]
fn shard_manager(
    model_id: String,
    revision: Option<String>,
334
    quantize: Option<Quantization>,
335
    dtype: Option<Dtype>,
336
    trust_remote_code: bool,
337
338
339
340
341
342
343
344
345
346
    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>,
347
    cuda_memory_fraction: f32,
Nicolas Patry's avatar
Nicolas Patry committed
348
349
    rope_scaling: Option<RopeScaling>,
    rope_factor: Option<f32>,
350
351
    otlp_endpoint: Option<String>,
    status_sender: mpsc::Sender<ShardStatus>,
352
    shutdown: Arc<AtomicBool>,
353
354
    _shutdown_sender: mpsc::Sender<()>,
) {
355
356
357
    // Enter shard-manager tracing span
    let _span = tracing::span!(tracing::Level::INFO, "shard-manager", rank = rank).entered();

358
359
360
361
    // Get UDS path
    let uds_string = format!("{uds_path}-{rank}");
    let uds = Path::new(&uds_string);
    // Clean previous runs
362
363
364
    if uds.exists() {
        fs::remove_file(uds).unwrap();
    }
365
366

    // Process args
OlivierDehaene's avatar
OlivierDehaene committed
367
    let mut shard_args = vec![
368
369
370
371
372
373
374
375
376
        "serve".to_string(),
        model_id,
        "--uds-path".to_string(),
        uds_path,
        "--logger-level".to_string(),
        "INFO".to_string(),
        "--json-output".to_string(),
    ];

377
378
    // Activate trust remote code
    if trust_remote_code {
OlivierDehaene's avatar
OlivierDehaene committed
379
        shard_args.push("--trust-remote-code".to_string());
380
381
    }

382
383
    // Activate tensor parallelism
    if world_size > 1 {
OlivierDehaene's avatar
OlivierDehaene committed
384
        shard_args.push("--sharded".to_string());
385
386
    }

387
    if let Some(quantize) = quantize {
OlivierDehaene's avatar
OlivierDehaene committed
388
389
        shard_args.push("--quantize".to_string());
        shard_args.push(quantize.to_string())
390
    }
391

392
    if let Some(dtype) = dtype {
OlivierDehaene's avatar
OlivierDehaene committed
393
394
        shard_args.push("--dtype".to_string());
        shard_args.push(dtype.to_string())
395
396
    }

397
398
    // Model optional revision
    if let Some(revision) = revision {
OlivierDehaene's avatar
OlivierDehaene committed
399
400
        shard_args.push("--revision".to_string());
        shard_args.push(revision)
401
    }
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
402

Nicolas Patry's avatar
Nicolas Patry committed
403
404
405
406
407
408
    let rope = match (rope_scaling, rope_factor) {
        (None, None) => None,
        (Some(scaling), None) => Some((scaling, 1.0)),
        (Some(scaling), Some(factor)) => Some((scaling, factor)),
        (None, Some(factor)) => Some((RopeScaling::Linear, factor)),
    };
409
410
    // OpenTelemetry
    if let Some(otlp_endpoint) = otlp_endpoint {
OlivierDehaene's avatar
OlivierDehaene committed
411
412
        shard_args.push("--otlp-endpoint".to_string());
        shard_args.push(otlp_endpoint);
413
414
415
    }

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

    // Torch Distributed Env vars
OlivierDehaene's avatar
OlivierDehaene committed
419
420
421
422
423
    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()));
424

425
426
427
428
429
430
    // CUDA memory fraction
    envs.push((
        "CUDA_MEMORY_FRACTION".into(),
        cuda_memory_fraction.to_string().into(),
    ));

431
    // Safetensors load fast
OlivierDehaene's avatar
OlivierDehaene committed
432
    envs.push(("SAFETENSORS_FAST_GPU".into(), "1".into()));
433
434
435

    // 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
436
    envs.push((
437
438
439
440
441
442
        "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
443
        envs.push(("HUGGING_FACE_HUB_TOKEN".into(), api_token.into()))
444
445
    };

Nicolas Patry's avatar
Nicolas Patry committed
446
447
448
449
450
451
452
453
454
    // Detect rope scaling
    // Sending as env instead of CLI args to not bloat everything
    // those only can be used by RoPE models, so passing information around
    // for all models will complexify code unnecessarily
    if let Some((scaling, factor)) = rope {
        envs.push(("ROPE_SCALING".into(), scaling.to_string().into()));
        envs.push(("ROPE_FACTOR".into(), factor.to_string().into()));
    }

455
456
457
    // 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
458
        envs.push(("HUGGINGFACE_HUB_CACHE".into(), huggingface_hub_cache.into()));
459
460
461
462
463
    };

    // 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
464
        envs.push((
465
466
467
468
469
470
471
            "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
472
        envs.push(("DISABLE_CUSTOM_KERNELS".into(), "True".into()))
473
474
475
476
    }

    // Watermark Gamma
    if let Some(watermark_gamma) = watermark_gamma {
OlivierDehaene's avatar
OlivierDehaene committed
477
        envs.push(("WATERMARK_GAMMA".into(), watermark_gamma.to_string().into()))
478
479
480
481
    }

    // Watermark Delta
    if let Some(watermark_delta) = watermark_delta {
OlivierDehaene's avatar
OlivierDehaene committed
482
        envs.push(("WATERMARK_DELTA".into(), watermark_delta.to_string().into()))
483
484
485
    }

    // Start process
486
    tracing::info!("Starting shard");
487
    let mut p = match Command::new("text-generation-server")
OlivierDehaene's avatar
OlivierDehaene committed
488
489
        .args(shard_args)
        .envs(envs)
490
491
492
493
494
        .stdout(Stdio::piped())
        .stderr(Stdio::piped())
        .process_group(0)
        .spawn()
    {
495
496
        Ok(p) => p,
        Err(err) => {
497
498
499
            if err.kind() == io::ErrorKind::NotFound {
                tracing::error!("text-generation-server not found in PATH");
                tracing::error!("Please install it with `make install-server`")
500
501
            }
            {
502
                tracing::error!("{}", err);
503
            }
504

505
            status_sender.send(ShardStatus::Failed(rank)).unwrap();
506
507
508
509
510
            return;
        }
    };

    // Redirect STDOUT to the console
511
    let shard_stdout_reader = BufReader::new(p.stdout.take().unwrap());
512
    let shard_stderr_reader = BufReader::new(p.stderr.take().unwrap());
513

514
    //stdout tracing thread
515
    thread::spawn(move || {
516
        log_lines(shard_stdout_reader.lines());
517
518
519
520
521
522
523
    });

    let mut ready = false;
    let start_time = Instant::now();
    let mut wait_time = Instant::now();
    loop {
        // Process exited
524
        if let Some(exit_status) = p.try_wait().unwrap() {
525
            // We read stderr in another thread as it seems that lines() can block in some cases
526
527
            let (err_sender, err_receiver) = mpsc::channel();
            thread::spawn(move || {
528
529
530
                for line in shard_stderr_reader.lines().flatten() {
                    err_sender.send(line).unwrap_or(());
                }
531
            });
532
533
534
535
            let mut err = String::new();
            while let Ok(line) = err_receiver.recv_timeout(Duration::from_millis(10)) {
                err = err + "\n" + &line;
            }
536

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

539
            if let Some(signal) = exit_status.signal() {
540
541
542
                tracing::error!("Shard process was signaled to shutdown with signal {signal}");
            }

543
            status_sender.send(ShardStatus::Failed(rank)).unwrap();
544
545
546
547
            return;
        }

        // We received a shutdown signal
548
        if shutdown.load(Ordering::SeqCst) {
549
            p.kill().unwrap();
550
            let _ = p.wait();
551
            tracing::info!("Shard terminated");
552
553
554
555
556
            return;
        }

        // Shard is ready
        if uds.exists() && !ready {
557
            tracing::info!("Shard ready in {:?}", start_time.elapsed());
558
559
560
            status_sender.send(ShardStatus::Ready).unwrap();
            ready = true;
        } else if !ready && wait_time.elapsed() > Duration::from_secs(10) {
561
            tracing::info!("Waiting for shard to be ready...");
562
563
564
565
566
567
            wait_time = Instant::now();
        }
        sleep(Duration::from_millis(100));
    }
}

568
fn shutdown_shards(shutdown: Arc<AtomicBool>, shutdown_receiver: &mpsc::Receiver<()>) {
569
570
571
    tracing::info!("Shutting down shards");
    // Update shutdown value to true
    // This will be picked up by the shard manager
572
    shutdown.store(true, Ordering::SeqCst);
573
574
575
576
577
578
579

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

fn num_cuda_devices() -> Option<usize> {
580
581
582
583
    let devices = match env::var("CUDA_VISIBLE_DEVICES") {
        Ok(devices) => devices,
        Err(_) => env::var("NVIDIA_VISIBLE_DEVICES").ok()?,
    };
584
585
    let n_devices = devices.split(',').count();
    Some(n_devices)
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
}

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

630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
impl TryFrom<&String> for PythonLogMessage {
    type Error = serde_json::Error;

    fn try_from(value: &String) -> Result<Self, Self::Error> {
        serde_json::from_str::<Self>(value)
    }
}

fn log_lines<S: Sized + BufRead>(lines: Lines<S>) {
    for line in lines.flatten() {
        match PythonLogMessage::try_from(&line) {
            Ok(log) => log.trace(),
            Err(_) => tracing::debug!("{line}"),
        }
    }
}

647
648
649
650
fn find_num_shards(
    sharded: Option<bool>,
    num_shard: Option<usize>,
) -> Result<usize, LauncherError> {
651
652
653
654
    // 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
655
656
657
            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");
658
            if n_devices <= 1 {
659
660
661
                return Err(LauncherError::NotEnoughCUDADevices(format!(
                    "`sharded` is true but only found {n_devices} CUDA devices"
                )));
662
            }
663
            n_devices
664
        }
665
666
667
        (Some(true), Some(num_shard)) => {
            // we can't have only one shard while sharded
            if num_shard <= 1 {
668
669
670
                return Err(LauncherError::ArgumentValidation(
                    "`sharded` is true but `num_shard` <= 1".to_string(),
                ));
671
672
            }
            num_shard
673
        }
674
675
676
677
        (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,
678
    };
679
    if num_shard < 1 {
680
681
682
        return Err(LauncherError::ArgumentValidation(
            "`num_shard` cannot be < 1".to_string(),
        ));
683
    }
684
    Ok(num_shard)
685
}
686

687
688
#[derive(Debug)]
enum LauncherError {
689
690
    ArgumentValidation(String),
    NotEnoughCUDADevices(String),
691
692
693
694
695
696
697
    DownloadError,
    ShardCannotStart,
    ShardDisconnected,
    ShardFailed,
    WebserverFailed,
    WebserverCannotStart,
}
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
698

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

OlivierDehaene's avatar
OlivierDehaene committed
703
    let mut download_args = vec![
704
705
706
707
708
709
710
711
        "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(),
    ];
712

713
714
    // Model optional revision
    if let Some(revision) = &args.revision {
OlivierDehaene's avatar
OlivierDehaene committed
715
716
        download_args.push("--revision".to_string());
        download_args.push(revision.to_string())
717
    }
718

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

722
    // If huggingface_hub_cache is set, pass it to the download process
723
724
    // Useful when running inside a docker container
    if let Some(ref huggingface_hub_cache) = args.huggingface_hub_cache {
OlivierDehaene's avatar
OlivierDehaene committed
725
        envs.push(("HUGGINGFACE_HUB_CACHE".into(), huggingface_hub_cache.into()));
726
    };
727

728
729
    // 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
730
    envs.push((
731
732
733
        "HF_HUB_ENABLE_HF_TRANSFER".into(),
        enable_hf_transfer.into(),
    ));
734

735
736
    // Parse Inference API token
    if let Ok(api_token) = env::var("HF_API_TOKEN") {
OlivierDehaene's avatar
OlivierDehaene committed
737
        envs.push(("HUGGING_FACE_HUB_TOKEN".into(), api_token.into()))
738
    };
739

740
741
742
    // 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
743
        envs.push((
744
745
746
747
748
            "WEIGHTS_CACHE_OVERRIDE".into(),
            weights_cache_override.into(),
        ));
    };

749
750
    // Start process
    tracing::info!("Starting download process.");
751
    let mut download_process = match Command::new("text-generation-server")
OlivierDehaene's avatar
OlivierDehaene committed
752
753
        .args(download_args)
        .envs(envs)
754
755
756
757
758
        .stdout(Stdio::piped())
        .stderr(Stdio::piped())
        .process_group(0)
        .spawn()
    {
759
760
        Ok(p) => p,
        Err(err) => {
761
762
763
            if err.kind() == io::ErrorKind::NotFound {
                tracing::error!("text-generation-server not found in PATH");
                tracing::error!("Please install it with `make install-server`")
764
765
            } else {
                tracing::error!("{}", err);
766
            }
767

768
769
770
            return Err(LauncherError::DownloadError);
        }
    };
771

772
773
    // Redirect STDOUT to the console
    let download_stdout = download_process.stdout.take().unwrap();
774
775
    let stdout = BufReader::new(download_stdout);

776
    thread::spawn(move || {
777
        log_lines(stdout.lines());
778
    });
779

780
    loop {
781
782
783
784
        if let Some(status) = download_process.try_wait().unwrap() {
            if status.success() {
                tracing::info!("Successfully downloaded weights.");
                break;
785
            }
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802

            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);
803
        }
804
        if !running.load(Ordering::SeqCst) {
OlivierDehaene's avatar
OlivierDehaene committed
805
            terminate("download", download_process, Duration::from_secs(10)).unwrap();
806
807
808
            return Ok(());
        }
        sleep(Duration::from_millis(100));
809
    }
810
811
    Ok(())
}
812

813
#[allow(clippy::too_many_arguments)]
814
815
816
fn spawn_shards(
    num_shard: usize,
    args: &Args,
817
    shutdown: Arc<AtomicBool>,
818
819
820
821
822
823
    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
824
825
    // Start shard processes
    for rank in 0..num_shard {
826
827
828
829
830
831
        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
832
833
834
        let status_sender = status_sender.clone();
        let shutdown = shutdown.clone();
        let shutdown_sender = shutdown_sender.clone();
835
        let otlp_endpoint = args.otlp_endpoint.clone();
836
        let quantize = args.quantize;
837
        let dtype = args.dtype;
838
        let trust_remote_code = args.trust_remote_code;
839
840
841
842
        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;
843
        let cuda_memory_fraction = args.cuda_memory_fraction;
Nicolas Patry's avatar
Nicolas Patry committed
844
845
        let rope_scaling = args.rope_scaling;
        let rope_factor = args.rope_factor;
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
846
847
        thread::spawn(move || {
            shard_manager(
848
                model_id,
849
                revision,
850
                quantize,
851
                dtype,
852
                trust_remote_code,
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
853
854
855
856
857
                uds_path,
                rank,
                num_shard,
                master_addr,
                master_port,
858
859
                huggingface_hub_cache,
                weights_cache_override,
860
                disable_custom_kernels,
861
862
                watermark_gamma,
                watermark_delta,
863
                cuda_memory_fraction,
Nicolas Patry's avatar
Nicolas Patry committed
864
865
                rope_scaling,
                rope_factor,
866
                otlp_endpoint,
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
                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));
            }
888
            Ok(ShardStatus::Failed(rank)) => {
889
                tracing::error!("Shard {rank} failed to start");
890
                shutdown_shards(shutdown, shutdown_receiver);
891
                return Err(LauncherError::ShardCannotStart);
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
892
893
894
            }
            Err(TryRecvError::Disconnected) => {
                tracing::error!("Shard status channel disconnected");
895
                shutdown_shards(shutdown, shutdown_receiver);
896
                return Err(LauncherError::ShardDisconnected);
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
897
898
899
            }
        }
    }
900
901
    Ok(())
}
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
902

903
904
fn spawn_webserver(
    args: Args,
905
    shutdown: Arc<AtomicBool>,
906
    shutdown_receiver: &mpsc::Receiver<()>,
907
) -> Result<Child, LauncherError> {
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
908
909
910
    // All shard started
    // Start webserver
    tracing::info!("Starting Webserver");
OlivierDehaene's avatar
OlivierDehaene committed
911
    let mut router_args = vec![
912
        "--max-concurrent-requests".to_string(),
913
        args.max_concurrent_requests.to_string(),
914
        "--max-best-of".to_string(),
915
        args.max_best_of.to_string(),
916
        "--max-stop-sequences".to_string(),
917
        args.max_stop_sequences.to_string(),
918
        "--max-input-length".to_string(),
919
        args.max_input_length.to_string(),
920
        "--max-total-tokens".to_string(),
921
        args.max_total_tokens.to_string(),
922
923
        "--max-batch-prefill-tokens".to_string(),
        args.max_batch_prefill_tokens.to_string(),
924
        "--waiting-served-ratio".to_string(),
925
        args.waiting_served_ratio.to_string(),
926
        "--max-waiting-tokens".to_string(),
927
        args.max_waiting_tokens.to_string(),
928
929
        "--validation-workers".to_string(),
        args.validation_workers.to_string(),
930
931
        "--hostname".to_string(),
        args.hostname.to_string(),
932
        "--port".to_string(),
933
        args.port.to_string(),
934
        "--master-shard-uds-path".to_string(),
935
        format!("{}-0", args.shard_uds_path),
936
        "--tokenizer-name".to_string(),
937
        args.model_id,
938
939
    ];

940
941
942
943
944
945
    // Model optional max batch total tokens
    if let Some(max_batch_total_tokens) = args.max_batch_total_tokens {
        router_args.push("--max-batch-total-tokens".to_string());
        router_args.push(max_batch_total_tokens.to_string());
    }

946
947
    // Model optional revision
    if let Some(ref revision) = args.revision {
OlivierDehaene's avatar
OlivierDehaene committed
948
949
        router_args.push("--revision".to_string());
        router_args.push(revision.to_string())
950
951
    }

952
    if args.json_output {
OlivierDehaene's avatar
OlivierDehaene committed
953
        router_args.push("--json-output".to_string());
954
955
    }

956
    // OpenTelemetry
957
    if let Some(otlp_endpoint) = args.otlp_endpoint {
OlivierDehaene's avatar
OlivierDehaene committed
958
959
        router_args.push("--otlp-endpoint".to_string());
        router_args.push(otlp_endpoint);
960
961
962
963
    }

    // CORS origins
    for origin in args.cors_allow_origin.into_iter() {
OlivierDehaene's avatar
OlivierDehaene committed
964
965
        router_args.push("--cors-allow-origin".to_string());
        router_args.push(origin);
966
967
    }

968
969
    // Ngrok
    if args.ngrok {
OlivierDehaene's avatar
OlivierDehaene committed
970
971
        router_args.push("--ngrok".to_string());
        router_args.push("--ngrok-authtoken".to_string());
972
973
974
        router_args.push(args.ngrok_authtoken.unwrap());
        router_args.push("--ngrok-edge".to_string());
        router_args.push(args.ngrok_edge.unwrap());
975
976
    }

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

980
981
    // Parse Inference API token
    if let Ok(api_token) = env::var("HF_API_TOKEN") {
OlivierDehaene's avatar
OlivierDehaene committed
982
        envs.push(("HUGGING_FACE_HUB_TOKEN".into(), api_token.into()))
983
    };
984

985
    let mut webserver = match Command::new("text-generation-router")
OlivierDehaene's avatar
OlivierDehaene committed
986
987
        .args(router_args)
        .envs(envs)
988
989
990
991
992
        .stdout(Stdio::piped())
        .stderr(Stdio::piped())
        .process_group(0)
        .spawn()
    {
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
993
994
        Ok(p) => p,
        Err(err) => {
995
            tracing::error!("Failed to start webserver: {}", err);
996
997
998
            if err.kind() == io::ErrorKind::NotFound {
                tracing::error!("text-generation-router not found in PATH");
                tracing::error!("Please install it with `make install-router`")
999
1000
            } else {
                tracing::error!("{}", err);
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1001
            }
1002

1003
            shutdown_shards(shutdown, shutdown_receiver);
1004
            return Err(LauncherError::WebserverCannotStart);
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1005
1006
1007
        }
    };

1008
1009
1010
    // Redirect STDOUT and STDERR to the console
    let webserver_stdout = webserver.stdout.take().unwrap();
    let webserver_stderr = webserver.stderr.take().unwrap();
1011
1012

    thread::spawn(move || {
1013
1014
        let stdout = BufReader::new(webserver_stdout);
        let stderr = BufReader::new(webserver_stderr);
1015
        for line in stdout.lines() {
1016
            println!("{}", line.unwrap());
1017
        }
1018
1019
        for line in stderr.lines() {
            println!("{}", line.unwrap());
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1020
        }
1021
1022
1023
    });
    Ok(webserver)
}
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1024

OlivierDehaene's avatar
OlivierDehaene committed
1025
1026
1027
1028
1029
1030
1031
1032
1033
1034
1035
1036
1037
1038
1039
1040
1041
1042
1043
1044
1045
1046
1047
1048
1049
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)
}

1050
1051
fn main() -> Result<(), LauncherError> {
    // Pattern match configuration
1052
    let args: Args = Args::parse();
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1053

1054
1055
1056
1057
    // Filter events with LOG_LEVEL
    let env_filter =
        EnvFilter::try_from_env("LOG_LEVEL").unwrap_or_else(|_| EnvFilter::new("info"));

1058
    if args.json_output {
1059
1060
1061
1062
        tracing_subscriber::fmt()
            .with_env_filter(env_filter)
            .json()
            .init();
1063
    } else {
1064
1065
1066
1067
        tracing_subscriber::fmt()
            .with_env_filter(env_filter)
            .compact()
            .init();
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1068
1069
    }

1070
1071
1072
1073
1074
    if args.env {
        let env_runtime = env_runtime::Env::new();
        tracing::info!("{}", env_runtime);
    }

1075
1076
    tracing::info!("{:?}", args);

1077
1078
1079
1080
1081
1082
1083
1084
1085
1086
1087
1088
    // 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
        )));
    }
1089

1090
1091
1092
1093
1094
    if args.validation_workers == 0 {
        return Err(LauncherError::ArgumentValidation(
            "`validation_workers` must be > 0".to_string(),
        ));
    }
1095
1096
1097
1098
1099
1100
    if args.trust_remote_code {
        tracing::warn!(
            "`trust_remote_code` is set. Trusting that model `{}` do not contain malicious code.",
            args.model_id
        );
    }
1101
1102

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

1107
1108
1109
1110
1111
1112
1113
1114
1115
1116
1117
1118
1119
1120
1121
    if let Some(ref max_batch_total_tokens) = args.max_batch_total_tokens {
        if args.max_batch_prefill_tokens > *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, max_batch_total_tokens
            )));
        }
        if args.max_total_tokens as u32 > *max_batch_total_tokens {
            return Err(LauncherError::ArgumentValidation(format!(
                "`max_total_tokens` must be <= `max_batch_total_tokens`. Given: {} and {}",
                args.max_total_tokens, max_batch_total_tokens
            )));
        }
    }

1122
1123
1124
1125
1126
1127
1128
1129
1130
1131
1132
1133
1134
1135
    if args.ngrok {
        if args.ngrok_authtoken.is_none() {
            return Err(LauncherError::ArgumentValidation(
                "`ngrok-authtoken` must be set when using ngrok tunneling".to_string(),
            ));
        }

        if args.ngrok_edge.is_none() {
            return Err(LauncherError::ArgumentValidation(
                "`ngrok-edge` must be set when using ngrok tunneling".to_string(),
            ));
        }
    }

1136
1137
1138
1139
1140
1141
1142
    // 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");
1143

1144
    // Download and convert model weights
1145
    download_convert_model(&args, running.clone())?;
1146

OlivierDehaene's avatar
OlivierDehaene committed
1147
1148
1149
1150
1151
    if !running.load(Ordering::SeqCst) {
        // Launcher was asked to stop
        return Ok(());
    }

1152
    // Shared shutdown bool
1153
    let shutdown = Arc::new(AtomicBool::new(false));
1154
1155
1156
    // Shared shutdown channel
    // When shutting down, the main thread will wait for all senders to be dropped
    let (shutdown_sender, shutdown_receiver) = mpsc::channel();
1157

1158
1159
    // Shared channel to track shard status
    let (status_sender, status_receiver) = mpsc::channel();
1160

1161
1162
1163
1164
1165
1166
1167
1168
1169
1170
    spawn_shards(
        num_shard,
        &args,
        shutdown.clone(),
        &shutdown_receiver,
        shutdown_sender,
        &status_receiver,
        status_sender,
        running.clone(),
    )?;
1171

1172
1173
1174
1175
1176
    // We might have received a termination signal
    if !running.load(Ordering::SeqCst) {
        shutdown_shards(shutdown, &shutdown_receiver);
        return Ok(());
    }
1177

OlivierDehaene's avatar
OlivierDehaene committed
1178
1179
1180
1181
1182
    let mut webserver =
        spawn_webserver(args, shutdown.clone(), &shutdown_receiver).map_err(|err| {
            shutdown_shards(shutdown.clone(), &shutdown_receiver);
            err
        })?;
1183
1184
1185
1186
1187

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

    while running.load(Ordering::SeqCst) {
1188
        if let Ok(ShardStatus::Failed(rank)) = status_receiver.try_recv() {
OlivierDehaene's avatar
OlivierDehaene committed
1189
            tracing::error!("Shard {rank} crashed");
1190
1191
1192
1193
            exit_code = Err(LauncherError::ShardFailed);
            break;
        };

1194
        match webserver.try_wait().unwrap() {
1195
1196
1197
1198
1199
1200
1201
1202
1203
            Some(_) => {
                tracing::error!("Webserver Crashed");
                shutdown_shards(shutdown, &shutdown_receiver);
                return Err(LauncherError::WebserverFailed);
            }
            None => {
                sleep(Duration::from_millis(100));
            }
        };
1204
    }
1205
1206

    // Graceful termination
OlivierDehaene's avatar
OlivierDehaene committed
1207
    terminate("webserver", webserver, Duration::from_secs(90)).unwrap();
1208
1209
1210
    shutdown_shards(shutdown, &shutdown_receiver);

    exit_code
1211
}