litdata 0.2.61__tar.gz → 0.2.64__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: litdata
-Version: 0.2.61
+Version: 0.2.64
 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.64}/src/litdata/__about__.py

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

{litdata-0.2.61 → litdata-0.2.64}/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.64/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.64/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.64/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.64}/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")
 

{litdata-0.2.61 → litdata-0.2.64}/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.61 → litdata-0.2.64}/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"
 

{litdata-0.2.61 → litdata-0.2.64}/src/litdata/streaming/config.py

@@ -302,6 +302,11 @@ class ChunksConfig:
 
         return local_chunkpath, begin, filesize_bytes
 
+    def download_filepath(self, chunk_index: int) -> str:
+        """The raw on-disk path that the chunk is downloaded to before any decompression."""
+        assert self._chunks is not None
+        return os.path.join(self._cache_dir, self._chunks[chunk_index]["filename"])
+
     def _get_chunk_index_from_filename(self, chunk_filename: str) -> int:
         """Retrieves the associated chunk_index for a given chunk filename."""
         assert self._chunks is not None

{litdata-0.2.61 → litdata-0.2.64}/src/litdata/streaming/dataloader.py

@@ -519,6 +519,43 @@ class StreamingDataLoaderCollateFn:
         return self.collate_fn(items)
 
 
+def _is_fork_context(multiprocessing_context: Any) -> bool:
+    import multiprocessing as mp
+
+    if multiprocessing_context is None:
+        try:
+            return mp.get_start_method(allow_none=False) == "fork"
+        except Exception:
+            return False
+
+    if isinstance(multiprocessing_context, str):
+        return multiprocessing_context == "fork"
+
+    if hasattr(multiprocessing_context, "get_start_method"):
+        try:
+            return multiprocessing_context.get_start_method() == "fork"
+        except Exception:
+            return False
+
+    if hasattr(multiprocessing_context, "start_method"):
+        return multiprocessing_context.start_method == "fork"
+
+    return False
+
+
+def _has_parquet_loader(dataset: Any) -> bool:
+    from litdata.streaming.dataset import StreamingDataset
+    from litdata.streaming.item_loader import ParquetLoader
+
+    if isinstance(dataset, StreamingDataset):
+        return isinstance(dataset.item_loader, ParquetLoader)
+
+    if hasattr(dataset, "_datasets") and isinstance(dataset._datasets, (list, tuple)):
+        return any(_has_parquet_loader(d) for d in dataset._datasets)
+
+    return False
+
+
 class StreamingDataLoader(DataLoader):
     r"""The StreamingDataLoader combines a dataset and a sampler, and provides an iterable over the given dataset.
 
@@ -589,6 +626,14 @@ class StreamingDataLoader(DataLoader):
         collate_fn: Callable | None = None,
         **kwargs: Any,
     ) -> None:  # pyright: ignore
+        if num_workers > 0 and _is_fork_context(kwargs.get("multiprocessing_context")) and _has_parquet_loader(dataset):
+            raise RuntimeError(
+                "The `ParquetLoader` uses Polars, which is not compatible with the `fork` multiprocessing context "
+                "used by PyTorch's DataLoader on Linux. Using `fork` will cause deadlocks due to Polars' "
+                "internal thread pool. Please pass `multiprocessing_context='spawn'` (or 'forkserver') to "
+                "`StreamingDataLoader`. Check thread: https://github.com/Lightning-AI/litData/issues/823"
+            )
+
         if not isinstance(dataset, (StreamingDataset, _BaseStreamingDatasetWrapper)):
             raise RuntimeError(
                 "The provided dataset should be either an instance of StreamingDataset, CombinedStreamingDataset or "

{litdata-0.2.61 → litdata-0.2.64}/src/litdata/streaming/downloader.py

@@ -268,7 +268,7 @@ class R2Downloader(Downloader):
         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

{litdata-0.2.61 → litdata-0.2.64}/src/litdata/streaming/fs_provider.py

@@ -20,7 +20,9 @@ from litdata.streaming.client import R2Client, S3Client
 
 
 class FsProvider(ABC):
-    def __init__(self, storage_options: dict[str, Any] | None = {}):
+    def __init__(self, storage_options: dict[str, Any] | None = None):
+        if storage_options is None:
+            storage_options = {}
         self.storage_options = storage_options
 
     @abstractmethod
@@ -50,7 +52,7 @@ class FsProvider(ABC):
 
 
 class GCPFsProvider(FsProvider):
-    def __init__(self, storage_options: dict[str, Any] | None = {}):
+    def __init__(self, storage_options: dict[str, Any] | None = None):
         if not _GOOGLE_STORAGE_AVAILABLE:
             raise ModuleNotFoundError(str(_GOOGLE_STORAGE_AVAILABLE))
         from google.cloud import storage
@@ -133,9 +135,9 @@ class GCPFsProvider(FsProvider):
 
 
 class S3FsProvider(FsProvider):
-    def __init__(self, storage_options: dict[str, Any] | None = {}):
+    def __init__(self, storage_options: dict[str, Any] | None = None):
         super().__init__(storage_options=storage_options)
-        self.client = S3Client(storage_options=storage_options)
+        self.client = S3Client(storage_options=self.storage_options)
 
     def upload_file(self, local_path: str, remote_path: str) -> None:
         bucket_name, blob_path = get_bucket_and_path(remote_path, "s3")
@@ -225,11 +227,11 @@ class S3FsProvider(FsProvider):
 
 
 class R2FsProvider(S3FsProvider):
-    def __init__(self, storage_options: dict[str, Any] | None = {}):
+    def __init__(self, storage_options: dict[str, Any] | None = None):
         super().__init__(storage_options=storage_options)
 
         # Create R2Client with refreshable credentials
-        self.client = R2Client(storage_options=storage_options)
+        self.client = R2Client(storage_options=self.storage_options)
 
     def upload_file(self, local_path: str, remote_path: str) -> None:
         bucket_name, blob_path = get_bucket_and_path(remote_path, "r2")
@@ -324,7 +326,7 @@ def get_bucket_and_path(remote_filepath: str, expected_scheme: str = "s3") -> tu
     return bucket_name, blob_path
 
 
-def _get_fs_provider(remote_filepath: str, storage_options: dict[str, Any] | None = {}) -> FsProvider:
+def _get_fs_provider(remote_filepath: str, storage_options: dict[str, Any] | None = None) -> FsProvider:
     obj = parse.urlparse(remote_filepath)
     if obj.scheme == "gs":
         return GCPFsProvider(storage_options=storage_options)

{litdata-0.2.61 → litdata-0.2.64}/src/litdata/streaming/item_loader.py

@@ -642,6 +642,7 @@ class ParquetLoader(BaseItemLoader):
         self._df: dict[int, Any] = {}
         self._chunk_row_groups: dict[int, Any] = {}
         self._chunk_row_group_item_read_count: dict[int, Any] = {}
+        self._chunk_row_group_offsets: dict[int, list[int]] = {}
 
     def generate_intervals(self) -> list[Interval]:
         intervals = []
@@ -712,18 +713,28 @@ class ParquetLoader(BaseItemLoader):
         Returns:
             Any: The dataframe row corresponding to the specified index.
         """
+        import bisect
+
         import polars as pl
         import pyarrow.parquet as pq
 
         # Load the Parquet file metadata if not already loaded
         if chunk_index not in self._df:
-            self._df[chunk_index] = pq.ParquetFile(chunk_filepath)
-
-        # Determine the row group and the row index within the row group
-        parquet_file = self._df[chunk_index]
-        num_rows_per_row_group = parquet_file.metadata.row_group(0).num_rows
-        row_group_index = row_index // num_rows_per_row_group
-        row_index_within_group = row_index % num_rows_per_row_group
+            parquet_file = pq.ParquetFile(chunk_filepath)
+            self._df[chunk_index] = parquet_file
+            # Precompute cumulative row offsets as a prefix-sum so lookup works for row groups of any size.
+            offsets = [0]
+            num_row_groups = parquet_file.metadata.num_row_groups
+            for i in range(num_row_groups):
+                num_rows = parquet_file.metadata.row_group(i).num_rows
+                offsets.append(offsets[-1] + num_rows)
+            self._chunk_row_group_offsets[chunk_index] = offsets
+
+        # Locate the row group containing row_index and the offset inside it.
+        offsets = self._chunk_row_group_offsets[chunk_index]
+        row_group_index = bisect.bisect_right(offsets, row_index) - 1
+        row_index_within_group = row_index - offsets[row_group_index]
+        row_group_size = offsets[row_group_index + 1] - offsets[row_group_index]
 
         # Check if the row group is already loaded
         if chunk_index in self._chunk_row_groups and row_group_index in self._chunk_row_groups[chunk_index]:
@@ -746,7 +757,7 @@ class ParquetLoader(BaseItemLoader):
 
         # Check if the row group has been fully read and release memory if necessary
         read_count = self._chunk_row_group_item_read_count[chunk_index][row_group_index]
-        if read_count >= num_rows_per_row_group:
+        if read_count >= row_group_size:
             # Release memory for the fully read row group
             del self._chunk_row_groups[chunk_index][row_group_index]
             del self._chunk_row_group_item_read_count[chunk_index][row_group_index]
@@ -797,6 +808,8 @@ class ParquetLoader(BaseItemLoader):
 
         if chunk_index in self._chunk_row_group_item_read_count:
             del self._chunk_row_group_item_read_count[chunk_index]
+        if chunk_index in self._chunk_row_group_offsets:
+            del self._chunk_row_group_offsets[chunk_index]
         if os.path.exists(chunk_filepath):
             os.remove(chunk_filepath)
         logger.debug(
@@ -820,5 +833,8 @@ class ParquetLoader(BaseItemLoader):
         if chunk_index in self._chunk_row_group_item_read_count:
             del self._chunk_row_group_item_read_count[chunk_index]
 
+        if chunk_index in self._chunk_row_group_offsets:
+            del self._chunk_row_group_offsets[chunk_index]
+
     def encode_data(self, data: list[bytes], sizes: list[int], flattened: list[Any]) -> Any:
         pass

{litdata-0.2.61 → litdata-0.2.64}/src/litdata/streaming/reader.py

@@ -239,25 +239,39 @@ class PrepareChunksThread(Thread):
 
     def _force_download(self) -> None:
         chunk_index = _get_from_queue(self._force_download_queue)
-        if chunk_index is not None:
-            # force apply deletion before redownload
-            self._apply_delete(chunk_index, skip_lock=True)
-            if _DEBUG:
-                chunk_filepath, _, _ = self._config[ChunkedIndex(index=-1, chunk_index=chunk_index)]
-                print(
-                    f"[Reader] Requested force download for {chunk_filepath} "
-                    f"by {self._rank} at {datetime.now().isoformat()}"
-                )
+        if chunk_index is None:
+            return
+
+        chunk_filepath, _, filesize_bytes = self._config[ChunkedIndex(index=-1, chunk_index=chunk_index)]
+        download_lock_path = self._config.download_filepath(chunk_index) + ".lock"
+        try:
+            with FileLock(download_lock_path, timeout=0):
+                # The chunk may have been fully downloaded by the time this
+                # request was processed, so double check that it still needs
+                # downloading before we delete it.
+                if os.path.exists(chunk_filepath) and os.stat(chunk_filepath).st_size >= filesize_bytes:
+                    return
+
+                # force apply deletion before redownload
+                self._apply_delete(chunk_index, skip_lock=True)
+                if _DEBUG:
+                    print(
+                        f"[Reader] Requested force download for {chunk_filepath} "
+                        f"by {self._rank} at {datetime.now().isoformat()}"
+                    )
 
             self._config.download_chunk_from_index(chunk_index, skip_lock=True)
+        except Timeout:
+            # Another worker is actively downloading this chunk. Defer to them.
+            return
 
-            # Preload item if possible to gain some time but only
-            # if this is one of the pre-downloaded chunk
-            if self._pre_download_counter > 0:
-                self._pre_load_chunk(chunk_index)
+        # Preload item if possible to gain some time but only
+        # if this is one of the pre-downloaded chunk
+        if self._pre_download_counter > 0:
+            self._pre_load_chunk(chunk_index)
 
-            # Avoid downloading too many chunks in advance at the risk of over using the disk space
-            self._pre_download_counter += 1
+        # Avoid downloading too many chunks in advance at the risk of over using the disk space
+        self._pre_download_counter += 1
 
     def run(self) -> None:
         while True:

{litdata-0.2.61 → litdata-0.2.64}/src/litdata/streaming/resolver.py

@@ -217,7 +217,19 @@ def _resolve_data_connection(dir_path: str) -> "V1DataConnection":
 def _resolve_s3_connections(dir_path: str) -> Dir:
     data_connection = _resolve_data_connection(dir_path)
 
-    return Dir(path=dir_path, url=os.path.join(data_connection.aws.source, *dir_path.split("/")[4:]))
+    # S3 connections flagged available on non-AWS providers can't rely on the FUSE mount or an
+    # instance profile off AWS, so thread the connection id through. That makes the S3 client mint
+    # temporary project-role credentials (same mechanism as R2 / lightning_storage below). Plain
+    # AWS-only S3 connections leave this unset and keep using the ambient credentials.
+    data_connection_id = (
+        data_connection.id if getattr(data_connection.aws, "available_in_non_aws_providers", False) else None
+    )
+
+    return Dir(
+        path=dir_path,
+        url=os.path.join(data_connection.aws.source, *dir_path.split("/")[4:]),
+        data_connection_id=data_connection_id,
+    )
 
 
 def _resolve_gcs_connections(dir_path: str) -> Dir:

{litdata-0.2.61 → litdata-0.2.64}/src/litdata/streaming/serializers.py

@@ -20,6 +20,7 @@ from abc import ABC, abstractmethod
 from collections import OrderedDict
 from contextlib import suppress
 from copy import deepcopy
+from dataclasses import asdict
 from itertools import chain
 from typing import Any
 
@@ -32,6 +33,7 @@ from litdata.constants import (
     _NUMPY_DTYPES_MAPPING,
     _PIL_AVAILABLE,
     _TORCH_DTYPES_MAPPING,
+    _TORCH_VISION_LESS_THAN_0_26,
 )
 
 
@@ -403,6 +405,13 @@ class VideoSerializer(Serializer):
             return f.read(), f"video:{file_extension}"
 
     def deserialize(self, data: bytes) -> Any:
+        # if using torchvision <=0.25, we will use torchvision.io to decode the video
+        # otherwise, we will use torchcodec to decode the video, which is faster and more robust
+        if _TORCH_VISION_LESS_THAN_0_26:
+            return self._deserialize_with_torchvision_io(data)
+        return self._deserialize_with_torchcodec(data)
+
+    def _deserialize_with_torchvision_io(self, data: bytes) -> Any:
         if not _AV_AVAILABLE:
             raise ModuleNotFoundError("av is required. Run `pip install av`")
 
@@ -416,6 +425,29 @@ class VideoSerializer(Serializer):
                 stream.write(data)
             return torchvision.io.read_video(fname, pts_unit="sec")
 
+    def _deserialize_with_torchcodec(self, data: bytes) -> Any:
+        try:
+            import torch
+            from torchcodec.decoders import AudioDecoder, VideoDecoder
+        except ImportError:
+            raise ModuleNotFoundError("torchcodec is required. Run `pip install torchcodec>0.11`")
+
+        dec = VideoDecoder(data, dimension_order="NHWC")  # NHWC → T,H,W,C after stacking
+        metadata = asdict(dec.metadata) if dec.metadata is not None else {}
+
+        # get_all_frames() returns a FrameBatch; .data is (N, C, H, W) or (N, H, W, C)
+        # depending on dimension_order above
+        frame_batch = dec.get_all_frames()
+        video = frame_batch.data  # shape: (T, H, W, C) with NHWC
+
+        try:
+            audio_dec = AudioDecoder(data)
+            audio = audio_dec.get_all_samples().data  # (num_channels, num_samples)
+        except ValueError:
+            audio = torch.zeros(1, 0)  # old torchvision path returns aframes with shape (1, 0) for no-audio videos.
+
+        return video, audio, metadata
+
     def can_serialize(self, data: Any) -> bool:
         return isinstance(data, str) and os.path.isfile(data) and any(data.endswith(ext) for ext in self._EXTENSIONS)
 

{litdata-0.2.61 → litdata-0.2.64}/src/litdata/streaming/writer.py

@@ -13,6 +13,7 @@
 
 import json
 import os
+import re
 import uuid
 import warnings
 from dataclasses import dataclass
@@ -499,7 +500,10 @@ class BinaryWriter:
             chunks_info.extend(existing_index["chunks"])
             config = existing_index["config"]
 
-        for index_filename in sorted(index_files):
+        def _natural_key(s: str) -> list:
+            return [int(t) if t.isdigit() else t for t in re.split(r"(\d+)", s)]
+
+        for index_filename in sorted(index_files, key=_natural_key):
             chunk_path = os.path.join(self._cache_dir, index_filename)
             with open(chunk_path) as f:
                 data = json.load(f)

{litdata-0.2.61 → litdata-0.2.64}/src/litdata/utilities/parquet.py

@@ -21,9 +21,12 @@ class ParquetDir(ABC):
         self,
         dir_path: str | Dir | None,
         cache_path: str | None = None,
-        storage_options: dict | None = {},
+        storage_options: dict | None = None,
         num_workers: int = 4,
     ):
+        if storage_options is None:
+            storage_options = {}
+
         self.dir = _resolve_dir(dir_path)
         self.cache_path = cache_path
         self.storage_options = storage_options
@@ -147,7 +150,7 @@ class CloudParquetDir(ParquetDir):
         for provider in _CLOUD_PROVIDER:
             if self.dir.url.startswith(provider):
                 # Initialize the cloud filesystem
-                self.fs = fsspec.filesystem(provider, *self.storage_options)
+                self.fs = fsspec.filesystem(provider, **self.storage_options)
                 print(f"using provider: {provider}")
                 break
 
@@ -323,7 +326,12 @@ def get_parquet_indexer_cls(
 
     obj = parse.urlparse(dir_path)
 
-    if obj.scheme in ("local", ""):
+    # On Windows, paths like "C:\\Users\\..." parse with scheme='c'. A single-letter
+    # scheme is never a valid URI scheme (RFC 3986 requires >=2 chars) so treat it
+    # as a Windows drive letter and dispatch to LocalParquetDir.
+    is_windows_drive = len(obj.scheme) == 1 and obj.scheme.isalpha()
+
+    if obj.scheme in ("local", "") or is_windows_drive:
         return LocalParquetDir(*args)
     if obj.scheme in ("gs", "s3"):
         return CloudParquetDir(*args)

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: litdata
-Version: 0.2.61
+Version: 0.2.64
 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.64}/src/litdata.egg-info/SOURCES.txt

@@ -22,6 +22,9 @@ src/litdata.egg-info/top_level.txt
 src/litdata/cli/__init__.py
 src/litdata/cli/commands.py
 src/litdata/cli/parser.py
+src/litdata/cli/handler/__init__.py
+src/litdata/cli/handler/cache.py
+src/litdata/cli/handler/optimize.py
 src/litdata/processing/__init__.py
 src/litdata/processing/data_processor.py
 src/litdata/processing/functions.py

{litdata-0.2.61 → litdata-0.2.64}/src/litdata.egg-info/requires.txt

@@ -11,9 +11,9 @@ obstore
 [extras]
 fsspec
 google-cloud-storage
-lightning-sdk==2025.12.17
+lightning-sdk==2026.2.6
 pillow
 polars
-pyarrow
+pyarrow<25.0.0
 tqdm
 viztracer