use rapids container

Signed-off-by: Sarah Yurick <sarahyurick@gmail.com>

.github/workflows/gpuci.yml

@@ -36,11 +36,11 @@ jobs:
 
     - name: Verify installations
       run: |
-        docker exec nemo-curator-container pip list
+        docker exec nemo-curator-container /bin/bash -c "source activate rapids && pip list"
 
     - name: Run PyTests with GPU mark
       run: |
-        docker exec nemo-curator-container pytest -m gpu --rootdir /opt/NeMo-Curator /opt/NeMo-Curator/tests
+        docker exec nemo-curator-container /bin/bash -c "source activate rapids && pytest -m gpu --rootdir /opt/NeMo-Curator /opt/NeMo-Curator/tests"
 
     - name: Cleanup
       run: |

Dockerfile

@@ -1,15 +1,27 @@
-FROM nvcr.io/nvidia/pytorch:24.07-py3
+FROM rapidsai/ci-conda:cuda12.5.1-ubuntu22.04-py3.10
 
 WORKDIR /opt
 
 ARG FORKED_REPO_URL
 ARG CURATOR_COMMIT
 
+conda create -y --name rapids -c conda-forge -c nvidia \
+  python=3.10 \
+  cuda-cudart \
+  libcufft \
+  libcublas \
+  libcurand \
+  libcusparse \
+  libcusolver
+
 RUN <<"EOF" bash -exu
 git clone $FORKED_REPO_URL
 cd NeMo-Curator
 git fetch origin $CURATOR_COMMIT --depth=1
 git checkout $CURATOR_COMMIT
-pip install cython pytest setuptools pip --upgrade
-pip install --extra-index-url https://pypi.nvidia.com ".[cuda12x]"
+RUN /bin/bash -c "source activate rapids && \
+    conda install -y -c conda-forge cython pytest setuptools && \
+    pip install --upgrade pip"
+RUN /bin/bash -c "source activate rapids && \
+    pip install --extra-index-url https://pypi.nvidia.com '.[cuda12x]'"
 EOF