skypilot-nightly 1.0.0.dev2024053101__py3-none-any.whl → 1.0.0.dev2025022801__py3-none-any.whl

sky/cloud_stores.py

@@ -7,15 +7,26 @@ TODO:
 * Better interface.
 * Better implementation (e.g., fsspec, smart_open, using each cloud's SDK).
 """
+import os
+import shlex
 import subprocess
+import time
 import urllib.parse
 
+from sky import exceptions as sky_exceptions
+from sky import sky_logging
 from sky.adaptors import aws
+from sky.adaptors import azure
 from sky.adaptors import cloudflare
 from sky.adaptors import ibm
+from sky.adaptors import oci
 from sky.clouds import gcp
 from sky.data import data_utils
 from sky.data.data_utils import Rclone
+from sky.skylet import constants
+from sky.utils import ux_utils
+
+logger = sky_logging.init_logger(__name__)
 
 
 class CloudStorage:
@@ -43,7 +54,8 @@ class S3CloudStorage(CloudStorage):
 
     # List of commands to install AWS CLI
     _GET_AWSCLI = [
-        'aws --version >/dev/null 2>&1 || pip3 install awscli',
+        'aws --version >/dev/null 2>&1 || '
+        f'{constants.SKY_UV_PIP_CMD} install awscli',
     ]
 
     def is_directory(self, url: str) -> bool:
@@ -73,7 +85,8 @@ class S3CloudStorage(CloudStorage):
         # AWS Sync by default uses 10 threads to upload files to the bucket.
         # To increase parallelism, modify max_concurrent_requests in your
         # aws config file (Default path: ~/.aws/config).
-        download_via_awscli = ('aws s3 sync --no-follow-symlinks '
+        download_via_awscli = (f'{constants.SKY_REMOTE_PYTHON_ENV}/bin/aws s3 '
+                               'sync --no-follow-symlinks '
                                f'{source} {destination}')
 
         all_commands = list(self._GET_AWSCLI)
@@ -82,7 +95,8 @@ class S3CloudStorage(CloudStorage):
 
     def make_sync_file_command(self, source: str, destination: str) -> str:
         """Downloads a file using AWS CLI."""
-        download_via_awscli = f'aws s3 cp {source} {destination}'
+        download_via_awscli = (f'{constants.SKY_REMOTE_PYTHON_ENV}/bin/aws s3 '
+                               f'cp {source} {destination}')
 
         all_commands = list(self._GET_AWSCLI)
         all_commands.append(download_via_awscli)
@@ -102,8 +116,16 @@ class GcsCloudStorage(CloudStorage):
     @property
     def _gsutil_command(self):
         gsutil_alias, alias_gen = data_utils.get_gsutil_command()
-        return (f'{alias_gen}; GOOGLE_APPLICATION_CREDENTIALS='
-                f'{gcp.DEFAULT_GCP_APPLICATION_CREDENTIAL_PATH} {gsutil_alias}')
+        return (
+            f'{alias_gen}; GOOGLE_APPLICATION_CREDENTIALS='
+            f'{gcp.DEFAULT_GCP_APPLICATION_CREDENTIAL_PATH}; '
+            # Explicitly activate service account. Unlike the gcp packages
+            # and other GCP commands, gsutil does not automatically pick up
+            # the default credential keys when it is a service account.
+            'gcloud auth activate-service-account '
+            '--key-file=$GOOGLE_APPLICATION_CREDENTIALS '
+            '2> /dev/null || true; '
+            f'{gsutil_alias}')
 
     def is_directory(self, url: str) -> bool:
         """Returns whether 'url' is a directory.
@@ -124,7 +146,7 @@ class GcsCloudStorage(CloudStorage):
         # 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:
+        if not key:
             return True
         # Otherwise, gsutil ls -d url will return:
         #   --> url.rstrip('/')          if url is not a directory
@@ -153,12 +175,190 @@ class GcsCloudStorage(CloudStorage):
         return ' && '.join(all_commands)
 
 
+class AzureBlobCloudStorage(CloudStorage):
+    """Azure Blob Storage."""
+    # AzCopy is utilized for downloading data from Azure Blob Storage
+    # containers to remote systems due to its superior performance compared to
+    # az-cli. While az-cli's `az storage blob sync` can synchronize data from
+    # local to container, it lacks support to sync from container to remote
+    # synchronization. Moreover, `az storage blob download-batch` in az-cli
+    # does not leverage AzCopy's efficient multi-threaded capabilities, leading
+    # to slower performance.
+    #
+    # AzCopy requires appending SAS tokens directly in commands, as it does not
+    # support using STORAGE_ACCOUNT_KEY, unlike az-cli, which can generate
+    # SAS tokens but lacks direct multi-threading support like AzCopy.
+    # Hence, az-cli for SAS token generation is ran on the local machine and
+    # AzCopy is installed at the remote machine for efficient data transfer
+    # from containers to remote systems.
+    # Note that on Azure instances, both az-cli and AzCopy are typically
+    # pre-installed. And installing both would be used with AZ container is
+    # used from non-Azure instances.
+
+    _GET_AZCOPY = [
+        'azcopy --version > /dev/null 2>&1 || '
+        '(mkdir -p /usr/local/bin; '
+        'curl -L https://aka.ms/downloadazcopy-v10-linux -o azcopy.tar.gz; '
+        'sudo tar -xvzf azcopy.tar.gz --strip-components=1 -C /usr/local/bin --exclude=*.txt; '  # pylint: disable=line-too-long
+        'sudo chmod +x /usr/local/bin/azcopy; '
+        'rm azcopy.tar.gz)'
+    ]
+
+    def is_directory(self, url: str) -> bool:
+        """Returns whether 'url' of the AZ Container is a directory.
+
+        In cloud object stores, a "directory" refers to a regular object whose
+        name is a prefix of other objects.
+
+        Args:
+            url: Endpoint url of the container/blob.
+
+        Returns:
+            True if the url is an endpoint of a directory and False if it
+            is a blob(file).
+
+        Raises:
+            azure.core.exceptions.HttpResponseError: If the user's Azure
+                Azure account does not have sufficient IAM role for the given
+                storage account.
+            StorageBucketGetError: Provided container name does not exist.
+            TimeoutError: If unable to determine the container path status
+                in time.
+        """
+        storage_account_name, container_name, path = data_utils.split_az_path(
+            url)
+
+        # If there are more, we need to check if it is a directory or a file.
+        container_url = data_utils.AZURE_CONTAINER_URL.format(
+            storage_account_name=storage_account_name,
+            container_name=container_name)
+        resource_group_name = azure.get_az_resource_group(storage_account_name)
+        role_assignment_start = time.time()
+        refresh_client = False
+        role_assigned = False
+
+        # 1. List blobs in the container_url to decide wether it is a directory
+        # 2. If it fails due to permission issues, try to assign a permissive
+        # role for the storage account to the current Azure account
+        # 3. Wait for the role assignment to propagate and retry.
+        while (time.time() - role_assignment_start <
+               constants.WAIT_FOR_STORAGE_ACCOUNT_ROLE_ASSIGNMENT):
+            container_client = data_utils.create_az_client(
+                client_type='container',
+                container_url=container_url,
+                storage_account_name=storage_account_name,
+                resource_group_name=resource_group_name,
+                refresh_client=refresh_client)
+
+            if not container_client.exists():
+                with ux_utils.print_exception_no_traceback():
+                    raise sky_exceptions.StorageBucketGetError(
+                        f'The provided container {container_name!r} from the '
+                        f'passed endpoint url {url!r} does not exist. Please '
+                        'check if the name is correct.')
+
+            # If there aren't more than just container name and storage account,
+            # that's a directory.
+            # Note: This must be ran after existence of the storage account is
+            # checked while obtaining container client.
+            if not path:
+                return True
+
+            num_objects = 0
+            try:
+                for blob in container_client.list_blobs(name_starts_with=path):
+                    if blob.name == path:
+                        return False
+                    num_objects += 1
+                    if num_objects > 1:
+                        return True
+                # A directory with few or no items
+                return True
+            except azure.exceptions().HttpResponseError as e:
+                # Handle case where user lacks sufficient IAM role for
+                # a private container in the same subscription. Attempt to
+                # assign appropriate role to current user.
+                if 'AuthorizationPermissionMismatch' in str(e):
+                    if not role_assigned:
+                        logger.info('Failed to list blobs in container '
+                                    f'{container_url!r}. This implies '
+                                    'insufficient IAM role for storage account'
+                                    f' {storage_account_name!r}.')
+                        azure.assign_storage_account_iam_role(
+                            storage_account_name=storage_account_name,
+                            resource_group_name=resource_group_name)
+                        role_assigned = True
+                        refresh_client = True
+                    else:
+                        logger.info(
+                            'Waiting due to the propagation delay of IAM '
+                            'role assignment to the storage account '
+                            f'{storage_account_name!r}.')
+                        time.sleep(
+                            constants.RETRY_INTERVAL_AFTER_ROLE_ASSIGNMENT)
+                    continue
+                raise
+        else:
+            raise TimeoutError(
+                'Failed to determine the container path status within '
+                f'{constants.WAIT_FOR_STORAGE_ACCOUNT_ROLE_ASSIGNMENT}'
+                'seconds.')
+
+    def _get_azcopy_source(self, source: str, is_dir: bool) -> str:
+        """Converts the source so it can be used as an argument for azcopy."""
+        storage_account_name, container_name, blob_path = (
+            data_utils.split_az_path(source))
+        storage_account_key = data_utils.get_az_storage_account_key(
+            storage_account_name)
+
+        if storage_account_key is None:
+            # public containers do not require SAS token for access
+            sas_token = ''
+        else:
+            if is_dir:
+                sas_token = azure.get_az_container_sas_token(
+                    storage_account_name, storage_account_key, container_name)
+            else:
+                sas_token = azure.get_az_blob_sas_token(storage_account_name,
+                                                        storage_account_key,
+                                                        container_name,
+                                                        blob_path)
+        # "?" is a delimiter character used when SAS token is attached to the
+        # container endpoint.
+        # Reference: https://learn.microsoft.com/en-us/azure/ai-services/translator/document-translation/how-to-guides/create-sas-tokens?tabs=Containers # pylint: disable=line-too-long
+        converted_source = f'{source}?{sas_token}' if sas_token else source
+
+        return shlex.quote(converted_source)
+
+    def make_sync_dir_command(self, source: str, destination: str) -> str:
+        """Fetches a directory using AZCOPY from storage to remote instance."""
+        source = self._get_azcopy_source(source, is_dir=True)
+        # destination is guaranteed to not have '/' at the end of the string
+        # by tasks.py::set_file_mounts(). It is necessary to add from this
+        # method due to syntax of azcopy.
+        destination = f'{destination}/'
+        download_command = (f'azcopy sync {source} {destination} '
+                            '--recursive --delete-destination=false')
+        all_commands = list(self._GET_AZCOPY)
+        all_commands.append(download_command)
+        return ' && '.join(all_commands)
+
+    def make_sync_file_command(self, source: str, destination: str) -> str:
+        """Fetches a file using AZCOPY from storage to remote instance."""
+        source = self._get_azcopy_source(source, is_dir=False)
+        download_command = f'azcopy copy {source} {destination}'
+        all_commands = list(self._GET_AZCOPY)
+        all_commands.append(download_command)
+        return ' && '.join(all_commands)
+
+
 class R2CloudStorage(CloudStorage):
     """Cloudflare Cloud Storage."""
 
     # List of commands to install AWS CLI
     _GET_AWSCLI = [
-        'aws --version >/dev/null 2>&1 || pip3 install awscli',
+        'aws --version >/dev/null 2>&1 || '
+        f'{constants.SKY_UV_PIP_CMD} install awscli',
     ]
 
     def is_directory(self, url: str) -> bool:
@@ -193,7 +393,8 @@ class R2CloudStorage(CloudStorage):
             source = source.replace('r2://', 's3://')
         download_via_awscli = ('AWS_SHARED_CREDENTIALS_FILE='
                                f'{cloudflare.R2_CREDENTIALS_PATH} '
-                               'aws s3 sync --no-follow-symlinks '
+                               f'{constants.SKY_REMOTE_PYTHON_ENV}/bin/aws s3 '
+                               'sync --no-follow-symlinks '
                                f'{source} {destination} '
                                f'--endpoint {endpoint_url} '
                                f'--profile={cloudflare.R2_PROFILE_NAME}')
@@ -209,7 +410,8 @@ class R2CloudStorage(CloudStorage):
             source = source.replace('r2://', 's3://')
         download_via_awscli = ('AWS_SHARED_CREDENTIALS_FILE='
                                f'{cloudflare.R2_CREDENTIALS_PATH} '
-                               f'aws s3 cp {source} {destination} '
+                               f'{constants.SKY_REMOTE_PYTHON_ENV}/bin/aws s3 '
+                               f'cp {source} {destination} '
                                f'--endpoint {endpoint_url} '
                                f'--profile={cloudflare.R2_PROFILE_NAME}')
 
@@ -218,16 +420,6 @@ class R2CloudStorage(CloudStorage):
         return ' && '.join(all_commands)
 
 
-def get_storage_from_path(url: str) -> CloudStorage:
-    """Returns a CloudStorage by identifying the scheme:// in a URL."""
-    result = urllib.parse.urlsplit(url)
-
-    if result.scheme not in _REGISTRY:
-        assert False, (f'Scheme {result.scheme} not found in'
-                       f' supported storage ({_REGISTRY.keys()}); path {url}')
-    return _REGISTRY[result.scheme]
-
-
 class IBMCosCloudStorage(CloudStorage):
     """IBM Cloud Storage."""
     # install rclone if package isn't already installed
@@ -294,10 +486,82 @@ class IBMCosCloudStorage(CloudStorage):
         return self.make_sync_dir_command(source, destination)
 
 
+class OciCloudStorage(CloudStorage):
+    """OCI Cloud Storage."""
+
+    def is_directory(self, url: str) -> bool:
+        """Returns whether OCI 'url' is a directory.
+        In cloud object stores, a "directory" refers to a regular object whose
+        name is a prefix of other objects.
+        """
+        bucket_name, path = data_utils.split_oci_path(url)
+
+        client = oci.get_object_storage_client()
+        namespace = client.get_namespace(
+            compartment_id=oci.get_oci_config()['tenancy']).data
+
+        objects = client.list_objects(namespace_name=namespace,
+                                      bucket_name=bucket_name,
+                                      prefix=path).data.objects
+
+        if len(objects) == 0:
+            # A directory with few or no items
+            return True
+
+        if len(objects) > 1:
+            # A directory with more than 1 items
+            return True
+
+        object_name = objects[0].name
+        if path.endswith(object_name):
+            # An object path
+            return False
+
+        # A directory with only 1 item
+        return True
+
+    @oci.with_oci_env
+    def make_sync_dir_command(self, source: str, destination: str) -> str:
+        """Downloads using OCI CLI."""
+        bucket_name, path = data_utils.split_oci_path(source)
+
+        download_via_ocicli = (f'oci os object sync --no-follow-symlinks '
+                               f'--bucket-name {bucket_name} '
+                               f'--prefix "{path}" --dest-dir "{destination}"')
+
+        return download_via_ocicli
+
+    @oci.with_oci_env
+    def make_sync_file_command(self, source: str, destination: str) -> str:
+        """Downloads a file using OCI CLI."""
+        bucket_name, path = data_utils.split_oci_path(source)
+        filename = os.path.basename(path)
+        destination = os.path.join(destination, filename)
+
+        download_via_ocicli = (f'oci os object get --bucket-name {bucket_name} '
+                               f'--name "{path}" --file "{destination}"')
+
+        return download_via_ocicli
+
+
+def get_storage_from_path(url: str) -> CloudStorage:
+    """Returns a CloudStorage by identifying the scheme:// in a URL."""
+    result = urllib.parse.urlsplit(url)
+    if result.scheme not in _REGISTRY:
+        assert False, (f'Scheme {result.scheme} not found in'
+                       f' supported storage ({_REGISTRY.keys()}); path {url}')
+    return _REGISTRY[result.scheme]
+
+
 # Maps bucket's URIs prefix(scheme) to its corresponding storage class
 _REGISTRY = {
     'gs': GcsCloudStorage(),
     's3': S3CloudStorage(),
     'r2': R2CloudStorage(),
     'cos': IBMCosCloudStorage(),
+    'oci': OciCloudStorage(),
+    # TODO: This is a hack, as Azure URL starts with https://, we should
+    # refactor the registry to be able to take regex, so that Azure blob can
+    # be identified with `https://(.*?)\.blob\.core\.windows\.net`
+    'https': AzureBlobCloudStorage()
 }

sky/clouds/__init__.py

@@ -3,26 +3,29 @@
 from sky.clouds.cloud import Cloud
 from sky.clouds.cloud import cloud_in_iterable
 from sky.clouds.cloud import CloudImplementationFeatures
+from sky.clouds.cloud import OpenPortsVersion
 from sky.clouds.cloud import ProvisionerVersion
 from sky.clouds.cloud import Region
 from sky.clouds.cloud import StatusVersion
 from sky.clouds.cloud import Zone
-from sky.clouds.cloud_registry import CLOUD_REGISTRY
 
 # NOTE: import the above first to avoid circular imports.
 # isort: split
 from sky.clouds.aws import AWS
 from sky.clouds.azure import Azure
 from sky.clouds.cudo import Cudo
+from sky.clouds.do import DO
 from sky.clouds.fluidstack import Fluidstack
 from sky.clouds.gcp import GCP
 from sky.clouds.ibm import IBM
 from sky.clouds.kubernetes import Kubernetes
 from sky.clouds.lambda_cloud import Lambda
+from sky.clouds.nebius import Nebius
 from sky.clouds.oci import OCI
 from sky.clouds.paperspace import Paperspace
 from sky.clouds.runpod import RunPod
 from sky.clouds.scp import SCP
+from sky.clouds.vast import Vast
 from sky.clouds.vsphere import Vsphere
 
 __all__ = [
@@ -33,19 +36,21 @@ __all__ = [
     'Cudo',
     'GCP',
     'Lambda',
+    'DO',
     'Paperspace',
     'SCP',
     'RunPod',
+    'Vast',
     'OCI',
     'Vsphere',
     'Kubernetes',
     'CloudImplementationFeatures',
     'Region',
     'Zone',
-    'CLOUD_REGISTRY',
     'ProvisionerVersion',
     'StatusVersion',
     'Fluidstack',
+    'Nebius',
     # Utility functions
     'cloud_in_iterable',
 ]