Update default images for AWS and GCP (skypilot-org#1608)

* Update images for AWS

* remnant

* comment

* refactor fetch_aws a bit

* fix tests

* fix

* wait longer for recovery on AWS

* fix

* fix no conda test

* longer recover time

* fix tests

* fix

* fix azure zones

* format

* update comments

* newline

Author: Michaelvll

examples/huggingface_glue_imdb_app.yaml

@@ -10,11 +10,11 @@ setup: |
   git clone https://github.com/huggingface/transformers/
   # checkout to the correct version
   cd transformers
-  git checkout v4.21.0
+  git checkout v4.25.1
   pip3 install .
   cd examples/pytorch/text-classification
   # SkyPilot's default image on AWS/GCP has CUDA 11.6 (Azure 11.5).
-  pip3 install -r requirements.txt torch==1.12.1+cu113 --extra-index-url https://download.pytorch.org/whl/cu113
+  pip3 install -r requirements.txt tensorboard torch==1.12.1+cu113 --extra-index-url https://download.pytorch.org/whl/cu113
 
 # The command to run.  Will be run under the working directory.
 run: |

examples/per_region_images.yaml

@@ -2,8 +2,8 @@ resources:
   cloud: aws
   instance_type: g4dn.xlarge
   image_id: 
-    us-west-2: ami-0fe5af21074ad2a10  # Deep learning AMI with CUDA 11.6 without conda installed
-    us-west-1: skypilot:gpu-ubuntu-1804
+    us-west-2: skypilot:gpu-ubuntu-1804
+    us-east-2: ami-0406ab83559331633 # AWS Deep Learning AMI GPU CUDA 11.4.3 (Ubuntu 20.04) 20220210 without conda installed
 
 
 setup: |

sky/authentication.py

@@ -245,7 +245,8 @@ def setup_gcp_authentication(config: Dict[str, Any]) -> Dict[str, Any]:
     # OS Login is not enabled for the project. Add the ssh key directly to the
     # metadata.
     # TODO(zhwu): Use cloud init to add ssh public key, to avoid the permission
-    # issue.
+    # issue. A blocker is that the cloud init is not installed in the debian
+    # image by default.
     project_keys = next(
         (item for item in project['commonInstanceMetadata'].get('items', [])
          if item['key'] == 'ssh-keys'), {}).get('value', '')

sky/clouds/service_catalog/data_fetchers/fetch_aws.py

@@ -251,29 +251,36 @@ def get_all_regions_instance_types_df(regions: Set[str]) -> pd.DataFrame:
 
 
 # Fetch Images
-_GPU_TO_IMAGE_DATE = {
-    # https://console.aws.amazon.com/ec2/v2/home?region=us-east-1#Images:visibility=public-images;v=3;search=:64,:Ubuntu%2020,:Deep%20Learning%20AMI%20GPU%20PyTorch # pylint: disable=line-too-long
-    # Current AMIs:
-    # Deep Learning AMI GPU PyTorch 1.10.0 (Ubuntu 20.04) 20220308
-    #   Nvidia driver: 510.47.03, CUDA Version: 11.6 (does not support torch==1.13.0+cu117)
-    #
-    # Use a list to fallback to newer AMI, as some regions like ap-southeast-3 does not have
-    # the older AMI.
-    'gpu': ['20220308', '20221101'],
-    # Deep Learning AMI GPU PyTorch 1.10.0 (Ubuntu 20.04) 20211208
-    # Downgrade the AMI for K80 due as it is only compatible with
-    # NVIDIA driver lower than 470.
-    'k80': ['20211208']
-}
-_UBUNTU_VERSION = ['18.04', '20.04']
-
-
-def _fetch_image_id(region: str, ubuntu_version: str,
-                    creation_date: str) -> Optional[str]:
+# https://console.aws.amazon.com/ec2/v2/home?region=us-east-1#Images:visibility=public-images;v=3;search=:64,:Ubuntu%2020,:Deep%20Learning%20AMI%20GPU%20PyTorch # pylint: disable=line-too-long
+# Current AMIs (we have to use different PyTorch versions for different OS as Ubuntu 18.04
+# does not have the latest PyTorch version):
+# GPU:
+# Deep Learning AMI GPU PyTorch 1.13.1 (Ubuntu 20.04) 20230103
+#   Nvidia driver: 515.65.01, CUDA Version: 11.7
+#
+# Deep Learning AMI GPU PyTorch 1.10.0 (Ubuntu 18.04) 20221114
+#   Nvidia driver: 510.47.03, CUDA Version: 11.6
+#
+# K80:
+# Deep Learning AMI GPU PyTorch 1.10.0 (Ubuntu 20.04) 20211208
+#   Nvidia driver: 470.57.02, CUDA Version: 11.4
+#
+# Deep Learning AMI GPU PyTorch 1.10.0 (Ubuntu 18.04) 20211208
+#   Nvidia driver: 470.57.02, CUDA Version: 11.4
+_GPU_UBUNTU_DATE_PYTORCH = [
+    ('gpu', '20.04', '20230103', '1.13.1'),
+    ('gpu', '18.04', '20221114', '1.10.0'),
+    ('k80', '20.04', '20211208', '1.10.0'),
+    ('k80', '18.04', '20211208', '1.10.0'),
+]
+
+
+def _fetch_image_id(region: str, ubuntu_version: str, creation_date: str,
+                    pytorch_version: str) -> Optional[str]:
     try:
         image = subprocess.check_output(f"""\
             aws ec2 describe-images --region {region} --owners amazon \\
-                --filters 'Name=name,Values="Deep Learning AMI GPU PyTorch 1.10.0 (Ubuntu {ubuntu_version}) {creation_date}"' \\
+                --filters 'Name=name,Values="Deep Learning AMI GPU PyTorch {pytorch_version} (Ubuntu {ubuntu_version}) {creation_date}"' \\
                     'Name=state,Values=available' --query 'Images[:1].ImageId' --output text
             """,
                                         shell=True)
@@ -290,33 +297,25 @@ def _fetch_image_id(region: str, ubuntu_version: str,
 
 @ray.remote
 def _get_image_row(
-        region: str, ubuntu_version: str,
-        cpu_or_gpu: str) -> Tuple[str, str, str, str, Optional[str], str]:
-    print(f'Getting image for {region}, {ubuntu_version}, {cpu_or_gpu}')
-    creation_date = _GPU_TO_IMAGE_DATE[cpu_or_gpu]
-    date = None
-    for date in creation_date:
-        image_id = _fetch_image_id(region, ubuntu_version, date)
-        if image_id:
-            break
-    else:
+        region: str, gpu: str, ubuntu_version: str, date: str,
+        pytorch_version) -> Tuple[str, str, str, str, Optional[str], str]:
+    print(f'Getting image for {region}, {ubuntu_version}, {gpu}')
+    image_id = _fetch_image_id(region, ubuntu_version, date, pytorch_version)
+    if image_id is None:
         # not found
-        print(
-            f'Failed to find image for {region}, {ubuntu_version}, {cpu_or_gpu}'
-        )
-    if date is None:
-        raise ValueError(f'Could not find the creation date for {cpu_or_gpu}.')
-    tag = f'skypilot:{cpu_or_gpu}-ubuntu-{ubuntu_version.replace(".", "")}'
+        print(f'Failed to find image for {region}, {ubuntu_version}, {gpu}')
+    tag = f'skypilot:{gpu}-ubuntu-{ubuntu_version.replace(".", "")}'
     return tag, region, 'ubuntu', ubuntu_version, image_id, date
 
 
 def get_all_regions_images_df(regions: Set[str]) -> pd.DataFrame:
     workers = []
-    for cpu_or_gpu in _GPU_TO_IMAGE_DATE:
-        for ubuntu_version in _UBUNTU_VERSION:
-            for region in regions:
-                workers.append(
-                    _get_image_row.remote(region, ubuntu_version, cpu_or_gpu))
+    for (gpu, ubuntu_version, date,
+         pytorch_version) in _GPU_UBUNTU_DATE_PYTORCH:
+        for region in regions:
+            workers.append(
+                _get_image_row.remote(region, gpu, ubuntu_version, date,
+                                      pytorch_version))
 
     results = ray.get(workers)
     results = pd.DataFrame(

tests/test_smoke.py

@@ -337,21 +337,21 @@ def test_aws_image_id_dict_region():
         'aws_image_id_dict_region',
         [
             # Use region to filter image_id dict.
-            f'sky launch -y -c {name} --region us-east-2 examples/per_region_images.yaml && exit 1 || true',
+            f'sky launch -y -c {name} --region us-east-1 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',
+            f'sky launch -y -c {name} --region us-west-2 examples/per_region_images.yaml',
             # Should success because the image id match for the region.
             f'sky launch -c {name} --image-id skypilot:gpu-ubuntu-1804 examples/minimal.yaml',
             f'sky exec {name} --image-id skypilot:gpu-ubuntu-1804 examples/minimal.yaml',
             f'sky exec {name} --image-id skypilot:gpu-ubuntu-2004 examples/minimal.yaml && exit 1 || true',
             f'sky logs {name} 1 --status',
             f'sky logs {name} 2 --status',
             f'sky logs {name} 3 --status',
-            f'sky status --all | grep {name} | grep us-west-1',  # Ensure the region is correct.
+            f'sky status --all | grep {name} | grep us-west-2',  # 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} --region us-west-2 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} --cloud aws --region us-west-2 "ls ~"',
             f'sky exec {name} "ls ~"',
             f'sky logs {name} 4 --status',
             f'sky logs {name} 5 --status',
@@ -403,9 +403,9 @@ def test_aws_image_id_dict_zone():
         'aws_image_id_dict_zone',
         [
             # Use zone to filter image_id dict.
-            f'sky launch -y -c {name} --zone us-east-2b examples/per_region_images.yaml && exit 1 || true',
+            f'sky launch -y -c {name} --zone us-east-1b 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',
+            f'sky launch -y -c {name} --zone us-west-2a examples/per_region_images.yaml',
             # Should success because the image id match for the zone.
             f'sky launch -y -c {name} --image-id skypilot:gpu-ubuntu-1804 examples/minimal.yaml',
             f'sky exec {name} --image-id skypilot:gpu-ubuntu-1804 examples/minimal.yaml',
@@ -414,11 +414,11 @@ def test_aws_image_id_dict_zone():
             f'sky logs {name} 1 --status',
             f'sky logs {name} 2 --status',
             f'sky logs {name} 3 --status',
-            f'sky status --all | grep {name} | grep us-west-1a',  # Ensure the zone is correct.
+            f'sky status --all | grep {name} | grep us-west-2a',  # 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} --zone us-west-2a 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} --cloud aws --region us-west-2 "ls ~"',
             f'sky exec {name} "ls ~"',
             f'sky logs {name} 4 --status',
             f'sky logs {name} 5 --status',
@@ -471,7 +471,7 @@ def test_image_no_conda():
         'image_no_conda',
         [
             # Use image id dict.
-            f'sky launch -y -c {name} --region us-west-2 examples/per_region_images.yaml',
+            f'sky launch -y -c {name} --region us-east-2 examples/per_region_images.yaml',
             f'sky logs {name} 1 --status',
             f'sky stop {name} -y',
             f'sky start {name} -y',
@@ -606,7 +606,7 @@ def test_aws_storage_mounts():
         file_path = f.name
         test_commands = [
             *storage_setup_commands,
-            f'sky launch -y -c {name}-aws --cloud aws {file_path}',
+            f'sky launch -y -c {name} --cloud aws {file_path}',
             f'sky logs {name} 1 --status',  # Ensure job succeeded.
             f'aws s3 ls {storage_name}/hello.txt',
         ]
@@ -1246,7 +1246,7 @@ def test_spot_recovery_multi_node_aws():
              '--output text)'),
             'sleep 50',
             f'{_SPOT_QUEUE_WAIT}| grep {name} | head -n1 | grep "RECOVERING"',
-            'sleep 420',
+            'sleep 500',
             f'{_SPOT_QUEUE_WAIT}| grep {name} | head -n1 | grep "RUNNING"',
             f'RUN_ID=$(cat /tmp/{name}-run-id); echo $RUN_ID; sky spot logs -n {name} --no-follow | grep SKYPILOT_JOB_ID | cut -d: -f2 | grep "$RUN_ID"',
         ],