diff --git a/sky/clouds/azure.py b/sky/clouds/azure.py index 5a696bf7e14..e2a15ec63bd 100644 --- a/sky/clouds/azure.py +++ b/sky/clouds/azure.py @@ -476,7 +476,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, ( diff --git a/sky/data/data_utils.py b/sky/data/data_utils.py index 8e994f3a78b..8c19598f620 100644 --- a/sky/data/data_utils.py +++ b/sky/data/data_utils.py @@ -63,13 +63,13 @@ 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) @@ -77,6 +77,7 @@ def split_az_path(az_path: str) -> Tuple[str, str, str]: 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,18 +158,20 @@ 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'. """ container_url = kwargs.pop('container_url', None) if client_type == 'container': - assert container_url is not None + assert container_url is not None, ('container_url must be provided for ' + 'container client') + subscription_id = azure.get_subscription_id() return azure.get_client(client_type, subscription_id, @@ -183,7 +186,7 @@ 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, @@ -199,11 +202,12 @@ def get_az_resource_group( """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 + storage_account_name: Name of the storage account + storage_client: Client object facing storage Returns: - Name of the resource group the given storage account belongs to. + Name of the resource group the given storage account belongs to, or + None if not found. """ if storage_client is None: storage_client = create_az_client('storage') @@ -232,14 +236,15 @@ 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') @@ -249,7 +254,7 @@ def get_az_storage_account_key( resource_group_name = 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,16 +291,16 @@ 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(/[^/]+)*$') diff --git a/sky/data/mounting_utils.py b/sky/data/mounting_utils.py index 7b1ecba1a76..5d4eb61156c 100644 --- a/sky/data/mounting_utils.py +++ b/sky/data/mounting_utils.py @@ -85,15 +85,14 @@ def get_az_mount_cmd(container_name: str, """Returns a command to mount an AZ Container using blobfuse2. Args: - container_name: str; Name of the mounting container - storage_account_name: str; Name of the storage account the given - container belongs to - mount_path: str; Path where the container will be mounting - storage_account_key: Optional[str]; Access key for the given storage - account + container_name: Name of the mounting container. + storage_account_name: Name of the storage account the given container + belongs to. + mount_path: Path where the container will be mounting. + storage_account_key: Access key for the given storage account. Returns: - str: command used to mount AZ container with blobfuse2 + str: Command used to mount AZ container with blobfuse2. """ # Storage_account_key is set to None when mounting public container, and # mounting public containers are not officially supported by blobfuse2 yet. @@ -163,10 +162,10 @@ def _get_mount_binary(mount_cmd: str) -> str: """Returns mounting binary in string given as the mount command. Args: - mount_cmd: str; command used to mount a cloud storage. + mount_cmd: Command used to mount a cloud storage. Returns: - str: name of the binary used to mount a cloud storage. + str: Name of the binary used to mount a cloud storage. """ if 'goofys' in mount_cmd: return 'goofys' diff --git a/sky/data/storage.py b/sky/data/storage.py index 55571dfc715..0ba353ee9d7 100644 --- a/sky/data/storage.py +++ b/sky/data/storage.py @@ -172,6 +172,15 @@ def store_prefix(self) -> str: @classmethod def get_endpoint_url(cls, store: 'AbstractStore', path: str) -> str: + """Generates the endpoint URL for a given store and path. + + Args: + store: Store object implementing AbstractStore. + path: Path within the store. + + Returns: + Endpoint URL of the bucket as a string. + """ store_type = cls.from_store(store) if store_type == StoreType.AZURE: assert isinstance(store, AzureBlobStore) @@ -1943,7 +1952,7 @@ class AzureBlobStore(AbstractStore): DEFAULT_RESOURCE_GROUP_NAME = 'sky{user_hash}' class AzureBlobStoreMetadata(AbstractStore.StoreMetadata): - """A pickle-able representation of Azure Blob Store + """A pickle-able representation of Azure Blob Store. Allows store objects to be written to and reconstructed from global_user_state. @@ -1992,11 +2001,17 @@ def __init__(self, @classmethod def from_metadata(cls, metadata: AbstractStore.StoreMetadata, - **override_args): - """Create a Store from a AzureBlobStoreMetadata object. + **override_args) -> 'AzureBlobStore': + """Creates AzureBlobStore from a AzureBlobStoreMetadata object. Used when reconstructing Storage and Store objects from global_user_state. + + Args: + metadata: Metadata object containing AzureBlobStore information. + + Returns: + An instance of AzureBlobStore. """ assert isinstance(metadata, AzureBlobStore.AzureBlobStoreMetadata) return cls(name=override_args.get('name', metadata.name), @@ -2072,10 +2087,10 @@ def validate_name(cls, name: str) -> str: Source for rules: https://learn.microsoft.com/en-us/rest/api/storageservices/Naming-and-Referencing-Containers--Blobs--and-Metadata#container-names # pylint: disable=line-too-long Args: - name: str; name of the container + name: Name of the container Returns: - name: str; name of the container + Name of the container Raises: StorageNameError: if the given container name does not follow the @@ -2171,6 +2186,10 @@ def _get_storage_account_and_resource_group( 4) If none of the above are true, default naming conventions are used to create the resource group and storage account for the users. + Returns: + str: The storage account name. + Optional[str]: The resource group name, or None if not found. + Raises: StorageBucketCreateError: If storage account attempted to be created already exists @@ -2265,9 +2284,9 @@ def _create_storage_account(self, resource_group_name: str, """Creates new storage account and assign Storage Blob Data Owner role. Args: - resource_group_name: str; Name of the resource group which the - storage account will be created under. - storage_account_name: str; Name of the storage account to be created + resource_group_name: Name of the resource group which the storage + account will be created under. + storage_account_name: Name of the storage account to be created. Raises: StorageBucketCreateError: If storage account attempted to be @@ -2607,7 +2626,7 @@ def mount_command(self, mount_path: str) -> str: Uses blobfuse2 to mount the container. Args: - mount_path: str; Path to mount the container to + mount_path: Path to mount the container to Returns: str: a heredoc used to setup the AZ Container mount @@ -2624,16 +2643,17 @@ def _create_az_bucket(self, container_name: str) -> StorageHandle: """Creates AZ Container. Args: - container_name: str; Name of bucket(container) - region: str; Region name, e.g. eastus, westus + container_name: Name of bucket(container) Returns: - StorageHandle: handle to interact with the container + StorageHandle: Handle to interact with the container Raises: StorageBucketCreateError: If container creation fails. """ try: + # Container is created under the region which the storage account + # belongs to. container = self.storage_client.blob_containers.create( self.resource_group_name, self.storage_account_name, @@ -2662,11 +2682,15 @@ def _delete_az_bucket(self, container_name: str) -> bool: """Deletes AZ Container, including all objects in Container. Args: - container_name: str; Name of bucket(container) + container_name: Name of bucket(container). Returns: - bool: True if container was deleted, False if it's - deleted externally. + bool: True if container was deleted, False if it's deleted + externally. + + Raises: + StorageBucketDeleteError: If deletion fails for reasons other than + the container not existing. """ try: with rich_utils.safe_status( @@ -2683,7 +2707,7 @@ def _delete_az_bucket(self, container_name: str) -> bool: container_name, ) except azure.exceptions().ResourceNotFoundError as e: - if 'Code: ContainerNotFound' in e.message: + if 'Code: ContainerNotFound' in str(e): logger.debug( _BUCKET_EXTERNALLY_DELETED_DEBUG_MESSAGE.format( bucket_name=container_name)) @@ -2692,7 +2716,7 @@ def _delete_az_bucket(self, container_name: str) -> bool: with ux_utils.print_exception_no_traceback(): raise exceptions.StorageBucketDeleteError( f'Failed to delete Azure container {container_name}. ' - f'Detailed error: {e.output}') + f'Detailed error: {e}') return True