Update the S3 client initialization to explicitly use boto3.Session() (#461)

* Refactor S3 client initialization to use boto3 session for improved configuration handling
* update client
* update tests to use boto3 session

Author: bhimrazy

src/litdata/streaming/client.py

@@ -38,17 +38,19 @@ def _create_client(self) -> None:
         )
 
         if has_shared_credentials_file or not _IS_IN_STUDIO or self._storage_options:
-            self._client = boto3.client(
+            session = boto3.Session()
+            self._client = session.client(
                 "s3",
                 **{
                     "config": botocore.config.Config(retries={"max_attempts": 1000, "mode": "adaptive"}),
-                    **self._storage_options,
+                    **self._storage_options,  # If additional options are provided
                 },
             )
         else:
             provider = InstanceMetadataProvider(iam_role_fetcher=InstanceMetadataFetcher(timeout=3600, num_attempts=5))
             credentials = provider.load()
-            self._client = boto3.client(
+            session = boto3.Session()
+            self._client = session.client(
                 "s3",
                 aws_access_key_id=credentials.access_key,
                 aws_secret_access_key=credentials.secret_key,

tests/streaming/test_client.py

@@ -8,12 +8,14 @@
 
 
 def test_s3_client_with_storage_options(monkeypatch):
-    boto3 = mock.MagicMock()
+    boto3_session = mock.MagicMock()
+    boto3 = mock.MagicMock(Session=boto3_session)
     monkeypatch.setattr(client, "boto3", boto3)
 
     botocore = mock.MagicMock()
     monkeypatch.setattr(client, "botocore", botocore)
 
+    # Create S3Client with storage options
     storage_options = {
         "region_name": "us-west-2",
         "endpoint_url": "https://custom.endpoint",
@@ -23,24 +25,27 @@ def test_s3_client_with_storage_options(monkeypatch):
 
     assert s3_client.client
 
-    boto3.client.assert_called_with(
+    boto3_session().client.assert_called_with(
         "s3",
         region_name="us-west-2",
         endpoint_url="https://custom.endpoint",
         config=botocore.config.Config(retries={"max_attempts": 100}),
     )
 
+    # Create S3Client without storage options
     s3_client = client.S3Client()
-
     assert s3_client.client
 
-    boto3.client.assert_called_with(
-        "s3", config=botocore.config.Config(retries={"max_attempts": 1000, "mode": "adaptive"})
+    # Verify that boto3.Session().client was called with the default parameters
+    boto3_session().client.assert_called_with(
+        "s3",
+        config=botocore.config.Config(retries={"max_attempts": 1000, "mode": "adaptive"}),
     )
 
 
 def test_s3_client_without_cloud_space_id(monkeypatch):
-    boto3 = mock.MagicMock()
+    boto3_session = mock.MagicMock()
+    boto3 = mock.MagicMock(Session=boto3_session)
     monkeypatch.setattr(client, "boto3", boto3)
 
     botocore = mock.MagicMock()
@@ -59,13 +64,14 @@ def test_s3_client_without_cloud_space_id(monkeypatch):
     assert s3.client
     assert s3.client
 
-    boto3.client.assert_called_once()
+    boto3_session().client.assert_called_once()
 
 
 @pytest.mark.skipif(sys.platform == "win32", reason="not supported on windows")
 @pytest.mark.parametrize("use_shared_credentials", [False, True, None])
 def test_s3_client_with_cloud_space_id(use_shared_credentials, monkeypatch):
-    boto3 = mock.MagicMock()
+    boto3_session = mock.MagicMock()
+    boto3 = mock.MagicMock(Session=boto3_session)
     monkeypatch.setattr(client, "boto3", boto3)
 
     botocore = mock.MagicMock()
@@ -85,14 +91,14 @@ def test_s3_client_with_cloud_space_id(use_shared_credentials, monkeypatch):
     s3 = client.S3Client(1)
     assert s3.client
     assert s3.client
-    boto3.client.assert_called_once()
+    boto3_session().client.assert_called_once()
     sleep(1 - (time() - s3._last_time))
     assert s3.client
     assert s3.client
-    assert len(boto3.client._mock_mock_calls) == 6
+    assert len(boto3_session().client._mock_mock_calls) == 6
     sleep(1 - (time() - s3._last_time))
     assert s3.client
     assert s3.client
-    assert len(boto3.client._mock_mock_calls) == 9
+    assert len(boto3_session().client._mock_mock_calls) == 9
 
     assert instance_metadata_provider._mock_call_count == 0 if use_shared_credentials else 3