Update tests and docs

s3torchconnector/src/s3torchconnector/s3iterable_dataset.py

@@ -66,6 +66,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.
+          s3_connector_client_config: Optional S3ConnectorClientConfig with parameters for S3 client.
 
         Returns:
             S3IterableDataset: An IterableStyle dataset created from S3 objects.
@@ -99,6 +100,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.
+          s3_connector_client_config: Optional S3ConnectorClientConfig with parameters for S3 client.
 
         Returns:
             S3IterableDataset: An IterableStyle dataset created from S3 objects.

s3torchconnector/src/s3torchconnector/s3map_dataset.py

@@ -75,6 +75,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.
+          s3_connector_client_config: Optional S3ConnectorClientConfig with parameters for S3 client.
 
         Returns:
             S3MapDataset: A Map-Style dataset created from S3 objects.
@@ -108,6 +109,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.
+          s3_connector_client_config: Optional S3ConnectorClientConfig with parameters for S3 client.
 
         Returns:
             S3MapDataset: A Map-Style dataset created from S3 objects.

s3torchconnector/tst/unit/test_s3_client.py

@@ -4,12 +4,15 @@
 import pytest
 
 from hypothesis import given
-from hypothesis.strategies import lists, text
+from hypothesis.strategies import lists, text, integers, floats
 from unittest.mock import MagicMock
 
+from s3torchconnectorclient._mountpoint_s3_client import S3Exception
+
 from s3torchconnector._user_agent import UserAgent
 from s3torchconnector._version import __version__
 from s3torchconnector._s3client import S3Client, MockS3Client
+from s3torchconnectorclient import S3ConnectorClientConfig
 
 TEST_BUCKET = "test-bucket"
 TEST_KEY = "test-key"
@@ -82,3 +85,34 @@ def test_user_agent_always_starts_with_package_version(comments):
     if comments_str:
         assert comments_str in s3_client.user_agent_prefix
         assert comments_str in s3_client._client.user_agent_prefix
+
+
+@given(
+    part_size=integers(min_value=5 * 1024, max_value=5 * 1024 * 1024),
+    throughput_target_gbps=floats(min_value=10.0, max_value=100.0),
+)
+def test_s3_client_custom_config(part_size: int, throughput_target_gbps: float):
+    # Part size must be at between 5MiB and 5GiB
+    part_size = part_size * 1024
+    s3_client = S3Client(
+        region=TEST_REGION,
+        s3_connector_client_config=S3ConnectorClientConfig(
+            part_size=part_size,
+            throughput_target_gbps=throughput_target_gbps,
+        ),
+    )
+    assert s3_client._client.part_size == part_size
+    assert abs(s3_client._client.throughput_target_gbps - throughput_target_gbps) < 1e-9
+
+
+def test_s3_client_invalid_part_size_config():
+    with pytest.raises(
+        S3Exception,
+        match="invalid configuration: part size must be at between 5MiB and 5GiB",
+    ):
+        s3_client = S3Client(
+            region=TEST_REGION,
+            s3_connector_client_config=S3ConnectorClientConfig(part_size=1),
+        )
+        # The client is lazily initialized
+        assert s3_client._client.part_size is not None

s3torchconnectorclient/python/src/s3torchconnectorclient/__init__.py

@@ -6,7 +6,7 @@
 from ._logger_patch import TRACE as LOG_TRACE
 from ._logger_patch import _install_trace_logging
 from ._mountpoint_s3_client import S3Exception, __version__
-from .s3_connector_client_config import S3ConnectorClientConfig
+from s3torchconnectorclient.s3_connector_client_config import S3ConnectorClientConfig
 
 _install_trace_logging()
 

s3torchconnectorclient/python/src/s3torchconnectorclient/s3_connector_client_config.py

@@ -1,10 +1,27 @@
 #  Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 #  // SPDX-License-Identifier: BSD
+from dataclasses import dataclass
 
 
+@dataclass
 class S3ConnectorClientConfig:
+    """A dataclass exposing configurable parameters for the S3 client."""
+
     def __init__(
-        self, throughput_target_gbps: float = 10.0, part_size: int = 8 * 1024 * 1024
+        self,
+        throughput_target_gbps: float = 10.0,
+        part_size: int = 8 * 1024 * 1024,
     ):
+        """Returns a config wrapper object.
+
+        Args:
+          throughput_target_gbps(float): Throughput target in Gigabits per second (Gbps) that we are trying to reach.
+          You can also use get_recommended_throughput_target_gbps() to get recommended value for your system.
+          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).
+          8MB by default (may change in future).
+        """
         self.throughput_target_gbps = throughput_target_gbps
         self.part_size = part_size

s3torchconnectorclient/python/tst/unit/test_s3_connector_client_config.py

@@ -0,0 +1,31 @@
+#  Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#  // SPDX-License-Identifier: BSD
+from hypothesis import given
+from hypothesis.strategies import integers, floats
+
+from s3torchconnectorclient import S3ConnectorClientConfig
+
+
+def test_default():
+    config = S3ConnectorClientConfig()
+    assert config is not None
+    assert config.part_size == 8 * 1024 * 1024
+    assert abs(config.throughput_target_gbps - 10.0) < 1e-9
+
+
+@given(part_size=integers(min_value=1, max_value=1e12))
+def test_part_size_setup(part_size: int):
+    config = S3ConnectorClientConfig(part_size=part_size)
+    assert config is not None
+    assert config.part_size == part_size
+    assert abs(config.throughput_target_gbps - 10.0) < 1e-9
+
+
+@given(throughput_target_gbps_setup=floats(min_value=1.0, max_value=100.0))
+def test_throughput_target_gbps_setup(throughput_target_gbps_setup: float):
+    config = S3ConnectorClientConfig(
+        throughput_target_gbps=throughput_target_gbps_setup
+    )
+    assert config is not None
+    assert config.part_size == 8 * 1024 * 1024
+    assert abs(config.throughput_target_gbps - throughput_target_gbps_setup) < 1e-9