Expose S3ClientConfig (#191)

* Expose S3ClientConfig

We expose the following configuration flags with performance impact:
throughput_target_gbps(float): Throughput target in Gigabits per second (Gbps).
part_size(int): Size, in bytes, of parts that files will be downloaded or uploaded in.

Author: dnanuti

CHANGELOG.md

@@ -1,3 +1,8 @@
+## Unreleased
+
+### New features
+* Expose a new class, S3ClientConfig, with `throughput_target_gbps` and `part_size` parameters of the inner S3 client.
+
 ## v1.2.1 (March 14, 2024)
 
 ### Breaking changes

s3torchbenchmarking/src/s3torchbenchmarking/benchmark_utils.py

@@ -75,7 +75,6 @@ def throughput(self):
 
 
 class ExperimentResultJsonEncoder(JSONEncoder):
-
     def default(self, o: Any) -> Any:
         if isinstance(o, ExperimentResult):
             o: ExperimentResult = o

s3torchconnector/src/s3torchconnector/__init__.py

@@ -11,6 +11,7 @@
 from .s3map_dataset import S3MapDataset
 from .s3checkpoint import S3Checkpoint
 from ._version import __version__
+from ._s3client import S3ClientConfig
 
 __all__ = [
     "S3IterableDataset",
@@ -19,5 +20,6 @@
     "S3Reader",
     "S3Writer",
     "S3Exception",
+    "S3ClientConfig",
     "__version__",
 ]

s3torchconnector/src/s3torchconnector/_s3client/__init__.py

@@ -1,7 +1,12 @@
 #  Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 #  // SPDX-License-Identifier: BSD
 
+from .s3client_config import S3ClientConfig
 from ._s3client import S3Client
 from ._mock_s3client import MockS3Client
 
-__all__ = ["S3Client", "MockS3Client"]
+__all__ = [
+    "S3ClientConfig",
+    "S3Client",
+    "MockS3Client",
+]

s3torchconnector/src/s3torchconnector/_s3client/_mock_s3client.py

@@ -10,6 +10,7 @@
 
 from . import S3Client
 from .._user_agent import UserAgent
+from .s3client_config import S3ClientConfig
 
 """
 _mock_s3client.py
@@ -22,14 +23,19 @@ def __init__(
         self,
         region: str,
         bucket: str,
-        part_size: int = 8 * 1024 * 1024,
         user_agent: Optional[UserAgent] = None,
+        s3client_config: Optional[S3ClientConfig] = None,
     ):
-        super().__init__(region, user_agent=user_agent)
+        super().__init__(
+            region,
+            user_agent=user_agent,
+            s3client_config=s3client_config,
+        )
         self._mock_client = MockMountpointS3Client(
             region,
             bucket,
-            part_size=part_size,
+            throughput_target_gbps=self.s3client_config.throughput_target_gbps,
+            part_size=self.s3client_config.part_size,
             user_agent_prefix=self.user_agent_prefix,
         )
 

s3torchconnector/src/s3torchconnector/_s3client/_s3client.py

@@ -7,6 +7,7 @@
 from typing import Optional, Any
 
 from s3torchconnector import S3Reader, S3Writer
+from .s3client_config import S3ClientConfig
 
 from s3torchconnectorclient._mountpoint_s3_client import (
     MountpointS3Client,
@@ -35,15 +36,18 @@ class S3Client:
     def __init__(
         self,
         region: str,
+        *,
         endpoint: Optional[str] = None,
         user_agent: Optional[UserAgent] = None,
+        s3client_config: Optional[S3ClientConfig] = None,
     ):
         self._region = region
         self._endpoint = endpoint
         self._real_client: Optional[MountpointS3Client] = None
         self._client_pid: Optional[int] = None
         user_agent = user_agent or UserAgent()
         self._user_agent_prefix = user_agent.prefix
+        self._s3client_config = s3client_config or S3ClientConfig()
 
     @property
     def _client(self) -> MountpointS3Client:
@@ -58,6 +62,10 @@ def _client(self) -> MountpointS3Client:
     def region(self) -> str:
         return self._region
 
+    @property
+    def s3client_config(self) -> S3ClientConfig:
+        return self._s3client_config
+
     @property
     def user_agent_prefix(self) -> str:
         return self._user_agent_prefix
@@ -67,6 +75,8 @@ def _client_builder(self) -> MountpointS3Client:
             region=self._region,
             endpoint=self._endpoint,
             user_agent_prefix=self._user_agent_prefix,
+            throughput_target_gbps=self._s3client_config.throughput_target_gbps,
+            part_size=self._s3client_config.part_size,
         )
 
     def get_object(

s3torchconnector/src/s3torchconnector/_s3client/s3client_config.py

@@ -0,0 +1,21 @@
+#  Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#  // SPDX-License-Identifier: BSD
+from dataclasses import dataclass
+
+
+@dataclass(frozen=True)
+class S3ClientConfig:
+    """A dataclass exposing configurable parameters for the S3 client.
+
+    Args:
+    throughput_target_gbps(float): Throughput target in Gigabits per second (Gbps) that we are trying to reach.
+        10.0 Gbps by default (may change in future).
+    part_size(int): Size, in bytes, of parts that files will be downloaded or uploaded in.
+        Note: for saving checkpoints, the inner client will adjust the part size to meet the service limits.
+        (max number of parts per upload is 10,000, minimum upload part size is 5 MiB).
+        Part size must have values between 5MiB and 5GiB.
+        8MB by default (may change in future).
+    """
+
+    throughput_target_gbps: float = 10.0
+    part_size: int = 8 * 1024 * 1024

s3torchconnector/src/s3torchconnector/lightning/s3_lightning_checkpoint.py

@@ -8,18 +8,26 @@
 
 from lightning.pytorch.plugins.io import CheckpointIO
 
-from .._s3client import S3Client
+from .._s3client import S3Client, S3ClientConfig
 from .._s3dataset_common import parse_s3_uri
 from .._user_agent import UserAgent
 
 
 class S3LightningCheckpoint(CheckpointIO):
     """A checkpoint manager for S3 using the :class:`CheckpointIO` interface."""
 
-    def __init__(self, region: str):
+    def __init__(
+        self,
+        region: str,
+        s3client_config: Optional[S3ClientConfig] = None,
+    ):
         self.region = region
         user_agent = UserAgent(["lightning", lightning.__version__])
-        self._client = S3Client(region, user_agent=user_agent)
+        self._client = S3Client(
+            region,
+            user_agent=user_agent,
+            s3client_config=s3client_config,
+        )
 
     def save_checkpoint(
         self,

s3torchconnector/src/s3torchconnector/s3checkpoint.py

@@ -3,7 +3,7 @@
 from typing import Optional
 
 from ._s3dataset_common import parse_s3_uri
-from ._s3client import S3Client
+from ._s3client import S3Client, S3ClientConfig
 from . import S3Reader, S3Writer
 
 
@@ -17,10 +17,17 @@ class S3Checkpoint:
     torch.load, and torch.save.
     """
 
-    def __init__(self, region: str, endpoint: Optional[str] = None):
+    def __init__(
+        self,
+        region: str,
+        endpoint: Optional[str] = None,
+        s3client_config: Optional[S3ClientConfig] = None,
+    ):
         self.region = region
         self.endpoint = endpoint
-        self._client = S3Client(region, endpoint)
+        self._client = S3Client(
+            region, endpoint=endpoint, s3client_config=s3client_config
+        )
 
     def reader(self, s3_uri: str) -> S3Reader:
         """Creates an S3Reader from a given s3_uri.

s3torchconnector/src/s3torchconnector/s3iterable_dataset.py

@@ -8,7 +8,7 @@
 
 from . import S3Reader
 from ._s3bucket_key_data import S3BucketKeyData
-from ._s3client import S3Client
+from ._s3client import S3Client, S3ClientConfig
 from ._s3dataset_common import (
     identity,
     get_objects_from_uris,
@@ -31,11 +31,13 @@ def __init__(
         get_dataset_objects: Callable[[S3Client], Iterable[S3BucketKeyData]],
         endpoint: Optional[str] = None,
         transform: Callable[[S3Reader], Any] = identity,
+        s3client_config: Optional[S3ClientConfig] = None,
     ):
         self._get_dataset_objects = get_dataset_objects
         self._transform = transform
         self._region = region
         self._endpoint = endpoint
+        self._s3client_config = s3client_config
         self._client = None
 
     @property
@@ -54,6 +56,7 @@ def from_objects(
         region: str,
         endpoint: Optional[str] = None,
         transform: Callable[[S3Reader], Any] = identity,
+        s3client_config: Optional[S3ClientConfig] = None,
     ):
         """Returns an instance of S3IterableDataset using the S3 URI(s) provided.
 
@@ -62,6 +65,7 @@ def from_objects(
           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.
+          s3client_config: Optional S3ClientConfig with parameters for S3 client.
 
         Returns:
             S3IterableDataset: An IterableStyle dataset created from S3 objects.
@@ -75,6 +79,7 @@ def from_objects(
             partial(get_objects_from_uris, object_uris),
             endpoint,
             transform=transform,
+            s3client_config=s3client_config,
         )
 
     @classmethod
@@ -85,6 +90,7 @@ def from_prefix(
         region: str,
         endpoint: Optional[str] = None,
         transform: Callable[[S3Reader], Any] = identity,
+        s3client_config: Optional[S3ClientConfig] = None,
     ):
         """Returns an instance of S3IterableDataset using the S3 URI provided.
 
@@ -93,6 +99,7 @@ def from_prefix(
           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.
+          s3client_config: Optional S3ClientConfig with parameters for S3 client.
 
         Returns:
             S3IterableDataset: An IterableStyle dataset created from S3 objects.
@@ -106,11 +113,16 @@ def from_prefix(
             partial(get_objects_from_prefix, s3_uri),
             endpoint,
             transform=transform,
+            s3client_config=s3client_config,
         )
 
     def _get_client(self):
         if self._client is None:
-            self._client = S3Client(self.region, self.endpoint)
+            self._client = S3Client(
+                self.region,
+                endpoint=self.endpoint,
+                s3client_config=self._s3client_config,
+            )
         return self._client
 
     def _get_transformed_object(self, bucket_key: S3BucketKeyData) -> Any: