From ad149c1a7506b314b75a550a73b0cfc270acf811 Mon Sep 17 00:00:00 2001
From: cblmemo <cblmemo@gmail.com>
Date: Sun, 17 Sep 2023 23:19:10 -0700
Subject: [PATCH] apply suggestions from code review

---
 sky/cli.py              |  5 +----
 sky/skylet/constants.py |  5 +++++
 sky/task.py             | 47 +++++++++++++++++++++++++----------------
 3 files changed, 35 insertions(+), 22 deletions(-)

diff --git a/sky/cli.py b/sky/cli.py
index e2461378fea..9e650fbb301 100644
--- a/sky/cli.py
+++ b/sky/cli.py
@@ -1132,16 +1132,13 @@ def _make_task_or_dag_from_entrypoint_with_overrides(
     old_resources = list(task.resources)[0]
     new_resources = old_resources.copy(**override_params)
 
-    # We need to update envs before we set resources, since the envs might
-    # contains docker login config, which will be used in setting up the
-    # resources.
-    task.update_envs(env)
     task.set_resources({new_resources})
 
     if num_nodes is not None:
         task.num_nodes = num_nodes
     if name is not None:
         task.name = name
+    task.update_envs(env)
     # TODO(wei-lin): move this validation into Python API.
     if new_resources.accelerators is not None:
         acc, _ = list(new_resources.accelerators.items())[0]
diff --git a/sky/skylet/constants.py b/sky/skylet/constants.py
index f657ec60264..f35c6ef3ba1 100644
--- 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
index 5951aac5e75..a0a4bf34375 100644
--- 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)
 
@@ -472,8 +483,6 @@ def update_envs(
         Raises:
           ValueError: if various invalid inputs errors are detected.
         """
-        # NOTE(dev): This should be called before task.set_resources() because
-        # of the docker login config.
         if envs is None:
             envs = {}
         if isinstance(envs, (list, tuple)):
@@ -496,6 +505,9 @@ def update_envs(
                     'envs must be List[Tuple[str, str]] or Dict[str, str]: '
                     f'{envs}')
         self._envs.update(envs)
+        if _check_docker_login_config(self._envs):
+            self.resources = _with_docker_login_config(self.resources,
+                                                       self._envs)
         return self
 
     @property
@@ -561,12 +573,11 @@ def set_resources(
         Returns:
           self: The current task, with resources set.
         """
-        # NOTE(dev): This should be called after task.update_envs() because of
-        # the docker login config.
         if isinstance(resources, sky.Resources):
             resources = {resources}
         # TODO(woosuk): Check if the resources are None.
-        self.resources = _with_docker_login_config(resources, self.envs)
+        if _check_docker_login_config(self.envs):
+            self.resources = _with_docker_login_config(resources, self.envs)
         return self
 
     def get_resources(self):