unify extras & minor CI cleaning: move env. var (#15942)

Co-authored-by: Carlos Mocholí <[email protected]>

.actions/assistant.py

@@ -60,9 +60,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'
@@ -79,6 +79,7 @@ def _augment_requirement(ln: str, comment_char: str = "#", unfreeze: str = "all"
     >>> _augment_requirement("arrow", unfreeze="major")
     'arrow'
     """
+    assert unfreeze in {"none", "major", "all"}
     # filer all comments
     if comment_char in ln:
         comment = ln[ln.index(comment_char) :]
@@ -88,7 +89,7 @@ def _augment_requirement(ln: str, comment_char: str = "#", unfreeze: str = "all"
         is_strict = False
     req = ln.strip()
     # skip directly installed dependencies
-    if not req or req.startswith("http") or "@" in req:
+    if not req or any(c in req for c in ["http:", "https:", "@"]):
         return ""
     # extract the major version from all listed versions
     if unfreeze == "major":
@@ -99,7 +100,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
@@ -121,6 +122,7 @@ def load_requirements(
     >>> load_requirements(path_req, "docs.txt", unfreeze="major")  # doctest: +ELLIPSIS +NORMALIZE_WHITESPACE
     ['sphinx>=4.0, <6.0  # strict', ...]
     """
+    assert unfreeze in {"none", "major", "all"}
     with open(os.path.join(path_dir, file_name)) as file:
         lines = [ln.strip() for ln in file.readlines()]
     reqs = [_augment_requirement(ln, comment_char=comment_char, unfreeze=unfreeze) for ln in lines]
@@ -206,17 +208,21 @@ def _download_frontend(pkg_path: str):
 
 
 def _load_aggregate_requirements(req_dir: str = "requirements", freeze_requirements: bool = False) -> None:
-    """Load all base requirements from all particular packages and prune duplicates."""
+    """Load all base requirements from all particular packages and prune duplicates.
+
+    >>> _load_aggregate_requirements(os.path.join(_PROJECT_ROOT, "requirements"))
+    """
     requires = [
-        load_requirements(d, file_name="base.txt", unfreeze=not freeze_requirements)
+        # TODO: consider passing unfreeze as string instead
+        load_requirements(d, file_name="base.txt", 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
         if os.path.isdir(d) and len(glob.glob(os.path.join(d, "*"))) > 0 and "__pycache__" not in d
     ]
     if not requires:
         return None
     # TODO: add some smarter version aggregation per each package
-    requires = list(chain(*requires))
+    requires = sorted(set(chain(*requires)))
     with open(os.path.join(req_dir, "base.txt"), "w") as fp:
         fp.writelines([ln + os.linesep for ln in requires] + [os.linesep])
 

.azure/app-cloud-e2e.yml

@@ -113,11 +113,15 @@ jobs:
       displayName: 'Info'
 
     # TODO: we are testing it as `lightning`, so add also version for `lightning_app`
-    - bash: git restore . && python -m pip install -e . -r requirements/app/test.txt -r requirements/app/ui.txt --find-links https://download.pytorch.org/whl/cpu/torch_stable.html
+    - bash: |
+        pip install -e .[test] \
+          -f https://download.pytorch.org/whl/cpu/torch_stable.html
+      env:
+        FREEZE_REQUIREMENTS: "1"
       displayName: 'Install Lightning & dependencies'
 
     - bash: |
-        python -m pip install playwright
+        pip install playwright
         python -m playwright install  # --with-deps
       displayName: 'Install Playwright system dependencies'
 
@@ -130,7 +134,6 @@ jobs:
       condition: eq(variables['name'], 'quick_start')
       displayName: 'Clone Quick start Repo'
 
-
     # - bash: |
     #     rm -rf examples/app_template_jupyterlab || true
     #     git clone https://github.com/Lightning-AI/LAI-lightning-template-jupyterlab-App examples/app_template_jupyterlab
@@ -144,7 +147,7 @@ jobs:
       condition: eq(variables['name'], 'template_react_ui')
       displayName: 'Clone Template React UI Repo'
 
-    - bash: python -m pip install -q -r .actions/requirements.txt
+    - bash: pip install -q -r .actions/requirements.txt
       displayName: 'Install assistant dependencies'
 
     # Fix imports to use `lightning` instead of `lightning_app` since we install lightning only ATM

.azure/gpu-tests-lite.yml

@@ -96,6 +96,7 @@ jobs:
     - bash: python -m coverage run --source lightning_lite -m pytest --ignore benchmarks -v --junitxml=$(Build.StagingDirectory)/test-results.xml --durations=50
       env:
         PL_RUN_CUDA_TESTS: "1"
+        CUDA_LAUNCH_BLOCKING: "1"
       workingDirectory: tests/tests_lite
       displayName: 'Testing: Lite standard'
       timeoutInMinutes: "10"

.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

@@ -110,8 +110,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
@@ -133,7 +134,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/__setup__.py

@@ -32,7 +32,7 @@ def _prepare_extras() -> Dict[str, Any]:
     # From remote, use like `pip install pytorch-lightning[dev, docs]`
     # From local copy of repo, use like `pip install ".[dev, docs]"`
     req_files = [Path(p) for p in glob.glob(os.path.join(_PATH_REQUIREMENTS, "*", "*.txt"))]
-    common_args = dict(unfreeze="major" if _FREEZE_REQUIREMENTS else "all")
+    common_args = dict(unfreeze="none" if _FREEZE_REQUIREMENTS else "major")
     extras = {
         f"{p.parent.name}-{p.stem}": _ASSISTANT.load_requirements(file_name=p.name, path_dir=p.parent, **common_args)
         for p in req_files
@@ -41,12 +41,11 @@ def _prepare_extras() -> Dict[str, Any]:
     for extra in list(extras):
         name = "-".join(extra.split("-")[1:])
         extras[name] = extras.get(name, []) + extras[extra]
-    # todo
-    # extras["extra"] = extras["cloud"] + extras["ui"]
-    # extras["dev"] = extras["extra"] + extras["test"]  # + extras['docs']
-    # extras["all"] = extras["dev"]
-    extras = {name: list(set(reqs)) for name, reqs in extras.items()}
-    print("The extras are", extras)
+    extras["extra"] += extras["cloud"] + extras["ui"] + extras["components"]
+    extras["all"] = extras["extra"]
+    extras["dev"] = extras["all"] + extras["test"]  # + extras['docs']
+    extras = {name: sorted(set(reqs)) for name, reqs in extras.items()}
+    print("The extras are: ", extras)
     return extras
 
 
@@ -84,7 +83,9 @@ def _setup_args() -> Dict[str, Any]:
             ],
         },
         setup_requires=[],
-        install_requires=_ASSISTANT.load_requirements(_PATH_REQUIREMENTS, unfreeze="all"),
+        install_requires=_ASSISTANT.load_requirements(
+            _PATH_REQUIREMENTS, unfreeze="none" if _FREEZE_REQUIREMENTS else "major"
+        ),
         extras_require=_prepare_extras(),
         project_urls={
             "Bug Tracker": "https://github.com/Lightning-AI/lightning/issues",

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
 
 
@@ -78,9 +80,9 @@ def _setup_args() -> Dict[str, Any]:
                 "lightning = lightning_app.cli.lightning_cli:main",
             ],
         },
-        setup_requires=["wheel"],
+        setup_requires=[],
         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={