Improve abort command (#49)

* Reset timings in reset command

* Improve abort command

* Add --reset option to abort command

* Close the shutter when aborting an exposure

* Additional tests

* Improve test

* Fix test

* Close shutter before aborting exposures

* Update changelog

CHANGELOG.md

@@ -1,5 +1,12 @@
 # Changelog
 
+## Next version
+
+### ✨ Improved
+
+* [#49](https://github.com/sdss/archon/pull/49) Improve the `abort` command. When an integration or readout is ongoing, the current task is stored in the delegate and cancelled if an abort happens. This should immediately cancel the command and any controller action. To perform a full abort, one should also run the `reset` command, which will reset the controller, including a reset timings, or call `abort --reset`.
+
+
 ## 0.13.4 - February 27, 2024
 
 ### ✨ Improved

archon/actor/commands/expose.py

@@ -191,15 +191,17 @@ async def expose(
         flavours = [flavour, "dark"] if with_dark else [flavour]
         for nf, this_flavour in enumerate(flavours):
             delegate.use_shutter = not no_shutter
-            exposure_result = await delegate.expose(
-                command,
-                selected_controllers,
-                flavour=this_flavour,
-                exposure_time=exposure_time,
-                readout=False,
-                window_mode=window_mode,
-                write=not no_write,
-                seqno=seqno,
+            exposure_result = await delegate.set_task(
+                delegate.expose(
+                    command,
+                    selected_controllers,
+                    flavour=this_flavour,
+                    exposure_time=exposure_time,
+                    readout=False,
+                    window_mode=window_mode,
+                    write=not no_write,
+                    seqno=seqno,
+                )
             )
 
             if not exposure_result:
@@ -213,7 +215,7 @@ async def expose(
                 if is_async:
                     command.finish("Returning while readout is ongoing.")
 
-                readout_task = asyncio.create_task(
+                readout_task = delegate.set_task(
                     delegate.readout(
                         command,
                         extra_header=extra_header,
@@ -267,10 +269,12 @@ async def read(
         command.warning("Waiting for image recovery to finish.")
         await command.actor.exposure_recovery.locker.wait()
 
-    result = await delegate.readout(
-        command,
-        extra_header=extra_header,
-        delay_readout=delay_readout,
+    result = await delegate.set_task(
+        delegate.readout(
+            command,
+            extra_header=extra_header,
+            delay_readout=delay_readout,
+        )
     )
 
     if result is True:
@@ -281,44 +285,52 @@ async def read(
 
 @parser.command()
 @click.option("--flush", is_flag=True, help="Flush the device after aborting.")
-@click.option("--force", is_flag=True, help="Forces abort.")
-@click.option("--all", "all_", is_flag=True, help="Aborts all the controllers.")
+@click.option(
+    "--reset",
+    is_flag=True,
+    help="Reset the controllers after aborting. This will discard any "
+    "ongoing exposures or pending readouts.",
+)
 async def abort(
     command: Command[archon.actor.actor.ArchonActor],
     controllers: dict[str, ArchonController],
-    flush: bool,
-    force: bool,
-    all_: bool,
+    flush: bool = False,
+    reset: bool = False,
 ):
-    """Aborts the exposure."""
+    """Aborts the current exposure."""
 
     assert command.actor
 
-    if all_:
-        force = True
-
-    expose_data = command.actor.exposure_delegate.expose_data
+    delegate = command.actor.exposure_delegate
+    expose_data = delegate.expose_data
 
-    if expose_data is None:
-        if force:
-            command.warning(error="No exposure found.")
-        else:
-            command.actor.exposure_delegate.reset()
-            return command.fail(error="No exposure found.")
+    if delegate is None or delegate._command is None:
+        return command.fail(error="Expose command is not running.")
 
     scontr: list[ArchonController]
-    if all_ or not expose_data:
+    if not expose_data:
         scontr = list(controllers.values())
     else:
         scontr = expose_data.controllers
 
     command.debug(text="Aborting exposures")
+
+    await delegate.shutter(False)  # Close the shutter.
+
     try:
         await asyncio.gather(*[contr.abort(readout=False) for contr in scontr])
     except ArchonError as err:
         return command.fail(error=f"Failed aborting exposures: {err}")
     finally:
-        command.actor.exposure_delegate.reset()
+        # This will also cancel any ongoing exposure or readout.
+        delegate.fail("Exposure was aborted")
+
+    if reset:
+        command.debug(text="Resetting controllers")
+        try:
+            await asyncio.gather(*[contr.reset(reset_timing=True) for contr in scontr])
+        except ArchonError as err:
+            return command.fail(error=f"Failed resetting devices: {err}")
 
     if flush:
         command.debug(text="Flushing devices")

archon/actor/commands/reset.py

@@ -25,7 +25,7 @@ async def reset(command: Command, controller: ArchonController):
     assert command.actor
 
     try:
-        await controller.reset()
+        await controller.reset(reset_timing=True)
         command.actor.expose_data = None
     except (ArchonControllerError, ArchonError) as err:
         return error_controller(command, controller, f"Failed resetting: {err}")

archon/actor/delegate.py

@@ -23,6 +23,7 @@
 from typing import (
     TYPE_CHECKING,
     Any,
+    Coroutine,
     Dict,
     Generic,
     List,
@@ -104,6 +105,8 @@ def __init__(self, actor: Actor_co):
         self._command: Command[Actor_co] | None = None
         self._expose_cotasks: asyncio.Task | None = None
 
+        self._current_task: asyncio.Task | None = None
+
         self._check_fitsio()
 
     @property
@@ -139,6 +142,10 @@ def reset(self):
     def fail(self, message: str | None = None):
         """Fails a command."""
 
+        if self._current_task:
+            self._current_task.cancel()
+            self._current_task = None
+
         if message:
             self.command.fail(error=message)
         else:
@@ -148,6 +155,13 @@ def fail(self, message: str | None = None):
 
         return False
 
+    def set_task(self, coro: Coroutine):
+        """Schedules a coroutine as a task. The current task is cancelled by `.fail`."""
+
+        self._current_task = asyncio.create_task(coro)
+
+        return self._current_task
+
     async def pre_expose(self, controllers: List[ArchonController]):
         """A routine that runs before integration begins."""
 

archon/controller/controller.py

@@ -752,7 +752,13 @@ async def set_autoflush(self, mode: bool):
 
         self.auto_flush = mode
 
-    async def reset(self, autoflush=True, release_timing=True, update_status=True):
+    async def reset(
+        self,
+        autoflush: bool = True,
+        release_timing: bool = True,
+        reset_timing: bool = False,
+        update_status: bool = True,
+    ):
         """Resets timing and discards current exposures."""
 
         self._parse_params()
@@ -774,9 +780,15 @@ async def reset(self, autoflush=True, release_timing=True, update_status=True):
             for param in default_parameters:
                 await self.set_param(param, default_parameters[param])
 
+        timing_commands: list[str] = []
+        if reset_timing:
+            release_timing = True
+            timing_commands.append("RESETTIMING")
+
         if release_timing:
             log.info(f"{self.name}: restarting timing .")
-            for cmd_str in ["RELEASETIMING"]:
+            timing_commands.append("RELEASETIMING")
+            for cmd_str in timing_commands:
                 cmd = await self.send_command(cmd_str, timeout=1)
                 if not cmd.succeeded():
                     self.update_status(ControllerStatus.ERROR)

tests/actor/test_command_expose.py

@@ -6,10 +6,12 @@
 # @Filename: test_command_expose.py
 # @License: BSD 3-clause (http://www.opensource.org/licenses/BSD-3-Clause)
 
+from __future__ import annotations
+
 import asyncio
 import os
 
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 from astropy.io import fits
 
@@ -19,6 +21,12 @@
 from archon.exceptions import ArchonError
 
 
+if TYPE_CHECKING:
+    from pytest_mock import MockFixture
+
+    from archon.controller.controller import ArchonController
+
+
 async def test_expose_start(delegate, actor: ArchonActor):
     command = await actor.invoke_mock_command("expose --no-readout 0.01")
     await command
@@ -133,47 +141,81 @@ async def fail(command, *args, **kwargs):
     assert command.status.did_fail
 
 
-async def test_expose_abort(delegate, actor: ArchonActor):
-    await actor.invoke_mock_command("expose --no-readout 1")
+async def test_expose_abort(
+    delegate,
+    actor: ArchonActor,
+    controller: ArchonController,
+    mocker: MockFixture,
+):
+    reset_mock = mocker.patch.object(controller, "reset", return_value=True)
+
+    expose_command = await actor.invoke_mock_command("expose --no-readout 1")
     await asyncio.sleep(0.5)
 
-    abort = await actor.invoke_mock_command("abort")
+    abort = await actor.invoke_mock_command("abort --reset")
     await abort
 
+    await asyncio.sleep(0.5)
+
     assert abort.status.did_succeed
 
+    assert expose_command.status.did_fail
+    assert expose_command.replies[-1].message["error"] == "Exposure was aborted"
 
-async def test_expose_abort_no_expose_data(delegate, actor: ArchonActor):
+    assert delegate._current_task is None
+
+    reset_mock.assert_called()
+    reset_mock.assert_called_with(reset_timing=True)
+
+
+async def test_expose_abort_reset_fails(
+    delegate,
+    actor: ArchonActor,
+    controller: ArchonController,
+    mocker: MockFixture,
+):
     await actor.invoke_mock_command("expose --no-readout 1")
-    await asyncio.sleep(0.05)
+    await asyncio.sleep(0.5)
 
-    actor.exposure_delegate.expose_data = None
+    mocker.patch.object(controller, "reset", side_effect=ArchonError)
 
-    abort = await actor.invoke_mock_command("abort")
+    abort = await actor.invoke_mock_command("abort --reset")
     await abort
 
     assert abort.status.did_fail
 
 
-async def test_expose_abort_no_expose_data_force(delegate, actor: ArchonActor):
+async def test_expose_abort_no_command(delegate, actor: ArchonActor):
+    delegate._command = None
+
+    abort = await actor.invoke_mock_command("abort --reset")
+    await abort
+
+    assert abort.status.did_fail
+
+    error = abort.replies[-1].message
+    assert error["error"] == "Expose command is not running."
+
+
+async def test_expose_abort_no_expose_data(delegate, actor: ArchonActor):
     await actor.invoke_mock_command("expose --no-readout 1")
     await asyncio.sleep(0.05)
 
     actor.exposure_delegate.expose_data = None
 
-    abort = await actor.invoke_mock_command("abort --force")
+    abort = await actor.invoke_mock_command("abort")
     await abort
 
     assert abort.status.did_succeed
 
 
-async def test_expose_abort_no_expose_data_all(delegate, actor: ArchonActor):
+async def test_expose_abort_no_expose_data_force(delegate, actor: ArchonActor):
     await actor.invoke_mock_command("expose --no-readout 1")
-    await asyncio.sleep(0.5)
+    await asyncio.sleep(0.05)
 
     actor.exposure_delegate.expose_data = None
 
-    abort = await actor.invoke_mock_command("abort --all")
+    abort = await actor.invoke_mock_command("abort")
     await abort
 
     assert abort.status.did_succeed