Add signal handling during shutdown operations

src/sorunlib/__init__.py

@@ -1,3 +1,5 @@
+import signal
+
 from . import acu, hwp, seq, smurf, stimulator, wiregrid
 
 from .commands import wait_until
@@ -27,6 +29,15 @@ def initialize(test_mode=False):
            "wait_until",
            "initialize"]
 
+
+# Treat SIGTERM like SIGINT and raise an exception
+def term_handler(sig, frame):
+    raise KeyboardInterrupt
+
+
+signal.signal(signal.SIGTERM, term_handler)
+
+
 # Define the variable '__version__':
 # This has the closest behavior to versioneer that I could find
 # https://github.com/maresb/hatch-vcs-footgun-example

src/sorunlib/_internal.py

@@ -5,9 +5,13 @@
 """
 
 import datetime as dt
+import signal
 import time
 
+from functools import wraps
+
 import ocs
+import sorunlib as run
 
 from sorunlib.commands import _timestamp_to_utc_datetime
 
@@ -180,3 +184,39 @@ def monitor_process(client, operation, stop_time, check_interval=10):
 
         # Recompute diff
         diff = _seconds_until_target(stop_time)
+
+
+def protect_shutdown(f):
+    """Decorator to install temporary signal handlers while operations required
+    to safely shutdown are handled.
+
+    This will catch and print the caught signals to ``stdout`` while shutdown
+    is happening. Currently handles only ``SIGINT`` and ``SIGTERM``.
+
+    """
+    @wraps(f)
+    def wrapper(*args, **kwds):
+        def handler(sig, frame):
+            print(f'Caught {signal.Signals(sig).name} during shutdown.')
+
+        int_handler = signal.signal(signal.SIGINT, handler)
+        term_handler = signal.signal(signal.SIGTERM, handler)
+
+        result = f(*args, **kwds)
+
+        signal.signal(signal.SIGINT, int_handler)
+        signal.signal(signal.SIGTERM, term_handler)
+        return result
+    return wrapper
+
+
+@protect_shutdown
+def stop_smurfs():
+    """Simple wrapper to shutdown all SMuRF systems and handle any errors that
+    occur.
+
+    """
+    try:
+        run.smurf.stream('off')
+    except RuntimeError as e:
+        print(f"Caught error while shutting down SMuRF streams: {e}")

src/sorunlib/hwp.py

@@ -1,5 +1,5 @@
 import sorunlib as run
-from sorunlib._internal import check_response
+from sorunlib._internal import check_response, stop_smurfs
 
 
 def _get_direction():
@@ -74,7 +74,7 @@ def spin_up(freq):
         check_response(hwp, resp)
         run.hwp.set_freq(freq=freq, timeout=1800)
     finally:
-        run.smurf.stream('off')
+        stop_smurfs()
 
 
 def spin_down(active=True, brake_voltage=None):
@@ -97,7 +97,7 @@ def spin_down(active=True, brake_voltage=None):
         resp = hwp.disable_driver_board()
         check_response(hwp, resp)
     finally:
-        run.smurf.stream('off')
+        stop_smurfs()
 
 
 def stop(active=True, brake_voltage=None):

src/sorunlib/seq.py

@@ -4,28 +4,22 @@
 import sorunlib as run
 
 from sorunlib.commands import _timestamp_to_utc_datetime
-from sorunlib._internal import check_response, check_started, monitor_process
+from sorunlib._internal import check_response, check_started, monitor_process, protect_shutdown, stop_smurfs
 
 
 OP_TIMEOUT = 60
 
 
-def _stop_smurfs():
-    # Stop SMuRF streams
-    try:
-        run.smurf.stream('off')
-    except RuntimeError as e:
-        print(f"Caught error while shutting down SMuRF streams: {e}")
-
-
+@protect_shutdown
 def _stop_scan():
     acu = run.CLIENTS['acu']
 
     print("Stopping scan.")
-    _stop_smurfs()
+    stop_smurfs()
 
     # Stop motion
     acu.generate_scan.stop()
+    print("Waiting for telescope motion to stop.")
     resp = acu.generate_scan.wait(timeout=OP_TIMEOUT)
     check_response(acu, resp)
     print("Scan finished.")
@@ -66,10 +60,10 @@ def scan(description, stop_time, width, az_drift=0, tag=None, subtype=None,
 
     acu = run.CLIENTS['acu']
 
-    # Enable SMuRF streams
-    run.smurf.stream('on', subtype=subtype, tag=tag)
-
     try:
+        # Enable SMuRF streams
+        run.smurf.stream('on', subtype=subtype, tag=tag)
+
         # Grab current telescope position
         resp = acu.monitor.status()
         az = resp.session['data']['StatusDetailed']['Azimuth current position']
@@ -110,10 +104,10 @@ def el_nod(el1, el2, num=5, pause=5):
     """
     acu = run.CLIENTS['acu']
 
-    # Enable SMuRF streams
-    run.smurf.stream('on', subtype='cal', tag='el_nods')
-
     try:
+        # Enable SMuRF streams
+        run.smurf.stream('on', subtype='cal', tag='el_nods')
+
         # Grab current telescope position
         resp = acu.monitor.status()
         init_az = resp.session['data']['StatusDetailed']['Azimuth current position']
@@ -129,4 +123,4 @@ def el_nod(el1, el2, num=5, pause=5):
             # Return to initial position
             run.acu.move_to(az=init_az, el=init_el)
     finally:
-        _stop_smurfs()
+        stop_smurfs()

src/sorunlib/stimulator.py

@@ -1,6 +1,6 @@
 import time
 import sorunlib as run
-from sorunlib._internal import check_response
+from sorunlib._internal import check_response, stop_smurfs
 
 ID_SHUTTER = 1
 
@@ -99,10 +99,7 @@ def calibrate_tau(duration_step=10,
 
             time.sleep(duration_step)
     finally:
-        try:
-            run.smurf.stream('off')
-        except RuntimeError as e:
-            print(f"Caught error while shutting down SMuRF streams: {e}")
+        stop_smurfs()
 
         if stop:
             _stop()
@@ -147,10 +144,7 @@ def calibrate_gain(duration=60, speed_rpm=90,
         # Data taking
         time.sleep(duration)
     finally:
-        try:
-            run.smurf.stream('off')
-        except RuntimeError as e:
-            print(f"Caught error while shutting down SMuRF streams: {e}")
+        stop_smurfs()
 
         if stop:
             _stop()

src/sorunlib/wiregrid.py

@@ -1,7 +1,7 @@
 import time
 
 import sorunlib as run
-from sorunlib._internal import check_response, check_running
+from sorunlib._internal import check_response, check_running, stop_smurfs
 
 EL_DIFF_THRESHOLD = 0.5  # deg diff from target that its ok to run calibration
 BORESIGHT_DIFF_THRESHOLD = 0.5  # deg
@@ -322,7 +322,7 @@ def calibrate(continuous=False, elevation_check=True, boresight_check=True,
         eject()
     finally:
         # Stop SMuRF streams
-        run.smurf.stream('off')
+        stop_smurfs()
 
 
 def time_constant(num_repeats=1):
@@ -381,7 +381,7 @@ def time_constant(num_repeats=1):
         insert()
         time.sleep(5)
     finally:
-        run.smurf.stream('off')
+        stop_smurfs()
 
     for i in range(num_repeats):
         if current_hwp_direction == 'ccw':
@@ -404,7 +404,7 @@ def time_constant(num_repeats=1):
             # Run stepwise rotation
             rotate(continuous=False)
         finally:
-            run.smurf.stream('off')
+            stop_smurfs()
 
         # Stop the HWP while streaming
         try:
@@ -413,7 +413,7 @@ def time_constant(num_repeats=1):
             run.smurf.stream('on', tag=stream_tag, subtype='cal')
             run.hwp.stop(active=True)
         finally:
-            run.smurf.stream('off')
+            stop_smurfs()
 
         # Reverse the HWP while streaming
         try:
@@ -430,7 +430,7 @@ def time_constant(num_repeats=1):
                 run.hwp.set_freq(freq=-2.0)
             current_hwp_direction = target_hwp_direction
         finally:
-            run.smurf.stream('off')
+            stop_smurfs()
 
     # Run stepwise rotation after changing the HWP rotation
     try:
@@ -441,7 +441,7 @@ def time_constant(num_repeats=1):
         # Run stepwise rotation
         rotate(continuous=False)
     finally:
-        run.smurf.stream('off')
+        stop_smurfs()
 
     # Bias step (the wire grid is on the window)
     # After changing the HWP rotation
@@ -458,7 +458,7 @@ def time_constant(num_repeats=1):
         eject()
         time.sleep(5)
     finally:
-        run.smurf.stream('off')
+        stop_smurfs()
 
     # Bias step (the wire grid is off the window)
     bs_tag = 'wiregrid, wg_time_constant, wg_ejected, ' + \