add an environment flag to raise an error if result is not cached (brain-score#8)

Author: mschrimpf

README.md

@@ -34,3 +34,4 @@ changed through the environment variable `RESULTCACHING_HOME`.
 | 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 |
 |                       | * `'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) |
+| RESULTCACHING_CACHEDONLY | If enabled, raises an error when trying to run a function that does not have its result already cached (follows the same matching rules as `RESULTCACHING_DISABLE`) |

result_caching/__init__.py

@@ -33,12 +33,25 @@ def get_function_identifier(function, call_args):
 
 def is_enabled(function_identifier):
     disable = os.getenv('RESULTCACHING_DISABLE', '0')
-    if disable == '1':
-        return False
-    if disable == '':
+    return not _match_identifier(function_identifier, disable)
+
+
+def cached_only(function_identifier):
+    cachedonly = os.getenv('RESULTCACHING_CACHEDONLY', '0')
+    return _match_identifier(function_identifier, cachedonly)
+
+
+def _match_identifier(function_identifier, match_value):
+    if match_value == '1':
         return True
-    disabled_modules = disable.split(',')
-    return not any(function_identifier.startswith(disabled_module) for disabled_module in disabled_modules)
+    if match_value == '':
+        return False
+    disabled_modules = match_value.split(',')
+    return any(function_identifier.startswith(disabled_module) for disabled_module in disabled_modules)
+
+
+class NotCachedError(Exception):
+    pass
 
 
 class _Storage(object):
@@ -58,6 +71,8 @@ def wrapper(*args, **kwargs):
             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)
+            if cached_only(function_identifier):
+                raise NotCachedError(f"No result stored for '{function_identifier}'")
             self._logger.debug("Running function: {}".format(function_identifier))
             result = function(*args, **kwargs)
             if is_enabled(function_identifier):
@@ -197,6 +212,9 @@ def wrapper(*args, **kwargs):
                     reduced_call_args = {**non_variable_call_args, **infile_missing_call_args}
                     self._logger.debug(f"Computing missing: {reduced_call_args}")
             if reduced_call_args:
+                if cached_only(function_identifier):
+                    raise NotCachedError(f"The following arguments for '{function_identifier}' "
+                                         f"are not stored: {reduced_call_args}")
                 # run function if some args are uncomputed
                 self._logger.debug(f"Running function: {function_identifier}")
                 result = function(**reduced_call_args)
@@ -283,6 +301,9 @@ def wrapper(*args, **kwargs):
                     reduced_call_args = {**non_variable_call_args, **missing_call_args}
                     self._logger.debug(f"Computing missing: {reduced_call_args}")
             if reduced_call_args:
+                if cached_only(function_identifier):
+                    raise NotCachedError(f"The following arguments for '{function_identifier}' "
+                                         f"are not stored: {reduced_call_args}")
                 self._logger.debug(f"Running function: {function_identifier}")
                 # run function if some args are uncomputed
                 result = function(**reduced_call_args)

tests/test___init__.py

@@ -6,7 +6,7 @@
 import pytest
 import xarray as xr
 
-from result_caching import store_xarray, store, cache, store_dict, get_function_identifier
+from result_caching import store_xarray, store, cache, store_dict, get_function_identifier, NotCachedError
 
 
 class TestFunctionIdentifier:
@@ -296,6 +296,35 @@ def func2(x):
 
         os.environ['RESULTCACHING_DISABLE'] = previous_disable_value
 
+    def test_cachedonly_specific(self):
+        previous_cached_value = os.getenv('RESULTCACHING_CACHEDONLY', '')
+        with tempfile.TemporaryDirectory() as storage_dir:
+            os.environ['RESULTCACHING_HOME'] = storage_dir
+
+            @store()
+            def func1(x):
+                return x
+
+            @store()
+            def func2(x):
+                return x
+
+            # when allowing only cached results from func2, func1 should work, but func2 should err
+            os.environ['RESULTCACHING_CACHEDONLY'] = 'test___init__.func2'
+            assert func1(1) == 1
+            with pytest.raises(NotCachedError):
+                func2(2)
+
+            # when allow reruns, func2 should work again
+            os.environ['RESULTCACHING_CACHEDONLY'] = ''
+            assert func2(2) == 2
+
+            # when now only allowing cached results again, func2 should work because results are already cached
+            os.environ['RESULTCACHING_CACHEDONLY'] = 'test___init__.func2'
+            assert func2(2) == 2
+
+        os.environ['RESULTCACHING_CACHEDONLY'] = previous_cached_value
+
 
 class TestCache:
     def test(self):