From 28abc9e69af0b9bf389aeaca9de174a011629352 Mon Sep 17 00:00:00 2001 From: allegroai <> Date: Mon, 12 Apr 2021 22:58:22 +0300 Subject: [PATCH] Fix Task.set_base_docker() Add docker bash script to Task.create() --- clearml/backend_interface/task/populate.py | 13 +++++++++++-- clearml/backend_interface/task/task.py | 14 ++++++++------ clearml/task.py | 7 ++++++- 3 files changed, 25 insertions(+), 9 deletions(-) diff --git a/clearml/backend_interface/task/populate.py b/clearml/backend_interface/task/populate.py index 2aa94a8c..c4eab3e2 100644 --- a/clearml/backend_interface/task/populate.py +++ b/clearml/backend_interface/task/populate.py @@ -27,6 +27,8 @@ class CreateAndPopulate(object): packages=None, # Optional[Union[bool, Sequence[str]]] requirements_file=None, # Optional[Union[str, Path]] docker=None, # Optional[str] + docker_args=None, # Optional[str] + docker_bash_setup_script=None, # Optional[str] base_task_id=None, # Optional[str] add_task_init_call=True, # 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. 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_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. 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. @@ -96,7 +101,7 @@ class CreateAndPopulate(object): else (packages or None) self.requirements_file = Path(requirements_file) if requirements_file else None 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.project_name = project_name self.task_name = task_name @@ -274,7 +279,11 @@ class CreateAndPopulate(object): # set base docker image if provided 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 task_state['script']['repository']: diff --git a/clearml/backend_interface/task/task.py b/clearml/backend_interface/task/task.py index 4a76f37e..be24e8fb 100644 --- a/clearml/backend_interface/task/task.py +++ b/clearml/backend_interface/task/task.py @@ -1153,26 +1153,28 @@ class Task(IdObjectBase, AccessMixin, SetupUploadMixin): """ image = docker_cmd.split(' ')[0] if docker_cmd else '' if not docker_arguments and docker_cmd: - 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)) \ - if docker_arguments else '' + docker_arguments = docker_cmd.split(' ')[1:] if len(docker_cmd.split(' ')) > 1 else '' + + arguments = (docker_arguments if isinstance(docker_arguments, str) else ' '.join(docker_arguments)) \ + if docker_arguments else '' + if 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) else: - setup_shell_script = None + setup_shell_script = '' with self._edit_lock: self.reload() if Session.check_min_api_version("2.13"): self.data.container = dict(image=image, arguments=arguments, setup_shell_script=setup_shell_script) + self._edit(container=self.data.container) else: if setup_shell_script: raise ValueError( "Your ClearML-server does not support docker bash script feature, please upgrade.") 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) def get_base_docker(self): diff --git a/clearml/task.py b/clearml/task.py index 37678f2d..a69e59d0 100644 --- a/clearml/task.py +++ b/clearml/task.py @@ -622,6 +622,8 @@ class Task(_Task): packages=None, # Optional[Union[bool, Sequence[str]]] requirements_file=None, # Optional[Union[str, Path]] docker=None, # Optional[str] + docker_args=None, # Optional[str] + docker_bash_setup_script=None, # Optional[str] argparse_args=None, # Optional[Sequence[Tuple[str, str]]] base_task_id=None, # Optional[str] 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. 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_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) 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. @@ -677,7 +682,7 @@ class Task(_Task): repo=repo, branch=branch, commit=commit, script=script, working_directory=working_directory, 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, add_task_init_call=add_task_init_call, raise_on_missing_entries=False,