Implementation:Predibase Lorax Router Client Lib

Knowledge Sources	Predibase_Lorax
Domains	Client, gRPC
Last Updated	2026-02-08 00:00 GMT

Overview

This is the root library crate for the LoRAX gRPC client, which re-exports protobuf-generated types, client and sharded client structs, error types, and a trait for converting chunked inputs to strings.

Description

The crate serves as the public API surface for the LoRAX router's gRPC client library. It declares three internal modules: client (the single-shard gRPC client), pb (protobuf-generated message types), and sharded_client (a client that fans out across multiple model shards). The crate re-exports a large set of protobuf types from pb::generate::v1 including Batch, CachedBatch, Generation, Request, AdapterParameters, Embedding, InputChunk, and many more. It defines a ClientError enum with three variants (Connection, Generation, EmptyResults) and provides From trait implementations to convert tonic::Status and tonic::transport::Error into ClientError. The ChunksToString trait provides backward-compatible string serialization for backends that do not support chunked inputs, handling both text and base64-encoded image chunks. A static WARMUP_IMAGE_BASE64 constant holds a small PNG image used during model warm-up.

Usage

This crate is used by the LoRAX router to communicate with the Python inference server shards over gRPC. The Client and ShardedClient types are the primary entry points for sending inference requests (generate, prefill, decode) and managing adapter lifecycle operations. The ChunksToString trait is used for backward-compatible multimodal input handling.

Code Reference

Source Location

• Repository: Predibase_Lorax
• File: router/client/src/lib.rs
• Lines: 1-83

Signature

#[derive(Error, Debug, Clone)]
pub enum ClientError {
    #[error("Could not connect to LoRAX server: {0}")]
    Connection(String),
    #[error("Server error: {0}")]
    Generation(String),
    #[error("Sharded results are empty")]
    EmptyResults,
}

pub trait ChunksToString {
    fn chunks_to_string(&self) -> String;
}

pub type Result<T> = std::result::Result<T, ClientError>;

Import

pub use client::Client;
pub use sharded_client::ShardedClient;
pub use pb::generate::v1::HealthResponse;
pub use pb::generate::v1::InfoResponse as ShardInfo;
pub use pb::generate::v1::{
    input_chunk, AdapterParameters, AlternativeTokens, Batch, CachedBatch,
    ClassifyPredictionList, DownloadAdapterResponse, Embedding, Entity,
    EntityList, FinishReason, GeneratedText, Generation, Image, InputChunk,
    MajoritySignMethod, MergeStrategy, NextTokenChooserParameters, NextTokens,
    PreloadedAdapter, Request, StoppingCriteriaParameters, TokenizedInputs,
};

I/O Contract

Inputs

Name	Type	Required	Description
chunks	Vec<InputChunk>	Yes	A vector of input chunks containing either text or image data, used by the ChunksToString trait
tonic::Status	tonic::Status	No	Converted into ClientError::Generation via the From trait
tonic::transport::Error	tonic::transport::Error	No	Converted into ClientError::Connection via the From trait

Outputs

Name	Type	Description
chunks_to_string	String	Concatenated text and base64-encoded image data from input chunks
ClientError	ClientError	Error type wrapping connection, generation, and empty result errors
Result<T>	std::result::Result<T, ClientError>	Convenience type alias for fallible client operations

Usage Examples

use lorax_client::{Client, ShardedClient, ClientError, ChunksToString, InputChunk};

// Create a client connection to a single shard
let client = Client::connect("http://localhost:8080".to_string()).await?;

// Use ChunksToString to convert multimodal input to a string
let chunks: Vec<InputChunk> = vec![/* ... */];
let text_repr = chunks.chunks_to_string();

// Handle client errors
match result {
    Ok(generation) => { /* process generation */ },
    Err(ClientError::Connection(msg)) => { /* handle connection error */ },
    Err(ClientError::Generation(msg)) => { /* handle generation error */ },
    Err(ClientError::EmptyResults) => { /* handle empty sharded results */ },
}

Related Pages

• Environment:Predibase_Lorax_CUDA_GPU_Runtime