Repository: skypilot-org/skypilot
Pull request: [AWS] Avoid key pair permission issue by using cloud-init for authorized keys
Closes #1425 

This PR uses the cloud-init to setup the authorized keys to avoid using `ec2.describe_key_pairs` reducing the permission requirement.

Backward compatibility should be guaranteed by our old yaml content restoring mechanism.

Tested:
`rm ~/.ssh/sky-key*`
- [x] `sky launch --cloud aws` and ssh into the cluster
- [x] `sky launch --cloud aws --num-nodes 3` and ssh into the head and workers.
- [x] `./tests/run_smoke_tests.sh` (except #1408 )
- [x] `tests/backward_comaptibility_tests.sh`

Problem statement:
[AWS] Permission issues with DescribeKeyPairs
We require `DescribeKeyPairs` permission in `auth. setup_aws_authentication()`, which is not checked in `sky check` and can fail the `sky launch` when not available.

```
botocore.exceptions.ClientError: An error occurred (UnauthorizedOperation) when calling the DescribeKeyPairs operation: You are not authorized to perform this operation.
```

TODO:
- [ ] Check if we can get around that permission.
- [ ] If not, add the checking in `sky check`.

Hints from issue discussion:

Commit message: address comments

Patch to review:
diff --git a/sky/authentication.py b/sky/authentication.py
index 8eff3e482b4..784d8c397d9 100644
--- a/sky/authentication.py
+++ b/sky/authentication.py
@@ -1,14 +1,13 @@
 """Module to enable a single SkyPilot key for all VMs in each cloud."""
 import copy
 import functools
-import hashlib
 import os
-import pathlib
 import re
 import subprocess
 import sys
+import textwrap
 import time
-import uuid
+from typing import Any, Dict, Tuple
 
 import colorama
 from cryptography.hazmat.primitives import serialization
@@ -17,7 +16,7 @@
 
 from sky import clouds
 from sky import sky_logging
-from sky.adaptors import aws, gcp
+from sky.adaptors import gcp
 from sky.utils import common_utils
 from sky.utils import subprocess_utils
 from sky.utils import ux_utils
@@ -30,10 +29,12 @@
 # development.
 
 MAX_TRIALS = 64
+# TODO(zhwu): Support user specified key pair.
 PRIVATE_SSH_KEY_PATH = '~/.ssh/sky-key'
+PUBLIC_SSH_KEY_PATH = '~/.ssh/sky-key.pub'
 
 
-def generate_rsa_key_pair():
+def _generate_rsa_key_pair() -> Tuple[str, str]:
     key = rsa.generate_private_key(backend=default_backend(),
                                    public_exponent=65537,
                                    key_size=2048)
@@ -50,7 +51,8 @@ def generate_rsa_key_pair():
     return public_key, private_key
 
 
-def save_key_pair(private_key_path, public_key_path, private_key, public_key):
+def _save_key_pair(private_key_path: str, public_key_path: str,
+                   private_key: str, public_key: str) -> None:
     private_key_dir = os.path.dirname(private_key_path)
     os.makedirs(private_key_dir, exist_ok=True)
 
@@ -65,92 +67,42 @@ def save_key_pair(private_key_path, public_key_path, private_key, public_key):
         f.write(public_key)
 
 
-def get_public_key_path(private_key_path):
-    if private_key_path.endswith('.pem'):
-        private_key_path, _ = private_key_path.rsplit('.', 1)
-    return private_key_path + '.pub'
-
-
-def get_or_generate_keys(private_key_path: str, public_key_path: str):
-    """Returns private and public keys from the given paths.
-
-    If the private_key_path is not provided or does not exist, then a new
-    keypair is generated and written to the path.
-    """
+def get_or_generate_keys() -> Tuple[str, str]:
+    """Returns the abosulte public and private key paths."""
+    private_key_path = os.path.expanduser(PRIVATE_SSH_KEY_PATH)
+    public_key_path = os.path.expanduser(PUBLIC_SSH_KEY_PATH)
     if private_key_path is None or not os.path.exists(private_key_path):
-        public_key, private_key = generate_rsa_key_pair()
-        save_key_pair(private_key_path, public_key_path, private_key,
-                      public_key)
+        public_key, private_key = _generate_rsa_key_pair()
+        _save_key_pair(private_key_path, public_key_path, private_key,
+                       public_key)
     else:
         assert os.path.exists(public_key_path)
-        public_key = open(public_key_path, 'rb').read().decode('utf-8')
-        private_key = open(private_key_path, 'rb').read().decode('utf-8')
-    return private_key, public_key
-
-
-# Snippets of code inspired from
-# https://github.com/ray-project/ray/blob/master/python/ray/autoscaler/_private/aws/config.py
-# Takes in config, a yaml dict and outputs a postprocessed dict
-def setup_aws_authentication(config):
-    from Crypto.PublicKey import RSA  # pylint: disable=import-outside-toplevel
-    config = copy.deepcopy(config)
-    private_key_path = config['auth'].get('ssh_private_key', None)
-    if private_key_path is None:
-        private_key_path = PRIVATE_SSH_KEY_PATH
-        config['auth']['ssh_private_key'] = private_key_path
-
-    private_key_path = os.path.expanduser(private_key_path)
-    public_key_path = get_public_key_path(private_key_path)
-
-    # Generating ssh key if it does not exist
-    _, public_key = get_or_generate_keys(private_key_path, public_key_path)
-
-    ec2 = aws.client('ec2', region_name=config['provider']['region'])
-    key_pairs = ec2.describe_key_pairs()['KeyPairs']
-    key_name = None
-    all_key_names = set()
-
-    def _get_fingerprint(public_key_path):
-        key = RSA.importKey(open(public_key_path).read())
-
-        def insert_char_every_n_chars(string, char='\n', every=2):
-            return char.join(
-                string[i:i + every] for i in range(0, len(string), every))
-
-        md5digest = hashlib.md5(key.exportKey('DER', pkcs=8)).hexdigest()
-        fingerprint = insert_char_every_n_chars(md5digest, ':', 2)
-        return fingerprint
-
-    for key in key_pairs:
-        # Compute Fingerprint of public key
-        aws_fingerprint = key['KeyFingerprint']
-        local_fingerprint = _get_fingerprint(public_key_path)
-        if aws_fingerprint == local_fingerprint:
-            key_name = key['KeyName']
-        # Add key name to key name list
-        all_key_names.add(key['KeyName'])
-
-    if key_name is None:
-        for fail_counter in range(MAX_TRIALS):
-            key_name = 'sky-key-' + uuid.uuid4().hex[:6]
-            if key_name not in all_key_names:
-                ec2.import_key_pair(KeyName=key_name,
-                                    PublicKeyMaterial=public_key)
-                break
-        if fail_counter == MAX_TRIALS - 1:
-            raise RuntimeError(
-                'Failed to generate a unique key pair ID for AWS')
-
-    node_types = config['available_node_types']
-
-    for node_type in node_types.values():
-        node_type['node_config']['KeyName'] = key_name
+    return private_key_path, public_key_path
+
+
+def setup_aws_authentication(config: Dict[str, Any]) -> Dict[str, Any]:
+    _, public_key_path = get_or_generate_keys()
+    with open(public_key_path, 'r') as f:
+        public_key = f.read()
+    # Use cloud init in UserData to set up the authorized_keys to get
+    # around the number of keys limit and permission issues with
+    # ec2.describe_key_pairs.
+    for node_type in config['available_node_types']:
+        config['available_node_types'][node_type]['node_config']['UserData'] = (
+            textwrap.dedent(f"""\
+            #cloud-config
+            users:
+            - name: {config['auth']['ssh_user']}
+              ssh-authorized-keys:
+                - {public_key}
+            """))
     return config
 
 
 # Reference:
 # https://github.com/ray-project/ray/blob/master/python/ray/autoscaler/_private/gcp/config.py
-def _wait_for_compute_global_operation(project_name, operation_name, compute):
+def _wait_for_compute_global_operation(project_name: str, operation_name: str,
+                                       compute: Any) -> None:
     """Poll for global compute operation until finished."""
     logger.debug('wait_for_compute_global_operation: '
                  'Waiting for operation {} to finish...'.format(operation_name))
@@ -180,15 +132,8 @@ def _wait_for_compute_global_operation(project_name, operation_name, compute):
 # Retry for the GCP as sometimes there will be connection reset by peer error.
 @common_utils.retry
 @gcp.import_package
-def setup_gcp_authentication(config):
-    config = copy.deepcopy(config)
-    private_key_path = config['auth'].get('ssh_private_key', None)
-    if private_key_path is None:
-        private_key_path = PRIVATE_SSH_KEY_PATH
-        config['auth']['ssh_private_key'] = private_key_path
-
-    private_key_path = os.path.expanduser(private_key_path)
-    public_key_path = get_public_key_path(private_key_path)
+def setup_gcp_authentication(config: Dict[str, Any]) -> Dict[str, Any]:
+    private_key_path, public_key_path = get_or_generate_keys()
     config = copy.deepcopy(config)
 
     project_id = config['provider']['project_id']
@@ -261,9 +206,6 @@ def setup_gcp_authentication(config):
                         'account information.')
         config['auth']['ssh_user'] = account.replace('@', '_').replace('.', '_')
 
-        # Generating ssh key if it does not exist
-        get_or_generate_keys(private_key_path, public_key_path)
-
         # Add ssh key to GCP with oslogin
         subprocess.run(
             'gcloud compute os-login ssh-keys add '
@@ -292,13 +234,16 @@ def setup_gcp_authentication(config):
 
     # OS Login is not enabled for the project. Add the ssh key directly to the
     # metadata.
+    # TODO(zhwu): Use cloud init to add ssh public key, to avoid the permission
+    # issue.
     project_keys = next(
         (item for item in project['commonInstanceMetadata'].get('items', [])
          if item['key'] == 'ssh-keys'), {}).get('value', '')
     ssh_keys = project_keys.split('\n') if project_keys else []
 
-    # Generating ssh key if it does not exist
-    _, public_key = get_or_generate_keys(private_key_path, public_key_path)
+    # Get public key from file.
+    with open(public_key_path, 'r') as f:
+        public_key = f.read()
 
     # Check if ssh key in Google Project's metadata
     public_key_token = public_key.split(' ')[1]
@@ -338,34 +283,14 @@ def setup_gcp_authentication(config):
     return config
 
 
-def _unexpand_user(path):
-    """Inverse of `os.path.expanduser`."""
-    return ('~' /
-            pathlib.Path(path).relative_to(pathlib.Path.home())).as_posix()
-
-
-# Takes in config, a yaml dict and outputs a postprocessed dict
-def setup_azure_authentication(config):
-    # Doesn't need special library calls!
-    config = copy.deepcopy(config)
-    private_key_path = config['auth'].get('ssh_private_key', None)
-    if private_key_path is None:
-        private_key_path = PRIVATE_SSH_KEY_PATH
-        config['auth']['ssh_private_key'] = private_key_path
-
-    private_key_path = os.path.expanduser(private_key_path)
-    public_key_path = get_public_key_path(private_key_path)
-
-    # Generating ssh key if it does not exist
-    _, _ = get_or_generate_keys(private_key_path, public_key_path)
-
-    # Need to convert /Users/<username> back to ~ because Ray uses the same
+def setup_azure_authentication(config: Dict[str, Any]) -> Dict[str, Any]:
+    get_or_generate_keys()
+    # Need to use ~ relative path because Ray uses the same
     # path for finding the public key path on both local and head node.
-    public_key_path = _unexpand_user(public_key_path)
-    config['auth']['ssh_public_key'] = public_key_path
+    config['auth']['ssh_public_key'] = PUBLIC_SSH_KEY_PATH
 
     file_mounts = config['file_mounts']
-    file_mounts[public_key_path] = public_key_path
+    file_mounts[PUBLIC_SSH_KEY_PATH] = PUBLIC_SSH_KEY_PATH
     config['file_mounts'] = file_mounts
 
     return config
diff --git a/sky/backends/backend_utils.py b/sky/backends/backend_utils.py
index ed91265fe2e..ce053b1a5a3 100644
--- a/sky/backends/backend_utils.py
+++ b/sky/backends/backend_utils.py
@@ -746,7 +746,7 @@ def write_cluster_config(
     credentials = sky_check.get_cloud_credential_file_mounts()
 
     ip_list = None
-    auth_config = None
+    auth_config = {'ssh_private_key': auth.PRIVATE_SSH_KEY_PATH}
     if isinstance(cloud, clouds.Local):
         ip_list = onprem_utils.get_local_ips(cluster_name)
         auth_config = onprem_utils.get_local_auth_config(cluster_name)
@@ -794,10 +794,7 @@ def write_cluster_config(
                 'head_ip': None if ip_list is None else ip_list[0],
                 'worker_ips': None if ip_list is None else ip_list[1:],
                 # Authentication (optional).
-                'ssh_user':
-                    (None if auth_config is None else auth_config['ssh_user']),
-                'ssh_private_key': (None if auth_config is None else
-                                    auth_config['ssh_private_key']),
+                **auth_config,
             }),
         output_path=tmp_yaml_path)
     config_dict['cluster_name'] = cluster_name
diff --git a/sky/templates/aws-ray.yml.j2 b/sky/templates/aws-ray.yml.j2
index 1ba5cc3ebaf..4646ff12241 100644
--- a/sky/templates/aws-ray.yml.j2
+++ b/sky/templates/aws-ray.yml.j2
@@ -19,6 +19,7 @@ provider:
 
 auth:
   ssh_user: ubuntu
+  ssh_private_key: {{ssh_private_key}}
 
 available_node_types:
   ray.head.default:
diff --git a/sky/templates/azure-ray.yml.j2 b/sky/templates/azure-ray.yml.j2
index ab3de7405c4..3d30ad7219a 100644
--- a/sky/templates/azure-ray.yml.j2
+++ b/sky/templates/azure-ray.yml.j2
@@ -22,6 +22,7 @@ provider:
 
 auth:
     ssh_user: azureuser
+    ssh_private_key: {{ssh_private_key}}
 
 available_node_types:
     ray.head.default:
diff --git a/sky/templates/gcp-ray.yml.j2 b/sky/templates/gcp-ray.yml.j2
index 620e2f008cb..255230ee8b7 100644
--- a/sky/templates/gcp-ray.yml.j2
+++ b/sky/templates/gcp-ray.yml.j2
@@ -23,6 +23,7 @@ provider:
 
 auth:
   ssh_user: gcpuser
+  ssh_private_key: {{ssh_private_key}}
 
 available_node_types:
   ray_head_default:
diff --git a/tests/backward_comaptibility_tests.sh b/tests/backward_comaptibility_tests.sh
index 228c4e4cee8..769a89b7055 100755
--- a/tests/backward_comaptibility_tests.sh
+++ b/tests/backward_comaptibility_tests.sh
@@ -8,25 +8,22 @@ source ~/.bashrc
 CLUSTER_NAME="test-back-compat-$USER"
 source $(conda info --base 2> /dev/null)/etc/profile.d/conda.sh
 
-git clone https://github.com/skypilot-org/skypilot.git sky-master || true
+git clone https://github.com/skypilot-org/skypilot.git ../sky-master || true
 
 
 # Create environment for compatibility tests
-mamba > /dev/null 2>&1 || conda install -n base -c conda-forge mamba -y
-mamba init
-source ~/.bashrc
-mamba env list | grep sky-back-compat-master || mamba create -n sky-back-compat-master -y python=3.9
+conda env list | grep sky-back-compat-master || conda create -n sky-back-compat-master -y python=3.9
 
-mamba activate sky-back-compat-master
-mamba install -c conda-forge google-cloud-sdk -y
+conda activate sky-back-compat-master
+conda install -c conda-forge google-cloud-sdk -y
 rm -r  ~/.sky/wheels || true
-cd sky-master
+cd ../sky-master
 git pull origin master
 pip install -e ".[all]"
 cd -
 
-mamba env list | grep sky-back-compat-current || mamba create -n sky-back-compat-current -y --clone sky-back-compat-master
-mamba activate sky-back-compat-current
+conda env list | grep sky-back-compat-current || conda create -n sky-back-compat-current -y python=3.9
+conda activate sky-back-compat-current
 rm -r  ~/.sky/wheels || true
 pip uninstall -y skypilot
 pip install -e ".[all]"
@@ -34,18 +31,18 @@ pip install -e ".[all]"
 
 # exec + launch
 if [ "$start_from" -le 1 ]; then
-mamba activate sky-back-compat-master
+conda activate sky-back-compat-master
 rm -r  ~/.sky/wheels || true
 which sky
-sky launch --cloud gcp -y -c ${CLUSTER_NAME} examples/minimal.yaml
+sky launch -y -c ${CLUSTER_NAME} examples/minimal.yaml
 
-mamba activate sky-back-compat-current
+conda activate sky-back-compat-current
 rm -r  ~/.sky/wheels || true
 if [ "$need_launch" -eq "1" ]; then
-  sky launch --cloud gcp -y -c ${CLUSTER_NAME}
+  sky launch -y -c ${CLUSTER_NAME}
 fi
-sky exec --cloud gcp ${CLUSTER_NAME} examples/minimal.yaml
-s=$(sky launch --cloud gcp -d -c ${CLUSTER_NAME} examples/minimal.yaml)
+sky exec ${CLUSTER_NAME} examples/minimal.yaml
+s=$(sky launch -d -c ${CLUSTER_NAME} examples/minimal.yaml)
 echo $s
 # remove color and find the job id
 echo $s | sed -r "s/\x1B\[([0-9]{1,3}(;[0-9]{1,2})?)?[mGK]//g" | grep "Job ID: 3" || exit 1
@@ -54,24 +51,24 @@ fi
 
 # sky stop + sky start + sky exec
 if [ "$start_from" -le 2 ]; then
-mamba activate sky-back-compat-master
+conda activate sky-back-compat-master
 rm -r  ~/.sky/wheels || true
-sky launch --cloud gcp -y -c ${CLUSTER_NAME}-2 examples/minimal.yaml
-mamba activate sky-back-compat-current
+sky launch -y -c ${CLUSTER_NAME}-2 examples/minimal.yaml
+conda activate sky-back-compat-current
 rm -r  ~/.sky/wheels || true
 sky stop -y ${CLUSTER_NAME}-2
 sky start -y ${CLUSTER_NAME}-2
-s=$(sky exec --cloud gcp -d ${CLUSTER_NAME}-2 examples/minimal.yaml)
+s=$(sky exec -d ${CLUSTER_NAME}-2 examples/minimal.yaml)
 echo $s
 echo $s | sed -r "s/\x1B\[([0-9]{1,3}(;[0-9]{1,2})?)?[mGK]//g" | grep "Job ID: 2" || exit 1
 fi
 
 # `sky autostop` + `sky status -r`
 if [ "$start_from" -le 3 ]; then
-mamba activate sky-back-compat-master
+conda activate sky-back-compat-master
 rm -r  ~/.sky/wheels || true
-sky launch --cloud gcp -y -c ${CLUSTER_NAME}-3 examples/minimal.yaml
-mamba activate sky-back-compat-current
+sky launch -y -c ${CLUSTER_NAME}-3 examples/minimal.yaml
+conda activate sky-back-compat-current
 rm -r  ~/.sky/wheels || true
 sky autostop -y -i0 ${CLUSTER_NAME}-3
 sleep 120
@@ -81,13 +78,13 @@ fi
 
 # (1 node) sky launch + sky exec + sky queue + sky logs
 if [ "$start_from" -le 4 ]; then
-mamba activate sky-back-compat-master
+conda activate sky-back-compat-master
 rm -r  ~/.sky/wheels || true
-sky launch --cloud gcp -y -c ${CLUSTER_NAME}-4 examples/minimal.yaml
+sky launch -y -c ${CLUSTER_NAME}-4 examples/minimal.yaml
 sky stop -y ${CLUSTER_NAME}-4
-mamba activate sky-back-compat-current
+conda activate sky-back-compat-current
 rm -r  ~/.sky/wheels || true
-sky launch --cloud gcp -y -c ${CLUSTER_NAME}-4 examples/minimal.yaml
+sky launch -y -c ${CLUSTER_NAME}-4 examples/minimal.yaml
 sky queue ${CLUSTER_NAME}-4
 sky logs ${CLUSTER_NAME}-4 1 --status
 sky logs ${CLUSTER_NAME}-4 2 --status
@@ -97,17 +94,17 @@ fi
 
 # (1 node) sky start + sky exec + sky queue + sky logs
 if [ "$start_form" -le 5 ]; then
-mamba activate sky-back-compat-master
+conda activate sky-back-compat-master
 rm -r  ~/.sky/wheels || true
-sky launch --cloud gcp -y -c ${CLUSTER_NAME}-5 examples/minimal.yaml
+sky launch -y -c ${CLUSTER_NAME}-5 examples/minimal.yaml
 sky stop -y ${CLUSTER_NAME}-5
-mamba activate sky-back-compat-current
+conda activate sky-back-compat-current
 rm -r  ~/.sky/wheels || true
 sky start -y ${CLUSTER_NAME}-5
 sky queue ${CLUSTER_NAME}-5
 sky logs ${CLUSTER_NAME}-5 1 --status
 sky logs ${CLUSTER_NAME}-5 1
-sky launch --cloud gcp -y -c ${CLUSTER_NAME}-5 examples/minimal.yaml
+sky launch -y -c ${CLUSTER_NAME}-5 examples/minimal.yaml
 sky queue ${CLUSTER_NAME}-5
 sky logs ${CLUSTER_NAME}-5 2 --status
 sky logs ${CLUSTER_NAME}-5 2
@@ -115,17 +112,17 @@ fi
 
 # (2 nodes) sky launch + sky exec + sky queue + sky logs
 if [ "$start_from" -le 6 ]; then
-mamba activate sky-back-compat-master
+conda activate sky-back-compat-master
 rm -r  ~/.sky/wheels || true
-sky launch --cloud gcp -y -c ${CLUSTER_NAME}-6 examples/multi_hostname.yaml
+sky launch -y -c ${CLUSTER_NAME}-6 examples/multi_hostname.yaml
 sky stop -y ${CLUSTER_NAME}-6
-mamba activate sky-back-compat-current
+conda activate sky-back-compat-current
 rm -r  ~/.sky/wheels || true
 sky start -y ${CLUSTER_NAME}-6
 sky queue ${CLUSTER_NAME}-6
 sky logs ${CLUSTER_NAME}-6 1 --status
 sky logs ${CLUSTER_NAME}-6 1
-sky exec --cloud gcp ${CLUSTER_NAME}-6 examples/multi_hostname.yaml
+sky exec ${CLUSTER_NAME}-6 examples/multi_hostname.yaml
 sky queue ${CLUSTER_NAME}-6
 sky logs ${CLUSTER_NAME}-6 2 --status
 sky logs ${CLUSTER_NAME}-6 2


Code review comments:
diff_hunk: @@ -65,92 +67,42 @@ def save_key_pair(private_key_path, public_key_path, private_key, public_key):
         f.write(public_key)
 
 
-def get_public_key_path(private_key_path):
-    if private_key_path.endswith('.pem'):
-        private_key_path, _ = private_key_path.rsplit('.', 1)
-    return private_key_path + '.pub'
-
-
-def get_or_generate_keys(private_key_path: str, public_key_path: str):
-    """Returns private and public keys from the given paths.
-
-    If the private_key_path is not provided or does not exist, then a new
-    keypair is generated and written to the path.
-    """
+def get_or_generate_keys() -> Tuple[str, str]:
+    """Returns the abosulte public and private key paths."""
+    private_key_path = os.path.expanduser(PRIVATE_SSH_KEY_PATH)
+    public_key_path = os.path.expanduser(PUBLIC_SSH_KEY_PATH)
     if private_key_path is None or not os.path.exists(private_key_path):
path: sky/authentication.py
text: @user1:
nit: `private_key_path is None` can't seem to happen.

@author:
Good catch! Fixed. Thanks!

diff_hunk: @@ -65,27 +67,23 @@ def save_key_pair(private_key_path, public_key_path, private_key, public_key):
         f.write(public_key)
 
 
-def get_or_generate_keys():
-    """Returns private and public keys from the given paths.
-
-    If the private_key_path is not provided or does not exist, then a new
-    keypair is generated and written to the path.
-    """
+def get_or_generate_keys() -> Tuple[str, str]:
+    """Returns the abosulte public and private key paths."""
path: sky/authentication.py
text: @user1:
nit: typo - `absolute`

@author:
Good catch! Fixed. Thanks!