Repository: skypilot-org/skypilot
Pull request #1792 (2023-03-17 22:27:50)

[Spot] Spot jobs only failover through the regions in ssh_proxy_command
<!-- Describe the changes in this PR -->

To reproduce the problem:
1. Have a `~/.sky/config.yaml`:
```
aws:
  vpc_name: skypilot-test-vpc
  use_internal_ips: True
  ssh_proxy_command:
    us-east-2: ssh -W %h:%p -i ~/.ssh/sky-key -o StrictHostKeyChecking=no ec2-user@jump-ip
```
2. `sky spot launch --gpus A100:8 --cloud aws echo hi`

In the current the master, the failover of the spot cluster to be launched will go through all the regions, while the expected behavior should be similar as `sky launch --gpus A100:8 --cloud aws`, i.e. only trying the zones in the us-east-1 specified in `ssh_proxy_command`

This PR fixes the problem mentioned above.

Future TODOs:
- [ ] Add test for the failover.

<!-- Describe the tests ran -->
<!-- Unit tests (tests/test_*.py) are part of GitHub CI; below are tests that launch on the cloud. -->

Tested (run the relevant ones):
All the tests are run with `~/.sky/config.yaml` set.
- [x] Any manual or new tests for this PR (please specify below)
  - [x] snippets above
  - [x] ` sky launch --cloud aws --gpus A100:8 echo hi --use-spot`
  - [x] ` sky launch --cloud aws --gpus A100:8 echo hi`
  - [x] ` sky launch --cloud aws --gpus A100:8 --region us-east-2 echo hi` (should fail with `FAILED_NO_RESOURCES`)
  - [x] ` sky launch --cloud aws --gpus A100:8 --region us-west-2 echo hi` (should fail with `FAILED_PRECHECKS`)
- [x] All smoke tests: `pytest tests/test_smoke.py --managed-spot --aws` 



Patch:
diff --git a/sky/backends/backend_utils.py b/sky/backends/backend_utils.py
index 513615f7e..0594c4cf2 100644
--- a/sky/backends/backend_utils.py
+++ b/sky/backends/backend_utils.py
@@ -818,13 +818,23 @@ def write_cluster_config(
     else:
         # ssh_proxy_command_config: Dict[str, str], region_name -> command
         # This type check is done by skypilot_config at config load time.
-        if region_name not in ssh_proxy_command_config:
-            # Skip this region. The upper layer will handle the failover to
-            # other regions.
-            raise exceptions.ResourcesUnavailableError(
-                f'No ssh_proxy_command provided for region {region_name}. Skipped.'
-            )
-        ssh_proxy_command = ssh_proxy_command_config[region_name]
+
+        # There are two cases:
+        if keep_launch_fields_in_existing_config:
+            # (1) We're re-provisioning an existing cluster.
+            #
+            # We use None for ssh_proxy_command, which will be restored to the
+            # cluster's original value later by _replace_yaml_dicts().
+            ssh_proxy_command = None
+        else:
+            # (2) We're launching a new cluster.
+            #
+            # Resources.get_valid_regions_for_launchable() respects the keys (regions)
+            # in ssh_proxy_command in skypilot_config. So here we add an assert.
+            assert region_name in ssh_proxy_command_config, (
+                region_name, ssh_proxy_command_config)
+            ssh_proxy_command = ssh_proxy_command_config[region_name]
+
     logger.debug(f'Using ssh_proxy_command: {ssh_proxy_command!r}')
 
     # Use a tmp file path to avoid incomplete YAML file being re-used in the
diff --git a/sky/backends/cloud_vm_ray_backend.py b/sky/backends/cloud_vm_ray_backend.py
index 50dfee00a..95401d48e 100644
--- a/sky/backends/cloud_vm_ray_backend.py
+++ b/sky/backends/cloud_vm_ray_backend.py
@@ -68,7 +68,7 @@ _PATH_SIZE_MEGABYTES_WARN_THRESHOLD = 256
 _NODES_LAUNCHING_PROGRESS_TIMEOUT = {
     clouds.AWS: 90,
     clouds.Azure: 90,
-    clouds.GCP: 90,
+    clouds.GCP: 120,
     clouds.Lambda: 120,
     clouds.Local: 90,
 }
diff --git a/sky/execution.py b/sky/execution.py
index 599ea10b7..6991eacd0 100644
--- a/sky/execution.py
+++ b/sky/execution.py
@@ -595,6 +595,7 @@ def spot_launch(
             'uuid': task_uuid,
             'gcloud_installation_commands': gcp.GCLOUD_INSTALLATION_COMMAND,
             'is_dev': env_options.Options.IS_DEVELOPER.get(),
+            'is_debug': env_options.Options.SHOW_DEBUG_INFO.get(),
             'disable_logging': env_options.Options.DISABLE_LOGGING.get(),
             'logging_user_hash': common_utils.get_user_hash(),
             'retry_until_up': retry_until_up,
@@ -605,11 +606,11 @@ def spot_launch(
             # 'skypilot_config' module. Don't simply read the on-disk file as
             # it may have changed since this process started.
             #
-            # Pop any proxy command, because the controller would've been
-            # launched behind the proxy, and in general any nodes we launch may
-            # not have or need the proxy setup. (If the controller needs to
-            # launch spot clusters in another region/VPC, the user should
-            # properly set up VPC peering, which will allow the
+            # Set any proxy command to None, because the controller would've
+            # been launched behind the proxy, and in general any nodes we
+            # launch may not have or need the proxy setup. (If the controller
+            # needs to launch spot clusters in another region/VPC, the user
+            # should properly set up VPC peering, which will allow the
             # cross-region/VPC communication. The proxy command is orthogonal
             # to this scenario.)
             #
@@ -628,11 +629,23 @@ def spot_launch(
             # controller VPC (or name) == the spot job's VPC (or name). It may
             # not be a sufficient check (as it's always possible that peering
             # is not set up), but it may catch some obvious errors.
-            # TODO(zhwu): hacky. We should pop the proxy command of the cloud
-            # where the controller is launched (currently, only aws user uses
-            # proxy_command).
-            config_dict = skypilot_config.pop_nested(
-                ('aws', 'ssh_proxy_command'))
+            # TODO(zhwu): hacky. We should only set the proxy command of the
+            # cloud where the controller is launched (currently, only aws user
+            # uses proxy_command).
+            proxy_command_key = ('aws', 'ssh_proxy_command')
+            ssh_proxy_command = skypilot_config.get_nested(
+                proxy_command_key, None)
+            if isinstance(ssh_proxy_command, str):
+                config_dict = skypilot_config.set_nested(
+                    proxy_command_key, None)
+            elif isinstance(ssh_proxy_command, dict):
+                # Instead of removing the key, we set the value to empty string
+                # so that the controller will only try the regions specified by
+                # the keys.
+                ssh_proxy_command = {k: None for k in ssh_proxy_command}
+                config_dict = skypilot_config.set_nested(
+                    proxy_command_key, ssh_proxy_command)
+
             with tempfile.NamedTemporaryFile(mode='w', delete=False) as tmpfile:
                 common_utils.dump_yaml(tmpfile.name, config_dict)
                 vars_to_fill.update({
diff --git a/sky/optimizer.py b/sky/optimizer.py
index af5bf3195..aaa679d41 100644
--- a/sky/optimizer.py
+++ b/sky/optimizer.py
@@ -260,8 +260,17 @@ class Optimizer:
             num_resources = len(node.get_resources())
             for orig_resources, launchable_list in launchable_resources.items():
                 if not launchable_list:
+                    specified_resources_str = ''
+                    if node.get_resources():
+                        specified_resources = list(node.get_resources())[0]
+                        specified_resources_str = f' ({specified_resources})'
+                        if specified_resources.region:
+                            specified_resources_str = (
+                                f' ({specified_resources} '
+                                f'in {specified_resources.region})')
                     error_msg = (
-                        f'No launchable resource found for task {node}. '
+                        'Optimizer: No launchable resource found for task '
+                        f'{node}{specified_resources_str}. '
                         'To fix: relax its resource requirements.\n'
                         'Hint: \'sky show-gpus --all\' '
                         'to list available accelerators.\n'
diff --git a/sky/skypilot_config.py b/sky/skypilot_config.py
index eb3a70ccc..098fbb45d 100644
--- a/sky/skypilot_config.py
+++ b/sky/skypilot_config.py
@@ -11,9 +11,9 @@ To read a nested-key config:
 
   >> skypilot_config.get_nested(('auth', 'some_auth_config'), default_value)
 
-To pop a nested-key config:
+To set a value in the nested-key config:
 
-  >> config_dict = skypilot_config.pop_nested(('auth', 'some_key'))
+  >> config_dict = skypilot_config.set_nested(('auth', 'some_key'), value)
 
 This operation returns a deep-copy dict, and is safe in that any key not found
 will not raise an error.
@@ -93,8 +93,8 @@ def get_nested(keys: Sequence[str], default_value: Any) -> Any:
     return curr
 
 
-def pop_nested(keys: Sequence[str]) -> Dict[str, Any]:
-    """Returns a deep-copied config with the nested key popped.
+def set_nested(keys: Sequence[str], value: Any) -> Dict[str, Any]:
+    """Returns a deep-copied config with the nested key set to value.
 
     Like get_nested(), if any key is not found, this will not raise an error.
     """
@@ -105,15 +105,15 @@ def pop_nested(keys: Sequence[str]) -> Dict[str, Any]:
     to_return = curr
     prev = None
     for i, key in enumerate(keys):
-        if key in curr:
-            prev = curr
-            curr = curr[key]
-            if i == len(keys) - 1:
-                prev.pop(key)
-                logger.debug(f'Popped {keys}. Returning conf: {to_return}')
-        else:
-            # If any key not found, simply return.
-            return to_return
+        if key not in curr:
+            curr[key] = {}
+        prev = curr
+        curr = curr[key]
+        if i == len(keys) - 1:
+            prev_value = prev[key]
+            prev[key] = value
+            logger.debug(f'Set the value of {keys} to {value} (previous: '
+                         f'{prev_value}). Returning conf: {to_return}')
     return to_return
 
 
@@ -126,7 +126,7 @@ def _syntax_check_for_ssh_proxy_command(cloud: str) -> None:
 
     if isinstance(ssh_proxy_command_config, dict):
         for region, cmd in ssh_proxy_command_config.items():
-            if not isinstance(cmd, str):
+            if cmd and not isinstance(cmd, str):
                 raise ValueError(
                     f'Invalid ssh_proxy_command config for region {region!r} '
                     f'(expected a str): {cmd!r}')
diff --git a/sky/templates/spot-controller.yaml.j2 b/sky/templates/spot-controller.yaml.j2
index 7931706ce..ba22458ad 100644
--- a/sky/templates/spot-controller.yaml.j2
+++ b/sky/templates/spot-controller.yaml.j2
@@ -51,6 +51,9 @@ envs:
 {% if is_dev %}
   SKYPILOT_DEV: 1
 {% endif %}
+{% if is_debug %}
+  SKYPILOT_DEBUG: 1
+{% endif %}
 {% if disable_logging %}
   SKYPILOT_DISABLE_USAGE_COLLECTION: 1
 {% endif %}


Test patch:
diff --git a/.github/workflows/pytest.yml b/.github/workflows/pytest.yml
index 77f91927c..43fddd51e 100644
--- a/.github/workflows/pytest.yml
+++ b/.github/workflows/pytest.yml
@@ -17,6 +17,7 @@ jobs:
         python-version: [3.6]
         test-path:
           - tests/test_cli.py
+          - tests/test_config.py
           - tests/test_global_user_state.py
           - tests/test_list_accelerators.py
           - tests/test_optimizer_dryruns.py
diff --git a/tests/conftest.py b/tests/conftest.py
index 9c20348fd..ffb3dbfe5 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -171,3 +171,15 @@ def enable_all_clouds(monkeypatch):
         prefix='tmp_backup_config_default', delete=False)
     monkeypatch.setattr('sky.clouds.gcp.GCP_CONFIG_SKY_BACKUP_PATH',
                         config_file_backup.name)
+
+
+@pytest.fixture
+def aws_config_region(monkeypatch) -> str:
+    from sky import skypilot_config
+    region = 'us-west-2'
+    if skypilot_config.loaded():
+        ssh_proxy_command = skypilot_config.get_nested(
+            ('aws', 'ssh_proxy_command'), None)
+        if isinstance(ssh_proxy_command, dict) and ssh_proxy_command:
+            region = list(ssh_proxy_command.keys())[0]
+    return region
diff --git a/tests/test_config.py b/tests/test_config.py
new file mode 100644
index 000000000..dffd4843f
--- /dev/null
+++ b/tests/test_config.py
@@ -0,0 +1,121 @@
+import copy
+import pathlib
+import textwrap
+
+import pytest
+
+from sky import skypilot_config
+from sky.utils import common_utils
+
+VPC_NAME = 'vpc-12345678'
+PROXY_COMMAND = 'ssh -W %h:%p -i ~/.ssh/id_rsa -o StrictHostKeyChecking=no'
+
+
+def _reload_config() -> None:
+    skypilot_config._dict = None
+    skypilot_config._try_load_config()
+
+
+def _check_empty_config() -> None:
+    """Check that the config is empty."""
+    assert not skypilot_config.loaded()
+    assert skypilot_config.get_nested(
+        ('aws', 'ssh_proxy_command'), None) is None
+    assert skypilot_config.get_nested(('aws', 'ssh_proxy_command'),
+                                      'default') == 'default'
+    with pytest.raises(RuntimeError):
+        skypilot_config.set_nested(('aws', 'ssh_proxy_command'), 'value')
+
+
+def _create_config_file(config_file_path: pathlib.Path) -> None:
+    config_file_path.open('w').write(
+        textwrap.dedent(f"""\
+            aws:
+                vpc_name: {VPC_NAME}
+                use_internal_ips: true
+                ssh_proxy_command: {PROXY_COMMAND}
+            """))
+
+
+def test_no_config(monkeypatch) -> None:
+    """Test that the config is not loaded if the config file does not exist."""
+    monkeypatch.setattr(skypilot_config, 'CONFIG_PATH', '/tmp/does_not_exist')
+    _reload_config()
+    _check_empty_config()
+
+
+def test_empty_config(monkeypatch, tmp_path) -> None:
+    """Test that the config is not loaded if the config file is empty."""
+    with open(tmp_path / 'empty.yaml', 'w') as f:
+        f.write('')
+    monkeypatch.setattr(skypilot_config, 'CONFIG_PATH', tmp_path / 'empty.yaml')
+    _reload_config()
+    _check_empty_config()
+
+
+def test_config_get_set_nested(monkeypatch, tmp_path) -> None:
+    """Test that set_nested(), get_nested() works."""
+
+    # Load from a config file
+    config_path = tmp_path / 'config.yaml'
+    _create_config_file(config_path)
+    monkeypatch.setattr(skypilot_config, 'CONFIG_PATH', config_path)
+    _reload_config()
+    # Check that the config is loaded with the expected values
+    assert skypilot_config.loaded()
+    assert skypilot_config.get_nested(('aws', 'vpc_name'), None) == VPC_NAME
+    assert skypilot_config.get_nested(('aws', 'use_internal_ips'), None)
+    assert skypilot_config.get_nested(('aws', 'ssh_proxy_command'),
+                                      None) == PROXY_COMMAND
+
+    # Check set_nested() will copy the config dict and return a new dict
+    new_config = skypilot_config.set_nested(('aws', 'ssh_proxy_command'),
+                                            'new_value')
+    assert new_config['aws']['ssh_proxy_command'] == 'new_value'
+    assert skypilot_config.get_nested(('aws', 'ssh_proxy_command'),
+                                      None) == PROXY_COMMAND
+
+    # Check that dumping the config to a file with the new None can be reloaded
+    new_config2 = skypilot_config.set_nested(('aws', 'ssh_proxy_command'), None)
+    new_config_path = tmp_path / 'new_config.yaml'
+    common_utils.dump_yaml(new_config_path, new_config2)
+    monkeypatch.setattr(skypilot_config, 'CONFIG_PATH', new_config_path)
+    _reload_config()
+    assert skypilot_config.get_nested(('aws', 'vpc_name'), None) == VPC_NAME
+    assert skypilot_config.get_nested(('aws', 'use_internal_ips'), None)
+    assert skypilot_config.get_nested(
+        ('aws', 'ssh_proxy_command'), None) is None
+
+    # Check config with only partial keys still works
+    new_config3 = copy.copy(new_config2)
+    del new_config3['aws']['ssh_proxy_command']
+    del new_config3['aws']['use_internal_ips']
+    new_config_path = tmp_path / 'new_config3.yaml'
+    common_utils.dump_yaml(new_config_path, new_config3)
+    monkeypatch.setattr(skypilot_config, 'CONFIG_PATH', new_config_path)
+    _reload_config()
+    assert skypilot_config.get_nested(('aws', 'vpc_name'), None) == VPC_NAME
+    assert skypilot_config.get_nested(('aws', 'use_internal_ips'), None) is None
+    assert skypilot_config.get_nested(
+        ('aws', 'ssh_proxy_command'), None) is None
+    # set_nested() should still work
+    new_config4 = skypilot_config.set_nested(('aws', 'ssh_proxy_command'),
+                                             'new_value')
+    assert new_config4['aws']['ssh_proxy_command'] == 'new_value'
+    assert skypilot_config.get_nested(
+        ('aws', 'ssh_proxy_command'), None) is None
+
+
+def test_config_with_env(monkeypatch, tmp_path) -> None:
+    """Test that the config is loaded with environment variables."""
+    config_path = tmp_path / 'config.yaml'
+    _create_config_file(config_path)
+    monkeypatch.setenv(skypilot_config.ENV_VAR_SKYPILOT_CONFIG, config_path)
+    monkeypatch.setattr(skypilot_config, 'CONFIG_PATH',
+                        tmp_path / 'does_not_exist')
+    _reload_config()
+    assert skypilot_config.loaded()
+    assert skypilot_config.get_nested(('aws', 'vpc_name'), None) == VPC_NAME
+    assert skypilot_config.get_nested(('aws', 'use_internal_ips'), None)
+    assert skypilot_config.get_nested(('aws', 'ssh_proxy_command'),
+                                      None) == PROXY_COMMAND
diff --git a/tests/test_smoke.py b/tests/test_smoke.py
index d7638a895..520344fc0 100644
--- a/tests/test_smoke.py
+++ b/tests/test_smoke.py
@@ -1231,10 +1231,10 @@ def test_spot_failed_setup(generic_cloud: str):
 # ---------- Testing managed spot recovery ----------
 @pytest.mark.aws
 @pytest.mark.managed_spot
-def test_spot_recovery_aws():
+def test_spot_recovery_aws(aws_config_region):
     """Test managed spot recovery."""
     name = _get_cluster_name()
-    region = 'us-west-2'
+    region = aws_config_region
     test = Test(
         'spot_recovery_aws',
         [
@@ -1312,10 +1312,10 @@ def test_spot_recovery_default_resources(generic_cloud: str):
 
 @pytest.mark.aws
 @pytest.mark.managed_spot
-def test_spot_recovery_multi_node_aws():
+def test_spot_recovery_multi_node_aws(aws_config_region):
     """Test managed spot recovery."""
     name = _get_cluster_name()
-    region = 'us-west-2'
+    region = aws_config_region
     test = Test(
         'spot_recovery_multi_node_aws',
         [
@@ -1337,7 +1337,7 @@ def test_spot_recovery_multi_node_aws():
             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"',
         ],
         _SPOT_CANCEL_WAIT.format(job_name=name),
-        timeout=25 * 60,
+        timeout=30 * 60,
     )
     run_one_test(test)
 
@@ -1378,9 +1378,9 @@ def test_spot_recovery_multi_node_gcp():
 
 @pytest.mark.aws
 @pytest.mark.managed_spot
-def test_spot_cancellation_aws():
+def test_spot_cancellation_aws(aws_config_region):
     name = _get_cluster_name()
-    region = 'us-east-2'
+    region = aws_config_region
     test = Test(
         'spot_cancellation_aws',
         [
@@ -1421,7 +1421,7 @@ def test_spot_cancellation_aws():
              f'--filters Name=tag:ray-cluster-name,Values={name}-3-* '
              f'--query Reservations[].Instances[].InstanceId '
              '--output text)'),
-            'sleep 100',
+            'sleep 120',
             f'{_SPOT_QUEUE_WAIT}| grep {name}-3 | head -n1 | grep "RECOVERING"',
             _SPOT_CANCEL_WAIT.format(job_name=f'{name}-3'),
             'sleep 5',