diff --git a/build.py b/build.py index cb3490bf0d..557f55ae77 100755 --- a/build.py +++ b/build.py @@ -972,6 +972,13 @@ def create_dockerfile_buildbase_rhel(ddir, dockerfile_name, argmap): xz-devel \\ zlib-devel """ + if argmap["NVIDIA_BUILD_ID"] is not None: + df += """ +ENV BUILD_NUMBER={} +""".format( + argmap["NVIDIA_BUILD_ID"] + ) + if os.getenv("CCACHE_REMOTE_ONLY") and os.getenv("CCACHE_REMOTE_STORAGE"): df += """ RUN curl -k -s -L https://github.com/ccache/ccache/archive/refs/tags/v4.10.2.tar.gz -o /tmp/ccache.tar.gz \\ @@ -1049,6 +1056,12 @@ def create_dockerfile_buildbase(ddir, dockerfile_name, argmap): ARG TRITON_CONTAINER_VERSION ENV PIP_BREAK_SYSTEM_PACKAGES=1 CMAKE_POLICY_VERSION_MINIMUM=3.5 """ + if argmap["NVIDIA_BUILD_ID"] is not None: + df += """ +ENV BUILD_NUMBER={} +""".format( + argmap["NVIDIA_BUILD_ID"] + ) # Install the windows- or linux-specific buildbase dependencies if target_platform() == "windows": df += """ diff --git a/src/python/build_wheel.py b/src/python/build_wheel.py index 875dd32a70..e3a597c89d 100755 --- a/src/python/build_wheel.py +++ b/src/python/build_wheel.py @@ -1,5 +1,5 @@ #!/usr/bin/env python3 -# Copyright 2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # Redistribution and use in source and binary forms, with or without # modification, are permitted provided that the following conditions @@ -32,7 +32,6 @@ import shutil import subprocess import sys -from distutils.dir_util import copy_tree from tempfile import mkstemp @@ -51,7 +50,7 @@ def touch(path): def cpdir(src, dest): - copy_tree(src, dest, preserve_symlinks=1) + shutil.copytree(src, dest, symlinks=True, dirs_exist_ok=True) def sed(pattern, replace, source, dest=None): @@ -70,6 +69,125 @@ def sed(pattern, replace, source, dest=None): shutil.copyfile(name, source) +def _detect_cuda_version(): + """Detect the CUDA toolkit version visible to the build. + + Prefers the CUDA_VERSION env var (set by official NVIDIA base + images); falls back to parsing /usr/local/cuda/version.json which + is the canonical location for the installed toolkit. Returns the + raw string (e.g. "13.2.1") or None when CUDA is not available. + """ + v = os.environ.get("CUDA_VERSION") + if v: + return v + try: + import json as _json + + with open("/usr/local/cuda/version.json") as f: + data = _json.load(f) + return data.get("cuda", {}).get("version") + except (OSError, ValueError, KeyError): + return None + + +def _compose_version(base_version): + """Compose the full wheel version string. + + Appends a PEP 440 local-version segment describing the NVIDIA + container release and CUDA toolkit so consumers can tell an + nv26.04 wheel from an nv26.05 wheel and a cu132 wheel from a + cu128 wheel. All sources are optional; local non-CI builds return + the version unchanged. + """ + nv = ( + os.environ.get("NVIDIA_UPSTREAM_VERSION") + or os.environ.get("NVIDIA_TRITON_SERVER_VERSION") + or os.environ.get("TRITON_CONTAINER_VERSION") + ) + cuda = _detect_cuda_version() + print( + f"=== Wheel local-version inputs: " + f"NVIDIA_UPSTREAM_VERSION={os.environ.get('NVIDIA_UPSTREAM_VERSION')!r} " + f"NVIDIA_TRITON_SERVER_VERSION={os.environ.get('NVIDIA_TRITON_SERVER_VERSION')!r} " + f"TRITON_CONTAINER_VERSION={os.environ.get('TRITON_CONTAINER_VERSION')!r} " + f"-> nv={nv!r}, cuda={cuda!r}", + file=sys.stderr, + ) + local = [] + if nv: + local.append(f"nv{nv}") + if cuda: + parts = cuda.split(".") + if len(parts) >= 2 and parts[0].isdigit() and parts[1].isdigit(): + local.append(f"cu{parts[0]}{parts[1]}") + if local: + return f"{base_version}+{'.'.join(local)}" + return base_version + + +def _repair_wheel_with_auditwheel(whl_dir, dest_dir): + """Upgrade a linux_ wheel to manylinux_2_X_. + + Ports the pattern established for tritonclient in TRI-286: + 1. auditwheel repair — auto-discovers the minimum manylinux tag + by inspecting glibc symbol requirements of the embedded .so. + 2. python -m wheel tags fallback — used when auditwheel reports + "no ELF" (the wheel has no native extension, e.g. a downstream + build disabled bindings). Mirrors the documented fallback. + 3. No-op with warning — when auditwheel is not installed in the + build image, keep the linux_ wheel as-is so the build + does not regress. + """ + if shutil.which("auditwheel") is None: + print( + "=== WARNING: auditwheel not found on PATH; keeping linux_ " + "wheel as-is. Install auditwheel in the build image to produce " + "PyPI-acceptable manylinux_2_X_ wheels.", + file=sys.stderr, + ) + shutil.copytree(os.path.join(whl_dir, "dist"), dest_dir, dirs_exist_ok=True) + return + + dist_dir = os.path.join(whl_dir, "dist") + wheels = [ + os.path.join(dist_dir, w) for w in os.listdir(dist_dir) if w.endswith(".whl") + ] + fail_if(not wheels, "no wheel produced by the build") + + for wheel_path in wheels: + print(f"=== Running auditwheel repair on {wheel_path}") + r = subprocess.run( + ["auditwheel", "repair", wheel_path, "--wheel-dir", dest_dir], + capture_output=True, + text=True, + ) + if r.returncode != 0 and "no ELF" in r.stderr: + arch = os.uname().machine + manylinux_tag = f"manylinux_2_28_{arch}" + print( + f"=== Pure-Python wheel detected; falling back to wheel tags " + f"({manylinux_tag})" + ) + copied = os.path.join(dest_dir, os.path.basename(wheel_path)) + shutil.copy(wheel_path, copied) + r2 = subprocess.run( + [ + "python3", + "-m", + "wheel", + "tags", + "--platform-tag", + manylinux_tag, + "--remove", + copied, + ] + ) + fail_if(r2.returncode != 0, "wheel tags fallback failed") + elif r.returncode != 0: + sys.stderr.write(r.stderr) + fail_if(True, "auditwheel repair failed") + + def main(): parser = argparse.ArgumentParser() @@ -117,15 +235,37 @@ def main(): os.chdir(FLAGS.whl_dir) print("=== Building wheel") args = ["python3", "setup.py", "bdist_wheel"] + # PEP 427 build tag: lets two wheels of the same version coexist + # (e.g. reruns of the same CI pipeline). Sources, first non-empty + # and usable wins: + # CI_PIPELINE_ID - GitLab pipeline-scoped ID (preferred). + # NVIDIA_BUILD_ID - from build.py's --build-id flag. + # BUILD_NUMBER - generic CI systems. + # PEP 427 requires the build tag to start with a digit. + build_tag = ( + os.environ.get("CI_PIPELINE_ID") + or os.environ.get("NVIDIA_BUILD_ID") + or os.environ.get("BUILD_NUMBER") + ) + print( + f"=== Wheel build-tag inputs: " + f"CI_PIPELINE_ID={os.environ.get('CI_PIPELINE_ID')!r} " + f"NVIDIA_BUILD_ID={os.environ.get('NVIDIA_BUILD_ID')!r} " + f"BUILD_NUMBER={os.environ.get('BUILD_NUMBER')!r} " + f"-> build-tag={build_tag!r}", + file=sys.stderr, + ) + if build_tag and build_tag != "" and build_tag[:1].isdigit(): + args += [f"--build-number={build_tag}"] wenv = os.environ.copy() - wenv["VERSION"] = FLAGS.triton_version + wenv["VERSION"] = _compose_version(FLAGS.triton_version) wenv["TRITON_PYBIND"] = PYBIND_LIB p = subprocess.Popen(args, env=wenv) p.wait() fail_if(p.returncode != 0, "setup.py failed") - cpdir("dist", FLAGS.dest_dir) + _repair_wheel_with_auditwheel(FLAGS.whl_dir, FLAGS.dest_dir) print(f"=== Output wheel file is in: {FLAGS.dest_dir}") touch(os.path.join(FLAGS.dest_dir, "stamp.whl")) diff --git a/src/python/setup.py b/src/python/setup.py index 2c7c12a9ee..12cf7b9afb 100755 --- a/src/python/setup.py +++ b/src/python/setup.py @@ -1,5 +1,5 @@ #!/usr/bin/env python3 -# Copyright 2024-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # Redistribution and use in source and binary forms, with or without # modification, are permitted provided that the following conditions @@ -26,34 +26,23 @@ # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. import os -import sys -from setuptools import find_packages, setup - -if "--plat-name" in sys.argv: - PLATFORM_FLAG = sys.argv[sys.argv.index("--plat-name") + 1] -else: - PLATFORM_FLAG = "any" +from setuptools import Distribution, find_packages, setup if "VERSION" not in os.environ: raise Exception("envvar VERSION must be specified") VERSION = os.environ["VERSION"] -try: - from wheel.bdist_wheel import bdist_wheel as _bdist_wheel - - class bdist_wheel(_bdist_wheel): - def finalize_options(self): - _bdist_wheel.finalize_options(self) - self.root_is_pure = False - def get_tag(self): - pyver, abi, plat = "py3", "none", PLATFORM_FLAG - return pyver, abi, plat +# The wheel ships an arch-specific pybind11 extension bundled via +# package_data. Without has_ext_modules()=True setuptools marks the +# wheel pure-Python (py3-none-any), which auditwheel rejects. +# See TRI-983. +class BinaryDistribution(Distribution): + def has_ext_modules(self): + return True -except ImportError: - bdist_wheel = None this_directory = os.path.abspath(os.path.dirname(__file__)) @@ -105,7 +94,7 @@ def get_tag(self): "": platform_package_data, }, zip_safe=False, - cmdclass={"bdist_wheel": bdist_wheel}, + distclass=BinaryDistribution, data_files=data_files, install_requires=["tritonserver", "pydantic==2.10.6"], extras_require={"GPU": gpu_extras, "test": test_extras, "all": all_extras},