diff --git a/src/sorunlib/__init__.py b/src/sorunlib/__init__.py index c07ec0e..ce0b9ee 100644 --- a/src/sorunlib/__init__.py +++ b/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 diff --git a/src/sorunlib/_internal.py b/src/sorunlib/_internal.py index c59c876..94e3d46 100644 --- a/src/sorunlib/_internal.py +++ b/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}") diff --git a/src/sorunlib/hwp.py b/src/sorunlib/hwp.py index 7d5d495..04bf2dd 100644 --- a/src/sorunlib/hwp.py +++ b/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): diff --git a/src/sorunlib/seq.py b/src/sorunlib/seq.py index b43988a..e039ff4 100644 --- a/src/sorunlib/seq.py +++ b/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() diff --git a/src/sorunlib/stimulator.py b/src/sorunlib/stimulator.py index bd9bcc2..c9b72df 100644 --- a/src/sorunlib/stimulator.py +++ b/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() diff --git a/src/sorunlib/wiregrid.py b/src/sorunlib/wiregrid.py index f8b993c..b98cead 100644 --- a/src/sorunlib/wiregrid.py +++ b/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, ' + \