feat: server side streaming rather than a single response

Cargo.toml

@@ -11,6 +11,7 @@ tracing = "0.1"
 tracing-subscriber = "0.3"
 metrics = "0.22"
 metrics-exporter-prometheus = "0.13"
+tokio-stream = "0.1"
 
 [build-dependencies]
 tonic-build = "0.11"

proto/inference.proto

@@ -5,6 +5,7 @@ package Inference;
 // gRPC service and endpoints
 service InferenceService {
   rpc Predict (PredictRequest) returns (PredictResponse);
+  rpc PredictStream (PredictRequest) returns (stream StreamChunk);
 }
 
 // incoming request / tensor data
@@ -16,3 +17,9 @@ message PredictRequest {
 message PredictResponse{
   repeated float tensor_data = 1;
 }
+
+// stream message
+message StreamChunk {
+  repeated float tensor_data = 1;
+  bool is_final = 2;
+}

src/bin/client.rs

@@ -22,12 +22,12 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
         tensor_data: my_data,
     });
 
-    let response = client.predict(request).await?;
+    let mut stream = client.predict_stream(request).await?.into_inner();
 
-    println!(
-        "Received response from the server: {:?}",
-        response.into_inner().tensor_data
-    );
+    println!("Receiving streaming response...");
+    while let Some(chunk) = stream.message().await? {
+        println!("Got chunk: {:?} (is_final: {})", chunk.tensor_data, chunk.is_final);
+    }
 
     Ok(())
 }

src/bin/worker.rs

@@ -1,10 +1,14 @@
 use tonic::{Request, Response, Status, transport::Server};
+use std::pin::Pin;
+use tokio_stream::wrappers::ReceiverStream;
+use tokio_stream::Stream;
+use std::time::Duration;
 
 pub mod inference {
     tonic::include_proto!("inference");
 }
 use inference::inference_service_server::{InferenceService, InferenceServiceServer};
-use inference::{PredictRequest, PredictResponse};
+use inference::{PredictRequest, PredictResponse, StreamChunk};
 
 #[derive(Default)]
 pub struct GpuWorker {}
@@ -32,8 +36,45 @@ impl InferenceService for GpuWorker {
 
         Ok(Response::new(response))
     }
+
+    type PredictStreamStream = Pin<Box<dyn Stream<Item = Result<StreamChunk, Status>> + Send>>;
+
+    async fn predict_stream(
+        &self,
+        request: Request<PredictRequest>,
+    ) -> Result<Response<Self::PredictStreamStream>, Status> {
+        let req_data = request.into_inner();
+        let chunk_size = 2;
+
+        let (tx, rx) = tokio::sync::mpsc::channel(32);
+
+        tokio::spawn(async move {
+            let chunks: Vec<&[f32]> = req_data.tensor_data.chunks(chunk_size).collect();
+            let total_chunks = chunks.len();
+
+            for (i, chunk) in chunks.into_iter().enumerate() {
+                let processed: Vec<f32> = chunk.iter().map(|x| x * 2.0).collect();
+                let is_final = i == total_chunks - 1;
+
+                tokio::time::sleep(Duration::from_millis(500)).await;
+
+                let stream_chunk = StreamChunk {
+                    tensor_data: processed,
+                    is_final,
+                };
+
+                if tx.send(Ok(stream_chunk)).await.is_err() {
+                    break;
+                }
+            }
+        });
+
+        let stream = ReceiverStream::new(rx);
+        Ok(Response::new(Box::pin(stream)))
+    }
 }
 
+
 #[tokio::main]
 async fn main() -> Result<(), Box<dyn std::error::Error>> {
     let args: Vec<String> = std::env::args().collect();

src/main.rs

@@ -76,11 +76,12 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
     // initialize mpsc channel for batching
     let (tx, rx) = mpsc::channel(100);
 
+    let proxy_workers = worker_clients.clone();
     tokio::spawn(async move {
         run_batcher(rx, worker_clients).await;
     });
 
-    let proxy = MyInferenceProxy::new(tx);
+    let proxy = MyInferenceProxy::new(tx, proxy_workers);
 
     info!("Inference Proxy listening on {}", addr);
 

src/proxy/mod.rs

@@ -1,4 +1,4 @@
-use tokio::sync::oneshot;
+use tokio::sync::{oneshot, mpsc};
 
 pub mod batcher;
 pub mod server;
@@ -7,3 +7,8 @@ pub struct BatchRequest {
     pub tensor_data: Vec<f32>,
     pub response_channel: oneshot::Sender<Vec<f32>>,
 }
+
+pub struct StreamRequest {
+    pub tensor_data: Vec<f32>,
+    pub stream_tx: mpsc::Sender<Vec<f32>>,
+}

src/proxy/server.rs

@@ -2,21 +2,30 @@ use super::BatchRequest;
 use metrics::{counter, histogram};
 use std::time::Instant;
 use tokio::sync::{mpsc, oneshot};
-use tonic::{Request, Response, Status};
+use tonic::{Request, Response, Status, transport::Channel};
+use std::sync::atomic::{AtomicUsize, Ordering};
+use std::sync::Arc;
 use tracing::info;
 
 use crate::inference::inference_service_server::InferenceService;
-use crate::inference::{PredictRequest, PredictResponse};
+use crate::inference::inference_service_client::InferenceServiceClient;
+use crate::inference::{PredictRequest, PredictResponse, StreamChunk};
 
 #[derive(Clone)]
 pub struct MyInferenceProxy {
     // channel to send batch requests to the batcher
     request_tx: mpsc::Sender<BatchRequest>,
+    worker_clients: Vec<InferenceServiceClient<Channel>>,
+    next_worker_index: Arc<AtomicUsize>,
 }
 
 impl MyInferenceProxy {
-    pub fn new(request_tx: mpsc::Sender<BatchRequest>) -> Self {
-        Self { request_tx }
+    pub fn new(request_tx: mpsc::Sender<BatchRequest>, worker_clients: Vec<InferenceServiceClient<Channel>>) -> Self {
+        Self { 
+            request_tx, 
+            worker_clients,
+            next_worker_index: Arc::new(AtomicUsize::new(0))
+        }
     }
 }
 #[tonic::async_trait]
@@ -66,4 +75,26 @@ impl InferenceService for MyInferenceProxy {
 
         Ok(Response::new(response))
     }
+
+    type PredictStreamStream = tonic::Streaming<StreamChunk>;
+
+    async fn predict_stream(
+        &self,
+        request: Request<PredictRequest>,
+    ) -> Result<Response<Self::PredictStreamStream>, Status> {
+        info!("Received a STREAMING request from a client.");
+
+        if self.worker_clients.is_empty() {
+            return Err(Status::internal("No workers available"));
+        }
+
+        let idx = self.next_worker_index.fetch_add(1, Ordering::Relaxed) % self.worker_clients.len();
+        let mut worker_client = self.worker_clients[idx].clone();
+
+        info!("Forwarding streaming request directly to worker #{}...", idx + 1);
+
+        let response = worker_client.predict_stream(request).await?;
+
+        Ok(response)
+    }
 }