ClearML/clearml-agent
1
0
Fork 0
mirror of https://github.com/clearml/clearml-agent synced 2025-02-07 05:19:17 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

Pass correct GPU limit when skipping gpus flag in docker mode

Browse Source
This commit is contained in:
allegroai 2020-03-05 14:07:44 +02:00
parent 0a3a8a1c52
commit 47bcd3839a
1 changed files with 10 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

12
trains_agent/commands/worker.py
View File

@ -1868,11 +1868,19 @@ class Worker(ServiceCommandSection):
docker = 'docker'
base_cmd = [docker, 'run', '-t']
dockers_nvidia_visible_devices = 'all'
gpu_devices = os.environ.get('NVIDIA_VISIBLE_DEVICES', None)
if gpu_devices is None or gpu_devices.lower().strip() == 'all':
base_cmd += ['--gpus', 'all', ]
if os.environ.get('TRAINS_DOCKER_SKIP_GPUS_FLAG', None):
dockers_nvidia_visible_devices = os.environ.get('NVIDIA_VISIBLE_DEVICES') or \
dockers_nvidia_visible_devices
else:
base_cmd += ['--gpus', 'all', ]
elif gpu_devices.strip() and gpu_devices.strip() != 'none':
base_cmd += ['--gpus', 'device='+gpu_devices, ]
if os.environ.get('TRAINS_DOCKER_SKIP_GPUS_FLAG', None):
dockers_nvidia_visible_devices = gpu_devices
else:
base_cmd += ['--gpus', 'device='+gpu_devices, ]
# We are using --gpu, so we should not pass NVIDIA_VISIBLE_DEVICES, I think.
# base_cmd += ['-e', 'NVIDIA_VISIBLE_DEVICES=' + gpu_devices, ]