[Storage] Azure blob storage support

sky/adaptors/common.py

@@ -60,6 +60,10 @@ def wrapper(*args, **kwargs):
                 m.load_module()
             return func(*args, **kwargs)
 
+        # Preserve lru_cache methods
+        if hasattr(func, 'cache_clear'):
+            wrapper.cache_clear = func.cache_clear
+
         return wrapper
 
     return decorator

sky/backends/cloud_vm_ray_backend.py

@@ -4508,23 +4508,35 @@ def _execute_file_mounts(self, handle: CloudVmRayResourceHandle,
                 continue
 
             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:
-                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)}')
+            try:
+                if storage.is_directory(src):
+                    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_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)}')
+            except Exception as e:  # pylint: disable=broad-except
+                logger.error(
+                    f'Failed to fetch from the bucket {src!r} to '
+                    f'remote instance at {dst!r}.\n'
+                    'Error details: '
+                    f'{common_utils.format_exception(e, use_bracket=True)}.')
+                # If 'cmd' was appended to 'symlink_commands' for this sync, we
+                # remove as it failed to sync.
+                if not dst.startswith('~/') and not dst.startswith('/tmp/'):
+                    symlink_commands.pop()
+                continue
 
             download_target_commands = [
                 # 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

@@ -9,15 +9,20 @@
 """
 import shlex
 import subprocess
+import time
 import urllib.parse
 
+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
+
+logger = sky_logging.init_logger(__name__)
 
 
 class CloudStorage:
@@ -189,30 +194,84 @@ def is_directory(self, url: str) -> bool:
 
         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.
         """
-        # split the url using split_az_path
         storage_account_name, container_name, 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 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)
-        # 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=container_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
+        resource_group_name = azure.get_az_resource_group(storage_account_name)
+        role_assignment_start = time.time()
+        refresh_client = False
+        role_assigned = False
+
+        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)
+
+            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
+            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
+            # A directory with few or no items
+            return True
+        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."""
@@ -408,5 +467,8 @@ def get_storage_from_path(url: str) -> CloudStorage:
     '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

@@ -481,7 +481,7 @@ def check_credentials(cls) -> Tuple[bool, Optional[str]]:
         # Check if the azure blob storage dependencies are installed.
         try:
             # pylint: disable=redefined-outer-name, import-outside-toplevel, unused-import
-            import azure.storage.blob
+            from azure.storage import blob
             import msgraph
         except ImportError as e:
             return False, (

sky/data/data_utils.py

@@ -63,20 +63,21 @@ def split_az_path(az_path: str) -> Tuple[str, str, str]:
     """Splits Path into Storage account and Container names and Relative Path
 
     Args:
-        az_path: str; Container Path,
+        az_path: Container Path,
           e.g. https://azureopendatastorage.blob.core.windows.net/nyctlc
 
     Returns:
-        str; Name of the storage account
-        str; Name of the container
-        str; Paths of the file/directory defined within the container
+        str: Name of the storage account
+        str: Name of the container
+        str: Paths of the file/directory defined within the container
     """
     path_parts = az_path.replace('https://', '').split('/')
     service_endpoint = path_parts.pop(0)
     service_endpoint_parts = service_endpoint.split('.')
     storage_account_name = service_endpoint_parts[0]
     container_name = path_parts.pop(0)
     path = '/'.join(path_parts)
+
     return storage_account_name, container_name, path
 
 
@@ -157,22 +158,37 @@ def verify_gcs_bucket(name: str) -> bool:
         return False
 
 
-def create_az_client(client_type: str, **kwargs) -> Client:
+def create_az_client(client_type: str, **kwargs: Any) -> Client:
     """Helper method that connects to AZ client for diverse Resources.
 
     Args:
       client_type: str; specify client type, e.g. storage, resource, container
 
     Returns:
-      Client object facing AZ Resource of the 'client_type'.
+        Client object facing AZ Resource of the 'client_type'.
     """
+    resource_group_name = kwargs.pop('resource_group_name', None)
     container_url = kwargs.pop('container_url', None)
+    storage_account_name = kwargs.pop('storage_account_name', None)
+    refresh_client = kwargs.pop('refresh_client', False)
     if client_type == 'container':
-        assert container_url is not None
+        # We do not assert on resource_group_name as it is set to None when the
+        # container_url is for public container with user access.
+        assert container_url is not None, ('container_url must be provided for '
+                                           'container client')
+        assert storage_account_name is not None, ('storage_account_name must '
+                                                  'be provided for container '
+                                                  'client')
+
     subscription_id = azure.get_subscription_id()
-    return azure.get_client(client_type,
-                            subscription_id,
-                            container_url=container_url)
+    if refresh_client:
+        azure.get_client.cache_clear()
+    client = azure.get_client(client_type,
+                              subscription_id,
+                              container_url=container_url,
+                              storage_account_name=storage_account_name,
+                              resource_group_name=resource_group_name)
+    return client
 
 
 def verify_az_bucket(storage_account_name: str, container_name: str) -> bool:
@@ -183,46 +199,20 @@ def verify_az_bucket(storage_account_name: str, container_name: str) -> bool:
       container_name: str; Name of the container
 
     Returns:
-      boolean; Shows either or not the container exists.
+      True if the container exists, False otherwise.
     """
     container_url = AZURE_CONTAINER_URL.format(
         storage_account_name=storage_account_name,
         container_name=container_name)
-    container_client = create_az_client(client_type='container',
-                                        container_url=container_url)
+    resource_group_name = azure.get_az_resource_group(storage_account_name)
+    container_client = create_az_client(
+        client_type='container',
+        container_url=container_url,
+        storage_account_name=storage_account_name,
+        resource_group_name=resource_group_name)
     return container_client.exists()
 
 
-def get_az_resource_group(
-        storage_account_name: str,
-        storage_client: Optional[Client] = None) -> Optional[str]:
-    """Returns the resource group name the given storage account belongs to.
-
-    Args:
-        storage_account_name: str; Name of the storage account
-        storage_client: Optional[Client]; Client object facing storage
-
-    Returns:
-        Name of the resource group the given storage account belongs to.
-    """
-    if storage_client is None:
-        storage_client = create_az_client('storage')
-    for account in storage_client.storage_accounts.list():
-        if account.name == storage_account_name:
-            # Extract the resource group name from the account ID
-            # An example of account.id would be the following:
-            # /subscriptions/{subscription_id}/resourceGroups/{resource_group_name}/providers/Microsoft.Storage/storageAccounts/{container_name} # pylint: disable=line-too-long
-            split_account_id = account.id.split('/')
-            assert len(split_account_id) == 9
-            resource_group_name = split_account_id[4]
-            return resource_group_name
-    # resource group cannot be found when using container not created
-    # under the user's subscription id, i.e. public container, or
-    # private containers not belonging to the user or when the storage account
-    # does not exist.
-    return None
-
-
 def get_az_storage_account_key(
     storage_account_name: str,
     resource_group_name: Optional[str] = None,
@@ -232,24 +222,25 @@ def get_az_storage_account_key(
     """Returns access key of the given name of storage account.
 
     Args:
-        storage_account_name: str; Name of the storage account
-        resource_group_name: Optional[str]; Name of the resource group the
+        storage_account_name: Name of the storage account
+        resource_group_name: Name of the resource group the
             passed storage account belongs to.
-        storage_clent: Optional[Client]; Client object facing Storage
-        resource_client: Optional[Client]; Client object facing Resource
+        storage_clent: Client object facing Storage
+        resource_client: Client object facing Resource
 
     Returns:
-        One of the few access keys to the given storage account
+        One of the two access keys to the given storage account, or None if
+        the account is not found.
     """
     if resource_client is None:
         resource_client = create_az_client('resource')
     if storage_client is None:
         storage_client = create_az_client('storage')
     if resource_group_name is None:
-        resource_group_name = get_az_resource_group(storage_account_name,
-                                                    storage_client)
+        resource_group_name = azure.get_az_resource_group(
+            storage_account_name, storage_client)
     # resource_group_name is None when using a public container or
-    # a private containers not belonging to the user.
+    # a private container not belonging to the user.
     if resource_group_name is None:
         return None
 
@@ -286,23 +277,23 @@ def get_az_storage_account_key(
 
 
 def is_az_container_endpoint(endpoint_url: str) -> bool:
-    """Checks if the provided url is valid container endpoint
+    """Checks if provided url follows a valid container endpoint naming format.
 
     Args:
-      endpoint_url: str; Url of container endpoint.
+      endpoint_url: Url of container endpoint.
         e.g. https://azureopendatastorage.blob.core.windows.net/nyctlc
 
     Returns:
-      boolean; Shows either or not the container endpoint is valid.
+      bool: True if the endpoint is valid, False otherwise.
     """
-    # storage account must be length of 3-24
+    # Storage account must be length of 3-24
     # Reference: https://learn.microsoft.com/en-us/azure/azure-resource-manager/management/resource-name-rules#microsoftstorage # pylint: disable=line-too-long
     pattern = re.compile(
         r'^https://([a-z0-9]{3,24})\.blob\.core\.windows\.net(/[^/]+)*$')
     match = pattern.match(endpoint_url)
-    if match:
-        return True
-    return False
+    if match is None:
+        return False
+    return True
 
 
 def create_r2_client(region: str = 'auto') -> Client: