pytorch-lightning

1
#!/usr/bin/env python
2
# Copyright The Lightning AI team.
3
#
4
# Licensed under the Apache License, Version 2.0 (the "License");
5
# you may not use this file except in compliance with the License.
6
# You may obtain a copy of the License at
7
#
8
#     http://www.apache.org/licenses/LICENSE-2.0
9
#
10
# Unless required by applicable law or agreed to in writing, software
11
# distributed under the License is distributed on an "AS IS" BASIS,
12
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13
# See the License for the specific language governing permissions and
14
# limitations under the License.
15
"""This is the main and only one setup entry point for installing each package as stand-alone as well as joint
16
installation for all packages.
17

18
There are considered three main scenarios for installing this project:
19

20
1. Using PyPI registry when you can install `pytorch-lightning`, `lightning-app`, etc. or `lightning` for all.
21

22
2. Installation from source code after cloning repository.
23
    In such case we recommend to use command `pip install .` or `pip install -e .` for development version
24
     (development ver. do not copy python files to your pip file system, just create links, so you can edit here)
25
    In case you want to install just one package you need to export env. variable before calling `pip`
26

27
     - for `pytorch-lightning` use `export PACKAGE_NAME=pytorch ; pip install .`
28
     - for `lightning-fabric` use `export PACKAGE_NAME=fabric ; pip install .`
29
     - for `lightning-app` use `export PACKAGE_NAME=app ; pip install .`
30

31
3. Building packages as sdist or binary wheel and installing or publish to PyPI afterwords you use command
32
    `python setup.py sdist` or `python setup.py bdist_wheel` accordingly.
33
   In case you want to build just a particular package you want to set an environment variable:
34
   `PACKAGE_NAME=lightning|pytorch|app|fabric python setup.py sdist|bdist_wheel`
35

36
4. Automated releasing with GitHub action is natural extension of 3) is composed of three consecutive steps:
37
    a) determine which packages shall be released based on version increment in `__version__.py` and eventually
38
     compared against PyPI registry
39
    b) with a parameterization build desired packages in to standard `dist/` folder
40
    c) validate packages and publish to PyPI
41

42
"""
43

44
import contextlib
45
import glob
46
import logging
47
import os
48
import tempfile
49
from importlib.util import module_from_spec, spec_from_file_location
50
from types import ModuleType
51
from typing import Generator, Mapping, Optional
52

53
import setuptools
54
import setuptools.command.egg_info
55

56
_PACKAGE_NAME = os.environ.get("PACKAGE_NAME")
57
_PACKAGE_MAPPING = {
58
    "lightning": "lightning",
59
    "pytorch": "pytorch_lightning",
60
    "app": "lightning_app",
61
    "fabric": "lightning_fabric",
62
}
63
# https://packaging.python.org/guides/single-sourcing-package-version/
64
# http://blog.ionelmc.ro/2014/05/25/python-packaging/
65
_PATH_ROOT = os.path.dirname(__file__)
66
_PATH_SRC = os.path.join(_PATH_ROOT, "src")
67
_PATH_REQUIRE = os.path.join(_PATH_ROOT, "requirements")
68
_FREEZE_REQUIREMENTS = os.environ.get("FREEZE_REQUIREMENTS", "0").lower() in ("1", "true")
69

70

71
def _load_py_module(name: str, location: str) -> ModuleType:
72
    spec = spec_from_file_location(name, location)
73
    assert spec, f"Failed to load module {name} from {location}"
74
    py = module_from_spec(spec)
75
    assert spec.loader, f"ModuleSpec.loader is None for {name} from {location}"
76
    spec.loader.exec_module(py)
77
    return py
78

79

80
def _named_temporary_file(directory: Optional[str] = None) -> str:
81
    # `tempfile.NamedTemporaryFile` has issues in Windows
82
    # https://github.com/deepchem/deepchem/issues/707#issuecomment-556002823
83
    if directory is None:
84
        directory = tempfile.gettempdir()
85
    return os.path.join(directory, os.urandom(24).hex())
86

87

88
@contextlib.contextmanager
89
def _set_manifest_path(manifest_dir: str, aggregate: bool = False, mapping: Mapping = _PACKAGE_MAPPING) -> Generator:
90
    if aggregate:
91
        # aggregate all MANIFEST.in contents into a single temporary file
92
        manifest_path = _named_temporary_file(manifest_dir)
93
        lines = []
94
        # load manifest and aggregated all manifests
95
        for pkg in mapping.values():
96
            pkg_manifest = os.path.join(_PATH_SRC, pkg, "MANIFEST.in")
97
            if os.path.isfile(pkg_manifest):
98
                with open(pkg_manifest) as fh:
99
                    lines.extend(fh.readlines())
100
        # convert lightning_foo to lightning/foo
101
        for new, old in mapping.items():
102
            if old == "lightning":
103
                continue  # avoid `lightning` -> `lightning/lightning`
104
            lines = [ln.replace(old, f"lightning/{new}") for ln in lines]
105
        lines = sorted(set(filter(lambda ln: not ln.strip().startswith("#"), lines)))
106
        logging.debug(f"aggregated manifest consists of: {lines}")
107
        with open(manifest_path, mode="w") as fp:
108
            fp.writelines(lines)
109
    else:
110
        manifest_path = os.path.join(manifest_dir, "MANIFEST.in")
111
        assert os.path.exists(manifest_path)
112
    # avoid error: setup script specifies an absolute path
113
    manifest_path = os.path.relpath(manifest_path, _PATH_ROOT)
114
    logging.info("Set manifest path to", manifest_path)
115
    setuptools.command.egg_info.manifest_maker.template = manifest_path
116
    yield
117
    # cleanup
118
    setuptools.command.egg_info.manifest_maker.template = "MANIFEST.in"
119
    if aggregate:
120
        os.remove(manifest_path)
121

122

123
if __name__ == "__main__":
124
    assistant = _load_py_module(name="assistant", location=os.path.join(_PATH_ROOT, ".actions", "assistant.py"))
125

126
    if os.path.isdir(_PATH_SRC):
127
        # copy the version information to all packages
128
        assistant.distribute_version(_PATH_SRC)
129
    print(f"Requested package: '{_PACKAGE_NAME}'")  # requires `-v` to appear
130

131
    local_pkgs = [
132
        os.path.basename(p)
133
        for p in glob.glob(os.path.join(_PATH_SRC, "*"))
134
        if os.path.isdir(p) and not p.endswith(".egg-info")
135
    ]
136
    print(f"Local package candidates: {local_pkgs}")
137
    is_source_install = len(local_pkgs) > 2
138
    print(f"Installing from source: {is_source_install}")
139
    if is_source_install:
140
        if _PACKAGE_NAME is not None and _PACKAGE_NAME not in _PACKAGE_MAPPING:
141
            raise ValueError(
142
                f"Unexpected package name: {_PACKAGE_NAME}. Possible choices are: {list(_PACKAGE_MAPPING)}"
143
            )
144
        package_to_install = _PACKAGE_MAPPING.get(_PACKAGE_NAME, "lightning")
145
        if package_to_install == "lightning":
146
            # merge all requirements files
147
            assistant._load_aggregate_requirements(_PATH_REQUIRE, _FREEZE_REQUIREMENTS)
148
        else:
149
            # replace imports and copy the code
150
            assistant.create_mirror_package(_PATH_SRC, _PACKAGE_MAPPING)
151
    else:
152
        assert len(local_pkgs) > 0
153
        # PL as a package is distributed together with Fabric, so in such case there are more than one candidate
154
        package_to_install = "pytorch_lightning" if "pytorch_lightning" in local_pkgs else local_pkgs[0]
155
    print(f"Installing package: {package_to_install}")
156

157
    # going to install with `setuptools.setup`
158
    pkg_path = os.path.join(_PATH_SRC, package_to_install)
159
    pkg_setup = os.path.join(pkg_path, "__setup__.py")
160
    if not os.path.exists(pkg_setup):
161
        raise RuntimeError(f"Something's wrong, no package was installed. Package name: {_PACKAGE_NAME}")
162
    setup_module = _load_py_module(name=f"{package_to_install}_setup", location=pkg_setup)
163
    setup_args = setup_module._setup_args()
164
    is_main_pkg = package_to_install == "lightning"
165
    print(f"Installing as the main package: {is_main_pkg}")
166
    if is_source_install:
167
        # we are installing from source, set the correct manifest path
168
        with _set_manifest_path(pkg_path, aggregate=is_main_pkg):
169
            setuptools.setup(**setup_args)
170
    else:
171
        setuptools.setup(**setup_args)
172
    print("Finished setup configuration.")
173