moolib/setup.py at main · facebookresearch/moolib · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
#!/usr/bin/env python
#
# Copyright (c) Facebook, Inc. and its affiliates.
#
# To install: pip install .
#
# For debug builds: python setup.py build --debug install
#
# The environment variable USE_CUDA can be set to "OFF" (or 0).
#

import os
import pathlib
import subprocess
import sys

import setuptools
from setuptools.command import build_ext
from distutils import spawn


class CMakeBuild(build_ext.build_ext):
    def run(self):  # Necessary for pip install -e.
        for ext in self.extensions:
            self.build_extension(ext)

    def build_extension(self, ext):
        source_path = pathlib.Path(__file__).parent.resolve()
        output_path = pathlib.Path(self.get_ext_fullpath(ext.name)).parent.absolute()

        os.makedirs(self.build_temp, exist_ok=True)

        build_type = "Debug" if self.debug else "RelWithDebInfo"

        generator = "Ninja" if spawn.find_executable("ninja") else "Unix Makefiles"

        cmake_cmd = [
            "cmake",
            str(source_path),
            "-G%s" % generator,
            "-DCMAKE_BUILD_TYPE=%s" % build_type,
            "-DCMAKE_LIBRARY_OUTPUT_DIRECTORY=%s" % output_path,
        ]

        use_cuda = os.environ.get("USE_CUDA", True)
        if use_cuda == "OFF":
            use_cuda = False
        if not int(use_cuda):
            cmake_cmd.append("-DUSE_CUDA=OFF")

        build_cmd = ["cmake", "--build", ".", "--parallel"]

        # pip install (but not python setup.py install) runs with a modified PYTHONPATH.
        # This can prevent cmake from finding the torch libraries.
        env = os.environ.copy()
        if "PYTHONPATH" in env:
            del env["PYTHONPATH"]
        try:
            subprocess.check_call(cmake_cmd, cwd=self.build_temp, env=env)
            subprocess.check_call(build_cmd, cwd=self.build_temp, env=env)
        except subprocess.CalledProcessError:
            # Don't obscure the error with a setuptools backtrace.
            sys.exit(1)


def main():
    with open("README.md") as f:
        long_description = f.read()

    setuptools.setup(
        name="moolib",
        version="0.2.0",
        description=("A library for distributed ML training with PyTorch"),
        long_description=long_description,
        long_description_content_type="text/markdown",
        author="tscmoo & the moolib dev team",
        url="https://github.com/facebookresearch/moolib",
        classifiers=[
            "Programming Language :: C++",
            "Programming Language :: Python :: 3",
            "License :: OSI Approved :: MIT License",
            "Operating System :: POSIX :: Linux",
            "Operating System :: MacOS :: MacOS X",
            "Environment :: GPU :: NVIDIA CUDA",
        ],
        packages=["moolib", "moolib.examples.common", "moolib.examples.vtrace"],
        package_dir={"": "py", "moolib.examples": "examples"},
        ext_modules=[setuptools.Extension("moolib._C", sources=[])],
        install_requires=["torch>=1.6.0"],
        cmdclass={"build_ext": CMakeBuild},
        zip_safe=False,
    )


if __name__ == "__main__":
    main()