Optimize building and packaging (#53)

ggeorgakoudis · web-flow · commit b5c38e33596c · 2026-02-13T13:43:44.000-08:00
* Remove custom bdist_wheel and prepare once OpenMP sources

* Optimize openmp source extraction and building
  - Extract and apply patches once per build
  - Remove symlinks
  - Change pyomp to load the versioned libomptarget library

* Clean up packaged files
  - Remove unnecessary manifest
  - Do not include package data, wheel builds will collect binaries from
the build dir

* Update conda test for correct filename
diff --git a/MANIFEST.in b/MANIFEST.in
diff --git a/buildscripts/conda-recipes/pyomp/meta.yaml b/buildscripts/conda-recipes/pyomp/meta.yaml
@@ -1,5 +1,7 @@
 {% set version = environ.get('GIT_DESCRIBE_TAG', '0.0.0').lstrip('v') %}
 {% set LLVM_VERSION = environ.get('LLVM_VERSION', '20.1.8') %}
+{% set LLVM_VERSION_MAJOR = LLVM_VERSION.split('.')[0] %}
+{% set LLVM_VERSION_MINOR = LLVM_VERSION.split('.')[1] %}
 
 package:
   name: pyomp
@@ -64,11 +66,11 @@ requirements:
 
 test:
   commands:
-    - test -f $SP_DIR/numba/openmp/libs/pass/libIntrinsicsOpenMP.dylib                 # [osx]
-    - test -f $SP_DIR/numba/openmp/libs/pass/libIntrinsicsOpenMP.so                    # [linux]
-    - test -f $SP_DIR/numba/openmp/libs/openmp/lib/libomptarget-amdgpu.bc        # [linux]
-    - test -f $SP_DIR/numba/openmp/libs/openmp/lib/libomptarget-nvptx.bc         # [linux]
-    - test -f $SP_DIR/numba/openmp/libs/openmp/lib/libomptarget.so               # [linux]
+    - test -f $SP_DIR/numba/openmp/libs/pass/libIntrinsicsOpenMP.dylib                                                # [osx]
+    - test -f $SP_DIR/numba/openmp/libs/pass/libIntrinsicsOpenMP.so                                                   # [linux]
+    - test -f $SP_DIR/numba/openmp/libs/openmp/lib/libomptarget-amdgpu.bc                                             # [linux]
+    - test -f $SP_DIR/numba/openmp/libs/openmp/lib/libomptarget-nvptx.bc                                              # [linux]
+    - test -f $SP_DIR/numba/openmp/libs/openmp/lib/libomptarget.so.{{ LLVM_VERSION_MAJOR }}.{{ LLVM_VERSION_MINOR }}  # [linux]
 
 about:
   home: https://github.com/Python-for-HPC/PyOMP
diff --git a/pyproject.toml b/pyproject.toml
@@ -27,17 +27,12 @@ Homepage = "https://github.com/Python-for-HPC/PyOMP"
 Issues = "https://github.com/Python-for-HPC/PyOMP/issues"
 
 [tool.setuptools]
-include-package-data = true
+include-package-data = false
 package-dir = { "" = "src" }
 
-# Use discovery for the numba.* namespace.
 [tool.setuptools.packages.find]
 where = ["src"]
-include = ["numba.openmp*"]
-
-# Bundle the CMake-installed artifacts into the wheel.
-[tool.setuptools.package-data]
-"numba.openmp.libs" = ["pass/*", "openmp/**/*"]
+include = ["numba*"]
 
 # setuptools-scm config
 [tool.setuptools_scm]
diff --git a/setup.py b/setup.py
@@ -8,7 +8,6 @@
 from setuptools import setup, Extension
 from setuptools import Command
 from setuptools.command.build_ext import build_ext
-from setuptools.command.bdist_wheel import bdist_wheel
 
 
 class CleanCommand(Command):
@@ -29,34 +28,21 @@ def run(self):
             shutil.rmtree(egg_info, ignore_errors=True)
 
 
-class CustomBdistWheel(bdist_wheel):
-    def run(self):
-        # Ensure all build steps are run before bdist_wheel
-        self.run_command("build_ext")
-        super().run()
-
-
 class CMakeExtension(Extension):
     def __init__(
         self,
         name,
         *,
+        setup=None,
         source_dir=None,
         install_dir=None,
-        url=None,
-        sha256=None,
         cmake_args=[],
     ):
         # Don't invoke the original build_ext for this special extension.
         super().__init__(name, sources=[])
-        if source_dir and url:
-            raise ValueError(
-                "CMakeExtension should have either a source_dir or a url, not both."
-            )
+        self.setup = setup
         self.source_dir = source_dir
         self.install_dir = install_dir
-        self.url = url
-        self.sha256 = sha256
         self.cmake_args = cmake_args
 
 
@@ -81,9 +67,17 @@ def _build_cmake(self, ext: CMakeExtension):
         shutil.rmtree(build_dir, ignore_errors=True)
         build_dir.mkdir(parents=True, exist_ok=True)
 
-        lib_dir = Path(
-            self.get_finalized_command("build_py").get_package_dir("numba.openmp.libs")
-        )
+        if ext.setup:
+            ext.setup.setup()
+
+        if self.inplace:
+            lib_dir = Path(
+                self.get_finalized_command("build_py").get_package_dir(
+                    "numba.openmp.libs"
+                )
+            )
+        else:
+            lib_dir = Path(self.build_lib) / "numba/openmp/libs"
 
         extra_cmake_args = self._env_toolchain_args(ext)
         # Set RPATH.
@@ -135,6 +129,10 @@ def _build_cmake(self, ext: CMakeExtension):
             include_dir = install_dir / "lib/cmake"
             if include_dir.exists():
                 shutil.rmtree(include_dir)
+        # Remove symlinks in the install directory to avoid copies.
+        for file in install_dir.rglob("*"):
+            if file.is_symlink():
+                file.unlink()
 
     def _env_toolchain_args(self, ext):
         args = []
@@ -149,79 +147,93 @@ def _env_toolchain_args(self, ext):
         return args
 
 
-def _prepare_source_openmp(sha256=None):
+class PrepareOpenMP:
+    setup_done = False
     LLVM_VERSION = os.environ.get("LLVM_VERSION", None)
-    assert LLVM_VERSION is not None, "LLVM_VERSION environment variable must be set."
-    url = f"https://github.com/llvm/llvm-project/releases/download/llvmorg-{LLVM_VERSION}/llvm-project-{LLVM_VERSION}.src.tar.xz"
-
-    tmp = Path("_downloads/libomp") / f"llvm-project-{LLVM_VERSION}.tar.gz"
-    tmp.parent.mkdir(parents=True, exist_ok=True)
-
-    # Download the source tarball if it does not exist.
-    if not tmp.exists():
-        print(f"Downloading llvm-project version {LLVM_VERSION} url:", url)
-        with urllib.request.urlopen(url) as r:
-            with tmp.open("wb") as f:
-                f.write(r.read())
-    else:
-        print(f"Using downloaded llvm-project at {tmp}")
-
-    if sha256:
-        import hashlib
-
-        hasher = hashlib.sha256()
-        with tmp.open("rb") as f:
-            hasher.update(f.read())
-        if hasher.hexdigest() != sha256:
-            raise ValueError(f"SHA256 mismatch for {url}")
-
-    print("Extracting llvm-project...")
-    with tarfile.open(tmp) as tf:
-        # The root dir llvm-project-20.1.8.src
-        root_name = tf.getnames()[0]
-
-        # Extract only needed subdirectories
-        members = [
-            m
-            for m in tf.getmembers()
-            if m.name.startswith(f"{root_name}/openmp/")
-            or m.name.startswith(f"{root_name}/offload/")
-            or m.name.startswith(f"{root_name}/runtimes/")
-            or m.name.startswith(f"{root_name}/cmake/")
-            or m.name.startswith(f"{root_name}/llvm/cmake/")
-            or m.name.startswith(f"{root_name}/llvm/utils/")
-            or m.name.startswith(f"{root_name}/libc/")
-        ]
-
-        parentdir = tmp.parent
-        # Base arguments for extractall.
-        kwargs = {"path": parentdir, "members": members}
-
-        # Check if data filter is available.
-        if hasattr(tarfile, "data_filter"):
-            # If this exists, the 'filter' argument is guaranteed to work
-            kwargs["filter"] = "data"
-
-        tf.extractall(**kwargs)
-
-        source_dir = parentdir / root_name
-        print("Extracted llvm-project to:", source_dir)
-
-    print("Applying patches to llvm-project...")
-    for patch in sorted(
-        Path(f"src/numba/openmp/libs/openmp/patches/{LLVM_VERSION}")
-        .absolute()
-        .glob("*.patch")
-    ):
-        print("applying patch", patch)
-        subprocess.run(
-            ["patch", "-p1", "-i", str(patch)],
-            cwd=source_dir,
-            check=True,
-            stdin=subprocess.DEVNULL,
+
+    @classmethod
+    def setup(cls):
+        if not cls.setup_done:
+            cls._prepare_source_openmp()
+            cls.setup_done = True
+
+    @classmethod
+    def get_source_dir(cls):
+        return Path(
+            f"_downloads/libomp/llvm-project-{cls.LLVM_VERSION}.src/runtimes"
+        ).absolute()
+
+    @classmethod
+    def _prepare_source_openmp(cls):
+        assert cls.LLVM_VERSION is not None, (
+            "LLVM_VERSION environment variable must be set."
         )
+        url = f"https://github.com/llvm/llvm-project/releases/download/llvmorg-{cls.LLVM_VERSION}/llvm-project-{cls.LLVM_VERSION}.src.tar.xz"
 
-    return f"{source_dir}/runtimes"
+        tmp = Path("_downloads/libomp") / f"llvm-project-{cls.LLVM_VERSION}.tar.gz"
+        tmp.parent.mkdir(parents=True, exist_ok=True)
+
+        # Download the source tarball if it does not exist.
+        if not tmp.exists():
+            print(
+                f"Downloading llvm-project version {cls.LLVM_VERSION} url:",
+                url,
+                flush=True,
+            )
+            with urllib.request.urlopen(url) as r:
+                with tmp.open("wb") as f:
+                    f.write(r.read())
+        else:
+            print(f"Using downloaded llvm-project at {tmp}", flush=True)
+
+        print("Extracting llvm-project...", flush=True)
+        root_dir = Path(f"_downloads/libomp/llvm-project-{cls.LLVM_VERSION}.src")
+        if not root_dir.exists():
+            with tarfile.open(tmp) as tf:
+                # Extract only needed subdirectories
+                root_name = f"llvm-project-{cls.LLVM_VERSION}.src"
+
+                # Use prefix tuple + any() for faster membership testing
+                prefixes = (
+                    f"{root_name}/openmp/",
+                    f"{root_name}/offload/",
+                    f"{root_name}/runtimes/",
+                    f"{root_name}/cmake/",
+                    f"{root_name}/llvm/cmake/",
+                    f"{root_name}/llvm/utils/",
+                    f"{root_name}/libc/",
+                )
+
+                include_members = [
+                    m
+                    for m in tf.getmembers()
+                    if any(m.name.startswith(p) for p in prefixes)
+                ]
+
+                parentdir = tmp.parent
+                # Base arguments for extractall.
+                kwargs = {"path": parentdir, "members": include_members}
+
+                # Check if data filter is available.
+                if hasattr(tarfile, "data_filter"):
+                    # If this exists, the 'filter' argument is guaranteed to work
+                    kwargs["filter"] = "data"
+
+                tf.extractall(**kwargs)
+
+            print("Applying patches to llvm-project...", flush=True)
+            for patch in sorted(
+                Path(f"src/numba/openmp/libs/openmp/patches/{cls.LLVM_VERSION}")
+                .absolute()
+                .glob("*.patch")
+            ):
+                print("applying patch", patch, flush=True)
+                subprocess.run(
+                    ["patch", "-p1", "-N", "-i", str(patch)],
+                    cwd=root_dir,
+                    check=True,
+                    stdin=subprocess.DEVNULL,
+                )
 
 
 def _check_true(env_var):
@@ -234,18 +246,15 @@ def _check_true(env_var):
 ext_modules = [CMakeExtension("pass", source_dir="src/numba/openmp/libs/pass")]
 
 
-# Prepare source directory if either bundled libomp or libomptarget is enabled.
-if _check_true("ENABLE_BUNDLED_LIBOMP") or _check_true("ENABLE_BUNDLED_LIBOMPTARGET"):
-    openmp_source_dir = _prepare_source_openmp()
-
 # Optionally enable bundled libomp build via ENABLE_BUNDLED_LIBOMP=1.  We want
 # to avoid bundling for conda builds to avoid duplicate OpenMP runtime conflicts
 # (e.g., numba 0.62+ and libopenblas already require llvm-openmp).
 if _check_true("ENABLE_BUNDLED_LIBOMP"):
     ext_modules.append(
         CMakeExtension(
             "libomp",
-            source_dir=openmp_source_dir,
+            setup=PrepareOpenMP,
+            source_dir=PrepareOpenMP.get_source_dir(),
             install_dir="openmp",
             cmake_args=[
                 "-DOPENMP_STANDALONE_BUILD=ON",
@@ -265,7 +274,8 @@ def _check_true(env_var):
     ext_modules.append(
         CMakeExtension(
             "libomptarget",
-            source_dir=openmp_source_dir,
+            setup=PrepareOpenMP,
+            source_dir=PrepareOpenMP.get_source_dir(),
             install_dir="openmp",
             cmake_args=[
                 "-DOPENMP_STANDALONE_BUILD=ON",
@@ -283,6 +293,5 @@ def _check_true(env_var):
     cmdclass={
         "clean": CleanCommand,
         "build_ext": BuildCMakeExt,
-        **({"bdist_wheel": CustomBdistWheel} if CustomBdistWheel else {}),
     },
 )
diff --git a/src/numba/openmp/__init__.py b/src/numba/openmp/__init__.py
@@ -92,7 +92,10 @@ def _init():
     if sys_platform.startswith("darwin") or sys_platform.startswith("win32"):
         return
 
-    omptargetlib = libpath / "openmp" / "lib" / "libomptarget.so"
+    llvm_major, llvm_minor, _ = ll.llvm_version_info
+    omptargetlib = (
+        libpath / "openmp" / "lib" / f"libomptarget.so.{llvm_major}.{llvm_minor}"
+    )
     if omptargetlib.exists():
         if DEBUG_OPENMP >= 1:
             print("Found OpenMP target runtime library at", omptargetlib)
