Unverified Commit c33f2489 authored by Tan SU's avatar Tan SU Committed by GitHub
Browse files

[Enhancement] Fix collect_env on Windows (#1789)

* [fix] fix collect_env() on Windows

* fix(utils): add back 'Cuda compilation tools' to nvcc info

* feat(utils): find default ccompiler

* fix(utils, test): parse cl.exe version instead of path

* fix: code style

* fix(tools): use distutils

* [fix] collect_env() returns GCC or MSVC

* [fix] unitest for collect_env()

* add comments

* test: add GCC to expected_keys

* fix: code style

* Update mmcv/utils/env.py
parent e0eebebf
......@@ -26,6 +26,7 @@ def collect_env():
- CUDA_HOME (optional): The env var ``CUDA_HOME``.
- NVCC (optional): NVCC version.
- GCC: GCC version, "n/a" if GCC is not installed.
- MSVC: Microsoft Virtual C++ Compiler version, Windows only.
- PyTorch: PyTorch version.
- PyTorch compiling details: The output of \
``torch.__config__.show()``.
......@@ -56,18 +57,40 @@ def collect_env():
if CUDA_HOME is not None and osp.isdir(CUDA_HOME):
try:
nvcc = osp.join(CUDA_HOME, 'bin/nvcc')
nvcc = subprocess.check_output(
f'"{nvcc}" -V | tail -n1', shell=True)
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
try:
gcc = subprocess.check_output('gcc --version | head -n1', shell=True)
gcc = gcc.decode('utf-8').strip()
env_info['GCC'] = gcc
except subprocess.CalledProcessError: # gcc is unavailable
# Check C++ Compiler.
# For Unix-like, sysconfig has 'CC' variable like 'gcc -pthread ...',
# indicating the compiler used, we use this to get the compiler name
import sysconfig
cc = sysconfig.get_config_var('CC')
if cc:
cc = osp.basename(cc.split()[0])
cc_info = subprocess.check_output(f'{cc} --version', shell=True)
env_info['GCC'] = cc_info.decode('utf-8').partition(
'\n')[0].strip()
else:
# on Windows, cl.exe is not in PATH. We need to find the path.
# distutils.ccompiler.new_compiler() returns a msvccompiler
# object and after initialization, path to cl.exe is found.
import locale
from distutils.ccompiler import new_compiler
ccompiler = new_compiler()
ccompiler.initialize()
cc = subprocess.check_output(
f'{ccompiler.cc}', stderr=subprocess.STDOUT, shell=True)
encoding = locale.getdefaultlocale()[1]
env_info['MSVC'] = cc.decode(encoding).partition('\n')[0].strip()
env_info['GCC'] = 'n/a'
except subprocess.CalledProcessError:
env_info['GCC'] = 'n/a'
env_info['PyTorch'] = torch.__version__
......
......@@ -16,7 +16,7 @@ def test_collect_env():
env_info = collect_env()
expected_keys = [
'sys.platform', 'Python', 'CUDA available', 'PyTorch',
'PyTorch compiling details', 'OpenCV', 'MMCV', 'MMCV Compiler',
'PyTorch compiling details', 'OpenCV', 'MMCV', 'MMCV Compiler', 'GCC',
'MMCV CUDA Compiler'
]
for key in expected_keys:
......@@ -26,8 +26,8 @@ def test_collect_env():
for key in ['CUDA_HOME', 'NVCC']:
assert key in env_info
if sys.platform != 'win32':
assert 'GCC' in env_info
if sys.platform == 'win32':
assert 'MSVC' in env_info
assert env_info['sys.platform'] == sys.platform
assert env_info['Python'] == sys.version.replace('\n', '')
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment