From 3bf04aea3f24817a879a6fd1dd5cf9dd663584ab Mon Sep 17 00:00:00 2001
From: Nicolas Patry <patry.nicolas@protonmail.com>
Date: Mon, 1 May 2023 15:40:03 +0200
Subject: [PATCH] Adding the env cli (in launcher).

---
 launcher/build.rs        | 14 +++++++++++++
 launcher/src/env_cli.rs  | 43 ++++++++++++++++++++++++++++++++++++++++
 launcher/src/main.rs     | 13 ++++++++++++
 launcher/src/versions.rs |  1 +
 4 files changed, 71 insertions(+)
 create mode 100644 launcher/build.rs
 create mode 100644 launcher/src/env_cli.rs
 create mode 100644 launcher/src/versions.rs

diff --git a/launcher/build.rs b/launcher/build.rs
new file mode 100644
index 00000000..3d7b7c2d
--- /dev/null
+++ b/launcher/build.rs
@@ -0,0 +1,14 @@
+use std::fs::File;
+use std::io::Write;
+use std::process::Command;
+
+fn main() {
+    let output = Command::new("git")
+        .args(&["rev-parse", "HEAD"])
+        .output()
+        .unwrap();
+    let git_hash = String::from_utf8(output.stdout).unwrap().trim().to_string();
+    let mut file = File::create("src/versions.rs").unwrap();
+    file.write(format!(r#"pub static GIT_HASH: &str = "{git_hash}";"#).as_bytes())
+        .unwrap();
+}
diff --git a/launcher/src/env_cli.rs b/launcher/src/env_cli.rs
new file mode 100644
index 00000000..de78cbb5
--- /dev/null
+++ b/launcher/src/env_cli.rs
@@ -0,0 +1,43 @@
+use std::process::Command;
+
+pub fn nvidia_smi() -> Option<String> {
+    let output = Command::new("nvidia-smi").output().ok()?;
+    let nvidia_smi = String::from_utf8(output.stdout).ok()?;
+    let output = nvidia_smi.replace("\n", "\n   ");
+    Some(output.trim().to_string())
+}
+
+pub fn cargo_version() -> Option<String> {
+    let output = Command::new("cargo").args(&["version"]).output().ok()?;
+    let output = String::from_utf8(output.stdout).ok()?;
+    Some(output.trim().to_string())
+}
+
+pub fn docker_image() -> Option<String> {
+    let output = Command::new("docker")
+        .args(&[
+            "image",
+            "inspect",
+            "--format",
+            "{{.RepoDigests}}",
+            "ghcr.io/huggingface/text-generation-inference:latest",
+        ])
+        .output()
+        .ok()?;
+    let output = String::from_utf8(output.stdout).ok()?;
+    Some(output.trim().to_string())
+}
+
+pub fn print_env() {
+    println!(
+        "Cargo version: {}",
+        cargo_version().unwrap_or("N/A".to_string())
+    );
+    println!("Commit SHA: {}", crate::versions::GIT_HASH);
+    println!(
+        "Docker image sha: {}",
+        docker_image().unwrap_or("N/A".to_string())
+    );
+    let nvidia_smi = nvidia_smi().unwrap_or("N/A".to_string());
+    println!("Nvidia-smi:\n     {}", nvidia_smi);
+}
diff --git a/launcher/src/main.rs b/launcher/src/main.rs
index 77b0cf8a..63325f67 100644
--- a/launcher/src/main.rs
+++ b/launcher/src/main.rs
@@ -14,6 +14,9 @@ use std::time::{Duration, Instant};
 use std::{fs, io};
 use subprocess::{ExitStatus, Popen, PopenConfig, PopenError, Redirection};
 
+mod env_cli;
+mod versions;
+
 /// App Configuration
 #[derive(Parser, Debug)]
 #[clap(author, version, about, long_about = None)]
@@ -200,6 +203,11 @@ struct Args {
     watermark_gamma: Option<f32>,
     #[clap(long, env)]
     watermark_delta: Option<f32>,
+
+    /// Print a lot of information about your environment
+    /// and exits.
+    #[clap(long, short, action)]
+    env: bool,
 }
 
 #[derive(Debug)]
@@ -823,6 +831,11 @@ fn main() -> Result<(), LauncherError> {
     // Pattern match configuration
     let args = Args::parse();
 
+    if args.env {
+        env_cli::print_env();
+        return Ok(());
+    }
+
     if args.json_output {
         tracing_subscriber::fmt().json().init();
     } else {
diff --git a/launcher/src/versions.rs b/launcher/src/versions.rs
new file mode 100644
index 00000000..669e58ec
--- /dev/null
+++ b/launcher/src/versions.rs
@@ -0,0 +1 @@
+pub static GIT_HASH: &str = "61fd8c3c014d9ea31bc1df736fd276606ce29d66";
\ No newline at end of file