worker: cleaner termination on exception in user code, improve unittest

m-labs · Oct 12, 2015 · 1d14975 · 1d14975
1 parent fbf94f9
commit 1d14975
Showing 3 changed files with 35 additions and 2 deletions.
diff --git a/artiq/master/worker.py b/artiq/master/worker.py
@@ -21,6 +21,10 @@ class WorkerWatchdogTimeout(Exception):
     pass
 
 
+class WorkerException(Exception):
+    pass
+
+
 class WorkerError(Exception):
     pass
 
@@ -159,6 +163,8 @@ async def _handle_worker_requests(self):
                 return True
             elif action == "pause":
                 return False
+            elif action == "exception":
+                raise WorkerException
             del obj["action"]
             if action == "create_watchdog":
                 func = self.create_watchdog

diff --git a/artiq/master/worker_impl.py b/artiq/master/worker_impl.py
@@ -1,6 +1,7 @@
 import sys
 import time
 import os
+import traceback
 
 from artiq.protocols import pyon
 from artiq.tools import file_import
@@ -214,6 +215,9 @@ def main():
                 put_object({"action": "completed"})
             elif action == "terminate":
                 break
+    except:
+        traceback.print_exc()
+        put_object({"action": "exception"})
     finally:
         device_mgr.close_devices()
 

diff --git a/artiq/test/worker.py b/artiq/test/worker.py
@@ -8,6 +8,22 @@
 from artiq.master.worker import *
 
 
+class SimpleExperiment(EnvExperiment):
+    def build(self):
+        pass
+
+    def run(self):
+        pass
+
+
+class ExceptionTermination(EnvExperiment):
+    def build(self):
+        pass
+
+    def run(self):
+        raise TypeError
+
+
 class WatchdogNoTimeout(EnvExperiment):
     def build(self):
         pass
@@ -53,18 +69,25 @@ def _run_experiment(class_name):
         "arguments": dict()
     }
     loop = asyncio.get_event_loop()
-    worker = Worker()
+    worker = Worker(handlers={"log": lambda message: None})
     loop.run_until_complete(_call_worker(worker, expid))
 
 
-class WatchdogCase(unittest.TestCase):
+class WorkerCase(unittest.TestCase):
     def setUp(self):
         if os.name == "nt":
             self.loop = asyncio.ProactorEventLoop()
         else:
             self.loop = asyncio.new_event_loop()
         asyncio.set_event_loop(self.loop)
 
+    def test_simple_run(self):
+        _run_experiment("SimpleExperiment")
+
+    def test_exception(self):
+        with self.assertRaises(WorkerException):
+            _run_experiment("ExceptionTermination")
+
     def test_watchdog_no_timeout(self):
         _run_experiment("WatchdogNoTimeout")