Make azure-start-stop use 1 node to speed up tests. (skypilot-org#619)

* Make azure-start-stop use 1 node to speed up tests.

* Add --num_nodes to both launch and exec.

Tested:
- sky launch --num_nodes=2 'echo $(hostname)' -c test
- sky exec test 'echo $(hostname)' --num_nodes=2
- sky exec test 'echo $(hostname)'

* Support test_azure_start_stop_two_nodes() under --runslow.

* Support passing a test name.

* Make Task.num_nodes a property and validate it.

Author: concretevitamin

examples/azure_start_stop.yaml

@@ -4,7 +4,10 @@ name: azure-start-stop
 resources:
   cloud: azure
 
-num_nodes: 2
+# Optimizing for smoke tests
+#   2 nodes: smoke tests ~37 mins
+#   1 node: smoke tests ~19 mins
+# num_nodes: 2
 
 # The setup command.  Will be run under the working directory.
 setup: 'echo "azure-start-stop [setup]"'

sky/cli.py

@@ -500,6 +500,12 @@ def cli():
           'resources and is used for scheduling the task. '
           'Overrides the "accelerators" '
           'config in the YAML if both are supplied.'))
+@click.option('--num_nodes',
+              required=False,
+              type=int,
+              help=('Number of nodes to launch and to execute the task on. '
+                    'Overrides the "num_nodes" config in the YAML if both are '
+                    'supplied.'))
 @click.option(
     '--use-spot/--no-use-spot',
     required=False,
@@ -523,10 +529,21 @@ def cli():
               default=False,
               required=False,
               help='Skip confirmation prompt.')
-def launch(entrypoint: str, cluster: Optional[str], dryrun: bool,
-           detach_run: bool, backend_name: str, workdir: Optional[str],
-           cloud: Optional[str], gpus: Optional[str], use_spot: Optional[bool],
-           name: Optional[str], disk_size: Optional[int], yes: bool):
+def launch(
+    entrypoint: str,
+    cluster: Optional[str],
+    dryrun: bool,
+    detach_run: bool,
+    backend_name: str,
+    workdir: Optional[str],
+    cloud: Optional[str],
+    gpus: Optional[str],
+    num_nodes: Optional[int],
+    use_spot: Optional[bool],
+    name: Optional[str],
+    disk_size: Optional[int],
+    yes: bool,
+):
     """Launch a task from a YAML or a command (rerun setup if cluster exists).
 
     If ENTRYPOINT points to a valid YAML file, it is read in as the task
@@ -570,6 +587,8 @@ def launch(entrypoint: str, cluster: Optional[str], dryrun: bool,
         if disk_size is not None:
             new_resources.disk_size = disk_size
         task.set_resources({new_resources})
+        if num_nodes is not None:
+            task.num_nodes = num_nodes
         if name is not None:
             task.name = name
 
@@ -623,21 +642,34 @@ def launch(entrypoint: str, cluster: Optional[str], dryrun: bool,
     '--gpus',
     required=False,
     type=str,
-    help=('Task demands: Type and number of GPUs to use. Example values: '
+    help=('Task demand: Type and number of GPUs to use. Example values: '
           '"V100:8", "V100" (short for a count of 1), or "V100:0.5" '
           '(fractional counts are supported by the scheduling framework). '
           'This is used for scheduling the task, so it must fit the '
           'cluster\'s total resources. Overrides the "accelerators" '
           'config in the YAML if both are supplied.'))
+@click.option('--num_nodes',
+              required=False,
+              type=int,
+              help=('Task demand: Number of nodes to execute the task on. '
+                    'Overrides the "num_nodes" config in the YAML if both are '
+                    'supplied.'))
 @click.option('--name',
               '-n',
               required=False,
               type=str,
               help=('Task name. Overrides the "name" '
                     'config in the YAML if both are supplied.'))
 # pylint: disable=redefined-builtin
-def exec(cluster: str, entrypoint: str, detach_run: bool,
-         workdir: Optional[str], gpus: Optional[str], name: Optional[str]):
+def exec(
+    cluster: str,
+    entrypoint: str,
+    detach_run: bool,
+    workdir: Optional[str],
+    gpus: Optional[str],
+    num_nodes: Optional[int],
+    name: Optional[str],
+):
     """Execute a task or a command on a cluster (skip setup).
 
     If ENTRYPOINT points to a valid YAML file, it is read in as the task
@@ -646,13 +678,15 @@ def exec(cluster: str, entrypoint: str, detach_run: bool,
     \b
     Execution and scheduling behavior:
     \b
-    - If ENTRYPOINT is a YAML, or if it is a command with `--gpus` specified:
-      it is treated as a proper task that will undergo job queue scheduling,
-      respecting its resource requirement. It can be executed on any node of th
-      cluster with enough resources.
-    - Otherwise (if ENTRYPOINT is a command and no `--gpus` specified), it is
-      treated as an inline command, to be executed only on the head node of the
-      cluster.
+    - If ENTRYPOINT is a YAML, or if it is a command with a resource demand
+      flag specified (`--gpus` or `--num_nodes`): it is treated as a proper
+      task that will undergo job queue scheduling, respecting its resource
+      requirement. It can be executed on any node of th cluster with enough
+      resources.
+    - Otherwise (if ENTRYPOINT is a command and no resource demand flag
+      specified), it is treated as an inline command, to be executed only on
+      the head node of the cluster. This is useful for monitoring commands
+      (e.g., gpustat, htop).
 
     In both cases, the commands are run under the task's workdir (if specified).
 
@@ -704,6 +738,7 @@ def exec(cluster: str, entrypoint: str, detach_run: bool,
         raise click.BadParameter(f'Cluster \'{cluster}\' not found.  '
                                  'Use `sky launch` to provision first.')
     backend = backend_utils.get_backend_from_handle(handle)
+    resource_demand_specified = gpus is not None or num_nodes is not None
 
     with sky.Dag() as dag:
         if _check_yaml(entrypoint):
@@ -722,7 +757,7 @@ def exec(cluster: str, entrypoint: str, detach_run: bool,
                 # Run inline commands directly on head node if the resources are
                 # not set. User should take the responsibility to not overload
                 # the cluster.
-                if gpus is None:
+                if not resource_demand_specified:
                     if workdir is not None:
                         backend.sync_workdir(handle, workdir)
                     backend.run_on_head(
@@ -743,6 +778,8 @@ def exec(cluster: str, entrypoint: str, detach_run: bool,
             copied = copy.deepcopy(list(task.resources)[0])
             copied.accelerators = _parse_accelerator_options(gpus)
             task.set_resources({copied})
+        if num_nodes is not None:
+            task.num_nodes = num_nodes
         if name is not None:
             task.name = name
 

sky/task.py

@@ -110,10 +110,9 @@ def __init__(
         self.storage_plans = {}
         self.setup = setup
         self.workdir = workdir
-        self.docker_image = docker_image if docker_image \
-            else 'gpuci/miniconda-cuda:11.4-runtime-ubuntu18.04'
-        self._explicit_num_nodes = num_nodes  # Used as a scheduling constraint.
-        self.num_nodes = 1 if num_nodes is None else num_nodes
+        self.docker_image = (docker_image if docker_image else
+                             'gpuci/miniconda-cuda:11.4-runtime-ubuntu18.04')
+        self.num_nodes = num_nodes
         self.inputs = None
         self.outputs = None
         self.estimated_inputs_size_gigabytes = None
@@ -323,9 +322,18 @@ def from_yaml(yaml_path):
         task.set_resources({resources})
         return task
 
-    def validate_config(self):
-        if self.num_nodes <= 0:
-            raise ValueError('Must set Task.num_nodes to >0.')
+    @property
+    def num_nodes(self) -> int:
+        return self._num_nodes
+
+    @num_nodes.setter
+    def num_nodes(self, num_nodes: Optional[int]) -> None:
+        if num_nodes is None:
+            num_nodes = 1
+        if not isinstance(num_nodes, int) or num_nodes <= 0:
+            raise ValueError(
+                f'num_nodes should be a positive int. Got: {num_nodes}')
+        self._num_nodes = num_nodes
 
     # E.g., 's3://bucket', 'gs://bucket', or None.
     def set_inputs(self, inputs, estimated_size_gigabytes):

tests/conftest.py

@@ -0,0 +1,27 @@
+import pytest
+
+# Usage: use
+#   @pytest.mark.slow
+# to mark a test as slow and to skip by default.
+# https://docs.pytest.org/en/latest/example/simple.html#control-skipping-of-tests-according-to-command-line-option
+
+
+def pytest_addoption(parser):
+    parser.addoption('--runslow',
+                     action='store_true',
+                     default=False,
+                     help='run slow tests')
+
+
+def pytest_configure(config):
+    config.addinivalue_line('markers', 'slow: mark test as slow to run')
+
+
+def pytest_collection_modifyitems(config, items):
+    if config.getoption('--runslow'):
+        # --runslow given in cli: do not skip slow tests
+        return
+    skip_slow = pytest.mark.skip(reason='need --runslow option to run')
+    for item in items:
+        if 'slow' in item.keywords:
+            item.add_marker(skip_slow)

tests/run_smoke_tests.sh

@@ -1,2 +1,22 @@
 #!/bin/bash
-time pytest -s -n 16 -q --tb=short --disable-warnings tests/test_smoke.py
+# Usage:
+#
+#   # Run everything
+#   bash tests/run_smoke_tests.sh
+#
+#   # Re-run a failed test
+#   bash tests/run_smoke_tests.sh test_azure_start_stop
+#
+
+test=${1:-""}
+if [ -z "$test" ]
+then
+    test_spec=tests/test_smoke.py
+else
+    test_spec=tests/test_smoke.py::"$test"
+fi
+
+pytest -s -n 16 -q --tb=short --disable-warnings "$test_spec"
+
+# To run all tests including the slow ones, add the --runslow flag:
+# pytest --runslow -s -n 16 -q --tb=short --disable-warnings tests/test_smoke.py

tests/test_optimizer_dryruns.py

@@ -174,3 +174,12 @@ def test_parse_accelerators_from_yaml():
     with pytest.raises(ValueError) as e:
         _test_parse_accelerators(spec, None)
         assert 'The "accelerators" field as a str ' in str(e.value)
+
+
+def test_invalid_num_nodes():
+    for invalid_value in (-1, 2.2, 1.0):
+        with pytest.raises(ValueError) as e:
+            with sky.Dag():
+                task = sky.Task()
+                task.num_nodes = invalid_value
+            assert 'num_nodes should be a positive int' in str(e.value)

tests/test_smoke.py

@@ -4,6 +4,7 @@
 from typing import List, Optional, Tuple, NamedTuple
 
 import colorama
+import pytest
 
 from sky.backends import backend_utils
 
@@ -33,16 +34,12 @@ def run_one_test(test: Test) -> Tuple[int, str, str]:
                                            delete=False)
     test.echo('Test started.')
     for command in test.commands:
-        # test.echo(f'  {command}')
         proc = subprocess.Popen(
             command,
             stdout=log_file,
             stderr=subprocess.STDOUT,
             shell=True,
         )
-        message = (f'Test command exited with non-zero status: '
-                   f'{command!r}')
-        error_occurred = (f'{message}')
         try:
             proc.wait(timeout=test.timeout)
         except subprocess.TimeoutExpired as e:
@@ -56,7 +53,8 @@ def run_one_test(test: Test) -> Tuple[int, str, str]:
 
     style = colorama.Style
     fore = colorama.Fore
-    outcome = f'{fore.RED}Failed{style.RESET_ALL}' if proc.returncode else f'{fore.GREEN}Passed{style.RESET_ALL}'
+    outcome = (f'{fore.RED}Failed{style.RESET_ALL}'
+               if proc.returncode else f'{fore.GREEN}Passed{style.RESET_ALL}')
     reason = f'\nReason: {command!r}' if proc.returncode else ''
     test.echo(f'{outcome}.'
               f'{reason}'
@@ -231,6 +229,25 @@ def test_distributed_tf():
     run_one_test(test)
 
 
+# ---------- Testing GCP start and stop instances ----------
+def test_gcp_start_stop():
+    test = Test(
+        'gcp-start-stop',
+        [
+            'sky launch -y -c test-gcp-start-stop examples/gcp_start_stop.yaml',
+            'sky logs test-gcp-start-stop 1 --status',  # Ensure the job succeeded.
+            'sky exec test-gcp-start-stop examples/gcp_start_stop.yaml',
+            'sky logs test-gcp-start-stop 2 --status',  # Ensure the job succeeded.
+            'sky stop -y test-gcp-start-stop',
+            'sky start -y test-gcp-start-stop',
+            'sky exec test-gcp-start-stop examples/gcp_start_stop.yaml',
+            'sky logs test-gcp-start-stop 3 --status',  # Ensure the job succeeded.
+        ],
+        'sky down -y test-gcp-start-stop',
+    )
+    run_one_test(test)
+
+
 # ---------- Testing Azure start and stop instances ----------
 def test_azure_start_stop():
     test = Test(
@@ -245,25 +262,25 @@ def test_azure_start_stop():
             'sky logs test-azure-start-stop 2 --status',  # Ensure the job succeeded.
         ],
         'sky down -y test-azure-start-stop',
-        timeout=30 * 60,  # 30 mins  (it takes around ~23 mins)
+        timeout=30 * 60,  # 30 mins
     )
     run_one_test(test)
 
 
-# ---------- Testing GCP start and stop instances ----------
-def test_gcp_start_stop():
+@pytest.mark.slow
+def test_azure_start_stop_two_nodes():
     test = Test(
-        'gcp-start-stop',
+        'azure-start-stop-two-nodes',
         [
-            'sky launch -y -c test-gcp-start-stop examples/gcp_start_stop.yaml',
-            'sky logs test-gcp-start-stop 1 --status',  # Ensure the job succeeded.
-            'sky exec test-gcp-start-stop examples/gcp_start_stop.yaml',
-            'sky logs test-gcp-start-stop 2 --status',  # Ensure the job succeeded.
-            'sky stop -y test-gcp-start-stop',
-            'sky start -y test-gcp-start-stop',
-            'sky exec test-gcp-start-stop examples/gcp_start_stop.yaml',
-            'sky logs test-gcp-start-stop 3 --status',  # Ensure the job succeeded.
+            'sky launch --num_nodes=2 -y -c test-azure-start-stop examples/azure_start_stop.yaml',
+            'sky exec --num_nodes=2 test-azure-start-stop examples/azure_start_stop.yaml',
+            'sky logs test-azure-start-stop 1 --status',  # Ensure the job succeeded.
+            'sky stop -y test-azure-start-stop',
+            'sky start -y test-azure-start-stop',
+            'sky exec --num_nodes=2 test-azure-start-stop examples/azure_start_stop.yaml',
+            'sky logs test-azure-start-stop 2 --status',  # Ensure the job succeeded.
         ],
-        'sky down -y test-gcp-start-stop',
+        'sky down -y test-azure-start-stop',
+        timeout=30 * 60,  # 30 mins  (it takes around ~23 mins)
     )
     run_one_test(test)