Repository: Lightning-AI/litData

Issue:
s3 client expire after 1 hour: an error occurred (400) when calling the headobject operation: bad request" due to download_chunk_from_index(chunk_index)
## 🐛 Bug

<!-- A clear and concise description of what the bug is. -->
I am training with a large tabular dataset (~510m+) rows using PyTorch DDP. The dataset has been optimized into chunks and store in s3. We used `StreamingDataset `and `StreamingDataLoader` to load chunks. It was fine at early stage and model is training. But after 1 hour, I guess the credential is expired, which leads `"botocore.exceptions.ClientError: An error occurred (400) when calling the headobject operation: bad request" due to download_chunk_from_index(chunk_index). I would not say my chunks are wrong, as I have tested in our validation dataset and both ends to this error whenever if is a long-running (>1 hour).
`
### To Reproduce

Steps to reproduce the behavior...

<!-- If you have a code sample, error messages, stack traces, please provide it here as well -->

```
def _create_client(self) -> None:
        has_shared_credentials_file = (
            os.getenv("AWS_SHARED_CREDENTIALS_FILE") == os.getenv("AWS_CONFIG_FILE") == "/.credentials/.aws_credentials"
        )

        if has_shared_credentials_file or not _IS_IN_STUDIO or self._storage_options:
            self._client = boto3.client(
                "s3",
                **{
                    "config": botocore.config.Config(retries={"max_attempts": 1000, "mode": "adaptive"}),
                    **self._storage_options,
                },
            )
        else:
            provider = InstanceMetadataProvider(iam_role_fetcher=InstanceMetadataFetcher(timeout=3600, num_attempts=5))
            credentials = provider.load()
            self._client = boto3.client(
                "s3",
                aws_access_key_id=credentials.access_key,
                aws_secret_access_key=credentials.secret_key,
                aws_session_token=credentials.token,
                config=botocore.config.Config(retries={"max_attempts": 1000, "mode": "adaptive"}),
```
For this part, since the s3 token will automatically expire after 1 hour. Is there any way I can customize this function using monkey patch? 


<details>
  <summary>Code sample</summary>

```python
# Ideally attach a minimal code sample to reproduce the decried issue.
# Minimal means having the shortest code but still preserving the bug.
```

</details>

### Expected behavior

<!-- A clear and concise description of what you expected to happen. -->

### Additional context

<!-- Add any other context about the problem here. -->

<details>
  <summary>Environment detail</summary>

- PyTorch Version (e.g., 1.0):
- OS (e.g., Linux):
- How you installed PyTorch (`conda`, `pip`, source):
- Build command you used (if compiling from source):
- Python version:
- CUDA/cuDNN version:
- GPU models and configuration:
- Any other relevant information:

</details>

Hints from issue comments:
@tchaton: Hey @leonzyzy In theory, it should re-create it: https://github.com/Lightning-AI/litdata/blob/main/src/litdata/streaming/client.py#L66@leonzyzy: Hi @tchaton, thanks for the help! Yes it re-create, but it seems did not create a session and only call the client() in boto3. https://github.com/Lightning-AI/litdata/blob/main/src/litdata/streaming/client.py#L40C8-L47C14@tchaton: Hey @leonzyzy feel free to make a PR to fix it. @tchaton: @leonzyzy Any updates ?@bhimrazy: ```
session = boto3.Session()
self._client = session.client(
    "s3",
    config=botocore.config.Config(retries={"max_attempts": 1000, "mode": "adaptive"}),
    **self._storage_options,
)
```
@tchaton, could this be a more reliable approach?@tchaton: Yes, possible. @bhimrazy: Hi @leonzyzy, feel free to try the updates from [PR #461](https://github.com/Lightning-AI/litdata/pull/461) and let us know how it goes! 

Solution patch:
diff --git a/src/litdata/streaming/client.py b/src/litdata/streaming/client.py
index 68c1491..e9bb5be 100644
--- a/src/litdata/streaming/client.py
+++ b/src/litdata/streaming/client.py
@@ -38,17 +38,19 @@ class S3Client:
         )
 
         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,


Test patch:
diff --git a/tests/streaming/test_client.py b/tests/streaming/test_client.py
index b0f95d1..636734a 100644
--- a/tests/streaming/test_client.py
+++ b/tests/streaming/test_client.py
@@ -8,12 +8,14 @@ from litdata.streaming import client
 
 
 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