Remove s5cmd from the R2 downloader (#714)

Author: pwgardipee

src/litdata/streaming/downloader.py

@@ -273,12 +273,9 @@ def __init__(
         **kwargs: Any,
     ):
         super().__init__(remote_dir, cache_dir, chunks, storage_options)
-        self._s5cmd_available = os.system("s5cmd > /dev/null 2>&1") == 0
         # check if kwargs contains session_options
         self.session_options = kwargs.get("session_options", {})
-
-        if not self._s5cmd_available or _DISABLE_S5CMD:
-            self._client = R2Client(storage_options=self._storage_options, session_options=self.session_options)
+        self._client = R2Client(storage_options=self._storage_options, session_options=self.session_options)
 
     def download_file(self, remote_filepath: str, local_filepath: str) -> None:
         obj = parse.urlparse(remote_filepath)
@@ -293,61 +290,19 @@ def download_file(self, remote_filepath: str, local_filepath: str) -> None:
             suppress(Timeout, FileNotFoundError),
             FileLock(local_filepath + ".lock", timeout=1 if obj.path.endswith(_INDEX_FILENAME) else 0),
         ):
-            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(
-                    cmd,
-                    shell=True,
-                    stdout=subprocess.PIPE,
-                    stderr=subprocess.PIPE,
-                    env=env,
+            from boto3.s3.transfer import TransferConfig
+
+            extra_args: dict[str, Any] = {}
+
+            if not os.path.exists(local_filepath):
+                # Issue: https://github.com/boto/boto3/issues/3113
+                self._client.client.download_file(
+                    obj.netloc,
+                    obj.path.lstrip("/"),
+                    local_filepath,
+                    ExtraArgs=extra_args,
+                    Config=TransferConfig(use_threads=False),
                 )
-                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
-
-                extra_args: dict[str, Any] = {}
-
-                if not os.path.exists(local_filepath):
-                    # Issue: https://github.com/boto/boto3/issues/3113
-                    self._client.client.download_file(
-                        obj.netloc,
-                        obj.path.lstrip("/"),
-                        local_filepath,
-                        ExtraArgs=extra_args,
-                        Config=TransferConfig(use_threads=False),
-                    )
 
     def download_bytes(self, remote_filepath: str, offset: int, length: int, local_chunkpath: str) -> bytes:
         obj = parse.urlparse(remote_filepath)

tests/streaming/test_downloader.py

@@ -248,116 +248,32 @@ def test_s3_downloader_s5cmd_error_handling(popen_mock, system_mock, tmpdir):
     process_mock.stderr.read.assert_called_once()
 
 
-def test_r2_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 = R2Downloader(tmpdir, tmpdir, [])
-    downloader.download_file("r2://random_bucket/a.txt", os.path.join(tmpdir, "a.txt"))
-    popen_mock.wait.assert_called()
-
-
-@patch("os.system")
-@patch("subprocess.Popen")
-def test_r2_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 R2Downloader without storage options
-    downloader = R2Downloader("r2://random_bucket", str(tmpdir), [])
-
-    # Action: Call the download_file method
-    remote_filepath = "r2://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 called with correct arguments and no env variables
-    popen_mock.assert_called_once_with(
-        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("litdata.streaming.downloader.R2Client")
-def test_r2_downloader_s5cmd_available_but_disabled(r2_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
-
+def test_r2_downloader_fast(r2_client_mock, tmpdir):
     # Mock the R2Client
     r2_client_instance = MagicMock()
     r2_client_mock.return_value = r2_client_instance
 
     # Mock the download_file method to avoid credential errors
     r2_client_instance.client.download_file = MagicMock()
 
-    # Initialize the R2Downloader
     downloader = R2Downloader("r2://random_bucket", str(tmpdir), [])
+    downloader.download_file("r2://random_bucket/a.txt", os.path.join(tmpdir, "a.txt"))
 
-    # Action: Call the download_file method
-    remote_filepath = "r2://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 R2Client download_file was called
+    # Verify R2Client download_file was called
     r2_client_instance.client.download_file.assert_called_once()
 
 
-@patch("os.system")
-@patch("subprocess.Popen")
-def test_r2_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"}
-
-    # Initialize the R2Downloader with storage options
-    downloader = R2Downloader("r2://random_bucket", str(tmpdir), [], storage_options)
-
-    # Action: Call the download_file method
-    remote_filepath = "r2://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 cp {remote_filepath} {local_filepath}",
-        shell=True,
-        stdout=subprocess.PIPE,
-        stderr=subprocess.PIPE,
-        env=expected_env,
-    )
-    process_mock.wait.assert_called_once()
-
+@mock.patch("litdata.streaming.downloader.R2Client")
+def test_r2_downloader_with_storage_options(r2_client_mock, tmpdir):
+    storage_options = {"data_connection_id": "test_connection_id"}
 
-@patch("os.system")
-@patch("subprocess.Popen")
-def test_r2_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
+    # Mock the R2Client
+    r2_client_instance = MagicMock()
+    r2_client_mock.return_value = r2_client_instance
 
-    storage_options = {"AWS_NO_SIGN_REQUEST": "Yes"}
+    # Mock the download_file method to avoid credential errors
+    r2_client_instance.client.download_file = MagicMock()
 
     # Initialize the R2Downloader with storage options
     downloader = R2Downloader("r2://random_bucket", str(tmpdir), [], storage_options)
@@ -367,51 +283,34 @@ def test_r2_downloader_with_s5cmd_with_storage_options_unsigned(popen_mock, syst
     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 R2Client was initialized with storage options
+    r2_client_mock.assert_called_once_with(storage_options=storage_options, session_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()
+    # Assertion: Verify R2Client download_file was called
+    r2_client_instance.client.download_file.assert_called_once()
 
 
-@patch("os.system")
-@patch("subprocess.Popen")
-def test_r2_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 R2 error message"
-    popen_mock.return_value = process_mock
+@mock.patch("litdata.streaming.downloader.R2Client")
+def test_r2_downloader_error_handling(r2_client_mock, tmpdir):
+    # Mock the R2Client to raise an exception
+    r2_client_instance = MagicMock()
+    r2_client_mock.return_value = r2_client_instance
+
+    # Mock the download_file method to raise an exception
+    r2_client_instance.client.download_file.side_effect = Exception("Simulated R2 error")
 
-    # Initialize the R2Downloader without storage options
+    # Initialize the R2Downloader
     downloader = R2Downloader("r2://random_bucket", str(tmpdir), [])
 
-    # Action: Call the download_file method and expect a RuntimeError
+    # Action: Call the download_file method and expect an exception
     remote_filepath = "r2://random_bucket/sample_file.txt"
     local_filepath = os.path.join(tmpdir, "sample_file.txt")
 
-    with pytest.raises(RuntimeError, match="Failed to execute command"):
+    with pytest.raises(Exception, match="Simulated R2 error"):
         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()
+    # Assertion: Verify R2Client download_file was called
+    r2_client_instance.client.download_file.assert_called_once()
 
 
 @mock.patch("litdata.streaming.downloader._GOOGLE_STORAGE_AVAILABLE", True)