[Storage] Azure blob storage support

sky/adaptors/azure.py

@@ -1,10 +1,15 @@
 """Azure cli adaptor"""
 
 # pylint: disable=import-outside-toplevel
+import datetime
 import functools
+import logging
 import threading
 
+from sky import exceptions as sky_exceptions
 from sky.adaptors import common
+from sky.utils import common_utils
+from sky.utils import ux_utils
 
 azure = common.LazyImport(
     'azure',
@@ -38,7 +43,7 @@ def exceptions():
 
 @functools.lru_cache()
 @common.load_lazy_modules(modules=_LAZY_MODULES)
-def get_client(name: str, subscription_id: str):
+def get_client(name: str, subscription_id: str, **kwargs):
     # Sky only supports Azure CLI credential for now.
     # Increase the timeout to fix the Azure get-access-token timeout issue.
     # Tracked in
@@ -55,10 +60,140 @@ def get_client(name: str, subscription_id: str):
         elif name == 'resource':
             from azure.mgmt.resource import ResourceManagementClient
             return ResourceManagementClient(credential, subscription_id)
+        elif name == 'storage':
+            from azure.mgmt.storage import StorageManagementClient
+            return StorageManagementClient(credential, subscription_id)
+        elif name == 'authorization':
+            from azure.mgmt.authorization import AuthorizationManagementClient
+            return AuthorizationManagementClient(credential, subscription_id)
+        elif name == 'graph':
+            from msgraph import GraphServiceClient
+            return GraphServiceClient(credential)
+        elif name == 'container':
+            # Currently, there is no way to check if the given endpoint url
+            # of a container is private or public. Hence, we approach with
+            # try-except block by first assuming the url is a public container.
+            # If it turns out to be a private container, an error will be
+            # caught and treated as private container.
+            # Reference: https://github.com/Azure/azure-sdk-for-python/issues/35770 # pylint: disable=line-too-long
+            from azure.storage.blob import ContainerClient
+            container_url = kwargs.pop('container_url', None)
+            assert container_url is not None, ('Must provide "container_url"'
+                                               ' keyword arguments for '
+                                               'container client.')
+            container_client = ContainerClient.from_container_url(container_url)
+            try:
+                container_client.exists()
+            except exceptions().ClientAuthenticationError:
+                # Caught when credential is not provided to the private
+                # container url or wrong container name is provided as the
+                # public container url. We reattempt with credentials assuming
+                # user is using private container with access rights.
+                container_client = ContainerClient.from_container_url(
+                    container_url, credential)
+                try:
+                    # Suppress noisy logs from Azure SDK when attempting
+                    # to run exists() on public container with credentials.
+                    # Reference:
+                    # https://github.com/Azure/azure-sdk-for-python/issues/9422 # pylint: disable=line-too-long
+                    azure_logger = logging.getLogger('azure')
+                    original_level = azure_logger.getEffectiveLevel()
+                    azure_logger.setLevel(logging.CRITICAL)
+                    container_client.exists()
+                    azure_logger.setLevel(original_level)
+                except exceptions().ClientAuthenticationError as error:
+                    # Caught when user attempted to use private container
+                    # without access rights.
+                    # Reference:
+                    # https://learn.microsoft.com/en-us/troubleshoot/azure/entra/entra-id/app-integration/error-code-aadsts50020-user-account-identity-provider-does-not-exist # pylint: disable=line-too-long
+                    if 'ERROR: AADSTS50020' in error.message:
+                        with ux_utils.print_exception_no_traceback():
+                            raise sky_exceptions.StorageBucketGetError(
+                                'Attempted to fetch a non-existant public '
+                                'container name: '
+                                f'{container_client.container_name}. '
+                                'Please check if the name is correct.')
+                    else:
+                        with ux_utils.print_exception_no_traceback():
+                            raise sky_exceptions.StorageBucketGetError(
+                                'Failed to retreive the container client '
+                                'for the container: '
+                                f'{container_client.container_name!r}. '
+                                f'Details: '
+                                f'{common_utils.format_exception(error, use_bracket=True)}'
+                            )
+
+            return container_client
         else:
             raise ValueError(f'Client not supported: "{name}"')
 
 
+@common.load_lazy_modules(modules=_LAZY_MODULES)
+def get_az_container_sas_token(
+    storage_account_name: str,
+    storage_account_key: str,
+    container_name: str,
+) -> str:
+    """Returns SAS token used to access container.
+
+    Args:
+        storage_account_name: str; Name of the storage account
+        storage_account_key: str; Access key for the given storage
+            account
+        container_name: str; The name of the mounting container
+
+    Returns:
+        SAS token prepended with the delimiter character, "?"
+    """
+    from azure.storage.blob import ContainerSasPermissions
+    from azure.storage.blob import generate_container_sas
+    sas_token = generate_container_sas(
+        account_name=storage_account_name,
+        container_name=container_name,
+        account_key=storage_account_key,
+        permission=ContainerSasPermissions(read=True,
+                                           write=True,
+                                           list=True,
+                                           create=True),
+        expiry=datetime.datetime.now(datetime.timezone.utc) +
+        datetime.timedelta(hours=1))
+    # "?" 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
+    return f'?{sas_token}'
+
+
+@common.load_lazy_modules(modules=_LAZY_MODULES)
+def get_az_blob_sas_token(storage_account_name: str, storage_account_key: str,
+                          container_name: str, blob_name: str) -> str:
+    """Returns SAS token used to access a blob.
+
+    Args:
+        storage_account_name: str; Name of the storage account
+        storage_account_key: str; access key for the given storage
+            account
+        container_name: str; name of the mounting container
+        blob_name: str; path to the blob(file)
+
+    Returns:
+        SAS token prepended with the delimiter character, "?"
+    """
+    from azure.storage.blob import BlobSasPermissions
+    from azure.storage.blob import generate_blob_sas
+    sas_token = generate_blob_sas(
+        account_name=storage_account_name,
+        container_name=container_name,
+        blob_name=blob_name,
+        account_key=storage_account_key,
+        permission=BlobSasPermissions(read=True,
+                                      write=True,
+                                      list=True,
+                                      create=True),
+        expiry=datetime.datetime.now(datetime.timezone.utc) +
+        datetime.timedelta(hours=1))
+    return f'{sas_token}'
+
+
 @common.load_lazy_modules(modules=_LAZY_MODULES)
 def create_security_rule(**kwargs):
     from azure.mgmt.network.models import SecurityRule

sky/cloud_stores.py

@@ -7,10 +7,12 @@
 * Better interface.
 * Better implementation (e.g., fsspec, smart_open, using each cloud's SDK).
 """
+import shlex
 import subprocess
 import urllib.parse
 
 from sky.adaptors import aws
+from sky.adaptors import azure
 from sky.adaptors import cloudflare
 from sky.adaptors import ibm
 from sky.clouds import gcp
@@ -153,6 +155,113 @@ def make_sync_file_command(self, source: str, destination: str) -> str:
         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.
+        """
+        # split the url using split_az_path
+        _, _, path = data_utils.split_az_path(url)
+        # If there aren't more than just container name and storage account,
+        # that's a directory.
+        if not path:
+            return True
+        # If there's more, we'd need to check if it's a directory or a file.
+        container_client = data_utils.create_az_client(client_type='container',
+                                                       container_url=url)
+        num_objects = 0
+        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
+
+    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))
+        resource_group_name = data_utils.get_az_resource_group(
+            storage_account_name)
+        storage_account_key = data_utils.get_az_storage_account_key(
+            storage_account_name, resource_group_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)
+        if is_dir:
+            converted_source = f'{source}/{sas_token}'
+        else:
+            # "?" is a delimiter character used when SAS token is attached to
+            # the blob endpoint.
+            converted_source = f'{source}/{blob_path}?{sas_token}'
+
+        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."""
 
@@ -218,16 +327,6 @@ def make_sync_file_command(self, source: str, destination: str) -> str:
         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 +393,20 @@ def make_sync_file_command(self, source: str, destination: str) -> str:
         return self.make_sync_dir_command(source, destination)
 
 
+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(),
+    'https': AzureBlobCloudStorage()
 }

sky/core.py

@@ -828,7 +828,7 @@ def storage_delete(name: str) -> None:
     if handle is None:
         raise ValueError(f'Storage name {name!r} not found.')
     else:
-        store_object = data.Storage(name=handle.storage_name,
-                                    source=handle.source,
-                                    sync_on_reconstruction=False)
-        store_object.delete()
+        storage_object = data.Storage(name=handle.storage_name,
+                                      source=handle.source,
+                                      sync_on_reconstruction=False)
+        storage_object.delete()