Merge pull request #503 from moorepants/cache-binary

Enables caching of the compiled code if tmp_dir is set.

Author: moorepants

opty/direct_collocation.py

@@ -1149,7 +1149,12 @@ def __init__(self, equations_of_motion, state_symbols,
         tmp_dir : string, optional
             If you want to see the generated Cython and C code for the
             constraint and constraint Jacobian evaluations, pass in a path to a
-            directory here.
+            directory here. Additionally, if this temporary directory is set to
+            an existing populated directory and the equations of motion have
+            not changed relative to prior instantiations of this class, the
+            compilation step will be skipped if equivalent compiled modules are
+            already present and cached. This may save significant computational
+            time when repeatedly using the same set of equations of motion.
         integration_method : string, optional
             The integration method to use, either ``backward euler`` or
             ``midpoint``.

opty/tests/test_direct_collocation.py

@@ -1,3 +1,6 @@
+import os
+import shutil
+import tempfile
 from collections import OrderedDict
 
 import numpy as np
@@ -408,43 +411,81 @@ def test_pendulum():
     np.testing.assert_allclose(prob._upp_con_bounds, expected_upp_con_bounds)
 
 
-def test_Problem():
+def TestProblem():
 
-    m, c, k, t = sym.symbols('m, c, k, t')
-    x, v, f = [s(t) for s in sym.symbols('x, v, f', cls=sym.Function)]
+    def setup_method(self):
+        self.tmp_dir = tempfile.mkdtemp("opty_cache_test")
+        if os.path.exists(self.tmp_dir):
+            shutil.rmtree(self.tmp_dir)
+        os.mkdir(self.tmp_dir)
+
+    def teardown_method(self):
+        if os.path.exists(self.tmp_dir):
+            shutil.rmtree(self.tmp_dir)
+
+    def test_problem(self):
 
-    state_symbols = (x, v)
+        m, c, k, t = sym.symbols('m, c, k, t')
+        x, v, f = [s(t) for s in sym.symbols('x, v, f', cls=sym.Function)]
 
-    interval_value = 0.01
+        state_symbols = (x, v)
 
-    eom = sym.Matrix([x.diff() - v,
-                      m * v.diff() + c * v + k * x - f])
+        interval_value = 0.01
 
-    prob = Problem(lambda x: 1.0,
-                   lambda x: x,
-                   eom,
-                   state_symbols,
-                   2,
-                   interval_value,
-                   time_symbol=t,
-                   bounds={x: (-10.0, 10.0),
-                           f: (-8.0, 8.0),
-                           m: (-1.0, 1.0),
-                           c: (-0.5, 0.5)})
+        eom = sym.Matrix([x.diff() - v,
+                          m * v.diff() + c * v + k * x - f])
 
-    INF = 10e19
-    expected_lower = np.array([-10.0, -10.0,
-                               -INF, -INF,
-                               -8.0, -8.0,
-                               -0.5, -INF, -1.0])
-    np.testing.assert_allclose(prob.lower_bound, expected_lower)
-    expected_upper = np.array([10.0, 10.0,
-                               INF, INF,
-                               8.0, 8.0,
-                               0.5, INF, 1.0])
-    np.testing.assert_allclose(prob.upper_bound, expected_upper)
+        prob = Problem(
+            lambda x: 1.0,
+            lambda x: x,
+            eom,
+            state_symbols,
+            2,
+            interval_value,
+            time_symbol=t,
+            bounds={
+                x: (-10.0, 10.0),
+                f: (-8.0, 8.0),
+                m: (-1.0, 1.0),
+                c: (-0.5, 0.5),
+            },
+            tmp_dir=self.tmp_dir)
+
+        # Only two modules should be generated
+        c_file_list = [f for f in os.listdir(self.tmp_dir) if
+                       f.endswith('_c.c')]
+        assert len(c_file_list) == 2
+
+        INF = 10e19
+        expected_lower = np.array([-10.0, -10.0,
+                                   -INF, -INF,
+                                   -8.0, -8.0,
+                                   -0.5, -INF, -1.0])
+        np.testing.assert_allclose(prob.lower_bound, expected_lower)
+        expected_upper = np.array([10.0, 10.0,
+                                   INF, INF,
+                                   8.0, 8.0,
+                                   0.5, INF, 1.0])
+        np.testing.assert_allclose(prob.upper_bound, expected_upper)
+
+        assert prob.collocator.num_instance_constraints == 0
+
+        # run Problem again to see if the cache worked.
+        prob = Problem(
+            lambda x: 1.0,
+            lambda x: x,
+            eom,
+            state_symbols,
+            4,
+            interval_value,
+            time_symbol=t,
+            tmp_dir=self.tmp_dir,
+        )
 
-    assert prob.collocator.num_instance_constraints == 0
+        # no more C files should have been generated
+        c_file_list = [f for f in os.listdir(self.tmp_dir) if
+                       f.endswith('_c.c')]
+        assert len(c_file_list) == 2
 
 
 class TestConstraintCollocator():

opty/utils.py

@@ -15,6 +15,7 @@
 from timeit import default_timer as timer
 import logging
 import locale
+import hashlib
 
 import numpy as np
 import sympy as sm
@@ -79,6 +80,11 @@ def ccode(expr, assign_to=None, **settings):
 
 def _forward_jacobian(expr, wrt):
 
+    # NOTE : free_symbols are sets and are not guaranteed to be in the same
+    # order, so sympy.ordered() is used throughout to ensure a deterministic
+    # behavior. This is important for the binary caching to work as it hashes
+    # the generated code string.
+
     def add_to_cache(node):
         if node in expr_to_replacement_cache:
             replacement_symbol = expr_to_replacement_cache[node]
@@ -120,6 +126,8 @@ def add_to_cache(node):
     replacement_symbols = numbered_symbols(
         prefix='z',
         cls=sm.Symbol,
+        # TODO : free symbols should be able to be passed in to save time in
+        # recomputing
         exclude=expr.free_symbols,
         real=True,
     )
@@ -158,7 +166,7 @@ def add_to_cache(node):
     start = timer()
     zeros = sm.ImmutableDenseMatrix.zeros(1, len(wrt))
     for symbol, subexpr in replacements:
-        free_symbols = subexpr.free_symbols
+        free_symbols = sm.ordered(subexpr.free_symbols)
         absolute_derivative = zeros
         for free_symbol in free_symbols:
             replacement_symbol, partial_derivative = add_to_cache(
@@ -182,8 +190,8 @@ def add_to_cache(node):
         entry = stack.pop()
         if entry in required_replacement_symbols or entry in wrt:
             continue
-        children = list(
-            replacement_to_reduced_expr_cache.get(entry, entry).free_symbols)
+        children = list(sm.ordered(
+            replacement_to_reduced_expr_cache.get(entry, entry).free_symbols))
         for child in children:
             if child not in required_replacement_symbols and child not in wrt:
                 stack.append(child)
@@ -198,9 +206,9 @@ def add_to_cache(node):
         if replacement_symbol in required_replacement_symbols
     }
 
-    counter = Counter(replaced_jacobian.free_symbols)
+    counter = Counter(sm.ordered(replaced_jacobian.free_symbols))
     for replaced_subexpr in required_replacements_dense.values():
-        counter.update(replaced_subexpr.free_symbols)
+        counter.update(sm.ordered(replaced_subexpr.free_symbols))
 
     logger.info('Substituting required replacements...')
     required_replacements = {}
@@ -459,6 +467,7 @@ def sort_sympy(seq):
 
 
 _c_template = """\
+// opty_code_hash={eval_code_hash}
 {win_math_def}
 #include <math.h>
 #include "{file_prefix}_h.h"
@@ -620,34 +629,38 @@ def ufuncify_matrix(args, expr, const=None, tmp_dir=None, parallel=False,
     args : iterable of sympy.Symbol
         A list of all symbols in expr in the desired order for the output
         function.
-    expr : sympy.Matrix
-        A matrix of expressions.
+    expr : sympy.Matrix or 2-tuple
+        A matrix of expressions or the output of ``cse()`` of a matrix of
+        expressions.
     const : tuple, optional
-        This should include any of the symbols in args that should be
-        constant with respect to the loop.
+        This should include any of the symbols in args that should be constant
+        with respect to the evaluation loop.
     tmp_dir : string, optional
-        The path to a directory in which to store the generated files. If
-        None then the files will be not be retained after the function is
-        compiled.
+        The path to a directory in which to store the generated files. If None
+        then the files will be not be retained after the function is compiled.
+        If this temporary directory is set to an existing populated directory
+        and ``expr`` has not changed relative to prior executions of this
+        function, the compilation step will be skipped if equivalent compiled
+        modules are already present and cached.
     parallel : boolean, optional
         If True and openmp is installed, the generated code will be
-        parallelized across threads. This is only useful when expr are
+        parallelized across threads. This is only useful when ``expr`` are
         extremely large.
     show_compile_output : boolean, optional
         If True, STDOUT and STDERR of the Cython compilation call will be
         shown.
 
     """
 
-    # TODO : This is my first ever global variable in Python. It'd probably
-    # be better if this was a class attribute of a Ufuncifier class. And I'm
-    # not sure if this current version counts sequentially.
+    # TODO : This is my first ever global variable in Python. It'd probably be
+    # better if this was a class attribute of a Ufuncifier class. And I'm not
+    # sure if this current version counts sequentially.
     global module_counter
 
     if hasattr(expr, 'shape'):
         num_rows = expr.shape[0]
         num_cols = expr.shape[1]
-    else:
+    else:  # output of cse()
         num_rows = expr[1][0].shape[0]
         num_cols = expr[1][0].shape[1]
 
@@ -675,6 +688,8 @@ def ufuncify_matrix(args, expr, const=None, tmp_dir=None, parallel=False,
             file_prefix = '{}_{}'.format(file_prefix_base, module_counter)
             module_counter += 1
 
+    prior_module_number = module_counter - 1
+
     d = {'routine_name': 'eval_matrix',
          'file_prefix': file_prefix,
          'matrix_output_size': matrix_size,
@@ -723,6 +738,19 @@ def ufuncify_matrix(args, expr, const=None, tmp_dir=None, parallel=False,
     d['eval_code'] = '    ' + '\n    '.join((sub_expr_code + '\n' +
                                              matrix_code).split('\n'))
 
+    # NOTE : It is very unlikely that the contents of evaluation code can be
+    # identical for two different sets of differential equations, so we hash it
+    # and store the hash value in the C file that contains the evaluation code.
+    # TODO : Maybe we should only do this if tmp_dir is not None, as it could
+    # have an undesired computational cost.
+    logger.debug('Calculating cache hash.')
+    hasher = hashlib.new('sha256')
+    const_str = 'const=None' if const is None else 'const={}'.format(const)
+    parallel_str = 'parallel={}'.format(parallel)
+    hasher.update((const_str + parallel_str + d['eval_code']).encode())
+    d['eval_code_hash'] = str(hasher.hexdigest())
+    logger.debug('Done calculating cache hash: {}'.format(d['eval_code_hash']))
+
     c_indent = len('void {routine_name}('.format(**d))
     c_arg_spacer = ',\n' + ' ' * c_indent
 
@@ -734,6 +762,7 @@ def ufuncify_matrix(args, expr, const=None, tmp_dir=None, parallel=False,
     memory_views = []
     for a in args:
         if const is not None and a in const:
+            # TODO : Should these be declared const in C?
             typ = 'double'
             idexy = '{}'
             cython_input_args.append('{} {}'.format(typ, ccode(a)))
@@ -772,6 +801,42 @@ def ufuncify_matrix(args, expr, const=None, tmp_dir=None, parallel=False,
 
     workingdir = os.getcwd()
     os.chdir(codedir)
+    logger.info('Changed directory to {}'.format(codedir))
+
+    # NOTE : If there are other files present in the directory (will only occur
+    # if a tmp_dir is set) then search through them starting with the most
+    # recent and see if it has a matching hash to the evaluation code generated
+    # here. If a match is found, store the module number.
+    matching_module_num = None
+    for prior_num in reversed(range(prior_module_number + 1)):
+        old_file_prefix = '{}_{}'.format(file_prefix_base, prior_num)
+        logger.info(f'Checking {old_file_prefix} for cached code.')
+        try:
+            with open(old_file_prefix + '_c.c', 'r') as f:
+                hash_line = f.readline()
+                logger.debug(hash_line.strip())
+                if 'opty_code_hash={}'.format(d['eval_code_hash']) in hash_line:
+                    matching_module_num = prior_num
+                    logger.info(f'{old_file_prefix} matches!')
+                    break
+        except FileNotFoundError:
+            logger.debug(f'{old_file_prefix} not found.')
+            pass
+
+    # NOTE : If we found a matching C file, then try to simply load that module
+    # instead of compiling a new one. This lets us skip the compile step if we
+    # have not changed anything about the model.
+    if matching_module_num is not None:
+        try:
+            cython_module = importlib.import_module(old_file_prefix)
+        except ImportError:  # false positive, so compile a new one
+            logger.info(f'Failed to import {old_file_prefix}.')
+            pass
+        else:
+            logger.info(f'Skipped compile, {old_file_prefix} module loaded.')
+            os.chdir(workingdir)
+            logger.info(f'Changed directory to {workingdir}.')
+            return getattr(cython_module, d['routine_name'] + '_loop')
 
     try:
         sys.path.append(codedir)
@@ -797,6 +862,7 @@ def ufuncify_matrix(args, expr, const=None, tmp_dir=None, parallel=False,
         else:
             encoding = None
         try:
+            logger.info('Compiling matrix evaluation.')
             proc = subprocess.run(cmd, capture_output=True, text=True,
                                   encoding=encoding)
         # On Windows this can raise a UnicodeDecodeError, but only in the
@@ -811,8 +877,13 @@ def ufuncify_matrix(args, expr, const=None, tmp_dir=None, parallel=False,
         if show_compile_output:
             print(stdout)
             print(stderr)
+        else:
+            logger.debug(stdout)
+            logger.debug(stderr)
+
         try:
             cython_module = importlib.import_module(d['file_prefix'])
+            logger.info("Loaded {} module".format(d['file_prefix']))
         except ImportError as error:
             msg = ('Unable to import the compiled Cython module {}, '
                    'compilation likely failed. STDERR output from '
@@ -827,6 +898,7 @@ def ufuncify_matrix(args, expr, const=None, tmp_dir=None, parallel=False,
             # so I don't delete the directory on Windows.
             if sys.platform != "win32":
                 shutil.rmtree(codedir)
+                logger.info('Removed directory {}'.format(codedir))
 
     return getattr(cython_module, d['routine_name'] + '_loop')