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 9eb45ec1ea7..448c2339e23 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) @@ -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