98 lines
3.0 KiB
Rust
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(())
|
|
}
|