Skip to content
Snippets Groups Projects
Unverified Commit 29390688 authored by Zaida Zhou's avatar Zaida Zhou Committed by GitHub
Browse files

[Feats]: Add ROCm info when collecting env (#633)

* Add ROCm info when collecting env

* minor fix
parent 06e077d7
No related branches found
No related tags found
No related merge requests found
...@@ -67,15 +67,27 @@ def collect_env(): ...@@ -67,15 +67,27 @@ def collect_env():
env_info['CUDA_HOME'] = CUDA_HOME env_info['CUDA_HOME'] = CUDA_HOME
if CUDA_HOME is not None and osp.isdir(CUDA_HOME): if CUDA_HOME is not None and osp.isdir(CUDA_HOME):
try: if CUDA_HOME == '/opt/rocm':
nvcc = osp.join(CUDA_HOME, 'bin/nvcc') try:
nvcc = subprocess.check_output(f'"{nvcc}" -V', shell=True) nvcc = osp.join(CUDA_HOME, 'hip/bin/hipcc')
nvcc = nvcc.decode('utf-8').strip() nvcc = subprocess.check_output(
release = nvcc.rfind('Cuda compilation tools') f'"{nvcc}" --version', shell=True)
build = nvcc.rfind('Build ') nvcc = nvcc.decode('utf-8').strip()
nvcc = nvcc[release:build].strip() release = nvcc.rfind('HIP version:')
except subprocess.SubprocessError: build = nvcc.rfind('')
nvcc = 'Not Available' nvcc = nvcc[release:build].strip()
except subprocess.SubprocessError:
nvcc = 'Not Available'
else:
try:
nvcc = osp.join(CUDA_HOME, 'bin/nvcc')
nvcc = subprocess.check_output(f'"{nvcc}" -V', shell=True)
nvcc = nvcc.decode('utf-8').strip()
release = nvcc.rfind('Cuda compilation tools')
build = nvcc.rfind('Build ')
nvcc = nvcc[release:build].strip()
except subprocess.SubprocessError:
nvcc = 'Not Available'
env_info['NVCC'] = nvcc env_info['NVCC'] = nvcc
try: try:
......
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment