ci: adjust version in all requirements (#16100)

Author: Borda

.actions/assistant.py

@@ -218,7 +218,7 @@ def _load_aggregate_requirements(req_dir: str = "requirements", freeze_requireme
     # TODO: add some smarter version aggregation per each package
     requires = list(chain(*requires))
     with open(os.path.join(req_dir, "base.txt"), "w") as fp:
-        fp.writelines([ln + os.linesep for ln in requires])
+        fp.writelines([ln + os.linesep for ln in requires] + [os.linesep])
 
 
 def _retrieve_files(directory: str, *ext: str) -> List[str]:
@@ -333,7 +333,7 @@ def _prune_packages(req_file: str, packages: Sequence[str]) -> None:
             if req.name not in packages:
                 final.append(line)
         print(final)
-        path.write_text("\n".join(final))
+        path.write_text("\n".join(final) + "\n")
 
     @staticmethod
     def _replace_min(fname: str) -> None:

.azure/gpu-tests-lite.yml

@@ -74,8 +74,9 @@ jobs:
 
     - bash: |
         PYTORCH_VERSION=$(python -c "import torch; print(torch.__version__.split('+')[0])")
-        python ./requirements/pytorch/adjust-versions.py requirements/lite/base.txt ${PYTORCH_VERSION}
-        python ./requirements/pytorch/adjust-versions.py requirements/lite/examples.txt ${PYTORCH_VERSION}
+        for fpath in `ls requirements/**/*.txt`; do \
+          python ./requirements/pytorch/adjust-versions.py $fpath ${PYTORCH_VERSION}; \
+        done
       displayName: 'Adjust dependencies'
 
     - bash: |

.azure/gpu-tests-pytorch.yml

@@ -87,14 +87,12 @@ jobs:
       displayName: 'Image info & NVIDIA'
 
     - bash: |
-        python -c "fname = 'requirements/pytorch/strategies.txt' ; lines = [line for line in open(fname).readlines() if 'horovod' not in line] ; open(fname, 'w').writelines(lines)"
-        python -c "fname = 'requirements/pytorch/strategies.txt' ; lines = [line for line in open(fname).readlines() if 'bagua' not in line] ; open(fname, 'w').writelines(lines)"
-        python -c "fname = 'requirements/pytorch/strategies.txt' ; lines = [line for line in open(fname).readlines() if 'colossalai' not in line] ; open(fname, 'w').writelines(lines)"
+        python .actions/assistant.py requirements_prune_pkgs --packages="[horovod,bagua,colossalai]"
 
         PYTORCH_VERSION=$(python -c "import torch; print(torch.__version__.split('+')[0])")
-        python ./requirements/pytorch/adjust-versions.py requirements/pytorch/base.txt ${PYTORCH_VERSION}
-        python ./requirements/pytorch/adjust-versions.py requirements/pytorch/extra.txt ${PYTORCH_VERSION}
-        python ./requirements/pytorch/adjust-versions.py requirements/pytorch/examples.txt ${PYTORCH_VERSION}
+        for fpath in `ls requirements/**/*.txt`; do \
+          python ./requirements/pytorch/adjust-versions.py $fpath ${PYTORCH_VERSION}; \
+        done
       displayName: 'Adjust dependencies'
 
     - bash: pip install -e .[dev,examples] --find-links ${TORCH_URL}

.azure/ipu-tests.yml

@@ -68,8 +68,9 @@ jobs:
       displayName: "Reset IPU devices"
 
     - bash: |
-        python ./requirements/pytorch/adjust-versions.py requirements/pytorch/extra.txt
-        python ./requirements/pytorch/adjust-versions.py requirements/pytorch/examples.txt
+        for fpath in `ls requirements/**/*.txt`; do \
+          python ./requirements/pytorch/adjust-versions.py $fpath; \
+        done
         pip install -e .[dev]
         pip list
       env:

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

@@ -87,7 +87,9 @@ jobs:
     - name: Adjust PyTorch versions in requirements files
       if: ${{ matrix.requires != 'oldest' }}
       run: |
-        python ./requirements/pytorch/adjust-versions.py requirements/lite/base.txt ${{ matrix.pytorch-version }}
+        for fpath in `ls requirements/**/*.txt`; do \
+          python ./requirements/pytorch/adjust-versions.py $fpath ${{ matrix.pytorch-version }}; \
+        done
         cat requirements/lite/base.txt
 
     - name: Get pip cache dir

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

@@ -104,9 +104,9 @@ jobs:
     - name: Adjust PyTorch versions in requirements files
       if: ${{ matrix.requires != 'oldest' }}
       run: |
-        python ./requirements/pytorch/adjust-versions.py requirements/lite/base.txt ${{ matrix.pytorch-version }}
-        python ./requirements/pytorch/adjust-versions.py requirements/pytorch/base.txt ${{ matrix.pytorch-version }}
-        python ./requirements/pytorch/adjust-versions.py requirements/pytorch/examples.txt ${{ matrix.pytorch-version }}
+        for fpath in `ls requirements/**/*.txt`; do \
+          python ./requirements/pytorch/adjust-versions.py $fpath ${{ matrix.pytorch-version }}; \
+        done
         cat requirements/pytorch/base.txt
 
     - name: Get pip cache dir

dockers/base-cuda/Dockerfile

@@ -91,10 +91,10 @@ RUN \
     # Disable cache \
     export CUDA_VERSION_MM=$(python -c "print(''.join('$CUDA_VERSION'.split('.')[:2]))") && \
     pip config set global.cache-dir false && \
-    # set particular PyTorch version
-    python ./requirements/pytorch/adjust-versions.py requirements/pytorch/base.txt ${PYTORCH_VERSION} && \
-    python ./requirements/pytorch/adjust-versions.py requirements/pytorch/extra.txt ${PYTORCH_VERSION} && \
-    python ./requirements/pytorch/adjust-versions.py requirements/pytorch/examples.txt ${PYTORCH_VERSION} && \
+    # set particular PyTorch version \
+    for fpath in `ls requirements/**/*.txt`; do \
+      python ./requirements/pytorch/adjust-versions.py $fpath ${PYTORCH_VERSION}; \
+    done && \
 
     # Install base requirements \
     pip install -r requirements/pytorch/base.txt --no-cache-dir --find-links https://download.pytorch.org/whl/cu${CUDA_VERSION_MM}/torch_stable.html && \

dockers/release/Dockerfile

@@ -24,6 +24,8 @@ ARG LIGHTNING_VERSION=""
 
 COPY ./ /home/lightning/
 
+ENV PACKAGE_NAME=pytorch
+
 # install dependencies
 RUN \
     cd /home && \
@@ -39,9 +41,8 @@ RUN \
     fi && \
     # otherwise there is collision with folder name ans pkg name on Pypi
     cd lightning && \
-    # remove colossalai from requirements since they are installed separately
-    python -c "fname = 'requirements/pytorch/strategies.txt' ; lines = [line for line in open(fname).readlines() if 'colossalai' not in line] ; open(fname, 'w').writelines(lines)" ; \
-    PACKAGE_NAME=pytorch pip install '.[extra,loggers,strategies]' --no-cache-dir && \
+    pip install setuptools==59.5.0 && \
+    pip install '.[extra,loggers,strategies]' --no-cache-dir && \
     cd .. && \
     rm -rf lightning
 

requirements/pytorch/adjust-versions.py

@@ -32,7 +32,7 @@ def find_latest(ver: str) -> Dict[str, str]:
     raise ValueError(f"Missing {ver} in {VERSIONS}")
 
 
-def main(req: str, torch_version: Optional[str] = None) -> str:
+def replace(req: str, torch_version: Optional[str] = None) -> str:
     if not torch_version:
         import torch
 
@@ -50,34 +50,7 @@ def main(req: str, torch_version: Optional[str] = None) -> str:
     return req
 
 
-def test_check():
-    requirements = """
-    torch>=1.2.*
-    torch==1.2.3
-    torch==1.4
-    torch
-    future>=0.17.1
-    pytorch==1.5.6+123dev0
-    torchvision
-    torchmetrics>=0.4.1
-    """
-    expected = """
-    torch==1.12.1
-    torch==1.12.1
-    torch==1.12.1
-    torch==1.12.1
-    future>=0.17.1
-    pytorch==1.5.6+123dev0
-    torchvision==0.13.1
-    torchmetrics>=0.4.1
-    """.strip()
-    actual = main(requirements, "1.12")
-    assert actual == expected, (actual, expected)
-
-
 if __name__ == "__main__":
-    test_check()  # sanity check
-
     if len(sys.argv) == 3:
         requirements_path, torch_version = sys.argv[1:]
     else:
@@ -86,7 +59,7 @@ def test_check():
 
     with open(requirements_path) as fp:
         requirements = fp.read()
-    requirements = main(requirements, torch_version)
+    requirements = replace(requirements, torch_version)
     print(requirements)  # on purpose - to debug
     with open(requirements_path, "w") as fp:
         fp.write(requirements)