20231088/vllm
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[CI/Build] update default number of jobs and nvcc threads to avoid overloading the system (#3675)

Browse Source
This commit is contained in:
youkaichao 2024-03-27 21:18:54 -07:00 committed by GitHub
parent 098e1776ba
commit 3492859b68
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 15 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

20
setup.py
View File

@ -56,6 +56,8 @@ class cmake_build_ext(build_ext):
# Determine number of compilation jobs and optionally nvcc compile threads.
#
def compute_num_jobs(self):
# `num_jobs` is either the value of the MAX_JOBS environment variable
# (if defined) or the number of CPUs available.
num_jobs = os.environ.get("MAX_JOBS", None)
if num_jobs is not None:
num_jobs = int(num_jobs)
@ -69,11 +71,19 @@ class cmake_build_ext(build_ext):
num_jobs = os.cpu_count()
nvcc_threads = None
if _is_cuda():
nvcc_cuda_version = get_nvcc_cuda_version()
if nvcc_cuda_version >= Version("11.2"):
nvcc_threads = int(os.getenv("NVCC_THREADS", 8))
num_jobs = max(1, round(num_jobs / (nvcc_threads / 4)))
if _is_cuda() and get_nvcc_cuda_version() >= Version("11.2"):
# `nvcc_threads` is either the value of the NVCC_THREADS
# environment variable (if defined) or 1.
# when it is set, we reduce `num_jobs` to avoid
# overloading the system.
nvcc_threads = os.getenv("NVCC_THREADS", None)
if nvcc_threads is not None:
nvcc_threads = int(nvcc_threads)
logger.info(f"Using NVCC_THREADS={nvcc_threads} as the number"
" of nvcc threads.")
else:
nvcc_threads = 1
num_jobs = max(1, num_jobs // nvcc_threads)
return num_jobs, nvcc_threads