Add zone support in YAML (skypilot-org#1014)

* Add zone support

* Zone validation and acc validation

* Remove CLI overwrite

* yapf

* fix

* address comments

* yapf

* address comments

* yapf

* fix

* fix and remove test

* fix

* fix test

* fix huggingface smoke test

* handle multi-node

* Replace ValueError with assert

* fix multi-node

* replace sky.cloud with clouds.cloud

* Add invalid zone test

* add smoke test

* yapf

* Support region/zone close matching

* Don't overwrite ray yaml

* add comments

* add zone info to usage_lib

* update docstring

* roll back sky.clouds

* refactor

* missed

* Fix

* yapf

* add zone to cli

* update test

* fix

* Add acc_count constraint

* remove unused func

* shorten

* comments

* fix

Author: infwinston

docs/source/reference/yaml-spec.rst

@@ -34,6 +34,9 @@ describe all fields available.
       # if this is specified.
       region: us-east-1
 
+      # The zone to use (optional).
+      region: us-east-1a
+
       # Accelerator name and count per node (optional).
       #
       # Use `sky show-gpus` to view available accelerator configurations.

examples/horovod_distributed_tf_app.py

@@ -4,7 +4,6 @@
 
 import sky
 import time_estimators
-from sky import clouds
 
 IPAddr = str
 
@@ -55,7 +54,7 @@ def run_fn(ip_list: List[IPAddr]) -> Dict[IPAddr, str]:
                      estimated_size_gigabytes=70)
     train.set_outputs('resnet-model-dir', estimated_size_gigabytes=0.1)
     train.set_resources({
-        sky.Resources(clouds.AWS(), 'p3.2xlarge'),
+        sky.Resources(sky.AWS(), 'p3.2xlarge'),
     })
 
 dag = sky.Optimizer.optimize(dag)

examples/ray_tune_app.py

@@ -2,7 +2,6 @@
 from typing import List, Optional
 
 import sky
-from sky import clouds
 
 with sky.Dag() as dag:
     # Total Nodes, INCLUDING Head Node
@@ -31,7 +30,7 @@ def run_fn(node_rank: int, ip_list: List[str]) -> Optional[str]:
     )
 
     train.set_resources({
-        sky.Resources(clouds.AWS(), 'p3.2xlarge'),
+        sky.Resources(sky.AWS(), 'p3.2xlarge'),
     })
 
 sky.launch(dag)

sky/backends/backend_utils.py

@@ -1075,7 +1075,7 @@ def _process_cli_query(
                  f'{stdout}\n'
                  '**** STDERR ****\n'
                  f'{stderr}')
-    if (cloud == str(sky.Azure()) and returncode == 2 and
+    if (cloud == str(clouds.Azure()) and returncode == 2 and
             'argument --ids: expected at least one argument' in stderr):
         # Azure CLI has a returncode 2 when the cluster is not found, as
         # --ids <empty> is passed to the query command. In that case, the

sky/backends/cloud_vm_ray_backend.py

@@ -686,25 +686,32 @@ def _yield_region_zones(self, to_provision: resources_lib.Resources,
                 prev_resources = handle.launched_resources
                 if prev_resources is not None and cloud.is_same_cloud(
                         prev_resources.cloud):
-                    if cloud.is_same_cloud(sky.GCP()) or cloud.is_same_cloud(
-                            sky.AWS()):
+                    if cloud.is_same_cloud(clouds.GCP()) or cloud.is_same_cloud(
+                            clouds.AWS()):
                         region = config['provider']['region']
                         zones = config['provider']['availability_zone']
-                    elif cloud.is_same_cloud(sky.Azure()):
+                    elif cloud.is_same_cloud(clouds.Azure()):
                         region = config['provider']['location']
                         zones = None
-                    elif cloud.is_same_cloud(sky.Local()):
+                    elif cloud.is_same_cloud(clouds.Local()):
                         local_regions = clouds.Local.regions()
                         region = local_regions[0].name
                         zones = None
                     else:
                         assert False, cloud
-                    if region != prev_resources.region:
-                        raise ValueError(
-                            f'Region mismatch. The region in '
-                            f'{handle.cluster_yaml} '
-                            'has been changed from '
-                            f'{prev_resources.region} to {region}.')
+                    assert region == prev_resources.region, (
+                        f'Region mismatch. The region in '
+                        f'{handle.cluster_yaml} '
+                        'has been changed from '
+                        f'{prev_resources.region} to {region}.')
+                    assert (zones is None or prev_resources.zone is None or
+                            prev_resources.zone
+                            in zones), (f'{prev_resources.zone} not found in '
+                                        f'zones of {handle.cluster_yaml}.')
+                    # Note that we don't overwrite the zone field in Ray YAML
+                    # even if prev_resources.zone != zones.
+                    # This is because Ray will consider the YAML hash changed
+                    # and not reuse the existing cluster.
             except FileNotFoundError:
                 # Happens if no previous cluster.yaml exists.
                 pass
@@ -798,10 +805,15 @@ def _yield_region_zones(self, to_provision: resources_lib.Resources,
                 accelerators=to_provision.accelerators,
                 use_spot=to_provision.use_spot,
         ):
-            # Do not retry on region if it's not in the requested region.
+            # Only retry requested region/zones or all if not specified.
             if (to_provision.region is not None and
                     region.name != to_provision.region):
                 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)
 
     def _try_provision_tpu(self, to_provision: resources_lib.Resources,
@@ -1410,12 +1422,12 @@ def _update_cluster_region(self):
             config = common_utils.read_yaml(self.cluster_yaml)
             provider = config['provider']
             cloud = self.launched_resources.cloud
-            if cloud.is_same_cloud(sky.Azure()):
+            if cloud.is_same_cloud(clouds.Azure()):
                 region = provider['location']
-            elif cloud.is_same_cloud(sky.GCP()) or cloud.is_same_cloud(
-                    sky.AWS()):
+            elif cloud.is_same_cloud(clouds.GCP()) or cloud.is_same_cloud(
+                    clouds.AWS()):
                 region = provider['region']
-            elif cloud.is_same_cloud(sky.Local()):
+            elif cloud.is_same_cloud(clouds.Local()):
                 # There is only 1 region for Local cluster, 'Local'.
                 local_regions = clouds.Local.regions()
                 region = local_regions[0].name
@@ -1495,6 +1507,16 @@ def check_resources_fit_cluster(self, handle: ResourceHandle,
                         'Task requested resources in region '
                         f'{task_resources.region!r}, but the existing cluster '
                         f'is in region {launched_resources.region!r}.')
+            if (task_resources.zone is not None and
+                    task_resources.zone != launched_resources.zone):
+                zone_str = (f'is in zone {launched_resources.zone!r}.'
+                            if launched_resources.zone is not None else
+                            'does not have zone specified.')
+                with ux_utils.print_exception_no_traceback():
+                    raise exceptions.ResourcesMismatchError(
+                        'Task requested resources in zone '
+                        f'{task_resources.zone!r}, but the existing cluster '
+                        f'{zone_str}')
             with ux_utils.print_exception_no_traceback():
                 raise exceptions.ResourcesMismatchError(
                     'Requested resources do not match the existing cluster.\n'
@@ -1620,6 +1642,21 @@ def _provision(self,
                 # TPU.
                 tpu_create_script=config_dict.get('tpu-create-script'),
                 tpu_delete_script=config_dict.get('tpu-delete-script'))
+
+            # Get actual zone info and save it into handle
+            get_zone_cmd = handle.launched_resources.cloud.get_zone_shell_cmd()
+            if get_zone_cmd is not None:
+                # We leave the zone field to None for multi-node cases
+                # if zone is not specified because head and worker nodes
+                # can be launched in different zones.
+                if (task.num_nodes == 1 or
+                        handle.launched_resources.zone is not None):
+                    returncode, stdout, _ = self.run_on_head(
+                        handle, get_zone_cmd, require_outputs=True)
+                    # zone will be checked during Resources cls initialization.
+                    handle.launched_resources = handle.launched_resources.copy(
+                        zone=stdout.strip())
+
             usage_lib.messages.usage.update_cluster_resources(
                 handle.launched_nodes, handle.launched_resources)
             usage_lib.messages.usage.update_final_cluster_status(

sky/cli.py

@@ -249,6 +249,12 @@ def _parse_env_var(env_var: str) -> Tuple[str, str]:
         type=str,
         help=('The region to use. If specified, overrides the '
               '"resources.region" config. Passing "none" resets the config.')),
+    click.option(
+        '--zone',
+        required=False,
+        type=str,
+        help=('The zone to use. If specified, overrides the '
+              '"resources.zone" config. Passing "none" resets the config.')),
     click.option(
         '--num-nodes',
         required=False,
@@ -328,6 +334,7 @@ def _add_options(func):
 
 def _parse_override_params(cloud: Optional[str] = None,
                            region: Optional[str] = None,
+                           zone: Optional[str] = None,
                            gpus: Optional[str] = None,
                            instance_type: Optional[str] = None,
                            use_spot: Optional[bool] = None,
@@ -345,6 +352,11 @@ def _parse_override_params(cloud: Optional[str] = None,
             override_params['region'] = None
         else:
             override_params['region'] = region
+    if zone is not None:
+        if zone.lower() == 'none':
+            override_params['zone'] = None
+        else:
+            override_params['zone'] = zone
     if gpus is not None:
         if gpus.lower() == 'none':
             override_params['accelerators'] = None
@@ -684,6 +696,7 @@ def _make_dag_from_entrypoint_with_overrides(
     workdir: Optional[str] = None,
     cloud: Optional[str] = None,
     region: Optional[str] = None,
+    zone: Optional[str] = None,
     gpus: Optional[str] = None,
     instance_type: Optional[str] = None,
     num_nodes: Optional[int] = None,
@@ -695,7 +708,6 @@ def _make_dag_from_entrypoint_with_overrides(
     spot_recovery: Optional[str] = None,
 ) -> sky.Dag:
     entrypoint = ' '.join(entrypoint)
-
     with sky.Dag() as dag:
         is_yaml, yaml_config = _check_yaml(entrypoint)
         if is_yaml:
@@ -726,6 +738,7 @@ def _make_dag_from_entrypoint_with_overrides(
 
         override_params = _parse_override_params(cloud=cloud,
                                                  region=region,
+                                                 zone=zone,
                                                  gpus=gpus,
                                                  instance_type=instance_type,
                                                  use_spot=use_spot,
@@ -877,6 +890,7 @@ def launch(
     workdir: Optional[str],
     cloud: Optional[str],
     region: Optional[str],
+    zone: Optional[str],
     gpus: Optional[str],
     instance_type: Optional[str],
     num_nodes: Optional[int],
@@ -908,6 +922,7 @@ def launch(
         workdir=workdir,
         cloud=cloud,
         region=region,
+        zone=zone,
         gpus=gpus,
         instance_type=instance_type,
         num_nodes=num_nodes,
@@ -956,6 +971,7 @@ def exec(
     name: Optional[str],
     cloud: Optional[str],
     region: Optional[str],
+    zone: Optional[str],
     workdir: Optional[str],
     gpus: Optional[str],
     instance_type: Optional[str],
@@ -1038,6 +1054,7 @@ def exec(
         workdir=workdir,
         cloud=cloud,
         region=region,
+        zone=zone,
         gpus=gpus,
         instance_type=instance_type,
         use_spot=use_spot,
@@ -2239,6 +2256,7 @@ def spot_launch(
     workdir: Optional[str],
     cloud: Optional[str],
     region: Optional[str],
+    zone: Optional[str],
     gpus: Optional[str],
     instance_type: Optional[str],
     num_nodes: Optional[int],
@@ -2262,6 +2280,7 @@ def spot_launch(
         workdir=workdir,
         cloud=cloud,
         region=region,
+        zone=zone,
         gpus=gpus,
         instance_type=instance_type,
         num_nodes=num_nodes,
@@ -2510,6 +2529,7 @@ def benchmark_launch(
     workdir: Optional[str],
     cloud: Optional[str],
     region: Optional[str],
+    zone: Optional[str],
     gpus: Optional[str],
     num_nodes: Optional[int],
     use_spot: Optional[bool],
@@ -2556,6 +2576,9 @@ def benchmark_launch(
         if region is not None:
             if any('region' in candidate for candidate in candidates):
                 raise click.BadParameter(f'region {message}')
+        if zone is not None:
+            if any('zone' in candidate for candidate in candidates):
+                raise click.BadParameter(f'zone {message}')
         if gpus is not None:
             if any('accelerators' in candidate for candidate in candidates):
                 raise click.BadParameter(f'gpus (accelerators) {message}')
@@ -2605,6 +2628,7 @@ def benchmark_launch(
         config['num_nodes'] = num_nodes
     override_params = _parse_override_params(cloud=cloud,
                                              region=region,
+                                             zone=zone,
                                              gpus=gpus,
                                              use_spot=use_spot,
                                              image_id=image_id,
@@ -2617,6 +2641,9 @@ def benchmark_launch(
     if 'region' in resources_config:
         if resources_config['region'] is None:
             resources_config.pop('region')
+    if 'zone' in resources_config:
+        if resources_config['zone'] is None:
+            resources_config.pop('zone')
     if 'accelerators' in resources_config:
         if resources_config['accelerators'] is None:
             resources_config.pop('accelerators')

sky/clouds/aws.py

@@ -120,6 +120,16 @@ def get_default_ami(cls, region_name: str, instance_type: str) -> str:
         assert region_name in amis, region_name
         return amis[region_name]
 
+    @classmethod
+    def get_zone_shell_cmd(cls) -> Optional[str]:
+        # The command for getting the current zone is from:
+        # https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/instance-identity-documents.html  # pylint: disable=line-too-long
+        command_str = (
+            'curl -s http://169.254.169.254/latest/dynamic/instance-identity/document'  # pylint: disable=line-too-long
+            ' | python3 -u -c "import sys, json; '
+            'print(json.load(sys.stdin)[\'availabilityZone\'])"')
+        return command_str
+
     #### Normal methods ####
 
     def instance_type_to_hourly_cost(self, instance_type: str, use_spot: bool):
@@ -309,5 +319,13 @@ def get_credential_file_mounts(self) -> Dict[str, str]:
     def instance_type_exists(self, instance_type):
         return service_catalog.instance_type_exists(instance_type, clouds='aws')
 
-    def region_exists(self, region: str) -> bool:
-        return service_catalog.region_exists(region, 'aws')
+    def validate_region_zone(self, region: Optional[str], zone: Optional[str]):
+        return service_catalog.validate_region_zone(region, zone, clouds='aws')
+
+    def accelerator_in_region_or_zone(self,
+                                      accelerator: str,
+                                      acc_count: int,
+                                      region: Optional[str] = None,
+                                      zone: Optional[str] = None) -> bool:
+        return service_catalog.accelerator_in_region_or_zone(
+            accelerator, acc_count, region, zone, 'aws')

sky/clouds/azure.py

@@ -159,6 +159,10 @@ def get_accelerators_from_instance_type(
         return service_catalog.get_accelerators_from_instance_type(
             instance_type, clouds='azure')
 
+    @classmethod
+    def get_zone_shell_cmd(cls) -> Optional[str]:
+        return None
+
     def make_deploy_resources_variables(
             self, resources: 'resources.Resources',
             region: Optional['clouds.Region'],
@@ -281,8 +285,18 @@ def instance_type_exists(self, instance_type):
         return service_catalog.instance_type_exists(instance_type,
                                                     clouds='azure')
 
-    def region_exists(self, region: str) -> bool:
-        return service_catalog.region_exists(region, 'azure')
+    def validate_region_zone(self, region: Optional[str], zone: Optional[str]):
+        return service_catalog.validate_region_zone(region,
+                                                    zone,
+                                                    clouds='azure')
+
+    def accelerator_in_region_or_zone(self,
+                                      accelerator: str,
+                                      acc_count: int,
+                                      region: Optional[str] = None,
+                                      zone: Optional[str] = None) -> bool:
+        return service_catalog.accelerator_in_region_or_zone(
+            accelerator, acc_count, region, zone, 'azure')
 
     @classmethod
     def get_project_id(cls, dryrun: bool = False) -> str: