Merge pull request #10 from saturncloud/bugfix/s3anon

Adding option to set anon=True which will let user access S3 bucket anonymously

Author: skirmer

dask_pytorch_ddp/data.py

@@ -34,13 +34,19 @@ def _list_all_files(bucket: str, prefix: str, s3_client=None) -> List[str]:
     return all_files
 
 
-def _read_s3_fileobj(bucket, path, fileobj):
+def _read_s3_fileobj(bucket, path, fileobj, anon=False):
     """
     read an obj from s3 to a file like object
     """
     import boto3  # pylint: disable=import-outside-toplevel
+    from botocore import UNSIGNED  # pylint: disable=import-outside-toplevel
+    from botocore.client import Config  # pylint: disable=import-outside-toplevel
+
+    if anon:
+        s3 = boto3.resource("s3", config=Config(signature_version=UNSIGNED))
+    else:
+        s3 = boto3.resource("s3")
 
-    s3 = boto3.resource("s3")
     bucket = s3.Bucket(bucket)
     bucket.download_fileobj(path, fileobj)
     fileobj.seek(0)
@@ -59,12 +65,16 @@ class S3ImageFolder(Dataset):
     An image folder that lives in S3.  Directories containing the image are classes.
     """
 
+    # pylint: disable=too-many-instance-attributes
+    # pylint: disable=too-many-arguments
+
     def __init__(
         self,
         s3_bucket: str,
         s3_prefix: str,
         transform: Optional[Callable] = None,
         target_transform: Optional[Callable] = None,
+        anon: Optional[bool] = False,
     ):
         self.s3_bucket = s3_bucket
         self.s3_prefix = s3_prefix
@@ -73,6 +83,7 @@ def __init__(
         self.class_to_idx = {k: idx for idx, k in enumerate(self.classes)}
         self.transform = transform
         self.target_transform = target_transform
+        self.anon = anon
 
     @classmethod
     def _get_class(cls, path):
@@ -88,7 +99,7 @@ def __getitem__(self, idx):
         path = self.all_files[idx]
         label = self.class_to_idx[self._get_class(path)]
         with tempfile.TemporaryFile() as f:
-            f = _read_s3_fileobj(self.s3_bucket, path, f)
+            f = _read_s3_fileobj(self.s3_bucket, path, f, self.anon)
             img = _load_image_obj(f)
         if self.transform is not None:
             img = self.transform(img)

tests/test_data.py

@@ -40,7 +40,7 @@ def test_image_folder_getitem():
         read_s3_fileobj.return_value = Mock()
         load_image_obj.return_value = Mock()
         val, label = folder[0]
-        read_s3_fileobj.assert_called_once_with("fake-bucket", fake_file_list[0], ANY)
+        read_s3_fileobj.assert_called_once_with("fake-bucket", fake_file_list[0], ANY, False)
         load_image_obj.assert_called_once_with(read_s3_fileobj())
         assert val == load_image_obj.return_value
         assert label == 1