From c52d28c6699cdeebf6803ecafdf98a9bf2bfc87f Mon Sep 17 00:00:00 2001 From: Abel Tesfa Date: Wed, 29 Jul 2026 09:54:04 +0300 Subject: [PATCH 1/6] refactor: updated base to accept env for including docker --- src/bin/client.rs | 2 +- src/bin/worker.rs | 2 +- src/main.rs | 24 ++++++++++++++++-------- 3 files changed, 18 insertions(+), 10 deletions(-) diff --git a/src/bin/client.rs b/src/bin/client.rs index c807592..3e457db 100644 --- a/src/bin/client.rs +++ b/src/bin/client.rs @@ -9,7 +9,7 @@ use inference::PredictRequest; async fn main() -> Result<(), Box> { println!("Connecting to the Inference Proxy server..."); - let mut client = InferenceServiceClient::connect("http://[::1]:50051").await?; + let mut client = InferenceServiceClient::connect("http://0.0.0.0:50051").await?; // dummy data to be sent to the server for inference let my_data = vec![1.0, 2.0, 3.0, 4.0]; diff --git a/src/bin/worker.rs b/src/bin/worker.rs index 08a2f92..f5e0a29 100644 --- a/src/bin/worker.rs +++ b/src/bin/worker.rs @@ -42,7 +42,7 @@ async fn main() -> Result<(), Box> { "50052".to_string() }; - let addr_str = format!("[::1]:{}", port); + let addr_str = format!("0.0.0.0:{}", port); let addr = addr_str.parse()?; let worker = GpuWorker::default(); diff --git a/src/main.rs b/src/main.rs index af7a0a3..2d7065f 100644 --- a/src/main.rs +++ b/src/main.rs @@ -25,29 +25,37 @@ async fn main() -> Result<(), Box> { info!("Metrics server running on http://localhost:9090/metrics"); - let addr = "[::1]:50051".parse()?; + // Read from environment variable, fallback to localhost for local dev + let addr = std::env::var("PROXY_LISTEN_ADDR") + .unwrap_or_else(|_| "[::1]:50051".to_string()) + .parse()?; info!("Attempting to connect to worker nodes..."); - let worker_ports = vec!["50052", "50053", "50054"]; + // Read worker addresses from env var (comma-separated), fallback for local dev + let worker_addrs_str = std::env::var("WORKER_ADDRESSES") + .unwrap_or_else(|_| "http://[::1]:50052,http://[::1]:50053,http://[::1]:50054".to_string()); + + let worker_addrs: Vec<&str> = worker_addrs_str.split(',').collect(); let mut worker_clients = Vec::new(); - for port in worker_ports { - let worker_addr = format!("http://[::1]:{}", port); - match InferenceServiceClient::connect(worker_addr).await { + for addr_str in worker_addrs { + let addr_str = addr_str.trim(); + match InferenceServiceClient::connect(addr_str.to_string()).await { Ok(client) => { - info!("Connected to worker node on port {}", port); + info!("Connected to worker node at {}", addr_str); worker_clients.push(client); } Err(e) => { - info!("Failed to connect to worker node on port {}: {}", port, e); + warn!("Failed to connect to worker at {}: {}", addr_str, e); } } } // safety check: if all workers are offline, crash the proxy if worker_clients.is_empty() { - panic!("FATAL: No worker nodes are available. Exiting..."); + error!("FATAL: No worker nodes are available. Exiting..."); + panic!("No workers online"); } // initialize mpsc channel for batching let (tx, rx) = mpsc::channel(100); From 2dc90c03286d2fe5cb77d145dbfc183ddf27c2fd Mon Sep 17 00:00:00 2001 From: Abel Tesfa Date: Wed, 29 Jul 2026 10:30:11 +0300 Subject: [PATCH 2/6] feat: dockerized inference proxy and worker nodes --- .dockerignore | 2 ++ Dockerfile | 17 +++++++++++++++++ docker-compose.yml | 32 ++++++++++++++++++++++++++++++++ src/bin/worker.rs | 6 ++++-- 4 files changed, 55 insertions(+), 2 deletions(-) create mode 100644 .dockerignore create mode 100644 Dockerfile create mode 100644 docker-compose.yml diff --git a/.dockerignore b/.dockerignore new file mode 100644 index 0000000..64083d1 --- /dev/null +++ b/.dockerignore @@ -0,0 +1,2 @@ +target/ +.git \ No newline at end of file diff --git a/Dockerfile b/Dockerfile new file mode 100644 index 0000000..61c9731 --- /dev/null +++ b/Dockerfile @@ -0,0 +1,17 @@ +FROM rust:bookworm AS builder +RUN apt-get update && apt-get install -y protobuf-compiler && rm -rf /var/lib/apt/lists/* + +WORKDIR /app +COPY . . + +RUN cargo build --release --bin inference_proxy --bin worker + +FROM debian:bookworm-slim + +RUN apt-get update && apt-get install -y ca-certificates && rm -rf /var/lib/apt/lists/* + +# Copy only the compiled binaries from the builder stage +COPY --from=builder /app/target/release/inference_proxy /usr/local/bin/inference_proxy +COPY --from=builder /app/target/release/worker /usr/local/bin/worker + +EXPOSE 50051 50052 9090 diff --git a/docker-compose.yml b/docker-compose.yml new file mode 100644 index 0000000..5093ba9 --- /dev/null +++ b/docker-compose.yml @@ -0,0 +1,32 @@ +services: + worker-1: + build: . + command: ["worker"] + environment: + WORKER_LISTEN_ADDR: "0.0.0.0:50052" + + worker-2: + build: . + command: ["worker"] + environment: + WORKER_LISTEN_ADDR: "0.0.0.0:50052" + + worker-3: + build: . + command: ["worker"] + environment: + WORKER_LISTEN_ADDR: "0.0.0.0:50052" + + proxy: + build: . + command: ["inference_proxy"] + environment: + PROXY_LISTEN_ADDR: "0.0.0.0:50051" + WORKER_ADDRESSES: "http://worker-1:50052,http://worker-2:50052,http://worker-3:50052" + ports: + - "50051:50051" + - "9090:9090" + depends_on: + - worker-1 + - worker-2 + - worker-3 \ No newline at end of file diff --git a/src/bin/worker.rs b/src/bin/worker.rs index f5e0a29..ff92903 100644 --- a/src/bin/worker.rs +++ b/src/bin/worker.rs @@ -42,8 +42,10 @@ async fn main() -> Result<(), Box> { "50052".to_string() }; - let addr_str = format!("0.0.0.0:{}", port); - let addr = addr_str.parse()?; + let listen_addr = std::env::var("WORKER_LISTEN_ADDR") + .unwrap_or_else(|_| format!("[::1]:{}", port)); + + let addr = listen_addr.parse()?; let worker = GpuWorker::default(); From 92788273f281051b888f8b3617ac6b21a2103874 Mon Sep 17 00:00:00 2001 From: Abel Tesfa Date: Wed, 29 Jul 2026 10:38:35 +0300 Subject: [PATCH 3/6] refactor: edgecase fixes --- src/bin/client.rs | 4 +++- src/main.rs | 32 +++++++++++++++++++++++--------- 2 files changed, 26 insertions(+), 10 deletions(-) diff --git a/src/bin/client.rs b/src/bin/client.rs index 3e457db..35e7754 100644 --- a/src/bin/client.rs +++ b/src/bin/client.rs @@ -9,7 +9,9 @@ use inference::PredictRequest; async fn main() -> Result<(), Box> { println!("Connecting to the Inference Proxy server..."); - let mut client = InferenceServiceClient::connect("http://0.0.0.0:50051").await?; + let proxy_addr = std::env::var("PROXY_ADDRESS") + .unwrap_or_else(|_| "http://[::1]:50051".to_string()); + let mut client = InferenceServiceClient::connect(proxy_addr).await?; // dummy data to be sent to the server for inference let my_data = vec![1.0, 2.0, 3.0, 4.0]; diff --git a/src/main.rs b/src/main.rs index 2d7065f..b575c35 100644 --- a/src/main.rs +++ b/src/main.rs @@ -36,18 +36,32 @@ async fn main() -> Result<(), Box> { let worker_addrs_str = std::env::var("WORKER_ADDRESSES") .unwrap_or_else(|_| "http://[::1]:50052,http://[::1]:50053,http://[::1]:50054".to_string()); - let worker_addrs: Vec<&str> = worker_addrs_str.split(',').collect(); + let worker_addrs: Vec<&str> = worker_addrs_str + .split(',') + .map(|s| s.trim()) + .filter(|s| !s.is_empty()) + .collect(); + let mut worker_clients = Vec::new(); for addr_str in worker_addrs { - let addr_str = addr_str.trim(); - match InferenceServiceClient::connect(addr_str.to_string()).await { - Ok(client) => { - info!("Connected to worker node at {}", addr_str); - worker_clients.push(client); - } - Err(e) => { - warn!("Failed to connect to worker at {}: {}", addr_str, e); + let mut attempts = 0; + loop { + match InferenceServiceClient::connect(addr_str.to_string()).await { + Ok(client) => { + info!("Connected to worker node at {}", addr_str); + worker_clients.push(client); + break; + } + Err(e) => { + attempts += 1; + if attempts >= 5 { + warn!("Gave up connecting to worker at {}: {}", addr_str, e); + break; + } + warn!("Worker {} not ready yet. Retrying in 2 seconds...", addr_str); + tokio::time::sleep(std::time::Duration::from_secs(2)).await; + } } } } From 41541a3107aa525195201a8a417b3d632a09bf1b Mon Sep 17 00:00:00 2001 From: Abel Tesfa Date: Wed, 29 Jul 2026 10:57:45 +0300 Subject: [PATCH 4/6] feat: github workflow and some clippy fixes --- .github/workflows/ci.yml | 40 ++++++++++++++++++++++++++++++++++++++++ src/bin/client.rs | 13 ++++++++----- src/bin/worker.rs | 20 +++++++++++--------- src/main.rs | 24 +++++++++++++----------- src/proxy/batcher.rs | 39 ++++++++++++++++++++++++++++----------- src/proxy/mod.rs | 4 ++-- src/proxy/server.rs | 26 ++++++++++++-------------- 7 files changed, 114 insertions(+), 52 deletions(-) create mode 100644 .github/workflows/ci.yml diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml new file mode 100644 index 0000000..79019a6 --- /dev/null +++ b/.github/workflows/ci.yml @@ -0,0 +1,40 @@ +name: Rust CI + +on: + push: + branches: [ "main" ] + pull_request: + branches: [ "main" ] + +env: + CARGO_TERM_COLOR: always + +jobs: + build_and_test: + name: Build, Lint, & Test + runs-on: ubuntu-latest + + steps: + - name: Checkout Repository + uses: actions/checkout@v3 + + - name: Install Protobuf Compiler + run: sudo apt-get install -y protobuf-compiler + + - name: Setup Rust Toolchain + uses: dtolnay/rust-toolchain@stable + with: + components: rustfmt, clippy + + - name: Rust Cache + uses: Swatinem/rust-cache@v2 + + - name: Check Formatting: + run: cargo fmt --all -- --check + + - name: Run Clippy (Linter) + run: cargo clippy --all-targets --all-features + + - name: Run Tests + run: cargo test --all-features -- --nocapture + \ No newline at end of file diff --git a/src/bin/client.rs b/src/bin/client.rs index 35e7754..db379ed 100644 --- a/src/bin/client.rs +++ b/src/bin/client.rs @@ -2,15 +2,15 @@ pub mod inference { tonic::include_proto!("inference"); } -use inference::inference_service_client::InferenceServiceClient; use inference::PredictRequest; +use inference::inference_service_client::InferenceServiceClient; #[tokio::main] async fn main() -> Result<(), Box> { println!("Connecting to the Inference Proxy server..."); - let proxy_addr = std::env::var("PROXY_ADDRESS") - .unwrap_or_else(|_| "http://[::1]:50051".to_string()); + let proxy_addr = + std::env::var("PROXY_ADDRESS").unwrap_or_else(|_| "http://[::1]:50051".to_string()); let mut client = InferenceServiceClient::connect(proxy_addr).await?; // dummy data to be sent to the server for inference @@ -24,7 +24,10 @@ async fn main() -> Result<(), Box> { let response = client.predict(request).await?; - println!("Received response from the server: {:?}", response.into_inner().tensor_data); + println!( + "Received response from the server: {:?}", + response.into_inner().tensor_data + ); Ok(()) -} \ No newline at end of file +} diff --git a/src/bin/worker.rs b/src/bin/worker.rs index ff92903..9fa152c 100644 --- a/src/bin/worker.rs +++ b/src/bin/worker.rs @@ -1,4 +1,4 @@ -use tonic::{transport::Server, Request, Response, Status}; +use tonic::{Request, Response, Status, transport::Server}; pub mod inference { tonic::include_proto!("inference"); @@ -15,15 +15,17 @@ impl InferenceService for GpuWorker { &self, request: Request, ) -> Result, Status> { - let req_data = request.into_inner(); - println!("Worker Node received a batch of {} items to process!", req_data.tensor_data.len()); + println!( + "Worker Node received a batch of {} items to process!", + req_data.tensor_data.len() + ); // here you would call your GPU inference code / dummy for now - let result_data: Vec = req_data.tensor_data.iter().map(|x| x * 2.0).collect(); + let result_data: Vec = req_data.tensor_data.iter().map(|x| x * 2.0).collect(); - // generate a response + // generate a response let response = PredictResponse { tensor_data: result_data, }; @@ -35,15 +37,15 @@ impl InferenceService for GpuWorker { #[tokio::main] async fn main() -> Result<(), Box> { let args: Vec = std::env::args().collect(); - + let port = if args.len() > 1 { args[1].clone() } else { "50052".to_string() }; - let listen_addr = std::env::var("WORKER_LISTEN_ADDR") - .unwrap_or_else(|_| format!("[::1]:{}", port)); + let listen_addr = + std::env::var("WORKER_LISTEN_ADDR").unwrap_or_else(|_| format!("[::1]:{}", port)); let addr = listen_addr.parse()?; @@ -57,4 +59,4 @@ async fn main() -> Result<(), Box> { .await?; Ok(()) -} \ No newline at end of file +} diff --git a/src/main.rs b/src/main.rs index b575c35..f6722ef 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,28 +1,27 @@ -use tonic::transport::Server; -use tokio::sync::mpsc; -use tracing::{info, warn, error}; use metrics_exporter_prometheus::PrometheusBuilder; +use tokio::sync::mpsc; +use tonic::transport::Server; +use tracing::{error, info, warn}; pub mod proxy; pub mod inference { tonic::include_proto!("inference"); } -use inference::inference_service_server::InferenceServiceServer; use inference::inference_service_client::InferenceServiceClient; -use proxy::server::MyInferenceProxy; +use inference::inference_service_server::InferenceServiceServer; use proxy::batcher::run_batcher; +use proxy::server::MyInferenceProxy; #[tokio::main] async fn main() -> Result<(), Box> { - tracing_subscriber::fmt::init(); PrometheusBuilder::new() .with_http_listener(([0, 0, 0, 0], 9090)) .install() .expect("failed to install Prometheus recorder"); - + info!("Metrics server running on http://localhost:9090/metrics"); // Read from environment variable, fallback to localhost for local dev @@ -35,13 +34,13 @@ async fn main() -> Result<(), Box> { // Read worker addresses from env var (comma-separated), fallback for local dev let worker_addrs_str = std::env::var("WORKER_ADDRESSES") .unwrap_or_else(|_| "http://[::1]:50052,http://[::1]:50053,http://[::1]:50054".to_string()); - + let worker_addrs: Vec<&str> = worker_addrs_str .split(',') .map(|s| s.trim()) .filter(|s| !s.is_empty()) .collect(); - + let mut worker_clients = Vec::new(); for addr_str in worker_addrs { @@ -59,7 +58,10 @@ async fn main() -> Result<(), Box> { warn!("Gave up connecting to worker at {}: {}", addr_str, e); break; } - warn!("Worker {} not ready yet. Retrying in 2 seconds...", addr_str); + warn!( + "Worker {} not ready yet. Retrying in 2 seconds...", + addr_str + ); tokio::time::sleep(std::time::Duration::from_secs(2)).await; } } @@ -86,6 +88,6 @@ async fn main() -> Result<(), Box> { .add_service(InferenceServiceServer::new(proxy)) .serve(addr) .await?; - + Ok(()) } diff --git a/src/proxy/batcher.rs b/src/proxy/batcher.rs index 5632f07..4000709 100644 --- a/src/proxy/batcher.rs +++ b/src/proxy/batcher.rs @@ -18,22 +18,18 @@ pub async fn run_batcher( info!("Continuous Batch processor ready. Connected to {} workers.", worker_clients.len()); loop { - // 1. Block until the FIRST request arrives let first_req = match request_rx.recv().await { Some(req) => req, None => break, // Channel closed }; - let mut current_requests = vec![first_req]; - // 2. We got a request! Wait a tiny micro-window (50ms) to let other concurrent requests pile up tokio::time::sleep(Duration::from_millis(50)).await; - // 3. Drain everything else that arrived in that 50ms window instantly while let Ok(req) = request_rx.try_recv() { current_requests.push(req); - // Optional: Cap the batch size to protect the GPU (e.g., max 32 items) + // Optional: Cap the batch size to protect the GPU if current_requests.len() >= 32 { break; } @@ -42,17 +38,13 @@ pub async fn run_batcher( let batch_size = current_requests.len(); histogram!("proxy_batch_size").record(batch_size as f64); - // 4. Round-Robin Worker Selection let worker_idx = next_worker_index; next_worker_index = (next_worker_index + 1) % worker_clients.len(); - // Tonic clients are just Arc wrappers around HTTP/2 multiplexers, so cloning them is incredibly cheap! let mut worker_client = worker_clients[worker_idx].clone(); info!("Dispatching batch of {} requests to worker node #{}", batch_size, worker_idx + 1); - // 5. THE MAGIC: Spawn a background task to handle the network call to the worker! - // This means the batcher loop instantly goes back to catching new requests! tokio::spawn(async move { let mut aggregated_tensor = Vec::new(); let mut input_sizes = Vec::new(); @@ -74,7 +66,7 @@ pub async fn run_batcher( let processed_batch = worker_response.into_inner().tensor_data; let mut start_index = 0; - for (channel, size) in response_channels.into_iter().zip(input_sizes.into_iter()) { + for (channel, size) in response_channels.into_iter().zip(input_sizes) { let client_result = processed_batch[start_index..start_index + size].to_vec(); let _ = channel.send(client_result); start_index += size; @@ -87,4 +79,29 @@ pub async fn run_batcher( } }); } -} \ No newline at end of file +} + +#[cfg(test)] +mod tests { + + #[tokio::test] + async fn test_tensor_aggregation_logic() { + let input_tensors = vec![ + vec![1.0, 2.0], + vec![3.0, 4.0, 5.0], + vec![6.0], + ]; + + let mut aggregated_tensor = Vec::new(); + let mut input_sizes = Vec::new(); + + for tensor in input_tensors { + input_sizes.push(tensor.len()); + aggregated_tensor.extend(tensor); + } + + assert_eq!(aggregated_tensor, vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]); + assert_eq!(input_sizes, vec![2, 3, 1]); + + } +} diff --git a/src/proxy/mod.rs b/src/proxy/mod.rs index 38aa06a..5ebb2f3 100644 --- a/src/proxy/mod.rs +++ b/src/proxy/mod.rs @@ -1,9 +1,9 @@ use tokio::sync::oneshot; -pub mod server; pub mod batcher; +pub mod server; pub struct BatchRequest { pub tensor_data: Vec, pub response_channel: oneshot::Sender>, -} \ No newline at end of file +} diff --git a/src/proxy/server.rs b/src/proxy/server.rs index a988038..1f57d98 100644 --- a/src/proxy/server.rs +++ b/src/proxy/server.rs @@ -1,9 +1,9 @@ -use tonic::{Request, Response, Status}; -use tokio::sync::{mpsc, oneshot}; -use tracing::{info, error}; -use metrics::{counter, histogram}; -use std::time::{Duration, Instant}; use super::BatchRequest; +use metrics::{counter, histogram}; +use std::time::Instant; +use tokio::sync::{mpsc, oneshot}; +use tonic::{Request, Response, Status}; +use tracing::info; use crate::inference::inference_service_server::InferenceService; use crate::inference::{PredictRequest, PredictResponse}; @@ -16,9 +16,7 @@ pub struct MyInferenceProxy { impl MyInferenceProxy { pub fn new(request_tx: mpsc::Sender) -> Self { - Self { - request_tx, - } + Self { request_tx } } } #[tonic::async_trait] @@ -28,7 +26,6 @@ impl InferenceService for MyInferenceProxy { &self, request: Request, ) -> Result, Status> { - info!("Received a request from a client."); counter!("proxy_requests_total").increment(1); @@ -43,7 +40,7 @@ impl InferenceService for MyInferenceProxy { response_channel: tx, }; - if let Err(_) = self.request_tx.send(batch_req).await { + if self.request_tx.send(batch_req).await.is_err() { tracing::error!("Failed to send request to batcher. Batcher might be down."); return Err(Status::internal("Internal Server Error: Router Offline.")); } @@ -52,7 +49,9 @@ impl InferenceService for MyInferenceProxy { Ok(data) => data, Err(_) => { tracing::error!("Failed to receive response: GPU Worker failed to process batch."); - return Err(Status::internal("Internal Server Error: GPU Processing Failed.")); + return Err(Status::internal( + "Internal Server Error: GPU Processing Failed.", + )); } }; @@ -60,12 +59,11 @@ impl InferenceService for MyInferenceProxy { histogram!("proxy_request_duration_seconds").record(duration); info!("Request processed in {:.3} seconds.", duration); - // generate a response + // generate a response let response = PredictResponse { tensor_data: result_data, }; Ok(Response::new(response)) - } -} \ No newline at end of file +} From a58619b56e8de43cec402af98cebc3578682e887 Mon Sep 17 00:00:00 2001 From: Abel Tesfa Date: Wed, 29 Jul 2026 11:09:50 +0300 Subject: [PATCH 5/6] refactor: fixed ci --- .github/workflows/ci.yml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 79019a6..838e119 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -29,11 +29,11 @@ jobs: - name: Rust Cache uses: Swatinem/rust-cache@v2 - - name: Check Formatting: + - name: Check Formatting run: cargo fmt --all -- --check - name: Run Clippy (Linter) - run: cargo clippy --all-targets --all-features + run: cargo clippy --all-targets --all-features -- -D warnings - name: Run Tests run: cargo test --all-features -- --nocapture From 18f0f49f7894f76f7b4f1e5138f596a1549f5e93 Mon Sep 17 00:00:00 2001 From: Abel Tesfa Date: Wed, 29 Jul 2026 11:12:31 +0300 Subject: [PATCH 6/6] style: run cargo fmt --- src/proxy/batcher.rs | 41 ++++++++++++++++++++++------------------- 1 file changed, 22 insertions(+), 19 deletions(-) diff --git a/src/proxy/batcher.rs b/src/proxy/batcher.rs index 4000709..9bf181a 100644 --- a/src/proxy/batcher.rs +++ b/src/proxy/batcher.rs @@ -1,8 +1,8 @@ -use tokio::sync::mpsc; -use std::time::Duration; -use tracing::{info, error}; use metrics::histogram; +use std::time::Duration; +use tokio::sync::mpsc; use tonic::transport::Channel; +use tracing::{error, info}; use super::BatchRequest; use crate::inference::PredictRequest; @@ -11,12 +11,15 @@ use crate::inference::inference_service_client::InferenceServiceClient; // background task that aggregates and processes requests pub async fn run_batcher( mut request_rx: mpsc::Receiver, - worker_clients: Vec> + worker_clients: Vec>, ) { let mut next_worker_index = 0; - - info!("Continuous Batch processor ready. Connected to {} workers.", worker_clients.len()); - + + info!( + "Continuous Batch processor ready. Connected to {} workers.", + worker_clients.len() + ); + loop { let first_req = match request_rx.recv().await { Some(req) => req, @@ -28,8 +31,8 @@ pub async fn run_batcher( while let Ok(req) = request_rx.try_recv() { current_requests.push(req); - - // Optional: Cap the batch size to protect the GPU + + // Optional: Cap the batch size to protect the GPU if current_requests.len() >= 32 { break; } @@ -40,10 +43,14 @@ pub async fn run_batcher( let worker_idx = next_worker_index; next_worker_index = (next_worker_index + 1) % worker_clients.len(); - + let mut worker_client = worker_clients[worker_idx].clone(); - info!("Dispatching batch of {} requests to worker node #{}", batch_size, worker_idx + 1); + info!( + "Dispatching batch of {} requests to worker node #{}", + batch_size, + worker_idx + 1 + ); tokio::spawn(async move { let mut aggregated_tensor = Vec::new(); @@ -64,10 +71,11 @@ pub async fn run_batcher( match worker_client.predict(request).await { Ok(worker_response) => { let processed_batch = worker_response.into_inner().tensor_data; - + let mut start_index = 0; for (channel, size) in response_channels.into_iter().zip(input_sizes) { - let client_result = processed_batch[start_index..start_index + size].to_vec(); + let client_result = + processed_batch[start_index..start_index + size].to_vec(); let _ = channel.send(client_result); start_index += size; } @@ -86,11 +94,7 @@ mod tests { #[tokio::test] async fn test_tensor_aggregation_logic() { - let input_tensors = vec![ - vec![1.0, 2.0], - vec![3.0, 4.0, 5.0], - vec![6.0], - ]; + let input_tensors = vec![vec![1.0, 2.0], vec![3.0, 4.0, 5.0], vec![6.0]]; let mut aggregated_tensor = Vec::new(); let mut input_sizes = Vec::new(); @@ -102,6 +106,5 @@ mod tests { assert_eq!(aggregated_tensor, vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]); assert_eq!(input_sizes, vec![2, 3, 1]); - } }