ClearML/clearml-agent
1
0
Fork 0
mirror of https://github.com/clearml/clearml-agent synced 2025-01-31 09:06:52 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

Allow to specifying cudatoolkit version in "installed packages" when using Conda as package manager (trains issue #229)

Browse Source
This commit is contained in:
allegroai 2020-10-30 10:06:02 +02:00
parent 293a92f486
commit 216b3e2179
1 changed files with 4 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
trains_agent/helper/package/conda_api.py
View File

@ -378,7 +378,8 @@ class CondaAPI(PackageManager):
line = 'tensorflow={}'.format(line.split('=')[1])
elif name == 'tensorflow' and cuda_version > 0:
line = 'tensorflow-gpu={}'.format(line.split('=')[1])
elif name in ('cudatoolkit', 'cupti', 'cudnn'):
elif name in ('cupti', 'cudnn'):
# cudatoolkit should pull them based on the cudatoolkit version
continue
elif name.startswith('_'):
continue
@ -543,7 +544,8 @@ class CondaAPI(PackageManager):
reqs.append(MarkerRequirement(Requirement.parse('kiwisolver')))
# remove specific cudatoolkit, it should have being preinstalled.
reqs = [r for r in reqs if r.name not in ('cudatoolkit', 'cudnn', 'cupti')]
# allow to override default cudatoolkit, but not the derivative packages, cudatoolkit should pull them
reqs = [r for r in reqs if r.name not in ('cudnn', 'cupti')]
if has_torch and cuda_version == 0:
reqs.append(MarkerRequirement(Requirement.parse('cpuonly')))