From fc8e1340521fc585ea52514748509ad2f4a28b5b Mon Sep 17 00:00:00 2001 From: bhimrazy Date: Tue, 14 Jul 2026 00:00:48 +0545 Subject: [PATCH] fix(r2-downloader): reuse client for byte-range reads --- src/litdata/streaming/downloader.py | 2 +- tests/streaming/test_downloader.py | 21 +++++++++++++++++++++ 2 files changed, 22 insertions(+), 1 deletion(-) diff --git a/src/litdata/streaming/downloader.py b/src/litdata/streaming/downloader.py index 53558beec..ae335a17c 100644 --- a/src/litdata/streaming/downloader.py +++ b/src/litdata/streaming/downloader.py @@ -268,7 +268,7 @@ def download_bytes(self, remote_filepath: str, offset: int, length: int, local_c if obj.scheme != "r2": raise ValueError(f"Expected obj.scheme to be `r2`, instead, got {obj.scheme} for remote={remote_filepath}") - if not hasattr(self, "client"): + if not hasattr(self, "_client"): self._client = R2Client(storage_options=self._storage_options, session_options=self.session_options) bucket = obj.netloc diff --git a/tests/streaming/test_downloader.py b/tests/streaming/test_downloader.py index 4461a728e..6c1b7eaf3 100644 --- a/tests/streaming/test_downloader.py +++ b/tests/streaming/test_downloader.py @@ -116,6 +116,27 @@ def test_r2_downloader_error_handling(r2_client_mock, tmpdir): r2_client_instance.client.download_file.assert_called_once() +@mock.patch("litdata.streaming.downloader.R2Client") +def test_r2_downloader_download_bytes_reuses_client(r2_client_mock, tmpdir): + r2_client_instance = MagicMock() + r2_client_mock.return_value = r2_client_instance + + body = MagicMock() + body.read.return_value = b"hello" + r2_client_instance.client.get_object.return_value = {"Body": body} + + downloader = R2Downloader("r2://random_bucket", str(tmpdir), []) + + assert downloader.download_bytes("r2://random_bucket/a.txt", 0, 5, os.path.join(tmpdir, "a.txt")) == b"hello" + assert downloader.download_bytes("r2://random_bucket/a.txt", 5, 5, os.path.join(tmpdir, "a.txt")) == b"hello" + + r2_client_mock.assert_called_once_with(storage_options={}, session_options={}) + assert r2_client_instance.client.get_object.call_args_list == [ + mock.call(Bucket="random_bucket", Key="a.txt", Range="bytes=0-4"), + mock.call(Bucket="random_bucket", Key="a.txt", Range="bytes=5-9"), + ] + + @mock.patch("litdata.streaming.downloader._GOOGLE_STORAGE_AVAILABLE", True) def test_gcp_downloader(tmpdir, monkeypatch, google_mock): # Create mock objects