[Storage] Azure blob storage support

sky/backends/cloud_vm_ray_backend.py

@@ -4423,13 +4423,13 @@ def _execute_file_mounts(self, handle: CloudVmRayResourceHandle,
 
             storage = cloud_stores.get_storage_from_path(src)
             if storage.is_directory(src):
-                sync = storage.make_sync_dir_command(source=src,
-                                                     destination=wrapped_dst)
+                sync_cmd = (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:
-                sync = storage.make_sync_file_command(source=src,
-                                                      destination=wrapped_dst)
+                sync_cmd = (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)}')
@@ -4438,7 +4438,7 @@ def _execute_file_mounts(self, handle: CloudVmRayResourceHandle,
                 # Ensure sync can write to wrapped_dst (e.g., '/data/').
                 mkdir_for_wrapped_dst,
                 # Both the wrapped and the symlink dir exist; sync.
-                sync,
+                sync_cmd,
             ]
             command = ' && '.join(download_target_commands)
             # dst is only used for message printing.

sky/cloud_stores.py

@@ -7,15 +7,24 @@
 * Better interface.
 * Better implementation (e.g., fsspec, smart_open, using each cloud's SDK).
 """
+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.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:
@@ -153,6 +162,183 @@ 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.
+
+        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."""
 
@@ -218,16 +404,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 +470,23 @@ 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(),
+    # 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/azure.py

@@ -475,6 +475,19 @@ def check_credentials(cls) -> Tuple[bool, Optional[str]]:
             return False, (f'Getting user\'s Azure identity failed.{help_str}\n'
                            f'{cls._INDENT_PREFIX}Details: '
                            f'{common_utils.format_exception(e)}')
+
+        # Check if the azure blob storage dependencies are installed.
+        try:
+            # pylint: disable=redefined-outer-name, import-outside-toplevel, unused-import
+            from azure.storage import blob
+            import msgraph
+        except ImportError as e:
+            return False, (
+                f'Azure blob storage depdencies are not installed. '
+                'Run the following commands:'
+                f'\n{cls._INDENT_PREFIX}  $ pip install skypilot[azure]'
+                f'\n{cls._INDENT_PREFIX}Details: '
+                f'{common_utils.format_exception(e)}')
         return True, None
 
     def get_credential_file_mounts(self) -> Dict[str, str]:

sky/core.py

@@ -831,7 +831,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()