Sky init and cloud credentials for storage

prototype/README.md

@@ -53,7 +53,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/backends/cloud_vm_ray_backend.py

@@ -981,17 +981,17 @@ def sync_file_mounts(
             else:
                 wrapped_dst = backend_utils.FileMountHelper.wrap_file_mount(dst)
             storage = cloud_stores.get_storage_from_path(src)
-            if storage.is_directory(src):
-                sync = storage.make_sync_dir_command(source=src,
-                                                     destination=wrapped_dst)
-                # It is a directory so make sure it exists.
-                mkdir_for_wrapped_dst = f'mkdir -p {wrapped_dst}'
-            else:
+            if storage.is_file(src):
                 sync = storage.make_sync_file_command(source=src,
                                                       destination=wrapped_dst)
                 # It is a file so make sure *its parent dir* exists.
                 mkdir_for_wrapped_dst = \
                     f'mkdir -p {os.path.dirname(wrapped_dst)}'
+            else:
+                sync = storage.make_sync_dir_command(source=src,
+                                                     destination=wrapped_dst)
+                # It is a directory so make sure it exists.
+                mkdir_for_wrapped_dst = f'mkdir -p {wrapped_dst}'
             download_target_commands = [
                 # Ensure sync can write to wrapped_dst (e.g., '/data/').
                 mkdir_for_wrapped_dst,

prototype/sky/cli.py

@@ -640,6 +640,38 @@ def tpunode(cluster: str, port_forward: Optional[List[int]], screen):
     )
 
 
+@cli.command()
+def init():
+    """Determines a set of clouds that Sky will use.
+
+    It checks access credentials for AWS, Azure and GCP. Sky jobs will only
+    run in clouds that you have access to. After configuring access for a
+    cloud, rerun `sky init` to reflect the changes.
+    """
+    click.echo('Sky will use the following clouds to run jobs. '
+               'To change this, configure cloud access credentials,'
+               ' and rerun ' + click.style('sky init', bold=True) + '.\n')
+
+    enabled_clouds = []
+    for cloud in [sky.AWS(), sky.Azure(), sky.GCP()]:
+        click.echo(f'  Checking {cloud}...', nl=False)
+        ok, reason = cloud.check_credentials()
+        click.echo('\r', nl=False)
+        status_msg = 'enabled' if ok else 'disabled'
+        status_color = 'green' if ok else 'red'
+        click.echo(
+            '  ' +
+            click.style(f'{cloud}: {status_msg}', fg=status_color, bold=True) +
+            ' ' * 10)
+        if ok:
+            enabled_clouds.append(str(cloud))
+        else:
+            click.echo(f'    Reason: {reason}')
+    click.echo()
+
+    global_user_state.set_enabled_clouds(enabled_clouds)
+
+
 def main():
     return cli()
 

prototype/sky/cloud_stores.py

@@ -6,13 +6,11 @@
 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
 
 import boto3
+import botocore
 
 from sky.backends import backend_utils
 from sky.data import data_utils
@@ -21,11 +19,11 @@
 class CloudStorage(object):
     """Interface for a cloud object store."""
 
-    def is_directory(self, url: str) -> bool:
-        """Returns whether 'url' is a directory.
+    def is_file(self, url: str) -> bool:
+        """Returns whether 'url' is a regular file.
 
-        In cloud object stores, a "directory" refers to a regular object whose
-        name is a prefix of other objects.
+        Returning false means <url> is a directory, or does not exist. Useful
+        for deciding whether to use cp or sync/rsync to download.
         """
         raise NotImplementedError
 
@@ -46,27 +44,17 @@ class S3CloudStorage(CloudStorage):
         'pip install awscli',
     ]
 
-    def is_directory(self, url: str) -> bool:
-        """Returns whether S3 'url' is a directory.
-
-        In cloud object stores, a "directory" refers to a regular object whose
-        name is a prefix of other objects.
-        """
-        s3 = boto3.resource('s3')
+    def is_file(self, url: str) -> bool:
+        """Returns whether 'url' is a regular file."""
         bucket_name, path = data_utils.split_s3_path(url)
-        bucket = s3.Bucket(bucket_name)
-
-        num_objects = 0
-        for obj in bucket.objects.filter(Prefix=path):
-            num_objects += 1
-            if obj.key == path:
-                return False
-            # If there are more than 1 object in filter, then it is a directory
-            if num_objects == 3:
-                return True
-
-        # A directory with few or no items
-        return True
+        if len(path) == 0:
+            return False
+        s3 = boto3.client('s3')
+        try:
+            s3.head_object(Bucket=bucket_name, Key=path)
+            return True
+        except botocore.errorfactory.ClientError:
+            return False
 
     def make_sync_dir_command(self, source: str, destination: str) -> str:
         """Downloads using AWS CLI."""
@@ -97,55 +85,41 @@ 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.
+    def is_file(self, url: str) -> bool:
+        """Returns whether 'url' is a regular file."""
+        command = ' && '.join(self._GET_GSUTIL)
+        backend_utils.run(command)
 
-        In cloud object stores, a "directory" refers to a regular object whose
-        name is a prefix of other objects.
-        """
-        commands = list(self._GET_GSUTIL)
-        commands.append(f'{self._GSUTIL} ls -d {url}')
-        command = ' && '.join(commands)
-        p = backend_utils.run(command, stdout=subprocess.PIPE)
-        out = p.stdout.decode().strip()
-        # gsutil ls -d url
-        #   --> url.rstrip('/')          if url is not a directory
-        #   --> url with an ending '/'   if url is a directory
-        if not out.endswith('/'):
-            assert out == url.rstrip('/'), (out, url)
-            return False
-        url = url if url.endswith('/') else (url + '/')
-        assert out == url, (out, url)
-        return True
+        # https://cloud.google.com/storage/docs/gsutil/commands/stat
+        gsutil_cmd = f'{self._GSUTIL} -q stat {url}'
+        p = backend_utils.run(gsutil_cmd, check=False)
+        rc = p.returncode
+        assert rc in [0, 1], command
+        return rc == 0
 
     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

@@ -14,3 +14,13 @@
     'Region',
     'Zone',
 ]
+
+__CLOUD_DICT__ = {
+    'AWS': AWS,
+    'Azure': Azure,
+    'GCP': GCP,
+}
+
+
+def cloud_factory(name):
+    return __CLOUD_DICT__[name]

prototype/sky/clouds/aws.py

@@ -1,12 +1,19 @@
 """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, stdout=subprocess.PIPE)
+    return proc.stdout.decode('ascii')
+
+
 class AWS(clouds.Cloud):
     """Amazon Web Services."""
 
@@ -160,3 +167,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."""
+        try:
+            assert os.path.isfile(os.path.expanduser('~/.aws/credentials'))
+            output = _run_output('aws configure list')
+        except (AssertionError, subprocess.CalledProcessError):
+            return False, 'AWS CLI not installed properly.'
+        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`.'

prototype/sky/clouds/azure.py

@@ -1,12 +1,18 @@
 """Azure."""
 import copy
 import json
+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, stdout=subprocess.PIPE)
+    return proc.stdout.decode('ascii')
+
+
 class Azure(clouds.Cloud):
     """Azure."""
 
@@ -128,3 +134,13 @@ 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."""
+        try:
+            output = _run_output('az account show --output=json')
+        except subprocess.CalledProcessError:
+            return False, 'Azure CLI returned error.'
+        if output.startswith('{'):
+            return True, None
+        return False, 'Azure credentials not set. Run `az login`.'

prototype/sky/clouds/cloud.py

@@ -110,3 +110,7 @@ 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."""
+        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
 
 
@@ -276,3 +279,17 @@ 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:
+            for file in [
+                    '~/.config/gcloud/access_tokens.db',
+                    '~/.config/gcloud/credentials.db'
+            ]:
+                assert os.path.isfile(os.path.expanduser(file))
+            auth.default()
+        except (AssertionError, auth.exceptions.DefaultCredentialsError):
+            return False, ('GCP credentials not set.'
+                           ' Run `gcloud auth application-default login`.')
+        return True, None

prototype/sky/data/storage.py

@@ -174,6 +174,11 @@ def __init__(self,
         # from existing ones
         self.stores = {} if stores is None else stores
 
+        if 's3://' in self.source:
+            self.get_or_copy_to_s3()
+        elif 'gs://' in self.source:
+            self.get_or_copy_to_gcs()
+
     def get_or_copy_to_s3(self):
         """Adds AWS S3 Store to Storage
         """