From 76323055d2564c7dab9f401b8a169271bafc29ab Mon Sep 17 00:00:00 2001 From: Abel Tesfa Date: Fri, 31 Jul 2026 02:55:04 +0300 Subject: [PATCH] feat: server side streaming rather than a single response --- Cargo.lock | 1 + Cargo.toml | 1 + proto/inference.proto | 7 +++++++ src/bin/client.rs | 10 +++++----- src/bin/worker.rs | 43 ++++++++++++++++++++++++++++++++++++++++++- src/main.rs | 3 ++- src/proxy/mod.rs | 7 ++++++- src/proxy/server.rs | 39 +++++++++++++++++++++++++++++++++++---- 8 files changed, 99 insertions(+), 12 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 7addb2d..f810a09 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -475,6 +475,7 @@ dependencies = [ "metrics-exporter-prometheus", "prost", "tokio", + "tokio-stream", "tonic", "tonic-build", "tracing", diff --git a/Cargo.toml b/Cargo.toml index 3a01647..85f0dfa 100644 --- a/Cargo.toml +++ b/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" diff --git a/proto/inference.proto b/proto/inference.proto index e281625..733b129 100644 --- a/proto/inference.proto +++ b/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; +} diff --git a/src/bin/client.rs b/src/bin/client.rs index db379ed..d4288d6 100644 --- a/src/bin/client.rs +++ b/src/bin/client.rs @@ -22,12 +22,12 @@ async fn main() -> Result<(), Box> { 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(()) } diff --git a/src/bin/worker.rs b/src/bin/worker.rs index 9fa152c..4674bfb 100644 --- a/src/bin/worker.rs +++ b/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> + Send>>; + + async fn predict_stream( + &self, + request: Request, + ) -> Result, 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 = 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> { let args: Vec = std::env::args().collect(); diff --git a/src/main.rs b/src/main.rs index f6722ef..6e56698 100644 --- a/src/main.rs +++ b/src/main.rs @@ -76,11 +76,12 @@ async fn main() -> Result<(), Box> { // 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); diff --git a/src/proxy/mod.rs b/src/proxy/mod.rs index 5ebb2f3..d37c5a4 100644 --- a/src/proxy/mod.rs +++ b/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, pub response_channel: oneshot::Sender>, } + +pub struct StreamRequest { + pub tensor_data: Vec, + pub stream_tx: mpsc::Sender>, +} \ No newline at end of file diff --git a/src/proxy/server.rs b/src/proxy/server.rs index 1f57d98..6306ffa 100644 --- a/src/proxy/server.rs +++ b/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, + worker_clients: Vec>, + next_worker_index: Arc, } impl MyInferenceProxy { - pub fn new(request_tx: mpsc::Sender) -> Self { - Self { request_tx } + pub fn new(request_tx: mpsc::Sender, worker_clients: Vec>) -> 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; + + async fn predict_stream( + &self, + request: Request, + ) -> Result, 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) + } }