Merge pull request #2769 from devitocodes/JDBetteridge/remove_switchenv

misc: Un-cargo-cult switchenv

Author: FabioLuporini

devito/parameters.py

@@ -1,7 +1,7 @@
 """The parameters dictionary contains global parameter settings."""
-
+from abc import ABC, abstractmethod
 from collections import OrderedDict
-from os import environ
+import os
 from functools import wraps
 
 from devito.logger import info, warning
@@ -170,24 +170,24 @@ def _signature_items(self):
 def init_configuration(configuration=configuration, env_vars_mapper=env_vars_mapper,
                        env_vars_deprecated=env_vars_deprecated):
     # Populate `configuration` with user-provided options
-    if environ.get('DEVITO_CONFIG') is None:
+    if os.environ.get('DEVITO_CONFIG') is None:
         # It is important to configure `platform`, `compiler`, and the rest, in this order
         process_order = filter_ordered(['platform', 'compiler'] +
                                        list(env_vars_mapper.values()))
         queue = sorted(env_vars_mapper.items(), key=lambda i: process_order.index(i[1]))
-        unprocessed = OrderedDict([(v, environ.get(k, configuration._defaults[v]))
+        unprocessed = OrderedDict([(v, os.environ.get(k, configuration._defaults[v]))
                                    for k, v in queue])
 
         # Handle deprecated env vars
         mapper = dict(queue)
         for k, (v, msg) in env_vars_deprecated.items():
-            if environ.get(k):
+            if os.environ.get(k):
                 warning(f"`{k}` is deprecated. {msg}")
-                if environ.get(v):
+                if os.environ.get(v):
                     warning(f"Both `{k}` and `{v}` set. Ignoring `{k}`")
                 else:
-                    warning(f"Setting `{v}={environ[k]}`")
-                    unprocessed[mapper[v]] = environ[k]
+                    warning(f"Setting `{v}={os.environ[k]}`")
+                    unprocessed[mapper[v]] = os.environ[k]
     else:
         # Attempt reading from the specified configuration file
         raise NotImplementedError("Devito doesn't support configuration via file yet.")
@@ -224,13 +224,25 @@ def init_configuration(configuration=configuration, env_vars_mapper=env_vars_map
     configuration.initialize()
 
 
-class abstractswitch:
-
+class SwitchDecorator(ABC):
     """
-    Abstract class for switch(whatever) decorators.
+    Abstract base class that turns a context manager class into a decorator.
     """
+    @abstractmethod
+    def __init__(self, *args, **kwargs):
+        pass
+
+    @abstractmethod
+    def __enter__(self):
+        pass
+
+    @abstractmethod
+    def __exit__(self, exc_type, exc_val, traceback):
+        pass
 
     def __call__(self, func, *args, **kwargs):
+        """ The call method turns the context manager class into a decorator
+        """
         @wraps(func)
         def wrapper(*args, **kwargs):
             with self:
@@ -239,26 +251,24 @@ def wrapper(*args, **kwargs):
         return wrapper
 
 
-class switchconfig(abstractswitch):
-
+class switchconfig(SwitchDecorator):
     """
     Decorator or context manager to temporarily change `configuration` parameters.
     """
-
     def __init__(self, condition=True, **params):
         if condition:
             self.params = {k.replace('_', '-'): v for k, v in params.items()}
         else:
             self.params = {}
         self.previous = {}
 
-    def __enter__(self, condition=True, **params):
+    def __enter__(self):
         self.previous = {}
         for k, v in self.params.items():
             self.previous[k] = configuration[k]
             configuration[k] = v
 
-    def __exit__(self, exc_type, exc_val, exc_tb):
+    def __exit__(self, exc_type, exc_val, traceback):
         for k, v in self.params.items():
             try:
                 configuration[k] = self.previous[k]
@@ -267,27 +277,23 @@ def __exit__(self, exc_type, exc_val, exc_tb):
                 super(Parameters, configuration).__setitem__(k, self.previous[k])
 
 
-class switchenv(abstractswitch):
+class switchenv(SwitchDecorator):
     """
-    Decorator to temporarily change environment variables.
-    Adapted from https://stackoverflow.com/questions/2059482/
-    """
-
-    def __init__(self, condition=True, **params):
-        self.previous = dict(environ)
+    Temporarily set environment variables from a dictionary
 
-        if condition:
-            # Environment variables are essentially always uppercase
-            self.params = {k.upper(): v for k, v in params.items()}
-        else:
-            self.params = params
+    Note: This does not propagate any environment variables that change inside
+    the context manager, so should be used cautiously.
+    """
+    def __init__(self, params):
+        self.previous = dict(os.environ)
+        self.params = params
 
-    def __enter__(self, condition=True, **params):
-        environ.update(self.params)
+    def __enter__(self):
+        os.environ.update(self.params)
 
-    def __exit__(self, exc_type, exc_val, exc_tb):
-        environ.clear()
-        environ.update(self.previous)
+    def __exit__(self, exc_type, exc_val, traceback):
+        os.environ.clear()
+        os.environ.update(self.previous)
 
 
 def print_defaults():

tests/test_gpu_common.py

@@ -1,7 +1,6 @@
 import cloudpickle as pickle
 
 import pytest
-import os
 import numpy as np
 import sympy
 import scipy.sparse
@@ -81,11 +80,11 @@ class TestDeviceID:
     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"}])
+    @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
@@ -96,19 +95,12 @@ 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):
+        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)
 
@@ -131,12 +123,10 @@ def test_visible_devices_mpi(self, visible_devices, mode):
 
         eq = Eq(u, u+1)
 
-        with switchenv(cuda_visible_devices=visible_devices):
+        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
@@ -151,7 +141,7 @@ def test_visible_devices_with_devito_deviceid(self):
 
         eq = Eq(u, u+1)
 
-        with switchenv(cuda_visible_devices="1,3"), switchconfig(deviceid=1):
+        with switchenv({'CUDA_VISIBLE_DEVICES': "1,3"}), switchconfig(deviceid=1):
             op = Operator(eq)
 
             argmap = op.arguments()

tests/test_tools.py

@@ -1,10 +1,11 @@
+import os
 import numpy as np
 import pytest
 from sympy.abc import a, b, c, d, e
 
 import time
 
-from devito import Operator, Eq
+from devito import Operator, Eq, switchenv
 from devito.tools import (UnboundedMultiTuple, ctypes_to_cstr, toposort,
                           filter_ordered, transitive_closure, UnboundTuple,
                           CacheInstances)
@@ -209,3 +210,18 @@ def __init__(self, value: int):
         # Cache should be cleared after Operator construction
         cache_size = Object._instance_cache.cache_info()[-1]
         assert cache_size == 0
+
+
+def test_switchenv():
+    # Save previous environment
+    previous_environ = dict(os.environ)
+
+    # Check a temporary variable is set inside the context manager
+    with switchenv({'TEST_VAR': 'foo'}):
+        assert os.environ['TEST_VAR'] == 'foo'
+
+    # Check a temporary variable is unset inside the context manager
+    assert os.environ.get('TEST_VAR') is None
+
+    # Make sure the switchenv does not persist to verify switchenv works as intended
+    assert dict(os.environ) == previous_environ