Skip to content

Commit

Permalink
Finished state de-abstraction
Browse files Browse the repository at this point in the history
  • Loading branch information
unkcpz committed Nov 28, 2024
1 parent 684cc07 commit 103193c
Show file tree
Hide file tree
Showing 2 changed files with 63 additions and 5 deletions.
23 changes: 22 additions & 1 deletion src/plumpy/base/state_machine.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,7 @@
import os
import sys
from types import TracebackType
from typing import Any, Callable, Dict, Hashable, Iterable, List, Optional, Sequence, Set, Type, Union, cast
from typing import Any, Callable, ClassVar, Dict, Hashable, Iterable, List, Optional, Protocol, Sequence, Set, Type, Union, cast, runtime_checkable

from plumpy.futures import Future

Expand Down Expand Up @@ -165,6 +165,27 @@ def do_exit(self) -> None:
self.in_state = False


@runtime_checkable
class StateP(Protocol):
LABEL: ClassVar[str]

# FIXME: fix the LABEL_TYPE
ALLOWED: ClassVar[set[LABEL_TYPE]]

def do_enter(self) -> None:
...

def do_exit(self) -> None:
...

def create_state(self, state_label: Hashable, *args: Any, **kwargs: Any) -> 'State':
...

@classmethod
def is_terminal(cls) -> bool:
...


class StateEventHook(enum.Enum):
"""
Hooks that can be used to register callback at various points in the state transition
Expand Down
45 changes: 41 additions & 4 deletions src/plumpy/process_states.py
Original file line number Diff line number Diff line change
Expand Up @@ -267,7 +267,6 @@ def _action_command(self, command: Union[Kill, Stop, Wait, Continue]) -> State:
return cast(State, state) # casting from base.State to process.State


# class Waiting(state_machine.State):
class Waiting(state_machine.State, persistence.Savable):
"""The basic waiting state."""

Expand All @@ -292,6 +291,7 @@ def __str__(self) -> str:
state_info += f' ({self.msg})'
return state_info

# FIXME: fully get rid of state_machine.State as parent class (as a protocol with contract)
def __init__(
self,
process: 'Process',
Expand All @@ -314,6 +314,7 @@ def process(self) -> state_machine.StateMachine:
"""
return self._process

# FIXME: this is shared by all states
def create_state(self, state_label: Hashable, *args: Any, **kwargs: Any) -> 'state_machine.State':
return self._process.create_state(state_label, *args, **kwargs)

Expand Down Expand Up @@ -424,15 +425,51 @@ def get_exc_info(self) -> Tuple[Optional[Type[BaseException]], Optional[BaseExce
return type(self.exception) if self.exception else None, self.exception, self.traceback


@auto_persist('result', 'successful')
class Finished(State):
class Finished(state_machine.State, persistence.Savable):
LABEL = ProcessState.FINISHED
is_terminal_state = True

_auto_persist = {'result', 'successful'}

def __init__(self, process: 'Process', result: Any, successful: bool) -> None:
super().__init__(process)
self._process = process
self.in_state: bool = False
self.result = result
self.successful = successful

@property
def process(self) -> state_machine.StateMachine:
"""
:return: The process
"""
return self._process

# FIXME: this is shared by all states
def create_state(self, state_label: Hashable, *args: Any, **kwargs: Any) -> 'state_machine.State':
return self._process.create_state(state_label, *args, **kwargs)

# FIXME: this is shared
def save_instance_state(self, out_state: SAVED_STATE_TYPE, save_context: persistence.LoadSaveContext) -> None:
super().save_instance_state(out_state, save_context)

# FIXME: this is shared
def load_instance_state(self, saved_state: SAVED_STATE_TYPE, load_context: persistence.LoadSaveContext) -> None:
super().load_instance_state(saved_state, load_context)

def do_enter(self) -> None:
self.in_state = True

def do_exit(self) -> None:
if self.is_terminal():
raise exceptions.InvalidStateError(f'Cannot exit a terminal state {self.LABEL}')

self.in_state = False

@classmethod
def is_terminal(cls) -> bool:
# deprecated using class attribute `is_terminal_state` directly.
return cls.is_terminal_state


@auto_persist('msg')
class Killed(State):
Expand Down

0 comments on commit 103193c

Please sign in to comment.