huggingface/text-generation-inference
5
0
Fork 0
mirror of https://github.com/huggingface/text-generation-inference.git synced 2025-09-11 12:24:53 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix: dynamically set routes when using google feature

Browse Source
This commit is contained in:
drbh 2024-01-26 11:49:05 -05:00
parent 5e38c4bfda
commit 7abb5ae4fa
1 changed files with 16 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

19
router/src/server.rs
View File

@ -1041,7 +1041,6 @@ pub async fn run(
.route("/generate", post(generate))
.route("/generate_stream", post(generate_stream))
.route("/v1/chat/completions", post(chat_completions))
.route("/vertex", post(vertex_compatibility))
.route("/tokenize", post(tokenize))
.route("/health", get(health))
.route("/ping", get(health))
@ -1058,10 +1057,24 @@ pub async fn run(
ComputeType(std::env::var("COMPUTE_TYPE").unwrap_or("gpu+optimized".to_string()));
// Combine routes and layers
let app = Router::new()
let mut app = Router::new()
.merge(swagger_ui)
.merge(base_routes)
.merge(aws_sagemaker_route)
.merge(aws_sagemaker_route);
if cfg!(feature = "google") {
// in the google feature case throw and error if any of the env vars are not set
let env_predict_route = std::env::var("AIP_PREDICT_ROUTE")
.expect("AIP_PREDICT_ROUTE must be set when building with `google` feature");
app = app.route(&env_predict_route, post(vertex_compatibility));
let env_health_route = std::env::var("AIP_HEALTH_ROUTE")
.expect("AIP_HEALTH_ROUTE must be set when building with `google` feature");
app = app.route(&env_health_route, get(health));
}
// add layers after routes
app = app
.layer(Extension(info))
.layer(Extension(health_ext.clone()))
.layer(Extension(compat_return_full_text))