Support for parallel blob downloads

[martinResearch]

Is your feature request related to a problem? Please describe.

I have a large list of blobs I want to download using the azure SDK. I can do a loop over the list and call the client's download_blob for each blob sequentially but this is very slow.

I implemented a class derived from azure.storage.blob.ContainerClient that uses ThreadPoolExecutor to do the downloads in parallel, with a new method with this interface:

    def download_blobs_to_files(
        self,
        blob_filename_pairs: Iterable[Tuple[str, str]],
        concurrency_limit: int = 1000,
        verbose: bool = False,
    ) -> int:
        """Downloads a list of files from an azure blob container.

        Args:
            blob_filename_pairs: List[Tuple[str, str]]:List of blob and local path pairs
            concurrency_limit: Maximum number of threads.
            verbose: controls verbosity of the function.

It works is it a bit brittle and it is not clear how to automatically choose the right number of threads (concurrency_limit). Ideally this would be a feature supported by the Azure SDK. It seems to me to be a frequent user need.

[github-actions]

Thanks for the feedback! We are routing this to the appropriate team for follow-up. cc @jalauzon-msft @vincenttran-msft.

[github-actions]

Hello @martinResearch. I'm an AI assistant for the azure-sdk-for-python repository. I have some suggestions that you can try out while the team gets back to you.

• The SDK supports parallel processing for blob downloads via the built-in max_concurrency parameter that you can pass to the download_blob method. This controls parallel chunk downloads per blob. Source

• For multiple blobs, you can issue parallel calls to download_blob (using threading, async or a thread pool), but optimal concurrency depends on factors such as network resources, blob size, and machine capabilities. Performance tuning provides some guidance.

• While your custom method works, a recommendation is to leverage the SDK's built-in parallelism for chunk downloads (by setting max_concurrency) and then control parallelism at the application level. A formula (e.g. based on CPU cores) or dynamic concurrent limit may help choose the right number.

• Overall, your observation is valid; improved native support for parallel downloads (for multiple blobs) might be a useful enhancement.

The team will get back to you shortly, hopefully this helps in the meantime.

[martinResearch]

It seems that download_blob with max_concurrency allows to use multiple threads to download a single blob, which seems useful when downloading a large blob. My use case is different: I have many small blobs and what to parallelize over blobs

[jalauzon-msft]

Hi @martinResearch, thanks for reaching out. Unfortunately, like #40270, this is also not something we are likely to support in the SDK anytime soon. Again, in a generalized fashion, this is a pretty complex feature that we don't really have the resourcing for. Additionally, like the sync function, if we did ever do this, it would be in a separate package likely. Again, if you have something currently implemented that is working for you, I recommend sticking with that.

We do offer this functionality programmatically via the .NET library Azure.Storage.DataMovement but obliviously that is .NET so it likely does not meet your needs. This is the type of library where we would introduce features like this and sync if we were to ever do them in Python.

[martinResearch]

My current implementation:

"""ContainerClientExtended - a helper class for accessing Azure Blob Storage.

ContainerClientExtended is a wrapper around azure.storage.blob.ContainerClient which
provides few additional helper methods.
"""

from concurrent.futures import ThreadPoolExecutor
from datetime import datetime, timezone
import hashlib
import os
from pathlib import Path
import time
from typing import Iterable, Tuple, Union, Dict, Optional
import uuid

from azure.storage.blob import BlobProperties, ContainerClient
from azure.storage.blob._container_client import AzureNamedKeyCredential, AzureSasCredential, TokenCredential


class ContainerClientExtended(ContainerClient):
    """An extension of azure.storage.blob.ContainerClient which automatically performs authentication.

    Args:
        account_url: The URI to the storage account, e.g. 'https://mydata.blob.core.windows.net'.
        container_name: The name of the blob container.
    """

    def __init__(
        self,
        account_url: str,
        container_name: str,
        timeout: int = 20,
        copy_dates: bool = True,
        credential: Optional[
            Union[str, Dict[str, str], AzureNamedKeyCredential, AzureSasCredential, "TokenCredential"]
        ] = None,
    ):
        """Performs authentication and initializes BlobStorageClient.

        Args:
            account_url: URL of the Azure storage account.
            container_name: Name of the Azure blob storage container.
            timeout: Timeout in seconds for connections/reads from the container.
            copy_dates: controls if we copy last_modified date from the azure blobs when downloading files
        """
        super().__init__(
            account_url,
            container_name,
            credential=credential,
            connection_timeout=timeout,
            read_timeout=timeout,
        )
        self._copy_dates = copy_dates

    @classmethod
    def _local_copy_exists_md5(cls, blob_props: BlobProperties, local_file: Union[Path, str]) -> bool:
        """Checks if the blob already exists locally based on MD5 hash.

        Args:
            blob_props: Blob properties of the blob to check.
            local_file: Location of the local file.

        Returns:
            true if the local file exists and contains the same data.
        """
        local_file = Path(local_file)
        if local_file.exists():
            hash_md5 = hashlib.md5()
            with open(local_file, "rb") as f:
                for chunk in iter(lambda: f.read(4096), b""):
                    hash_md5.update(chunk)
            md5_local = hash_md5.digest()
            md5_remote = blob_props.content_settings.content_md5
            return md5_local == md5_remote
        else:
            return False

    @classmethod
    def _local_copy_exists_dates_and_size(
        cls, blob_props: BlobProperties, local_file: Union[Path, str], verbose: bool = False
    ) -> bool:
        """Checks if the blob already exists locally base on dates and file size.

        This is used a fallback when the file is too large for azure storage to compute the md5.
        """
        local_file = Path(local_file)
        if not local_file.exists():
            return False

        local_file_stats = os.stat(local_file)

        local_file_last_modified_utc = datetime.fromtimestamp(local_file_stats.st_mtime, tz=timezone.utc)
        remote_file_last_modified_utc = blob_props["last_modified"]
        remote_file_size = blob_props["size"]
        local_file_size = local_file_stats.st_size
        if remote_file_last_modified_utc != local_file_last_modified_utc:
            if verbose:
                print(
                    f"The local and remote last modified dates are different for file {local_file}:\n"
                    f"Local last modified date utc: {local_file_last_modified_utc}\n"
                    f"Remote last modified date utc: {remote_file_last_modified_utc})"
                )
            return False
        elif remote_file_size != local_file_size:
            if verbose:
                print(
                    f"The local and remote file sizes are different for file {local_file}:\n"
                    f"Local file size:  {local_file_size}B\n"
                    f"Remote file size: {remote_file_size}B"
                )
            return False
        else:
            return True

    @classmethod
    def _local_copy_exists(cls, blob_props: BlobProperties, local_file: Union[Path, str]) -> bool:
        if blob_props.content_settings.content_md5:
            return cls._local_copy_exists_md5(blob_props=blob_props, local_file=local_file)
        else:
            # Azure storage does not compute md5 files for large files.
            # Fall back using the modification date and file size to detect if local copy already exist
            return cls._local_copy_exists_dates_and_size(blob_props=blob_props, local_file=local_file)

    def local_copy_exists(self, blob_name: str, file_path: Union[Path, str]) -> bool:
        """Check if the local copy of the file already exists."""
        file_path = Path(file_path)
        return self._local_copy_exists(self.get_blob_client(blob_name).get_blob_properties(), file_path)

    def download_blob_to_file(
        self, blob_name: str, file_path: Union[Path, str], overwrite: bool = False, verbose: bool = False
    ) -> None:
        """Downloads a blob and saves it to a file.

        Args:
            blob_name: Blob name, possibly including multiple path segments.
            file_path: Local file path where the blob will be saved.
            overwrite: Forces re-download if the local file already exists.
            verbose:   controls the verbosity of the function.
        """
        file_path = Path(file_path)

        blob_properties = self.get_blob_client(blob_name).get_blob_properties()

        if overwrite or not self._local_copy_exists(blob_properties, file_path):
            # Downloading the data first in a temporary file then renaming the file to the in order
            # to avoid getting a corrupted file if one interrupts a process while a file
            # is being downloaded in that process.
            tmp_file = file_path.with_suffix(file_path.suffix + "." + uuid.uuid4().hex[:10] + ".tmp")
            file_path.parent.mkdir(exist_ok=True, parents=True)
            if verbose:
                print(f"    {blob_name} -> {file_path}")
            with open(tmp_file, "wb") as file:
                super().download_blob(blob_name).readinto(file)

            if self._copy_dates:
                # copy last modified data into local file last modified date
                last_modified = blob_properties["last_modified"].timestamp()
                last_accessed = last_modified
                os.utime(tmp_file, (last_modified, last_accessed))

            if file_path.exists():
                file_path.unlink()
            tmp_file.rename(file_path)
        else:
            if verbose:
                print(f"    {blob_name} present locally, not overwritten")

    def upload_file_to_blob(self, blob_name: str, file_path: Union[Path, str], overwrite: bool = False) -> None:
        """Uploads file to blob.

        Args:
            blob_name: Blob name, possibly including multiple path segments.
            file_path: Local file path.
            overwrite: Whether the blob to be uploaded should overwrite the current data.
                       If True, upload_blob will overwrite the existing data. If set to False, the
                       operation will fail with ResourceExistsError
        """
        with open(file_path, "rb") as data:
            super().upload_blob(blob_name, data, overwrite=overwrite)

    def download_blobs_to_files(
        self,
        blob_filename_pairs: Iterable[Tuple[str, str]],
        concurrency_limit: int = 1000,
        verbose: bool = False,
    ) -> int:
        """Downloads a list of files from an azure blob container.

        Args:
            blob_filename_pairs: List[Tuple[str, str]]:List of blob and local path pairs
            concurrency_limit: Maximum number of threads.
            verbose: controls verbosity of the function.
        """
        if verbose:
            print("Checking files exist locally...", end="")

        start_time = time.perf_counter()

        with ThreadPoolExecutor(concurrency_limit) as executor:
            futures_local_exists = [
                (blob_path, local_file_path, executor.submit(self.local_copy_exists, blob_path, local_file_path))
                for blob_path, local_file_path in blob_filename_pairs
            ]

            files_to_download = [(future[0], future[1]) for future in futures_local_exists if not future[2].result()]

        finished_time = time.perf_counter()

        if verbose:
            print(f"done in {finished_time-start_time} seconds")

        # Downloading with overwrite set to true as we already checked these file do not exist locally.

        if verbose:
            print(f"Downloading {len(files_to_download)} files...", end="")

        start_time = time.perf_counter()

        with ThreadPoolExecutor(concurrency_limit) as executor:
            futures_download = [
                executor.submit(self.download_blob_to_file, blob_path, local_file_path, overwrite=True, verbose=False)
                for blob_path, local_file_path in files_to_download
            ]

            # Getting the results in order to wait for all download to complete
            # and raise the potential errors that are thrown in any subtask.
            for future in futures_download:
                future.result()

        nb_downloaded = len(files_to_download)

        finished_time = time.perf_counter()

        if verbose:
            print(f"done in {finished_time-start_time} seconds")

        return nb_downloaded