Repository: mosaicml/streaming (language: python)

Issue:
Support for azure Data Lake Gen2 type storage
## 🚀 Feature Request
The current support for Azure only covers the blob type storage and not the Data Lake Gen2 type storage which is more generally used for AI workloads. These support hierarchical namespaces in comparison to the flat namespaces in the case of blob storage. Also these type of storage are optimized for Big Data analytics and training on large sets of data.

## Motivation

With PR #256, we can use Streaming with azure blob type storage. I tried it out and it worked seamlessly. But when I tried it with Data Lake type storage it crashed, mainly due to differences in the way these two are handled by the sdk. 

## Implementation
 The code additions to allow their usage is minimal:

```
   # Create a new session per thread
    service = DataLakeServiceClient(
        account_url=f"https://{os.environ['AZURE_ACCOUNT_NAME']}.dfs.core.windows.net",
        credential= os.environ['AZURE_ACCOUNT_ACCESS_KEY'])
    try:
        file_client = service.get_file_client(file_system=obj.netloc, file_path=obj.path.lstrip('/'))
        with open(local, 'wb') as my_blob:
            file_data = file_client.download_file()
            file_data.readinto(my_blob)
    except ResourceNotFoundError:
        raise FileNotFoundError(f'Object {remote} not found.')
    except Exception:
        raise
```

We need only add a new service client for Data Lakes and change the corresponding  _get_file_client_ method. Also  a corresponding uploader needs to be created.

The only reason I haven't created a PR yet is that it is not yet clear to me what may be the best way to differentiate between calling either of the clients. Can we have a flag for the storage type ( storage_type = blob  or storage_type = datalake), hat can be set by the user? Which other files may need changing because of this? Please suggest and I can contribute.
Support for azure Data Lake Gen2 type storage
## 🚀 Feature Request
The current support for Azure only covers the blob type storage and not the Data Lake Gen2 type storage which is more generally used for AI workloads. These support hierarchical namespaces in comparison to the flat namespaces in the case of blob storage. Also these type of storage are optimized for Big Data analytics and training on large sets of data.

## Motivation

With PR #256, we can use Streaming with azure blob type storage. I tried it out and it worked seamlessly. But when I tried it with Data Lake type storage it crashed, mainly due to differences in the way these two are handled by the sdk. 

## Implementation
 The code additions to allow their usage is minimal:

```
   # Create a new session per thread
    service = DataLakeServiceClient(
        account_url=f"https://{os.environ['AZURE_ACCOUNT_NAME']}.dfs.core.windows.net",
        credential= os.environ['AZURE_ACCOUNT_ACCESS_KEY'])
    try:
        file_client = service.get_file_client(file_system=obj.netloc, file_path=obj.path.lstrip('/'))
        with open(local, 'wb') as my_blob:
            file_data = file_client.download_file()
            file_data.readinto(my_blob)
    except ResourceNotFoundError:
        raise FileNotFoundError(f'Object {remote} not found.')
    except Exception:
        raise
```

We need only add a new service client for Data Lakes and change the corresponding  _get_file_client_ method. Also  a corresponding uploader needs to be created.

The only reason I haven't created a PR yet is that it is not yet clear to me what may be the best way to differentiate between calling either of the clients. Can we have a flag for the storage type ( storage_type = blob  or storage_type = datalake), hat can be set by the user? Which other files may need changing because of this? Please suggest and I can contribute.


Discussion / hints:
Hi @shivshandilya , thanks for showing interest in the streaming dataset and interest in contribution. When someone uses Azure Data Lake Gen2 type storage, is the data location prefix path starts from `azure://` like blob type or is it different ? 
To answer simply, its the same. But there is a catch. In azure the data can be accessed in many ways. Simple URIs like **azure://datastores/<-data_store_name->/path/** can be used to access the data, but they need account details like account_url and key. They can be URIs of different types which differ in the way the data is accessed and how easily(all dependent on the type of data asset we create) . The example above is just one way. [Here](https://learn.microsoft.com/en-us/azure/machine-learning/how-to-create-data-assets?tabs=Python-SDK&view=azureml-api-2) are more ways. The main info that we need is always the container name and the path to the data is the container, which a URI or path like **azure://container_name/path-to-data**  is sufficient to convey along with the account details. This is the way the PR #256 did it. Now if we want to access data lakes instead of blobs they can have same or different URIs depending on the access type we want but to make it simple we can stick to the **azure://** format, otherwise we would have to handle more prefixes. I suggest that we add support for a custom prefix like **azure-dl** (for azure Data Lake). This way we can make minimalistic changes and also respect the way azure present the data(in some way, only the prefix is different, everything else in the URI or path is the same).


Hi @shivshandilya, thanks for the detailed explanation. I liked the approach of a custom prefix like `azure-dl` (for azure Data Lake). Also, if we keep the same prefix `azure`, I am guessing there is no way to identify at runtime whether this is a data lake or a blob storage right ? Since the credential setup is the same and both need an environment variables setup AZURE_ACCOUNT_NAME and AZURE_ACCOUNT_ACCESS_KEY. 
Yes, @karan6181. There is no way to know the storage type programmatically during runtime up to my knowledge. and yes they both need the same variables.

Hi @shivshandilya , thanks for confirming. I am okay going with a custom prefix like `azure-dl` (for azure Data Lake). If you have a bandwidth, do you mind creating the PR for the same ? I think you already know where and what to change in the code base, but feel free to reach out to me in case you need any help! Happy to guide you! Thank you for your contribution!  
Sure @karan6181. I will create a PR for the same.
Hi @shivshandilya , thanks for showing interest in the streaming dataset and interest in contribution. When someone uses Azure Data Lake Gen2 type storage, is the data location prefix path starts from `azure://` like blob type or is it different ? 
To answer simply, its the same. But there is a catch. In azure the data can be accessed in many ways. Simple URIs like **azure://datastores/<-data_store_name->/path/** can be used to access the data, but they need account details like account_url and key. They can be URIs of different types which differ in the way the data is accessed and how easily(all dependent on the type of data asset we create) . The example above is just one way. [Here](https://learn.microsoft.com/en-us/azure/machine-learning/how-to-create-data-assets?tabs=Python-SDK&view=azureml-api-2) are more ways. The main info that we need is always the container name and the path to the data is the container, which a URI or path like **azure://container_name/path-to-data**  is sufficient to convey along with the account details. This is the way the PR #256 did it. Now if we want to access data lakes instead of blobs they can have same or different URIs depending on the access type we want but to make it simple we can stick to the **azure://** format, otherwise we would have to handle more prefixes. I suggest that we add support for a custom prefix like **azure-dl** (for azure Data Lake). This way we can make minimalistic changes and also respect the way azure present the data(in some way, only the prefix is different, everything else in the URI or path is the same).


Hi @shivshandilya, thanks for the detailed explanation. I liked the approach of a custom prefix like `azure-dl` (for azure Data Lake). Also, if we keep the same prefix `azure`, I am guessing there is no way to identify at runtime whether this is a data lake or a blob storage right ? Since the credential setup is the same and both need an environment variables setup AZURE_ACCOUNT_NAME and AZURE_ACCOUNT_ACCESS_KEY. 
Yes, @karan6181. There is no way to know the storage type programmatically during runtime up to my knowledge. and yes they both need the same variables.

Hi @shivshandilya , thanks for confirming. I am okay going with a custom prefix like `azure-dl` (for azure Data Lake). If you have a bandwidth, do you mind creating the PR for the same ? I think you already know where and what to change in the code base, but feel free to reach out to me in case you need any help! Happy to guide you! Thank you for your contribution!  
Sure @karan6181. I will create a PR for the same.

Fix patch:
diff --git a/docs/source/how_to_guides/configure_cloud_storage_credentials.md b/docs/source/how_to_guides/configure_cloud_storage_credentials.md
index 46909ad6a..9f4c0ec5f 100644
--- a/docs/source/how_to_guides/configure_cloud_storage_credentials.md
+++ b/docs/source/how_to_guides/configure_cloud_storage_credentials.md
@@ -145,7 +145,7 @@ region=us-ashburn-1
 
 The key file (`~/.oci/oci_api_key.pem`) is a PEM file that would look like a typical RSA private key file. The streaming dataset authenticates the credentials by reading the `~/.oci/config` and `~/.oci/oci_api_key.pem`.
 
-## Azure Blob Storage
+## Azure Blob Storage and Azure DataLake
 
 If you wish to create a new storage account, you can use the [Azure Portal](https://docs.microsoft.com/azure/storage/common/storage-quickstart-create-account?tabs=azure-portal), [Azure PowerShell](https://docs.microsoft.com/azure/storage/common/storage-quickstart-create-account?tabs=azure-powershell), or [Azure CLI](https://docs.microsoft.com/azure/storage/common/storage-quickstart-create-account?tabs=azure-cli):
 
diff --git a/setup.py b/setup.py
index f8894ca96..408ba5799 100644
--- a/setup.py
+++ b/setup.py
@@ -56,6 +56,8 @@
     'zstd>=1.5.2.5,<2',
     'oci>=2.88,<3',
     'azure-storage-blob>=12.0.0,<13',
+    'azure-storage-file-datalake>=12.11.0,<13',
+    'azure-identity>=1.13.0',
 ]
 
 extra_deps = {}
diff --git a/streaming/base/storage/__init__.py b/streaming/base/storage/__init__.py
index cb83bc5c6..fffdecbd8 100644
--- a/streaming/base/storage/__init__.py
+++ b/streaming/base/storage/__init__.py
@@ -4,10 +4,10 @@
 """Base module for downloading/uploading files from/to cloud storage."""
 
 from streaming.base.storage.download import download_file, download_or_wait
-from streaming.base.storage.upload import (AzureUploader, CloudUploader, GCSUploader,
-                                           LocalUploader, OCIUploader, S3Uploader)
+from streaming.base.storage.upload import (AzureDataLakeUploader, AzureUploader, CloudUploader,
+                                           GCSUploader, LocalUploader, OCIUploader, S3Uploader)
 
 __all__ = [
     'download_file', 'download_or_wait', 'CloudUploader', 'S3Uploader', 'GCSUploader',
-    'OCIUploader', 'LocalUploader', 'AzureUploader'
+    'OCIUploader', 'LocalUploader', 'AzureUploader', 'AzureDataLakeUploader'
 ]
diff --git a/streaming/base/storage/download.py b/streaming/base/storage/download.py
index f1499094b..bdcf55378 100644
--- a/streaming/base/storage/download.py
+++ b/streaming/base/storage/download.py
@@ -248,6 +248,38 @@ def download_from_azure(remote: str, local: str) -> None:
         raise
 
 
+def download_from_azure_datalake(remote: str, local: str) -> None:
+    """Download a file from remote Microsoft Azure Data Lake to local.
+
+    Args:
+        remote (str): Remote path (azure).
+        local (str): Local path (local filesystem).
+    """
+    from azure.core.exceptions import ResourceNotFoundError
+    from azure.storage.filedatalake import DataLakeServiceClient
+
+    obj = urllib.parse.urlparse(remote)
+    if obj.scheme != 'azure-dl':
+        raise ValueError(
+            f'Expected obj.scheme to be `azure-dl`, got {obj.scheme} for remote={remote}')
+
+    # Create a new session per thread
+    service = DataLakeServiceClient(
+        account_url=f"https://{os.environ['AZURE_ACCOUNT_NAME']}.dfs.core.windows.net",
+        credential=os.environ['AZURE_ACCOUNT_ACCESS_KEY'],
+    )
+    try:
+        file_client = service.get_file_client(file_system=obj.netloc,
+                                              file_path=obj.path.lstrip('/'))
+        with open(local, 'wb') as my_file:
+            file_data = file_client.download_file()
+            file_data.readinto(my_file)
+    except ResourceNotFoundError:
+        raise FileNotFoundError(f'Object {remote} not found.')
+    except Exception:
+        raise
+
+
 def download_from_local(remote: str, local: str) -> None:
     """Download a file from remote to local.
 
@@ -294,6 +326,8 @@ def download_file(remote: Optional[str], local: str, timeout: float):
         download_from_oci(remote, local)
     elif remote.startswith('azure://'):
         download_from_azure(remote, local)
+    elif remote.startswith('azure-dl://'):
+        download_from_azure_datalake(remote, local)
     else:
         download_from_local(remote, local)
 
diff --git a/streaming/base/storage/upload.py b/streaming/base/storage/upload.py
index b7d496a98..2ffe4e327 100644
--- a/streaming/base/storage/upload.py
+++ b/streaming/base/storage/upload.py
@@ -16,7 +16,8 @@
 from streaming.base.storage.download import BOTOCORE_CLIENT_ERROR_CODES
 
 __all__ = [
-    'CloudUploader', 'S3Uploader', 'GCSUploader', 'OCIUploader', 'AzureUploader', 'LocalUploader'
+    'CloudUploader', 'S3Uploader', 'GCSUploader', 'OCIUploader', 'AzureUploader',
+    'AzureDataLakeUploader', 'LocalUploader'
 ]
 
 logger = logging.getLogger(__name__)
@@ -26,6 +27,7 @@
     'gs': 'GCSUploader',
     'oci': 'OCIUploader',
     'azure': 'AzureUploader',
+    'azure-dl': 'AzureDataLakeUploader',
     '': 'LocalUploader',
 }
 
@@ -475,6 +477,81 @@ def check_bucket_exists(self, remote: str):
                 f'or check the bucket permission.',)
 
 
+class AzureDataLakeUploader(CloudUploader):
+    """Upload file from local machine to Microsoft Azure DataLake.
+
+    Args:
+        out (str | Tuple[str, str]): Output dataset directory to save shard files.
+
+            1. If ``out`` is a local directory, shard files are saved locally.
+            2. If ``out`` is a remote directory, a local temporary directory is created to
+               cache the shard files and then the shard files are uploaded to a remote
+               location. At the end, the temp directory is deleted once shards are uploaded.
+            3. If ``out`` is a tuple of ``(local_dir, remote_dir)``, shard files are saved in
+               the `local_dir` and also uploaded to a remote location.
+        keep_local (bool): If the dataset is uploaded, whether to keep the local dataset
+            shard file or remove it after uploading. Defaults to ``False``.
+        progress_bar (bool): Display TQDM progress bars for uploading output dataset files to
+            a remote location. Default to ``False``.
+    """
+
+    def __init__(self,
+                 out: Union[str, Tuple[str, str]],
+                 keep_local: bool = False,
+                 progress_bar: bool = False) -> None:
+        super().__init__(out, keep_local, progress_bar)
+
+        from azure.storage.filedatalake import DataLakeServiceClient
+
+        # Create a session and use it to make our client. Unlike Resources and Sessions,
+        # clients are generally thread-safe.
+        self.azure_service = DataLakeServiceClient(
+            account_url=f"https://{os.environ['AZURE_ACCOUNT_NAME']}.dfs.core.windows.net",
+            credential=os.environ['AZURE_ACCOUNT_ACCESS_KEY'])
+        self.check_container_exists(self.remote)  # pyright: ignore
+
+    def upload_file(self, filename: str):
+        """Upload file from local instance to Azure DataLalke container.
+
+        Args:
+            filename (str): File to upload.
+        """
+        local_filename = os.path.join(self.local, filename)
+        local_filename = local_filename.replace('\\', '/')
+        remote_filename = os.path.join(self.remote, filename)  # pyright: ignore
+        remote_filename = remote_filename.replace('\\', '/')
+        obj = urllib.parse.urlparse(remote_filename)
+        logger.debug(f'Uploading to {remote_filename}')
+        file_size = os.stat(local_filename).st_size
+        file_client = self.azure_service.get_file_client(file_system=obj.netloc,
+                                                         file_path=obj.path.lstrip('/'))
+
+        with tqdm.tqdm(total=file_size,
+                       unit='B',
+                       unit_scale=True,
+                       desc=f'Uploading to {remote_filename}',
+                       disable=(not self.progress_bar)) as pbar:
+            with open(local_filename, 'rb') as data:
+                file_client.upload_data(data=data, overwrite=True)
+                pbar.update(file_size)
+        self.clear_local(local=local_filename)
+
+    def check_container_exists(self, remote: str):
+        """Raise an exception if the container does not exist.
+
+        Args:
+            remote (str): azure container path.
+
+        Raises:
+            error: Container does not exist.
+        """
+        container_name = urllib.parse.urlparse(remote).netloc
+        if self.azure_service.get_file_system_client(file_system=container_name).exists() is False:
+            raise FileNotFoundError(
+                f'Either container `{container_name}` does not exist! ' +
+                f'or check the container permission.',)
+
+
 class LocalUploader(CloudUploader):
     """Copy file from one local directory to another local directory.
 


Test patch:
diff --git a/tests/test_download.py b/tests/test_download.py
index fc6ccbbf6..419bfc789 100644
--- a/tests/test_download.py
+++ b/tests/test_download.py
@@ -10,7 +10,8 @@
 import pytest
 from botocore.exceptions import ClientError
 
-from streaming.base.storage.download import (download_file, download_from_azure, download_from_gcs,
+from streaming.base.storage.download import (download_file, download_from_azure,
+                                             download_from_azure_datalake, download_from_gcs,
                                              download_from_local, download_from_s3,
                                              download_or_wait)
 from tests.conftest import GCS_URL, MY_BUCKET, R2_URL
@@ -42,6 +43,7 @@ def test_invalid_cloud_prefix(self, remote_local_file: Any):
             mock_remote_filepath, mock_local_filepath = remote_local_file(
                 cloud_prefix='aaazure://')
             download_from_azure(mock_remote_filepath, mock_local_filepath)
+            download_from_azure_datalake(mock_remote_filepath, mock_local_filepath)
 
 
 class TestS3Client:
@@ -149,6 +151,15 @@ def test_download_from_azure_gets_called(self, mocked_requests: Mock, remote_loc
         mocked_requests.assert_called_once()
         mocked_requests.assert_called_once_with(mock_remote_filepath, mock_local_filepath)
 
+    @patch('streaming.base.storage.download.download_from_azure_datalake')
+    @pytest.mark.usefixtures('remote_local_file')
+    def test_download_from_azure_datalake_gets_called(self, mocked_requests: Mock,
+                                                      remote_local_file: Any):
+        mock_remote_filepath, mock_local_filepath = remote_local_file(cloud_prefix='azure-dl://')
+        download_file(mock_remote_filepath, mock_local_filepath, 60)
+        mocked_requests.assert_called_once()
+        mocked_requests.assert_called_once_with(mock_remote_filepath, mock_local_filepath)
+
     @patch('streaming.base.storage.download.download_from_sftp')
     @pytest.mark.usefixtures('remote_local_file')
     def test_download_from_sftp_gets_called(self, mocked_requests: Mock, remote_local_file: Any):
diff --git a/tests/test_upload.py b/tests/test_upload.py
index 979c9df63..ed52acedb 100644
--- a/tests/test_upload.py
+++ b/tests/test_upload.py
@@ -9,8 +9,8 @@
 
 import pytest
 
-from streaming.base.storage.upload import (AzureUploader, CloudUploader, GCSUploader,
-                                           LocalUploader, S3Uploader)
+from streaming.base.storage.upload import (AzureDataLakeUploader, AzureUploader, CloudUploader,
+                                           GCSUploader, LocalUploader, S3Uploader)
 from tests.conftest import R2_URL
 
 
@@ -224,6 +224,39 @@ def test_local_directory_is_empty(self, local_remote_dir: Tuple[str, str]):
             _ = AzureUploader(out=local)
 
 
+class TestAzureDataLakeUploader:
+
+    @patch('streaming.base.storage.upload.AzureDataLakeUploader.check_container_exists')
+    @pytest.mark.usefixtures('azure_credentials')
+    @pytest.mark.parametrize('out',
+                             ['azure://container/dir', ('./dir1', 'azure://container/dir/')])
+    def test_instantiation(self, mocked_requests: Mock, out: Any):
+        mocked_requests.side_effect = None
+        _ = AzureDataLakeUploader(out=out)
+        if not isinstance(out, str):
+            shutil.rmtree(out[0])
+
+    @pytest.mark.parametrize('out', ['ss4://container/dir'])
+    def test_invalid_remote_str(self, out: str):
+        with pytest.raises(ValueError, match=f'Invalid Cloud provider prefix.*'):
+            _ = AzureDataLakeUploader(out=out)
+
+    @pytest.mark.parametrize('out', ['ss4://container/dir', ('./dir1', 'gcs://container/dir/')])
+    def test_invalid_remote_list(self, out: Any):
+        with pytest.raises(ValueError, match=f'Invalid Cloud provider prefix.*'):
+            _ = AzureDataLakeUploader(out=out)
+
+    def test_local_directory_is_empty(self, local_remote_dir: Tuple[str, str]):
+        with pytest.raises(FileExistsError, match=f'Directory is not empty.*'):
+            local, _ = local_remote_dir
+            os.makedirs(local, exist_ok=True)
+            local_file_path = os.path.join(local, 'file.txt')
+            # Creating an empty file at specified location
+            with open(local_file_path, 'w') as _:
+                pass
+            _ = AzureDataLakeUploader(out=local)
+
+
 class TestLocalUploader:
 
     def test_upload_file(self, local_remote_dir: Tuple[str, str]):