Fix Task.set_base_docker()

Add docker bash script to Task.create()
This commit is contained in:
allegroai 2021-04-12 22:58:22 +03:00
parent 50d620cc9a
commit 28abc9e69a
3 changed files with 25 additions and 9 deletions

View File

@ -27,6 +27,8 @@ class CreateAndPopulate(object):
packages=None, # Optional[Union[bool, Sequence[str]]] packages=None, # Optional[Union[bool, Sequence[str]]]
requirements_file=None, # Optional[Union[str, Path]] requirements_file=None, # Optional[Union[str, Path]]
docker=None, # Optional[str] docker=None, # Optional[str]
docker_args=None, # Optional[str]
docker_bash_setup_script=None, # Optional[str]
base_task_id=None, # Optional[str] base_task_id=None, # Optional[str]
add_task_init_call=True, # bool add_task_init_call=True, # bool
raise_on_missing_entries=False, # bool raise_on_missing_entries=False, # bool
@ -59,6 +61,9 @@ class CreateAndPopulate(object):
:param requirements_file: Specify requirements.txt file to install when setting the session. :param requirements_file: Specify requirements.txt file to install when setting the session.
If not provided, the requirements.txt from the repository will be used. If not provided, the requirements.txt from the repository will be used.
:param docker: Select the docker image to be executed in by the remote session :param docker: Select the docker image to be executed in by the remote session
:param docker_args: Add docker arguments, pass a single string
:param docker_bash_setup_script: Add bash script to be executed
inside the docker before setting up the Task's environement
:param base_task_id: Use a pre-existing task in the system, instead of a local repo/script. :param base_task_id: Use a pre-existing task in the system, instead of a local repo/script.
Essentially clones an existing task and overrides arguments/requirements. Essentially clones an existing task and overrides arguments/requirements.
:param add_task_init_call: If True, a 'Task.init()' call is added to the script entry point in remote execution. :param add_task_init_call: If True, a 'Task.init()' call is added to the script entry point in remote execution.
@ -96,7 +101,7 @@ class CreateAndPopulate(object):
else (packages or None) else (packages or None)
self.requirements_file = Path(requirements_file) if requirements_file else None self.requirements_file = Path(requirements_file) if requirements_file else None
self.base_task_id = base_task_id self.base_task_id = base_task_id
self.docker = docker self.docker = dict(image=docker, args=docker_args, bash_script=docker_bash_setup_script)
self.add_task_init_call = add_task_init_call self.add_task_init_call = add_task_init_call
self.project_name = project_name self.project_name = project_name
self.task_name = task_name self.task_name = task_name
@ -274,7 +279,11 @@ class CreateAndPopulate(object):
# set base docker image if provided # set base docker image if provided
if self.docker: if self.docker:
task.set_base_docker(self.docker) task.set_base_docker(
docker_cmd=self.docker.get('image'),
docker_arguments=self.docker.get('args'),
docker_setup_bash_script=self.docker.get('bash_script'),
)
if self.verbose: if self.verbose:
if task_state['script']['repository']: if task_state['script']['repository']:

View File

@ -1153,26 +1153,28 @@ class Task(IdObjectBase, AccessMixin, SetupUploadMixin):
""" """
image = docker_cmd.split(' ')[0] if docker_cmd else '' image = docker_cmd.split(' ')[0] if docker_cmd else ''
if not docker_arguments and docker_cmd: if not docker_arguments and docker_cmd:
arguments = docker_cmd.split(' ')[1:] if len(docker_cmd.split(' ')) > 1 else '' docker_arguments = docker_cmd.split(' ')[1:] if len(docker_cmd.split(' ')) > 1 else ''
else:
arguments = (docker_arguments if isinstance(docker_arguments, str) else ' '.join(docker_arguments)) \ arguments = (docker_arguments if isinstance(docker_arguments, str) else ' '.join(docker_arguments)) \
if docker_arguments else '' if docker_arguments else ''
if docker_setup_bash_script: if docker_setup_bash_script:
setup_shell_script = docker_setup_bash_script \ setup_shell_script = docker_setup_bash_script \
if isinstance(docker_setup_bash_script, str) else '\n'.join(docker_setup_bash_script) if isinstance(docker_setup_bash_script, str) else '\n'.join(docker_setup_bash_script)
else: else:
setup_shell_script = None setup_shell_script = ''
with self._edit_lock: with self._edit_lock:
self.reload() self.reload()
if Session.check_min_api_version("2.13"): if Session.check_min_api_version("2.13"):
self.data.container = dict(image=image, arguments=arguments, setup_shell_script=setup_shell_script) self.data.container = dict(image=image, arguments=arguments, setup_shell_script=setup_shell_script)
self._edit(container=self.data.container)
else: else:
if setup_shell_script: if setup_shell_script:
raise ValueError( raise ValueError(
"Your ClearML-server does not support docker bash script feature, please upgrade.") "Your ClearML-server does not support docker bash script feature, please upgrade.")
execution = self.data.execution execution = self.data.execution
execution.docker_cmd = docker_cmd + (' {}'.format(arguments) if arguments else '') execution.docker_cmd = image + (' {}'.format(arguments) if arguments else '')
self._edit(execution=execution) self._edit(execution=execution)
def get_base_docker(self): def get_base_docker(self):

View File

@ -622,6 +622,8 @@ class Task(_Task):
packages=None, # Optional[Union[bool, Sequence[str]]] packages=None, # Optional[Union[bool, Sequence[str]]]
requirements_file=None, # Optional[Union[str, Path]] requirements_file=None, # Optional[Union[str, Path]]
docker=None, # Optional[str] docker=None, # Optional[str]
docker_args=None, # Optional[str]
docker_bash_setup_script=None, # Optional[str]
argparse_args=None, # Optional[Sequence[Tuple[str, str]]] argparse_args=None, # Optional[Sequence[Tuple[str, str]]]
base_task_id=None, # Optional[str] base_task_id=None, # Optional[str]
add_task_init_call=True, # bool add_task_init_call=True, # bool
@ -658,6 +660,9 @@ class Task(_Task):
:param requirements_file: Specify requirements.txt file to install when setting the session. :param requirements_file: Specify requirements.txt file to install when setting the session.
If not provided, the requirements.txt from the repository will be used. If not provided, the requirements.txt from the repository will be used.
:param docker: Select the docker image to be executed in by the remote session :param docker: Select the docker image to be executed in by the remote session
:param docker_args: Add docker arguments, pass a single string
:param docker_bash_setup_script: Add bash script to be executed
inside the docker before setting up the Task's environement
:param argparse_args: Arguments to pass to the remote execution, list of string pairs (argument, value) :param argparse_args: Arguments to pass to the remote execution, list of string pairs (argument, value)
Notice, only supported if the codebase itself uses argparse.ArgumentParser Notice, only supported if the codebase itself uses argparse.ArgumentParser
:param base_task_id: Use a pre-existing task in the system, instead of a local repo/script. :param base_task_id: Use a pre-existing task in the system, instead of a local repo/script.
@ -677,7 +682,7 @@ class Task(_Task):
repo=repo, branch=branch, commit=commit, repo=repo, branch=branch, commit=commit,
script=script, working_directory=working_directory, script=script, working_directory=working_directory,
packages=packages, requirements_file=requirements_file, packages=packages, requirements_file=requirements_file,
docker=docker, docker=docker, docker_args=docker_args, docker_bash_setup_script=docker_bash_setup_script,
base_task_id=base_task_id, base_task_id=base_task_id,
add_task_init_call=add_task_init_call, add_task_init_call=add_task_init_call,
raise_on_missing_entries=False, raise_on_missing_entries=False,