Assistant for Unified Package (#15207)

* Update assistant and workflow files
* Update .actions/assistant.py

Co-authored-by: otaj <6065855+otaj@users.noreply.github.com>
Co-authored-by: Jirka Borovec <Borda@users.noreply.github.com>
Co-authored-by: otaj <ota@lightning.ai>

.actions/assistant.py

@@ -14,7 +14,7 @@ from typing import List, Optional, Sequence
 from urllib import request
 from urllib.request import Request, urlopen
 
-import fire
+import jsonargparse
 import pkg_resources
 from packaging.version import parse as version_parse
 
@@ -32,7 +32,7 @@ REQUIREMENT_FILES = {
         "requirements/app/cloud.txt",
     ),
 }
-REQUIREMENT_FILES_ALL = tuple(chain(*REQUIREMENT_FILES.values()))
+REQUIREMENT_FILES_ALL = list(chain(*REQUIREMENT_FILES.values()))
 PACKAGE_MAPPING = {"app": "lightning-app", "pytorch": "pytorch-lightning"}
 
 
@@ -60,6 +60,16 @@ def _load_py_module(name: str, location: str) -> ModuleType:
     return py
 
 
+def _retrieve_files(directory: str, *ext: str) -> List[str]:
+    all_files = []
+    for root, _, files in os.walk(directory):
+        for fname in files:
+            if not ext or any(os.path.split(fname)[1].lower().endswith(e) for e in ext):
+                all_files.append(os.path.join(root, fname))
+
+    return all_files
+
+
 class AssistantCLI:
     _PATH_ROOT = str(Path(__file__).parent.parent)
     _PATH_SRC = os.path.join(_PATH_ROOT, "src")
@@ -93,7 +103,16 @@ class AssistantCLI:
         path = Path(req_file)
         assert path.exists()
         text = path.read_text()
-        final = [str(req) for req in pkg_resources.parse_requirements(text) if req.name not in packages]
+        lines = text.splitlines()
+        final = []
+        for line in lines:
+            ln_ = line.strip()
+            if not ln_ or ln_.startswith("#"):
+                final.append(line)
+                continue
+            req = list(pkg_resources.parse_requirements(ln_))[0]
+            if req.name not in packages:
+                final.append(line)
         pprint(final)
         path.write_text("\n".join(final))
 
@@ -168,6 +187,47 @@ class AssistantCLI:
                 shutil.rmtree(py_dir2, ignore_errors=True)
                 shutil.copytree(py_dir, py_dir2)
 
+    @staticmethod
+    def copy_replace_imports(
+        source_dir: str, source_import: str, target_import: str, target_dir: Optional[str] = None
+    ) -> None:
+        """Recursively replace imports in given folder."""
+
+        source_imports = source_import.strip().split(",")
+        target_imports = target_import.strip().split(",")
+        assert len(source_imports) == len(target_imports), (
+            "source and target imports must have the same length, "
+            f"source: {len(source_import)}, target: {len(target_import)}"
+        )
+
+        if target_dir is None:
+            target_dir = source_dir
+
+        ls = _retrieve_files(source_dir)
+
+        for fp in ls:
+            if fp.endswith(".py"):
+                with open(fp, encoding="utf-8") as fo:
+                    py = fo.readlines()
+
+                for source_import, target_import in zip(source_imports, target_imports):
+                    for i, ln in enumerate(py):
+                        py[i] = re.sub(rf"(?!_){source_import}(?!_)", target_import, ln)
+
+                if target_dir:
+                    fp_new = fp.replace(source_dir, target_dir)
+                    os.makedirs(os.path.dirname(fp_new), exist_ok=True)
+                else:
+                    fp_new = fp
+
+                with open(fp_new, "w", encoding="utf-8") as fo:
+                    fo.writelines(py)
+            elif not fp.endswith(".pyc"):
+                fp_new = fp.replace(source_dir, target_dir)
+                os.makedirs(os.path.dirname(fp_new), exist_ok=True)
+                if os.path.abspath(fp) != os.path.abspath(fp_new):
+                    shutil.copy2(fp, fp_new)
+
 
 if __name__ == "__main__":
-    fire.Fire(AssistantCLI)
+    jsonargparse.CLI(AssistantCLI)

.actions/requirements.txt

@@ -1,3 +1,4 @@
-fire
+jsonargparse>=4.15.0
 packaging
 requests
+typing_extensions

.azure/hpu-tests.yml

@@ -63,7 +63,7 @@ jobs:
         pip --version
         sudo pip uninstall -y lightning pytorch-lightning
         pip install -q -r .actions/requirements.txt
-        python .actions/assistant.py requirements-prune-pkgs torch,torchvision
+        python .actions/assistant.py requirements_prune_pkgs "[torch,torchvision]"
         pip install ".[extra,test]"
         pip list
       env:

.github/workflows/ci-lite-test-full.yml

@@ -59,7 +59,7 @@ jobs:
     - name: Setup Windows
       if: runner.os == 'windows'
       run: |
-        python .actions/assistant.py requirements_prune_pkgs horovod
+        python .actions/assistant.py requirements_prune_pkgs "[horovod]"
 
     - name: Set min. dependencies
       if: matrix.requires == 'oldest'

.github/workflows/ci-pkg-install.yml

@@ -125,7 +125,7 @@ jobs:
         run: |
           pip install -q -r .actions/requirements.txt
           for pkg in 'app' 'pytorch' ; do
-            python .actions/assistant.py download-package "$pkg" --folder pypi
+            python .actions/assistant.py download_package "$pkg" --folder pypi
           done
           ls -lh pypi/
 
@@ -141,7 +141,7 @@ jobs:
       - name: Miror source
         run: |
           pip install -q -r .actions/requirements.txt
-          python .actions/assistant.py mirror-pkg2source pypi src
+          python .actions/assistant.py mirror_pkg2source pypi src
           ls -R src/
 
       - uses: ./.github/actions/pkg-check

.github/workflows/events-nightly.yml

@@ -27,7 +27,7 @@ jobs:
       - name: Build packages
         run: |
           pip install -q -r .actions/requirements.txt
-          python .actions/assistant.py prepare-nightly-version
+          python .actions/assistant.py prepare_nightly_version
           python setup.py sdist bdist_wheel
           ls -lh dist/
 

.github/workflows/release-pypi.yml

@@ -41,13 +41,13 @@ jobs:
     - name: Find changed packages
       id: candidate
       run: |
-        echo "::set-output name=pkgs::{include: $(python .actions/assistant.py determine-releasing-pkgs 2>&1)}"
+        echo "::set-output name=pkgs::{include: $(python .actions/assistant.py determine_releasing_pkgs 2>&1)}"
     - run: echo "${{ steps.candidate.outputs.pkgs }}"
 
     - name: Inverse packages to pull
       id: download
       run: |
-        echo "::set-output name=pkgs::{include: $(python .actions/assistant.py determine-releasing-pkgs --inverse 2>&1)}"
+        echo "::set-output name=pkgs::{include: $(python .actions/assistant.py determine_releasing_pkgs --inverse 2>&1)}"
     - run: echo "${{ steps.download.outputs.pkgs }}"
 
   # based on https://github.com/pypa/gh-action-pypi-publish
@@ -106,7 +106,7 @@ jobs:
     - name: Dowload package
       run: |
         pip install -q -r .actions/requirements.txt
-        python .actions/assistant.py download-package ${{ matrix.pkg }} --folder pypi
+        python .actions/assistant.py download_package ${{ matrix.pkg }} --folder pypi
 
     - uses: actions/upload-artifact@v3
       with:
@@ -148,7 +148,7 @@ jobs:
     - name: Miror source
       run: |
         pip install -q -r .actions/requirements.txt
-        python .actions/assistant.py mirror-pkg2source pypi src
+        python .actions/assistant.py mirror_pkg2source pypi src
         ls -R src/
 
     - name: Build packages

dockers/base-xla/Dockerfile

@@ -91,13 +91,13 @@ COPY ./ ./pytorch-lightning/
 RUN \
     python --version && \
     cd pytorch-lightning && \
-    pip install -q fire && \
+    pip install -q -r .actions/requirements.txt && \
     # Pin mkl version to avoid OSError on torch import
     # OSError: libmkl_intel_lp64.so.1: cannot open shared object file: No such file or directory
     # https://github.com/pytorch/xla/issues/1666
     pip install mkl==2021.4.0 && \
     # drop packages installed with XLA
-    python .actions/assistant.py requirements_prune_pkgs torch,torchvision && \
+    python .actions/assistant.py requirements_prune_pkgs "[torch,torchvision]" && \
     # drop unnecessary packages
     python ./requirements/pytorch/adjust-versions.py ./requirements/pytorch/extra.txt && \
     # install PL dependencies