implement disabling specific modules (brain-score#6)

* implement disabling specific modules

* make example more specific wrt. brain-score

Author: mschrimpf

README.md

@@ -32,4 +32,5 @@ changed through the environment variable `RESULTCACHING_HOME`.
 | Variable | description |
 |-----------------------|----------------------------------------------------------------------------------|
 | RESULTCACHING_HOME | directory to cache results (benchmark ceilings) in, `~/.result_caching` by default |
-| RESULTCACHING_DISABLE | `'1'` to disable loading and saving of results, functions will be called directly |
+| RESULTCACHING_DISABLE | * `'1'` to disable loading and saving of results, functions will be called directly |
+|                       | * `'candidate_models.score_model,model_tools.activations`' to disable loading and saving of function identifiers starting with one of the specifiers separated by a comma (e.g. any package or function inside `model_tools.activations` will not be considered) |

result_caching/__init__.py

@@ -31,8 +31,14 @@ def get_function_identifier(function, call_args):
     return function_identifier
 
 
-def is_enabled():
-    return os.getenv('RESULTCACHING_DISABLE', '0') != '1'
+def is_enabled(function_identifier):
+    disable = os.getenv('RESULTCACHING_DISABLE', '0')
+    if disable == '1':
+        return False
+    if disable == '':
+        return True
+    disabled_modules = disable.split(',')
+    return not any(function_identifier.startswith(disabled_module) for disabled_module in disabled_modules)
 
 
 class _Storage(object):
@@ -49,12 +55,12 @@ def __call__(self, function):
         def wrapper(*args, **kwargs):
             call_args = self.getcallargs(function, *args, **kwargs)
             function_identifier = self.get_function_identifier(function, call_args)
-            if is_enabled() and self.is_stored(function_identifier):
+            if is_enabled(function_identifier) and self.is_stored(function_identifier):
                 self._logger.debug("Loading from storage: {}".format(function_identifier))
                 return self.load(function_identifier)
             self._logger.debug("Running function: {}".format(function_identifier))
             result = function(*args, **kwargs)
-            if is_enabled():
+            if is_enabled(function_identifier):
                 self._logger.debug("Saving to storage: {}".format(function_identifier))
                 self.save(result, function_identifier)
             return result
@@ -258,7 +264,7 @@ def wrapper(*args, **kwargs):
                                 if key in self._combine_fields}
             function_identifier = self.get_function_identifier(function, call_args)
             stored_result, reduced_call_args = None, call_args
-            if is_enabled() and self.is_stored(function_identifier):
+            if is_enabled(function_identifier) and self.is_stored(function_identifier):
                 self._logger.debug(f"Loading from storage: {function_identifier}")
                 stored_result = self.load(function_identifier)
                 missing_call_args = self.filter_coords(infile_call_args, stored_result) if not self._sub_fields \

tests/test___init__.py

@@ -1,5 +1,6 @@
 import os
 import tempfile
+from collections import defaultdict
 
 import numpy as np
 import pytest
@@ -210,8 +211,8 @@ def __repr__(self):
             # second call returns same thing and doesn't actually call function again
             assert c.f(1) == 1
 
-    def test_disable_store(self):
-        previous_disable_value = os.getenv('RESULTCACHING_DISABLE', '0')
+    def test_disable_all(self):
+        previous_disable_value = os.getenv('RESULTCACHING_DISABLE', '')
         with tempfile.TemporaryDirectory() as storage_dir:
             os.environ['RESULTCACHING_HOME'] = storage_dir
             os.environ['RESULTCACHING_DISABLE'] = '1'
@@ -232,6 +233,69 @@ def func(x):
 
         os.environ['RESULTCACHING_DISABLE'] = previous_disable_value
 
+    def test_disable_specific(self):
+        previous_disable_value = os.getenv('RESULTCACHING_DISABLE', '')
+        with tempfile.TemporaryDirectory() as storage_dir:
+            os.environ['RESULTCACHING_HOME'] = storage_dir
+            os.environ['RESULTCACHING_DISABLE'] = 'test___init__.func1'
+
+            function_calls = defaultdict(lambda: 0)
+
+            @store()
+            def func1(x):
+                nonlocal function_calls
+                function_calls[1] += 1
+                return x
+
+            @store()
+            def func2(x):
+                nonlocal function_calls
+                function_calls[2] += 1
+                return x
+
+            assert func1(1) == 1
+            assert function_calls[1] == 1
+            assert not os.listdir(storage_dir)
+            assert func1(1) == 1
+            assert function_calls[1] == 2
+
+            assert func2(1) == 1
+            assert function_calls[2] == 1
+            assert func2(1) == 1
+            assert function_calls[2] == 1
+
+        os.environ['RESULTCACHING_DISABLE'] = previous_disable_value
+
+    def test_disable_module(self):
+        previous_disable_value = os.getenv('RESULTCACHING_DISABLE', '')
+        with tempfile.TemporaryDirectory() as storage_dir:
+            os.environ['RESULTCACHING_HOME'] = storage_dir
+            os.environ['RESULTCACHING_DISABLE'] = 'test___init__'
+
+            function_calls = defaultdict(lambda: 0)
+
+            @store()
+            def func1(x):
+                nonlocal function_calls
+                function_calls[1] += 1
+                return x
+
+            @store()
+            def func2(x):
+                nonlocal function_calls
+                function_calls[2] += 1
+                return x
+
+            assert func1(1) == 1
+            assert func1(1) == 1
+            assert function_calls[1] == 2
+            assert func2(1) == 1
+            assert func2(1) == 1
+            assert function_calls[2] == 2
+            assert not os.listdir(storage_dir)
+
+        os.environ['RESULTCACHING_DISABLE'] = previous_disable_value
+
 
 class TestCache:
     def test(self):