|
| 1 | +# -*- coding: utf-8 -*- |
| 2 | +import os |
| 3 | +import sys |
| 4 | +import subprocess |
| 5 | +import textwrap |
| 6 | +import inspect |
| 7 | +from pathlib import Path |
| 8 | +from setuptools import setup, Extension |
| 9 | +from setuptools.command.build_ext import build_ext |
| 10 | +from distutils.dep_util import newer_group |
| 11 | + |
| 12 | +from paddle.utils.cpp_extension import load_op_meta_info_and_register_op |
| 13 | +from paddle.utils.cpp_extension.extension_utils import _jit_compile, _import_module_from_library |
| 14 | +from paddle.utils.cpp_extension.cpp_extension import ( |
| 15 | + CUDA_HOME, CppExtension, BuildExtension as PaddleBuildExtension) |
| 16 | +from paddlenlp.utils.env import PPNLP_HOME |
| 17 | +from paddlenlp.utils.log import logger |
| 18 | + |
| 19 | +if not os.path.exists(CUDA_HOME): |
| 20 | + # CUDA_HOME is only None when `core.is_compiled_with_cuda()` is True in |
| 21 | + # find_cuda_home. Clear it for paddle cpu version. |
| 22 | + CUDA_HOME = None |
| 23 | + |
| 24 | + |
| 25 | +class CMakeExtension(Extension): |
| 26 | + def __init__(self, name, source_dir=None): |
| 27 | + # A CMakeExtension needs a source_dir instead of a file list. |
| 28 | + Extension.__init__(self, name, sources=[]) |
| 29 | + if source_dir is None: |
| 30 | + self.source_dir = Path(__file__).parent.resolve() |
| 31 | + else: |
| 32 | + self.source_dir = os.path.abspath(os.path.expanduser(source_dir)) |
| 33 | + self.sources = [ |
| 34 | + os.path.join(self.source_dir, f) |
| 35 | + for f in os.listdir(self.source_dir) |
| 36 | + ] |
| 37 | + |
| 38 | + def build_with_command(self, ext_builder): |
| 39 | + """ |
| 40 | + Custom `build_ext.build_extension` in `Extension` instead of `Command`. |
| 41 | + `ext_builder` is the instance of `build_ext` command. |
| 42 | + """ |
| 43 | + # refer to https://github.com/pybind/cmake_example/blob/master/setup.py |
| 44 | + if ext_builder.compiler.compiler_type == "msvc": |
| 45 | + raise NotImplementedError |
| 46 | + cmake_args = getattr(self, "cmake_args", []) + [ |
| 47 | + "-DCMAKE_BUILD_TYPE={}".format("Debug" |
| 48 | + if ext_builder.debug else "Release"), |
| 49 | + "-DCMAKE_LIBRARY_OUTPUT_DIRECTORY={}".format(ext_builder.build_lib), |
| 50 | + ] |
| 51 | + build_args = [] |
| 52 | + |
| 53 | + # Set CMAKE_BUILD_PARALLEL_LEVEL to control the parallel build level |
| 54 | + # across all generators. |
| 55 | + if "CMAKE_BUILD_PARALLEL_LEVEL" not in os.environ: |
| 56 | + # self.parallel is a Python 3 only way to set parallel jobs by hand |
| 57 | + # using -j in the build_ext call, not supported by pip or PyPA-build. |
| 58 | + if hasattr(ext_builder, "parallel") and ext_builder.parallel: |
| 59 | + # CMake 3.12+ only. |
| 60 | + build_args += ["-j{}".format(ext_builder.parallel)] |
| 61 | + |
| 62 | + if not os.path.exists(ext_builder.build_temp): |
| 63 | + os.makedirs(ext_builder.build_temp) |
| 64 | + |
| 65 | + # Redirect stdout/stderr to mute, especially when allowing errors |
| 66 | + stdout = getattr(self, "_std_out_handle", None) |
| 67 | + subprocess.check_call( |
| 68 | + ["cmake", self.source_dir] + cmake_args, |
| 69 | + cwd=ext_builder.build_temp, |
| 70 | + stdout=stdout, |
| 71 | + stderr=stdout) |
| 72 | + subprocess.check_call( |
| 73 | + ["cmake", "--build", "."] + build_args, |
| 74 | + cwd=ext_builder.build_temp, |
| 75 | + stdout=stdout, |
| 76 | + stderr=stdout) |
| 77 | + |
| 78 | + def get_target_filename(self): |
| 79 | + raise NotImplementedError |
| 80 | + |
| 81 | + |
| 82 | +class FasterTransformerExtension(CMakeExtension): |
| 83 | + def __init__(self, name, source_dir=None): |
| 84 | + super(FasterTransformerExtension, self).__init__(name, source_dir) |
| 85 | + self._std_out_handle = None |
| 86 | + # Env variable may not work as expected, since jit compile by `load` |
| 87 | + # would not re-built if source code is not update. |
| 88 | + # self.sm = os.environ.get("PPNLP_GENERATE_CODE", None) |
| 89 | + |
| 90 | + def build_with_command(self, ext_builder): |
| 91 | + if CUDA_HOME is None: # GPU only |
| 92 | + # TODO(guosheng): should we touch a dummy file or add a quick exit |
| 93 | + # method to avoid meaningless process in `load` |
| 94 | + logger.warning( |
| 95 | + "FasterTransformer is not available because CUDA can not be found." |
| 96 | + ) |
| 97 | + raise NotImplementedError |
| 98 | + # TODO(guosheng): Multiple -std seems be passed in FasterTransformer, |
| 99 | + # which is not allowed by NVCC. Fix it later. |
| 100 | + self.cmake_args = [f"-DPY_CMD={sys.executable}"] |
| 101 | + # `GetCUDAComputeCapability` is not exposed yet, and detect CUDA/GPU |
| 102 | + # version in cmake file. |
| 103 | + # self.cmake_args += [f"-DSM={self.sm}"] if self.sm is not None else [] |
| 104 | + self.cmake_args = [f"-DWITH_GPT=ON"] |
| 105 | + try: |
| 106 | + super(FasterTransformerExtension, |
| 107 | + self).build_with_command(ext_builder) |
| 108 | + # FasterTransformer cmake file resets `CMAKE_LIBRARY_OUTPUT_DIRECTORY` |
| 109 | + # to `CMAKE_BINARY_DIR/lib`, thus copy the lib back to `build_ext.build_lib`. |
| 110 | + # Maybe move this copy to CMakeList. |
| 111 | + # `copy_tree` or `copy_file`, boost lib might be included |
| 112 | + ext_builder.copy_tree( |
| 113 | + os.path.join(ext_builder.build_temp, "lib"), |
| 114 | + ext_builder.build_lib) |
| 115 | + except Exception as e: |
| 116 | + logger.warning( |
| 117 | + "FasterTransformer is not available due to build errors.") |
| 118 | + raise e |
| 119 | + |
| 120 | + def get_target_filename(self): |
| 121 | + # CMake file has fixed the name of lib, maybe we can copy it as the name |
| 122 | + # returned by `BuildExtension.get_ext_filename` after build. |
| 123 | + return "libdecoding_op.so" |
| 124 | + |
| 125 | + |
| 126 | +class BuildExtension(PaddleBuildExtension): |
| 127 | + """ |
| 128 | + Support both `CppExtention` of Paddle and custom extensions of PaddleNLP. |
| 129 | + """ |
| 130 | + |
| 131 | + def build_extensions(self): |
| 132 | + custom_exts = [] # for |
| 133 | + no_custom_exts = [] # for normal extentions paddle.utils.cpp_extension |
| 134 | + for ext in self.extensions: |
| 135 | + if hasattr(ext, "build_with_command"): |
| 136 | + # custom build in Extension |
| 137 | + ext.build_with_command(self) |
| 138 | + custom_exts.append(ext) |
| 139 | + else: |
| 140 | + no_custom_exts.append(ext) |
| 141 | + if no_custom_exts: |
| 142 | + # Build CppExtentio/CUDAExtension with `PaddleBuildExtension` |
| 143 | + self.extensions = no_custom_exts |
| 144 | + super(BuildExtension, self).build_extensions() |
| 145 | + self.extensions = custom_exts + no_custom_exts |
| 146 | + |
| 147 | + |
| 148 | +EXTENSIONS = {"FasterTransformer": FasterTransformerExtension} |
| 149 | + |
| 150 | + |
| 151 | +def get_extension_maker(name): |
| 152 | + # Use `paddle.utils.cpp_extension.CppExtension` as the default |
| 153 | + # TODO(guosheng): Maybe register extension classes into `Extensions`. |
| 154 | + return EXTENSIONS.get(name, CppExtension) |
| 155 | + |
| 156 | + |
| 157 | +def _write_setup_file(name, file_path, build_dir, **kwargs): |
| 158 | + """ |
| 159 | + Automatically generate setup.py and write it into build directory. |
| 160 | + `kwargws` is arguments for the corresponding Extension initialization. |
| 161 | + Any type extension can be jit build. |
| 162 | + """ |
| 163 | + template = textwrap.dedent(""" |
| 164 | + from setuptools import setup |
| 165 | + from paddlenlp.ops.ext_utils import get_extension_maker, BuildExtension |
| 166 | +
|
| 167 | + setup( |
| 168 | + name='{name}', |
| 169 | + ext_modules=[ |
| 170 | + get_extension_maker('{name}')( |
| 171 | + name='{name}', |
| 172 | + {kwargs_str})], |
| 173 | + cmdclass={{'build_ext' : BuildExtension.with_options( |
| 174 | + output_dir=r'{build_dir}') |
| 175 | + }})""").lstrip() |
| 176 | + kwargs_str = "" |
| 177 | + for key, value in kwargs.items(): |
| 178 | + kwargs_str += key + "=" + (f"'{value}'" |
| 179 | + if isinstance(value, str) else value) + "," |
| 180 | + content = template.format( |
| 181 | + name=name, kwargs_str=kwargs_str, build_dir=build_dir) |
| 182 | + |
| 183 | + with open(file_path, 'w') as f: |
| 184 | + f.write(content) |
| 185 | + |
| 186 | + |
| 187 | +def load(name, build_dir=None, force=False, verbose=False, **kwargs): |
| 188 | + # TODO(guosheng): Need better way to resolve unsupported such as CPU. Currently, |
| 189 | + # raise NotImplementedError and skip `_jit_compile`. Otherwise, `_jit_compile` |
| 190 | + # will output the error to stdout (when verbose is True) and raise `RuntimeError`, |
| 191 | + # which is not friendly for users though no other bad effect. |
| 192 | + if CUDA_HOME is None: |
| 193 | + logger.warning("%s is not available because CUDA can not be found." % |
| 194 | + name) |
| 195 | + raise NotImplementedError |
| 196 | + if build_dir is None: |
| 197 | + build_dir = os.path.join(PPNLP_HOME, 'extenstions') |
| 198 | + build_base_dir = os.path.abspath( |
| 199 | + os.path.expanduser(os.path.join(build_dir, name))) |
| 200 | + if not os.path.exists(build_base_dir): |
| 201 | + os.makedirs(build_base_dir) |
| 202 | + |
| 203 | + extension = get_extension_maker(name)(name, **kwargs) |
| 204 | + # Check if 'target' is out-of-date with respect to any file to avoid rebuild |
| 205 | + if isinstance(extension, CMakeExtension): |
| 206 | + # `CppExtention/CUDAExtension `has version manager by `PaddleBuildExtension` |
| 207 | + # Maybe move this to CMakeExtension later. |
| 208 | + # TODO(guosheng): flags/args changes may also trigger build, and maybe |
| 209 | + # need version manager like `PaddleBuildExtension`. |
| 210 | + ext_filename = extension.get_target_filename() |
| 211 | + ext_filepath = os.path.join(build_base_dir, ext_filename) |
| 212 | + if not force: |
| 213 | + ext_sources = extension.sources |
| 214 | + if os.path.exists(ext_filepath) and not newer_group( |
| 215 | + ext_sources, ext_filepath, 'newer'): |
| 216 | + logger.debug("skipping '%s' extension (up-to-date) build" % |
| 217 | + name) |
| 218 | + return load_op_meta_info_and_register_op(ext_filepath) |
| 219 | + |
| 220 | + # write setup file and jit compile |
| 221 | + file_path = os.path.join(build_dir, "{}_setup.py".format(name)) |
| 222 | + _write_setup_file(name, file_path, build_base_dir, **kwargs) |
| 223 | + _jit_compile(file_path, verbose) |
| 224 | + if isinstance(extension, CMakeExtension): |
| 225 | + # Load a shared library (if exists) only to register op. |
| 226 | + if os.path.exists(ext_filepath): |
| 227 | + load_op_meta_info_and_register_op(ext_filepath) |
| 228 | + else: |
| 229 | + # Import as callable python api |
| 230 | + return _import_module_from_library(name, build_base_dir, verbose) |
0 commit comments