Add a regression test for StreamingDataset using cloud providers

regression/cloud_providers.yaml

@@ -0,0 +1,27 @@
+name: streaming-regression-test-cloud-providers
+compute:
+  cluster: r1z1
+gpu_num: 8
+command: |-
+  pip uninstall -y mosaicml-streaming
+  cd streaming
+  pip install -e '.[dev]'
+  composer -n 1 regression/synthetic_dataset.py --create --cloud gs
+  composer -n 2 regression/iterate_data.py --cloud gs --check_download
+  composer -n 1 regression/synthetic_dataset.py --delete --cloud gs
+  composer -n 1 regression/synthetic_dataset.py --create --cloud s3
+  composer -n 4 regression/iterate_data.py --cloud s3 --check_download
+  composer -n 1 regression/synthetic_dataset.py --delete --cloud s3
+  composer -n 1 regression/synthetic_dataset.py --create --cloud oci
+  composer -n 8 regression/iterate_data.py --cloud oci --check_download
+  composer -n 1 regression/synthetic_dataset.py --delete --cloud oci
+
+image: mosaicml/composer:0.15.0
+scheduling:
+  resumable: true
+  priority: medium
+integrations:
+  - integration_type: git_repo
+    git_repo: b-chu/streaming
+    git_branch: test_cloud_providers
+    ssh_clone: false

regression/iterate_data.py

@@ -6,12 +6,15 @@
 import os
 import shutil
 import tempfile
+import urllib.parse
 from argparse import ArgumentParser, Namespace
 
+import utils
+from torch import distributed as dist
 from torch.utils.data import DataLoader
 
 from streaming import StreamingDataset
-from streaming.base.distributed import barrier
+from streaming.base.distributed import barrier, maybe_init_dist
 
 _TRAIN_EPOCHS = 2
 
@@ -35,6 +38,8 @@ def parse_args() -> tuple[Namespace, dict[str, str]]:
         tuple(Namespace, dict[str, str]): Command-line arguments and named arguments.
     """
     args = ArgumentParser()
+    args.add_argument('--cloud', type=str)
+    args.add_argument('--check_download', default=False, action='store_true')
     args.add_argument('--local', default=False, action='store_true')
     args.add_argument(
         '--keep_zip',
@@ -56,18 +61,57 @@ def parse_args() -> tuple[Namespace, dict[str, str]]:
     return args, kwargs
 
 
+def get_file_count(cloud: str) -> int:
+    """Get the number of files in a remote directory.
+
+    Args:
+        cloud (str): Cloud provider.
+    """
+    remote_dir = utils.get_remote_dir(cloud)
+    obj = urllib.parse.urlparse(remote_dir)
+    files = []
+    if cloud == 'gs':
+        from google.cloud.storage import Bucket, Client
+
+        service_account_path = os.environ['GOOGLE_APPLICATION_CREDENTIALS']
+        gcs_client = Client.from_service_account_json(service_account_path)
+
+        bucket = Bucket(gcs_client, obj.netloc)
+        files = bucket.list_blobs(prefix=obj.path.lstrip('/'))
+    elif cloud == 's3':
+        import boto3
+
+        s3 = boto3.resource('s3')
+        bucket = s3.Bucket(obj.netloc)
+        files = bucket.objects.filter(Prefix=obj.path.lstrip('/'))
+    elif cloud == 'oci':
+        import oci
+
+        config = oci.config.from_file()
+        oci_client = oci.object_storage.ObjectStorageClient(
+            config=config, retry_strategy=oci.retry.DEFAULT_RETRY_STRATEGY)
+        namespace = oci_client.get_namespace().data
+        objects = oci_client.list_objects(namespace, obj.netloc, prefix=obj.path.lstrip('/'))
+
+        files = objects.data.objects
+
+    return sum(1 for _ in files)
+
+
 def main(args: Namespace, kwargs: dict[str, str]) -> None:
     """Benchmark time taken to generate the epoch for a given dataset.
 
     Args:
         args (Namespace): Command-line arguments.
         kwargs (dict): Named arguments.
     """
+    # Initialize torch dist ourselves, if necessary.
+    destroy_dist = maybe_init_dist()
+
     tmp_dir = tempfile.gettempdir()
-    tmp_remote_dir = os.path.join(tmp_dir, 'regression_remote')
     tmp_download_dir = os.path.join(tmp_dir, 'test_iterate_data_download')
     dataset = StreamingDataset(
-        remote=tmp_remote_dir,
+        remote=utils.get_remote_dir(args.cloud),
         local=tmp_download_dir if args.local else None,
         split=kwargs.get('split'),
         download_retry=int(kwargs.get('download_retry', 2)),
@@ -92,12 +136,23 @@ def main(args: Namespace, kwargs: dict[str, str]) -> None:
         for _ in dataloader:
             pass
 
+    if args.check_download and args.cloud is not None:
+        num_cloud_files = get_file_count(args.cloud)
+        local_dir = dataset.streams[0].local
+        num_local_files = len([
+            name for name in os.listdir(local_dir) if os.path.isfile(os.path.join(local_dir, name))
+        ])
+        assert num_cloud_files == num_local_files
+
     barrier()
     # Clean up directories
     for stream in dataset.streams:
         shutil.rmtree(stream.local, ignore_errors=True)
     shutil.rmtree(tmp_download_dir, ignore_errors=True)
 
+    if destroy_dist:
+        dist.destroy_process_group()
+
 
 if __name__ == '__main__':
     args, kwargs = parse_args()

regression/iterate_data.yaml

@@ -23,7 +23,6 @@ command: |-
   composer -n 8 regression/iterate_data.py --batch_size 1000
   composer -n 8 regression/iterate_data.py --shuffle --shuffle_algo py1b --shuffle_seed 12 --shuffle_block_size 10000
   composer -n 1 regression/synthetic_dataset.py --delete
-
 image: mosaicml/composer:0.15.0
 scheduling:
   resumable: true

regression/synthetic_dataset.py

@@ -5,11 +5,12 @@
 
 import os
 import shutil
-import tempfile
+import urllib.parse
 from argparse import ArgumentParser, Namespace
 from typing import Union
 
 import numpy as np
+import utils
 
 from streaming import MDSWriter
 
@@ -32,6 +33,7 @@ def parse_args() -> Namespace:
         Namespace: Command-line arguments.
     """
     args = ArgumentParser()
+    args.add_argument('--cloud', type=str)
     args.add_argument('--create', default=False, action='store_true')
     args.add_argument('--delete', default=False, action='store_true')
     args.add_argument(
@@ -98,19 +100,71 @@ def get_dataset(num_samples: int) -> list[dict[str, Union[int, str]]]:
     return samples
 
 
+def delete_gcs(remote_dir: str) -> None:
+    """Delete a remote directory from gcs.
+
+    Args:
+        remote_dir (str): Location of the remote directory.
+    """
+    from google.cloud.storage import Bucket, Client
+
+    service_account_path = os.environ['GOOGLE_APPLICATION_CREDENTIALS']
+    gcs_client = Client.from_service_account_json(service_account_path)
+    obj = urllib.parse.urlparse(remote_dir)
+
+    bucket = Bucket(gcs_client, obj.netloc)
+    blobs = bucket.list_blobs(prefix=obj.path.lstrip('/'))
+
+    for blob in blobs:
+        blob.delete()
+
+
+def delete_s3(remote_dir: str) -> None:
+    """Delete a remote directory from s3.
+
+    Args:
+        remote_dir (str): Location of the remote directory.
+    """
+    import boto3
+
+    obj = urllib.parse.urlparse(remote_dir)
+
+    s3 = boto3.resource('s3')
+    bucket = s3.Bucket(obj.netloc)
+    bucket.objects.filter(Prefix=obj.path.lstrip('/')).delete()
+
+
+def delete_oci(remote_dir: str) -> None:
+    """Delete a remote directory from oci.
+
+    Args:
+        remote_dir (str): Location of the remote directory.
+    """
+    import oci
+
+    obj = urllib.parse.urlparse(remote_dir)
+
+    config = oci.config.from_file()
+    oci_client = oci.object_storage.ObjectStorageClient(
+        config=config, retry_strategy=oci.retry.DEFAULT_RETRY_STRATEGY)
+    namespace = oci_client.get_namespace().data
+    objects = oci_client.list_objects(namespace, obj.netloc, prefix=obj.path.lstrip('/'))
+
+    for filenames in objects.data.objects:
+        oci_client.delete_object(namespace, obj.netloc, filenames.name)
+
+
 def main(args: Namespace) -> None:
     """Benchmark time taken to generate the epoch for a given dataset.
 
     Args:
         args (Namespace): Command-line arguments.
     """
-    tmp_dir = tempfile.gettempdir()
-    tmp_remote_dir = os.path.join(tmp_dir, 'regression_remote')
-
+    remote_dir = utils.get_remote_dir(args.cloud)
     if args.create:
         dataset = get_dataset(_NUM_SAMPLES)
         with MDSWriter(
-                out=tmp_remote_dir,
+                out=remote_dir,
                 columns=_COLUMNS,
                 compression=args.compression,
                 hashes=args.hashes,
@@ -119,7 +173,14 @@ def main(args: Namespace) -> None:
             for sample in dataset:
                 out.write(sample)
     if args.delete:
-        shutil.rmtree(tmp_remote_dir, ignore_errors=True)
+        if args.cloud is None:
+            shutil.rmtree(remote_dir, ignore_errors=True)
+        elif args.cloud == 'gs':
+            delete_gcs(remote_dir)
+        elif args.cloud == 's3':
+            delete_s3(remote_dir)
+        elif args.cloud == 'oci':
+            delete_oci(remote_dir)
 
 
 if __name__ == '__main__':

regression/utils.py

@@ -0,0 +1,36 @@
+# Copyright 2023 MosaicML Streaming authors
+# SPDX-License-Identifier: Apache-2.0
+
+"""Utility and helper functions for regression testing."""
+
+import os
+import tempfile
+from typing import Optional
+
+_CLOUD_REMOTE_LOCATIONS = {
+    'gs': 'gs://mosaicml-composer-tests/streaming/regression/',
+    's3': 's3://streaming-upload-test-bucket/streaming/regression/',
+    'oci': 'oci://streaming-test/regression'
+}
+
+
+def get_remote_dir(storage: Optional[str]) -> str:
+    """Get an remote directory.
+
+    Args:
+        storage (str): Type of storage to use.
+
+    Returns:
+        str: Remote directory.
+    """
+    if storage is None:
+        return get_local_remote_dir()
+    else:
+        return _CLOUD_REMOTE_LOCATIONS[storage]
+
+
+def get_local_remote_dir() -> str:
+    """Get a local remote directory."""
+    tmp_dir = tempfile.gettempdir()
+    tmp_remote_dir = os.path.join(tmp_dir, 'regression_remote')
+    return tmp_remote_dir