Dev

.dockerignore

@@ -0,0 +1,2 @@
+target/
+.git

.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
+

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

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

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<dyn std::error::Error>> {
     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<dyn std::error::Error>> {
 
     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(())
-}
+}

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<PredictRequest>,
     ) -> Result<Response<PredictResponse>, 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<f32> = req_data.tensor_data.iter().map(|x| x * 2.0).collect(); 
+        let result_data: Vec<f32> = 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<dyn std::error::Error>> {
     let args: Vec<String> = 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<dyn std::error::Error>> {
         .await?;
 
     Ok(())
-}
+}

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<dyn std::error::Error>> {
-
     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<dyn std::error::Error>> {
         .add_service(InferenceServiceServer::new(proxy))
         .serve(addr)
         .await?;
-    
+
     Ok(())
 }