Fix Task.set_base_docker()

Add docker bash script to Task.create()

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']:

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 ''
+            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)) \
             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):

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,