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

Add support for updating back docker on new API v2.13

Browse Source
This commit is contained in:
allegroai 2021-04-07 18:46:58 +03:00
parent 8eace6d57b
commit 45009e6cc2
1 changed files with 41 additions and 20 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

57
clearml_agent/commands/worker.py
View File

@ -249,6 +249,29 @@ def get_task_container(session, task_id):
return container
def set_task_container(session, task_id, docker_image=None, docker_arguments=None, docker_bash_script=None):
if session.check_min_api_version("2.13"):
container = dict(
image=docker_image or None,
arguments=' '.join(docker_arguments) if docker_arguments else None,
setup_shell_script=docker_bash_script or None,
)
result = session.send_request(
service='tasks',
action='edit',
version='2.13',
json={'task': task_id, 'container': container, 'force': True},
method='get',
async_enable=False,
)
return result.ok
else:
return session.send_api(
tasks_api.EditRequest(task_id, force=True, execution=dict( # noqa
docker_cmd=' '.join([docker_image] + docker_arguments)
if docker_arguments else str(docker_image or ''))))
class TaskStopSignal(object):
"""
Follow task status and signals when it should be stopped
@ -541,21 +564,18 @@ class Worker(ServiceCommandSection):
except Exception:
task_container = {}
if task_container:
self.send_logs(task_id=task_id,
lines=['Running Task {} inside docker: {}\n'.format(task_id, str(task_container))],
level="INFO")
docker_image = task_container.get('image')
docker_arguments = task_container.get('arguments')
docker_image = task_container.get('image') or self._docker_image
docker_arguments = task_container.get('arguments') or self._docker_arguments
docker_setup_script = task_container.get('setup_shell_script')
else:
self.send_logs(task_id=task_id,
lines=['running Task {} inside default docker image: {} {}\n'.format(
task_id, self._docker_image, self._docker_arguments or '')],
self.send_logs(
task_id=task_id,
lines=
['Running Task {} inside {}docker: {} arguments: {}\n'.format(
task_id, "default " if not task_container.get('image') else '',
docker_image, docker_arguments)]
+ ['custom_setup_bash_script:\n{}'.format(docker_setup_script)] if docker_setup_script else [],
level="INFO")
docker_image = self._docker_image
docker_arguments = self._docker_arguments
docker_setup_script = None
# Update docker command
if self._services_mode:
@ -572,13 +592,14 @@ class Worker(ServiceCommandSection):
docker_bash_setup_script=docker_setup_script)
# if we are using the default docker, update back the Task:
if not task_container:
if not task_container.get('image') or not task_container.get('arguments'):
# noinspection PyBroadException
try:
self._session.send_api(
tasks_api.EditRequest(task_id, force=True, execution=dict( # noqa
docker_cmd=' '.join([docker_image] + docker_arguments)
if docker_arguments else str(docker_image))))
set_task_container(
self._session, task_id,
docker_image=docker_image,
docker_arguments=docker_arguments,
docker_bash_script=docker_setup_script)
except Exception:
pass