ClearML/clearml-agent
1
0
Fork 0
mirror of https://github.com/clearml/clearml-agent synced 2025-04-16 21:41:55 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

Add support for running daemon inside a K8s pod in daemon mode

Browse Source
This commit is contained in:
allegroai 2020-03-05 13:03:36 +02:00
parent 8f95eecf2e
commit 231a907cff
2 changed files with 13 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

11
trains_agent/commands/worker.py
View File

@ -1882,6 +1882,17 @@ class Worker(ServiceCommandSection):
if isinstance(extra_docker_arguments, six.string_types) else extra_docker_arguments
base_cmd += [str(a) for a in extra_docker_arguments if a]
# check if running inside a kubernetes
if os.environ.get('KUBERNETES_SERVICE_HOST') and os.environ.get('KUBERNETES_PORT'):
# map network to sibling docker
try:
network_mode = get_bash_output(
'docker inspect --format=\'{{.HostConfig.NetworkMode}}\' $(basename $(cat /proc/1/cpuset))')
base_cmd += ['--network', network_mode]
except:
pass
base_cmd += ['-e', 'NVIDIA_VISIBLE_DEVICES={}'.format(dockers_nvidia_visible_devices)]
base_cmd += ['-e', 'TRAINS_WORKER_ID='+worker_id, ]
if host_ssh_cache:

3
trains_agent/session.py
View File

@ -75,7 +75,8 @@ class Session(_Session):
cpu_only = kwargs.get('cpu_only')
if cpu_only:
os.environ['CUDA_VISIBLE_DEVICES'] = os.environ['NVIDIA_VISIBLE_DEVICES'] = 'none'
if kwargs.get('gpus'):
if kwargs.get('gpus') and not os.environ.get('KUBERNETES_SERVICE_HOST') \
and not os.environ.get('KUBERNETES_PORT'):
os.environ['CUDA_VISIBLE_DEVICES'] = os.environ['NVIDIA_VISIBLE_DEVICES'] = kwargs.get('gpus')
if kwargs.get('only_load_config'):
from trains_agent.backend_api.config import load