cache source on import of modules that may contain kernels. Closes #416

artiq/compiler/import_cache.py

@@ -0,0 +1,55 @@
+import sys
+import builtins
+import linecache
+import tokenize
+import logging
+import importlib.machinery as im
+
+from artiq.experiment import kernel, portable
+
+
+__all__ = ["install_hook"]
+
+
+logger = logging.getLogger(__name__)
+
+
+cache = dict()
+im_exec_module = None
+linecache_getlines = None
+
+
+def hook_exec_module(self, module):
+    im_exec_module(self, module)
+    if (hasattr(module, "__file__")
+            # Heuristic to determine if the module may contain ARTIQ kernels.
+            # This breaks if kernel is not imported the usual way.
+            and ((getattr(module, "kernel", None) is kernel)
+                 or (getattr(module, "portable", None) is portable))):
+        fn = module.__file__
+        try:
+            with tokenize.open(fn) as fp:
+                cache[fn] = fp.readlines()
+        except:
+            logger.warning("failed to add '%s' to cache", fn, exc_info=True)
+        else:
+            logger.debug("added '%s' to cache", fn)
+
+
+def hook_getlines(filename, module_globals=None):
+    if filename in cache:
+        return cache[filename]
+    else:
+        return linecache_getlines(filename, module_globals)
+
+
+def install_hook():
+    global im_exec_module, linecache_getlines
+
+    im_exec_module = im.SourceFileLoader.exec_module
+    im.SourceFileLoader.exec_module = hook_exec_module
+
+    linecache_getlines = linecache.getlines
+    linecache.getlines = hook_getlines
+
+    logger.debug("hook installed")

artiq/frontend/artiq_run.py

@@ -19,6 +19,7 @@ from artiq.master.worker_db import DeviceManager, DatasetManager
 from artiq.coredevice.core import CompileError, host_only
 from artiq.compiler.embedding import EmbeddingMap
 from artiq.compiler.targets import OR1KTarget
+from artiq.compiler import import_cache
 from artiq.tools import *
 
 
@@ -158,6 +159,7 @@ def _build_experiment(device_mgr, dataset_mgr, args):
         elif is_bc:
             return LLVMBitcodeRunner(device_mgr, dataset_mgr, file=args.file)
         else:
+            import_cache.install_hook()
             module = file_import(args.file, prefix="artiq_run_")
         file = args.file
     else:

artiq/language/core.py

@@ -43,6 +43,9 @@ def kernel(arg=None, flags={}):
 
     The decorator takes an optional parameter that defaults to ``core`` and
     specifies the name of the attribute to use as core device driver.
+
+    This decorator must be present in the global namespace of all modules using
+    it for the import cache to work properly.
     """
     if isinstance(arg, str):
         def inner_decorator(function):
@@ -70,6 +73,9 @@ def portable(arg=None, flags={}):
     host will be executed on the host (no compilation and execution on the
     core device). A decorated function called from a kernel will be executed
     on the core device (no RPC).
+
+    This decorator must be present in the global namespace of all modules using
+    it for the import cache to work properly.
     """
     if arg is None:
         def inner_decorator(function):

artiq/master/worker_impl.py

@@ -16,6 +16,7 @@ from artiq.language.environment import (is_experiment, TraceArgumentManager,
                                         ProcessArgumentManager)
 from artiq.language.core import set_watchdog_factory, TerminationRequested
 from artiq.language.types import TBool
+from artiq.compiler import import_cache
 from artiq.coredevice.core import CompileError, host_only, _render_diagnostic
 from artiq import __version__ as artiq_version
 
@@ -191,6 +192,8 @@ def main():
                                virtual_devices={"scheduler": Scheduler()})
     dataset_mgr = DatasetManager(ParentDatasetDB)
 
+    import_cache.install_hook()
+
     try:
         while True:
             obj = get_object()