Lightning Dataset (including optimized dataloading of s3 buckets)

.actions/assistant.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import glob
+import logging
 import os
 import pathlib
 import re
@@ -43,6 +44,11 @@
         "requirements/fabric/base.txt",
         "requirements/fabric/strategies.txt",
     ),
+    "data": (
+        "requirements/data/data.txt",
+        "requirements/data/cloud.txt",
+        "requirements/data/examples.txt",
+    ),
 }
 REQUIREMENT_FILES_ALL = list(chain(*REQUIREMENT_FILES.values()))
 
@@ -146,6 +152,9 @@ def load_requirements(path_dir: str, file_name: str = "base.txt", unfreeze: str
     """
     assert unfreeze in {"none", "major", "all"}
     path = Path(path_dir) / file_name
+    if not path.exists():
+        logging.warning(f"Folder {path_dir} does not have any base requirements.")
+        return []
     assert path.exists(), (path_dir, file_name, path)
     text = path.read_text()
     return [req.adjust(unfreeze) for req in _parse_requirements(text)]
@@ -240,7 +249,7 @@ def _load_aggregate_requirements(req_dir: str = "requirements", freeze_requireme
     requires = [
         load_requirements(d, unfreeze="none" if freeze_requirements else "major")
         for d in glob.glob(os.path.join(req_dir, "*"))
-        # skip empty folder as git artefacts, and resolving Will's special issue
+        # skip empty folder (git artifacts), and resolving Will's special issue
         if os.path.isdir(d) and len(glob.glob(os.path.join(d, "*"))) > 0 and not os.path.basename(d).startswith("_")
     ]
     if not requires:
@@ -404,6 +413,7 @@ def _replace_min(fname: str) -> None:
     def replace_oldest_ver(requirement_fnames: Sequence[str] = REQUIREMENT_FILES_ALL) -> None:
         """Replace the min package version by fixed one."""
         for fname in requirement_fnames:
+            print(fname)
             AssistantCLI._replace_min(fname)
 
     @staticmethod

.github/CODEOWNERS

@@ -41,6 +41,11 @@
 /src/lightning/pytorch/core/hooks.py         @williamfalcon @tchaton @awaelchli @carmocca
 /src/lightning/pytorch/core/module.py        @williamfalcon @tchaton @awaelchli @carmocca
 
+# Data Utilities
+/examples/data/           @nohalon @justusschock
+/src/lightning/data/      @nohalon @justusschock
+/tests/tests_data         @nohalon @justusschock
+
 # Lightning Fabric
 /src/lightning/fabric                       @awaelchli @carmocca @justusschock
 /src/lightning_fabric                       @awaelchli @carmocca @justusschock

.github/checkgroup.yml

@@ -150,6 +150,26 @@ subprojects:
       - "build-pl (3.9, 1.13, 11.7.1)"
       - "build-pl (3.10, 2.0, 11.7.1)"
 
+  # SECTIONS: lightning_data
+
+  - id: "lightning_data: CPU workflow"
+    paths:
+      - ".actions/**"
+      - "requirements/data/**"
+      - "src/lightning/data/**"
+      - "src/lightning_data/*"
+      - "tests/tests_data/**"
+      - "examples/data/**"
+      - "pyproject.toml"  # includes pytest config
+      - ".github/workflows/ci-tests-data.yml"
+      - "!requirements/*/docs.txt"
+      - "!*.md"
+      - "!**/*.md"
+    checks:
+      - "data-cpu (macOS-11, lightning, 3.10, 2.0)"
+      - "data-cpu (ubuntu-20.04, lightning, 3.10, 2.0)"
+      - "data-cpu (windows-2022, lightning, 3.10, 2.0)"
+
   # SECTION: lightning_fabric
 
   - id: "lightning_fabric: CPU workflow"

.github/workflows/ci-tests-data.yml

@@ -0,0 +1,118 @@
+name: Test Data
+
+# see: https://help.github.com/en/actions/reference/events-that-trigger-workflows
+on:
+  push:
+    branches: [master, "release/*"]
+  pull_request:
+    branches: [master, "release/*"]
+    types: [opened, reopened, ready_for_review, synchronize]  # added `ready_for_review` since draft is skipped
+    paths:
+      - ".actions/**"
+      - "requirements/data/**"
+      - "src/lightning/data/**"
+      - "tests/tests_data/**"
+      - "pyproject.toml"  # includes pytest config
+      - ".github/workflows/ci-tests-data.yml"
+      - "!requirements/*/docs.txt"
+      - "!*.md"
+      - "!**/*.md"
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.ref }}-${{ github.head_ref }}
+  cancel-in-progress: ${{ ! (github.ref == 'refs/heads/master' || startsWith(github.ref, 'refs/heads/release/')) }}
+
+defaults:
+  run:
+    shell: bash
+
+jobs:
+  data-cpu:
+    runs-on: ${{ matrix.os }}
+    if: github.event.pull_request.draft == false
+    strategy:
+      fail-fast: false
+      matrix:
+        include:
+          - {os: "macOS-11", pkg-name: "lightning", python-version: "3.10", pytorch-version: "2.0"}
+          - {os: "ubuntu-20.04", pkg-name: "lightning", python-version: "3.10", pytorch-version: "2.0"}
+          - {os: "windows-2022", pkg-name: "lightning", python-version: "3.10", pytorch-version: "2.0"}
+          # "oldest" versions tests, only on minimum Python
+          # - {os: "macOS-11", pkg-name: "lightning", python-version: "3.8", pytorch-version: "2.0", requires: "oldest"}
+          # - {os: "ubuntu-20.04", pkg-name: "lightning", python-version: "3.8", pytorch-version: "2.0", requires: "oldest"}
+          # - {os: "windows-2022", pkg-name: "lightning", python-version: "3.8", pytorch-version: "2.0", requires: "oldest"}
+    timeout-minutes: 25  # because of building grpcio on Mac
+    env:
+      PACKAGE_NAME: ${{ matrix.pkg-name }}
+      FREEZE_REQUIREMENTS: ${{ ! (github.ref == 'refs/heads/master' || startsWith(github.ref, 'refs/heads/release/')) }}
+      # PYPI_CACHE_DIR: "_pip-wheels"
+      TORCH_URL_STABLE: "https://download.pytorch.org/whl/cpu/torch_stable.html"
+      TORCH_URL_TEST: "https://download.pytorch.org/whl/test/cpu/torch_test.html"
+    steps:
+    - uses: actions/checkout@v3
+
+    - name: Set up Python ${{ matrix.python-version }}
+      uses: actions/setup-python@v4
+      with:
+        python-version: ${{ matrix.python-version }}
+
+    - name: basic setup
+      run: pip install -q -r .actions/requirements.txt
+
+    - name: Set min. dependencies
+      if: ${{ matrix.requires == 'oldest' }}
+      run: |
+        python .actions/assistant.py replace_oldest_ver
+
+    - name: Adjust PyTorch versions in requirements files
+      if: ${{ matrix.requires != 'oldest' && matrix.release != 'pre' }}
+      run: |
+        pip install -q wget packaging
+        python -m wget https://raw.githubusercontent.com/Lightning-AI/utilities/main/scripts/adjust-torch-versions.py
+        for fpath in `ls requirements/data/*.txt`; do \
+          python ./adjust-torch-versions.py $fpath ${{ matrix.pytorch-version }}; \
+        done
+        cat requirements/data/data.txt
+        cat requirements/data/cloud.txt
+
+    # - name: pip wheels cache
+    #   uses: actions/cache/restore@v3
+    #   with:
+    #     path: ${{ env.PYPI_CACHE_DIR }}
+    #     key: pypi_wheels
+    # - run: |
+    #     mkdir -p $PYPI_CACHE_DIR
+    #     ls -lh $PYPI_CACHE_DIR
+
+    # removing torch stable line:
+    # pip install -e ".[${extra}test]" "pytest-timeout" -U -f ${TORCH_URL} ${TORCH_PREINSTALL} -f ${PYPI_CACHE_DIR} --prefer-binary
+    - name: Install package & dependencies
+      run: |
+        python -m pip install -q pip -U
+        pip install -e ".[data-dev]" "pytest-timeout" -U -f ${TORCH_URL} --prefer-binary
+        pip list
+
+    - name: Testing Data
+      working-directory: tests/tests_data
+      # NOTE: do not include coverage report here, see: https://github.com/nedbat/coveragepy/issues/1003
+      run: |
+        python -m coverage run --source lightning \
+          -m pytest -v --timeout=30 --durations=50
+
+    - name: Statistics
+      if: success()
+      working-directory: tests/tests_data
+      run: |
+        coverage report
+        coverage xml
+
+    - name: Upload coverage to Codecov
+      uses: codecov/codecov-action@v3
+      # see: https://github.com/actions/toolkit/issues/399
+      continue-on-error: true
+      with:
+        token: ${{ secrets.CODECOV_TOKEN }}
+        file: tests/tests_data/coverage.xml
+        flags: lightning,cpu,pytest,python${{ matrix.python-version }}
+        name: CPU-coverage
+        fail_ci_if_error: false

.gitignore

@@ -189,6 +189,10 @@ our_model.tar
 test.png
 saved_models
 data/
+!src/lightning/data/
+!examples/data/
+!tests/tests_pytorch/utilities/data/
+!requirements/data/
 .shared
 .lightning
 node_modules/