import inspect
import time
import numpy as np
from artiq.experiment import kernel, rpc, delay, now_mu, portable
from artiq.language.environment import EnvExperiment, EnumerationValue, BooleanValue, NumberValue, StringValue
from artiq.language.scan import Scannable, NoScan, MultiScanManager
from artiq.language.units import ms
from artiq.language.core import TerminationRequested
from artiq.coredevice.exceptions import RTIOUnderflow
from artiq.master.worker_db import DummyDevice, DeviceManager
from atomiq.components import ComponentFactory
from atomiq.components.primitives import Component
from atomiq.arguments import NativeArgumentProvider
from atomiq.hooks import hooks
from atomiq.heros import heros, AtomiqHERO
from atomiq.helper import (
component_dict,
component_data,
block_dict,
rec_getattr,
random_ascii_string,
master_ipc_call,
normalize_component_name,
)
import logging
default_argument_provider = NativeArgumentProvider()
logging.basicConfig()
logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)
[docs]
class KernelChunk:
"""Represents a chunk of scan points optimized for kernel execution.
This class behave similar to :py:class:`Chunk` but for use inside an ARTIQ kernel with reduced python capabilities.
The main change is that this class is not iterable but :py:meth:`KernelChunk.to_list()` must be used to obtain
an iterable chunk list.
"""
def __init__(self, points):
self._points = points
self.kernel_invariants = {
"_variable_args",
}
self._variable_args = []
compare_dict = {
key: getattr(self._points[0], key)
for key in self._points[0].__dict__
if key not in ("identifier", "step_counter")
}
for point in self._points:
drop = [key for key, val in compare_dict.items() if val != getattr(point, key)]
self._variable_args.extend(drop)
for key in drop:
compare_dict.pop(key)
for key, val in compare_dict.items():
setattr(self, key, val)
self.kernel_invariants.add(key)
self.kernel_invariants |= set(compare_dict.keys())
[docs]
def __getattribute__(self, name):
if name == "_variable_args":
return super().__getattribute__(name)
elif name in self._variable_args:
msg = "Argument is used per chunk but changes in a chunk:"
raise RuntimeError(msg, name)
return super().__getattribute__(name)
[docs]
@portable
def to_list(self):
return self._points
[docs]
def __repr__(self) -> str:
return f"KernelChunk({self._points})"
[docs]
class Chunk(list):
"""Extended list class for managing chunks of scan points.
Provides convenience methods for accessing common attributes across all points in the chunk.
"""
[docs]
def __getattr__(self, attr):
return self.get(attr)
[docs]
@portable
def get(self, attr):
iterator = iter(self)
try:
first = getattr(next(iterator), attr)
except StopIteration:
return None
if not all(first == getattr(x, attr) for x in iterator):
msg = f"Argument {attr} is used per chunk but changes in a chunk."
raise RuntimeError(msg)
return getattr(self[0], attr)
[docs]
def get_kernel_chunk(self):
"""Get a KernelChunk copy of the chunk for ARTIQ kernel execution."""
return KernelChunk(list(self))
[docs]
class AtomiqExperiment(EnvExperiment):
CHUNKSIZE = 10
TAGS = []
components = ["log"]
heros_realm = "heros"
log = logger
arg_provider = default_argument_provider
max_step_counter: int
scan_argument_order: list[str]
def __init__(self, managers_or_parent, name=None, arg_provider=None, component_map=None, *args, **kwargs):
# propagate heros realm and check if heros wanted before building components and blocks
if heros is not None:
hero_condition = True
no_realm_msg = "Can not infer heros realm from experiment, falling back to default `heros`. This can happen if the experiment was not queued through the scheduler"
if isinstance(managers_or_parent, tuple):
if isinstance(managers_or_parent[0], DeviceManager):
# only do heros if not in scanning loop, first manager is a dummy class while scanning
try:
self.heros_realm = managers_or_parent[0].virtual_devices["scheduler"].expid["heros_realm"]
except KeyError:
logger.info(no_realm_msg)
else:
hero_condition = False
else:
if not isinstance(managers_or_parent.scheduler, DummyDevice):
# only do heros if not in scanning loop
try:
self.heros_realm = managers_or_parent.heros_realm
except AttributeError:
logger.info(no_realm_msg)
else:
hero_condition = False
else:
hero_condition = False
# combine components from parent class with our own and make entries unique
self._components = []
for cls in self._getmro():
if hasattr(cls, "components"):
self._components += cls.components
self._components = list(set(self._components))
self._blocks = []
for cls in self._getmro():
if hasattr(cls, "blocks"):
self._blocks += cls.blocks
self.name = name if name is not None else type(self).__name__
self.component_map = component_map
if arg_provider is not None:
self.arg_provider = arg_provider
self.scannable_arguments = {}
self.non_scannable_arguments = {}
self.scan_argument_order = []
self.clock_at_start = np.int64(-1)
self.identifier = np.int64(-1)
self.step_counter = 0
self.run_id = 0
# if we are an atomiq block, leave the experiment attribute untouched
if not hasattr(self, "experiment"):
self.experiment = self
self.__terminate_asap__ = False
# latch to the core device logger to handle kernel sequence and collision errors.
def _sequence_error_filter(record):
if "reported during kernel execution" in record.msg.lower():
self._sequence_error = True
return True
self._sequence_error = False
logging.getLogger("artiq.coredevice.comm_kernel").addFilter(_sequence_error_filter)
super().__init__(managers_or_parent)
# only attach HERO if not in examine phase and if HEROS is loaded
# also, do not attach the HERO to a block
if hero_condition and not isinstance(self, AtomiqBlock):
self._create_atomiq_hero()
[docs]
def _create_atomiq_hero(self):
self.hero = AtomiqHERO(self, realm=self.heros_realm)
# Emit starting signal via master process if heros is available
metadata = {"name": self.__class__.__name__, "hero_name": self.hero._name}
master_ipc_call(self, "run_created")(self.scheduler.rid, metadata=metadata)
[docs]
def _getmro(self):
return inspect.getmro(type(self))
[docs]
def _import_from_block(self, blkdata: dict):
"""
This uses monkey patching of the class to make members available to
the current experiment
Args:
blkdata: Block data with the keys "class", "map", "alias"
"""
# get proper name of the step function after importing
alias = blkdata["alias"]
cls_name = blkdata["class"].__name__ if alias is None else alias
# propagate our component mappings if necessary
if self.component_map is not None and len(self.component_map) > 0:
prop_map = {
key: (self.component_map[val] if val in self.component_map else val)
for key, val in blkdata["map"].items()
}
else:
prop_map = blkdata["map"]
# instantiate the block object and attach it to self
obj = blkdata["class"](self, arg_provider=self.arg_provider, component_map=prop_map, name=cls_name)
setattr(type(self), cls_name, obj)
# check if we need any components from the attached block object
if hasattr(obj, "_components"):
cmp = [
component_data(dict(comp, **{"id": prop_map[comp["id"]]}))
if comp["id"] in prop_map
else component_data(comp)
for comp in map(component_dict, obj._components)
]
logger.info(f"Transferring components {cmp} to parent")
self._components += cmp
self._components = list(set(self._components))
[docs]
def _add_to_hook_handler(self, obj_list):
"""
Add the hooks the objects in obj_list to the experiment hooks such that they are called when processing the
hooks. Only if the child object posses a method with the according name of the hook, it will be called.
Args:
obj_list: list of child objects relative to self (e.g. 'modulator.rfsource.switch') for which the hook
should be called
"""
for hook, no_args in hooks:
args = [f"arg{i}" for i in range(no_args)]
code = f"def _{hook}({', '.join(['self'] + args)}):" + "\n"
code += f" self.log.debug('Doing {hook} for {{0}}', [self])\n"
for obj_path in obj_list:
# account for HEROS
obj_path = normalize_component_name(obj_path)
obj = rec_getattr(self, obj_path)
if hasattr(obj, f"_{hook}"):
wrapped_hook = f"do_{hook}" if hasattr(obj, f"_do_{hook}") else hook
if isinstance(obj, Component):
code += f" if not self.{obj_path}._{hook}_done:" + "\n"
code += " self." + f"{obj_path}._{wrapped_hook}({', '.join(args)})" + "\n"
code += f" self.{obj_path}._{hook}_done = True" + "\n"
else:
code += " self." + f"{obj_path}._{wrapped_hook}({', '.join(args)})" + "\n"
if hasattr(self, f"_{hook}"):
# we need to set an ephemeral name for the existing hook, other wise we would recurse into the method
# we are about to build
rnd = random_ascii_string(8)
setattr(self, f"_{hook}_{rnd}", getattr(self, f"_{hook}"))
code += f" self._{hook}_{rnd}({', '.join(args)})" + "\n"
else:
code += f" self.{hook}({', '.join(args)})" + "\n"
loc = locals()
exec(code, globals(), loc)
f = loc[f"_{hook}"]
if not hook.endswith("_host"):
f = kernel(f)
# Save source code for the compiler to pick up later.
f.artiq_embedded = f.artiq_embedded._replace(function=code)
setattr(self.__class__, f"_{hook}", f)
setattr(self, f"_{hook}", f.__get__(self))
[docs]
def _build_blocks(self):
for blkdata in map(block_dict, self._blocks):
self._import_from_block(blkdata)
blks = [
blkdata["alias"] if "alias" in blkdata and blkdata["alias"] is not None else blkdata["class"].__name__
for blkdata in map(block_dict, self._blocks)
]
self._add_to_hook_handler(blks)
[docs]
def _build_components(self):
recursive_components = []
for comp in map(component_dict, self._components):
if self.component_map is not None and comp["id"] in self.component_map:
target_id = self.component_map[comp["id"]]
else:
target_id = comp["id"]
comp_obj = ComponentFactory.produce(target_id, self)
if isinstance(comp_obj, comp["type"]) or isinstance(self.scheduler, DummyDevice):
# in the case of repo scan, do not check for correct type
setattr(self, normalize_component_name(comp["id"]), comp_obj)
else:
raise TypeError(
f"Component {comp['id']} does not have specified type {comp['type']},"
f"instead it has {type(comp_obj)}"
)
recursive_components += comp_obj.required_components(ancestors=[comp["id"]])
# In the following we add all hooks of the components to the experiment hooks
# remove duplicates
dedup_components = []
for obj, path in recursive_components:
if obj not in [obj for obj, _ in dedup_components]:
dedup_components.append((obj, path))
prerun_components = [
(obj, path) for obj, path in dedup_components if obj.__class__._prerun != Component._prerun
]
self._add_to_hook_handler([".".join(path) for obj, path in prerun_components])
[docs]
def setattr_argument(self, key, processor=None, group=None, tooltip=None):
"""Set an argument as an attribute.
This extends the ARTIQ setattr_argument function by checking if there are incompatible collisions between
arguments of different blocks and the experiment.
"""
try:
EnvExperiment.setattr_argument(self, key, processor=processor, group=group, tooltip=tooltip)
except TypeError as e:
if "int() argument must be" in str(e):
msg = (
f"Argument '{key}' is defined as a non-scannable in '{self.__class__.__name__}' but is "
+ "supplied with a scannable (float). Make sure you don't shadow the argument in "
+ "another Block or the main experiment."
)
raise TypeError(msg) from e
elif "object is not subscriptable" in str(e):
msg = (
f"Argument '{key}' is defined as a scannable (float) in '{self.__class__.__name__}' but is "
+ "supplied with a non-scannable. Make sure you don't shadow the argument in "
+ "another Block or the main experiment."
)
raise TypeError(msg) from e
else:
# something else gone wrong
raise
if isinstance(processor, Scannable):
self.scannable_arguments[key] = getattr(self, key)
else:
self.non_scannable_arguments[key] = getattr(self, key)
[docs]
def _build_arguments(self):
for argname, argdict in self.arg_provider.get_arguments(self.__class__).items():
group = argdict["group"] if "group" in argdict else self.name
if "options" in argdict:
default = argdict["default"] if "default" in argdict else argdict["options"][0]
self.setattr_argument(argname, EnumerationValue(argdict["options"], default=default), group=group)
else:
if "default" in argdict:
if isinstance(argdict["default"], bool):
self.setattr_argument(argname, BooleanValue(argdict["default"]), group=group)
elif isinstance(argdict["default"], str):
self.setattr_argument(argname, StringValue(argdict["default"]), group=group)
elif (
("scannable" in argdict and argdict["scannable"] is False)
or ("type" in argdict and argdict["type"] in ("int", int))
or (
"step" in argdict
and argdict["step"] == int(argdict["step"])
and "precision" in argdict
and argdict["precision"] == 0
and "scale" in argdict
and argdict["scale"] == 1
)
):
if "type" in argdict and argdict["type"] in ("int", int):
# set default compatible int values so we can define args with just "type"
argdict["step"] = argdict.get("step", 1)
argdict["precision"] = argdict.get("precision", 0)
argdict["scale"] = argdict.get("scale", 1)
self.setattr_argument(
argname,
NumberValue(
argdict["default"],
**{
k: argdict[k]
for k in ("unit", "scale", "precision", "step", "min", "max", "type")
if k in argdict
},
),
group=group,
)
else:
self.setattr_argument(
argname,
Scannable(
default=NoScan(argdict["default"], 1),
**{k: argdict[k] for k in ("unit", "scale", "precision") if k in argdict},
),
group=group,
)
else:
logger.error(f"no default value set for argument {argname}")
[docs]
def get_multiscanmanager(self) -> tuple[list[str], MultiScanManager]:
"""
Get the MultiScanManager of all arguments.
"""
scannable_arguments_list = [item for item in self.scannable_arguments.items()]
all_arguments = self._HasEnvironment__argument_mgr.unprocessed_arguments
scannable_arguments_list.sort(
key=lambda x: all_arguments[x[0]]["scan_priority"] if "scan_priority" in all_arguments[x[0]] else 100,
)
scan_argument_order = [varname for varname, var in scannable_arguments_list if len(var) > 1]
scan_order_msg = "scan argument order (inner to outer loop): " + " -> ".join(scan_argument_order)
self.log.info(scan_order_msg)
# reverse order to it consistent with dashboard and print
return scan_argument_order, MultiScanManager(*reversed(scannable_arguments_list))
[docs]
def prepare(self):
"""
Prepares components and structure.
Called by ARTIQ in the prepare phase, see ARTIQ documentation for more information on experiment phases.
Note:
If you overwrite this method in your experiment, make sure to call `super().prepare()`.
"""
for component in self._components:
getattr(self, normalize_component_name(component_dict(component)["id"]))._recursive_prepare()
# check if we need any attributes from the attached block object
for blkdata in map(block_dict, self._blocks):
alias = blkdata["alias"]
cls_name = blkdata["class"].__name__ if alias is None else alias
obj = getattr(type(self), cls_name)
obj.prepare()
# include child arguments with dedup using dict update
if hasattr(obj, "scannable_arguments"):
self.scannable_arguments.update(obj.scannable_arguments)
if hasattr(obj, "non_scannable_arguments"):
self.non_scannable_arguments.update(obj.non_scannable_arguments)
[docs]
def _build_core(self):
self.setattr_device("scheduler")
self.setattr_device("core")
self.setattr_device("core_dma")
[docs]
def build(self):
"""
Initializes arguments.
Called by ARTIQ in the build phase, see ARTIQ documentation for more information on experiment phases.
Note:
If you overwrite this method in your experiment, make sure to call `super().build()`
"""
self._build_core()
# It's important to do this before the components are built since
# the blocks modify the component list
self._build_blocks()
self._build_components()
# maybe components themselves need something being done in build phase?
for component in self._components:
getattr(self, normalize_component_name(component_dict(component)["id"]))._recursive_build()
self._build_arguments()
self._build()
[docs]
@kernel
def _prerun_core(self, reset=False):
self.clock_at_start = self.core.get_rtio_counter_mu()
self.log.info("setting up coredevice")
if reset:
delay(500 * ms)
self.core.reset()
self.core.break_realtime()
self.prerun()
[docs]
@rpc(flags={"async"})
def _set_identifiers_in_master(self, identifier, run_id, step_counter):
"""
makes an rpc call to set the identifier in the master
"""
self.identifier = identifier
self.run_id = run_id
self.step_counter = step_counter
self.set_dataset("step_counter", step_counter, broadcast=True)
[docs]
def chunker(self, mult_scan, size: np.int32 = 100):
"""
Generator to call a kernel with chunks of scan points.
"""
chunk = Chunk()
i = 1
counter = 0
# add max_step_counter to experiment
for _p in mult_scan:
# broadcast the max step counter
# we do this inside the loop to update the dataset in case of
# multiple concurrent runs (pause after chunk)
self.set_dataset("max_step_counter", self.max_step_counter, broadcast=True)
step_counter = counter
_p.__dict__.update(self.non_scannable_arguments)
_p.attr |= self.non_scannable_arguments.keys()
setattr(_p, "step_counter", step_counter)
setattr(_p, "run_id", self.scheduler.rid)
setattr(_p, "identifier", np.int64(-1))
chunk.append(_p)
counter += 1
if len(chunk) >= size:
logger.info(f"starting kernel with chunk {i}")
yield chunk
chunk = Chunk()
i += 1
if len(chunk) > 0:
logger.info(f"starting kernel with chunk {i}")
yield chunk
else:
return
[docs]
def _loop_chunks(self, chunk_list: list):
"""
Loop over a list of chunks in the host, i.e. iterate over starting kernels holding CHUNKSIZE experiments
:param chunks: (list) list of chunks. Each entry is again an iterable with CHUNKSIZE Scanpoint objects inside
"""
current_chunk = 0
for points in chunk_list:
try:
self._prechunk_host(points)
self._runchunk(points.get_kernel_chunk())
self._postchunk_host(points)
except RTIOUnderflow as ex:
self.log.warning(f"{ex} during step {self.step_counter} in chunk {current_chunk}")
self.log.warning(ex.artiq_core_exception)
self._prerun_core(reset=True)
finally:
if self._sequence_error:
self.log.warning(f"Sequence error reported during chunk {current_chunk}, resetting kernel...")
self._prerun_core(reset=True)
self._sequence_error = False
current_chunk += 1
if self.scheduler.check_pause():
self.core.comm.close()
self.scheduler.pause()
# emergency stop if required
if self.__terminate_asap__:
logger.warning("Terminating NOW")
break
self.log.info(f"Come with me if you want to live - RID {self.scheduler.rid}")
[docs]
def _run_init(self):
# generate and log run id
exp_identifier = f"{self.__class__.__name__}_{self.scheduler.rid}"
self.log.info(f"starting run {exp_identifier}")
# get timestamp of the run
self.run_timestamp = float(time.time())
self._prerun_core()
self._prerun_host()
self._prerun()
logger.info(f"running with block list {self._blocks}")
logger.info(f"running with component list {self._components}")
[docs]
def run(self):
"""
Run entry point for ARTIQ, see ARTIQ documentation for more information on experiment phases.
Warning:
Do not implement this entry point in your experiment. Use the provided sub-phases (`prerun`, `step`, etc.)
More information can be found in the :ref:`phases_chunking` documentation.
"""
# Emit starting signal via master process if heros is available
if heros is not None:
metadata = {"name": self.__class__.__name__, "hero_name": self.hero._name}
master_ipc_call(self, "run_started")(self.scheduler.rid, metadata=metadata)
try:
self.scan_argument_order, mult_scan = self.get_multiscanmanager()
# first count the number of points for time estimation
self.max_step_counter = len(list(mult_scan)) - 1
self._run_init()
try:
self._loop_chunks(self.chunker(mult_scan, self.CHUNKSIZE))
except TerminationRequested:
logger.info(f"I'll be back - RID {self.scheduler.rid}")
self._postrun()
self._postrun_host()
except Exception as e_run:
logger.exception("Run sub-phase failed with %s - RID %i ", e_run, self.scheduler.rid)
try:
self._postfail()
except Exception as e_postfail_kernel:
logger.exception(
"Post-fail kernel phase failed with %s - RID %i ", e_postfail_kernel, self.scheduler.rid
)
try:
self._postfail_host()
except Exception as e_postfail_host:
logger.exception("Post-fail host phase failed with %s - RID %i ", e_postfail_host, self.scheduler.rid)
raise e_run
# Emit ending signal via master process if heros is available
if heros is not None:
metadata = {"name": self.__class__.__name__, "hero_name": self.hero._name}
master_ipc_call(self, "run_ended")(self.scheduler.rid, metadata=metadata)
[docs]
@kernel
def _runchunk(self, points):
# make everything settle
self.core.wait_until_mu(now_mu())
self.core.reset()
self._prechunk(points)
# execute the step function for every point in the chunk.
for point in points.to_list():
# set the identifier in the master and the running core device
self.identifier = np.int64(self.run_timestamp * 1e6) + np.int64(
(self.core.get_rtio_counter_mu() - self.clock_at_start) * self.core.ref_period * 1e6
)
point.identifier = self.identifier
self.step_counter = point.step_counter
self.run_id = point.run_id
self._set_identifiers_in_master(self.identifier, self.run_id, self.step_counter)
self.core.break_realtime()
# delay to have some slack when the step starts
delay(0.5 * ms)
# update run progress
self.set_dataset(
"run_progress_pct", 100 * ((self.step_counter + 1) / (self.max_step_counter + 1)), broadcast=True
)
# do what the user wants to be done...
self._prestep(point)
self.step(point)
self._poststep(point)
self.core.wait_until_mu(now_mu())
self._postchunk(points)
self.core.wait_until_mu(now_mu())
[docs]
def terminate(self):
"""
Terminate the running experiment as soon as possible. Normally this means after the current chunk.
"""
self.__terminate_asap__ = True
[docs]
@kernel
def prerun(self):
"""
Kernel entry point, run once at the beginning of the run phase of an experiment.
This method can be overloaded by the user.
Details can be found in the :ref:`phases_chunking` documentation.
"""
[docs]
def prerun_host(self):
"""
Host entry point, run once at the beginning of the run phase of an experiment.
This method can be overloaded by the user.
Details can be found in the :ref:`phases_chunking` documentation.
"""
[docs]
@kernel
def postfail(self):
pass
[docs]
def postfail_host(self):
pass
[docs]
@kernel
def postrun(self):
"""
Kernel entry point, run once at the end of the run phase of an experiment.
This method can be overloaded by the user.
Details can be found in the :ref:`phases_chunking` documentation.
"""
[docs]
def postrun_host(self):
"""
Host entry point, run once at the end of the run phase of an experiment.
This method can be overloaded by the user.
Details can be found in the :ref:`phases_chunking` documentation.
"""
[docs]
@kernel
def prestep(self, point):
"""
Kernel entry point, run before every step.
This method can be overloaded by the user.
Details can be found in the :ref:`phases_chunking` documentation.
"""
[docs]
@kernel
def poststep(self, point):
"""
Kernel entry point, run after every step.
This method can be overloaded by the user.
Details can be found in the :ref:`phases_chunking` documentation.
"""
[docs]
@kernel
def prechunk(self, points):
"""
Kernel entry point, run once at the beginning of a chunk.
This method can be overloaded by the user.
Details can be found in the :ref:`phases_chunking` documentation.
"""
[docs]
@kernel
def postchunk(self, points):
"""
Kernel entry point, run once at the end of a chunk.
This method can be overloaded by the user.
Details can be found in the :ref:`phases_chunking` documentation.
"""
[docs]
def prechunk_host(self, points):
"""
Host entry point, run once at the beginning of a chunk.
This method can be overloaded by the user.
Details can be found in the :ref:`phases_chunking` documentation.
"""
[docs]
def postchunk_host(self, points):
"""
Host entry point, run once at the end of a chunk.
This method can be overloaded by the user.
Details can be found in the :ref:`phases_chunking` documentation.
"""
[docs]
@kernel
def step(self, point):
"""
Kernel entry point, for the main experiment sequence code.
This method **must** be overloaded by the user.
Details can be found in the :ref:`phases_chunking` documentation.
"""
raise NotImplementedError
[docs]
class AtomiqBlock(AtomiqExperiment):
def __new__(cls, *args, **kwargs):
# if we get a name in the constructor make a class with that name and instantiate it
if "name" in kwargs and kwargs["name"] is not None:
return super().__new__(type(kwargs["name"], (cls,), {}))
else:
return super().__new__(cls)
def __init__(self, *args, **kwargs):
self.experiment = args[0].experiment
super().__init__(*args, **kwargs)
[docs]
def _getmro(self):
# since we dynamically add the dummy class with the alias name, we need to exclude that
# from the mro here again.
return inspect.getmro(type(self))[1:]