setup.py

import os
import re
import subprocess
import sys

from typing import List
from setuptools import Extension, setup, find_packages
from setuptools.command.build_ext import build_ext

ROOT_DIR = os.path.dirname(__file__)
# A CMakeExtension needs a sourcedir instead of a file list.
# The name must be the _single_ output extension from the CMake build.
# If you need multiple extensions, see scikit-build.
class CMakeExtension(Extension):
    def __init__(self, name, target="C", sourcedir=""):
        Extension.__init__(self, name, sources=[])
        self.target = target
        self.sourcedir = os.path.abspath(sourcedir)


class CMakeBuild(build_ext):
    def build_extension(self, ext):
        extdir = os.path.abspath(os.path.dirname(self.get_ext_fullpath(ext.name)))

        # required for auto-detection & inclusion of auxiliary "native" libs
        if not extdir.endswith(os.path.sep):
            extdir += os.path.sep

        debug = int(os.environ.get("DEBUG", 0)) if self.debug is None else self.debug
        cfg = "RelWithDebInfo" if debug else "Release"
        testing = int(os.environ.get("TESTING", 1))
        testing_cfg = "ON" if testing else "OFF"

        # CMake lets you override the generator - we need to check this.
        # Can be set with Conda-Build, for example.
        cmake_generator = os.environ.get("CMAKE_GENERATOR", "")

        # Set Python_EXECUTABLE instead if you use PYBIND11_FINDPYTHON
        # EXAMPLE_VERSION_INFO shows you how to pass a value into the C++ code
        # from Python.
        cmake_args = [
            f"-DCMAKE_CXX_STANDARD=17",
            f"-DCMAKE_LIBRARY_OUTPUT_DIRECTORY={extdir}",
            f"-DPYTHON_EXECUTABLE={sys.executable}",
            f"-DPYTHON_VERSION={sys.version_info.major}.{sys.version_info.minor}",
            f"-DCMAKE_BUILD_TYPE={cfg}",  # not used on MSVC, but no harm
            f"-DWITH_TESTING={testing_cfg}",
        ]

        # zhilight can be compiled with various versions of g++ and CUDA,
        # only use standard toolchain in packaging container.
        if os.path.exists("/opt/rh/devtoolset-7/root/bin/gcc"):
            cmake_args.extend(
                [
                    f"-DCMAKE_C_COMPILER=/opt/rh/devtoolset-7/root/bin/gcc",
                    f"-DCMAKE_CXX_COMPILER=/opt/rh/devtoolset-7/root/bin/g++",
                    f"-DCMAKE_CUDA_COMPILER=/usr/local/cuda-11.1/bin/nvcc",
                ]
            )
        build_args = [f"--target {ext.target}"]
        # Adding CMake arguments set as environment variable
        # (needed e.g. to build for ARM OSx on conda-forge)
        if "CMAKE_ARGS" in os.environ:
            cmake_args += [item for item in os.environ["CMAKE_ARGS"].split(" ") if item]

        # In this example, we pass in the version to C++. You might not need to.
        cmake_args += [f"-DEXAMPLE_VERSION_INFO={self.distribution.get_version()}"]

        # Using Ninja-build since it a) is available as a wheel and b)
        # multithreads automatically. MSVC would require all variables be
        # exported for Ninja to pick it up, which is a little tricky to do.
        # Users can override the generator with CMAKE_GENERATOR in CMake
        # 3.15+.
        if not cmake_generator or cmake_generator == "Ninja":
            try:
                import ninja  # noqa: F401

                ninja_executable_path = os.path.join(ninja.BIN_DIR, "ninja")
                cmake_args += [
                    "-GNinja",
                    f"-DCMAKE_MAKE_PROGRAM:FILEPATH={ninja_executable_path}",
                ]
            except ImportError:
                pass

        # Set CMAKE_BUILD_PARALLEL_LEVEL to control the parallel build level
        # across all generators.
        if "CMAKE_BUILD_PARALLEL_LEVEL" not in os.environ:
            # self.parallel is a Python 3 only way to set parallel jobs by hand
            # using -j in the build_ext call, not supported by pip or PyPA-build.
            if hasattr(self, "parallel") and self.parallel:
                # CMake 3.12+ only.
                build_args += [f"-j{self.parallel}"]

        build_temp = os.path.join(self.build_temp, ext.name)
        if not os.path.exists(build_temp):
            os.makedirs(build_temp)

        cmake_args += [
            "-DPython_ROOT_DIR=" + os.path.dirname(os.path.dirname(sys.executable))
        ]
        subprocess.check_call(["cmake", ext.sourcedir] + cmake_args, cwd=build_temp)
        subprocess.check_call(["cmake", "--build", "."] + build_args, cwd=build_temp)


ext_modules = [
    CMakeExtension("zhilight.C", "C"),
]

testing = int(os.environ.get("TESTING", 1))

if testing:
    ext_modules.append(CMakeExtension("zhilight.C", "internals_"))

# 移步version.py里去修改__version__
from version import __version__

def get_path(*filepath) -> str:
    return os.path.join(ROOT_DIR, *filepath)

def get_requirements() -> List[str]:
    """Get Python package dependencies from requirements.txt."""

    def _read_requirements(filename: str) -> List[str]:
        with open(get_path(filename)) as f:
            requirements = f.read().strip().split("\n")
        resolved_requirements = []
        for line in requirements:
            if line.startswith("-r "):
                resolved_requirements += _read_requirements(line.split()[1])
            elif line.startswith("--"):
                continue
            else:
                resolved_requirements.append(line)
        return resolved_requirements
    return _read_requirements("requirements.txt")

setup(
    name="zhilight",
    version=__version__,
    author="Zhihu and ModelBest Teams",
    description="Optimized inference engine for llama and similar models",
    long_description="",
    ext_modules=ext_modules,
    cmdclass={"build_ext": CMakeBuild},
    zip_safe=False,
    packages= find_packages(exclude=("tests", )),
    python_requires=">=3.9",
    include_package_data=True,
    install_requires=get_requirements(),
)