diff --git a/artiq/language/environment.py b/artiq/language/environment.py index 0e1b2b918..2709643c5 100644 --- a/artiq/language/environment.py +++ b/artiq/language/environment.py @@ -207,10 +207,12 @@ class HasEnvironment: self.__device_mgr = managers_or_parent[0] self.__dataset_mgr = managers_or_parent[1] self.__argument_mgr = managers_or_parent[2] + self.__scheduler_defaults = managers_or_parent[3] else: self.__device_mgr = managers_or_parent.__device_mgr self.__dataset_mgr = managers_or_parent.__dataset_mgr self.__argument_mgr = managers_or_parent.__argument_mgr + self.__scheduler_defaults = {} managers_or_parent.register_child(self) self.__in_build = True @@ -364,6 +366,21 @@ class HasEnvironment: dataset and of the attribute are the same.""" setattr(self, key, self.get_dataset(key, default, archive)) + def set_default_scheduling(self, priority=None, pipeline_name=None, flush=None): + """Sets the default scheduling options. + + This function should only be called from ``build``.""" + if not self.__in_build: + raise TypeError("set_default_scheduling() should only " + "be called from build()") + + if priority is not None: + self.__scheduler_defaults["priority"] = int(priority) + if pipeline_name is not None: + self.__scheduler_defaults["pipeline_name"] = pipeline_name + if flush is not None: + self.__scheduler_defaults["flush"] = flush + class Experiment: """Base class for top-level experiments. diff --git a/artiq/master/experiments.py b/artiq/master/experiments.py index dc17915fc..892d6d30a 100644 --- a/artiq/master/experiments.py +++ b/artiq/master/experiments.py @@ -46,7 +46,8 @@ class _RepoScanner: entry = { "file": filename, "class_name": class_name, - "arginfo": arginfo + "arginfo": arginfo, + "scheduler_defaults": class_desc["scheduler_defaults"] } entry_dict[name] = entry @@ -115,7 +116,6 @@ class ExperimentDB: t1 = time.monotonic() new_explist = await _RepoScanner(self.worker_handlers).scan(wd) logger.info("repository scan took %d seconds", time.monotonic()-t1) - update_from_dict(self.explist, new_explist) finally: self._scanning = False diff --git a/artiq/master/worker.py b/artiq/master/worker.py index ac12ba4b1..293d0f180 100644 --- a/artiq/master/worker.py +++ b/artiq/master/worker.py @@ -303,8 +303,8 @@ class Worker: await self._create_process(logging.WARNING) r = dict() - def register(class_name, name, arginfo): - r[class_name] = {"name": name, "arginfo": arginfo} + def register(class_name, name, arginfo, scheduler_defaults): + r[class_name] = {"name": name, "arginfo": arginfo, "scheduler_defaults": scheduler_defaults} self.register_experiment = register await self._worker_action({"action": "examine", "file": file}, timeout) diff --git a/artiq/master/worker_impl.py b/artiq/master/worker_impl.py index 6a2df4cfd..b9fa404de 100644 --- a/artiq/master/worker_impl.py +++ b/artiq/master/worker_impl.py @@ -174,11 +174,12 @@ def examine(device_mgr, dataset_mgr, file): if name[-1] == ".": name = name[:-1] argument_mgr = TraceArgumentManager() - exp_class((device_mgr, dataset_mgr, argument_mgr)) + scheduler_defaults = {} + cls = exp_class((device_mgr, dataset_mgr, argument_mgr, scheduler_defaults)) arginfo = OrderedDict( (k, (proc.describe(), group, tooltip)) for k, (proc, group, tooltip) in argument_mgr.requested_args.items()) - register_experiment(class_name, name, arginfo) + register_experiment(class_name, name, arginfo, scheduler_defaults) finally: new_keys = set(sys.modules.keys()) for key in new_keys - previous_keys: @@ -277,7 +278,7 @@ def main(): os.makedirs(dirname, exist_ok=True) os.chdir(dirname) argument_mgr = ProcessArgumentManager(expid["arguments"]) - exp_inst = exp((device_mgr, dataset_mgr, argument_mgr)) + exp_inst = exp((device_mgr, dataset_mgr, argument_mgr, {})) put_object({"action": "completed"}) elif action == "prepare": exp_inst.prepare()