fixing install from src package & CI release (#16027)

* debug install
* hotfix
* local
* refactor
* lit swap
* pruning

.github/actions/pkg-install/action.yml

@@ -33,11 +33,10 @@ runs:
       python -c "import ${{ env.PKG_IMPORT }}; print(${{ env.PKG_IMPORT }}.__version__)"
     shell: bash
 
-# FIXME !!!
-#  - name: Install package - archive
-#    working-directory: ${{ inputs.pkg-folder }}
-#    run: |
-#      pip install ${PKG_SOURCE} ${{ inputs.pip-flags }}
-#      pip list | grep lightning
-#      python -c "import ${{ env.PKG_IMPORT }}; print(${{ env.PKG_IMPORT }}.__version__)"
-#    shell: bash
+  - name: Install package - archive
+    working-directory: ${{ inputs.pkg-folder }}
+    run: |
+      pip install ${PKG_SOURCE} ${{ inputs.pip-flags }}
+      pip list | grep lightning
+      python -c "import ${{ env.PKG_IMPORT }}; print(${{ env.PKG_IMPORT }}.__version__)"
+    shell: bash

.github/workflows/_build-packages.yml

@@ -50,7 +50,8 @@ jobs:
         python-version: 3.9
 
     - run: python -c "print('NB_DIRS=' + str(2 if '${{ matrix.pkg-name }}' == 'pytorch' else 1))" >> $GITHUB_ENV
-    - uses: ./.github/actions/pkg-check
+    - name: Build & check package
+      uses: ./.github/actions/pkg-check
       with:
         pkg-name: ${{ matrix.pkg-name }}
         nb-dirs: ${{ env.NB_DIRS }}

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

@@ -58,7 +58,8 @@ jobs:
 
       - run: |
           python -c "print('PKG_DIR=' + {'notset': 'lightning'}.get('${{matrix.pkg-name}}', '${{matrix.pkg-name}}'))" >> $GITHUB_ENV
-      - uses: ./.github/actions/pkg-install
+      - name: Install package - wheel & archive
+        uses: ./.github/actions/pkg-install
         with:
           pkg-folder: dist/${{ env.PKG_DIR }}
           pkg-name: ${{ matrix.pkg-name }}

.github/workflows/release-pypi.yml

@@ -122,6 +122,7 @@ jobs:
     needs: build-packages
     if: startsWith(github.event.ref, 'refs/tags') || github.event_name == 'release'
     steps:
+    - uses: actions/checkout@v3  # needed for local action bellow
     - uses: actions/download-artifact@v3
       with:
         name: dist-packages-${{ github.sha }}
@@ -152,6 +153,7 @@ jobs:
     needs: [build-packages, waiting]
     if: startsWith(github.event.ref, 'refs/tags') || github.event_name == 'release'
     steps:
+    - uses: actions/checkout@v3  # needed for local action bellow
     - uses: actions/download-artifact@v3
       with:
         name: dist-packages-${{ github.sha }}

setup.py

@@ -40,6 +40,7 @@
     c) validate packages and publish to PyPI
 """
 import contextlib
+import glob
 import os
 import tempfile
 from importlib.util import module_from_spec, spec_from_file_location
@@ -49,7 +50,7 @@
 import setuptools
 import setuptools.command.egg_info
 
-_PACKAGE_NAME = os.environ.get("PACKAGE_NAME", "lightning")
+_PACKAGE_NAME = os.environ.get("PACKAGE_NAME")
 _PACKAGE_MAPPING = {
     "lightning": "lightning",
     "pytorch": "pytorch_lightning",
@@ -87,22 +88,17 @@ def _set_manifest_path(manifest_dir: str, aggregate: bool = False) -> Generator:
         # aggregate all MANIFEST.in contents into a single temporary file
         manifest_path = _named_temporary_file(manifest_dir)
         mapping = _PACKAGE_MAPPING.copy()
-        del mapping["lightning"]
-        lines = ["include src/lightning/version.info\n"]
+        lines = ["include src/lightning/version.info\n", "include requirements/base.txt\n"]
+        # load manifest and aggregated all manifests
         for pkg in mapping.values():
-            pkg_path = os.path.join(_PATH_SRC, pkg)
-            if not os.path.exists(pkg_path):
-                # this function is getting called with `pip install .` and `pip install *.tar.gz`, however, it only
-                # should be called with the former. i haven't found a way to differentiate the two so this is the hacky
-                # solution to avoid an error
-                print(f"{pkg_path!r} does not exist")
-                yield
-                return
-            with open(os.path.join(pkg_path, "MANIFEST.in")) as fh:
-                lines.extend(fh.readlines())
+            pkg_manifest = os.path.join(_PATH_SRC, pkg, "MANIFEST.in")
+            if os.path.isfile(pkg_manifest):
+                with open(pkg_manifest) as fh:
+                    lines.extend(fh.readlines())
         # convert lightning_foo to lightning/foo
         for new, old in mapping.items():
-            lines = [line.replace(old, f"lightning/{new}") for line in lines]
+            lines += [ln.replace(old, f"lightning/{new}") for ln in lines]
+        lines = sorted(set(filter(lambda ln: not ln.strip().startswith("#"), lines)))
         with open(manifest_path, mode="w") as fp:
             fp.writelines(lines)
     else:
@@ -122,41 +118,49 @@ def _set_manifest_path(manifest_dir: str, aggregate: bool = False) -> Generator:
 if __name__ == "__main__":
     assistant = _load_py_module(name="assistant", location=os.path.join(_PATH_ROOT, ".actions", "assistant.py"))
 
-    if os.path.exists(_PATH_SRC):
+    if os.path.isdir(_PATH_SRC):
         # copy the version information to all packages
         assistant.distribute_version(_PATH_SRC)
-
-    package_to_install = _PACKAGE_NAME or "lightning"
-    print(f"Installing the {package_to_install} package")  # requires `-v` to appear
-    is_wheel_install = "PEP517_BUILD_BACKEND" in os.environ
-    print("is_wheel_install:", is_wheel_install)
-    if package_to_install not in _PACKAGE_MAPPING or (not is_wheel_install and _PACKAGE_NAME is None):
-        raise ValueError(f"Unexpected package name: {_PACKAGE_NAME}. Possible choices are: {list(_PACKAGE_MAPPING)}")
-    is_wheel_install &= _PACKAGE_NAME is None
-
-    if package_to_install == "lightning":  # install everything
-        # merge all requirements files
-        assistant._load_aggregate_requirements(_PATH_REQUIRE, _FREEZE_REQUIREMENTS)
-        # replace imports and copy the code
-        assistant.create_mirror_package(_PATH_SRC, _PACKAGE_MAPPING)
-
-    # if it's a wheel install (hence _PACKAGE_NAME should not be set), iterate over all possible packages until we find
-    # one that can be installed. the wheel should have included only the relevant files of the package to install
-    possible_packages = _PACKAGE_MAPPING.values() if is_wheel_install else [_PACKAGE_MAPPING[_PACKAGE_NAME]]
-    for pkg in possible_packages:
-        pkg_path = os.path.join(_PATH_SRC, pkg)
-        pkg_setup = os.path.join(pkg_path, "__setup__.py")
-        if os.path.exists(pkg_setup):
-            print(f"{pkg_setup} exists. Running `setuptools.setup`")
-            setup_module = _load_py_module(name=f"{pkg}_setup", location=pkg_setup)
-            setup_args = setup_module._setup_args()
-            is_main_pkg = pkg == "lightning"
-            if is_wheel_install and not is_main_pkg:
-                setuptools.setup(**setup_args)
-            else:
-                # we are installing from source, set the correct manifest path
-                with _set_manifest_path(pkg_path, aggregate=is_main_pkg):
-                    setuptools.setup(**setup_args)
-            break
+    print(f"Requested package: '{_PACKAGE_NAME}'")  # requires `-v` to appear
+
+    local_pkgs = [
+        os.path.basename(p)
+        for p in glob.glob(os.path.join(_PATH_SRC, "*"))
+        if os.path.isdir(p) and not p.endswith(".egg-info")
+    ]
+    print(f"Local package candidates: {local_pkgs}")
+    is_source_install = len(local_pkgs) > 2
+    print(f"Installing from source: {is_source_install}")
+    if is_source_install:
+        if _PACKAGE_NAME is not None and _PACKAGE_NAME not in _PACKAGE_MAPPING:
+            raise ValueError(
+                f"Unexpected package name: {_PACKAGE_NAME}. Possible choices are: {list(_PACKAGE_MAPPING)}"
+            )
+        package_to_install = _PACKAGE_MAPPING.get(_PACKAGE_NAME, "lightning")
+        if package_to_install == "lightning":  # install everything
+            # merge all requirements files
+            assistant._load_aggregate_requirements(_PATH_REQUIRE, _FREEZE_REQUIREMENTS)
+            # replace imports and copy the code
+            assistant.create_mirror_package(_PATH_SRC, _PACKAGE_MAPPING)
     else:
+        assert len(local_pkgs) > 0
+        # PL as a package is distributed together with Lite, so in such case there are more than one candidate
+        package_to_install = "pytorch_lightning" if "pytorch_lightning" in local_pkgs else local_pkgs[0]
+    print(f"Installing package: {package_to_install}")
+
+    # going to install with `setuptools.setup`
+    pkg_path = os.path.join(_PATH_SRC, package_to_install)
+    pkg_setup = os.path.join(pkg_path, "__setup__.py")
+    if not os.path.exists(pkg_setup):
         raise RuntimeError(f"Something's wrong, no package was installed. Package name: {_PACKAGE_NAME}")
+    setup_module = _load_py_module(name=f"{package_to_install}_setup", location=pkg_setup)
+    setup_args = setup_module._setup_args()
+    is_main_pkg = package_to_install == "lightning"
+    print(f"Installing as the main package: {is_main_pkg}")
+    if is_source_install:
+        # we are installing from source, set the correct manifest path
+        with _set_manifest_path(pkg_path, aggregate=is_main_pkg):
+            setuptools.setup(**setup_args)
+    else:
+        setuptools.setup(**setup_args)
+    print("Finished setup configuration.")