add docker label in build

.github/workflows/build.yaml

@@ -97,6 +97,7 @@ jobs:
           platforms: 'linux/amd64'
           build-args: |
             GIT_SHA=${{ env.GITHUB_SHA }}
+            DOCKER_LABEL=sha-${{ env.GITHUB_SHA_SHORT }}
           tags: ${{ steps.meta.outputs.tags }}
           labels: ${{ steps.meta.outputs.labels }}
           cache-from: type=registry,ref=registry.internal.huggingface.tech/api-inference/community/text-generation-inference:cache,mode=max

Dockerfile

@@ -13,6 +13,7 @@ RUN cargo chef prepare --recipe-path recipe.json
 FROM chef AS builder
 
 ARG GIT_SHA
+ARG DOCKER_LABEL
 
 RUN PROTOC_ZIP=protoc-21.12-linux-x86_64.zip && \
     curl -OL https://github.com/protocolbuffers/protobuf/releases/download/v21.12/$PROTOC_ZIP && \

launcher/build.rs

@@ -2,11 +2,28 @@ use std::error::Error;
 use vergen::EmitBuilder;
 
 fn main() -> Result<(), Box<dyn Error>> {
-    // Emit the instructions
-    EmitBuilder::builder()
-        .all_cargo()
-        .all_git()
-        .all_rustc()
-        .emit()?;
+    // Emit cargo and rustc compile time values
+    EmitBuilder::builder().all_cargo().all_rustc().emit()?;
+
+    // Try to get the git sha from the local git repository
+    if EmitBuilder::builder()
+        .fail_on_error()
+        .git_sha(false)
+        .emit()
+        .is_err()
+    {
+        // Unable to get the git sha
+        if let Ok(sha) = std::env::var("GIT_SHA") {
+            // Set it from an env var
+            println!("cargo:rustc-env=VERGEN_GIT_SHA={sha}");
+        }
+    }
+
+    // Set docker label if present
+    if let Ok(label) = std::env::var("DOCKER_LABEL") {
+        // Set it from an env var
+        println!("cargo:rustc-env=DOCKER_LABEL={label}");
+    }
+
     Ok(())
 }

launcher/src/env_cli.rs

@@ -1,37 +0,0 @@
-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 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!("Target: {}", env!("VERGEN_CARGO_TARGET_TRIPLE"));
-    println!("Cargo version: {}", env!("VERGEN_RUSTC_SEMVER"));
-    println!("Commit SHA: {}", env!("VERGEN_GIT_SHA"));
-    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);
-    println!("Command line used: <FILL IN>");
-    println!("OS: <FILL IN>");
-}

launcher/src/env_runtime.rs

@@ -0,0 +1,45 @@
+use std::fmt;
+use std::process::Command;
+
+pub(crate) struct Env {
+    cargo_target: &'static str,
+    cargo_version: &'static str,
+    git_sha: &'static str,
+    docker_label: &'static str,
+    nvidia_env: String,
+}
+
+impl Env {
+    pub fn new() -> Self {
+        let nvidia_env = nvidia_smi();
+
+        Self {
+            nvidia_env: nvidia_env.unwrap_or("N/A".to_string()),
+            cargo_target: env!("VERGEN_CARGO_TARGET_TRIPLE"),
+            cargo_version: env!("VERGEN_RUSTC_SEMVER"),
+            git_sha: option_env!("VERGEN_GIT_SHA").unwrap_or("N/A"),
+            docker_label: option_env!("DOCKER_LABEL").unwrap_or("N/A"),
+        }
+    }
+}
+
+impl fmt::Display for Env {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        writeln!(f, "Runtime environment:")?;
+
+        writeln!(f, "Target: {}", self.cargo_target)?;
+        writeln!(f, "Cargo version: {}", self.cargo_version)?;
+        writeln!(f, "Commit sha: {}", self.git_sha)?;
+        writeln!(f, "Docker label: {}", self.docker_label)?;
+        write!(f, "nvidia-smi:\n{}", self.nvidia_env)?;
+
+        Ok(())
+    }
+}
+
+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())
+}

launcher/src/main.rs

@@ -14,7 +14,7 @@ use std::time::{Duration, Instant};
 use std::{fs, io};
 use subprocess::{ExitStatus, Popen, PopenConfig, PopenError, Redirection};
 
-mod env_cli;
+mod env_runtime;
 
 /// App Configuration
 #[derive(Parser, Debug)]
@@ -203,8 +203,7 @@ struct Args {
     #[clap(long, env)]
     watermark_delta: Option<f32>,
 
-    /// Print a lot of information about your environment
-    /// and exits.
+    /// Display a lot of information about your runtime environment
     #[clap(long, short, action)]
     env: bool,
 }
@@ -830,17 +829,17 @@ 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 {
         tracing_subscriber::fmt().compact().init();
     }
 
+    if args.env {
+        let env_runtime = env_runtime::Env::new();
+        tracing::info!("{}", env_runtime);
+    }
+
     tracing::info!("{:?}", args);
 
     let num_shard = find_num_shards(args.sharded, args.num_shard);

router/build.rs

@@ -15,5 +15,12 @@ fn main() -> Result<(), Box<dyn Error>> {
             println!("cargo:rustc-env=VERGEN_GIT_SHA={sha}");
         }
     }
+
+    // Set docker label if present
+    if let Ok(label) = std::env::var("DOCKER_LABEL") {
+        // Set it from an env var
+        println!("cargo:rustc-env=DOCKER_LABEL={label}");
+    }
+
     Ok(())
 }

router/src/lib.rs

@@ -57,6 +57,8 @@ pub struct Info {
     pub version: &'static str,
     #[schema(nullable = true, example = "null")]
     pub sha: Option<&'static str>,
+    #[schema(nullable = true, example = "null")]
+    pub docker_label: Option<&'static str>,
 }
 
 #[derive(Clone, Debug, Deserialize, ToSchema)]

router/src/server.rs

@@ -635,6 +635,7 @@ pub async fn run(
         validation_workers,
         version: env!("CARGO_PKG_VERSION"),
         sha: option_env!("VERGEN_GIT_SHA"),
+        docker_label: option_env!("DOCKER_LABEL"),
     };
 
     // Create router