diff --git a/.ci/aarch64_linux/aarch64_wheel_ci_build.py b/.ci/aarch64_linux/aarch64_wheel_ci_build.py index 1cce2836974dc..a182f4d36bd98 100755 --- a/.ci/aarch64_linux/aarch64_wheel_ci_build.py +++ b/.ci/aarch64_linux/aarch64_wheel_ci_build.py @@ -55,9 +55,22 @@ def build_ArmComputeLibrary() -> None: shutil.copytree(f"{acl_checkout_dir}/{d}", f"{acl_install_dir}/{d}") -def update_wheel(wheel_path, desired_cuda) -> None: +def replace_tag(filename) -> None: + with open(filename) as f: + lines = f.readlines() + for i, line in enumerate(lines): + if line.startswith("Tag:"): + lines[i] = line.replace("-linux_", "-manylinux_2_28_") + print(f"Updated tag from {line} to {lines[i]}") + break + + with open(filename, "w") as f: + f.writelines(lines) + + +def package_cuda_wheel(wheel_path, desired_cuda) -> None: """ - Update the cuda wheel libraries + Package the cuda wheel libraries """ folder = os.path.dirname(wheel_path) wheelname = os.path.basename(wheel_path) @@ -88,30 +101,19 @@ def update_wheel(wheel_path, desired_cuda) -> None: "/usr/lib64/libgfortran.so.5", "/acl/build/libarm_compute.so", "/acl/build/libarm_compute_graph.so", + "/usr/local/lib/libnvpl_lapack_lp64_gomp.so.0", + "/usr/local/lib/libnvpl_blas_lp64_gomp.so.0", + "/usr/local/lib/libnvpl_lapack_core.so.0", + "/usr/local/lib/libnvpl_blas_core.so.0", ] - if enable_cuda: - libs_to_copy += [ - "/usr/local/lib/libnvpl_lapack_lp64_gomp.so.0", - "/usr/local/lib/libnvpl_blas_lp64_gomp.so.0", - "/usr/local/lib/libnvpl_lapack_core.so.0", - "/usr/local/lib/libnvpl_blas_core.so.0", - ] - if "126" in desired_cuda: - libs_to_copy += [ - "/usr/local/cuda/lib64/libnvrtc-builtins.so.12.6", - "/usr/local/cuda/lib64/libcufile.so.0", - "/usr/local/cuda/lib64/libcufile_rdma.so.1", - ] - elif "128" in desired_cuda: - libs_to_copy += [ - "/usr/local/cuda/lib64/libnvrtc-builtins.so.12.8", - "/usr/local/cuda/lib64/libcufile.so.0", - "/usr/local/cuda/lib64/libcufile_rdma.so.1", - ] - else: + + if "128" in desired_cuda: libs_to_copy += [ - "/opt/OpenBLAS/lib/libopenblas.so.0", + "/usr/local/cuda/lib64/libnvrtc-builtins.so.12.8", + "/usr/local/cuda/lib64/libcufile.so.0", + "/usr/local/cuda/lib64/libcufile_rdma.so.1", ] + # Copy libraries to unzipped_folder/a/lib for lib_path in libs_to_copy: lib_name = os.path.basename(lib_path) @@ -120,6 +122,13 @@ def update_wheel(wheel_path, desired_cuda) -> None: f"cd {folder}/tmp/torch/lib/; " f"patchelf --set-rpath '$ORIGIN' --force-rpath {folder}/tmp/torch/lib/{lib_name}" ) + + # Make sure the wheel is tagged with manylinux_2_28 + for f in os.scandir(f"{folder}/tmp/"): + if f.is_dir() and f.name.endswith(".dist-info"): + replace_tag(f"{f.path}/WHEEL") + break + os.mkdir(f"{folder}/cuda_wheel") os.system(f"cd {folder}/tmp/; zip -r {folder}/cuda_wheel/{wheelname} *") shutil.move( @@ -242,6 +251,6 @@ def parse_arguments(): print("Updating Cuda Dependency") filename = os.listdir("/pytorch/dist/") wheel_path = f"/pytorch/dist/{filename[0]}" - update_wheel(wheel_path, desired_cuda) + package_cuda_wheel(wheel_path, desired_cuda) pytorch_wheel_name = complete_wheel("/pytorch/") print(f"Build Complete. Created {pytorch_wheel_name}..")