apply suggestion from code review

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 ' +

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}