nvibrant/hatch_build.py at main · Tremeschin/nvibrant · 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
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
import hashlib
import os
import subprocess
import sys
from pathlib import Path
from tempfile import TemporaryDirectory

import packaging.tags
from hatchling.builders.hooks.plugin.interface import BuildHookInterface

# ---------------------------------------------------------------------------- #

_workdir = TemporaryDirectory(prefix="nvibrant-")

class Dirs:

    repository: Path = Path(__file__).parent
    """Repository root"""

    opengpu: Path = repository.joinpath("open-gpu")
    """Open GPU Kernel Modules submodule"""

    dist: Path = repository.joinpath("dist")
    """Distribution output directory"""

    workdir: Path = Path(_workdir.name)
    """Temporary directory for building"""

    build: Path = workdir.joinpath("build")
    """Meson build directory"""

# ---------------------------------------------------------------------------- #

class BuildHook(BuildHookInterface):
    def initialize(self, version: str, build_data: dict) -> None:

        # Make wheels strictly for the host platform
        # https://packaging.python.org/en/latest/specifications/platform-compatibility-tags/
        for tag in packaging.tags.sys_tags():

            # Skip generic linux not allowed in PyPI
            # https://github.com/pypa/packaging/issues/160
            if tag.platform.startswith("linux_"):
                continue
            if "local" in tag.platform:
                continue

            # Mark broader compatibility than host
            for arch in ("x86_64", "aarch64"):
                if arch not in tag.platform:
                    continue
                if "manylinux" in tag.platform:
                    build_data["tag"] = f"py3-none-manylinux_2_17_{arch}"
                elif "musllinux" in tag.platform:
                    build_data["tag"] = f"py3-none-musllinux_1_1_{arch}"
            break

        build_data["pure_python"] = False

        # Ensure submodule on git main
        subprocess.check_call(
            ("git", "submodule", "update", "--init", "--remote"),
            cwd=Dirs.repository,
        )

        # Configure the project
        subprocess.check_call((
            sys.executable, "-m", "mesonbuild.mesonmain",
            "setup", Dirs.build,
            "--buildtype", "release",
            "--reconfigure", "--wipe",
        ), cwd=Dirs.repository)

        # Intended operation
        subprocess.check_call(
            ("git", "config", "advice.detachedHead", "false"),
            cwd=Dirs.opengpu,
        )

        # Keep track of breaking changes
        hashes: dict[str, str] = dict()

        # Make binaries for all known driver version
        for driver in sorted(subprocess.check_output(
            args=("git", "tag"),
            cwd=Dirs.opengpu
        ).decode().strip().splitlines()):

            # Checkout driver version
            subprocess.check_call(
                ("git", "checkout", "-f", driver),
                cwd=Dirs.opengpu,
            )

            # Compile an executable
            subprocess.check_call((
                sys.executable, "-m", "ninja",
                "-C", Dirs.build,
            ))

            # Find and vendor the binary for this version
            binary = Dirs.build.joinpath("nvibrant")
            target = Dirs.workdir/f"nvibrant-{driver}"
            target.write_bytes(binary.read_bytes())
            target.chmod(0o755)
            binary.unlink()

            # Include in the wheel
            build_data["force_include"][str(target)] = f"nvibrant/resources/{target.name}"
            hashes.setdefault(hashlib.md5(target.read_bytes()).hexdigest(), driver)

        # Revert back main branch
        subprocess.check_call(
            ("git", "checkout", "-f", "main"),
            cwd=Dirs.opengpu
        )

        print("\nBreaking changes across driver versions:")

        for hashsum, driver in hashes.items():
            print(f"• {hashsum}  {driver}")
        print()