Add cache option to CLinker methods and clean up type hints

aesara/link/c/basic.py

@@ -7,7 +7,7 @@ import sys
 from collections import defaultdict
 from copy import copy
 from io import StringIO
-from typing import Dict
+from typing import TYPE_CHECKING, Any, Dict, Optional
 
 import numpy as np
 
@@ -21,6 +21,7 @@ from aesara.graph.basic import (
     vars_between,
 )
 from aesara.link.basic import Container, Linker, LocalLinker, PerformLinker
+from aesara.link.c import cutils
 from aesara.link.c.cmodule import (
     METH_VARARGS,
     DynamicModule,
@@ -34,13 +35,14 @@ from aesara.link.utils import gc_helper, map_storage, raise_with_op, streamline
 from aesara.utils import difference, uniq
 
 
-_logger = logging.getLogger("aesara.link.c.basic")
-
+if TYPE_CHECKING:
+    from aesara.graph.fg import FunctionGraph
+    from aesara.link.c.cmodule import ModuleCache
 
-run_cthunk = None  # Will be imported only when needed.
+_logger = logging.getLogger("aesara.link.c.basic")
 
 
-def get_module_cache(init_args=None):
+def get_module_cache(init_args: Optional[Dict[str, Any]] = None) -> "ModuleCache":
     """
 
     Parameters
@@ -580,7 +582,9 @@ class CLinker(Linker):
         self.fgraph = None
         super().__init__(scheduler=schedule)
 
-    def accept(self, fgraph, no_recycling=None, profile=None):
+    def accept(
+        self, fgraph: "FunctionGraph", no_recycling=None, profile=None
+    ) -> "CLinker":
         r"""Associate this `Linker` with `fgraph`.
 
         The `no_recycling` argument can contain a list of `Variable`\s that
@@ -1080,18 +1084,18 @@ class CLinker(Linker):
         input_storage=None,
         output_storage=None,
         storage_map=None,
+        cache: Optional["ModuleCache"] = None,
     ):
-        """
-        Compiles this linker's fgraph.
+        """Compile `self.fgraph`.
 
         Parameters
         ----------
         input_storage: list or None
             List of lists of length 1. In order to use the thunk returned
-            by __compile__, the inputs must be put in that storage.
+            by this method, the inputs must be put in that storage.
             If None, storage will be allocated.
         output_storage: list of lists of length 1
-            The thunk returned by __compile__ will put the variables of the
+            The thunk returned by this method will put the variables of the
             computation in these lists. If None, storage will be allocated.
 
         Returns
@@ -1121,6 +1125,7 @@ class CLinker(Linker):
             input_storage,
             output_storage,
             storage_map,
+            cache,
         )
         return (
             thunk,
@@ -1152,37 +1157,51 @@ class CLinker(Linker):
             id += 2
         return init_tasks, tasks
 
-    def make_thunk(self, input_storage=None, output_storage=None, storage_map=None):
-        """
-        Compiles this linker's fgraph and returns a function to perform the
-        computations, as well as lists of storage cells for both the inputs
-        and outputs.
+    def make_thunk(
+        self,
+        input_storage=None,
+        output_storage=None,
+        storage_map=None,
+        cache: Optional["ModuleCache"] = None,
+        **kwargs,
+    ):
+        """Compile this linker's `self.fgraph` and return a function that performs the computations.
+
+        The return values can be used as follows:
+
+        .. code-block::
+
+            f, istor, ostor = clinker.make_thunk()
+            istor[0].data = first_input
+            istor[1].data = second_input
+            f()
+            first_output = ostor[0].data
+
 
         Parameters
         ----------
         input_storage: list or None
             List of lists of length 1. In order to use
-            the thunk returned by __compile__, the inputs must be put in
+            the thunk returned by `CLinker.__compile__`, the inputs must be put in
             that storage. If None, storage will be allocated.
         output_storage: list of lists of length 1.
-            The thunk returned by __compile__ will put the variables
+            The thunk returned by `CLinker.__compile__` will put the variables
             of the computation in these lists. If None, storage will
             be allocated.
         storage_map: dict that map variables to storages.
             This is used when you need to customize the storage of
             this thunk
-        Returns: thunk, input_storage, output_storage
+        cache
+            A cache in which to store the compilation results.
+
+        Returns
+        -------
+        thunk, input_storage, output_storage
 
-        The return values can be used as follows:
-          f, istor, ostor = clinker.make_thunk()
-          istor[0].data = first_input
-          istor[1].data = second_input
-          f()
-          first_output = ostor[0].data
         """
         init_tasks, tasks = self.get_init_tasks()
         cthunk, module, in_storage, out_storage, error_storage = self.__compile__(
-            input_storage, output_storage, storage_map
+            input_storage, output_storage, storage_map, cache
         )
 
         res = _CThunk(cthunk, init_tasks, tasks, error_storage, module)
@@ -1584,7 +1603,14 @@ class CLinker(Linker):
             self._mod = mod
         return self._mod
 
-    def cthunk_factory(self, error_storage, in_storage, out_storage, storage_map=None):
+    def cthunk_factory(
+        self,
+        error_storage,
+        in_storage,
+        out_storage,
+        storage_map=None,
+        cache: Optional["ModuleCache"] = None,
+    ):
         """
         Returns a thunk that points to an instance of a C struct that
         can carry on the computation of this linker's fgraph
@@ -1605,6 +1631,7 @@ class CLinker(Linker):
             key = self.cmodule_key()
         except KeyError:
             key = None
+
         if key is None:
             # If we can't get a key, then forget the cache mechanism.
             module = self.compile_cmodule()
@@ -1612,7 +1639,9 @@ class CLinker(Linker):
             # Set compute_map as None as clinker do not support lazy evaluation
             for node in self.node_order:
                 node.op.prepare_node(node, storage_map, None, "c")
-            module = get_module_cache().module_from_key(key=key, lnk=self)
+            if cache is None:
+                cache = get_module_cache()
+            module = cache.module_from_key(key=key, lnk=self)
 
         vars = self.inputs + self.outputs + self.orphans
         # List of indices that should be ignored when passing the arguments
@@ -1703,15 +1732,12 @@ class _CThunk:
     """
 
     def __init__(self, cthunk, init_tasks, tasks, error_storage, module):
-        global run_cthunk
-        if run_cthunk is None:
-            # Lazy import to avoid compilation when importing aesara.
-            from aesara.link.c.cutils import run_cthunk  # noqa
         self.cthunk = cthunk
         self.init_tasks = init_tasks
         self.tasks = tasks
         self.error_storage = error_storage
         self.module = module
+        self.nodes = None
 
     def find_task(self, failure_code):
         """
@@ -1727,7 +1753,7 @@ class _CThunk:
             return self.tasks[failure_code - n]
 
     def __call__(self):
-        failure = run_cthunk(self.cthunk)
+        failure = cutils.run_cthunk(self.cthunk)
         if failure:
             task, taskname, id = self.find_task(failure)
             try:

aesara/link/c/cmodule.py

@@ -1612,18 +1612,18 @@ def _rmtree(
                 )
 
 
-_module_cache = None
+_module_cache: Optional[ModuleCache] = None
 
 
-def get_module_cache(dirname, init_args=None):
-    """
-    Create a new module_cache with the (k, v) pairs in this dictionary
+def get_module_cache(dirname: str, init_args=None) -> ModuleCache:
+    """Create a new module_cache.
 
     Parameters
     ----------
+    dirname
+        The name of the directory used by the cache.
     init_args
-        If not None, the (k, v) pairs in this dictionary will be forwarded to
-        the ModuleCache constructor as keyword arguments.
+        Keyword arguments passed to the `ModuleCache` constructor.
 
     """
     global _module_cache
@@ -1639,8 +1639,8 @@ def get_module_cache(dirname, init_args=None):
         )
     if _module_cache.dirname != dirname:
         _logger.warning(
-            "Returning module cache instance with different "
-            f"dirname ({_module_cache.dirname}) than you requested ({dirname})"
+            "Returning a module cache instance with a different "
+            f"dirname ({_module_cache.dirname}) than requested ({dirname})"
         )
     return _module_cache