[Test/Azure] Fix the torch version in examples for smoke test and change the credential for Azure (skypilot-org#1330)

* Upgrade images for three clouds

* Fix cuda version

* pin cuda version for torch

* Fix torch version

* fix comments

* Fix azure provider

* fix credential

* revert back to previous azure image

* switch back to cuda 11.3 for pytorch due to azure's image

* fix torch installation

* increase the multi-node timeout

* Update sky/clouds/azure.py

Co-authored-by: Zongheng Yang <[email protected]>

* revert aws image version

* pin cu113 for huggingface

* Add comment

* format

* Update sky/clouds/aws.py

Co-authored-by: Zongheng Yang <[email protected]>

* Update sky/clouds/gcp.py

Co-authored-by: Zongheng Yang <[email protected]>

* revert gcp image

* Fix doc

Co-authored-by: Zongheng Yang <[email protected]>

Author: Michaelvll

docs/source/examples/distributed-jobs.rst

@@ -20,8 +20,9 @@ For example, here is a simple PyTorch Distributed training example:
   setup: |
     pip3 install --upgrade pip
     git clone https://github.com/michaelzhiluo/pytorch-distributed-resnet
+    cd pytorch-distributed-resnet
     # SkyPilot's default image on AWS/GCP has CUDA 11.6 (Azure 11.5).
-    cd pytorch-distributed-resnet && pip3 install -r requirements.txt torch==1.12.1+cu113 --extra-index-url https://download.pytorch.org/whl/cu113
+    pip3 install -r requirements.txt torch==1.12.1+cu113 --extra-index-url https://download.pytorch.org/whl/cu113
     mkdir -p data  && mkdir -p saved_models && cd data && \
       wget -c --quiet https://www.cs.toronto.edu/~kriz/cifar-10-python.tar.gz
     tar -xvzf cifar-10-python.tar.gz

examples/huggingface_glue_imdb_app.yaml

@@ -13,7 +13,8 @@ setup: |
   git checkout v4.21.0
   pip3 install .
   cd examples/pytorch/text-classification
-  pip3 install -r requirements.txt
+  # 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
 
 # The command to run.  Will be run under the working directory.
 run: |

examples/resnet_distributed_torch.yaml

@@ -9,7 +9,9 @@ num_nodes: 2
 setup: |
     pip3 install --upgrade pip
     git clone https://github.com/michaelzhiluo/pytorch-distributed-resnet
-    cd pytorch-distributed-resnet && pip3 install -r requirements.txt torch==1.12.1+cu113 --extra-index-url https://download.pytorch.org/whl/cu113
+    cd pytorch-distributed-resnet
+    # 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
     mkdir -p data  && mkdir -p saved_models && cd data && \
     wget -c --quiet https://www.cs.toronto.edu/~kriz/cifar-10-python.tar.gz
     tar -xvzf cifar-10-python.tar.gz

examples/resnet_distributed_torch_app.py

@@ -9,7 +9,7 @@
 setup = 'echo \"alias python=python3\" >> ~/.bashrc && pip3 install --upgrade pip && \
     [ -d pytorch-distributed-resnet ] || \
     (git clone https://github.com/michaelzhiluo/pytorch-distributed-resnet && \
-    cd pytorch-distributed-resnet && pip3 install -r requirements.txt && \
+    cd pytorch-distributed-resnet && pip3 install -r requirements.txt torch==1.12.1+cu113 --extra-index-url https://download.pytorch.org/whl/cu113 && \
     mkdir -p data  && mkdir -p saved_models && cd data && \
     wget -c --quiet https://www.cs.toronto.edu/~kriz/cifar-10-python.tar.gz && \
     tar -xvzf cifar-10-python.tar.gz)'

examples/resnet_distributed_torch_scripts/setup.sh

@@ -7,10 +7,11 @@ if [ $? -eq 0 ]; then
     echo "conda env exists"
 else
     echo "conda env does not exist"
-    conda create -n resnet python=3.6 -y
+    conda create -n resnet python=3.7 -y
     conda activate resnet
 fi
-pip install -r requirements.txt
+# SkyPilot's default image on AWS/GCP has CUDA 11.6 (Azure 11.5).
+pip install -r requirements.txt torch==1.12.1+cu113 --extra-index-url https://download.pytorch.org/whl/cu113
 
 mkdir -p data
 mkdir -p saved_models

sky/clouds/aws.py

@@ -110,9 +110,20 @@ def get_default_ami(cls, region_name: str, instance_type: str) -> str:
                 }
                 assert region_name in amis, region_name
                 return amis[region_name]
-        # Deep Learning AMI GPU PyTorch 1.10.0 (Ubuntu 20.04) 20220308
         # 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
-        # Nvidia driver: 510.47.03, CUDA Version: 11.6
+
+        # Commented below are newer AMIs, but as other clouds do not support
+        # torch==1.13.0+cu117 we do not use these AMIs to avoid frequent updates:
+        # Deep Learning AMI GPU PyTorch 1.12.1 (Ubuntu 20.04) 20221025
+        #   Nvidia driver: 510.47.03, CUDA Version: 11.6 (supports torch==1.13.0+cu117)
+        #     'us-east-1': 'ami-0eb1f91977a3fcc1b'
+        #     'us-east-2': 'ami-0274a6db2e19b7cc6'
+        #     'us-west-1': 'ami-0fb299af41d32cfd3'
+        #     'us-west-2': 'ami-04ba15f9bd464eb20'
+        #
+        # 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)
         amis = {
             'us-east-1': 'ami-0729d913a335efca7',
             'us-east-2': 'ami-070f4af81c19b41bf',

sky/clouds/azure.py

@@ -82,6 +82,13 @@ def get_default_instance_type(cls) -> str:
         return 'Standard_D8_v4'
 
     def _get_image_config(self, gen_version, instance_type):
+        # az vm image list \
+        #  --publisher microsoft-dsvm --all --output table
+        # nvidia-driver: 495.29.05, cuda: 11.5
+
+        # The latest image 2022.09.14/2022.08.11/22.06.10/22.05.11/
+        # 22.04.27/22.04.05 has even older nvidia driver 470.57.02,
+        # cuda: 11.4
         image_config = {
             'image_publisher': 'microsoft-dsvm',
             'image_offer': 'ubuntu-2004',

sky/clouds/gcp.py

@@ -214,6 +214,9 @@ def make_deploy_resources_variables(
         region_name = region.name
         zones = [zones[0].name]
 
+        # gcloud compute images list \
+        # --project deeplearning-platform-release \
+        # --no-standard-images
         image_id = _IMAGE_ID_PREFIX + 'common-cpu-v20220806'
 
         r = resources
@@ -255,10 +258,15 @@ def make_deploy_resources_variables(
                         acc.lower())
                 resources_vars['gpu_count'] = acc_count
                 if acc == 'K80':
+                    # Though the image is called cu113, it actually has later
+                    # versions of CUDA as noted below.
                     # CUDA driver version 470.57.02, CUDA Library 11.4
                     image_id = _IMAGE_ID_PREFIX + 'common-cu113-v20220701'
                 else:
+                    # Though the image is called cu113, it actually has later
+                    # versions of CUDA as noted below.
                     # CUDA driver version 510.47.03, CUDA Library 11.6
+                    # Does not support torch==1.13.0 with cu117
                     image_id = _IMAGE_ID_PREFIX + 'common-cu113-v20220806'
 
         if resources.image_id is not None:

sky/skylet/providers/azure/node_provider.py

@@ -4,7 +4,7 @@
 from threading import RLock
 from uuid import uuid4
 
-from azure.identity import DefaultAzureCredential
+from azure.identity import AzureCliCredential
 from azure.mgmt.compute import ComputeManagementClient
 from azure.mgmt.network import NetworkManagementClient
 from azure.mgmt.resource import ResourceManagementClient
@@ -65,12 +65,8 @@ def __init__(self, provider_config, cluster_name):
         _configure_resource_group({"provider": provider_config})
         subscription_id = provider_config["subscription_id"]
         self.cache_stopped_nodes = provider_config.get("cache_stopped_nodes", True)
-        # AWS provides managed identity for Azure, but it is not setup properly by
-        # default. This interferes with azure-cli credentials and causes failures,
-        # when using sky to launch Azure on AWS ec2 instances. We disable it to give
-        # way to azure-cli credentials.
-        credential = DefaultAzureCredential(exclude_shared_token_cache_credential=True,
-                                            exclude_managed_identity_credential=True)
+        # Sky only supports Azure CLI credential for now.
+        credential = AzureCliCredential()
         self.compute_client = ComputeManagementClient(credential, subscription_id)
         self.network_client = NetworkManagementClient(credential, subscription_id)
         self.resource_client = ResourceManagementClient(credential, subscription_id)