mirror of
https://github.com/NVIDIA/TensorRT-LLM.git
synced 2026-01-14 06:27:45 +08:00
89 lines
2.8 KiB
Python
89 lines
2.8 KiB
Python
import argparse
|
|
import contextlib
|
|
import logging
|
|
import os
|
|
import platform
|
|
import shutil
|
|
import subprocess
|
|
from os import PathLike
|
|
from pathlib import Path
|
|
|
|
|
|
@contextlib.contextmanager
|
|
def working_directory(path: PathLike):
|
|
"""Changes working directory and returns to previous on exit."""
|
|
prev_cwd = Path.cwd()
|
|
os.chdir(path)
|
|
try:
|
|
yield
|
|
finally:
|
|
os.chdir(prev_cwd)
|
|
|
|
|
|
def build_cpp_examples(build_dir: PathLike, trt_dir: PathLike,
|
|
enable_multi_device: str, loglevel: int) -> None:
|
|
logging.basicConfig(level=loglevel,
|
|
format='%(asctime)s - %(levelname)s - %(message)s')
|
|
# Convert input paths to pathlib.Path objects
|
|
build_dir = Path(build_dir)
|
|
trt_dir = Path(trt_dir)
|
|
|
|
assert trt_dir.is_dir()
|
|
|
|
def cmake_parse(path: PathLike) -> str:
|
|
return str(path).replace("\\", "/")
|
|
|
|
# Remove the build directory if it exists
|
|
if build_dir.exists():
|
|
logging.info(f"Removed directory: {build_dir}")
|
|
shutil.rmtree(build_dir)
|
|
|
|
# Create the build directory
|
|
build_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
# Change to the build directory
|
|
with working_directory(build_dir):
|
|
# Run CMake with the specified TensorRT directories
|
|
generator = ["-GNinja"] if platform.system() == "Windows" else []
|
|
generate_command = [
|
|
'cmake',
|
|
'-S',
|
|
'..',
|
|
'-B',
|
|
'.',
|
|
f'-DTensorRT_ROOT={cmake_parse(trt_dir)}',
|
|
f'-DENABLE_MULTI_DEVICE={enable_multi_device}',
|
|
] + generator
|
|
logging.info(f"Executing {generate_command}")
|
|
subprocess.run(generate_command, check=True)
|
|
|
|
# Build the project using make
|
|
build_command = ["cmake", "--build", ".", "--config", "Release"]
|
|
logging.info(f"Executing {build_command}")
|
|
subprocess.run(build_command, check=True)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
parser = argparse.ArgumentParser(description='Build C++ examples')
|
|
parser.add_argument('--build-dir',
|
|
default='examples/cpp/executor/build',
|
|
help='Build directory path')
|
|
parser.add_argument('--trt-dir',
|
|
default='/usr/local/tensorrt',
|
|
help='TensorRT directory path')
|
|
parser.add_argument('--enable-multi-device',
|
|
default='ON',
|
|
help='Enable multi device support (requires MPI)')
|
|
parser.add_argument('-v',
|
|
'--verbose',
|
|
help="verbose",
|
|
action="store_const",
|
|
dest="loglevel",
|
|
const=logging.DEBUG,
|
|
default=logging.INFO)
|
|
cli = parser.parse_args()
|
|
|
|
args = vars(cli)
|
|
print(args) # Log on Jenkins instance.
|
|
build_cpp_examples(**args)
|