Add distributed open telemetry tracing (#6)

We would like to add distributed open-telemetry tracing to the fmaas-router. This requires 3 changes on the fmaas-router: 1. Enabling open-telemetry tracing in the fmaas-router (this is done by the `init_logging` function in the newly created `tracing.rs` file. 2. Propagating trace context information in the GRPC metadata to the TGIS router. This is done with helper functions `inject_context_span` etc. defined in the same file and used in the `server.rs` file for the `/generate`, `/generateStream` and `/tokenize` endpoints. Span tags have been added following the [Semantic Convention for RPC calls](https://opentelemetry.io/docs/specs/semconv/rpc/rpc-spans/) 3. Receives trace context information from wx-inference-proxy. This is done with helper function `extract_context_spa Co-authored-by: Kavya <[email protected]>
IBM · Mar 28, 2024 · abf31c0 · abf31c0
1 parent 8efd754
commit abf31c0
Show file tree

Hide file tree

Showing 6 changed files with 388 additions and 29 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/fmaas-router/Cargo.toml b/fmaas-router/Cargo.toml
@@ -27,6 +27,10 @@ prost = "^0.12.3"
 prost-types = "^0.12.3"
 serde_yaml = "^0.9.33"
 serde = { version = "^1.0.197", features = ["derive"] }
+opentelemetry = { version = "0.22", features = ["trace"] }
+opentelemetry_sdk = {version = "0.22", features = ["rt-tokio"]}
+opentelemetry-otlp = "0.15.0"
+tracing-opentelemetry = "0.23.0"
 
 mio = "^0.8.11" # Override to address CVE-2024-27308
 rustls-webpki = "^0.102.2" # Override to address WS-2023-0305, CVE-2018-16875

diff --git a/fmaas-router/src/lib.rs b/fmaas-router/src/lib.rs
@@ -11,6 +11,7 @@ use tracing::info;
 mod pb;
 pub mod rpc;
 pub mod server;
+pub mod tracing_utils;
 
 #[derive(Debug, Clone, Deserialize)]
 pub struct ServiceAddr {
@@ -118,4 +119,4 @@ async fn create_clients<C>(
         .expect("Error creating upstream service clients")
         .into_iter()
         .collect()
-}
+}
diff --git a/fmaas-router/src/main.rs b/fmaas-router/src/main.rs
@@ -1,8 +1,7 @@
 use std::net::{IpAddr, Ipv4Addr, SocketAddr};
 
 use clap::Parser;
-use fmaas_router::{server, ModelMap};
-use tracing_subscriber::EnvFilter;
+use fmaas_router::{server, tracing_utils::init_logging, ModelMap};
 
 /// App Configuration
 #[derive(Parser, Debug)]
@@ -28,30 +27,16 @@ struct Args {
     upstream_tls: bool,
     #[clap(long, env)]
     upstream_tls_ca_cert_path: Option<String>,
+    #[clap(long, env = "OTEL_EXPORTER_OTLP_ENDPOINT")]
+    otlp_endpoint: Option<String>,
+    #[clap(long, env = "OTEL_SERVICE_NAME", default_value = "fmaas-router")]
+    otlp_service_name: String,
 }
 
 fn main() -> Result<(), std::io::Error> {
     //Get args
     let args = Args::parse();
 
-    // Configure log level; use info by default
-    let filter_layer = EnvFilter::try_from_default_env()
-        .or_else(|_| EnvFilter::try_new("info"))
-        .unwrap();
-
-    if args.json_output {
-        tracing_subscriber::fmt()
-            .json()
-            .with_env_filter(filter_layer)
-            .with_current_span(false)
-            .init();
-    } else {
-        tracing_subscriber::fmt()
-            .compact()
-            .with_env_filter(filter_layer)
-            .init();
-    }
-
     if args.tls_key_path.is_some() != args.tls_cert_path.is_some() {
         panic!("tls: must provide both cert and key")
     }
@@ -71,6 +56,8 @@ fn main() -> Result<(), std::io::Error> {
             let grpc_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0)), args.grpc_port);
             let http_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0)), args.probe_port);
 
+            init_logging(args.otlp_service_name, args.json_output, args.otlp_endpoint);
+
             server::run(
                 grpc_addr,
                 http_addr,