Add script to inject libnuma into published ROCm torch wheels

[atalman]

Hotfix path for pytorch/pytorch#195670, so the already-published ROCm 7.14 wheels can be repaired without a rebuild or a cherry-pick.

Why

rocSHMEM's NUMAWrapper global ctor does dlopen("libnuma.so"), and torch/lib is on libtorch_rocshmem.so's RPATH via $ORIGIN, so a bundled copy there satisfies it. The AlmaLinux manywheel builder had no numactl installed, so repair_wheel.py's rocm_os_deps() skipped libnuma silently and the published wheels ship without it. Every import torch prints:

E-001h rocSHMEM Could not open libnuma. Returning	NUMAWrapper@src/gda/numa_wrapper.cpp:48

Per pytorch/pytorch#189110 the same failed dlopen can make rocSHMEM exit() at load, tripping a rocprofiler-sdk atexit deadlock that hangs import torch on hosts with no GPU or kfd — so this is not purely cosmetic.

pytorch/pytorch#195672 fixes the builder for future builds.

What it does

Follows repackage_torchaudio_cu130_to_cu132.py: discover from the channel index, download, modify, regenerate RECORD, upload back to S3 and R2 with x-amz-meta-checksum-sha256 so the PEP 503 index picks up the new hash.

python release/inject_libnuma_rocm_wheels.py --version 2.14.0 --rocm rocm7.14 \
    --libnuma ./libnuma.so.1 --dry-run

Both libnuma.so.1 and a bare libnuma.so are written as real files. A correct build makes the second a symlink, but the repack doesn't round-trip symlinks and the two are equivalent to dlopen. Wheels that already contain libnuma are skipped, so re-running is safe.

libnuma isn't vendored here — --libnuma takes a path, and the docstring gives a docker run almalinux:8 one-liner to extract one matching the builder's glibc so the manylinux_2_28 tag stays honest.

Repacking: auditwheel, and no new dependency surface

wheel pack emits an invalid ZIP64 header above 4GB (pypa/wheel#692), and ROCm wheels are ~1.4GB compressed / >4GB unpacked. pytorch#189748 tracked exactly that: wheels installing under pip but failing under uv. So this uses auditwheel, matching .ci/manywheel/repair_wheel.py and the torchaudio script here.

Only InWheelCtx is used — unpack, edit, repack, regenerate RECORD. The repair path is never invoked, so no patchelf, no RPATH rewriting, no library bundling, no platform retagging. Verified by inspection: InWheelCtx's source has zero references to patchelf, repair or lib bundling, and one to RECORD.

Dependencies: auditwheel needs only packaging and pyelftools, both pure Python. The patchelf binary is a requirement of auditwheel repair, which this doesn't call. Module-level imports are stdlib only; auditwheel and boto3 are imported at their call sites.

And rather than just assert that, verify_unchanged() enforces it: the repacked wheel must differ from the original by exactly the two added members, with every other member's CRC and size identical. RECORD is exempt (it legitimately lists the new files) and zip directory entries are skipped (the repack writes explicit ones the original may lack).

Test plan

Built a synthetic torch-2.14.0+rocm7.14-cp312-cp312-manylinux_2_28_x86_64.whl containing torch/lib/libtorch_rocshmem.so and a valid RECORD, plus a real x86-64 libnuma:

• inject → both libs added → repack; filename preserved
• RECORD regenerated with correct hashes and sizes:

  torch/lib/libnuma.so,sha256=ScpOG8SgwQOYmqyXm3qpMPlhUXEFqSUBPYITRJDf0dQ,62064
  torch/lib/libnuma.so.1,sha256=ScpOG8SgwQOYmqyXm3qpMPlhUXEFqSUBPYITRJDf0dQ,62064
  

• verify_unchanged: "only 2 members added, nothing else changed"
• libtorch_rocshmem.so preserved; second run skips (idempotent)
• guard rejects an injected extra member: repack added unexpected members: ['torch/lib/libsneaky.so']
• guard rejects a modified existing member: torch/version.py changed during repack (crc/size ...)
• input validation: a non-x86-64 ELF is rejected, as is anything without the libnuma.so.1 SONAME string

Two bugs found while testing, both fixed:

• InWheelCtx chdirs into its unpack directory, so a relative out_wheel resolved against the wrong cwd and repacking died with FileNotFoundError. Both paths are now resolved before entering the context.
• The first verify_unchanged was too strict and flagged the repack's own directory entries as unexpected additions.

Not yet run against real 1.4GB wheels — --dry-run covers that without uploading. Each wheel unpacks in full, so budget ~6GB of scratch per wheel and point TMPDIR at real storage rather than a tmpfs.

[vercel]

@atalman is attempting to deploy a commit to the Meta Open Source Team on Vercel.

A member of the Team first needs to authorize it.