diff --git a/python/build_wheel.py b/python/build_wheel.py index 91f98e9b9..423273890 100755 --- a/python/build_wheel.py +++ b/python/build_wheel.py @@ -32,9 +32,24 @@ import shutil import subprocess import sys +import zipfile from distutils.dir_util import copy_tree from tempfile import mkstemp +# ANSI colors for CI log readability (rendered by GitLab CI, harmlessly +# inert in non-ANSI viewers). Suppressed when stderr is redirected to a +# file by piping through a tool that strips them, or set NO_COLOR=1. +if os.environ.get("NO_COLOR") or not sys.stderr.isatty() and not os.environ.get("CI"): + _GREEN = _YELLOW = _CYAN = _RED = _RESET = "" +else: + _GREEN, _YELLOW, _CYAN, _RED, _RESET = ( + "\033[32m", + "\033[33m", + "\033[36m", + "\033[31m", + "\033[0m", + ) + def fail_if(p, msg): if p: @@ -98,89 +113,137 @@ def _detect_cuda_version() -> str | None: return None -def _compose_version(base_version): - """Compose the full wheel version string. - - The base version comes from TRITON_VERSION and may already include a - PEP 440 pre-release suffix (e.g. "2.69.0.dev0"). Append a PEP 440 - local-version segment describing the NVIDIA container release and - CUDA toolkit the wheel was built against, so consumers can tell an - nv26.04 wheel from an nv26.05 wheel (same upstream Triton version) - and a cu132 wheel from a cu128 wheel. The local-version segment is - primarily for distinguishing these builds; while it does not change - the public upstream version, it can still affect version ordering - and candidate selection among wheels with the same base version. - - Sources for NVIDIA upstream version (first non-empty wins): - NVIDIA_UPSTREAM_VERSION - propagated by build.py via - `docker run -e` from - FLAGS.upstream_container_version. - NVIDIA_TRITON_SERVER_VERSION - set as ENV in the buildbase image - at image-build time from the - TRITON_CONTAINER_VERSION ARG - (survives even if the docker-run - `-e` forwarding is not applied). - TRITON_CONTAINER_VERSION - set as ENV in some downstream - images; same value as above in CI. - Source for CUDA toolkit version: - CUDA_VERSION / toolkit - discovered by _detect_cuda_version() - - All sources are optional; if none is present the version is returned - unchanged so local non-CI builds stay stable. Each detection - outcome is logged to stderr so any future gap is self-announcing - in the build log rather than surfacing only as a missing suffix in - the wheel filename. - """ +def _compose_variant_label(): + """PEP 817 variant label 'nv.cu'. Returns None + if neither input is detectable or the label violates ^[a-z0-9._]{1,16}$.""" 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 = [] + parts = [] if nv: - local.append(f"nv{nv}") + parts.append(f"nv{nv}") if cuda: - # "13.2" / "13.2.0" / "13.2.1" -> "cu132" - 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 + cu = cuda.split(".") + if len(cu) >= 2 and cu[0].isdigit() and cu[1].isdigit(): + parts.append(f"cu{cu[0]}{cu[1]}") + if not parts: + return None + label = ".".join(parts) + if len(label) > 16 or not re.fullmatch(r"[a-z0-9._]+", label): + print( + f"{_RED}=== Variant label {label!r} violates PEP 817; skipping{_RESET}", + file=sys.stderr, + ) + return None + return label -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. +def _normalize_to_highest_manylinux(dist_dir): + """Collapse compressed manylinux tag sets to the highest version. + + auditwheel may emit a wheel whose PEP 425 platform-tag component + contains multiple manylinux entries joined by `.`, e.g. + `manylinux_2_27_x86_64.manylinux_2_28_x86_64`. Per project policy + (TRI-1118), keep only the highest version -- the strictest glibc + baseline -- so each variant's tag reflects its build container's + manylinux policy rather than the binary's symbol-minimum. + + No-op for wheels already carrying a single platform tag. Other + non-manylinux entries in the compressed set are preserved. """ - if shutil.which("auditwheel") is None: + manylinux_re = re.compile(r"^manylinux_(\d+)_(\d+)_(.+)$") + for fname in os.listdir(dist_dir): + if not fname.endswith(".whl"): + continue + parts = fname[:-4].split("-") + if len(parts) < 5: + continue + plat = parts[-1] + if "." not in plat: + continue + tags = plat.split(".") + manylinux_tags = [] + other_tags = [] + for t in tags: + m = manylinux_re.match(t) + if m: + manylinux_tags.append(((int(m.group(1)), int(m.group(2))), t)) + else: + other_tags.append(t) + if len(manylinux_tags) <= 1 and not other_tags: + continue + if not manylinux_tags: + continue + manylinux_tags.sort() + highest = manylinux_tags[-1][1] + new_plat = ".".join([highest] + other_tags) if other_tags else highest + if new_plat == plat: + continue + wheel_path = os.path.join(dist_dir, fname) 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.", + f"{_CYAN}=== Compressed platform tag in {fname!r}: " + f"{plat!r} -> {new_plat!r} (highest manylinux){_RESET}", file=sys.stderr, ) - cpdir("dist", dest_dir) - return + r = subprocess.run( + [ + "python3", + "-m", + "wheel", + "tags", + "--platform-tag", + new_plat, + "--remove", + wheel_path, + ] + ) + fail_if(r.returncode != 0, "wheel tags normalization failed") + + +def _wheel_has_so(wheel_path): + """True if the wheel zip contains a native shared library. + Detects both unversioned (`libfoo.so`) and versioned (`libfoo.so.1.2`) + SONAMEs by checking for `.so` extension or `.so.` infix in any member + basename. ELF inspection is not necessary -- the file-extension test + is what auditwheel and pip both use to classify wheels. + """ + with zipfile.ZipFile(wheel_path) as zf: + for name in zf.namelist(): + base = os.path.basename(name) + if base.endswith(".so") or ".so." in base: + return True + return False + + +def _repair_wheel_with_auditwheel(whl_dir, dest_dir): + """Apply the correct PEP 425 platform-compatibility tag to each wheel. + + Routing rules (per the relevant PEPs): + - Has native `.so` -> PEP 513 / PEP 599 / PEP 600 `manylinux___` + via `auditwheel repair`. auditwheel inspects the .so's glibc + symbol requirements and picks the lowest manylinux policy that + covers them, then bundles any non-allowlisted dynamic deps. + Original linux_ wheel is removed on success. + - No native `.so` -> PEP 425 pure-Python tag `py3-none-any`. + The manylinux platform tag is OMITTED -- claiming manylinux on + a wheel with no glibc-bound code would be a false compatibility + promise. python_tag/abi_tag are also reset to py3/none because + the project has no compiled extensions. + + Notes: + - PEP 656 musllinux is not produced here: build containers are + glibc-based (Ubuntu/RHEL). `auditwheel-musl` would be required + on musl-based distros. + - PEP 440 version normalization is handled by the dev-counter + rewrite in `__main__`, before this function runs. + - If a wheel has a `.so` but `auditwheel` is missing from PATH, + the linux_ wheel is kept as-is and a warning is logged + rather than mis-tagging it as manylinux. + """ 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") @@ -188,42 +251,72 @@ def _repair_wheel_with_auditwheel(whl_dir, dest_dir): 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, - ) - # `auditwheel` logs via Python's logging module, which writes to - # stderr — the "no ELF" sentinel only appears there, not in - # stdout. See TRI-286 root-cause write-up. - if r.returncode != 0 and "no ELF" in r.stderr: - arch = os.uname().machine - manylinux_tag = f"manylinux_2_28_{arch}" + fname = os.path.basename(wheel_path) + # Skip wheels that already carry a manylinux/musllinux platform + # tag. Re-running auditwheel on an already-repaired wheel produces + # a compressed PEP 425 tag set + # (e.g. manylinux_2_27_x86_64.manylinux_2_28_x86_64) -- valid but + # noisy. This guards against CMake invoking this custom command + # twice (build + install phases) and finding stale wheels in dist/. + if "manylinux" in fname or "musllinux" in fname: + print( + f"{_CYAN}=== Skipping already-tagged wheel: {fname}{_RESET}", + file=sys.stderr, + ) + continue + if _wheel_has_so(wheel_path): + if shutil.which("auditwheel") is None: + print( + f"{_RED}=== WARNING: native .so found in " + f"{os.path.basename(wheel_path)} but auditwheel not on " + f"PATH; keeping linux_ wheel as-is. Install " + f"auditwheel in the build image to produce " + f"PyPI-acceptable manylinux wheels (PEP 513/599/600).{_RESET}", + file=sys.stderr, + ) + continue + print( + f"{_CYAN}=== Native extension in {os.path.basename(wheel_path)}: " + f"auditwheel repair -> PEP 513/599/600 manylinux{_RESET}", + file=sys.stderr, + ) + r = subprocess.run( + ["auditwheel", "repair", wheel_path, "--wheel-dir", dist_dir], + capture_output=True, + text=True, + ) + if r.returncode != 0: + sys.stderr.write(r.stderr) + fail_if(True, "auditwheel repair failed") + os.remove(wheel_path) + else: print( - f"=== Pure-Python wheel detected; falling back to wheel tags " - f"({manylinux_tag})" + f"{_CYAN}=== No native extension in " + f"{os.path.basename(wheel_path)}: retagging as PEP 425 " + f"pure-Python (py3-none-any); manylinux tag omitted{_RESET}", + file=sys.stderr, ) - copied = os.path.join(dest_dir, os.path.basename(wheel_path)) - shutil.copy(wheel_path, copied) - # `wheel tags --remove` replaces the linux_ wheel in - # dest_dir with the correctly-tagged manylinux one. - r2 = subprocess.run( + r = subprocess.run( [ "python3", "-m", "wheel", "tags", + "--python-tag", + "py3", + "--abi-tag", + "none", "--platform-tag", - manylinux_tag, + "any", "--remove", - copied, + wheel_path, ] ) - fail_if(r2.returncode != 0, "wheel tags fallback failed") - elif r.returncode != 0: - sys.stderr.write(r.stderr) - fail_if(True, "auditwheel repair failed") + fail_if(r.returncode != 0, "wheel tags retag failed for pure-Python wheel") + + # Post-process: if any resulting wheel carries a compressed manylinux + # tag set, collapse it to the highest version (project policy). + _normalize_to_highest_manylinux(dist_dir) if __name__ == "__main__": @@ -235,12 +328,58 @@ def _repair_wheel_with_auditwheel(whl_dir, dest_dir): parser.add_argument( "--binding-path", type=str, required=True, help="Path to Triton Python binding." ) + parser.add_argument( + "--release-version", + type=str, + required=False, + default=None, + help=( + "Base PEP 440 release version (e.g. '2.70.0'). Overrides the " + "TRITON_RELEASE_VERSION env var and the in-tree TRITON_VERSION file. " + "Precedence: --release-version > TRITON_RELEASE_VERSION > TRITON_VERSION file." + ), + ) FLAGS = parser.parse_args() - FLAGS.triton_version = None - with open("TRITON_VERSION", "r") as vfile: - FLAGS.triton_version = vfile.readline().strip() + # Base release version source — explicit precedence so CI can pin a + # release tag without editing the in-tree TRITON_VERSION file: + # 1. --release-version CLI flag + # 2. TRITON_RELEASE_VERSION env var + # 3. TRITON_VERSION file in CWD (legacy behaviour) + env_release_version = os.environ.get("TRITON_RELEASE_VERSION") + if FLAGS.release_version: + FLAGS.triton_version = FLAGS.release_version + base_source = "--release-version" + elif env_release_version: + FLAGS.triton_version = env_release_version + base_source = "TRITON_RELEASE_VERSION env" + else: + with open("TRITON_VERSION", "r") as vfile: + FLAGS.triton_version = vfile.readline().strip() + base_source = "TRITON_VERSION file" + print( + f"=== Wheel base version: {FLAGS.triton_version!r} (source: {base_source})", + file=sys.stderr, + ) + + # Replace the PEP 440 dev counter with CI_PIPELINE_ID when present, so + # each CI rebuild gets a monotonic, PyPI-uploadable, naturally-sortable + # version (e.g. 2.70.0.dev0 + CI_PIPELINE_ID=12345 -> 2.70.0.dev12345). + # Replaces the legacy PEP 427 build-tag scheme which PyPI rejects and + # which embedded the pipeline id in a separate filename segment. + # Regex tolerates both 2.70.0.dev0 (canonical PEP 440) and 2.71.0dev + # (legacy in-tree shape with no period and no counter). + _pipeline = os.environ.get("CI_PIPELINE_ID", "") + _dev_m = re.match(r"^(\d+\.\d+\.\d+)\.?dev\d*$", FLAGS.triton_version) + if _dev_m and _pipeline.isdigit(): + _new = f"{_dev_m.group(1)}.dev{_pipeline}" + print( + f"{_CYAN}=== PEP 440 dev counter: {FLAGS.triton_version!r} -> " + f"{_new!r} (from CI_PIPELINE_ID={_pipeline}){_RESET}", + file=sys.stderr, + ) + FLAGS.triton_version = _new FLAGS.whl_dir = os.path.join(FLAGS.dest_dir, "wheel") @@ -264,50 +403,37 @@ def _repair_wheel_with_auditwheel(whl_dir, dest_dir): shutil.copyfile("setup.py", os.path.join(FLAGS.whl_dir, "setup.py")) shutil.copyfile("pyproject.toml", os.path.join(FLAGS.whl_dir, "pyproject.toml")) # pyproject.toml resolves the wheel version from the TRITON_VERSION file - # next to it (see [tool.setuptools.dynamic]). Write the *composed* version - # (which appends the +nv…cu… local segment) into the wheel build root so - # that the full version — not just the bare release number — is embedded - # in the wheel filename. Do NOT modify the source-tree TRITON_VERSION. - composed_version = _compose_version(FLAGS.triton_version) + # next to it. Write the chosen version into the wheel build root; do NOT + # modify the source-tree TRITON_VERSION. with open(os.path.join(FLAGS.whl_dir, "TRITON_VERSION"), "w") as vf: - vf.write(composed_version) - print(f"=== Wheel TRITON_VERSION set to: {composed_version!r}", file=sys.stderr) + vf.write(FLAGS.triton_version) os.chdir(FLAGS.whl_dir) + # Clean dist/ to prevent accumulating wheels from prior runs. CMake may + # invoke this custom command twice (build + install phases); without + # this, dist/ would end up with the linux_ wheel just produced + # AND the manylinux___ wheel left over from the previous + # run. _repair_wheel_with_auditwheel would then process both, producing + # wheels with compressed PEP 425 tag sets like + # manylinux_2_26_x86_64.manylinux_2_28_x86_64.whl. + _dist = os.path.join(FLAGS.whl_dir, "dist") + if os.path.isdir(_dist): + shutil.rmtree(_dist) print("=== Building wheel") args = ["python3", "-m", "build"] - # PEP 427 "build tag": an optional segment between version and - # python-tag that 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, matches the - # identifier used in RHEL .zip artifact - # naming (.gitlab-ci.yml). Preferred so all - # wheels in a pipeline share one build tag. - # NVIDIA_BUILD_ID - set from build.py's --build-id flag - # (CI feeds ${CI_JOB_ID}); falls back for - # non-CI builds that pass --build-id. - # BUILD_NUMBER - generic CI systems that set this instead. - # PEP 427 requires the build tag to start with a digit. Skip the - # slot when the value does not satisfy that constraint or is the - # "" default emitted for local builds without --build-id. - # The value is forwarded through `python -m build` to the setuptools - # backend's `bdist_wheel --build=` (alias for --build-number). - build_tag = ( - os.environ.get("CI_PIPELINE_ID") - or os.environ.get("NVIDIA_BUILD_ID") - or os.environ.get("BUILD_NUMBER") - ) + + # Release-semantic X.Y.Z -> PyPI-clean (no variant label). + # Anything else -> PEP 817 variant label. The pipeline id is already + # encoded as the PEP 440 .dev counter above, so no separate + # PEP 427 build tag is needed. + is_release = bool(re.match(r"^\d+\.\d+\.\d+$", FLAGS.triton_version)) 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}", + f"{_GREEN if is_release else _YELLOW}" + f"=== Version {FLAGS.triton_version!r} -> " + f"{'PEP 440 release (PyPI-clean)' if is_release else 'PEP 817 variant'}" + f"{_RESET}", file=sys.stderr, ) - if build_tag and build_tag != "" and build_tag[:1].isdigit(): - args += [f"-C--build-option=--build={build_tag}"] wenv = os.environ.copy() wenv["TRITON_PYBIND"] = PYBIND_LIB @@ -315,13 +441,26 @@ def _repair_wheel_with_auditwheel(whl_dir, dest_dir): p.wait() fail_if(p.returncode != 0, "Building wheel failed") - # Post-process with auditwheel so the wheel is tagged with a proper - # manylinux_2_X_ platform (required by canonical PyPI). When - # auditwheel is unavailable in the build image we keep the - # linux_ wheel and emit a warning; the Poetry/pip lock-file - # problem is already solved by the distinct filename, and the tag can - # be fixed up in a follow-up publish step if needed. _repair_wheel_with_auditwheel(FLAGS.whl_dir, FLAGS.dest_dir) + if not is_release: + label = _compose_variant_label() + if label: + print( + f"{_CYAN}=== PEP 817 variant label: {label!r}{_RESET}", file=sys.stderr + ) + for fname in os.listdir(FLAGS.dest_dir): + if fname.endswith(".whl"): + os.rename( + os.path.join(FLAGS.dest_dir, fname), + os.path.join(FLAGS.dest_dir, fname[:-4] + f"-{label}.whl"), + ) + else: + print( + f"{_RED}=== PEP 817 variant: no nv/cu inputs detected; " + f"wheel emitted unlabeled{_RESET}", + file=sys.stderr, + ) + print("=== Output wheel file is in: {}".format(FLAGS.dest_dir)) touch(os.path.join(FLAGS.dest_dir, "stamp.whl"))