Skip to content

Commit

Permalink
cache source on import of modules that may contain kernels. Closes #416
Browse files Browse the repository at this point in the history
sbourdeauducq committed Aug 6, 2016
1 parent d51b27e commit 84f4725
Showing 4 changed files with 66 additions and 0 deletions.
55 changes: 55 additions & 0 deletions artiq/compiler/import_cache.py
Original file line number Diff line number Diff line change
@@ -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")
2 changes: 2 additions & 0 deletions artiq/frontend/artiq_run.py
Original file line number Diff line number Diff line change
@@ -19,6 +19,7 @@
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:
6 changes: 6 additions & 0 deletions artiq/language/core.py
Original file line number Diff line number Diff line change
@@ -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):
3 changes: 3 additions & 0 deletions artiq/master/worker_impl.py
Original file line number Diff line number Diff line change
@@ -16,6 +16,7 @@
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()

0 comments on commit 84f4725

Please sign in to comment.