Skip to content

Commit e9151ec

Browse files
zhzhcookiezhzhcookie
authored
[BUILD] Add unpack .triton shell for offline-building on triton_v3.3.x (#366)
* [BUILD] Add unpack .triton shell for offline-building * [Build] Print color when downloading
1 parent 38760e5 commit e9151ec

3 files changed

Lines changed: 125 additions & 3 deletions

File tree

‎python/scripts/unpack_triton_build_deps.sh‎

Lines changed: 103 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,103 @@
1+
#!/bin/bash
2+
3+
set -e
4+
5+
# Print
6+
if [ "$OS" = "Windows_NT" ]; then
7+
RED=''
8+
GREEN=''
9+
YELLOW=''
10+
NC=''
11+
else
12+
RED='\033[1;31m'
13+
GREEN='\033[0;32m'
14+
YELLOW='\033[1;33m'
15+
NC='\033[0m'
16+
fi
17+
INFO="${GREEN}[INFO]${NC}"
18+
NOTE="${YELLOW}[NOTE]${NC}"
19+
ERROR="${RED}[ERROR]${NC}"
20+
21+
printfln() {
22+
printf "%b
23+
" "$@"
24+
}
25+
26+
# Input
27+
if [ $# -ge 1 ] && [ -f "$1" ]; then
28+
input_tar_gz="$1"
29+
printfln "${INFO} Use ${input_tar_gz} as input packed .tar.gz file"
30+
else
31+
printfln "${ERROR} No input .tar.gz file specified"
32+
printfln "${INFO} Usage: sh $0 [input_tar_gz]"
33+
exit 1
34+
fi
35+
36+
# Check system architecture
37+
file_arch="${input_tar_gz##*-}" # x64.tar.gz
38+
file_arch="${file_arch%%.tar.gz}" # x64
39+
sys_arch="$(uname -m)"
40+
case "${sys_arch}" in
41+
x86_64|amd64) sys_arch="x64" ;;
42+
aarch64|arm64) sys_arch="aarch64" ;;
43+
esac
44+
if [ "${file_arch}" != "${sys_arch}" ]; then
45+
printfln "${ERROR} Arch mismatch: input_file=${RED}${file_arch}${NC}, system=${RED}${sys_arch}${NC}"
46+
exit 1
47+
fi
48+
49+
# Output
50+
if [ "${TRITON_HOME}" != "" ]; then
51+
output_dir="${TRITON_HOME}"
52+
printfln "${INFO} Use ${output_dir} as output directory because TRITON_HOME is set"
53+
else
54+
output_dir="${HOME}/.triton"
55+
printfln "${INFO} Use ${output_dir} as default output directory"
56+
fi
57+
58+
if [ -d "${output_dir}" ]; then
59+
last_output_dir=${output_dir}.$(date +%Y%m%d_%H%M%S)
60+
if [ -d "${last_output_dir}" ]; then
61+
printfln "${ERROR} Backup directory ${RED}${last_output_dir}${NC} already exists, retrying will resolve it"
62+
exit 1
63+
fi
64+
printfln "${NOTE} Output directory ${YELLOW}${output_dir}${NC} already exists, will mv to ${YELLOW}${last_output_dir}${NC}"
65+
fi
66+
67+
# Check unpack dirs
68+
printfln "${NOTE} Will unpack following dirs to ${YELLOW}${output_dir}${NC} (will be created):"
69+
tar tzf "${input_tar_gz}" \
70+
| awk -F'/' '{
71+
sub(/^\.triton\/?/, ""); if ($0 == "") next
72+
if ($1 == "nvidia") {
73+
if (NF >= 2 && $2 != "") {
74+
print "nvidia/"$2"/"
75+
} else {
76+
print "nvidia/"
77+
}
78+
} else {
79+
print $1"/"
80+
}
81+
}' \
82+
| uniq
83+
printfln "${NOTE} Press any key to confirm and continue, or Ctrl+C to cancel ..."
84+
read dummy
85+
86+
# Create output dir
87+
if [ -d "${output_dir}" ]; then
88+
set -x
89+
mv "${output_dir}" "${last_output_dir}"
90+
{ set +x; } 2>/dev/null
91+
fi
92+
set -x
93+
mkdir -p "${output_dir}"
94+
{ set +x; } 2>/dev/null
95+
printfln ""
96+
97+
# Unpack
98+
printfln "${NOTE} Unpacking ${YELLOW}${input_tar_gz}${NC} into ${YELLOW}${output_dir}${NC}"
99+
set -x
100+
tar zxf "${input_tar_gz}" -C "${output_dir}" --strip-components=1
101+
{ set +x; } 2>/dev/null
102+
103+
printfln "${INFO} Finished successfully."

‎python/setup.py‎

Lines changed: 12 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,14 @@
3131
from build_helpers import get_base_dir, get_cmake_dir
3232
from setup_tools import setup_helper as helper
3333

34+
# flagtree setup print config
35+
if platform.system() == "Windows":
36+
YELLOW = ""
37+
RESET = ""
38+
else:
39+
YELLOW = "\033[1;33m"
40+
NC = "\033[0m"
41+
3442

3543
@dataclass
3644
class Backend:
@@ -301,7 +309,8 @@ def get_thirdparty_packages(packages: list):
301309
with contextlib.suppress(Exception):
302310
shutil.rmtree(package_root_dir)
303311
os.makedirs(package_root_dir, exist_ok=True)
304-
print(f'downloading and extracting {p.url} ...')
312+
print(f'{YELLOW}downloading and extracting {p.url} to {package_root_dir} ... {NC}', file=sys.stderr,
313+
flush=True)
305314
with open_url(p.url) as response:
306315
if p.url.endswith(".zip"):
307316
file_bytes = BytesIO(response.read())
@@ -346,11 +355,11 @@ def download_and_copy(name, src_func, dst_path, variable, version, url_func):
346355
assert curr_version is not None, f"No version information for {dst_path}"
347356
download = download or curr_version.group(1) != version
348357
if download and not is_offline_build():
349-
print(f'downloading and extracting {url} ...')
358+
print(f'{YELLOW}downloading and extracting {url} ... {NC}', file=sys.stderr, flush=True)
350359
file = tarfile.open(fileobj=open_url(url), mode="r|*")
351360
file.extractall(path=tmp_path)
352361
os.makedirs(os.path.split(dst_path)[0], exist_ok=True)
353-
print(f'copy {src_path} to {dst_path} ...')
362+
print(f'copy {src_path} to {dst_path} ...', file=sys.stderr, flush=True)
354363
if os.path.isdir(src_path):
355364
shutil.copytree(src_path, dst_path, dirs_exist_ok=True)
356365
else:

‎unittest/googletest.cmake‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,16 @@ include(FetchContent)
22

33
set(GOOGLETEST_DIR "" CACHE STRING "Location of local GoogleTest repo to build against")
44

5+
if(NOT GOOGLETEST_DIR)
6+
set(_default_googletest "$ENV{HOME}/.triton/googletest-release-1.12.1")
7+
if(EXISTS "${_default_googletest}")
8+
set(GOOGLETEST_DIR "${_default_googletest}" CACHE STRING "Location of local GoogleTest repo to build against" FORCE)
9+
else()
10+
message("Dowloading https://github.com/google/googletest/archive/refs/tags/release-1.12.1.tar.gz")
11+
endif()
12+
unset(_default_googletest)
13+
endif()
14+
515
if(GOOGLETEST_DIR)
616
set(FETCHCONTENT_SOURCE_DIR_GOOGLETEST ${GOOGLETEST_DIR} CACHE STRING "GoogleTest source directory override")
717
endif()

0 commit comments

Comments
 (0)