Enable custom image for AWS and GCP (skypilot-org#931)

* add image id

* format

* Add schema

* fix resources equality

* add gcp image id

* retry ray up

* Install pip

* Remove output of pip installation

* Fix conda

* Add docs

* fix typo

* Fix aws worker image

* Add tests

* Fix azure and add customized image example

* format

* format

* Address comments

* fix typo

* Fix output

* fix ux_utils

* revert minimal

* fix template

* Address comments

* Reverted to python3

* revert parens

* fix conda env

Author: Michaelvll

docs/source/reference/yaml-spec.rst

@@ -79,6 +79,16 @@ describe all fields available.
         tpu_name: mytpu
         tpu_vm: False  # False to use TPU nodes (the default); True to use TPU VMs.
 
+      # Custom image id (optional, advanced). The image id used to boot the
+      # instances. Only supported for AWS and GCP. If not specified, sky will use
+      # the default debian-based image suitable for machine learning tasks.
+      # To find AWS AMI ids: https://leaherb.com/how-to-find-an-aws-marketplace-ami-image-id
+      # AWS
+      image_id: ami-0868a20f5a3bf9702
+      # To find GCP images: https://cloud.google.com/compute/docs/images
+      # GCP
+      # image_id: projects/deeplearning-platform-release/global/images/family/tf2-ent-2-1-cpu-ubuntu-2004
+
     file_mounts:
       # Uses rsync to copy local files to all nodes of the cluster.
       #

examples/custom_image.yaml

@@ -0,0 +1,13 @@
+resources:
+  cloud: aws
+  region: us-east-2
+  # Nvidia image from
+  # https://aws.amazon.com/marketplace/pp/prodview-rf7na2b2ttvdg
+  image_id: ami-062ddd90fb6f8267a
+  accelerators: V100
+
+setup: |
+  echo "running setup"
+
+run: |
+  echo "hello sky"

sky/backends/backend_utils.py

@@ -558,32 +558,10 @@ def write_cluster_config(to_provision: 'resources.Resources',
     # task.best_resources may not be equal to to_provision if the user
     # is running a job with less resources than the cluster has.
     cloud = to_provision.cloud
-    resources_vars = cloud.make_deploy_resources_variables(to_provision)
+    resources_vars = cloud.make_deploy_resources_variables(
+        to_provision, region, zones)
     config_dict = {}
 
-    if region is None:
-        assert zones is None, 'Set either both or neither for: region, zones.'
-        region = cloud.get_default_region()
-        zones = region.zones
-    else:
-        assert isinstance(
-            cloud, clouds.Azure
-        ) or zones is not None, 'Set either both or neither for: region, zones.'
-    region = region.name
-    if isinstance(cloud, clouds.AWS):
-        # Only AWS supports multiple zones in the 'availability_zone' field.
-        zones = [zone.name for zone in zones]
-    elif isinstance(cloud, clouds.Azure):
-        # Azure does not support specific zones.
-        zones = []
-    else:
-        zones = [zones[0].name]
-
-    aws_default_ami = None
-    if isinstance(cloud, clouds.AWS):
-        instance_type = resources_vars['instance_type']
-        aws_default_ami = cloud.get_default_ami(region, instance_type)
-
     azure_subscription_id = None
     if isinstance(cloud, clouds.Azure):
         azure_subscription_id = cloud.get_project_id(dryrun=dryrun)
@@ -594,6 +572,7 @@ def write_cluster_config(to_provision: 'resources.Resources',
 
     assert cluster_name is not None
     credentials = sky_check.get_cloud_credential_file_mounts()
+    region_name = resources_vars['region']
     yaml_path = fill_template(
         cluster_config_template,
         dict(
@@ -602,11 +581,6 @@ def write_cluster_config(to_provision: 'resources.Resources',
                 'cluster_name': cluster_name,
                 'num_nodes': num_nodes,
                 'disk_size': to_provision.disk_size,
-                # Region/zones.
-                'region': region,
-                'zones': ','.join(zones),
-                # AWS only.
-                'aws_default_ami': aws_default_ami,
                 # Temporary measure, as deleting per-cluster SGs is too slow.
                 # See https://github.com/sky-proj/sky/pull/742.
                 # Generate the name of the security group we're looking for.
@@ -616,7 +590,7 @@ def write_cluster_config(to_provision: 'resources.Resources',
                 'security_group': f'sky-sg-{user_and_hostname_hash()}',
                 # Azure only.
                 'azure_subscription_id': azure_subscription_id,
-                'resource_group': f'{cluster_name}-{region}',
+                'resource_group': f'{cluster_name}-{region_name}',
                 # GCP only.
                 'gcp_project_id': gcp_project_id,
                 # Ray version.
@@ -645,7 +619,6 @@ def write_cluster_config(to_provision: 'resources.Resources',
                 template_name,
                 dict(
                     resources_vars, **{
-                        'zones': ','.join(zones),
                         'tpu_name': tpu_name,
                         'gcp_project_id': gcp_project_id,
                     }),

sky/backends/cloud_vm_ray_backend.py

@@ -1019,6 +1019,16 @@ def ray_up():
                     f'{region_name}{colorama.Style.RESET_ALL} ({zone_str})')
         start = time.time()
         returncode, stdout, stderr = ray_up()
+        if (returncode != 0 and 'Processing file mounts' in stdout and
+                'Running setup commands' not in stdout):
+            # Retry ray up if it failed due to file mounts, because it is
+            # probably due to too many ssh connections issue and can be fixed
+            # by retrying.
+            # This is required when using custom image for GCP.
+            logger.info(
+                'Retrying sky runtime setup due to ssh connection issue.')
+            returncode, stdout, stderr = ray_up()
+
         logger.debug(f'Ray up takes {time.time() - start} seconds.')
 
         # Only 1 node or head node provisioning failure.

sky/cli.py

@@ -261,6 +261,11 @@ def _parse_env_var(env_var: str) -> Tuple[str, str]:
         default=None,
         help=('Whether to request spot instances. If specified, overrides the '
               '"resources.use_spot" config.')),
+    click.option('--image-id',
+                 required=False,
+                 default=None,
+                 help=('Custom image id for launching the instances. '
+                       'Passing "none" resets the config.')),
     click.option(
         '--env',
         required=False,
@@ -552,6 +557,7 @@ def _make_dag_from_entrypoint_with_overrides(
     gpus: Optional[str] = None,
     num_nodes: Optional[int] = None,
     use_spot: Optional[bool] = None,
+    image_id: Optional[str] = None,
     disk_size: Optional[int] = None,
     env: List[Dict[str, str]] = None,
     # spot launch specific
@@ -563,16 +569,17 @@ def _make_dag_from_entrypoint_with_overrides(
         if _check_yaml(entrypoint):
             # Treat entrypoint as a yaml.
             click.secho('Task from YAML spec: ', fg='yellow', nl=False)
+            click.secho(entrypoint, bold=True)
             task = sky.Task.from_yaml(entrypoint)
         else:
             if not entrypoint:
                 entrypoint = None
             else:
                 # Treat entrypoint as a bash command.
                 click.secho('Task from command: ', fg='yellow', nl=False)
+                click.secho(entrypoint, bold=True)
             task = sky.Task(name='sky-cmd', run=entrypoint)
             task.set_resources({sky.Resources()})
-        click.secho(entrypoint, bold=True)
         # Override.
         if workdir is not None:
             task.workdir = workdir
@@ -597,6 +604,11 @@ def _make_dag_from_entrypoint_with_overrides(
             override_params['use_spot'] = use_spot
         if disk_size is not None:
             override_params['disk_size'] = disk_size
+        if image_id is not None:
+            if image_id.lower() == 'none':
+                override_params['image_id'] = None
+            else:
+                override_params['image_id'] = image_id
 
         # Spot launch specific.
         if spot_recovery is not None:
@@ -735,6 +747,7 @@ def launch(
     gpus: Optional[str],
     num_nodes: Optional[int],
     use_spot: Optional[bool],
+    image_id: Optional[str],
     env: List[Dict[str, str]],
     disk_size: Optional[int],
     idle_minutes_to_autostop: Optional[int],
@@ -763,6 +776,7 @@ def launch(
         gpus=gpus,
         num_nodes=num_nodes,
         use_spot=use_spot,
+        image_id=image_id,
         env=env,
         disk_size=disk_size,
     )
@@ -808,6 +822,7 @@ def exec(
     gpus: Optional[str],
     num_nodes: Optional[int],
     use_spot: Optional[bool],
+    image_id: Optional[str],
     env: List[Dict[str, str]],
 ):
     """Execute a task or a command on a cluster (skip setup).
@@ -882,6 +897,7 @@ def exec(
         region=region,
         gpus=gpus,
         use_spot=use_spot,
+        image_id=image_id,
         num_nodes=num_nodes,
         env=env,
     )
@@ -2064,6 +2080,7 @@ def spot_launch(
     gpus: Optional[str],
     num_nodes: Optional[int],
     use_spot: Optional[bool],
+    image_id: Optional[str],
     spot_recovery: Optional[str],
     env: List[Dict[str, str]],
     disk_size: Optional[int],
@@ -2085,6 +2102,7 @@ def spot_launch(
         gpus=gpus,
         num_nodes=num_nodes,
         use_spot=use_spot,
+        image_id=image_id,
         env=env,
         disk_size=disk_size,
         spot_recovery=spot_recovery,

sky/clouds/aws.py

@@ -178,19 +178,42 @@ def get_accelerators_from_instance_type(
         return service_catalog.get_accelerators_from_instance_type(
             instance_type, clouds='aws')
 
-    def make_deploy_resources_variables(self,
-                                        resources: 'resources_lib.Resources'):
+    def make_deploy_resources_variables(
+            self, resources: 'resources_lib.Resources',
+            region: Optional['clouds.Region'],
+            zones: Optional[List['clouds.Zone']]) -> Dict[str, str]:
+        if region is None:
+            assert zones is None, (
+                'Set either both or neither for: region, zones.')
+            region = self._get_default_region()
+            zones = region.zones
+        else:
+            assert zones is not None, (
+                'Set either both or neither for: region, zones.')
+
+        region_name = region.name
+        zones = [zone.name for zone in zones]
+
         r = resources
         # r.accelerators is cleared but .instance_type encodes the info.
         acc_dict = self.get_accelerators_from_instance_type(r.instance_type)
         if acc_dict is not None:
             custom_resources = json.dumps(acc_dict, separators=(',', ':'))
         else:
             custom_resources = None
+
+        if r.image_id is not None:
+            image_id = r.image_id
+        else:
+            image_id = self.get_default_ami(region_name, r.instance_type)
+
         return {
             'instance_type': r.instance_type,
             'custom_resources': custom_resources,
             'use_spot': r.use_spot,
+            'region': region_name,
+            'zones': ','.join(zones),
+            'image_id': image_id,
         }
 
     def get_feasible_launchable_resources(self,

sky/clouds/azure.py

@@ -2,12 +2,16 @@
 import json
 import os
 import subprocess
+import typing
 from typing import Dict, Iterator, List, Optional, Tuple
 
 from sky import clouds
 from sky.adaptors import azure
 from sky.clouds import service_catalog
 
+if typing.TYPE_CHECKING:
+    from sky import resources
+
 # Minimum set of files under ~/.azure that grant Azure access.
 _CREDENTIAL_FILES = [
     'azureProfile.json',
@@ -157,7 +161,19 @@ def get_accelerators_from_instance_type(
         return service_catalog.get_accelerators_from_instance_type(
             instance_type, clouds='azure')
 
-    def make_deploy_resources_variables(self, resources):
+    def make_deploy_resources_variables(
+            self, resources: 'resources.Resources',
+            region: Optional['clouds.Region'],
+            zones: Optional[List['clouds.Zone']]) -> Dict[str, str]:
+        if region is None:
+            assert zones is None, (
+                'Set either both or neither for: region, zones.')
+            region = self._get_default_region()
+
+        region_name = region.name
+        # Azure does not support specific zones.
+        zones = []
+
         r = resources
         assert not r.use_spot, \
             'Our subscription offer ID does not support spot instances.'
@@ -175,6 +191,8 @@ def make_deploy_resources_variables(self, resources):
             'instance_type': r.instance_type,
             'custom_resources': custom_resources,
             'use_spot': r.use_spot,
+            'region': region_name,
+            'zones': zones,
             **image_config
         }
 

sky/clouds/cloud.py

@@ -1,7 +1,11 @@
 """Interfaces: clouds, regions, and zones."""
 import collections
+import typing
 from typing import Dict, Iterator, List, Optional, Tuple
 
+if typing.TYPE_CHECKING:
+    from sky import resources
+
 
 class Region(collections.namedtuple('Region', ['name'])):
     """A region."""
@@ -107,7 +111,12 @@ def get_egress_cost(self, num_gigabytes):
     def is_same_cloud(self, other):
         raise NotImplementedError
 
-    def make_deploy_resources_variables(self, resources):
+    def make_deploy_resources_variables(
+        self,
+        resources: 'resources.Resources',
+        region: Optional['Region'],
+        zones: Optional[List['Zone']],
+    ) -> Dict[str, str]:
         """Converts planned sky.Resources to cloud-specific resource variables.
 
         These variables are used to fill the node type section (instance type,
@@ -136,7 +145,7 @@ def get_default_instance_type(cls,
         raise NotImplementedError
 
     @classmethod
-    def get_default_region(cls) -> Region:
+    def _get_default_region(cls) -> Region:
         raise NotImplementedError
 
     def get_feasible_launchable_resources(self, resources):