ClearML/clearml-agent
1
0
Fork 0
mirror of https://github.com/clearml/clearml-agent synced 2025-06-26 18:16:15 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

Upgrade pynvml add detect CUDA version from driver level

Browse Source
This commit is contained in:
allegroai
2021-02-17 00:03:16 +02:00
parent 22d5892b12
commit 58cb344ee6
3 changed files with 1938 additions and 44 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
clearml_agent/helper/package/requirements.py
View File

@@ -17,6 +17,7 @@ import six
from clearml_agent.definitions import PIP_EXTRA_INDICES
from clearml_agent.helper.base import warning, is_conda, which, join_lines, is_windows_platform
from clearml_agent.helper.process import Argv, PathLike
from clearml_agent.helper.gpu.gpustat import get_driver_cuda_version
from clearml_agent.session import Session, normalize_cuda_version
from clearml_agent.external.requirements_parser import parse
from clearml_agent.external.requirements_parser.requirement import Requirement
@@ -537,6 +538,9 @@ class RequirementsManager(object):
if cuda_version and cudnn_version:
return normalize_cuda_version(cuda_version), normalize_cuda_version(cudnn_version)
if not cuda_version:
cuda_version = get_driver_cuda_version()
if not cuda_version and is_windows_platform():
try:
cuda_vers = [int(k.replace('CUDA_PATH_V', '').replace('_', '')) for k in os.environ.keys()