From 40330c73f0ad6daee75ead64d6815b44c371a5ad Mon Sep 17 00:00:00 2001 From: Morgan Funtowicz Date: Mon, 22 Jul 2024 14:14:57 +0000 Subject: [PATCH] align all the linker search dependency --- Cargo.lock | 1 + backends/trtllm/Cargo.toml | 2 +- backends/trtllm/Dockerfile | 19 ++++++--- backends/trtllm/build.rs | 82 ++++++++++++++++++++++++-------------- 4 files changed, 68 insertions(+), 36 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index cbcd1955..6970b73d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3499,6 +3499,7 @@ dependencies = [ "cxx", "cxx-build", "log", + "pkg-config", "text-generation-router", "thiserror", "tokenizers", diff --git a/backends/trtllm/Cargo.toml b/backends/trtllm/Cargo.toml index a5f09c11..870a3c88 100644 --- a/backends/trtllm/Cargo.toml +++ b/backends/trtllm/Cargo.toml @@ -22,5 +22,5 @@ log = { version = "0.4", features = [] } [build-dependencies] cmake = "0.1" -cxx-build = "1.0" +cxx-build = { version = "1.0", features = ["parallel"] } pkg-config = "0.3" \ No newline at end of file diff --git a/backends/trtllm/Dockerfile b/backends/trtllm/Dockerfile index 5eb04815..24ba48da 100644 --- a/backends/trtllm/Dockerfile +++ b/backends/trtllm/Dockerfile @@ -1,4 +1,6 @@ ARG CUDA_ARCH_LIST="75-real;80-real;86-real;89-real;90-real" +ARG INSTALL_PREFIX="/usr/local/tgi +ARG TENSORRT_ROOT_DIR="/usr/local/tensorrt # Build dependencies resolver stage FROM lukemathwalker/cargo-chef:latest as chef @@ -29,21 +31,28 @@ RUN curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | bash -s -- -y && chmod -R a+w $HOME/.rustup && \ chmod -R a+w $HOME/.cargo -ENV PATH="$HOME/.cargo/bin:$PATH" -RUN $HOME/.cargo/bin/cargo install cargo-chef +ENV PATH="/root/.cargo/bin:$PATH" +RUN cargo install cargo-chef # Backend build step WORKDIR /usr/src/text-generation-inference # Cache dependencies COPY --from=planner /usr/src/text-generation-inference/recipe.json . -RUN $HOME/.cargo/bin/cargo chef cook --release --recipe-path recipe.json +RUN cargo chef cook --release --recipe-path recipe.json + +ENV LD_LIBRARY_PATH="/usr/local/mpi/lib:$LD_LIBRARY_PATH" # Build actual TGI +ENV PKG_CONFIG_PATH="/usr/local/mpi/lib:$PKG_CONFIG_PATH" +ENV CMAKE_INSTALL_PREFIX="/usr/local/tgi" COPY . . -RUN $HOME/.cargo/bin/cargo build --release --bin text-generation-backends-trtllm +RUN mkdir /usr/local/tgi && mkdir /usr/local/tgi/include && mkdir /usr/local/tgi/lib && \ + cargo build --release --bin text-generation-backends-trtllm FROM nvcr.io/nvidia/pytorch:24.05-py3 WORKDIR /opt + COPY --from=cuda-builder /usr/local/tensorrt /usr/local/tensorrt -COPY --from=cuda-builder /usr/src/text-generation-inference/target/release/text-generation-backends-trtllm /opt/text-generation-launcher \ No newline at end of file +COPY --from=cuda-builder /usr/local/tgi /usr/local/tgi +COPY --from=cuda-builder /usr/src/text-generation-inference/target/release/text-generation-backends-trtllm /usr/local/tgi/bin/text-generation-launcher diff --git a/backends/trtllm/build.rs b/backends/trtllm/build.rs index c1c25ae7..d69565c9 100644 --- a/backends/trtllm/build.rs +++ b/backends/trtllm/build.rs @@ -1,5 +1,5 @@ use std::env; -use std::path::PathBuf; +use std::path::{absolute, PathBuf}; use cxx_build::CFG; use pkg_config; @@ -14,23 +14,24 @@ const TENSORRT_ROOT_DIR: Option<&str> = option_env!("TENSORRT_ROOT_DIR"); macro_rules! probe { ($name: literal, $version: expr) => { if let Err(_) = pkg_config::probe_library($name) { - pkg_config::probe_library(&format!("cuda-{}", $version)) + pkg_config::probe_library(&format!("{}-{}", $name, $version)) .expect(&format!("Failed to locate {}", $name)); } }; } -fn main() { - // Misc variables - let out_dir = PathBuf::from(env::var("OUT_DIR").unwrap()); - let build_profile = env::var("PROFILE").unwrap(); - let (is_debug, opt_level) = match build_profile.as_ref() { - "debug" => (true, "0"), - _ => (false, "3"), - }; - +fn build_backend(is_debug: bool, opt_level: &str, out_dir: &PathBuf) -> (PathBuf, PathBuf) { // Build the backend implementation through CMake - let backend_path = cmake::Config::new(".") + let install_path = INSTALL_PREFIX.unwrap_or("/usr/local/tgi"); + let tensorrt_path = TENSORRT_ROOT_DIR.unwrap_or("/usr/local/tensorrt"); + let cuda_arch_list = CUDA_ARCH_LIST.unwrap_or("90-real"); // Hopper by default + + let mut install_path = PathBuf::from(install_path); + if !install_path.is_absolute() { + install_path = absolute(out_dir).expect("cannot happen").join(install_path); + } + + let _ = cmake::Config::new(".") .uses_cxx11() .generator("Ninja") .profile(match is_debug { @@ -38,31 +39,37 @@ fn main() { false => "Release", }) .env("OPT_LEVEL", opt_level) - .out_dir(INSTALL_PREFIX.unwrap_or("/usr/local/tgi")) + .define("CMAKE_INSTALL_PREFIX", &install_path) .define("CMAKE_CUDA_COMPILER", "/usr/local/cuda/bin/nvcc") - .define( - "TGI_TRTLLM_BACKEND_TARGET_CUDA_ARCH_LIST", - CUDA_ARCH_LIST.unwrap_or("90-real"), // Hopper by default - ) - .define( - "TGI_TRTLLM_BACKEND_TRT_ROOT", - TENSORRT_ROOT_DIR.unwrap_or("/usr/local/tensorrt"), - ) + .define("TGI_TRTLLM_BACKEND_TARGET_CUDA_ARCH_LIST", cuda_arch_list) + .define("TGI_TRTLLM_BACKEND_TRT_ROOT", tensorrt_path) .build(); // Additional transitive CMake dependencies let deps_folder = out_dir.join("build").join("_deps"); for dependency in ADDITIONAL_BACKEND_LINK_LIBRARIES { - let dep_name = match build_profile.as_ref() { - "debug" => format!("{}d", dependency), - _ => String::from(dependency), + let dep_name = match is_debug { + true => format!("{}d", dependency), + false => String::from(dependency), }; let dep_path = deps_folder.join(format!("{}-build", dependency)); println!("cargo:rustc-link-search={}", dep_path.display()); println!("cargo:rustc-link-lib=static={}", dep_name); } - // Build the FFI layer calling the backend above + // Emit linkage information from the artifacts we just built + let install_lib_path = install_path.join("lib"); + + println!( + r"cargo:warning=Adding link search path: {}", + install_lib_path.display() + ); + println!(r"cargo:rustc-link-search={}", install_lib_path.display()); + + (PathBuf::from(install_path), deps_folder) +} + +fn build_ffi_layer(deps_folder: &PathBuf) { CFG.include_prefix = "backends/trtllm"; cxx_build::bridge("src/lib.rs") .static_flag(true) @@ -81,6 +88,22 @@ fn main() { println!("cargo:rerun-if-changed=lib/backend.cpp"); println!("cargo:rerun-if-changed=include/ffi.h"); println!("cargo:rerun-if-changed=src/ffi.cpp"); +} + +fn main() { + // Misc variables + let out_dir = PathBuf::from(env::var("OUT_DIR").unwrap()); + let build_profile = env::var("PROFILE").unwrap(); + let (is_debug, opt_level) = match build_profile.as_ref() { + "debug" => (true, "0"), + _ => (false, "3"), + }; + + // Build the backend + let (_backend_path, deps_folder) = build_backend(is_debug, opt_level, &out_dir); + + // Build the FFI layer calling the backend above + build_ffi_layer(&deps_folder); // Emit linkage search path probe!("ompi", MPI_REQUIRED_VERSION); @@ -92,14 +115,13 @@ fn main() { probe!("nvidia-ml", CUDA_REQUIRED_VERSION); // TensorRT - println!(r"cargo:rustc-link-search=native=/usr/local/tensorrt/lib"); + println!( + r"cargo:rustc-link-search=native={}", + TENSORRT_ROOT_DIR.unwrap_or("/usr/local/tensorrt/lib") + ); println!("cargo:rustc-link-lib=dylib=nvinfer"); // TensorRT-LLM - println!( - r"cargo:rustc-link-search=native={}", - backend_path.join("lib").display() - ); println!("cargo:rustc-link-lib=dylib=tensorrt_llm"); println!("cargo:rustc-link-lib=static=tensorrt_llm_executor_static"); println!("cargo:rustc-link-lib=dylib=nvinfer_plugin_tensorrt_llm");