lit extras (#15793)

Co-authored-by: Carlos Mocholí <[email protected]>
(cherry picked from commit 8ee889b)

Author: Borda

.github/workflows/docs-checks.yml

@@ -76,7 +76,7 @@ jobs:
         # This is needed as App docs is heavily using/referring to lightning package
         if: ${{ matrix.pkg-name == 'app' }}
         run: |
-          pip install -e . -U  -f https://download.pytorch.org/whl/cpu/torch_stable.html -f pypi
+          pip install -e . -U -v  -f https://download.pytorch.org/whl/cpu/torch_stable.html -f pypi
           git checkout -- .
 
       - name: Adjust docs refs

requirements/pytorch/strategies.txt

@@ -6,4 +6,4 @@ fairscale>=0.4.5, <=0.4.6
 deepspeed>=0.6.0, <=0.7.0
 # no need to install with [pytorch] as pytorch is already installed
 horovod>=0.21.2, !=0.24.0, <=0.26.1
-hivemind>=1.0.1, <=1.0.1; sys_platform == 'linux'
+hivemind==1.0.1; sys_platform == 'linux'

src/lightning/__setup__.py

@@ -1,5 +1,7 @@
+import glob
 import os.path
 from importlib.util import module_from_spec, spec_from_file_location
+from pathlib import Path
 from types import ModuleType
 from typing import Any, Dict
 
@@ -8,7 +10,7 @@
 _PROJECT_ROOT = "."
 _SOURCE_ROOT = os.path.join(_PROJECT_ROOT, "src")
 _PACKAGE_ROOT = os.path.join(_SOURCE_ROOT, "lightning")
-_PATH_REQUIREMENTS = os.path.join("requirements")
+_PATH_REQUIREMENTS = os.path.join(_PROJECT_ROOT, "requirements")
 _FREEZE_REQUIREMENTS = bool(int(os.environ.get("FREEZE_REQUIREMENTS", 0)))
 
 
@@ -24,6 +26,30 @@ def _load_py_module(name: str, location: str) -> ModuleType:
 _SETUP_TOOLS = _load_py_module("setup_tools", os.path.join(_PROJECT_ROOT, ".actions", "setup_tools.py"))
 
 
+def _prepare_extras() -> Dict[str, Any]:
+    # https://setuptools.readthedocs.io/en/latest/setuptools.html#declaring-extras
+    # Define package extras. These are only installed if you specify them.
+    # From remote, use like `pip install pytorch-lightning[dev, docs]`
+    # From local copy of repo, use like `pip install ".[dev, docs]"`
+    req_files = [Path(p) for p in glob.glob(os.path.join(_PATH_REQUIREMENTS, "*", "*.txt"))]
+    common_args = dict(unfreeze="major" if _FREEZE_REQUIREMENTS else "all")
+    extras = {
+        f"{p.parent.name}-{p.stem}": _SETUP_TOOLS.load_requirements(file_name=p.name, path_dir=p.parent, **common_args)
+        for p in req_files
+        if p.name not in ("docs.txt", "devel.txt", "base.txt")
+    }
+    for extra in list(extras):
+        name = "-".join(extra.split("-")[1:])
+        extras[name] = extras.get(name, []) + extras[extra]
+    # todo
+    # extras["extra"] = extras["cloud"] + extras["ui"]
+    # extras["dev"] = extras["extra"] + extras["test"]  # + extras['docs']
+    # extras["all"] = extras["dev"]
+    extras = {name: list(set(reqs)) for name, reqs in extras.items()}
+    print("The extras are", extras)
+    return extras
+
+
 def _adjust_manifest(**kwargs: Any) -> None:
     # todo: consider rather aggregation of particular manifest adjustments
     manifest_path = os.path.join(_PROJECT_ROOT, "MANIFEST.in")
@@ -82,7 +108,7 @@ def _setup_args(**kwargs: Any) -> Dict[str, Any]:
         },
         setup_requires=[],
         install_requires=_SETUP_TOOLS.load_requirements(_PATH_REQUIREMENTS, unfreeze="all"),
-        extras_require={},  # todo: consider porting all other packages extras with prefix
+        extras_require=_prepare_extras(),
         project_urls={
             "Bug Tracker": "https://github.com/Lightning-AI/lightning/issues",
             "Documentation": "https://lightning.ai/lightning-docs",