main.rs 40.7 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
720
721
722
723
    // Trust remote code for automatic peft fusion
    if args.trust_remote_code {
        download_args.push("--trust-remote-code".to_string());
    }

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

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

733
734
    // 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
735
    envs.push((
736
737
738
        "HF_HUB_ENABLE_HF_TRANSFER".into(),
        enable_hf_transfer.into(),
    ));
739

740
741
    // Parse Inference API token
    if let Ok(api_token) = env::var("HF_API_TOKEN") {
OlivierDehaene's avatar
OlivierDehaene committed
742
        envs.push(("HUGGING_FACE_HUB_TOKEN".into(), api_token.into()))
743
    };
744

745
746
747
    // 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
748
        envs.push((
749
750
751
752
753
            "WEIGHTS_CACHE_OVERRIDE".into(),
            weights_cache_override.into(),
        ));
    };

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

773
774
775
            return Err(LauncherError::DownloadError);
        }
    };
776

777
778
    // Redirect STDOUT to the console
    let download_stdout = download_process.stdout.take().unwrap();
779
780
    let stdout = BufReader::new(download_stdout);

781
    thread::spawn(move || {
782
        log_lines(stdout.lines());
783
    });
784

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

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

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

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

945
946
947
948
949
950
    // 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());
    }

951
952
    // Model optional revision
    if let Some(ref revision) = args.revision {
OlivierDehaene's avatar
OlivierDehaene committed
953
954
        router_args.push("--revision".to_string());
        router_args.push(revision.to_string())
955
956
    }

957
    if args.json_output {
OlivierDehaene's avatar
OlivierDehaene committed
958
        router_args.push("--json-output".to_string());
959
960
    }

961
    // OpenTelemetry
962
    if let Some(otlp_endpoint) = args.otlp_endpoint {
OlivierDehaene's avatar
OlivierDehaene committed
963
964
        router_args.push("--otlp-endpoint".to_string());
        router_args.push(otlp_endpoint);
965
966
967
968
    }

    // CORS origins
    for origin in args.cors_allow_origin.into_iter() {
OlivierDehaene's avatar
OlivierDehaene committed
969
970
        router_args.push("--cors-allow-origin".to_string());
        router_args.push(origin);
971
972
    }

973
974
    // Ngrok
    if args.ngrok {
OlivierDehaene's avatar
OlivierDehaene committed
975
976
        router_args.push("--ngrok".to_string());
        router_args.push("--ngrok-authtoken".to_string());
977
978
979
        router_args.push(args.ngrok_authtoken.unwrap());
        router_args.push("--ngrok-edge".to_string());
        router_args.push(args.ngrok_edge.unwrap());
980
981
    }

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

985
986
    // Parse Inference API token
    if let Ok(api_token) = env::var("HF_API_TOKEN") {
OlivierDehaene's avatar
OlivierDehaene committed
987
        envs.push(("HUGGING_FACE_HUB_TOKEN".into(), api_token.into()))
988
    };
989

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

1008
            shutdown_shards(shutdown, shutdown_receiver);
1009
            return Err(LauncherError::WebserverCannotStart);
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1010
1011
1012
        }
    };

1013
1014
1015
    // Redirect STDOUT and STDERR to the console
    let webserver_stdout = webserver.stdout.take().unwrap();
    let webserver_stderr = webserver.stderr.take().unwrap();
1016
1017

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

OlivierDehaene's avatar
OlivierDehaene committed
1030
1031
1032
1033
1034
1035
1036
1037
1038
1039
1040
1041
1042
1043
1044
1045
1046
1047
1048
1049
1050
1051
1052
1053
1054
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)
}

1055
1056
fn main() -> Result<(), LauncherError> {
    // Pattern match configuration
1057
    let args: Args = Args::parse();
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1058

1059
1060
1061
1062
    // Filter events with LOG_LEVEL
    let env_filter =
        EnvFilter::try_from_env("LOG_LEVEL").unwrap_or_else(|_| EnvFilter::new("info"));

1063
    if args.json_output {
1064
1065
1066
1067
        tracing_subscriber::fmt()
            .with_env_filter(env_filter)
            .json()
            .init();
1068
    } else {
1069
1070
1071
1072
        tracing_subscriber::fmt()
            .with_env_filter(env_filter)
            .compact()
            .init();
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1073
1074
    }

1075
1076
1077
1078
1079
    if args.env {
        let env_runtime = env_runtime::Env::new();
        tracing::info!("{}", env_runtime);
    }

1080
1081
    tracing::info!("{:?}", args);

1082
1083
1084
1085
1086
1087
1088
1089
1090
1091
1092
1093
    // 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
        )));
    }
1094

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

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

1112
1113
1114
1115
1116
1117
1118
1119
1120
1121
1122
1123
1124
1125
1126
    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
            )));
        }
    }

1127
1128
1129
1130
1131
1132
1133
1134
1135
1136
1137
1138
1139
1140
    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(),
            ));
        }
    }

1141
1142
1143
1144
1145
1146
1147
    // 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");
1148

1149
    // Download and convert model weights
1150
    download_convert_model(&args, running.clone())?;
1151

OlivierDehaene's avatar
OlivierDehaene committed
1152
1153
1154
1155
1156
    if !running.load(Ordering::SeqCst) {
        // Launcher was asked to stop
        return Ok(());
    }

1157
    // Shared shutdown bool
1158
    let shutdown = Arc::new(AtomicBool::new(false));
1159
1160
1161
    // Shared shutdown channel
    // When shutting down, the main thread will wait for all senders to be dropped
    let (shutdown_sender, shutdown_receiver) = mpsc::channel();
1162

1163
1164
    // Shared channel to track shard status
    let (status_sender, status_receiver) = mpsc::channel();
1165

1166
1167
1168
1169
1170
1171
1172
1173
1174
1175
    spawn_shards(
        num_shard,
        &args,
        shutdown.clone(),
        &shutdown_receiver,
        shutdown_sender,
        &status_receiver,
        status_sender,
        running.clone(),
    )?;
1176

1177
1178
1179
1180
1181
    // We might have received a termination signal
    if !running.load(Ordering::SeqCst) {
        shutdown_shards(shutdown, &shutdown_receiver);
        return Ok(());
    }
1182

OlivierDehaene's avatar
OlivierDehaene committed
1183
1184
1185
1186
1187
    let mut webserver =
        spawn_webserver(args, shutdown.clone(), &shutdown_receiver).map_err(|err| {
            shutdown_shards(shutdown.clone(), &shutdown_receiver);
            err
        })?;
1188
1189
1190
1191
1192

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

    while running.load(Ordering::SeqCst) {
1193
        if let Ok(ShardStatus::Failed(rank)) = status_receiver.try_recv() {
OlivierDehaene's avatar
OlivierDehaene committed
1194
            tracing::error!("Shard {rank} crashed");
1195
1196
1197
1198
            exit_code = Err(LauncherError::ShardFailed);
            break;
        };

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

    // Graceful termination
OlivierDehaene's avatar
OlivierDehaene committed
1212
    terminate("webserver", webserver, Duration::from_secs(90)).unwrap();
1213
1214
1215
    shutdown_shards(shutdown, &shutdown_receiver);

    exit_code
1216
}