Test Running NVTabular GPU tests with rapids runner

NVIDIA-Merlin · Jun 2, 2023 · 3064bc4 · 3064bc4
1 parent 6dd49c9
commit 3064bc4
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 16 deletions.
diff --git a/.github/workflows/gpu-tests.yml b/.github/workflows/gpu-tests.yml
@@ -16,19 +16,19 @@ concurrency:
 
 jobs:
   gpu-tests:
-    runs-on: 2GPU
-
+    runs-on: linux-amd64-gpu-v100-latest-1
+    container:
+      image: nvidia/cuda:11.8.0-base-ubuntu22.04
+      env:
+        NVIDIA_VISIBLE_DEVICES: ${{ env.NVIDIA_VISIBLE_DEVICES }}
     steps:
       - uses: actions/checkout@v3
         with:
           fetch-depth: 0
+      - name: Get Branch name
+        id: get-branch-name
+        uses: NVIDIA-Merlin/.github/actions/branch-name@main
       - 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
-          cd ${{ github.workspace }}; tox -e test-gpu -- $branch
+          merlin_branch="${{ steps.get-branch-name.outputs.branch }}"
+          RAPIDS_VERSION=23.04 MERLIN_BRANCH="$merlin_branch" tox -e test-gpu
diff --git a/tox.ini b/tox.ini
@@ -28,9 +28,8 @@ commands =
 setenv =
     TF_GPU_ALLOCATOR=cuda_malloc_async
 passenv =
-    OPAL_PREFIX
-    NR_USER
     CUDA_VISIBLE_DEVICES
+    NVIDIA_VISIBLE_DEVICES
 sitepackages=true
 ; Runs in: Internal Jenkins
 ; Runs GPU-based tests.
@@ -42,10 +41,12 @@ deps =
     pytest
     pytest-cov
 commands =
-    python -m pip install --upgrade git+https://github.com/NVIDIA-Merlin/core.git
-    python -m pip install --upgrade git+https://github.com/NVIDIA-Merlin/dataloader.git
-    python -m pip install --upgrade git+https://github.com/NVIDIA-Merlin/models.git
-    python -m pip install --upgrade git+https://github.com/NVIDIA-Merlin/core.git@{posargs:main}
+    python -m pip install rmm-cu11=={env:RAPIDS_VERSION} cudf-cu11=={env:RAPIDS_VERSION} dask-cudf-cu11=={env:RAPIDS_VERSION} --extra-index-url=https://pypi.nvidia.com
+    python -m pip install --upgrade \
+           git+https://github.com/NVIDIA-Merlin/core.git@{env:MERLIN_BRANCH:main} \
+           git+https://github.com/NVIDIA-Merlin/dataloader.git@{env:MERLIN_BRANCH:main} \
+           git+https://github.com/NVIDIA-Merlin/models.git@{env:MERLIN_BRANCH:main} \
+           git+https://github.com/NVIDIA-Merlin/core.git@{env:MERLIN_BRANCH:main} \
     python -m pytest --cov-report term --cov merlin -rxs tests/unit
 
 [testenv:test-merlin]