litdata 0.2.71__tar.gz → 0.2.73__tar.gz

{litdata-0.2.71/src/litdata.egg-info → litdata-0.2.73}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: litdata
-Version: 0.2.71
+Version: 0.2.73
 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
@@ -34,6 +34,7 @@ Requires-Dist: filelock
 Requires-Dist: numpy
 Requires-Dist: boto3
 Requires-Dist: requests
+Requires-Dist: urllib3>=1.26
 Requires-Dist: tifffile
 Requires-Dist: obstore
 Provides-Extra: extras

{litdata-0.2.71 → litdata-0.2.73}/requirements.txt

@@ -5,5 +5,6 @@ filelock
 numpy
 boto3
 requests
+urllib3 >=1.26  # Retry(allowed_methods=...) in streaming/client.py
 tifffile
 obstore

{litdata-0.2.71 → litdata-0.2.73}/src/litdata/__about__.py

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

{litdata-0.2.71 → litdata-0.2.73}/src/litdata/processing/data_processor.py

@@ -1523,7 +1523,16 @@ class DataChunkRecipe(DataRecipe):
 
         merge_cache = Cache(cache_dir, chunk_bytes=1)
         node_rank = _get_node_rank()
-        merge_cache._merge_no_wait(node_rank if num_nodes > 1 else None, getattr(self, "existing_index", None))
+        # With more than one node, every node would otherwise fold the existing
+        # index into its own {node_rank}-index.json, and the final cross-node
+        # merge below then repeats those chunks once per node. Add the existing
+        # index only once, either here for the single-node case or in the final
+        # merge for the multi-node case.
+        existing_index = getattr(self, "existing_index", None)
+        merge_cache._merge_no_wait(
+            node_rank if num_nodes > 1 else None,
+            existing_index if num_nodes == 1 else None,
+        )
 
         self._merge_and_upload_keys(output_dir, cache_dir, num_nodes, node_rank)
         self._upload_index(output_dir, cache_dir, num_nodes, node_rank)
@@ -1734,7 +1743,9 @@ class DataChunkRecipe(DataRecipe):
                     shutil.copyfile(remote_filepath, node_index_filepath)
 
             merge_cache = Cache(merge_dir, chunk_bytes=1)
-            merge_cache._merge_no_wait()
+            # The per-node index files hold only their own new chunks now, so the
+            # existing index is folded in here, once, as the node files are merged.
+            merge_cache._merge_no_wait(existing_index=getattr(self, "existing_index", None))
             self._upload_index(output_dir, merge_dir, 1, None)
 
 

{litdata-0.2.71 → litdata-0.2.73}/src/litdata/processing/functions.py

@@ -147,6 +147,17 @@ class LambdaMapRecipe(MapRecipe):
         self._contains_device = "device" in params
         self._contains_is_last = "is_last" in params
 
+    def __getstate__(self) -> dict[str, Any]:
+        """Drop the full input sequence when pickling into spawn workers.
+
+        The parent already shards items onto per-worker queues. Workers only need
+        ``prepare_item``; keeping ``_inputs`` in the pickle would duplicate the list
+        once per process.
+        """
+        state = self.__dict__.copy()
+        state["_inputs"] = None
+        return state
+
     def prepare_structure(self, _: str | None) -> Any:
         return self._inputs
 
@@ -210,6 +221,17 @@ class LambdaDataChunkRecipe(DataChunkRecipe):
 
         self.prepare_item = self._prepare_item_generator if self.is_generator else self._prepare_item  # type: ignore
 
+    def __getstate__(self) -> dict[str, Any]:
+        """Drop the full input sequence when pickling into spawn workers.
+
+        The parent already shards items onto per-worker queues. Workers only need
+        ``prepare_item``; keeping ``_inputs`` in the pickle would duplicate the list
+        once per process.
+        """
+        state = self.__dict__.copy()
+        state["_inputs"] = None
+        return state
+
     def check_fn(self) -> None:
         if (
             isinstance(self._fn, (partial, FunctionType))

{litdata-0.2.71 → litdata-0.2.73}/src/litdata/streaming/client.py

@@ -12,9 +12,11 @@
 # limitations under the License.
 
 import json
+import logging
 import os
+import random
 import threading
-from time import time
+from time import sleep, time
 from typing import Any
 
 import boto3
@@ -27,14 +29,60 @@ from urllib3.util.retry import Retry
 
 from litdata.constants import _IS_IN_STUDIO
 
+logger = logging.getLogger("litdata.streaming.client")
+
 # Constants for the retry adapter. Docs: https://urllib3.readthedocs.io/en/stable/reference/urllib3.util.html
-# Maximum number of total connection retry attempts (e.g., 2880 retries = 24 hours with 30s timeout per request)
-_CONNECTION_RETRY_TOTAL = 2880
+# Retries per request. Deliberately small: a refresh runs inline on whichever DataLoader worker
+# thread crossed the interval and holds the client lock for its whole duration, so it has to
+# return in seconds, not hours. `S3Client._refresh_client` is what rides out a longer outage.
+_CONNECTION_RETRY_TOTAL = 4
 # Backoff factor for connection retries (wait time increases by this factor after each failure)
 _CONNECTION_RETRY_BACKOFF_FACTOR = 0.5
 # Default timeout for each HTTP request in seconds
 _DEFAULT_REQUEST_TIMEOUT = 30  # seconds
 
+# The control plane mints credentials with a 1 hour TTL for S3 project-role data connections, and
+# the response carries no expiry for us to read, so refresh well inside it. The remaining time is
+# the window in which a failed refresh can be retried while the credentials in hand still work.
+_DEFAULT_REFETCH_INTERVAL = 2700  # seconds
+# How long past the refetch interval we keep serving existing credentials while refreshes fail.
+# Sized against the TTL, not comfort: 2700 + 600 leaves ~5 minutes before the 1 hour S3 expiry,
+# so we stop before reads start failing as unexplained S3 403s.
+_REFRESH_GRACE_PERIOD = 600  # seconds
+# How long to wait for the control plane when there are no credentials yet. No TTL constrains
+# this one — nothing is being served — so it can be more patient than the refresh grace.
+_INITIAL_RETRY_BUDGET = 900  # seconds
+# Spacing between refresh attempts once one has failed. Without it every subsequent read
+# re-requests credentials, turning one control-plane blip into a request storm per worker.
+_REFRESH_RETRY_INTERVAL = 60  # seconds
+# Fraction of the interval by which each process refreshes early. DataLoader workers are forked
+# together, so without jitter they all reach the interval — and stampede — in the same instant.
+_REFETCH_JITTER_RATIO = 0.1
+
+
+class _CredentialsError(RuntimeError):
+    """Raised when credentials could not be obtained from the control plane or IMDS.
+
+    Only these are retried. Anything else out of ``_create_client`` — a bad ``storage_options``
+    key, a malformed endpoint — is a local mistake that no amount of waiting fixes, and must
+    reach the caller immediately rather than stalling the job for the retry budget.
+    """
+
+
+class _CredentialsUnavailableError(_CredentialsError):
+    """A credential failure that may clear on its own: unreachable, timed out, 5xx, 429."""
+
+
+class _CredentialsConfigurationError(_CredentialsError):
+    """A credential failure that retrying cannot fix: missing configuration, or rejected auth."""
+
+
+def _credentials_error(status_code: int, message: str) -> _CredentialsError:
+    """Classify an HTTP failure. A 4xx other than 408/429 will not fix itself on a retry."""
+    if 400 <= status_code < 500 and status_code not in (408, 429):
+        return _CredentialsConfigurationError(message)
+    return _CredentialsUnavailableError(message)
+
 
 class _CustomRetryAdapter(HTTPAdapter):
     def __init__(self, *args: Any, **kwargs: Any) -> None:
@@ -42,7 +90,11 @@ class _CustomRetryAdapter(HTTPAdapter):
         super().__init__(*args, **kwargs)
 
     def send(self, request: Any, *args: Any, **kwargs: Any) -> Any:
-        kwargs["timeout"] = kwargs.get("timeout", self.timeout)
+        # requests always passes `timeout` explicitly, as None when the caller gave none, so a
+        # `kwargs.get("timeout", self.timeout)` default never applied and every request without
+        # an explicit timeout could hang forever.
+        if kwargs.get("timeout") is None:
+            kwargs["timeout"] = self.timeout
         return super().send(request, **kwargs)
 
 
@@ -57,6 +109,10 @@ def _login_and_get_temp_bucket_credentials(data_connection_id: str) -> dict[str,
     retry_strategy = Retry(
         total=_CONNECTION_RETRY_TOTAL,
         backoff_factor=_CONNECTION_RETRY_BACKOFF_FACTOR,
+        # urllib3 leaves POST out of its default allowed_methods, which left the login below
+        # unretried: a single 429 or 502 there failed the whole refresh even though the
+        # credentials GET would have retried it. Minting a token is safe to repeat.
+        allowed_methods=Retry.DEFAULT_ALLOWED_METHODS | {"POST"},
         status_forcelist=[
             408,  # Request Timeout
             429,  # Too Many Requests
@@ -77,17 +133,27 @@ def _login_and_get_temp_bucket_credentials(data_connection_id: str) -> dict[str,
     project_id = os.getenv("LIGHTNING_CLOUD_PROJECT_ID")
 
     if not all([api_key, username, project_id]):
-        raise RuntimeError("Missing required environment variables")
+        raise _CredentialsConfigurationError("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")
+    try:
+        response = session.post(login_url, data=json.dumps(payload))
+    except requests.exceptions.RequestException as e:
+        raise _CredentialsUnavailableError(f"Could not reach the Lightning Cloud API to log in: {e}") from e
+
+    # Check the status before the body: a proxy in front of the API answers with HTML, which
+    # would otherwise surface as a JSONDecodeError with nothing pointing back at the login call.
+    if response.status_code != 200:
+        raise _credentials_error(
+            response.status_code, f"Failed to log in to the Lightning Cloud API: {response.status_code}"
+        )
 
-    token = response.json()["token"]
+    try:
+        token = response.json()["token"]
+    except (ValueError, KeyError) as e:
+        raise RuntimeError("Failed to get authentication token") from e
 
     # Get temporary bucket credentials
     headers = {"Authorization": f"Bearer {token}", "Content-Type": "application/json"}
@@ -95,10 +161,15 @@ def _login_and_get_temp_bucket_credentials(data_connection_id: str) -> dict[str,
         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)
+    try:
+        credentials_response = session.get(credentials_url, headers=headers, timeout=10)
+    except requests.exceptions.RequestException as e:
+        raise _CredentialsUnavailableError(f"Could not reach the Lightning Cloud API for credentials: {e}") from e
 
     if credentials_response.status_code != 200:
-        raise RuntimeError(f"Failed to get credentials: {credentials_response.status_code}")
+        raise _credentials_error(
+            credentials_response.status_code, f"Failed to get credentials: {credentials_response.status_code}"
+        )
 
     return credentials_response.json()
 
@@ -108,17 +179,33 @@ class S3Client:
 
     def __init__(
         self,
-        refetch_interval: int = 3300,
+        refetch_interval: int = _DEFAULT_REFETCH_INTERVAL,
         storage_options: dict | None = {},
         session_options: dict | None = {},
     ) -> None:
         self._refetch_interval = refetch_interval
         self._last_time: float | None = None
-        self._client: Any | None = None
         self._storage_options: dict = storage_options or {}
         self._session_options: dict = session_options or {}
+        self._reset_process_state()
+
+    def _reset_process_state(self) -> None:
+        """Reset the state that must not be shared across a fork.
+
+        Re-rolls the refresh jitter too: DataLoader workers inherit the parent's timings, so
+        without a fresh roll per process they would all refresh in the same instant.
+        """
+        self._client: Any | None = None
         # Guards lazy create + credential refresh (range GETs hit .client from many threads).
         self._client_lock = threading.Lock()
+        self._owner_pid = os.getpid()
+        self._refetch_deadline = self._jittered_refetch_interval()
+        self._refresh_retry_time: float | None = None
+
+    def _jittered_refetch_interval(self) -> float:
+        # Only ever early, never late: callers set the interval as an upper bound on how long
+        # a set of credentials is held, and the grace period below is measured against it.
+        return self._refetch_interval * (1.0 - random.uniform(0.0, _REFETCH_JITTER_RATIO))  # noqa: S311
 
     def __getstate__(self) -> dict[str, Any]:
         state = self.__dict__.copy()
@@ -130,9 +217,7 @@ class S3Client:
 
     def __setstate__(self, state: dict[str, Any]) -> None:
         self.__dict__.update(state)
-        self._client = None
-        self._client_lock = threading.Lock()
-        self._owner_pid = os.getpid()
+        self._reset_process_state()
 
     def _create_client(self) -> None:
         # S3 data connections marked available on non-AWS providers can't reach the bucket via the
@@ -159,7 +244,10 @@ class S3Client:
             )
         else:
             provider = InstanceMetadataProvider(iam_role_fetcher=InstanceMetadataFetcher(timeout=3600, num_attempts=5))
-            credentials = provider.load()
+            try:
+                credentials = provider.load()
+            except Exception as e:
+                raise _CredentialsUnavailableError(f"Could not load instance metadata credentials: {e}") from e
             session = boto3.Session()
             self._client = session.client(
                 "s3",
@@ -192,24 +280,91 @@ class S3Client:
             },
         )
 
+    def _mark_refreshed(self) -> None:
+        self._last_time = time()
+        self._refetch_deadline = self._jittered_refetch_interval()
+        self._refresh_retry_time = None
+
+    def _create_initial_client(self) -> None:
+        """Create the first client, waiting out a control plane that is briefly unavailable.
+
+        Unlike a refresh there are no credentials to fall back on, so this has to keep trying
+        rather than hand the caller an error. Bounded by ``_INITIAL_RETRY_BUDGET``: a job that
+        cannot reach the control plane for that long should fail with a clear reason.
+        """
+        started = time()
+        attempt = 0
+        while True:
+            attempt += 1
+            try:
+                self._create_client()
+            # Only a control plane that might come back is worth waiting for. A configuration
+            # error, or any local failure building the client, propagates on the first attempt.
+            except _CredentialsUnavailableError as e:
+                waited = time() - started
+                if waited + _REFRESH_RETRY_INTERVAL >= _INITIAL_RETRY_BUDGET:
+                    raise RuntimeError(f"Could not get credentials after {waited:.0f}s of retrying: {e}") from e
+                logger.warning(
+                    "Could not get credentials (attempt %d, %.0fs elapsed), so data loading is blocked; "
+                    "retrying in %ds: %s",
+                    attempt,
+                    waited,
+                    _REFRESH_RETRY_INTERVAL,
+                    e,
+                )
+                sleep(_REFRESH_RETRY_INTERVAL)
+            else:
+                self._mark_refreshed()
+                return
+
+    def _refresh_client(self) -> None:
+        """Re-mint credentials, tolerating a control plane that is briefly unavailable.
+
+        Credentials are refreshed before they expire, so a failed refresh is not immediately
+        fatal — the ones already in hand still work. Keep serving those and retry on a timer
+        until the grace period runs out, rather than failing the read on the first error, which
+        kills the DataLoader worker and with it the run.
+        """
+        now = time()
+        if self._refresh_retry_time is not None and now < self._refresh_retry_time:
+            return
+
+        try:
+            self._create_client()
+        # Both kinds get the grace period here, unlike initial creation. The credentials in hand
+        # still work, so there is nothing to gain by failing fast on a 403 that might be a proxy
+        # misbehaving mid-deploy — and if it is a real revocation, the deadline still catches it.
+        except _CredentialsError as e:
+            held_for = 0.0 if self._last_time is None else now - self._last_time
+            if held_for > self._refetch_deadline + _REFRESH_GRACE_PERIOD:
+                raise RuntimeError(
+                    f"Failed to refresh credentials for {held_for:.0f}s, so they are assumed expired: {e}"
+                ) from e
+            self._refresh_retry_time = now + _REFRESH_RETRY_INTERVAL
+            logger.warning(
+                "Could not refresh credentials (%.0fs since the last successful refresh); reusing the current "
+                "ones and retrying in %ds: %s",
+                held_for,
+                _REFRESH_RETRY_INTERVAL,
+                e,
+            )
+            return
+
+        self._mark_refreshed()
+
     @property
     def client(self) -> Any:
         # boto3 clients are thread-safe for requests; construction/refresh is not.
-        pid = os.getpid()
-        if getattr(self, "_owner_pid", None) != pid:
-            # DataLoader fork: drop the inherited client and lock.
-            self._client = None
-            self._client_lock = threading.Lock()
-            self._owner_pid = pid
+        if getattr(self, "_owner_pid", None) != os.getpid():
+            # DataLoader fork: drop the inherited client, lock and refresh schedule.
+            self._reset_process_state()
+
         with self._client_lock:
             if self._client is None:
-                self._create_client()
-                self._last_time = time()
-
+                self._create_initial_client()
             # Re-generate credentials for EC2 / temporary Studio creds
-            if self._last_time is None or (time() - self._last_time) > self._refetch_interval:
-                self._create_client()
-                self._last_time = time()
+            elif self._last_time is None or (time() - self._last_time) > self._refetch_deadline:
+                self._refresh_client()
 
             return self._client
 
@@ -219,7 +374,7 @@ class R2Client(S3Client):
 
     def __init__(
         self,
-        refetch_interval: int = 3300,
+        refetch_interval: int = _DEFAULT_REFETCH_INTERVAL,
         storage_options: dict | None = {},
         session_options: dict | None = {},
     ) -> None:
@@ -248,17 +403,17 @@ class R2Client(S3Client):
                 "endpoint_url": endpoint_url,
             }
 
+        except _CredentialsError:
+            raise
         except Exception as e:
-            # Fallback to hardcoded credentials if API call fails
-            print(f"Failed to get R2 credentials from API: {e}. Using fallback credentials.")
-            raise RuntimeError(f"Failed to get R2 credentials and no fallback available: {e}")
+            raise RuntimeError(f"Failed to get R2 credentials: {e}") from e
 
     def _create_client(self) -> None:
         """Create a new R2 client with fresh credentials."""
         # Get data connection ID from storage options
         data_connection_id = self._base_storage_options.get("data_connection_id")
         if not data_connection_id:
-            raise RuntimeError("data_connection_id is required in storage_options for R2 client")
+            raise _CredentialsConfigurationError("data_connection_id is required in storage_options for R2 client")
 
         # Get fresh R2 credentials
         r2_credentials = self.get_r2_bucket_credentials(data_connection_id)

{litdata-0.2.71 → litdata-0.2.73}/src/litdata/streaming/dataloader.py

@@ -558,25 +558,30 @@ class _StreamingMultiProcessingDataLoaderIter(_MultiProcessingDataLoaderIter):
         )
         self._num_workers = loader.num_workers
         self._cprofile: cProfile.Profile | None = None
-        self._orig_worker_loop: Any = None
 
         distributed_env = _DistributedEnv.detect()
         profile_cprofile = bool(getattr(self._loader, "_profile_cprofile", False))
+        from torch.utils.data._utils import worker
 
-        if distributed_env.global_rank == 0:
-            from torch.utils.data._utils import worker
+        original_worker_loop: Any = None
 
+        if distributed_env.global_rank == 0:
             if self._loader._profile_batches and _VIZ_TRACKER_AVAILABLE:
+                original_worker_loop = worker._worker_loop
                 worker._worker_loop = _ProfileWorkerLoop(
                     self._loader._profile_batches, self._loader._profile_skip_batches, self._loader._profile_dir
                 )
             elif profile_cprofile:
-                self._orig_worker_loop = worker._worker_loop
+                original_worker_loop = worker._worker_loop
                 worker._worker_loop = _CProfileWorkerLoop(_cprofile_output_dir(self._loader._profile_dir))
 
         # Workers fork/spawn here. Enable the parent profiler only after that so
         # the child does not inherit an active cProfile (one profiler per process).
-        super().__init__(loader)
+        try:
+            super().__init__(loader)
+        finally:
+            if original_worker_loop is not None:
+                worker._worker_loop = original_worker_loop
 
         if profile_cprofile and distributed_env.global_rank == 0:
             self._cprofile = cProfile.Profile()
@@ -587,11 +592,6 @@ class _StreamingMultiProcessingDataLoaderIter(_MultiProcessingDataLoaderIter):
             stem = os.path.join(_cprofile_output_dir(self._loader._profile_dir), _CPROFILE_MAIN_STEM)
             _dump_cprofile(self._cprofile, stem)
             self._cprofile = None
-        if self._orig_worker_loop is not None:
-            from torch.utils.data._utils import worker
-
-            worker._worker_loop = self._orig_worker_loop
-            self._orig_worker_loop = None
         super()._shutdown_workers()
 
     def _try_put_index(self) -> None:

{litdata-0.2.71 → litdata-0.2.73}/src/litdata/streaming/serializers.py

@@ -477,7 +477,7 @@ class NumpySerializer(Serializer):
             shape.append(np.frombuffer(data[8 + 4 * shape_idx : 8 + 4 * (shape_idx + 1)], np.uint32).item())
 
         # deserialize the numpy array bytes
-        tensor = np.frombuffer(data[8 + 4 * shape_size : len(data)], dtype=dtype)
+        tensor = np.frombuffer(data[8 + 4 * shape_size : len(data)], dtype=dtype).copy()
         if tensor.shape == shape:
             return tensor
         return np.reshape(tensor, shape)
@@ -503,7 +503,7 @@ class NoHeaderNumpySerializer(Serializer):
 
     def deserialize(self, data: bytes) -> np.ndarray:
         assert self._dtype
-        return np.frombuffer(data, dtype=self._dtype)
+        return np.frombuffer(data, dtype=self._dtype).copy()
 
     def can_serialize(self, item: np.ndarray) -> bool:
         return isinstance(item, np.ndarray) and len(item.shape) == 1

{litdata-0.2.71 → litdata-0.2.73/src/litdata.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: litdata
-Version: 0.2.71
+Version: 0.2.73
 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
@@ -34,6 +34,7 @@ Requires-Dist: filelock
 Requires-Dist: numpy
 Requires-Dist: boto3
 Requires-Dist: requests
+Requires-Dist: urllib3>=1.26
 Requires-Dist: tifffile
 Requires-Dist: obstore
 Provides-Extra: extras

{litdata-0.2.71 → litdata-0.2.73}/src/litdata.egg-info/requires.txt

@@ -5,6 +5,7 @@ filelock
 numpy
 boto3
 requests
+urllib3>=1.26
 tifffile
 obstore