main.rs 12 KB
Newer Older
1
/// Text Generation Inference webserver entrypoint
2
use axum::http::HeaderValue;
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
3
use clap::Parser;
4
5
6
7
8
9
use opentelemetry::sdk::propagation::TraceContextPropagator;
use opentelemetry::sdk::trace;
use opentelemetry::sdk::trace::Sampler;
use opentelemetry::sdk::Resource;
use opentelemetry::{global, KeyValue};
use opentelemetry_otlp::WithExportConfig;
Olivier Dehaene's avatar
Olivier Dehaene committed
10
use std::net::{IpAddr, Ipv4Addr, SocketAddr};
11
use std::path::Path;
12
use std::time::Duration;
13
use text_generation_client::{ClientError, ShardedClient};
14
use text_generation_router::{server, HubModelInfo};
15
use thiserror::Error;
16
use tokenizers::{FromPretrainedParameters, Tokenizer};
17
use tower_http::cors::AllowOrigin;
18
19
20
use tracing_subscriber::layer::SubscriberExt;
use tracing_subscriber::util::SubscriberInitExt;
use tracing_subscriber::{EnvFilter, Layer};
Olivier Dehaene's avatar
Olivier Dehaene committed
21
22
23
24
25

/// App Configuration
#[derive(Parser, Debug)]
#[clap(author, version, about, long_about = None)]
struct Args {
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
26
27
    #[clap(default_value = "128", long, env)]
    max_concurrent_requests: usize,
28
29
    #[clap(default_value = "2", long, env)]
    max_best_of: usize,
30
31
    #[clap(default_value = "4", long, env)]
    max_stop_sequences: usize,
32
    #[clap(default_value = "1024", long, env)]
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
33
    max_input_length: usize,
34
    #[clap(default_value = "2048", long, env)]
35
    max_total_tokens: usize,
36
37
    #[clap(default_value = "1.2", long, env)]
    waiting_served_ratio: f32,
38
39
    #[clap(default_value = "4096", long, env)]
    max_batch_prefill_tokens: u32,
40
    #[clap(default_value = "16000", long, env)]
41
    max_batch_total_tokens: u32,
42
43
    #[clap(default_value = "20", long, env)]
    max_waiting_tokens: usize,
44
45
    #[clap(default_value = "0.0.0.0", long, env)]
    hostname: String,
Olivier Dehaene's avatar
Olivier Dehaene committed
46
47
    #[clap(default_value = "3000", long, short, env)]
    port: u16,
48
    #[clap(default_value = "/tmp/text-generation-server-0", long, env)]
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
49
    master_shard_uds_path: String,
Olivier Dehaene's avatar
Olivier Dehaene committed
50
51
    #[clap(default_value = "bigscience/bloom", long, env)]
    tokenizer_name: String,
52
53
    #[clap(default_value = "main", long, env)]
    revision: String,
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
54
55
    #[clap(default_value = "2", long, env)]
    validation_workers: usize,
56
57
    #[clap(long, env)]
    json_output: bool,
58
59
    #[clap(long, env)]
    otlp_endpoint: Option<String>,
60
61
    #[clap(long, env)]
    cors_allow_origin: Option<Vec<String>>,
62
63
64
65
66
67
68
69
70
71
    #[clap(long, env)]
    ngrok: bool,
    #[clap(long, env)]
    ngrok_authtoken: Option<String>,
    #[clap(long, env)]
    ngrok_domain: Option<String>,
    #[clap(long, env)]
    ngrok_username: Option<String>,
    #[clap(long, env)]
    ngrok_password: Option<String>,
Olivier Dehaene's avatar
Olivier Dehaene committed
72
}
Olivier Dehaene's avatar
Init  
Olivier Dehaene committed
73

74
fn main() -> Result<(), RouterError> {
Olivier Dehaene's avatar
Olivier Dehaene committed
75
76
    // Get args
    let args = Args::parse();
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
77
    // Pattern match configuration
Olivier Dehaene's avatar
Olivier Dehaene committed
78
    let Args {
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
79
        max_concurrent_requests,
80
        max_best_of,
81
        max_stop_sequences,
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
82
        max_input_length,
83
        max_total_tokens,
84
        waiting_served_ratio,
85
86
        max_batch_prefill_tokens,
        max_batch_total_tokens,
87
        max_waiting_tokens,
88
        hostname,
Olivier Dehaene's avatar
Olivier Dehaene committed
89
        port,
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
90
        master_shard_uds_path,
Olivier Dehaene's avatar
Olivier Dehaene committed
91
        tokenizer_name,
92
        revision,
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
93
        validation_workers,
94
        json_output,
95
        otlp_endpoint,
96
        cors_allow_origin,
97
98
99
100
101
        ngrok,
        ngrok_authtoken,
        ngrok_domain,
        ngrok_username,
        ngrok_password,
Olivier Dehaene's avatar
Olivier Dehaene committed
102
103
    } = args;

104
105
106
107
108
109
110
111
112
113
    // Validate args
    if max_input_length as u32 > max_batch_prefill_tokens {
        panic!("{}", format!("`max_batch_prefill_tokens` must be >= `max_input_length`. Given: {max_batch_prefill_tokens} and {max_input_length}"));
    }
    if max_batch_prefill_tokens > max_batch_total_tokens {
        panic!("{}", format!("`max_batch_prefill_tokens` must be <= `max_batch_total_tokens`. Given: {max_batch_prefill_tokens} and {max_batch_total_tokens}"));
    }
    if max_total_tokens as u32 > max_batch_total_tokens {
        panic!("{}", format!("`max_total_tokens` must be <= `max_batch_total_tokens`. Given: {max_total_tokens} and {max_batch_total_tokens}"));
    }
114
    if validation_workers == 0 {
115
        panic!("`validation_workers` must be > 0");
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
116
117
    }

118
119
120
121
122
123
124
125
126
127
128
    // CORS allowed origins
    // map to go inside the option and then map to parse from String to HeaderValue
    // Finally, convert to AllowOrigin
    let cors_allow_origin: Option<AllowOrigin> = cors_allow_origin.map(|cors_allow_origin| {
        AllowOrigin::list(
            cors_allow_origin
                .iter()
                .map(|origin| origin.parse::<HeaderValue>().unwrap()),
        )
    });

129
130
131
    // Parse Huggingface hub token
    let authorization_token = std::env::var("HUGGING_FACE_HUB_TOKEN").ok();

132
    // Tokenizer instance
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
133
    // This will only be used to validate payloads
134
    let local_path = Path::new(&tokenizer_name);
135
136
137
138
139
140
141
142
143
    let local_model = local_path.exists() && local_path.is_dir();
    let tokenizer = if local_model {
        // Load local tokenizer
        Tokenizer::from_file(local_path.join("tokenizer.json")).ok()
    } else {
        // Download and instantiate tokenizer
        // We need to download it outside of the Tokio runtime
        let params = FromPretrainedParameters {
            revision: revision.clone(),
144
            auth_token: authorization_token.clone(),
145
            ..Default::default()
146
        };
147
148
        Tokenizer::from_pretrained(tokenizer_name.clone(), Some(params)).ok()
    };
Olivier Dehaene's avatar
Olivier Dehaene committed
149

Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
150
    // Launch Tokio runtime
Olivier Dehaene's avatar
Olivier Dehaene committed
151
152
    tokio::runtime::Builder::new_multi_thread()
        .enable_all()
153
        .build()?
Olivier Dehaene's avatar
Olivier Dehaene committed
154
        .block_on(async {
OlivierDehaene's avatar
OlivierDehaene committed
155
156
            init_logging(otlp_endpoint, json_output);

157
158
159
160
161
162
163
            if tokenizer.is_none() {
                tracing::warn!(
                    "Could not find a fast tokenizer implementation for {tokenizer_name}"
                );
                tracing::warn!("Rust input length validation and truncation is disabled");
            }

164
165
            // Get Model info
            let model_info = match local_model {
166
                true => HubModelInfo {
167
168
169
170
                    model_id: tokenizer_name.clone(),
                    sha: None,
                    pipeline_tag: None,
                },
171
172
173
174
175
176
177
178
179
180
                false => get_model_info(&tokenizer_name, &revision, authorization_token)
                    .await
                    .unwrap_or_else(|| {
                        tracing::warn!("Could not retrieve model info from the Hugging Face hub.");
                        HubModelInfo {
                            model_id: tokenizer_name.to_string(),
                            sha: None,
                            pipeline_tag: None,
                        }
                    }),
181
            };
182
183

            // if pipeline-tag == text-generation we default to return_full_text = true
184
            let compat_return_full_text = match &model_info.pipeline_tag {
185
186
187
188
                None => {
                    tracing::warn!("no pipeline tag found for model {tokenizer_name}");
                    false
                }
189
                Some(pipeline_tag) => pipeline_tag.as_str() == "text-generation",
190
191
            };

Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
192
            // Instantiate sharded client from the master unix socket
193
            let mut sharded_client = ShardedClient::connect_uds(master_shard_uds_path)
Olivier Dehaene's avatar
Olivier Dehaene committed
194
                .await
195
                .map_err(RouterError::Connection)?;
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
196
            // Clear the cache; useful if the webserver rebooted
Olivier Dehaene's avatar
Olivier Dehaene committed
197
            sharded_client
198
                .clear_cache(None)
Olivier Dehaene's avatar
Olivier Dehaene committed
199
                .await
200
                .map_err(RouterError::Cache)?;
201
            // Get info from the shard
202
            let shard_info = sharded_client.info().await.map_err(RouterError::Info)?;
203
204
205
206
207
208
209
210
211
212

            // Warmup model
            tracing::info!("Warming up model");
            sharded_client
                .warmup(
                    max_input_length as u32,
                    max_batch_prefill_tokens,
                    max_batch_total_tokens,
                )
                .await
213
                .map_err(RouterError::Warmup)?;
Olivier Dehaene's avatar
Olivier Dehaene committed
214
            tracing::info!("Connected");
Olivier Dehaene's avatar
Init  
Olivier Dehaene committed
215

216
217
218
219
220
221
222
            let addr = match hostname.parse() {
                Ok(ip) => SocketAddr::new(ip, port),
                Err(_) => {
                    tracing::warn!("Invalid hostname, defaulting to 0.0.0.0");
                    SocketAddr::new(IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0)), port)
                }
            };
Olivier Dehaene's avatar
Init  
Olivier Dehaene committed
223

Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
224
225
            // Run server
            server::run(
226
                model_info,
227
                shard_info,
228
                compat_return_full_text,
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
229
                max_concurrent_requests,
230
                max_best_of,
231
                max_stop_sequences,
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
232
                max_input_length,
233
                max_total_tokens,
234
                waiting_served_ratio,
235
                max_batch_prefill_tokens,
236
                max_batch_total_tokens,
237
                max_waiting_tokens,
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
238
239
240
241
                sharded_client,
                tokenizer,
                validation_workers,
                addr,
242
                cors_allow_origin,
243
244
245
246
247
                ngrok,
                ngrok_authtoken,
                ngrok_domain,
                ngrok_username,
                ngrok_password,
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
248
            )
249
            .await?;
Olivier Dehaene's avatar
Olivier Dehaene committed
250
            Ok(())
Olivier Dehaene's avatar
Olivier Dehaene committed
251
        })
Olivier Dehaene's avatar
Init  
Olivier Dehaene committed
252
}
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307

/// Init logging using env variables LOG_LEVEL and LOG_FORMAT:
///     - otlp_endpoint is an optional URL to an Open Telemetry collector
///     - LOG_LEVEL may be TRACE, DEBUG, INFO, WARN or ERROR (default to INFO)
///     - LOG_FORMAT may be TEXT or JSON (default to TEXT)
fn init_logging(otlp_endpoint: Option<String>, json_output: bool) {
    let mut layers = Vec::new();

    // STDOUT/STDERR layer
    let fmt_layer = tracing_subscriber::fmt::layer()
        .with_file(true)
        .with_line_number(true);

    let fmt_layer = match json_output {
        true => fmt_layer.json().flatten_event(true).boxed(),
        false => fmt_layer.boxed(),
    };
    layers.push(fmt_layer);

    // OpenTelemetry tracing layer
    if let Some(otlp_endpoint) = otlp_endpoint {
        global::set_text_map_propagator(TraceContextPropagator::new());

        let tracer = opentelemetry_otlp::new_pipeline()
            .tracing()
            .with_exporter(
                opentelemetry_otlp::new_exporter()
                    .tonic()
                    .with_endpoint(otlp_endpoint),
            )
            .with_trace_config(
                trace::config()
                    .with_resource(Resource::new(vec![KeyValue::new(
                        "service.name",
                        "text-generation-inference.router",
                    )]))
                    .with_sampler(Sampler::AlwaysOn),
            )
            .install_batch(opentelemetry::runtime::Tokio);

        if let Ok(tracer) = tracer {
            layers.push(tracing_opentelemetry::layer().with_tracer(tracer).boxed());
            axum_tracing_opentelemetry::init_propagator().unwrap();
        };
    }

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

    tracing_subscriber::registry()
        .with(env_filter)
        .with(layers)
        .init();
}
308
309

/// get model info from the Huggingface Hub
310
311
312
313
314
pub async fn get_model_info(
    model_id: &str,
    revision: &str,
    token: Option<String>,
) -> Option<HubModelInfo> {
315
    let client = reqwest::Client::new();
316
    // Poor man's urlencode
317
    let revision = revision.replace('/', "%2F");
318
    let url = format!("https://huggingface.co/api/models/{model_id}/revision/{revision}");
319
    let mut builder = client.get(url).timeout(Duration::from_secs(5));
320
321
322
323
    if let Some(token) = token {
        builder = builder.bearer_auth(token);
    }

324
325
326
327
328
329
    let response = builder.send().await.ok()?;

    if response.status().is_success() {
        return serde_json::from_str(&response.text().await.ok()?).ok();
    }
    None
330
}
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346

#[derive(Debug, Error)]
enum RouterError {
    #[error("Unable to connect to the Python model shards: {0}")]
    Connection(ClientError),
    #[error("Unable to clear the Python model shards cache: {0}")]
    Cache(ClientError),
    #[error("Unable to get the Python model shards info: {0}")]
    Info(ClientError),
    #[error("Unable to warmup the Python model shards: {0}")]
    Warmup(ClientError),
    #[error("Tokio runtime failed to start: {0}")]
    Tokio(#[from] std::io::Error),
    #[error("Axum webserver failed: {0}")]
    Axum(#[from] axum::BoxError),
}