WIP: create_state refact

src/plumpy/base/state_machine.py

@@ -147,7 +147,6 @@ def interrupt(self, reason: Exception) -> None: ...
 
 @runtime_checkable
 class Proceedable(Protocol):
-
     def execute(self) -> State | None:
         """
         Execute the state, performing the actions that this state is responsible for.
@@ -354,17 +353,13 @@ def get_debug(self) -> bool:
     def set_debug(self, enabled: bool) -> None:
         self._debug: bool = enabled
 
-    def create_state(self, state_label: Hashable, *args: Any, **kwargs: Any) -> State:
-        # XXX: this method create state from label, which is duplicate as _create_state_instance and less generic
-        # because the label is defined after the state and required to be know before calling this function.
-        # This method should be replaced by `_create_state_instance`.
-        # aiida-core using this method for its Waiting state override.
-        try:
-            state_cls = self.get_states_map()[state_label]
-            return state_cls(self, *args, **kwargs)
-        except KeyError:
+    def create_state(self, state_label: Hashable, *args, **kwargs: Any) -> State:
+        if state_label not in self.get_states_map():
             raise ValueError(f'{state_label} is not a valid state')
 
+        state_cls = self.get_states_map()[state_label]
+        return state_cls(*args, **kwargs)
+
     def _exit_current_state(self, next_state: State) -> None:
         """Exit the given state"""
 
@@ -387,9 +382,3 @@ def _enter_next_state(self, next_state: State) -> None:
         next_state.enter()
         self._state = next_state
         self._fire_state_event(StateEventHook.ENTERED_STATE, last_state)
-
-    def _create_state_instance(self, state_cls: type[State], **kwargs: Any) -> State:
-        if state_cls.LABEL not in self.get_states_map():
-            raise ValueError(f'{state_cls.LABEL} is not a valid state')
-
-        return state_cls(self, **kwargs)

src/plumpy/process_states.py

@@ -138,6 +138,7 @@ class ProcessState(Enum):
     The possible states that a :class:`~plumpy.processes.Process` can be in.
     """
 
+    # FIXME: see LSP error of return a exception, the type is Literal[str] which is invariant, tricky
     CREATED = 'created'
     RUNNING = 'running'
     WAITING = 'waiting'
@@ -173,7 +174,9 @@ def load_instance_state(self, saved_state: SAVED_STATE_TYPE, load_context: persi
         self.run_fn = getattr(self.process, saved_state[self.RUN_FN])
 
     async def execute(self) -> state_machine.State:
-        return self.process.create_state(ProcessState.RUNNING, self.run_fn, *self.args, **self.kwargs)
+        return self.process.create_state(
+            ProcessState.RUNNING, process=self.process, run_fn=self.run_fn, *self.args, **self.kwargs
+        )
 
     def enter(self) -> None: ...
 
@@ -241,8 +244,11 @@ async def execute(self) -> state_machine.State:
                 # Let this bubble up to the caller
                 raise
             except Exception:
-                excepted = self.process.create_state(ProcessState.EXCEPTED, *sys.exc_info()[1:])
-                return cast(state_machine.State, excepted)
+                # state_cls: Excepted = self.process.get_states_map()[ProcessState.EXCEPTED]
+                _, exception, traceback = sys.exc_info()
+                # excepted = state_cls(exception=exception, traceback=traceback)
+                excepted = Excepted(exception=exception, traceback=traceback)
+                return excepted
             else:
                 if not isinstance(result, Command):
                     if isinstance(result, exceptions.UnsuccessfulResult):
@@ -258,15 +264,25 @@ async def execute(self) -> state_machine.State:
 
     def _action_command(self, command: Union[Kill, Stop, Wait, Continue]) -> state_machine.State:
         if isinstance(command, Kill):
-            state = self.process.create_state(ProcessState.KILLED, command.msg)
+            state = self.process.create_state(ProcessState.KILLED, msg=command.msg)
         # elif isinstance(command, Pause):
         #     self.pause()
         elif isinstance(command, Stop):
-            state = self.process.create_state(ProcessState.FINISHED, command.result, command.successful)
+            state = self.process.create_state(
+                ProcessState.FINISHED, result=command.result, successful=command.successful
+            )
         elif isinstance(command, Wait):
-            state = self.process.create_state(ProcessState.WAITING, command.continue_fn, command.msg, command.data)
+            state = self.process.create_state(
+                ProcessState.WAITING,
+                process=self.process,
+                done_callback=command.continue_fn,
+                msg=command.msg,
+                data=command.data,
+            )
         elif isinstance(command, Continue):
-            state = self.process.create_state(ProcessState.RUNNING, command.continue_fn, *command.args)
+            state = self.process.create_state(
+                ProcessState.RUNNING, process=self.process, run_fn=command.continue_fn, *command.args, **command.kwargs
+            )
         else:
             raise ValueError('Unrecognised command')
 
@@ -345,9 +361,13 @@ async def execute(self) -> state_machine.State:  # type: ignore
             raise
 
         if result == NULL:
-            next_state = self.process.create_state(ProcessState.RUNNING, self.done_callback)
+            next_state = self.process.create_state(
+                ProcessState.RUNNING, process=self.process, done_callback=self.done_callback
+            )
         else:
-            next_state = self.process.create_state(ProcessState.RUNNING, self.done_callback, result)
+            next_state = self.process.create_state(
+                ProcessState.RUNNING, process=self.process, done_callback=self.done_callback, *result
+            )
 
         return cast(state_machine.State, next_state)  # casting from base.State to process.State
 
@@ -367,10 +387,10 @@ def exit(self) -> None: ...
 @final
 class Excepted(persistence.Savable):
     """
-    Excepted state, can optionally provide exception and trace_back
+    Excepted state, can optionally provide exception and traceback
 
     :param exception: The exception instance
-    :param trace_back: An optional exception traceback
+    :param traceback: An optional exception traceback
     """
 
     LABEL = ProcessState.EXCEPTED
@@ -383,18 +403,15 @@ class Excepted(persistence.Savable):
 
     def __init__(
         self,
-        process: 'Process',
         exception: Optional[BaseException],
-        trace_back: Optional[TracebackType] = None,
+        traceback: Optional[TracebackType] = None,
     ):
         """
-        :param process: The associated process
         :param exception: The exception instance
-        :param trace_back: An optional exception traceback
+        :param traceback: An optional exception traceback
         """
-        self.process = process
         self.exception = exception
-        self.traceback = trace_back
+        self.traceback = traceback
 
     def __str__(self) -> str:
         exception = traceback.format_exception_only(type(self.exception) if self.exception else None, self.exception)[0]
@@ -408,7 +425,6 @@ def save_instance_state(self, out_state: SAVED_STATE_TYPE, save_context: persist
 
     def load_instance_state(self, saved_state: SAVED_STATE_TYPE, load_context: persistence.LoadSaveContext) -> None:
         super().load_instance_state(saved_state, load_context)
-        self.process = load_context.process
 
         self.exception = yaml.load(saved_state[self.EXC_VALUE], Loader=Loader)
         if _HAS_TBLIB:
@@ -450,14 +466,12 @@ class Finished(persistence.Savable):
 
     is_terminal = True
 
-    def __init__(self, process: 'Process', result: Any, successful: bool) -> None:
-        self.process = process
+    def __init__(self, result: Any, successful: bool) -> None:
         self.result = result
         self.successful = successful
 
     def load_instance_state(self, saved_state: SAVED_STATE_TYPE, load_context: persistence.LoadSaveContext) -> None:
         super().load_instance_state(saved_state, load_context)
-        self.process = load_context.process
 
     def enter(self) -> None: ...
 
@@ -481,17 +495,14 @@ class Killed(persistence.Savable):
 
     is_terminal = True
 
-    def __init__(self, process: 'Process', msg: Optional[MessageType]):
+    def __init__(self, msg: Optional[MessageType]):
         """
-        :param process: The associated process
         :param msg: Optional kill message
         """
-        self.process = process
         self.msg = msg
 
     def load_instance_state(self, saved_state: SAVED_STATE_TYPE, load_context: persistence.LoadSaveContext) -> None:
         super().load_instance_state(saved_state, load_context)
-        self.process = load_context.process
 
     def enter(self) -> None: ...
 

src/plumpy/processes.py

@@ -1074,8 +1074,8 @@ def transition_failed(
         if final_state == process_states.ProcessState.CREATED:
             raise exception.with_traceback(trace)
 
-        state_class = self.get_states_map()[process_states.ProcessState.EXCEPTED]
-        new_state = self._create_state_instance(state_class, exception=exception, trace_back=trace)
+        # state_class = self.get_states_map()[process_states.ProcessState.EXCEPTED]
+        new_state = self.create_state(process_states.ProcessState.EXCEPTED, exception=exception, traceback=trace)
         self.transition_to(new_state)
 
     def pause(self, msg: Union[str, None] = None) -> Union[bool, futures.CancellableAction]:
@@ -1148,10 +1148,12 @@ def _create_interrupt_action(self, exception: process_states.Interruption) -> fu
 
             def do_kill(_next_state: state_machine.State) -> Any:
                 try:
-                    state_class = self.get_states_map()[process_states.ProcessState.KILLED]
-                    new_state = self._create_state_instance(state_class, msg=exception.msg)
+                    # state_class = self.get_states_map()[process_states.ProcessState.KILLED]
+                    new_state = self.create_state(process_states.ProcessState.KILLED, msg=exception.msg)
                     self.transition_to(new_state)
                     return True
+                    # FIXME: if try block except, will hit deadlock in event loop
+                    # need to know how to debug it, and where to set a timeout.
                 finally:
                     self._killing = None
 
@@ -1196,14 +1198,14 @@ def resume(self, *args: Any) -> None:
         return self._state.resume(*args)  # type: ignore
 
     @event(to_states=process_states.Excepted)
-    def fail(self, exception: Optional[BaseException], trace_back: Optional[TracebackType]) -> None:
+    def fail(self, exception: Optional[BaseException], traceback: Optional[TracebackType]) -> None:
         """
         Fail the process in response to an exception
         :param exception: The exception that caused the failure
-        :param trace_back: Optional exception traceback
+        :param traceback: Optional exception traceback
         """
-        state_class = self.get_states_map()[process_states.ProcessState.EXCEPTED]
-        new_state = self._create_state_instance(state_class, exception=exception, trace_back=trace_back)
+        # state_class = self.get_states_map()[process_states.ProcessState.EXCEPTED]
+        new_state = self.create_state(process_states.ProcessState.EXCEPTED, exception=exception, traceback=trace_back)
         self.transition_to(new_state)
 
     def kill(self, msg: Optional[MessageType] = None) -> Union[bool, asyncio.Future]:
@@ -1232,8 +1234,8 @@ def kill(self, msg: Optional[MessageType] = None) -> Union[bool, asyncio.Future]
             self._state.interrupt(interrupt_exception)
             return cast(futures.CancellableAction, self._interrupt_action)
 
-        state_class = self.get_states_map()[process_states.ProcessState.KILLED]
-        new_state = self._create_state_instance(state_class, msg=msg)
+        # state_class = self.get_states_map()[process_states.ProcessState.KILLED]
+        new_state = self.create_state(process_states.ProcessState.KILLED, msg=msg)
         self.transition_to(new_state)
         return True
 
@@ -1325,7 +1327,10 @@ async def step(self) -> None:
                 raise
             except Exception:
                 # Overwrite the next state to go to excepted directly
-                next_state = self.create_state(process_states.ProcessState.EXCEPTED, *sys.exc_info()[1:])
+                _, exception, traceback = sys.exc_info()
+                next_state = self.create_state(
+                    process_states.ProcessState.EXCEPTED, exception=exception, traceback=traceback
+                )
                 self._set_interrupt_action(None)
 
             if self._interrupt_action: