Merge pull request #2746 from devitocodes/deviceid

compiler: Make nbytes available mapper aware of visible devices environment variables

Author: FabioLuporini

devito/arch/archinfo.py

@@ -16,10 +16,11 @@
 from devito.logger import warning
 from devito.tools import as_tuple, all_equal, memoized_func
 
-__all__ = ['platform_registry', 'get_cpu_info', 'get_gpu_info', 'get_nvidia_cc',
-           'get_cuda_path', 'get_hip_path', 'check_cuda_runtime', 'get_m1_llvm_path',
-           'get_advisor_path', 'Platform', 'Cpu64', 'Intel64', 'IntelSkylake', 'Amd',
-           'Arm', 'Power', 'Device', 'NvidiaDevice', 'AmdDevice', 'IntelDevice',
+__all__ = ['platform_registry', 'get_cpu_info', 'get_gpu_info', 'get_visible_devices',
+           'get_nvidia_cc', 'get_cuda_path', 'get_hip_path', 'check_cuda_runtime',
+           'get_m1_llvm_path', 'get_advisor_path', 'Platform', 'Cpu64', 'Intel64',
+           'IntelSkylake', 'Amd', 'Arm', 'Power', 'Device', 'NvidiaDevice',
+           'AmdDevice', 'IntelDevice',
            # Brand-agnostic
            'ANYCPU', 'ANYGPU',
            # Intel CPUs
@@ -488,6 +489,27 @@ def parse_product_arch():
     return None
 
 
+def get_visible_devices():
+    device_vars = (
+        'CUDA_VISIBLE_DEVICES',
+        'ROCR_VISIBLE_DEVICES',
+        'HIP_VISIBLE_DEVICES'
+    )
+    for v in device_vars:
+        try:
+            return 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"
+                    " identifiers is currently unsupported: environment variable"
+                    f" {v}={os.environ[v]} ignored.")
+        except KeyError:
+            # Environment variable not set
+            continue
+
+    return None
+
+
 @memoized_func
 def get_nvidia_cc():
     libnames = ('libcuda.so', 'libcuda.dylib', 'cuda.dll')

devito/operator/operator.py

@@ -10,7 +10,8 @@
 import sympy
 import numpy as np
 
-from devito.arch import ANYCPU, Device, compiler_registry, platform_registry
+from devito.arch import (ANYCPU, Device, compiler_registry, platform_registry,
+                         get_visible_devices)
 from devito.data import default_allocator
 from devito.exceptions import (CompilationError, ExecutionError, InvalidArgument,
                                InvalidOperator)
@@ -1388,6 +1389,23 @@ def _get_nbytes(self, i):
 
         return nbytes
 
+    @cached_property
+    def _physical_deviceid(self):
+        if isinstance(self.platform, Device):
+            # Get the physical device ID (as CUDA_VISIBLE_DEVICES may be set)
+            logical_deviceid = self.get('deviceid', -1)
+            if logical_deviceid < 0:
+                rank = self.comm.Get_rank() if self.comm != MPI.COMM_NULL else 0
+                logical_deviceid = rank
+
+            visible_devices = get_visible_devices()
+            if visible_devices is None:
+                return logical_deviceid
+            else:
+                return visible_devices[logical_deviceid]
+        else:
+            return None
+
     @cached_property
     def nbytes_avail_mapper(self):
         """
@@ -1402,8 +1420,8 @@ def nbytes_avail_mapper(self):
 
         # The amount of space available on the device
         if isinstance(self.platform, Device):
-            deviceid = max(self.get('deviceid', 0), 0)
-            mapper[device_layer] = self.platform.memavail(deviceid=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.
@@ -251,13 +266,28 @@ def __exit__(self, exc_type, exc_val, exc_tb):
                 # 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,111 @@ def test_autopad_with_platform_switch(self):
             assert f.shape_allocated[1] == 64
 
 
+class TestDeviceID:
+    """
+    Test that device IDs and associated environment variables such as
+    CUDA_VISIBLE_DEVICES 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)
+
+        # Save previous environment to verify switchenv works as intended
+        previous_environ = dict(os.environ)
+
+        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
+        assert dict(os.environ) == previous_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
+
+    @pytest.mark.parallel(mode=2)
+    def test_deviceid_per_rank(self, mode):
+        """
+        Test that Device IDs set by the user on a per-rank basis do not
+        get modifed.
+        """
+        # Reversed order to ensure it is different to default
+        user_set_deviceids = (1, 0)
+
+        grid = Grid(shape=(10, 10))
+        u = Function(name='u', grid=grid)
+
+        rank = grid.distributor.myrank
+        deviceid = user_set_deviceids[rank]
+
+        op = Operator(Eq(u, u+1))
+
+        argmap = op.arguments(deviceid=deviceid)
+        assert argmap._physical_deviceid == deviceid
+
+
 class TestCodeGeneration:
 
     def test_maxpar_option(self):