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/.github/workflows/ci.yml b/.github/workflows/ci.yml new file mode 100644 index 0000000..838e119 --- /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 -- -D warnings + + - name: Run Tests + run: cargo test --all-features -- --nocapture + \ 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/client.rs b/src/bin/client.rs index c807592..db379ed 100644 --- a/src/bin/client.rs +++ b/src/bin/client.rs @@ -2,14 +2,16 @@ 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 mut client = InferenceServiceClient::connect("http://[::1]: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]; @@ -22,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 08a2f92..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,17 @@ 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 addr_str = format!("[::1]:{}", 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(); @@ -55,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 af7a0a3..f6722ef 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,53 +1,77 @@ -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"); - 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(',') + .map(|s| s.trim()) + .filter(|s| !s.is_empty()) + .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 { - Ok(client) => { - info!("Connected to worker node on port {}", port); - worker_clients.push(client); - } - Err(e) => { - info!("Failed to connect to worker node on port {}: {}", port, e); + for addr_str in worker_addrs { + 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; + } } } } // 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); @@ -64,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..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,29 +11,28 @@ 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 { - // 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 +41,17 @@ 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); + 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(); @@ -72,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.into_iter()) { - let client_result = processed_batch[start_index..start_index + size].to_vec(); + 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 +87,24 @@ 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 +}