Repository: skypilot-org/skypilot
Pull request: [Test/Azure] Fix the torch version in examples for smoke test and change the credential for Azure 
Closes #1327 

PyTorch bumps up the version to `1.13.0` yesterday, which causes the cuda version mismatch problem, when directly installing the latest `PyTorch==1.13.0` `pip install torch`.

The backward compatibility should be guaranteed by our `_RAY_YAML_KEYS_TO_RESTORE_FOR_BACK_COMPATIBILITY` in `backend_utils.py`, as the ami is set in the `node_config` section.
https://github.com/skypilot-org/skypilot/blob/855d4b46800d70597808d4c773e1f53e146bdbf8/sky/backends/backend_utils.py#L126

Update: Have to revert the aws/gcp/azure image, since failed `test_cancel_azure` (The latest `22.09.14` azure image has an older nvidia driver `470.141.03` cuda version `11.3` than our previous azure image `21.11.04` with nvidia-driver `495.29.05`, cuda `11.5`).

Future TODO:
- [ ] Ask the user to use older version of the pytorch to be able to run with CUDA lower than 11.6 #1333 

Tested:
- [x] `./tests/run_smoke_tests.sh`

Problem statement:
[Smoke tests] Cuda error for training
Our smoke tests with `Pytorch` fails (`test_cancel_pytorch` and `test_huggingface`, due to the following cuda issue:
```
(node-1 pid=31582) Traceback (most recent call last):
(node-1 pid=31582)   File "/opt/conda/lib/python3.7/runpy.py", line 183, in _run_module_as_main
(node-1 pid=31582)     mod_name, mod_spec, code = _get_module_details(mod_name, _Error)
(node-1 pid=31582)   File "/opt/conda/lib/python3.7/runpy.py", line 109, in _get_module_details
(node-1 pid=31582)     __import__(pkg_name)
(node-1 pid=31582)   File "/opt/conda/lib/python3.7/site-packages/torch/__init__.py", line 191, in <module>
(node-1 pid=31582)     _load_global_deps()
(node-1 pid=31582)   File "/opt/conda/lib/python3.7/site-packages/torch/__init__.py", line 153, in _load_global_deps
(node-1 pid=31582)     ctypes.CDLL(lib_path, mode=ctypes.RTLD_GLOBAL)
(node-1 pid=31582)   File "/opt/conda/lib/python3.7/ctypes/__init__.py", line 364, in __init__
(node-1 pid=31582)     self._handle = _dlopen(self._name, mode)
(node-1 pid=31582) OSError: /opt/conda/lib/python3.7/site-packages/torch/lib/../../nvidia/cublas/lib/libcublas.so.11: symbol cublasLtGetStatusString version libcublasLt.so.11 not defined in file libcublasLt.so.11 with link time reference
(node-0 pid=21417, ip=10.128.15.237) Traceback (most recent call last):
(node-0 pid=21417, ip=10.128.15.237)   File "/opt/conda/lib/python3.7/runpy.py", line 183, in _run_module_as_main
(node-0 pid=21417, ip=10.128.15.237)     mod_name, mod_spec, code = _get_module_details(mod_name, _Error)
(node-0 pid=21417, ip=10.128.15.237)   File "/opt/conda/lib/python3.7/runpy.py", line 109, in _get_module_details
(node-0 pid=21417, ip=10.128.15.237)     __import__(pkg_name)
(node-0 pid=21417, ip=10.128.15.237)   File "/opt/conda/lib/python3.7/site-packages/torch/__init__.py", line 191, in <module>
(node-0 pid=21417, ip=10.128.15.237)     _load_global_deps()
(node-0 pid=21417, ip=10.128.15.237)   File "/opt/conda/lib/python3.7/site-packages/torch/__init__.py", line 153, in _load_global_deps
(node-0 pid=21417, ip=10.128.15.237)     ctypes.CDLL(lib_path, mode=ctypes.RTLD_GLOBAL)
(node-0 pid=21417, ip=10.128.15.237)   File "/opt/conda/lib/python3.7/ctypes/__init__.py", line 364, in __init__
(node-0 pid=21417, ip=10.128.15.237)     self._handle = _dlopen(self._name, mode)
(node-0 pid=21417, ip=10.128.15.237) OSError: /opt/conda/lib/python3.7/site-packages/torch/lib/../../nvidia/cublas/lib/libcublas.so.11: symbol cublasLtGetStatusString version libcublasLt.so.11 not defined in file libcublasLt.so.11 with link time referenc
e
```
```
(huggingface pid=15893) Traceback (most recent call last):
(huggingface pid=15893)   File "run_glue.py", line 31, in <module>
(huggingface pid=15893)     from transformers import (
(huggingface pid=15893)   File "<frozen importlib._bootstrap>", line 1032, in _handle_fromlist
(huggingface pid=15893)   File "/opt/conda/lib/python3.7/site-packages/transformers/utils/import_utils.py", line 992, in __getattr__
(huggingface pid=15893)     module = self._get_module(self._class_to_module[name])
(huggingface pid=15893)   File "/opt/conda/lib/python3.7/site-packages/transformers/utils/import_utils.py", line 1007, in _get_module
(huggingface pid=15893)     ) from e
(huggingface pid=15893) RuntimeError: Failed to import transformers.data.data_collator because of the following error (look up to see its traceback):
(huggingface pid=15893) /opt/conda/lib/python3.7/site-packages/torch/lib/../../nvidia/cublas/lib/libcublas.so.11: symbol cublasLtGetStatusString version libcublasLt.so.11 not defined in file libcublasLt.so.11 with link time reference
```


On the cluster, importing the `pytorch` will fail:
```
gcpuser@ray-test-huggingface-3446f302-84-head-d164904e-compute:~$ python
Python 3.7.12 | packaged by conda-forge | (default, Oct 26 2021, 06:08:53)
[GCC 9.4.0] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import torch
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/opt/conda/lib/python3.7/site-packages/torch/__init__.py", line 191, in <module>
    _load_global_deps()
  File "/opt/conda/lib/python3.7/site-packages/torch/__init__.py", line 153, in _load_global_deps
    ctypes.CDLL(lib_path, mode=ctypes.RTLD_GLOBAL)
  File "/opt/conda/lib/python3.7/ctypes/__init__.py", line 364, in __init__
    self._handle = _dlopen(self._name, mode)
OSError: /opt/conda/lib/python3.7/site-packages/torch/lib/../../nvidia/cublas/lib/libcublas.so.11: symbol cublasLtGetStatusString version libcublasLt.so.11 not defined in file libcublasLt.so.11 with link time reference
```


[Azure] Launching the azure cluster can cause nondeterministic error
`sky launch -y -c launch-azure examples/azure_start_stop.yaml` sometimes getting the following issue:
```
Exception in thread Thread-1:
Traceback (most recent call last):
  File "/home/ubuntu/.conda/envs/sky/lib/python3.10/site-packages/ray/autoscaler/_private/updater.py", line 153, in run
    self.do_update()
  File "/home/ubuntu/.conda/envs/sky/lib/python3.10/site-packages/ray/autoscaler/_private/updater.py", line 327, in do_update
    self.wait_ready(deadline)
  File "/home/ubuntu/.conda/envs/sky/lib/python3.10/site-packages/ray/autoscaler/_private/updater.py", line 271, in wait_ready
    if self.provider.is_terminated(self.node_id):
  File "/home/ubuntu/skypilot/sky/skylet/providers/azure/node_provider.py", line 166, in is_terminated
    node = self._get_node(node_id=node_id)
  File "/home/ubuntu/skypilot/sky/skylet/providers/azure/node_provider.py", line 361, in _get_node
    return self.cached_nodes[node_id]
KeyError: 'ray-test-azure-start-stop-3446f302-c9-head-492094d10'

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/home/ubuntu/.conda/envs/sky/lib/python3.10/threading.py", line 1016, in _bootstrap_inner
    self.run()
  File "/home/ubuntu/.conda/envs/sky/lib/python3.10/site-packages/ray/autoscaler/_private/updater.py", line 155, in run
    self.provider.set_node_tags(
  File "/home/ubuntu/skypilot/sky/skylet/providers/azure/node_provider.py", line 38, in wrapper
    return f(self, *args, **kwargs)
  File "/home/ubuntu/skypilot/sky/skylet/providers/azure/node_provider.py", line 266, in set_node_tags
    node_tags = self._get_cached_node(node_id)["tags"]
  File "/home/ubuntu/skypilot/sky/skylet/providers/azure/node_provider.py", line 366, in _get_cached_node
    return self._get_node(node_id=node_id)
  File "/home/ubuntu/skypilot/sky/skylet/providers/azure/node_provider.py", line 361, in _get_node
    return self.cached_nodes[node_id]
KeyError: 'ray-test-azure-start-stop-3446f302-c9-head-492094d10'
2022-10-30 04:21:36,033 PANIC commands.py:813 -- Failed to setup head node.
Error: Failed to setup head node.
```

@suquark Do you think this is a bug in our `node_provider`?

I tried to get the `resources_group` with the code in `self.provider.is_terminated(self.node_id)` above to get an existing `resource_group` but it fails. However, `az vm list --resource-group test-azure-start-stop-3446f302-c9-eastus` can get the correct cluster. Do you have any idea of the issue @suquark ?
```
>>> with open('/home/ubuntu//.sky/generated/test-azure-start-stop-3446f302-c9.yml', 'r') as f:
...     config = yaml.safe_load(f.read())
...
>>> from sky.skylet.providers.azure.node_provider import AzureNodeProvider
>>> provider = AzureNodeProvider(config['provider'], config['cluster_name'])
>>> vms = provider.compute_client.virtual_machines.list(resource_group_name=provider.provider_config['resource_group'])
>>> list(vms)
[]
```

Hints from issue discussion:
Just identified the issue, I was launching the azure cluster from an AWS ec2 VM, where there are some empty Azure credential services set up by default. We can limit the credential used by the `node_provider` to be only `AzureCliCredential` instead of a chain of different kind of credentials (fixed in #1330)

Update: Seems changing the credential can only partially fix the `KeyError` issue mentioned above. The problem still exists.

Commit message: increase the multi-node timeout

Patch to review:
diff --git a/examples/huggingface_glue_imdb_app.yaml b/examples/huggingface_glue_imdb_app.yaml
index ba114fcd0b8..8eb45dda016 100644
--- a/examples/huggingface_glue_imdb_app.yaml
+++ b/examples/huggingface_glue_imdb_app.yaml
@@ -13,7 +13,7 @@ setup: |
   git checkout v4.21.0
   pip3 install .
   cd examples/pytorch/text-classification
-  pip3 install -r requirements.txt
+  pip3 install -r requirements.txt --extra-index-url https://download.pytorch.org/whl/cu116
 
 # The command to run.  Will be run under the working directory.
 run: |
diff --git a/examples/resnet_distributed_torch.yaml b/examples/resnet_distributed_torch.yaml
index 8df3810c25e..69f9f384242 100644
--- a/examples/resnet_distributed_torch.yaml
+++ b/examples/resnet_distributed_torch.yaml
@@ -9,7 +9,7 @@ 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
+    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
diff --git a/examples/resnet_distributed_torch_app.py b/examples/resnet_distributed_torch_app.py
index 25a92228437..1bc38886536 100644
--- a/examples/resnet_distributed_torch_app.py
+++ b/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)'
diff --git a/examples/resnet_distributed_torch_scripts/setup.sh b/examples/resnet_distributed_torch_scripts/setup.sh
index a87a0194041..3b74f82f701 100644
--- a/examples/resnet_distributed_torch_scripts/setup.sh
+++ b/examples/resnet_distributed_torch_scripts/setup.sh
@@ -7,10 +7,10 @@ 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
+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
diff --git a/sky/backends/cloud_vm_ray_backend.py b/sky/backends/cloud_vm_ray_backend.py
index c52375cb6e6..0b66b956d47 100644
--- a/sky/backends/cloud_vm_ray_backend.py
+++ b/sky/backends/cloud_vm_ray_backend.py
@@ -62,7 +62,7 @@
 
 # Timeout (seconds) for provision progress: if in this duration no new nodes
 # are launched, abort and failover.
-_NODES_LAUNCHING_PROGRESS_TIMEOUT = 60
+_NODES_LAUNCHING_PROGRESS_TIMEOUT = 90
 
 # Time gap between retries after failing to provision in all possible places.
 # Used only if --retry-until-up is set.
diff --git a/sky/clouds/aws.py b/sky/clouds/aws.py
index 1f6f4e21a6f..f0f4485cd9f 100644
--- a/sky/clouds/aws.py
+++ b/sky/clouds/aws.py
@@ -110,15 +110,14 @@ 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
+        # Deep Learning AMI GPU PyTorch 1.12.1 (Ubuntu 20.04) 20221025
         # 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
+        # Nvidia driver: 510.47.03, CUDA Version: 11.6 (supports torch==1.13.0+cu117)
         amis = {
-            'us-east-1': 'ami-0729d913a335efca7',
-            'us-east-2': 'ami-070f4af81c19b41bf',
-            # This AMI is 20210623 as aws does not provide a newer one.
-            'us-west-1': 'ami-0b3c34d643904a734',
-            'us-west-2': 'ami-050814f384259894c',
+            'us-east-1': 'ami-0eb1f91977a3fcc1b',
+            'us-east-2': 'ami-0274a6db2e19b7cc6',
+            'us-west-1': 'ami-0fb299af41d32cfd3',
+            'us-west-2': 'ami-04ba15f9bd464eb20',
         }
         assert region_name in amis, region_name
         return amis[region_name]
diff --git a/sky/clouds/azure.py b/sky/clouds/azure.py
index 9a22bab943d..594224d27d2 100644
--- a/sky/clouds/azure.py
+++ b/sky/clouds/azure.py
@@ -82,6 +82,12 @@ 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',
diff --git a/sky/clouds/gcp.py b/sky/clouds/gcp.py
index 3b288362878..7874502c1bd 100644
--- a/sky/clouds/gcp.py
+++ b/sky/clouds/gcp.py
@@ -214,7 +214,10 @@ def make_deploy_resources_variables(
         region_name = region.name
         zones = [zones[0].name]
 
-        image_id = _IMAGE_ID_PREFIX + 'common-cpu-v20220806'
+        # gcloud compute images list \
+        # --project deeplearning-platform-release \
+        # --no-standard-images
+        image_id = _IMAGE_ID_PREFIX + 'common-cpu-v20221026'
 
         r = resources
         # Find GPU spec, if any.
@@ -259,7 +262,8 @@ def make_deploy_resources_variables(
                     image_id = _IMAGE_ID_PREFIX + 'common-cu113-v20220701'
                 else:
                     # CUDA driver version 510.47.03, CUDA Library 11.6
-                    image_id = _IMAGE_ID_PREFIX + 'common-cu113-v20220806'
+                    # Does not support torch==1.13.0 with cu117
+                    image_id = _IMAGE_ID_PREFIX + 'common-cu113-v20221026'
 
         if resources.image_id is not None:
             image_id = resources.image_id
diff --git a/sky/skylet/providers/azure/node_provider.py b/sky/skylet/providers/azure/node_provider.py
index 630b5993474..29ba743fc31 100644
--- a/sky/skylet/providers/azure/node_provider.py
+++ b/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)


Code review comments:
diff_hunk: @@ -110,15 +110,14 @@ 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
+        # Deep Learning AMI GPU PyTorch 1.12.1 (Ubuntu 20.04) 20221025
         # 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
+        # Nvidia driver: 510.47.03, CUDA Version: 11.6 (supports torch==1.13.0+cu117)
path: sky/clouds/aws.py
text: @user1:
Q: Just to understand, if the old image also comes with CUDA 11.6, this image upgrade is not necessary to fix the pytorch issue? 

(I'm weary of us frequently updating images. If there's no strong reasons or user demands, I think keeping the images stable is better.)

@author:
Good point!

The latest version of the AMI supports the PyTorch with the CUDA 11.7, which is the one installed by default with `pip install torch`, but the previous version does not.

I agree that we should not frequently update the images, and since the other cloud does not support PyTorch with CUDA 11.7. I am now leaning toward reverting it.

@user2:
have you checked if the AMI has an older version of CUDA under /usr/local (or other places)? then we only need to softlink /usr/local/cuda to the older version. this is how you are supposed to handle backwards compatibility

@user2:
e.g., in setup we can have `sudo rm /usr/local/cuda && sudo ln -s /usr/local/cudaxxxx /usr/local/cuda`

@author:
I think the backward compatibility of cuda is handled well even without re-do the soft link, as the `pip install torch==1.12.1+cu113` works well on the VM with cuda 11.6 installed.