main.rs 42.9 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
#[derive(Clone, Copy, Debug, ValueEnum)]
enum Quantization {
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
    /// 4 bit quantization. Requires a specific GTPQ quantized model:
    ///   https://hf.co/models?search=awq.
    /// Should replace GPTQ models whereever possible because of the better latency
    Awq,
    /// 8 bit quantization, doesn't require specific model.
    /// Should be a drop-in replacement to bitsandbytes with much better performance.
    /// Kernels are from https://github.com/NetEase-FuXi/EETQ.git
    Eetq,
    /// 4 bit quantization. Requires a specific GTPQ quantized model: https://hf.co/models?search=gptq.
    /// text-generation-inference will use exllama (faster) kernels whereever possible, and use
    /// triton kernel (wider support) when it's not.
    /// AWQ has faster kernels.
    Gptq,
    /// Bitsandbytes 8bit. Can be applied on any model, will cut the memory requirement in half,
    /// but it is known that the model will be much slower to run than the native f16.
    #[deprecated(
        since = "1.1.0",
        note = "Use `eetq` instead, which provides better latencies overall and is drop-in in most cases"
    )]
43
    Bitsandbytes,
44
45
    /// Bitsandbytes 4bit. Can be applied on any model, will cut the memory requirement by 4x,
    /// but it is known that the model will be much slower to run than the native f16.
Nicolas Patry's avatar
Nicolas Patry committed
46
    BitsandbytesNF4,
47
48
    /// Bitsandbytes 4bit. nf4 should be preferred in most cases but maybe this one has better
    /// perplexity performance for you model
Nicolas Patry's avatar
Nicolas Patry committed
49
    BitsandbytesFP4,
50
51
52
53
54
55
56
57
58
}

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")
            }
Nicolas Patry's avatar
Nicolas Patry committed
59
60
61
62
63
64
            Quantization::BitsandbytesNF4 => {
                write!(f, "bitsandbytes-nf4")
            }
            Quantization::BitsandbytesFP4 => {
                write!(f, "bitsandbytes-fp4")
            }
65
66
67
            Quantization::Gptq => {
                write!(f, "gptq")
            }
68
69
70
            Quantization::Awq => {
                write!(f, "awq")
            }
71
72
73
            Quantization::Eetq => {
                write!(f, "eetq")
            }
74
75
76
77
        }
    }
}

78
79
80
#[derive(Clone, Copy, Debug, ValueEnum)]
enum Dtype {
    Float16,
81
    #[clap(name = "bfloat16")]
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
    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
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
#[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
119
120
121
122
/// App Configuration
#[derive(Parser, Debug)]
#[clap(author, version, about, long_about = None)]
struct Args {
123
124
125
126
127
    /// 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
128
    #[clap(default_value = "bigscience/bloom-560m", long, env)]
129
    model_id: String,
130
131
132

    /// 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
133
    #[clap(long, env)]
134
    revision: Option<String>,
135

136
137
138
139
140
    /// The number of tokenizer workers used for payload validation and truncation inside the
    /// router.
    #[clap(default_value = "2", long, env)]
    validation_workers: usize,

141
    /// Whether to shard the model across multiple GPUs
142
143
    /// By default text-generation-inference will use all available GPUs to run
    /// the model. Setting it to `false` deactivates `num_shard`.
144
145
    #[clap(long, env)]
    sharded: Option<bool>,
146
147

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

154
    /// Whether you want the model to be quantized.
155
156
    #[clap(long, env, value_enum)]
    quantize: Option<Quantization>,
157

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

162
163
164
165
166
167
    /// 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,

168
169
170
    /// 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
171
172
    #[clap(default_value = "128", long, env)]
    max_concurrent_requests: usize,
173
174
175
176

    /// 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
177
178
    #[clap(default_value = "2", long, env)]
    max_best_of: usize,
179
180
181
182
183
184

    /// 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.
185
186
    #[clap(default_value = "4", long, env)]
    max_stop_sequences: usize,
187

Nicolas Patry's avatar
Nicolas Patry committed
188
189
190
191
192
193
194
195
    /// This is the maximum allowed value for clients to set `top_n_tokens`.
    /// `top_n_tokens is used to return information about the the `n` most likely
    /// tokens at each generation step, instead of just the sampled token. This
    /// information can be used for downstream tasks like for classification or
    /// ranking.
    #[clap(default_value = "5", long, env)]
    max_top_n_tokens: u32,

196
197
198
199
    /// 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.
200
    #[clap(default_value = "1024", long, env)]
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
201
    max_input_length: usize,
202
203
204
205
206
207
208
209
210

    /// 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.
211
    #[clap(default_value = "2048", long, env)]
212
    max_total_tokens: usize,
213
214
215
216
217
218
219
220
221
222
223

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

227
228
229
230
231
232
    /// 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,

233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
    /// **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.
250
251
    #[clap(long, env)]
    max_batch_total_tokens: Option<u32>,
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269

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

273
274
275
276
    /// The IP address to listen on
    #[clap(default_value = "0.0.0.0", long, env)]
    hostname: String,

277
    /// The port to listen on.
278
    #[clap(default_value = "3000", long, short, env)]
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
279
    port: u16,
280
281
282

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

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

    /// The address the master port will listen on. (setting used by torch distributed)
291
    #[clap(default_value = "29500", long, env)]
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
292
    master_port: usize,
293
294
295

    /// The location of the huggingface hub cache.
    /// Used to override the location if you want to provide a mounted disk for instance
296
    #[clap(long, env)]
297
    huggingface_hub_cache: Option<String>,
298
299
300

    /// The location of the huggingface hub cache.
    /// Used to override the location if you want to provide a mounted disk for instance
301
302
    #[clap(long, env)]
    weights_cache_override: Option<String>,
303
304
305
306
307

    /// 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.
308
    #[clap(long, env)]
309
    disable_custom_kernels: bool,
310

311
312
313
314
315
    /// 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
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
    /// 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>,

336
    /// Outputs the logs in JSON format (useful for telemetry)
337
    #[clap(long, env)]
338
    json_output: bool,
339

340
341
    #[clap(long, env)]
    otlp_endpoint: Option<String>,
342

343
344
    #[clap(long, env)]
    cors_allow_origin: Vec<String>,
345
346
347
348
    #[clap(long, env)]
    watermark_gamma: Option<f32>,
    #[clap(long, env)]
    watermark_delta: Option<f32>,
349

350
351
352
353
354
355
356
357
    /// Enable ngrok tunneling
    #[clap(long, env)]
    ngrok: bool,

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

358
    /// ngrok edge
359
    #[clap(long, env)]
360
    ngrok_edge: Option<String>,
361

362
363
364
    /// 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
365
366
}

367
368
369
#[derive(Debug)]
enum ShardStatus {
    Ready,
370
    Failed(usize),
371
}
372

373
374
375
376
#[allow(clippy::too_many_arguments)]
fn shard_manager(
    model_id: String,
    revision: Option<String>,
377
    quantize: Option<Quantization>,
378
    dtype: Option<Dtype>,
379
    trust_remote_code: bool,
380
381
382
383
384
385
386
387
388
389
    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>,
390
    cuda_memory_fraction: f32,
Nicolas Patry's avatar
Nicolas Patry committed
391
392
    rope_scaling: Option<RopeScaling>,
    rope_factor: Option<f32>,
393
394
    otlp_endpoint: Option<String>,
    status_sender: mpsc::Sender<ShardStatus>,
395
    shutdown: Arc<AtomicBool>,
396
397
    _shutdown_sender: mpsc::Sender<()>,
) {
398
399
400
    // Enter shard-manager tracing span
    let _span = tracing::span!(tracing::Level::INFO, "shard-manager", rank = rank).entered();

401
402
403
404
    // Get UDS path
    let uds_string = format!("{uds_path}-{rank}");
    let uds = Path::new(&uds_string);
    // Clean previous runs
405
406
407
    if uds.exists() {
        fs::remove_file(uds).unwrap();
    }
408
409

    // Process args
OlivierDehaene's avatar
OlivierDehaene committed
410
    let mut shard_args = vec![
411
412
413
414
415
416
417
418
419
        "serve".to_string(),
        model_id,
        "--uds-path".to_string(),
        uds_path,
        "--logger-level".to_string(),
        "INFO".to_string(),
        "--json-output".to_string(),
    ];

420
421
    // Activate trust remote code
    if trust_remote_code {
OlivierDehaene's avatar
OlivierDehaene committed
422
        shard_args.push("--trust-remote-code".to_string());
423
424
    }

425
426
    // Activate tensor parallelism
    if world_size > 1 {
OlivierDehaene's avatar
OlivierDehaene committed
427
        shard_args.push("--sharded".to_string());
428
429
    }

430
    if let Some(quantize) = quantize {
OlivierDehaene's avatar
OlivierDehaene committed
431
432
        shard_args.push("--quantize".to_string());
        shard_args.push(quantize.to_string())
433
    }
434

435
    if let Some(dtype) = dtype {
OlivierDehaene's avatar
OlivierDehaene committed
436
437
        shard_args.push("--dtype".to_string());
        shard_args.push(dtype.to_string())
438
439
    }

440
441
    // Model optional revision
    if let Some(revision) = revision {
OlivierDehaene's avatar
OlivierDehaene committed
442
443
        shard_args.push("--revision".to_string());
        shard_args.push(revision)
444
    }
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
445

Nicolas Patry's avatar
Nicolas Patry committed
446
447
448
449
450
451
    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)),
    };
452
453
    // OpenTelemetry
    if let Some(otlp_endpoint) = otlp_endpoint {
OlivierDehaene's avatar
OlivierDehaene committed
454
455
        shard_args.push("--otlp-endpoint".to_string());
        shard_args.push(otlp_endpoint);
456
457
458
    }

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

    // Torch Distributed Env vars
OlivierDehaene's avatar
OlivierDehaene committed
462
463
464
465
466
    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()));
467

468
469
470
471
472
473
    // CUDA memory fraction
    envs.push((
        "CUDA_MEMORY_FRACTION".into(),
        cuda_memory_fraction.to_string().into(),
    ));

474
    // Safetensors load fast
OlivierDehaene's avatar
OlivierDehaene committed
475
    envs.push(("SAFETENSORS_FAST_GPU".into(), "1".into()));
476
477
478

    // 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
479
    envs.push((
480
481
482
483
484
485
        "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
486
        envs.push(("HUGGING_FACE_HUB_TOKEN".into(), api_token.into()))
487
488
    };

Nicolas Patry's avatar
Nicolas Patry committed
489
490
491
492
493
494
495
496
497
    // 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()));
    }

498
499
500
    // 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
501
        envs.push(("HUGGINGFACE_HUB_CACHE".into(), huggingface_hub_cache.into()));
502
503
504
505
506
    };

    // 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
507
        envs.push((
508
509
510
511
512
513
514
            "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
515
        envs.push(("DISABLE_CUSTOM_KERNELS".into(), "True".into()))
516
517
518
519
    }

    // Watermark Gamma
    if let Some(watermark_gamma) = watermark_gamma {
OlivierDehaene's avatar
OlivierDehaene committed
520
        envs.push(("WATERMARK_GAMMA".into(), watermark_gamma.to_string().into()))
521
522
523
524
    }

    // Watermark Delta
    if let Some(watermark_delta) = watermark_delta {
OlivierDehaene's avatar
OlivierDehaene committed
525
        envs.push(("WATERMARK_DELTA".into(), watermark_delta.to_string().into()))
526
527
528
    }

    // Start process
529
    tracing::info!("Starting shard");
530
    let mut p = match Command::new("text-generation-server")
OlivierDehaene's avatar
OlivierDehaene committed
531
532
        .args(shard_args)
        .envs(envs)
533
534
535
536
537
        .stdout(Stdio::piped())
        .stderr(Stdio::piped())
        .process_group(0)
        .spawn()
    {
538
539
        Ok(p) => p,
        Err(err) => {
540
541
542
            if err.kind() == io::ErrorKind::NotFound {
                tracing::error!("text-generation-server not found in PATH");
                tracing::error!("Please install it with `make install-server`")
543
544
            }
            {
545
                tracing::error!("{}", err);
546
            }
547

548
            status_sender.send(ShardStatus::Failed(rank)).unwrap();
549
550
551
552
553
            return;
        }
    };

    // Redirect STDOUT to the console
554
    let shard_stdout_reader = BufReader::new(p.stdout.take().unwrap());
555
    let shard_stderr_reader = BufReader::new(p.stderr.take().unwrap());
556

557
    //stdout tracing thread
558
    thread::spawn(move || {
559
        log_lines(shard_stdout_reader.lines());
560
561
562
563
564
565
566
    });

    let mut ready = false;
    let start_time = Instant::now();
    let mut wait_time = Instant::now();
    loop {
        // Process exited
567
        if let Some(exit_status) = p.try_wait().unwrap() {
568
            // We read stderr in another thread as it seems that lines() can block in some cases
569
570
            let (err_sender, err_receiver) = mpsc::channel();
            thread::spawn(move || {
571
572
573
                for line in shard_stderr_reader.lines().flatten() {
                    err_sender.send(line).unwrap_or(());
                }
574
            });
575
576
577
578
            let mut err = String::new();
            while let Ok(line) = err_receiver.recv_timeout(Duration::from_millis(10)) {
                err = err + "\n" + &line;
            }
579

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

582
            if let Some(signal) = exit_status.signal() {
583
584
585
                tracing::error!("Shard process was signaled to shutdown with signal {signal}");
            }

586
            status_sender.send(ShardStatus::Failed(rank)).unwrap();
587
588
589
590
            return;
        }

        // We received a shutdown signal
591
        if shutdown.load(Ordering::SeqCst) {
592
            p.kill().unwrap();
593
            let _ = p.wait();
594
            tracing::info!("Shard terminated");
595
596
597
598
599
            return;
        }

        // Shard is ready
        if uds.exists() && !ready {
600
            tracing::info!("Shard ready in {:?}", start_time.elapsed());
601
602
603
            status_sender.send(ShardStatus::Ready).unwrap();
            ready = true;
        } else if !ready && wait_time.elapsed() > Duration::from_secs(10) {
604
            tracing::info!("Waiting for shard to be ready...");
605
606
607
608
609
610
            wait_time = Instant::now();
        }
        sleep(Duration::from_millis(100));
    }
}

611
fn shutdown_shards(shutdown: Arc<AtomicBool>, shutdown_receiver: &mpsc::Receiver<()>) {
612
613
614
    tracing::info!("Shutting down shards");
    // Update shutdown value to true
    // This will be picked up by the shard manager
615
    shutdown.store(true, Ordering::SeqCst);
616
617
618
619
620
621
622

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

fn num_cuda_devices() -> Option<usize> {
623
624
625
626
    let devices = match env::var("CUDA_VISIBLE_DEVICES") {
        Ok(devices) => devices,
        Err(_) => env::var("NVIDIA_VISIBLE_DEVICES").ok()?,
    };
627
628
    let n_devices = devices.split(',').count();
    Some(n_devices)
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
}

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

673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
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}"),
        }
    }
}

690
691
692
693
fn find_num_shards(
    sharded: Option<bool>,
    num_shard: Option<usize>,
) -> Result<usize, LauncherError> {
694
695
696
697
    // 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
698
699
700
            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");
701
            if n_devices <= 1 {
702
703
704
                return Err(LauncherError::NotEnoughCUDADevices(format!(
                    "`sharded` is true but only found {n_devices} CUDA devices"
                )));
705
            }
706
            n_devices
707
        }
708
709
710
        (Some(true), Some(num_shard)) => {
            // we can't have only one shard while sharded
            if num_shard <= 1 {
711
712
713
                return Err(LauncherError::ArgumentValidation(
                    "`sharded` is true but `num_shard` <= 1".to_string(),
                ));
714
715
            }
            num_shard
716
        }
717
718
719
720
        (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,
721
    };
722
    if num_shard < 1 {
723
724
725
        return Err(LauncherError::ArgumentValidation(
            "`num_shard` cannot be < 1".to_string(),
        ));
726
    }
727
    Ok(num_shard)
728
}
729

730
731
#[derive(Debug)]
enum LauncherError {
732
733
    ArgumentValidation(String),
    NotEnoughCUDADevices(String),
734
735
736
737
738
739
740
    DownloadError,
    ShardCannotStart,
    ShardDisconnected,
    ShardFailed,
    WebserverFailed,
    WebserverCannotStart,
}
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
741

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

OlivierDehaene's avatar
OlivierDehaene committed
746
    let mut download_args = vec![
747
748
749
750
751
752
753
754
        "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(),
    ];
755

756
757
    // Model optional revision
    if let Some(revision) = &args.revision {
OlivierDehaene's avatar
OlivierDehaene committed
758
759
        download_args.push("--revision".to_string());
        download_args.push(revision.to_string())
760
    }
761

762
763
764
765
766
    // Trust remote code for automatic peft fusion
    if args.trust_remote_code {
        download_args.push("--trust-remote-code".to_string());
    }

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

770
    // If huggingface_hub_cache is set, pass it to the download process
771
772
    // Useful when running inside a docker container
    if let Some(ref huggingface_hub_cache) = args.huggingface_hub_cache {
OlivierDehaene's avatar
OlivierDehaene committed
773
        envs.push(("HUGGINGFACE_HUB_CACHE".into(), huggingface_hub_cache.into()));
774
    };
775

776
777
    // 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
778
    envs.push((
779
780
781
        "HF_HUB_ENABLE_HF_TRANSFER".into(),
        enable_hf_transfer.into(),
    ));
782

783
784
    // Parse Inference API token
    if let Ok(api_token) = env::var("HF_API_TOKEN") {
OlivierDehaene's avatar
OlivierDehaene committed
785
        envs.push(("HUGGING_FACE_HUB_TOKEN".into(), api_token.into()))
786
    };
787

788
789
790
    // 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
791
        envs.push((
792
793
794
795
796
            "WEIGHTS_CACHE_OVERRIDE".into(),
            weights_cache_override.into(),
        ));
    };

797
798
    // Start process
    tracing::info!("Starting download process.");
799
    let mut download_process = match Command::new("text-generation-server")
OlivierDehaene's avatar
OlivierDehaene committed
800
801
        .args(download_args)
        .envs(envs)
802
803
804
805
806
        .stdout(Stdio::piped())
        .stderr(Stdio::piped())
        .process_group(0)
        .spawn()
    {
807
808
        Ok(p) => p,
        Err(err) => {
809
810
811
            if err.kind() == io::ErrorKind::NotFound {
                tracing::error!("text-generation-server not found in PATH");
                tracing::error!("Please install it with `make install-server`")
812
813
            } else {
                tracing::error!("{}", err);
814
            }
815

816
817
818
            return Err(LauncherError::DownloadError);
        }
    };
819

820
821
    // Redirect STDOUT to the console
    let download_stdout = download_process.stdout.take().unwrap();
822
823
    let stdout = BufReader::new(download_stdout);

824
    thread::spawn(move || {
825
        log_lines(stdout.lines());
826
    });
827

828
    loop {
829
830
831
832
        if let Some(status) = download_process.try_wait().unwrap() {
            if status.success() {
                tracing::info!("Successfully downloaded weights.");
                break;
833
            }
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850

            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);
851
        }
852
        if !running.load(Ordering::SeqCst) {
OlivierDehaene's avatar
OlivierDehaene committed
853
            terminate("download", download_process, Duration::from_secs(10)).unwrap();
854
855
856
            return Ok(());
        }
        sleep(Duration::from_millis(100));
857
    }
858
859
    Ok(())
}
860

861
#[allow(clippy::too_many_arguments)]
862
863
864
fn spawn_shards(
    num_shard: usize,
    args: &Args,
865
    shutdown: Arc<AtomicBool>,
866
867
868
869
870
871
    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
872
873
    // Start shard processes
    for rank in 0..num_shard {
874
875
876
877
878
879
        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
880
881
882
        let status_sender = status_sender.clone();
        let shutdown = shutdown.clone();
        let shutdown_sender = shutdown_sender.clone();
883
        let otlp_endpoint = args.otlp_endpoint.clone();
884
        let quantize = args.quantize;
885
        let dtype = args.dtype;
886
        let trust_remote_code = args.trust_remote_code;
887
888
889
890
        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;
891
        let cuda_memory_fraction = args.cuda_memory_fraction;
Nicolas Patry's avatar
Nicolas Patry committed
892
893
        let rope_scaling = args.rope_scaling;
        let rope_factor = args.rope_factor;
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
894
895
        thread::spawn(move || {
            shard_manager(
896
                model_id,
897
                revision,
898
                quantize,
899
                dtype,
900
                trust_remote_code,
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
901
902
903
904
905
                uds_path,
                rank,
                num_shard,
                master_addr,
                master_port,
906
907
                huggingface_hub_cache,
                weights_cache_override,
908
                disable_custom_kernels,
909
910
                watermark_gamma,
                watermark_delta,
911
                cuda_memory_fraction,
Nicolas Patry's avatar
Nicolas Patry committed
912
913
                rope_scaling,
                rope_factor,
914
                otlp_endpoint,
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
                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));
            }
936
            Ok(ShardStatus::Failed(rank)) => {
937
                tracing::error!("Shard {rank} failed to start");
938
                shutdown_shards(shutdown, shutdown_receiver);
939
                return Err(LauncherError::ShardCannotStart);
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
940
941
942
            }
            Err(TryRecvError::Disconnected) => {
                tracing::error!("Shard status channel disconnected");
943
                shutdown_shards(shutdown, shutdown_receiver);
944
                return Err(LauncherError::ShardDisconnected);
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
945
946
947
            }
        }
    }
948
949
    Ok(())
}
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
950

951
952
fn spawn_webserver(
    args: Args,
953
    shutdown: Arc<AtomicBool>,
954
    shutdown_receiver: &mpsc::Receiver<()>,
955
) -> Result<Child, LauncherError> {
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
956
957
958
    // All shard started
    // Start webserver
    tracing::info!("Starting Webserver");
OlivierDehaene's avatar
OlivierDehaene committed
959
    let mut router_args = vec![
960
        "--max-concurrent-requests".to_string(),
961
        args.max_concurrent_requests.to_string(),
962
        "--max-best-of".to_string(),
963
        args.max_best_of.to_string(),
964
        "--max-stop-sequences".to_string(),
965
        args.max_stop_sequences.to_string(),
Nicolas Patry's avatar
Nicolas Patry committed
966
967
        "--max-top-n-tokens".to_string(),
        args.max_top_n_tokens.to_string(),
968
        "--max-input-length".to_string(),
969
        args.max_input_length.to_string(),
970
        "--max-total-tokens".to_string(),
971
        args.max_total_tokens.to_string(),
972
973
        "--max-batch-prefill-tokens".to_string(),
        args.max_batch_prefill_tokens.to_string(),
974
        "--waiting-served-ratio".to_string(),
975
        args.waiting_served_ratio.to_string(),
976
        "--max-waiting-tokens".to_string(),
977
        args.max_waiting_tokens.to_string(),
978
979
        "--validation-workers".to_string(),
        args.validation_workers.to_string(),
980
981
        "--hostname".to_string(),
        args.hostname.to_string(),
982
        "--port".to_string(),
983
        args.port.to_string(),
984
        "--master-shard-uds-path".to_string(),
985
        format!("{}-0", args.shard_uds_path),
986
        "--tokenizer-name".to_string(),
987
        args.model_id,
988
989
    ];

990
991
992
993
994
995
    // 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());
    }

996
997
    // Model optional revision
    if let Some(ref revision) = args.revision {
OlivierDehaene's avatar
OlivierDehaene committed
998
999
        router_args.push("--revision".to_string());
        router_args.push(revision.to_string())
1000
1001
    }

1002
    if args.json_output {
OlivierDehaene's avatar
OlivierDehaene committed
1003
        router_args.push("--json-output".to_string());
1004
1005
    }

1006
    // OpenTelemetry
1007
    if let Some(otlp_endpoint) = args.otlp_endpoint {
OlivierDehaene's avatar
OlivierDehaene committed
1008
1009
        router_args.push("--otlp-endpoint".to_string());
        router_args.push(otlp_endpoint);
1010
1011
1012
1013
    }

    // CORS origins
    for origin in args.cors_allow_origin.into_iter() {
OlivierDehaene's avatar
OlivierDehaene committed
1014
1015
        router_args.push("--cors-allow-origin".to_string());
        router_args.push(origin);
1016
1017
    }

1018
1019
    // Ngrok
    if args.ngrok {
OlivierDehaene's avatar
OlivierDehaene committed
1020
1021
        router_args.push("--ngrok".to_string());
        router_args.push("--ngrok-authtoken".to_string());
1022
1023
1024
        router_args.push(args.ngrok_authtoken.unwrap());
        router_args.push("--ngrok-edge".to_string());
        router_args.push(args.ngrok_edge.unwrap());
1025
1026
    }

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

1030
1031
    // Parse Inference API token
    if let Ok(api_token) = env::var("HF_API_TOKEN") {
OlivierDehaene's avatar
OlivierDehaene committed
1032
        envs.push(("HUGGING_FACE_HUB_TOKEN".into(), api_token.into()))
1033
    };
1034

1035
    let mut webserver = match Command::new("text-generation-router")
OlivierDehaene's avatar
OlivierDehaene committed
1036
1037
        .args(router_args)
        .envs(envs)
1038
1039
1040
1041
1042
        .stdout(Stdio::piped())
        .stderr(Stdio::piped())
        .process_group(0)
        .spawn()
    {
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1043
1044
        Ok(p) => p,
        Err(err) => {
1045
            tracing::error!("Failed to start webserver: {}", err);
1046
1047
1048
            if err.kind() == io::ErrorKind::NotFound {
                tracing::error!("text-generation-router not found in PATH");
                tracing::error!("Please install it with `make install-router`")
1049
1050
            } else {
                tracing::error!("{}", err);
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1051
            }
1052

1053
            shutdown_shards(shutdown, shutdown_receiver);
1054
            return Err(LauncherError::WebserverCannotStart);
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1055
1056
1057
        }
    };

1058
1059
1060
    // Redirect STDOUT and STDERR to the console
    let webserver_stdout = webserver.stdout.take().unwrap();
    let webserver_stderr = webserver.stderr.take().unwrap();
1061
1062

    thread::spawn(move || {
1063
1064
        let stdout = BufReader::new(webserver_stdout);
        let stderr = BufReader::new(webserver_stderr);
1065
        for line in stdout.lines() {
1066
            println!("{}", line.unwrap());
1067
        }
1068
1069
        for line in stderr.lines() {
            println!("{}", line.unwrap());
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1070
        }
1071
1072
1073
    });
    Ok(webserver)
}
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1074

OlivierDehaene's avatar
OlivierDehaene committed
1075
1076
1077
1078
1079
1080
1081
1082
1083
1084
1085
1086
1087
1088
1089
1090
1091
1092
1093
1094
1095
1096
1097
1098
1099
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)
}

1100
1101
fn main() -> Result<(), LauncherError> {
    // Pattern match configuration
1102
    let args: Args = Args::parse();
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1103

1104
1105
1106
1107
    // Filter events with LOG_LEVEL
    let env_filter =
        EnvFilter::try_from_env("LOG_LEVEL").unwrap_or_else(|_| EnvFilter::new("info"));

1108
    if args.json_output {
1109
1110
1111
1112
        tracing_subscriber::fmt()
            .with_env_filter(env_filter)
            .json()
            .init();
1113
    } else {
1114
1115
1116
1117
        tracing_subscriber::fmt()
            .with_env_filter(env_filter)
            .compact()
            .init();
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1118
1119
    }

1120
1121
1122
1123
1124
    if args.env {
        let env_runtime = env_runtime::Env::new();
        tracing::info!("{}", env_runtime);
    }

1125
1126
    tracing::info!("{:?}", args);

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

1140
1141
1142
1143
1144
    if args.validation_workers == 0 {
        return Err(LauncherError::ArgumentValidation(
            "`validation_workers` must be > 0".to_string(),
        ));
    }
1145
1146
1147
1148
1149
1150
    if args.trust_remote_code {
        tracing::warn!(
            "`trust_remote_code` is set. Trusting that model `{}` do not contain malicious code.",
            args.model_id
        );
    }
1151
1152

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

1157
1158
1159
1160
1161
1162
1163
1164
1165
1166
1167
1168
1169
1170
1171
    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
            )));
        }
    }

1172
1173
1174
1175
1176
1177
1178
1179
1180
1181
1182
1183
1184
1185
    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(),
            ));
        }
    }

1186
1187
1188
1189
1190
1191
1192
    // 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");
1193

1194
    // Download and convert model weights
1195
    download_convert_model(&args, running.clone())?;
1196

OlivierDehaene's avatar
OlivierDehaene committed
1197
1198
1199
1200
1201
    if !running.load(Ordering::SeqCst) {
        // Launcher was asked to stop
        return Ok(());
    }

1202
    // Shared shutdown bool
1203
    let shutdown = Arc::new(AtomicBool::new(false));
1204
1205
1206
    // Shared shutdown channel
    // When shutting down, the main thread will wait for all senders to be dropped
    let (shutdown_sender, shutdown_receiver) = mpsc::channel();
1207

1208
1209
    // Shared channel to track shard status
    let (status_sender, status_receiver) = mpsc::channel();
1210

1211
1212
1213
1214
1215
1216
1217
1218
1219
1220
    spawn_shards(
        num_shard,
        &args,
        shutdown.clone(),
        &shutdown_receiver,
        shutdown_sender,
        &status_receiver,
        status_sender,
        running.clone(),
    )?;
1221

1222
1223
1224
1225
1226
    // We might have received a termination signal
    if !running.load(Ordering::SeqCst) {
        shutdown_shards(shutdown, &shutdown_receiver);
        return Ok(());
    }
1227

OlivierDehaene's avatar
OlivierDehaene committed
1228
1229
1230
1231
1232
    let mut webserver =
        spawn_webserver(args, shutdown.clone(), &shutdown_receiver).map_err(|err| {
            shutdown_shards(shutdown.clone(), &shutdown_receiver);
            err
        })?;
1233
1234
1235
1236
1237

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

    while running.load(Ordering::SeqCst) {
1238
        if let Ok(ShardStatus::Failed(rank)) = status_receiver.try_recv() {
OlivierDehaene's avatar
OlivierDehaene committed
1239
            tracing::error!("Shard {rank} crashed");
1240
1241
1242
1243
            exit_code = Err(LauncherError::ShardFailed);
            break;
        };

1244
        match webserver.try_wait().unwrap() {
1245
1246
1247
1248
1249
1250
1251
1252
1253
            Some(_) => {
                tracing::error!("Webserver Crashed");
                shutdown_shards(shutdown, &shutdown_receiver);
                return Err(LauncherError::WebserverFailed);
            }
            None => {
                sleep(Duration::from_millis(100));
            }
        };
1254
    }
1255
1256

    // Graceful termination
OlivierDehaene's avatar
OlivierDehaene committed
1257
    terminate("webserver", webserver, Duration::from_secs(90)).unwrap();
1258
1259
1260
    shutdown_shards(shutdown, &shutdown_receiver);

    exit_code
1261
}