feat: add OpenTelemetry tracing support to router (#55)

This PR enables OpenTelemetry tracing support at the router level for
`prefill`, `generate`, and `generate_stream` methods.

Signed-off-by: Daniel Clark <[email protected]>
Signed-off-by: Joe Runde <[email protected]>
Co-authored-by: PRATIBHA MOOGI <[email protected]>
Co-authored-by: Daniel Clark <[email protected]>
Co-authored-by: Joe Runde <[email protected]>

Author: 3 people

Cargo.lock

@@ -11,5 +11,5 @@ ctrlc = { version = "3.4.2", features = ["termination"] }
 nix = { version = "0.28.0", features = ["process", "signal"] }
 serde_json = "^1.0.114"
 tracing = "0.1.40"
-tracing-subscriber = { version = "0.3.18", features = ["json"] }
-uuid = { version = "1.7.0", features = ["v4", "fast-rng"] }
+tracing-subscriber = { version = "0.3.18", features = ["json", "env-filter"] }
+uuid = { version = "1.7.0", features = ["v4", "fast-rng"] }

launcher/Cargo.toml

@@ -89,6 +89,8 @@ struct Args {
     // Default for default_include_stop_seqs is true for now, for backwards compatibility
     #[clap(default_value = "true", long, env, action = clap::ArgAction::Set)]
     default_include_stop_seqs: bool,
+    #[clap(long, env)]
+    otlp_endpoint: Option<String>,
 }
 
 fn main() -> ExitCode {
@@ -107,7 +109,6 @@ fn main() -> ExitCode {
 
     // Pattern match configuration
     let args = Args::parse();
-
     if args.json_output {
         tracing_subscriber::fmt()
             .json()
@@ -326,6 +327,12 @@ fn main() -> ExitCode {
         argv.push("--json-output".to_string());
     }
 
+    // OpenTelemetry
+    if let Some(otlp_endpoint) = args.otlp_endpoint {
+        argv.push("--otlp-endpoint".to_string());
+        argv.push(otlp_endpoint);
+    }
+
     if args.output_special_tokens {
         argv.push("--output-special-tokens".into());
     }
@@ -400,6 +407,7 @@ fn main() -> ExitCode {
         };
     }
 
+    // Graceful termination
     terminate_gracefully(&mut webserver, shutdown.clone(), shutdown_receiver);
 
     exit_code

launcher/src/main.rs

@@ -15,6 +15,7 @@ path = "src/main.rs"
 
 [dependencies]
 axum = { version = "0.6.20", features = ["json"] }
+axum-tracing-opentelemetry = "0.10.0"
 text-generation-client = { path = "client" }
 clap = { version = "^4.5.2", features = ["derive", "env"] }
 futures = "^0.3.30"
@@ -38,12 +39,15 @@ tokio = { version = "1.36.0", features = ["rt", "rt-multi-thread", "parking_lot"
 tokio-rustls = "^0.25.0"
 rustls = "0.22.2"
 tracing = "^0.1.40"
-tracing-subscriber = { version = "0.3.18", features = ["json"] }
 prost = "^0.12.3"
 tonic = { version = "^0.11.0", features = ["tls"] }
+tracing-subscriber = { version = "0.3.18", features = ["json", "env-filter"] }
+tracing-opentelemetry = "0.19.0"
 tokio-stream ="^0.1.14"
 unicode-segmentation = "^1.11.0"
 unicode-truncate = "^0.2.0"
+opentelemetry = { version = "0.19.0", features = ["rt-tokio"] }
+opentelemetry-otlp = "0.12.0"
 
 [build-dependencies]
-tonic-build = "^0.11.0"
+tonic-build = "^0.11.0"

router/Cargo.toml

@@ -52,15 +52,12 @@ impl Client {
         })
     }
 
-    /// Returns a list of uris or unix sockets of all shards
-    #[instrument(skip(self))]
+    // Returns a list of uris or unix sockets of all shards
+    //#[instrument(skip(self))]
+    // Below function is a method only used once during pod startup and not tied to any external requests/transactions, disabling otel instrumentation
     pub async fn service_discovery(&mut self) -> Result<Vec<String>> {
         let request = tonic::Request::new(ServiceDiscoveryRequest {});
-        let response = self
-            .stub
-            .service_discovery(request)
-            .instrument(info_span!("service_discovery"))
-            .await?;
+        let response = self.stub.service_discovery(request).await?;
         let urls = response
             .into_inner()
             .urls
@@ -75,26 +72,20 @@ impl Client {
     }
 
     /// Clear the past generations cache
-    #[instrument(skip(self))]
+    //#[instrument(skip(self))]
+    //Below function is a method only used once during pod startup and not tied to any external requests/transactions, disabling otel instrumentation
     pub async fn clear_cache(&mut self) -> Result<()> {
         let request = tonic::Request::new(ClearCacheRequest {});
-        self.stub
-            .clear_cache(request)
-            .instrument(info_span!("clear_cache"))
-            .await?;
+        self.stub.clear_cache(request).await?;
         Ok(())
     }
 
     /// Get shard model info
-    #[instrument(skip(self))]
+    // Below function is a method only used once during pod startup and not tied to any external requests/transactions, disabling otel instrumentation
+    //#[instrument(skip(self))]
     pub async fn model_info(&mut self) -> Result<(ModelType, u32, bool, MemoryScalingModel)> {
         let request = tonic::Request::new(ModelInfoRequest {});
-        let response = self
-            .stub
-            .model_info(request)
-            .instrument(info_span!("model_info"))
-            .await?
-            .into_inner();
+        let response = self.stub.model_info(request).await?.into_inner();
         ModelType::try_from(response.model_type)
             .map(|mt| {
                 (
@@ -108,7 +99,7 @@ impl Client {
     }
 
     /// Get model health
-    #[instrument(skip(self))]
+    //#[instrument(skip(self))]
     pub async fn health(&mut self) -> Result<HealthResponse> {
         let request = tonic::Request::new(HealthRequest {});
         let response = self.stub.health(request).await?.into_inner();
@@ -120,20 +111,15 @@ impl Client {
     pub async fn prefix_lookup(&mut self, prefix_id: String) -> Result<u32> {
         let mut request = tonic::Request::new(PrefixLookupRequest { prefix_id });
         request.set_timeout(PREFIX_LOOKUP_TIMEOUT);
-        let response = self
-            .stub
-            .prefix_lookup(request)
-            .instrument(info_span!("prefix_lookup"))
-            .await?
-            .into_inner();
+        let response = self.stub.prefix_lookup(request).await?.into_inner();
         Ok(response.prefix_length)
     }
 
     /// Generate one token for each request in the given batch
     ///
     /// Returns first generated token for each request in the batch, id of the next cached batch,
     /// and input token info if requested
-    #[instrument(skip(self))]
+    #[instrument(skip_all, fields(batch_id = &batch.id))]
     pub async fn prefill(
         &mut self,
         batch: Batch,
@@ -143,12 +129,7 @@ impl Client {
             batch: Some(batch),
             to_prune,
         });
-        let response = self
-            .stub
-            .prefill(request)
-            .instrument(info_span!("generate"))
-            .await?
-            .into_inner();
+        let response = self.stub.prefill(request).await?.into_inner();
         let result = response
             .result
             .ok_or_else(|| ClientError::Generation("Unexpected empty response".into()))?;
@@ -164,18 +145,13 @@ impl Client {
     /// Generate one token for each request in the given cached batch(es)
     ///
     /// Returns next generated token of each request in the batches and id of the next cached batch
-    #[instrument(skip(self))]
+    //#[instrument(skip(self))] <You can uncomment it for getting traces at each next_token() level
     pub async fn next_token(
         &mut self,
         batches: Vec<CachedBatch>,
     ) -> Result<Option<GenerateTokenResponse>> {
         let request = tonic::Request::new(NextTokenRequest { batches });
-        let response = self
-            .stub
-            .next_token(request)
-            .instrument(info_span!("generate_with_cache"))
-            .await?
-            .into_inner();
+        let response = self.stub.next_token(request).await?.into_inner();
         Ok(response.result.map(|result| {
             (
                 result.output_tokens,

router/client/src/client.rs

@@ -11,7 +11,7 @@ use tonic::{
     transport::{Certificate, Identity, Server, ServerTlsConfig},
     Request, Response, Status,
 };
-use tracing::{info_span, instrument, Span};
+use tracing::{instrument, Span};
 use unicode_truncate::UnicodeTruncateStr;
 
 use crate::{
@@ -99,11 +99,18 @@ impl GenerationService for GenerationServicer {
     #[instrument(
         skip_all,
         fields(
+            request_id=tracing::field::Empty,
             input=?request.get_ref().requests.iter().map(|r| truncate(&r.text, 32)).collect::<Vec<Cow<'_,str>>>(),
             prefix_id=?request.get_ref().prefix_id,
             correlation_id=?request.metadata().get("x-correlation-id").map(|mv| mv.to_str().unwrap_or("<non-ascii>")).unwrap_or("<none>"),
             input_bytes=?request.get_ref().requests.iter().map(|r| r.text.len()).collect::<Vec<usize>>(),
             params=?request.get_ref().params,
+            validation_time=tracing::field::Empty,
+            queue_time=tracing::field::Empty,
+            inference_time=tracing::field::Empty,
+            time_per_token=tracing::field::Empty,
+            total_time=tracing::field::Empty,
+            input_toks=tracing::field::Empty,
         )
     )]
     async fn generate(
@@ -224,11 +231,18 @@ impl GenerationService for GenerationServicer {
     #[instrument(
         skip_all,
         fields(
+            request_id=tracing::field::Empty,
             input=?truncate(request.get_ref().request.as_ref().map(|r| &*r.text).unwrap_or(""), 32),
             prefix_id=?request.get_ref().prefix_id,
             correlation_id=?request.metadata().get("x-correlation-id").map(|mv| mv.to_str().unwrap_or("<non-ascii>")).unwrap_or("<none>"),
             input_bytes=?request.get_ref().request.as_ref().map(|r| r.text.len()).unwrap_or(0),
             params=?request.get_ref().params,
+            validation_time=tracing::field::Empty,
+            queue_time=tracing::field::Empty,
+            inference_time=tracing::field::Empty,
+            time_per_token=tracing::field::Empty,
+            total_time=tracing::field::Empty,
+            input_toks=tracing::field::Empty,
         )
     )]
     async fn generate_stream(
@@ -416,8 +430,6 @@ fn log_response(
     kind_log: &str,
     request_id: Option<u64>,
 ) {
-    let span;
-    let _enter;
     // Timings
     let total_time = Instant::now() - start_time;
     if let Some(times) = times.as_ref() {
@@ -429,17 +441,14 @@ fn log_response(
             .unwrap_or_else(|| Duration::new(0, 0));
 
         // Tracing metadata
-        span = info_span!(
-            "",
-            validation_time = ?validation_time,
-            queue_time = ?queue_time,
-            inference_time = ?inference_time,
-            time_per_token = ?time_per_token,
-            total_time = ?total_time,
-            input_toks = input_tokens,
-            request_id = request_id,
-        );
-        _enter = span.enter();
+        let span = Span::current();
+        span.record("request_id", request_id.unwrap_or_default());
+        span.record("validation_time", format!("{validation_time:?}"));
+        span.record("queue_time", format!("{queue_time:?}"));
+        span.record("inference_time", format!("{inference_time:?}"));
+        span.record("time_per_token", format!("{time_per_token:?}"));
+        span.record("total_time", format!("{total_time:?}"));
+        span.record("input_toks", input_tokens);
 
         metrics::histogram!(
             "tgi_request_inference_duration",

router/src/grpc_server.rs

@@ -7,10 +7,17 @@ use std::{
 
 /// Text Generation Inference external gRPC server entrypoint
 use clap::Parser;
+use opentelemetry::{
+    global,
+    sdk::{propagation::TraceContextPropagator, trace, trace::Sampler, Resource},
+    KeyValue,
+};
+use opentelemetry_otlp::WithExportConfig;
 use text_generation_client::ShardedClient;
 use text_generation_router::{server, server::ServerRunArgs};
 use tokenizers::Tokenizer;
 use tracing::warn;
+use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt, EnvFilter, Layer};
 
 /// App Configuration
 #[derive(Parser, Debug)]
@@ -50,6 +57,8 @@ struct Args {
     output_special_tokens: bool,
     #[clap(long, env)]
     default_include_stop_seqs: bool,
+    #[clap(long, env)]
+    otlp_endpoint: Option<String>,
 }
 
 fn main() -> Result<(), std::io::Error> {
@@ -69,15 +78,6 @@ fn main() -> Result<(), std::io::Error> {
     // Get args
     let args = Args::parse();
 
-    if args.json_output {
-        tracing_subscriber::fmt()
-            .json()
-            .with_current_span(false)
-            .init();
-    } else {
-        tracing_subscriber::fmt().compact().init();
-    }
-
     // Validate args
     validate_args(&args);
 
@@ -104,6 +104,7 @@ fn main() -> Result<(), std::io::Error> {
         .build()
         .unwrap()
         .block_on(async {
+            init_logging(args.otlp_endpoint, args.json_output);
             // Instantiate sharded client from the master unix socket
             let mut sharded_client = ShardedClient::connect_uds(args.master_shard_uds_path)
                 .await
@@ -202,3 +203,54 @@ fn write_termination_log(msg: &str) -> Result<(), io::Error> {
     writeln!(f, "{}", msg)?;
     Ok(())
 }
+
+fn init_logging(otlp_endpoint: Option<String>, json_output: bool) {
+    let mut layers = Vec::new();
+
+    // STDOUT/STDERR layer
+    let fmt_layer = tracing_subscriber::fmt::layer()
+        .with_file(true)
+        .with_line_number(true);
+
+    let fmt_layer = match json_output {
+        true => fmt_layer.json().flatten_event(true).boxed(),
+        false => fmt_layer.boxed(),
+    };
+    layers.push(fmt_layer);
+
+    // OpenTelemetry tracing layer
+    if let Some(otlp_endpoint) = otlp_endpoint {
+        global::set_text_map_propagator(TraceContextPropagator::new());
+
+        let tracer = opentelemetry_otlp::new_pipeline()
+            .tracing()
+            .with_exporter(
+                opentelemetry_otlp::new_exporter()
+                    .tonic()
+                    .with_endpoint(otlp_endpoint),
+            )
+            .with_trace_config(
+                trace::config()
+                    .with_resource(Resource::new(vec![KeyValue::new(
+                        "service.name",
+                        "text-generation-inference.router",
+                    )]))
+                    .with_sampler(Sampler::AlwaysOn),
+            )
+            .install_batch(opentelemetry::runtime::Tokio);
+
+        if let Ok(tracer) = tracer {
+            layers.push(tracing_opentelemetry::layer().with_tracer(tracer).boxed());
+            axum_tracing_opentelemetry::init_propagator().unwrap();
+        };
+    }
+
+    // Filter events with LOG_LEVEL
+    let env_filter =
+        EnvFilter::try_from_env("LOG_LEVEL").unwrap_or_else(|_| EnvFilter::new("info"));
+
+    tracing_subscriber::registry()
+        .with(env_filter)
+        .with(layers)
+        .init();
+}