Merges pytorch#88125 and uses container from pytorch/builder#1196

syed-ahmed · Dec 16, 2022 · a29c6b8 · a29c6b8
1 parent 815f460
commit a29c6b8
Show file tree

Hide file tree

Showing 3 changed files with 43 additions and 38 deletions.
diff --git a/.github/scripts/generate_binary_build_matrix.py b/.github/scripts/generate_binary_build_matrix.py
@@ -31,7 +31,10 @@ def arch_type(arch_version: str) -> str:
 WHEEL_CONTAINER_IMAGES = {
     **{
         gpu_arch: f"pytorch/manylinux-builder:cuda{gpu_arch}"
-        for gpu_arch in CUDA_ARCHES
+        for gpu_arch in CUDA_ARCHES if gpu_arch == "11.6"
+    },
+    **{
+        "11.7": "tousif111/manylinux-builder:cuda11.7"
     },
     **{
         gpu_arch: f"pytorch/manylinux-builder:rocm{gpu_arch}"
@@ -219,14 +222,16 @@ def generate_wheels_matrix(os: str,
                         "container_image": WHEEL_CONTAINER_IMAGES[arch_version],
                         "package_type": package_type,
                         "pytorch_extra_install_requirements":
-                        "nvidia-cuda-runtime-cu11; platform_system == 'Linux' | "
+                        "nvidia-cuda-nvrtc-cu11==11.7.99; platform_system == 'Linux' | "
+                        "nvidia-cuda-runtime-cu11==11.7.99; platform_system == 'Linux' | "
                         "nvidia-cuda-cupti-cu11==11.7.101; platform_system == 'Linux' | "
                         "nvidia-cudnn-cu11==8.5.0.96; platform_system == 'Linux' | "
                         "nvidia-cublas-cu11==11.10.3.66; platform_system == 'Linux' | "
                         "nvidia-cufft-cu11==10.9.0.58; platform_system == 'Linux' | "
                         "nvidia-curand-cu11==10.2.10.91; platform_system == 'Linux' | "
                         "nvidia-cusolver-cu11==11.4.0.1; platform_system == 'Linux' | "
                         "nvidia-cusparse-cu11==11.7.4.91; platform_system == 'Linux' | "
+                        "nvidia-nccl-cu11==2.14.3; platform_system == 'Linux' | "
                         "nvidia-nvtx-cu11==11.7.91; platform_system == 'Linux'",
                         "build_name":
                         f"{package_type}-py{python_version}-{gpu_arch_type}{gpu_arch_version}-with-pypi-cudnn"

diff --git a/.github/workflows/_binary-build-linux.yml b/.github/workflows/_binary-build-linux.yml
@@ -160,7 +160,7 @@ jobs:
         with:
           ref: main
           submodules: recursive
-          repository: pytorch/builder
+          repository: syed-ahmed/builder
           path: builder
           quiet-checkout: true
       - name: Clean pytorch/builder checkout

diff --git a/.github/workflows/generated-linux-binary-manywheel-nightly.yml b/.github/workflows/generated-linux-binary-manywheel-nightly.yml