From 9fc54cd91c4e37aa8d298febf9b60004463f76b2 Mon Sep 17 00:00:00 2001
From: erikkaum <erik.kaum@gmail.com>
Date: Tue, 16 Jul 2024 14:14:36 +0200
Subject: [PATCH] more robust way of checking if is in container

---
 router/src/main.rs        |  6 +++---
 router/src/usage_stats.rs | 18 ++++++++++++++++++
 2 files changed, 21 insertions(+), 3 deletions(-)

diff --git a/router/src/main.rs b/router/src/main.rs
index f6e04123..ee02d676 100644
--- a/router/src/main.rs
+++ b/router/src/main.rs
@@ -382,10 +382,10 @@ async fn main() -> Result<(), RouterError> {
         }
     };
 
-    // Only send usage stats when TGI is run in docker
-    let is_docker = option_env!("DOCKER_LABEL").is_some();
+    // Only send usage stats when TGI is run in container
+    let is_container = usage_stats::is_container();
 
-    let user_agent = if !disable_usage_stats && is_docker {
+    let user_agent = if !disable_usage_stats && is_container {
         let reduced_args = usage_stats::Args::new(
             config.clone(),
             tokenizer_class,
diff --git a/router/src/usage_stats.rs b/router/src/usage_stats.rs
index dda58bd8..929765c8 100644
--- a/router/src/usage_stats.rs
+++ b/router/src/usage_stats.rs
@@ -233,3 +233,21 @@ fn xpu_smi() -> Option<String> {
     let output = xpu_smi.replace('\n', "\n   ");
     Some(output.trim().to_string())
 }
+
+pub fn is_container() -> io::Result<bool> {
+    let path = Path::new("/proc/self/cgroup");
+    let file = File::open(&path)?;
+    let reader = io::BufReader::new(file);
+
+    for line in reader.lines() {
+        let line = line?;
+        // Check for common container runtimes
+        if line.contains("/docker/") || line.contains("/docker-") ||
+           line.contains("/kubepods/") || line.contains("/kubepods-") ||
+           line.contains("containerd") || line.contains("crio") ||
+           line.contains("podman") {
+            return Ok(true);
+        }
+    }
+    Ok(false)
+}