Sky init and cloud credentials for storage

prototype/README.md

@@ -55,7 +55,7 @@ ray attach config/gcp.yml
 ray down config/gcp.yml
 ```
 
-**Azure**. Install the Azure CLI (`pip install azure-cli`) then login using `az login`. Set the subscription to use from the command line (`az account set -s <subscription_id>`) or by modifying the provider section of the Azure template (`config/azure.yml.j2`). Ray Autoscaler does not work with the latest version of `azure-cli`. Hotfix: `pip install azure-cli-core==2.22.0` (this will make Ray work but at the cost of making the `az` CLI tool unusable).
+**Azure**. Install the Azure CLI (`pip install azure-cli==2.22.0`) then login using `az login`. Set the subscription to use from the command line (`az account set -s <subscription_id>`). Ray Autoscaler does not work with the latest version of `azure-cli` as of 1.9.1, hence the fixed Azure version.
 
 ## Open issues
 

prototype/requirements.txt

@@ -14,6 +14,6 @@ tabulate
 docker
 
 awscli==1.22.17
-azure-cli
+azure-cli==2.22.0
 google-api-python-client
 google-cloud-storage

prototype/sky/__init__.py

@@ -9,7 +9,6 @@
 from sky.execution import execute
 from sky.resources import Resources
 from sky.task import ParTask, Task
-from sky.registry import fill_in_launchable_resources
 from sky.optimizer import Optimizer, OptimizeTarget
 from sky.data import Storage, StorageType
 
@@ -34,7 +33,6 @@
     'Task',
     'backends',
     'execute',
-    'fill_in_launchable_resources',
     'list_accelerators',
     '__root_dir__',
     'Storage',

prototype/sky/cli.py

@@ -40,6 +40,7 @@
 import sky
 from sky import backends
 from sky import global_user_state
+from sky import init as sky_init
 from sky import task as task_lib
 from sky.backends import backend as backend_lib
 from sky.backends import backend_utils
@@ -208,6 +209,7 @@ def _create_and_ssh_into_node(
     dag = sky.optimize(dag)
     task = dag.tasks[0]
     backend.register_info(dag=dag)
+    task.update_file_mounts(sky_init.get_cloud_credential_file_mounts())
 
     handle = global_user_state.get_handle_from_cluster_name(cluster_name)
     if handle is None or handle.head_ip is None:
@@ -793,6 +795,17 @@ def tpunode(cluster: str, port_forward: Optional[List[int]],
     )
 
 
+@cli.command()
+def init():
+    """Determines a set of clouds that Sky will use.
+
+    It checks access credentials for AWS, Azure and GCP. Sky tasks will only
+    run in clouds that you have access to. After configuring access for a
+    cloud, rerun `sky init` to reflect the changes.
+    """
+    sky_init.init()
+
+
 def main():
     return cli()
 

prototype/sky/cloud_stores.py

@@ -6,8 +6,6 @@
 TODO:
 * Better interface.
 * Better implementation (e.g., fsspec, smart_open, using each cloud's SDK).
-  The full-blown impl should handle authentication so each user's private
-  datasets can be accessed.
 """
 import subprocess
 import urllib.parse
@@ -97,15 +95,18 @@ class GcsCloudStorage(CloudStorage):
     # multi-threaded download is nice, which frees us from implementing
     # parellel workers on our end.
     _GET_GSUTIL = [
+        # Skip if gsutil already exists.
         'pushd /tmp &>/dev/null',
-        # Skip if /tmp/gsutil already exists.
-        '(test -f /tmp/gsutil/gsutil || (wget --quiet '
-        'https://storage.googleapis.com/pub/gsutil.tar.gz && '
-        'tar xzf gsutil.tar.gz))',
+        '(test -f ~/google-cloud-sdk/bin/gsutil || (wget --quiet '
+        'https://dl.google.com/dl/cloudsdk/channels/rapid/downloads/'
+        'google-cloud-sdk-367.0.0-linux-x86_64.tar.gz && '
+        'tar xzf google-cloud-sdk-367.0.0-linux-x86_64.tar.gz && '
+        'mv google-cloud-sdk ~/ && '
+        '~/google-cloud-sdk/install.sh -q ))',
         'popd &>/dev/null',
     ]
 
-    _GSUTIL = '/tmp/gsutil/gsutil'
+    _GSUTIL = '~/google-cloud-sdk/bin/gsutil'
 
     def is_directory(self, url: str) -> bool:
         """Returns whether 'url' is a directory.
@@ -118,7 +119,12 @@ def is_directory(self, url: str) -> bool:
         command = ' && '.join(commands)
         p = backend_utils.run(command, stdout=subprocess.PIPE)
         out = p.stdout.decode().strip()
-        # gsutil ls -d url
+        # If <url> is a bucket root, then we only need `gsutil` to succeed
+        # to make sure the bucket exists. It is already a directory.
+        _, key = data_utils.split_gcs_path(url)
+        if len(key) == 0:
+            return True
+        # Otherwise, gsutil ls -d url will return:
         #   --> url.rstrip('/')          if url is not a directory
         #   --> url with an ending '/'   if url is a directory
         if not out.endswith('/'):
@@ -129,23 +135,15 @@ def is_directory(self, url: str) -> bool:
         return True
 
     def make_sync_dir_command(self, source: str, destination: str) -> str:
-        """Downloads a directory using gsutil.
-
-        Limitation: no authentication support; 'source' is assumed to in a
-        publicly accessible bucket.
-        """
+        """Downloads a directory using gsutil."""
         download_via_gsutil = (
             f'{self._GSUTIL} -m rsync -d -r {source} {destination}')
         all_commands = list(self._GET_GSUTIL)
         all_commands.append(download_via_gsutil)
         return ' && '.join(all_commands)
 
     def make_sync_file_command(self, source: str, destination: str) -> str:
-        """Downloads a file using gsutil.
-
-        Limitation: no authentication support; 'source' is assumed to in a
-        publicly accessible bucket.
-        """
+        """Downloads a file using gsutil."""
         download_via_gsutil = f'{self._GSUTIL} -m cp {source} {destination}'
         all_commands = list(self._GET_GSUTIL)
         all_commands.append(download_via_gsutil)

prototype/sky/clouds/__init__.py

@@ -7,10 +7,24 @@
 from sky.clouds.gcp import GCP
 
 __all__ = [
+    'ALL_CLOUDS',
     'AWS',
     'Azure',
     'Cloud',
     'GCP',
     'Region',
     'Zone',
+    'from_str',
 ]
+
+__CLOUD_DICT__ = {
+    'AWS': AWS(),
+    'Azure': Azure(),
+    'GCP': GCP(),
+}
+
+ALL_CLOUDS = list(__CLOUD_DICT__.values())
+
+
+def from_str(name: str) -> 'Cloud':
+    return __CLOUD_DICT__[name]

prototype/sky/clouds/aws.py

@@ -1,12 +1,23 @@
 """Amazon Web Services."""
 import copy
 import json
+import os
+import subprocess
 from typing import Dict, Iterator, List, Optional, Tuple
 
 from sky import clouds
 from sky.clouds.service_catalog import aws_catalog
 
 
+def _run_output(cmd):
+    proc = subprocess.run(cmd,
+                          shell=True,
+                          check=True,
+                          stderr=subprocess.PIPE,
+                          stdout=subprocess.PIPE)
+    return proc.stdout.decode('ascii')
+
+
 class AWS(clouds.Cloud):
     """Amazon Web Services."""
 
@@ -176,3 +187,41 @@ def _make(instance_type):
         if instance_type is None:
             return []
         return _make(instance_type)
+
+    def check_credentials(self) -> Tuple[bool, Optional[str]]:
+        """Checks if the user has access credentials to this cloud."""
+        # This file is required because it will be synced to remote VMs for
+        # `aws` to access private storage buckets.
+        # `aws configure list` does not guarantee this file exists.
+        if not os.path.isfile(os.path.expanduser('~/.aws/credentials')):
+            return (False,
+                    '~/.aws/credentials does not exist. Run `aws configure`.')
+        try:
+            output = _run_output('aws configure list')
+        except subprocess.CalledProcessError:
+            return False, 'AWS CLI not installed properly.'
+        # Configured correctly, the AWS output should look like this:
+        #   ...
+        #   access_key     ******************** shared-credentials-file
+        #   secret_key     ******************** shared-credentials-file
+        #   ...
+        # Otherwise, one or both keys will show as '<not set>'.
+        lines = output.split('\n')
+        if len(lines) < 2:
+            return False, 'AWS CLI output invalid.'
+        access_key_ok = False
+        secret_key_ok = False
+        for line in lines[2:]:
+            line = line.lstrip()
+            if line.startswith('access_key'):
+                if '<not set>' not in line:
+                    access_key_ok = True
+            elif line.startswith('secret_key'):
+                if '<not set>' not in line:
+                    secret_key_ok = True
+        if access_key_ok and secret_key_ok:
+            return True, None
+        return False, 'AWS credentials not set. Run `aws configure`.'
+
+    def get_credential_file_mounts(self) -> Dict[str, str]:
+        return {'~/.aws': '~/.aws'}

prototype/sky/clouds/azure.py

@@ -1,12 +1,23 @@
 """Azure."""
 import copy
 import json
+import os
+import subprocess
 from typing import Dict, Iterator, List, Optional, Tuple
 
 from sky import clouds
 from sky.clouds.service_catalog import azure_catalog
 
 
+def _run_output(cmd):
+    proc = subprocess.run(cmd,
+                          shell=True,
+                          check=True,
+                          stderr=subprocess.PIPE,
+                          stdout=subprocess.PIPE)
+    return proc.stdout.decode('ascii')
+
+
 class Azure(clouds.Cloud):
     """Azure."""
 
@@ -141,3 +152,27 @@ def _make(instance_type):
         if instance_type is None:
             return []
         return _make(instance_type)
+
+    def check_credentials(self) -> Tuple[bool, Optional[str]]:
+        """Checks if the user has access credentials to this cloud."""
+        # This file is required because it will be synced to remote VMs for
+        # `az` to access private storage buckets.
+        # `az account show` does not guarantee this file exists.
+        if not os.path.isfile(os.path.expanduser('~/.azure/accessTokens.json')):
+            return (
+                False,
+                '~/.azure/accessTokens.json does not exist. Run `az login`.')
+        try:
+            output = _run_output('az account show --output=json')
+        except subprocess.CalledProcessError:
+            return False, 'Azure CLI returned error.'
+        # If Azure is properly logged in, this will return something like:
+        #   {"id": ..., "user": ...}
+        # and if not, it will return:
+        #   Please run 'az login' to setup account.
+        if output.startswith('{'):
+            return True, None
+        return False, 'Azure credentials not set. Run `az login`.'
+
+    def get_credential_file_mounts(self) -> Dict[str, str]:
+        return {'~/.azure': '~/.azure'}

prototype/sky/clouds/cloud.py

@@ -125,3 +125,17 @@ def get_feasible_launchable_resources(self, resources):
         Launchable resources require a cloud and an instance type be assigned.
         """
         raise NotImplementedError
+
+    def check_credentials(self) -> Tuple[bool, Optional[str]]:
+        """Checks if the user has access credentials to this cloud.
+
+        Returns a boolean of whether the user can access this cloud, and a
+        string describing the reason if the user cannot access.
+        """
+        raise NotImplementedError
+
+    def get_credential_file_mounts(self) -> Dict[str, str]:
+        """Returns the files necessary to access this cloud.
+
+        Returns a dictionary that will be added to a task's file mounts."""
+        raise NotImplementedError

prototype/sky/clouds/gcp.py

@@ -1,8 +1,11 @@
 """Google Cloud Platform."""
 import copy
 import json
+import os
 from typing import Dict, Iterator, List, Optional, Tuple
 
+from google import auth
+
 from sky import clouds
 
 
@@ -285,3 +288,32 @@ def get_accelerators_from_instance_type(
             instance_type: str,
     ) -> Optional[Dict[str, int]]:
         return None
+
+    def check_credentials(self) -> Tuple[bool, Optional[str]]:
+        """Checks if the user has access credentials to this cloud."""
+        try:
+            # These files are required because they will be synced to remote
+            # VMs for `gsutil` to access private storage buckets.
+            # `auth.default()` does not guarantee these files exist.
+            for file in [
+                    '~/.config/gcloud/access_tokens.db',
+                    '~/.config/gcloud/credentials.db'
+            ]:
+                assert os.path.isfile(os.path.expanduser(file))
+            # Calling `auth.default()` ensures the GCP client library works,
+            # which is used by Ray Autoscaler to launch VMs.
+            auth.default()
+        except (AssertionError, auth.exceptions.DefaultCredentialsError):
+            return False, (
+                'GCP credentials not set. Run the following commands:\n    '
+                # This authenticates the CLI to make `gsutil` work:
+                '$ gcloud auth login\n    '
+                '$ gcloud config set project <proj>\n    '
+                # These two commands setup the client library to make
+                # Ray Autoscaler work:
+                '$ gcloud auth application-default login\n    '
+                '$ gcloud auth application-default set-quota-project <proj>')
+        return True, None
+
+    def get_credential_file_mounts(self) -> Dict[str, str]:
+        return {'~/.config/gcloud': '~/.config/gcloud'}

prototype/sky/execution.py

@@ -19,6 +19,7 @@
 
 import sky
 from sky import backends
+from sky import init
 from sky import logging
 from sky import optimizer
 
@@ -86,6 +87,8 @@ def execute(dag: sky.Dag,
     backend = backend if backend is not None else backends.CloudVmRayBackend()
     backend.register_info(dag=dag, optimize_target=optimize_target)
 
+    task.update_file_mounts(init.get_cloud_credential_file_mounts())
+
     if task.storage_mounts is not None:
         # Optimizer should eventually choose where to store bucket
         task.add_storage_mounts()