litdata 0.2.61__tar.gz → 0.2.65__tar.gz

{litdata-0.2.61/src/litdata.egg-info → litdata-0.2.65}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: litdata
-Version: 0.2.61
+Version: 0.2.65
 Summary: The Deep Learning framework to train, deploy, and ship AI products Lightning fast.
 Home-page: https://github.com/Lightning-AI/litdata
 Download-URL: https://github.com/Lightning-AI/litdata
@@ -39,10 +39,10 @@ Requires-Dist: obstore
 Provides-Extra: extras
 Requires-Dist: fsspec; extra == "extras"
 Requires-Dist: google-cloud-storage; extra == "extras"
-Requires-Dist: lightning-sdk==2025.12.17; extra == "extras"
+Requires-Dist: lightning-sdk==2026.2.6; extra == "extras"
 Requires-Dist: pillow; extra == "extras"
 Requires-Dist: polars; extra == "extras"
-Requires-Dist: pyarrow; extra == "extras"
+Requires-Dist: pyarrow<25.0.0; extra == "extras"
 Requires-Dist: tqdm; extra == "extras"
 Requires-Dist: viztracer; extra == "extras"
 Dynamic: author

{litdata-0.2.61 → litdata-0.2.65}/src/litdata/__about__.py

@@ -14,7 +14,7 @@
 
 import time
 
-__version__ = "0.2.61"
+__version__ = "0.2.65"
 __author__ = "Lightning AI et al."
 __author_email__ = "pytorch@lightning.ai"
 __license__ = "Apache-2.0"

{litdata-0.2.61 → litdata-0.2.65}/src/litdata/cli/commands.py

@@ -13,9 +13,9 @@
 
 from argparse import _SubParsersAction
 
+from litdata.cli import LitFormatter
 from litdata.cli.handler.cache import clear_cache, show_cache_path
 from litdata.cli.handler.optimize import optimize_dataset
-from litdata.cli.parser import LitFormatter
 
 
 def register_cache_subcommand(subparser: _SubParsersAction) -> None:

litdata-0.2.65/src/litdata/cli/handler/__init__.py

@@ -0,0 +1,21 @@
+# Copyright The Lightning AI team.
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from litdata.cli.handler.cache import clear_cache, show_cache_path
+from litdata.cli.handler.optimize import optimize_dataset
+
+__all__ = [
+    "clear_cache",
+    "show_cache_path",
+    "optimize_dataset",
+]

litdata-0.2.65/src/litdata/cli/handler/cache.py

@@ -0,0 +1,32 @@
+# Copyright The Lightning AI team.
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import shutil
+from argparse import Namespace
+
+from litdata.utilities.dataset_utilities import get_default_cache_dir
+
+
+def clear_cache(args: Namespace) -> None:
+    """Clear default cache used for StreamingDataset and other utilities."""
+    streaming_default_cache_dir = get_default_cache_dir()
+
+    shutil.rmtree(streaming_default_cache_dir, ignore_errors=True)
+
+    print(f"Cache directory '{streaming_default_cache_dir}' cleared.")
+
+
+def show_cache_path(args: Namespace) -> None:
+    """Show the path to the cache directory."""
+    streaming_default_cache_dir = get_default_cache_dir()
+    print(f"Default cache directory: {streaming_default_cache_dir}")

litdata-0.2.65/src/litdata/cli/handler/optimize.py

@@ -0,0 +1,19 @@
+# Copyright The Lightning AI team.
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from argparse import Namespace
+
+
+def optimize_dataset(args: Namespace) -> None:
+    """Handle the optimize command."""
+    print(f"Optimizing dataset at {args.dataset}...")

{litdata-0.2.61 → litdata-0.2.65}/src/litdata/constants.py

@@ -49,6 +49,7 @@ _PYARROW_AVAILABLE = RequirementCache("pyarrow")
 _POLARS_AVAILABLE = RequirementCache("polars>1.0.0")
 _PIL_AVAILABLE = RequirementCache("PIL")
 _TORCH_VISION_AVAILABLE = RequirementCache("torchvision")
+_TORCH_VISION_LESS_THAN_0_26 = RequirementCache("torchvision<0.26.0")
 _AV_AVAILABLE = RequirementCache("av")
 _OBSTORE_AVAILABLE = RequirementCache("obstore")
 
@@ -58,6 +59,8 @@ _PRINT_DEBUG_LOGS = bool(int(os.getenv("PRINT_DEBUG_LOGS", "0")))
 _MAX_WAIT_TIME = int(os.getenv("MAX_WAIT_TIME", "120"))
 _FORCE_DOWNLOAD_TIME = int(os.getenv("FORCE_DOWNLOAD_TIME", "30"))
 _LITDATA_DISABLE_VERSION_CHECK = int(os.getenv("LITDATA_DISABLE_VERSION_CHECK", "0"))
+# Experimental async chunk prefetch: set env ``LITDATA_ASYNC_CHUNK_PREFETCH=1``
+# (see ``litdata.streaming.async_prefetch``). Off by default.
 
 # DON'T CHANGE ORDER
 _TORCH_DTYPES_MAPPING = {

{litdata-0.2.61 → litdata-0.2.65}/src/litdata/processing/data_processor.py

@@ -160,7 +160,7 @@ def _download_data_target(
                 if input_dir.path:
                     local_path = path.replace(input_dir.path, cache_dir)
 
-                if input_dir.url and input_dir.path:
+                if input_dir.url and input_dir.path and not os.path.isfile(path):
                     path = path.replace(input_dir.path, input_dir.url)
 
                 obj = parse.urlparse(path)

litdata-0.2.65/src/litdata/streaming/async_prefetch.py

@@ -0,0 +1,256 @@
+# Copyright The Lightning AI team.
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Experimental asyncio helpers for overlapping **remote chunk downloads**.
+
+This is intentionally **not** an async ``StreamingDataLoader``. Training stays on
+the sync ``for batch in loader`` API; decode stays on process workers. Asyncio is
+only useful where we wait on network IO.
+
+On by default for remote datasets; force off with ``LITDATA_ASYNC_CHUNK_PREFETCH=0``,
+or force on locally with ``LITDATA_ASYNC_CHUNK_PREFETCH=1``.
+
+Strategy:
+  * Prefer ``Downloader.adownload_file`` (streaming to disk) when overridden.
+  * Else ``Downloader.adownload_fileobj`` + atomic write.
+  * Otherwise run sync ``download_chunk_from_index`` in ``asyncio.to_thread`` and
+    ``gather`` several chunk indexes — still overlaps latency for blocking cloud SDKs.
+
+Real-S3 note (Studio benches on ~67MB ImageNet chunks): gather concurrency needs
+``max_pre_download >= 4`` to matter; the prepare thread raises the floor when
+async prefetch is enabled (see :func:`async_prefetch_min_pre_download`).
+"""
+
+from __future__ import annotations
+
+import asyncio
+import contextlib
+import logging
+import os
+import threading
+from typing import TYPE_CHECKING
+
+from litdata.streaming.downloader import Downloader
+
+if TYPE_CHECKING:
+    from litdata.streaming.config import ChunksConfig
+
+logger = logging.getLogger("litdata.streaming.async_prefetch")
+
+# Thread-local event loop so PrepareChunksThread does not pay asyncio.run()
+# startup on every drain batch.
+_THREAD_LOOPS = threading.local()
+
+# Empirically, async gather on real S3 is bottlenecked when max_pre_download==2
+# (only 1–2 in-flight). Floor to 4 when the feature is enabled unless overridden.
+_DEFAULT_ASYNC_MIN_PRE_DOWNLOAD = 4
+
+
+def async_chunk_prefetch_enabled(remote_dir: str | None = None) -> bool:
+    """Return True when async chunk prefetch should run.
+
+    Env ``LITDATA_ASYNC_CHUNK_PREFETCH`` wins when set (``1``/``0``). When unset,
+    default **on** for remote datasets (real-S3: process+async is the fast path)
+    and **off** for local-only caches.
+    """
+    raw = os.getenv("LITDATA_ASYNC_CHUNK_PREFETCH")
+    if raw is not None:
+        return bool(int(raw))
+    return bool(remote_dir)
+
+
+def async_prefetch_min_pre_download() -> int:
+    """Minimum ``max_pre_download`` applied when async chunk prefetch is on.
+
+    Override with ``LITDATA_ASYNC_MIN_PRE_DOWNLOAD`` (default 4). Set to ``0`` to
+    disable the floor and keep the caller's ``max_pre_download`` unchanged.
+    """
+    raw = os.getenv("LITDATA_ASYNC_MIN_PRE_DOWNLOAD")
+    if raw is None:
+        return _DEFAULT_ASYNC_MIN_PRE_DOWNLOAD
+    return max(0, int(raw))
+
+
+def apply_async_pre_download_floor(max_pre_download: int, remote_dir: str | None = None) -> int:
+    """Raise ``max_pre_download`` when async prefetch needs gather width."""
+    if not async_chunk_prefetch_enabled(remote_dir):
+        return max_pre_download
+    floor = async_prefetch_min_pre_download()
+    if floor <= 0 or max_pre_download >= floor:
+        return max_pre_download
+    logger.info(
+        "Async chunk prefetch: raising max_pre_download %s → %s so "
+        "asyncio.gather can overlap remote chunk downloads. "
+        "Set LITDATA_ASYNC_MIN_PRE_DOWNLOAD=0 to keep the original value.",
+        max_pre_download,
+        floor,
+    )
+    return floor
+
+
+def downloader_supports_adownload(downloader: Downloader | None) -> bool:
+    """True when an async download path is available on ``downloader``."""
+    if downloader is None:
+        return False
+    cls = type(downloader)
+    return (
+        cls.adownload_file is not Downloader.adownload_file or cls.adownload_fileobj is not Downloader.adownload_fileobj
+    )
+
+
+def _remote_join(remote_dir: str, filename: str) -> str:
+    """Join cloud URLs without ``os.path.join`` quirks on schemes."""
+    return remote_dir.rstrip("/") + "/" + filename.lstrip("/")
+
+
+async def _adownload_file_to_path(downloader: Downloader, remote_filepath: str, local_filepath: str) -> None:
+    """Fetch ``remote_filepath`` asynchronously and publish atomically."""
+    if os.path.exists(local_filepath):
+        return
+    # Prefer streaming-to-disk when the backend overrides adownload_file.
+    if type(downloader).adownload_file is not Downloader.adownload_file:
+        await downloader.adownload_file(remote_filepath, local_filepath)
+        return
+    data = await downloader.adownload_fileobj(remote_filepath)
+    if data is None:
+        raise NotImplementedError(
+            f"{type(downloader).__name__}.adownload_fileobj returned None; "
+            "cannot use async chunk prefetch for this backend."
+        )
+    tmp_path = downloader._temp_download_path(local_filepath)
+    try:
+        os.makedirs(os.path.dirname(local_filepath) or ".", exist_ok=True)
+        with open(tmp_path, "wb") as f:
+            f.write(data)
+        downloader._atomic_replace(tmp_path, local_filepath)
+    except Exception:
+        with contextlib.suppress(FileNotFoundError, PermissionError):
+            os.remove(tmp_path)
+        raise
+
+
+async def _adownload_chunk_index(config: ChunksConfig, chunk_index: int) -> None:
+    """Async download + decompress for one chunk index (mirrors sync config path)."""
+    assert config._chunks is not None
+    downloader = config._downloader
+    if downloader is None:
+        return
+
+    chunk_filename = config._chunks[chunk_index]["filename"]
+    local_chunkpath = os.path.join(config._cache_dir, chunk_filename)
+    remote_chunkpath = _remote_join(downloader._remote_dir, chunk_filename)
+    lazily_ref_counted = chunk_index not in config._shared_chunk_indexes
+    lock_path = (
+        local_chunkpath.replace(f".{config._compressor_name}", "") if config._compressor_name else local_chunkpath
+    )
+
+    if os.path.exists(local_chunkpath):
+        config.try_decompress(local_chunkpath)
+        if lazily_ref_counted:
+            downloader._increment_local_lock(lock_path, chunk_index)
+        return
+
+    if lazily_ref_counted:
+        downloader._increment_local_lock(lock_path, chunk_index)
+
+    if downloader_supports_adownload(downloader):
+        await _adownload_file_to_path(downloader, remote_chunkpath, local_chunkpath)
+    else:
+        # Overlap blocking SDK calls across threads when native async is unavailable.
+        await asyncio.to_thread(downloader.download_chunk_from_index, chunk_index)
+
+    config.try_decompress(local_chunkpath)
+
+
+async def adownload_chunk_indexes(config: ChunksConfig, chunk_indexes: list[int]) -> None:
+    """Download several chunk indexes concurrently (gather)."""
+    if not chunk_indexes:
+        return
+    if len(chunk_indexes) == 1:
+        await _adownload_chunk_index(config, chunk_indexes[0])
+        return
+    await asyncio.gather(*[_adownload_chunk_index(config, idx) for idx in chunk_indexes])
+
+
+def _thread_event_loop() -> asyncio.AbstractEventLoop:
+    """Return a reusable event loop for the current thread."""
+    loop = getattr(_THREAD_LOOPS, "loop", None)
+    if loop is None or loop.is_closed():
+        loop = asyncio.new_event_loop()
+        _THREAD_LOOPS.loop = loop
+    return loop
+
+
+def close_thread_event_loop() -> None:
+    """Shut down the thread-local loop and its default executor.
+
+    ``asyncio.to_thread`` (used when a downloader has no native async path) parks
+    workers named ``asyncio_N`` on the loop's default executor. Leaving the loop
+    open leaks those threads and trips the session thread-police on Windows.
+
+    Executor shutdown is non-blocking (``wait=False``) so prepare-thread ``finally``
+    never joins forever if a download worker is stuck. Threads are marked daemon and
+    futures cancelled so orphans cannot keep the process alive or poison later
+    DataLoader forks under pytest-xdist.
+    """
+    loop = getattr(_THREAD_LOOPS, "loop", None)
+    _THREAD_LOOPS.loop = None
+    if loop is None or loop.is_closed():
+        return
+    try:
+        if not loop.is_running():
+            with contextlib.suppress(Exception):
+                loop.run_until_complete(loop.shutdown_asyncgens())
+            # Prefer non-blocking teardown over ``shutdown_default_executor()``,
+            # which joins executor threads and can hang prepare-thread exit.
+            executor = getattr(loop, "_default_executor", None)
+            if executor is not None:
+                with contextlib.suppress(Exception):
+                    loop._default_executor = None
+                    for thread in getattr(executor, "_threads", set()):
+                        thread.daemon = True
+                    # cancel_futures is 3.9+; litdata already requires newer Python.
+                    executor.shutdown(wait=False, cancel_futures=True)
+    finally:
+        with contextlib.suppress(Exception):
+            loop.close()
+
+
+def download_chunk_indexes_concurrently(config: ChunksConfig, chunk_indexes: list[int]) -> None:
+    """Sync entry point for ``PrepareChunksThread``: run ``adownload_chunk_indexes``."""
+    if not chunk_indexes:
+        return
+    if len(chunk_indexes) == 1:
+        config.download_chunk_from_index(chunk_indexes[0])
+        return
+    # Reuse a per-thread loop instead of asyncio.run() (which creates+closes a loop
+    # on every prefetch batch — measurable overhead under high chunk churn).
+    try:
+        running = asyncio.get_running_loop()
+    except RuntimeError:
+        running = None
+    if running is not None:
+        # Already inside an event loop (tests / nested callers): block via a bridge.
+        import concurrent.futures
+
+        def _run() -> None:
+            try:
+                _thread_event_loop().run_until_complete(adownload_chunk_indexes(config, chunk_indexes))
+            finally:
+                close_thread_event_loop()
+
+        with concurrent.futures.ThreadPoolExecutor(max_workers=1) as pool:
+            pool.submit(_run).result()
+        return
+    loop = _thread_event_loop()
+    loop.run_until_complete(adownload_chunk_indexes(config, chunk_indexes))

{litdata-0.2.61 → litdata-0.2.65}/src/litdata/streaming/client.py

@@ -45,6 +45,63 @@ class _CustomRetryAdapter(HTTPAdapter):
         return super().send(request, **kwargs)
 
 
+def _login_and_get_temp_bucket_credentials(data_connection_id: str) -> dict[str, Any]:
+    """Mint temporary bucket credentials for a data connection via the Lightning Cloud API.
+
+    Shared by R2 (lightning storage) connections and by S3 connections marked
+    ``available_in_non_aws_providers``: both bypass the FUSE mount and need short-lived creds
+    from the same control-plane ``temp-bucket-credentials`` endpoint. Returns the raw response
+    (accessKeyId/secretAccessKey/sessionToken, plus accountId for R2).
+    """
+    retry_strategy = Retry(
+        total=_CONNECTION_RETRY_TOTAL,
+        backoff_factor=_CONNECTION_RETRY_BACKOFF_FACTOR,
+        status_forcelist=[
+            408,  # Request Timeout
+            429,  # Too Many Requests
+            500,  # Internal Server Error
+            502,  # Bad Gateway
+            503,  # Service Unavailable
+            504,  # Gateway Timeout
+        ],
+    )
+    adapter = _CustomRetryAdapter(max_retries=retry_strategy, timeout=_DEFAULT_REQUEST_TIMEOUT)
+    session = requests.Session()
+    session.mount("http://", adapter)
+    session.mount("https://", adapter)
+
+    cloud_url = os.getenv("LIGHTNING_CLOUD_URL", "https://lightning.ai")
+    api_key = os.getenv("LIGHTNING_API_KEY")
+    username = os.getenv("LIGHTNING_USERNAME")
+    project_id = os.getenv("LIGHTNING_CLOUD_PROJECT_ID")
+
+    if not all([api_key, username, project_id]):
+        raise RuntimeError("Missing required environment variables")
+
+    # Login to get token
+    payload = {"apiKey": api_key, "username": username}
+    login_url = f"{cloud_url}/v1/auth/login"
+    response = session.post(login_url, data=json.dumps(payload))
+
+    if "token" not in response.json():
+        raise RuntimeError("Failed to get authentication token")
+
+    token = response.json()["token"]
+
+    # Get temporary bucket credentials
+    headers = {"Authorization": f"Bearer {token}", "Content-Type": "application/json"}
+    credentials_url = (
+        f"{cloud_url}/v1/projects/{project_id}/data-connections/{data_connection_id}/temp-bucket-credentials"
+    )
+
+    credentials_response = session.get(credentials_url, headers=headers, timeout=10)
+
+    if credentials_response.status_code != 200:
+        raise RuntimeError(f"Failed to get credentials: {credentials_response.status_code}")
+
+    return credentials_response.json()
+
+
 class S3Client:
     # TODO: Generalize to support more cloud providers.
 
@@ -61,6 +118,15 @@ class S3Client:
         self._session_options: dict = session_options or {}
 
     def _create_client(self) -> None:
+        # S3 data connections marked available on non-AWS providers can't reach the bucket via the
+        # FUSE mount or an instance profile off AWS, so mint temporary project-role credentials from
+        # the control plane instead (mirrors R2Client). Gated on data_connection_id being threaded
+        # through storage_options, so plain S3 access is unaffected.
+        data_connection_id = self._storage_options.get("data_connection_id")
+        if data_connection_id:
+            self._create_client_from_temp_credentials(data_connection_id)
+            return
+
         has_shared_credentials_file = (
             os.getenv("AWS_SHARED_CREDENTIALS_FILE") == os.getenv("AWS_CONFIG_FILE") == "/.credentials/.aws_credentials"
         )
@@ -86,6 +152,29 @@ class S3Client:
                 config=botocore.config.Config(retries={"max_attempts": 1000, "mode": "adaptive"}),
             )
 
+    def _create_client_from_temp_credentials(self, data_connection_id: str) -> None:
+        """Create an S3 client backed by temporary project-role credentials for a data connection.
+
+        Used for S3 connections available on non-AWS providers. Unlike R2 there is no custom
+        endpoint — this is real AWS S3, so botocore resolves the bucket region on first use.
+        """
+        temp_credentials = _login_and_get_temp_bucket_credentials(data_connection_id)
+
+        # data_connection_id is our own metadata; drop it before handing options to boto3.
+        storage_options = {k: v for k, v in self._storage_options.items() if k != "data_connection_id"}
+
+        session = boto3.Session(**self._session_options)
+        self._client = session.client(
+            "s3",
+            **{
+                "aws_access_key_id": temp_credentials["accessKeyId"],
+                "aws_secret_access_key": temp_credentials["secretAccessKey"],
+                "aws_session_token": temp_credentials["sessionToken"],
+                "config": botocore.config.Config(retries={"max_attempts": 1000, "mode": "adaptive"}),
+                **storage_options,
+            },
+        )
+
     @property
     def client(self) -> Any:
         if self._client is None:
@@ -121,56 +210,8 @@ class R2Client(S3Client):
 
     def get_r2_bucket_credentials(self, data_connection_id: str) -> dict[str, str]:
         """Fetch temporary R2 credentials for the current lightning storage connection."""
-        # Create session with retry logic
-        retry_strategy = Retry(
-            total=_CONNECTION_RETRY_TOTAL,
-            backoff_factor=_CONNECTION_RETRY_BACKOFF_FACTOR,
-            status_forcelist=[
-                408,  # Request Timeout
-                429,  # Too Many Requests
-                500,  # Internal Server Error
-                502,  # Bad Gateway
-                503,  # Service Unavailable
-                504,  # Gateway Timeout
-            ],
-        )
-        adapter = _CustomRetryAdapter(max_retries=retry_strategy, timeout=_DEFAULT_REQUEST_TIMEOUT)
-        session = requests.Session()
-        session.mount("http://", adapter)
-        session.mount("https://", adapter)
-
         try:
-            # Get Lightning Cloud API token
-            cloud_url = os.getenv("LIGHTNING_CLOUD_URL", "https://lightning.ai")
-            api_key = os.getenv("LIGHTNING_API_KEY")
-            username = os.getenv("LIGHTNING_USERNAME")
-            project_id = os.getenv("LIGHTNING_CLOUD_PROJECT_ID")
-
-            if not all([api_key, username, project_id]):
-                raise RuntimeError("Missing required environment variables")
-
-            # Login to get token
-            payload = {"apiKey": api_key, "username": username}
-            login_url = f"{cloud_url}/v1/auth/login"
-            response = session.post(login_url, data=json.dumps(payload))
-
-            if "token" not in response.json():
-                raise RuntimeError("Failed to get authentication token")
-
-            token = response.json()["token"]
-
-            # Get temporary bucket credentials
-            headers = {"Authorization": f"Bearer {token}", "Content-Type": "application/json"}
-            credentials_url = (
-                f"{cloud_url}/v1/projects/{project_id}/data-connections/{data_connection_id}/temp-bucket-credentials"
-            )
-
-            credentials_response = session.get(credentials_url, headers=headers, timeout=10)
-
-            if credentials_response.status_code != 200:
-                raise RuntimeError(f"Failed to get credentials: {credentials_response.status_code}")
-
-            temp_credentials = credentials_response.json()
+            temp_credentials = _login_and_get_temp_bucket_credentials(data_connection_id)
 
             endpoint_url = f"https://{temp_credentials['accountId']}.r2.cloudflarestorage.com"