Fix RPATH of the _k2 Python module

[csukuangfj]

No description provided.

[gemini-code-assist]

Code Review

This pull request updates the GitHub Actions build scripts to restrict setuptools versions for PyTorch < 2.0, adds PyTorch 2.12.1 to the build matrix, and introduces a step to patch the RPATH of the _k2 shared library inside the built wheels using patchelf. The review feedback highlights several robustness issues in the RPATH patching logic across the build scripts—such as fragile path slicing, unsafe bash interpolation, and potential failures when locating the shared library or handling multiple wheels—and provides a more robust implementation using a loop, find, and os.path.relpath.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

The current implementation of patching the RPATH has several potential issues and robustness concerns:

1. Multiple Wheels / No Wheels: Using whl=$(ls *.whl) can fail or behave unexpectedly if there are multiple wheels or no wheels in the directory. Using a for loop with a file existence check is much safer.
2. Finding the Shared Library: Using ls _k2.cpython-*.so assumes the .so file is at the root of the unzipped directory. If it is located in a subdirectory (e.g., k2/), ls will fail. Using find is much more robust.
3. Bash Interpolation in Python: Passing $whl directly inside the double-quoted Python command string can lead to syntax or execution errors if the filename contains special characters or spaces. Passing it via sys.argv is the standard and safe way.
4. Path Slicing Hack: Using path[2:] to strip the leading ./ is a fragile hack. Using os.path.relpath is the standard, cross-platform, and robust way to get the relative path for the zip archive.

pushd /var/www/wheelhouse
for whl in *.whl; do
  [ -f "$whl" ] || continue
  _tmp_whl="_tmp_${whl%.whl}"
  mkdir -p "$_tmp_whl"
  pushd "$_tmp_whl"
  unzip -o "../$whl"
  so_file=$(find . -name "_k2.cpython-*.so" -print -quit)
  if [ -n "$so_file" ]; then
    echo "Patching rpath for $so_file"
    current_rpath=$(patchelf --print-rpath "$so_file")
    echo "Current rpath: $current_rpath"
    new_rpath="\$ORIGIN/nvidia/nvtx/lib:\$ORIGIN/nvidia/cuda_runtime/lib:\$ORIGIN/nvidia/cuda_nvrtc/lib:${current_rpath}"
    echo "New rpath: $new_rpath"
    patchelf --set-rpath "$new_rpath" "$so_file"
    echo "Verified rpath:"
    patchelf --print-rpath "$so_file"
    python3 -c "
import zipfile, os, sys
whl = sys.argv[1]
with zipfile.ZipFile(os.path.join('..', whl), 'w', zipfile.ZIP_DEFLATED) as zf:
    for root, dirs, files in os.walk('.'):
        for f in files:
            path = os.path.join(root, f)
            zf.write(path, os.path.relpath(path, '.'))
" "$whl"
  else
    echo "Warning: No _k2 shared library found in $whl"
  fi
  popd
  rm -rf "$_tmp_whl"
done
popd

[gemini-code-assist]

The current implementation of patching the RPATH has several potential issues and robustness concerns:

1. Multiple Wheels / No Wheels: Using whl=$(ls *.whl) can fail or behave unexpectedly if there are multiple wheels or no wheels in the directory. Using a for loop with a file existence check is much safer.
2. Finding the Shared Library: Using ls _k2.cpython-*.so assumes the .so file is at the root of the unzipped directory. If it is located in a subdirectory (e.g., k2/), ls will fail. Using find is much more robust.
3. Bash Interpolation in Python: Passing $whl directly inside the double-quoted Python command string can lead to syntax or execution errors if the filename contains special characters or spaces. Passing it via sys.argv is the standard and safe way.
4. Path Slicing Hack: Using path[2:] to strip the leading ./ is a fragile hack. Using os.path.relpath is the standard, cross-platform, and robust way to get the relative path for the zip archive.

pushd /var/www/wheelhouse
for whl in *.whl; do
  [ -f "$whl" ] || continue
  _tmp_whl="_tmp_${whl%.whl}"
  mkdir -p "$_tmp_whl"
  pushd "$_tmp_whl"
  unzip -o "../$whl"
  so_file=$(find . -name "_k2.cpython-*.so" -print -quit)
  if [ -n "$so_file" ]; then
    echo "Patching rpath for $so_file"
    current_rpath=$(patchelf --print-rpath "$so_file")
    echo "Current rpath: $current_rpath"
    new_rpath="\$ORIGIN/nvidia/nvtx/lib:\$ORIGIN/nvidia/cuda_runtime/lib:\$ORIGIN/nvidia/cuda_nvrtc/lib:${current_rpath}"
    echo "New rpath: $new_rpath"
    patchelf --set-rpath "$new_rpath" "$so_file"
    echo "Verified rpath:"
    patchelf --print-rpath "$so_file"
    python3 -c "
import zipfile, os, sys
whl = sys.argv[1]
with zipfile.ZipFile(os.path.join('..', whl), 'w', zipfile.ZIP_DEFLATED) as zf:
    for root, dirs, files in os.walk('.'):
        for f in files:
            path = os.path.join(root, f)
            zf.write(path, os.path.relpath(path, '.'))
" "$whl"
  else
    echo "Warning: No _k2 shared library found in $whl"
  fi
  popd
  rm -rf "$_tmp_whl"
done
popd

[gemini-code-assist]

The current implementation of patching the RPATH has several potential issues and robustness concerns:

1. Multiple Wheels / No Wheels: Using whl=$(ls *.whl) can fail or behave unexpectedly if there are multiple wheels or no wheels in the directory. Using a for loop with a file existence check is much safer.
2. Finding the Shared Library: Using ls _k2.cpython-*.so assumes the .so file is at the root of the unzipped directory. If it is located in a subdirectory (e.g., k2/), ls will fail. Using find is much more robust.
3. Bash Interpolation in Python: Passing $whl directly inside the double-quoted Python command string can lead to syntax or execution errors if the filename contains special characters or spaces. Passing it via sys.argv is the standard and safe way.
4. Path Slicing Hack: Using path[2:] to strip the leading ./ is a fragile hack. Using os.path.relpath is the standard, cross-platform, and robust way to get the relative path for the zip archive.

pushd /var/www/wheelhouse
for whl in *.whl; do
  [ -f "$whl" ] || continue
  _tmp_whl="_tmp_${whl%.whl}"
  mkdir -p "$_tmp_whl"
  pushd "$_tmp_whl"
  unzip -o "../$whl"
  so_file=$(find . -name "_k2.cpython-*.so" -print -quit)
  if [ -n "$so_file" ]; then
    echo "Patching rpath for $so_file"
    current_rpath=$(patchelf --print-rpath "$so_file")
    echo "Current rpath: $current_rpath"
    new_rpath="\$ORIGIN/nvidia/nvtx/lib:\$ORIGIN/nvidia/cuda_runtime/lib:\$ORIGIN/nvidia/cuda_nvrtc/lib:${current_rpath}"
    echo "New rpath: $new_rpath"
    patchelf --set-rpath "$new_rpath" "$so_file"
    echo "Verified rpath:"
    patchelf --print-rpath "$so_file"
    python3 -c "
import zipfile, os, sys
whl = sys.argv[1]
with zipfile.ZipFile(os.path.join('..', whl), 'w', zipfile.ZIP_DEFLATED) as zf:
    for root, dirs, files in os.walk('.'):
        for f in files:
            path = os.path.join(root, f)
            zf.write(path, os.path.relpath(path, '.'))
" "$whl"
  else
    echo "Warning: No _k2 shared library found in $whl"
  fi
  popd
  rm -rf "$_tmp_whl"
done
popd