Repository: Lightning-AI/litData

Issue:
storage_options endpoint_url works for S3Client but not s5cmd in S3Downloader
## 🐛 Bug

When using a StreamingDataset stored in Wasabi S3 Compatible Storage, `storage_options` that work with the S3Client don't work with s5cmd, and I don't think it's possible to pass the required arguments to s5cmd.

  <summary>Code sample</summary>

```python
from litdata import StreamingDataset
from botocore import UNSIGNED
from botocore.client import Config

storage_options = {
    "endpoint_url": "https://s3.wasabisys.com",
    "config": Config(signature_version=UNSIGNED)
}

# succeeds with S3Client, fails with s5cmd, can't figure out a workaround...
dataset = StreamingDataset("s3://visionlab-litdata/mnist/val", 
                           storage_options=storage_options)
```

### To Reproduce

The above code succeeds when using an environment without s5cmd installed, or when I edit `litdata/streaming/downloader.py` to set self._s5cmd_available to False
```
nano +72 /home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/litdata/streaming/downloader.py
```
pasting in:
```
self._s5cmd_available = False
```

Without that edit, s5cmd is used and throws an error, and as far as I can tell, there's no way to pass the required arguments through the StreamingDataset storage_options or system environment variables. 

For example, this command works:
```
s5cmd --no-sign-request --endpoint-url https://s3.wasabisys.com cp s3://visionlab-litdata/mnist/val/index.json index.json
```

But s5cmd does not appear to use environment variables for either `--no-sign-request` or `--endpoint-url `, so for instance the following doesn't work.
```
export AWS_ENDPOINT_URL=https://s3.wasabisys.com 
s5cmd --no-sign-request cp s3://visionlab-litdata/mnist/val/index.json index.json
```

Whereas the aws cli would use that environment variable, so the following succeeds
```
export AWS_ENDPOINT_URL=https://s3.wasabisys.com 
aws s3 --no-sign-request cp s3://visionlab-litdata/mnist/val/index.json /tmp/index.json
```

I think the command in the S3Downloader class might need to be constructed differently to handle storage_options, hopefully in a way that's unform across S3Client and  s5cmd? Or perhaps there needs to be a separate s5cmd_options argument?

```python
if self._s5cmd_available:
                env = None
                if self._storage_options:
                    env = os.environ.copy()
                    env.update(self._storage_options)
                proc = subprocess.Popen(
                    f"s5cmd cp {remote_filepath} {local_filepath}",
                    shell=True,
                    stdout=subprocess.PIPE,
                    env=env,
                )
                proc.wait()
```

Hopefully there's a workaround for passing the required args to s5cmd that I missed.

I suppose the minimal fix would be to allow the option to force using the S3Client, but that seems less desirable. 

Otherwise the construction of the s5cmd command might need to be modified, something like the following 

```python
from botocore import UNSIGNED
from botocore.client import Config

if self._s5cmd_available:
    env = os.environ.copy()
    extra_args = ''
    if self._storage_options:
        unsigned = self._storage_options.get('unsigned')
        config = self._storage_options.get('config', {})
        signature_version = getattr(config, 'signature_version')
        endpoint_url = self._storage_options.get('endpoint_url')
        if unsigned or signature_version==UNSIGNED:
            extra_args += f'--no-sign-request '
        if endpoint_url:
            extra_args += f'--endpoint-url {endpoint_url} '

    cmd = f"s5cmd {extra_args.strip()} cp {remote_filepath} {local_filepath}"
    proc = subprocess.Popen(
        cmd,
        shell=True,
        stdout=subprocess.PIPE,
        env=env,
    )
    proc.wait()
```

But this might be a headache to maintain if there other input arguments need to be taken into account that I'm not anticipating (bucket region?).

It might just be simpler to add a s5cmd_options argument to StreamingDataset, which in this case would have been `s5cmd_options = "--no-sign-request --endpoint-url https://s3.wasabisys.com"` I guess the downside with this option is that as a user, you have to know what your backend is (S3Client or s5cmd), and might want to support both.

### Expected behavior

I would expect a given set of storage_options to work for both supported backends (S3Client and s5cmd), or for there to be possible configuration options to pass the required args to s5cmd. e.g., it would be great if something like this would work with either S3Client or s5cmd.

```python
from litdata import StreamingDataset

storage_options = {
    "endpoint_url": "https://s3.wasabisys.com",
    "unsigned": True,
}

dataset = StreamingDataset("s3://visionlab-litdata/mnist/val",  storage_options=storage_options)
```

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

- Litdata Version:  latest (0.2.42)
- 5cmd Version: 0.2.0
- OS: Linux lightning.ai studio

</details>

Discussion:
@bhimrazy: Hi @grez72,  
Thank you for raising this issue.  

For **s5cmd**, the storage options can be passed as environment variables in storage options. 
Could you please try using capitalized variables?  

### **Signed Request**  
```python
storage_options = {
    "AWS_ACCESS_KEY_ID": "your-access-key",
    "AWS_SECRET_ACCESS_KEY": "your-secret-key",
    "S3_ENDPOINT_URL": "your-endpoint-url",
}
```

### **Unsigned Request**  
```python
storage_options = {
    "AWS_NO_SIGN_REQUEST": "yes",
    "S3_ENDPOINT_URL": "your-endpoint-url",
}
```

Let us know if this works for you.  

For other supported environment variables, please refer to the s5cmd repository:
🔗 [s5cmd GitHub Repo](https://github.com/peak/s5cmd)

We will also update the README to clarify the correct storage options for **s5cmd**.  
Thank you, and we appreciate your feedback!  

@tchaton: Hey @grez72. Can you confirm that it works with the right name for the env variables.

@bhimrazy Do you think we could add those examples on the README for s5cmd and let's introduce an env variable to disable s5cmd ? @bhimrazy: > Do you think we could add those examples to the README for s5cmd, and introduce an environment variable to disable s5cmd?

Yes @tchaton.
btw I have prepared the draft for the README update and am waiting for confirmation from the user. 
So far, I have only tested locally with MinIO.

And sure, I will create the PR with the README and environment variable update.
@grez72: Thanks for your quick help with this. 

I tried the following and it still didn't work. I think the culprit is AWS_NO_SIGN_REQUEST not being recognized.
```python
storage_options = {
    "AWS_NO_SIGN_REQUEST": "yes",
    "S3_ENDPOINT_URL": "https://s3.wasabisys.com",
}

dataset = StreamingDataset("s3://visionlab-litdata/mnist/val", 
                           storage_options=storage_options)
```

A quick test shows that setting environment variable S3_ENDPOINT_URL allows downloading files, but setting AWS_NO_SIGN_REQUEST as an environment variable does not.
```
# these work (download index.json)
s5cmd --no-sign-request --endpoint-url https://s3.wasabisys.com cp s3://visionlab-litdata/mnist/val/index.json index.json

S3_ENDPOINT_URL=https://s3.wasabisys.com s5cmd --no-sign-request cp s3://visionlab-litdata/mnist/val/index.json index.json

# these don't don't work
# throws The AWS Access Key Id you provided does not exist in our records
AWS_NO_SIGN_REQUEST=yes s5cmd --endpoint-url https://s3.wasabisys.com cp s3://visionlab-litdata/mnist/val/index.json index.json
AWS_NO_SIGN_REQUEST='yes' s5cmd --endpoint-url https://s3.wasabisys.com cp s3://visionlab-litdata/mnist/val/index.json index.json
AWS_NO_SIGN_REQUEST=1 s5cmd --endpoint-url https://s3.wasabisys.com cp s3://visionlab-litdata/mnist/val/index.json index.json
```

@tchaton: Hey @grez72 

Would you mind making a PR to document this on the README ?

Best,
Thomas Chaton.

Solution patch:
diff --git a/README.md b/README.md
index fc80dc6..4d85fd8 100644
--- a/README.md
+++ b/README.md
@@ -216,15 +216,29 @@ Additionally, you can inject client connection settings for [S3](https://boto3.a
 ```python
 from litdata import StreamingDataset
 
+# boto3 compatible storage options for a custom S3-compatible endpoint
 storage_options = {
     "endpoint_url": "your_endpoint_url",
     "aws_access_key_id": "your_access_key_id",
     "aws_secret_access_key": "your_secret_access_key",
 }
 
+dataset = StreamingDataset('s3://my-bucket/my-data', storage_options=storage_options)
+
+# s5cmd compatible storage options for a custom S3-compatible endpoint
+# Note: If s5cmd is installed, it will be used by default for S3 operations. If you prefer not to use s5cmd, you can disable it by setting the environment variable: `DISABLE_S5CMD=1`
+storage_options = {
+    "AWS_ACCESS_KEY_ID": "your_access_key_id",
+    "AWS_SECRET_ACCESS_KEY": "your_secret_access_key",
+    "S3_ENDPOINT_URL": "your_endpoint_url",  # Required only for custom endpoints
+}
+
+
 dataset = StreamingDataset('s3://my-bucket/my-data', storage_options=storage_options)
 ```
 
+Alternative: Using `s5cmd` for S3 Operations
+
 
 Also, you can specify a custom cache directory when initializing your dataset. This is useful when you want to store the cache in a specific location.
 ```python
@@ -366,19 +380,37 @@ for batch in val_dataloader:
 
 &nbsp;
 
-The StreamingDataset supports reading optimized datasets from common cloud providers. 
+The `StreamingDataset` provides support for reading optimized datasets from common cloud storage providers like AWS S3, Google Cloud Storage (GCS), and Azure Blob Storage. Below are examples of how to use StreamingDataset with each cloud provider.
 
 ```python
 import os
 import litdata as ld
 
-# Read data from AWS S3
+# Read data from AWS S3 using s5cmd
+# Note: If s5cmd is installed, it will be used by default for S3 operations. If you prefer not to use s5cmd, you can disable it by setting the environment variable: `DISABLE_S5CMD=1`
 aws_storage_options={
     "AWS_ACCESS_KEY_ID": os.environ['AWS_ACCESS_KEY_ID'],
     "AWS_SECRET_ACCESS_KEY": os.environ['AWS_SECRET_ACCESS_KEY'],
+    "S3_ENDPOINT_URL": os.environ['AWS_ENDPOINT_URL'],  # Required only for custom endpoints
+}
+dataset = ld.StreamingDataset("s3://my-bucket/my-data", storage_options=aws_storage_options)
+
+# Read Data from AWS S3 with Unsigned Request using s5cmd
+aws_storage_options={
+  "AWS_NO_SIGN_REQUEST": "Yes" # Required for unsigned requests
+  "S3_ENDPOINT_URL": os.environ['AWS_ENDPOINT_URL'],  # Required only for custom endpoints
+}
+dataset = ld.StreamingDataset("s3://my-bucket/my-data", storage_options=aws_storage_options)
+
+# Read data from AWS S3 using boto3
+os.environ["DISABLE_S5CMD"] = "1"
+aws_storage_options={
+    "aws_access_key_id": os.environ['AWS_ACCESS_KEY_ID'],
+    "aws_secret_access_key": os.environ['AWS_SECRET_ACCESS_KEY'],
 }
 dataset = ld.StreamingDataset("s3://my-bucket/my-data", storage_options=aws_storage_options)
 
+
 # Read data from GCS
 gcp_storage_options={
     "project": os.environ['PROJECT_ID'],
diff --git a/src/litdata/constants.py b/src/litdata/constants.py
index ada4a05..89ed7cd 100644
--- a/src/litdata/constants.py
+++ b/src/litdata/constants.py
@@ -44,6 +44,7 @@ _DEBUG = bool(int(os.getenv("DEBUG", "0")))
 
 _MAX_WAIT_TIME = int(os.getenv("MAX_WAIT_TIME", "120"))
 _FORCE_DOWNLOAD_TIME = int(os.getenv("FORCE_DOWNLOAD_TIME", "30"))
+_DISABLE_S5CMD = bool(int(os.getenv("DISABLE_S5CMD", "0")))
 
 # DON'T CHANGE ORDER
 _TORCH_DTYPES_MAPPING = {
diff --git a/src/litdata/streaming/downloader.py b/src/litdata/streaming/downloader.py
index d06336d..02a7904 100644
--- a/src/litdata/streaming/downloader.py
+++ b/src/litdata/streaming/downloader.py
@@ -24,6 +24,7 @@ from filelock import FileLock, Timeout
 
 from litdata.constants import (
     _AZURE_STORAGE_AVAILABLE,
+    _DISABLE_S5CMD,
     _GOOGLE_STORAGE_AVAILABLE,
     _HF_HUB_AVAILABLE,
     _INDEX_FILENAME,
@@ -70,7 +71,7 @@ class S3Downloader(Downloader):
         super().__init__(remote_dir, cache_dir, chunks, storage_options)
         self._s5cmd_available = os.system("s5cmd > /dev/null 2>&1") == 0
 
-        if not self._s5cmd_available:
+        if not self._s5cmd_available or _DISABLE_S5CMD:
             self._client = S3Client(storage_options=self._storage_options)
 
     def download_file(self, remote_filepath: str, local_filepath: str, remote_chunk_filename: str = "") -> None:
@@ -91,18 +92,47 @@ class S3Downloader(Downloader):
             if os.path.exists(local_filepath):
                 return
 
-            if self._s5cmd_available:
+            if self._s5cmd_available and not _DISABLE_S5CMD:
                 env = None
                 if self._storage_options:
                     env = os.environ.copy()
                     env.update(self._storage_options)
+
+                aws_no_sign_request = self._storage_options.get("AWS_NO_SIGN_REQUEST", "no").lower() == "yes"
+                # prepare the s5cmd command
+                no_signed_option = "--no-sign-request" if aws_no_sign_request else None
+                cmd_parts = ["s5cmd", no_signed_option, "cp", remote_filepath, local_filepath]
+                cmd = " ".join(part for part in cmd_parts if part)
+
                 proc = subprocess.Popen(
-                    f"s5cmd cp {remote_filepath} {local_filepath}",
+                    cmd,
                     shell=True,
                     stdout=subprocess.PIPE,
+                    stderr=subprocess.PIPE,
                     env=env,
                 )
-                proc.wait()
+                return_code = proc.wait()
+
+                if return_code != 0:
+                    stderr_output = proc.stderr.read().decode().strip() if proc.stderr else ""
+                    error_message = (
+                        f"Failed to execute command `{cmd}` (exit code: {return_code}). "
+                        "This might be due to an incorrect file path, insufficient permissions, or network issues. "
+                        "To resolve this issue, you can either:\n"
+                        "- Pass `storage_options` with the necessary credentials and endpoint. \n"
+                        "- Example:\n"
+                        "  storage_options = {\n"
+                        '      "AWS_ACCESS_KEY_ID": "your-key",\n'
+                        '      "AWS_SECRET_ACCESS_KEY": "your-secret",\n'
+                        '      "S3_ENDPOINT_URL": "https://s3.example.com" (Optional if using AWS)\n'
+                        "  }\n"
+                        "- or disable `s5cmd` by setting `DISABLE_S5CMD=1` if `storage_options` do not work.\n"
+                    )
+                    if stderr_output:
+                        error_message += (
+                            f"For further debugging, please check the command output below:\n{stderr_output}"
+                        )
+                    raise RuntimeError(error_message)
             else:
                 from boto3.s3.transfer import TransferConfig
 


Test patch:
diff --git a/tests/streaming/test_downloader.py b/tests/streaming/test_downloader.py
index 50f754a..52bace5 100644
--- a/tests/streaming/test_downloader.py
+++ b/tests/streaming/test_downloader.py
@@ -1,5 +1,6 @@
 # ruff: noqa: S604
 import os
+import sys
 from unittest import mock
 from unittest.mock import MagicMock, patch
 
@@ -52,6 +53,7 @@ def test_get_downloader(tmpdir):
 def test_s3_downloader_fast(tmpdir, monkeypatch):
     monkeypatch.setattr(os, "system", MagicMock(return_value=0))
     popen_mock = MagicMock()
+    popen_mock.wait.return_value = 0  # Simulate a successful download
     monkeypatch.setattr(subprocess, "Popen", MagicMock(return_value=popen_mock))
     downloader = S3Downloader(tmpdir, tmpdir, [])
     downloader.download_file("s3://random_bucket/a.txt", os.path.join(tmpdir, "a.txt"))
@@ -63,6 +65,7 @@ def test_s3_downloader_fast(tmpdir, monkeypatch):
 def test_s3_downloader_with_s5cmd_no_storage_options(popen_mock, system_mock, tmpdir):
     system_mock.return_value = 0  # Simulates s5cmd being available
     process_mock = MagicMock()
+    process_mock.wait.return_value = 0  # Simulate a successful download
     popen_mock.return_value = process_mock
 
     # Initialize the S3Downloader without storage options
@@ -78,16 +81,53 @@ def test_s3_downloader_with_s5cmd_no_storage_options(popen_mock, system_mock, tm
         f"s5cmd cp {remote_filepath} {local_filepath}",
         shell=True,
         stdout=subprocess.PIPE,
+        stderr=subprocess.PIPE,
         env=None,
     )
     process_mock.wait.assert_called_once()
 
 
+@patch("os.system")
+@patch("subprocess.Popen")
+@mock.patch("litdata.streaming.downloader._DISABLE_S5CMD", True)
+@mock.patch("boto3.client")
+def test_s3_downloader_s5cmd_available_but_disabled(boto3_client_mock, popen_mock, system_mock, tmpdir):
+    system_mock.return_value = 0  # Simulates s5cmd being available
+    process_mock = MagicMock()
+    popen_mock.return_value = process_mock
+
+    # Mock the boto3 client
+    boto3_client_instance = MagicMock()
+    boto3_client_mock.return_value = boto3_client_instance
+
+    # Mock the download_file method to avoid NoCredentialsError
+    boto3_client_instance.download_file = MagicMock()
+
+    # Mock the S3Client class to avoid creating a real boto3 client
+    with mock.patch("litdata.streaming.downloader.S3Client") as S3ClientMock:
+        S3ClientMock.return_value.client = boto3_client_instance
+
+        # Initialize the S3Downloader
+        downloader = S3Downloader("s3://random_bucket", str(tmpdir), [])
+
+        # Action: Call the download_file method
+        remote_filepath = "s3://random_bucket/sample_file.txt"
+        local_filepath = os.path.join(tmpdir, "sample_file.txt")
+        downloader.download_file(remote_filepath, local_filepath)
+
+        # Assertion: Verify subprocess.Popen was not called
+        popen_mock.assert_not_called()
+
+        # Assertion: Verify boto3 download_file was called
+        boto3_client_instance.download_file.assert_called_once()
+
+
 @patch("os.system")
 @patch("subprocess.Popen")
 def test_s3_downloader_with_s5cmd_with_storage_options(popen_mock, system_mock, tmpdir):
     system_mock.return_value = 0  # Simulates s5cmd being available
     process_mock = MagicMock()
+    process_mock.wait.return_value = 0  # Simulate a successful download
     popen_mock.return_value = process_mock
 
     storage_options = {"AWS_ACCESS_KEY_ID": "dummy_key", "AWS_SECRET_ACCESS_KEY": "dummy_secret"}
@@ -109,11 +149,101 @@ def test_s3_downloader_with_s5cmd_with_storage_options(popen_mock, system_mock,
         f"s5cmd cp {remote_filepath} {local_filepath}",
         shell=True,
         stdout=subprocess.PIPE,
+        stderr=subprocess.PIPE,
+        env=expected_env,
+    )
+    process_mock.wait.assert_called_once()
+
+
+@patch("os.system")
+@patch("subprocess.Popen")
+def test_s3_downloader_with_s5cmd_with_storage_options_unsigned(popen_mock, system_mock, tmpdir):
+    system_mock.return_value = 0  # Simulates s5cmd being available
+    process_mock = MagicMock()
+    process_mock.wait.return_value = 0  # Simulate a successful download
+    popen_mock.return_value = process_mock
+
+    storage_options = {"AWS_NO_SIGN_REQUEST": "Yes"}
+
+    # Initialize the S3Downloader with storage options
+    downloader = S3Downloader("s3://random_bucket", str(tmpdir), [], storage_options)
+
+    # Action: Call the download_file method
+    remote_filepath = "s3://random_bucket/sample_file.txt"
+    local_filepath = os.path.join(tmpdir, "sample_file.txt")
+    downloader.download_file(remote_filepath, local_filepath)
+
+    # Create expected environment variables by merging the current env with storage_options
+    expected_env = os.environ.copy()
+    expected_env.update(storage_options)
+
+    # Assertion: Verify subprocess.Popen was called with the correct arguments and environment variables
+    popen_mock.assert_called_once_with(
+        f"s5cmd --no-sign-request cp {remote_filepath} {local_filepath}",
+        shell=True,
+        stdout=subprocess.PIPE,
+        stderr=subprocess.PIPE,
         env=expected_env,
     )
     process_mock.wait.assert_called_once()
 
 
+@pytest.mark.skipif(
+    shutil.which("s5cmd") is None or sys.platform == "win32", reason="s5cmd is not available or running on Windows"
+)
+def test_s3_downloader_with_s5cmd_with_storage_options_unsigned_pl(tmpdir):
+    # Set up the test environment
+    remote_filepath = "s3://pl-flash-data/optimized_tiny_imagenet/index.json"
+    local_filepath = os.path.join(tmpdir, "index.json")
+
+    storage_options = {"AWS_NO_SIGN_REQUEST": "Yes"}
+    # Initialize the S3Downloader
+    downloader = S3Downloader("s3://pl-flash-data", str(tmpdir), [], storage_options)
+
+    # Download the file
+    downloader.download_file(remote_filepath, local_filepath)
+
+    # Verify the download
+    assert os.path.exists(local_filepath), "The index.json file was not downloaded successfully."
+
+    # verify the contents of the file
+    with open(local_filepath) as f:
+        content = f.read()
+        assert content.startswith("{"), "The downloaded file does not appear to be a valid JSON file."
+
+
+@patch("os.system")
+@patch("subprocess.Popen")
+def test_s3_downloader_s5cmd_error_handling(popen_mock, system_mock, tmpdir):
+    system_mock.return_value = 0  # Simulates s5cmd being available
+    process_mock = MagicMock()
+    process_mock.wait.return_value = 1  # Simulate a non-zero return code
+    process_mock.stderr.read.return_value = b"Simulated error message"
+    popen_mock.return_value = process_mock
+
+    # Initialize the S3Downloader without storage options
+    downloader = S3Downloader("s3://random_bucket", str(tmpdir), [])
+
+    # Action: Call the download_file method and expect a RuntimeError
+    remote_filepath = "s3://random_bucket/sample_file.txt"
+    local_filepath = os.path.join(tmpdir, "sample_file.txt")
+
+    with pytest.raises(RuntimeError, match="Failed to execute command"):
+        downloader.download_file(remote_filepath, local_filepath)
+
+    # Assertion: Verify subprocess.Popen was called with the correct arguments
+    popen_mock.assert_called_once_with(
+        f"s5cmd cp {remote_filepath} {local_filepath}",
+        shell=True,
+        stdout=subprocess.PIPE,
+        stderr=subprocess.PIPE,
+        env=None,
+    )
+
+    # Assertion: Verify the error message includes the simulated stderr output
+    process_mock.stderr.read.assert_called_once()
+
+
 @mock.patch("litdata.streaming.downloader._GOOGLE_STORAGE_AVAILABLE", True)
 def test_gcp_downloader(tmpdir, monkeypatch, google_mock):
     # Create mock objects