Merge pull request #181 from mrava87/ci-gputests

ci: added nccl tests to cupy GA

Author: mrava87

.github/workflows/build_cupy.yml

@@ -3,7 +3,7 @@ name: PyLops-MPI (CuPy)
 on: [push, pull_request]
 
 jobs:
-  build:
+  build-mpi:
     strategy:
       matrix:
         rank: ['2','4']
@@ -32,6 +32,39 @@ jobs:
       - name: Tests with pytest 
         run: |
           export CUPY_PYLOPS=1; export TEST_CUPY_PYLOPS=1; export PYLOPS_MPI_CUDA_AWARE=0
-          export PYTEST=.venv-pylops/bin/pytest
-          srun --account=yuxilab -n ${{matrix.rank}} -N 1 --gres=gpu:L40S:${{matrix.rank}} $PYTEST
+          export PYTEST=".venv-pylops/bin/pytest tests/ --with-mpi"
+          srun --account=yuxilab --mpi=pmix -n ${{matrix.rank}} -N 1 --gres=gpu:L40S:${{matrix.rank}} $PYTEST
+          echo "done!"
+
+  build-nccl:
+    strategy:
+      matrix:
+        rank: ['2','4']
+    runs-on: self-hosted
+    steps:
+      - name: Check out source repository
+        uses: actions/checkout@v4
+      - name: Set up Python environment and Install dependencies and pylops
+        run: |
+          python3 -m venv .venv-pylops
+          # write install scripts
+          cat << 'EOF' > pylops_tests.sh
+          #!/bin/bash
+          PYTHON=.venv-pylops/bin/python3
+          PIP=.venv-pylops/bin/pip
+          ls $PYTHON
+          ls $PIP
+          $PIP install --upgrade pip setuptools
+          $PIP install flake8 pytest setuptools-scm
+          $PIP install -r requirements-dev.txt
+          $PIP install cupy-cuda12x nvidia-nccl-cu12
+          $PYTHON -m setuptools_scm
+          $PIP install .
+          EOF
+          srun --account=yuxilab -n 1 -N 1 --gres=gpu:L40S:1 bash pylops_tests.sh
+      - name: Tests with pytest 
+        run: |
+          export CUPY_PYLOPS=1;
+          export PYTEST=".venv-pylops/bin/pytest tests_nccl/ --with-mpi"
+          srun --account=yuxilab --mpi=pmix -n ${{matrix.rank}} -N 1 --gres=gpu:L40S:${{matrix.rank}} $PYTEST
           echo "done!"

tests_nccl/test_ncclutils_nccl.py

@@ -10,6 +10,7 @@
 
 from pylops_mpi.utils._nccl import initialize_nccl_comm, nccl_allgather
 from pylops_mpi.utils._common import _prepare_allgather_inputs, _unroll_allgather_recv
+from pylops_mpi.utils.deps import nccl_enabled
 
 np.random.seed(42)
 
@@ -18,6 +19,11 @@
 par1 = {'n': 3, 'dtype': np.float64}
 
 
+def test_nccl_enabled():
+    """Test nccl_enabled is True (ensures that the tests are effectively
+    run with NCCL and not with the MPI fallack)"""
+    assert nccl_enabled
+
 @pytest.mark.mpi(min_size=2)
 @pytest.mark.parametrize("par", [(par1), ])
 def test_allgather_samesize(par):