From 07b3d3e9e72d89c139ba4a39d325a718abe352c0 Mon Sep 17 00:00:00 2001
From: Julio Perez <jperez@nvidia.com>
Date: Thu, 12 Oct 2023 12:01:58 -0400
Subject: [PATCH] use rapids infra to run testing

---
 .github/workflows/gpu.yml | 114 +++++++++++++++++++++++---------------
 1 file changed, 69 insertions(+), 45 deletions(-)

diff --git a/.github/workflows/gpu.yml b/.github/workflows/gpu.yml
index 694ff16a9e..4ed77bb095 100644
--- a/.github/workflows/gpu.yml
+++ b/.github/workflows/gpu.yml
@@ -1,58 +1,82 @@
-name: gpu-ci
+name: GPU CI
 
 on:
   workflow_dispatch:
   push:
-    branches: [main]
+    branches: 
+      - main
+      - "pull-request/[0-9]+"
     tags:
       - "v[0-9]+.[0-9]+.[0-9]+"
-  pull_request:
-    branches: [main]
-    types: [opened, synchronize, reopened]
 
-concurrency:
-  group: ${{ github.workflow }}-${{ github.ref }}
-  cancel-in-progress: true
+jobs:
+  gpu-ci-tensorflow:
+    runs-on: linux-amd64-gpu-p100-latest-1
+    container:
+          image: nvcr.io/nvidia/merlin/merlin-tensorflow:nightly
+          env:
+            NVIDIA_VISIBLE_DEVICES: ${{ env.NVIDIA_VISIBLE_DEVICES }}
+          options: --shm-size=1G
+    steps:
+    - uses: actions/checkout@v3
+      with:
+        fetch-depth: 0
+    - name: Run tests
+      run: |
+        pip install tox
+        PYTEST_MARKERS="unit and not (examples or integration or notebook) and (singlegpu or not multigpu) $extra_pytest_markers" MERLIN_BRANCH=$branch COMPARE_BRANCH=${{ github.base_ref }} tox -e gpu
 
 jobs:
-  gpu-ci:
-    runs-on: 1GPU
+  gpu-ci-pytorch:
+    runs-on: linux-amd64-gpu-p100-latest-1
+    container:
+          image: nvcr.io/nvidia/merlin/merlin-pytorch:nightly
+          env:
+            NVIDIA_VISIBLE_DEVICES: ${{ env.NVIDIA_VISIBLE_DEVICES }}
+          options: --shm-size=1G
+    steps:
+    - uses: actions/checkout@v3
+      with:
+        fetch-depth: 0
+    - name: Run tests
+      run: |
+        pip install tox
+        PYTEST_MARKERS="unit and not (examples or integration or notebook) and (singlegpu or not multigpu) $extra_pytest_markers" MERLIN_BRANCH=$branch COMPARE_BRANCH=${{ github.base_ref }} tox -e gpu
+
 
+jobs:
+  gpu-ci-tensorflow-examples:
+    runs-on: linux-amd64-gpu-p100-latest-1
+    container:
+          image: nvcr.io/nvidia/merlin/merlin-tensorflow:nightly
+          env:
+            NVIDIA_VISIBLE_DEVICES: ${{ env.NVIDIA_VISIBLE_DEVICES }}
+          options: --shm-size=1G
     steps:
-      - uses: actions/checkout@v3
-        with:
-          fetch-depth: 0
-      - 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\///')
-          fi
-          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) and (singlegpu or not multigpu) $extra_pytest_markers" MERLIN_BRANCH=$branch COMPARE_BRANCH=${{ github.base_ref }} tox -e gpu
-
-  tests-examples:
-    runs-on: 1GPU
+    - uses: actions/checkout@v3
+      with:
+        fetch-depth: 0
+    - name: Run tests
+      run: |
+        pip install tox
+        PYTEST_MARKERS="(examples or notebook) $extra_pytest_markers" tox -e gpu
+
 
+jobs:
+  gpu-ci-pytorch-examples:
+    runs-on: linux-amd64-gpu-p100-latest-1
+    container:
+          image: nvcr.io/nvidia/merlin/merlin-pytorch:nightly
+          env:
+            NVIDIA_VISIBLE_DEVICES: ${{ env.NVIDIA_VISIBLE_DEVICES }}
+          options: --shm-size=1G
     steps:
-      - uses: actions/checkout@v3
-        with:
-          fetch-depth: 0
-      - 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\///')
-          fi
-          if [[ "${{ github.ref }}" != 'refs/heads/main' ]]; then
-              extra_pytest_markers="and changed"
-          fi
-          cd ${{ github.workspace }}; PYTEST_MARKERS="(examples or notebook) $extra_pytest_markers" MERLIN_BRANCH=$branch COMPARE_BRANCH=${{ github.base_ref }} tox -e gpu
+    - uses: actions/checkout@v3
+      with:
+        fetch-depth: 0
+    - name: Run tests
+      run: |
+        pip install tox
+        PYTEST_MARKERS="(examples or notebook) $extra_pytest_markers" tox -e gpu
+
+        
\ No newline at end of file