hf_text-generation-inference/backends/trtllm/build.rs

98 lines
3.0 KiB
Rust

use std::env;
use std::path::{Path, PathBuf};
use anyhow::{anyhow, Error};
use git2::{Error as GitError, Repository, SubmoduleUpdateOptions};
use git2::build::CheckoutBuilder;
// const ENV_TGI_TENSORRT_LLM_VERSION: &str = "TGI_TRTLLM_VERSION";
const ENV_TGI_TENSORRT_LLM_VERSION_COMMIT: &str = "TGI_TRTLLM_VERSION_COMMIT";
const TENSORRT_LLM_REPOSITORY_URL: &str = "https://github.com/nvidia/tensorrt-llm";
const TENSORRT_LLM_REPOSITORY_COMMIT_HASH: &str = "9691e12bce7ae1c126c435a049eb516eb119486c";
macro_rules! log {
($fmt:expr) => (println!(concat!("trtllm-sys/build.rs:{}: ", $fmt), line!()));
($fmt:expr, $($arg:tt)*) => (println!(concat!("trtllm-sys/build.rs:{}: ", $fmt),
line!(), $($arg)*));
}
fn get_trtllm_snapshot<P: AsRef<Path>>(
origin: &str,
version: &str,
dest: P,
) -> Result<(), GitError> {
let dest = dest.as_ref();
let repo = if dest.join(".git").exists() {
Repository::open(dest)
} else {
Repository::init(dest)
}?;
let mut remote = repo.remote_anonymous(origin)?;
remote.fetch(&["main"], None, None)?;
log!("Successfully fetched main on remote");
let mut fetch_head = repo.find_reference("FETCH_HEAD")?;
let commit = repo.find_commit_by_prefix(version)?;
log!("Targeting commit {:?}", commit);
let target = fetch_head.set_target(
commit.id(),
&format!("Fast-Forward to {TENSORRT_LLM_REPOSITORY_COMMIT_HASH}"),
)?;
repo.set_head(target.name().unwrap())?;
repo.checkout_head(Some(CheckoutBuilder::default().force()))?;
for mut submodule in repo.submodules()? {
submodule.sync()?;
submodule.update(true, Some(SubmoduleUpdateOptions::new().allow_fetch(true)))?
}
Ok(())
}
fn cmake_backend_build<P: AsRef<Path>>(trtllm_sources_dir: P, dest: P, profile: &str) {
cmake::Config::new(PathBuf::from("csrc"))
.profile(match profile {
"release" => "Release",
_ => "DebWithRelInfo",
})
.generator("Ninja")
.define(
"TGI_TRTLLM_BACKEND_TRTLLM_ROOT_DIR",
trtllm_sources_dir.as_ref().as_os_str(),
)
.uses_cxx11()
.build_target("tgi_trtllm_backend")
.out_dir(dest)
.build();
}
fn main() -> Result<(), Error> {
env::set_var(
ENV_TGI_TENSORRT_LLM_VERSION_COMMIT,
TENSORRT_LLM_REPOSITORY_COMMIT_HASH,
);
let tensorrt_llm_sources_dir = PathBuf::from(env::var("OUT_DIR").unwrap()).join("tensorrt-llm");
let backend_out_dir = PathBuf::from(env::var("OUT_DIR").unwrap()).join("tgi-trtllm-backend");
let build_profile: String = env::var("PROFILE").unwrap();
// First we need to retrieve TensorRT-LLM sources to build the library to link against
get_trtllm_snapshot(
TENSORRT_LLM_REPOSITORY_URL,
TENSORRT_LLM_REPOSITORY_COMMIT_HASH,
&tensorrt_llm_sources_dir,
)
.map_err(|e| anyhow!(e))?;
cmake_backend_build(&tensorrt_llm_sources_dir, &backend_out_dir, &build_profile);
Ok(())
}