huggingface/text-generation-inference
5
0
Fork 0
mirror of https://github.com/huggingface/text-generation-inference.git synced 2025-09-10 03:44:54 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(router): add ngrok integration

Browse Source
This commit is contained in:
OlivierDehaene 2023-06-14 18:06:33 +02:00
parent 5ce89059f8
commit 51b2d4edca
3 changed files with 240 additions and 9 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

197
Cargo.lock generated
View File

@ -95,6 +95,23 @@ version = "1.0.70"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "7de8ce5e0f9f8d88245311066a578d72b7af3e7088f32783804676302df237e4"
[[package]]
name = "arc-swap"
version = "1.6.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "bddcadddf5e9015d310179a59bb28c4d4b9920ad0f11e8e14dbadf654890c9a6"
[[package]]
name = "async-rustls"
version = "0.3.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "93b21a03b7c21702a0110f9f8d228763a533570deb376119042dabf33c37a01a"
dependencies = [
"futures-io",
"rustls",
"webpki",
]
[[package]]
name = "async-stream"
version = "0.3.5"
@ -145,6 +162,18 @@ dependencies = [
"num-traits",
]
[[package]]
name = "awaitdrop"
version = "0.1.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "771051cdc7eec2dc1b23fbf870bb7fbb89136fe374227c875e377f1eed99a429"
dependencies = [
"futures",
"generational-arena",
"parking_lot",
"slotmap",
]
[[package]]
name = "axum"
version = "0.6.15"
@ -758,7 +787,7 @@ dependencies = [
"futures-sink",
"nanorand",
"pin-project",
"spin",
"spin 0.9.8",
]
[[package]]
@ -890,6 +919,15 @@ dependencies = [
"slab",
]
[[package]]
name = "generational-arena"
version = "0.2.9"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "877e94aff08e743b651baaea359664321055749b398adff8740a7399af7796e7"
dependencies = [
"cfg-if",
]
[[package]]
name = "generic-array"
version = "0.14.7"
@ -987,6 +1025,17 @@ dependencies = [
"digest",
]
[[package]]
name = "hostname"
version = "0.3.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "3c731c3e10504cc8ed35cfe2f1db4c9274c3d35fa486e3b31df46f068ef3e867"
dependencies = [
"libc",
"match_cfg",
"winapi",
]
[[package]]
name = "http"
version = "0.2.9"
@ -1284,6 +1333,12 @@ version = "0.1.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "58093314a45e00c77d5c508f76e77c3396afbbc0d01506e7fae47b018bac2b1d"
[[package]]
name = "match_cfg"
version = "0.1.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "ffbee8634e0d45d258acb448e7eaab3fce7a0a467395d4d9f228e3c1f01fb2e4"
[[package]]
name = "matchers"
version = "0.1.0"
@ -1442,6 +1497,25 @@ version = "0.8.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "e5ce46fe64a9d73be07dcbe690a38ce1b293be448fd8ce1e6c1b8062c9f72c6a"
[[package]]
name = "muxado"
version = "0.4.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "e92b89ac3127251efde6f5a9586e5aae99468d06fcf9f133b377f58d5ed66446"
dependencies = [
"async-trait",
"awaitdrop",
"bitflags",
"bytes",
"futures",
"pin-project",
"rand",
"thiserror",
"tokio",
"tokio-util",
"tracing",
]
[[package]]
name = "nanorand"
version = "0.7.0"
@ -1469,6 +1543,37 @@ dependencies = [
"tempfile",
]
[[package]]
name = "ngrok"
version = "0.12.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "98ce3514eec7338e2d4663e3efb4429e08d8f3662996be4b9585350e7d8ad728"
dependencies = [
"arc-swap",
"async-rustls",
"async-trait",
"awaitdrop",
"axum",
"base64 0.13.1",
"bytes",
"futures",
"hostname",
"hyper",
"muxado",
"once_cell",
"parking_lot",
"regex",
"rustls-pemfile",
"serde",
"serde_json",
"thiserror",
"tokio",
"tokio-retry",
"tokio-util",
"tracing",
"windows-sys 0.45.0",
]
[[package]]
name = "nix"
version = "0.26.2"
@ -2140,6 +2245,21 @@ dependencies = [
"winreg",
]
[[package]]
name = "ring"
version = "0.16.20"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "3053cf52e236a3ed746dfc745aa9cacf1b791d846bdaf412f60a8d7d6e17c8fc"
dependencies = [
"cc",
"libc",
"once_cell",
"spin 0.5.2",
"untrusted",
"web-sys",
"winapi",
]
[[package]]
name = "rust-embed"
version = "6.6.1"
@ -2198,6 +2318,27 @@ dependencies = [
"windows-sys 0.48.0",
]
[[package]]
name = "rustls"
version = "0.20.8"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "fff78fc74d175294f4e83b28343315ffcfb114b156f0185e9741cb5570f50e2f"
dependencies = [
"log",
"ring",
"sct",
"webpki",
]
[[package]]
name = "rustls-pemfile"
version = "1.0.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "d194b56d58803a43635bdc398cd17e383d6f71f9182b9a192c127ca42494a59b"
dependencies = [
"base64 0.21.0",
]
[[package]]
name = "rustversion"
version = "1.0.12"
@ -2234,6 +2375,16 @@ version = "1.1.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "d29ab0c6d3fc0ee92fe66e2d99f700eab17a8d57d1c1d3b748380fb20baa78cd"
[[package]]
name = "sct"
version = "0.7.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "d53dcdb7c9f8158937a7981b48accfd39a43af418591a5d008c7b22b5e1b7ca4"
dependencies = [
"ring",
"untrusted",
]
[[package]]
name = "security-framework"
version = "2.8.2"
@ -2400,6 +2551,15 @@ dependencies = [
"autocfg",
]
[[package]]
name = "slotmap"
version = "1.0.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "e1e08e261d0e8f5c43123b7adf3e4ca1690d655377ac93a03b2c9d3e98de1342"
dependencies = [
"version_check",
]
[[package]]
name = "smallvec"
version = "1.10.0"
@ -2416,6 +2576,12 @@ dependencies = [
"winapi",
]
[[package]]
name = "spin"
version = "0.5.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "6e63cff320ae2c57904679ba7cb63280a3dc4613885beafb148ee7bf9aa9042d"
[[package]]
name = "spin"
version = "0.9.8"
@ -2619,6 +2785,7 @@ dependencies = [
"futures",
"metrics",
"metrics-exporter-prometheus",
"ngrok",
"nohash-hasher",
"opentelemetry",
"opentelemetry-otlp",
@ -2797,6 +2964,17 @@ dependencies = [
"tokio",
]
[[package]]
name = "tokio-retry"
version = "0.3.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "7f57eb36ecbe0fc510036adff84824dd3c24bb781e21bfa67b69d556aa85214f"
dependencies = [
"pin-project",
"rand",
"tokio",
]
[[package]]
name = "tokio-stream"
version = "0.1.12"
@ -2816,6 +2994,7 @@ checksum = "5427d89453009325de0d8f342c9490009f76e999cb7672d77e46267448f7e6b2"
dependencies = [
"bytes",
"futures-core",
"futures-io",
"futures-sink",
"pin-project-lite",
"tokio",
@ -3105,6 +3284,12 @@ version = "0.1.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "39ec24b3121d976906ece63c9daad25b85969647682eee313cb5779fdd69e14e"
[[package]]
name = "untrusted"
version = "0.7.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "a156c684c91ea7d62626509bce3cb4e1d9ed5c4d978f7b4352658f96a4c26b4a"
[[package]]
name = "url"
version = "2.3.1"
@ -3301,6 +3486,16 @@ dependencies = [
"wasm-bindgen",
]
[[package]]
name = "webpki"
version = "0.22.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "f095d78192e208183081cc07bc5515ef55216397af48b873e5edcd72637fa1bd"
dependencies = [
"ring",
"untrusted",
]
[[package]]
name = "which"
version = "4.4.0"

1
router/Cargo.toml
View File

@ -40,6 +40,7 @@ tracing-opentelemetry = "0.18.0"
tracing-subscriber = { version = "0.3.16", features = ["json", "env-filter"] }
utoipa = { version = "3.0.1", features = ["axum_extras"] }
utoipa-swagger-ui = { version = "3.0.2", features = ["axum"] }
ngrok = { version = "0.12.3", features = ["axum"] }
[build-dependencies]
vergen = { version = "8.0.0", features = ["build", "git", "gitcl"] }

37
router/src/server.rs
View File

@ -1,5 +1,5 @@
use crate::health::Health;
/// HTTP Server logic
use crate::health::Health;
use crate::infer::{InferError, InferResponse, InferStreamResponse};
use crate::validation::ValidationError;
use crate::{
@ -17,7 +17,10 @@ use axum_tracing_opentelemetry::opentelemetry_tracing_layer;
use futures::stream::StreamExt;
use futures::Stream;
use metrics_exporter_prometheus::{Matcher, PrometheusBuilder, PrometheusHandle};
use ngrok::config::TunnelBuilder;
use ngrok::tunnel::UrlTunnel;
use std::convert::Infallible;
use std::env;
use std::net::SocketAddr;
use std::sync::atomic::AtomicBool;
use std::sync::Arc;
@ -683,6 +686,37 @@ pub async fn run(
.layer(opentelemetry_tracing_layer())
.layer(cors_layer);
if let Ok(ngrok_token) = env::var("NGROK_AUTHTOKEN") {
let _ = addr;
let mut tunnel = ngrok::Session::builder()
.authtoken(ngrok_token)
.connect()
.await
.unwrap()
.http_endpoint();
if let Ok(domain) = env::var("NGROK_DOMAIN") {
tunnel = tunnel.domain(domain);
}
if let (Ok(username), Ok(password)) =
(env::var("NGROK_USERNAME"), env::var("NGROK_PASSWORD"))
{
tunnel = tunnel.basic_auth(username, password);
}
let listener = tunnel.listen().await.unwrap();
// Run server
tracing::info!("Ingress URL: {:?}", listener.url());
axum::Server::builder(listener)
.serve(app.into_make_service())
//Wait until all requests are finished to shut down
.with_graceful_shutdown(shutdown_signal())
.await
.unwrap();
} else {
// Run server
axum::Server::bind(&addr)
.serve(app.into_make_service())
@ -691,6 +725,7 @@ pub async fn run(
.await
.unwrap();
}
}
/// Shutdown signal handler
async fn shutdown_signal() {