Merge branch 'main' into fix-parallel-streaming-dset-resume-with-load-state

Author: bhimrazy

.github/benchmark/benchmark.py

@@ -2,7 +2,6 @@
 
 import argparse
 from dataclasses import dataclass
-from typing import Optional
 
 from lightning_sdk import Machine, Studio
 
@@ -19,8 +18,8 @@ class BenchmarkArgs:
 
     pr_number: int
     branch: str
-    org: Optional[str]
-    user: Optional[str]
+    org: str | None
+    user: str | None
     teamspace: str
     machine: Machine
     make_args: str
@@ -75,7 +74,7 @@ def __init__(self, config: BenchmarkArgs):
         self.org = config.org
         self.machine = config.machine
         self.make_args = config.make_args
-        self.studio: Optional[Studio] = None
+        self.studio: Studio | None = None
 
     def run(self) -> None:
         """Run the LitData benchmark."""

.github/workflows/ci-benchmark.yml

@@ -23,7 +23,7 @@ jobs:
       )
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/checkout@v5
+      - uses: actions/checkout@v6
       - name: Extract make args
         id: extract
         uses: actions/github-script@v8

.github/workflows/ci-testing.yml

@@ -30,7 +30,7 @@ jobs:
       UV_TORCH_BACKEND: "cpu"
 
     steps:
-      - uses: actions/checkout@v5
+      - uses: actions/checkout@v6
       - name: Install uv and setup python ${{ matrix.python-version }}
         uses: astral-sh/setup-uv@v7
         with:

.github/workflows/cleanup-caches.yml

@@ -9,7 +9,7 @@ jobs:
     runs-on: ubuntu-latest
     steps:
       - name: Check out code
-        uses: actions/checkout@v5
+        uses: actions/checkout@v6
 
       - name: Cleanup
         run: |

.github/workflows/docs-build.yml

@@ -19,15 +19,15 @@ jobs:
     if: github.event_name != 'pull_request'
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/checkout@v5 # deploy needs git credentials
+      - uses: actions/checkout@v6 # deploy needs git credentials
       - name: Download prepared docs
-        uses: actions/download-artifact@v6
+        uses: actions/download-artifact@v7
         with:
           name: docs-html-${{ github.sha }}
           path: docs/build/html
 
       - name: Deploy 🚀
-        uses: JamesIves/github-pages-deploy-action@v4.7.4
+        uses: JamesIves/github-pages-deploy-action@v4.7.6
         if: ${{ github.event_name == 'push' }}
         with:
           token: ${{ secrets.GITHUB_TOKEN }}

.github/workflows/release-pypi.yml

@@ -18,7 +18,7 @@ jobs:
     runs-on: ubuntu-22.04
 
     steps:
-      - uses: actions/checkout@v5
+      - uses: actions/checkout@v6
       - uses: actions/setup-python@v6
         with:
           python-version: "3.10"

.pre-commit-config.yaml

@@ -55,7 +55,7 @@ repos:
         #args: ["--write-changes"] # uncomment if you want to get automatic fixing
 
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.13.3
+    rev: v0.14.10
     hooks:
       - id: ruff
         args: ["--fix"]
@@ -85,7 +85,7 @@ repos:
         args: ["--print-width=120"]
 
   - repo: https://github.com/tox-dev/pyproject-fmt
-    rev: v2.7.0
+    rev: v2.11.1
     hooks:
       - id: pyproject-fmt
         additional_dependencies: [tox]

benchmarks/litdata/optimize_imagenet.py

@@ -10,7 +10,6 @@
 import os
 import time
 from functools import lru_cache, partial
-from typing import Union
 
 import numpy as np
 import requests
@@ -126,7 +125,7 @@ def main():
     seed_everything(args.seed)
 
     # Handle resize_size: if two ints are given, treat as tuple, else int or None
-    resize_size: Union[int, tuple[int, int], None] = None
+    resize_size: int | tuple[int, int] | None = None
     if args.resize_size is not None:
         if isinstance(args.resize_size, list):
             if len(args.resize_size) == 1:

docs/source/conf.py

@@ -175,6 +175,10 @@ def _convert_markdown(path_in: str, path_out: str) -> None:
     "logo_only": False,
 }
 
+html_context = {
+    "style": "pytorch",
+}
+
 html_favicon = "_static/images/icon.svg"
 
 # Add any paths that contain custom static files (such as style sheets) here,

examples/multi_modal/dataloader.py

@@ -2,7 +2,7 @@
 
 import logging
 import os
-from typing import Any, Union
+from typing import Any
 
 import joblib
 import lightning as pl
@@ -44,7 +44,7 @@ def load_tokenizer(self):
 class DocumentClassificationDataset(StreamingDataset):
     """Streaming dataset class."""
 
-    def __init__(self, input_dir: Union[str, Any], hyperparameters: Union[dict, Any] = None) -> None:
+    def __init__(self, input_dir: str | Any, hyperparameters: dict | Any = None) -> None:
         super().__init__(input_dir, shuffle=True, max_cache_size=hyperparameters["max_cache_size"])
         self.hyperparameters = hyperparameters
         self.image_transform = transforms.Compose(