Added support for passing S3 endpoint URL (#124)

* Added support for passing S3 endpoint as an argument to dataset construction API to access dataset stored in AWS S3 compatible storage systems
* Make the endpoint argument optional by setting empty string as default value to S3IterableDataset constructor

Author: bharatendra

s3torchconnector/src/s3torchconnector/_s3client/_s3client.py

@@ -32,8 +32,9 @@ def _identity(obj: Any) -> Any:
 
 
 class S3Client:
-    def __init__(self, region: str):
+    def __init__(self, region: str, endpoint: str = None):
         self._region = region
+        self._endpoint = endpoint
         self._real_client = None
         self._client_pid = None
 
@@ -51,7 +52,9 @@ def region(self) -> str:
 
     def _client_builder(self) -> MountpointS3Client:
         return MountpointS3Client(
-            region=self._region, user_agent_prefix=user_agent_prefix
+            region=self._region,
+            endpoint=self._endpoint,
+            user_agent_prefix=user_agent_prefix,
         )
 
     def get_object(self, bucket: str, key: str) -> S3Reader:

s3torchconnector/src/s3torchconnector/s3checkpoint.py

@@ -16,9 +16,10 @@ class S3Checkpoint:
     torch.load, and torch.save.
     """
 
-    def __init__(self, region: str):
+    def __init__(self, region: str, endpoint: str = None):
         self.region = region
-        self._client = S3Client(region)
+        self.endpoint = endpoint
+        self._client = S3Client(region, endpoint)
 
     def reader(self, s3_uri: str) -> S3Reader:
         """Creates an S3Reader from a given s3_uri.

s3torchconnector/src/s3torchconnector/s3iterable_dataset.py

@@ -29,30 +29,38 @@ def __init__(
         self,
         region: str,
         get_dataset_objects: Callable[[S3Client], Iterable[S3BucketKey]],
+        endpoint: str = None,
         transform: Callable[[S3Reader], Any] = identity,
     ):
         self._get_dataset_objects = get_dataset_objects
         self._transform = transform
         self._region = region
+        self._endpoint = endpoint
         self._client = None
 
     @property
     def region(self):
         return self._region
 
+    @property
+    def endpoint(self):
+        return self._endpoint
+
     @classmethod
     def from_objects(
         cls,
         object_uris: Union[str, Iterable[str]],
         *,
         region: str,
+        endpoint: str = None,
         transform: Callable[[S3Reader], Any] = identity,
     ):
         """Returns an instance of S3IterableDataset using the S3 URI(s) provided.
 
         Args:
           object_uris(str | Iterable[str]): S3 URI of the object(s) desired.
           region(str): AWS region of the S3 bucket where the objects are stored.
+          endpoint(str): AWS endpoint of the S3 bucket where the objects are stored.
           transform: Optional callable which is used to transform an S3Reader into the desired type.
 
         Returns:
@@ -63,7 +71,10 @@ def from_objects(
         """
         log.info(f"Building {cls.__name__} from_objects")
         return cls(
-            region, partial(get_objects_from_uris, object_uris), transform=transform
+            region,
+            partial(get_objects_from_uris, object_uris),
+            endpoint,
+            transform=transform,
         )
 
     @classmethod
@@ -72,13 +83,15 @@ def from_prefix(
         s3_uri: str,
         *,
         region: str,
+        endpoint: str = None,
         transform: Callable[[S3Reader], Any] = identity,
     ):
         """Returns an instance of S3IterableDataset using the S3 URI provided.
 
         Args:
           s3_uri(str): An S3 URI (prefix) of the object(s) desired. Objects matching the prefix will be included in the returned dataset.
           region(str): AWS region of the S3 bucket where the objects are stored.
+          endpoint(str): AWS endpoint of the S3 bucket where the objects are stored.
           transform: Optional callable which is used to transform an S3Reader into the desired type.
 
         Returns:
@@ -89,12 +102,15 @@ def from_prefix(
         """
         log.info(f"Building {cls.__name__} from_prefix {s3_uri=}")
         return cls(
-            region, partial(get_objects_from_prefix, s3_uri), transform=transform
+            region,
+            partial(get_objects_from_prefix, s3_uri),
+            endpoint,
+            transform=transform,
         )
 
     def _get_client(self):
         if self._client is None:
-            self._client = S3Client(self.region)
+            self._client = S3Client(self.region, self.endpoint)
         return self._client
 
     def _get_transformed_object(self, bucket_key: S3BucketKey) -> Any:

s3torchconnector/src/s3torchconnector/s3map_dataset.py

@@ -30,18 +30,24 @@ def __init__(
         self,
         region: str,
         get_dataset_objects: Callable[[S3Client], Iterable[S3BucketKey]],
+        endpoint: str = None,
         transform: Callable[[S3Reader], Any] = identity,
     ):
         self._get_dataset_objects = get_dataset_objects
         self._transform = transform
         self._region = region
+        self._endpoint = endpoint
         self._client = None
         self._bucket_key_pairs = None
 
     @property
     def region(self):
         return self._region
 
+    @property
+    def endpoint(self):
+        return self._endpoint
+
     @property
     def _dataset_bucket_key_pairs(self) -> List[S3BucketKey]:
         if self._bucket_key_pairs is None:
@@ -54,13 +60,15 @@ def from_objects(
         object_uris: Union[str, Iterable[str]],
         *,
         region: str,
+        endpoint: str = None,
         transform: Callable[[S3Reader], Any] = identity,
     ):
         """Returns an instance of S3MapDataset using the S3 URI(s) provided.
 
         Args:
           object_uris(str | Iterable[str]): S3 URI of the object(s) desired.
           region(str): AWS region of the S3 bucket where the objects are stored.
+          endpoint(str): AWS endpoint of the S3 bucket where the objects are stored.
           transform: Optional callable which is used to transform an S3Reader into the desired type.
 
         Returns:
@@ -71,7 +79,10 @@ def from_objects(
         """
         log.info(f"Building {cls.__name__} from_objects")
         return cls(
-            region, partial(get_objects_from_uris, object_uris), transform=transform
+            region,
+            partial(get_objects_from_uris, object_uris),
+            endpoint,
+            transform=transform,
         )
 
     @classmethod
@@ -80,13 +91,15 @@ def from_prefix(
         s3_uri: str,
         *,
         region: str,
+        endpoint: str = None,
         transform: Callable[[S3Reader], Any] = identity,
     ):
         """Returns an instance of S3MapDataset using the S3 URI provided.
 
         Args:
           s3_uri(str): An S3 URI (prefix) of the object(s) desired. Objects matching the prefix will be included in the returned dataset.
           region(str): AWS region of the S3 bucket where the objects are stored.
+          endpoint(str): AWS endpoint of the S3 bucket where the objects are stored.
           transform: Optional callable which is used to transform an S3Reader into the desired type.
 
         Returns:
@@ -97,12 +110,15 @@ def from_prefix(
         """
         log.info(f"Building {cls.__name__} from_prefix {s3_uri=}")
         return cls(
-            region, partial(get_objects_from_prefix, s3_uri), transform=transform
+            region,
+            partial(get_objects_from_prefix, s3_uri),
+            endpoint,
+            transform=transform,
         )
 
     def _get_client(self):
         if self._client is None:
-            self._client = S3Client(self.region)
+            self._client = S3Client(self.region, self.endpoint)
         return self._client
 
     def _get_object(self, i: int) -> S3Reader:

s3torchconnector/tst/unit/test_checkpointing.py

@@ -37,7 +37,7 @@
 TEST_BUCKET = "test-bucket"
 TEST_KEY = "test-key"
 TEST_REGION = "us-east-1"
-
+TEST_ENDPOINT = "https://s3.us-east-1.amazonaws.com"
 
 scalars = (
     none()
@@ -147,6 +147,12 @@ def test_general_checkpointing_untyped_storage_loads_no_modern_pytorch_format(
     )
 
 
+def test_checkpoint_creation_with_region_and_endpoint():
+    checkpoint = S3Checkpoint(TEST_REGION, endpoint=TEST_ENDPOINT)
+    assert isinstance(checkpoint, S3Checkpoint)
+    assert checkpoint.endpoint == TEST_ENDPOINT
+
+
 def test_checkpoint_seek_logging(caplog):
     checkpoint = S3Checkpoint(TEST_REGION)
 

s3torchconnector/tst/unit/test_s3dataset_common.py

@@ -26,6 +26,7 @@
 TEST_KEY = "test-key"
 TEST_REGION = "us-east-1"
 S3_PREFIX = f"s3://{TEST_BUCKET}"
+TEST_ENDPOINT = "https://s3.us-east-1.amazonaws.com"
 
 
 @pytest.mark.parametrize(

s3torchconnector/tst/unit/test_s3iterable_dataset.py

@@ -12,6 +12,7 @@
     _create_mock_client_with_dummy_objects,
     S3_PREFIX,
     TEST_REGION,
+    TEST_ENDPOINT,
 )
 
 
@@ -191,6 +192,14 @@ def test_iteration_multiple_times(
     _verify_dataset(dataset, expected_keys, lambda data: data._object_info is not None)
 
 
+def test_dataset_creation_from_prefix_with_region_and_endpoint():
+    dataset = S3IterableDataset.from_prefix(
+        S3_PREFIX, region=TEST_REGION, endpoint=TEST_ENDPOINT
+    )
+    assert isinstance(dataset, S3IterableDataset)
+    assert dataset.endpoint == TEST_ENDPOINT
+
+
 def _verify_dataset(
     dataset: S3IterableDataset,
     expected_keys: Sequence[str],

s3torchconnector/tst/unit/test_s3mapdataset.py

@@ -12,6 +12,7 @@
     TEST_REGION,
     _create_mock_client_with_dummy_objects,
     S3_PREFIX,
+    TEST_ENDPOINT,
 )
 
 
@@ -178,6 +179,14 @@ def test_call_len_twice(keys: Sequence[str], length: int):
     assert len(dataset) == length
 
 
+def test_dataset_creation_from_prefix_with_region_and_endpoint():
+    dataset = S3MapDataset.from_prefix(
+        S3_PREFIX, region=TEST_REGION, endpoint=TEST_ENDPOINT
+    )
+    assert isinstance(dataset, S3MapDataset)
+    assert dataset.endpoint == TEST_ENDPOINT
+
+
 def verify_item(dataset: S3MapDataset, index: int, expected_key: str):
     data = dataset[index]
 

s3torchconnectorclient/python/src/s3torchconnectorclient/_mountpoint_s3_client.pyi

@@ -7,6 +7,7 @@ from typing import List, Optional
 
 class MountpointS3Client:
     region: str
+    endpoint: str
     user_agent_prefix: str
     throughput_target_gbps: float
     part_size: int
@@ -16,6 +17,7 @@ class MountpointS3Client:
     def __init__(
         self,
         region: str,
+        endpoint: str = "",
         user_agent_prefix: str = "",
         throughput_target_gbps: float = 10.0,
         part_size: int = 8 * 1024 * 1024,
@@ -37,6 +39,7 @@ class MockMountpointS3Client:
         self,
         region: str,
         bucket: str,
+        endpoint: str = "",
         throughput_target_gbps: float = 10.0,
         part_size: int = 8 * 1024 * 1024,
     ): ...

s3torchconnectorclient/python/tst/integration/test_mountpoint_s3_integration.py

@@ -36,6 +36,20 @@ def test_get_object(sample_directory):
     assert full_data == HELLO_WORLD_DATA
 
 
+def test_get_object_with_endpoint(sample_directory):
+    client = MountpointS3Client(
+        sample_directory.region,
+        TEST_USER_AGENT_PREFIX,
+        endpoint="http://s3.amazonaws.com",
+    )
+    stream = client.get_object(
+        sample_directory.bucket, f"{sample_directory.prefix}hello_world.txt"
+    )
+
+    full_data = b"".join(stream)
+    assert full_data == HELLO_WORLD_DATA
+
+
 def test_get_object_with_unpickled_client(sample_directory):
     original_client = MountpointS3Client(
         sample_directory.region, TEST_USER_AGENT_PREFIX