From 8659c769cb2cf9a5c5787b0f266faac9b011e2fa Mon Sep 17 00:00:00 2001
From: Chris Ballance <chris.ballance@physics.ox.ac.uk>
Date: Sat, 2 Feb 2019 22:47:57 +0000
Subject: [PATCH] master/language: add methods to set experiment
 pipeline/priority/flush defaults

---
 artiq/language/environment.py | 17 +++++++++++++++++
 artiq/master/experiments.py   |  4 ++--
 artiq/master/worker.py        |  4 ++--
 artiq/master/worker_impl.py   |  7 ++++---
 4 files changed, 25 insertions(+), 7 deletions(-)

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()