Merge pull request #2791 from devitocodes/deviceid-errors

test: Enhance tests for interaction between deviceid and visible devices

Author: FabioLuporini

devito/arch/archinfo.py

@@ -497,7 +497,7 @@ def get_visible_devices():
     )
     for v in device_vars:
         try:
-            return tuple(int(i) for i in os.environ[v].split(','))
+            return v, tuple(int(i) for i in os.environ[v].split(','))
         except ValueError:
             # Visible devices set via UUIDs or other non-integer identifiers.
             warning("Setting visible devices via UUIDs or other non-integer"
@@ -507,7 +507,7 @@ def get_visible_devices():
             # Environment variable not set
             continue
 
-    return None
+    return None, None
 
 
 @memoized_func

devito/operator/operator.py

@@ -1398,11 +1398,18 @@ def _physical_deviceid(self):
                 rank = self.comm.Get_rank() if self.comm != MPI.COMM_NULL else 0
                 logical_deviceid = rank
 
-            visible_devices = get_visible_devices()
+            visible_device_var, visible_devices = get_visible_devices()
             if visible_devices is None:
                 return logical_deviceid
             else:
-                return visible_devices[logical_deviceid]
+                try:
+                    return visible_devices[logical_deviceid]
+                except IndexError:
+                    errmsg = (f"A deviceid value of {logical_deviceid} is not valid "
+                              f"with {visible_device_var}={visible_devices}. Note that "
+                              "deviceid corresponds to the logical index within the "
+                              "visible devices, not the physical device index.")
+                    raise ValueError(errmsg)
         else:
             return None
 

tests/test_gpu_common.py

@@ -149,6 +149,12 @@ def test_visible_devices_with_devito_deviceid(self):
             # So should be the second of the two visible devices specified (3)
             assert argmap._physical_deviceid == 3
 
+        with switchenv({'CUDA_VISIBLE_DEVICES': "1"}), switchconfig(deviceid=0):
+            op1 = Operator(eq)
+
+            argmap1 = op1.arguments()
+            assert argmap1._physical_deviceid == 1
+
     @pytest.mark.parallel(mode=2)
     def test_deviceid_per_rank(self, mode):
         """