From dfa2d3079f681a6bfafd783a6397fbb59cc8f0b1 Mon Sep 17 00:00:00 2001
From: Huanyu Zhang <huanyuzhang@meta.com>
Date: Thu, 15 Feb 2024 10:23:28 -0800
Subject: [PATCH] Fix github action CI for GPU (#625)

Summary: Pull Request resolved: https://github.com/pytorch/opacus/pull/625

Differential Revision: D53292685
---
 .github/workflows/ci_gpu.yml | 52 +++++++++++++++++++++++-------------
 scripts/install_via_pip.sh   |  7 ++---
 2 files changed, 38 insertions(+), 21 deletions(-)

diff --git a/.github/workflows/ci_gpu.yml b/.github/workflows/ci_gpu.yml
index 90dca200..4fb203b5 100644
--- a/.github/workflows/ci_gpu.yml
+++ b/.github/workflows/ci_gpu.yml
@@ -8,48 +8,64 @@ on:
     branches:
       - main
 
+jobs:
   unittest_multi_gpu:
-    runs-on: linux.4xlarge.nvidia.gpu
+    runs-on: ${{ matrix.os }}
+    strategy:
+      matrix:
+        os: [linux.4xlarge.nvidia.gpu]
+        python-version: [3.8]
+        cuda-tag: ["cu11"]
     steps:
       - name: Checkout
         uses: actions/checkout@v2
 
-      - name: Set up Python
-        uses: actions/setup-python@v2
-        with:
-          python-version: 3.9
+     - name: Display Python version
+        run: python -c "import sys; print(sys.version)"
+
+      # - name: Set up Python
+      #   uses: actions/setup-python@v2
+      #   with:
+      #     python-version: '3.9'
 
       - name: Install dependencies
         run: |
+          python3 -m pip install --upgrade pip
           ./scripts/install_via_pip.sh -c
 
       - name: Run multi-GPU unit tests
         run: |
           nvidia-smi
           nvcc --version
-          python -m unittest opacus.tests.multigpu_gradcheck.GradientComputationTest.test_gradient_correct
+          python3 -m unittest opacus.tests.multigpu_gradcheck.GradientComputationTest.test_gradient_correct
 
 
   integrationtest_py39_torch_release_cuda:
-    runs-on: ubuntu-latest
-    container:
-      # https://hub.docker.com/r/nvidia/cuda
-      image: nvidia/cuda:12.3.1-base-ubuntu22.04
-      options: --gpus all
-    env:
-      TZ: 'UTC'
+    # runs-on: linux.4xlarge.nvidia.gpu
+    runs-on: ${{ matrix.os }}
+    strategy:
+      matrix:
+        os: [linux.4xlarge.nvidia.gpu]
+        python-version: [3.8]
+        cuda-tag: ["cu11"]
+    # container:
+    #   # https://hub.docker.com/r/nvidia/cuda
+    #   image: nvidia/cuda:12.3.1-base-ubuntu22.04
+    #   options: --gpus all
+    # env:
+    #   TZ: 'UTC'
     steps:
       - name: Checkout
         uses: actions/checkout@v2
 
-      - name: Set up Python
-        uses: actions/setup-python@v2
-        with:
-          python-version: 3.9
+      # - name: Set up Python
+      #   uses: actions/setup-python@v2
+      #   with:
+      #     python-version: '3.10.13'
 
       - name: Install dependencies
         run: |
-          python -m pip install --upgrade pip
+          python3 -m pip install --upgrade pip
           pip install pytest coverage coveralls
           ./scripts/install_via_pip.sh -c
 
diff --git a/scripts/install_via_pip.sh b/scripts/install_via_pip.sh
index d906094d..17581c1e 100755
--- a/scripts/install_via_pip.sh
+++ b/scripts/install_via_pip.sh
@@ -33,9 +33,10 @@ while getopts 'ncdv:' flag; do
 # NOTE: Only Debian variants are supported, since this script is only
 # used by our tests on CircleCI.
 
-curl -sL https://dl.yarnpkg.com/debian/pubkey.gpg | sudo apt-key add -
-echo "deb https://dl.yarnpkg.com/debian/ stable main" | sudo tee /etc/apt/sources.list.d/yarn.list
-sudo apt-get update && sudo apt-get install yarn
+# curl -sL https://dl.yarnpkg.com/debian/pubkey.gpg | sudo apt-key add -
+# echo "deb https://dl.yarnpkg.com/debian/ stable main" | sudo tee /etc/apt/sources.list.d/yarn.list
+# sudo apt-get update && sudo apt-get install yarn
+# sudo yum install yarn
 
 # yarn needs terminal info
 export TERM=xterm