Implementation:Predibase Lorax GRPC Metadata Lib

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

Overview

This crate provides OpenTelemetry context propagation for gRPC requests by extracting and injecting tracing span context into tonic metadata maps.

Description

The crate defines two internal structs, MetadataExtractor and MetadataInjector, which wrap a tonic::metadata::MetadataMap reference and implement the OpenTelemetry Extractor and Injector traits respectively. MetadataExtractor reads key-value pairs from the gRPC metadata for context extraction, while MetadataInjector writes propagation headers into the metadata. The private inject function uses the global OpenTelemetry text map propagator to inject the current tracing span's context into a metadata map. The public trait InjectTelemetryContext is implemented for tonic::Request<T>, providing an inject_context method that enriches an outgoing gRPC request with the current distributed tracing context.

Usage

This crate is used by the LoRAX router's gRPC client layer to propagate distributed tracing context from the router to the Python inference server shards. When the router sends a gRPC request to a shard, it calls inject_context() on the request to include the OpenTelemetry trace and span IDs in the gRPC metadata, enabling end-to-end tracing across service boundaries.

Code Reference

Source Location

• Repository: Predibase_Lorax
• File: router/grpc-metadata/src/lib.rs
• Lines: 1-63

Signature

pub trait InjectTelemetryContext {
    fn inject_context(self) -> Self;
}

impl<T> InjectTelemetryContext for tonic::Request<T> {
    fn inject_context(mut self) -> Self;
}

struct MetadataExtractor<'a>(pub &'a tonic::metadata::MetadataMap);

impl<'a> Extractor for MetadataExtractor<'a> {
    fn get(&self, key: &str) -> Option<&str>;
    fn keys(&self) -> Vec<&str>;
}

struct MetadataInjector<'a>(pub &'a mut tonic::metadata::MetadataMap);

impl<'a> Injector for MetadataInjector<'a> {
    fn set(&mut self, key: &str, value: String);
}

Import

use opentelemetry::global;
use opentelemetry::propagation::{Extractor, Injector};
use tracing_opentelemetry::OpenTelemetrySpanExt;

I/O Contract

Inputs

Name	Type	Required	Description
self	tonic::Request<T>	Yes	The outgoing gRPC request whose metadata will be enriched with tracing context
key	&str	Yes	The metadata key to read (for MetadataExtractor::get) or write (for MetadataInjector::set)
value	String	Yes	The propagation header value to inject (for MetadataInjector::set)

Outputs

Name	Type	Description
inject_context	tonic::Request<T>	The same request with OpenTelemetry context injected into its metadata
get	Option<&str>	The metadata value for the given key, or None if not found or not convertible to a string
keys	Vec<&str>	All metadata key names in the map

Usage Examples

use grpc_metadata::InjectTelemetryContext;

// Inject the current OpenTelemetry span context into a gRPC request
let request = tonic::Request::new(my_payload);
let request_with_context = request.inject_context();

// The request metadata now contains trace propagation headers
// (e.g., traceparent, tracestate) that the server can extract
client.generate(request_with_context).await?;

Related Pages

• Environment:Predibase_Lorax_CUDA_GPU_Runtime