CI: fix requirements freeze

.azure/gpu-benchmark.yml

@@ -40,6 +40,7 @@ jobs:
         pip list
       env:
         PACKAGE_NAME: pytorch
+        FREEZE_REQUIREMENTS: 1
       displayName: 'Install package'
 
     - bash: python -m pytest benchmarks -v --durations=0

.azure/gpu-tests.yml

@@ -56,6 +56,7 @@ jobs:
         pip list
       env:
         PACKAGE_NAME: pytorch
+        FREEZE_REQUIREMENTS: 1
       displayName: 'Install dependencies'
 
     - bash: |

.azure/hpu-tests.yml

@@ -33,6 +33,7 @@ jobs:
         pip install -e .[extra] -r requirements/pytorch/test.txt
       env:
         PACKAGE_NAME: pytorch
+        FREEZE_REQUIREMENTS: 1
       displayName: 'Install dependencies'
 
     - bash: |

.azure/ipu-tests.yml

@@ -57,6 +57,7 @@ jobs:
         pip list
       env:
         PACKAGE_NAME: pytorch
+        FREEZE_REQUIREMENTS: 1
       displayName: 'Install dependencies'
 
     - bash: |

.github/workflows/ci-pytorch_test-conda.yml

@@ -40,6 +40,7 @@ jobs:
     - name: Update base dependencies
       env:
         PACKAGE_NAME: pytorch
+        FREEZE_REQUIREMENTS: 1
       run: |
         conda info
         conda list

.github/workflows/ci-pytorch_test-full.yml

@@ -81,6 +81,7 @@ jobs:
     - name: Install dependencies
       env:
         PACKAGE_NAME: pytorch
+        FREEZE_REQUIREMENTS: 1
       run: |
         flag=$(python -c "print('--pre' if '${{matrix.release}}' == 'pre' else '')" 2>&1)
         url=$(python -c "print('test/cpu/torch_test.html' if '${{matrix.release}}' == 'pre' else 'cpu/torch_stable.html')" 2>&1)

.github/workflows/ci-pytorch_test-slow.yml

@@ -52,6 +52,7 @@ jobs:
     - name: Install dependencies
       env:
         PACKAGE_NAME: pytorch
+        FREEZE_REQUIREMENTS: 1
       run: |
         # adjust versions according installed Torch version
         python ./requirements/pytorch/adjust-versions.py requirements.txt ${{ matrix.pytorch-version }}

.github/workflows/docs-checks.yml

@@ -38,6 +38,7 @@ jobs:
       - name: Install dependencies
         env:
           PACKAGE_NAME: pytorch
+          FREEZE_REQUIREMENTS: 1
         run: |
           sudo apt-get update
           sudo apt-get install -y cmake pandoc
@@ -82,6 +83,7 @@ jobs:
       - name: Install dependencies
         env:
           PACKAGE_NAME: pytorch
+          FREEZE_REQUIREMENTS: 1
         run: |
           sudo apt-get update
           sudo apt-get install -y cmake pandoc

dockers/base-ipu/Dockerfile

@@ -73,20 +73,18 @@ ENV \
     CONDA_DEFAULT_ENV=${CONDA_ENV} \
     MKL_THREADING_LAYER=GNU
 
-COPY ./requirements/pytorch/extra.txt requirements-extra.txt
-COPY ./requirements/pytorch/test.txt requirements-test.txt
-COPY ./requirements/pytorch/adjust-versions.py adjust_versions.py
+COPY ./requirements/pytorch/* requirements/
 COPY ./.actions/assistant.py assistant.py
 
 RUN \
     pip list | grep torch && \
     python -c "import torch; print(torch.__version__)" && \
-    python adjust_versions.py requirements-extra.txt && \
     pip install -q fire && \
+    python requirements/adjust-versions.py requirements/extra.txt && \
     # Install remaining requirements
-    pip install -r requirements-extra.txt --no-cache-dir && \
-    pip install -r requirements-test.txt --no-cache-dir && \
-    rm requirements*
+    pip install -r requirements/extra.txt --no-cache-dir && \
+    pip install -r requirements/test.txt --no-cache-dir && \
+    rm -rf requirements/
 
 RUN \
     # Show what we have

dockers/tpu-tests/tpu_test_cases.jsonnet

@@ -29,6 +29,8 @@ local tputests = base.BaseTest {
       git ls-remote --refs origin
       git fetch origin "refs/pull/{PR_NUMBER}/head:pr/{PR_NUMBER}" && git checkout "pr/{PR_NUMBER}"
       git checkout {SHA}
+      export PACKAGE_NAME=pytorch
+      export FREEZE_REQUIREMENTS=1
       pip install -e .
       echo $KUBE_GOOGLE_CLOUD_TPU_ENDPOINTS
       export XRT_TPU_CONFIG="tpu_worker;0;${KUBE_GOOGLE_CLOUD_TPU_ENDPOINTS:7}"

requirements/pytorch/base.txt

@@ -1,5 +1,5 @@
 numpy>=1.17.2, <1.23.1
-torch>=1.9.*, <=1.11.0
+torch>=1.9.*, <=1.11.0  # strict
 tqdm>=4.57.0, <=4.63.0
 PyYAML>=5.4, <=6.0
 fsspec[http]>=2021.05.0, !=2021.06.0, <2022.6.0
@@ -8,4 +8,4 @@ torchmetrics>=0.7.0, <0.9.2  # needed for using fixed compare_version
 pyDeprecate>=0.3.1, <=0.3.2
 packaging>=17.0, <=21.3
 typing-extensions>=4.0.0, <4.2.1
-protobuf<=3.20.1  # strict. TODO: Remove after tensorboard gets compatible https://github.com/tensorflow/tensorboard/issues/5708
+protobuf<=3.20.1  # strict TODO: Remove after tensorboard gets compatible https://github.com/tensorflow/tensorboard/issues/5708

src/lightning/__setup__.py

@@ -10,6 +10,7 @@
 _PROJECT_ROOT = "."
 _SOURCE_ROOT = os.path.join(_PROJECT_ROOT, "src")
 _PACKAGE_ROOT = os.path.join(_SOURCE_ROOT, "lightning")
+_FREEZE_REQUIREMENTS = bool(int(os.environ.get("FREEZE_REQUIREMENTS", 0)))
 
 
 def _load_py_module(name: str, location: str) -> ModuleType:
@@ -56,7 +57,9 @@ def _setup_args(**kwargs: Any) -> Dict[str, Any]:
     else:
         _include_pkgs = ["*"]
         _requires = [
-            _setup_tools.load_requirements(d) for d in glob.glob(os.path.join("requirements", "*")) if os.path.isdir(d)
+            _setup_tools.load_requirements(d, unfreeze=not _FREEZE_REQUIREMENTS)
+            for d in glob.glob(os.path.join("requirements", "*"))
+            if os.path.isdir(d)
         ]
         _requires = list(chain(*_requires))
     # todo: consider invaliding some additional arguments from packages, for example if include data or safe to zip

src/lightning_app/__setup__.py

@@ -9,6 +9,7 @@
 _SOURCE_ROOT = os.path.join(_PROJECT_ROOT, "src")
 _PACKAGE_ROOT = os.path.join(_SOURCE_ROOT, "lightning_app")
 _PATH_REQUIREMENTS = os.path.join("requirements", "app")
+_FREEZE_REQUIREMENTS = bool(int(os.environ.get("FREEZE_REQUIREMENTS", 0)))
 
 
 def _load_py_module(name: str, location: str) -> ModuleType:
@@ -25,11 +26,12 @@ def _prepare_extras(**kwargs: Any) -> 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=not _FREEZE_REQUIREMENTS)
     extras = {
         # 'docs': load_requirements(file_name='docs.txt'),
-        "cloud": _setup_tools.load_requirements(path_dir=_PATH_REQUIREMENTS, file_name="cloud.txt"),
-        "ui": _setup_tools.load_requirements(path_dir=_PATH_REQUIREMENTS, file_name="ui.txt"),
-        "test": _setup_tools.load_requirements(path_dir=_PATH_REQUIREMENTS, file_name="test.txt"),
+        "cloud": _setup_tools.load_requirements(file_name="cloud.txt", **common_args),
+        "ui": _setup_tools.load_requirements(file_name="ui.txt", **common_args),
+        "test": _setup_tools.load_requirements(file_name="test.txt", **common_args),
     }
     extras["dev"] = extras["cloud"] + extras["ui"] + extras["test"]  # + extras['docs']
     extras["all"] = extras["cloud"] + extras["ui"]
@@ -83,7 +85,7 @@ def _setup_args(**__: Any) -> Dict[str, Any]:
             ],
         },
         setup_requires=["wheel"],
-        install_requires=_setup_tools.load_requirements(_PATH_REQUIREMENTS),
+        install_requires=_setup_tools.load_requirements(_PATH_REQUIREMENTS, unfreeze=not _FREEZE_REQUIREMENTS),
         extras_require=_prepare_extras(),
         project_urls={
             "Bug Tracker": "https://github.com/Lightning-AI/lightning/issues",

src/pytorch_lightning/__setup__.py

@@ -10,6 +10,7 @@
 _SOURCE_ROOT = os.path.join(_PROJECT_ROOT, "src")
 _PACKAGE_ROOT = os.path.join(_SOURCE_ROOT, "pytorch_lightning")
 _PATH_REQUIREMENTS = os.path.join("requirements", "pytorch")
+_FREEZE_REQUIREMENTS = bool(int(os.environ.get("FREEZE_REQUIREMENTS", 0)))
 
 
 def _load_py_module(name: str, location: str) -> ModuleType:
@@ -26,13 +27,14 @@ def _prepare_extras(**kwargs: Any) -> 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=not _FREEZE_REQUIREMENTS)
     extras = {
         # 'docs': load_requirements(file_name='docs.txt'),
-        "examples": _setup_tools.load_requirements(path_dir=_PATH_REQUIREMENTS, file_name="examples.txt"),
-        "loggers": _setup_tools.load_requirements(path_dir=_PATH_REQUIREMENTS, file_name="loggers.txt"),
-        "extra": _setup_tools.load_requirements(path_dir=_PATH_REQUIREMENTS, file_name="extra.txt"),
-        "strategies": _setup_tools.load_requirements(path_dir=_PATH_REQUIREMENTS, file_name="strategies.txt"),
-        "test": _setup_tools.load_requirements(path_dir=_PATH_REQUIREMENTS, file_name="test.txt"),
+        "examples": _setup_tools.load_requirements(file_name="examples.txt", **common_args),
+        "loggers": _setup_tools.load_requirements(file_name="loggers.txt", **common_args),
+        "extra": _setup_tools.load_requirements(file_name="extra.txt", **common_args),
+        "strategies": _setup_tools.load_requirements(file_name="strategies.txt", **common_args),
+        "test": _setup_tools.load_requirements(file_name="test.txt", **common_args),
     }
     for req in parse_requirements(extras["strategies"]):
         extras[req.key] = [str(req)]
@@ -83,7 +85,7 @@ def _setup_args(**__: Any) -> Dict[str, Any]:
         keywords=["deep learning", "pytorch", "AI"],
         python_requires=">=3.7",
         setup_requires=[],
-        install_requires=_setup_tools.load_requirements(_PATH_REQUIREMENTS),
+        install_requires=_setup_tools.load_requirements(_PATH_REQUIREMENTS, unfreeze=not _FREEZE_REQUIREMENTS),
         extras_require=_prepare_extras(),
         project_urls={
             "Bug Tracker": "https://github.com/Lightning-AI/lightning/issues",