align all the linker search dependency

This commit is contained in:
Morgan Funtowicz 2024-07-22 14:14:57 +00:00
parent 6a9e925ec1
commit 40330c73f0
4 changed files with 68 additions and 36 deletions

1
Cargo.lock generated
View File

@ -3499,6 +3499,7 @@ dependencies = [
"cxx",
"cxx-build",
"log",
"pkg-config",
"text-generation-router",
"thiserror",
"tokenizers",

View File

@ -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"

View File

@ -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
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

View File

@ -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");