From 9fd583f7f09b5feea79336ecb614f333c4fb1168 Mon Sep 17 00:00:00 2001
From: henrylhtsang <91030427+henrylhtsang@users.noreply.github.com>
Date: Mon, 29 Jan 2024 10:48:30 -0800
Subject: [PATCH] Update release_build.yml

---
 .github/workflows/release_build.yml | 42 +++++++++++++++++++++--------
 1 file changed, 31 insertions(+), 11 deletions(-)

diff --git a/.github/workflows/release_build.yml b/.github/workflows/release_build.yml
index dc064e658..64d94e094 100644
--- a/.github/workflows/release_build.yml
+++ b/.github/workflows/release_build.yml
@@ -6,10 +6,6 @@ name: Push Binary Release
 on:
   workflow_call:
     secrets:
-      AWS_PYTORCH_UPLOADER_ACCESS_KEY_ID:
-        required: true
-      AWS_PYTORCH_UPLOADER_SECRET_ACCESS_KEY:
-        required: true
       PYPI_TOKEN:
         required: false
   workflow_dispatch:
@@ -25,15 +21,23 @@ jobs:
          - os: linux.2xlarge
            python-version: 3.8
            python-tag: "py38"
-           cuda-tag: "cu11"
+           cuda-tag: "cu121"
          - os: linux.2xlarge
            python-version: 3.9
            python-tag: "py39"
-           cuda-tag: "cu11"
+           cuda-tag: "cu121"
          - os: linux.2xlarge
            python-version: '3.10'
            python-tag: "py310"
-           cuda-tag: "cu11"
+           cuda-tag: "cu121"
+         - os: linux.2xlarge
+           python-version: '3.11'
+           python-tag: "py311"
+           cuda-tag: "cu121"
+         - os: linux.2xlarge
+           python-version: '3.12'
+           python-tag: "py312"
+           cuda-tag: "cu121"
     steps:
     # Checkout the repository to the GitHub Actions runner
     - name: Check ldd --version
@@ -69,7 +73,12 @@ jobs:
     - name: Install PyTorch and CUDA
       shell: bash
       run: |
-        conda install -n build_binary -y pytorch pytorch-cuda=11.7 -c pytorch-test -c nvidia
+        conda run -n build_binary pip install torch --index-url https://download.pytorch.org/whl/test/cu121
+    - name: Install fbgemm
+      shell: bash
+      run: |
+        conda run -n build_binary pip install numpy
+        conda run -n build_binary pip install fbgemm-gpu --index-url https://download.pytorch.org/whl/test/cu121
     - name: Install Dependencies
       shell: bash
       run: |
@@ -102,8 +111,8 @@ jobs:
     strategy:
       matrix:
         os: [linux.4xlarge.nvidia.gpu]
-        python-version: [3.8, 3.9, "3.10"]
-        cuda-tag: ["cu11"]
+        python-version: [3.8, 3.9, "3.10", "3.11", "3.12"]
+        cuda-tag: ["cu121"]
     needs: build_on_cpu
     # the glibc version should match the version of the one we used to build the binary
     # for this case, it's 2.26
@@ -168,8 +177,17 @@ jobs:
     - name: Install PyTorch and CUDA
       shell: bash
       run: |
-        conda install -n build_binary -y pytorch pytorch-cuda=11.7 -c pytorch-test -c nvidia
+        conda run -n build_binary pip install torch --index-url https://download.pytorch.org/whl/test/cu121
     # download wheel from GHA
+    - name: Install fbgemm
+      shell: bash
+      run: |
+        conda run -n build_binary pip install numpy
+        conda run -n build_binary pip install fbgemm-gpu --index-url https://download.pytorch.org/whl/test/cu121
+    - name: Install torchmetrics
+      shell: bash
+      run: |
+        conda run -n build_binary pip install torchmetrics==1.0.3
     - name: Download wheel
       uses: actions/download-artifact@v2
       with:
@@ -183,6 +201,8 @@ jobs:
     - name: Test fbgemm_gpu and torchrec installation
       shell: bash
       run: |
+        conda run -n build_binary \
+          python -c "import fbgemm_gpu"
         conda run -n build_binary \
           python -c "import torchrec"
     # Push to Pypi