Repository: skypilot-org/skypilot
Pull request: [Image] Support passing AMIs for different regions
Closes #1345 

Main changes:
1. Support a dict for `image_id`
2. Added 3 smoke tests for the image_id dict.

Notes:
1. The failover will only go through the regions listed in the `image_id`.
3. If a region or zone is specified and not in the `image_ids`, it will throw an error.
4. If a task to be run on a existing cluster has a different `image_id` (even if it is the same in the same region the cluster is in), we will just error out.

Tested:
- [x] `sky launch -c test-image-dict test.yaml`
```
# test.yaml
resources:
  cloud: aws
  image_id:
    us-east-1: ami-0fb70a87623506e3e
    us-west-2: ami-03e40ee6f58eea284
```
- [x] `sky launch -c test-image-dict-failover test.yaml`
```
# test.yaml
resources:
  cloud: aws
  image_id:
    us-east-1: ami-0fb70a87623506e3e
    us-west-2: invalid
```
- [x] `tests/run_smoke_tests.sh test_image_id_dict`
- [x] `tests/run_smoke_tests.sh test_image_id_dict_with_region`
- [x] `tests/run_smoke_tests.sh test_image_id_dict_with_zone`

Problem statement:
Support passing multiple `image_id`, one per region
From user, whose app requires a specific Ubuntu 18.04 version, and does not support the latest 20.04 we use:

> An alternative could be to allow specifying image_id as a mapping from region to AMI so that at least the user won’t be limited to specifying a fixed region. This would be more flexible (esp if user wants to use custom AMIs) but a bit more work on the user’s side.

Previously, @pounde also requested this feature.

-----

Proposal:
```
resources:
  image_id:
    us-east-1: <find your us-east-1 AMI>
    us-west-1: <find your us-west-1 AMI>
```
where SkyPilot will restrict itself to provision in these regions only.

Hints from issue discussion:
The API for having a dict for `image_id` leaks the things each field can control. Previously, the region controls the regions to retry on, and the `image_id` specifies the image only. However, with the current specification, there are now two ways to limit the regions to failover through:
```
resources:
  region: us-west-1
  image_id: xxx
```
```
resources:
  image_id:
    us-west-1: xxx
```
This makes the API not as clean, though I could not think of a better way to do it.

Commit message: add tests

Patch to review:
diff --git a/docs/source/reference/yaml-spec.rst b/docs/source/reference/yaml-spec.rst
index fa1e29a8b33..586e0366cbc 100644
--- a/docs/source/reference/yaml-spec.rst
+++ b/docs/source/reference/yaml-spec.rst
@@ -96,6 +96,11 @@ Available fields:
       #   image_id: skypilot:k80-ubuntu-2004
       #   image_id: skypilot:gpu-ubuntu-1804
       #   image_id: skypilot:k80-ubuntu-1804
+      # It is also possible to specify a per-region image id (failover will only go through the regions sepcified as keys; 
+      # useful when you have the custom images in multiple regions):
+      #   image_id:
+      #     us-east-1: ami-0729d913a335efca7
+      #     us-west-2: ami-050814f384259894c
       image_id: ami-0868a20f5a3bf9702
       # GCP
       # To find GCP images: https://cloud.google.com/compute/docs/images
diff --git a/sky/backends/cloud_vm_ray_backend.py b/sky/backends/cloud_vm_ray_backend.py
index 39a7302e67c..85d26bb5ca8 100644
--- a/sky/backends/cloud_vm_ray_backend.py
+++ b/sky/backends/cloud_vm_ray_backend.py
@@ -904,13 +904,10 @@ def _yield_region_zones(self, to_provision: resources_lib.Resources,
                 use_spot=to_provision.use_spot,
         ):
             # Only retry requested region/zones or all if not specified.
-            if (to_provision.region is not None and
-                    region.name != to_provision.region):
+            zone_names = [zone.name for zone in zones]
+            if not to_provision.valid_on_region_zones(region.name, zone_names):
                 continue
             if to_provision.zone is not None:
-                zones_name = [zone.name for zone in zones]
-                if to_provision.zone not in zones_name:
-                    continue
                 zones = [clouds.Zone(name=to_provision.zone)]
             yield (region, zones)
 
diff --git a/sky/cli.py b/sky/cli.py
index 7f1c5b3c6ad..11d6e13c84a 100644
--- a/sky/cli.py
+++ b/sky/cli.py
@@ -216,11 +216,10 @@ def _interactive_node_cli_command(cli_func):
                                  help=('Automatically stop the cluster after '
                                        'this many minutes of idleness, i.e. '
                                        'no running or pending jobs in the '
-                                       'cluster\'s job queue. Idleness starts '
-                                       'counting after setup/file_mounts are '
-                                       'done; the clock gets reset whenever '
-                                       'there are running/pending jobs in the '
-                                       'job queue. If not set, the cluster '
+                                       'cluster\'s job queue. Idleness gets '
+                                       'reset whenever setting-up/running/'
+                                       'pending jobs are found in the job '
+                                       'queue. If not set, the cluster '
                                        'will not be auto-stopped.'))
     autodown = click.option('--down',
                             default=False,
@@ -763,10 +762,9 @@ def _create_and_ssh_into_node(
         user_requested_resources: If true, user requested resources explicitly.
         no_confirm: If true, skips confirmation prompt presented to user.
         idle_minutes_to_autostop: Automatically stop the cluster after
-                                  specified minutes of idleness. Idleness
-                                  starts counting after setup/file_mounts are
-                                  done; the clock gets reset whenever there
-                                  are running/pending jobs in the job queue.
+                                  specified minutes of idleness. Idleness gets
+                                  reset whenever setting-up/running/pending
+                                  jobs are found in the job queue.
         down: If true, autodown the cluster after all jobs finish. If
               idle_minutes_to_autostop is also set, the cluster will be torn
               down after the specified idle time.
@@ -1096,9 +1094,8 @@ def cli():
     required=False,
     help=('Automatically stop the cluster after this many minutes '
           'of idleness, i.e., no running or pending jobs in the cluster\'s job '
-          'queue. Idleness starts counting after setup/file_mounts are done; '
-          'the clock gets reset whenever there are running/pending jobs in the '
-          'job queue. '
+          'queue. Idleness gets reset whenever setting-up/running/pending jobs '
+          'are found in the job queue. '
           'Setting this flag is equivalent to '
           'running ``sky launch -d ...`` and then ``sky autostop -i <minutes>``'
           '. If not set, the cluster will not be autostopped.'))
@@ -1760,9 +1757,8 @@ def autostop(
     required=False,
     help=('Automatically stop the cluster after this many minutes '
           'of idleness, i.e., no running or pending jobs in the cluster\'s job '
-          'queue. Idleness starts counting after setup/file_mounts are done; '
-          'the clock gets reset whenever there are running/pending jobs in the '
-          'job queue. '
+          'queue. Idleness gets reset whenever setting-up/running/pending jobs '
+          'are found in the job queue. '
           'Setting this flag is equivalent to '
           'running ``sky launch -d ...`` and then ``sky autostop -i <minutes>``'
           '. If not set, the cluster will not be autostopped.'))
diff --git a/sky/clouds/aws.py b/sky/clouds/aws.py
index 2d6443f4a71..e39f580ac15 100644
--- a/sky/clouds/aws.py
+++ b/sky/clouds/aws.py
@@ -115,8 +115,13 @@ def get_default_ami(cls, region_name: str, instance_type: str) -> str:
 
     @classmethod
     def _get_image_id(cls, region_name: str, instance_type: str,
-                      image_id: Optional[str]) -> str:
+                      image_id: Optional[Dict[str, str]]) -> str:
         if image_id is not None:
+            if None in image_id:
+                image_id = image_id[None]
+            else:
+                assert region_name in image_id, image_id
+                image_id = image_id[region_name]
             if image_id.startswith('skypilot:'):
                 image_id = service_catalog.get_image_id_from_tag(image_id,
                                                                  region_name,
diff --git a/sky/clouds/gcp.py b/sky/clouds/gcp.py
index 8d571a526d2..e176e20d2e8 100644
--- a/sky/clouds/gcp.py
+++ b/sky/clouds/gcp.py
@@ -273,7 +273,11 @@ def make_deploy_resources_variables(
                         'skypilot:gpu-debian-10', clouds='gcp')
 
         if resources.image_id is not None:
-            image_id = resources.image_id
+            if None in resources.image_id:
+                image_id = resources.image_id[None]
+            else:
+                assert region_name in resources.image_id, resources.image_id
+                image_id = resources.image_id[region_name]
 
         assert image_id is not None, (image_id, r)
         resources_vars['image_id'] = image_id
diff --git a/sky/clouds/service_catalog/__init__.py b/sky/clouds/service_catalog/__init__.py
index d84676aaa6c..f259dd4f268 100644
--- a/sky/clouds/service_catalog/__init__.py
+++ b/sky/clouds/service_catalog/__init__.py
@@ -98,9 +98,10 @@ def instance_type_exists(instance_type: str,
     return _map_clouds_catalog(clouds, 'instance_type_exists', instance_type)
 
 
-def validate_region_zone(region_name: Optional[str],
-                         zone_name: Optional[str],
-                         clouds: CloudFilter = None) -> bool:
+def validate_region_zone(
+        region_name: Optional[str],
+        zone_name: Optional[str],
+        clouds: CloudFilter = None) -> Tuple[Optional[str], Optional[str]]:
     """Returns the zone by name."""
     return _map_clouds_catalog(clouds, 'validate_region_zone', region_name,
                                zone_name)
diff --git a/sky/clouds/service_catalog/aws_catalog.py b/sky/clouds/service_catalog/aws_catalog.py
index 67265f586f0..172cd9afd33 100644
--- a/sky/clouds/service_catalog/aws_catalog.py
+++ b/sky/clouds/service_catalog/aws_catalog.py
@@ -19,7 +19,9 @@ def instance_type_exists(instance_type: str) -> bool:
     return common.instance_type_exists_impl(_df, instance_type)
 
 
-def validate_region_zone(region: Optional[str], zone: Optional[str]):
+def validate_region_zone(
+        region: Optional[str],
+        zone: Optional[str]) -> Tuple[Optional[str], Optional[str]]:
     return common.validate_region_zone_impl(_df, region, zone)
 
 
diff --git a/sky/clouds/service_catalog/azure_catalog.py b/sky/clouds/service_catalog/azure_catalog.py
index 7c571dfd33f..930c6499db9 100644
--- a/sky/clouds/service_catalog/azure_catalog.py
+++ b/sky/clouds/service_catalog/azure_catalog.py
@@ -16,7 +16,9 @@ def instance_type_exists(instance_type: str) -> bool:
     return common.instance_type_exists_impl(_df, instance_type)
 
 
-def validate_region_zone(region: Optional[str], zone: Optional[str]):
+def validate_region_zone(
+        region: Optional[str],
+        zone: Optional[str]) -> Tuple[Optional[str], Optional[str]]:
     if zone is not None:
         with ux_utils.print_exception_no_traceback():
             raise ValueError('Azure does not support zones.')
diff --git a/sky/clouds/service_catalog/common.py b/sky/clouds/service_catalog/common.py
index 2fcffd6afb2..e176fb0563a 100644
--- a/sky/clouds/service_catalog/common.py
+++ b/sky/clouds/service_catalog/common.py
@@ -100,8 +100,9 @@ def instance_type_exists_impl(df: pd.DataFrame, instance_type: str) -> bool:
     return instance_type in df['InstanceType'].unique()
 
 
-def validate_region_zone_impl(df: pd.DataFrame, region: Optional[str],
-                              zone: Optional[str]):
+def validate_region_zone_impl(
+        df: pd.DataFrame, region: Optional[str],
+        zone: Optional[str]) -> Tuple[Optional[str], Optional[str]]:
     """Validates whether region and zone exist in the catalog."""
 
     def _get_candidate_str(loc: str, all_loc: List[str]) -> List[str]:
@@ -113,30 +114,32 @@ def _get_candidate_str(loc: str, all_loc: List[str]) -> List[str]:
             candidate_strs = f'\nDid you mean one of these: {candidate_strs!r}?'
         return candidate_strs
 
+    validated_region, validated_zone = region, zone
+
+    filter_df = df
     if region is not None:
-        all_regions = df['Region'].unique()
-        if region not in all_regions:
+        filter_df = filter_df[filter_df['Region'] == region]
+        if len(filter_df) == 0:
             with ux_utils.print_exception_no_traceback():
                 error_msg = (f'Invalid region {region!r}')
-                error_msg += _get_candidate_str(region, all_regions)
+                error_msg += _get_candidate_str(region, df['Region'].unique())
                 raise ValueError(error_msg)
 
     if zone is not None:
-        all_zones = df['AvailabilityZone'].unique()
-        if zone not in all_zones:
-            with ux_utils.print_exception_no_traceback():
-                error_msg = (f'Invalid zone {zone!r}')
-                error_msg += _get_candidate_str(zone, all_zones)
-                raise ValueError(error_msg)
-
-    if region is not None and zone is not None:
-        if zone not in df[df['Region'] == region]['AvailabilityZone'].unique():
+        maybe_region_df = filter_df
+        filter_df = filter_df[filter_df['AvailabilityZone'] == zone]
+        if len(filter_df) == 0:
+            region_str = f' for region {region!r}' if region else ''
+            df = maybe_region_df if region else df
             with ux_utils.print_exception_no_traceback():
-                error_msg = (f'Invalid zone {zone!r} for region {region!r}')
+                error_msg = (f'Invalid zone {zone!r}{region_str}')
                 error_msg += _get_candidate_str(
-                    zone,
-                    df[df['Region'] == region]['AvailabilityZone'].unique())
+                    zone, maybe_region_df['AvailabilityZone'].unique())
                 raise ValueError(error_msg)
+        region_df = filter_df['Region'].unique()
+        assert len(region_df) == 1, 'Zone should be unique across regions.'
+        validated_region = region_df[0]
+    return validated_region, validated_zone
 
 
 def get_hourly_cost_impl(
diff --git a/sky/clouds/service_catalog/gcp_catalog.py b/sky/clouds/service_catalog/gcp_catalog.py
index 663058167bc..79eecd28c9f 100644
--- a/sky/clouds/service_catalog/gcp_catalog.py
+++ b/sky/clouds/service_catalog/gcp_catalog.py
@@ -202,7 +202,9 @@ def get_instance_type_for_accelerator(
     return [f'{_DEFAULT_HOST_VM_FAMILY}-{mem_type}-{num_cpus}'], []
 
 
-def validate_region_zone(region: Optional[str], zone: Optional[str]):
+def validate_region_zone(
+        region: Optional[str],
+        zone: Optional[str]) -> Tuple[Optional[str], Optional[str]]:
     return common.validate_region_zone_impl(_df, region, zone)
 
 
diff --git a/sky/core.py b/sky/core.py
index 1db8ee8fad8..59189267a1b 100644
--- a/sky/core.py
+++ b/sky/core.py
@@ -131,12 +131,12 @@ def start(
         cluster_name: name of the cluster to start.
         idle_minutes_to_autostop: automatically stop the cluster after this
             many minute of idleness, i.e., no running or pending jobs in the
-            cluster's job queue. Idleness starts counting after
-            setup/file_mounts are done; the clock gets reset whenever there
-            are running/pending jobs in the job queue. Setting this flag is
-            equivalent to running ``sky.launch(..., detach_run=True, ...)``
-            and then ``sky.autostop(idle_minutes=<minutes>)``. If not set,
-            the cluster will not be autostopped.
+            cluster's job queue. Idleness gets reset whenever setting-up/
+            running/pending jobs are found in the job queue. Setting this
+            flag is equivalent to running
+            ``sky.launch(..., detach_run=True, ...)`` and then
+            ``sky.autostop(idle_minutes=<minutes>)``. If not set, the
+            cluster will not be autostopped.
         retry_until_up: whether to retry launching the cluster until it is
             up.
         down: Autodown the cluster: tear down the cluster after specified
diff --git a/sky/execution.py b/sky/execution.py
index 8b6ce6baf91..9639965b28b 100644
--- a/sky/execution.py
+++ b/sky/execution.py
@@ -318,12 +318,12 @@ def launch(
             up.
         idle_minutes_to_autostop: automatically stop the cluster after this
             many minute of idleness, i.e., no running or pending jobs in the
-            cluster's job queue. Idleness starts counting after
-            setup/file_mounts are done; the clock gets reset whenever there
-            are running/pending jobs in the job queue. Setting this flag is
-            equivalent to running ``sky.launch(..., detach_run=True, ...)``
-            and then ``sky.autostop(idle_minutes=<minutes>)``. If not set,
-            the cluster will not be autostopped.
+            cluster's job queue. Idleness gets reset whenever setting-up/
+            running/pending jobs are found in the job queue. Setting this
+            flag is equivalent to running
+            ``sky.launch(..., detach_run=True, ...)`` and then
+            ``sky.autostop(idle_minutes=<minutes>)``. If not set, the cluster
+            will not be autostopped.
         down: Tear down the cluster after all jobs finish (successfully or
             abnormally). If --idle-minutes-to-autostop is also set, the
             cluster will be torn down after the specified idle time.
diff --git a/sky/resources.py b/sky/resources.py
index ab9ea9161c2..99653624a8b 100644
--- a/sky/resources.py
+++ b/sky/resources.py
@@ -57,7 +57,7 @@ def __init__(
         disk_size: Optional[int] = None,
         region: Optional[str] = None,
         zone: Optional[str] = None,
-        image_id: Optional[str] = None,
+        image_id: Union[Dict[str, str], str, None] = None,
     ):
         self._version = self._VERSION
         self._cloud = cloud
@@ -88,6 +88,8 @@ def __init__(
             self._disk_size = _DEFAULT_DISK_SIZE_GB
 
         self._image_id = image_id
+        if isinstance(image_id, str):
+            self._image_id = {None: image_id}
 
         self._set_accelerators(accelerators, accelerator_args)
 
@@ -114,7 +116,10 @@ def __repr__(self) -> str:
 
         image_id = ''
         if self.image_id is not None:
-            image_id = f', image_id={self.image_id!r}'
+            if None in self.image_id:
+                image_id = f', image_id={self.image_id[None]}'
+            else:
+                image_id = f', image_id={self.image_id!r}'
 
         disk_size = ''
         if self.disk_size != _DEFAULT_DISK_SIZE_GB:
@@ -175,7 +180,7 @@ def disk_size(self) -> int:
         return self._disk_size
 
     @property
-    def image_id(self) -> Optional[str]:
+    def image_id(self) -> Optional[Dict[str, str]]:
         return self._image_id
 
     def _set_accelerators(
@@ -263,11 +268,10 @@ def _set_region_zone(self, region: Optional[str],
                 raise ValueError(
                     'Cloud must be specified when region/zone are specified.')
 
-        # Validate whether region and zone exist in the catalog.
-        self._cloud.validate_region_zone(region, zone)
-
-        self._region = region
-        self._zone = zone
+        # Validate whether region and zone exist in the catalog, and set the
+        # region if zone is specified.
+        self._region, self._zone = self._cloud.validate_region_zone(
+            region, zone)
 
     def _try_validate_instance_type(self):
         if self.instance_type is None:
@@ -404,23 +408,29 @@ def _try_validate_image_id(self) -> None:
                     'image_id is only supported for AWS and GCP, please '
                     'explicitly specify the cloud.')
 
-        if (self._image_id.startswith('skypilot:') and
-                not self._cloud.is_image_tag_valid(self._image_id,
-                                                   self._region)):
-            region_or_zone = self._region or self._zone
-            region_str = f' ({region_or_zone})' if region_or_zone else ''
+        if (self._region is not None and None not in self._image_id and
+                self._region not in self._image_id):
             with ux_utils.print_exception_no_traceback():
                 raise ValueError(
-                    f'Image tag {self._image_id} is not valid, please make sure'
-                    f' the tag exists in {self._cloud}{region_str}.')
+                    'image_id should contain the image for the specified '
+                    f'region {self._region}.')
+
+        for region, image_id in self._image_id.items():
+            region = region or self._region
+            if (image_id.startswith('skypilot:') and
+                    not self._cloud.is_image_tag_valid(image_id, region)):
+                region_str = f' ({region})' if region else ''
+                with ux_utils.print_exception_no_traceback():
+                    raise ValueError(
+                        f'Image tag {image_id} is not valid, please make sure'
+                        f' the tag exists in {self._cloud}{region_str}.')
 
-        if (self._cloud.is_same_cloud(clouds.AWS()) and
-                not self._image_id.startswith('skypilot:') and
-                self._region is None):
-            with ux_utils.print_exception_no_traceback():
-                raise ValueError(
-                    'image_id is only supported for AWS in a specific '
-                    'region, please explicitly specify the region.')
+            if (self._cloud.is_same_cloud(clouds.AWS()) and
+                    not image_id.startswith('skypilot:') and region is None):
+                with ux_utils.print_exception_no_traceback():
+                    raise ValueError(
+                        'image_id is only supported for AWS in a specific '
+                        'region, please explicitly specify the region.')
 
     def get_cost(self, seconds: float) -> float:
         """Returns cost in USD for the runtime in seconds."""
@@ -584,6 +594,17 @@ def copy(self, **override) -> 'Resources':
         assert len(override) == 0
         return resources
 
+    def valid_on_region_zones(self, region: str, zones: List[str]) -> bool:
+        """Returns whether this Resources is valid on given region and zones"""
+        if self.region is not None and self.region != region:
+            return False
+        if self.zone is not None and self.zone not in zones:
+            return False
+        if self.image_id is not None:
+            if None not in self.image_id and region not in self.image_id:
+                return False
+        return True
+
     @classmethod
     def from_yaml_config(cls, config: Optional[Dict[str, str]]) -> 'Resources':
         if config is None:
diff --git a/sky/utils/schemas.py b/sky/utils/schemas.py
index 3493ebaff7e..928505957bf 100644
--- a/sky/utils/schemas.py
+++ b/sky/utils/schemas.py
@@ -66,7 +66,13 @@ def get_resources_schema():
                 }
             },
             'image_id': {
-                'type': 'string',
+                'anyOf': [{
+                    'type': 'string',
+                }, {
+                    'type': 'object',
+                    'required': [],
+                    'additionalProperties': True,
+                }]
             }
         }
     }
diff --git a/tests/test_smoke.py b/tests/test_smoke.py
index 9f4a736bfb2..064c624c4b2 100644
--- a/tests/test_smoke.py
+++ b/tests/test_smoke.py
@@ -155,8 +155,9 @@ def test_region():
             f'sky launch -y -c {name} --region us-west-2 examples/minimal.yaml',
             f'sky exec {name} examples/minimal.yaml',
             f'sky logs {name} 1 --status',  # Ensure the job succeeded.
+            f'sky status --all | grep {name} | grep us-west-2',  # Ensure the region is correct.
         ],
-        f'sky down -y {name}',
+        f'sky down -y {name} {name}-2 {name}-3',
     )
     run_one_test(test)
 
@@ -172,6 +173,82 @@ def test_zone():
             f'sky logs {name} 1 --status',  # Ensure the job succeeded.
             f'sky status --all | grep {name} | grep us-west-2b',  # Ensure the zone is correct.
         ],
+        f'sky down -y {name} {name}-2 {name}-3',
+    )
+    run_one_test(test)
+
+
+def test_image_id_dict():
+    name = _get_cluster_name()
+    test = Test(
+        'image_id_dict',
+        [
+            # Use image id dict.
+            f'sky launch -y -c {name} examples/per_region_images.yaml',
+            f'sky exec {name} examples/per_region_images.py',
+            f'sky exec {name} "ls ~"',
+            f'sky logs {name} 1 --status',
+            f'sky logs {name} 2 --status',
+            f'sky logs {name} 3 --status',
+        ],
+        f'sky down -y {name} {name}-2 {name}-3',
+    )
+    run_one_test(test)
+
+
+def test_image_id_dict_with_region():
+    name = _get_cluster_name()
+    test = Test(
+        'image_id_dict_with_region',
+        [
+            # Use region to filter image_id dict.
+            f'sky launch -y -c {name} --region us-west-2 examples/per_region_images.yaml && exit 1 || true',
+            f'sky status | grep {name} && exit 1 || true',  # Ensure the cluster is not created.
+            f'sky launch -y -c {name} --region us-west-1 examples/per_region_images.yaml',
+            # Fail the task if even if the image id match for the region (image_id should be the same dict).
+            f'sky launch -y -c {name} --image-id skypilot:gpu-ubuntu-1804 examples/minimal.yaml && exit 1 || true',
+            f'sky exec {name} --image-id skypilot:gpu-ubuntu-1804 examples/minimal.yaml && exit 1 || true',
+            f'sky status --all | grep {name} | grep us-west-1',  # Ensure the region is correct.
+            # Ensure exec works.
+            f'sky exec {name} --region us-west-1 examples/per_region_images.yaml',
+            f'sky exec {name} examples/per_region_images.yaml',
+            f'sky exec {name} --cloud aws --region us-west-1 "ls ~"',
+            f'sky exec {name} "ls ~"',
+            f'sky logs {name} 1 --status',
+            f'sky logs {name} 2 --status',
+            f'sky logs {name} 3 --status',
+            f'sky logs {name} 4 --status',
+            f'sky logs {name} 5 --status',
+        ],
+        f'sky down -y {name}',
+    )
+    run_one_test(test)
+
+
+def test_image_id_dict_with_zone():
+    name = _get_cluster_name()
+    test = Test(
+        'image_id_dict_with_region',
+        [
+            # Use zone to filter image_id dict.
+            f'sky launch -y -c {name} --zone us-west-2b examples/per_region_images.yaml && exit 1 || true',
+            f'sky status | grep {name} && exit 1 || true',  # Ensure the cluster is not created.
+            f'sky launch -y -c {name} --zone us-west-1a examples/per_region_images.yaml',
+            # Fail the task if even if the image id match for the region (image_id should be the same dict).
+            f'sky launch -y -c {name} --image-id skypilot:gpu-ubuntu-1804 examples/minimal.yaml && exit 1 || true',
+            f'sky exec {name} --image-id skypilot:gpu-ubuntu-1804 examples/minimal.yaml && exit 1 || true',
+            f'sky status --all | grep {name} | grep us-west-1a',  # Ensure the zone is correct.
+            # Ensure exec works.
+            f'sky exec {name} --zone us-west-1a examples/per_region_images.yaml',
+            f'sky exec {name} examples/per_region_images.yaml',
+            f'sky exec {name} --cloud aws --region us-west-1 "ls ~"',
+            f'sky exec {name} "ls ~"',
+            f'sky logs {name} 1 --status',
+            f'sky logs {name} 2 --status',
+            f'sky logs {name} 3 --status',
+            f'sky logs {name} 4 --status',
+            f'sky logs {name} 5 --status',
+        ],
         f'sky down -y {name}',
     )
     run_one_test(test)


Code review comments:
diff_hunk: @@ -66,7 +66,13 @@ def get_resources_schema():
                 }
             },
             'image_id': {
-                'type': 'string',
+                'anyOf': [{
+                    'type': 'string',
+                }, {
+                    'type': 'object',
+                    'required': [],
+                    'additionalProperties': True,
path: sky/utils/schemas.py
text: @user1:
nit: `additionalProperties` defaults to True if not specified, maybe remove?

@author:
Good point! Removed. Thanks!

diff_hunk: @@ -155,8 +155,9 @@ def test_region():
             f'sky launch -y -c {name} --region us-west-2 examples/minimal.yaml',
             f'sky exec {name} examples/minimal.yaml',
             f'sky logs {name} 1 --status',  # Ensure the job succeeded.
+            f'sky status --all | grep {name} | grep us-west-2',  # Ensure the region is correct.
         ],
-        f'sky down -y {name}',
+        f'sky down -y {name} {name}-2 {name}-3',
path: tests/test_smoke.py
text: @user1:
Is this change required? Not sure how `{name}-2` and `{name}-3` clusters are being created. Same at L176

@author:
Oops, this is a remnant. Removed. Thanks!

diff_hunk: @@ -172,6 +173,82 @@ def test_zone():
             f'sky logs {name} 1 --status',  # Ensure the job succeeded.
             f'sky status --all | grep {name} | grep us-west-2b',  # Ensure the zone is correct.
         ],
+        f'sky down -y {name} {name}-2 {name}-3',
+    )
+    run_one_test(test)
+
+
+def test_image_id_dict():
+    name = _get_cluster_name()
+    test = Test(
+        'image_id_dict',
+        [
+            # Use image id dict.
+            f'sky launch -y -c {name} examples/per_region_images.yaml',
+            f'sky exec {name} examples/per_region_images.py',
+            f'sky exec {name} "ls ~"',
+            f'sky logs {name} 1 --status',
+            f'sky logs {name} 2 --status',
+            f'sky logs {name} 3 --status',
+        ],
+        f'sky down -y {name} {name}-2 {name}-3',
+    )
+    run_one_test(test)
+
+
+def test_image_id_dict_with_region():
+    name = _get_cluster_name()
+    test = Test(
+        'image_id_dict_with_region',
+        [
+            # Use region to filter image_id dict.
+            f'sky launch -y -c {name} --region us-west-2 examples/per_region_images.yaml && exit 1 || true',
+            f'sky status | grep {name} && exit 1 || true',  # Ensure the cluster is not created.
+            f'sky launch -y -c {name} --region us-west-1 examples/per_region_images.yaml',
+            # Fail the task if even if the image id match for the region (image_id should be the same dict).
path: tests/test_smoke.py
text: @user1:
nit: typo: "Fail the task ~if~ even if"