Apply rustfmt, clippy suggestions, update .gitignore

declark1 · prashantgupta24 · commit 506fdf0c2274 · 2024-03-28T13:43:14.000-07:00
diff --git a/fmaas-router/src/lib.rs b/fmaas-router/src/lib.rs
@@ -1,5 +1,5 @@
-extern crate lazy_static;
 extern crate core;
+extern crate lazy_static;
 
+mod pb;
 pub mod server;
-mod pb;
diff --git a/fmaas-router/src/main.rs b/fmaas-router/src/main.rs
@@ -1,7 +1,8 @@
-use clap::Parser;
 use std::net::{IpAddr, Ipv4Addr, SocketAddr};
-use tracing_subscriber::{EnvFilter};
+
+use clap::Parser;
 use fmaas_router::server;
+use tracing_subscriber::EnvFilter;
 
 /// App Configuration
 #[derive(Parser, Debug)]
@@ -27,7 +28,6 @@ struct Args {
     upstream_tls_ca_cert_path: Option<String>,
 }
 
-
 fn main() -> Result<(), std::io::Error> {
     //Get args
     let args = Args::parse();
@@ -66,22 +66,20 @@ fn main() -> Result<(), std::io::Error> {
         .block_on(async {
             //TODO initialize clients
 
-            let grpc_addr = SocketAddr::new(
-                IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0)), args.grpc_port
-            );
-            let http_addr = SocketAddr::new(
-                IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0)), args.probe_port
-            );
+            let grpc_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0)), args.grpc_port);
+            let http_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0)), args.probe_port);
 
             server::run(
                 grpc_addr,
                 http_addr,
-                args.tls_cert_path.map(|cp| (cp, args.tls_key_path.unwrap())),
+                args.tls_cert_path
+                    .map(|cp| (cp, args.tls_key_path.unwrap())),
                 args.tls_client_ca_cert_path,
                 args.default_upstream_port,
                 args.upstream_tls,
-                args.upstream_tls_ca_cert_path
-            ).await;
+                args.upstream_tls_ca_cert_path,
+            )
+            .await;
 
             Ok(())
         })
diff --git a/fmaas-router/src/server.rs b/fmaas-router/src/server.rs
@@ -1,27 +1,24 @@
-use std::collections::HashMap;
-use std::net::SocketAddr;
-use std::time::Duration;
-use anyhow::Context;
+use std::{collections::HashMap, net::SocketAddr, time::Duration};
 
-use axum::Router;
-use axum::routing::get;
+use anyhow::Context;
+use axum::{routing::get, Router};
 use futures::future::try_join_all;
 use ginepro::LoadBalancedChannel;
 use lazy_static::lazy_static;
-use tokio::fs::read;
-use tokio::signal;
-use tokio::time::sleep;
-use tonic::{Request, Response, Status, Streaming};
-use tonic::transport::{Certificate, ClientTlsConfig, Identity, Server, ServerTlsConfig};
+use tokio::{fs::read, signal, time::sleep};
+use tonic::{
+    transport::{Certificate, ClientTlsConfig, Identity, Server, ServerTlsConfig},
+    Request, Response, Status, Streaming,
+};
 use tracing::instrument;
+
 use crate::pb::fmaas::{
-    BatchedGenerationRequest, BatchedGenerationResponse, GenerationResponse,
-    SingleGenerationRequest, BatchedTokenizeRequest, BatchedTokenizeResponse,
-    ModelInfoRequest, ModelInfoResponse,
+    generation_service_client::GenerationServiceClient,
+    generation_service_server::{GenerationService, GenerationServiceServer},
+    BatchedGenerationRequest, BatchedGenerationResponse, BatchedTokenizeRequest,
+    BatchedTokenizeResponse, GenerationResponse, ModelInfoRequest, ModelInfoResponse,
+    SingleGenerationRequest,
 };
-use crate::pb::fmaas::generation_service_client::GenerationServiceClient;
-use crate::pb::fmaas::generation_service_server::{GenerationService, GenerationServiceServer};
-
 
 const MODEL_MAP_ENV_VAR_NAME: &str = "MODEL_MAP_CONFIG";
 
@@ -40,8 +37,8 @@ lazy_static! {
                 }
             };
         }
-        let map: HashMap<&'static str, &'static str> = serde_yaml::from_str(&MODEL_MAP_STR)
-            .expect("Failed to parse the model mapping config");
+        let map: HashMap<&'static str, &'static str> =
+            serde_yaml::from_str(&MODEL_MAP_STR).expect("Failed to parse the model mapping config");
         tracing::info!("{} model mappings configured", map.len());
         map
     };
@@ -54,7 +51,7 @@ pub async fn run(
     tls_client_ca_cert: Option<String>,
     default_target_port: u16,
     upstream_tls: bool,
-    upstream_tls_ca_cert: Option<String>
+    upstream_tls_ca_cert: Option<String>,
 ) {
     let mut builder = Server::builder();
 
@@ -82,40 +79,47 @@ pub async fn run(
             let ca_cert_pem = load_pem(ca_cert_path, "client ca cert").await;
             tls_config = tls_config.client_ca_root(Certificate::from_pem(ca_cert_pem));
         }
-        builder = builder.tls_config(tls_config).expect("tls configuration error");
+        builder = builder
+            .tls_config(tls_config)
+            .expect("tls configuration error");
     } else if upstream_tls {
         panic!("Upstream TLS enabled without any certificates");
     }
 
     // Set up clients
-    let clients = try_join_all(
-        MODEL_MAP.iter().map(|(model, service)| async {
-            tracing::info!("Configuring client for model name: [{}]", *model);
-            // Parse hostname and optional port from target service name
-            let mut service_parts = service.split(":");
-            let hostname = service_parts.next().unwrap();
-            let port = service_parts.next().map_or(
-                default_target_port,
-                |p| p.parse::<u16>().expect(
-                    &*format!("Invalid port in configured service name: {}", p)
-                ),
-            );
-            if service_parts.next().is_some() {
-                panic!("Configured service name contains more than one : character");
-            }
-            // Build a load-balanced channel given a service name and a port.
-            let mut builder = LoadBalancedChannel::builder((hostname, port));
-                //.dns_probe_interval(Duration::from_secs(10))
-            if let Some(tls_config) = &client_tls {
-                builder = builder.with_tls(tls_config.clone());
-            }
-            let channel = builder.channel().await
-                .context(format!("Channel failed for service {}", *service))?;
-            Ok((*model, GenerationServiceClient::new(channel))) as Result<
-                (&'static str, GenerationServiceClient<LoadBalancedChannel>), anyhow::Error
-            >
-    })).await.expect("Error creating upstream service clients").into_iter().collect();
-    tracing::info!("{} upstream gRPC clients created successfully", grpc_addr.port());
+    let clients = try_join_all(MODEL_MAP.iter().map(|(model, service)| async {
+        tracing::info!("Configuring client for model name: [{}]", *model);
+        // Parse hostname and optional port from target service name
+        let mut service_parts = service.split(':');
+        let hostname = service_parts.next().unwrap();
+        let port = service_parts.next().map_or(default_target_port, |p| {
+            p.parse::<u16>()
+                .unwrap_or_else(|_| panic!("Invalid port in configured service name: {}", p))
+        });
+        if service_parts.next().is_some() {
+            panic!("Configured service name contains more than one : character");
+        }
+        // Build a load-balanced channel given a service name and a port.
+        let mut builder = LoadBalancedChannel::builder((hostname, port));
+        //.dns_probe_interval(Duration::from_secs(10))
+        if let Some(tls_config) = &client_tls {
+            builder = builder.with_tls(tls_config.clone());
+        }
+        let channel = builder
+            .channel()
+            .await
+            .context(format!("Channel failed for service {}", *service))?;
+        Ok((*model, GenerationServiceClient::new(channel)))
+            as Result<(&'static str, GenerationServiceClient<LoadBalancedChannel>), anyhow::Error>
+    }))
+    .await
+    .expect("Error creating upstream service clients")
+    .into_iter()
+    .collect();
+    tracing::info!(
+        "{} upstream gRPC clients created successfully",
+        grpc_addr.port()
+    );
 
     // Build and start gRPC server in background task
     let grpc_service = GenerationServicer { clients };
@@ -132,13 +136,15 @@ pub async fn run(
     // fail before starting
     sleep(Duration::from_secs(2)).await;
     if grpc_server_handle.is_finished() {
-        grpc_server_handle.await.unwrap().expect("gRPC server startup failed");
+        grpc_server_handle
+            .await
+            .unwrap()
+            .expect("gRPC server startup failed");
         panic!(); // should not reach here
     }
 
     // Build and await on the HTTP server
-    let app = Router::new()
-      .route("/health", get(health));
+    let app = Router::new().route("/health", get(health));
 
     let server = axum::Server::bind(&http_addr)
         .serve(app.into_make_service())
@@ -147,7 +153,10 @@ pub async fn run(
     tracing::info!("HTTP server started on port {}", http_addr.port());
     server.await.expect("HTTP server crashed!");
 
-    grpc_server_handle.await.unwrap().expect("gRPC server crashed");
+    grpc_server_handle
+        .await
+        .unwrap()
+        .expect("gRPC server crashed");
 }
 
 async fn health() -> &'static str {
@@ -156,42 +165,47 @@ async fn health() -> &'static str {
 }
 
 async fn load_pem(path: String, name: &str) -> Vec<u8> {
-    read(&path).await.expect(&*format!("couldn't load {name} from {path}"))
+    read(&path)
+        .await
+        .unwrap_or_else(|_| panic!("couldn't load {name} from {path}"))
 }
 
 /*
   TODO:
   - Log errors/timings
 */
 
-
 #[derive(Debug, Default)]
 pub struct GenerationServicer {
-    clients: HashMap<&'static str, GenerationServiceClient<LoadBalancedChannel>>
+    clients: HashMap<&'static str, GenerationServiceClient<LoadBalancedChannel>>,
 }
 
 impl GenerationServicer {
-    async fn client(&self, model_id: &String)
-        -> Result<GenerationServiceClient<LoadBalancedChannel>, Status> {
-        Ok(
-            self.clients.get(&**model_id)
-                .ok_or_else(|| Status::not_found(
-                    format!("Unrecognized model_id: {model_id}"))
-                )?
-                .clone()
-        )
+    async fn client(
+        &self,
+        model_id: &String,
+    ) -> Result<GenerationServiceClient<LoadBalancedChannel>, Status> {
+        Ok(self
+            .clients
+            .get(&**model_id)
+            .ok_or_else(|| Status::not_found(format!("Unrecognized model_id: {model_id}")))?
+            .clone())
     }
 }
 
 #[tonic::async_trait]
 impl GenerationService for GenerationServicer {
     #[instrument(skip_all)]
-    async fn generate(&self, request: Request<BatchedGenerationRequest>)
-        -> Result<Response<BatchedGenerationResponse>, Status> {
+    async fn generate(
+        &self,
+        request: Request<BatchedGenerationRequest>,
+    ) -> Result<Response<BatchedGenerationResponse>, Status> {
         //let start_time = Instant::now();
         let br = request.get_ref();
         if br.requests.is_empty() {
-            return Ok(Response::new(BatchedGenerationResponse{ responses: vec![] }));
+            return Ok(Response::new(BatchedGenerationResponse {
+                responses: vec![],
+            }));
         }
         tracing::debug!("Routing generation request for Model ID {}", &br.model_id);
         self.client(&br.model_id).await?.generate(request).await
@@ -200,32 +214,50 @@ impl GenerationService for GenerationServicer {
     type GenerateStreamStream = Streaming<GenerationResponse>;
 
     #[instrument(skip_all)]
-    async fn generate_stream(&self, request: Request<SingleGenerationRequest>)
-        -> Result<Response<Self::GenerateStreamStream>, Status> {
+    async fn generate_stream(
+        &self,
+        request: Request<SingleGenerationRequest>,
+    ) -> Result<Response<Self::GenerateStreamStream>, Status> {
         let sr = request.get_ref();
         if sr.request.is_none() {
             return Err(Status::invalid_argument("missing request"));
         }
-        tracing::debug!("Routing streaming generation request for Model ID {}", &sr.model_id);
-        self.client(&sr.model_id).await?.generate_stream(request).await
+        tracing::debug!(
+            "Routing streaming generation request for Model ID {}",
+            &sr.model_id
+        );
+        self.client(&sr.model_id)
+            .await?
+            .generate_stream(request)
+            .await
     }
 
     #[instrument(skip_all)]
-    async fn tokenize(&self, request: Request<BatchedTokenizeRequest>)
-        -> Result<Response<BatchedTokenizeResponse>, Status> {
+    async fn tokenize(
+        &self,
+        request: Request<BatchedTokenizeRequest>,
+    ) -> Result<Response<BatchedTokenizeResponse>, Status> {
         let br = request.get_ref();
         if br.requests.is_empty() {
-            return Ok(Response::new(BatchedTokenizeResponse{ responses: vec![] }));
+            return Ok(Response::new(BatchedTokenizeResponse { responses: vec![] }));
         }
         tracing::debug!("Routing tokenization request for Model ID {}", &br.model_id);
         self.client(&br.model_id).await?.tokenize(request).await
     }
 
     #[instrument(skip_all)]
-    async fn model_info(&self, request: Request<ModelInfoRequest>)
-        -> Result<Response<ModelInfoResponse>, Status> {
-        tracing::debug!("Routing model info request for Model ID {}", &request.get_ref().model_id);
-        self.client(&request.get_ref().model_id).await?.model_info(request).await
+    async fn model_info(
+        &self,
+        request: Request<ModelInfoRequest>,
+    ) -> Result<Response<ModelInfoResponse>, Status> {
+        tracing::debug!(
+            "Routing model info request for Model ID {}",
+            &request.get_ref().model_id
+        );
+        self.client(&request.get_ref().model_id)
+            .await?
+            .model_info(request)
+            .await
     }
 }
 
@@ -238,15 +270,15 @@ async fn shutdown_signal() {
     };
 
     #[cfg(unix)]
-        let terminate = async {
+    let terminate = async {
         signal::unix::signal(signal::unix::SignalKind::terminate())
             .expect("failed to install signal handler")
             .recv()
             .await;
     };
 
     #[cfg(not(unix))]
-        let terminate = std::future::pending::<()>();
+    let terminate = std::future::pending::<()>();
 
     tokio::select! {
         _ = ctrl_c => {},
diff --git a/rustfmt.toml b/rustfmt.toml
@@ -0,0 +1,2 @@
+group_imports = "StdExternalCrate"
+imports_granularity = "Crate"

Original file line number	Diff line number	Diff line change
`@@ -0,0 +1,2 @@`
	`1`	`+group_imports = "StdExternalCrate"`
	`2`	`+imports_granularity = "Crate"`