[AWS] Bring-your-own-VPC that disables public IPs for all SkyPilot nodes.

sky/authentication.py

@@ -77,7 +77,10 @@ def get_or_generate_keys() -> Tuple[str, str]:
         _save_key_pair(private_key_path, public_key_path, private_key,
                        public_key)
     else:
-        assert os.path.exists(public_key_path)
+        # FIXME(skypilot): sky-launched instances fail this assert (has private
+        # key; but has not generated public key).
+        #   AssertionError: /home/ubuntu/.ssh/sky-key.pub
+        assert os.path.exists(public_key_path), public_key_path
     return private_key_path, public_key_path
 
 

sky/backends/backend_utils.py

@@ -26,7 +26,7 @@
 from requests import adapters
 from requests.packages.urllib3.util import retry as retry_lib
 from ray.autoscaler._private import commands as ray_commands
-from ray.autoscaler._private import util as ray_util
+from ray.autoscaler._private import util as ray_autoscaler_private_util
 import rich.console as rich_console
 import rich.progress as rich_progress
 import yaml
@@ -38,6 +38,7 @@
 from sky import clouds
 from sky import exceptions
 from sky import global_user_state
+from sky import sky_config
 from sky import sky_logging
 from sky import spot as spot_lib
 from sky.backends import onprem_utils
@@ -361,7 +362,12 @@ class SSHConfigHelper(object):
 
     @classmethod
     def _get_generated_config(cls, autogen_comment: str, host_name: str,
-                              ip: str, username: str, ssh_key_path: str):
+                              ip: str, username: str, ssh_key_path: str,
+                              proxy_command: Optional[str]):
+        if proxy_command is not None:
+            proxy = f'ProxyCommand {proxy_command}'
+        else:
+            proxy = ''
         codegen = textwrap.dedent(f"""\
             {autogen_comment}
             Host {host_name}
@@ -372,6 +378,7 @@ def _get_generated_config(cls, autogen_comment: str, host_name: str,
               ForwardAgent yes
               StrictHostKeyChecking no
               Port 22
+              {proxy}
             """)
         return codegen
 
@@ -437,8 +444,9 @@ def add_cluster(
                 f.writelines(config)
             os.chmod(config_path, 0o644)
 
+        proxy_command = auth_config.get('ssh_proxy_command', None)
         codegen = cls._get_generated_config(sky_autogen_comment, host_name, ip,
-                                            username, key_path)
+                                            username, key_path, proxy_command)
 
         # Add (or overwrite) the new config.
         if overwrite:
@@ -524,6 +532,8 @@ def _add_multinode_config(
         with open(config_path) as f:
             config = f.readlines()
 
+        proxy_command = auth_config.get('ssh_proxy_command', None)
+
         # Check if ~/.ssh/config contains existing names
         host_lines = [f'Host {c_name}' for c_name in worker_names]
         for i, line in enumerate(config):
@@ -536,7 +546,7 @@ def _add_multinode_config(
                 logger.warning(f'Using {host_name} to identify host instead.')
                 codegens[idx] = cls._get_generated_config(
                     sky_autogen_comment, host_name, external_worker_ips[idx],
-                    username, key_path)
+                    username, key_path, proxy_command)
 
         # All workers go to SKY_USER_FILE_PATH/ssh/{cluster_name}
         for i, line in enumerate(extra_config):
@@ -549,14 +559,14 @@ def _add_multinode_config(
                     overwrite_begin_idxs[idx] = i - 1
                 codegens[idx] = cls._get_generated_config(
                     sky_autogen_comment, host_name, external_worker_ips[idx],
-                    username, key_path)
+                    username, key_path, proxy_command)
 
         # This checks if all codegens have been created.
         for idx, ip in enumerate(external_worker_ips):
             if not codegens[idx]:
                 codegens[idx] = cls._get_generated_config(
                     sky_autogen_comment, worker_names[idx], ip, username,
-                    key_path)
+                    key_path, proxy_command)
 
         for idx in range(len(external_worker_ips)):
             # Add (or overwrite) the new config.
@@ -789,14 +799,25 @@ def write_cluster_config(
                 # See https://github.com/skypilot-org/skypilot/pull/742.
                 # Generate the name of the security group we're looking for.
                 # (username, last 4 chars of hash of hostname): for uniquefying
-                # users on shared-account cloud providers. Using uuid.getnode()
-                # is incorrect; observed to collide on Macs.
-                'security_group': f'sky-sg-{common_utils.user_and_hostname_hash()}',
-                # Azure only.
+                # users on shared-account scenarios.
+                'security_group': sky_config.get_nested(
+                    ('aws', 'security_group_name'),
+                    f'sky-sg-{common_utils.user_and_hostname_hash()}'),
+                'vpc_name': sky_config.get_nested(('aws', 'vpc_name'), None),
+                'use_internal_ips': sky_config.get_nested(
+                    ('aws', 'use_internal_ips'), False),
+                # Not exactly AWS only, but we only test it's supported on AWS
+                # for now:
+                'ssh_proxy_command': sky_config.get_nested(
+                    ('auth', 'ssh_proxy_command'), None),
+
+                # Azure only:
                 'azure_subscription_id': azure_subscription_id,
                 'resource_group': f'{cluster_name}-{region_name}',
-                # GCP only.
+
+                # GCP only:
                 'gcp_project_id': gcp_project_id,
+
                 # Ray version.
                 'ray_version': constants.SKY_REMOTE_RAY_VERSION,
                 # Cloud credentials for cloud storage.
@@ -1032,10 +1053,12 @@ def ssh_credential_from_yaml(cluster_yaml: str) -> Dict[str, str]:
     ssh_user = auth_section['ssh_user'].strip()
     ssh_private_key = auth_section.get('ssh_private_key')
     ssh_control_name = config.get('cluster_name', '__default__')
+    ssh_proxy_command = auth_section.get('ssh_proxy_command')
     return {
         'ssh_user': ssh_user,
         'ssh_private_key': ssh_private_key,
-        'ssh_control_name': ssh_control_name
+        'ssh_control_name': ssh_control_name,
+        'ssh_proxy_command': ssh_proxy_command,
     }
 
 
@@ -1473,7 +1496,7 @@ def _ray_launch_hash(cluster_name: str, ray_config: Dict[str, Any]) -> Set[str]:
     assert metadata is not None, cluster_name
     ray_launch_hashes = metadata.get('ray_launch_hashes', None)
     if ray_launch_hashes is not None:
-        logger.debug('Using cached launch_caches')
+        logger.debug('Using cached launch_hashes.')
         return set(ray_launch_hashes)
     with ux_utils.suppress_output():
         ray_config = ray_commands._bootstrap_config(ray_config)  # pylint: disable=protected-access
@@ -1484,14 +1507,27 @@ def _ray_launch_hash(cluster_name: str, ray_config: Dict[str, Any]) -> Set[str]:
     for node_type, node_config in ray_config['available_node_types'].items():
         if node_type == head_node_type:
             launch_config = ray_config.get('head_node', {})
+            auth_config = ray_config['auth']
         else:
             launch_config = ray_config.get('worker_nodes', {})
+            auth_config = dict(ray_config['auth'])
+            # Ray uses the head node to launch worker nodes. On the head node,
+            # ~/ray_bootstrap_config.yaml, which has any ssh_proxy_command
+            # field removed (see Ray's autoscaler/_private/commands.py), is
+            # passed to the autoscaler. Therefore when Ray calculates the hash
+            # for workers, ssh_proxy_command is not included. Here we follow
+            # this (otherwise our hash here would not match with what's on the
+            # console for workers).
+            #
+            # Note that head node is launched from the local client, whose
+            # launch has *does* include ssh_proxy_command.
+            auth_config.pop('ssh_proxy_command', None)
         launch_config = copy.deepcopy(launch_config)
 
         launch_config.update(node_config['node_config'])
         with ux_utils.suppress_output():
-            current_hash = ray_util.hash_launch_conf(launch_config,
-                                                     ray_config['auth'])
+            current_hash = ray_autoscaler_private_util.hash_launch_conf(
+                launch_config, auth_config)
         launch_hashes.add(current_hash)
     # Cache the launch hashes for the cluster.
     metadata['ray_launch_hashes'] = list(launch_hashes)
@@ -1690,18 +1726,22 @@ def _update_cluster_status_no_lock(
     # where the cluster is terminated by the user manually through the UI.
     to_terminate = not node_statuses
 
-    # A cluster is considered "abnormal", if not all nodes are TERMINATED or not all
-    # nodes are STOPPED. We check that with the following logic:
-    #   * not all nodes are terminated and there's at least one node terminated; or
-    #   * any of the non-TERMINATED nodes is in a non-STOPPED status.
+    # A cluster is considered "abnormal", if not all nodes are TERMINATED or
+    # not all nodes are STOPPED. We check that with the following logic:
+    #   * Not all nodes are terminated and there's at least one node
+    #     terminated; or
+    #   * Any of the non-TERMINATED nodes is in a non-STOPPED status.
     #
     # This includes these special cases:
-    # All stopped are considered normal and will be cleaned up at the end of the function.
-    # Some of the nodes UP should be considered abnormal, because the ray cluster is
-    # probably down.
-    # The cluster is partially terminated or stopped should be considered abnormal.
+    #   * All stopped are considered normal and will be cleaned up at the end
+    #     of the function.
+    #   * Some of the nodes UP should be considered abnormal, because the ray
+    #     cluster is probably down.
+    #   * The cluster is partially terminated or stopped should be considered
+    #     abnormal.
     #
-    # An abnormal cluster will transition to INIT and have any autostop setting reset.
+    # An abnormal cluster will transition to INIT and have any autostop setting
+    # reset.
     is_abnormal = ((0 < len(node_statuses) < handle.launched_nodes) or
                    any(status != global_user_state.ClusterStatus.STOPPED
                        for status in node_statuses))