ClearML/clearml-agent
1
0
Fork 0
mirror of https://github.com/clearml/clearml-agent synced 2025-03-13 15:08:44 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

Improve k8s nvidia container integration

Browse Source
This commit is contained in:
allegroai 2020-11-26 01:15:49 +02:00
parent ea63e4f66e
commit 0cf485f7a9
1 changed files with 60 additions and 39 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

97
trains_agent/glue/k8s.py
View File

@ -12,8 +12,6 @@ import json
from time import sleep from time import sleep
from typing import Text, List from typing import Text, List
from pyhocon import HOCONConverter
from trains_agent.commands.events import Events from trains_agent.commands.events import Events
from trains_agent.commands.worker import Worker from trains_agent.commands.worker import Worker
from trains_agent.errors import APIError from trains_agent.errors import APIError
@ -40,34 +38,34 @@ class K8sIntegration(Worker):
"--field-selector=status.phase!=Pending,status.phase!=Running " \ "--field-selector=status.phase!=Pending,status.phase!=Running " \
"--namespace=trains" "--namespace=trains"
BASH_INSTALL_SSH_CMD = \ BASH_INSTALL_SSH_CMD = [
"apt-get install -y openssh-server ; " \ "apt-get install -y openssh-server",
"mkdir -p /var/run/sshd ; " \ "mkdir -p /var/run/sshd",
"echo 'root:training' | chpasswd ; " \ "echo 'root:training' | chpasswd",
"echo 'PermitRootLogin yes' >> /etc/ssh/sshd_config ; " \ "echo 'PermitRootLogin yes' >> /etc/ssh/sshd_config",
"sed -i 's/PermitRootLogin prohibit-password/PermitRootLogin yes/' /etc/ssh/sshd_config ; " \ "sed -i 's/PermitRootLogin prohibit-password/PermitRootLogin yes/' /etc/ssh/sshd_config",
"sed 's@session\s*required\s*pam_loginuid.so@session optional pam_loginuid.so@g' -i /etc/pam.d/sshd ; " \ r"sed 's@session\s*required\s*pam_loginuid.so@session optional pam_loginuid.so@g' -i /etc/pam.d/sshd",
"echo 'AcceptEnv TRAINS_API_ACCESS_KEY TRAINS_API_SECRET_KEY' >> /etc/ssh/sshd_config ; " \ "echo 'AcceptEnv TRAINS_API_ACCESS_KEY TRAINS_API_SECRET_KEY' >> /etc/ssh/sshd_config",
'echo "export VISIBLE=now" >> /etc/profile ; ' \ 'echo "export VISIBLE=now" >> /etc/profile',
'echo "export PATH=$PATH" >> /etc/profile ; ' \ 'echo "export PATH=$PATH" >> /etc/profile',
'echo "ldconfig" >> /etc/profile ; ' \ 'echo "ldconfig" >> /etc/profile',
"/usr/sbin/sshd -p {port} & ; " "/usr/sbin/sshd -p {port}"]
CONTAINER_BASH_SCRIPT = \ CONTAINER_BASH_SCRIPT = [
"export DEBIAN_FRONTEND='noninteractive'; " \ "export DEBIAN_FRONTEND='noninteractive'",
"echo 'Binary::apt::APT::Keep-Downloaded-Packages \"true\";' > /etc/apt/apt.conf.d/docker-clean ; " \ "echo 'Binary::apt::APT::Keep-Downloaded-Packages \"true\";' > /etc/apt/apt.conf.d/docker-clean",
"chown -R root /root/.cache/pip ; " \ "chown -R root /root/.cache/pip",
"apt-get update ; " \ "apt-get update",
"apt-get install -y git libsm6 libxext6 libxrender-dev libglib2.0-0 ; " \ "apt-get install -y git libsm6 libxext6 libxrender-dev libglib2.0-0",
"declare LOCAL_PYTHON ; " \ "declare LOCAL_PYTHON",
"for i in {{10..5}}; do which python3.$i && " \ "for i in {{10..5}}; do which python3.$i && python3.$i -m pip --version && "
"python3.$i -m pip --version && " \ "export LOCAL_PYTHON=$(which python3.$i) && break ; done",
"export LOCAL_PYTHON=$(which python3.$i) && break ; done ; " \ "[ ! -z $LOCAL_PYTHON ] || apt-get install -y python3-pip",
"[ ! -z $LOCAL_PYTHON ] || apt-get install -y python3-pip ; " \ "[ ! -z $LOCAL_PYTHON ] || export LOCAL_PYTHON=python3",
"[ ! -z $LOCAL_PYTHON ] || export LOCAL_PYTHON=python3 ; " \ "$LOCAL_PYTHON -m pip install trains-agent",
"$LOCAL_PYTHON -m pip install trains-agent ; " \ "{extra_bash_init_cmd}",
"{extra_bash_init_cmd} ; " \ "$LOCAL_PYTHON -m trains_agent execute --full-monitoring --require-queue --id {task_id}"
"$LOCAL_PYTHON -m trains_agent execute --full-monitoring --require-queue --id {task_id} ; " ]
AGENT_LABEL = "TRAINS=agent" AGENT_LABEL = "TRAINS=agent"
LIMIT_POD_LABEL = "ai.allegro.agent.serial=pod-{pod_number}" LIMIT_POD_LABEL = "ai.allegro.agent.serial=pod-{pod_number}"
@ -131,7 +129,7 @@ class K8sIntegration(Worker):
self.template_dict = None self.template_dict = None
self.extra_bash_init_script = extra_bash_init_script or None self.extra_bash_init_script = extra_bash_init_script or None
if self.extra_bash_init_script and not isinstance(self.extra_bash_init_script, str): if self.extra_bash_init_script and not isinstance(self.extra_bash_init_script, str):
self.extra_bash_init_script = ' ; '.join(self.extra_bash_init_script) self.extra_bash_init_script = ' ; '.join(self.extra_bash_init_script) # noqa
self.pod_limits = [] self.pod_limits = []
self.pod_requests = [] self.pod_requests = []
if overrides_yaml: if overrides_yaml:
@ -163,6 +161,8 @@ class K8sIntegration(Worker):
if trains_conf_file: if trains_conf_file:
with open(os.path.expandvars(os.path.expanduser(str(trains_conf_file))), 'rt') as f: with open(os.path.expandvars(os.path.expanduser(str(trains_conf_file))), 'rt') as f:
self.trains_conf_file = f.read() self.trains_conf_file = f.read()
# make sure we use system packages!
self.trains_conf_file += '\nagent.package_manager.system_site_packages=true\n'
def _set_task_user_properties(self, task_id: str, **properties: str): def _set_task_user_properties(self, task_id: str, **properties: str):
if self._edit_hyperparams_support is not True: if self._edit_hyperparams_support is not True:
@ -223,7 +223,7 @@ class K8sIntegration(Worker):
# get the trains.conf encoded file # get the trains.conf encoded file
# noinspection PyProtectedMember # noinspection PyProtectedMember
hocon_config_encoded = (self.trains_conf_file or self._session._config_file).encode('ascii') hocon_config_encoded = (self.trains_conf_file or self._session._config_file).encode('ascii')
create_trains_conf = "echo '{}' | base64 --decode >> ~/trains.conf && ".format( create_trains_conf = "echo '{}' | base64 --decode >> ~/trains.conf".format(
base64.b64encode( base64.b64encode(
hocon_config_encoded hocon_config_encoded
).decode('ascii') ).decode('ascii')
@ -326,13 +326,28 @@ class K8sIntegration(Worker):
template['metadata'].setdefault('labels', {}) template['metadata'].setdefault('labels', {})
template['metadata']['labels'].update(labels_dict) template['metadata']['labels'].update(labels_dict)
container = self._parse_docker_args(docker_args) container = self._parse_docker_args(docker_args)
container_bash_script = [self.container_bash_script] if isinstance(self.container_bash_script, str) \
else self.container_bash_script
script_encoded = '\n'.join(
['#!/bin/bash', ] +
[line.format(extra_bash_init_cmd=self.extra_bash_init_script or '', task_id=task_id)
for line in container_bash_script])
create_init_script = \
"echo '{}' | base64 --decode >> ~/__start_agent__.sh ; " \
"/bin/bash ~/__start_agent__.sh".format(
base64.b64encode(
script_encoded.encode('ascii')
).decode('ascii'))
container = merge_dicts( container = merge_dicts(
container, container,
dict(name=name, image=docker_image, dict(name=name, image=docker_image,
command=['/bin/bash'], command=['/bin/bash'],
args=['-c', create_trains_conf + self.container_bash_script.format( args=['-c', '{} ; {}'.format(create_trains_conf, create_init_script)])
extra_bash_init_cmd=self.extra_bash_init_script, )
task_id=task_id)]))
if template['spec']['containers']: if template['spec']['containers']:
template['spec']['containers'][0] = merge_dicts(template['spec']['containers'][0], container) template['spec']['containers'][0] = merge_dicts(template['spec']['containers'][0], container)
@ -386,14 +401,18 @@ class K8sIntegration(Worker):
if self.pod_requests: if self.pod_requests:
kubectl_cmd += ['--requests', ",".join(self.pod_requests)] kubectl_cmd += ['--requests', ",".join(self.pod_requests)]
container_bash_script = [self.container_bash_script] if isinstance(self.container_bash_script, str) \
else self.container_bash_script
container_bash_script = ' ; '.join(container_bash_script)
kubectl_cmd += [ kubectl_cmd += [
"--labels=" + ",".join(labels), "--labels=" + ",".join(labels),
"--command", "--command",
"--", "--",
"/bin/sh", "/bin/sh",
"-c", "-c",
create_trains_conf + self.container_bash_script.format( "{} ; {}".format(create_trains_conf, container_bash_script.format(
extra_bash_init_cmd=self.extra_bash_init_script, task_id=task_id), extra_bash_init_cmd=self.extra_bash_init_script, task_id=task_id)),
] ]
process = subprocess.Popen(kubectl_cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE) process = subprocess.Popen(kubectl_cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
output, error = process.communicate() output, error = process.communicate()
@ -413,6 +432,7 @@ class K8sIntegration(Worker):
events_service = self.get_service(Events) events_service = self.get_service(Events)
# make sure we have a k8s pending queue # make sure we have a k8s pending queue
# noinspection PyBroadException
try: try:
self._session.api_client.queues.create(self.k8s_pending_queue_name) self._session.api_client.queues.create(self.k8s_pending_queue_name)
except Exception: except Exception:
@ -471,8 +491,9 @@ class K8sIntegration(Worker):
:param list(str) queue: queue name to pull from :param list(str) queue: queue name to pull from
""" """
return self.daemon(queues=[ObjectID(name=queue)] if queue else None, log_level=logging.INFO, foreground=True, docker=False) return self.daemon(queues=[ObjectID(name=queue)] if queue else None,
log_level=logging.INFO, foreground=True, docker=False)
@classmethod @classmethod
def get_ssh_server_bash(cls, ssh_port_number): def get_ssh_server_bash(cls, ssh_port_number):
return cls.BASH_INSTALL_SSH_CMD.format(port=ssh_port_number) return ' ; '.join(line.format(port=ssh_port_number) for line in cls.BASH_INSTALL_SSH_CMD)