compiler: Refine handling of device environment variables and associated tests

Author: EdCaunt

devito/operator/operator.py

@@ -1435,7 +1435,8 @@ def nbytes_avail_mapper(self):
 
         # The amount of space available on the device
         if isinstance(self.platform, Device):
-            mapper[device_layer] = self.platform.memavail(deviceid=self._physical_deviceid)
+            mapper[device_layer] = \
+                self.platform.memavail(deviceid=self._physical_deviceid)
 
         # The amount of space available on the host
         try:

devito/parameters.py

@@ -265,7 +265,7 @@ 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])
-    
+
 
 class switchenv(abstractswitch):
     """
@@ -279,7 +279,7 @@ def __init__(self, condition=True, **params):
         if condition:
             # Environment variables are essentially always uppercase
             self.params = {k.upper(): v for k, v in params.items()}
-        else:   
+        else:
             self.params = params
 
     def __enter__(self, condition=True, **params):

tests/test_gpu_common.py

@@ -95,6 +95,9 @@ def test_visible_devices(self, env_variables):
 
         eq = Eq(u, u+1)
 
+        # Save previous environment to verify switchenv works as intended
+        previous_environ = dict(os.environ)
+
         with switchenv(**env_variables):
             op1 = Operator(eq)
 
@@ -103,8 +106,7 @@ def test_visible_devices(self, env_variables):
             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
+        assert dict(os.environ) == previous_environ
 
         # Check that physical deviceid is 0 when no environment variables set 
         op2 = Operator(eq)
@@ -113,7 +115,6 @@ def test_visible_devices(self, env_variables):
         # 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):