refactored the download module to have reusable clients

mosaicml · Oct 28, 2024 · 571b056 · 571b056
1 parent ecca820
commit 571b056
Show file tree

Hide file tree

Showing 7 changed files with 857 additions and 644 deletions.
diff --git a/streaming/base/storage/__init__.py b/streaming/base/storage/__init__.py
@@ -3,16 +3,16 @@
 
 """Base module for downloading/uploading files from/to cloud storage."""
 # isort: off
-from streaming.base.storage.download import (
-    download_file, download_from_alipan, download_from_azure, download_from_azure_datalake,
-    download_from_databricks_unity_catalog, download_from_dbfs, download_from_gcs,
-    download_from_hf, download_from_local, download_from_oci, download_from_s3, download_from_sftp)
+from streaming.base.storage.download import (CloudDownloader, S3Downloader, SFTPDownloader,
+                                             GCSDownloader, OCIDownloader, AzureDownloader,
+                                             AzureDataLakeDownloader, HFDownloader,
+                                             DatabricksUnityCatalogDownloader, DBFSDownloader,
+                                             AlipanDownloader, LocalDownloader)
 from streaming.base.storage.upload import (AzureDataLakeUploader, AzureUploader, CloudUploader,
                                            GCSUploader, HFUploader, LocalUploader, OCIUploader,
                                            S3Uploader)
 
 __all__ = [
-    'download_file',
     'CloudUploader',
     'S3Uploader',
     'GCSUploader',
@@ -21,15 +21,16 @@
     'AzureUploader',
     'AzureDataLakeUploader',
     'HFUploader',
-    'download_from_s3',
-    'download_from_sftp',
-    'download_from_gcs',
-    'download_from_oci',
-    'download_from_azure',
-    'download_from_azure_datalake',
-    'download_from_databricks_unity_catalog',
-    'download_from_dbfs',
-    'download_from_alipan',
-    'download_from_local',
-    'download_from_hf',
+    'CloudDownloader',
+    'S3Downloader',
+    'SFTPDownloader',
+    'GCSDownloader',
+    'OCIDownloader',
+    'AzureDownloader',
+    'AzureDataLakeDownloader',
+    'HFDownloader',
+    'DatabricksUnityCatalogDownloader',
+    'DBFSDownloader',
+    'AlipanDownloader',
+    'LocalDownloader',
 ]