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>


Discussion:
codecov[bot]: ## [Codecov](https://app.codecov.io/gh/Lightning-AI/litdata/pull/461?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Lightning-AI) Report
Attention: Patch coverage is `50.00000%` with `2 lines` in your changes missing coverage. Please review.
> Project coverage is 71%. Comparing base [(`7cbb3ef`)](https://app.codecov.io/gh/Lightning-AI/litdata/commit/7cbb3ef7588246d38ca2b0e3ab68d1976d34ad14?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Lightning-AI) to head [(`53e2f9b`)](https://app.codecov.io/gh/Lightning-AI/litdata/commit/53e2f9ba4d73ac69d423fd514dfda4b4d06596c1?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Lightning-AI).

> :exclamation:  There is a different number of reports uploaded between BASE (7cbb3ef) and HEAD (53e2f9b). Click for more details.
> 
> <details><summary>HEAD has 18 uploads less than BASE</summary>
>
>| Flag | BASE (7cbb3ef) | HEAD (53e2f9b) |
>|------|------|------|
>|unittests|7|1|
>|3.10|4|1|
>|ubuntu-22.04|2|0|
>|3.9|3|0|
>|windows-2022|2|1|
>|macos-13|3|0|
></details>

<details><summary>Additional details and impacted files</summary>


```diff
@@         Coverage Diff          @@
##           main   #461    +/-   ##
====================================
- Coverage    78%    71%    -7%     
====================================
  Files        37     37            
  Lines      5365   5367     +2     
====================================
- Hits       4182   3832   -350     
- Misses     1183   1535   +352     
```

</details>

Fix (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 changes:
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