From 2c7f3376c5bb83716db86c4af75d3b38a7c0f0be Mon Sep 17 00:00:00 2001
From: Oliver Holworthy <1216955+oliverholworthy@users.noreply.github.com>
Date: Fri, 30 Jun 2023 14:20:40 +0100
Subject: [PATCH] Run single GPU tests in nvidia/tensorflow and nvidia/cuda
 images

---
 .github/workflows/gpu.yml | 62 +++++++++++++++++++++++++++++++++------
 tox.ini                   |  2 ++
 2 files changed, 55 insertions(+), 9 deletions(-)

diff --git a/.github/workflows/gpu.yml b/.github/workflows/gpu.yml
index 90478661c9..28b5bc797c 100644
--- a/.github/workflows/gpu.yml
+++ b/.github/workflows/gpu.yml
@@ -16,25 +16,69 @@ concurrency:
 
 jobs:
   gpu-ci:
-    runs-on: 1GPU
-
+    runs-on: linux-amd64-gpu-p100-latest-1
+    container:
+      image: nvcr.io/nvidia/tensorflow:23.06-tf2-py3
+      env:
+        NVIDIA_VISIBLE_DEVICES: ${{ env.NVIDIA_VISIBLE_DEVICES }}
     steps:
       - uses: actions/checkout@v3
         with:
           fetch-depth: 0
+      - name: Install Ubuntu packages
+        run: |
+          apt-get update -y
+          apt-get install -y lsb-release
+      - name: Install and upgrade python packages
+        run: |
+          python -m pip install --upgrade pip tox
+      - name: Get Branch name
+        id: get-branch-name
+        uses: NVIDIA-Merlin/.github/actions/branch-name@branch-name-pull-request
       - name: Run tests
         run: |
-          ref_type=${{ github.ref_type }}
-          branch=main
-          if [[ $ref_type == "tag"* ]]
-          then
-            git -c protocol.version=2 fetch --no-tags --prune --progress --no-recurse-submodules --depth=1 origin +refs/heads/release*:refs/remotes/origin/release*
-            branch=$(git branch -r --contains ${{ github.ref_name }} --list '*release*' --format "%(refname:short)" | sed -e 's/^origin\///')
+          if [[ "${{ github.ref }}" != 'refs/heads/main' ]]; then
+              extra_pytest_markers="and changed"
           fi
+          merlin_branch="${{ steps.get-branch-name.outputs.branch }}"
+          MERLIN_BRANCH=$merlin_branch \
+            PYTEST_MARKERS="unit and not (examples or integration or notebook) $extra_pytest_markers" \
+            tox -e gpu
+
+  gpu-cu11:
+    runs-on: linux-amd64-gpu-p100-latest-1
+    container:
+      image: nvidia/cuda:11.8.0-devel-ubuntu22.04
+      env:
+        NVIDIA_VISIBLE_DEVICES: ${{ env.NVIDIA_VISIBLE_DEVICES }}
+    steps:
+      - uses: actions/checkout@v3
+        with:
+          fetch-depth: 0
+      - name: Install Ubuntu packages
+        run: |
+          apt-get update -y
+          # libcudnn8 installed for tensorflow GPU support
+          apt-get install -y git lsb-release 'libcudnn8=*cuda11.8'
+      - name: Set up Python 3.8
+        uses: actions/setup-python@v4
+        with:
+          python-version: 3.8
+      - name: Install and upgrade python packages
+        run: |
+          python -m pip install --upgrade pip tox
+      - name: Get Branch name
+        id: get-branch-name
+        uses: NVIDIA-Merlin/.github/actions/branch-name@branch-name-pull-request
+      - name: Run tests
+        run: |
           if [[ "${{ github.ref }}" != 'refs/heads/main' ]]; then
               extra_pytest_markers="and changed"
           fi
-          cd ${{ github.workspace }}; PYTEST_MARKERS="unit and not (examples or integration or notebook) $extra_pytest_markers" MERLIN_BRANCH=$branch COMPARE_BRANCH=${{ github.base_ref }} tox -e gpu
+          merlin_branch="${{ steps.get-branch-name.outputs.branch }}"
+          RAPIDS_VERSION=23.04 MERLIN_BRANCH=$merlin_branch \
+            PYTEST_MARKERS="unit and not (examples or integration or notebook) $extra_pytest_markers" \
+            tox -e gpu-cu11
 
   tests-examples:
     runs-on: 1GPU
diff --git a/tox.ini b/tox.ini
index 3d624c88c5..1bfd4b2e17 100644
--- a/tox.ini
+++ b/tox.ini
@@ -37,6 +37,7 @@ commands =
 ; Runs GPU-based tests.
 allowlist_externals =
     bash
+    cp
 deps =
     -rrequirements/test.txt
     git+https://github.com/NVIDIA-Merlin/core.git@{env:MERLIN_BRANCH}
@@ -49,6 +50,7 @@ setenv =
     TF_GPU_ALLOCATOR=cuda_malloc_async
 sitepackages=true
 commands =
+    bash -c 'cp $(python -c "import sys; print(sys.base_prefix)")/lib/*.so* $(python -c "import sys; print(sys.prefix)")/lib'
     bash -c 'python -m pytest --cov-report term --cov merlin -m "{env:PYTEST_MARKERS}" -rxs {posargs:tests} || ([ $? = 5 ] && exit 0 || exit $?)'