From 07be5c79e197520ec3cc0f50d694e0b10f93f96d Mon Sep 17 00:00:00 2001 From: Tobias Wolf Date: Thu, 9 May 2024 17:15:01 +0200 Subject: [PATCH 1/7] Add "transitions" based execution machine This adapts the logic how modules are loaded and sorted for execution. Signed-off-by: Tobias Wolf --- requirements.txt | 1 + src/rookify/__main__.py | 62 ++------------ src/rookify/modules/__init__.py | 140 +++++++++----------------------- src/rookify/modules/machine.py | 28 +++++++ src/rookify/modules/module.py | 84 +++++++++++++------ 5 files changed, 132 insertions(+), 183 deletions(-) create mode 100644 src/rookify/modules/machine.py diff --git a/requirements.txt b/requirements.txt index 9246059..0464477 100644 --- a/requirements.txt +++ b/requirements.txt @@ -26,6 +26,7 @@ requests-oauthlib==1.3.1 rsa==4.9 six==1.16.0 structlog==24.1.0 +transitions==0.9.0 urllib3==2.2.1 yamale==5.1.0 websocket-client==1.7.0 diff --git a/src/rookify/__main__.py b/src/rookify/__main__.py index 89c8853..6474da3 100644 --- a/src/rookify/__main__.py +++ b/src/rookify/__main__.py @@ -1,10 +1,9 @@ # -*- coding: utf-8 -*- -import os -import rookify.modules -from types import MappingProxyType +from .modules import load_modules +from .modules.machine import Machine from .logger import configure_logging, get_logger -from .yaml import load_config, load_module_data, save_module_data +from .yaml import load_config def main() -> None: @@ -23,56 +22,7 @@ def main() -> None: log = get_logger() log.debug("Executing Rookify") - preflight_modules, migration_modules = rookify.modules.load_modules( - config["migration_modules"] - ) + machine = Machine() + load_modules(machine, config) - module_data = dict() - - try: - module_data.update(load_module_data(config["general"]["module_data_file"])) - except FileNotFoundError: - pass - - # Get absolute path of the rookify instance - rookify_path = os.path.dirname(__file__) - - # Run preflight requirement modules - for preflight_module in preflight_modules: - module_path = os.path.join( - rookify_path, "modules", preflight_module.MODULE_NAME - ) - handler = preflight_module.HANDLER_CLASS( - config=MappingProxyType(config), - data=MappingProxyType(module_data), - module_path=module_path, - ) - result = handler.run() - module_data[preflight_module.MODULE_NAME] = result - - # Run preflight and append handlers to list - handlers = list() - for migration_module in migration_modules: - module_path = os.path.join( - rookify_path, "modules", migration_module.MODULE_NAME - ) - handler = migration_module.HANDLER_CLASS( - config=MappingProxyType(config), - data=MappingProxyType(module_data), - module_path=module_path, - ) - handler.preflight() - handlers.append((migration_module, handler)) - - # Run migration modules - for migration_module, handler in handlers: - result = handler.run() - module_data[migration_module.MODULE_NAME] = result - - save_module_data(config["general"]["module_data_file"], module_data) - - log.info("Data was updated to module_data_file.") - - -if __name__ == "__main__": - main() + machine.execute() diff --git a/src/rookify/modules/__init__.py b/src/rookify/modules/__init__.py index 3167e96..40993f7 100644 --- a/src/rookify/modules/__init__.py +++ b/src/rookify/modules/__init__.py @@ -1,11 +1,8 @@ # -*- coding: utf-8 -*- import importlib -import types -from typing import List - -from collections import OrderedDict -from .module import ModuleHandler +from typing import Any, Dict +from .machine import Machine class ModuleLoadException(Exception): @@ -24,9 +21,39 @@ def __init__(self, module_name: str, message: str): self.message = message -def load_modules( - module_names: List[str], -) -> tuple[List[types.ModuleType], List[types.ModuleType]]: +def _load_module(machine: Machine, config: Dict[str, Any], module_name: str) -> None: + """ + Dynamically loads a module from the 'rookify.modules' package. + + :param module_names: The module names to load + :return: returns tuple of preflight_modules, modules + """ + + module = importlib.import_module("rookify.modules.{0}".format(module_name)) + additional_modules = [] + + if not hasattr(module, "ModuleHandler") or not callable( + getattr(module.ModuleHandler, "register_state") + ): + raise ModuleLoadException(module_name, "Module structure is invalid") + + if hasattr(module.ModuleHandler, "PREFLIGHT_REQUIRES"): + assert isinstance(module.ModuleHandler.PREFLIGHT_REQUIRES, list) + additional_modules += module.ModuleHandler.PREFLIGHT_REQUIRES + + if hasattr(module.ModuleHandler, "REQUIRES"): + assert isinstance(module.ModuleHandler.REQUIRES, list) + for module_name in module.ModuleHandler.REQUIRES: + if module_name not in additional_modules: + additional_modules.append(module_name) + + for module_name in additional_modules: + _load_module(machine, config, module_name) + + module.ModuleHandler.register_state(machine, config) + + +def load_modules(machine: Machine, config: Dict[str, Any]) -> None: """ Dynamically loads modules from the 'modules' package. @@ -34,97 +61,6 @@ def load_modules( :return: returns tuple of preflight_modules, modules """ - # Sanity checks for modules - def check_module_sanity(module_name: str, module: types.ModuleType) -> None: - for attr_type, attr_name in ( - (ModuleHandler, "HANDLER_CLASS"), - (str, "MODULE_NAME"), - (list, "REQUIRES"), - (list, "AFTER"), - (list, "PREFLIGHT_REQUIRES"), - ): - if not hasattr(module, attr_name): - raise ModuleLoadException( - module_name, f"Module has no attribute {attr_name}" - ) - - attr = getattr(module, attr_name) - if not isinstance(attr, attr_type) and not issubclass(attr, attr_type): - raise ModuleLoadException( - module_name, f"Attribute {attr_name} is not type {attr_type}" - ) - - # Load the modules in the given list and recursivley load required modules - required_modules: OrderedDict[str, types.ModuleType] = OrderedDict() - - def load_required_modules( - modules_out: OrderedDict[str, types.ModuleType], module_names: List[str] - ) -> None: - for module_name in module_names: - if module_name in modules_out: - continue - - module = importlib.import_module(f".{module_name}", "rookify.modules") - check_module_sanity(module_name, module) - - load_required_modules(modules_out, module.REQUIRES) - module.AFTER.extend(module.REQUIRES) - - modules_out[module_name] = module - - load_required_modules(required_modules, module_names) - - # Recursively load the modules in the PREFLIGHT_REQUIRES attribute of the given modules - preflight_modules: OrderedDict[str, types.ModuleType] = OrderedDict() - - def load_preflight_modules( - modules_in: OrderedDict[str, types.ModuleType], - modules_out: OrderedDict[str, types.ModuleType], - module_names: List[str], - ) -> None: - for module_name in module_names: - if module_name in modules_out: - continue - - module = importlib.import_module(f".{module_name}", "rookify.modules") - check_module_sanity(module_name, module) - - # We have to check, if the preflight_requires list is already loaded as migration requirement - for preflight_requirement in module.PREFLIGHT_REQUIRES: - if preflight_requirement in modules_in: - raise ModuleLoadException( - module_name, - f"Module {preflight_requirement} is already loaded as migration requirement", - ) - - load_preflight_modules(modules_in, modules_out, module.PREFLIGHT_REQUIRES) - if module_name not in modules_in: - modules_out[module_name] = module - - load_preflight_modules( - required_modules, preflight_modules, list(required_modules.keys()) - ) - - # Sort the modules by the AFTER keyword - modules: OrderedDict[str, types.ModuleType] = OrderedDict() - - def sort_modules( - modules_in: OrderedDict[str, types.ModuleType], - modules_out: OrderedDict[str, types.ModuleType], - module_names: List[str], - ) -> None: - for module_name in module_names: - if module_name not in modules_in: - continue - - if module_name in modules_out: - continue - - after_modules_name = modules_in[module_name].AFTER - sort_modules(modules_in, modules_out, after_modules_name) - - modules_out[module_name] = modules_in[module_name] - - sort_modules(required_modules, modules, list(required_modules.keys())) - - return list(preflight_modules.values()), list(modules.values()) + for entry in importlib.resources.files("rookify.modules").iterdir(): + if entry.is_dir() and entry.name in config["migration_modules"]: + _load_module(machine, config, entry.name) diff --git a/src/rookify/modules/machine.py b/src/rookify/modules/machine.py new file mode 100644 index 0000000..5b7340c --- /dev/null +++ b/src/rookify/modules/machine.py @@ -0,0 +1,28 @@ +# -*- coding: utf-8 -*- + +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 + + +@add_state_features(Tags, Timeout) +class Machine(_Machine): # type: ignore + def __init__(self) -> None: + _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 execute(self) -> None: + self.add_state("migrated") + self.add_ordered_transitions(loop=False) + + try: + while True: + self.next_state() + except MachineError: + if self.state != "migrated": + raise diff --git a/src/rookify/modules/module.py b/src/rookify/modules/module.py index 8642717..5b4c514 100644 --- a/src/rookify/modules/module.py +++ b/src/rookify/modules/module.py @@ -9,9 +9,11 @@ import fabric import jinja2 import structlog -from rookify.logger import get_logger from typing import Any, Dict, List, Optional +from ..logger import get_logger +from .machine import Machine + class ModuleException(Exception): pass @@ -152,7 +154,7 @@ def yaml(self) -> Any: self.__result_yaml = yaml.safe_load(self.raw) return self.__result_yaml - def __init__(self, config: Dict[str, Any], data: Dict[str, Any], module_path: str): + def __init__(self, machine: Machine, config: Dict[str, Any]): """ Construct a new 'ModuleHandler' object. @@ -162,31 +164,13 @@ def __init__(self, config: Dict[str, Any], data: Dict[str, Any], module_path: st :return: returns nothing """ self._config = config - self._data = data - self.__module_path = module_path + self._machine = machine + self.__ceph: Optional[ModuleHandler.__Ceph] = None self.__k8s: Optional[ModuleHandler.__K8s] = None self.__ssh: Optional[ModuleHandler.__SSH] = None self.__logger = get_logger() - self.__logger.debug("Executing {0}".format(self.__class__.__name__)) - - @abc.abstractmethod - def preflight(self) -> None: - """ - Run the modules preflight check - """ - pass - - @abc.abstractmethod - def run(self) -> Dict[str, Any]: - """ - Run the modules tasks - - :return: returns result - """ - pass - @property def ceph(self) -> __Ceph: if self.__ceph is None: @@ -194,8 +178,8 @@ def ceph(self) -> __Ceph: return self.__ceph @property - def logger(self) -> structlog.getLogger: - return self.__logger + def machine(self) -> Machine: + return self._machine @property def k8s(self) -> __K8s: @@ -203,14 +187,64 @@ def k8s(self) -> __K8s: self.__k8s = ModuleHandler.__K8s(self._config["kubernetes"]) return self.__k8s + @property + def logger(self) -> structlog.getLogger: + return self.__logger + @property def ssh(self) -> __SSH: if self.__ssh is None: self.__ssh = ModuleHandler.__SSH(self._config["ssh"]) return self.__ssh + @abc.abstractmethod + def preflight(self) -> None: + """ + Run the modules preflight check + """ + pass + + @abc.abstractmethod + def run(self) -> Dict[str, Any]: + """ + Run the modules tasks + + :return: returns result + """ + pass + def load_template(self, filename: str, **variables: Any) -> __Template: - template_path = os.path.join(self.__module_path, "templates", filename) + template_path = os.path.join(os.path.dirname(__file__), "templates", filename) template = ModuleHandler.__Template(template_path) template.render(**variables) return template + + @classmethod + def register_state( + cls, machine: Machine, config: Dict[str, Any], **kwargs: Any + ) -> None: + """ + Register state for transitions + """ + + state_name = cls.STATE_NAME if hasattr(cls, "STATE_NAME") else cls.__name__ + + handler = cls(machine, config) + + if hasattr(cls, "preflight") and not getattr( + cls.preflight, "__isabstractmethod__", False + ): + kwargs["on_enter"] = handler.preflight + + if hasattr(cls, "run") and not getattr(cls.run, "__isabstractmethod__", False): + kwargs["on_exit"] = handler.run + + if len(kwargs) > 0: + get_logger().debug("Registering state {0}".format(state_name)) + machine.add_migrating_state(state_name, **kwargs) + else: + get_logger().warn( + "Not registering state {0} because ModuleHandler has no expected callables".format( + state_name + ) + ) From 61013ca5619aab82450af3d358d7c99ef0a9b9ad Mon Sep 17 00:00:00 2001 From: Tobias Wolf Date: Thu, 9 May 2024 17:30:00 +0200 Subject: [PATCH 2/7] Adapt existing modules for "transitions" Signed-off-by: Tobias Wolf --- src/rookify/modules/analyze_ceph/__init__.py | 9 +--- src/rookify/modules/analyze_ceph/main.py | 35 +++++++++++----- .../modules/cephx_auth_config/__init__.py | 9 +--- src/rookify/modules/cephx_auth_config/main.py | 41 +++++++++++++++---- src/rookify/modules/example/__init__.py | 13 +----- src/rookify/modules/example/main.py | 6 ++- .../modules/migrate_monitors/__init__.py | 9 +--- src/rookify/modules/migrate_monitors/main.py | 4 +- src/rookify/modules/migrate_osds/__init__.py | 9 +--- src/rookify/modules/migrate_osds/main.py | 18 ++++---- tests/modules/test_example.py | 3 +- 11 files changed, 79 insertions(+), 77 deletions(-) diff --git a/src/rookify/modules/analyze_ceph/__init__.py b/src/rookify/modules/analyze_ceph/__init__.py index d821800..58914f8 100644 --- a/src/rookify/modules/analyze_ceph/__init__.py +++ b/src/rookify/modules/analyze_ceph/__init__.py @@ -1,10 +1,3 @@ # -*- coding: utf-8 -*- -# type: ignore -from .main import AnalyzeCephHandler - -MODULE_NAME = "analyze_ceph" -HANDLER_CLASS = AnalyzeCephHandler -REQUIRES = [] -AFTER = [] -PREFLIGHT_REQUIRES = [] +from .main import AnalyzeCephHandler as ModuleHandler # noqa diff --git a/src/rookify/modules/analyze_ceph/main.py b/src/rookify/modules/analyze_ceph/main.py index 594b5f8..04ff418 100644 --- a/src/rookify/modules/analyze_ceph/main.py +++ b/src/rookify/modules/analyze_ceph/main.py @@ -1,29 +1,44 @@ # -*- coding: utf-8 -*- -from ..module import ModuleHandler from typing import Any, Dict +from ..machine import Machine +from ..module import ModuleHandler class AnalyzeCephHandler(ModuleHandler): def run(self) -> Any: commands = ["mon dump", "osd dump", "device ls", "fs dump", "node ls"] - results: Dict[str, Any] = dict() + state = self.machine.get_state("AnalyzeCephHandler") + state.data: Dict[str, Any] = {} # type: ignore + for command in commands: parts = command.split(" ") - leaf = results + leaf = state.data for idx, part in enumerate(parts): if idx < len(parts) - 1: - leaf[part] = dict() + leaf[part] = {} else: leaf[part] = self.ceph.mon_command(command) leaf = leaf[part] - self.logger.info("Dictionary created") - results["ssh"] = dict() - results["ssh"]["osd"] = dict() - for node, values in results["node"]["ls"]["osd"].items(): + self.logger.debug("AnalyzeCephHandler commands executed") + + state.data["ssh"] = {} + state.data["ssh"]["osd"] = {} + + for node, values in state.data["node"]["ls"]["osd"].items(): devices = self.ssh.command(node, "find /dev/ceph-*/*").stdout.splitlines() - results["ssh"]["osd"][node] = {"devices": devices} + state.data["ssh"]["osd"][node] = {"devices": devices} + self.logger.info("AnalyzeCephHandler ran successfully.") - return results + + @classmethod + def register_state( + _, machine: Machine, config: Dict[str, Any], **kwargs: Any + ) -> None: + """ + Register state for transitions + """ + + super().register_state(machine, config, tags=["data"]) diff --git a/src/rookify/modules/cephx_auth_config/__init__.py b/src/rookify/modules/cephx_auth_config/__init__.py index 6070741..1fdf5ee 100644 --- a/src/rookify/modules/cephx_auth_config/__init__.py +++ b/src/rookify/modules/cephx_auth_config/__init__.py @@ -1,10 +1,3 @@ # -*- coding: utf-8 -*- -# type: ignore -from .main import CephXAuthHandler - -MODULE_NAME = "cephx_auth_config" -HANDLER_CLASS = CephXAuthHandler -REQUIRES = [] -AFTER = [] -PREFLIGHT_REQUIRES = [] +from .main import CephXAuthHandler as ModuleHandler # noqa diff --git a/src/rookify/modules/cephx_auth_config/main.py b/src/rookify/modules/cephx_auth_config/main.py index 45248d6..7780931 100644 --- a/src/rookify/modules/cephx_auth_config/main.py +++ b/src/rookify/modules/cephx_auth_config/main.py @@ -1,16 +1,39 @@ # -*- coding: utf-8 -*- -from ..module import ModuleHandler -from typing import Any +from typing import Any, Dict +from ..machine import Machine +from ..module import ModuleException, ModuleHandler class CephXAuthHandler(ModuleHandler): - def run(self) -> Any: - self.logger.debug("Reconfiguring Ceph to expect cephx auth") + def preflight(self) -> Any: + if not self.is_cephx_set(self.ceph.conf_get("auth_cluster_required")): + raise ModuleException( + "Ceph config value auth_cluster_required does not contain cephx" + ) - self.ceph.conf_set("auth_cluster_required", "cephx") - self.ceph.conf_set("auth_service_required", "cephx") - self.ceph.conf_set("auth_client_required", "cephx") + if not self.is_cephx_set(self.ceph.conf_get("auth_service_required")): + raise ModuleException( + "Ceph config value auth_service_required does not contain cephx" + ) - self.logger.info("Reconfigured Ceph to expect cephx auth") - return {"reconfigured": True} + if not self.is_cephx_set(self.ceph.conf_get("auth_client_required")): + raise ModuleException( + "Ceph config value auth_client_required does not contain cephx" + ) + + self.machine.get_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 + ) -> None: + """ + Register state for transitions + """ + + super().register_state(machine, config, tags=["verified"]) diff --git a/src/rookify/modules/example/__init__.py b/src/rookify/modules/example/__init__.py index ff0dcf5..5d3c104 100644 --- a/src/rookify/modules/example/__init__.py +++ b/src/rookify/modules/example/__init__.py @@ -1,14 +1,3 @@ # -*- coding: utf-8 -*- -# type: ignore -from .main import ExampleHandler - -MODULE_NAME = "example" # Name of the module -HANDLER_CLASS = ExampleHandler # Define the handler class for this module -REQUIRES = [] # A list of modules that are required to run before this module. Modules in this list will be imported, even if they are not configured -AFTER = [ - "migrate_monitors" -] # A list of modules that should be run before this module, if they are defined in config -PREFLIGHT_REQUIRES = [ - "analyze_ceph" -] # A list of modules that are required to run the preflight_check of this module. Modules in this list will be imported and run in preflight stage. +from .main import ExampleHandler as ModuleHandler # noqa diff --git a/src/rookify/modules/example/main.py b/src/rookify/modules/example/main.py index e62394c..bd0fde7 100644 --- a/src/rookify/modules/example/main.py +++ b/src/rookify/modules/example/main.py @@ -1,11 +1,13 @@ # -*- coding: utf-8 -*- -from ..module import ModuleHandler, ModuleException - from typing import Any +from ..module import ModuleHandler, ModuleException class ExampleHandler(ModuleHandler): + # A list of modules that are required to run the preflight_check of this module. Modules in this list will be imported and run in preflight stage. + PREFLIGHT_REQUIRES = ["analyze_ceph"] + 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!") diff --git a/src/rookify/modules/migrate_monitors/__init__.py b/src/rookify/modules/migrate_monitors/__init__.py index ff3c337..65c6808 100644 --- a/src/rookify/modules/migrate_monitors/__init__.py +++ b/src/rookify/modules/migrate_monitors/__init__.py @@ -1,10 +1,3 @@ # -*- coding: utf-8 -*- -# type: ignore -from .main import MigrateMonitorsHandler - -MODULE_NAME = "migrate_monitors" -HANDLER_CLASS = MigrateMonitorsHandler -REQUIRES = [] -AFTER = [] -PREFLIGHT_REQUIRES = ["analyze_ceph"] +from .main import MigrateMonitorsHandler as ModuleHandler # noqa diff --git a/src/rookify/modules/migrate_monitors/main.py b/src/rookify/modules/migrate_monitors/main.py index f057695..d4e75f9 100644 --- a/src/rookify/modules/migrate_monitors/main.py +++ b/src/rookify/modules/migrate_monitors/main.py @@ -1,10 +1,12 @@ # -*- coding: utf-8 -*- -from ..module import ModuleHandler from typing import Dict, Any +from ..module import ModuleHandler class MigrateMonitorsHandler(ModuleHandler): + REQUIRES = ["analyze_ceph"] + def run(self) -> Dict[str, Any]: self.logger.info("MigrateMonitorsHandler ran successfully.") return {} diff --git a/src/rookify/modules/migrate_osds/__init__.py b/src/rookify/modules/migrate_osds/__init__.py index 86041ae..8aa151c 100644 --- a/src/rookify/modules/migrate_osds/__init__.py +++ b/src/rookify/modules/migrate_osds/__init__.py @@ -1,10 +1,3 @@ # -*- coding: utf-8 -*- -# type: ignore -from .main import MigrateOSDsHandler - -MODULE_NAME = "migrate_osds" -HANDLER_CLASS = MigrateOSDsHandler -REQUIRES = [] -AFTER = ["migrate_monitors"] -PREFLIGHT_REQUIRES = ["analyze_ceph"] +from .main import MigrateOSDsHandler as ModuleHandler # noqa diff --git a/src/rookify/modules/migrate_osds/main.py b/src/rookify/modules/migrate_osds/main.py index 8f89c42..b9df296 100644 --- a/src/rookify/modules/migrate_osds/main.py +++ b/src/rookify/modules/migrate_osds/main.py @@ -1,23 +1,22 @@ # -*- coding: utf-8 -*- -from ..module import ModuleHandler from typing import Any, Dict +from ..module import ModuleHandler class MigrateOSDsHandler(ModuleHandler): - def preflight(self) -> None: - pass - # result = self.ceph.mon_command("osd dump") - # raise ModuleException('test error') + REQUIRES = ["analyze_ceph"] def run(self) -> Any: - osd_config: Dict[str, Any] = dict() - for node, osds in self._data["analyze_ceph"]["node"]["ls"]["osd"].items(): + osd_config: Dict[str, Any] = {} + state_data = self.machine.get_state("AnalyzeCephHandler").data + + for node, osds in state_data["node"]["ls"]["osd"].items(): osd_config[node] = {"osds": {}} for osd in osds: osd_config[node]["osds"][osd] = dict() - for osd in self._data["analyze_ceph"]["osd"]["dump"]["osds"]: + for osd in state_data["osd"]["dump"]["osds"]: number = osd["osd"] uuid = osd["uuid"] for host in osd_config.values(): @@ -26,7 +25,7 @@ def run(self) -> Any: break for node, values in osd_config.items(): - devices = self._data["analyze_ceph"]["ssh"]["osd"][node]["devices"] + devices = state_data["ssh"]["osd"][node]["devices"] for osd in values["osds"].values(): for device in devices: if osd["uuid"] in device: @@ -34,4 +33,3 @@ def run(self) -> Any: break self.logger.info(osd_config) - return {} diff --git a/tests/modules/test_example.py b/tests/modules/test_example.py index 1ff294a..3615560 100644 --- a/tests/modules/test_example.py +++ b/tests/modules/test_example.py @@ -3,9 +3,10 @@ import pytest from rookify.modules.example.main import ExampleHandler +from rookify.modules.machine import Machine from rookify.modules.module import ModuleException def test_preflight() -> None: with pytest.raises(ModuleException): - ExampleHandler({}, {}, "").preflight() + ExampleHandler(Machine(), {}).preflight() From 32352710a93d32c76d421a8f28c33fb370762eb1 Mon Sep 17 00:00:00 2001 From: Tobias Wolf Date: Wed, 15 May 2024 11:32:05 +0200 Subject: [PATCH 3/7] 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 --- .gitignore | 2 +- requirements.txt | 1 + src/config.example.yaml | 2 +- src/rookify/__main__.py | 2 +- src/rookify/config.schema.yaml | 2 +- src/rookify/modules/__init__.py | 4 +- src/rookify/modules/analyze_ceph/main.py | 18 ++-- src/rookify/modules/cephx_auth_config/main.py | 18 ++-- src/rookify/modules/example/main.py | 2 +- src/rookify/modules/machine.py | 92 +++++++++++++++++-- src/rookify/modules/migrate_monitors/main.py | 4 +- src/rookify/modules/migrate_osds/main.py | 2 +- src/rookify/modules/module.py | 53 ++++++++--- src/rookify/yaml.py | 13 --- 14 files changed, 152 insertions(+), 63 deletions(-) diff --git a/.gitignore b/.gitignore index 8174be4..7863217 100644 --- a/.gitignore +++ b/.gitignore @@ -86,7 +86,7 @@ dmypy.json cython_debug/ # Project specific -data.yaml +data.pickle config.yaml .ceph .k8s diff --git a/requirements.txt b/requirements.txt index 0464477..1c8606a 100644 --- a/requirements.txt +++ b/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 diff --git a/src/config.example.yaml b/src/config.example.yaml index 771b2c8..57d2f16 100644 --- a/src/config.example.yaml +++ b/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 diff --git a/src/rookify/__main__.py b/src/rookify/__main__.py index 6474da3..fd9fad0 100644 --- a/src/rookify/__main__.py +++ b/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() diff --git a/src/rookify/config.schema.yaml b/src/rookify/config.schema.yaml index 3504eab..baba81c 100644 --- a/src/rookify/config.schema.yaml +++ b/src/rookify/config.schema.yaml @@ -1,5 +1,5 @@ general: - module_data_file: str() + machine_pickle_file: str(required=False) logging: level: str() diff --git a/src/rookify/modules/__init__.py b/src/rookify/modules/__init__.py index 40993f7..9c2e697 100644 --- a/src/rookify/modules/__init__.py +++ b/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: diff --git a/src/rookify/modules/analyze_ceph/main.py b/src/rookify/modules/analyze_ceph/main.py index 04ff418..273e90f 100644 --- a/src/rookify/modules/analyze_ceph/main.py +++ b/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"] + ) diff --git a/src/rookify/modules/cephx_auth_config/main.py b/src/rookify/modules/cephx_auth_config/main.py index 7780931..c7d8c58 100644 --- a/src/rookify/modules/cephx_auth_config/main.py +++ b/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"] + ) diff --git a/src/rookify/modules/example/main.py b/src/rookify/modules/example/main.py index bd0fde7..790605c 100644 --- a/src/rookify/modules/example/main.py +++ b/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 {} diff --git a/src/rookify/modules/machine.py b/src/rookify/modules/machine.py index 5b7340c..e0ec915 100644 --- a/src/rookify/modules/machine.py +++ b/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 + ) + ) diff --git a/src/rookify/modules/migrate_monitors/main.py b/src/rookify/modules/migrate_monitors/main.py index d4e75f9..ac54483 100644 --- a/src/rookify/modules/migrate_monitors/main.py +++ b/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 {} diff --git a/src/rookify/modules/migrate_osds/main.py b/src/rookify/modules/migrate_osds/main.py index b9df296..cb38a79 100644 --- a/src/rookify/modules/migrate_osds/main.py +++ b/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 diff --git a/src/rookify/modules/module.py b/src/rookify/modules/module.py index 5b4c514..d0466e8 100644 --- a/src/rookify/modules/module.py +++ b/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) diff --git a/src/rookify/yaml.py b/src/rookify/yaml.py index 6712ae4..b0c178b 100644 --- a/src/rookify/yaml.py +++ b/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) From 893beed2dbc04c17aee89e25e53a8c161724dada Mon Sep 17 00:00:00 2001 From: Tobias Wolf Date: Wed, 15 May 2024 11:42:47 +0200 Subject: [PATCH 4/7] Remove `PREFLIGHT_REQUIRES` logic Signed-off-by: Tobias Wolf --- src/rookify/modules/__init__.py | 8 +------- src/rookify/modules/example/main.py | 2 +- 2 files changed, 2 insertions(+), 8 deletions(-) diff --git a/src/rookify/modules/__init__.py b/src/rookify/modules/__init__.py index 9c2e697..fa5529b 100644 --- a/src/rookify/modules/__init__.py +++ b/src/rookify/modules/__init__.py @@ -37,15 +37,9 @@ def _load_module(machine: Machine, config: Dict[str, Any], module_name: str) -> ): raise ModuleLoadException(module_name, "Module structure is invalid") - if hasattr(module.ModuleHandler, "PREFLIGHT_REQUIRES"): - assert isinstance(module.ModuleHandler.PREFLIGHT_REQUIRES, list) - additional_modules += module.ModuleHandler.PREFLIGHT_REQUIRES - if hasattr(module.ModuleHandler, "REQUIRES"): assert isinstance(module.ModuleHandler.REQUIRES, list) - for module_name in module.ModuleHandler.REQUIRES: - if module_name not in additional_modules: - additional_modules.append(module_name) + additional_modules = module.ModuleHandler.REQUIRES for module_name in additional_modules: _load_module(machine, config, module_name) diff --git a/src/rookify/modules/example/main.py b/src/rookify/modules/example/main.py index 790605c..cddba34 100644 --- a/src/rookify/modules/example/main.py +++ b/src/rookify/modules/example/main.py @@ -6,7 +6,7 @@ class ExampleHandler(ModuleHandler): # A list of modules that are required to run the preflight_check of this module. Modules in this list will be imported and run in preflight stage. - PREFLIGHT_REQUIRES = ["analyze_ceph"] + REQUIRES = ["analyze_ceph"] def preflight(self) -> None: # Do something for checking if all needed preconditions are met else throw ModuleException From feec06163921bc68b0c4e74b944a429de260a6d1 Mon Sep 17 00:00:00 2001 From: Tobias Wolf Date: Wed, 15 May 2024 11:51:04 +0200 Subject: [PATCH 5/7] Fix execution if no `machine_pickle_file` is configured Signed-off-by: Tobias Wolf --- src/rookify/__main__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/rookify/__main__.py b/src/rookify/__main__.py index fd9fad0..9c2e3c4 100644 --- a/src/rookify/__main__.py +++ b/src/rookify/__main__.py @@ -22,7 +22,7 @@ def main() -> None: log = get_logger() log.debug("Executing Rookify") - machine = Machine(config["general"]["machine_pickle_file"]) + machine = Machine(config["general"].get("machine_pickle_file")) load_modules(machine, config) machine.execute() From 1d9d7771fd5daf186f9eecca7e66d1ebec77dd4d Mon Sep 17 00:00:00 2001 From: Tobias Wolf Date: Wed, 22 May 2024 09:19:14 +0200 Subject: [PATCH 6/7] Fix access to preflight state in `migrate_osds` Signed-off-by: Tobias Wolf --- src/rookify/modules/migrate_osds/main.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/rookify/modules/migrate_osds/main.py b/src/rookify/modules/migrate_osds/main.py index cb38a79..170ddbe 100644 --- a/src/rookify/modules/migrate_osds/main.py +++ b/src/rookify/modules/migrate_osds/main.py @@ -9,7 +9,7 @@ class MigrateOSDsHandler(ModuleHandler): def execute(self) -> Any: osd_config: Dict[str, Any] = {} - state_data = self.machine.get_state("AnalyzeCephHandler").data + state_data = self.machine.get_preflight_state("AnalyzeCephHandler").data for node, osds in state_data["node"]["ls"]["osd"].items(): osd_config[node] = {"osds": {}} From af699c71863d9a8163d63d6851238911774913ba Mon Sep 17 00:00:00 2001 From: Tobias Wolf Date: Wed, 29 May 2024 11:28:10 +0200 Subject: [PATCH 7/7] Fix docstring for class `ModuleHandler` Signed-off-by: Tobias Wolf --- src/rookify/modules/module.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/src/rookify/modules/module.py b/src/rookify/modules/module.py index d0466e8..1de3997 100644 --- a/src/rookify/modules/module.py +++ b/src/rookify/modules/module.py @@ -158,9 +158,8 @@ def __init__(self, machine: Machine, config: Dict[str, Any]): """ Construct a new 'ModuleHandler' object. + :param machine: Machine object :param config: The global config file - :param data: The output of modules required by this module - :param module_path: The filesystem path of this module :return: returns nothing """ self._config = config