From 9aa1636b862aa35fd0f05bf5e141578110032930 Mon Sep 17 00:00:00 2001
From: cblmemo <cblmemo@gmail.com>
Date: Mon, 18 Sep 2023 00:23:06 -0700
Subject: [PATCH] apply suggestion from code review

---
 sky/skylet/providers/command_runner.py |  4 +---
 sky/task.py                            | 12 ++++++++++--
 2 files changed, 11 insertions(+), 5 deletions(-)

diff --git a/sky/skylet/providers/command_runner.py b/sky/skylet/providers/command_runner.py
index e12305882b9..02eea737727 100644
--- a/sky/skylet/providers/command_runner.py
+++ b/sky/skylet/providers/command_runner.py
@@ -143,9 +143,7 @@ def run_init(self, *, as_head: bool, file_mounts: Dict[str, str],
                 docker_login_config.password,
                 docker_login_config.server,
             ))
-            server_prefix = f'{docker_login_config.server}/'
-            if not specific_image.startswith(server_prefix):
-                specific_image = f'{server_prefix}{specific_image}'
+            specific_image = f'{docker_login_config.server}/{specific_image}'
 
         if self.docker_config.get('pull_before_run', True):
             assert specific_image, ('Image must be included in config if ' +
diff --git a/sky/task.py b/sky/task.py
index 9eb45ec1ea7..9e9850f83af 100644
--- a/sky/task.py
+++ b/sky/task.py
@@ -126,13 +126,21 @@ def _with_docker_login_config(
         task_envs)
 
     def _add_docker_login_config(resources: 'resources_lib.Resources'):
-        if resources.extract_docker_image() is None:
+        docker_image = resources.extract_docker_image()
+        if 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}')
             return resources
-        return resources.copy(_docker_login_config=docker_login_config)
+        # Already checked in extract_docker_image
+        assert len(resources.image_id) == 1
+        region = list(resources.image_id.keys())[0]
+        server_prefix = f'{docker_login_config.server}'
+        if not docker_image.startswith(server_prefix):
+            docker_image = f'{server_prefix}{docker_image}'
+        return resources.copy(image_id={region: 'docker:' + docker_image},
+                              _docker_login_config=docker_login_config)
 
     return {_add_docker_login_config(resources) for resources in resources_set}