main.rs 43.5 KB
Newer Older
1
use clap::{Parser, ValueEnum};
2
3
use nix::sys::signal::{self, Signal};
use nix::unistd::Pid;
4
use serde::Deserialize;
Nicolas Patry's avatar
Nicolas Patry committed
5
use std::env;
6
use std::ffi::OsString;
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

Nicolas Patry's avatar
Nicolas Patry committed
158
159
160
161
162
163
164
    /// The number of input_ids to speculate on
    /// If using a medusa model, the heads will be picked up automatically
    /// Other wise, it will use n-gram speculation which is relatively free
    /// in terms of compute, but the speedup heavily depends on the task.
    #[clap(long, env)]
    speculate: Option<usize>,

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

169
170
171
172
173
174
    /// 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,

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

    /// 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
184
185
    #[clap(default_value = "2", long, env)]
    max_best_of: usize,
186
187
188
189
190
191

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

Nicolas Patry's avatar
Nicolas Patry committed
195
196
197
198
199
200
201
202
    /// 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,

203
204
205
206
    /// 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.
207
    #[clap(default_value = "1024", long, env)]
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
208
    max_input_length: usize,
209
210
211
212
213
214
215
216
217

    /// 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.
218
    #[clap(default_value = "2048", long, env)]
219
    max_total_tokens: usize,
220
221
222
223
224
225
226
227
228
229
230

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

234
235
236
237
238
239
    /// 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,

240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
    /// **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.
257
258
    #[clap(long, env)]
    max_batch_total_tokens: Option<u32>,
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276

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

280
281
282
283
    /// The IP address to listen on
    #[clap(default_value = "0.0.0.0", long, env)]
    hostname: String,

284
    /// The port to listen on.
285
    #[clap(default_value = "3000", long, short, env)]
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
286
    port: u16,
287
288
289

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

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

    /// The address the master port will listen on. (setting used by torch distributed)
298
    #[clap(default_value = "29500", long, env)]
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
299
    master_port: usize,
300
301
302

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

    /// The location of the huggingface hub cache.
    /// Used to override the location if you want to provide a mounted disk for instance
308
309
    #[clap(long, env)]
    weights_cache_override: Option<String>,
310
311
312
313
314

    /// 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.
315
    #[clap(long, env)]
316
    disable_custom_kernels: bool,
317

318
319
320
321
322
    /// 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
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
    /// 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>,

343
    /// Outputs the logs in JSON format (useful for telemetry)
344
    #[clap(long, env)]
345
    json_output: bool,
346

347
348
    #[clap(long, env)]
    otlp_endpoint: Option<String>,
349

350
351
    #[clap(long, env)]
    cors_allow_origin: Vec<String>,
352
353
354
355
    #[clap(long, env)]
    watermark_gamma: Option<f32>,
    #[clap(long, env)]
    watermark_delta: Option<f32>,
356

357
358
359
360
361
362
363
364
    /// Enable ngrok tunneling
    #[clap(long, env)]
    ngrok: bool,

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

365
    /// ngrok edge
366
    #[clap(long, env)]
367
    ngrok_edge: Option<String>,
368

369
370
371
    /// 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
372
373
}

374
375
376
#[derive(Debug)]
enum ShardStatus {
    Ready,
377
    Failed(usize),
378
}
379

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

409
410
411
412
    // Get UDS path
    let uds_string = format!("{uds_path}-{rank}");
    let uds = Path::new(&uds_string);
    // Clean previous runs
413
414
415
    if uds.exists() {
        fs::remove_file(uds).unwrap();
    }
416
417

    // Process args
OlivierDehaene's avatar
OlivierDehaene committed
418
    let mut shard_args = vec![
419
420
421
422
423
424
425
426
427
        "serve".to_string(),
        model_id,
        "--uds-path".to_string(),
        uds_path,
        "--logger-level".to_string(),
        "INFO".to_string(),
        "--json-output".to_string(),
    ];

428
429
    // Activate trust remote code
    if trust_remote_code {
OlivierDehaene's avatar
OlivierDehaene committed
430
        shard_args.push("--trust-remote-code".to_string());
431
432
    }

433
434
    // Activate tensor parallelism
    if world_size > 1 {
OlivierDehaene's avatar
OlivierDehaene committed
435
        shard_args.push("--sharded".to_string());
436
437
    }

438
    if let Some(quantize) = quantize {
OlivierDehaene's avatar
OlivierDehaene committed
439
440
        shard_args.push("--quantize".to_string());
        shard_args.push(quantize.to_string())
441
    }
442

Nicolas Patry's avatar
Nicolas Patry committed
443
444
445
446
447
    if let Some(speculate) = speculate {
        shard_args.push("--speculate".to_string());
        shard_args.push(speculate.to_string())
    }

448
    if let Some(dtype) = dtype {
OlivierDehaene's avatar
OlivierDehaene committed
449
450
        shard_args.push("--dtype".to_string());
        shard_args.push(dtype.to_string())
451
452
    }

453
454
    // Model optional revision
    if let Some(revision) = revision {
OlivierDehaene's avatar
OlivierDehaene committed
455
456
        shard_args.push("--revision".to_string());
        shard_args.push(revision)
457
    }
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
458

Nicolas Patry's avatar
Nicolas Patry committed
459
460
461
462
463
464
    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)),
    };
465
466
    // OpenTelemetry
    if let Some(otlp_endpoint) = otlp_endpoint {
OlivierDehaene's avatar
OlivierDehaene committed
467
468
        shard_args.push("--otlp-endpoint".to_string());
        shard_args.push(otlp_endpoint);
469
470
471
    }

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

    // Torch Distributed Env vars
OlivierDehaene's avatar
OlivierDehaene committed
475
476
477
478
479
    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()));
480

481
482
483
484
485
486
    // CUDA memory fraction
    envs.push((
        "CUDA_MEMORY_FRACTION".into(),
        cuda_memory_fraction.to_string().into(),
    ));

487
    // Safetensors load fast
OlivierDehaene's avatar
OlivierDehaene committed
488
    envs.push(("SAFETENSORS_FAST_GPU".into(), "1".into()));
489
490
491

    // 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
492
    envs.push((
493
494
495
496
497
498
        "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
499
        envs.push(("HUGGING_FACE_HUB_TOKEN".into(), api_token.into()))
500
501
    };

Nicolas Patry's avatar
Nicolas Patry committed
502
503
504
505
506
507
508
509
510
    // 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()));
    }

511
512
513
    // 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
514
        envs.push(("HUGGINGFACE_HUB_CACHE".into(), huggingface_hub_cache.into()));
515
516
517
518
519
    };

    // 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
520
        envs.push((
521
522
523
524
525
526
527
            "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
528
        envs.push(("DISABLE_CUSTOM_KERNELS".into(), "True".into()))
529
530
531
532
    }

    // Watermark Gamma
    if let Some(watermark_gamma) = watermark_gamma {
OlivierDehaene's avatar
OlivierDehaene committed
533
        envs.push(("WATERMARK_GAMMA".into(), watermark_gamma.to_string().into()))
534
535
536
537
    }

    // Watermark Delta
    if let Some(watermark_delta) = watermark_delta {
OlivierDehaene's avatar
OlivierDehaene committed
538
        envs.push(("WATERMARK_DELTA".into(), watermark_delta.to_string().into()))
539
540
541
    }

    // Start process
542
    tracing::info!("Starting shard");
543
    let mut p = match Command::new("text-generation-server")
OlivierDehaene's avatar
OlivierDehaene committed
544
545
        .args(shard_args)
        .envs(envs)
546
547
548
549
550
        .stdout(Stdio::piped())
        .stderr(Stdio::piped())
        .process_group(0)
        .spawn()
    {
551
552
        Ok(p) => p,
        Err(err) => {
553
554
555
            if err.kind() == io::ErrorKind::NotFound {
                tracing::error!("text-generation-server not found in PATH");
                tracing::error!("Please install it with `make install-server`")
556
557
            }
            {
558
                tracing::error!("{}", err);
559
            }
560

561
            status_sender.send(ShardStatus::Failed(rank)).unwrap();
562
563
564
565
566
            return;
        }
    };

    // Redirect STDOUT to the console
567
    let shard_stdout_reader = BufReader::new(p.stdout.take().unwrap());
568
    let shard_stderr_reader = BufReader::new(p.stderr.take().unwrap());
569

570
    //stdout tracing thread
571
    thread::spawn(move || {
572
        log_lines(shard_stdout_reader.lines());
573
574
575
576
577
578
579
    });

    let mut ready = false;
    let start_time = Instant::now();
    let mut wait_time = Instant::now();
    loop {
        // Process exited
580
        if let Some(exit_status) = p.try_wait().unwrap() {
581
            // We read stderr in another thread as it seems that lines() can block in some cases
582
583
            let (err_sender, err_receiver) = mpsc::channel();
            thread::spawn(move || {
584
585
586
                for line in shard_stderr_reader.lines().flatten() {
                    err_sender.send(line).unwrap_or(());
                }
587
            });
588
589
590
591
            let mut err = String::new();
            while let Ok(line) = err_receiver.recv_timeout(Duration::from_millis(10)) {
                err = err + "\n" + &line;
            }
592

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

595
            if let Some(signal) = exit_status.signal() {
596
597
598
                tracing::error!("Shard process was signaled to shutdown with signal {signal}");
            }

599
            status_sender.send(ShardStatus::Failed(rank)).unwrap();
600
601
602
603
            return;
        }

        // We received a shutdown signal
604
        if shutdown.load(Ordering::SeqCst) {
605
            p.kill().unwrap();
606
            let _ = p.wait();
607
            tracing::info!("Shard terminated");
608
609
610
611
612
            return;
        }

        // Shard is ready
        if uds.exists() && !ready {
613
            tracing::info!("Shard ready in {:?}", start_time.elapsed());
614
615
616
            status_sender.send(ShardStatus::Ready).unwrap();
            ready = true;
        } else if !ready && wait_time.elapsed() > Duration::from_secs(10) {
617
            tracing::info!("Waiting for shard to be ready...");
618
619
620
621
622
623
            wait_time = Instant::now();
        }
        sleep(Duration::from_millis(100));
    }
}

624
fn shutdown_shards(shutdown: Arc<AtomicBool>, shutdown_receiver: &mpsc::Receiver<()>) {
625
626
627
    tracing::info!("Shutting down shards");
    // Update shutdown value to true
    // This will be picked up by the shard manager
628
    shutdown.store(true, Ordering::SeqCst);
629
630
631
632
633
634
635

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

fn num_cuda_devices() -> Option<usize> {
636
637
638
639
    let devices = match env::var("CUDA_VISIBLE_DEVICES") {
        Ok(devices) => devices,
        Err(_) => env::var("NVIDIA_VISIBLE_DEVICES").ok()?,
    };
640
641
    let n_devices = devices.split(',').count();
    Some(n_devices)
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
673
674
675
676
677
678
679
680
681
682
683
684
685
}

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

686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
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}"),
        }
    }
}

703
704
705
706
fn find_num_shards(
    sharded: Option<bool>,
    num_shard: Option<usize>,
) -> Result<usize, LauncherError> {
707
708
709
710
    // 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
711
712
713
            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");
714
            if n_devices <= 1 {
715
716
717
                return Err(LauncherError::NotEnoughCUDADevices(format!(
                    "`sharded` is true but only found {n_devices} CUDA devices"
                )));
718
            }
719
            n_devices
720
        }
721
722
723
        (Some(true), Some(num_shard)) => {
            // we can't have only one shard while sharded
            if num_shard <= 1 {
724
725
726
                return Err(LauncherError::ArgumentValidation(
                    "`sharded` is true but `num_shard` <= 1".to_string(),
                ));
727
728
            }
            num_shard
729
        }
730
731
732
733
        (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,
734
    };
735
    if num_shard < 1 {
736
737
738
        return Err(LauncherError::ArgumentValidation(
            "`num_shard` cannot be < 1".to_string(),
        ));
739
    }
740
    Ok(num_shard)
741
}
742

743
744
#[derive(Debug)]
enum LauncherError {
745
746
    ArgumentValidation(String),
    NotEnoughCUDADevices(String),
747
748
749
750
751
752
753
    DownloadError,
    ShardCannotStart,
    ShardDisconnected,
    ShardFailed,
    WebserverFailed,
    WebserverCannotStart,
}
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
754

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

OlivierDehaene's avatar
OlivierDehaene committed
759
    let mut download_args = vec![
760
761
762
763
764
765
766
767
        "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(),
    ];
768

769
770
    // Model optional revision
    if let Some(revision) = &args.revision {
OlivierDehaene's avatar
OlivierDehaene committed
771
772
        download_args.push("--revision".to_string());
        download_args.push(revision.to_string())
773
    }
774

775
776
777
778
779
    // Trust remote code for automatic peft fusion
    if args.trust_remote_code {
        download_args.push("--trust-remote-code".to_string());
    }

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

783
    // If huggingface_hub_cache is set, pass it to the download process
784
785
    // Useful when running inside a docker container
    if let Some(ref huggingface_hub_cache) = args.huggingface_hub_cache {
OlivierDehaene's avatar
OlivierDehaene committed
786
        envs.push(("HUGGINGFACE_HUB_CACHE".into(), huggingface_hub_cache.into()));
787
    };
788

789
790
    // 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
791
    envs.push((
792
793
794
        "HF_HUB_ENABLE_HF_TRANSFER".into(),
        enable_hf_transfer.into(),
    ));
795

796
797
    // Parse Inference API token
    if let Ok(api_token) = env::var("HF_API_TOKEN") {
OlivierDehaene's avatar
OlivierDehaene committed
798
        envs.push(("HUGGING_FACE_HUB_TOKEN".into(), api_token.into()))
799
    };
800

801
802
803
    // 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
804
        envs.push((
805
806
807
808
809
            "WEIGHTS_CACHE_OVERRIDE".into(),
            weights_cache_override.into(),
        ));
    };

810
811
    // Start process
    tracing::info!("Starting download process.");
812
    let mut download_process = match Command::new("text-generation-server")
OlivierDehaene's avatar
OlivierDehaene committed
813
814
        .args(download_args)
        .envs(envs)
815
816
817
818
819
        .stdout(Stdio::piped())
        .stderr(Stdio::piped())
        .process_group(0)
        .spawn()
    {
820
821
        Ok(p) => p,
        Err(err) => {
822
823
824
            if err.kind() == io::ErrorKind::NotFound {
                tracing::error!("text-generation-server not found in PATH");
                tracing::error!("Please install it with `make install-server`")
825
826
            } else {
                tracing::error!("{}", err);
827
            }
828

829
830
831
            return Err(LauncherError::DownloadError);
        }
    };
832

833
834
    // Redirect STDOUT to the console
    let download_stdout = download_process.stdout.take().unwrap();
835
836
    let stdout = BufReader::new(download_stdout);

837
    thread::spawn(move || {
838
        log_lines(stdout.lines());
839
    });
840

841
    loop {
842
843
844
845
        if let Some(status) = download_process.try_wait().unwrap() {
            if status.success() {
                tracing::info!("Successfully downloaded weights.");
                break;
846
            }
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863

            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);
864
        }
865
        if !running.load(Ordering::SeqCst) {
OlivierDehaene's avatar
OlivierDehaene committed
866
            terminate("download", download_process, Duration::from_secs(10)).unwrap();
867
868
869
            return Ok(());
        }
        sleep(Duration::from_millis(100));
870
    }
871
872
    Ok(())
}
873

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

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

1005
1006
1007
1008
1009
1010
    // 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());
    }

1011
1012
    // Model optional revision
    if let Some(ref revision) = args.revision {
OlivierDehaene's avatar
OlivierDehaene committed
1013
1014
        router_args.push("--revision".to_string());
        router_args.push(revision.to_string())
1015
1016
    }

1017
    if args.json_output {
OlivierDehaene's avatar
OlivierDehaene committed
1018
        router_args.push("--json-output".to_string());
1019
1020
    }

1021
    // OpenTelemetry
1022
    if let Some(otlp_endpoint) = args.otlp_endpoint {
OlivierDehaene's avatar
OlivierDehaene committed
1023
1024
        router_args.push("--otlp-endpoint".to_string());
        router_args.push(otlp_endpoint);
1025
1026
1027
1028
    }

    // CORS origins
    for origin in args.cors_allow_origin.into_iter() {
OlivierDehaene's avatar
OlivierDehaene committed
1029
1030
        router_args.push("--cors-allow-origin".to_string());
        router_args.push(origin);
1031
1032
    }

1033
1034
    // Ngrok
    if args.ngrok {
OlivierDehaene's avatar
OlivierDehaene committed
1035
1036
        router_args.push("--ngrok".to_string());
        router_args.push("--ngrok-authtoken".to_string());
1037
1038
1039
        router_args.push(args.ngrok_authtoken.unwrap());
        router_args.push("--ngrok-edge".to_string());
        router_args.push(args.ngrok_edge.unwrap());
1040
1041
    }

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

1045
1046
    // Parse Inference API token
    if let Ok(api_token) = env::var("HF_API_TOKEN") {
OlivierDehaene's avatar
OlivierDehaene committed
1047
        envs.push(("HUGGING_FACE_HUB_TOKEN".into(), api_token.into()))
1048
    };
1049

1050
    let mut webserver = match Command::new("text-generation-router")
OlivierDehaene's avatar
OlivierDehaene committed
1051
1052
        .args(router_args)
        .envs(envs)
1053
1054
1055
1056
1057
        .stdout(Stdio::piped())
        .stderr(Stdio::piped())
        .process_group(0)
        .spawn()
    {
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1058
1059
        Ok(p) => p,
        Err(err) => {
1060
            tracing::error!("Failed to start webserver: {}", err);
1061
1062
1063
            if err.kind() == io::ErrorKind::NotFound {
                tracing::error!("text-generation-router not found in PATH");
                tracing::error!("Please install it with `make install-router`")
1064
1065
            } else {
                tracing::error!("{}", err);
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1066
            }
1067

1068
            shutdown_shards(shutdown, shutdown_receiver);
1069
            return Err(LauncherError::WebserverCannotStart);
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1070
1071
1072
        }
    };

1073
1074
1075
    // Redirect STDOUT and STDERR to the console
    let webserver_stdout = webserver.stdout.take().unwrap();
    let webserver_stderr = webserver.stderr.take().unwrap();
1076
1077

    thread::spawn(move || {
1078
1079
        let stdout = BufReader::new(webserver_stdout);
        let stderr = BufReader::new(webserver_stderr);
1080
        for line in stdout.lines() {
1081
            println!("{}", line.unwrap());
1082
        }
1083
1084
        for line in stderr.lines() {
            println!("{}", line.unwrap());
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1085
        }
1086
1087
1088
    });
    Ok(webserver)
}
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1089

OlivierDehaene's avatar
OlivierDehaene committed
1090
1091
1092
1093
1094
1095
1096
1097
1098
1099
1100
1101
1102
1103
1104
1105
1106
1107
1108
1109
1110
1111
1112
1113
1114
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)
}

1115
1116
fn main() -> Result<(), LauncherError> {
    // Pattern match configuration
1117
    let args: Args = Args::parse();
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1118

1119
1120
1121
1122
    // Filter events with LOG_LEVEL
    let env_filter =
        EnvFilter::try_from_env("LOG_LEVEL").unwrap_or_else(|_| EnvFilter::new("info"));

1123
    if args.json_output {
1124
1125
1126
1127
        tracing_subscriber::fmt()
            .with_env_filter(env_filter)
            .json()
            .init();
1128
    } else {
1129
1130
1131
1132
        tracing_subscriber::fmt()
            .with_env_filter(env_filter)
            .compact()
            .init();
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1133
1134
    }

1135
1136
1137
1138
1139
    if args.env {
        let env_runtime = env_runtime::Env::new();
        tracing::info!("{}", env_runtime);
    }

1140
1141
    tracing::info!("{:?}", args);

1142
1143
1144
1145
1146
1147
1148
1149
1150
1151
1152
1153
    // 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
        )));
    }
1154

1155
1156
1157
1158
1159
    if args.validation_workers == 0 {
        return Err(LauncherError::ArgumentValidation(
            "`validation_workers` must be > 0".to_string(),
        ));
    }
1160
1161
1162
1163
1164
1165
    if args.trust_remote_code {
        tracing::warn!(
            "`trust_remote_code` is set. Trusting that model `{}` do not contain malicious code.",
            args.model_id
        );
    }
1166
1167

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

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

1187
1188
1189
1190
1191
1192
1193
1194
1195
1196
1197
1198
1199
1200
    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(),
            ));
        }
    }

1201
1202
1203
1204
1205
1206
1207
    // 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");
1208

1209
    // Download and convert model weights
1210
    download_convert_model(&args, running.clone())?;
1211

OlivierDehaene's avatar
OlivierDehaene committed
1212
1213
1214
1215
1216
    if !running.load(Ordering::SeqCst) {
        // Launcher was asked to stop
        return Ok(());
    }

1217
    // Shared shutdown bool
1218
    let shutdown = Arc::new(AtomicBool::new(false));
1219
1220
1221
    // Shared shutdown channel
    // When shutting down, the main thread will wait for all senders to be dropped
    let (shutdown_sender, shutdown_receiver) = mpsc::channel();
1222

1223
1224
    // Shared channel to track shard status
    let (status_sender, status_receiver) = mpsc::channel();
1225

1226
1227
1228
1229
1230
1231
1232
1233
1234
1235
    spawn_shards(
        num_shard,
        &args,
        shutdown.clone(),
        &shutdown_receiver,
        shutdown_sender,
        &status_receiver,
        status_sender,
        running.clone(),
    )?;
1236

1237
1238
1239
1240
1241
    // We might have received a termination signal
    if !running.load(Ordering::SeqCst) {
        shutdown_shards(shutdown, &shutdown_receiver);
        return Ok(());
    }
1242

OlivierDehaene's avatar
OlivierDehaene committed
1243
1244
1245
1246
1247
    let mut webserver =
        spawn_webserver(args, shutdown.clone(), &shutdown_receiver).map_err(|err| {
            shutdown_shards(shutdown.clone(), &shutdown_receiver);
            err
        })?;
1248
1249
1250
1251
1252

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

    while running.load(Ordering::SeqCst) {
1253
        if let Ok(ShardStatus::Failed(rank)) = status_receiver.try_recv() {
OlivierDehaene's avatar
OlivierDehaene committed
1254
            tracing::error!("Shard {rank} crashed");
1255
1256
1257
1258
            exit_code = Err(LauncherError::ShardFailed);
            break;
        };

1259
        match webserver.try_wait().unwrap() {
1260
1261
1262
1263
1264
1265
1266
1267
1268
            Some(_) => {
                tracing::error!("Webserver Crashed");
                shutdown_shards(shutdown, &shutdown_receiver);
                return Err(LauncherError::WebserverFailed);
            }
            None => {
                sleep(Duration::from_millis(100));
            }
        };
1269
    }
1270
1271

    // Graceful termination
OlivierDehaene's avatar
OlivierDehaene committed
1272
    terminate("webserver", webserver, Duration::from_secs(90)).unwrap();
1273
1274
1275
    shutdown_shards(shutdown, &shutdown_receiver);

    exit_code
1276
}