[Tests] Add options for test to select clouds (skypilot-org#1587)

Michaelvll · concretevitamin · web-flow · commit 36181a6ae3bb · 2023-01-19T10:49:25.000-08:00
* Add options for test to select clouds

* strip image_id

* strip

* not error out if azure status not exist

* fix gcp images

* reduce test name

* format

* fix

* fix gcp tests

* name length

* Fix smoke test

* fix gcp region

* use west3

* fix some more tests

* use ubuntu instead

* fix name

* fix

* Only refresh status for cluster

* remnant

* add spot queue waiting

* revert gcp

* use ":" instead for gcloud

* fix launch hash filter str

* Use our own lsof to avoid not installed issue

* instruction for rerun failed tests

* spot queue left

* fix templates

* wait longer for autostop

* Fix autodown test

* adopt the mylsof to azure

* lint

* change zone

* format

* reduce test name length

* wait longer for autostop

* longer autostop

* fix format exception

* shorter cluster name

* Update CONTRIBUTING.md

Co-authored-by: Zongheng Yang &lt;zongheng.y@gmail.com&gt;

* address comments

* format

* contributing

* Document

* pyproject

* revert gcloud filter

* fix

* fix

* format

* fix test_smoke matching

* fix

* reduce length

* reduce cluster name length

* increase cancel wait time

* Reset the meaning of --cloud

* fix comment

* format

* address comments

* yapf

Co-authored-by: Zongheng Yang &lt;zongheng.y@gmail.com&gt;
diff --git a/.github/pull_request_template.md b/.github/pull_request_template.md
@@ -8,6 +8,6 @@
 Tested (run the relevant ones):
 
 - [ ] Any manual or new tests for this PR (please specify below)
-- [ ] All smoke tests: `bash tests/run_smoke_tests.sh` 
-- [ ] Relevant individual smoke tests: `bash tests/run_smoke_tests.sh test_fill_in_the_name` 
+- [ ] All smoke tests: `pytest tests/test_smoke.py` 
+- [ ] Relevant individual smoke tests: `pytest tests/test_smoke.py::test_fill_in_the_name` 
 - [ ] Backward compatibility tests: `bash tests/backward_comaptibility_tests.sh`
diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md
@@ -26,10 +26,20 @@ pip install -r requirements-dev.txt
 ### Testing
 To run smoke tests (NOTE: Running all smoke tests launches ~20 clusters):
 ```
-bash tests/run_smoke_tests.sh
+# Run all tests except for AWS
+pytest tests/test_smoke.py
+
+# Re-run last failed tests
+pytest --lf
 
 # Run one of the smoke tests
-bash tests/run_smoke_tests.sh test_minimal
+pytest tests/test_smoke.py::test_minimal
+
+# Only run test for AWS + generic tests
+pytest tests/test_smoke.py --aws
+
+# Change cloud for generic tests to aws
+pytest tests/test_smoke.py --generic-cloud aws
 ```
 
 For profiling code, use:
diff --git a/examples/env_check.yaml b/examples/env_check.yaml
@@ -1,6 +1,3 @@
-resources:
-  cloud: aws
-
 num_nodes: 2  
 
 workdir: .
diff --git a/examples/job_queue/cluster.yaml b/examples/job_queue/cluster.yaml
@@ -7,5 +7,4 @@
 #   sky exec jq job.yaml
 
 resources:
-  cloud: aws
   accelerators: K80
diff --git a/examples/job_queue/cluster_multinode.yaml b/examples/job_queue/cluster_multinode.yaml
@@ -8,7 +8,6 @@
 #   sky exec mjq job.yaml
 
 resources:
-  cloud: aws
   accelerators: T4
 
 num_nodes: 2
diff --git a/examples/job_queue/job_multinode.yaml b/examples/job_queue/job_multinode.yaml
@@ -21,7 +21,7 @@ setup: |
 run: |
   timestamp=$(date +%s)
   conda env list
-  for i in {1..240}; do
+  for i in {1..360}; do
     echo "$timestamp $i"
     sleep 1
   done
diff --git a/examples/managed_spot.yaml b/examples/managed_spot.yaml
@@ -1,10 +1,5 @@
 name: minimal
 
-resources:
-  cloud: aws
-  use_spot: true
-  spot_recovery: failover
-
 setup: |
   echo "running setup"
   pip install tqdm
diff --git a/examples/multi_echo.py b/examples/multi_echo.py
@@ -8,7 +8,7 @@
 import sky
 
 
-def run(cluster: Optional[str] = None):
+def run(cluster: Optional[str] = None, cloud: Optional[str] = None):
     if cluster is None:
         # (username, last 4 chars of hash of hostname): for uniquefying users on
         # shared-account cloud providers.
@@ -17,9 +17,13 @@ def run(cluster: Optional[str] = None):
         _user_and_host = f'{getpass.getuser()}-{hostname_hash}'
         cluster = f'test-multi-echo-{_user_and_host}'
 
+    if cloud is None:
+        cloud = 'gcp'
+    cloud = sky.clouds.CLOUD_REGISTRY.from_str(cloud)
+
     # Create the cluster.
     with sky.Dag() as dag:
-        cluster_resources = sky.Resources(sky.AWS(), accelerators={'K80': 1})
+        cluster_resources = sky.Resources(cloud, accelerators={'K80': 1})
         task = sky.Task(num_nodes=2).set_resources(cluster_resources)
     # `detach_run` will only detach the `run` command. The provision and
     # `setup` are still blocking.
@@ -38,7 +42,10 @@ def _exec(i):
 
 if __name__ == '__main__':
     cluster = None
+    cloud = None
     if len(sys.argv) > 1:
         # For smoke test passing in a cluster name.
         cluster = sys.argv[1]
-    run(cluster)
+    if len(sys.argv) > 2:
+        cloud = sys.argv[2]
+    run(cluster, cloud)
diff --git a/examples/multi_hostname.yaml b/examples/multi_hostname.yaml
@@ -1,8 +1,5 @@
 name: multi_hostname
 
-resources:
-  cloud: gcp
-
 num_nodes: 2
 
 # The run command will be run on *all* nodes.
diff --git a/examples/per_region_images.yaml b/examples/per_region_images.yaml
@@ -1,7 +1,8 @@
 resources:
   cloud: aws
+  instance_type: g4dn.xlarge
   image_id: 
-    us-east-1: ami-0729d913a335efca7 # Ubuntu 20.04 
+    us-west-2: ami-0fe5af21074ad2a10  # Deep learning AMI with CUDA 11.6 without conda installed
     us-west-1: skypilot:gpu-ubuntu-1804
 
 
@@ -10,3 +11,4 @@ setup: |
 
 run: |
   conda env list
+  nvidia-smi
diff --git a/examples/resnet_distributed_tf_app.py b/examples/resnet_distributed_tf_app.py
@@ -7,7 +7,7 @@
 import sky
 
 
-def run(cluster: Optional[str] = None):
+def run(cluster: Optional[str] = None, cloud: Optional[str] = None):
     if cluster is None:
         # (username, last 4 chars of hash of hostname): for uniquefying users on
         # shared-account cloud providers.
@@ -75,14 +75,19 @@ def run_fn(node_rank: int, ip_list: List[str]) -> Optional[str]:
         train.set_inputs('gs://cloud-tpu-test-datasets/fake_imagenet',
                          estimated_size_gigabytes=70)
         train.set_outputs('resnet-model-dir', estimated_size_gigabytes=0.1)
-        train.set_resources(sky.Resources(sky.AWS(), accelerators='V100'))
+        train.set_resources(
+            sky.Resources(sky.clouds.CLOUD_REGISTRY.from_str(cloud),
+                          accelerators='V100'))
 
     sky.launch(dag, cluster_name=cluster, retry_until_up=True)
 
 
 if __name__ == '__main__':
     cluster = None
+    cloud = None
     if len(sys.argv) > 1:
         # For smoke test passing in a cluster name.
         cluster = sys.argv[1]
-    run(cluster)
+    if len(sys.argv) > 2:
+        cloud = sys.argv[2]
+    run(cluster, cloud)
diff --git a/examples/using_file_mounts.yaml b/examples/using_file_mounts.yaml
@@ -81,7 +81,7 @@ file_mounts:
 
 setup: |
   sudo apt update
-  sudo apt install tree
+  sudo apt install -y tree
 
 run: |
   set -ex
diff --git a/pyproject.toml b/pyproject.toml
@@ -16,6 +16,7 @@ env = [
     "SKYPILOT_DEBUG=1",
     "SKYPILOT_DISABLE_USAGE_COLLECTION=1"
 ]
+addopts = "-s -n 16 -q --tb=short --disable-warnings"
 
 [tool.mypy]
 python_version = "3.8"
diff --git a/sky/backends/backend_utils.py b/sky/backends/backend_utils.py
@@ -1355,8 +1355,8 @@ def _get_tpu_vm_pod_ips(ray_config: Dict[str, Any],
     cluster_name = ray_config['cluster_name']
     zone = ray_config['provider']['availability_zone']
     query_cmd = (f'gcloud compute tpus tpu-vm list --filter='
-                 f'\\(labels.ray-cluster-name={cluster_name}\\) '
-                 f'--zone={zone} --format=value\\(name\\)')
+                 f'"(labels.ray-cluster-name={cluster_name})" '
+                 f'--zone={zone} --format="value(name)"')
     returncode, stdout, stderr = log_lib.run_with_log(query_cmd,
                                                       '/dev/null',
                                                       shell=True,
@@ -1678,6 +1678,9 @@ def _query_status_gcp(
     cluster: str,
     ray_config: Dict[str, Any],
 ) -> List[global_user_state.ClusterStatus]:
+    # Note: we use ":" for filtering labels for gcloud, as the latest gcloud (v393.0)
+    # fails to filter labels with "=".
+    # Reference: https://cloud.google.com/sdk/gcloud/reference/topic/filters
     launch_hashes = _ray_launch_hash(cluster, ray_config)
     assert launch_hashes is not None
     hash_filter_str = ' '.join(launch_hashes)
diff --git a/sky/backends/cloud_vm_ray_backend.py b/sky/backends/cloud_vm_ray_backend.py
@@ -2897,10 +2897,9 @@ def teardown_no_lock(self,
                     terminate_cmd = tpu_utils.terminate_tpu_vm_cluster_cmd(
                         cluster_name, zone, log_abs_path)
                 else:
-                    query_cmd = (
-                        f'gcloud compute instances list --filter='
-                        f'\\(labels.ray-cluster-name={cluster_name}\\) '
-                        f'--zones={zone} --format=value\\(name\\)')
+                    query_cmd = (f'gcloud compute instances list --filter='
+                                 f'"(labels.ray-cluster-name={cluster_name})" '
+                                 f'--zones={zone} --format=value\\(name\\)')
                     terminate_cmd = (
                         f'gcloud compute instances delete --zone={zone}'
                         f' --quiet $({query_cmd})')
@@ -2956,8 +2955,14 @@ def teardown_no_lock(self,
             #   never launched and the errors are related to pre-launch
             #   configurations (such as VPC not found). So it's safe & good UX
             #   to not print a failure message.
+            #
+            # '(ResourceGroupNotFound)': this indicates the resource group on
+            #   Azure is not found. That means the cluster is already deleted
+            #   on the cloud. So it's safe & good UX to not print a failure
+            #   message.
             elif ('TPU must be specified.' not in stderr and
-                  'SKYPILOT_ERROR_NO_NODES_LAUNCHED: ' not in stderr):
+                  'SKYPILOT_ERROR_NO_NODES_LAUNCHED: ' not in stderr and
+                  '(ResourceGroupNotFound)' not in stderr):
                 logger.error(
                     _TEARDOWN_FAILURE_MESSAGE.format(
                         extra_reason='',
diff --git a/sky/cli.py b/sky/cli.py
@@ -1368,11 +1368,18 @@ def exec(
     is_flag=True,
     required=False,
     help='Query the latest cluster statuses from the cloud provider(s).')
+@click.argument('clusters',
+                required=False,
+                type=str,
+                nargs=-1,
+                **_get_shell_complete_args(_complete_cluster_name))
 @usage_lib.entrypoint
-def status(all: bool, refresh: bool):  # pylint: disable=redefined-builtin
+def status(all: bool, refresh: bool, clusters: List[str]):  # pylint: disable=redefined-builtin
     # NOTE(dev): Keep the docstring consistent between the Python API and CLI.
     """Show clusters.
 
+    If CLUSTERS is given, show those clusters. Otherwise, show all clusters.
+
     The following fields for each cluster are recorded: cluster name, time
     since last launch, resources, region, zone, hourly price, status, autostop,
     command.
@@ -1417,7 +1424,11 @@ def status(all: bool, refresh: bool):  # pylint: disable=redefined-builtin
       or for autostop-enabled clusters, use ``--refresh`` to query the latest
       cluster statuses from the cloud providers.
     """
-    cluster_records = core.status(refresh=refresh)
+    if clusters:
+        clusters = _get_glob_clusters(clusters)
+    else:
+        clusters = None
+    cluster_records = core.status(cluster_names=clusters, refresh=refresh)
     nonreserved_cluster_records = []
     reserved_clusters = dict()
     for cluster_record in cluster_records:
diff --git a/sky/clouds/gcp.py b/sky/clouds/gcp.py
@@ -308,8 +308,7 @@ def make_deploy_resources_variables(
         # --no-standard-images
         # We use the debian image, as the ubuntu image has some connectivity
         # issue when first booted.
-        image_id = service_catalog.get_image_id_from_tag(
-            'skypilot:cpu-debian-10', clouds='gcp')
+        image_id = 'skypilot:cpu-debian-10'
 
         r = resources
         # Find GPU spec, if any.
@@ -353,22 +352,23 @@ def make_deploy_resources_variables(
                     # 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 = service_catalog.get_image_id_from_tag(
-                        'skypilot:k80-debian-10', clouds='gcp')
+                    image_id = 'skypilot:k80-debian-10'
                 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 = service_catalog.get_image_id_from_tag(
-                        'skypilot:gpu-debian-10', clouds='gcp')
+                    image_id = 'skypilot:gpu-debian-10'
 
         if resources.image_id is not None:
             if None in resources.image_id:
                 image_id = resources.image_id[None]
             else:
                 assert region_name in resources.image_id, resources.image_id
                 image_id = resources.image_id[region_name]
+        if image_id.startswith('skypilot:'):
+            image_id = service_catalog.get_image_id_from_tag(image_id,
+                                                             clouds='gcp')
 
         assert image_id is not None, (image_id, r)
         resources_vars['image_id'] = image_id
diff --git a/sky/clouds/service_catalog/gcp_catalog.py b/sky/clouds/service_catalog/gcp_catalog.py
@@ -460,6 +460,8 @@ def get_image_id_from_tag(tag: str, region: Optional[str]) -> Optional[str]:
     return common.get_image_id_from_tag_impl(_image_df, tag, region)
 
 
-def validate_image_tag(tag: str, region: Optional[str]) -> bool:
+def is_image_tag_valid(tag: str, region: Optional[str]) -> bool:
     """Returns whether the image tag is valid."""
-    return common.is_image_tag_valid_impl(_image_df, tag, region)
+    # GCP images are not region-specific.
+    del region  # Unused.
+    return common.is_image_tag_valid_impl(_image_df, tag, None)
diff --git a/sky/core.py b/sky/core.py
@@ -33,7 +33,10 @@
 def status(cluster_names: Optional[Union[str, Sequence[str]]] = None,
            refresh: bool = False) -> List[Dict[str, Any]]:
     # NOTE(dev): Keep the docstring consistent between the Python API and CLI.
-    """Get all cluster statuses.
+    """Get cluster statuses.
+
+    If cluster_names is given, return those clusters. Otherwise, return all
+    clusters.
 
     Each returned value has the following fields:
 
diff --git a/sky/resources.py b/sky/resources.py
@@ -87,9 +87,14 @@ def __init__(
         # The key is None if the same image_id applies for all regions.
         self._image_id = image_id
         if isinstance(image_id, str):
-            self._image_id = {self._region: image_id}
-        elif isinstance(image_id, dict) and None in image_id:
-            self._image_id = {self._region: image_id[None]}
+            self._image_id = {self._region: image_id.strip()}
+        elif isinstance(image_id, dict):
+            if None in image_id:
+                self._image_id = {self._region: image_id[None].strip()}
+            else:
+                self._image_id = {
+                    k.strip(): v.strip() for k, v in image_id.items()
+                }
 
         self._set_accelerators(accelerators, accelerator_args)
 
@@ -486,7 +491,7 @@ def _try_validate_image_id(self) -> None:
                 region_str = f' ({region})' if region else ''
                 with ux_utils.print_exception_no_traceback():
                     raise ValueError(
-                        f'Image tag {image_id} is not valid, please make sure'
+                        f'Image tag {image_id!r} is not valid, please make sure'
                         f' the tag exists in {self._cloud}{region_str}.')
 
             if (self._cloud.is_same_cloud(clouds.AWS()) and
diff --git a/sky/spot/recovery_strategy.py b/sky/spot/recovery_strategy.py
@@ -169,7 +169,9 @@ def _launch(self, max_retry=3, raise_on_failure=True) -> Optional[float]:
                 # If the launch fails, it will be recovered by the following
                 # code.
                 logger.info('Failed to launch the spot cluster with error: '
-                            f'{type(e)}: {e}')
+                            f'{common_utils.format_exception(e)})')
+                import traceback  # pylint: disable=import-outside-toplevel
+                logger.info(f'  Traceback: {traceback.format_exc()}')
                 retry_launch = True
                 exception = e
 
diff --git a/sky/templates/aws-ray.yml.j2 b/sky/templates/aws-ray.yml.j2
@@ -118,13 +118,16 @@ setup_commands:
   # Line 'sudo grep ..': set the number of threads per process to unlimited to avoid ray job submit stucking issue when the number of running ray jobs increase.
   # Line 'mkdir -p ..': disable host key check
   # Line 'python3 -c ..': patch the buggy ray files and enable `-o allow_other` option for `goofys`
-  - sudo systemctl stop unattended-upgrades || true;
+  - function mylsof { p=$(for pid in /proc/{0..9}*; do i=$(basename "$pid"); for file in "$pid"/fd/*; do link=$(readlink -e "$file"); if [ "$link" = "$1" ]; then echo "$i"; fi; done; done); echo "$p"; };
+    sudo systemctl stop unattended-upgrades || true;
     sudo systemctl disable unattended-upgrades || true;
     sudo sed -i 's/Unattended-Upgrade "1"/Unattended-Upgrade "0"/g' /etc/apt/apt.conf.d/20auto-upgrades || true;
-    sudo kill -9 `sudo lsof /var/lib/dpkg/lock-frontend | awk '{print $2}' | tail -n 1` || true;
-    sudo pkill -9 apt-get;
+    p=$(mylsof "/var/lib/dpkg/lock-frontend"); echo "$p";
+    sudo kill -9 `echo "$p" | tail -n 1` || true;
+    sudo rm /var/lib/dpkg/lock-frontend;
     sudo pkill -9 dpkg;
-    sudo dpkg --configure -a;
+    sudo pkill -9 apt-get;
+    sudo dpkg --configure --force-overwrite -a;
     mkdir -p ~/.ssh; touch ~/.ssh/config;
     pip3 --version > /dev/null 2>&1 || (curl -sSL https://bootstrap.pypa.io/get-pip.py -o get-pip.py && python3 get-pip.py && echo "PATH=$HOME/.local/bin:$PATH" >> ~/.bashrc);
     (type -a python | grep -q python3) || echo 'alias python=python3' >> ~/.bashrc;
diff --git a/sky/templates/azure-ray.yml.j2 b/sky/templates/azure-ray.yml.j2
diff --git a/sky/templates/gcp-ray.yml.j2 b/sky/templates/gcp-ray.yml.j2
diff --git a/sky/utils/tpu_utils.py b/sky/utils/tpu_utils.py
diff --git a/tests/conftest.py b/tests/conftest.py
diff --git a/tests/run_smoke_tests.sh b/tests/run_smoke_tests.sh
diff --git a/tests/test_smoke.py b/tests/test_smoke.py
diff --git a/tests/test_yamls/gcp_per_region_images.yaml b/tests/test_yamls/gcp_per_region_images.yaml
diff --git a/tests/test_yamls/minimal.yaml b/tests/test_yamls/minimal.yaml
diff --git a/tests/test_yamls/no_conda_ami.yaml b/tests/test_yamls/no_conda_ami.yaml
