From e24f2ba7fb40306672768f403b271652630970d8 Mon Sep 17 00:00:00 2001
From: Oliver Holworthy <oholworthy@nvidia.com>
Date: Mon, 17 Apr 2023 13:11:45 +0100
Subject: [PATCH 1/2] Add job to run NVTabular unit tests on GPU

---
 .github/workflows/merlin.yml | 9 ++++++++-
 tox.ini                      | 9 +++++++--
 2 files changed, 15 insertions(+), 3 deletions(-)

diff --git a/.github/workflows/merlin.yml b/.github/workflows/merlin.yml
index cbb68af28..534d65773 100644
--- a/.github/workflows/merlin.yml
+++ b/.github/workflows/merlin.yml
@@ -32,4 +32,11 @@ jobs:
     name: "NVTabular (CPU)"
     uses: ./.github/workflows/tox.yml
     with:
-      env: test-nvtabular-cpu
+      env: test-nvtabular
+
+  nvtabular-gpu:
+    runs-on: 1GPU
+    name: "NVTabular (GPU)"
+    uses: ./.github/workflows/tox.yml
+    with:
+      env: test-nvtabular
diff --git a/tox.ini b/tox.ini
index 91ce69f1c..2ae67e1ee 100644
--- a/tox.ini
+++ b/tox.ini
@@ -52,8 +52,13 @@ commands =
     python -m pytest --cov-report term --cov merlin -rxs tests/unit
 
 
-[testenv:test-nvtabular-cpu]
-passenv=GIT_COMMIT
+[testenv:test-nvtabular]
+setenv =
+    TF_GPU_ALLOCATOR=cuda_malloc_async
+passenv =
+    GIT_COMMIT
+    OPAL_PREFIX
+    CUDA_VISIBLE_DEVICES
 sitepackages=true
 allowlist_externals = git
 deps =

From b020fe58e690fac54229341f8a106e53662bdf9a Mon Sep 17 00:00:00 2001
From: Oliver Holworthy <oholworthy@nvidia.com>
Date: Mon, 17 Apr 2023 13:26:41 +0100
Subject: [PATCH 2/2] Add separate reusable workflow for running tox in gpu
 runner

---
 .github/workflows/merlin.yml  |  3 +--
 .github/workflows/tox-gpu.yml | 26 ++++++++++++++++++++++++++
 2 files changed, 27 insertions(+), 2 deletions(-)
 create mode 100644 .github/workflows/tox-gpu.yml

diff --git a/.github/workflows/merlin.yml b/.github/workflows/merlin.yml
index 534d65773..db6c31954 100644
--- a/.github/workflows/merlin.yml
+++ b/.github/workflows/merlin.yml
@@ -35,8 +35,7 @@ jobs:
       env: test-nvtabular
 
   nvtabular-gpu:
-    runs-on: 1GPU
     name: "NVTabular (GPU)"
-    uses: ./.github/workflows/tox.yml
+    uses: ./.github/workflows/tox-gpu.yml
     with:
       env: test-nvtabular
diff --git a/.github/workflows/tox-gpu.yml b/.github/workflows/tox-gpu.yml
new file mode 100644
index 000000000..a94ac42e2
--- /dev/null
+++ b/.github/workflows/tox-gpu.yml
@@ -0,0 +1,26 @@
+name: "Reusable Workflow - Run Tox Env (1GPU)"
+
+on:
+  workflow_call:
+    inputs:
+      env:
+        description: "The name of the tox environment to run"
+        required: true
+        type: string
+
+jobs:
+  check:
+    runs-on: 1GPU
+
+    steps:
+      - uses: actions/checkout@v3
+      - name: Run tests
+        run: |
+          ref_type=${{ github.ref_type }}
+          branch=main
+          if [[ $ref_type == "tag"* ]]
+          then
+            raw=$(git branch -r --contains ${{ github.ref_name }})
+            branch=${raw/origin\/}
+          fi
+          GIT_COMMIT=`git rev-parse HEAD` tox -e ${{ inputs.env }} -- $branch