fix(core): do not discard operation outputs when shutdown is set (#219)

* fix(core): operation output typing

Update typing to reflect that None can be returned by an operation

* fix(core): save operation output when shutdown set

The code assumed that if shutdown was set when an operation returned there would be no output, but this may not be the case #202

Also the logs assumed that shutdown could only be set by SIGTERM. However because of #200 we cannot for sure if this is the case. Hence the logs are temporarily changed here to reflect this.

Author: michael-johnston

orchestrator/modules/operators/_explore_orchestration.py

@@ -141,7 +141,7 @@ def handler(sig, frame):
 
 def run_explore_operation_core_closure(
     operator: "OperatorActor", state: "DiscoverySpaceManagerActor"
-) -> typing.Callable[[], OperationOutput]:
+) -> typing.Callable[[], OperationOutput | None]:
 
     def _run_explore_operation_core() -> OperationOutput:
         import numpy as np

orchestrator/modules/operators/_general_orchestration.py

@@ -37,7 +37,7 @@ def run_general_operation_core_closure(
             FunctionOperationInfo,
             ...,
         ],
-        OperationOutput,
+        OperationOutput | None,
     ],
     discovery_space: DiscoverySpace,
     operationInfo: FunctionOperationInfo,
@@ -47,7 +47,7 @@ def run_general_operation_core_closure(
     def _run_general_operation_core() -> OperationOutput:
         return operation_function(
             discovery_space, operationInfo, **operation_parameters
-        )  # type: OperationOutput
+        )  # type: OperationOutput | None
 
     return _run_general_operation_core
 

orchestrator/modules/operators/_orchestrate_core.py

@@ -46,7 +46,7 @@ def log_space_details(discovery_space: "DiscoverySpace"):
 
 
 def _run_operation_harness(
-    run_closure: typing.Callable[[], OperationOutput],
+    run_closure: typing.Callable[[], OperationOutput | None],
     base_operation_configuration: BaseOperationRunConfiguration,
     discovery_space: DiscoverySpace,
     operation_identifier: str | None = None,
@@ -85,7 +85,7 @@ def _run_operation_harness(
             OperationResourceStatus(event=OperationResourceEventEnum.STARTED)
         )
         discovery_space.metadataStore.updateResource(operation_resource)
-        operation_output = run_closure()
+        operation_output: OperationOutput | None = run_closure()
     except KeyboardInterrupt:
         sys.stdout.flush()
         moduleLog.warning("Caught keyboard interrupt - initiating graceful shutdown")
@@ -124,12 +124,19 @@ def _run_operation_harness(
         time.sleep(1)
         sys.stdout.flush()
         if shutdown:
-            moduleLog.warning("Operation exited normally but a signal was sent")
-            operation_output = None
-            operationStatus = OperationResourceStatus(
-                event=OperationResourceEventEnum.FINISHED,
-                exit_state=OperationExitStateEnum.ERROR,
-                message="Operation exited due to SIGTERM)",
+            moduleLog.warning(
+                "Operation exited normally but an external event e.g. SIGTERM, has already initiated shutdown"
+            )
+            if operation_output:
+                moduleLog.info("Operation returned output - will save")
+
+            operationStatus = (
+                OperationResourceStatus(
+                    event=OperationResourceEventEnum.FINISHED,
+                    exit_state=OperationExitStateEnum.ERROR,
+                    message="An external event e.g. SIGTERM, initiated shutdown. "
+                    "This may have caused the operation to exit early",
+                ),
             )
         else:
             if not operation_output: