[Provisioner] Fix bug when using docker login config and launching wi…

…thout yaml (#2572) * fix * apply suggestions from code review * fix * fix * add comment * Update sky/task.py Co-authored-by: Zhanghao Wu <[email protected]> --------- Co-authored-by: Zhanghao Wu <[email protected]>
skypilot-org · Sep 18, 2023 · 17ede26 · 17ede26
1 parent e5e400b
commit 17ede26
Show file tree

Hide file tree

Showing 2 changed files with 35 additions and 13 deletions.
diff --git a/sky/skylet/constants.py b/sky/skylet/constants.py
@@ -53,6 +53,11 @@
 DOCKER_USERNAME_ENV_VAR = 'SKYPILOT_DOCKER_USERNAME'
 DOCKER_PASSWORD_ENV_VAR = 'SKYPILOT_DOCKER_PASSWORD'
 DOCKER_SERVER_ENV_VAR = 'SKYPILOT_DOCKER_SERVER'
+DOCKER_LOGIN_ENV_VARS = {
+    DOCKER_USERNAME_ENV_VAR,
+    DOCKER_PASSWORD_ENV_VAR,
+    DOCKER_SERVER_ENV_VAR,
+}
 
 # Install conda on the remote cluster if it is not already installed.
 # We do not install the latest conda with python 3.11 because ray has not

diff --git a/sky/task.py b/sky/task.py
@@ -105,32 +105,43 @@ def replace_var(match):
     return json.loads(file_mounts_str)
 
 
-def _with_docker_login_config(
-    resources_set: Set['resources_lib.Resources'],
-    task_envs: Dict[str, str],
-) -> Set['resources_lib.Resources']:
-    all_keys = {
-        constants.DOCKER_USERNAME_ENV_VAR,
-        constants.DOCKER_PASSWORD_ENV_VAR,
-        constants.DOCKER_SERVER_ENV_VAR,
-    }
+def _check_docker_login_config(task_envs: Dict[str, str]) -> bool:
+    """Checks if there is a valid docker login config in task_envs.
+
+    If any of the docker login env vars is set, all of them must be set.
+
+    Raises:
+        ValueError: if any of the docker login env vars is set, but not all of
+            them are set.
+    """
+    all_keys = constants.DOCKER_LOGIN_ENV_VARS
     existing_keys = all_keys & set(task_envs.keys())
     if not existing_keys:
-        return resources_set
+        return False
     if len(existing_keys) != len(all_keys):
         with ux_utils.print_exception_no_traceback():
             raise ValueError(
                 f'If any of {", ".join(all_keys)} is set, all of them must '
                 f'be set. Missing envs: {all_keys - existing_keys}')
+    return True
+
+
+def _with_docker_login_config(
+    resources_set: Set['resources_lib.Resources'],
+    task_envs: Dict[str, str],
+) -> Set['resources_lib.Resources']:
+    if not _check_docker_login_config(task_envs):
+        return resources_set
     docker_login_config = command_runner.DockerLoginConfig.from_env_vars(
         task_envs)
 
     def _add_docker_login_config(resources: 'resources_lib.Resources'):
         if resources.extract_docker_image() is None:
             logger.warning(f'{colorama.Fore.YELLOW}Docker login configs '
-                           f'{", ".join(all_keys)} are provided, but no docker '
-                           'image is specified in `image_id`. The login configs'
-                           f' will be ignored.{colorama.Style.RESET_ALL}')
+                           f'{", ".join(constants.DOCKER_LOGIN_ENV_VARS)} '
+                           'are provided, but no docker image is specified '
+                           'in `image_id`. The login configs will be '
+                           f'ignored.{colorama.Style.RESET_ALL}')
             return resources
         return resources.copy(_docker_login_config=docker_login_config)
 
@@ -494,6 +505,12 @@ def update_envs(
                     'envs must be List[Tuple[str, str]] or Dict[str, str]: '
                     f'{envs}')
         self._envs.update(envs)
+        # If the update_envs() is called after set_resources(), we need to
+        # manually update docker login config in task resources, in case the
+        # docker login envs are newly added.
+        if _check_docker_login_config(self._envs):
+            self.resources = _with_docker_login_config(self.resources,
+                                                       self._envs)
         return self
 
     @property