forked from M-Labs/artiq
cache source on import of modules that may contain kernels. Closes #416
This commit is contained in:
parent
d51b27e0aa
commit
84f4725015
|
@ -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")
|
|
@ -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:
|
||||
|
|
|
@ -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):
|
||||
|
|
|
@ -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()
|
||||
|
|
Loading…
Reference in New Issue