From ce394f5d69544f4cf77ca456d24fabf02a265663 Mon Sep 17 00:00:00 2001 From: OlivierDehaene <23298448+OlivierDehaene@users.noreply.github.com> Date: Tue, 2 May 2023 15:20:27 +0200 Subject: [PATCH] add docker label in build --- .github/workflows/build.yaml | 1 + Dockerfile | 1 + launcher/build.rs | 29 ++++++++++++++++++----- launcher/src/env_cli.rs | 37 ----------------------------- launcher/src/env_runtime.rs | 45 ++++++++++++++++++++++++++++++++++++ launcher/src/main.rs | 15 ++++++------ router/build.rs | 7 ++++++ router/src/lib.rs | 2 ++ router/src/server.rs | 1 + 9 files changed, 87 insertions(+), 51 deletions(-) delete mode 100644 launcher/src/env_cli.rs create mode 100644 launcher/src/env_runtime.rs diff --git a/.github/workflows/build.yaml b/.github/workflows/build.yaml index 67d1c730..289e4f67 100644 --- a/.github/workflows/build.yaml +++ b/.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 diff --git a/Dockerfile b/Dockerfile index c65d3dee..d59acc10 100644 --- a/Dockerfile +++ b/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 && \ diff --git a/launcher/build.rs b/launcher/build.rs index e777b049..71d2c0c5 100644 --- a/launcher/build.rs +++ b/launcher/build.rs @@ -2,11 +2,28 @@ use std::error::Error; use vergen::EmitBuilder; fn main() -> Result<(), Box> { - // 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(()) } diff --git a/launcher/src/env_cli.rs b/launcher/src/env_cli.rs deleted file mode 100644 index 29243293..00000000 --- a/launcher/src/env_cli.rs +++ /dev/null @@ -1,37 +0,0 @@ -use std::process::Command; - -pub fn nvidia_smi() -> Option { - 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 { - 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: "); - println!("OS: "); -} diff --git a/launcher/src/env_runtime.rs b/launcher/src/env_runtime.rs new file mode 100644 index 00000000..9dbc83f7 --- /dev/null +++ b/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 { + 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()) +} diff --git a/launcher/src/main.rs b/launcher/src/main.rs index 87ff47b8..736aa5a7 100644 --- a/launcher/src/main.rs +++ b/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, - /// 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); diff --git a/router/build.rs b/router/build.rs index 1b1fdc86..f5eb8a26 100644 --- a/router/build.rs +++ b/router/build.rs @@ -15,5 +15,12 @@ fn main() -> Result<(), Box> { 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(()) } diff --git a/router/src/lib.rs b/router/src/lib.rs index c2ff669b..080dc4f4 100644 --- a/router/src/lib.rs +++ b/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)] diff --git a/router/src/server.rs b/router/src/server.rs index f25fa2b3..c5fc41ef 100644 --- a/router/src/server.rs +++ b/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