unify extras & minor CI cleaning: move env. var

.actions/assistant.py

@@ -27,6 +27,11 @@
 
 from pkg_resources import parse_requirements
 
+try:
+    from typing import Literal
+except ImportError:
+    from typing_extensions import Literal
+
 REQUIREMENT_FILES = {
     "pytorch": (
         "requirements/pytorch/base.txt",
@@ -49,7 +54,7 @@
 _PROJECT_ROOT = os.path.dirname(os.path.dirname(__file__))
 
 
-def _augment_requirement(ln: str, comment_char: str = "#", unfreeze: str = "all") -> str:
+def _augment_requirement(ln: str, comment_char: str = "#", unfreeze: Literal["none", "major", "all"] = "all") -> str:
     """Adjust the upper version contrains.
 
     Args:
@@ -60,9 +65,9 @@ def _augment_requirement(ln: str, comment_char: str = "#", unfreeze: str = "all"
     Returns:
         adjusted requirement
 
-    >>> _augment_requirement("arrow<=1.2.2,>=1.2.0  # anything", unfreeze="")
+    >>> _augment_requirement("arrow<=1.2.2,>=1.2.0  # anything", unfreeze="none")
     'arrow<=1.2.2,>=1.2.0'
-    >>> _augment_requirement("arrow<=1.2.2,>=1.2.0  # strict", unfreeze="")
+    >>> _augment_requirement("arrow<=1.2.2,>=1.2.0  # strict", unfreeze="none")
     'arrow<=1.2.2,>=1.2.0  # strict'
     >>> _augment_requirement("arrow<=1.2.2,>=1.2.0  # my name", unfreeze="all")
     'arrow>=1.2.0'
@@ -99,7 +104,7 @@ def _augment_requirement(ln: str, comment_char: str = "#", unfreeze: str = "all"
         ver_major = None
 
     # remove version restrictions unless they are strict
-    if unfreeze and "<" in req and not is_strict:
+    if unfreeze != "none" and "<" in req and not is_strict:
         req = re.sub(r",? *<=? *[\d\.\*]+,? *", "", req).strip()
     if ver_major is not None and not is_strict:
         # add , only if there are already some versions
@@ -113,7 +118,10 @@ def _augment_requirement(ln: str, comment_char: str = "#", unfreeze: str = "all"
 
 
 def load_requirements(
-    path_dir: str, file_name: str = "base.txt", comment_char: str = "#", unfreeze: str = "all"
+    path_dir: str,
+    file_name: str = "base.txt",
+    comment_char: str = "#",
+    unfreeze: Literal["none", "major", "all"] = "all",
 ) -> List[str]:
     """Loading requirements from a file.
 

.actions/requirements.txt

@@ -1,2 +1,3 @@
 jsonargparse>=4.16.0
+typing-extensions
 requests

.github/workflows/ci-examples-app.yml

@@ -72,10 +72,13 @@ jobs:
           key: ${{ runner.os }}-pip-py${{ matrix.python-version }}-${{ matrix.pkg-name }}-${{ matrix.requires }}-${{ hashFiles('requirements/app/base.txt') }}
           restore-keys: ${{ runner.os }}-pip-py${{ matrix.python-version }}-${{ matrix.pkg-name }}-${{ matrix.requires }}-
 
-      - name: Install dependencies
+      - name: Install Lightning package & dependencies
+        env:
+          PACKAGE_NAME: ${{ matrix.pkg-name }}
+          FREEZE_REQUIREMENTS: 1
         run: |
-          pip --version
-          pip install -r requirements/app/devel.txt --find-links https://download.pytorch.org/whl/cpu/torch_stable.html
+          # do not use `-e` because it will make both packages available since it adds `src` to `sys.path` automatically
+          pip install .[dev] --find-links https://download.pytorch.org/whl/cpu/torch_stable.html
           pip list
 
       - name: Setup Node.js
@@ -86,12 +89,6 @@ jobs:
       - name: Install Yarn
         run: npm install -g yarn
 
-      - name: Install Lightning package
-        env:
-          PACKAGE_NAME: ${{ matrix.pkg-name }}
-        # do not use -e because it will make both packages available since it adds `src` to `sys.path` automatically
-        run: pip install .
-
       - name: Adjust tests
         if: ${{ matrix.pkg-name == 'lightning' }}
         run: |
@@ -116,7 +113,10 @@ jobs:
           AWS_DEFAULT_REGION: us-east-1
           PYTEST_ARTIFACT: results-${{ runner.os }}-${{ matrix.python-version }}-${{ matrix.requires }}.xml
         run: |
-          coverage run --source ${COVERAGE_SCOPE} -m pytest -m "not cloud" tests_examples_app --timeout=300 -vvvv --junitxml=$PYTEST_ARTIFACT --durations=0
+          python -m coverage run --source ${COVERAGE_SCOPE} \
+            -m pytest -m "not cloud" tests_examples_app \
+            --timeout=300 --durations=0 -vvvv \
+            --junitxml=$PYTEST_ARTIFACT
 
       - name: Upload pytest test results
         uses: actions/upload-artifact@v3

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

@@ -22,9 +22,6 @@ concurrency:
   group: ${{ github.workflow }}-${{ github.ref }}-${{ github.head_ref }}
   cancel-in-progress: ${{ github.ref != 'refs/heads/master' }}
 
-env:
-  FREEZE_REQUIREMENTS: 1
-
 defaults:
   run:
     shell: bash
@@ -77,11 +74,12 @@ jobs:
       - name: Switch PyTorch URL
         run: python -c "print('TORCH_URL=https://download.pytorch.org/whl/' + str('test/cpu/torch_test.html' if '${{matrix.release}}' == 'pre' else 'cpu/torch_stable.html'))" >> $GITHUB_ENV
 
-      - name: Install package & depenencies
+      - name: Install package & dependencies
         env:
           PACKAGE_NAME: ${{ matrix.pkg-name }}
+          FREEZE_REQUIREMENTS: 1
         run: |
-          pip install -e . -r requirements/app/devel.txt -U -q --find-links ${TORCH_URL}
+          pip install -e .[dev] --upgrade --find-links ${TORCH_URL}
           pip list
 
       - name: Setup Node.js
@@ -117,7 +115,10 @@ jobs:
           AWS_DEFAULT_REGION: us-east-1
           PYTEST_ARTIFACT: results-${{ runner.os }}-${{ matrix.python-version }}-${{ matrix.requires }}.xml
         run: |
-          coverage run --source ${COVERAGE_SCOPE} -m pytest -m "not cloud" tests_app --timeout=300 -vvvv --junitxml=$PYTEST_ARTIFACT --durations=50
+          python -m coverage run --source ${COVERAGE_SCOPE} \
+            -m pytest -m "not cloud" tests_app \
+            --timeout=300 -vvvv --durations=50 \
+            --junitxml=$PYTEST_ARTIFACT
 
       - name: Upload pytest test results
         uses: actions/upload-artifact@v3

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

@@ -108,8 +108,9 @@ jobs:
     - name: Install package & dependencies
       env:
         PACKAGE_NAME: ${{ matrix.pkg-name }}
+        FREEZE_REQUIREMENTS: 1
       run: |
-        pip install -e . "pytest-timeout" -r requirements/lite/devel.txt --upgrade --find-links ${TORCH_URL}
+        pip install -e .[test] "pytest-timeout" --upgrade --find-links ${TORCH_URL}
         pip list
 
     - name: Adjust tests
@@ -131,7 +132,10 @@ jobs:
     - name: Testing Lite
       working-directory: tests/tests_lite
       # NOTE: do not include coverage report here, see: https://github.com/nedbat/coveragepy/issues/1003
-      run: coverage run --source ${COVERAGE_SCOPE} -m pytest -v --timeout=30 --durations=50 --junitxml=results-${{ runner.os }}-py${{ matrix.python-version }}-${{ matrix.requires }}-${{ matrix.release }}.xml
+      run: |
+        python -m coverage run --source ${COVERAGE_SCOPE} \
+          -m pytest -v --timeout=30 --durations=50 \
+          --junitxml=results-${{ runner.os }}-py${{ matrix.python-version }}-${{ matrix.requires }}-${{ matrix.release }}.xml
 
     - name: Upload pytest results
       if: failure()

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

@@ -135,8 +135,11 @@ jobs:
     - name: Install package & dependencies
       env:
         PACKAGE_NAME: ${{ matrix.pkg-name }}
+        FREEZE_REQUIREMENTS: 1
       run: |
-        pip install -e . "pytest-timeout" -r requirements/pytorch/devel.txt --upgrade --find-links ${TORCH_URL}
+        pip install -e .[extra,test] "pytest-timeout" --upgrade --find-links ${TORCH_URL}
+        # TODO: installing the strategies from file as deepspeed expects already installed PyTorch and extras do not install base first
+        pip install -r requirements/pytorch/strategies.txt --find-links ${TORCH_URL}
         pip list
 
     - name: Reinstall Horovod if necessary
@@ -184,7 +187,12 @@ jobs:
     - name: Testing PyTorch
       working-directory: tests/tests_pytorch
       # NOTE: do not include coverage report here, see: https://github.com/nedbat/coveragepy/issues/1003
-      run: coverage run --source ${COVERAGE_SCOPE} -m pytest -v --timeout=${TEST_TIMEOUT} --durations=50 --junitxml=results-${{ runner.os }}-py${{ matrix.python-version }}-${{ matrix.requires }}-${{ matrix.release }}.xml
+      run: |
+        python -m coverage run --source ${COVERAGE_SCOPE} \
+          -m pytest . -v \
+          --timeout=${TEST_TIMEOUT} --durations=50 \
+          --reruns 3 --reruns-delay 1 \
+          --junitxml=results-${{ runner.os }}-py${{ matrix.python-version }}-${{ matrix.requires }}-${{ matrix.release }}.xml
 
     - name: Upload pytest results
       if: failure()

.github/workflows/docs-checks.yml

@@ -23,9 +23,6 @@ concurrency:
   group: ${{ github.workflow }}-${{ github.ref }}-${{ github.head_ref }}
   cancel-in-progress: ${{ ! (github.ref == 'refs/heads/master' || startsWith(github.ref, 'refs/heads/release/')) }}
 
-env:
-  FREEZE_REQUIREMENTS: "1"
-
 defaults:
   run:
     shell: bash
@@ -78,6 +75,8 @@ jobs:
       - name: Install LAI package
         # This is needed as App docs is heavily using/referring to lightning package
         if: ${{ matrix.pkg-name == 'app' }}
+        env:
+          FREEZE_REQUIREMENTS: 1
         run: |
           pip install -e . -U -v -f https://download.pytorch.org/whl/cpu/torch_stable.html -f pypi
 
@@ -92,6 +91,7 @@ jobs:
       - name: Install this package
         env:
           PACKAGE_NAME: ${{ matrix.pkg-name }}
+          FREEZE_REQUIREMENTS: 1
         run: |
           pip install -e .[extra,cloud,ui] -U -r requirements/${{ matrix.pkg-name }}/docs.txt -f pypi
           pip list
@@ -110,8 +110,6 @@ jobs:
       fail-fast: false
       matrix:
         pkg-name: ["app", "pytorch"]
-    env:
-      FREEZE_REQUIREMENTS: "1"
     steps:
       - uses: actions/checkout@v3
         with:
@@ -145,6 +143,7 @@ jobs:
       - name: Install package & dependencies
         env:
           PACKAGE_NAME: ${{ matrix.pkg-name }}
+          FREEZE_REQUIREMENTS: 1
         run: |
           sudo apt-get update
           sudo apt-get install -y cmake pandoc texlive-latex-extra dvipng texlive-pictures

.github/workflows/docs-deploy.yml

@@ -8,9 +8,6 @@ on:
       - ".github/workflows/docs-deploy.yml"
       # TODO: this workflow is just for debugging. extend the paths that should trigger it
 
-env:
-  FREEZE_REQUIREMENTS: 1
-
 defaults:
   run:
     shell: bash
@@ -48,6 +45,8 @@ jobs:
             ${{ runner.os }}-deploy-docs-pip-
 
       - name: Install package & dependencies
+        env:
+          FREEZE_REQUIREMENTS: 1
         run: |
           sudo apt-get update
           sudo apt-get install -y cmake pandoc

src/lightning_app/__setup__.py

@@ -1,5 +1,7 @@
+import glob
 import os
 from importlib.util import module_from_spec, spec_from_file_location
+from pathlib import Path
 from types import ModuleType
 from typing import Any, Dict
 
@@ -32,16 +34,16 @@ def _prepare_extras() -> Dict[str, Any]:
     # Define package extras. These are only installed if you specify them.
     # From remote, use like `pip install pytorch-lightning[dev, docs]`
     # From local copy of repo, use like `pip install ".[dev, docs]"`
-    common_args = dict(path_dir=_PATH_REQUIREMENTS, unfreeze="major" if _FREEZE_REQUIREMENTS else "all")
+    req_files = [Path(p) for p in glob.glob(os.path.join(_PATH_REQUIREMENTS, "*.txt"))]
+    common_args = dict(path_dir=_PATH_REQUIREMENTS, unfreeze="none" if _FREEZE_REQUIREMENTS else "major")
     extras = {
-        # 'docs': load_requirements(file_name='docs.txt'),
-        "cloud": assistant.load_requirements(file_name="cloud.txt", **common_args),
-        "ui": assistant.load_requirements(file_name="ui.txt", **common_args),
-        "test": assistant.load_requirements(file_name="test.txt", **common_args),
+        p.stem: assistant.load_requirements(file_name=p.name, **common_args)
+        for p in req_files
+        if p.name not in ("docs.txt", "devel.txt", "base.txt")
     }
-    extras["extra"] = extras["cloud"] + extras["ui"]
-    extras["dev"] = extras["extra"] + extras["test"]  # + extras['docs']
-    extras["all"] = extras["dev"]
+    extras["extra"] = extras["cloud"] + extras["ui"] + extras["components"]
+    extras["all"] = extras["extra"]
+    extras["dev"] = extras["all"] + extras["test"]  # + extras['docs']
     return extras
 
 
@@ -80,7 +82,7 @@ def _setup_args() -> Dict[str, Any]:
         },
         setup_requires=["wheel"],
         install_requires=assistant.load_requirements(
-            _PATH_REQUIREMENTS, unfreeze="major" if _FREEZE_REQUIREMENTS else "all"
+            _PATH_REQUIREMENTS, unfreeze="none" if _FREEZE_REQUIREMENTS else "major"
         ),
         extras_require=_prepare_extras(),
         project_urls={

src/lightning_lite/__setup__.py

@@ -1,5 +1,7 @@
+import glob
 import os
 from importlib.util import module_from_spec, spec_from_file_location
+from pathlib import Path
 from types import ModuleType
 from typing import Any, Dict
 
@@ -33,16 +35,17 @@ def _prepare_extras() -> Dict[str, Any]:
     # Define package extras. These are only installed if you specify them.
     # From remote, use like `pip install pytorch-lightning[dev, docs]`
     # From local copy of repo, use like `pip install ".[dev, docs]"`
-    common_args = dict(path_dir=_PATH_REQUIREMENTS, unfreeze="" if _FREEZE_REQUIREMENTS else "all")
+    common_args = dict(path_dir=_PATH_REQUIREMENTS, unfreeze="none" if _FREEZE_REQUIREMENTS else "all")
+    req_files = [Path(p) for p in glob.glob(os.path.join(_PATH_REQUIREMENTS, "*.txt"))]
     extras = {
-        "examples": assistant.load_requirements(file_name="examples.txt", **common_args),
-        "strategies": assistant.load_requirements(file_name="strategies.txt", **common_args),
-        "test": assistant.load_requirements(file_name="test.txt", **common_args),
+        p.stem: assistant.load_requirements(file_name=p.name, **common_args)
+        for p in req_files
+        if p.name not in ("docs.txt", "devel.txt", "base.txt")
     }
     for req in parse_requirements(extras["strategies"]):
         extras[req.key] = [str(req)]
-    extras["dev"] = extras["test"]
-    extras["all"] = extras["dev"] + extras["examples"] + extras["strategies"]
+    extras["all"] = extras["strategies"] + extras["examples"]
+    extras["dev"] = extras["all"] + extras["test"]
     return extras
 
 
@@ -72,7 +75,9 @@ def _setup_args() -> Dict[str, Any]:
         keywords=["deep learning", "pytorch", "AI"],
         python_requires=">=3.7",
         setup_requires=["wheel"],
-        install_requires=assistant.load_requirements(_PATH_REQUIREMENTS, unfreeze=not _FREEZE_REQUIREMENTS),
+        install_requires=assistant.load_requirements(
+            _PATH_REQUIREMENTS, unfreeze="none" if _FREEZE_REQUIREMENTS else "all"
+        ),
         extras_require=_prepare_extras(),
         project_urls={
             "Bug Tracker": "https://github.com/Lightning-AI/lightning/issues",

src/pytorch_lightning/__setup__.py

@@ -1,5 +1,7 @@
+import glob
 import os.path
 from importlib.util import module_from_spec, spec_from_file_location
+from pathlib import Path
 from types import ModuleType
 from typing import Any, Dict
 
@@ -33,18 +35,17 @@ def _prepare_extras() -> Dict[str, Any]:
     # Define package extras. These are only installed if you specify them.
     # From remote, use like `pip install pytorch-lightning[dev, docs]`
     # From local copy of repo, use like `pip install ".[dev, docs]"`
-    common_args = dict(path_dir=_PATH_REQUIREMENTS, unfreeze="" if _FREEZE_REQUIREMENTS else "all")
+    common_args = dict(path_dir=_PATH_REQUIREMENTS, unfreeze="none" if _FREEZE_REQUIREMENTS else "all")
+    req_files = [Path(p) for p in glob.glob(os.path.join(_PATH_REQUIREMENTS, "*.txt"))]
     extras = {
-        # 'docs': load_requirements(file_name='docs.txt'),
-        "examples": assistant.load_requirements(file_name="examples.txt", **common_args),
-        "extra": assistant.load_requirements(file_name="extra.txt", **common_args),
-        "strategies": assistant.load_requirements(file_name="strategies.txt", **common_args),
-        "test": assistant.load_requirements(file_name="test.txt", **common_args),
+        p.stem: assistant.load_requirements(file_name=p.name, **common_args)
+        for p in req_files
+        if p.name not in ("docs.txt", "devel.txt", "base.txt")
     }
     for req in parse_requirements(extras["strategies"]):
         extras[req.key] = [str(req)]
-    extras["dev"] = extras["extra"] + extras["test"]
-    extras["all"] = extras["dev"] + extras["examples"] + extras["strategies"]  # + extras['docs']
+    extras["all"] = extras["extra"] + extras["strategies"] + extras["examples"]
+    extras["dev"] = extras["all"] + extras["test"]  # + extras['docs']
     return extras
 
 
@@ -84,7 +85,7 @@ def _setup_args() -> Dict[str, Any]:
         # TODO: aggregate pytorch and lite requirements as we include its source code directly in this package.
         # this is not a problem yet because lite's base requirements are all included in pytorch's base requirements
         install_requires=assistant.load_requirements(
-            _PATH_REQUIREMENTS, unfreeze="" if _FREEZE_REQUIREMENTS else "all"
+            _PATH_REQUIREMENTS, unfreeze="none" if _FREEZE_REQUIREMENTS else "all"
         ),
         extras_require=_prepare_extras(),
         project_urls={