Split "preflight" and "execution" transition states and store state data

This changes the way `preflight()` and `execute()` (formerly `run()`) methods are executed by "transitions". Additionally use Python's "pickle" formatted file to store arbitrary state data.

Signed-off-by: Tobias Wolf <[email protected]>

.gitignore

@@ -86,7 +86,7 @@ dmypy.json
 cython_debug/
 
 # Project specific
-data.yaml
+data.pickle
 config.yaml
 .ceph
 .k8s

requirements.txt

@@ -4,6 +4,7 @@ certifi==2024.2.2
 cffi==1.16.0
 charset-normalizer==3.3.2
 cryptography==42.0.4
+dill==0.3.8
 decorator==5.1.1
 Deprecated==1.2.14
 fabric==3.2.2

src/config.example.yaml

@@ -1,5 +1,5 @@
 general:
-  module_data_file: data.yaml
+  machine_pickle_file: data.pickle
 
 logging:
   level: INFO # level at which logging should start

src/rookify/__main__.py

@@ -22,7 +22,7 @@ def main() -> None:
     log = get_logger()
     log.debug("Executing Rookify")
 
-    machine = Machine()
+    machine = Machine(config["general"]["machine_pickle_file"])
     load_modules(machine, config)
 
     machine.execute()

src/rookify/config.schema.yaml

@@ -1,5 +1,5 @@
 general:
-  module_data_file: str()
+  machine_pickle_file: str(required=False)
 
 logging:
   level: str()

src/rookify/modules/__init__.py

@@ -33,7 +33,7 @@ def _load_module(machine: Machine, config: Dict[str, Any], module_name: str) ->
     additional_modules = []
 
     if not hasattr(module, "ModuleHandler") or not callable(
-        getattr(module.ModuleHandler, "register_state")
+        getattr(module.ModuleHandler, "register_states")
     ):
         raise ModuleLoadException(module_name, "Module structure is invalid")
 
@@ -50,7 +50,7 @@ def _load_module(machine: Machine, config: Dict[str, Any], module_name: str) ->
     for module_name in additional_modules:
         _load_module(machine, config, module_name)
 
-    module.ModuleHandler.register_state(machine, config)
+    module.ModuleHandler.register_states(machine, config)
 
 
 def load_modules(machine: Machine, config: Dict[str, Any]) -> None:

src/rookify/modules/analyze_ceph/main.py

@@ -6,10 +6,10 @@
 
 
 class AnalyzeCephHandler(ModuleHandler):
-    def run(self) -> Any:
+    def preflight(self) -> Any:
         commands = ["mon dump", "osd dump", "device ls", "fs dump", "node ls"]
 
-        state = self.machine.get_state("AnalyzeCephHandler")
+        state = self.machine.get_preflight_state("AnalyzeCephHandler")
         state.data: Dict[str, Any] = {}  # type: ignore
 
         for command in commands:
@@ -33,12 +33,10 @@ def run(self) -> Any:
 
         self.logger.info("AnalyzeCephHandler ran successfully.")
 
-    @classmethod
-    def register_state(
-        _, machine: Machine, config: Dict[str, Any], **kwargs: Any
+    @staticmethod
+    def register_preflight_state(
+        machine: Machine, state_name: str, handler: ModuleHandler, **kwargs: Any
     ) -> None:
-        """
-        Register state for transitions
-        """
-
-        super().register_state(machine, config, tags=["data"])
+        ModuleHandler.register_preflight_state(
+            machine, state_name, handler, tags=["data"]
+        )

src/rookify/modules/cephx_auth_config/main.py

@@ -1,6 +1,6 @@
 # -*- coding: utf-8 -*-
 
-from typing import Any, Dict
+from typing import Any
 from ..machine import Machine
 from ..module import ModuleException, ModuleHandler
 
@@ -22,18 +22,16 @@ def preflight(self) -> Any:
                 "Ceph config value auth_client_required does not contain cephx"
             )
 
-        self.machine.get_state("CephXAuthHandler").verified = True
+        self.machine.get_preflight_state("CephXAuthHandler").verified = True
         self.logger.info("Validated Ceph to expect cephx auth")
 
     def is_cephx_set(self, values: str) -> Any:
         return "cephx" in [value.strip() for value in values.split(",")]
 
-    @classmethod
-    def register_state(
-        _, machine: Machine, config: Dict[str, Any], **kwargs: Any
+    @staticmethod
+    def register_preflight_state(
+        machine: Machine, state_name: str, handler: ModuleHandler, **kwargs: Any
     ) -> None:
-        """
-        Register state for transitions
-        """
-
-        super().register_state(machine, config, tags=["verified"])
+        ModuleHandler.register_preflight_state(
+            machine, state_name, handler, tags=["verified"]
+        )

src/rookify/modules/example/main.py

@@ -12,6 +12,6 @@ def preflight(self) -> None:
         # Do something for checking if all needed preconditions are met else throw ModuleException
         raise ModuleException("Example module was loaded, so aborting!")
 
-    def run(self) -> Any:
+    def execute(self) -> Any:
         # Run the migration tasks
         return {}

src/rookify/modules/machine.py

@@ -1,28 +1,108 @@
 # -*- coding: utf-8 -*-
 
+from dill import Pickler, Unpickler
 from transitions import MachineError
 from transitions import Machine as _Machine
 from transitions.extensions.states import add_state_features, Tags, Timeout
-from typing import Any, Dict
+from typing import Any, Dict, IO, Optional, List
+from ..logger import get_logger
 
 
 @add_state_features(Tags, Timeout)
 class Machine(_Machine):  # type: ignore
-    def __init__(self) -> None:
+    STATE_NAME_EXECUTION_PREFIX = "Execution"
+    STATE_NAME_PREFLIGHT_PREFIX = "Preflight"
+
+    def __init__(self, machine_pickle_file: Optional[str] = None) -> None:
+        self._machine_pickle_file = machine_pickle_file
+        self._execution_states: List[str] = []
+        self._preflight_states: List[str] = []
+
         _Machine.__init__(self, states=["uninitialized"], initial="uninitialized")
 
-    def add_migrating_state(self, name: Any, **kwargs: Dict[str, Any]) -> None:
-        if not isinstance(name, str):
-            raise MachineError("Migration state name must be string")
-        self.add_state(name, **kwargs)
+    def add_execution_state(self, name: str, **kwargs: Dict[str, Any]) -> None:
+        self._execution_states.append(self.__class__.state_cls(name, **kwargs))
+
+    def add_preflight_state(self, name: str, **kwargs: Dict[str, Any]) -> None:
+        self._preflight_states.append(self.__class__.state_cls(name, **kwargs))
 
     def execute(self) -> None:
+        for state in self._preflight_states + self._execution_states:
+            self.add_state(state)
+
         self.add_state("migrated")
         self.add_ordered_transitions(loop=False)
 
+        if self._machine_pickle_file is None:
+            self._execute()
+        else:
+            with open(self._machine_pickle_file, "ab+") as file:
+                self._execute(file)
+
+    def _execute(self, pickle_file: Optional[IO[Any]] = None) -> None:
+        states_data = {}
+
+        if pickle_file is not None and pickle_file.tell() > 0:
+            pickle_file.seek(0)
+
+            states_data = Unpickler(pickle_file).load()
+            self._restore_state_data(states_data)
+
         try:
             while True:
                 self.next_state()
+
+                if pickle_file is not None:
+                    state_data = self._get_state_tags_data(self.state)
+
+                    if len(state_data) > 0:
+                        states_data[self.state] = state_data
         except MachineError:
             if self.state != "migrated":
                 raise
+        finally:
+            if pickle_file is not None:
+                get_logger().debug("Storing state data: {0}".format(states_data))
+                pickle_file.truncate(0)
+
+                Pickler(pickle_file).dump(states_data)
+
+    def _get_state_tags_data(self, name: str) -> Dict[str, Any]:
+        data = {}
+        state = self.get_state(name)
+
+        if len(state.tags) > 0:
+            for tag in state.tags:
+                data[tag] = getattr(state, tag)
+
+        return data
+
+    def get_execution_state(self, name: Optional[str] = None) -> Any:
+        if name is None:
+            name = self.state
+        else:
+            name = self.__class__.STATE_NAME_EXECUTION_PREFIX + name
+
+        return self.get_state(name)
+
+    def get_preflight_state(self, name: Optional[str] = None) -> Any:
+        if name is None:
+            name = self.state
+        else:
+            name = self.__class__.STATE_NAME_PREFLIGHT_PREFIX + name
+
+        return self.get_state(name)
+
+    def _restore_state_data(self, data: Dict[str, Any]) -> None:
+        for state_name in data:
+            try:
+                state = self.get_state(state_name)
+
+                for tag in data[state_name]:
+                    setattr(state, tag, data[state_name][tag])
+            except Exception as exc:
+                get_logger().debug(
+                    "Restoring state data failed for '{0}': {1!r}".format(
+                        state_name, exc
+                    )
+                )

src/rookify/modules/migrate_monitors/main.py

@@ -1,12 +1,10 @@
 # -*- coding: utf-8 -*-
 
-from typing import Dict, Any
 from ..module import ModuleHandler
 
 
 class MigrateMonitorsHandler(ModuleHandler):
     REQUIRES = ["analyze_ceph"]
 
-    def run(self) -> Dict[str, Any]:
+    def execute(self) -> None:
         self.logger.info("MigrateMonitorsHandler ran successfully.")
-        return {}

src/rookify/modules/migrate_osds/main.py

@@ -7,7 +7,7 @@
 class MigrateOSDsHandler(ModuleHandler):
     REQUIRES = ["analyze_ceph"]
 
-    def run(self) -> Any:
+    def execute(self) -> Any:
         osd_config: Dict[str, Any] = {}
         state_data = self.machine.get_state("AnalyzeCephHandler").data
 

src/rookify/modules/module.py

@@ -205,9 +205,9 @@ def preflight(self) -> None:
         pass
 
     @abc.abstractmethod
-    def run(self) -> Dict[str, Any]:
+    def execute(self) -> None:
         """
-        Run the modules tasks
+        Executes the modules tasks
 
         :return: returns result
         """
@@ -220,31 +220,58 @@ def load_template(self, filename: str, **variables: Any) -> __Template:
         return template
 
     @classmethod
-    def register_state(
-        cls, machine: Machine, config: Dict[str, Any], **kwargs: Any
-    ) -> None:
+    def register_states(cls, machine: Machine, config: Dict[str, Any]) -> None:
         """
-        Register state for transitions
+        Register states for transitions
         """
 
         state_name = cls.STATE_NAME if hasattr(cls, "STATE_NAME") else cls.__name__
 
         handler = cls(machine, config)
+        preflight_state_name = None
+        execution_state_name = None
 
         if hasattr(cls, "preflight") and not getattr(
             cls.preflight, "__isabstractmethod__", False
         ):
-            kwargs["on_enter"] = handler.preflight
+            preflight_state_name = Machine.STATE_NAME_PREFLIGHT_PREFIX + state_name
 
-        if hasattr(cls, "run") and not getattr(cls.run, "__isabstractmethod__", False):
-            kwargs["on_exit"] = handler.run
+        if hasattr(cls, "execute") and not getattr(
+            cls.execute, "__isabstractmethod__", False
+        ):
+            execution_state_name = Machine.STATE_NAME_EXECUTION_PREFIX + state_name
 
-        if len(kwargs) > 0:
-            get_logger().debug("Registering state {0}".format(state_name))
-            machine.add_migrating_state(state_name, **kwargs)
-        else:
+        if preflight_state_name is None and execution_state_name is None:
             get_logger().warn(
                 "Not registering state {0} because ModuleHandler has no expected callables".format(
                     state_name
                 )
             )
+        else:
+            get_logger().debug("Registering states for {0}".format(state_name))
+
+            if preflight_state_name is not None:
+                cls.register_preflight_state(machine, preflight_state_name, handler)
+
+            if execution_state_name is not None:
+                cls.register_execution_state(machine, execution_state_name, handler)
+
+    @staticmethod
+    def register_preflight_state(
+        machine: Machine, state_name: str, handler: Any, **kwargs: Any
+    ) -> None:
+        """
+        Register state for transitions
+        """
+
+        machine.add_preflight_state(state_name, on_enter=handler.preflight, **kwargs)
+
+    @staticmethod
+    def register_execution_state(
+        machine: Machine, state_name: str, handler: Any, **kwargs: Any
+    ) -> None:
+        """
+        Register state for transitions
+        """
+
+        machine.add_execution_state(state_name, on_enter=handler.execute, **kwargs)

src/rookify/yaml.py

@@ -3,7 +3,6 @@
 import importlib.resources
 import importlib.resources.abc
 import yamale
-import yaml
 from pathlib import Path
 from typing import Any, Dict
 
@@ -24,15 +23,3 @@ def load_config(path: str) -> Dict[str, Any]:
 
     assert isinstance(data[0][0], dict)
     return data[0][0]
-
-
-def load_module_data(path: str) -> Dict[str, Any]:
-    with open(path, "r") as file:
-        data = yaml.safe_load(file)
-        assert isinstance(data, dict)
-        return data
-
-
-def save_module_data(path: str, data: Dict[str, Any]) -> None:
-    with open(path, "w") as file:
-        yaml.safe_dump(data, file)