tests: Add tests for using CUDA_VISIBLE_DEVICES and similar

Author: EdCaunt

devito/operator/operator.py

@@ -1390,21 +1390,36 @@ def _get_nbytes(self, i):
         return nbytes
 
     @cached_property
-    def visible_devices(self):
+    def _visible_devices(self):
         device_vars = (
             'CUDA_VISIBLE_DEVICES',
             'ROCR_VISIBLE_DEVICES',
             'HIP_VISIBLE_DEVICES'
         )
         for v in device_vars:
-            if v in os.environ:
-                try:
-                    return tuple(int(i) for i in os.environ[v].split(','))
-                except ValueError:
-                    # Visible devices set via UUIDs or other non-integer identifiers
-                    continue
+            try:
+                return tuple(int(i) for i in os.environ[v].split(','))
+            except (ValueError, KeyError):
+                # Environment variable not set or visible devices set via UUIDs
+                # or other non-integer identifiers.
+                continue
 
         return None
+    
+    @cached_property
+    def _physical_deviceid(self):
+        if isinstance(self.platform, Device):
+            # Get the physical device ID (as CUDA_VISIBLE_DEVICES may be set)
+            rank = self.comm.Get_rank() if self.comm != MPI.COMM_NULL else 0
+
+            logical_deviceid = max(self.get('deviceid', 0), 0) + rank
+            if self._visible_devices is not None:
+                return self._visible_devices[logical_deviceid]
+            else:
+                return logical_deviceid
+
+        else:
+            return None
 
     @cached_property
     def nbytes_avail_mapper(self):
@@ -1420,16 +1435,7 @@ def nbytes_avail_mapper(self):
 
         # The amount of space available on the device
         if isinstance(self.platform, Device):
-            # Get the physical device ID (as CUDA_VISIBLE_DEVICES may be set)
-            rank = self.comm.Get_rank() if self.comm != MPI.COMM_NULL else 0
-
-            logical_deviceid = max(self.get('deviceid', 0), 0) + rank
-            if self.visible_devices is not None:
-                physical_deviceid = self.visible_devices[logical_deviceid]
-            else:
-                physical_deviceid = logical_deviceid
-
-            mapper[device_layer] = self.platform.memavail(deviceid=physical_deviceid)
+            mapper[device_layer] = self.platform.memavail(deviceid=self._physical_deviceid)
 
         # The amount of space available on the host
         try:

devito/parameters.py

@@ -8,7 +8,7 @@
 from devito.tools import Signer, filter_ordered
 
 __all__ = ['configuration', 'init_configuration', 'print_defaults', 'print_state',
-           'switchconfig']
+           'switchconfig', 'switchenv']
 
 # Be EXTREMELY careful when writing to a Parameters dictionary
 # Read here for reference: http://wiki.c2.com/?GlobalVariablesAreBad
@@ -224,7 +224,22 @@ def init_configuration(configuration=configuration, env_vars_mapper=env_vars_map
     configuration.initialize()
 
 
-class switchconfig:
+class abstractswitch:
+
+    """
+    Abstract class for switch(whatever) decorators.
+    """
+
+    def __call__(self, func, *args, **kwargs):
+        @wraps(func)
+        def wrapper(*args, **kwargs):
+            with self:
+                result = func(*args, **kwargs)
+            return result
+        return wrapper
+
+
+class switchconfig(abstractswitch):
 
     """
     Decorator or context manager to temporarily change `configuration` parameters.
@@ -250,14 +265,29 @@ def __exit__(self, exc_type, exc_val, exc_tb):
             except ValueError:
                 # E.g., `platform` and `compiler` will end up here
                 super(Parameters, configuration).__setitem__(k, self.previous[k])
+    
 
-    def __call__(self, func, *args, **kwargs):
-        @wraps(func)
-        def wrapper(*args, **kwargs):
-            with self:
-                result = func(*args, **kwargs)
-            return result
-        return wrapper
+class switchenv(abstractswitch):
+    """
+    Decorator to temporarily change environment variables.
+    Adapted from https://stackoverflow.com/questions/2059482/
+    """
+
+    def __init__(self, condition=True, **params):
+        self.previous = dict(environ)
+
+        if condition:
+            # Environment variables are essentially always uppercase
+            self.params = {k.upper(): v for k, v in params.items()}
+        else:   
+            self.params = params
+
+    def __enter__(self, condition=True, **params):
+        environ.update(self.params)
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        environ.clear()
+        environ.update(self.previous)
 
 
 def print_defaults():

tests/test_gpu_common.py

@@ -1,6 +1,7 @@
 import cloudpickle as pickle
 
 import pytest
+import os
 import numpy as np
 import sympy
 import scipy.sparse
@@ -10,7 +11,7 @@
                     Dimension, MatrixSparseTimeFunction, SparseTimeFunction,
                     SubDimension, SubDomain, SubDomainSet, TimeFunction, exp,
                     Operator, configuration, switchconfig, TensorTimeFunction,
-                    Buffer, assign)
+                    Buffer, assign, switchenv)
 from devito.arch import get_gpu_info, get_cpu_info, Device, Cpu64
 from devito.exceptions import InvalidArgument
 from devito.ir import (Conditional, Expression, Section, FindNodes, FindSymbols,
@@ -74,6 +75,89 @@ def test_autopad_with_platform_switch(self):
             assert f.shape_allocated[1] == 64
 
 
+class TestEnvironmentVariables:
+    """
+    Test that environment variables are correctly handled.
+    """
+
+    @pytest.mark.parametrize('env_variables', [{"cuda_visible_devices": "1"},
+                                               {"cuda_visible_devices": "1,2"},
+                                               {"cuda_visible_devices": "1,0"},
+                                               {"rocr_visible_devices": "1"},
+                                               {"hip_visible_devices": " 1"}])
+    def test_visible_devices(self, env_variables):
+        """
+        Test that physical device IDs used for querying memory on a device via
+        nvidia-smi correctly account for visible-device environment variables.
+        """
+        grid = Grid(shape=(10, 10))
+        u = Function(name='u', grid=grid)
+
+        eq = Eq(u, u+1)
+
+        with switchenv(**env_variables):
+            op1 = Operator(eq)
+
+            argmap1 = op1.arguments()
+            # All variants in parameterisation should yield deviceid 1
+            assert argmap1._physical_deviceid == 1
+
+        # Make sure the switchenv doesn't somehow persist
+        for i in ("CUDA", "ROCR", "HIP"):
+            assert f"{i}_VISIBLE_DEVICES" not in os.environ
+
+        # Check that physical deviceid is 0 when no environment variables set 
+        op2 = Operator(eq)
+
+        argmap2 = op2.arguments()
+        # Default physical deviceid expected to be 0
+        assert argmap2._physical_deviceid == 0
+
+
+    @pytest.mark.parallel(mode=2)
+    @pytest.mark.parametrize('visible_devices', ["1,2", "1,0", "0,2,3"])
+    def test_visible_devices_mpi(self, visible_devices, mode):
+        """
+        Test that physical device IDs used for querying memory on a device via
+        nvidia-smi correctly account for visible-device environment variables
+        when using MPI.
+        """
+
+        grid = Grid(shape=(10, 10))
+        rank = grid.distributor.myrank
+        u = Function(name='u', grid=grid)
+
+        eq = Eq(u, u+1)
+
+        with switchenv(cuda_visible_devices=visible_devices):
+            op1 = Operator(eq)
+            argmap1 = op1.arguments()
+
+            devices = [int(i) for i in visible_devices.split(',')]
+
+            assert argmap1._physical_deviceid == devices[rank]
+
+        # In default case, physical deviceid will equal rank
+        op2 = Operator(eq)
+        argmap2 = op2.arguments()
+        assert argmap2._physical_deviceid == rank
+
+    def test_visible_devices_with_devito_deviceid(self):
+        """Test interaction between CUDA_VISIBLE_DEVICES and DEVITO_DEVICEID"""
+        grid = Grid(shape=(10, 10))
+        u = Function(name='u', grid=grid)
+
+        eq = Eq(u, u+1)
+
+        with switchenv(cuda_visible_devices="1,3"), switchconfig(deviceid=1):
+            op = Operator(eq)
+
+            argmap = op.arguments()
+            # deviceid should see the world from within CUDA_VISIBLE_DEVICES
+            # So should be the second of the two visible devices specified (3)
+            assert argmap._physical_deviceid == 3
+
+
 class TestCodeGeneration:
 
     def test_maxpar_option(self):