pytorch-lightning

1
import glob
2
import os.path
3
from importlib.util import module_from_spec, spec_from_file_location
4
from pathlib import Path
5
from types import ModuleType
6
from typing import Any, Dict
7

8
from pkg_resources import parse_requirements
9
from setuptools import find_packages
10

11
_PROJECT_ROOT = "."
12
_SOURCE_ROOT = os.path.join(_PROJECT_ROOT, "src")
13
_PACKAGE_ROOT = os.path.join(_SOURCE_ROOT, "pytorch_lightning")
14
_PATH_REQUIREMENTS = os.path.join("requirements", "pytorch")
15
_FREEZE_REQUIREMENTS = os.environ.get("FREEZE_REQUIREMENTS", "0").lower() in ("1", "true")
16

17

18
def _load_py_module(name: str, location: str) -> ModuleType:
19
    spec = spec_from_file_location(name, location)
20
    assert spec, f"Failed to load module {name} from {location}"
21
    py = module_from_spec(spec)
22
    assert spec.loader, f"ModuleSpec.loader is None for {name} from {location}"
23
    spec.loader.exec_module(py)
24
    return py
25

26

27
def _load_assistant() -> ModuleType:
28
    location = os.path.join(_PROJECT_ROOT, ".actions", "assistant.py")
29
    return _load_py_module("assistant", location)
30

31

32
def _prepare_extras() -> Dict[str, Any]:
33
    assistant = _load_assistant()
34
    # https://setuptools.readthedocs.io/en/latest/setuptools.html#declaring-extras
35
    # Define package extras. These are only installed if you specify them.
36
    # From remote, use like `pip install "pytorch-lightning[dev, docs]"`
37
    # From local copy of repo, use like `PACKAGE_NAME=pytorch pip install ".[dev, docs]"`
38
    common_args = {"path_dir": _PATH_REQUIREMENTS, "unfreeze": "none" if _FREEZE_REQUIREMENTS else "all"}
39
    req_files = [Path(p) for p in glob.glob(os.path.join(_PATH_REQUIREMENTS, "*.txt"))]
40
    extras = {
41
        p.stem: assistant.load_requirements(file_name=p.name, **common_args)
42
        for p in req_files
43
        if p.name not in ("docs.txt", "base.txt")
44
    }
45
    for req in parse_requirements(extras["strategies"]):
46
        extras[req.key] = [str(req)]
47
    extras["all"] = extras["extra"] + extras["strategies"] + extras["examples"]
48
    extras["dev"] = extras["all"] + extras["test"]  # + extras['docs']
49
    return extras
50

51

52
def _setup_args() -> Dict[str, Any]:
53
    assistant = _load_assistant()
54
    about = _load_py_module("about", os.path.join(_PACKAGE_ROOT, "__about__.py"))
55
    version = _load_py_module("version", os.path.join(_PACKAGE_ROOT, "__version__.py"))
56
    long_description = assistant.load_readme_description(
57
        _PACKAGE_ROOT, homepage=about.__homepage__, version=version.version
58
    )
59
    return {
60
        "name": "pytorch-lightning",
61
        "version": version.version,
62
        "description": about.__docs__,
63
        "author": about.__author__,
64
        "author_email": about.__author_email__,
65
        "url": about.__homepage__,
66
        "download_url": "https://github.com/Lightning-AI/lightning",
67
        "license": about.__license__,
68
        "packages": find_packages(
69
            where="src",
70
            include=[
71
                "pytorch_lightning",
72
                "pytorch_lightning.*",
73
                "lightning_fabric",
74
                "lightning_fabric.*",
75
            ],
76
        ),
77
        "package_dir": {"": "src"},
78
        "include_package_data": True,
79
        "long_description": long_description,
80
        "long_description_content_type": "text/markdown",
81
        "zip_safe": False,
82
        "keywords": ["deep learning", "pytorch", "AI"],
83
        "python_requires": ">=3.8",
84
        "setup_requires": ["wheel"],
85
        # TODO: aggregate pytorch and lite requirements as we include its source code directly in this package.
86
        # this is not a problem yet because lite's base requirements are all included in pytorch's base requirements
87
        "install_requires": assistant.load_requirements(
88
            _PATH_REQUIREMENTS, unfreeze="none" if _FREEZE_REQUIREMENTS else "all"
89
        ),
90
        "extras_require": _prepare_extras(),
91
        "project_urls": {
92
            "Bug Tracker": "https://github.com/Lightning-AI/lightning/issues",
93
            "Documentation": "https://pytorch-lightning.rtfd.io/en/latest/",
94
            "Source Code": "https://github.com/Lightning-AI/lightning",
95
        },
96
        "classifiers": [
97
            "Environment :: Console",
98
            "Natural Language :: English",
99
            "Development Status :: 5 - Production/Stable",
100
            # Indicate who your project is intended for
101
            "Intended Audience :: Developers",
102
            "Topic :: Scientific/Engineering :: Artificial Intelligence",
103
            "Topic :: Scientific/Engineering :: Image Recognition",
104
            "Topic :: Scientific/Engineering :: Information Analysis",
105
            # Pick your license as you wish
106
            "License :: OSI Approved :: Apache Software License",
107
            "Operating System :: OS Independent",
108
            # Specify the Python versions you support here.
109
            "Programming Language :: Python :: 3",
110
            "Programming Language :: Python :: 3.8",
111
            "Programming Language :: Python :: 3.9",
112
            "Programming Language :: Python :: 3.10",
113
            "Programming Language :: Python :: 3.11",
114
        ],
115
    }
116