Skip to content

Repository files navigation

Aegis Inference Proxy

A high-throughput, asynchronous gRPC proxy server designed to act as the front door for Machine Learning inference workloads. Built completely in Rust using tokio and tonic.

Architecture & Concept

In production AI environments, sending individual requests to a GPU is highly inefficient. GPUs are designed for massive parallel processing and require batched data.

This proxy solves that by decoupling incoming user requests from the execution hardware:

  1. gRPC Interface: Clients send tensor data (lists of floats) via fast, compressed Protocol Buffers.
  2. Dynamic Request Batching: The proxy catches concurrent requests and routes them into an mpsc (Multi-Producer, Single-Consumer) channel.
  3. Async Background Loop: A dedicated tokio::spawn task aggregates incoming tensors until a max batch size is met, OR a latency timeout (e.g., 3 seconds) is reached.
  4. Result Dispatching: The concatenated batch is processed, and the results are perfectly sliced and returned to the exact clients who requested them via oneshot channels.

Technology Stack

  • Language: Rust (Edition 2024)
  • Async Runtime: tokio
  • RPC Framework: tonic (gRPC)
  • Serialization: Protocol Buffers (prost)

Getting Started

Prerequisites

Because this project compiles Protocol Buffers, you must have the protoc compiler installed on your system.

Debian/Ubuntu:

sudo apt update
sudo apt install protobuf-compiler

Running the Server

To start the Inference Proxy server, run:

cargo run

The server will initialize the background batching processor and begin listening on [::1]:50051.

Running the Test Client

This repository includes a built-in dummy client to test the gRPC connection and batching logic.

While the server is running, open a separate terminal and run:

cargo run --bin client

To see the Dynamic Batching in action, open two separate terminals and run the client command in both simultaneously. The server will catch both requests, aggregate them, process them as a single batch, and return the sliced results to both terminals.

About

A high-throughput, dynamically batching gRPC inference proxy written in Rust.

Resources

Stars

1 star

Watchers

0 watching

Forks

Releases

Packages

Contributors

Languages