fix(gke): accept null values for vertex env vars (huggingface#243)

Author: OlivierDehaene

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(())

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())

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,32 +278,24 @@ 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")]
     {
         // 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(())