remove slurm-related env in test scripts, fix Makefile type, add tests_gpu

Author: tharittk

Makefile

@@ -29,7 +29,7 @@ dev-install:
 
 dev-install_nccl:
 	make pipcheck
-	$(PIP) install -r requirements-dev.txt && $(PIP) install cupy-cuda12x nvidia-nccl-cu12  $(PIP) install -e .
+	$(PIP) install -r requirements-dev.txt && $(PIP) install cupy-cuda12x nvidia-nccl-cu12 && $(PIP) install -e .
 
 install_conda:
 	conda env create -f environment.yml && conda activate pylops_mpi && pip install .
@@ -49,6 +49,10 @@ lint:
 tests:
 	mpiexec -n $(NUM_PROCESSES) pytest tests/ --with-mpi
 
+# assuming NUM_PROCESSES <= number of gpus available
+tests_gpu:
+	export TEST_CUPY_PYLOPS=1 && mpiexec -n $(NUM_PROCESSES) pytest tests/ --with-mpi
+
 # assuming NUM_PROCESSES <= number of gpus available
 tests_nccl:	
 	mpiexec -n $(NUM_PROCESSES) pytest tests_nccl/ --with-mpi

tests/test_blockdiag.py

@@ -29,11 +29,7 @@
 np.random.seed(42)
 rank = MPI.COMM_WORLD.Get_rank()
 if backend == "cupy":
-    device_count = np.cuda.runtime.getDeviceCount()
-    device_id = int(
-        os.environ.get("OMPI_COMM_WORLD_LOCAL_RANK")
-        or rank % np.cuda.runtime.getDeviceCount()
-    )
+    device_id = rank % np.cuda.runtime.getDeviceCount()
     np.cuda.Device(device_id).use()
 
 

tests/test_derivative.py

@@ -26,11 +26,7 @@
 rank = MPI.COMM_WORLD.Get_rank()
 size = MPI.COMM_WORLD.Get_size()
 if backend == "cupy":
-    device_count = np.cuda.runtime.getDeviceCount()
-    device_id = int(
-        os.environ.get("OMPI_COMM_WORLD_LOCAL_RANK")
-        or rank % np.cuda.runtime.getDeviceCount()
-    )
+    device_id = rank % np.cuda.runtime.getDeviceCount()
     np.cuda.Device(device_id).use()
 
 

tests/test_distributedarray.py

@@ -23,11 +23,7 @@
 np.random.seed(42)
 rank = MPI.COMM_WORLD.Get_rank()
 if backend == "cupy":
-    device_count = np.cuda.runtime.getDeviceCount()
-    device_id = int(
-        os.environ.get("OMPI_COMM_WORLD_LOCAL_RANK")
-        or rank % np.cuda.runtime.getDeviceCount()
-    )
+    device_id = rank % np.cuda.runtime.getDeviceCount()
     np.cuda.Device(device_id).use()
 
 par1 = {'global_shape': (500, 501),

tests/test_fredholm.py

@@ -30,11 +30,7 @@
 rank = MPI.COMM_WORLD.Get_rank()
 size = MPI.COMM_WORLD.Get_size()
 if backend == "cupy":
-    device_count = np.cuda.runtime.getDeviceCount()
-    device_id = int(
-        os.environ.get("OMPI_COMM_WORLD_LOCAL_RANK")
-        or rank % np.cuda.runtime.getDeviceCount()
-    )
+    device_id = rank % np.cuda.runtime.getDeviceCount()
     np.cuda.Device(device_id).use()
 
 par1 = {

tests/test_linearop.py

@@ -32,11 +32,7 @@
 rank = MPI.COMM_WORLD.Get_rank()
 size = MPI.COMM_WORLD.Get_size()
 if backend == "cupy":
-    device_count = np.cuda.runtime.getDeviceCount()
-    device_id = int(
-        os.environ.get("OMPI_COMM_WORLD_LOCAL_RANK")
-        or rank % np.cuda.runtime.getDeviceCount()
-    )
+    device_id = rank % np.cuda.runtime.getDeviceCount()
     np.cuda.Device(device_id).use()
 
 par1 = {'ny': 101, 'nx': 101, 'dtype': np.float64}

tests/test_matrixmult.py

@@ -28,12 +28,9 @@
 size = base_comm.Get_size()
 rank = MPI.COMM_WORLD.Get_rank()
 if backend == "cupy":
-    device_count = np.cuda.runtime.getDeviceCount()
-    device_id = int(
-        os.environ.get("OMPI_COMM_WORLD_LOCAL_RANK")
-        or rank % np.cuda.runtime.getDeviceCount()
-    )
+    device_id = rank % np.cuda.runtime.getDeviceCount()
     np.cuda.Device(device_id).use()
+
 # Define test cases: (N, K, M, dtype_str)
 # M, K, N are matrix dimensions A(N,K), B(K,M)
 # P_prime will be ceil(sqrt(size)).

tests/test_solver.py

@@ -39,11 +39,7 @@
 size = MPI.COMM_WORLD.Get_size()
 rank = MPI.COMM_WORLD.Get_rank()
 if backend == "cupy":
-    device_count = np.cuda.runtime.getDeviceCount()
-    device_id = int(
-        os.environ.get("OMPI_COMM_WORLD_LOCAL_RANK")
-        or rank % np.cuda.runtime.getDeviceCount()
-    )
+    device_id = rank % np.cuda.runtime.getDeviceCount()
     np.cuda.Device(device_id).use()
 
 par1 = {

tests/test_stack.py

@@ -23,11 +23,7 @@
 
 rank = MPI.COMM_WORLD.Get_rank()
 if backend == "cupy":
-    device_count = np.cuda.runtime.getDeviceCount()
-    device_id = int(
-        os.environ.get("OMPI_COMM_WORLD_LOCAL_RANK")
-        or rank % np.cuda.runtime.getDeviceCount()
-    )
+    device_id = rank % np.cuda.runtime.getDeviceCount()
     np.cuda.Device(device_id).use()
 
 par1 = {'ny': 101, 'nx': 101, 'imag': 0, 'dtype': np.float64}

tests/test_stackedarray.py

@@ -23,11 +23,7 @@
 np.random.seed(42)
 rank = MPI.COMM_WORLD.Get_rank()
 if backend == "cupy":
-    device_count = np.cuda.runtime.getDeviceCount()
-    device_id = int(
-        os.environ.get("OMPI_COMM_WORLD_LOCAL_RANK")
-        or rank % np.cuda.runtime.getDeviceCount()
-    )
+    device_id = rank % np.cuda.runtime.getDeviceCount()
     np.cuda.Device(device_id).use()