Fix requirements/adjust_versions.py (#7149)

Co-authored-by: jirka <jirka.borovec@seznam.cz>

.github/workflows/ci_dockers.yml

@@ -10,7 +10,7 @@ on: # Trigger the workflow on push or pull request, but only for the master bran
     paths:
       - "dockers/**"
       - "!dockers/README.md"
-      - "requirements/*.txt"
+      - "requirements/*"
       - "environment.yml"
       - "requirements.txt"
       - ".github/workflows/*docker*.yml"

dockers/base-xla/Dockerfile

@@ -29,7 +29,7 @@ ENV \
     DEBIAN_FRONTEND=noninteractive \
     CONDA_ENV=lightning
 
-# show system inforation
+# show system info
 RUN lsb_release -a && cat /etc/*-release
 
 RUN apt-get update -qq && \
@@ -42,13 +42,13 @@ RUN apt-get update -qq && \
         ca-certificates \
         libomp5 \
     && \
-# Install conda and python.
-# NOTE new Conda does not forward the exit status... https://github.com/conda/conda/issues/8385
+    # Install conda and python.
+    # NOTE new Conda does not forward the exit status... https://github.com/conda/conda/issues/8385
     curl -o ~/miniconda.sh https://repo.anaconda.com/miniconda/Miniconda3-py38_${CONDA_VERSION}-Linux-x86_64.sh && \
     chmod +x ~/miniconda.sh && \
     ~/miniconda.sh -b && \
     rm ~/miniconda.sh && \
-# Cleaning
+    # Cleaning
     apt-get autoremove -y && \
     apt-get clean && \
     rm -rf /root/.cache && \
@@ -79,7 +79,7 @@ ENV \
 RUN pip --version && \
     pip config set global.cache-dir false && \
     conda remove pytorch torchvision && \
-# Install Pytorch XLA
+    # Install Pytorch XLA
     py_version=${PYTHON_VERSION/./} && \
     # Python 3.7 wheels are available. Replace cp36-cp36m with cp37-cp37m
     gsutil cp "gs://tpu-pytorch/wheels/torch-${XLA_VERSION}-cp${py_version}-cp${py_version}m-linux_x86_64.whl" . && \
@@ -91,20 +91,17 @@ RUN pip --version && \
 # Get package
 COPY ./ ./pytorch-lightning/
 
-# Install pytorch-lightning dependencies.
 RUN \
     python --version && \
-# Install PL dependencies
     cd pytorch-lightning && \
-    # drop Torch as it was installed with XLA
+    # drop packages installed with XLA
     python -c "fname = 'requirements.txt' ; lines = [line for line in open(fname).readlines() if not line.startswith('torch')] ; open(fname, 'w').writelines(lines)" && \
-    # drop Horovod as it is not needed
-    python -c "fname = 'requirements/extra.txt' ; lines = [line for line in open(fname).readlines() if not line.startswith('horovod')] ; open(fname, 'w').writelines(lines)" && \
-    # drop fairscale as it is not needed
-    python -c "fname = 'requirements/extra.txt' ; lines = [line for line in open(fname).readlines() if 'fairscale' not in line] ; open(fname, 'w').writelines(lines)" && \
-    # drop TorchVision as it was installed with XLA
     python -c "fname = 'requirements/examples.txt' ; lines = [line for line in open(fname).readlines() if not line.startswith('torchvision')] ; open(fname, 'w').writelines(lines)" && \
+    # drop unnecessary packages
+    python -c "fname = 'requirements/extra.txt' ; lines = [line for line in open(fname).readlines() if not line.startswith('horovod')] ; open(fname, 'w').writelines(lines)" && \
+    python -c "fname = 'requirements/extra.txt' ; lines = [line for line in open(fname).readlines() if 'fairscale' not in line] ; open(fname, 'w').writelines(lines)" && \
     python ./requirements/adjust_versions.py ./requirements/extra.txt && \
+    # install PL dependencies
     pip install --requirement ./requirements/devel.txt --no-cache-dir && \
     cd .. && \
     rm -rf pytorch-lightning && \

requirements/adjust_versions.py

@@ -1,47 +1,50 @@
 import os
 import re
 import sys
-from typing import Any, Dict
+from typing import Dict, Optional
 
-VERSIONS_LUT: Dict[str, Dict[str, Any]] = {
-    "1.4.0": dict(torchvision="0.5.0", torchtext="0.5"),
-    "1.5.0": dict(torchvision="0.6.0", torchtext="0.6"),
-    "1.5.1": dict(torchvision="0.6.1", torchtext="0.6"),
-    "1.6.0": dict(torchvision="0.7.0", torchtext="0.7"),
-    "1.7.0": dict(torchvision="0.8.1", torchtext="0.8"),
-    "1.7.1": dict(torchvision="0.8.2", torchtext="0.8.1"),
-    "1.8.0": dict(torchvision="0.9.0", torchtext="0.9"),
-    "1.8.1": dict(torchvision="0.9.0", torchtext="0.9"),
-}
+VERSIONS = [
+    dict(torch="1.9.0", torchvision="", torchtext=""),  # nightly
+    dict(torch="1.8.1", torchvision="0.9.1", torchtext="0.9.1"),
+    dict(torch="1.8.0", torchvision="0.9.0", torchtext="0.9.0"),
+    dict(torch="1.7.1", torchvision="0.8.2", torchtext="0.8.1"),
+    dict(torch="1.7.0", torchvision="0.8.1", torchtext="0.8.0"),
+    dict(torch="1.6.0", torchvision="0.7.0", torchtext="0.7"),
+    dict(torch="1.5.1", torchvision="0.6.1", torchtext="0.6"),
+    dict(torch="1.5.0", torchvision="0.6.0", torchtext="0.6"),
+    dict(torch="1.4.0", torchvision="0.5.0", torchtext="0.5"),
+]
+VERSIONS.sort(key=lambda v: v["torch"], reverse=True)
 
 
-def find_latest(ver: str, versions_all: list) -> str:
+def find_latest(ver: str) -> Dict[str, str]:
     # drop all except semantic version
     ver = re.search(r'([\.\d]+)', ver).groups()[0]
-    # find candidates, by starting version pattern
-    options = [v for v in versions_all if v.startswith(ver)]
-    assert options, f"missing {ver} among {versions_all}"
-    # take the last one...
-    return sorted(options)[-1]
+
+    # find first match
+    for option in VERSIONS:
+        if option["torch"].startswith(ver):
+            return option
+
+    raise ValueError(f"Missing {ver} in {VERSIONS}")
 
 
-def main(path_req: str, torch_version: str = None) -> None:
-    with open(path_req, "r") as fp:
-        req = fp.read()
-
+def main(path_req: str, torch_version: Optional[str] = None) -> None:
     if not torch_version:
         import torch
         torch_version = torch.__version__
-    assert torch_version, f"invalid/missing Torch: {torch_version}"
+    assert torch_version, f"invalid torch: {torch_version}"
 
-    torch_version = find_latest(torch_version, list(VERSIONS_LUT.keys()))
-    dep_versions = VERSIONS_LUT[torch_version]
-    dep_versions["torch"] = torch_version
-    for lib in dep_versions:
-        version = dep_versions[lib]
-        replace = f"{lib}=={version}\n"
+    with open(path_req, "r") as fp:
+        req = fp.read()
+
+    latest = find_latest(torch_version)
+    for lib, version in latest.items():
+        replace = f"{lib}=={version}" if version else lib
+        replace += os.linesep
         req = re.sub(rf"{lib}[>=]*[\d\.]*{os.linesep}", replace, req)
 
+    print(req)  # on purpose - to debug
     with open(path_req, "w") as fp:
         fp.write(req)