update internal CLI prune_pkgs

.actions/assistant.py

@@ -35,7 +35,15 @@ def prepare_nightly_version(proj_root: str = _PATH_ROOT) -> None:
             fp.write(init)
 
     @staticmethod
-    def requirements_prune_pkgs(req_file: str, packages: Sequence[str]) -> None:
+    def requirements_prune_pkgs(packages: Sequence[str], req_files: Sequence[str] = REQUIREMENT_FILES) -> None:
+        """Remove some packages from given requirement files."""
+        if isinstance(req_files, str):
+            req_files = [req_files]
+        for req in req_files:
+            AssistantCLI._prune_packages(req, packages)
+
+    @staticmethod
+    def _prune_packages(req_file: str, packages: Sequence[str]) -> None:
         """Remove some packages from given requirement files."""
         with open(req_file) as fp:
             lines = fp.readlines()

.azure-pipelines/hpu-tests.yml

@@ -28,5 +28,6 @@ jobs:
 
     steps:
     - bash: |
+        apt-get install -y hwinfo
         hwinfo --short
       displayName: 'Instance HW info'

.github/workflows/ci_test-full.yml

@@ -57,7 +57,7 @@ jobs:
     - name: Setup Windows
       if: runner.os == 'windows'
       run: |
-        python .actions/assistant.py requirements_prune_pkgs requirements/extra.txt --packages horovod
+        python .actions/assistant.py requirements_prune_pkgs horovod
     - name: Set min. dependencies
       if: matrix.requires == 'oldest'
       run: |

.github/workflows/docs-checks.yml

@@ -43,7 +43,7 @@ jobs:
           pip --version
           pip install -q fire
           # remove Horovod from requirements
-          python .actions/assistant.py requirements_prune_pkgs requirements/extra.txt --packages horovod
+          python .actions/assistant.py requirements_prune_pkgs horovod
           # python -m pip install --upgrade --user pip
           pip install --requirement requirements.txt --upgrade-strategy only-if-needed --find-links https://download.pytorch.org/whl/cpu/torch_stable.html --quiet
           pip install --requirement requirements/extra.txt

dockers/base-conda/Dockerfile

@@ -96,7 +96,7 @@ RUN \
     python requirements_adjust_versions.py requirements-extra.txt && \
     python -c "print(' '.join([ln for ln in open('requirements-extra.txt').readlines() if 'horovod' in ln]))" > requirements_horovod.txt && \
     pip install -q fire && \
-    python assistant.py requirements_prune_pkgs requirements-extra.txt "horovod" && \
+    python assistant.py requirements_prune_pkgs "horovod" --req_files requirements-extra.txt && \
     python requirements_adjust_versions.py requirements-examples.txt && \
     # Install remaining requirements
     pip install -r requirements.txt --no-cache-dir --find-links https://download.pytorch.org/whl/test/torch_test.html && \

dockers/base-cuda/Dockerfile

@@ -83,7 +83,7 @@ RUN \
     python ./requirements/adjust-versions.py requirements/extra.txt ${PYTORCH_VERSION} && \
     python ./requirements/adjust-versions.py requirements/examples.txt ${PYTORCH_VERSION} && \
     python -c "print(' '.join([ln for ln in open('requirements/extra.txt').readlines() if 'horovod' in ln]))" > ./requirements/horovod.txt && \
-    python assistant.py requirements_prune_pkgs requirements/extra.txt "horovod" && \
+    python assistant.py requirements_prune_pkgs "horovod" && \
     # Install all requirements \
     pip install -r requirements/devel.txt --no-cache-dir --find-links https://download.pytorch.org/whl/cu${CUDA_VERSION_MM}/torch_stable.html && \
     rm -rf requirements.* && \

dockers/base-ipu/Dockerfile

@@ -83,7 +83,7 @@ RUN \
     python -c "import torch; print(torch.__version__)" && \
     python adjust_versions.py requirements-extra.txt && \
     pip install -q fire && \
-    python assistant.py requirements_prune_pkgs requirements-extra.txt --packages fairscale,horovod && \
+    python assistant.py requirements_prune_pkgs fairscale,horovod --req_files requirements-extra.txt && \
     # Install remaining requirements
     pip install -r requirements-extra.txt --no-cache-dir && \
     pip install -r requirements-test.txt --no-cache-dir && \

dockers/base-xla/Dockerfile

@@ -93,10 +93,9 @@ RUN \
     cd pytorch-lightning && \
     pip install -q fire && \
     # drop packages installed with XLA
-    python .actions/assistant.py requirements_prune_pkgs requirements.txt --packages torch && \
-    python .actions/assistant.py requirements_prune_pkgs requirements/examples.txt --packages torchvision && \
+    python .actions/assistant.py requirements_prune_pkgs torch,torchvision && \
     # drop unnecessary packages
-    python .actions/assistant.py requirements_prune_pkgs requirements/extra.txt --packages fairscale,horovod && \
+    python .actions/assistant.py requirements_prune_pkgs fairscale,horovod && \
     python ./requirements/adjust-versions.py ./requirements/extra.txt && \
     # install PL dependencies
     pip install --requirement ./requirements/devel.txt --no-cache-dir && \

dockers/nvidia/Dockerfile

@@ -42,7 +42,7 @@ RUN \
 
 # Installations \
     pip install -q fire && \
-    python ./pytorch-lightning/.actions/assistant.py requirements_prune_pkgs ./pytorch-lightning/requirements/extra.txt --packages horovod && \
+    python ./pytorch-lightning/.actions/assistant.py requirements_prune_pkgs horovod --req_files ./pytorch-lightning/requirements/extra.txt  && \
     pip install "Pillow>=8.2, !=8.3.0" "cryptography>=3.4" "py>=1.10" --no-cache-dir --upgrade-strategy only-if-needed && \
     pip install -r ./pytorch-lightning/requirements/extra.txt --no-cache-dir --upgrade-strategy only-if-needed && \
     pip install -r ./pytorch-lightning/requirements/examples.txt --no-cache-dir --upgrade-strategy only-if-needed && \

dockers/tpu-tests/Dockerfile

@@ -30,7 +30,7 @@ RUN cd pytorch-lightning && \
 RUN \
     pip install -q fire && \
     # drop unnecessary packages
-    python .actions/assistant.py requirements_prune_pkgs ./pytorch-lightning/requirements/extra.txt --packages fairscale,horovod && \
+    python .actions/assistant.py requirements_prune_pkgs fairscale,horovod --req_files ./pytorch-lightning/requirements/extra.txt && \
     pip install -r pytorch-lightning/requirements/devel.txt --no-cache-dir
 
 COPY ./dockers/tpu-tests/docker-entrypoint.sh /usr/local/bin/