From 433373b2ec91a1a1fa0bd4059d23e2e90c8188c2 Mon Sep 17 00:00:00 2001
From: Hugo Larcher <hugo.larcher@huggingface.co>
Date: Tue, 28 Jan 2025 11:59:50 +0100
Subject: [PATCH] feat: Add backend name to telemetry

---
 backends/trtllm/src/looper.rs | 4 ++++
 backends/v2/src/backend.rs    | 4 ++++
 backends/v3/src/backend.rs    | 4 ++++
 router/src/infer/mod.rs       | 2 ++
 router/src/server.rs          | 1 +
 router/src/usage_stats.rs     | 3 +++
 6 files changed, 18 insertions(+)

diff --git a/backends/trtllm/src/looper.rs b/backends/trtllm/src/looper.rs
index b148ec30..5fed954f 100644
--- a/backends/trtllm/src/looper.rs
+++ b/backends/trtllm/src/looper.rs
@@ -340,4 +340,8 @@ impl Backend for TensorRtLlmBackendV2 {
     async fn health(&self, _: bool) -> bool {
         true
     }
+
+    fn name(&self) -> &'static str {
+        "TensorRT-LLM"
+    }
 }
diff --git a/backends/v2/src/backend.rs b/backends/v2/src/backend.rs
index cfe87f98..adca3d5d 100644
--- a/backends/v2/src/backend.rs
+++ b/backends/v2/src/backend.rs
@@ -108,6 +108,10 @@ impl Backend for BackendV2 {
     fn start_health(&self) -> bool {
         true
     }
+
+    fn name(&self) -> &'static str {
+        "tgi-v2"
+    }
 }
 
 /// Batching logic
diff --git a/backends/v3/src/backend.rs b/backends/v3/src/backend.rs
index 736301b3..98e8d76f 100644
--- a/backends/v3/src/backend.rs
+++ b/backends/v3/src/backend.rs
@@ -115,6 +115,10 @@ impl Backend for BackendV3 {
     fn start_health(&self) -> bool {
         true
     }
+
+    fn name(&self) -> &'static str {
+        "tgi-v3"
+    }
 }
 
 /// Batching logic
diff --git a/router/src/infer/mod.rs b/router/src/infer/mod.rs
index 6497d857..7eb8a41b 100644
--- a/router/src/infer/mod.rs
+++ b/router/src/infer/mod.rs
@@ -40,6 +40,8 @@ pub trait Backend {
     fn start_health(&self) -> bool {
         false
     }
+
+    fn name(&self) -> &'static str;
 }
 
 /// Inference struct
diff --git a/router/src/server.rs b/router/src/server.rs
index 2781f9fb..9e57af27 100644
--- a/router/src/server.rs
+++ b/router/src/server.rs
@@ -1898,6 +1898,7 @@ pub async fn run(
                 disable_grammar_support,
                 max_client_batch_size,
                 usage_stats_level,
+                backend.name(),
             );
             Some(usage_stats::UserAgent::new(reduced_args))
         }
diff --git a/router/src/usage_stats.rs b/router/src/usage_stats.rs
index 4139c4c5..bc0bede2 100644
--- a/router/src/usage_stats.rs
+++ b/router/src/usage_stats.rs
@@ -97,6 +97,7 @@ pub struct Args {
     disable_grammar_support: bool,
     max_client_batch_size: usize,
     usage_stats_level: UsageStatsLevel,
+    backend_name: &'static str
 }
 
 impl Args {
@@ -120,6 +121,7 @@ impl Args {
         disable_grammar_support: bool,
         max_client_batch_size: usize,
         usage_stats_level: UsageStatsLevel,
+        backend_name: &'static str
     ) -> Self {
         Self {
             model_config,
@@ -140,6 +142,7 @@ impl Args {
             disable_grammar_support,
             max_client_batch_size,
             usage_stats_level,
+            backend_name
         }
     }
 }