From d6409a5643a8b0e1e8511438f7d0590e6dfa67b9 Mon Sep 17 00:00:00 2001
From: ramilbakhshyiev <7597197+ramilbakhshyiev@users.noreply.github.com>
Date: Thu, 15 Aug 2024 16:36:19 -0700
Subject: [PATCH] Enable Python tests for ROCm (#3763)

Summary:
Pull Request resolved: https://github.com/facebookresearch/faiss/pull/3763

This change enables Python version of tests for ROCm builds.

Differential Revision: D61366282
---
 .github/actions/build_cmake/action.yml | 13 ++++++++++---
 1 file changed, 10 insertions(+), 3 deletions(-)

diff --git a/.github/actions/build_cmake/action.yml b/.github/actions/build_cmake/action.yml
index 5213034cb7..93c327ccfd 100644
--- a/.github/actions/build_cmake/action.yml
+++ b/.github/actions/build_cmake/action.yml
@@ -58,7 +58,10 @@ runs:
 
         # install test packages
         conda install -y pytest
-        if [ "${{ inputs.gpu }}" = "ON" ]; then
+        if [ "${{ inputs.rocm }}" = "ON" ]; then
+          : # skip torch install via conda, we need to install via pip to get
+            #  ROCm-enabled version until it's supported in conda by PyTorch
+        elif [ "${{ inputs.gpu }}" = "ON" ]; then
           conda install -y -q pytorch pytorch-cuda=12.4 -c pytorch -c nvidia/label/cuda-12.4.0
         else
           conda install -y -q pytorch -c pytorch
@@ -138,6 +141,11 @@ runs:
       working-directory: build/faiss/python
       run: |
         $CONDA/bin/python setup.py install
+    - name: ROCm - install ROCm-enabled torch via pip
+      if: inputs.rocm == 'ON'
+      shell: bash
+      run: |
+        pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/rocm6.1
     - name: Python tests (CPU only)
       if: inputs.gpu == 'OFF'
       shell: bash
@@ -145,7 +153,7 @@ runs:
         pytest --junitxml=test-results/pytest/results.xml tests/test_*.py
         pytest --junitxml=test-results/pytest/results-torch.xml tests/torch_*.py
     - name: Python tests (CPU + GPU)
-      if: inputs.gpu == 'ON' && inputs.rocm == 'OFF'
+      if: inputs.gpu == 'ON'
       shell: bash
       run: |
         pytest --junitxml=test-results/pytest/results.xml tests/test_*.py
@@ -160,7 +168,6 @@ runs:
         FAISS_DISABLE_CPU_FEATURES=AVX2 LD_DEBUG=libs $CONDA/bin/python -c "import faiss" 2>&1 | grep faiss.so
         LD_DEBUG=libs $CONDA/bin/python -c "import faiss" 2>&1 | grep faiss_avx2.so
     - name: Upload test results
-      if: inputs.rocm == 'OFF'
       uses: actions/upload-artifact@v4
       with:
         name: test-results-arch=${{ runner.arch }}-opt=${{ inputs.opt_level }}-gpu=${{ inputs.gpu }}-raft=${{ inputs.raft }}-rocm=${{ inputs.rocm }}