|
| 1 | +""" |
| 2 | +Custom build backend that downloads cubins before building the package. |
| 3 | +""" |
| 4 | + |
| 5 | +import os |
| 6 | +import sys |
| 7 | +from pathlib import Path |
| 8 | +from setuptools import build_meta as _orig |
| 9 | +from setuptools.build_meta import * |
| 10 | + |
| 11 | +# Add parent directory to path to import artifacts module |
| 12 | +sys.path.insert(0, str(Path(__file__).parent.parent)) |
| 13 | + |
| 14 | +# add flashinfer._build_meta if not there, it should exist in Path(__file__).parent.parent / "flashinfer" / "_build_meta.py" |
| 15 | +build_meta_file = Path(__file__).parent.parent / "flashinfer" / "_build_meta.py" |
| 16 | +if not build_meta_file.exists(): |
| 17 | + version_file = Path(__file__).parent.parent / "version.txt" |
| 18 | + if version_file.exists(): |
| 19 | + with open(version_file, "r") as f: |
| 20 | + version = f.read().strip() |
| 21 | + with open(build_meta_file, "w") as f: |
| 22 | + f.write('"""Build metadata for flashinfer package."""\n') |
| 23 | + f.write(f'__version__ = "{version}"\n') |
| 24 | + |
| 25 | + |
| 26 | +def _download_cubins(): |
| 27 | + """Download cubins to the source directory before building.""" |
| 28 | + from flashinfer.artifacts import download_artifacts |
| 29 | + |
| 30 | + # Create cubins directory in the source tree |
| 31 | + cubin_dir = Path(__file__).parent / "flashinfer_cubin" / "cubins" |
| 32 | + cubin_dir.mkdir(parents=True, exist_ok=True) |
| 33 | + |
| 34 | + # Set environment variable to download to our package directory |
| 35 | + original_cubin_dir = os.environ.get("FLASHINFER_CUBIN_DIR") |
| 36 | + os.environ["FLASHINFER_CUBIN_DIR"] = str(cubin_dir) |
| 37 | + |
| 38 | + try: |
| 39 | + print(f"Downloading cubins to {cubin_dir}...") |
| 40 | + download_artifacts() |
| 41 | + print(f"Successfully downloaded cubins to {cubin_dir}") |
| 42 | + |
| 43 | + # Count the downloaded files |
| 44 | + cubin_files = list(cubin_dir.rglob("*.cubin")) |
| 45 | + print(f"Downloaded {len(cubin_files)} cubin files") |
| 46 | + |
| 47 | + finally: |
| 48 | + # Restore original environment variable |
| 49 | + if original_cubin_dir: |
| 50 | + os.environ["FLASHINFER_CUBIN_DIR"] = original_cubin_dir |
| 51 | + else: |
| 52 | + os.environ.pop("FLASHINFER_CUBIN_DIR", None) |
| 53 | + |
| 54 | + |
| 55 | +def _create_build_metadata(): |
| 56 | + """Create build metadata file with version information.""" |
| 57 | + version_file = Path(__file__).parent.parent / "version.txt" |
| 58 | + if version_file.exists(): |
| 59 | + with open(version_file, "r") as f: |
| 60 | + version = f.read().strip() |
| 61 | + else: |
| 62 | + version = "0.0.0+unknown" |
| 63 | + |
| 64 | + # Create build metadata in the source tree |
| 65 | + package_dir = Path(__file__).parent / "flashinfer_cubin" |
| 66 | + build_meta_file = package_dir / "_build_meta.py" |
| 67 | + |
| 68 | + with open(build_meta_file, "w") as f: |
| 69 | + f.write('"""Build metadata for flashinfer-cubin package."""\n') |
| 70 | + f.write(f'__version__ = "{version}"\n') |
| 71 | + |
| 72 | + print(f"Created build metadata file with version {version}") |
| 73 | + return version |
| 74 | + |
| 75 | + |
| 76 | +def build_wheel(wheel_directory, config_settings=None, metadata_directory=None): |
| 77 | + """Build a wheel, downloading cubins first.""" |
| 78 | + _download_cubins() |
| 79 | + _create_build_metadata() |
| 80 | + return _orig.build_wheel(wheel_directory, config_settings, metadata_directory) |
| 81 | + |
| 82 | + |
| 83 | +def build_sdist(sdist_directory, config_settings=None): |
| 84 | + """Build a source distribution, downloading cubins first.""" |
| 85 | + _download_cubins() |
| 86 | + _create_build_metadata() |
| 87 | + return _orig.build_sdist(sdist_directory, config_settings) |
| 88 | + |
| 89 | + |
| 90 | +def build_editable(wheel_directory, config_settings=None, metadata_directory=None): |
| 91 | + """Build an editable install, downloading cubins first.""" |
| 92 | + _download_cubins() |
| 93 | + _create_build_metadata() |
| 94 | + return _orig.build_editable(wheel_directory, config_settings, metadata_directory) |
| 95 | + |
| 96 | + |
| 97 | +# Pass through all other hooks |
| 98 | +get_requires_for_build_wheel = _orig.get_requires_for_build_wheel |
| 99 | +get_requires_for_build_sdist = _orig.get_requires_for_build_sdist |
| 100 | +get_requires_for_build_editable = _orig.get_requires_for_build_editable |
| 101 | +prepare_metadata_for_build_wheel = _orig.prepare_metadata_for_build_wheel |
| 102 | +prepare_metadata_for_build_editable = _orig.prepare_metadata_for_build_editable |
0 commit comments