From ede0b37c6e0753aa8ca9ade9261b5fd5723d8825 Mon Sep 17 00:00:00 2001
From: Sebastien Bourdeauducq <sb@m-labs.hk>
Date: Wed, 13 Dec 2023 14:27:04 +0800
Subject: [PATCH] devices: introduce notify_run_end API

---
 artiq/frontend/artiq_run.py           | 1 +
 artiq/frontend/artiq_sinara_tester.py | 1 +
 artiq/master/worker_db.py             | 7 +++++++
 artiq/master/worker_impl.py           | 2 ++
 artiq/test/hardware_testbench.py      | 1 +
 5 files changed, 12 insertions(+)

diff --git a/artiq/frontend/artiq_run.py b/artiq/frontend/artiq_run.py
index f0a28e91c..ae3e5f1c7 100755
--- a/artiq/frontend/artiq_run.py
+++ b/artiq/frontend/artiq_run.py
@@ -222,6 +222,7 @@ def run(with_file=False):
             exp_inst = _build_experiment(device_mgr, dataset_mgr, args)
             exp_inst.prepare()
             exp_inst.run()
+            device_mgr.notify_run_end()
             exp_inst.analyze()
         except CompileError as error:
             return
diff --git a/artiq/frontend/artiq_sinara_tester.py b/artiq/frontend/artiq_sinara_tester.py
index eaee07441..e2d48b903 100755
--- a/artiq/frontend/artiq_sinara_tester.py
+++ b/artiq/frontend/artiq_sinara_tester.py
@@ -804,6 +804,7 @@ def main():
         experiment = SinaraTester((device_mgr, None, None, None))
         experiment.prepare()
         experiment.run(tests)
+        device_mgr.notify_run_end()
         experiment.analyze()
     finally:
         device_mgr.close_devices()
diff --git a/artiq/master/worker_db.py b/artiq/master/worker_db.py
index 18e51f726..a52b63a1f 100644
--- a/artiq/master/worker_db.py
+++ b/artiq/master/worker_db.py
@@ -92,6 +92,13 @@ class DeviceManager:
         self.active_devices.append((desc, dev))
         return dev
 
+    def notify_run_end(self):
+        """Sends a "end of Experiment run stage" notification to
+        all active devices."""
+        for _desc, dev in self.active_devices:
+            if hasattr(dev, "notify_run_end"):
+                dev.notify_run_end()
+
     def close_devices(self):
         """Closes all active devices, in the opposite order as they were
         requested."""
diff --git a/artiq/master/worker_impl.py b/artiq/master/worker_impl.py
index 7c5ad2c2f..7ecb8da10 100644
--- a/artiq/master/worker_impl.py
+++ b/artiq/master/worker_impl.py
@@ -343,6 +343,8 @@ def main():
                     # for end of analyze stage.
                     write_results()
                     raise
+                finally:
+                    device_mgr.notify_run_end()
                 put_completed()
             elif action == "analyze":
                 try:
diff --git a/artiq/test/hardware_testbench.py b/artiq/test/hardware_testbench.py
index 05e925d0c..55a4d33cb 100644
--- a/artiq/test/hardware_testbench.py
+++ b/artiq/test/hardware_testbench.py
@@ -55,6 +55,7 @@ class ExperimentCase(unittest.TestCase):
         try:
             exp = self.create(cls, *args, **kwargs)
             exp.run()
+            self.device_mgr.notify_run_end()
             exp.analyze()
             return exp
         except CompileError as error: