diff --git a/router/src/grpc/server.rs b/router/src/grpc/server.rs index df76275e..913455b3 100644 --- a/router/src/grpc/server.rs +++ b/router/src/grpc/server.rs @@ -1415,6 +1415,7 @@ pub async fn run( }; tracing::info!("Starting gRPC server: {}", &addr); + tracing::info!("Ready"); server.await?; Ok(()) diff --git a/router/src/http/server.rs b/router/src/http/server.rs index d7dc1a6c..6c48ad3e 100644 --- a/router/src/http/server.rs +++ b/router/src/http/server.rs @@ -1479,12 +1479,16 @@ pub async fn run( #[cfg(feature = "google")] { tracing::info!("Built with `google` feature"); - let env_predict_route = std::env::var("AIP_PREDICT_ROUTE") - .context("`AIP_PREDICT_ROUTE` env var must be set for Google Vertex deployments")?; - app = app.route(&env_predict_route, post(vertex_compatibility)); - let env_health_route = std::env::var("AIP_HEALTH_ROUTE") - .context("`AIP_HEALTH_ROUTE` env var must be set for Google Vertex deployments")?; - app = app.route(&env_health_route, get(health)); + + if let Ok(env_predict_route) = std::env::var("AIP_PREDICT_ROUTE") { + tracing::info!("Serving Vertex compatible route on {env_predict_route}"); + app = app.route(&env_predict_route, post(vertex_compatibility)); + } + + if let Ok(env_health_route) = std::env::var("AIP_HEALTH_ROUTE") { + tracing::info!("Serving Vertex compatible health route on {env_health_route}"); + app = app.route(&env_health_route, get(health)); + } } #[cfg(not(feature = "google"))] { @@ -1546,6 +1550,9 @@ pub async fn run( // Run server let listener = tokio::net::TcpListener::bind(&addr).await.unwrap(); + tracing::info!("Starting HTTP server: {}", &addr); + tracing::info!("Ready"); + axum::serve(listener, app) // Wait until all requests are finished to shut down .with_graceful_shutdown(shutdown::shutdown_signal()) diff --git a/router/src/lib.rs b/router/src/lib.rs index df0fad4b..3801af8a 100644 --- a/router/src/lib.rs +++ b/router/src/lib.rs @@ -248,7 +248,11 @@ pub async fn run( std::env::var("AIP_HTTP_PORT") .ok() .and_then(|p| p.parse().ok()) - .context("`AIP_HTTP_PORT` env var must be set for Google Vertex deployments")? + .map(|p| { + tracing::info!("`AIP_HTTP_PORT` is set: overriding port {port} by port {p}"); + p + }) + .unwrap_or(port) } else { port }; @@ -274,20 +278,16 @@ pub async fn run( #[cfg(feature = "http")] { - let server = tokio::spawn(async move { - http::server::run( - infer, - info, - addr, - prom_builder, - payload_limit, - api_key, - cors_allow_origin, - ) - .await - }); - tracing::info!("Ready"); - server.await??; + http::server::run( + infer, + info, + addr, + prom_builder, + payload_limit, + api_key, + cors_allow_origin, + ) + .await?; } #[cfg(feature = "grpc")] @@ -295,11 +295,7 @@ pub async fn run( // cors_allow_origin and payload_limit are not used for gRPC servers let _ = cors_allow_origin; let _ = payload_limit; - let server = tokio::spawn(async move { - grpc::server::run(infer, info, addr, prom_builder, api_key).await - }); - tracing::info!("Ready"); - server.await??; + grpc::server::run(infer, info, addr, prom_builder, api_key).await?; } Ok(())