From e64c26ea3d9c82c79db43a77f094b38943b6eb3d Mon Sep 17 00:00:00 2001 From: Miruna Serian Date: Thu, 4 Jun 2026 09:02:08 +0100 Subject: [PATCH 01/20] Make mypy compliant - no actions subdir yet --- pyproject.toml | 5 +- src/drunc/fsm/action_factory.py | 41 ++++++++++++----- src/drunc/fsm/configuration.py | 75 +++++++++++++++++------------- src/drunc/fsm/core.py | 82 ++++++++++++++++++++++----------- src/drunc/fsm/exceptions.py | 48 +++++++++---------- src/drunc/fsm/transition.py | 21 +++++++-- src/drunc/fsm/utils.py | 41 +++++++++++++---- 7 files changed, 201 insertions(+), 112 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index be11c1c90..ff6a6af8b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -109,9 +109,6 @@ show_error_codes = true # These overrides are because the library stubs dont exist and not on typeshed [[tool.mypy.overrides]] -module = ["google.rpc.*"] +module = "google.rpc.*,conffwk.*" ignore_missing_imports = true -[[tool.mypy.overrides]] -module = ["conffwk"] -ignore_missing_imports = true \ No newline at end of file diff --git a/src/drunc/fsm/action_factory.py b/src/drunc/fsm/action_factory.py index 5d23a1218..a14f5ce63 100644 --- a/src/drunc/fsm/action_factory.py +++ b/src/drunc/fsm/action_factory.py @@ -1,4 +1,7 @@ +from __future__ import annotations + import inspect +from typing import TYPE_CHECKING, Dict, Protocol, Type import conffwk @@ -15,28 +18,43 @@ from drunc.fsm.actions.usvc_elisa_logbook import ElisaLogbook from drunc.fsm.actions.usvc_provided_run_number import UsvcProvidedRunNumber +if TYPE_CHECKING: + # for mypy treat FSMaction_t as a local objext + FSMaction_t = object +else: + # at runtime use the actual type from conffwk.dal + FSMaction_t = conffwk.dal.FSMaction + +class ActionMethod(Protocol): + def __call__(self, *args: object, **kwargs: object) -> object: ... + +class FSMActionProtocol(Protocol): + name: str + class FSMActionFactory: - def __init__(self): + _instance: FSMActionFactory | None = None + + def __init__(self) -> None: raise DruncSetupException("Call get() instead") - def _get_pre_transitions(self, action): - retr = {} + def _get_pre_transitions(self, action: FSMActionProtocol) -> Dict[str, ActionMethod]: + retr: Dict[str, ActionMethod] = {} for name, method in inspect.getmembers(action): if inspect.ismethod(method): if name.startswith("pre_"): retr[name] = method return retr - def _get_post_transitions(self, action): - retr = {} + def _get_post_transitions(self, action: FSMActionProtocol) -> Dict[str, ActionMethod]: + retr: Dict[str, ActionMethod] = {} for name, method in inspect.getmembers(action): if inspect.ismethod(method): if name.startswith("post_"): retr[name] = method return retr - def _validate_signature(self, name, method, action): + def _validate_signature(self, name:str , method: ActionMethod, action: str) -> None: sig = inspect.signature(method) if ( @@ -53,7 +71,7 @@ def _validate_signature(self, name, method, action): if p.annotation is inspect._empty: raise fsme.MethodSignatureMissingAnnotation(action, name, pname) - def _validate_action(self, action): + def _validate_action(self, action: FSMActionProtocol) -> None: pre_transition = self._get_pre_transitions(action) post_transition = self._get_post_transitions(action) @@ -67,8 +85,8 @@ def _validate_action(self, action): self._validate_signature(k, v, action.name) def get_action( - self, action_name: str, action_configuration: "conffwk.dal.FSMaction" - ): + self, action_name: str, action_configuration: FSMaction_t + ) -> FSMActionProtocol: """ Construct the action interface for the given action name and configuration. @@ -84,7 +102,7 @@ def get_action( fsme.InvalidAction: If the constructed action does not have valid pre/post transition methods. """ - iface = None + iface: FSMActionProtocol | None = None match action_name: case "user-provided-run-number": iface = UserProvidedRunNumber(action_configuration) @@ -116,10 +134,9 @@ def get_action( return iface - _instance = None @classmethod - def get(cls): + def get(cls: Type[FSMActionFactory]) -> FSMActionFactory: if cls._instance is None: cls._instance = cls.__new__(cls) diff --git a/src/drunc/fsm/configuration.py b/src/drunc/fsm/configuration.py index 50760ae55..cbfec6a20 100644 --- a/src/drunc/fsm/configuration.py +++ b/src/drunc/fsm/configuration.py @@ -1,19 +1,41 @@ -import conffwk +import logging +from typing import TYPE_CHECKING, Dict, List from druncschema.controller_pb2 import FSMSequence -from drunc.fsm.action_factory import FSMActionFactory +from drunc.fsm.action_factory import FSMActionFactory, FSMActionProtocol from drunc.fsm.core import PreOrPostTransitionSequence from drunc.fsm.transition import Transition -from drunc.utils.configuration import ConfHandler +from drunc.utils.configuration import ConfHandler, OKSKey from drunc.utils.utils import get_logger +if TYPE_CHECKING: + FSMaction_t = object + FSMxTransition_t = object + FSMSequence_t = object +else: + import conffwk + FSMaction_t = conffwk.dal.FSMAction + FSMxTransition_t = conffwk.dal.FSMxTransition + FSMSequence_t = conffwk.dal.FSMSequence class FSMConfHandler(ConfHandler): + + def __init__(self, key: OKSKey) -> None: + super().__init__(key) + self.log: logging.Logger = get_logger("controller.core.FSMConfHandler") + self.pre_transitions: Dict[Transition, PreOrPostTransitionSequence] = {} + self.post_transitions: Dict[Transition, PreOrPostTransitionSequence] = {} + self.actions: Dict[str, FSMActionProtocol] = {} # (e.g. "thread_pinning": thread_pinning FSM action object) + self.transitions: List[Transition] = [] + self.sequences: List[FSMSequence] = [] + self.states: List[str] = [] + self.initial_state: str = "" + def _fill_pre_post_transition_sequence_oks( self, prefix: str, transition: Transition, - data: list("conffwk.dal.FSMxTransition"), + data: List[FSMxTransition_t] | None, ) -> PreOrPostTransitionSequence: """ Fill the pre or post transition sequence for a given transition. @@ -42,16 +64,17 @@ def _fill_pre_post_transition_sequence_oks( # transition. There is one FSMxTransition per transition. The FSMxTransition # contains the list of actions to execute for the given transition. for fsm_x_transition in data: - if fsm_x_transition.transition == transition.name: - for action_name in fsm_x_transition.order: + if getattr(fsm_x_transition, "transition", None) == transition.name: + for action_name in getattr(fsm_x_transition, "order", []): + mandatory_list = getattr(fsm_x_transition, "mandatory", []) seq.add_callback( action=self.actions[action_name], - mandatory=action_name in fsm_x_transition.mandatory, + mandatory=action_name in mandatory_list, ) break return seq - def _post_process_oks(self): + def _post_process_oks(self) -> None: """ Post-process the configuration data after it has been loaded and validated. @@ -70,26 +93,15 @@ def _post_process_oks(self): Raises: None """ - # Define the data structures to store the FSM configuration - self.log.debug("_post_process_oks configuration") - self.pre_transitions: dict(Transition, PreOrPostTransitionSequence) = {} - self.post_transitions: dict(Transition, PreOrPostTransitionSequence) = {} - self.actions: dict( - str, "conffwk.dal.FSMAction" - ) = {} # (e.g. "thread_pinning": thread_pinning FSM action object) - self.transitions: list(Transition) = [] - self.sequences: list["conffwk.dal.FSMSequence"] = [] - self.states: list(str) = self.data.states - self.initial_state: str = self.data.initial_state # Fill the actions dictionary with the FSMAction objects corresponding to the # action names defined in the configuration - for action in self.data.actions: # type: 'conffwk.dal.FSMAction' + for action in self.data.actions: self.actions[action.id] = FSMActionFactory.get().get_action( action.id, action ) - for transition in self.data.transitions: # type: 'conffwk.dal.FSMTransition' + for transition in self.data.transitions: tr = Transition( name=transition.id, source=transition.source, @@ -126,26 +138,23 @@ def _post_process_oks(self): cmd_ids = [cmd.id for cmd in sequence.sequence] self.sequences.append(FSMSequence(id=seq_id, command_ids=cmd_ids)) - # def _parse_dict(self, data): - # pass - - def get_actions(self): + def get_actions(self) -> Dict[str, FSMActionProtocol]: return self.actions - def get_initial_state(self): - return self.data.initial_state + def get_initial_state(self) -> str: + return str(self.data.initial_state) - def get_states(self): - return self.data.states + def get_states(self) -> List[str]: + return list(self.data.states) - def get_transitions(self): + def get_transitions(self) -> List[Transition]: return self.transitions - def get_pre_transitions_sequences(self): + def get_pre_transitions_sequences(self) -> Dict[Transition, PreOrPostTransitionSequence]: return self.pre_transitions - def get_post_transitions_sequences(self): + def get_post_transitions_sequences(self) -> Dict[Transition, PreOrPostTransitionSequence]: return self.post_transitions - def get_sequences(self): + def get_sequences(self) -> List[FSMSequence] : return self.sequences diff --git a/src/drunc/fsm/core.py b/src/drunc/fsm/core.py index 7ebb2c7e5..ab93e3499 100644 --- a/src/drunc/fsm/core.py +++ b/src/drunc/fsm/core.py @@ -1,14 +1,33 @@ +from __future__ import annotations +from typing import TYPE_CHECKING, Optional, Union, Dict, List, Set, Protocol, Any + +if TYPE_CHECKING: + FSMaction_t = object +else: + import conffwk + FSMaction_t = conffwk.dal.FSMAction + +class ActionProtocol(Protocol): + name: str # Every action must have a name string + +class ActionMethodProtocol(Protocol): + __name__: str + __module__: str + __self__: object + def __call__(self, _input_data: Dict[str, object], _context: object, **kwargs: object) -> Dict[str, object]: ... + + # Define the abcs first to avoid circular imports class FSMAction: """Abstract class defining a generic action""" - def __init__(self, name): + def __init__(self, name: str) -> None: self.name = name class Callback: - def __init__(self, method: "conffwk.dal.FSMAction", mandatory: bool = True): - self.method: "conffwk.dal.FSMAction" = method + def __init__(self, method: ActionMethodProtocol, mandatory: bool = True) -> None: + self.method: ActionMethodProtocol = method self.mandatory: bool = mandatory @@ -17,7 +36,6 @@ def __init__(self, method: "conffwk.dal.FSMAction", mandatory: bool = True): from dataclasses import dataclass from enum import Enum from inspect import Parameter, signature -from typing import Optional, Union import conffwk from druncschema.controller_pb2 import Argument, FSMSequence @@ -31,8 +49,11 @@ def __init__(self, method: "conffwk.dal.FSMAction", mandatory: bool = True): from drunc.utils.utils import get_logger, regex_match +class ContextProtocol(Protocol): + runinfo: Dict[str, object] + class PreOrPostTransitionSequence: - def __init__(self, transition: Transition, pre_or_post: str = "pre"): + def __init__(self, transition: Transition, pre_or_post: str = "pre") -> None: self.transition: Transition = transition if pre_or_post not in ["pre", "post"]: raise DruncSetupException( @@ -41,11 +62,11 @@ def __init__(self, transition: Transition, pre_or_post: str = "pre"): self.prefix: str = pre_or_post - self.sequence: list(Callback) = [] + self.sequence: list[Callback] = [] self.log = get_logger("controller.core.PreOrPostTransitionSequence") def add_callback( - self, action: "conffwk.dal.FSMaction", mandatory: bool = True + self, action: ActionProtocol, mandatory: bool = True ) -> None: """ Add a callback to the sequence. The method to be called will be determined by @@ -83,7 +104,7 @@ def add_callback( ) ] - def __str__(self): + def __str__(self) -> str: return ", ".join( [ f"{cb.method.__self__.__class__.__name__} (mandatory={cb.mandatory})" @@ -91,7 +112,7 @@ def __str__(self): ] ) - def execute(self, transition_data, transition_args, ctx=None): + def execute(self, transition_data: str | None, transition_args: Dict[str, object], ctx: ContextProtocol) -> str: self.log.debug(f"{transition_data=}, {transition_args=}") if not transition_data: transition_data = "{}" @@ -149,11 +170,11 @@ def get_arguments(self) -> list[Argument]: """ # Construct the list of arguments by looking at the signature of the methods - arguments: list(Argument) = [] + arguments: list[Argument] = [] # Check that there are no duplicate parameter names across the callbacks # otherwise, we won't know which one to use when executing the sequence - all_sequence_arguments: set(str) = set() # set(Argument names) + all_sequence_arguments: set[str] = set() # set(Argument names) # Iterate over the callbacks, construct the list of arguments for callback in self.sequence: @@ -184,7 +205,7 @@ def get_arguments(self) -> list[Argument]: # Determine the type of the argument, and set the default value if it is # optional. If the type is not one of the supported types, raise an # error - t: Argument.Type = Argument.Type.INT + t: int = Argument.Type.INT if p.annotation in (str, Optional[str], Union[str, None]): t = Argument.Type.STRING @@ -255,8 +276,17 @@ class FSMDestinationResult: destination_type: FSMDestinationType +class ConfigProtocol(Protocol): + def get_initial_state(self) -> str: ... + def get_states(self) -> List[str]: ... + def get_transitions(self) -> List[Transition]: ... + def get_sequences(self) -> List[FSMSequence]: ... + def get_pre_transitions_sequences(self) -> Dict[Transition, PreOrPostTransitionSequence]: ... + def get_post_transitions_sequences(self) -> Dict[Transition, PreOrPostTransitionSequence]: ... + + class FSM: - def __init__(self, conf): + def __init__(self, conf: ConfigProtocol) -> None: self.log = get_logger("controller.core.FSM") self.configuration = conf @@ -280,7 +310,7 @@ def __init__(self, conf): self.log.debug(f"Pre transition: {self.pre_transition_sequences[t]}") self.log.debug(f"Post transition: {self.post_transition_sequences[t]}") - def _enusure_unique_transition(self, transitions): + def _enusure_unique_transition(self, transitions: list[Transition]) -> None: a_set = set() for t in transitions: if t.name in a_set: @@ -299,10 +329,10 @@ def get_all_sequences(self) -> list[FSMSequence]: """Grab all the transitions""" return self.sequences - def is_destination_of_this_transition(self, state, transition) -> bool: - return transition.destination == state + def is_destination_of_this_transition(self, state: str, transition: Transition) -> bool: + return bool(transition.destination == state) - def get_destination_state(self, source_state, transition) -> FSMDestinationResult: + def get_destination_state(self, source_state: str, transition: Transition) -> FSMDestinationResult: """Tells us where a particular transition will take us, given the source_state""" right_name = [t for t in self.transitions if t == transition] @@ -336,7 +366,7 @@ def get_destination_state(self, source_state, transition) -> FSMDestinationResul destination_type=FSMDestinationType.TRANSITION_NOT_VALID, ) - def get_executable_transitions(self, source_state) -> list[Transition]: + def get_executable_transitions(self, source_state: str) -> list[Transition]: valid_transitions = [] for tr in self.transitions: @@ -349,7 +379,7 @@ def get_executable_transitions(self, source_state) -> list[Transition]: return valid_transitions - def get_executable_sequences(self, source_state) -> list[FSMSequence]: + def get_executable_sequences(self, source_state: str) -> list[FSMSequence]: valid_sequences = [] for seq in self.sequences: @@ -366,7 +396,7 @@ def get_executable_sequences(self, source_state) -> list[FSMSequence]: return valid_sequences - def get_transition(self, transition_name) -> Transition: + def get_transition(self, transition_name: str) -> Transition: self.log.debug(f"Searching for transition {transition_name}") transition = [t for t in self.transitions if t.name == transition_name] self.log.debug(f"Found transition {transition}") @@ -374,22 +404,22 @@ def get_transition(self, transition_name) -> Transition: raise fsme.NoTransitionOfName(transition_name) return transition[0] - def can_execute_transition(self, source_state, transition) -> bool: + def can_execute_transition(self, source_state: str, transition: Transition) -> bool: """Check that this transition is allowed given the source_state""" self.log.debug(f"can_execute_transition {transition.source!s} {source_state}") - return regex_match(transition.source, source_state) + return bool(regex_match(transition.source, source_state)) def prepare_transition( - self, transition, transition_data, transition_args, ctx=None - ): + self, transition: Transition, transition_data: str, transition_args: Dict[str, object], ctx: ContextProtocol + ) -> str: transition_data = self.pre_transition_sequences[transition].execute( transition_data, transition_args, ctx ) return transition_data def finalise_transition( - self, transition, transition_data, transition_args, ctx=None - ): + self, transition: Transition, transition_data: str, transition_args: Dict[str, object], ctx: ContextProtocol + ) -> str: transition_data = self.post_transition_sequences[transition].execute( transition_data, transition_args, ctx ) diff --git a/src/drunc/fsm/exceptions.py b/src/drunc/fsm/exceptions.py index 66148eaa5..a7e113d86 100644 --- a/src/drunc/fsm/exceptions.py +++ b/src/drunc/fsm/exceptions.py @@ -6,15 +6,15 @@ class FSMException(DruncCommandException): class NoTransitionOfName(FSMException): - def __init__(self, transition_name): + def __init__(self, transition_name: str) -> None: self.message = f'Transition "{transition_name}" does not exist' - super(NoTransitionOfName, self).__init__(self.message) + super().__init__(self.message) class DuplicateTransition(FSMException): """If a transition has the same name as another""" - def __init__(self, transition_name): + def __init__(self, transition_name: str) -> None: self.message = f'Transition "{transition_name}" is a duplicate' super().__init__(self.message) @@ -22,7 +22,7 @@ def __init__(self, transition_name): class InvalidTransition(FSMException): """Raised when the transition isn't in the list of currently accessible transitions""" - def __init__(self, transition, state): + def __init__(self, transition: str, state: str) -> None: self.message = f"Transition {transition} is not allowed from the state {state}." super().__init__(self.message) @@ -30,7 +30,7 @@ def __init__(self, transition, state): class UnregisteredTransition(FSMException): """Raised when the transition is allowed, but we can't find the implementation""" - def __init__(self, transition): + def __init__(self, transition: str) -> None: self.message = f"Implementation of {transition} not found." super().__init__(self.message) @@ -38,7 +38,7 @@ def __init__(self, transition): class UnknownAction(FSMException): """Raised when a plugin name is provided that does not correspond to any files in /plugins""" - def __init__(self, name): + def __init__(self, name: str) -> None: self.message = f'"{name}" is not a known plugin.' super().__init__(self.message) @@ -46,7 +46,7 @@ def __init__(self, name): class MissingArgument(FSMException): """Raised when a mandatory argument is not provided for a transition""" - def __init__(self, param, name): + def __init__(self, param: str, name: str) -> None: self.message = f'The mandatory argument "{param}" was not provided to the transition {name}' super().__init__(self.message) @@ -54,7 +54,7 @@ def __init__(self, param, name): class MissingArgumentValue(FSMException): """Raised when a mandatory argument is not provided for a transition""" - def __init__(self): + def __init__(self) -> None: self.message = "A passed argument does not have an associated value, arguments are key-value pairs." super().__init__(self.message) @@ -62,13 +62,13 @@ def __init__(self): class DoubleArgument(FSMException): """Raised when an argument is provided more than once""" - def __init__(self, txt): + def __init__(self, txt: str) -> None: self.message = txt super().__init__(self.message) class UnhandledArgumentType(FSMException): - def __init__(self, annotation): + def __init__(self, annotation: str) -> None: self.message = f'The argument "{annotation}" cannot be handled' super().__init__(self.message) @@ -76,7 +76,7 @@ def __init__(self, annotation): class UnknownArgument(FSMException): """Raised when an unwanted argument is given to a transition""" - def __init__(self, param, name): + def __init__(self, param: str, name: str) -> None: self.message = ( f'The mandatory argument "{param}" is not required by transition {name}' ) @@ -86,7 +86,7 @@ def __init__(self, param, name): class InvalidAction(FSMException): """Raised when an action doesn't have pre/post transitions""" - def __init__(self, iface): + def __init__(self, iface: str) -> None: self.message = ( f'The action "{iface}" does not have any pre or post transition method' ) @@ -96,48 +96,48 @@ def __init__(self, iface): class InvalidActionMethod(FSMException): """Raised when an action doesn't have the pre/post transitions arguments""" - def __init__(self, iface, method): + def __init__(self, iface: str, method: str) -> None: self.message = f'The action "{iface}" method {method} does not have the correct arguements, each one should have at least "_input_data" and "**kwargs", and have type annotations' super().__init__(self.message) class MethodSignatureMissingAnnotation(FSMException): - def __init__(self, iface, method, pname): + def __init__(self, iface: str, method: str, pname: str) -> None: self.message = f'The action "{iface}" method {method} does not have the correct arguement annotation for "{pname}", provide "argument:int" to your pre/post methods.' super().__init__(self.message) class TransitionDataOfIncorrectFormat(FSMException): - def __init__(self, data): + def __init__(self, data: object) -> None: self.message = f'The data "{data}" could not be interpreted as json' - super(MethodSignatureMissingAnnotation, self).__init__(self.message) + super().__init__(self.message) class CannotGetRunNumber(FSMException): - def __init__(self, data): + def __init__(self, data: object) -> None: self.message = f"Could not get Run Number because {data}" super().__init__(self.message) class DBRunRegistryConfigurationError(FSMException): - def __init__(self, data): + def __init__(self, data: object) -> None: self.message = f"RunRegistryDB configuration error: {data}" super().__init__(self.message) class CannotInsertRunNumber(FSMException): - def __init__(self, data): + def __init__(self, data: object) -> None: self.message = f"Could not insert Run into RunRegistryDB because {data}" super().__init__(self.message) class CannotUpdateStopTime(FSMException): - def __init__(self, data): + def __init__(self, data: object) -> None: self.message = f"Could not update stop time in RunRegistryDB because {data}" super().__init__(self.message) class InvalidDataReturnByFSMAction(FSMException): - def __init__(self, data): + def __init__(self, data: object) -> None: self.message = ( f"The action returns an incorrect object which isn't serialisable: {data}" ) @@ -145,19 +145,19 @@ def __init__(self, data): class ThreadPinningFailed(FSMException): - def __init__(self, host): + def __init__(self, host: str) -> None: self.message = f'The thread pinning on "{host}" failed' super().__init__(self.message) class CannotGetSoftwareVersion(FSMException): - def __init__(self): + def __init__(self) -> None: self.message = "RunRegistryDB: dunedaq version not in the variable env DUNE_DAQ_BASE_RELEASE! Exit drunc and export DUNE_DAQ_BASE_RELEASE=dunedaq-vX.XX.XX\n" super().__init__(self.message) class CannotSendElisaMessage(FSMException): - def __init__(self, data): + def __init__(self, data: object) -> None: self.message = f"Cannot send message to ELisA because {data}. Do it manually at https://np-vd-coldbox-elog.app.cern.ch or https://pdsp-elog.app.cern.ch!" super().__init__(self.message) diff --git a/src/drunc/fsm/transition.py b/src/drunc/fsm/transition.py index 4faafba37..15b70f3af 100644 --- a/src/drunc/fsm/transition.py +++ b/src/drunc/fsm/transition.py @@ -1,12 +1,25 @@ +from __future__ import annotations + +from typing import List, Optional + +from druncschema.controller_pb2 import Argument + + class Transition: - def __init__(self, name, source, destination, arguments=[], help: str = ""): + def __init__(self, + name: str, + source: str, + destination: str, + arguments: Optional[List[Argument]] = None, + help: str = "" + ) -> None: self.source = source self.destination = destination self.name = name self.arguments = arguments self.help = help - def __eq__(self, another): + def __eq__(self, another: object) -> bool: same_name = hasattr(another, "name") and self.name == another.name same_destination = ( hasattr(another, "destination") and self.destination == another.destination @@ -14,8 +27,8 @@ def __eq__(self, another): same_source = hasattr(another, "source") and self.source == another.source return same_name and same_destination and same_source - def __hash__(self): + def __hash__(self) -> int: return hash(self.__str__()) - def __str__(self): + def __str__(self) -> str: return f'"{self.name}": "{self.source}" → "{self.destination}"' diff --git a/src/drunc/fsm/utils.py b/src/drunc/fsm/utils.py index 6c8c0c44c..4c3aca92c 100644 --- a/src/drunc/fsm/utils.py +++ b/src/drunc/fsm/utils.py @@ -1,4 +1,4 @@ -from typing import MutableMapping +from typing import MutableMapping, Optional from druncschema.controller_pb2 import ( Argument, @@ -12,6 +12,22 @@ from drunc.fsm.transition import Transition from drunc.utils.grpc_utils import unpack_any +from typing import MutableMapping, Optional + +from druncschema.controller_pb2 import ( + Argument, + FSMCommandDescription, + FSMCommandsDescription, +) +from druncschema.generic_pb2 import bool_msg, float_msg, int_msg, string_msg + +# Alias Protobuf's Any to PbAny to prevent confusion with typing.Any +from google.protobuf.any_pb2 import Any as PbAny + +import drunc.fsm.exceptions as fsme +from drunc.fsm.transition import Transition +from drunc.utils.grpc_utils import unpack_any + def convert_fsm_transition(transitions: list[Transition]) -> FSMCommandsDescription: """ @@ -28,7 +44,7 @@ def convert_fsm_transition(transitions: list[Transition]) -> FSMCommandsDescript None. """ commands_description = FSMCommandsDescription() - for t in transitions: # type: Transition + for t in transitions: commands_description.commands.append( FSMCommandDescription( name=t.name, @@ -41,7 +57,10 @@ def convert_fsm_transition(transitions: list[Transition]) -> FSMCommandsDescript return commands_description -def decode_fsm_arguments(arguments: MutableMapping[str, Any], arguments_format: list[Argument]) -> dict[str, str | int | float | bool]: +def decode_fsm_arguments( + arguments: MutableMapping[str, PbAny], + arguments_format: list[Argument] +) -> dict[str, str | int | float | bool]: """ Decodes the arguments of a FSM command. @@ -50,7 +69,7 @@ def decode_fsm_arguments(arguments: MutableMapping[str, Any], arguments_format: drunc operating modes. Args: - arguments (MutableMapping[str, Any]): The arguments to decode. + arguments (MutableMapping[str, PbAny]): The arguments to decode. arguments_format (list[Argument]): The format of the arguments. Returns: @@ -61,15 +80,18 @@ def decode_fsm_arguments(arguments: MutableMapping[str, Any], arguments_format: fsme.UnhandledArgumentType: If an argument type is not handled. """ - def get_argument(name, arguments): - for n, k in arguments.items(): + # Added explicit type hints to the nested function + def get_argument(name: str, args: MutableMapping[str, PbAny]) -> Optional[PbAny]: + for n, k in args.items(): if n == name: return k return None - out_dict = {} + # Explicitly type out_dict so Mypy doesn't infer dict[Any, Any] + out_dict: dict[str, str | int | float | bool] = {} + for arg in arguments_format: - arg_value = get_argument(arg.name, arguments) + arg_value: Optional[PbAny] = get_argument(arg.name, arguments) if arg.presence == Argument.Presence.MANDATORY and arg_value is None: raise fsme.MissingArgument(arg.name, "") @@ -88,4 +110,5 @@ def get_argument(name, arguments): out_dict[arg.name] = unpack_any(arg_value, bool_msg).value case _: raise fsme.UnhandledArgumentType(arg.type) - return out_dict + + return out_dict \ No newline at end of file From c8a08a7c6b4812abe934d6d5f8cb983468b8b07a Mon Sep 17 00:00:00 2001 From: Miruna Serian Date: Tue, 9 Jun 2026 19:27:08 +0100 Subject: [PATCH 02/20] Move protocols to separate file --- src/drunc/fsm/_protocols.py | 87 +++++++++++++++++++ src/drunc/fsm/action_factory.py | 31 +++---- src/drunc/fsm/actions/db_run_registry.py | 23 +++-- src/drunc/fsm/actions/file_logbook.py | 20 +++-- src/drunc/fsm/actions/file_run_registry.py | 8 +- src/drunc/fsm/actions/some_test_action.py | 11 +-- src/drunc/fsm/actions/thread_pinning.py | 17 ++-- .../actions/timing/master_send_fl_command.py | 12 +-- .../fsm/actions/trigger_rate_specifier.py | 7 +- .../fsm/actions/user_provided_run_number.py | 14 +-- src/drunc/fsm/actions/usvc_elisa_logbook.py | 16 ++-- .../fsm/actions/usvc_provided_run_number.py | 19 ++-- src/drunc/fsm/actions/utils.py | 8 +- src/drunc/fsm/configuration.py | 22 ++--- src/drunc/fsm/core.py | 32 +------ src/drunc/fsm/utils.py | 18 +--- 16 files changed, 204 insertions(+), 141 deletions(-) create mode 100644 src/drunc/fsm/_protocols.py diff --git a/src/drunc/fsm/_protocols.py b/src/drunc/fsm/_protocols.py new file mode 100644 index 000000000..6cac35532 --- /dev/null +++ b/src/drunc/fsm/_protocols.py @@ -0,0 +1,87 @@ +from typing import Dict, Iterable, List, Optional, Protocol + +from druncschema.controller_pb2 import FSMSequence + +from drunc.fsm.core import PreOrPostTransitionSequence +from drunc.fsm.transition import Transition + + +class ParameterProtocol(Protocol): + name: str + value: str + +class InitConfigurationProtocol(Protocol): + parameters: Iterable[ParameterProtocol] + + +class ActorProtocol(Protocol): + def get_user_name(self) -> str: ... + + +class DetConfigProtocol(Protocol): + id: str + +class DALProtocol(Protocol): + detector_configuration: DetConfigProtocol + +class DBProtocol(Protocol): + def get_dal(self, class_name: str, uid: str) -> DALProtocol: ... + +class OksKeyProtocol(Protocol): + session: str + +class RuntimeConfigurationProtocol(Protocol): + initial_data: str + oks_key: OksKeyProtocol + + +class ConfigurationProtocol(Protocol): + id: str + db: DBProtocol + oks_key: OksKeyProtocol + initial_data: str + parameters: Iterable[ParameterProtocol] + name: str + + +class ContextProtocol(Protocol): + actor: ActorProtocol + configuration: ConfigurationProtocol + runinfo: Dict[str, object] + + +class SessionDalProtocol(Protocol): + segment: object + rte_script: Optional[str] + + +class SSHCommandProtocol(Protocol): + def __call__(self, *args: str, _err_to_out: bool = ...) -> object: ... + + +class ShErrorProtocol(Protocol): + stdout: bytes + stderr: bytes + +class ActionMethodProtocol(Protocol): + __name__: str + __module__: str + __self__: object + def __call__(self, *args: object, **kwargs: object) -> object: ... + + +class FSMActionProtocol(Protocol): + name: str + + +class ConfigProtocol(Protocol): + def get_initial_state(self) -> str: ... + def get_states(self) -> List[str]: ... + def get_transitions(self) -> List[Transition]: ... + def get_sequences(self) -> List[FSMSequence]: ... + def get_pre_transitions_sequences(self) -> Dict[Transition, PreOrPostTransitionSequence]: ... + def get_post_transitions_sequences(self) -> Dict[Transition, PreOrPostTransitionSequence]: ... + + +class ActionMethod(Protocol): + def __call__(self, *args: object, **kwargs: object) -> object: ... diff --git a/src/drunc/fsm/action_factory.py b/src/drunc/fsm/action_factory.py index a14f5ce63..805c78ec0 100644 --- a/src/drunc/fsm/action_factory.py +++ b/src/drunc/fsm/action_factory.py @@ -1,7 +1,7 @@ from __future__ import annotations import inspect -from typing import TYPE_CHECKING, Dict, Protocol, Type +from typing import TYPE_CHECKING, Dict, Protocol, Type, cast import conffwk @@ -18,18 +18,7 @@ from drunc.fsm.actions.usvc_elisa_logbook import ElisaLogbook from drunc.fsm.actions.usvc_provided_run_number import UsvcProvidedRunNumber -if TYPE_CHECKING: - # for mypy treat FSMaction_t as a local objext - FSMaction_t = object -else: - # at runtime use the actual type from conffwk.dal - FSMaction_t = conffwk.dal.FSMaction - -class ActionMethod(Protocol): - def __call__(self, *args: object, **kwargs: object) -> object: ... - -class FSMActionProtocol(Protocol): - name: str +from drunc.fsm._protocols import ActionMethodProtocol, ConfigurationProtocol, FSMActionProtocol class FSMActionFactory: @@ -38,23 +27,23 @@ class FSMActionFactory: def __init__(self) -> None: raise DruncSetupException("Call get() instead") - def _get_pre_transitions(self, action: FSMActionProtocol) -> Dict[str, ActionMethod]: - retr: Dict[str, ActionMethod] = {} + def _get_pre_transitions(self, action: FSMActionProtocol) -> Dict[str, ActionMethodProtocol]: + retr: Dict[str, ActionMethodProtocol] = {} for name, method in inspect.getmembers(action): if inspect.ismethod(method): if name.startswith("pre_"): - retr[name] = method + retr[name] = cast(ActionMethodProtocol, method) return retr - def _get_post_transitions(self, action: FSMActionProtocol) -> Dict[str, ActionMethod]: - retr: Dict[str, ActionMethod] = {} + def _get_post_transitions(self, action: FSMActionProtocol) -> Dict[str, ActionMethodProtocol]: + retr: Dict[str, ActionMethodProtocol] = {} for name, method in inspect.getmembers(action): if inspect.ismethod(method): if name.startswith("post_"): - retr[name] = method + retr[name] = cast(ActionMethodProtocol, method) return retr - def _validate_signature(self, name:str , method: ActionMethod, action: str) -> None: + def _validate_signature(self, name:str , method: ActionMethodProtocol, action: str) -> None: sig = inspect.signature(method) if ( @@ -85,7 +74,7 @@ def _validate_action(self, action: FSMActionProtocol) -> None: self._validate_signature(k, v, action.name) def get_action( - self, action_name: str, action_configuration: FSMaction_t + self, action_name: str, action_configuration: ConfigurationProtocol ) -> FSMActionProtocol: """ Construct the action interface for the given action name and configuration. diff --git a/src/drunc/fsm/actions/db_run_registry.py b/src/drunc/fsm/actions/db_run_registry.py index 79b8e6152..ac37d9a92 100644 --- a/src/drunc/fsm/actions/db_run_registry.py +++ b/src/drunc/fsm/actions/db_run_registry.py @@ -3,10 +3,11 @@ import tempfile import requests -from daqconf.consolidate import consolidate_db -from daqconf.jsonify import jsonify_xml_data -from daqconf.validate import validate_session +from daqconf.consolidate import consolidate_db # type: ignore[import-untyped] +from daqconf.jsonify import jsonify_xml_data # type: ignore[import-untyped] +from daqconf.validate import validate_session # type: ignore[import-untyped] +from drunc.fsm._protocols import ContextProtocol from drunc.fsm.actions.utils import get_dotdrunc_json from drunc.fsm.core import FSMAction from drunc.fsm.exceptions import ( @@ -20,7 +21,7 @@ class DBRunRegistry(FSMAction): - def __init__(self, configuration): + def __init__(self, configuration: object) -> None: super().__init__(name="db-run-registry") self.log = get_logger("controller.iface.usvc_db_run_registry") @@ -36,7 +37,12 @@ def __init__(self, configuration): ) from exc self.timeout = 2 - def pre_start(self, _input_data: dict, _context, **kwargs): + def pre_start( + self, + _input_data: dict[str, object], + _context: ContextProtocol, + **kwargs: object + ) -> dict[str, object]: """ Upload a copy of the XML and JSON configurations used to take the current run to the Run Registry, and insert a new run number entry. @@ -162,7 +168,12 @@ def pre_start(self, _input_data: dict, _context, **kwargs): return _input_data - def post_drain_dataflow(self, _input_data, _context, **kwargs): + def post_drain_dataflow( + self, + _input_data: dict[str, object], + _context: ContextProtocol, + **kwargs: object + ) -> None: try: requests.get( self.API_SOCKET + "/runregistry/updateStopTime/" + str(self.run_number), diff --git a/src/drunc/fsm/actions/file_logbook.py b/src/drunc/fsm/actions/file_logbook.py index b5cd530a9..3bab1e26a 100644 --- a/src/drunc/fsm/actions/file_logbook.py +++ b/src/drunc/fsm/actions/file_logbook.py @@ -1,18 +1,24 @@ from typing import Optional +from drunc.fsm._protocols import ContextProtocol, ConfigurationProtocol + from drunc.fsm.core import FSMAction from drunc.utils.utils import now_str class FileLogbook(FSMAction): - def __init__(self, configuration): + def __init__(self, configuration: ConfigurationProtocol) -> None: super().__init__(name="file-logbook") self.conf_dict = {p.name: p.value for p in configuration.parameters} self.file = self.conf_dict["file_name"] def post_start( - self, _input_data, _context, file_logbook_post: Optional[str] = None, **kwargs - ): + self, + _input_data: dict[str, object], + _context: ContextProtocol, + file_logbook_post: Optional[str] = None, + **kwargs: object + ) -> dict[str, object]: with open(self.file, "a") as f: f.write( f"Run {_input_data['run']} started by {_context.actor.get_user_name()} at {now_str()}\n" @@ -24,8 +30,12 @@ def post_start( return _input_data def post_drain_dataflow( - self, _input_data, _context, file_logbook_post: str = "", **kwargs - ): + self, + _input_data: dict[str, object], + _context: ContextProtocol, + file_logbook_post: str = "", + **kwargs: object + ) -> dict[str, object]: with open(self.file, "a") as f: f.write( f"Current run stopped by {_context.actor.get_user_name()} at {now_str()}\n" diff --git a/src/drunc/fsm/actions/file_run_registry.py b/src/drunc/fsm/actions/file_run_registry.py index 4254f0a8b..c3c739615 100644 --- a/src/drunc/fsm/actions/file_run_registry.py +++ b/src/drunc/fsm/actions/file_run_registry.py @@ -1,16 +1,18 @@ import os +from typing import Dict, Protocol -from daqconf.consolidate import consolidate_db +from daqconf.consolidate import consolidate_db # type: ignore[import-untyped] +from drunc.fsm._protocols import ContextProtocol from drunc.fsm.core import FSMAction class FileRunRegistry(FSMAction): - def __init__(self, configuration): + def __init__(self, configuration: object) -> None: super().__init__(name="file-run-registry") self.configuration = configuration - def pre_start(self, _input_data, _context, **kwargs): + def pre_start(self, _input_data: Dict[str, object], _context: ContextProtocol, **kwargs: object) -> Dict[str, object]: run_number = _input_data["run"] dest = os.getcwd() + "/run_conf" + str(run_number) + ".data.xml" consolidate_db(_context.configuration.initial_data.split(":")[1], f"{dest}") diff --git a/src/drunc/fsm/actions/some_test_action.py b/src/drunc/fsm/actions/some_test_action.py index 669b5c44f..2a42e1de0 100644 --- a/src/drunc/fsm/actions/some_test_action.py +++ b/src/drunc/fsm/actions/some_test_action.py @@ -1,5 +1,6 @@ from enum import Enum +from drunc.fsm._protocols import ConfigurationProtocol, ContextProtocol from drunc.fsm.core import FSMAction @@ -9,18 +10,18 @@ class an_enum(Enum): class SomeTestAction(FSMAction): - def __init__(self, configuration): + def __init__(self, configuration: ConfigurationProtocol) -> None: super().__init__(name="test-action") def pre_conf( self, - _input_data: dict, - _context, + _input_data: dict[str, object], + _context: ContextProtocol, some_int: int, some_str: str, some_float: float = 0.2, - **kwargs, - ) -> dict: + **kwargs: object, + ) -> dict[str, object]: print(f"Running pre_conf of {self.name}") _input_data["some_int"] = some_int _input_data["some_str"] = some_str diff --git a/src/drunc/fsm/actions/thread_pinning.py b/src/drunc/fsm/actions/thread_pinning.py index b1fa976e5..8a3dca809 100644 --- a/src/drunc/fsm/actions/thread_pinning.py +++ b/src/drunc/fsm/actions/thread_pinning.py @@ -1,8 +1,11 @@ import getpass +from asyncio import Protocol from os import environ +from typing import Iterable, Optional import conffwk -from sh import Command, ErrorReturnCode +from drunc.fsm._protocols import ContextProtocol, InitConfigurationProtocol +from sh import Command, ErrorReturnCode # type: ignore[import-untyped] from drunc.exceptions import DruncSetupException from drunc.fsm.core import FSMAction @@ -13,12 +16,12 @@ class ThreadPinning(FSMAction): - def __init__(self, configuration): + def __init__(self, configuration: InitConfigurationProtocol) -> None: super().__init__(name="thread-pinning") self.log = get_logger("controller.iface.thread-pinning") self.conf_dict = {p.name: p.value for p in configuration.parameters} - def pin_thread(self, thread_pinning_file, configuration, session): + def pin_thread(self, thread_pinning_file: str, configuration: str, session: str) -> None: db = conffwk.Configuration(configuration) session_dal = db.get_dal(class_name="Session", uid=session) @@ -27,7 +30,7 @@ def pin_thread(self, thread_pinning_file, configuration, session): session_name=session, session_dal_obj=session_dal, segment_obj=session_dal.segment, - env=environ, + env=dict(environ), tree_prefix=[], ) @@ -86,7 +89,7 @@ def pin_thread(self, thread_pinning_file, configuration, session): if failed_hosts: raise ThreadPinningFailed(failed_hosts_error_str) - def post_conf(self, _input_data, _context, **kwargs): + def post_conf(self, _input_data: dict[str, object], _context: ContextProtocol, **kwargs: object) -> dict[str, object]: if "post_conf" in self.conf_dict: self.pin_thread( self.conf_dict["post_conf"], @@ -95,7 +98,7 @@ def post_conf(self, _input_data, _context, **kwargs): ) return _input_data - def post_start(self, _input_data, _context, **kwargs): + def post_start(self, _input_data: dict[str, object], _context: ContextProtocol, **kwargs: object) -> dict[str, object]: if "post_start" in self.conf_dict: self.pin_thread( self.conf_dict["post_start"], @@ -104,7 +107,7 @@ def post_start(self, _input_data, _context, **kwargs): ) return _input_data - def pre_conf(self, _input_data, _context, **kwargs): + def pre_conf(self, _input_data: dict[str, object], _context: ContextProtocol, **kwargs: object) -> dict[str, object]: if "pre_conf" in self.conf_dict: self.pin_thread( self.conf_dict["pre_conf"], diff --git a/src/drunc/fsm/actions/timing/master_send_fl_command.py b/src/drunc/fsm/actions/timing/master_send_fl_command.py index a289af2a2..524034738 100644 --- a/src/drunc/fsm/actions/timing/master_send_fl_command.py +++ b/src/drunc/fsm/actions/timing/master_send_fl_command.py @@ -1,19 +1,19 @@ from drunc.fsm.core import FSMAction - +from drunc.fsm._protocols import ContextProtocol class MasterSendFLCommand(FSMAction): - def __init__(self, configuration): + def __init__(self, configuration: object) -> None: super().__init__(name="master-send-fl-command") def pre_master_send_fl_command( self, - _input_data, - _context, + _input_data: dict[str, object], + _context: ContextProtocol, fl_cmd_id: int, channel: int, number_of_commands_to_send: int, - **kwargs, - ): + **kwargs: object, + ) -> dict[str, object]: # parse fl_cmd_id... _input_data["fl_cmd_id"] = fl_cmd_id _input_data["channel"] = channel diff --git a/src/drunc/fsm/actions/trigger_rate_specifier.py b/src/drunc/fsm/actions/trigger_rate_specifier.py index 189da9002..19c0ada03 100644 --- a/src/drunc/fsm/actions/trigger_rate_specifier.py +++ b/src/drunc/fsm/actions/trigger_rate_specifier.py @@ -1,12 +1,13 @@ +from drunc.fsm._protocols import ConfigurationProtocol, ContextProtocol from drunc.fsm.core import FSMAction class TriggerRateSpecifier(FSMAction): - def __init__(self, configuration): + def __init__(self, configuration: ConfigurationProtocol) -> None: super().__init__(name="trigger-rate-specifier") def pre_change_rate( - self, _input_data: dict, _context, trigger_rate: float, **kwargs - ): + self, _input_data: dict[str, object], _context: ContextProtocol, trigger_rate: float, **kwargs: object + ) -> dict[str, object]: _input_data["trigger_rate"] = trigger_rate return _input_data diff --git a/src/drunc/fsm/actions/user_provided_run_number.py b/src/drunc/fsm/actions/user_provided_run_number.py index ea97dec5f..3e4aca02f 100644 --- a/src/drunc/fsm/actions/user_provided_run_number.py +++ b/src/drunc/fsm/actions/user_provided_run_number.py @@ -6,20 +6,22 @@ class UserProvidedRunNumber(FSMAction): - def __init__(self, configuration): + def __init__(self, configuration: object) -> None: super().__init__(name="run-number") def pre_start( self, - _input_data: dict, - _context, + _input_data: dict[str, object], + _context: object, run_number: int, run_type: Optional[str] = "TEST", disable_data_storage: bool = False, trigger_rate: Optional[float] = None, - **kwargs, - ): - run_type = validate_run_type(run_type.upper()) + **kwargs: object, + ) -> dict[str, object]: + + safe_run_type = run_type.upper() if run_type is not None else "TEST" + run_type = validate_run_type(safe_run_type) _input_data["production_vs_test"] = run_type _input_data["run"] = run_number _input_data["disable_data_storage"] = disable_data_storage diff --git a/src/drunc/fsm/actions/usvc_elisa_logbook.py b/src/drunc/fsm/actions/usvc_elisa_logbook.py index 7dfcb4716..207e57e99 100644 --- a/src/drunc/fsm/actions/usvc_elisa_logbook.py +++ b/src/drunc/fsm/actions/usvc_elisa_logbook.py @@ -4,6 +4,7 @@ import requests +from drunc.fsm._protocols import ContextProtocol from drunc.fsm.actions.utils import get_dotdrunc_json from drunc.fsm.core import FSMAction from drunc.fsm.exceptions import ( @@ -15,7 +16,7 @@ class ElisaLogbook(FSMAction): - def __init__(self): + def __init__(self) -> None: super().__init__(name="elisa-logbook") self.log = get_logger("controller.iface.elisa-logbook") @@ -65,7 +66,6 @@ def __init__(self): ) self.log.warning(warn_msg) else: - warn_msg: str = "" if default_elisa_logbook: warn_msg = ( "You need to update your ~/.drunc.json: The default ELisA logbook " @@ -83,10 +83,10 @@ def __init__(self): self.timeout = 5 def post_start( - self, _input_data: dict, _context, elisa_post: Optional[str] = None, **kwargs - ): + self, _input_data: dict[str, object], _context: ContextProtocol, elisa_post: Optional[str] = None, **kwargs: object + ) -> dict[str, object]: if self.elisa_hardware in self.no_publish_hardware: - return + return {} text = "" self.thread_id = None # Clear this value here, so that if it fails stop can't reply to an old message @@ -142,10 +142,10 @@ def post_start( return _input_data def post_drain_dataflow( - self, _input_data, _context, elisa_post: Optional[str] = None, **kwargs - ): + self, _input_data: dict[str, object], _context: ContextProtocol, elisa_post: Optional[str] = None, **kwargs: object + ) -> dict[str, object]: if self.elisa_hardware in self.no_publish_hardware: - return + return {} text = "" if elisa_post is not None: self.log.info( diff --git a/src/drunc/fsm/actions/usvc_provided_run_number.py b/src/drunc/fsm/actions/usvc_provided_run_number.py index a3553c966..335c269c5 100644 --- a/src/drunc/fsm/actions/usvc_provided_run_number.py +++ b/src/drunc/fsm/actions/usvc_provided_run_number.py @@ -1,5 +1,5 @@ import time -from typing import Optional +from typing import Optional, cast import requests @@ -7,10 +7,10 @@ from drunc.fsm.core import FSMAction from drunc.fsm.exceptions import CannotGetRunNumber, DotDruncJsonIncorrectFormat from drunc.utils.utils import get_logger - +from drunc.fsm._protocols import ContextProtocol class UsvcProvidedRunNumber(FSMAction): - def __init__(self, configuration): + def __init__(self, configuration: object) -> None: self.log = get_logger("controller.iface.usvc_run_number") super().__init__(name="usvc-provided-run-number") dotdrunc = get_dotdrunc_json() @@ -27,13 +27,13 @@ def __init__(self, configuration): def pre_start( self, - _input_data: dict, - _context, + _input_data: dict[str, object], + _context: ContextProtocol, run_type: str, disable_data_storage: bool = False, trigger_rate: Optional[float] = None, - **kwargs, - ): + **kwargs: object, + ) -> dict[str, object]: run_type = validate_run_type(run_type.upper()) _input_data["production_vs_test"] = run_type _input_data["run"] = self._getnew_run_number() @@ -45,7 +45,7 @@ def pre_start( return _input_data - def _getnew_run_number(self): + def _getnew_run_number(self) -> int: try: req = requests.get( self.API_SOCKET + "/runnumber/getnew", @@ -68,5 +68,6 @@ def _getnew_run_number(self): self.log.error(error) raise CannotGetRunNumber(error) from exc - self.run = req.json()[0][0][0] + self.run = cast(int, req.json()[0][0][0]) + return self.run diff --git a/src/drunc/fsm/actions/utils.py b/src/drunc/fsm/actions/utils.py index b32d5eed3..5f9a1ac59 100644 --- a/src/drunc/fsm/actions/utils.py +++ b/src/drunc/fsm/actions/utils.py @@ -11,8 +11,10 @@ def validate_run_type(run_type: str) -> str: """Validate the run type - :param run_type: the run type - :return: the validated run type + Args: + run_type: the run type + Returns: + the validated run type """ RUN_TYPES = ["PROD", "TEST"] if run_type not in RUN_TYPES: @@ -22,7 +24,7 @@ def validate_run_type(run_type: str) -> str: return run_type -def get_dotdrunc_json(path: str | None = None): +def get_dotdrunc_json(path: str | None = None): # type: ignore[no-untyped-def] # Resolution order: DOTDRUNC env var -> provided path -> default path file_path = os.getenv("DOTDRUNC") or path or "~/.drunc.json" try: diff --git a/src/drunc/fsm/configuration.py b/src/drunc/fsm/configuration.py index cbfec6a20..e3e7c918b 100644 --- a/src/drunc/fsm/configuration.py +++ b/src/drunc/fsm/configuration.py @@ -1,22 +1,15 @@ import logging -from typing import TYPE_CHECKING, Dict, List +from typing import Dict, List + from druncschema.controller_pb2 import FSMSequence +from drunc.fsm._protocols import ConfigurationProtocol from drunc.fsm.action_factory import FSMActionFactory, FSMActionProtocol from drunc.fsm.core import PreOrPostTransitionSequence from drunc.fsm.transition import Transition from drunc.utils.configuration import ConfHandler, OKSKey from drunc.utils.utils import get_logger -if TYPE_CHECKING: - FSMaction_t = object - FSMxTransition_t = object - FSMSequence_t = object -else: - import conffwk - FSMaction_t = conffwk.dal.FSMAction - FSMxTransition_t = conffwk.dal.FSMxTransition - FSMSequence_t = conffwk.dal.FSMSequence class FSMConfHandler(ConfHandler): @@ -35,7 +28,7 @@ def _fill_pre_post_transition_sequence_oks( self, prefix: str, transition: Transition, - data: List[FSMxTransition_t] | None, + data: List[ConfigurationProtocol] | None, ) -> PreOrPostTransitionSequence: """ Fill the pre or post transition sequence for a given transition. @@ -121,9 +114,10 @@ def _post_process_oks(self) -> None: ) ) - # Add the pre and post transition sequence arguments to the transition - tr.arguments += pre_transitions.get_arguments() - tr.arguments += post_transitions.get_arguments() + if tr.arguments is not None: + # Add the pre and post transition sequence arguments to the transition + tr.arguments += pre_transitions.get_arguments() + tr.arguments += post_transitions.get_arguments() # Store the pre and post transition sequences for the transition self.pre_transitions[tr] = pre_transitions diff --git a/src/drunc/fsm/core.py b/src/drunc/fsm/core.py index ab93e3499..5c37ebb46 100644 --- a/src/drunc/fsm/core.py +++ b/src/drunc/fsm/core.py @@ -1,21 +1,7 @@ from __future__ import annotations -from typing import TYPE_CHECKING, Optional, Union, Dict, List, Set, Protocol, Any - -if TYPE_CHECKING: - FSMaction_t = object -else: - import conffwk - FSMaction_t = conffwk.dal.FSMAction - -class ActionProtocol(Protocol): - name: str # Every action must have a name string - -class ActionMethodProtocol(Protocol): - __name__: str - __module__: str - __self__: object - def __call__(self, _input_data: Dict[str, object], _context: object, **kwargs: object) -> Dict[str, object]: ... +from typing import Optional, Union, Dict, List, Protocol +from drunc.fsm._protocols import ActionMethodProtocol, FSMActionProtocol, ContextProtocol, ConfigProtocol # Define the abcs first to avoid circular imports class FSMAction: @@ -49,9 +35,6 @@ def __init__(self, method: ActionMethodProtocol, mandatory: bool = True) -> None from drunc.utils.utils import get_logger, regex_match -class ContextProtocol(Protocol): - runinfo: Dict[str, object] - class PreOrPostTransitionSequence: def __init__(self, transition: Transition, pre_or_post: str = "pre") -> None: self.transition: Transition = transition @@ -66,7 +49,7 @@ def __init__(self, transition: Transition, pre_or_post: str = "pre") -> None: self.log = get_logger("controller.core.PreOrPostTransitionSequence") def add_callback( - self, action: ActionProtocol, mandatory: bool = True + self, action: FSMActionProtocol, mandatory: bool = True ) -> None: """ Add a callback to the sequence. The method to be called will be determined by @@ -275,15 +258,6 @@ class FSMDestinationResult: destination_state: str | None destination_type: FSMDestinationType - -class ConfigProtocol(Protocol): - def get_initial_state(self) -> str: ... - def get_states(self) -> List[str]: ... - def get_transitions(self) -> List[Transition]: ... - def get_sequences(self) -> List[FSMSequence]: ... - def get_pre_transitions_sequences(self) -> Dict[Transition, PreOrPostTransitionSequence]: ... - def get_post_transitions_sequences(self) -> Dict[Transition, PreOrPostTransitionSequence]: ... - class FSM: def __init__(self, conf: ConfigProtocol) -> None: diff --git a/src/drunc/fsm/utils.py b/src/drunc/fsm/utils.py index 4c3aca92c..e2ddf6a8f 100644 --- a/src/drunc/fsm/utils.py +++ b/src/drunc/fsm/utils.py @@ -6,22 +6,8 @@ FSMCommandsDescription, ) from druncschema.generic_pb2 import bool_msg, float_msg, int_msg, string_msg -from google.protobuf.any_pb2 import Any -import drunc.fsm.exceptions as fsme -from drunc.fsm.transition import Transition -from drunc.utils.grpc_utils import unpack_any - -from typing import MutableMapping, Optional - -from druncschema.controller_pb2 import ( - Argument, - FSMCommandDescription, - FSMCommandsDescription, -) -from druncschema.generic_pb2 import bool_msg, float_msg, int_msg, string_msg - -# Alias Protobuf's Any to PbAny to prevent confusion with typing.Any +# prevent confusion with typing.Any from google.protobuf.any_pb2 import Any as PbAny import drunc.fsm.exceptions as fsme @@ -109,6 +95,6 @@ def get_argument(name: str, args: MutableMapping[str, PbAny]) -> Optional[PbAny] case Argument.Type.BOOL: out_dict[arg.name] = unpack_any(arg_value, bool_msg).value case _: - raise fsme.UnhandledArgumentType(arg.type) + raise fsme.UnhandledArgumentType(str(arg.type)) return out_dict \ No newline at end of file From 26c6606d3ba3b4e243ce4856dc8e56c56fabe6bf Mon Sep 17 00:00:00 2001 From: Miruna Serian Date: Tue, 9 Jun 2026 19:37:24 +0100 Subject: [PATCH 03/20] Fix circular import error --- src/drunc/fsm/_protocols.py | 5 +++-- src/drunc/fsm/action_factory.py | 11 ++++++----- src/drunc/fsm/actions/file_logbook.py | 3 +-- src/drunc/fsm/actions/file_run_registry.py | 2 +- src/drunc/fsm/actions/thread_pinning.py | 6 ++---- .../fsm/actions/timing/master_send_fl_command.py | 3 ++- src/drunc/fsm/actions/usvc_provided_run_number.py | 3 ++- src/drunc/fsm/core.py | 12 +++++++++--- 8 files changed, 26 insertions(+), 19 deletions(-) diff --git a/src/drunc/fsm/_protocols.py b/src/drunc/fsm/_protocols.py index 6cac35532..5429ab8a3 100644 --- a/src/drunc/fsm/_protocols.py +++ b/src/drunc/fsm/_protocols.py @@ -1,8 +1,9 @@ -from typing import Dict, Iterable, List, Optional, Protocol +from typing import TYPE_CHECKING, Dict, Iterable, List, Optional, Protocol from druncschema.controller_pb2 import FSMSequence -from drunc.fsm.core import PreOrPostTransitionSequence +if TYPE_CHECKING: + from drunc.fsm.core import PreOrPostTransitionSequence from drunc.fsm.transition import Transition diff --git a/src/drunc/fsm/action_factory.py b/src/drunc/fsm/action_factory.py index 805c78ec0..50fe20814 100644 --- a/src/drunc/fsm/action_factory.py +++ b/src/drunc/fsm/action_factory.py @@ -1,12 +1,15 @@ from __future__ import annotations import inspect -from typing import TYPE_CHECKING, Dict, Protocol, Type, cast - -import conffwk +from typing import Dict, Type, cast import drunc.fsm.exceptions as fsme from drunc.exceptions import DruncSetupException +from drunc.fsm._protocols import ( + ActionMethodProtocol, + ConfigurationProtocol, + FSMActionProtocol, +) from drunc.fsm.actions.db_run_registry import DBRunRegistry from drunc.fsm.actions.file_logbook import FileLogbook from drunc.fsm.actions.file_run_registry import FileRunRegistry @@ -18,8 +21,6 @@ from drunc.fsm.actions.usvc_elisa_logbook import ElisaLogbook from drunc.fsm.actions.usvc_provided_run_number import UsvcProvidedRunNumber -from drunc.fsm._protocols import ActionMethodProtocol, ConfigurationProtocol, FSMActionProtocol - class FSMActionFactory: _instance: FSMActionFactory | None = None diff --git a/src/drunc/fsm/actions/file_logbook.py b/src/drunc/fsm/actions/file_logbook.py index 3bab1e26a..032e7a254 100644 --- a/src/drunc/fsm/actions/file_logbook.py +++ b/src/drunc/fsm/actions/file_logbook.py @@ -1,7 +1,6 @@ from typing import Optional -from drunc.fsm._protocols import ContextProtocol, ConfigurationProtocol - +from drunc.fsm._protocols import ConfigurationProtocol, ContextProtocol from drunc.fsm.core import FSMAction from drunc.utils.utils import now_str diff --git a/src/drunc/fsm/actions/file_run_registry.py b/src/drunc/fsm/actions/file_run_registry.py index c3c739615..6191487c5 100644 --- a/src/drunc/fsm/actions/file_run_registry.py +++ b/src/drunc/fsm/actions/file_run_registry.py @@ -1,5 +1,5 @@ import os -from typing import Dict, Protocol +from typing import Dict from daqconf.consolidate import consolidate_db # type: ignore[import-untyped] diff --git a/src/drunc/fsm/actions/thread_pinning.py b/src/drunc/fsm/actions/thread_pinning.py index 8a3dca809..360ddf043 100644 --- a/src/drunc/fsm/actions/thread_pinning.py +++ b/src/drunc/fsm/actions/thread_pinning.py @@ -1,13 +1,11 @@ import getpass -from asyncio import Protocol from os import environ -from typing import Iterable, Optional -import conffwk -from drunc.fsm._protocols import ContextProtocol, InitConfigurationProtocol +import conffwk # ignore[import-untyped] from sh import Command, ErrorReturnCode # type: ignore[import-untyped] from drunc.exceptions import DruncSetupException +from drunc.fsm._protocols import ContextProtocol, InitConfigurationProtocol from drunc.fsm.core import FSMAction from drunc.fsm.exceptions import ThreadPinningFailed from drunc.process_manager.oks_parser import collect_apps diff --git a/src/drunc/fsm/actions/timing/master_send_fl_command.py b/src/drunc/fsm/actions/timing/master_send_fl_command.py index 524034738..ee7858316 100644 --- a/src/drunc/fsm/actions/timing/master_send_fl_command.py +++ b/src/drunc/fsm/actions/timing/master_send_fl_command.py @@ -1,5 +1,6 @@ -from drunc.fsm.core import FSMAction from drunc.fsm._protocols import ContextProtocol +from drunc.fsm.core import FSMAction + class MasterSendFLCommand(FSMAction): def __init__(self, configuration: object) -> None: diff --git a/src/drunc/fsm/actions/usvc_provided_run_number.py b/src/drunc/fsm/actions/usvc_provided_run_number.py index 335c269c5..d63ef072b 100644 --- a/src/drunc/fsm/actions/usvc_provided_run_number.py +++ b/src/drunc/fsm/actions/usvc_provided_run_number.py @@ -3,11 +3,12 @@ import requests +from drunc.fsm._protocols import ContextProtocol from drunc.fsm.actions.utils import get_dotdrunc_json, validate_run_type from drunc.fsm.core import FSMAction from drunc.fsm.exceptions import CannotGetRunNumber, DotDruncJsonIncorrectFormat from drunc.utils.utils import get_logger -from drunc.fsm._protocols import ContextProtocol + class UsvcProvidedRunNumber(FSMAction): def __init__(self, configuration: object) -> None: diff --git a/src/drunc/fsm/core.py b/src/drunc/fsm/core.py index 5c37ebb46..0d9b74b13 100644 --- a/src/drunc/fsm/core.py +++ b/src/drunc/fsm/core.py @@ -1,7 +1,14 @@ from __future__ import annotations -from typing import Optional, Union, Dict, List, Protocol -from drunc.fsm._protocols import ActionMethodProtocol, FSMActionProtocol, ContextProtocol, ConfigProtocol +from typing import TYPE_CHECKING, Dict, Optional, Union + +if TYPE_CHECKING: + from drunc.fsm._protocols import ( + ActionMethodProtocol, + ConfigProtocol, + ContextProtocol, + FSMActionProtocol, + ) # Define the abcs first to avoid circular imports class FSMAction: @@ -23,7 +30,6 @@ def __init__(self, method: ActionMethodProtocol, mandatory: bool = True) -> None from enum import Enum from inspect import Parameter, signature -import conffwk from druncschema.controller_pb2 import Argument, FSMSequence from druncschema.generic_pb2 import bool_msg, float_msg, int_msg, string_msg from google.protobuf import any_pb2 From efe2ad849d814c165941411a4277036a893fef04 Mon Sep 17 00:00:00 2001 From: Miruna Serian Date: Wed, 10 Jun 2026 09:17:29 +0100 Subject: [PATCH 04/20] Add stubs and fix circular import --- pyproject.toml | 3 +++ src/drunc/fsm/_protocols.py | 14 +++++++++++--- src/drunc/fsm/actions/db_run_registry.py | 6 +++--- src/drunc/fsm/actions/file_run_registry.py | 2 +- src/drunc/fsm/actions/thread_pinning.py | 2 +- src/drunc/fsm/configuration.py | 2 +- src/drunc/fsm/core.py | 1 + src/drunc/fsm/utils.py | 2 +- typings/conffwk/__init__.py | 12 ++++++++++++ typings/daqconf/__init__.py | 0 typings/daqconf/consolidate.pyi | 7 +++++++ typings/daqconf/jsonify.pyi | 4 ++++ typings/daqconf/validate.pyi | 4 ++++ 13 files changed, 49 insertions(+), 10 deletions(-) create mode 100644 typings/conffwk/__init__.py create mode 100644 typings/daqconf/__init__.py create mode 100644 typings/daqconf/consolidate.pyi create mode 100644 typings/daqconf/jsonify.pyi create mode 100644 typings/daqconf/validate.pyi diff --git a/pyproject.toml b/pyproject.toml index ff6a6af8b..ab0ec2264 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -107,6 +107,9 @@ disallow_any_unimported = true warn_unused_configs = true show_error_codes = true +# Tell mypy where to find stubs +mypy_path = "src:typings" + # These overrides are because the library stubs dont exist and not on typeshed [[tool.mypy.overrides]] module = "google.rpc.*,conffwk.*" diff --git a/src/drunc/fsm/_protocols.py b/src/drunc/fsm/_protocols.py index 5429ab8a3..81f65740d 100644 --- a/src/drunc/fsm/_protocols.py +++ b/src/drunc/fsm/_protocols.py @@ -1,16 +1,19 @@ -from typing import TYPE_CHECKING, Dict, Iterable, List, Optional, Protocol +from __future__ import annotations -from druncschema.controller_pb2 import FSMSequence +from typing import TYPE_CHECKING, Dict, Iterable, List, Optional, Protocol if TYPE_CHECKING: + from druncschema.controller_pb2 import FSMSequence + from drunc.fsm.core import PreOrPostTransitionSequence -from drunc.fsm.transition import Transition + from drunc.fsm.transition import Transition class ParameterProtocol(Protocol): name: str value: str + class InitConfigurationProtocol(Protocol): parameters: Iterable[ParameterProtocol] @@ -22,15 +25,19 @@ def get_user_name(self) -> str: ... class DetConfigProtocol(Protocol): id: str + class DALProtocol(Protocol): detector_configuration: DetConfigProtocol + class DBProtocol(Protocol): def get_dal(self, class_name: str, uid: str) -> DALProtocol: ... + class OksKeyProtocol(Protocol): session: str + class RuntimeConfigurationProtocol(Protocol): initial_data: str oks_key: OksKeyProtocol @@ -64,6 +71,7 @@ class ShErrorProtocol(Protocol): stdout: bytes stderr: bytes + class ActionMethodProtocol(Protocol): __name__: str __module__: str diff --git a/src/drunc/fsm/actions/db_run_registry.py b/src/drunc/fsm/actions/db_run_registry.py index ac37d9a92..cc8982e06 100644 --- a/src/drunc/fsm/actions/db_run_registry.py +++ b/src/drunc/fsm/actions/db_run_registry.py @@ -3,9 +3,9 @@ import tempfile import requests -from daqconf.consolidate import consolidate_db # type: ignore[import-untyped] -from daqconf.jsonify import jsonify_xml_data # type: ignore[import-untyped] -from daqconf.validate import validate_session # type: ignore[import-untyped] +from daqconf.consolidate import consolidate_db +from daqconf.jsonify import jsonify_xml_data +from daqconf.validate import validate_session from drunc.fsm._protocols import ContextProtocol from drunc.fsm.actions.utils import get_dotdrunc_json diff --git a/src/drunc/fsm/actions/file_run_registry.py b/src/drunc/fsm/actions/file_run_registry.py index 6191487c5..7a1f76553 100644 --- a/src/drunc/fsm/actions/file_run_registry.py +++ b/src/drunc/fsm/actions/file_run_registry.py @@ -1,7 +1,7 @@ import os from typing import Dict -from daqconf.consolidate import consolidate_db # type: ignore[import-untyped] +from daqconf.consolidate import consolidate_db from drunc.fsm._protocols import ContextProtocol from drunc.fsm.core import FSMAction diff --git a/src/drunc/fsm/actions/thread_pinning.py b/src/drunc/fsm/actions/thread_pinning.py index 360ddf043..9a5f82d35 100644 --- a/src/drunc/fsm/actions/thread_pinning.py +++ b/src/drunc/fsm/actions/thread_pinning.py @@ -1,7 +1,7 @@ import getpass from os import environ -import conffwk # ignore[import-untyped] +import conffwk from sh import Command, ErrorReturnCode # type: ignore[import-untyped] from drunc.exceptions import DruncSetupException diff --git a/src/drunc/fsm/configuration.py b/src/drunc/fsm/configuration.py index e3e7c918b..77fb78318 100644 --- a/src/drunc/fsm/configuration.py +++ b/src/drunc/fsm/configuration.py @@ -89,7 +89,7 @@ def _post_process_oks(self) -> None: # Fill the actions dictionary with the FSMAction objects corresponding to the # action names defined in the configuration - for action in self.data.actions: + for action in self.data.actions: self.actions[action.id] = FSMActionFactory.get().get_action( action.id, action ) diff --git a/src/drunc/fsm/core.py b/src/drunc/fsm/core.py index 0d9b74b13..2e22b9a4c 100644 --- a/src/drunc/fsm/core.py +++ b/src/drunc/fsm/core.py @@ -10,6 +10,7 @@ FSMActionProtocol, ) + # Define the abcs first to avoid circular imports class FSMAction: """Abstract class defining a generic action""" diff --git a/src/drunc/fsm/utils.py b/src/drunc/fsm/utils.py index e2ddf6a8f..8a70951e9 100644 --- a/src/drunc/fsm/utils.py +++ b/src/drunc/fsm/utils.py @@ -73,7 +73,7 @@ def get_argument(name: str, args: MutableMapping[str, PbAny]) -> Optional[PbAny] return k return None - # Explicitly type out_dict so Mypy doesn't infer dict[Any, Any] + # Explicitly type out_dict out_dict: dict[str, str | int | float | bool] = {} for arg in arguments_format: diff --git a/typings/conffwk/__init__.py b/typings/conffwk/__init__.py new file mode 100644 index 000000000..2a158262e --- /dev/null +++ b/typings/conffwk/__init__.py @@ -0,0 +1,12 @@ +from typing import Any + + +class Configuration: + def __init__(self, connection: str = "oksconflibs:") -> None: ... + + def get_dal( + self, + class_name: str, + uid: str, + ) -> Any: ... + \ No newline at end of file diff --git a/typings/daqconf/__init__.py b/typings/daqconf/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/typings/daqconf/consolidate.pyi b/typings/daqconf/consolidate.pyi new file mode 100644 index 000000000..5afc5dadc --- /dev/null +++ b/typings/daqconf/consolidate.pyi @@ -0,0 +1,7 @@ +from typing import Optional + +def consolidate_db( + oksfile: str, + output_file: str, + session_id: Optional[str] = None, +) -> None: ... \ No newline at end of file diff --git a/typings/daqconf/jsonify.pyi b/typings/daqconf/jsonify.pyi new file mode 100644 index 000000000..f05b78fa6 --- /dev/null +++ b/typings/daqconf/jsonify.pyi @@ -0,0 +1,4 @@ +def jsonify_xml_data( + oksfile: str, + output: str, +) -> None: ... \ No newline at end of file diff --git a/typings/daqconf/validate.pyi b/typings/daqconf/validate.pyi new file mode 100644 index 000000000..a96125cf9 --- /dev/null +++ b/typings/daqconf/validate.pyi @@ -0,0 +1,4 @@ +def validate_session( + oksfile: str, + session_name: str, +) -> None: ... \ No newline at end of file From ab76d90cfb991d260f4b11e328404e211beadd34 Mon Sep 17 00:00:00 2001 From: Miruna Serian Date: Tue, 16 Jun 2026 18:47:59 +0100 Subject: [PATCH 05/20] add pybind11-stubgen for conffwk stubs --- clean_stubs.py | 103 +++++ pyproject.toml | 11 +- src/drunc/fsm/_protocols.py | 19 +- src/drunc/fsm/action_factory.py | 15 +- src/drunc/fsm/actions/db_run_registry.py | 25 +- src/drunc/fsm/actions/file_logbook.py | 20 +- src/drunc/fsm/actions/file_run_registry.py | 7 +- src/drunc/fsm/actions/thread_pinning.py | 25 +- .../fsm/actions/trigger_rate_specifier.py | 6 +- .../fsm/actions/user_provided_run_number.py | 2 +- src/drunc/fsm/actions/usvc_elisa_logbook.py | 12 +- src/drunc/fsm/actions/utils.py | 2 +- src/drunc/fsm/configuration.py | 70 ++-- src/drunc/fsm/core.py | 33 +- src/drunc/fsm/exceptions.py | 1 + src/drunc/fsm/transition.py | 15 +- src/drunc/fsm/utils.py | 15 +- src/drunc/utils/grpc_utils.py | 2 +- typings/conffwk/__init__.py | 12 - typings/conffwk/__init__.pyi | 21 + typings/conffwk/_daq_conffwk_py.pyi | 370 ++++++++++++++++++ typings/conffwk/dal.pyi | 52 +++ typings/conffwk/dalproperty.pyi | 12 + typings/conffwk/proxy.pyi | 46 +++ typings/conffwk/schema.pyi | 106 +++++ 25 files changed, 889 insertions(+), 113 deletions(-) create mode 100644 clean_stubs.py delete mode 100644 typings/conffwk/__init__.py create mode 100644 typings/conffwk/__init__.pyi create mode 100644 typings/conffwk/_daq_conffwk_py.pyi create mode 100644 typings/conffwk/dal.pyi create mode 100644 typings/conffwk/dalproperty.pyi create mode 100644 typings/conffwk/proxy.pyi create mode 100644 typings/conffwk/schema.pyi diff --git a/clean_stubs.py b/clean_stubs.py new file mode 100644 index 000000000..05fbc12c5 --- /dev/null +++ b/clean_stubs.py @@ -0,0 +1,103 @@ +""" +Manual overrides for conffwk stubs. + +The pybind11-stubgen auto-generator marks types it can't understand as 'Any' +or misses fields entirely. This adds those classes manually. +""" + +import os +import re +from pathlib import Path + + +def process_file(filepath: Path) -> None: + with open(filepath, 'r', encoding='utf-8') as f: + content = f.read() + + # Fix parameter types + content = re.sub(r': \.\.\.', ': object', content) + + # Fix return types + content = re.sub(r'-> \.\.\.:', '-> object:', content) + + # Fix untyped arguments + content = re.sub(r'= \.\.\.', '= object()', content) + + with open(filepath, 'w', encoding='utf-8') as f: + f.write(content) + + +def add_dal_classes(stubs_dir: Path) -> None: + dal_file = stubs_dir / "dal.pyi" + + dal_classes = """ +from typing import List + +from drunc.fsm._protocols import ( + ConfigurationProtocol, + DBProtocol, + OksKeyProtocol, + ParameterProtocol, +) + +class FSMParameter(ParameterProtocol): + name: str + value: str + +class FSMAction(ConfigurationProtocol): + id: str + name: str + parameters: List[FSMParameter] + db: DBProtocol + oks_key: OksKeyProtocol + initial_data: str + +class FSMxTransition: + transition: str + order: List[str] + mandatory: List[str] + +class FSMTransitionConfig: + id: str + source: str + dest: str + +class FSMCommand: + id: str + +class FSMCommandSequence: + id: str + sequence: List[FSMCommand] + +class FSMData: + states: List[str] + initial_state: str + actions: List[FSMAction] + transitions: List[FSMTransitionConfig] + pre_transitions: List[FSMxTransition] + post_transitions: List[FSMxTransition] + command_sequences: List[FSMCommandSequence] +""" + with open(dal_file, 'w', encoding='utf-8') as f: + f.write(dal_classes) + print("[+] Done. dal.pyi is now strictly typed.") + +def main() -> None: + stubs_dir = Path("typings/conffwk") + + if not stubs_dir.exists(): + print("Typings/conffwk directory not found. Run 'pybind11-stubgen conffwk --output-dir=typings' first.") + return + + # Fix syntax errors in all generated files + for root, _, files in os.walk(stubs_dir): + for file in files: + if file.endswith('.pyi'): + process_file(Path(root) / file) + + # Add missing classes into the DAL file + add_dal_classes(stubs_dir) + + +if __name__ == "__main__": + main() diff --git a/pyproject.toml b/pyproject.toml index ab0ec2264..5b831d130 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -34,7 +34,7 @@ dependencies = [ [project.optional-dependencies] prod = ["paramiko[gssapi]"] -dev = ["ruff", "pre-commit", "pytest", "pytest-cov", "grpcio-testing", "grpcio==1.75", "grpcio-tools==1.75", "grpcio-status==1.75"] +dev = ["ruff", "pre-commit", "pytest", "pytest-cov", "grpcio-testing", "grpcio==1.75", "grpcio-tools==1.75", "grpcio-status==1.75", "pybind11-stubgen"] test = ["pytest", "pytest-cov", "grpcio-testing", "grpcio==1.75", "grpcio-tools==1.75", "grpcio-status==1.75"] [project.scripts] @@ -112,6 +112,13 @@ mypy_path = "src:typings" # These overrides are because the library stubs dont exist and not on typeshed [[tool.mypy.overrides]] -module = "google.rpc.*,conffwk.*" +module = "google.rpc.*" ignore_missing_imports = true +[[tool.mypy.overrides]] +module = "conffwk.*" +disallow_untyped_defs = false +disallow_incomplete_defs = false +check_untyped_defs = false +ignore_missing_imports = true +ignore_errors = true \ No newline at end of file diff --git a/src/drunc/fsm/_protocols.py b/src/drunc/fsm/_protocols.py index 81f65740d..d2431b58a 100644 --- a/src/drunc/fsm/_protocols.py +++ b/src/drunc/fsm/_protocols.py @@ -42,7 +42,7 @@ class RuntimeConfigurationProtocol(Protocol): initial_data: str oks_key: OksKeyProtocol - + class ConfigurationProtocol(Protocol): id: str db: DBProtocol @@ -76,20 +76,31 @@ class ActionMethodProtocol(Protocol): __name__: str __module__: str __self__: object + def __call__(self, *args: object, **kwargs: object) -> object: ... class FSMActionProtocol(Protocol): name: str - + + +class FSMxTransitionProtocol(Protocol): + transition: str + order: list[str] + mandatory: list[str] + class ConfigProtocol(Protocol): def get_initial_state(self) -> str: ... def get_states(self) -> List[str]: ... def get_transitions(self) -> List[Transition]: ... def get_sequences(self) -> List[FSMSequence]: ... - def get_pre_transitions_sequences(self) -> Dict[Transition, PreOrPostTransitionSequence]: ... - def get_post_transitions_sequences(self) -> Dict[Transition, PreOrPostTransitionSequence]: ... + def get_pre_transitions_sequences( + self, + ) -> Dict[Transition, PreOrPostTransitionSequence]: ... + def get_post_transitions_sequences( + self, + ) -> Dict[Transition, PreOrPostTransitionSequence]: ... class ActionMethod(Protocol): diff --git a/src/drunc/fsm/action_factory.py b/src/drunc/fsm/action_factory.py index 50fe20814..d4024f075 100644 --- a/src/drunc/fsm/action_factory.py +++ b/src/drunc/fsm/action_factory.py @@ -25,10 +25,12 @@ class FSMActionFactory: _instance: FSMActionFactory | None = None - def __init__(self) -> None: + def __init__(self) -> None: raise DruncSetupException("Call get() instead") - def _get_pre_transitions(self, action: FSMActionProtocol) -> Dict[str, ActionMethodProtocol]: + def _get_pre_transitions( + self, action: FSMActionProtocol + ) -> Dict[str, ActionMethodProtocol]: retr: Dict[str, ActionMethodProtocol] = {} for name, method in inspect.getmembers(action): if inspect.ismethod(method): @@ -36,7 +38,9 @@ def _get_pre_transitions(self, action: FSMActionProtocol) -> Dict[str, ActionMet retr[name] = cast(ActionMethodProtocol, method) return retr - def _get_post_transitions(self, action: FSMActionProtocol) -> Dict[str, ActionMethodProtocol]: + def _get_post_transitions( + self, action: FSMActionProtocol + ) -> Dict[str, ActionMethodProtocol]: retr: Dict[str, ActionMethodProtocol] = {} for name, method in inspect.getmembers(action): if inspect.ismethod(method): @@ -44,7 +48,9 @@ def _get_post_transitions(self, action: FSMActionProtocol) -> Dict[str, ActionMe retr[name] = cast(ActionMethodProtocol, method) return retr - def _validate_signature(self, name:str , method: ActionMethodProtocol, action: str) -> None: + def _validate_signature( + self, name: str, method: ActionMethodProtocol, action: str + ) -> None: sig = inspect.signature(method) if ( @@ -124,7 +130,6 @@ def get_action( return iface - @classmethod def get(cls: Type[FSMActionFactory]) -> FSMActionFactory: if cls._instance is None: diff --git a/src/drunc/fsm/actions/db_run_registry.py b/src/drunc/fsm/actions/db_run_registry.py index cc8982e06..3767b07b7 100644 --- a/src/drunc/fsm/actions/db_run_registry.py +++ b/src/drunc/fsm/actions/db_run_registry.py @@ -38,10 +38,10 @@ def __init__(self, configuration: object) -> None: self.timeout = 2 def pre_start( - self, - _input_data: dict[str, object], - _context: ContextProtocol, - **kwargs: object + self, + _input_data: dict[str, object], + _context: ContextProtocol, + **kwargs: object, ) -> dict[str, object]: """ Upload a copy of the XML and JSON configurations used to take the current run to @@ -63,9 +63,7 @@ def pre_start( # Seems like run_number isn't in _input_data in post_drain_dataflow so need to # initialise it here - self.run_number = _input_data[ - "run" - ] + self.run_number = _input_data["run"] # Get the environment variables for upload software_version = os.getenv("DUNE_DAQ_BASE_RELEASE") @@ -125,10 +123,11 @@ def pre_start( with tarfile.open(fileobj=tarball, mode="w:gz") as tar: tar.add(xml_filename, arcname=os.path.basename(xml_filename)) tar.add(json_filename, arcname=os.path.basename(json_filename)) - tar.add(entry_point_filename, arcname=os.path.basename(entry_point_filename)) + tar.add( + entry_point_filename, arcname=os.path.basename(entry_point_filename) + ) # f_tar.close() - # Publish to the run registry with open(tarball_name, "rb") as f: files = {"file": f} @@ -169,10 +168,10 @@ def pre_start( return _input_data def post_drain_dataflow( - self, - _input_data: dict[str, object], - _context: ContextProtocol, - **kwargs: object + self, + _input_data: dict[str, object], + _context: ContextProtocol, + **kwargs: object, ) -> None: try: requests.get( diff --git a/src/drunc/fsm/actions/file_logbook.py b/src/drunc/fsm/actions/file_logbook.py index 032e7a254..f776160e6 100644 --- a/src/drunc/fsm/actions/file_logbook.py +++ b/src/drunc/fsm/actions/file_logbook.py @@ -12,11 +12,11 @@ def __init__(self, configuration: ConfigurationProtocol) -> None: self.file = self.conf_dict["file_name"] def post_start( - self, - _input_data: dict[str, object], - _context: ContextProtocol, - file_logbook_post: Optional[str] = None, - **kwargs: object + self, + _input_data: dict[str, object], + _context: ContextProtocol, + file_logbook_post: Optional[str] = None, + **kwargs: object, ) -> dict[str, object]: with open(self.file, "a") as f: f.write( @@ -29,11 +29,11 @@ def post_start( return _input_data def post_drain_dataflow( - self, - _input_data: dict[str, object], - _context: ContextProtocol, - file_logbook_post: str = "", - **kwargs: object + self, + _input_data: dict[str, object], + _context: ContextProtocol, + file_logbook_post: str = "", + **kwargs: object, ) -> dict[str, object]: with open(self.file, "a") as f: f.write( diff --git a/src/drunc/fsm/actions/file_run_registry.py b/src/drunc/fsm/actions/file_run_registry.py index 7a1f76553..b5b7cd145 100644 --- a/src/drunc/fsm/actions/file_run_registry.py +++ b/src/drunc/fsm/actions/file_run_registry.py @@ -12,7 +12,12 @@ def __init__(self, configuration: object) -> None: super().__init__(name="file-run-registry") self.configuration = configuration - def pre_start(self, _input_data: Dict[str, object], _context: ContextProtocol, **kwargs: object) -> Dict[str, object]: + def pre_start( + self, + _input_data: Dict[str, object], + _context: ContextProtocol, + **kwargs: object, + ) -> Dict[str, object]: run_number = _input_data["run"] dest = os.getcwd() + "/run_conf" + str(run_number) + ".data.xml" consolidate_db(_context.configuration.initial_data.split(":")[1], f"{dest}") diff --git a/src/drunc/fsm/actions/thread_pinning.py b/src/drunc/fsm/actions/thread_pinning.py index 9a5f82d35..0dd5323d1 100644 --- a/src/drunc/fsm/actions/thread_pinning.py +++ b/src/drunc/fsm/actions/thread_pinning.py @@ -19,7 +19,9 @@ def __init__(self, configuration: InitConfigurationProtocol) -> None: self.log = get_logger("controller.iface.thread-pinning") self.conf_dict = {p.name: p.value for p in configuration.parameters} - def pin_thread(self, thread_pinning_file: str, configuration: str, session: str) -> None: + def pin_thread( + self, thread_pinning_file: str, configuration: str, session: str + ) -> None: db = conffwk.Configuration(configuration) session_dal = db.get_dal(class_name="Session", uid=session) @@ -87,7 +89,12 @@ def pin_thread(self, thread_pinning_file: str, configuration: str, session: str) if failed_hosts: raise ThreadPinningFailed(failed_hosts_error_str) - def post_conf(self, _input_data: dict[str, object], _context: ContextProtocol, **kwargs: object) -> dict[str, object]: + def post_conf( + self, + _input_data: dict[str, object], + _context: ContextProtocol, + **kwargs: object, + ) -> dict[str, object]: if "post_conf" in self.conf_dict: self.pin_thread( self.conf_dict["post_conf"], @@ -96,7 +103,12 @@ def post_conf(self, _input_data: dict[str, object], _context: ContextProtocol, * ) return _input_data - def post_start(self, _input_data: dict[str, object], _context: ContextProtocol, **kwargs: object) -> dict[str, object]: + def post_start( + self, + _input_data: dict[str, object], + _context: ContextProtocol, + **kwargs: object, + ) -> dict[str, object]: if "post_start" in self.conf_dict: self.pin_thread( self.conf_dict["post_start"], @@ -105,7 +117,12 @@ def post_start(self, _input_data: dict[str, object], _context: ContextProtocol, ) return _input_data - def pre_conf(self, _input_data: dict[str, object], _context: ContextProtocol, **kwargs: object) -> dict[str, object]: + def pre_conf( + self, + _input_data: dict[str, object], + _context: ContextProtocol, + **kwargs: object, + ) -> dict[str, object]: if "pre_conf" in self.conf_dict: self.pin_thread( self.conf_dict["pre_conf"], diff --git a/src/drunc/fsm/actions/trigger_rate_specifier.py b/src/drunc/fsm/actions/trigger_rate_specifier.py index 19c0ada03..cd05cf10a 100644 --- a/src/drunc/fsm/actions/trigger_rate_specifier.py +++ b/src/drunc/fsm/actions/trigger_rate_specifier.py @@ -7,7 +7,11 @@ def __init__(self, configuration: ConfigurationProtocol) -> None: super().__init__(name="trigger-rate-specifier") def pre_change_rate( - self, _input_data: dict[str, object], _context: ContextProtocol, trigger_rate: float, **kwargs: object + self, + _input_data: dict[str, object], + _context: ContextProtocol, + trigger_rate: float, + **kwargs: object, ) -> dict[str, object]: _input_data["trigger_rate"] = trigger_rate return _input_data diff --git a/src/drunc/fsm/actions/user_provided_run_number.py b/src/drunc/fsm/actions/user_provided_run_number.py index 3e4aca02f..dba128c22 100644 --- a/src/drunc/fsm/actions/user_provided_run_number.py +++ b/src/drunc/fsm/actions/user_provided_run_number.py @@ -19,7 +19,7 @@ def pre_start( trigger_rate: Optional[float] = None, **kwargs: object, ) -> dict[str, object]: - + safe_run_type = run_type.upper() if run_type is not None else "TEST" run_type = validate_run_type(safe_run_type) _input_data["production_vs_test"] = run_type diff --git a/src/drunc/fsm/actions/usvc_elisa_logbook.py b/src/drunc/fsm/actions/usvc_elisa_logbook.py index 207e57e99..2b2d747d5 100644 --- a/src/drunc/fsm/actions/usvc_elisa_logbook.py +++ b/src/drunc/fsm/actions/usvc_elisa_logbook.py @@ -83,7 +83,11 @@ def __init__(self) -> None: self.timeout = 5 def post_start( - self, _input_data: dict[str, object], _context: ContextProtocol, elisa_post: Optional[str] = None, **kwargs: object + self, + _input_data: dict[str, object], + _context: ContextProtocol, + elisa_post: Optional[str] = None, + **kwargs: object, ) -> dict[str, object]: if self.elisa_hardware in self.no_publish_hardware: return {} @@ -142,7 +146,11 @@ def post_start( return _input_data def post_drain_dataflow( - self, _input_data: dict[str, object], _context: ContextProtocol, elisa_post: Optional[str] = None, **kwargs: object + self, + _input_data: dict[str, object], + _context: ContextProtocol, + elisa_post: Optional[str] = None, + **kwargs: object, ) -> dict[str, object]: if self.elisa_hardware in self.no_publish_hardware: return {} diff --git a/src/drunc/fsm/actions/utils.py b/src/drunc/fsm/actions/utils.py index 5f9a1ac59..e9e3efe52 100644 --- a/src/drunc/fsm/actions/utils.py +++ b/src/drunc/fsm/actions/utils.py @@ -24,7 +24,7 @@ def validate_run_type(run_type: str) -> str: return run_type -def get_dotdrunc_json(path: str | None = None): # type: ignore[no-untyped-def] +def get_dotdrunc_json(path: str | None = None): # type: ignore[no-untyped-def] # Resolution order: DOTDRUNC env var -> provided path -> default path file_path = os.getenv("DOTDRUNC") or path or "~/.drunc.json" try: diff --git a/src/drunc/fsm/configuration.py b/src/drunc/fsm/configuration.py index 77fb78318..bb54ee4a3 100644 --- a/src/drunc/fsm/configuration.py +++ b/src/drunc/fsm/configuration.py @@ -1,34 +1,26 @@ -import logging -from typing import Dict, List +from __future__ import annotations +from typing import List + +import conffwk from druncschema.controller_pb2 import FSMSequence -from drunc.fsm._protocols import ConfigurationProtocol -from drunc.fsm.action_factory import FSMActionFactory, FSMActionProtocol +from drunc.fsm._protocols import FSMActionProtocol +from drunc.fsm.action_factory import FSMActionFactory from drunc.fsm.core import PreOrPostTransitionSequence from drunc.fsm.transition import Transition -from drunc.utils.configuration import ConfHandler, OKSKey +from drunc.utils.configuration import ConfHandler from drunc.utils.utils import get_logger class FSMConfHandler(ConfHandler): + data: conffwk.dal.FSMData - def __init__(self, key: OKSKey) -> None: - super().__init__(key) - self.log: logging.Logger = get_logger("controller.core.FSMConfHandler") - self.pre_transitions: Dict[Transition, PreOrPostTransitionSequence] = {} - self.post_transitions: Dict[Transition, PreOrPostTransitionSequence] = {} - self.actions: Dict[str, FSMActionProtocol] = {} # (e.g. "thread_pinning": thread_pinning FSM action object) - self.transitions: List[Transition] = [] - self.sequences: List[FSMSequence] = [] - self.states: List[str] = [] - self.initial_state: str = "" - def _fill_pre_post_transition_sequence_oks( self, prefix: str, transition: Transition, - data: List[ConfigurationProtocol] | None, + data: List[conffwk.dal.FSMxTransition] | None, ) -> PreOrPostTransitionSequence: """ Fill the pre or post transition sequence for a given transition. @@ -57,12 +49,11 @@ def _fill_pre_post_transition_sequence_oks( # transition. There is one FSMxTransition per transition. The FSMxTransition # contains the list of actions to execute for the given transition. for fsm_x_transition in data: - if getattr(fsm_x_transition, "transition", None) == transition.name: - for action_name in getattr(fsm_x_transition, "order", []): - mandatory_list = getattr(fsm_x_transition, "mandatory", []) + if fsm_x_transition.transition == transition.name: + for action_name in fsm_x_transition.order: seq.add_callback( action=self.actions[action_name], - mandatory=action_name in mandatory_list, + mandatory=action_name in fsm_x_transition.mandatory, ) break return seq @@ -86,15 +77,26 @@ def _post_process_oks(self) -> None: Raises: None """ + # Define the data structures to store the FSM configuration + self.log.debug("_post_process_oks configuration") + self.pre_transitions: dict[Transition, PreOrPostTransitionSequence] = {} + self.post_transitions: dict[Transition, PreOrPostTransitionSequence] = {} + self.actions: dict[ + str, FSMActionProtocol + ] = {} # (e.g. "thread_pinning": thread_pinning FSM action object) + self.transitions: list[Transition] = [] + self.sequences: list[FSMSequence] = [] + self.states: list[str] = self.data.states + self.initial_state: str = self.data.initial_state # Fill the actions dictionary with the FSMAction objects corresponding to the # action names defined in the configuration - for action in self.data.actions: + for action in self.data.actions: self.actions[action.id] = FSMActionFactory.get().get_action( action.id, action ) - for transition in self.data.transitions: + for transition in self.data.transitions: tr = Transition( name=transition.id, source=transition.source, @@ -114,8 +116,8 @@ def _post_process_oks(self) -> None: ) ) + # Add the pre and post transition sequence arguments to the transition if tr.arguments is not None: - # Add the pre and post transition sequence arguments to the transition tr.arguments += pre_transitions.get_arguments() tr.arguments += post_transitions.get_arguments() @@ -132,23 +134,27 @@ def _post_process_oks(self) -> None: cmd_ids = [cmd.id for cmd in sequence.sequence] self.sequences.append(FSMSequence(id=seq_id, command_ids=cmd_ids)) - def get_actions(self) -> Dict[str, FSMActionProtocol]: + def get_actions(self) -> dict[str, FSMActionProtocol]: return self.actions def get_initial_state(self) -> str: - return str(self.data.initial_state) + return self.data.initial_state - def get_states(self) -> List[str]: - return list(self.data.states) + def get_states(self) -> list[str]: + return self.data.states - def get_transitions(self) -> List[Transition]: + def get_transitions(self) -> list[Transition]: return self.transitions - def get_pre_transitions_sequences(self) -> Dict[Transition, PreOrPostTransitionSequence]: + def get_pre_transitions_sequences( + self, + ) -> dict[Transition, PreOrPostTransitionSequence]: return self.pre_transitions - def get_post_transitions_sequences(self) -> Dict[Transition, PreOrPostTransitionSequence]: + def get_post_transitions_sequences( + self, + ) -> dict[Transition, PreOrPostTransitionSequence]: return self.post_transitions - def get_sequences(self) -> List[FSMSequence] : + def get_sequences(self) -> list[FSMSequence]: return self.sequences diff --git a/src/drunc/fsm/core.py b/src/drunc/fsm/core.py index 2e22b9a4c..e9611ce7a 100644 --- a/src/drunc/fsm/core.py +++ b/src/drunc/fsm/core.py @@ -55,9 +55,7 @@ def __init__(self, transition: Transition, pre_or_post: str = "pre") -> None: self.sequence: list[Callback] = [] self.log = get_logger("controller.core.PreOrPostTransitionSequence") - def add_callback( - self, action: FSMActionProtocol, mandatory: bool = True - ) -> None: + def add_callback(self, action: FSMActionProtocol, mandatory: bool = True) -> None: """ Add a callback to the sequence. The method to be called will be determined by the name of the transition and the prefix (pre or post). @@ -102,7 +100,12 @@ def __str__(self) -> str: ] ) - def execute(self, transition_data: str | None, transition_args: Dict[str, object], ctx: ContextProtocol) -> str: + def execute( + self, + transition_data: str | None, + transition_args: Dict[str, object], + ctx: ContextProtocol, + ) -> str: self.log.debug(f"{transition_data=}, {transition_args=}") if not transition_data: transition_data = "{}" @@ -265,7 +268,7 @@ class FSMDestinationResult: destination_state: str | None destination_type: FSMDestinationType - + class FSM: def __init__(self, conf: ConfigProtocol) -> None: self.log = get_logger("controller.core.FSM") @@ -310,10 +313,14 @@ def get_all_sequences(self) -> list[FSMSequence]: """Grab all the transitions""" return self.sequences - def is_destination_of_this_transition(self, state: str, transition: Transition) -> bool: + def is_destination_of_this_transition( + self, state: str, transition: Transition + ) -> bool: return bool(transition.destination == state) - def get_destination_state(self, source_state: str, transition: Transition) -> FSMDestinationResult: + def get_destination_state( + self, source_state: str, transition: Transition + ) -> FSMDestinationResult: """Tells us where a particular transition will take us, given the source_state""" right_name = [t for t in self.transitions if t == transition] @@ -391,7 +398,11 @@ def can_execute_transition(self, source_state: str, transition: Transition) -> b return bool(regex_match(transition.source, source_state)) def prepare_transition( - self, transition: Transition, transition_data: str, transition_args: Dict[str, object], ctx: ContextProtocol + self, + transition: Transition, + transition_data: str, + transition_args: Dict[str, object], + ctx: ContextProtocol, ) -> str: transition_data = self.pre_transition_sequences[transition].execute( transition_data, transition_args, ctx @@ -399,7 +410,11 @@ def prepare_transition( return transition_data def finalise_transition( - self, transition: Transition, transition_data: str, transition_args: Dict[str, object], ctx: ContextProtocol + self, + transition: Transition, + transition_data: str, + transition_args: Dict[str, object], + ctx: ContextProtocol, ) -> str: transition_data = self.post_transition_sequences[transition].execute( transition_data, transition_args, ctx diff --git a/src/drunc/fsm/exceptions.py b/src/drunc/fsm/exceptions.py index a7e113d86..acffab9ee 100644 --- a/src/drunc/fsm/exceptions.py +++ b/src/drunc/fsm/exceptions.py @@ -118,6 +118,7 @@ def __init__(self, data: object) -> None: self.message = f"Could not get Run Number because {data}" super().__init__(self.message) + class DBRunRegistryConfigurationError(FSMException): def __init__(self, data: object) -> None: self.message = f"RunRegistryDB configuration error: {data}" diff --git a/src/drunc/fsm/transition.py b/src/drunc/fsm/transition.py index 15b70f3af..cae2a6ab1 100644 --- a/src/drunc/fsm/transition.py +++ b/src/drunc/fsm/transition.py @@ -6,20 +6,21 @@ class Transition: - def __init__(self, - name: str, - source: str, - destination: str, + def __init__( + self, + name: str, + source: str, + destination: str, arguments: Optional[List[Argument]] = None, - help: str = "" - ) -> None: + help: str = "", + ) -> None: self.source = source self.destination = destination self.name = name self.arguments = arguments self.help = help - def __eq__(self, another: object) -> bool: + def __eq__(self, another: object) -> bool: same_name = hasattr(another, "name") and self.name == another.name same_destination = ( hasattr(another, "destination") and self.destination == another.destination diff --git a/src/drunc/fsm/utils.py b/src/drunc/fsm/utils.py index 8a70951e9..80f7e26d6 100644 --- a/src/drunc/fsm/utils.py +++ b/src/drunc/fsm/utils.py @@ -23,9 +23,9 @@ def convert_fsm_transition(transitions: list[Transition]) -> FSMCommandsDescript transitions (list[Transition]): A list of FSM transitions. Returns: - FSMCommandsDescription: A FSMCommandsDescription containing the converted + FSMCommandsDescription: A FSMCommandsDescription containing the converted transitions. - + Raises: None. """ @@ -44,13 +44,12 @@ def convert_fsm_transition(transitions: list[Transition]) -> FSMCommandsDescript def decode_fsm_arguments( - arguments: MutableMapping[str, PbAny], - arguments_format: list[Argument] + arguments: MutableMapping[str, PbAny], arguments_format: list[Argument] ) -> dict[str, str | int | float | bool]: """ Decodes the arguments of a FSM command. - Note there is separate logic to validate whether the required arguments are all + Note there is separate logic to validate whether the required arguments are all present at the click.core.Command level, this is a safeguard to support the multiple drunc operating modes. @@ -75,7 +74,7 @@ def get_argument(name: str, args: MutableMapping[str, PbAny]) -> Optional[PbAny] # Explicitly type out_dict out_dict: dict[str, str | int | float | bool] = {} - + for arg in arguments_format: arg_value: Optional[PbAny] = get_argument(arg.name, arguments) @@ -96,5 +95,5 @@ def get_argument(name: str, args: MutableMapping[str, PbAny]) -> Optional[PbAny] out_dict[arg.name] = unpack_any(arg_value, bool_msg).value case _: raise fsme.UnhandledArgumentType(str(arg.type)) - - return out_dict \ No newline at end of file + + return out_dict diff --git a/src/drunc/utils/grpc_utils.py b/src/drunc/utils/grpc_utils.py index 6ecedc48c..426ca0070 100644 --- a/src/drunc/utils/grpc_utils.py +++ b/src/drunc/utils/grpc_utils.py @@ -318,7 +318,7 @@ def extract_grpc_rich_error(grpc_error: grpc.RpcError) -> GrpcErrorDetails: # Fallback to simple error if no rich status if status is None: - return GrpcErrorDetails(code=code, message="No message", details=[]) + return GrpcErrorDetails(code=code, message="No message ", details=[]) # Extract all error details error_details = [] diff --git a/typings/conffwk/__init__.py b/typings/conffwk/__init__.py deleted file mode 100644 index 2a158262e..000000000 --- a/typings/conffwk/__init__.py +++ /dev/null @@ -1,12 +0,0 @@ -from typing import Any - - -class Configuration: - def __init__(self, connection: str = "oksconflibs:") -> None: ... - - def get_dal( - self, - class_name: str, - uid: str, - ) -> Any: ... - \ No newline at end of file diff --git a/typings/conffwk/__init__.pyi b/typings/conffwk/__init__.pyi new file mode 100644 index 000000000..14716e12a --- /dev/null +++ b/typings/conffwk/__init__.pyi @@ -0,0 +1,21 @@ +import __future__ +from __future__ import annotations +from conffwk.ConfigObject import ConfigObject +from conffwk.Configuration import Configuration +from . import _daq_conffwk_py +from . import dal +from . import dalproperty +from . import proxy +from . import schema +__all__: list[str] = ['ConfigObject', 'Configuration', 'absolute_import', 'dal', 'dalproperty', 'proxy', 'reset_updated_dals', 'schema', 'updated_dals'] +def reset_updated_dals(): + """ + Reset the set keeping track of modified DAL objects + + """ +def updated_dals(): + """ + Returns a set of DAL objects that were modified in this DB session + + """ +absolute_import: __future__._Feature # value = _Feature((2, 5, 0, 'alpha', 1), (3, 0, 0, 'alpha', 0), 262144) diff --git a/typings/conffwk/_daq_conffwk_py.pyi b/typings/conffwk/_daq_conffwk_py.pyi new file mode 100644 index 000000000..28d4c8e4f --- /dev/null +++ b/typings/conffwk/_daq_conffwk_py.pyi @@ -0,0 +1,370 @@ +""" +Python interface to the conffwk package +""" +from __future__ import annotations +import typing +__all__: list[str] = list() +class _ConfigObject: + def UID(self) -> str: + """ + Return object identity + """ + @typing.overload + def __init__(self) -> None: + ... + @typing.overload + def __init__(self, arg0: _ConfigObject) -> None: + ... + def class_name(self) -> str: + """ + Return object's class name + """ + def contained_in(self) -> str: + """ + Return the name of the database file this object belongs to. + """ + def full_name(self) -> str: + """ + Return full object name + """ + def get_bool(self, attr: str) -> bool: + """ + Simple getter function + """ + def get_bool_vec(self, attr: str) -> list[bool]: + """ + Getter function for a list + """ + def get_double(self, attr: str) -> float: + """ + Simple getter function + """ + def get_double_vec(self, attr: str) -> list[float]: + """ + Getter function for a list + """ + def get_float(self, attr: str) -> float: + """ + Simple getter function + """ + def get_float_vec(self, attr: str) -> list[float]: + """ + Getter function for a list + """ + def get_obj(self, attrname: str) -> _ConfigObject: + """ + Get a copy of an object + """ + def get_objs(self, attr: str) -> list[_ConfigObject]: + """ + Getter function for a list + """ + def get_s16(self, attr: str) -> int: + """ + Simple getter function + """ + def get_s16_vec(self, attr: str) -> list[int]: + """ + Getter function for a list + """ + def get_s32(self, attr: str) -> int: + """ + Simple getter function + """ + def get_s32_vec(self, attr: str) -> list[int]: + """ + Getter function for a list + """ + def get_s64(self, attr: str) -> int: + """ + Simple getter function + """ + def get_s64_vec(self, attr: str) -> list[int]: + """ + Getter function for a list + """ + def get_s8(self, attr: str) -> int: + """ + Simple getter function + """ + def get_s8_vec(self, attr: str) -> list[int]: + """ + Getter function for a list + """ + def get_string(self, attr: str) -> str: + """ + Simple getter function + """ + def get_string_vec(self, attr: str) -> list[str]: + """ + Getter function for a list + """ + def get_u16(self, attr: str) -> int: + """ + Simple getter function + """ + def get_u16_vec(self, attr: str) -> list[int]: + """ + Getter function for a list + """ + def get_u32(self, attr: str) -> int: + """ + Simple getter function + """ + def get_u32_vec(self, attr: str) -> list[int]: + """ + Getter function for a list + """ + def get_u64(self, attr: str) -> int: + """ + Simple getter function + """ + def get_u64_vec(self, attr: str) -> list[int]: + """ + Getter function for a list + """ + def get_u8(self, attr: str) -> int: + """ + Simple getter function + """ + def get_u8_vec(self, attr: str) -> list[int]: + """ + Getter function for a list + """ + def rename(self, new_id: str) -> None: + """ + Rename object + """ + def set_bool(self, name: str, value: bool) -> None: + """ + Simple setter function + """ + def set_bool_vec(self, attrname: str, l: list[bool]) -> None: + """ + Setter function for list + """ + def set_class(self, name: str, value: str) -> None: + """ + Set the class name + """ + def set_class_vec(self, attrname: str, l: list[str]) -> None: + """ + Set list of classes + """ + def set_date(self, name: str, value: str) -> None: + """ + Set the date + """ + def set_date_vec(self, attrname: str, l: list[str]) -> None: + """ + Set list of dates + """ + def set_double(self, name: str, value: float) -> None: + """ + Simple setter function + """ + def set_double_vec(self, attrname: str, l: list[float]) -> None: + """ + Setter function for list + """ + def set_enum(self, name: str, value: str) -> None: + """ + Set the enum + """ + def set_enum_vec(self, attrname: str, l: list[str]) -> None: + """ + Set list of enums + """ + def set_float(self, name: str, value: float) -> None: + """ + Simple setter function + """ + def set_float_vec(self, attrname: str, l: list[float]) -> None: + """ + Setter function for list + """ + def set_obj(self, name: str, o: _ConfigObject, skip_non_null_check: bool = False) -> None: + """ + Set relationship single-value + """ + def set_objs(self, name: str, o: list[_ConfigObject], skip_non_null_check: bool = False) -> None: + """ + Set relationship multi-value. + """ + def set_s16(self, name: str, value: int) -> None: + """ + Simple setter function + """ + def set_s16_vec(self, attrname: str, l: list[int]) -> None: + """ + Setter function for list + """ + def set_s32(self, name: str, value: int) -> None: + """ + Simple setter function + """ + def set_s32_vec(self, attrname: str, l: list[int]) -> None: + """ + Setter function for list + """ + def set_s64(self, name: str, value: int) -> None: + """ + Simple setter function + """ + def set_s64_vec(self, attrname: str, l: list[int]) -> None: + """ + Setter function for list + """ + def set_s8(self, name: str, value: int) -> None: + """ + Simple setter function + """ + def set_s8_vec(self, attrname: str, l: list[int]) -> None: + """ + Setter function for list + """ + def set_string(self, name: str, value: str) -> None: + """ + Simple setter function + """ + def set_string_vec(self, attrname: str, l: list[str]) -> None: + """ + Set list of strings + """ + def set_time(self, name: str, value: str) -> None: + """ + Set the time + """ + def set_time_vec(self, attrname: str, l: list[str]) -> None: + """ + Set list of times + """ + def set_u16(self, name: str, value: int) -> None: + """ + Simple setter function + """ + def set_u16_vec(self, attrname: str, l: list[int]) -> None: + """ + Setter function for list + """ + def set_u32(self, name: str, value: int) -> None: + """ + Simple setter function + """ + def set_u32_vec(self, attrname: str, l: list[int]) -> None: + """ + Setter function for list + """ + def set_u64(self, name: str, value: int) -> None: + """ + Simple setter function + """ + def set_u64_vec(self, attrname: str, l: list[int]) -> None: + """ + Setter function for list + """ + def set_u8(self, name: str, value: int) -> None: + """ + Simple setter function + """ + def set_u8_vec(self, attrname: str, l: list[int]) -> None: + """ + Setter function for list + """ +class _Configuration: + @typing.overload + def __init__(self) -> None: + ... + @typing.overload + def __init__(self, arg0: str) -> None: + ... + def add_include(self, db_name: str, include: str) -> None: + """ + Add include file to existing database. + """ + def attributes(self, class_name: str, all: bool) -> dict[str, dict[str, str]]: + """ + Get the properties of each attribute in a given class + """ + def classes(self) -> list[str]: + """ + Get the names of the superclasses for each class + """ + def commit(self, log_message: str = '') -> None: + """ + Commit database changes. + """ + def create_db(self, db_name: str, includes: list[str]) -> None: + """ + Create a database from a list of files + """ + @typing.overload + def create_obj(self, at: str, class_name: str, id: str) -> object: + """ + Create new object by class name and object id. + """ + @typing.overload + def create_obj(self, at: object, class_name: str, id: str) -> object: + """ + Create new object by class name and object id. + """ + def destroy_obj(self, object: object) -> None: + """ + The method tries to destroy given object. + """ + def get_impl_param(self) -> str: + """ + Get implementation plug-in parameter used to build conffwk object + """ + def get_impl_spec(self) -> str: + """ + Get implementation plug-in and its parameter used to build conffwk object + """ + def get_includes(self, db_name: str) -> list[str]: + """ + Returns list of files included by given database. + """ + def get_obj(self, class_name: str, id: str) -> object: + """ + Create a configuration object containing the desired entity from the database + """ + def get_objs(self, class_name: str, query: str = '') -> list[...]: + """ + Create a list of configuration objects of a given class from the database + """ + def get_schema_path(self, class_name: str) -> str: + """ + Get path to schema file with definition of the given class + """ + def load(self, db_name: str) -> None: + """ + Load database according to the name. + """ + def loaded(self) -> bool: + """ + Check if database is correctly loaded. + """ + def relations(self, class_name: str, all: bool) -> dict[str, dict[str, str]]: + """ + Get the properties of each relation in a given class + """ + def remove_include(self, db_name: str, include: str) -> None: + """ + Remove include file. + """ + def subclasses(self, class_name: str, all: bool) -> list[str]: + """ + Get the subclasses of a single class + """ + def superclasses(self, class_name: str, all: bool) -> list[str]: + """ + Get the superclasses of a single class + """ + def test_object(self, class_name: str, id: str, rlevel: int, rclasses: list[str]) -> bool: + """ + Test the existence of the object + """ + def unload(self) -> None: + """ + Unload previously-loaded database + """ diff --git a/typings/conffwk/dal.pyi b/typings/conffwk/dal.pyi new file mode 100644 index 000000000..e14bd856b --- /dev/null +++ b/typings/conffwk/dal.pyi @@ -0,0 +1,52 @@ +from typing import List + +from drunc.fsm._protocols import ( + ConfigurationProtocol, + DBProtocol, + OksKeyProtocol, + ParameterProtocol, +) + +class FSMParameter(ParameterProtocol): + name: str + value: str + + +class FSMAction(ConfigurationProtocol): + id: str + name: str + parameters: List[FSMParameter] + db: DBProtocol + oks_key: OksKeyProtocol + initial_data: str + + +class FSMxTransition: + transition: str + order: List[str] + mandatory: List[str] + + +class FSMTransitionConfig: + id: str + source: str + dest: str + + +class FSMCommand: + id: str + + +class FSMCommandSequence: + id: str + sequence: List[FSMCommand] + + +class FSMData: + states: List[str] + initial_state: str + actions: List[FSMAction] + transitions: List[FSMTransitionConfig] + pre_transitions: List[FSMxTransition] + post_transitions: List[FSMxTransition] + command_sequences: List[FSMCommandSequence] diff --git a/typings/conffwk/dalproperty.pyi b/typings/conffwk/dalproperty.pyi new file mode 100644 index 000000000..9143ae5da --- /dev/null +++ b/typings/conffwk/dalproperty.pyi @@ -0,0 +1,12 @@ +import __future__ +from __future__ import annotations +__all__: list[str] = ['absolute_import'] +def _assign_attribute(attribute): + ... +def _assign_relation(relation): + ... +def _return_attribute(attribute, dalobj = None, value = None): + ... +def _return_relation(relation, multi = False, data = None, dalobj = None, cache = None): + ... +absolute_import: __future__._Feature # value = _Feature((2, 5, 0, 'alpha', 1), (3, 0, 0, 'alpha', 0), 262144) diff --git a/typings/conffwk/proxy.pyi b/typings/conffwk/proxy.pyi new file mode 100644 index 000000000..81469fe94 --- /dev/null +++ b/typings/conffwk/proxy.pyi @@ -0,0 +1,46 @@ +""" +Proxing/Delegation tools + +Provide several tools to implement proxying/delegation of objects. The proxying +instances expose the same public interface of the proxied object, but avoiding +inheritance. This allows to control the reference counts of the proxied object. + +""" +from __future__ import annotations +__all__: list[str] = ['Proxy', 'make_proxy_class'] +class Proxy: + """ + A very basic holder class + + Just holds the reference to a provided object. + + + """ + def __init__(self, obj): + ... +def _DelegateMetaFunction(clsName, bases, atts): + """ + Implements a delegation pattern using a metaclass approach + + A class using this meta mechanism should have 'memberclass' class attribute + initialized at the class of the instance to proxied. The metaclass will + make sure the delegate class will expose all the public methods of the + proxied one. + Moreover, the metaclass will provide the delegate class with a '__init__' + function instantiating a 'memberclass' object, storing it in 'self._obj'. + The delegate class constructor method will therefore accept all the + arguments accepted by the proxied class constructor. + The delegate class uses slots + + + """ +def make_proxy_class(theclass): + """ + Builds a delegation class out of a given type. + + Uses the Proxy class to generate a new proxy class exposing the same + interface of the provided class and delegating the method calls to + the hosted object instance. + + + """ diff --git a/typings/conffwk/schema.pyi b/typings/conffwk/schema.pyi new file mode 100644 index 000000000..390a3765b --- /dev/null +++ b/typings/conffwk/schema.pyi @@ -0,0 +1,106 @@ +""" +A set of utilities to simplify OKS instrospection. +""" +from __future__ import annotations +from conffwk import ConfigObject +import logging as logging +import re as re +import sys as sys +__all__: list[str] = ['Cache', 'ConfigObject', 'check_cardinality', 'check_range', 'check_relation', 'coerce', 'decode_range', 'logging', 'map_coercion', 'oks_types', 'range_regexp', 're', 'str2integer', 'sys', 'to_int', 'to_long'] +class Cache: + """ + Defines a cache for all known schemas at a certain time. + + """ + def __getitem__(self, key): + """ + Gets the description of a certain class. + """ + def __init__(self, conffwk, all = True): + """ + Initializes the cache with information from the Configuration + object. + + This method will browse for all declared classes in the Configuration + object given as input and will setup the schema for all known classes. + After this you can still update the cache using the update() method. + + Keyword parameters: + + conffwk -- The conffwk.Configuration object to use as base for the + current cache. + + all -- A boolean indicating if I should store all the attributes and + relations from a certain class or just the ones directly associated + with a class. + + """ + def __str__(self): + """ + Prints a nice display of myself + """ + def update(self, conffwk): + """ + Updates this cache with information from the Configuration object. + + This method will add new classes not yet know to this cache. Classes + with existing names will not be added. No warning is generated (this + should be done by the OKS layer in any case. + + """ + def update_dal(self, conffwk): + """ + Updates this cache with information for DAL. + + This method will add new DAL classes not yet know to this cache. + Classes with existing DAL representations will not be touched. + + """ +def check_cardinality(v, prop): + """ + Checks the cardinality of a certain attribute or relationship. + """ +def check_range(v, range, range_re, pytype): + """ + Checks the range of the value 'v' to make sure it is inside. + """ +def check_relation(v, rel): + """ + Checks the value v against the relationship parameters in 'rel'. + """ +def coerce(v, attr): + """ + Coerces the input value 'v' in the way the attribute expects. + """ +def decode_range(s): + """ + Decodes a range string representation, returns a tuple with 2 values. + + This is the supported format in regexp representation: + '([-0x]*\\d+)\\D+-?\\d+' + + """ +def map_coercion(class_name, schema): + """ + Given a schema of a class, maps coercion functions from libpyconffwk. + """ +def str2integer(v, t, max): + """ + Converts a value v to integer, irrespectively of its formatting. + + If the number starts with a '0', we convert it using an octal + representation. Else, we try a decimal conversion. If any of these fail, + we try an hexa conversion before throwing a ValueError. + + Keyword arguments: + + v -- the value to be converted + t -- the python type (int or float) to use in the conversion + + """ +def to_int(v): + ... +def to_long(v): + ... +oks_types: dict = {'bool': ['bool'], 'integer': ['s8', 'u8', 's16', 'u16', 's32'], 'long': ['u32', 's64', 'u64'], 'float': ['float', 'double'], 'int-number': ['s8', 'u8', 's16', 'u16', 's32', 'u32', 's64', 'u64'], 'number': ['u32', 's64', 'u64', 's8', 'u8', 's16', 'u16', 's32', 'float', 'double'], 'time': ['date', 'time'], 'string': ['date', 'time', 'string', 'uid', 'enum', 'class']} +range_regexp: re.Pattern # value = re.compile('(?P-?0?x?[\\da-fA-F]+(\\.\\d+)?)-(?P-?0?x?[\\da-fA-F]+(\\.\\d+)?)') From 58b586d7d84dba1baec74b3510824329db91d317 Mon Sep 17 00:00:00 2001 From: Miruna Serian Date: Tue, 16 Jun 2026 19:06:36 +0100 Subject: [PATCH 06/20] fix ruff typings --- typings/conffwk/__init__.pyi | 34 +++++++---- typings/conffwk/_daq_conffwk_py.pyi | 33 ++++++----- typings/conffwk/dal.pyi | 6 -- typings/conffwk/dalproperty.pyi | 23 ++++---- typings/conffwk/proxy.pyi | 48 ++++++++-------- typings/conffwk/schema.pyi | 88 +++++++++++++++++++++-------- typings/daqconf/consolidate.pyi | 2 +- typings/daqconf/jsonify.pyi | 2 +- typings/daqconf/validate.pyi | 2 +- 9 files changed, 148 insertions(+), 90 deletions(-) diff --git a/typings/conffwk/__init__.pyi b/typings/conffwk/__init__.pyi index 14716e12a..a4eb04769 100644 --- a/typings/conffwk/__init__.pyi +++ b/typings/conffwk/__init__.pyi @@ -1,21 +1,35 @@ -import __future__ from __future__ import annotations +import __future__ + from conffwk.ConfigObject import ConfigObject from conffwk.Configuration import Configuration -from . import _daq_conffwk_py -from . import dal -from . import dalproperty -from . import proxy -from . import schema -__all__: list[str] = ['ConfigObject', 'Configuration', 'absolute_import', 'dal', 'dalproperty', 'proxy', 'reset_updated_dals', 'schema', 'updated_dals'] + +from . import dal, dalproperty, proxy, schema + +__all__: list[str] = [ + "ConfigObject", + "Configuration", + "absolute_import", + "dal", + "dalproperty", + "proxy", + "reset_updated_dals", + "schema", + "updated_dals", +] + def reset_updated_dals(): """ Reset the set keeping track of modified DAL objects - + """ + def updated_dals(): """ Returns a set of DAL objects that were modified in this DB session - + """ -absolute_import: __future__._Feature # value = _Feature((2, 5, 0, 'alpha', 1), (3, 0, 0, 'alpha', 0), 262144) + +absolute_import: ( + __future__._Feature +) # value = _Feature((2, 5, 0, 'alpha', 1), (3, 0, 0, 'alpha', 0), 262144) diff --git a/typings/conffwk/_daq_conffwk_py.pyi b/typings/conffwk/_daq_conffwk_py.pyi index 28d4c8e4f..1fcb0a6ca 100644 --- a/typings/conffwk/_daq_conffwk_py.pyi +++ b/typings/conffwk/_daq_conffwk_py.pyi @@ -1,20 +1,22 @@ """ Python interface to the conffwk package """ + from __future__ import annotations + import typing + __all__: list[str] = list() + class _ConfigObject: def UID(self) -> str: """ Return object identity """ @typing.overload - def __init__(self) -> None: - ... + def __init__(self) -> None: ... @typing.overload - def __init__(self, arg0: _ConfigObject) -> None: - ... + def __init__(self, arg0: _ConfigObject) -> None: ... def class_name(self) -> str: """ Return object's class name @@ -183,11 +185,15 @@ class _ConfigObject: """ Setter function for list """ - def set_obj(self, name: str, o: _ConfigObject, skip_non_null_check: bool = False) -> None: + def set_obj( + self, name: str, o: _ConfigObject, skip_non_null_check: bool = False + ) -> None: """ Set relationship single-value """ - def set_objs(self, name: str, o: list[_ConfigObject], skip_non_null_check: bool = False) -> None: + def set_objs( + self, name: str, o: list[_ConfigObject], skip_non_null_check: bool = False + ) -> None: """ Set relationship multi-value. """ @@ -271,13 +277,12 @@ class _ConfigObject: """ Setter function for list """ + class _Configuration: @typing.overload - def __init__(self) -> None: - ... + def __init__(self) -> None: ... @typing.overload - def __init__(self, arg0: str) -> None: - ... + def __init__(self, arg0: str) -> None: ... def add_include(self, db_name: str, include: str) -> None: """ Add include file to existing database. @@ -290,7 +295,7 @@ class _Configuration: """ Get the names of the superclasses for each class """ - def commit(self, log_message: str = '') -> None: + def commit(self, log_message: str = "") -> None: """ Commit database changes. """ @@ -328,7 +333,7 @@ class _Configuration: """ Create a configuration object containing the desired entity from the database """ - def get_objs(self, class_name: str, query: str = '') -> list[...]: + def get_objs(self, class_name: str, query: str = "") -> list[...]: """ Create a list of configuration objects of a given class from the database """ @@ -360,7 +365,9 @@ class _Configuration: """ Get the superclasses of a single class """ - def test_object(self, class_name: str, id: str, rlevel: int, rclasses: list[str]) -> bool: + def test_object( + self, class_name: str, id: str, rlevel: int, rclasses: list[str] + ) -> bool: """ Test the existence of the object """ diff --git a/typings/conffwk/dal.pyi b/typings/conffwk/dal.pyi index e14bd856b..03b68698e 100644 --- a/typings/conffwk/dal.pyi +++ b/typings/conffwk/dal.pyi @@ -11,7 +11,6 @@ class FSMParameter(ParameterProtocol): name: str value: str - class FSMAction(ConfigurationProtocol): id: str name: str @@ -20,28 +19,23 @@ class FSMAction(ConfigurationProtocol): oks_key: OksKeyProtocol initial_data: str - class FSMxTransition: transition: str order: List[str] mandatory: List[str] - class FSMTransitionConfig: id: str source: str dest: str - class FSMCommand: id: str - class FSMCommandSequence: id: str sequence: List[FSMCommand] - class FSMData: states: List[str] initial_state: str diff --git a/typings/conffwk/dalproperty.pyi b/typings/conffwk/dalproperty.pyi index 9143ae5da..56aecb9de 100644 --- a/typings/conffwk/dalproperty.pyi +++ b/typings/conffwk/dalproperty.pyi @@ -1,12 +1,13 @@ -import __future__ from __future__ import annotations -__all__: list[str] = ['absolute_import'] -def _assign_attribute(attribute): - ... -def _assign_relation(relation): - ... -def _return_attribute(attribute, dalobj = None, value = None): - ... -def _return_relation(relation, multi = False, data = None, dalobj = None, cache = None): - ... -absolute_import: __future__._Feature # value = _Feature((2, 5, 0, 'alpha', 1), (3, 0, 0, 'alpha', 0), 262144) +import __future__ + +__all__: list[str] = ["absolute_import"] + +def _assign_attribute(attribute): ... +def _assign_relation(relation): ... +def _return_attribute(attribute, dalobj=None, value=None): ... +def _return_relation(relation, multi=False, data=None, dalobj=None, cache=None): ... + +absolute_import: ( + __future__._Feature +) # value = _Feature((2, 5, 0, 'alpha', 1), (3, 0, 0, 'alpha', 0), 262144) diff --git a/typings/conffwk/proxy.pyi b/typings/conffwk/proxy.pyi index 81469fe94..a47b05295 100644 --- a/typings/conffwk/proxy.pyi +++ b/typings/conffwk/proxy.pyi @@ -6,41 +6,45 @@ instances expose the same public interface of the proxied object, but avoiding inheritance. This allows to control the reference counts of the proxied object. """ + from __future__ import annotations -__all__: list[str] = ['Proxy', 'make_proxy_class'] + +__all__: list[str] = ["Proxy", "make_proxy_class"] + class Proxy: """ A very basic holder class - + Just holds the reference to a provided object. - - + + """ - def __init__(self, obj): - ... + def __init__(self, obj): ... + def _DelegateMetaFunction(clsName, bases, atts): """ - Implements a delegation pattern using a metaclass approach - - A class using this meta mechanism should have 'memberclass' class attribute - initialized at the class of the instance to proxied. The metaclass will - make sure the delegate class will expose all the public methods of the - proxied one. - Moreover, the metaclass will provide the delegate class with a '__init__' - function instantiating a 'memberclass' object, storing it in 'self._obj'. - The delegate class constructor method will therefore accept all the - arguments accepted by the proxied class constructor. - The delegate class uses slots - - + Implements a delegation pattern using a metaclass approach + + A class using this meta mechanism should have 'memberclass' class attribute + initialized at the class of the instance to proxied. The metaclass will + make sure the delegate class will expose all the public methods of the + proxied one. + Moreover, the metaclass will provide the delegate class with a '__init__' + function instantiating a 'memberclass' object, storing it in 'self._obj'. + The delegate class constructor method will therefore accept all the + arguments accepted by the proxied class constructor. + The delegate class uses slots + + """ + def make_proxy_class(theclass): """ Builds a delegation class out of a given type. - + Uses the Proxy class to generate a new proxy class exposing the same interface of the provided class and delegating the method calls to the hosted object instance. - - + + """ diff --git a/typings/conffwk/schema.pyi b/typings/conffwk/schema.pyi index 390a3765b..4b7b1dd59 100644 --- a/typings/conffwk/schema.pyi +++ b/typings/conffwk/schema.pyi @@ -1,39 +1,61 @@ """ A set of utilities to simplify OKS instrospection. """ + from __future__ import annotations -from conffwk import ConfigObject + import logging as logging import re as re import sys as sys -__all__: list[str] = ['Cache', 'ConfigObject', 'check_cardinality', 'check_range', 'check_relation', 'coerce', 'decode_range', 'logging', 'map_coercion', 'oks_types', 'range_regexp', 're', 'str2integer', 'sys', 'to_int', 'to_long'] + +from conffwk import ConfigObject + +__all__: list[str] = [ + "Cache", + "ConfigObject", + "check_cardinality", + "check_range", + "check_relation", + "coerce", + "decode_range", + "logging", + "map_coercion", + "oks_types", + "range_regexp", + "re", + "str2integer", + "sys", + "to_int", + "to_long", +] + class Cache: """ Defines a cache for all known schemas at a certain time. - + """ def __getitem__(self, key): """ Gets the description of a certain class. """ - def __init__(self, conffwk, all = True): + def __init__(self, conffwk, all=True): """ Initializes the cache with information from the Configuration object. - + This method will browse for all declared classes in the Configuration object given as input and will setup the schema for all known classes. After this you can still update the cache using the update() method. - + Keyword parameters: - + conffwk -- The conffwk.Configuration object to use as base for the current cache. - + all -- A boolean indicating if I should store all the attributes and relations from a certain class or just the ones directly associated with a class. - + """ def __str__(self): """ @@ -42,65 +64,81 @@ class Cache: def update(self, conffwk): """ Updates this cache with information from the Configuration object. - + This method will add new classes not yet know to this cache. Classes with existing names will not be added. No warning is generated (this should be done by the OKS layer in any case. - + """ def update_dal(self, conffwk): """ Updates this cache with information for DAL. - + This method will add new DAL classes not yet know to this cache. Classes with existing DAL representations will not be touched. - + """ + def check_cardinality(v, prop): """ Checks the cardinality of a certain attribute or relationship. """ + def check_range(v, range, range_re, pytype): """ Checks the range of the value 'v' to make sure it is inside. """ + def check_relation(v, rel): """ Checks the value v against the relationship parameters in 'rel'. """ + def coerce(v, attr): """ Coerces the input value 'v' in the way the attribute expects. """ + def decode_range(s): """ Decodes a range string representation, returns a tuple with 2 values. - + This is the supported format in regexp representation: '([-0x]*\\d+)\\D+-?\\d+' - + """ + def map_coercion(class_name, schema): """ Given a schema of a class, maps coercion functions from libpyconffwk. """ + def str2integer(v, t, max): """ Converts a value v to integer, irrespectively of its formatting. - + If the number starts with a '0', we convert it using an octal representation. Else, we try a decimal conversion. If any of these fail, we try an hexa conversion before throwing a ValueError. - + Keyword arguments: - + v -- the value to be converted t -- the python type (int or float) to use in the conversion - - """ -def to_int(v): - ... -def to_long(v): - ... -oks_types: dict = {'bool': ['bool'], 'integer': ['s8', 'u8', 's16', 'u16', 's32'], 'long': ['u32', 's64', 'u64'], 'float': ['float', 'double'], 'int-number': ['s8', 'u8', 's16', 'u16', 's32', 'u32', 's64', 'u64'], 'number': ['u32', 's64', 'u64', 's8', 'u8', 's16', 'u16', 's32', 'float', 'double'], 'time': ['date', 'time'], 'string': ['date', 'time', 'string', 'uid', 'enum', 'class']} + + """ + +def to_int(v): ... +def to_long(v): ... + +oks_types: dict = { + "bool": ["bool"], + "integer": ["s8", "u8", "s16", "u16", "s32"], + "long": ["u32", "s64", "u64"], + "float": ["float", "double"], + "int-number": ["s8", "u8", "s16", "u16", "s32", "u32", "s64", "u64"], + "number": ["u32", "s64", "u64", "s8", "u8", "s16", "u16", "s32", "float", "double"], + "time": ["date", "time"], + "string": ["date", "time", "string", "uid", "enum", "class"], +} range_regexp: re.Pattern # value = re.compile('(?P-?0?x?[\\da-fA-F]+(\\.\\d+)?)-(?P-?0?x?[\\da-fA-F]+(\\.\\d+)?)') diff --git a/typings/daqconf/consolidate.pyi b/typings/daqconf/consolidate.pyi index 5afc5dadc..cd7c719f4 100644 --- a/typings/daqconf/consolidate.pyi +++ b/typings/daqconf/consolidate.pyi @@ -4,4 +4,4 @@ def consolidate_db( oksfile: str, output_file: str, session_id: Optional[str] = None, -) -> None: ... \ No newline at end of file +) -> None: ... diff --git a/typings/daqconf/jsonify.pyi b/typings/daqconf/jsonify.pyi index f05b78fa6..64766ceed 100644 --- a/typings/daqconf/jsonify.pyi +++ b/typings/daqconf/jsonify.pyi @@ -1,4 +1,4 @@ def jsonify_xml_data( oksfile: str, output: str, -) -> None: ... \ No newline at end of file +) -> None: ... diff --git a/typings/daqconf/validate.pyi b/typings/daqconf/validate.pyi index a96125cf9..d6d52894f 100644 --- a/typings/daqconf/validate.pyi +++ b/typings/daqconf/validate.pyi @@ -1,4 +1,4 @@ def validate_session( oksfile: str, session_name: str, -) -> None: ... \ No newline at end of file +) -> None: ... From bd3ae9f2fbd3f569b3815690325618a62ba1bfb7 Mon Sep 17 00:00:00 2001 From: James Paul Turner Date: Mon, 13 Jul 2026 17:22:02 +0100 Subject: [PATCH 07/20] Phased merge 1. --- .github/workflows/check_dirs_init_file.yml | 1 + .github/workflows/check_docs.yml | 20 ++++++++++++-------- 2 files changed, 13 insertions(+), 8 deletions(-) diff --git a/.github/workflows/check_dirs_init_file.yml b/.github/workflows/check_dirs_init_file.yml index 4b5eceea2..f34b86529 100644 --- a/.github/workflows/check_dirs_init_file.yml +++ b/.github/workflows/check_dirs_init_file.yml @@ -21,6 +21,7 @@ jobs: -path "src/tests*" -o \ -path "src/docs*" -o \ -path "src/config*" -o \ + -path "src/drunc/integtest*" -o \ -path "src/drunc/data*" \) -prune -o \ -type d -exec sh -c '[ ! -e "$1/__init__.py" ] && echo "$1"' _ {} \;) diff --git a/.github/workflows/check_docs.yml b/.github/workflows/check_docs.yml index 546c22a84..40ea29528 100644 --- a/.github/workflows/check_docs.yml +++ b/.github/workflows/check_docs.yml @@ -2,16 +2,16 @@ name: Check Links In Docs on: pull_request: - paths: - - 'docs/**' - workflow_dispatch: - workflow_call: - schedule: - - cron: "0 0 * * 6" # midnight every Saturday + workflow_dispatch: + workflow_call: + schedule: + - cron: "0 0 * * 6" # Runs every Saturday at midnight jobs: check-links: runs-on: ubuntu-latest + permissions: + contents: read steps: - name: Checkout code uses: actions/checkout@v4 @@ -19,5 +19,9 @@ jobs: - name: Run Lychee link checker uses: lycheeverse/lychee-action@v2 with: - args: docs/ - \ No newline at end of file + args: > + --exclude 'https://twiki.cern.ch/.*' + --exclude 'https://indico.*' + --exclude 'https://wiki.dunescience.org/wiki/.*' + --exclude 'https://github.com/DUNE-DAQ/drunc/wiki/Archive:.*' + docs/ \ No newline at end of file From 15bc3f894492c03f2e177aa0f21fdf1a78ebf3c0 Mon Sep 17 00:00:00 2001 From: James Paul Turner Date: Mon, 13 Jul 2026 17:25:36 +0100 Subject: [PATCH 08/20] Phased merge 2. --- docs/Messaging-format.md | 57 ------------------------------- docs/Process-manager-interface.md | 2 +- docs/Process-manager.md | 4 +-- docs/Unified-shell-reference.md | 1 - 4 files changed, 2 insertions(+), 62 deletions(-) diff --git a/docs/Messaging-format.md b/docs/Messaging-format.md index 9a048c783..3a3a196de 100644 --- a/docs/Messaging-format.md +++ b/docs/Messaging-format.md @@ -79,14 +79,12 @@ message Description { string name = 2; optional string session = 3; repeated CommandDescription commands = 4; - optional google.protobuf.Any broadcast = 5; } ``` * `type` can be `process_manager` or `controller` right now it is a string. * `name` is the name of the server. * `session` is the (optional) session name. * `commands` is a vector of acceptable commands to send to the endpoint. -* `broadcast` is a description of the broadcast service. `CommandDescription` is used to describe all the commands, it has the following format: ``` @@ -102,58 +100,3 @@ message CommandDescription { * `help` is a string providing help * `return_type` is the format of the `data` field that should be expected inside the `Response` if the command execution was successful. -For broadcasting, there is right now only one format that the `describe` command can fill, with a description of the Kafka service that is used to broadcast the log messages: -``` -message KafkaBroadcastHandlerConfiguration{ - string kafka_address = 1; - string topic = 2; -} -``` -* `kafka_address` is a bootstrap server -* `topic` is the Kafka topic on which this service is logging. - -More formats of broadcasting may be added in the future (potentially using [ERS's python binding](https://github.com/DUNE-DAQ/erskafka/tree/develop/python/erskafka)). - -# Broadcasting -As mentioned earlier, the only working solution for broadcasting is Kafka. All messages feed to Kafka by drunc have the [form](https://github.com/DUNE-DAQ/druncschema/blob/develop/schema/druncschema/broadcast.proto#L51): -``` -message BroadcastMessage{ - Emitter emitter = 1; - BroadcastType type = 2; - google.protobuf.Any data = 3; -} -``` -Where: -``` -message Emitter { - string process = 1; - string session = 2; -} -``` -and -``` -enum BroadcastType { - ACK = 0; - RECEIVER_REMOVED = 1; - RECEIVER_ADDED = 2; - SERVER_READY = 3; - SERVER_SHUTDOWN = 4; - TEXT_MESSAGE = 15; - COMMAND_EXECUTION_START = 5; - COMMAND_RECEIVED = 16; - COMMAND_EXECUTION_SUCCESS = 6; - EXCEPTION_RAISED = 7; - UNHANDLED_EXCEPTION_RAISED = 8; - STATUS_UPDATE = 9; - SUBPROCESS_STATUS_UPDATE = 10; - DEBUG = 11; - CHILD_COMMAND_EXECUTION_START = 12; - CHILD_COMMAND_EXECUTION_SUCCESS = 13; - CHILD_COMMAND_EXECUTION_FAILED = 14; - FSM_STATUS_UPDATE = 17; -} -``` - -Note that for now, the `data` field is only filled with [PlainText](https://github.com/DUNE-DAQ/druncschema/blob/develop/schema/druncschema/generic.proto#L7C1-L9C2) messages. - -In the future, this may change. diff --git a/docs/Process-manager-interface.md b/docs/Process-manager-interface.md index 90e03b7c1..bd97b814f 100644 --- a/docs/Process-manager-interface.md +++ b/docs/Process-manager-interface.md @@ -193,7 +193,7 @@ message ExceptionNotification { Each RPC call is described here. ### `describe` -Returns metadata about the process manager, including available commands, session info, and broadcast description. +Returns metadata about the process manager, including available commands, and session info. * input: `Request` * output: [Description](https://dune-daq-sw.readthedocs.io/en/latest/packages/drunc/Messaging-format) diff --git a/docs/Process-manager.md b/docs/Process-manager.md index b661f39d2..eeafdf59f 100644 --- a/docs/Process-manager.md +++ b/docs/Process-manager.md @@ -7,7 +7,7 @@ For a standalone `process_manager` you will need two shells - one shell to run t To boot a `process_manager`, you will need to choose the most appropriate configuration that applies to the use case. The configurations that are packaged with `drunc` are defined in `drunc/src/data/process_manager/`, which are * `ssh-standalone.json`: `ssh` based implementation without a `kafka` feed. Uses remote shell client processes to manage ssh connections. * `ssh-standalone-paramiko-client.json`: `ssh` based implementation without a `kafka` feed. Uses paramiko library to manage ssh connections from the python process. Currently NOT maintained but may be revisited in future - use `ssh-standalone` only for now. -* `ssh-pocket-kafka.json`: `ssh` based implementation with `pocket`'s `kafka` for message broadcasting. +* `ssh-pocket-kafka.json`: `ssh` based implementation with `pocket`'s `kafka`. * `ssh-CERN-kafka.json`: `ssh` based implementation with `kafka` service running at ENH1. * `ssh-CERN-kafka-OpMon.json`: `ssh` based implementation with `kafka` service running at ENH1, and with Opmon. * `k8s.json`: `kubernetes` implementation (not recommended nor working, so don't use this unless you are an working on getting it to work). @@ -26,8 +26,6 @@ To start the ssh version without kafka: drunc-process-manager ssh-standalone Using 'file://src/drunc/data/process_manager/ssh-standalone.json' as the ProcessManager configuration Starting 'SSHProcessManager' -[12:43:26] INFO "BroadcastSenderConfHandler": None configuration.py:25 - INFO "Controller": DummyAuthoriser ready dummy_authoriser.py:13 ProcessManager was started on np04-srv-019:10054 ``` Once this is done, you will not be able to send commands to the process from the current shell with the `process_manager` acting in the foreground. To interact with a standalone instance of `process_manager` you will need to connect to it (see below). diff --git a/docs/Unified-shell-reference.md b/docs/Unified-shell-reference.md index 2cf1c9c6a..bc2085672 100644 --- a/docs/Unified-shell-reference.md +++ b/docs/Unified-shell-reference.md @@ -414,7 +414,6 @@ logs -n root-controller --how-far 5 INFO "Controller": 'df-controller@localhost:5600' (type ChildNodeType.gRPC) controller.py:123 INFO "Controller": 'trg-controller@localhost:5700' (type ChildNodeType.gRPC) controller.py:123 INFO "Controller": 'hsi-controller@localhost:5800' (type ChildNodeType.gRPC) controller.py:123 - INFO "Broadcast": ready broadcast_sender.py:65 root-controller was started on localhost:3333 ───────────────────────────────────────────────────────────────────────────────────────────── End ────────────────────────────────────────────────────────────────────────────────────────────── ``` From 032b67b34e77067d01fe108384e76233bc07ce67 Mon Sep 17 00:00:00 2001 From: James Paul Turner Date: Mon, 13 Jul 2026 17:36:22 +0100 Subject: [PATCH 09/20] Phased merge 3. --- scripts/drunc_integtest_bundle.sh | 12 +++++++++++- .../drunc/integtest}/process_manager_test.py | 7 ++++++- 2 files changed, 17 insertions(+), 2 deletions(-) rename {integtest => src/drunc/integtest}/process_manager_test.py (97%) diff --git a/scripts/drunc_integtest_bundle.sh b/scripts/drunc_integtest_bundle.sh index 7cea68c87..e32a4aeb3 100755 --- a/scripts/drunc_integtest_bundle.sh +++ b/scripts/drunc_integtest_bundle.sh @@ -21,6 +21,7 @@ Options: -k -n -N + --verbosity --stop-on-failure : causes the script to stop when one of the integtests reports a failure --concise-output : suppresses run control and DAQApp messages in order to focus on test results --tmpdir : specifies a root directory to use for test output, e.g. a directory instead of '/tmp' @@ -49,7 +50,7 @@ CaptureOutput() { tee -a $1 } -GETOPT_TEMP=`getopt -o hs:f:l:k:n:N: --long help,stop-on-failure,concise-output,tmpdir: -- "$@"` +GETOPT_TEMP=`getopt -o hs:f:l:k:n:N: --long help,verbosity:,stop-on-failure,concise-output,tmpdir: -- "$@"` eval set -- "$GETOPT_TEMP" let first_test_index=0 @@ -57,6 +58,7 @@ let individual_test_requested_iterations=1 let full_set_requested_interations=1 let stop_on_failure=0 requested_test_names= +verbosity_level= PYTEST_COMMAND="pytest -c /dev/null -s --tb=short" # our core pytest command, with DAQ printout included and short pytest traceback while true; do @@ -85,6 +87,10 @@ while true; do let full_set_requested_interations=$2 shift 2 ;; + --verbosity) + verbosity_level=$2 + shift 2 + ;; --stop-on-failure) let stop_on_failure=1 PYTEST_COMMAND="${PYTEST_COMMAND} -x" # add the -x option to our pytest command to have it exit on first error @@ -106,6 +112,10 @@ while true; do esac done +if [[ "${verbosity_level}" != "" ]]; then + PYTEST_COMMAND="$PYTEST_COMMAND --integtest-verbosity ${verbosity_level}" +fi + # check if the numad daemon is running numad_grep_output=`ps -ef | grep numad | grep -v grep` if [[ "${numad_grep_output}" != "" ]]; then diff --git a/integtest/process_manager_test.py b/src/drunc/integtest/process_manager_test.py similarity index 97% rename from integtest/process_manager_test.py rename to src/drunc/integtest/process_manager_test.py index 39584c1a3..2a3640d23 100644 --- a/integtest/process_manager_test.py +++ b/src/drunc/integtest/process_manager_test.py @@ -32,6 +32,9 @@ "Worker with pid \\d+ was terminated due to signal", "Connection '.*' not found on the application registry", ], + "SSH_SHELL_process_manager": [ + "was terminated unexpectedly through the remote pid by a SIGKILL", + ], "connectivity-service": [ "errorlog: -", ], @@ -345,7 +348,9 @@ def test_restart_mlt_logs(run_dunerc) -> None: require_pattern_match( restart_text, - re.compile(r"Process 'mlt'.*?process exited\s+with exit code 0", re.DOTALL), + re.compile( + r"Process 'mlt' \(.*?\) was terminated by the process manager through the remote pid\. Reported exit code: 0\.", re.DOTALL + ), error_message="Did not find the mlt exit-code log line after graceful termination.", ) From cf5d919600e36d37a86327ac141ca60cb619c4bd Mon Sep 17 00:00:00 2001 From: James Paul Turner Date: Mon, 13 Jul 2026 17:46:15 +0100 Subject: [PATCH 10/20] Phased merge 4. --- clean_stubs.py | 103 ------- src/drunc/authoriser/configuration.py | 5 +- src/drunc/authoriser/dummy_authoriser.py | 23 +- src/drunc/broadcast/__init__.py | 0 src/drunc/broadcast/client/__init__.py | 0 .../broadcast/client/broadcast_handler.py | 38 --- .../broadcast_handler_implementation.py | 7 - src/drunc/broadcast/client/configuration.py | 43 --- .../client/grpc_stdout_broadcast_handler.py | 109 ------- .../client/kafka_stdout_broadcast_handler.py | 106 ------- src/drunc/broadcast/server/__init__.py | 0 .../broadcast/server/broadcast_sender.py | 128 -------- .../server/broadcast_sender_implementation.py | 17 -- src/drunc/broadcast/server/configuration.py | 34 --- src/drunc/broadcast/server/decorators.py | 50 ---- src/drunc/broadcast/server/grpc_servicer.py | 282 ------------------ src/drunc/broadcast/server/kafka_sender.py | 69 ----- src/drunc/broadcast/types.py | 15 - src/drunc/broadcast/utils.py | 30 -- 19 files changed, 11 insertions(+), 1048 deletions(-) delete mode 100644 clean_stubs.py delete mode 100644 src/drunc/broadcast/__init__.py delete mode 100644 src/drunc/broadcast/client/__init__.py delete mode 100644 src/drunc/broadcast/client/broadcast_handler.py delete mode 100644 src/drunc/broadcast/client/broadcast_handler_implementation.py delete mode 100644 src/drunc/broadcast/client/configuration.py delete mode 100644 src/drunc/broadcast/client/grpc_stdout_broadcast_handler.py delete mode 100644 src/drunc/broadcast/client/kafka_stdout_broadcast_handler.py delete mode 100644 src/drunc/broadcast/server/__init__.py delete mode 100644 src/drunc/broadcast/server/broadcast_sender.py delete mode 100644 src/drunc/broadcast/server/broadcast_sender_implementation.py delete mode 100644 src/drunc/broadcast/server/configuration.py delete mode 100644 src/drunc/broadcast/server/decorators.py delete mode 100644 src/drunc/broadcast/server/grpc_servicer.py delete mode 100644 src/drunc/broadcast/server/kafka_sender.py delete mode 100644 src/drunc/broadcast/types.py delete mode 100644 src/drunc/broadcast/utils.py diff --git a/clean_stubs.py b/clean_stubs.py deleted file mode 100644 index 05fbc12c5..000000000 --- a/clean_stubs.py +++ /dev/null @@ -1,103 +0,0 @@ -""" -Manual overrides for conffwk stubs. - -The pybind11-stubgen auto-generator marks types it can't understand as 'Any' -or misses fields entirely. This adds those classes manually. -""" - -import os -import re -from pathlib import Path - - -def process_file(filepath: Path) -> None: - with open(filepath, 'r', encoding='utf-8') as f: - content = f.read() - - # Fix parameter types - content = re.sub(r': \.\.\.', ': object', content) - - # Fix return types - content = re.sub(r'-> \.\.\.:', '-> object:', content) - - # Fix untyped arguments - content = re.sub(r'= \.\.\.', '= object()', content) - - with open(filepath, 'w', encoding='utf-8') as f: - f.write(content) - - -def add_dal_classes(stubs_dir: Path) -> None: - dal_file = stubs_dir / "dal.pyi" - - dal_classes = """ -from typing import List - -from drunc.fsm._protocols import ( - ConfigurationProtocol, - DBProtocol, - OksKeyProtocol, - ParameterProtocol, -) - -class FSMParameter(ParameterProtocol): - name: str - value: str - -class FSMAction(ConfigurationProtocol): - id: str - name: str - parameters: List[FSMParameter] - db: DBProtocol - oks_key: OksKeyProtocol - initial_data: str - -class FSMxTransition: - transition: str - order: List[str] - mandatory: List[str] - -class FSMTransitionConfig: - id: str - source: str - dest: str - -class FSMCommand: - id: str - -class FSMCommandSequence: - id: str - sequence: List[FSMCommand] - -class FSMData: - states: List[str] - initial_state: str - actions: List[FSMAction] - transitions: List[FSMTransitionConfig] - pre_transitions: List[FSMxTransition] - post_transitions: List[FSMxTransition] - command_sequences: List[FSMCommandSequence] -""" - with open(dal_file, 'w', encoding='utf-8') as f: - f.write(dal_classes) - print("[+] Done. dal.pyi is now strictly typed.") - -def main() -> None: - stubs_dir = Path("typings/conffwk") - - if not stubs_dir.exists(): - print("Typings/conffwk directory not found. Run 'pybind11-stubgen conffwk --output-dir=typings' first.") - return - - # Fix syntax errors in all generated files - for root, _, files in os.walk(stubs_dir): - for file in files: - if file.endswith('.pyi'): - process_file(Path(root) / file) - - # Add missing classes into the DAL file - add_dal_classes(stubs_dir) - - -if __name__ == "__main__": - main() diff --git a/src/drunc/authoriser/configuration.py b/src/drunc/authoriser/configuration.py index 2820db245..49f20ee2f 100644 --- a/src/drunc/authoriser/configuration.py +++ b/src/drunc/authoriser/configuration.py @@ -2,4 +2,7 @@ class DummyAuthoriserConfHandler(ConfHandler): - pass + """Handler for dummy authoriser configuration.""" + + def populate_from_dict(self, data: dict[str, object]) -> None: + pass diff --git a/src/drunc/authoriser/dummy_authoriser.py b/src/drunc/authoriser/dummy_authoriser.py index c093e8fb5..a7098f687 100644 --- a/src/drunc/authoriser/dummy_authoriser.py +++ b/src/drunc/authoriser/dummy_authoriser.py @@ -5,22 +5,21 @@ from drunc.utils.utils import get_logger -# TODO: Should be communicating over network -# The Rolls Royce of the authoriser systems class DummyAuthoriser: def __init__( - self, - system: SystemType, - configuration_handler: DummyAuthoriserConfHandler = None, + self, configuration: DummyAuthoriserConfHandler, system: SystemType.ValueType ): self.log = get_logger("utils.authorizer") self.log.debug("DummyAuthoriser ready") - self.configuration = configuration_handler - self.command_actions = {} # Dict[str, ActionType] + self.configuration = configuration self.system = system def is_authorised( - self, token: Token, action: ActionType, system: SystemType, cmd_name: str = "" + self, + token: Token, + action: ActionType.ValueType, + system: SystemType.ValueType, + cmd_name: str, ) -> bool: self.log.debug( f"Authorising {token.user_name} to {ActionType.Name(action)} ({cmd_name}) on {SystemType.Name(system)}" @@ -30,11 +29,3 @@ def is_authorised( def authorised_actions(self, token: Token) -> list[str]: self.log.info(f"Grabbing authorisations for {token.token}") return [] - - -def main(): - DummyAuthoriser() - - -if __name__ == "__main__": - main() diff --git a/src/drunc/broadcast/__init__.py b/src/drunc/broadcast/__init__.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/src/drunc/broadcast/client/__init__.py b/src/drunc/broadcast/client/__init__.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/src/drunc/broadcast/client/broadcast_handler.py b/src/drunc/broadcast/client/broadcast_handler.py deleted file mode 100644 index 1f144f582..000000000 --- a/src/drunc/broadcast/client/broadcast_handler.py +++ /dev/null @@ -1,38 +0,0 @@ -from drunc.broadcast.client.configuration import BroadcastClientConfHandler -from drunc.broadcast.types import BroadcastTypes - - -class BroadcastHandler: - def __init__(self, broadcast_configuration: BroadcastClientConfHandler): - super().__init__() - - from logging import getLogger - - self.log = getLogger("BroadcastHandler") - - self.configuration = broadcast_configuration - self.implementation = None - - match self.configuration.data.type: - # Being a bit sloppy here, having a Kafka sender doesn't mean we want to dump everything to stdout - # There could be cases where we want to do other things. - # For now, 1 server type <-> 1 client type... - # Maybe in the future some sort of callback-based functionality would be preferable. - case BroadcastTypes.Kafka: - from druncschema.broadcast_pb2 import BroadcastMessage - - from drunc.broadcast.client.kafka_stdout_broadcast_handler import ( - KafkaStdoutBroadcastHandler, - ) - - self.implementation = KafkaStdoutBroadcastHandler( - message_format=BroadcastMessage, conf=self.configuration - ) - case _: - self.log.info( - "Could not understand the BroadcastHandler technology you want to use, you will get no broadcast!" - ) - - def stop(self): - if self.implementation: - self.implementation.stop() diff --git a/src/drunc/broadcast/client/broadcast_handler_implementation.py b/src/drunc/broadcast/client/broadcast_handler_implementation.py deleted file mode 100644 index 99d0fa95f..000000000 --- a/src/drunc/broadcast/client/broadcast_handler_implementation.py +++ /dev/null @@ -1,7 +0,0 @@ -import abc - - -class BroadcastHandlerImplementation(abc.ABC): - @abc.abstractmethod - def stop(self): - pass diff --git a/src/drunc/broadcast/client/configuration.py b/src/drunc/broadcast/client/configuration.py deleted file mode 100644 index 4b5cdd7f1..000000000 --- a/src/drunc/broadcast/client/configuration.py +++ /dev/null @@ -1,43 +0,0 @@ -from drunc.broadcast.types import BroadcastTypes -from drunc.utils.configuration import ConfHandler - - -class BroadcastClientConfData: # OKSeroo - def __init__(self, type: BroadcastTypes, address: str, topic: str): - self.type = type - self.address = address - self.topic = topic - - -class BroadcastClientConfHandler(ConfHandler): - def __init__(self, *args, **kwargs): - super().__init__(*args, **kwargs) - - def get_impl_technology(self): - return self.impl_technology - - def _parse_pbany(self, data): - # potentially do something more complicated with different implementation technology here - # match data.format(): - # case KafkaBroadcastHandlerConfiguration - # ... - - from druncschema.broadcast_pb2 import KafkaBroadcastHandlerConfiguration - - from drunc.utils.grpc_utils import UnpackingError, unpack_any - - if not data.ByteSize(): - return BroadcastClientConfData(type=None, address=None, topic=None) - try: - data = unpack_any(data, KafkaBroadcastHandlerConfiguration) - return BroadcastClientConfData( - type=BroadcastTypes.Kafka, address=data.kafka_address, topic=data.topic - ) - - except UnpackingError as e: - from drunc.exceptions import DruncSetupException - - raise DruncSetupException( - f"Input configuration to configure the broadcast was not understood, could not setup the broadcast handler: {e}", - e, - ) diff --git a/src/drunc/broadcast/client/grpc_stdout_broadcast_handler.py b/src/drunc/broadcast/client/grpc_stdout_broadcast_handler.py deleted file mode 100644 index e0faa89b5..000000000 --- a/src/drunc/broadcast/client/grpc_stdout_broadcast_handler.py +++ /dev/null @@ -1,109 +0,0 @@ -import grpc -from druncschema.broadcast_pb2 import BroadcastMessage, BroadcastType -from druncschema.broadcast_pb2_grpc import BroadcastReceiverServicer -from druncschema.generic_pb2 import Empty - -from drunc.controller.configuration import ControllerConfHandler - - -class gRPCStdoutBroadcastHandler(BroadcastReceiverServicer): - def __init__(self, conf: ControllerConfHandler, token, **kwargs) -> None: - super(gRPCStdoutBroadcastHandler, self).__init__(**kwargs) - from drunc.exceptions import DruncSetupException - - raise DruncSetupException( - "gRPCStdoutBroadcastHandler is not handled it needs to be reworked!" - ) - self.ready = False - self.stub = None - self.token = token - from logging import getLogger - - self._log = getLogger("BroadcastReceiver") - self._address = None # f'[::]:{port}' - self._log.debug("Broadcast receiver initialised") - - def stop_receiving(self) -> None: - self._server.stop(0) - self._log.debug("Broadcast receiver stopped") - - def connect(self) -> None: - from druncschema.broadcast_pb2 import BroadcastRequest - - from drunc.utils.grpc_utils import send_command - - self._log.info(f"Connecting to {self.stub}") - try: - send_command( - controller=self.stub, - token=self.token, - command="add_to_broadcast_list", - data=BroadcastRequest(broadcast_receiver_address=self._address), - rethrow=True, - ) - - except Exception as e: - self._log.error("Could not connect to service to receive broadcast") - self._log.error(str(e)) - raise e - - def disconnect(self) -> None: - from druncschema.broadcast_receiver_pb2 import BroadcastRequest - - from drunc.utils.grpc_utils import send_command - - try: - send_command( - controller=self.stub, - token=self.token, - command="remove_from_broadcast_list", - data=BroadcastRequest(broadcast_receiver_address=self._address), - rethrow=True, - ) - - except Exception as e: - self._log.error( - "Could not disconnect from broadcaster (maybe it' dead and you won't receive any broadcast anyway)" - ) - self._log.error(str(e)) - - def terminate(self) -> None: - self.disconnect() - self.stop_receiving() - - def serve(self) -> None: - from concurrent import futures - - self._server = grpc.server(futures.ThreadPoolExecutor(max_workers=1)) - - from druncschema.broadcast_pb2_grpc import ( - add_BroadcastReceiverServicer_to_server, - ) - - add_BroadcastReceiverServicer_to_server(self, self._server) - - self._server.add_insecure_port(self._address) - - self._server.start() - self.ready = True - self._log.debug("Broadcast receiver server started") - self._server.wait_for_termination() - self.ready = False - - def handle_broadcast( - self, bm: BroadcastMessage, context: grpc.aio.ServicerContext = None - ) -> Empty: - from druncschema.generic_pb2 import PlainText - - from drunc.utils.grpc_utils import unpack_any - - type = bm.type - if type == BroadcastType.TEXT_MESSAGE: - pt = unpack_any(bm.data, PlainText) - self._log.info(f"{bm.emitter}: {pt}") - elif type == BroadcastType.ACK: - self._log.info(f"{bm.emitter}: Ack") - elif type == BroadcastType.SERVER_SHUTDOWN: - self._log.info(f"{bm.emitter} is shutting down") - - return Empty() diff --git a/src/drunc/broadcast/client/kafka_stdout_broadcast_handler.py b/src/drunc/broadcast/client/kafka_stdout_broadcast_handler.py deleted file mode 100644 index 36022ee96..000000000 --- a/src/drunc/broadcast/client/kafka_stdout_broadcast_handler.py +++ /dev/null @@ -1,106 +0,0 @@ -from drunc.broadcast.client.broadcast_handler_implementation import ( - BroadcastHandlerImplementation, -) - - -class KafkaStdoutBroadcastHandler(BroadcastHandlerImplementation): - def __init__(self, message_format, conf): - from drunc.broadcast.utils import broadcast_types_loglevels - - self.broadcast_types_loglevels = ( - broadcast_types_loglevels # in this case, we stick with default - ) - self.conf = conf - # import os - # drunc_shell_conf = os.getenv('DRUNC_SHELL_CONF', None) - # if drunc_shell_conf is not None: - - # with open(drunc_shell_conf) as f: - # import json - # self.global_kafka_stdout_conf = json.load(f).get('kafka_broadcast_handler', {}) - # if 'broadcast_types_loglevels' in self.global_kafka_stdout_conf: - # self.broadcast_types_loglevels.update(self.global_kafka_stdout_conf['broadcast_types_loglevels']) - - self.kafka_address = self.conf.data.address - self.topic = self.conf.data.topic - - # self.broadcast_types_loglevels.update(conf.data.get('broadcast_types_loglevels', {})) - - self.message_format = message_format - - import logging - - self._log = logging.getLogger("Broadcast") - - import getpass - - from drunc.utils.utils import get_random_string, now_str - - group_id = f"drunc-stdout-broadcasthandler-{getpass.getuser()}-{now_str(True)}-{get_random_string(5)}" - - from kafka import KafkaConsumer - - self.consumer = KafkaConsumer( - self.topic, - client_id="run_control", - bootstrap_servers=[self.kafka_address], - group_id=group_id, - ) - - self.run = True - import threading - - self.thread = threading.Thread(target=self.consume) - self.thread.start() - - def stop(self): - self._log.info(f"Stopping listening to '{self.topic}'") - self.run = False - self.thread.join() - - def consume(self): - from druncschema.broadcast_pb2 import BroadcastType - from druncschema.generic_pb2 import PlainText - from google.protobuf import text_format - - from drunc.utils.grpc_utils import unpack_any - - while self.run: - for messages in self.consumer.poll(timeout_ms=500).values(): - for message in messages: - decoded = "" - try: - decoded = self.message_format() - decoded.ParseFromString(message.value) - self._log.debug(f"{decoded=}, {type(decoded)=}") - except Exception as e: - self._log.error( - f"Unhandled broadcast message: {message} (error: {e!s})" - ) - pass - - try: - if decoded.data.Is(PlainText.DESCRIPTOR): - txt = unpack_any(decoded.data, PlainText).text - else: - txt = decoded.data - - from druncschema.broadcast_pb2 import BroadcastType - - from drunc.broadcast.utils import ( - get_broadcast_level_from_broadcast_type, - ) - - bt = BroadcastType.Name(decoded.type) - - get_broadcast_level_from_broadcast_type( - decoded.type, self._log, self.broadcast_types_loglevels - )(f"'{bt}' {txt}") - - except Exception as e: - self._log.error( - f"Weird broadcast message: {message} (error: {e!s})" - ) - text_proto = text_format.MessageToString(decoded) - self._log.info(text_proto) - pass diff --git a/src/drunc/broadcast/server/__init__.py b/src/drunc/broadcast/server/__init__.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/src/drunc/broadcast/server/broadcast_sender.py b/src/drunc/broadcast/server/broadcast_sender.py deleted file mode 100644 index f138a3f5a..000000000 --- a/src/drunc/broadcast/server/broadcast_sender.py +++ /dev/null @@ -1,128 +0,0 @@ -from drunc.broadcast.server.configuration import BroadcastSenderConfHandler - - -class BroadcastSender: - implementation = None - - def __init__( - self, - name: str, - configuration: BroadcastSenderConfHandler, - session: str = "no_session", - ): - super().__init__() - - self.configuration = configuration - - self.name = name - self.session = session - self.identifier = f"{self.session}.{self.name}" - - from logging import getLogger - - self.logger = getLogger("Broadcast") - - self.logger.info("Initialising broadcast") - - from drunc.broadcast.utils import broadcast_types_loglevels - - self.broadcast_types_loglevels = broadcast_types_loglevels - - # TODO - # self.broadcast_types_loglevels.update(self.configuration.get_raw('broadcast_types_loglevels', {})) - - self.impl_technology = self.configuration.get_impl_technology() - - self.implementation = None - - if self.impl_technology is None: - self.logger.info("There is no broadcasting service!") - return - - from drunc.broadcast.types import BroadcastTypes - - match self.impl_technology: - case BroadcastTypes.Kafka: - from drunc.broadcast.server.kafka_sender import KafkaSender - - self.implementation = KafkaSender( - self.configuration.data.address, - self.configuration.data.publish_timeout, - topic=f"control.{self.identifier}", - ) - case _: - from drunc.exceptions import DruncSetupException - - raise DruncSetupException( - f"Broadcaster cannot be {self.impl_technology}" - ) - - def describe_broadcast(self): - if self.implementation: - return self.implementation.describe_broadcast() - else: - return None - - def can_broadcast(self): - if not self.implementation: - return False - - return self.implementation.can_broadcast() - - def broadcast(self, message, btype): - if self.logger: - from drunc.broadcast.utils import get_broadcast_level_from_broadcast_type - - get_broadcast_level_from_broadcast_type( - btype, self.logger, self.broadcast_types_loglevels - )(message) - - if self.implementation is None: - # nice and easy case - return - - from druncschema.broadcast_pb2 import BroadcastMessage, Emitter - from druncschema.generic_pb2 import PlainText - - from drunc.utils.grpc_utils import pack_to_any - - any = pack_to_any(PlainText(text=message)) - emitter = Emitter( - process=self.name, - session=self.session, - ) - bm = BroadcastMessage( - emitter=emitter, - type=btype, - data=any, - ) - bm.type = btype - - self.implementation._send(bm) - - def _interrupt_with_exception(self, exception, context, stack=""): - from druncschema.broadcast_pb2 import BroadcastType - - txt = f"'{exception.__class__.__name__}' exception thrown: {exception}" - - from drunc.exceptions import DruncException - - self.broadcast( - btype=( - BroadcastType.DRUNC_EXCEPTION_RAISED - if isinstance(exception, DruncException) - else BroadcastType.UNHANDLED_EXCEPTION_RAISED - ), - message=txt, - ) - - if stack: - txt += "\n\n" + stack - - from google.rpc import code_pb2 - - error_code = getattr(exception, "grpc_error_code", code_pb2.INTERNAL) - context.abort( - code=error_code, - details=txt, - ) diff --git a/src/drunc/broadcast/server/broadcast_sender_implementation.py b/src/drunc/broadcast/server/broadcast_sender_implementation.py deleted file mode 100644 index f6ace3a2b..000000000 --- a/src/drunc/broadcast/server/broadcast_sender_implementation.py +++ /dev/null @@ -1,17 +0,0 @@ -import abc - -from druncschema.broadcast_pb2 import BroadcastMessage - - -class BroadcastSenderImplementation(abc.ABC): - @abc.abstractmethod - def _send(self, bm: BroadcastMessage): - pass - - @abc.abstractmethod - def describe_broadcast(self): - pass - - @abc.abstractmethod - def can_broadcast(self): - pass diff --git a/src/drunc/broadcast/server/configuration.py b/src/drunc/broadcast/server/configuration.py deleted file mode 100644 index fe485e4fa..000000000 --- a/src/drunc/broadcast/server/configuration.py +++ /dev/null @@ -1,34 +0,0 @@ -from drunc.utils.configuration import ConfHandler - - -class KafkaBroadcastSenderConfData: - def __init__(self, address=None, publish_timeout=None): - self.address = address - self.publish_timeout = publish_timeout - - @staticmethod - def from_dict(data: dict): - address = data.get("address") - if address is None: - address = data["kafka_address"] - - return KafkaBroadcastSenderConfData( - address=address, publish_timeout=data["publish_timeout"] - ) - - -class BroadcastSenderConfHandler(ConfHandler): - def _post_process_oks(self): - from drunc.broadcast.types import BroadcastTypes - - self.impl_technology = BroadcastTypes.Kafka if self.data else None - self.log.debug(self.data) - - def get_impl_technology(self): - return self.impl_technology - - def _parse_dict(self, data): - if data == {}: - self.impl_technology = None - return None - return KafkaBroadcastSenderConfData.from_dict(data) diff --git a/src/drunc/broadcast/server/decorators.py b/src/drunc/broadcast/server/decorators.py deleted file mode 100644 index ece3f95dd..000000000 --- a/src/drunc/broadcast/server/decorators.py +++ /dev/null @@ -1,50 +0,0 @@ - - -from drunc.utils.utils import get_logger - - -def broadcasted(cmd): - import functools - - @functools.wraps( - cmd - ) # this nifty decorator of decorator (!) is nicely preserving the cmd.__name__ (i.e. signature) - def wrap(obj, request, context): - log = get_logger("broadcasted_decorator", rich_handler=True) - - # hummmm I feel like creating a level myself, but... - # https://docs.python.org/3/howto/logging.html#custom-levels - # lets not - log.debug("Entering") - from druncschema.broadcast_pb2 import BroadcastType - - msg = f"User '{request.token.user_name}' executing '{cmd.__name__}'" - - log.debug(msg) - - obj.broadcast(message=msg, btype=BroadcastType.ACK) - - ret = None - try: - log.debug("Executing wrapped function") - ret = cmd(obj, request, context) - - except Exception as e: - log.exception(e) - - obj.broadcast( - message=f"Command '{cmd.__name__}' failed", - btype=BroadcastType.UNHANDLED_EXCEPTION_RAISED, - ) - # raise the exception to the client so the interceptor can handle it - raise e - - msg = f"User '{request.token.user_name}' successfully executed '{cmd.__name__}'" - - obj.broadcast(message=msg, btype=BroadcastType.COMMAND_EXECUTION_SUCCESS) - log.debug(msg) - - log.debug("Exiting") - return ret - - return wrap diff --git a/src/drunc/broadcast/server/grpc_servicer.py b/src/drunc/broadcast/server/grpc_servicer.py deleted file mode 100644 index eabed6282..000000000 --- a/src/drunc/broadcast/server/grpc_servicer.py +++ /dev/null @@ -1,282 +0,0 @@ -from queue import Queue -from threading import Lock, Thread - -import grpc -from druncschema.authoriser_pb2 import ActionType -from druncschema.broadcast_pb2 import BroadcastMessage, BroadcastRequest, BroadcastType -from druncschema.broadcast_pb2_grpc import BroadcastSenderServicer -from druncschema.generic_pb2 import PlainText, StringStringMap -from druncschema.request_response_pb2 import Request, Response -from google.protobuf.any_pb2 import Any - -import drunc.controller.exceptions as ctler_excpt -from drunc.utils.grpc_utils import unpack_any - - -class ListenerRepresentation: - def __init__(self, configuration): - self.address = configuration["address"] - self.channel = grpc.insecure_channel(self.address) - from druncschema.broadcast_pb2_grpc import BroadcastReceiverStub - - self.stub = BroadcastReceiverStub(self.channel) - - def handle_broadcast(self, message): - return self.stub.handle_broadcast(message) - - -class GRCPBroadcastSender(BroadcastSenderServicer): - def __init__(self): - from drunc.exceptions import DruncSetupException - - raise DruncSetupException("GRCPBroadcastSender not supported") - - from logging import getLogger - - self.name = "broadcast_sender" - self._log = getLogger("Broadcast Sender") - self._listeners = {} - self._listener_lock = Lock() - self._message_queue = Queue() - self._consumer_thread = Thread(target=self._consumer, name="broadcast_consumer") - self._consumer_thread.start() - self._log.info("Broadcaster started") - - def _send(self, bm: BroadcastMessage): - pass - - def get_listeners(self): - import copy as cp - - with self._listener_lock: - ret = cp.deepcopy(list(self._listeners.keys())) - return ret - - def broadcast(self, txt, type=BroadcastType.TEXT_MESSAGE): - from druncschema.broadcast_pb2 import BroadcastMessage - from google.protobuf import any_pb2 - - message = PlainText(text=txt) - data_detail = any_pb2.Any() - data_detail.Pack(message) - - bm = BroadcastMessage( - emitter=self.name, type=BroadcastType.TEXT_MESSAGE, data=data_detail - ) - - return self._message_queue.put(bm) - - def broadcast_exception(self, exception): - from druncschema.broadcast_pb2 import BroadcastMessage - from google.protobuf import any_pb2 - - message = PlainText(text=str(exception)) - data_detail = any_pb2.Any() - data_detail.Pack(message) - - bm = BroadcastMessage( - emitter=self.name, type=BroadcastType.EXCEPTION_RAISED, data=data_detail - ) - return bm - - def add_to_bl_logic(self, request: BroadcastRequest): - self.add_listener(request.broadcast_receiver_address) - - def execute_command(self, request, format, logic, action): - uname = request.token.user_name - if self.authoriser: - if not self.authoriser.is_authorised(request.token, action): - from drunc.authoriser.exceptions import Unauthorised - - raise Unauthorised(uname, action) - from drunc.utils.grpc_utils import unpack_any - - data = unpack_any(request.data, format) - - self.broadcast( - f"Executing {action} (user: {uname})", BroadcastType.COMMAND_EXECUTION_START - ) - ret = logic(data) - self.broadcast( - f"Finshed executing {action} (user: {uname})", - BroadcastType.COMMAND_EXECUTION_SUCCESS, - ) - - response = Response() - response.token.CopyFrom(request.token) - data = Any() - if ret: - data.Pack(ret) - response.data.CopyFrom(data) - - return response - - def add_to_broadcast_list(self, request: Request, context) -> Response: - try: - return self.execute_command( - request=request, - format=BroadcastRequest, - logic=self.add_to_bl_logic, - action=ActionType.CREATE, - ) - except Exception as e: - self.broadcast_exception(e) - - def remove_from_broadcast_list(self, request: Request, context) -> Response: - r = unpack_any(request, BroadcastRequest) - if not self.broadcaster.rm_listener(r.broadcast_receiver_address): - raise ctler_excpt.ControllerException( - f"Failed to remove {r.broadcast_receiver_address} from broadcast list" - ) - return PlainText( - text=f"Removed {r.broadcast_receiver_address} to broadcast list" - ) - - def get_broadcast_list(self, request: Request, context) -> Response: - # return self._generic_user_command(request, '_get_broadcast_list', context) - ret = StringStringMap() - listeners = self.broadcaster.get_listeners() - for k, v in listeners.items(): - ret[k] = v - return - - def ack(self, address): - if address not in self._listeners: - raise RuntimeError(f"Cannot send ack to {address}") - - stub = self._listeners[address] - self._log.debug(f"Ack to {address}") - - from druncschema.broadcast_pb2 import BroadcastMessage, BroadcastType - - message = BroadcastMessage(emitter=self.name, type=BroadcastType.ACK) - - try: - stub.handle_broadcast(message) - except Exception as e: - self._log.error(f"Could not Ack to {address}: {e!s}") - - def shutdown(self): - from druncschema.broadcast_pb2 import BroadcastMessage, BroadcastType - - bm = BroadcastMessage(emitter=self.name, type=BroadcastType.SERVER_SHUTDOWN) - - return self._message_queue.put(bm) - - def add_listener(self, address): - with self._listener_lock: - if address in self._listeners.keys(): - self._log.error(f"Listener {address} already exists") - self._listener_lock.release() - return False - self._log.info(f"Adding listener {address}") - self._listeners[address] = ListenerRepresentation(address) - self.ack(address) - - return True - - def rm_listener(self, address): - with self._listener_lock: - if address not in self._listeners.keys(): - self._log.error(f"Listener {address} does not exist") - self._listener_lock.release() - return False - self._log.info(f"Removing listener {address}") - del self._listeners[address] - return True - - def join(self): - return self._consumer_thread.join() - - def _consumer(self): - from druncschema.broadcast_pb2 import BroadcastType - - while True: - message = ( - self._message_queue.get() - ) # Wait for a message from the controller - self._log.debug("Received broadcast message: " + str(message)) - self._listener_lock.acquire() - for address, listener in self._listeners.items(): - self._log.debug(f"Broadcasting {message} to {address}") - try: - response = listener.handle_broadcast(message) - except Exception as e: - self._log.error(f"Could not broadcast to {address}: {e}") - self._log.debug(f"Received response from {address}: {response}") - - self._listener_lock.release() - - if message.type == BroadcastType.SERVER_SHUTDOWN: - break - - -def main(): - from druncschema.broadcast_pb2 import BroadcastMessage, BroadcastType - from druncschema.broadcast_pb2_grpc import BroadcastReceiver - from druncschema.generic_pb2 import Empty - - class StatusReceiver(BroadcastReceiver): - def __init__(self, port): - super(StatusReceiver, self).__init__() - self.port = port - - def handle_broadcast( - self, bm: BroadcastMessage, context: grpc.aio.ServicerContext = None - ): - from drunc.utils.grpc_utils import unpack_any - - if bm.type == BroadcastType.SERVER_SHUTDOWN: - print("End of broadcast") - - elif bm.type == BroadcastType.TEXT_MESSAGE: - pt = unpack_any(bm.data, PlainText) - print(f'Got broadcasted message: "{pt.text}" on port {self.port}') - - return Empty() - - def serve(port: int) -> None: - from concurrent import futures - - server = grpc.server(futures.ThreadPoolExecutor(max_workers=1)) - from druncschema.broadcast_pb2_grpc import ( - add_BroadcastReceiverServicer_to_server, - ) - - add_BroadcastReceiverServicer_to_server(StatusReceiver(port), server) - server.add_insecure_port(f"[::]:{port}") - server.start() - print(f"Status receiver started on {port}") - server.wait_for_termination() - - receiver_threads = [] - port_list = range(15000, 15010) - for port in port_list: - try: - server_thread = Thread( - target=serve, kwargs={"port": port}, name=f"serve_thread_{port}" - ) - server_thread.start() - receiver_threads.append(server_thread) - except: - pass - from drunc.broadcast.server.broadcast_sender import BroadcastSender - - broadcaster = BroadcastSender() - for port in port_list: - broadcaster.add_listener(f"[::]:{port}") - broadcaster.broadcast("Test message") - broadcaster.broadcast("How do you do?") - broadcaster.broadcast("Let's all go for coffee!") - broadcaster.broadcast("End of broadcast") - broadcaster.shutdown() - broadcaster.join() - print("Broadcasting done") - print("\n\nYou can now ctrl-c\n\n") - - for thread in receiver_threads: - thread.join() - - -if __name__ == "__main__": - main() diff --git a/src/drunc/broadcast/server/kafka_sender.py b/src/drunc/broadcast/server/kafka_sender.py deleted file mode 100644 index c1a861a71..000000000 --- a/src/drunc/broadcast/server/kafka_sender.py +++ /dev/null @@ -1,69 +0,0 @@ -from druncschema.broadcast_pb2 import BroadcastMessage - -from drunc.broadcast.server.broadcast_sender_implementation import ( - BroadcastSenderImplementation, -) - - -class KafkaSender(BroadcastSenderImplementation): - def __init__(self, kafka_address: str, publish_timeout: int, topic: str, **kwargs): - super(KafkaSender, self).__init__(**kwargs) - - import logging - - self._log = logging.getLogger(f"{topic}.KafkaSender") - - from kafka import KafkaProducer - from kafka import errors as Errors - - self.topic = topic - self._can_broadcast = False - - self.kafka_address = kafka_address - self.publish_timeout = publish_timeout - - try: - self.kafka = KafkaProducer( - bootstrap_servers=[self.kafka_address], - client_id="run_control", - ) - except Errors.NoBrokersAvailable as e: - t = f"{self.kafka_address} does not seem to point to a kafka broker." - self._log.critical(t) - from drunc.exceptions import DruncSetupException - - raise DruncSetupException(t) from e - - self._log.info( - f'Broadcasting to Kafka ({self.kafka_address}) client_id: "run_control", topic: "{self.topic}"' - ) - self._can_broadcast = True - - def can_broadcast(self): - return self._can_broadcast - - def _send(self, bm: BroadcastMessage): - from kafka.errors import KafkaError - - future = self.kafka.send(self.topic, bm.SerializeToString()) - - record_metadata = None - - try: - record_metadata = future.get(timeout=self.publish_timeout) - except KafkaError as e: - # Decide what to do if produce request failed... - self._log.error(f"Kafka exception sending message {bm}: {e!s}") - except Exception as e: - # Decide what to do if produce request failed... - self._log.error(f"Unhandled exception sending message {bm}: {e!s}") - else: - self._log.debug(f"{record_metadata} published") - - def describe_broadcast(self): - from druncschema.broadcast_pb2 import KafkaBroadcastHandlerConfiguration - - return KafkaBroadcastHandlerConfiguration( - topic=self.topic, - kafka_address=self.kafka_address, - ) diff --git a/src/drunc/broadcast/types.py b/src/drunc/broadcast/types.py deleted file mode 100644 index d91ce34be..000000000 --- a/src/drunc/broadcast/types.py +++ /dev/null @@ -1,15 +0,0 @@ -from enum import Enum - -from drunc.exceptions import DruncSetupException - - -class BroadcastTypes(Enum): - Unknown = 0 - Kafka = 1 - ERS = 2 - - -class BroadcastTypeNotHandled(DruncSetupException): - def __init__(self, btype): - message = f"{btype} not handled" - super(BroadcastTypeNotHandled, self).__init__(message) diff --git a/src/drunc/broadcast/utils.py b/src/drunc/broadcast/utils.py deleted file mode 100644 index 1492996ef..000000000 --- a/src/drunc/broadcast/utils.py +++ /dev/null @@ -1,30 +0,0 @@ -broadcast_types_loglevels = { - "ACK": "debug", - "RECEIVER_REMOVED": "info", - "RECEIVER_ADDED": "info", - "SERVER_READY": "info", - "SERVER_SHUTDOWN": "info", - "TEXT_MESSAGE": "info", - "COMMAND_EXECUTION_START": "info", - "COMMAND_EXECUTION_SUCCESS": "info", - "EXCEPTION_RAISED": "error", - "UNHANDLED_EXCEPTION_RAISED": "critical", - "STATUS_UPDATE": "info", - "SUBPROCESS_STATUS_UPDATE": "info", - "DEBUG": "debug", - "CHILD_COMMAND_EXECUTION_START": "info", - "CHILD_COMMAND_EXECUTION_SUCCESS": "info", - "CHILD_COMMAND_EXECUTION_FAILED": "error", -} - - -def get_broadcast_level_from_broadcast_type( - btype, logger, levels=broadcast_types_loglevels -): - from druncschema.broadcast_pb2 import BroadcastType - - bt = BroadcastType.Name(btype) - if bt not in levels: - return logger.info - else: - return getattr(logger, levels[bt].lower()) From 87f319e5a02e55ee58c960ab001e6f230a22d845 Mon Sep 17 00:00:00 2001 From: James Paul Turner Date: Mon, 13 Jul 2026 17:52:05 +0100 Subject: [PATCH 11/20] Phased merge 5. --- typings/conffwk/__init__.pyi | 35 --- typings/conffwk/_daq_conffwk_py.pyi | 377 ---------------------------- typings/conffwk/dal.pyi | 46 ---- typings/conffwk/dalproperty.pyi | 13 - typings/conffwk/proxy.pyi | 50 ---- typings/conffwk/schema.pyi | 144 ----------- 6 files changed, 665 deletions(-) delete mode 100644 typings/conffwk/__init__.pyi delete mode 100644 typings/conffwk/_daq_conffwk_py.pyi delete mode 100644 typings/conffwk/dal.pyi delete mode 100644 typings/conffwk/dalproperty.pyi delete mode 100644 typings/conffwk/proxy.pyi delete mode 100644 typings/conffwk/schema.pyi diff --git a/typings/conffwk/__init__.pyi b/typings/conffwk/__init__.pyi deleted file mode 100644 index a4eb04769..000000000 --- a/typings/conffwk/__init__.pyi +++ /dev/null @@ -1,35 +0,0 @@ -from __future__ import annotations -import __future__ - -from conffwk.ConfigObject import ConfigObject -from conffwk.Configuration import Configuration - -from . import dal, dalproperty, proxy, schema - -__all__: list[str] = [ - "ConfigObject", - "Configuration", - "absolute_import", - "dal", - "dalproperty", - "proxy", - "reset_updated_dals", - "schema", - "updated_dals", -] - -def reset_updated_dals(): - """ - Reset the set keeping track of modified DAL objects - - """ - -def updated_dals(): - """ - Returns a set of DAL objects that were modified in this DB session - - """ - -absolute_import: ( - __future__._Feature -) # value = _Feature((2, 5, 0, 'alpha', 1), (3, 0, 0, 'alpha', 0), 262144) diff --git a/typings/conffwk/_daq_conffwk_py.pyi b/typings/conffwk/_daq_conffwk_py.pyi deleted file mode 100644 index 1fcb0a6ca..000000000 --- a/typings/conffwk/_daq_conffwk_py.pyi +++ /dev/null @@ -1,377 +0,0 @@ -""" -Python interface to the conffwk package -""" - -from __future__ import annotations - -import typing - -__all__: list[str] = list() - -class _ConfigObject: - def UID(self) -> str: - """ - Return object identity - """ - @typing.overload - def __init__(self) -> None: ... - @typing.overload - def __init__(self, arg0: _ConfigObject) -> None: ... - def class_name(self) -> str: - """ - Return object's class name - """ - def contained_in(self) -> str: - """ - Return the name of the database file this object belongs to. - """ - def full_name(self) -> str: - """ - Return full object name - """ - def get_bool(self, attr: str) -> bool: - """ - Simple getter function - """ - def get_bool_vec(self, attr: str) -> list[bool]: - """ - Getter function for a list - """ - def get_double(self, attr: str) -> float: - """ - Simple getter function - """ - def get_double_vec(self, attr: str) -> list[float]: - """ - Getter function for a list - """ - def get_float(self, attr: str) -> float: - """ - Simple getter function - """ - def get_float_vec(self, attr: str) -> list[float]: - """ - Getter function for a list - """ - def get_obj(self, attrname: str) -> _ConfigObject: - """ - Get a copy of an object - """ - def get_objs(self, attr: str) -> list[_ConfigObject]: - """ - Getter function for a list - """ - def get_s16(self, attr: str) -> int: - """ - Simple getter function - """ - def get_s16_vec(self, attr: str) -> list[int]: - """ - Getter function for a list - """ - def get_s32(self, attr: str) -> int: - """ - Simple getter function - """ - def get_s32_vec(self, attr: str) -> list[int]: - """ - Getter function for a list - """ - def get_s64(self, attr: str) -> int: - """ - Simple getter function - """ - def get_s64_vec(self, attr: str) -> list[int]: - """ - Getter function for a list - """ - def get_s8(self, attr: str) -> int: - """ - Simple getter function - """ - def get_s8_vec(self, attr: str) -> list[int]: - """ - Getter function for a list - """ - def get_string(self, attr: str) -> str: - """ - Simple getter function - """ - def get_string_vec(self, attr: str) -> list[str]: - """ - Getter function for a list - """ - def get_u16(self, attr: str) -> int: - """ - Simple getter function - """ - def get_u16_vec(self, attr: str) -> list[int]: - """ - Getter function for a list - """ - def get_u32(self, attr: str) -> int: - """ - Simple getter function - """ - def get_u32_vec(self, attr: str) -> list[int]: - """ - Getter function for a list - """ - def get_u64(self, attr: str) -> int: - """ - Simple getter function - """ - def get_u64_vec(self, attr: str) -> list[int]: - """ - Getter function for a list - """ - def get_u8(self, attr: str) -> int: - """ - Simple getter function - """ - def get_u8_vec(self, attr: str) -> list[int]: - """ - Getter function for a list - """ - def rename(self, new_id: str) -> None: - """ - Rename object - """ - def set_bool(self, name: str, value: bool) -> None: - """ - Simple setter function - """ - def set_bool_vec(self, attrname: str, l: list[bool]) -> None: - """ - Setter function for list - """ - def set_class(self, name: str, value: str) -> None: - """ - Set the class name - """ - def set_class_vec(self, attrname: str, l: list[str]) -> None: - """ - Set list of classes - """ - def set_date(self, name: str, value: str) -> None: - """ - Set the date - """ - def set_date_vec(self, attrname: str, l: list[str]) -> None: - """ - Set list of dates - """ - def set_double(self, name: str, value: float) -> None: - """ - Simple setter function - """ - def set_double_vec(self, attrname: str, l: list[float]) -> None: - """ - Setter function for list - """ - def set_enum(self, name: str, value: str) -> None: - """ - Set the enum - """ - def set_enum_vec(self, attrname: str, l: list[str]) -> None: - """ - Set list of enums - """ - def set_float(self, name: str, value: float) -> None: - """ - Simple setter function - """ - def set_float_vec(self, attrname: str, l: list[float]) -> None: - """ - Setter function for list - """ - def set_obj( - self, name: str, o: _ConfigObject, skip_non_null_check: bool = False - ) -> None: - """ - Set relationship single-value - """ - def set_objs( - self, name: str, o: list[_ConfigObject], skip_non_null_check: bool = False - ) -> None: - """ - Set relationship multi-value. - """ - def set_s16(self, name: str, value: int) -> None: - """ - Simple setter function - """ - def set_s16_vec(self, attrname: str, l: list[int]) -> None: - """ - Setter function for list - """ - def set_s32(self, name: str, value: int) -> None: - """ - Simple setter function - """ - def set_s32_vec(self, attrname: str, l: list[int]) -> None: - """ - Setter function for list - """ - def set_s64(self, name: str, value: int) -> None: - """ - Simple setter function - """ - def set_s64_vec(self, attrname: str, l: list[int]) -> None: - """ - Setter function for list - """ - def set_s8(self, name: str, value: int) -> None: - """ - Simple setter function - """ - def set_s8_vec(self, attrname: str, l: list[int]) -> None: - """ - Setter function for list - """ - def set_string(self, name: str, value: str) -> None: - """ - Simple setter function - """ - def set_string_vec(self, attrname: str, l: list[str]) -> None: - """ - Set list of strings - """ - def set_time(self, name: str, value: str) -> None: - """ - Set the time - """ - def set_time_vec(self, attrname: str, l: list[str]) -> None: - """ - Set list of times - """ - def set_u16(self, name: str, value: int) -> None: - """ - Simple setter function - """ - def set_u16_vec(self, attrname: str, l: list[int]) -> None: - """ - Setter function for list - """ - def set_u32(self, name: str, value: int) -> None: - """ - Simple setter function - """ - def set_u32_vec(self, attrname: str, l: list[int]) -> None: - """ - Setter function for list - """ - def set_u64(self, name: str, value: int) -> None: - """ - Simple setter function - """ - def set_u64_vec(self, attrname: str, l: list[int]) -> None: - """ - Setter function for list - """ - def set_u8(self, name: str, value: int) -> None: - """ - Simple setter function - """ - def set_u8_vec(self, attrname: str, l: list[int]) -> None: - """ - Setter function for list - """ - -class _Configuration: - @typing.overload - def __init__(self) -> None: ... - @typing.overload - def __init__(self, arg0: str) -> None: ... - def add_include(self, db_name: str, include: str) -> None: - """ - Add include file to existing database. - """ - def attributes(self, class_name: str, all: bool) -> dict[str, dict[str, str]]: - """ - Get the properties of each attribute in a given class - """ - def classes(self) -> list[str]: - """ - Get the names of the superclasses for each class - """ - def commit(self, log_message: str = "") -> None: - """ - Commit database changes. - """ - def create_db(self, db_name: str, includes: list[str]) -> None: - """ - Create a database from a list of files - """ - @typing.overload - def create_obj(self, at: str, class_name: str, id: str) -> object: - """ - Create new object by class name and object id. - """ - @typing.overload - def create_obj(self, at: object, class_name: str, id: str) -> object: - """ - Create new object by class name and object id. - """ - def destroy_obj(self, object: object) -> None: - """ - The method tries to destroy given object. - """ - def get_impl_param(self) -> str: - """ - Get implementation plug-in parameter used to build conffwk object - """ - def get_impl_spec(self) -> str: - """ - Get implementation plug-in and its parameter used to build conffwk object - """ - def get_includes(self, db_name: str) -> list[str]: - """ - Returns list of files included by given database. - """ - def get_obj(self, class_name: str, id: str) -> object: - """ - Create a configuration object containing the desired entity from the database - """ - def get_objs(self, class_name: str, query: str = "") -> list[...]: - """ - Create a list of configuration objects of a given class from the database - """ - def get_schema_path(self, class_name: str) -> str: - """ - Get path to schema file with definition of the given class - """ - def load(self, db_name: str) -> None: - """ - Load database according to the name. - """ - def loaded(self) -> bool: - """ - Check if database is correctly loaded. - """ - def relations(self, class_name: str, all: bool) -> dict[str, dict[str, str]]: - """ - Get the properties of each relation in a given class - """ - def remove_include(self, db_name: str, include: str) -> None: - """ - Remove include file. - """ - def subclasses(self, class_name: str, all: bool) -> list[str]: - """ - Get the subclasses of a single class - """ - def superclasses(self, class_name: str, all: bool) -> list[str]: - """ - Get the superclasses of a single class - """ - def test_object( - self, class_name: str, id: str, rlevel: int, rclasses: list[str] - ) -> bool: - """ - Test the existence of the object - """ - def unload(self) -> None: - """ - Unload previously-loaded database - """ diff --git a/typings/conffwk/dal.pyi b/typings/conffwk/dal.pyi deleted file mode 100644 index 03b68698e..000000000 --- a/typings/conffwk/dal.pyi +++ /dev/null @@ -1,46 +0,0 @@ -from typing import List - -from drunc.fsm._protocols import ( - ConfigurationProtocol, - DBProtocol, - OksKeyProtocol, - ParameterProtocol, -) - -class FSMParameter(ParameterProtocol): - name: str - value: str - -class FSMAction(ConfigurationProtocol): - id: str - name: str - parameters: List[FSMParameter] - db: DBProtocol - oks_key: OksKeyProtocol - initial_data: str - -class FSMxTransition: - transition: str - order: List[str] - mandatory: List[str] - -class FSMTransitionConfig: - id: str - source: str - dest: str - -class FSMCommand: - id: str - -class FSMCommandSequence: - id: str - sequence: List[FSMCommand] - -class FSMData: - states: List[str] - initial_state: str - actions: List[FSMAction] - transitions: List[FSMTransitionConfig] - pre_transitions: List[FSMxTransition] - post_transitions: List[FSMxTransition] - command_sequences: List[FSMCommandSequence] diff --git a/typings/conffwk/dalproperty.pyi b/typings/conffwk/dalproperty.pyi deleted file mode 100644 index 56aecb9de..000000000 --- a/typings/conffwk/dalproperty.pyi +++ /dev/null @@ -1,13 +0,0 @@ -from __future__ import annotations -import __future__ - -__all__: list[str] = ["absolute_import"] - -def _assign_attribute(attribute): ... -def _assign_relation(relation): ... -def _return_attribute(attribute, dalobj=None, value=None): ... -def _return_relation(relation, multi=False, data=None, dalobj=None, cache=None): ... - -absolute_import: ( - __future__._Feature -) # value = _Feature((2, 5, 0, 'alpha', 1), (3, 0, 0, 'alpha', 0), 262144) diff --git a/typings/conffwk/proxy.pyi b/typings/conffwk/proxy.pyi deleted file mode 100644 index a47b05295..000000000 --- a/typings/conffwk/proxy.pyi +++ /dev/null @@ -1,50 +0,0 @@ -""" -Proxing/Delegation tools - -Provide several tools to implement proxying/delegation of objects. The proxying -instances expose the same public interface of the proxied object, but avoiding -inheritance. This allows to control the reference counts of the proxied object. - -""" - -from __future__ import annotations - -__all__: list[str] = ["Proxy", "make_proxy_class"] - -class Proxy: - """ - A very basic holder class - - Just holds the reference to a provided object. - - - """ - def __init__(self, obj): ... - -def _DelegateMetaFunction(clsName, bases, atts): - """ - Implements a delegation pattern using a metaclass approach - - A class using this meta mechanism should have 'memberclass' class attribute - initialized at the class of the instance to proxied. The metaclass will - make sure the delegate class will expose all the public methods of the - proxied one. - Moreover, the metaclass will provide the delegate class with a '__init__' - function instantiating a 'memberclass' object, storing it in 'self._obj'. - The delegate class constructor method will therefore accept all the - arguments accepted by the proxied class constructor. - The delegate class uses slots - - - """ - -def make_proxy_class(theclass): - """ - Builds a delegation class out of a given type. - - Uses the Proxy class to generate a new proxy class exposing the same - interface of the provided class and delegating the method calls to - the hosted object instance. - - - """ diff --git a/typings/conffwk/schema.pyi b/typings/conffwk/schema.pyi deleted file mode 100644 index 4b7b1dd59..000000000 --- a/typings/conffwk/schema.pyi +++ /dev/null @@ -1,144 +0,0 @@ -""" -A set of utilities to simplify OKS instrospection. -""" - -from __future__ import annotations - -import logging as logging -import re as re -import sys as sys - -from conffwk import ConfigObject - -__all__: list[str] = [ - "Cache", - "ConfigObject", - "check_cardinality", - "check_range", - "check_relation", - "coerce", - "decode_range", - "logging", - "map_coercion", - "oks_types", - "range_regexp", - "re", - "str2integer", - "sys", - "to_int", - "to_long", -] - -class Cache: - """ - Defines a cache for all known schemas at a certain time. - - """ - def __getitem__(self, key): - """ - Gets the description of a certain class. - """ - def __init__(self, conffwk, all=True): - """ - Initializes the cache with information from the Configuration - object. - - This method will browse for all declared classes in the Configuration - object given as input and will setup the schema for all known classes. - After this you can still update the cache using the update() method. - - Keyword parameters: - - conffwk -- The conffwk.Configuration object to use as base for the - current cache. - - all -- A boolean indicating if I should store all the attributes and - relations from a certain class or just the ones directly associated - with a class. - - """ - def __str__(self): - """ - Prints a nice display of myself - """ - def update(self, conffwk): - """ - Updates this cache with information from the Configuration object. - - This method will add new classes not yet know to this cache. Classes - with existing names will not be added. No warning is generated (this - should be done by the OKS layer in any case. - - """ - def update_dal(self, conffwk): - """ - Updates this cache with information for DAL. - - This method will add new DAL classes not yet know to this cache. - Classes with existing DAL representations will not be touched. - - """ - -def check_cardinality(v, prop): - """ - Checks the cardinality of a certain attribute or relationship. - """ - -def check_range(v, range, range_re, pytype): - """ - Checks the range of the value 'v' to make sure it is inside. - """ - -def check_relation(v, rel): - """ - Checks the value v against the relationship parameters in 'rel'. - """ - -def coerce(v, attr): - """ - Coerces the input value 'v' in the way the attribute expects. - """ - -def decode_range(s): - """ - Decodes a range string representation, returns a tuple with 2 values. - - This is the supported format in regexp representation: - '([-0x]*\\d+)\\D+-?\\d+' - - """ - -def map_coercion(class_name, schema): - """ - Given a schema of a class, maps coercion functions from libpyconffwk. - """ - -def str2integer(v, t, max): - """ - Converts a value v to integer, irrespectively of its formatting. - - If the number starts with a '0', we convert it using an octal - representation. Else, we try a decimal conversion. If any of these fail, - we try an hexa conversion before throwing a ValueError. - - Keyword arguments: - - v -- the value to be converted - t -- the python type (int or float) to use in the conversion - - """ - -def to_int(v): ... -def to_long(v): ... - -oks_types: dict = { - "bool": ["bool"], - "integer": ["s8", "u8", "s16", "u16", "s32"], - "long": ["u32", "s64", "u64"], - "float": ["float", "double"], - "int-number": ["s8", "u8", "s16", "u16", "s32", "u32", "s64", "u64"], - "number": ["u32", "s64", "u64", "s8", "u8", "s16", "u16", "s32", "float", "double"], - "time": ["date", "time"], - "string": ["date", "time", "string", "uid", "enum", "class"], -} -range_regexp: re.Pattern # value = re.compile('(?P-?0?x?[\\da-fA-F]+(\\.\\d+)?)-(?P-?0?x?[\\da-fA-F]+(\\.\\d+)?)') From 1b9925e5db357177e3687e65e59714b3d6d640fe Mon Sep 17 00:00:00 2001 From: James Paul Turner Date: Mon, 13 Jul 2026 17:57:11 +0100 Subject: [PATCH 12/20] Phased merge 6 --- typings/daqconf/__init__.py | 0 typings/daqconf/consolidate.pyi | 7 ------- typings/daqconf/jsonify.pyi | 4 ---- typings/daqconf/validate.pyi | 4 ---- 4 files changed, 15 deletions(-) delete mode 100644 typings/daqconf/__init__.py delete mode 100644 typings/daqconf/consolidate.pyi delete mode 100644 typings/daqconf/jsonify.pyi delete mode 100644 typings/daqconf/validate.pyi diff --git a/typings/daqconf/__init__.py b/typings/daqconf/__init__.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/typings/daqconf/consolidate.pyi b/typings/daqconf/consolidate.pyi deleted file mode 100644 index cd7c719f4..000000000 --- a/typings/daqconf/consolidate.pyi +++ /dev/null @@ -1,7 +0,0 @@ -from typing import Optional - -def consolidate_db( - oksfile: str, - output_file: str, - session_id: Optional[str] = None, -) -> None: ... diff --git a/typings/daqconf/jsonify.pyi b/typings/daqconf/jsonify.pyi deleted file mode 100644 index 64766ceed..000000000 --- a/typings/daqconf/jsonify.pyi +++ /dev/null @@ -1,4 +0,0 @@ -def jsonify_xml_data( - oksfile: str, - output: str, -) -> None: ... diff --git a/typings/daqconf/validate.pyi b/typings/daqconf/validate.pyi deleted file mode 100644 index d6d52894f..000000000 --- a/typings/daqconf/validate.pyi +++ /dev/null @@ -1,4 +0,0 @@ -def validate_session( - oksfile: str, - session_name: str, -) -> None: ... From cf717344d2e38c206b04442efc68d5905e266894 Mon Sep 17 00:00:00 2001 From: James Paul Turner Date: Mon, 13 Jul 2026 18:02:22 +0100 Subject: [PATCH 13/20] Phased merge 7 --- .../children_interface/grpc_child.py | 30 +-- .../children_interface/rest_api_child.py | 41 ++-- src/drunc/controller/configuration.py | 50 ++-- src/drunc/controller/controller.py | 62 +---- src/drunc/controller/interface/commands.py | 30 ++- src/drunc/controller/interface/context.py | 9 - src/drunc/controller/interface/controller.py | 7 +- src/drunc/controller/interface/shell_utils.py | 122 +++++++--- src/drunc/controller/stateful_node.py | 8 - src/drunc/controller/utils.py | 13 +- tests/broadcast/__init__.py | 0 tests/broadcast/client/__init__.py | 0 .../client/test_broadcast_handler.py | 0 .../test_broadcast_handler_implementation.py | 0 tests/broadcast/client/test_configuration.py | 0 .../test_grpc_stdout_broadcast_handler.py | 0 .../test_kafka_stdout_broadcast_handler.py | 0 tests/broadcast/server/__init__.py | 0 .../broadcast/server/test_broadcast_sender.py | 0 .../test_broadcast_sender_implementation.py | 0 tests/broadcast/server/test_configuration.py | 0 tests/broadcast/server/test_decorators.py | 72 ------ tests/broadcast/server/test_grpc_servicer.py | 0 tests/broadcast/server/test_kafka_sender.py | 93 -------- tests/broadcast/test_types.py | 0 tests/broadcast/test_utils.py | 0 tests/issues/test_issue309.py | 11 +- tests/issues/test_issue363.py | 13 +- .../process_manager_mock_impls.py | 6 +- tests/processes/conftest.py | 59 +++++ tests/processes/test_process_metadata.py | 48 ++++ ...est_ssh_process_lifetime_manager_common.py | 191 ++++++++++++--- ...h_process_lifetime_manager_forked_shell.py | 223 ++++++++++++++---- ...t_ssh_process_lifetime_manager_paramiko.py | 28 ++- ...test_ssh_process_lifetime_manager_shell.py | 30 ++- 35 files changed, 672 insertions(+), 474 deletions(-) delete mode 100644 tests/broadcast/__init__.py delete mode 100644 tests/broadcast/client/__init__.py delete mode 100644 tests/broadcast/client/test_broadcast_handler.py delete mode 100644 tests/broadcast/client/test_broadcast_handler_implementation.py delete mode 100644 tests/broadcast/client/test_configuration.py delete mode 100644 tests/broadcast/client/test_grpc_stdout_broadcast_handler.py delete mode 100644 tests/broadcast/client/test_kafka_stdout_broadcast_handler.py delete mode 100644 tests/broadcast/server/__init__.py delete mode 100644 tests/broadcast/server/test_broadcast_sender.py delete mode 100644 tests/broadcast/server/test_broadcast_sender_implementation.py delete mode 100644 tests/broadcast/server/test_configuration.py delete mode 100644 tests/broadcast/server/test_decorators.py delete mode 100644 tests/broadcast/server/test_grpc_servicer.py delete mode 100644 tests/broadcast/server/test_kafka_sender.py delete mode 100644 tests/broadcast/test_types.py delete mode 100644 tests/broadcast/test_utils.py create mode 100644 tests/processes/test_process_metadata.py diff --git a/src/drunc/controller/children_interface/grpc_child.py b/src/drunc/controller/children_interface/grpc_child.py index 54c7f6f0d..ee216e573 100644 --- a/src/drunc/controller/children_interface/grpc_child.py +++ b/src/drunc/controller/children_interface/grpc_child.py @@ -35,15 +35,13 @@ from druncschema.token_pb2 import Token from grpc_status import rpc_status -from drunc.broadcast.client.broadcast_handler import BroadcastHandler -from drunc.broadcast.client.configuration import BroadcastClientConfHandler from drunc.connectivity_service.exceptions import ( ApplicationLookupUnsuccessful, ) from drunc.controller.children_interface.child_node import ChildNode from drunc.exceptions import DruncSetupException from drunc.grpc_settings import CONTROLLER_CLIENT_GRPC_CONFIG -from drunc.utils.configuration import ConfHandler, ConfTypes +from drunc.utils.configuration import ConfHandler from drunc.utils.grpc_utils import ( ServerUnreachable, rethrow_if_unreachable_server, @@ -56,12 +54,17 @@ class gRCPChildConfHandler(ConfHandler): + """Handler for gRPC child node configuration.""" + + def _post_process_oks(self) -> None: + self.controller = self._raw_data.controller + def get_uri(self): - for service in self.data.controller.exposes_service: - if self.data.controller.id + "_control" in service.id: - return f"{service.protocol}://{self.data.controller.runs_on.runs_on.id}:{service.port}" + for service in self.controller.exposes_service: + if self.controller.id + "_control" in service.id: + return f"{service.protocol}://{self.controller.runs_on.runs_on.id}:{service.port}" raise DruncSetupException( - f"gRPC API child node {self.data.controller.id} does not expose a control service" + f"gRPC API child node {self.controller.id} does not expose a control service" ) @@ -116,7 +119,7 @@ def _setup_connection(self): tries_remaining -= 1 try: - response = self.stub.describe(request) + self.stub.describe(request) except grpc.RpcError as error: if tries_remaining == 0: @@ -131,7 +134,6 @@ def _setup_connection(self): else: self.log.info(f"Connected to the controller ({self.uri})!") - self.start_listening(response.description.broadcast) break def _attempt_reconnection(self, retry_call): @@ -185,9 +187,7 @@ def terminate(self) -> None: del self.channel if self.stub: del self.stub - self.channel = None - self.broadcast.stop() def check_connection(self) -> bool: """Probe child connectivity and retry once after reconnecting if needed. @@ -233,14 +233,6 @@ def check_connection(self) -> bool: return False - def start_listening(self, bdesc): - self.broadcast = BroadcastHandler( - BroadcastClientConfHandler( - data=bdesc, - type=ConfTypes.ProtobufAny, - ) - ) - def status( self, target: str = "", diff --git a/src/drunc/controller/children_interface/rest_api_child.py b/src/drunc/controller/children_interface/rest_api_child.py index 9eab6130d..aa28a5951 100644 --- a/src/drunc/controller/children_interface/rest_api_child.py +++ b/src/drunc/controller/children_interface/rest_api_child.py @@ -355,12 +355,24 @@ def check_response(self, timeout: int = 0) -> dict: class RESTAPIChildNodeConfHandler(ConfHandler): + """Handler for REST API child node configuration.""" + + def _post_process_oks(self) -> None: + raw = self._raw_data + # Flatten attributes that are always accessed from outside + self.id = getattr(raw, "id", None) + self.exposes_service = getattr(raw, "exposes_service", []) + self.runs_on = getattr(raw, "runs_on", None) + self.proxy = getattr(raw, "proxy", [None, None]) + # Pass through any other attributes callers may inspect dynamically + self._raw = raw + def get_host_port(self): - for service in self.data.exposes_service: - if self.data.id + "_control" in service.id: - return self.data.runs_on.runs_on.id, service.port + for service in self.exposes_service: + if self.id + "_control" in service.id: + return self.runs_on.runs_on.id, service.port raise DruncSetupException( - f"REST API child node {self.data.id} does not expose a control service" + f"REST API child node {self.id} does not expose a control service" ) @@ -379,7 +391,7 @@ def __init__( self.fsm_configuration = fsm_configuration self.connectivity_service = connectivity_service if fsm_configuration: - fsmch = FSMConfHandler(fsm_configuration) + fsmch = FSMConfHandler.from_pyobject(data=fsm_configuration) self.fsm = FSM(conf=fsmch) response_listener_host = socket.gethostname() @@ -392,7 +404,7 @@ def __init__( f"Application {name} does not expose a control service in the configuration, or has not advertised itself to the application registry service, or the application registry service is not reachable." ) - proxy_host, proxy_port = getattr(self.configuration.data, "proxy", [None, None]) + proxy_host, proxy_port = self.configuration.proxy proxy_port = int(proxy_port) if proxy_port is not None else None self.commander = AppCommander( @@ -542,16 +554,15 @@ def describe( if self.configuration is not None: if detector_name := get_detector_name(self.configuration): description.info = detector_name - if hasattr( - self.configuration.data, "application_name" - ): # Application nodes. - description.type = self.configuration.data.application_name - description.name = self.configuration.data.id - elif hasattr(self.configuration.data, "controller") and hasattr( - self.configuration.data.controller, "application_name" + raw = self.configuration._raw + if hasattr(raw, "application_name"): # Application nodes. + description.type = raw.application_name + description.name = raw.id + elif hasattr(raw, "controller") and hasattr( + raw.controller, "application_name" ): # Controller nodes. - description.type = self.configuration.data.controller.application_name - description.name = self.configuration.data.controller.id + description.type = raw.controller.application_name + description.name = raw.controller.id response.description.CopyFrom(description) diff --git a/src/drunc/controller/configuration.py b/src/drunc/controller/configuration.py index 64944fa8c..6d3c93513 100644 --- a/src/drunc/controller/configuration.py +++ b/src/drunc/controller/configuration.py @@ -21,7 +21,7 @@ ) from drunc.exceptions import DruncCommandException, DruncSetupException from drunc.process_manager.configuration import get_commandline_parameters -from drunc.utils.configuration import ConfHandler, ConfTypes +from drunc.utils.configuration import ConfHandler from drunc.utils.utils import ( ControlType, get_control_type_and_uri_from_cli, @@ -30,20 +30,9 @@ ) -class ControllerConfData: # the bastardised OKS - def __init__(self): - class id_able: - id = None - - class cler: - pass - - self.controller = cler() - self.controller.broadcaster = id_able() - self.controller.fsm = id_able() - - class ControllerConfHandler(ConfHandler): + """Handler for controller configuration.""" + @staticmethod def find_segment(segment, id_): if segment.controller.id == id_: @@ -67,30 +56,29 @@ def _grab_segment_conf_from_controller(self, configuration): ) return this_segment - def _post_process_oks(self, *args, **kwargs): + def _post_process_oks(self) -> None: self.authoriser = None - self.data = self._grab_segment_conf_from_controller(self.data) + segment = self._grab_segment_conf_from_controller(self._raw_data) + self.controller = segment.controller + self.segments = segment.segments + self.applications = segment.applications - self.this_host = self.data.controller.runs_on.runs_on.id + self.this_host = self.controller.runs_on.runs_on.id if self.this_host in ["localhost"] or self.this_host.startswith("127."): self.this_host = socket.gethostname() self.opmon_publisher = None self.opmon_conf = parse_opmon_conf( log=self.log, - conf=self.data.controller.opmon_conf, + conf=self.controller.opmon_conf, uri=self.session.opmon_uri, session=self.session_name, - application=self.data.controller.id, + application=self.controller.id, ) if self.opmon_conf.path == "./info.json": self.opmon_conf.path = ( - "./info." - + self.opmon_conf.session - + "." - + self.data.controller.id - + ".json" + "./info." + self.opmon_conf.session + "." + self.controller.id + ".json" ) self.log.debug("Initializing OpMon with configuration %s", self.opmon_conf) @@ -194,13 +182,13 @@ def process_application(app): # threading the children look up threads = [] - for segment in self.data.segments: + for segment in self.segments: self.log.debug(segment) t = threading.Thread(target=process_segment, args=(segment,)) threads.append(t) t.start() - for app in self.data.applications: + for app in self.applications: self.log.debug(app) t = threading.Thread(target=process_application, args=(app,)) threads.append(t) @@ -252,22 +240,22 @@ def child_node_factory( match ctype: case ControlType.gRPC: - grpc_conf_handler = gRCPChildConfHandler( - configuration, ConfTypes.PyObject + grpc_conf_handler = gRCPChildConfHandler.from_pyobject( + data=configuration ) return gRPCChildNode( name, grpc_conf_handler, uri, connectivity_service, init_token ) case ControlType.REST_API: - restapi_conf_handler = RESTAPIChildNodeConfHandler( - configuration, ConfTypes.PyObject + restapi_conf_handler = RESTAPIChildNodeConfHandler.from_pyobject( + data=configuration ) return RESTAPIChildNode( name, restapi_conf_handler, uri, - self.data.controller.fsm, + self.controller.fsm, connectivity_service=connectivity_service, ) diff --git a/src/drunc/controller/controller.py b/src/drunc/controller/controller.py index c73f8174b..7e076fbcc 100644 --- a/src/drunc/controller/controller.py +++ b/src/drunc/controller/controller.py @@ -6,7 +6,6 @@ from daqpytools.logging import LogHandlerConf, setup_daq_ers_logger from druncschema.authoriser_pb2 import ActionType, SystemType -from druncschema.broadcast_pb2 import BroadcastType from druncschema.controller_pb2 import ( DescribeFSMRequest, DescribeFSMResponse, @@ -47,9 +46,6 @@ from drunc.authoriser.configuration import DummyAuthoriserConfHandler from drunc.authoriser.decorators import authentified_and_authorised from drunc.authoriser.dummy_authoriser import DummyAuthoriser -from drunc.broadcast.server.broadcast_sender import BroadcastSender -from drunc.broadcast.server.configuration import BroadcastSenderConfHandler -from drunc.broadcast.server.decorators import broadcasted from drunc.connectivity_service.client import ConnectivityServiceClient from drunc.connectivity_service.exceptions import ApplicationLookupUnsuccessful from drunc.controller.children_interface.child_node import ChildNode @@ -85,7 +81,6 @@ def __init__(self, configuration, name: str, session: str, token: Token): self._previous_error_state = False self.name = name self.session = session - self.broadcast_service = None self.monitoring_metrics = ControllerMonitoringMetrics() self.handlerconf = LogHandlerConf(init_ers=True) self.log = get_logger(f"controller.core.{name}_ctrl") @@ -109,18 +104,9 @@ def __init__(self, configuration, name: str, session: str, token: Token): self.opmon_publisher = getattr(self.configuration, "opmon_publisher", None) self.stop_event: threading.Event | None = None self.thread: threading.Thread | None = None - bsch = BroadcastSenderConfHandler( - data=self.configuration.data.controller.broadcaster, - ) - self.broadcast_service = BroadcastSender( - name=name, - session=session, - configuration=bsch, - ) - - self.fsm_config = FSMConfHandler( - data=self.configuration.data.controller.fsm, + self.fsm_config = FSMConfHandler.from_pyobject( + data=self.configuration.controller.fsm, ) self.stateful_node = StatefulNode( @@ -132,10 +118,9 @@ def __init__(self, configuration, name: str, session: str, token: Token): top_segment_controller=self.top_segment_controller, ) - dach = DummyAuthoriserConfHandler( + dach = DummyAuthoriserConfHandler.from_pyobject( data=self.configuration.authoriser, ) - self.authoriser = DummyAuthoriser(dach, SystemType.CONTROLLER) self.actor = ControllerActor(token) @@ -245,7 +230,7 @@ def init_controller(self) -> None: log_init_controller.info(f"Taking control of {child.name}") child.take_control(execute_on_all_subsequent_children_in_path=True) - interval_s = getattr(self.configuration.data, "interval_s", 10.0) + interval_s = getattr(self.configuration, "interval_s", 10.0) if self.opmon_publisher is not None: self.stop_event = threading.Event() @@ -256,28 +241,9 @@ def init_controller(self) -> None: ) self.thread.start() - self.broadcast(message="ready", btype=BroadcastType.SERVER_READY) self.stateful_node.set_ready_state(True) log_init_controller.info("Controller ready") - """ - A couple of simple pass-through functions to the broadcasting service - """ - - def broadcast(self, *args, **kwargs): - return self.broadcast_service.broadcast(*args, **kwargs) - - def can_broadcast(self, *args, **kwargs): - if self.broadcast_service: - return self.broadcast_service.can_broadcast(*args, **kwargs) - return False - - def describe_broadcast(self, *args, **kwargs): - return self.broadcast_service.describe_broadcast(*args, **kwargs) - - def interrupt_with_exception(self, *args, **kwargs): - return self.broadcast_service._interrupt_with_exception(*args, **kwargs) - def controller_publisher(self, message, custom_origin: dict | None = None): if isinstance(message, FSMStatus) and message.in_error: if message.in_error and not self._previous_error_state: @@ -394,12 +360,6 @@ def terminate(self): self.log.info("Unregistering from the connectivity service") self.connectivity_service.retract(self.name + "_control", fail_quickly=True) - if self.can_broadcast(): - self.broadcast( - btype=BroadcastType.SERVER_SHUTDOWN, - message="over_and_out", - ) - self.log.info("Stopping children") for child in self.children_nodes: self.log.debug(f"Stopping {child.name}") @@ -625,7 +585,6 @@ def _partition_connected_children( ############# Status, description commands ############# ######################################################## - @broadcasted @authentified_and_authorised(action=ActionType.READ, system=SystemType.CONTROLLER) @publish_command_time def status( @@ -688,7 +647,6 @@ def status( return response - @broadcasted @authentified_and_authorised(action=ActionType.READ, system=SystemType.CONTROLLER) @publish_command_time def describe( @@ -717,8 +675,6 @@ def describe( session=self.session, commands=None, ) - if broadcast_description := self.describe_broadcast(): - description.broadcast.Pack(broadcast_description) response.description.CopyFrom(description) # Children nodes (ignore exclusion). @@ -754,7 +710,6 @@ def describe( return response - @broadcasted @authentified_and_authorised(action=ActionType.READ, system=SystemType.CONTROLLER) @publish_command_time def describe_fsm( @@ -837,7 +792,6 @@ def describe_fsm( ############# FSM commands ############# ######################################## - @broadcasted @authentified_and_authorised(action=ActionType.UPDATE, system=SystemType.CONTROLLER) @in_control @publish_command_time @@ -1016,7 +970,6 @@ def execute_fsm_command( return response - @broadcasted @authentified_and_authorised(action=ActionType.EXPERT, system=SystemType.CONTROLLER) @in_control @publish_command_time @@ -1076,7 +1029,6 @@ def execute_expert_command( return response - @broadcasted @authentified_and_authorised(action=ActionType.UPDATE, system=SystemType.CONTROLLER) @in_control @publish_command_time @@ -1142,7 +1094,6 @@ def include( return response - @broadcasted @authentified_and_authorised(action=ActionType.UPDATE, system=SystemType.CONTROLLER) @in_control @publish_command_time @@ -1208,7 +1159,6 @@ def exclude( return response - @broadcasted @authentified_and_authorised(action=ActionType.UPDATE, system=SystemType.CONTROLLER) @in_control @publish_command_time @@ -1337,7 +1287,6 @@ def recompute_status( ############# Actor commands ############# ########################################## - @broadcasted @authentified_and_authorised(action=ActionType.UPDATE, system=SystemType.CONTROLLER) @publish_command_time def take_control( @@ -1413,7 +1362,6 @@ def take_control( return response - @broadcasted @authentified_and_authorised(action=ActionType.UPDATE, system=SystemType.CONTROLLER) @in_control @publish_command_time @@ -1490,7 +1438,6 @@ def surrender_control( return response - @broadcasted @authentified_and_authorised(action=ActionType.READ, system=SystemType.CONTROLLER) @publish_command_time def who_is_in_charge( @@ -1552,7 +1499,6 @@ def who_is_in_charge( ####### Integration test commands ######## ########################################## - @broadcasted @authentified_and_authorised(action=ActionType.UPDATE, system=SystemType.CONTROLLER) @in_control @publish_command_time diff --git a/src/drunc/controller/interface/commands.py b/src/drunc/controller/interface/commands.py index f054a8e29..26f60972b 100644 --- a/src/drunc/controller/interface/commands.py +++ b/src/drunc/controller/interface/commands.py @@ -4,7 +4,7 @@ import click from drunc.controller.interface.context import ControllerContext -from drunc.controller.interface.shell_utils import controller_setup, get_status_table +from drunc.controller.interface.shell_utils import controller_setup, render_status_table from drunc.utils.utils import get_logger log = get_logger("controller.iface", rich_handler=True) @@ -70,25 +70,29 @@ def wait(obj: ControllerContext, sleep_time: int) -> None: help="Execute the command on all subsequent children in the path", default=True, ) +@click.option( + "--extended", + is_flag=True, + default=False, + help="Show additional columns, including the IP address of each endpoint.", +) @click.pass_obj def status( obj: ControllerContext, target: str, execute_along_path: bool, execute_on_all_subsequent_children_in_path: bool, + extended: bool, ) -> None: - statuses = obj.get_driver("controller").status( - target=target, - execute_along_path=execute_along_path, - execute_on_all_subsequent_children_in_path=execute_on_all_subsequent_children_in_path, - ) # Get the dynamic system information - descriptions = obj.get_driver("controller").describe( - target=target, - execute_along_path=execute_along_path, - execute_on_all_subsequent_children_in_path=execute_on_all_subsequent_children_in_path, - ) # Get the static system information - t = get_status_table(statuses, descriptions) - obj.print(t) + obj.print( + render_status_table( + obj, + target=target, + execute_along_path=execute_along_path, + execute_on_all_subsequent_children_in_path=execute_on_all_subsequent_children_in_path, + show_ip_address=extended, + ) + ) obj.print_status_summary() diff --git a/src/drunc/controller/interface/context.py b/src/drunc/controller/interface/context.py index 082c42a00..8da69f38e 100644 --- a/src/drunc/controller/interface/context.py +++ b/src/drunc/controller/interface/context.py @@ -2,10 +2,7 @@ from druncschema.token_pb2 import Token -from drunc.broadcast.client.broadcast_handler import BroadcastHandler -from drunc.broadcast.client.configuration import BroadcastClientConfHandler from drunc.controller.controller_driver import ControllerDriver -from drunc.utils.configuration import ConfTypes from drunc.utils.shell_utils import ( ShellContext, create_dummy_token_from_uname, @@ -33,12 +30,6 @@ def create_drivers(self, **kwargs) -> Mapping[str, object]: def create_token(self, **kwargs) -> Token: return create_dummy_token_from_uname() - def start_listening_controller(self, broadcaster_conf): - bcch = BroadcastClientConfHandler( - data=broadcaster_conf, type=ConfTypes.ProtobufAny - ) - self.status_receiver = BroadcastHandler(broadcast_configuration=bcch) - def terminate(self): if self.status_receiver: self.status_receiver.stop() diff --git a/src/drunc/controller/interface/controller.py b/src/drunc/controller/interface/controller.py index 4eadef8a2..a4b3ff059 100644 --- a/src/drunc/controller/interface/controller.py +++ b/src/drunc/controller/interface/controller.py @@ -14,7 +14,7 @@ CONTROLLER_SERVER_GRPC_CONFIG, CONTROLLER_SERVER_GRPC_MAX_WORKERS, ) -from drunc.utils.configuration import ConfTypes, OKSKey +from drunc.utils.configuration import OKSKey from drunc.utils.utils import ( get_logger, get_root_logger, @@ -86,9 +86,8 @@ def controller_cli( token="", ) - controller_configuration = ControllerConfHandler( - type=ConfTypes.OKSFileName, - data=configurationservice, + controller_configuration = ControllerConfHandler.from_oks( + url=configurationservice, oks_key=OKSKey( schema_file="schema/confmodel/dunedaq.schema.xml", class_name="RCApplication", diff --git a/src/drunc/controller/interface/shell_utils.py b/src/drunc/controller/interface/shell_utils.py index 3ffba7455..ec5529baa 100644 --- a/src/drunc/controller/interface/shell_utils.py +++ b/src/drunc/controller/interface/shell_utils.py @@ -38,6 +38,7 @@ ) from rich.table import Table +from drunc.controller.interface.context import ControllerContext from drunc.exceptions import DruncSetupException, DruncShellException from drunc.unified_shell.context import UnifiedShellContext, UnifiedShellMode from drunc.utils.grpc_utils import ( @@ -71,7 +72,10 @@ def match_children( def get_status_table( - status_response: StatusResponse, describe_response: DescribeResponse + status_response: StatusResponse, + describe_response: DescribeResponse, + display_host_overrides: dict[str, str] | None = None, + show_ip_address: bool = False, ): status = status_response.status description = describe_response.description @@ -90,6 +94,8 @@ def get_status_table( t.add_column("In error") t.add_column("Included") t.add_column("Endpoint") + if show_ip_address: + t.add_column("IP Address") def add_status_to_table( table: Table, @@ -102,34 +108,64 @@ def add_status_to_table( if status is None or description is None: return - def update_endpoint(endpoint: str) -> str: + def update_endpoint(endpoint: str, proc_name: str) -> tuple[str, str]: """ - Parses endpoint to a human readable hostname + Parses endpoint to a human readable hostname. Args: - endpoint: Process URI + endpoint: The endpoint to parse + proc_name: The name of the process to parse the endpoint for Returns: - str: URI with human readable hostname + tuple[str, str]: (display_endpoint, actual_endpoint) + display_endpoint: URI with human readable hostname + actual_endpoint: raw URI with actual IP/host (empty if same as display) """ if not endpoint: - return "" + return "", "" - ip_address = urlparse(endpoint).hostname - if not ip_address: - return "" - resolved_host = get_hostname_smart(ip_address) - return endpoint.replace(ip_address, resolved_host) + parsed = urlparse(endpoint) + raw_host = parsed.hostname + if not raw_host: + return "", "" - table.add_row( + scheme = parsed.scheme + port = parsed.port + + def make_uri(host: str) -> str: + uri = f"{scheme}://{host}" + if port is not None: + uri = f"{uri}:{port}" + return uri + + if display_host_overrides and proc_name in display_host_overrides: + display_host = get_hostname_smart(display_host_overrides[proc_name]) + pretty = make_uri(display_host) + if display_host != raw_host: + return pretty, make_uri(raw_host) + return pretty, "" + + resolved = get_hostname_smart(raw_host) + if resolved != raw_host: + return make_uri(resolved), endpoint + + return endpoint, "" + + display_ep, actual_ep = update_endpoint( + description.endpoint, status_response.name + ) + row = [ prefix + status_response.name, description.info, status.state, status.sub_state, format_bool(status.in_error, false_is_good=True), format_bool(status.included), - update_endpoint(description.endpoint), - ) + display_ep, + ] + if show_ip_address: + row.append(actual_ep) + table.add_row(*row) children = match_children(status_response.children, describe_response.children) children_list = sorted(list(children.keys())) @@ -180,6 +216,32 @@ def add_runinfo_to_table(table: Table, status: Status): return t +def render_status_table( + ctx: ControllerContext, + target: str = "", + execute_along_path: bool = True, + execute_on_all_subsequent_children_in_path: bool = True, + show_ip_address: bool = False, +): + statuses = ctx.get_driver("controller").status( + target=target, + execute_along_path=execute_along_path, + execute_on_all_subsequent_children_in_path=execute_on_all_subsequent_children_in_path, + ) + descriptions = ctx.get_driver("controller").describe( + target=target, + execute_along_path=execute_along_path, + execute_on_all_subsequent_children_in_path=execute_on_all_subsequent_children_in_path, + ) + display_host_overrides = ctx.get_endpoint_display_host_overrides() + return get_status_table( + statuses, + descriptions, + display_host_overrides=display_host_overrides, + show_ip_address=show_ip_address, + ) + + class StatusTableUpdater(Progress): def __init__(self, ctx, refresh_per_second=2, *args, **kwargs) -> None: self.ctx = ctx @@ -195,12 +257,7 @@ def __init__(self, ctx, refresh_per_second=2, *args, **kwargs) -> None: super().__init__(*args, refresh_per_second=refresh_per_second, **kwargs) def update_table(self): - # The following debug log line will be used in an integration test to validate - # that issue 817 does not appear again (rich table overriding the log entries) - self.ctx.log.debug("Updating the status table...") - statuses = self.ctx.get_driver("controller").status() - descriptions = self.ctx.get_driver("controller").describe() - self.table = get_status_table(statuses, descriptions) + self.table = render_status_table(self.ctx) def get_renderable(self) -> ConsoleRenderable | RichCast | str: renderable = Group(self.table, *self.get_renderables()) @@ -210,7 +267,6 @@ def get_renderable(self) -> ConsoleRenderable | RichCast | str: def controller_cleanup_wrapper(ctx): def controller_cleanup(): log = logging.getLogger("controller.shell_utils") - # remove the shell from the controller broadcast list dead = False who = "" @@ -291,8 +347,6 @@ def controller_setup(ctx, controller_address): f"{controller_address} is '{desc.name}.{desc.session}' (name.session), starting listening..." ) ctx.get_driver("controller").name = f"{desc.name}.{desc.session}" - if desc.HasField("broadcast"): - ctx.start_listening_controller(desc.broadcast) log.debug("Connected to the controller") @@ -689,10 +743,7 @@ def add_to_table(table, response, prefix=""): add_to_table(t, result) obj.print(t) # rich tables require console printing - statuses = obj.get_driver("controller").status() - descriptions = obj.get_driver("controller").describe() - t = get_status_table(statuses, descriptions) - obj.print(t) + obj.print(render_status_table(obj)) obj.print_status_summary() @@ -797,20 +848,27 @@ def generate_fsm_command(ctx, transition: FSMCommandDescription, controller_name @functools.lru_cache(maxsize=4096) def get_hostname_smart(ip_or_host: str, timeout_seconds: float = 0.2) -> str: """ - Resolves an IP to a hostname, with optimizations: + Resolves an IP or hostname to a human-readable hostname, with optimizations: 1. Caches all results. - 2. Immediately skips private/internal IPs (like K8s). - 3. Uses a short timeout for public IPs. + 2. Replaces localhost/loopback with the machine's actual hostname. + 3. Uses a short timeout for reverse DNS lookups on other IPs. """ if not ip_or_host: return "" + if ip_or_host == "localhost": + return socket.getfqdn() + try: ip_address = ipaddress.ip_address(ip_or_host) except ValueError: - return ip_or_host - # If public IP, try to resolve it. + fqdn = socket.getfqdn(ip_or_host) + return fqdn if fqdn else ip_or_host + + if ip_address.is_loopback: + return socket.getfqdn() + original_timeout = socket.getdefaulttimeout() try: socket.setdefaulttimeout(timeout_seconds) diff --git a/src/drunc/controller/stateful_node.py b/src/drunc/controller/stateful_node.py index 007bd200f..d503c63fa 100644 --- a/src/drunc/controller/stateful_node.py +++ b/src/drunc/controller/stateful_node.py @@ -20,14 +20,6 @@ def value(self): @value.setter def value(self, value): - # if self._broadcast_on_change is None or self._broadcast_key is None: - # self._value = value - # return - - # self._broadcast_on_change.broadcast( - # message = f'Changing {self._name} from {self._value} to {value}', - # btype = self._broadcast_key, - # ) self._value = value if self.stateful_node: self.stateful_node.log.info(f"{self._name} changed to {value}") diff --git a/src/drunc/controller/utils.py b/src/drunc/controller/utils.py index d9c0b9994..82970e6ea 100644 --- a/src/drunc/controller/utils.py +++ b/src/drunc/controller/utils.py @@ -49,17 +49,16 @@ def get_status_message(controller): def get_detector_name(configuration) -> str: detector_name = None log = get_logger("controller.core.get_detector_name") - if hasattr(configuration.data, "contains") and len(configuration.data.contains) > 0: - if len(configuration.data.contains) > 0: + raw = getattr(configuration, "_raw", None) + if raw is not None and hasattr(raw, "contains") and len(raw.contains) > 0: + if len(raw.contains) > 0: log.debug( - f"Application {configuration.data.id} has multiple contains, using the first one" + f"Application {raw.id} has multiple contains, using the first one" ) - detector_name = ( - configuration.data.contains[0].id.replace("-", "_").replace("_", " ") - ) + detector_name = raw.contains[0].id.replace("-", "_").replace("_", " ") else: log.debug( - f'Application {configuration.data.id} has no "contains" relation, hence no detector' + f'Application {getattr(raw, "id", "?")} has no "contains" relation, hence no detector' ) return detector_name diff --git a/tests/broadcast/__init__.py b/tests/broadcast/__init__.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/tests/broadcast/client/__init__.py b/tests/broadcast/client/__init__.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/tests/broadcast/client/test_broadcast_handler.py b/tests/broadcast/client/test_broadcast_handler.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/tests/broadcast/client/test_broadcast_handler_implementation.py b/tests/broadcast/client/test_broadcast_handler_implementation.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/tests/broadcast/client/test_configuration.py b/tests/broadcast/client/test_configuration.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/tests/broadcast/client/test_grpc_stdout_broadcast_handler.py b/tests/broadcast/client/test_grpc_stdout_broadcast_handler.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/tests/broadcast/client/test_kafka_stdout_broadcast_handler.py b/tests/broadcast/client/test_kafka_stdout_broadcast_handler.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/tests/broadcast/server/__init__.py b/tests/broadcast/server/__init__.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/tests/broadcast/server/test_broadcast_sender.py b/tests/broadcast/server/test_broadcast_sender.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/tests/broadcast/server/test_broadcast_sender_implementation.py b/tests/broadcast/server/test_broadcast_sender_implementation.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/tests/broadcast/server/test_configuration.py b/tests/broadcast/server/test_configuration.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/tests/broadcast/server/test_decorators.py b/tests/broadcast/server/test_decorators.py deleted file mode 100644 index 6dfbde783..000000000 --- a/tests/broadcast/server/test_decorators.py +++ /dev/null @@ -1,72 +0,0 @@ -from unittest.mock import MagicMock - -import pytest -from druncschema.broadcast_pb2 import BroadcastType -from druncschema.request_response_pb2 import Request -from druncschema.token_pb2 import Token - -from drunc.broadcast.server.decorators import broadcasted - - -class MockException(Exception): - pass - -@pytest.fixture -def mock_obj(): - """Mock the object that has the .broadcast() method.""" - obj = MagicMock() - obj.name = "test-node" - return obj - - -@pytest.fixture(scope="function") -def mock_request(): - return Request(token=Token(user_name="test", token="tets-token")) - - -@pytest.fixture -def mock_context(): - return MagicMock() - - -def test_broadcasted_success(mock_obj, mock_request, mock_context): - - @broadcasted - def dummy_command(obj, request, context): - return "Success" - - result = dummy_command(mock_obj, mock_request, mock_context) - - assert result == "Success" - - assert mock_obj.broadcast.call_count == 2 # ACK and COMMAND_EXECUTION_SUCCESS - - # first call - ACK - args, kwargs = mock_obj.broadcast.call_args_list[0] - assert kwargs['message'] == "User 'test' executing 'dummy_command'" - assert kwargs['btype'] == BroadcastType.ACK - - # second call - COMMAND_EXECUTION_SUCCESS - # check no missing or additional arguments - mock_obj.broadcast.assert_called_with( - message="User 'test' successfully executed 'dummy_command'", - btype=BroadcastType.COMMAND_EXECUTION_SUCCESS) - - -def test_broadcasted_failure(mock_obj, mock_request, mock_context): - - # command that raises an error - @broadcasted - def dummy_command(obj, request, context): - raise MockException("Test exception") - - with pytest.raises(MockException): - dummy_command(mock_obj, mock_request, mock_context) - - assert mock_obj.broadcast.call_count == 2 # ACK and Exception - - # check no missing or additional arguments passed to broadcast - mock_obj.broadcast.assert_called_with( - message="Command 'dummy_command' failed", - btype=BroadcastType.UNHANDLED_EXCEPTION_RAISED) - \ No newline at end of file diff --git a/tests/broadcast/server/test_grpc_servicer.py b/tests/broadcast/server/test_grpc_servicer.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/tests/broadcast/server/test_kafka_sender.py b/tests/broadcast/server/test_kafka_sender.py deleted file mode 100644 index f2e4df400..000000000 --- a/tests/broadcast/server/test_kafka_sender.py +++ /dev/null @@ -1,93 +0,0 @@ -from unittest.mock import MagicMock, patch - -import pytest -from kafka.errors import KafkaError, NoBrokersAvailable - -from drunc.broadcast.server.kafka_sender import KafkaSender -from drunc.exceptions import DruncSetupException - - -@pytest.fixture -def mock_kafka_producer(): - """Fixture to mock a KafkaProducer class and its instance.""" - with patch("kafka.KafkaProducer") as mock_class: - mock_instance = mock_class.return_value - yield mock_class, mock_instance - - -@pytest.fixture -def sender_with_mocked_producer(mock_kafka_producer): - """Fixture to create a KafkaSender instance with a mocked KafkaProducer.""" - kafka_sender = KafkaSender( - kafka_address="test-kafka-address", publish_timeout=5, topic="test-topic" - ) - kafka_sender._log = MagicMock() - return kafka_sender - - -def test_init_raises_drunc_setup_exception(): - """Test that KafkaSender raises DruncSetupException when KafkaProducer cannot connect to a broker.""" - mock_logger = MagicMock() - kafka_address = "test-address" - expected_exc_msg = f"{kafka_address} does not seem to point to a kafka broker." - - with ( - patch("logging.getLogger", return_value=mock_logger), - patch("kafka.KafkaProducer", side_effect=NoBrokersAvailable), - ): - with pytest.raises(DruncSetupException) as exc_info: - KafkaSender( - kafka_address=kafka_address, publish_timeout=5, topic="test-topic" - ) - - mock_logger.critical.assert_called_once() - log_call_args = mock_logger.critical.call_args[0][0] - assert expected_exc_msg in log_call_args - assert expected_exc_msg in str(exc_info.value) - - -def test_send_success(sender_with_mocked_producer, mock_kafka_producer): - """Test that KafkaSender._send successfully sends a message and logs the metadata.""" - _, mock_instance = mock_kafka_producer - sender = sender_with_mocked_producer - - mock_future = MagicMock() - mock_future.get.return_value = "test-metadata" - mock_instance.send.return_value = mock_future - - mock_broadcast_msg = MagicMock() - mock_broadcast_msg.SerializeToString.return_value = b"test-msg" - sender._send(mock_broadcast_msg) - - mock_instance.send.assert_called_with("test-topic", b"test-msg") - sender._log.debug.assert_called_with("test-metadata published") - - -def test_send_handle_exception(sender_with_mocked_producer, mock_kafka_producer): - """Test that KafkaSender._send handles KafkaError exceptions and logs the error.""" - _, mock_instance = mock_kafka_producer - sender = sender_with_mocked_producer - - mock_future = MagicMock() - mock_future.get.side_effect = KafkaError("Connection lost") - mock_instance.send.return_value = mock_future - - mock_broadcast_msg = MagicMock() - mock_broadcast_msg.SerializeToString.return_value = b"test-broadcast-msg" - - sender._log = MagicMock() - - sender._send(mock_broadcast_msg) - - sender._log.error.assert_called() - log_message = sender._log.error.call_args[0][0] - assert "Connection lost" in log_message - - -def test_describe_broadcast(sender_with_mocked_producer): - """Test that KafkaSender.describe_broadcast returns the correct information.""" - - result = sender_with_mocked_producer.describe_broadcast() - - assert result.topic == "test-topic" - assert result.kafka_address == "test-kafka-address" diff --git a/tests/broadcast/test_types.py b/tests/broadcast/test_types.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/tests/broadcast/test_utils.py b/tests/broadcast/test_utils.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/tests/issues/test_issue309.py b/tests/issues/test_issue309.py index dbd027628..d12070ad8 100644 --- a/tests/issues/test_issue309.py +++ b/tests/issues/test_issue309.py @@ -6,14 +6,13 @@ def test_issue309(load_test_config): get_root_logger("INFO") - from drunc.utils.configuration import OKSKey, parse_conf_url + from drunc.utils.configuration import OKSKey - conf_path, conf_type = parse_conf_url("oksconflibs:deep-segments-config.data.xml") + conf_path = "oksconflibs:deep-segments-config.data.xml" controller_id = "controller-3" - controller_configuration = ControllerConfHandler( - type=conf_type, - data=conf_path, + controller_configuration = ControllerConfHandler.from_oks( + url=conf_path, oks_key=OKSKey( schema_file="schema/confmodel/dunedaq.schema.xml", class_name="RCApplication", @@ -23,4 +22,4 @@ def test_issue309(load_test_config): session_name="test", ) - assert controller_configuration.data.controller.id == controller_id + assert controller_configuration.controller.id == controller_id diff --git a/tests/issues/test_issue363.py b/tests/issues/test_issue363.py index 02e2ec088..ff477a5e7 100644 --- a/tests/issues/test_issue363.py +++ b/tests/issues/test_issue363.py @@ -1,18 +1,17 @@ # https://github.com/DUNE-DAQ/drunc/issues/363 from drunc.controller.configuration import ControllerConfHandler -from drunc.utils.configuration import OKSKey, parse_conf_url +from drunc.utils.configuration import OKSKey from drunc.utils.utils import get_root_logger def test_issue363(load_test_config): get_root_logger("INFO") - conf_path, conf_type = parse_conf_url("oksconflibs:nestedConfig.data.xml") + conf_path = "oksconflibs:nestedConfig.data.xml" controller_id = "nested-segment-controller" - controller_configuration = ControllerConfHandler( - type=conf_type, - data=conf_path, + controller_configuration = ControllerConfHandler.from_oks( + url=conf_path, oks_key=OKSKey( schema_file="schema/confmodel/dunedaq.schema.xml", class_name="RCApplication", @@ -21,6 +20,6 @@ def test_issue363(load_test_config): ), session_name="test", ) - ids = [segment.id for segment in controller_configuration.data.segments] + ids = [segment.id for segment in controller_configuration.segments] assert ids == ["bottom-segment-1", "bottom-segment-2"] - assert controller_configuration.data.controller.id == controller_id + assert controller_configuration.controller.id == controller_id diff --git a/tests/process_manager/process_manager_mock_impls.py b/tests/process_manager/process_manager_mock_impls.py index b9b84cf23..831ac2c38 100644 --- a/tests/process_manager/process_manager_mock_impls.py +++ b/tests/process_manager/process_manager_mock_impls.py @@ -38,7 +38,8 @@ def __init__( """ all-default constructor for testing purposes. """ - configuration.get_data().opmon_publisher = None + configuration.opmon_conf = {"level": "info", "interval_s": 10.0} + configuration.opmon_publisher = None super().__init__(configuration, name, session, **kwargs) def _not_implemented_response(self): @@ -55,9 +56,6 @@ def _not_implemented_response(self): flag=ResponseFlag.NOT_EXECUTED_NOT_IMPLEMENTED, ) - def _create_broadcast_service(self, name, session): - self.broadcast_service = None - def _boot_impl(self, boot_request: BootRequest) -> ProcessInstanceList: """ Returns default not implemented response to indicate communication is working diff --git a/tests/processes/conftest.py b/tests/processes/conftest.py index 701ba3bbe..16f582c4f 100644 --- a/tests/processes/conftest.py +++ b/tests/processes/conftest.py @@ -215,3 +215,62 @@ def ssh_manager_forked() -> Generator[ logger.warning(f"Error during kill_all_processes in fixture cleanup: {e}") finally: manager.shutdown() + + +@pytest.fixture +def process_configs_flat(): + """ + Fixture providing process configurations for flat structure tests. + + Returns: + List of process config dicts with 'name', 'role', and 'tree_id' keys + """ + return [ + { + "name": "test_process_app_1", + "role": "application", + "tree_id": "0.1.2", + }, + { + "name": "test_process_app_2", + "role": "application", + "tree_id": "0.1.2", + }, + { + "name": "test_process_infra", + "role": "infrastructure-applications", + "tree_id": "infra.process", + }, + ] + + +@pytest.fixture +def process_configs_deep_nested(): + """ + Fixture providing process configurations for deeply nested structure tests. + + Returns: + List of process config dicts with 'name', 'role', and 'tree_id' keys + """ + return [ + { + "name": "bottom-segment-1-application", + "role": "application", + "tree_id": "0.1.2.3", + }, + { + "name": "bottom-segment-2-application", + "role": "application", + "tree_id": "0.2.3.4", + }, + { + "name": "nested-segment-application", + "role": "application", + "tree_id": "0.1.2", + }, + { + "name": "local-connection-server", + "role": "infrastructure-applications", + "tree_id": "1", + }, + ] diff --git a/tests/processes/test_process_metadata.py b/tests/processes/test_process_metadata.py new file mode 100644 index 000000000..65518449b --- /dev/null +++ b/tests/processes/test_process_metadata.py @@ -0,0 +1,48 @@ +""" +Tests for ProcessMetadata.compute_role_from_tree_id role classification logic. +""" + +import pytest + +from drunc.processes.process_metadata import ProcessMetadata + + +@pytest.mark.parametrize( + "tree_id, is_controller, expected_role", + [ + # Root roles + ("0", True, "root-controller"), + ("0", False, "infrastructure-applications"), + # Segment-controller (requires is_controller=True) + ("0.1", True, "segment-controller"), + ("0.1.2", True, "segment-controller"), + ("0.1.2.3", True, "segment-controller"), + # Application at various depths (requires 0. prefix, is_controller=False) + ("0.1", False, "application"), + ("0.1.2", False, "application"), + ("0.1.2.3", False, "application"), + ("0.2.3.4", False, "application"), + # Local connection server + ("1", False, "infrastructure-applications"), + ("1", True, "infrastructure-applications"), + # Other non-0-prefixed (infrastructure) + ("infra.process", False, "infrastructure-applications"), + ("infra.process", True, "infrastructure-applications"), + ("2.3", False, "infrastructure-applications"), + ("2.3", True, "infrastructure-applications"), + # Empty + ("", False, "unknown"), + ("", True, "unknown"), + ], +) +def test_compute_role_from_tree_id( + tree_id: str, is_controller: bool, expected_role: str +): + """Verify role classification matches expected value.""" + result = ProcessMetadata.compute_role_from_tree_id( + tree_id, is_controller=is_controller + ) + assert result == expected_role, ( + f"compute_role_from_tree_id({tree_id!r}, is_controller={is_controller}) " + f"returned {result!r}, expected {expected_role!r}" + ) diff --git a/tests/processes/test_ssh_process_lifetime_manager_common.py b/tests/processes/test_ssh_process_lifetime_manager_common.py index 33aef9359..0a0bab2b6 100644 --- a/tests/processes/test_ssh_process_lifetime_manager_common.py +++ b/tests/processes/test_ssh_process_lifetime_manager_common.py @@ -560,7 +560,7 @@ def boot_processes_and_kill_individually(ssh_manager, test_file_path): boot_request = create_boot_request( process_name=process_name, - tree_id="this.isan.application", + tree_id="0.1.2", log_file=log_file, test_file_path=test_file_path, ) @@ -624,7 +624,7 @@ def boot_processes_and_terminate_all_same_role(ssh_manager, test_file_path): boot_request = create_boot_request( process_name=process_name, - tree_id="this.isan.application", + tree_id="0.1.2", log_file=log_file, test_file_path=test_file_path, ) @@ -659,17 +659,20 @@ def boot_processes_and_terminate_all_same_role(ssh_manager, test_file_path): ) -def boot_processes_and_terminate_all_different_role(ssh_manager, test_file_path): +def boot_processes_and_terminate_all_different_role_flat( + ssh_manager, test_file_path, process_configs_flat +): """ - Execute SSH processes with different roles and verify priority-based termination. + Execute SSH processes with different roles (flat structure) and verify + priority-based termination. Tests that the role-based shutdown mechanism correctly terminates processes - in the expected order according to the defined shutdown sequence. Higher-priority - roles (earlier in shutdown order) are killed before lower-priority roles. + in the expected order. Args: ssh_manager: SSH process lifetime manager instance test_file_path: Path to test file (for locating simple_process.py) + process_configs_flat: Fixture providing flat process configurations """ import threading import time @@ -677,25 +680,7 @@ def boot_processes_and_terminate_all_different_role(ssh_manager, test_file_path) with tempfile.TemporaryDirectory() as temp_dir: log_dir = Path(temp_dir) - # Define processes with different roles based on shutdown order - # "application" is terminated before "segment-controller" - process_configs = [ - { - "name": "test_process_app_1", - "role": "application", - "tree_id": "this.isan.application", - }, - { - "name": "test_process_app_2", - "role": "application", - "tree_id": "this.isan.application", - }, - { - "name": "test_process_segment", - "role": "segment-controller", - "tree_id": "thisisa.segment-controller", - }, - ] + process_configs = process_configs_flat process_uuids = [] process_info = {} @@ -775,10 +760,10 @@ def monitor_termination(uuid_to_monitor): for uuid in process_uuids if process_info[uuid]["role"] == "application" ] - segment_processes = [ + infra_processes = [ uuid for uuid in process_uuids - if process_info[uuid]["role"] == "segment-controller" + if process_info[uuid]["role"] == "infrastructure-applications" ] # Find latest termination time among higher-priority processes @@ -788,30 +773,30 @@ def monitor_termination(uuid_to_monitor): latest_app_termination = max(app_termination_times) # Find earliest termination time among lower-priority processes - segment_termination_times = [ - process_info[uuid]["termination_time"] for uuid in segment_processes + infra_termination_times = [ + process_info[uuid]["termination_time"] for uuid in infra_processes ] - earliest_segment_termination = min(segment_termination_times) + earliest_infra_termination = min(infra_termination_times) print( f"Latest 'application' role termination: " f"{latest_app_termination - start_time:.3f}s" ) print( - f"Earliest 'segment-controller' role termination: " - f"{earliest_segment_termination - start_time:.3f}s" + f"Earliest 'infrastructure-applications' role termination: " + f"{earliest_infra_termination - start_time:.3f}s" ) # Verify shutdown order: all "application" processes must terminate - # before any "segment-controller" processes - assert latest_app_termination <= earliest_segment_termination, ( - f"Application processes should terminate before segment-controller. " + # before any "infrastructure-applications" processes + assert latest_app_termination <= earliest_infra_termination, ( + f"Application processes should terminate before infrastructure-applications. " f"Latest app: {latest_app_termination - start_time:.3f}s, " - f"Earliest segment: {earliest_segment_termination - start_time:.3f}s" + f"Earliest infra: {earliest_infra_termination - start_time:.3f}s" ) print( "✓ Role-based termination order verified: " - "'application' before 'segment-controller'" + "'application' before 'infrastructure-applications'" ) verify_cleanup_complete(ssh_manager, pid_snapshots=pid_snapshots) @@ -820,3 +805,135 @@ def monitor_termination(uuid_to_monitor): "\n✓ Test passed: Processes with different roles executed, logged, " "terminated in correct order, and cleaned up successfully" ) + + +def boot_processes_and_terminate_all_different_role_deep_nested( + ssh_manager, test_file_path, process_configs_deep_nested +): + """ + Execute SSH processes with deeply nested tree_ids and verify priority-based + termination. + + Tests that the role-based shutdown mechanism correctly handles applications + at arbitrary depth under "0." prefix (e.g. 0.1.2.3, 0.2.3.4), terminating + them before infrastructure-applications processes. + + Args: + ssh_manager: SSH process lifetime manager instance + test_file_path: Path to test file (for locating simple_process.py) + process_configs_deep_nested: Fixture providing deep nested process configurations + """ + import threading + import time + + with tempfile.TemporaryDirectory() as temp_dir: + log_dir = Path(temp_dir) + + process_configs = process_configs_deep_nested + + process_uuids = [] + process_info = {} + + print("\n=== Executing deeply nested processes ===") + for config in process_configs: + process_name = config["name"] + role = config["role"] + log_file = str(log_dir / f"{process_name}.log") + process_uuid = str(uuid.uuid4()) + process_uuids.append(process_uuid) + + boot_request = create_boot_request( + process_name=process_name, + tree_id=config["tree_id"], + log_file=log_file, + test_file_path=test_file_path, + ) + + ssh_manager.start_process(uuid=process_uuid, boot_request=boot_request) + + process_info[process_uuid] = { + "name": process_name, + "role": role, + "log_file": log_file, + "termination_time": None, + } + + print( + f"Executed {process_name} (tree_id={config['tree_id']!r}) " + f"with UUID {process_uuid} and role '{role}'" + ) + + verify_all_processes_alive(ssh_manager, process_uuids, len(process_configs)) + verify_log_output(ssh_manager, process_uuids, process_info) + + print("\n=== Terminating all processes (role-based shutdown, deep nested) ===") + + def monitor_termination(uuid_to_monitor): + while ssh_manager.is_process_alive(uuid_to_monitor): + time.sleep(0.05) + process_info[uuid_to_monitor]["termination_time"] = time.time() + + monitor_threads = [] + for process_uuid in process_uuids: + thread = threading.Thread( + target=monitor_termination, + args=(process_uuid,), + daemon=True, + ) + thread.start() + monitor_threads.append(thread) + + start_time = time.time() + exit_codes = ssh_manager.kill_processes( + process_uuids, process_timeouts={uuid: 10.0 for uuid in process_uuids} + ) + + for thread in monitor_threads: + thread.join(timeout=15.0) + + verify_all_processes_dead(ssh_manager, process_uuids, len(process_configs)) + verify_exit_codes(exit_codes, process_uuids, process_info) + + print("\n=== Verifying termination order (deep nested) ===") + app_processes = [ + u for u in process_uuids if process_info[u]["role"] == "application" + ] + infra_processes = [ + u + for u in process_uuids + if process_info[u]["role"] == "infrastructure-applications" + ] + + latest_app_termination = max( + process_info[u]["termination_time"] for u in app_processes + ) + earliest_infra_termination = min( + process_info[u]["termination_time"] for u in infra_processes + ) + + print( + f"Latest 'application' role termination: " + f"{latest_app_termination - start_time:.3f}s" + ) + print( + f"Earliest 'infrastructure-applications' termination: " + f"{earliest_infra_termination - start_time:.3f}s" + ) + + assert latest_app_termination <= earliest_infra_termination, ( + f"Deeply nested application processes should terminate before " + f"infrastructure-applications. " + f"Latest app: {latest_app_termination - start_time:.3f}s, " + f"Earliest infra: {earliest_infra_termination - start_time:.3f}s" + ) + print( + "✓ Role-based termination order verified (deep nested): " + "'application' before 'infrastructure-applications'" + ) + + verify_cleanup_complete(ssh_manager) + + print( + "\n✓ Test passed: Deeply nested processes correctly classified and " + "terminated in correct role-based order" + ) diff --git a/tests/processes/test_ssh_process_lifetime_manager_forked_shell.py b/tests/processes/test_ssh_process_lifetime_manager_forked_shell.py index d9d0d698c..a1e10e561 100644 --- a/tests/processes/test_ssh_process_lifetime_manager_forked_shell.py +++ b/tests/processes/test_ssh_process_lifetime_manager_forked_shell.py @@ -8,7 +8,6 @@ from tests.processes.test_ssh_process_lifetime_manager_common import ( boot_processes_and_kill_individually, boot_processes_and_terminate_all_same_role, - boot_processes_and_verify_exit_state_messages, capture_process_pid_snapshots, create_boot_request, verify_all_processes_alive, @@ -41,24 +40,18 @@ def test_ssh_terminate_all_same_role_forked(ssh_manager_forked): boot_processes_and_terminate_all_same_role(ssh_manager_forked, Path(__file__)) -def test_ssh_exit_status_messages_for_kill_paths_forked(ssh_manager_forked): +def boot_processes_and_terminate_all_different_role_flat_forked( + test_file_path, process_configs_flat +): """ - Verify forked shell manager emits correct exit-status messages for all kill paths. + Execute SSH processes with different roles (flat structure) via the forked manager + and verify priority-based termination order. - Boots one process per exit-status scenario, triggers all required kill modes - concurrently, and validates callback-delivered ExitStatus source and message. - """ - boot_processes_and_verify_exit_state_messages(ssh_manager_forked, Path(__file__)) - - -def boot_processes_and_terminate_all_different_role_forked(test_file_path): - """ - Execute SSH processes with different roles via the forked manager and verify - priority-based termination order. + Uses flat process configurations with termination order enforcement via callbacks. Args: - ssh_manager_forked: Forked SSH process lifetime manager instance test_file_path: Path to the test file (used to locate simple_process.py) + process_configs_flat: Fixture providing flat process configurations """ import threading import time @@ -75,23 +68,7 @@ def boot_processes_and_terminate_all_different_role_forked(test_file_path): termination_times: dict = {} callback_events: dict = {} - process_configs = [ - { - "name": "test_process_app_1", - "role": "application", - "tree_id": "this.isan.application", - }, - { - "name": "test_process_app_2", - "role": "application", - "tree_id": "this.isan.application", - }, - { - "name": "test_process_segment", - "role": "segment-controller", - "tree_id": "thisisa.segment-controller", - }, - ] + process_configs = process_configs_flat process_uuids = [] process_info = {} @@ -173,25 +150,174 @@ def on_exit(cb_uuid: str, exit_code, exception): app_processes = [ u for u in process_uuids if process_info[u]["role"] == "application" ] - segment_processes = [ + infra_processes = [ u for u in process_uuids - if process_info[u]["role"] == "segment-controller" + if process_info[u]["role"] == "infrastructure-applications" ] latest_app = max(termination_times[u] for u in app_processes) - earliest_segment = min(termination_times[u] for u in segment_processes) + earliest_infra = min(termination_times[u] for u in infra_processes) - print(f"Latest 'application' callback: {latest_app:.6f}") - print(f"Earliest 'segment-controller' callback: {earliest_segment:.6f}") + print(f"Latest 'application' callback: {latest_app:.6f}") + print( + f"Earliest 'infrastructure-applications' callback: {earliest_infra:.6f}" + ) - assert latest_app <= earliest_segment, ( - f"Application processes should terminate before segment-controller " - f"(delta: {earliest_segment - latest_app:.6f}s)" + assert latest_app <= earliest_infra, ( + f"Application processes should terminate before infrastructure-applications " + f"(delta: {earliest_infra - latest_app:.6f}s)" ) print( "✓ Role-based termination order verified via callbacks: " - "'application' before 'segment-controller'" + "'application' before 'infrastructure-applications'" + ) + + verify_cleanup_complete(manager, pid_snapshots=pid_snapshots) + + finally: + manager.shutdown() + + +def test_ssh_terminate_all_different_role_flat_forked(process_configs_flat): + """ + Test priority-based termination of flat-structure processes using the forked manager. + + Executes processes with varying role priorities, verifies callback-based + termination ordering, and confirms cleanup. + """ + boot_processes_and_terminate_all_different_role_flat_forked( + Path(__file__), process_configs_flat + ) + + +def boot_processes_and_terminate_all_different_role_deep_nested_forked( + test_file_path, process_configs_deep_nested +): + """ + Execute SSH processes with deeply nested tree_ids via the forked manager and verify + priority-based termination order. + + Tests that the role-based shutdown mechanism correctly handles applications + at arbitrary depth under "0." prefix, terminating them before + infrastructure-applications processes. + + Args: + test_file_path: Path to the test file (used to locate simple_process.py) + process_configs_deep_nested: Fixture providing deep nested process configurations + """ + import threading + import time + + from drunc.processes.ssh_process_lifetime_manager_from_forked_process import ( + SSHProcessLifetimeManagerShellOnForkedProcess, + ) + + with tempfile.TemporaryDirectory() as temp_dir: + log_dir = Path(temp_dir) + + termination_times: dict = {} + callback_events: dict = {} + + process_configs = process_configs_deep_nested + + process_uuids = [] + process_info = {} + + callback_lock = threading.Lock() + + def on_exit(cb_uuid: str, exit_code, exception): + with callback_lock: + termination_times[cb_uuid] = time.monotonic() + if cb_uuid in callback_events: + callback_events[cb_uuid].set() + + manager = SSHProcessLifetimeManagerShellOnForkedProcess( + disable_localhost_host_key_check=True, + on_process_exit=on_exit, + ) + + try: + print("\n=== Executing deeply nested processes (forked) ===") + for config in process_configs: + process_name = config["name"] + role = config["role"] + log_file = str(log_dir / f"{process_name}.log") + process_uuid = str(uuid.uuid4()) + process_uuids.append(process_uuid) + callback_events[process_uuid] = threading.Event() + + boot_request = create_boot_request( + process_name=process_name, + tree_id=config["tree_id"], + log_file=log_file, + test_file_path=test_file_path, + ) + + manager.start_process(uuid=process_uuid, boot_request=boot_request) + + process_info[process_uuid] = { + "name": process_name, + "role": role, + "log_file": log_file, + } + + print( + f"Executed {process_name} (tree_id={config['tree_id']!r}) " + f"with UUID {process_uuid} and role '{role}'" + ) + + verify_all_processes_alive(manager, process_uuids, len(process_configs)) + verify_log_output(manager, process_uuids, process_info) + + pid_snapshots = capture_process_pid_snapshots(manager, process_uuids) + + print( + "\n=== Terminating all processes " + "(role-based shutdown, deep nested, forked) ===" + ) + exit_codes = manager.kill_processes( + process_uuids, + process_timeouts={u: 10.0 for u in process_uuids}, + ) + + for process_uuid in process_uuids: + fired = callback_events[process_uuid].wait(timeout=15.0) + assert fired, ( + f"on_process_exit callback never fired for " + f"{process_info[process_uuid]['name']}" + ) + + verify_all_processes_dead(manager, process_uuids, len(process_configs)) + verify_exit_codes(exit_codes, process_uuids, process_info) + + print("\n=== Verifying termination order (deep nested, forked) ===") + + app_processes = [ + u for u in process_uuids if process_info[u]["role"] == "application" + ] + infra_processes = [ + u + for u in process_uuids + if process_info[u]["role"] == "infrastructure-applications" + ] + + latest_app = max(termination_times[u] for u in app_processes) + earliest_infra = min(termination_times[u] for u in infra_processes) + + print(f"Latest 'application' callback: {latest_app:.6f}") + print( + f"Earliest 'infrastructure-applications' callback: {earliest_infra:.6f}" + ) + + assert latest_app <= earliest_infra, ( + f"Deeply nested application processes should terminate before " + f"infrastructure-applications " + f"(delta: {earliest_infra - latest_app:.6f}s)" + ) + print( + "✓ Role-based termination order verified (deep nested, forked): " + "'application' before 'infrastructure-applications'" ) verify_cleanup_complete(manager, pid_snapshots=pid_snapshots) @@ -200,6 +326,21 @@ def on_exit(cb_uuid: str, exit_code, exception): manager.shutdown() +def test_ssh_terminate_all_different_role_deep_nested_forked( + process_configs_deep_nested, +): + """ + Test role classification and priority-based termination for deeply nested processes + using the forked manager. + + Tests that applications at arbitrary depth under "0." prefix are correctly + handled and terminate before infrastructure-applications processes. + """ + boot_processes_and_terminate_all_different_role_deep_nested_forked( + Path(__file__), process_configs_deep_nested + ) + + def test_forked_manager_worker_process_is_alive(ssh_manager_forked): """ Verify that the worker child process is alive immediately after construction. @@ -278,7 +419,7 @@ def on_exit(cb_uuid: str, exit_code, exception): boot_request = create_boot_request( process_name="callback_test_process", - tree_id="this.isan.application", + tree_id="0.1.2", log_file=log_file, test_file_path=Path(__file__), ) diff --git a/tests/processes/test_ssh_process_lifetime_manager_paramiko.py b/tests/processes/test_ssh_process_lifetime_manager_paramiko.py index 7e7430791..240bd1d42 100644 --- a/tests/processes/test_ssh_process_lifetime_manager_paramiko.py +++ b/tests/processes/test_ssh_process_lifetime_manager_paramiko.py @@ -4,7 +4,8 @@ from tests.processes.test_ssh_process_lifetime_manager_common import ( boot_processes_and_kill_individually, - boot_processes_and_terminate_all_different_role, + boot_processes_and_terminate_all_different_role_deep_nested, + boot_processes_and_terminate_all_different_role_flat, boot_processes_and_terminate_all_same_role, ) @@ -32,14 +33,31 @@ def test_ssh_terminate_all_same_role_paramiko(ssh_manager_paramiko): @pytest.mark.paramiko -def test_ssh_terminate_all_different_role_paramiko(ssh_manager_paramiko): +def test_ssh_terminate_all_different_role_flat_paramiko( + ssh_manager_paramiko, process_configs_flat +): """ - Test priority-based termination of processes with different roles using Paramiko. + Test priority-based termination of processes with different roles (flat) using Paramiko. Executes processes with varying role priorities via SSH, verifies log output, terminates all processes using role-based shutdown, verifies termination order, and confirms complete cleanup. """ - boot_processes_and_terminate_all_different_role( - ssh_manager_paramiko, Path(__file__) + boot_processes_and_terminate_all_different_role_flat( + ssh_manager_paramiko, Path(__file__), process_configs_flat + ) + + +@pytest.mark.paramiko +def test_ssh_terminate_all_different_role_deep_nested_paramiko( + ssh_manager_paramiko, process_configs_deep_nested +): + """ + Test role classification and priority-based termination for deeply nested processes using Paramiko. + + Exercises role classification for applications at arbitrary depth under "0." prefix, + and verifies they terminate before infrastructure-applications processes. + """ + boot_processes_and_terminate_all_different_role_deep_nested( + ssh_manager_paramiko, Path(__file__), process_configs_deep_nested ) diff --git a/tests/processes/test_ssh_process_lifetime_manager_shell.py b/tests/processes/test_ssh_process_lifetime_manager_shell.py index 27b091093..4d7e14fea 100644 --- a/tests/processes/test_ssh_process_lifetime_manager_shell.py +++ b/tests/processes/test_ssh_process_lifetime_manager_shell.py @@ -2,9 +2,9 @@ from tests.processes.test_ssh_process_lifetime_manager_common import ( boot_processes_and_kill_individually, - boot_processes_and_terminate_all_different_role, + boot_processes_and_terminate_all_different_role_deep_nested, + boot_processes_and_terminate_all_different_role_flat, boot_processes_and_terminate_all_same_role, - boot_processes_and_verify_exit_state_messages, ) @@ -28,26 +28,30 @@ def test_ssh_terminate_all_same_role_shell(ssh_manager_shell): boot_processes_and_terminate_all_same_role(ssh_manager_shell, Path(__file__)) -def test_ssh_terminate_all_different_role_shell(ssh_manager_shell): +def test_ssh_terminate_all_different_role_flat_shell( + ssh_manager_shell, process_configs_flat +): """ - Test priority-based termination of processes with different roles using shell. + Test priority-based termination of processes with different roles (flat) using shell. Executes processes with varying role priorities via SSH, verifies log output, terminates all processes using role-based shutdown, verifies termination order, and confirms complete cleanup. """ - boot_processes_and_terminate_all_different_role(ssh_manager_shell, Path(__file__)) + boot_processes_and_terminate_all_different_role_flat( + ssh_manager_shell, Path(__file__), process_configs_flat + ) -def test_ssh_exit_status_messages_for_kill_paths_shell(ssh_manager_shell): +def test_ssh_terminate_all_different_role_deep_nested_shell( + ssh_manager_shell, process_configs_deep_nested +): """ - Verify shell manager emits correct exit-status messages for all kill paths. + Test role classification and priority-based termination for deeply nested processes. - Boots one process per exit-status scenario, triggers all required kill modes - concurrently, and validates callback-delivered ExitStatus source and message. + Exercises role classification for applications at arbitrary depth under "0." prefix, + and verifies they terminate before infrastructure-applications processes. """ - - boot_processes_and_verify_exit_state_messages( - ssh_manager=ssh_manager_shell, - test_file_path=Path(__file__), + boot_processes_and_terminate_all_different_role_deep_nested( + ssh_manager_shell, Path(__file__), process_configs_deep_nested ) From 119bd600769af8a4670266a00610c5b94cf58c1a Mon Sep 17 00:00:00 2001 From: James Paul Turner Date: Mon, 13 Jul 2026 18:14:36 +0100 Subject: [PATCH 14/20] Phased merge 8 --- src/drunc/data/process_manager/k8s-CERN.json | 6 - src/drunc/data/process_manager/k8s.json | 5 - .../process-manager-k8s-pocket.json | 6 - .../schema/process_manager.schema.json | 10 - .../data/process_manager/ssh-CERN-kafka.json | 6 - .../process_manager/ssh-pocket-kafka.json | 9 +- src/drunc/utils/__init__.py | 2 + src/drunc/utils/configuration.py | 248 +++++++++---- src/drunc/utils/flask_manager.py | 153 ++++++-- src/drunc/utils/flask_typing.pyi | 29 ++ src/drunc/utils/grpc_utils.py | 156 ++++++-- src/drunc/utils/shell_utils.py | 250 +++++++++++-- src/drunc/utils/utils.py | 341 ++++++++++++++---- 13 files changed, 937 insertions(+), 284 deletions(-) create mode 100644 src/drunc/utils/flask_typing.pyi diff --git a/src/drunc/data/process_manager/k8s-CERN.json b/src/drunc/data/process_manager/k8s-CERN.json index 100dc2851..b59b10fb1 100644 --- a/src/drunc/data/process_manager/k8s-CERN.json +++ b/src/drunc/data/process_manager/k8s-CERN.json @@ -7,12 +7,6 @@ "authoriser": { "type": "dummy" }, - - "broadcaster": { - "type": "kafka", - "kafka_address": "monkafka.cern.ch:30092", - "publish_timeout": 2 - }, "environment": { "DUNEDAQ_ERS_ERROR": "erstrace,throttle,lstdout,protobufstream(monkafka.cern.ch:30092)", "DUNEDAQ_ERS_FATAL": "erstrace,lstdout,protobufstream(monkafka.cern.ch:30092)", diff --git a/src/drunc/data/process_manager/k8s.json b/src/drunc/data/process_manager/k8s.json index fa94a98ea..c46412eb6 100644 --- a/src/drunc/data/process_manager/k8s.json +++ b/src/drunc/data/process_manager/k8s.json @@ -6,11 +6,6 @@ "authoriser": { "type": "dummy" }, - "broadcaster": { - "type": "kafka", - "kafka_address": "monkafka.cern.ch:30092", - "publish_timeout": 2 - }, "environment": { "DUNEDAQ_ERS_ERROR": "erstrace,throttle,lstdout", "DUNEDAQ_ERS_FATAL": "erstrace,lstdout", diff --git a/src/drunc/data/process_manager/process-manager-k8s-pocket.json b/src/drunc/data/process_manager/process-manager-k8s-pocket.json index 54c295afd..069ec89d8 100644 --- a/src/drunc/data/process_manager/process-manager-k8s-pocket.json +++ b/src/drunc/data/process_manager/process-manager-k8s-pocket.json @@ -6,12 +6,6 @@ "authoriser": { "type": "dummy" }, - - "broadcaster": { - "type": "kafka", - "kafka_address": "localhost:30092", - "publish_timeout": 2 - }, "environment": { "DUNEDAQ_ERS_ERROR": "erstrace,throttle,lstdout,protobufstream(monkafka.cern.ch:30092)", "DUNEDAQ_ERS_FATAL": "erstrace,lstdout,protobufstream(monkafka.cern.ch:30092)", diff --git a/src/drunc/data/process_manager/schema/process_manager.schema.json b/src/drunc/data/process_manager/schema/process_manager.schema.json index 0ad199b8e..bcb906a7e 100644 --- a/src/drunc/data/process_manager/schema/process_manager.schema.json +++ b/src/drunc/data/process_manager/schema/process_manager.schema.json @@ -33,16 +33,6 @@ }, "additionalProperties": true }, - "broadcaster": { - "type": "object", - "properties": { - "type": {"type": "string"}, - "kafka_address": {"type": "string"}, - "publish_timeout": {"type": ["integer", "number"]} - }, - "required": ["type", "kafka_address", "publish_timeout"], - "additionalProperties": true - }, "command_address": {"type": "string"} }, "required": ["type", "name"], diff --git a/src/drunc/data/process_manager/ssh-CERN-kafka.json b/src/drunc/data/process_manager/ssh-CERN-kafka.json index 1158be0d9..0f62572e7 100644 --- a/src/drunc/data/process_manager/ssh-CERN-kafka.json +++ b/src/drunc/data/process_manager/ssh-CERN-kafka.json @@ -7,12 +7,6 @@ "authoriser": { "type": "dummy" }, - - "broadcaster": { - "type": "kafka", - "kafka_address": "monkafka.cern.ch:30092", - "publish_timeout": 2 - }, "environment": { "DUNEDAQ_ERS_ERROR": "erstrace,throttle,lstdout,protobufstream(monkafka.cern.ch:30092)", "DUNEDAQ_ERS_FATAL": "erstrace,lstdout,protobufstream(monkafka.cern.ch:30092)", diff --git a/src/drunc/data/process_manager/ssh-pocket-kafka.json b/src/drunc/data/process_manager/ssh-pocket-kafka.json index 9208772da..950083f7b 100644 --- a/src/drunc/data/process_manager/ssh-pocket-kafka.json +++ b/src/drunc/data/process_manager/ssh-pocket-kafka.json @@ -7,13 +7,6 @@ "authoriser": { "type": "dummy" }, - - "broadcaster": { - "type": "kafka", - "kafka_address": "localhost:31014", - "publish_timeout": 2 - }, - "environment": { "GRPC_ENABLE_FORK_SUPPORT": "false", "DUNEDAQ_ERS_ERROR": "erstrace,throttle,lstdout,protobufstream(monkafka.cern.ch:30092)", @@ -30,4 +23,4 @@ "level": "debug", "interval_s": 10.0 } -} \ No newline at end of file +} diff --git a/src/drunc/utils/__init__.py b/src/drunc/utils/__init__.py index 786dda6c3..4ebf79608 100644 --- a/src/drunc/utils/__init__.py +++ b/src/drunc/utils/__init__.py @@ -1,3 +1,5 @@ +"""drunc utilities module.""" + from drunc.utils.utils import get_logger # Initialise utils logger with Rich handler diff --git a/src/drunc/utils/configuration.py b/src/drunc/utils/configuration.py index 99d08db7b..9fb5b2199 100644 --- a/src/drunc/utils/configuration.py +++ b/src/drunc/utils/configuration.py @@ -1,6 +1,10 @@ +"""Configuration utilities for DRUNC.""" + import json +import logging import os from enum import Enum +from typing import cast import conffwk @@ -9,6 +13,8 @@ class ConfTypes(Enum): + """Enumeration of supported configuration types.""" + Unknown = 0 # End product @@ -21,6 +27,17 @@ class ConfTypes(Enum): def CLI_to_ConfTypes(scheme: str) -> ConfTypes: + """Convert a CLI scheme string to a ConfTypes enum. + + Args: + scheme: The scheme string ("file", "oksconflibs", or ""). + + Returns: + ConfTypes: The corresponding configuration type. + + Raises: + DruncSetupException: If the scheme is not recognized. + """ match scheme: case "file": return ConfTypes.JsonFileName @@ -31,20 +48,43 @@ def CLI_to_ConfTypes(scheme: str) -> ConfTypes: def parse_conf_url(url: str) -> tuple[str, ConfTypes]: + """Parse a configuration URL into scheme and type. + + Args: + url: The configuration URL (format: "scheme:filename"). + + Returns: + tuple[str, ConfTypes]: A tuple of (url, conf_type). + """ scheme, filename = url.split(":") t = CLI_to_ConfTypes(scheme) return url, t class ConfigurationNotFound(DruncSetupException): - def __init__(self, requested_path): + """Exception raised when configuration is not found.""" + + def __init__(self, requested_path: str) -> None: + """Initialize the ConfigurationNotFound exception. + + Args: + requested_path: The path to the configuration that was not found. + """ super().__init__( f"The configuration '{requested_path}' is not in $DUNEDAQ_DB_PATH, perhaps you forgot to 'dbt-workarea-env && dbt-build'?" ) class ConfTypeNotSupported(DruncSetupException): - def __init__(self, conf_type: ConfTypes, class_name: str): + """Exception raised when a configuration type is not supported.""" + + def __init__(self, conf_type: ConfTypes, class_name: str) -> None: + """Initialize the ConfTypeNotSupported exception. + + Args: + conf_type: The configuration type that is not supported. + class_name: The name of the class where this type is not supported. + """ if not isinstance(class_name, str): class_name = class_name.__class__.__name__ message = f"'{conf_type}' is not supported by '{class_name}'" @@ -52,7 +92,19 @@ def __init__(self, conf_type: ConfTypes, class_name: str): class OKSKey: - def __init__(self, schema_file: str, class_name: str, obj_uid: str, session: str): + """Key information for accessing OKS configuration objects.""" + + def __init__( + self, schema_file: str, class_name: str, obj_uid: str, session: str + ) -> None: + """Initialize an OKSKey. + + Args: + schema_file: The OKS schema file path. + class_name: The class name in the OKS schema. + obj_uid: The unique identifier for the object. + session: The session name. + """ self.schema_file = schema_file self.class_name = class_name self.obj_uid = obj_uid @@ -60,50 +112,138 @@ def __init__(self, schema_file: str, class_name: str, obj_uid: str, session: str class ConfHandler: - def __init__( - self, - data=None, - type=ConfTypes.PyObject, - oks_key: OKSKey = None, - *args, - **kwargs, - ): + """Handler for loading and parsing DRUNC configurations. + + Supports multiple configuration sources via from_* classmethods. + Subclasses override populate_from_dict / populate_from_pbany to handle + JSON and protobuf sources, and _post_process_oks to handle OKS/pyobject + sources (via self._raw_data). + """ + + type: ConfTypes + oks_key: OKSKey | None + class_name: str + log: logging.Logger + root_id: int + controller_id: int + process_id: int + process_id_infra: int + session_name: str | None + initial_data: object + oks_path: str + db: object + _raw_data: object # raw OKS/pyobject data, available during _post_process_oks + + @classmethod + def from_pyobject( + cls, data: object, session_name: str | None = None + ) -> "ConfHandler": + instance: ConfHandler = cls.__new__(cls) + instance._init_common(session_name) + instance.initial_data = data + instance._raw_data = data + instance.type = ConfTypes.PyObject + instance._post_process_oks() + return instance + + @classmethod + def from_pbany(cls, data: object, session_name: str | None = None) -> "ConfHandler": + instance: ConfHandler = cls.__new__(cls) + instance._init_common(session_name) + instance.initial_data = data + instance._raw_data = None + instance.populate_from_pbany(data) + instance.type = ConfTypes.PyObject + instance._post_process_oks() + return instance + + @classmethod + def from_json(cls, path: str, session_name: str | None = None) -> "ConfHandler": + instance: ConfHandler = cls.__new__(cls) + instance._init_common(session_name) + instance.initial_data = path + resolved = expand_path(path, True) + if not os.path.exists(expand_path(path)): + raise DruncSetupException(f"Location {resolved} ({path}) is empty!") + with open(resolved) as f: + json_data = json.load(f) + instance._raw_data = None + instance.populate_from_dict(cast(dict[str, object], json_data)) + instance.type = ConfTypes.PyObject + instance._post_process_oks() + return instance + + @classmethod + def from_oks( + cls, + url: str, + oks_key: OKSKey, + session_name: str | None = None, + ) -> "ConfHandler": + instance: ConfHandler = cls.__new__(cls) + instance._init_common(session_name) + instance.initial_data = url + instance.oks_key = oks_key + instance._raw_data = instance._parse_oks_file(url) + instance.type = ConfTypes.PyObject + instance._post_process_oks() + return instance + + def populate_from_dict(self, data: dict[str, object]) -> None: + """Populate from a dictionary (JSON source). + + Override in subclasses that support JSON configuration. + """ + raise ConfTypeNotSupported(ConfTypes.JsonFileName, self.__class__.__name__) + + def populate_from_pbany(self, pbany_data: object) -> None: + """Populate from a Protobuf Any message. + + Override in subclasses that support protobuf configuration. + """ + raise ConfTypeNotSupported(ConfTypes.ProtobufAny, self.__class__.__name__) + + def _init_common(self, session_name: str | None = None) -> None: + """Initialize common attributes. + + Args: + session_name: Optional session name. + """ self.class_name = self.__class__.__name__ self.log = get_logger("utils." + self.class_name) - self.initial_type = type - self.initial_data = data self.root_id = 0 self.controller_id = 0 self.process_id = 0 self.process_id_infra = 0 - self.session_name = kwargs.get("session_name") - - if type == ConfTypes.OKSFileName and oks_key is None: - raise DruncSetupException("Need to provide a key for the OKS file") + self.session_name = session_name + self.oks_key = None + self.type = ConfTypes.Unknown - self.oks_key = oks_key - self.validate_and_parse_configuration_location(*args, **kwargs) + def copy_oks_key(self) -> OKSKey | None: + """Get a copy of the OKS key if one exists. - def get_data(self): - return self.data - - def get_data_type_name(self): - return self.get_data().type._name_ + Returns: + OKSKey | None: The OKS key, or None if not using OKS configuration. + """ + return self.oks_key - def get_data_broadcaster(self): - return self.get_data().broadcaster + def _parse_oks_file(self, oks_path: str) -> object: + """Parse OKS configuration file. - def get_data_authoriser(self): - return self.get_data().authoriser + Args: + oks_path: Path to OKS database. - def copy_oks_key(self): - return self.oks_key + Returns: + object: The parsed DAL object. - def _parse_oks_file(self, oks_path): + Raises: + DruncSetupException: If OKS setup or parameters are missing. + """ try: self.oks_path = oks_path self.log.debug(f"Using {self.oks_path} to configure") self.db = conffwk.Configuration(self.oks_path) + assert self.oks_key is not None, "OKS key is required for OKS configuration" return self.db.get_dal( class_name=self.oks_key.class_name, uid=self.oks_key.obj_uid ) @@ -118,44 +258,10 @@ def _parse_oks_file(self, oks_path): "OKS params where not passed to this ConfigurationHandler, cannot parse OKS configurations" ) from e - def _post_process_oks(self): - pass + def _post_process_oks(self) -> None: + """Post-process configuration after loading. - def _parse_pbany(self, pbany_data): - raise ConfTypeNotSupported(ConfTypes.ProtobufAny, self) - - def _parse_dict(self, data): - raise ConfTypeNotSupported(ConfTypes.JsonFileName, self) - - def validate_and_parse_configuration_location(self, *args, **kwargs): - match self.initial_type: - case ConfTypes.PyObject: - self.data = self.initial_data - self.type = self.initial_type - self._post_process_oks(*args, **kwargs) - - case ConfTypes.JsonFileName: - resolved = expand_path(self.initial_data, True) - if not os.path.exists(expand_path(self.initial_data)): - raise DruncSetupException( - f"Location {resolved} ({self.initial_data}) is empty!" - ) - - with open(resolved) as f: - data = json.loads(f.read()) - self.data = self._parse_dict(data) - self.type = ConfTypes.PyObject - self._post_process_oks(*args, **kwargs) - - case ConfTypes.OKSFileName: - self.data = self._parse_oks_file(self.initial_data) - self.type = ConfTypes.PyObject - self._post_process_oks(*args, **kwargs) - - case ConfTypes.ProtobufAny: - self.data = self._parse_pbany(self.initial_data) - self.type = ConfTypes.PyObject - self._post_process_oks(*args, **kwargs) - - case _: - raise ConfTypeNotSupported(self.initial_type, self.class_name) + Override in subclasses to perform custom initialization. + For OKS/pyobject sources, self._raw_data holds the raw object. + """ + pass diff --git a/src/drunc/utils/flask_manager.py b/src/drunc/utils/flask_manager.py index 80c504e06..561820a3f 100644 --- a/src/drunc/utils/flask_manager.py +++ b/src/drunc/utils/flask_manager.py @@ -1,27 +1,51 @@ +"""Flask application manager utilities for DRUNC.""" + import os import signal import threading import time from multiprocessing import Process -from typing import NoReturn +from typing import TYPE_CHECKING -import gunicorn.app.base import psutil import requests from flask import Flask, jsonify, make_response, request -from flask_restful import Api, Resource + +if TYPE_CHECKING: + from drunc.utils.flask_typing import ( + Api, + _BaseApplication, + _Resource, + ) +else: + from flask_restful import Api + from flask_restful import Resource as _Resource + from gunicorn.app.base import BaseApplication as _BaseApplication from drunc.exceptions import DruncCommandException from drunc.utils.utils import get_logger, get_new_port -class GunicornStandaloneApplication(gunicorn.app.base.BaseApplication): - def __init__(self, app, options=None): +class GunicornStandaloneApplication(_BaseApplication): + """Standalone Gunicorn application wrapper.""" + + def __init__( + self, + app: Flask, + options: dict[str, object] | None = None, + ) -> None: + """Initialize a GunicornStandaloneApplication. + + Args: + app: The Flask application to run. + options: Configuration options for Gunicorn. Defaults to None. + """ self.options = options or {} self.application = app super().__init__() - def load_config(self): + def load_config(self) -> None: + """Load Gunicorn configuration from options.""" config = { key: value for key, value in self.options.items() @@ -30,23 +54,32 @@ def load_config(self): for key, value in config.items(): self.cfg.set(key.lower(), value) - def load(self): + def load(self) -> Flask: + """Load the Flask application. + + Returns: + Flask: The Flask application. + """ return self.application class CannotStartFlaskManager(DruncCommandException): + """Exception raised when the Flask manager cannot start.""" + pass class FlaskManager(threading.Thread): - """This class is a manager for flask. - It allows to have a Flask server under a thread, start and stop it. - Note that it creates another -trivial- endpoint accessible at the route /readystatus. - This is used to poll if the service is up, however the user can provide it, and + """Manager for Flask applications running in a separate thread. + + It allows to have a Flask server under a thread, + start and stop it. Note that it creates another endpoint accessible at the route + /readystatus. This is used to poll if the service is up, however the user can + provide it. To use this code, one can use the following example: - + ```python from flask import Flask from flask_restful import Api app = Flask('some-name') @@ -66,27 +99,44 @@ class FlaskManager(threading.Thread): while not manager.is_ready(): from time import sleep sleep(0.1) - + ``` Then, later on, to stop it: - + + ```python manager.stop() - + ``` """ - def __init__(self, name, app, port, workers=1, host="0.0.0.0"): + def __init__( + self, + name: str, + app: Flask, + port: int, + workers: int = 1, + host: str = "0.0.0.0", + ) -> None: + """Initialize a FlaskManager. + + Args: + name: The name of the Flask manager. + app: The Flask application to manage. + port: The port to run the Flask server on. + workers: The number of Gunicorn workers. Defaults to 1. + host: The host address to bind to. Defaults to "0.0.0.0". + """ super(FlaskManager, self).__init__(daemon=True) self.log = get_logger(f"{name}-flaskmanager", stream_handlers=True) self.name = name self.app = app - self.prod_app = None - self.flask = None + self.prod_app: GunicornStandaloneApplication | None = None + self.flask: Process | None = None self.host = host self.port = port self.workers = workers - self.gunicorn_pid = None + self.gunicorn_pid: int | None = None self.ready = False self.joined = False self.ready_lock = threading.Lock() @@ -97,7 +147,7 @@ def _create_flask(self) -> Process: if "get_ready_status" in rule.endpoint: need_ready = False - def get_ready_status(): + def get_ready_status() -> str: return "ready" if need_ready: @@ -112,8 +162,10 @@ def get_ready_status(): "workers": self.workers, }, ) + prod_app = self.prod_app + assert prod_app is not None, "GunicornStandaloneApplication creation failed" - def run_gunicorn_with_signal_handling(): + def run_gunicorn_with_signal_handling() -> None: """Run gunicorn with SIGHUP ignored to prevent reload on shutdown. This prevents gunicorn from reloading when the parent process receives SIGHUP. @@ -127,7 +179,7 @@ def run_gunicorn_with_signal_handling(): # May fail if already in a process group or on some systems, ignore pass - self.prod_app.run() + prod_app.run() thread_name = f"{self.name}_thread" flask_srv = Process( # Indeed, we've just forked this sucker @@ -185,21 +237,33 @@ def run_gunicorn_with_signal_handling(): return flask_srv - def __del__(self): + def __del__(self) -> None: + """Cleanup when the FlaskManager is destroyed.""" self.stop() - def stop(self) -> NoReturn: - # gunicorn is forked, so we need to now need send signal ourselves + def stop(self) -> None: + """Stop the Flask manager and terminate the Gunicorn process. + + Sends SIGTERM to the Gunicorn process and joins the Flask process thread. + """ if self.gunicorn_pid: gunicorn_proc = psutil.Process(self.gunicorn_pid) # https://github.com/benoitc/gunicorn/blob/ab9c8301cb9ae573ba597154ddeea16f0326fc15/docs/source/signals.rst#master-process # TOTAL DESTRUCTION gunicorn_proc.send_signal(signal.SIGTERM) - self.flask.terminate() + if self.flask is not None: + self.flask.terminate() self.join() - def restart_renew(self): + def restart_renew(self) -> "FlaskManager": + """Restart and renew the Flask manager. + + Stops the current instance and creates a new one with the same configuration. + + Returns: + FlaskManager: A new FlaskManager instance with the same settings. + """ # well, we cannot really do that. # we have to hack it a bit: # unfortunately, this means you need to do: @@ -217,15 +281,25 @@ def restart_renew(self): time.sleep(0.1) return fm - def is_ready(self): + def is_ready(self) -> bool: + """Check if the Flask manager is ready to serve requests. + + Returns: + bool: True if ready, False otherwise. + """ with self.ready_lock: return self.ready - def is_terminated(self): + def is_terminated(self) -> bool: + """Check if the Flask manager has been terminated. + + Returns: + bool: True if terminated, False otherwise. + """ with self.ready_lock: return self.joined - def _create_and_join_flask(self): + def _create_and_join_flask(self) -> None: with self.ready_lock: self.ready = False self.joined = False @@ -238,16 +312,25 @@ def _create_and_join_flask(self): self.log.info(f"{self.name}-flaskmanager terminated") - def run(self) -> NoReturn: + def run(self) -> None: + """Run the Flask server in the thread. + + This method is called when the thread is started. + """ self._create_and_join_flask() -def main(): - class DummyEndpoint(Resource): - def post(self): +def main() -> None: + """Main entry point for demonstrating the FlaskManager. + + Creates a simple Flask application with a dummy endpoint and starts it. + """ + + class DummyEndpoint(_Resource): + def post(self) -> None: print(request) - def get(self): + def get(self) -> object: return make_response(jsonify({"weeeee": "wooo"})) app = Flask("test-app") diff --git a/src/drunc/utils/flask_typing.pyi b/src/drunc/utils/flask_typing.pyi new file mode 100644 index 000000000..b0a46f1c5 --- /dev/null +++ b/src/drunc/utils/flask_typing.pyi @@ -0,0 +1,29 @@ +"""Typing stubs for Flask/Gunicorn-related types used by drunc.utils.flask_manager. + +This module is intended to be imported only under ``TYPE_CHECKING`` to +avoid importing runtime dependencies during normal execution. +""" +from typing import Protocol + +from flask import Flask + +class _GunicornConfig(Protocol): + settings: dict[str, object] + + def set(self, key: str, value: object) -> None: ... + + +class _BaseApplication: + cfg: _GunicornConfig + + def __init__(self, *args: object, **kwargs: object) -> None: ... + def run(self) -> None: ... + + +class _Resource: ... + + +class Api: + def __init__(self, app: Flask) -> None: ... + + def add_resource(self, resource: type[_Resource], *urls: str, **kwargs: object) -> None: ... diff --git a/src/drunc/utils/grpc_utils.py b/src/drunc/utils/grpc_utils.py index 426ca0070..a228de8a7 100644 --- a/src/drunc/utils/grpc_utils.py +++ b/src/drunc/utils/grpc_utils.py @@ -1,5 +1,9 @@ +"""gRPC utilities for DRUNC.""" + +from __future__ import annotations + from dataclasses import dataclass -from typing import List, NoReturn, Optional +from typing import Callable, NoReturn, cast import grpc from druncschema.generic_pb2 import PlainText @@ -20,7 +24,15 @@ class UnpackingError(DruncCommandException): - def __init__(self, data, format): + """Exception raised when unpacking gRPC messages fails.""" + + def __init__(self, data: object, format: type[Message]) -> None: + """Initialize the UnpackingError. + + Args: + data: The data that failed to unpack. + format: The expected format. + """ self.data = data self.format = format @@ -50,13 +62,33 @@ def unpack_error_response(name: str, text: str, token: Token) -> Response: ) -def pack_to_any(data) -> any_pb2.Any: +def pack_to_any(data: Message) -> any_pb2.Any: + """Pack a protobuf message into an Any message. + + Args: + data: The protobuf message to pack. + + Returns: + any_pb2.Any: The packed message. + """ any = any_pb2.Any() any.Pack(data) return any -def unpack_any(data, format): +def unpack_any(data: any_pb2.Any, format: type[Message]) -> Message: + """Unpack an Any message into a specific protobuf format. + + Args: + data: The Any message to unpack. + format: The protobuf message type to unpack into. + + Returns: + Message: The unpacked message. + + Raises: + UnpackingError: If the message cannot be unpacked into the specified format. + """ if not data.Is(format.DESCRIPTOR): raise UnpackingError(data, format) req = format() @@ -65,13 +97,27 @@ def unpack_any(data, format): class ServerUnreachable(DruncException): - def __init__(self, message): + """Exception raised when the gRPC server is unreachable.""" + + def __init__(self, message: str) -> None: + """Initialize the ServerUnreachable exception. + + Args: + message: The error message. + """ self.message = message super(ServerUnreachable, self).__init__(message) class ServerTimeout(DruncException): - def __init__(self, message): + """Exception raised when the gRPC server times out.""" + + def __init__(self, message: str) -> None: + """Initialize the ServerTimeout exception. + + Args: + message: The error message. + """ self.message = message super(ServerTimeout, self).__init__(message) @@ -97,7 +143,7 @@ def server_is_reachable(grpc_error: grpc.RpcError) -> bool: return True -def rethrow_if_unreachable_server(grpc_error: grpc.RpcError) -> NoReturn: +def rethrow_if_unreachable_server(grpc_error: grpc.RpcError) -> None: """ Raise a ServerUnreachable exception if the gRPC error indicates the server is unreachable. @@ -114,7 +160,7 @@ def rethrow_if_unreachable_server(grpc_error: grpc.RpcError) -> NoReturn: raise ServerUnreachable(grpc_error._details) from grpc_error -def rethrow_if_timeout(grpc_error: grpc.RpcError) -> NoReturn: +def rethrow_if_timeout(grpc_error: grpc.RpcError) -> None: """ Raise a ServerTimeout if timeout. @@ -135,6 +181,7 @@ def handle_grpc_error(error: grpc.RpcError) -> NoReturn: Args: error: The gRPC error to handle. + Raises: A custom exception if the error matches a known category, or the original gRPC error if no classification applies. @@ -144,12 +191,11 @@ def handle_grpc_error(error: grpc.RpcError) -> NoReturn: raise error -def interrupt_if_unreachable_server(grpc_error: grpc.RpcError) -> Optional[str]: - """ - Interrupt if server is not reachable and return the error details. +def interrupt_if_unreachable_server(grpc_error: grpc.RpcError) -> str | None: + """Interrupt if server is not reachable and return the error details. Args: - grpc_error (grpc.RpcError): The gRPC error + grpc_error: The gRPC error Returns: str | None: The internal error details if the server is unreachable and details are available; @@ -157,9 +203,10 @@ def interrupt_if_unreachable_server(grpc_error: grpc.RpcError) -> Optional[str]: """ if not server_is_reachable(grpc_error): if hasattr(grpc_error, "_state"): - return grpc_error._state.details + return str(grpc_error._state.details) elif hasattr(grpc_error, "_details"): - return grpc_error._details + return str(grpc_error._details) + return None def copy_token(token: Token) -> Token: @@ -176,10 +223,19 @@ def copy_token(token: Token) -> Token: return token_copy -def dict_to_grpc_proto(data: dict, proto_class_instance: Message) -> Message: - """ - Converts a Python dictionary into an instance of a gRPC Protobuf message. +def dict_to_grpc_proto( + data: dict[str, object], proto_class_instance: Message +) -> Message: + """Converts a Python dictionary into an instance of a gRPC Protobuf message. + 'proto_class_instance' should be an empty instance, e.g., Token() + + Args: + data: The dictionary to convert. + proto_class_instance: An empty instance of the target protobuf message type. + + Returns: + Message: The converted protobuf message. """ return json_format.ParseDict(data, proto_class_instance, ignore_unknown_fields=True) @@ -199,21 +255,19 @@ class GrpcErrorDetails: Attributes: code (str): The gRPC status code name (e.g., "NOT_FOUND") message (str): The error message from the gRPC status - details (List[str]): A list of formatted error detail strings + details: A list of formatted error detail strings or protobuf Messages. """ code: str message: str - details: List[str] + details: list[str | Message] - def __str__(self): - """ - Return a human-readable string representation of the error. - """ + def __str__(self) -> str: + """Return a human-readable string representation of the error.""" lines = [f"[{self.code}] {self.message}"] for detail in self.details: # If it's a Proto message format the error detail - if hasattr(detail, "DESCRIPTOR"): + if isinstance(detail, Message): lines.extend(format_error_details(detail)) else: lines.append(str(detail)) @@ -312,13 +366,13 @@ def extract_grpc_rich_error(grpc_error: grpc.RpcError) -> GrpcErrorDetails: """ code = grpc_error.code().name if grpc_error.code() else "UNKNOWN" try: - status = rpc_status.from_call(grpc_error) + status = rpc_status.from_call(cast(grpc.Call, grpc_error)) except NotImplementedError: return GrpcErrorDetails(code=code, message="No message", details=[]) # Fallback to simple error if no rich status if status is None: - return GrpcErrorDetails(code=code, message="No message ", details=[]) + return GrpcErrorDetails(code=code, message="No message", details=[]) # Extract all error details error_details = [] @@ -342,9 +396,9 @@ def extract_grpc_rich_error(grpc_error: grpc.RpcError) -> GrpcErrorDetails: def abort_with_rich_error_status( context: grpc.ServicerContext, - grpc_error_code: code_pb2.Code, + grpc_error_code: int, message: str, - error_obj: Message, + error_obj: object, ) -> NoReturn: """ Aborts the current gRPC call with a rich error status containing @@ -378,21 +432,34 @@ def abort_with_rich_error_status( class RichErrorServerInterceptor(grpc.ServerInterceptor): """ A gRPC server interceptor that catches exceptions and converts them into - rich error statuses with structured error details.""" + rich error statuses with structured error details. + """ - def intercept_service(self, continuation, handler_call_details): + def intercept_service( # type: ignore[override] + self, + continuation: Callable[ + [grpc.HandlerCallDetails], + grpc.RpcMethodHandler[object, object] | None, + ], + handler_call_details: grpc.HandlerCallDetails, + ) -> grpc.RpcMethodHandler[object, object] | None: """ Intercept gRPC service calls to handle exceptions and convert them into rich error statuses. """ handler = continuation(handler_call_details) + if handler is None: + return None - def error_wrapper(request, context): + def error_wrapper(request: object, context: grpc.ServicerContext) -> object: try: - return handler.unary_unary(request, context) + unary_unary = handler.unary_unary + if unary_unary is None: + return handler + return unary_unary(request, context) except DruncSetupException as e: - detail_obj = error_details_pb2.PreconditionFailure( + detail_obj_precondition = error_details_pb2.PreconditionFailure( violations=[ error_details_pb2.PreconditionFailure.Violation( type="MISSING OR INVALID", @@ -402,39 +469,48 @@ def error_wrapper(request, context): ] ) abort_with_rich_error_status( - context, e.grpc_error_code, str(e), detail_obj + context, + int(e.grpc_error_code), + str(e), + detail_obj_precondition, ) except DruncNotImplementedException as e: - detail_obj = error_details_pb2.ErrorInfo( + detail_obj_not_implemented = error_details_pb2.ErrorInfo( reason="NOT_IMPLEMENTED", domain="server", metadata={}, ) abort_with_rich_error_status( - context, e.grpc_error_code, str(e), detail_obj + context, + int(e.grpc_error_code), + str(e), + detail_obj_not_implemented, ) except DruncCommandException as e: exception_data = e.detail_kwargs - detail_obj = error_details_pb2.ErrorInfo( + detail_obj_command = error_details_pb2.ErrorInfo( reason=str(e.message), domain=str( exception_data.get("domain", ""), ), ) abort_with_rich_error_status( - context, e.grpc_error_code, str(e), detail_obj + context, + int(e.grpc_error_code), + str(e), + detail_obj_command, ) except Exception as e: # Fallback - detail_obj = error_details_pb2.ErrorInfo( + detail_obj_fallback = error_details_pb2.ErrorInfo( reason="Unexpected error", domain="server", metadata={"original_error": str(type(e))}, ) abort_with_rich_error_status( - context, code_pb2.INTERNAL, str(e), detail_obj + context, int(code_pb2.INTERNAL), str(e), detail_obj_fallback ) if handler.unary_unary: diff --git a/src/drunc/utils/shell_utils.py b/src/drunc/utils/shell_utils.py index d1443968a..02c9ab46f 100644 --- a/src/drunc/utils/shell_utils.py +++ b/src/drunc/utils/shell_utils.py @@ -1,6 +1,9 @@ +"""Shell utilities for DRUNC.""" + import abc import getpass -from collections.abc import Mapping +from collections.abc import MutableMapping +from typing import Callable, ParamSpec, Protocol, TypeVar, cast import click from druncschema.token_pb2 import Token @@ -10,13 +13,80 @@ from drunc.utils.utils import get_logger +class CommandLike(Protocol): + """Protocol for command-like objects.""" + + name: str + + +class SequenceLike(Protocol): + """Protocol for sequence-like objects.""" + + id: str + + +class FSMDescriptionLike(Protocol): + """Protocol for FSM description-like objects.""" + + commands: list[CommandLike] + sequences: list[SequenceLike] + + +class DescribeFSMReplyLike(Protocol): + """Protocol for describe FSM reply-like objects.""" + + description: FSMDescriptionLike + + +class StatusLike(Protocol): + """Protocol for status-like objects.""" + + state: str + in_error: bool + + +class StatusReplyLike(Protocol): + """Protocol for status reply-like objects.""" + + status: StatusLike + + +class ControllerDriverProtocol(Protocol): + """Protocol for controller driver objects.""" + + def status(self) -> StatusReplyLike: + """Get the current status. + + Returns: + StatusReplyLike: The current status. + """ + ... + + def describe_fsm(self) -> DescribeFSMReplyLike: + """Describe the FSM. + + Returns: + DescribeFSMReplyLike: The FSM description. + """ + ... + + +P = ParamSpec("P") +R = TypeVar("R") + + class InterruptedCommand(DruncShellException): - """This exception gets thrown if we don't want to have a full stack, but still want to interrupt a **shell** command""" + """Exception thrown to interrupt a shell command without a full stack trace.""" pass def create_dummy_token_from_uname() -> Token: + """Create a dummy token from the current username. + + Returns: + Token: A dummy token with the current username. + """ user = getpass.getuser() return ( Token( # fake token, but should be figured out from the environment/authoriser @@ -25,8 +95,14 @@ def create_dummy_token_from_uname() -> Token: ) -def add_traceback_flag(): - def wrapper(f0): +def add_traceback_flag() -> Callable[[Callable[P, R]], Callable[P, R]]: + """Add a traceback flag to a command. + + Returns: + Callable: A decorator that adds the traceback flag. + """ + + def wrapper(f0: Callable[P, R]) -> Callable[P, R]: f1 = click.option( "-t/-nt", "--traceback/--no-traceback", @@ -39,14 +115,35 @@ def wrapper(f0): class DecodedResponse: - ## Warning! This should be kept in sync with druncschema/request_response.proto/Response class + """Decoded response object. + + Warning: This should be kept in sync with + druncschema/request_response.proto/Response class + """ + name = None token = None data = None flag = None - children = [] + children: list["DecodedResponse"] = [] + + def __init__( + self, + name: str, + token: Token, + flag: object, + data: object | None = None, + children: list["DecodedResponse"] | None = None, + ) -> None: + """Initialize a DecodedResponse. - def __init__(self, name, token, flag, data=None, children=None): + Args: + name: The name of the response. + token: The token associated with the response. + flag: The response flag. + data: The response data. Defaults to None. + children: Child responses. Defaults to None. + """ self.name = name self.token = token self.flag = flag @@ -57,29 +154,56 @@ def __init__(self, name, token, flag, data=None, children=None): self.children = children @staticmethod - def str(obj, prefix=""): + def to_string(obj: "DecodedResponse", prefix: str = "") -> str: + """Convert a DecodedResponse to a string representation. + + Args: + obj: The DecodedResponse to convert. + prefix: A prefix to add to the string. Defaults to empty string. + + Returns: + str: The string representation of the response. + """ text = ( f"{prefix} {obj.name} -> response flag={obj.flag} type={type(obj.data)}\n" ) for v in obj.children: if v is None: continue - text += DecodedResponse.str(v, prefix + " ") + text += DecodedResponse.to_string(v, prefix + " ") return text - def __str__(self): - return DecodedResponse.str(self) + def __str__(self) -> str: + """Return string representation of the DecodedResponse. + + Returns: + str: The string representation. + """ + return DecodedResponse.to_string(self) class ShellContext: - def _reset(self, name: str, token_args: dict = {}, driver_args: dict = {}): + """Base class for shell contexts.""" + + def _reset( + self, + name: str, + token_args: dict[str, object] = {}, + driver_args: dict[str, object] = {}, + ) -> None: self._console = Console() self._token = self.create_token(**token_args) - self._drivers: Mapping[str, object] = self.create_drivers(**driver_args) + self._drivers: MutableMapping[str, object] = self.create_drivers(**driver_args) + + def __init__(self, *args: object, **kwargs: object) -> None: + """Initialize the shell context. - def __init__(self, *args, **kwargs): + Args: + *args: Additional positional arguments. + **kwargs: Additional keyword arguments. + """ log = get_logger("utils.ShellContext") - self.dynamic_commands = set() + self.dynamic_commands: set[str] = set() try: self.reset(*args, **kwargs) except Exception as e: @@ -87,27 +211,71 @@ def __init__(self, *args, **kwargs): exit(1) @abc.abstractmethod - def reset(self, **kwargs): + def reset(self, **kwargs: object) -> None: + """Reset the shell context. + + Args: + **kwargs: Additional keyword arguments. + """ pass @abc.abstractmethod - def create_drivers(self, **kwargs) -> Mapping[str, object]: + def create_drivers(self, **kwargs: object) -> MutableMapping[str, object]: + """Create drivers for the context. + + Args: + **kwargs: Additional keyword arguments. + + Returns: + MutableMapping[str, object]: A mapping of driver names to driver objects. + """ pass @abc.abstractmethod - def create_token(self, **kwargs) -> Token: + def create_token(self, **kwargs: object) -> Token: + """Create a token for the context. + + Args: + **kwargs: Additional keyword arguments. + + Returns: + Token: A token object. + """ pass @abc.abstractmethod def terminate(self) -> None: + """Terminate the shell context.""" pass def set_driver(self, name: str, driver: object) -> None: + """Set a driver in the context. + + Args: + name: The name of the driver. + driver: The driver object. + + Raises: + DruncShellException: If a driver with the same name already exists. + """ if name in self._drivers: raise DruncShellException(f"Driver {name} already present in this context") self._drivers[name] = driver - def get_driver(self, name: str = None, quiet_fail: bool = False) -> object: + def get_driver(self, name: str | None = None, quiet_fail: bool = False) -> object: + """Get a driver from the context. + + Args: + name: The name of the driver. If None, returns the only driver if there is exactly one. + quiet_fail: If True, return None on failure instead of raising an exception. + + Returns: + object: The driver object, or None if quiet_fail is True and the driver is not found. + + Raises: + DruncShellException: If there are multiple drivers and no name is specified. + SystemExit: If the driver is not found and quiet_fail is False. + """ try: if name: return self._drivers[name] @@ -127,9 +295,22 @@ def get_driver(self, name: str = None, quiet_fail: bool = False) -> object: ) # used to avoid having to catch multiple Attribute errors when this function gets called def has_driver(self, name: str) -> bool: + """Check if a driver exists in the context. + + Args: + name: The name of the driver. + + Returns: + bool: True if the driver exists, False otherwise. + """ return name in self._drivers def delete_driver(self, name: str) -> None: + """Delete a driver from the context. + + Args: + name: The name of the driver to delete. + """ log = get_logger("utils.ShellContext") if name in self._drivers: log.info(f"You will not be able to issue commands to the {name} anymore.") @@ -137,18 +318,37 @@ def delete_driver(self, name: str) -> None: log.info(f"{name.capitalize()} driver has been deleted.") def get_token(self) -> Token: + """Get the token from the context. + + Returns: + Token: The token object. + """ return self._token - def print(self, *args, **kwargs) -> None: - self._console.print(*args, **kwargs) # rich tables require console printing + def print(self, *args: object, **kwargs: object) -> None: + """Print to the console. + + Args: + *args: Positional arguments to pass to the console. + **kwargs: Keyword arguments to pass to the console. + """ + self._console.print(*args, **kwargs) # type: ignore[arg-type] + + def rule(self, *args: object, **kwargs: object) -> None: + """Print a rule to the console. - def rule(self, *args, **kwargs) -> None: - self._console.rule(*args, **kwargs) + Args: + *args: Positional arguments to pass to the console. + **kwargs: Keyword arguments to pass to the console. + """ + self._console.rule(*args, **kwargs) # type: ignore[arg-type] def print_status_summary(self) -> None: + """Print a summary of the FSM status and available transitions.""" log = get_logger("utils.ShellContext") - status = self.get_driver("controller").status().status - describe_fsm = self.get_driver("controller").describe_fsm().description + controller = cast(ControllerDriverProtocol, self.get_driver("controller")) + status = controller.status().status + describe_fsm = controller.describe_fsm().description current_state = status.state if status.in_error: log.error( diff --git a/src/drunc/utils/utils.py b/src/drunc/utils/utils.py index 3ba58dbd7..71b07ad58 100644 --- a/src/drunc/utils/utils.py +++ b/src/drunc/utils/utils.py @@ -1,3 +1,5 @@ +"""A set of utility functions for drunc.""" + import ctypes import logging import os @@ -11,11 +13,13 @@ from contextlib import closing from datetime import datetime from enum import Enum -from urllib.parse import urlparse +from typing import Protocol, cast +from urllib.parse import ParseResult, urlparse -from click import BadParameter +from click import BadParameter, Context, Parameter from daqpytools.logging import get_daq_logger, setup_root_logger -from requests import delete, get, patch, post +from requests import Response, delete, get, patch, post +from rich.console import Console from rich.logging import RichHandler from rich.progress import ( BarColumn, @@ -35,8 +39,8 @@ def get_root_logger(log_level: str) -> logging.Logger: - """ - Set up the base logger which all other loggers will inherit. + """Set up the base logger which all other loggers will inherit. + This base logger is named the 'drunc' logger, and functions similarly to the root logger. It should have no handlers attached to it. @@ -50,19 +54,32 @@ def get_root_logger(log_level: str) -> logging.Logger: return setup_root_logger("drunc", log_level) -def get_logger(logger_name: str, *args, **kwargs) -> logging.Logger: - """Returns / constructs default logging instances. Prepends all loggers with 'drunc' - to inherit from the root 'drunc' logger. - Wraps to the daqpytools implementation, see for more details - - Args: - logger_name (str): Name of the logger - args, kwargs: Passed without modification to the daqpytools implementation - """ - return get_daq_logger(f"drunc.{logger_name}", *args, **kwargs) +def get_logger( + logger_name: str, + log_level: int | str = logging.NOTSET, + use_parent_handlers: bool = True, + rich_handler: bool = False, + file_handler_path: str | None = None, + stream_handlers: bool = False, + ers_kafka_session: str | None = None, + throttle: bool = False, + **extras: object, +) -> logging.Logger: + """Get a logger instance for the given logger name.""" + return get_daq_logger( + f"drunc.{logger_name}", + log_level, + use_parent_handlers, + rich_handler, + file_handler_path, + stream_handlers, + ers_kafka_session, + throttle, + **extras, + ) -def get_shared_rich_console(logger: logging.Logger): +def get_shared_rich_console(logger: logging.Logger) -> Console | None: """ Traverses logger hierarchy to find a FormattedRichHandler's console. @@ -102,8 +119,7 @@ def get_shared_rich_console(logger: logging.Logger): def strip_non_drunc_loggers() -> None: - """ - Strip out all the basicConfig handlers from other repositories, which define + """Strip out all the basicConfig handlers from other repositories, which define handlers with the root logger. """ root = logging.getLogger() @@ -111,37 +127,92 @@ def strip_non_drunc_loggers() -> None: root.handlers.clear() -def get_random_string(length): +def get_random_string(length: int) -> str: + """Generate a random string of lowercase ASCII letters. + + Args: + length (int): The desired length of the random string. + + Returns: + str: A random string of the specified length. + """ letters = string.ascii_lowercase return "".join(random.choice(letters) for i in range(length)) -def regex_match(regex, string): +def regex_match(regex: str, string: str) -> bool: + """Check if a regex pattern matches a string. + + Args: + regex (str): The regular expression pattern. + string (str): The string to match against. + + Returns: + bool: True if the pattern matches, False otherwise. + """ return re.match(regex, string) is not None -def get_new_port(): +def get_new_port() -> int: + """Get an available port number. + + Returns: + int: An available port number. + """ with closing(socket.socket(socket.AF_INET, socket.SOCK_STREAM)) as s: s.bind(("", 0)) s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) - return s.getsockname()[1] + return int(s.getsockname()[1]) -def now_str(posix_friendly=False): +def now_str(posix_friendly: bool = False) -> str: + """Get the current time as a formatted string. + + Args: + posix_friendly (bool): If True, use POSIX-friendly format. Defaults to False. + + Returns: + str: The current time as a formatted string. + """ if not posix_friendly: return datetime.now().strftime("%m/%d/%Y,%H:%M:%S") else: return datetime.now().strftime("%Y-%m-%d-%H-%M-%S") -def expand_path(path, turn_to_abs_path=False): +def expand_path(path: str, turn_to_abs_path: bool = False) -> str: + """Expand a path with user and environment variables. + + Args: + path (str): The path to expand. + turn_to_abs_path (bool): If True, also convert to absolute path. + Defaults to False. + + Returns: + str: The expanded path. + """ if turn_to_abs_path: return os.path.abspath(os.path.expanduser(os.path.expandvars(path))) return os.path.expanduser(os.path.expandvars(path)) -def validate_command_facility(ctx, param, value): - parsed = "" +def validate_command_facility( + ctx: Context | None, param: Parameter | None, value: str +) -> str: + """Validate a command facility parameter. + + Args: + ctx (Any): Click context. + param (Any): Click parameter. + value (str): The value to validate. + + Returns: + str: The validated netloc. + + Raises: + BadParameter: If the value is invalid. + """ + parsed: ParseResult try: parsed = urlparse(value) except Exception as e: @@ -166,8 +237,7 @@ def validate_command_facility(ctx, param, value): def address_regex(address: str, hostname_or_ip: str) -> str: - """ - Replace 127.x.x.x and 0.x.x.x IPs with the provided hostname + """Replace 127.x.x.x and 0.x.x.x IPs with the provided hostname. This is useful when a service binds to localhost or 127.x.x.x, but we want to access it using the hostname or network IP. @@ -180,7 +250,7 @@ def address_regex(address: str, hostname_or_ip: str) -> str: str: The address with 127.x.x.x and 0.x.x.x replaced by the hostname or IP. """ - ip_match: re.Match = re.search( + ip_match: re.Match[str] | None = re.search( r"((25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)\.){3}(25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)", address, ) @@ -241,7 +311,15 @@ def resolve_localhost_and_127_ip_to_network_ip(address: str) -> str: return address_regex(address, this_ip) -def host_is_local(host): +def host_is_local(host: str) -> bool: + """Check if a host is local. + + Args: + host (str): The hostname or IP to check. + + Returns: + bool: True if the host is local, False otherwise. + """ if host in [ "localhost", socket.gethostname(), @@ -255,15 +333,21 @@ def host_is_local(host): return False -def pid_info_str(): +def pid_info_str() -> str: + """Get a string with process ID information. + + Returns: + str: A string containing the parent and current process IDs. + """ return f"Parent's PID: {os.getppid()} | This PID: {os.getpid()}" -def ignore_sigint_sighandler(): +def ignore_sigint_sighandler() -> None: + """Ignore SIGINT (Ctrl+C) signals.""" signal.signal(signal.SIGINT, signal.SIG_IGN) -def parent_death_pact(signal=signal.SIGHUP): +def parent_death_pact(signal: int = signal.SIGHUP) -> None: """Commit to kill current process when parent process dies. Each time you spawn a new process, run this to set signal handler appropriately (e.g put it at the beginning of each @@ -280,36 +364,82 @@ def parent_death_pact(signal=signal.SIGHUP): class IncorrectAddress(DruncException): + """Exception raised when an address is invalid.""" + pass -def https_or_http_present(address: str): +def https_or_http_present(address: str) -> None: + """Validate that an address starts with http:// or https://. + + Args: + address (str): The address to validate. + + Raises: + IncorrectAddress: If the address does not start with http:// or https://. + """ if not address.startswith("https://") and not address.startswith("http://"): raise IncorrectAddress("Endpoint should start with http:// or https://") -def http_post(address, data, as_json=True, ignore_errors=False, **post_kwargs): +def http_post( + address: str, + data: object, + as_json: bool = True, + ignore_errors: bool = False, + **post_kwargs: object, +) -> Response: + """Send an HTTP POST request. + + Args: + address (str): The URL to send the request to. + data (Any): The data to send in the request body. + as_json (bool): If True, send data as JSON. Defaults to True. + ignore_errors (bool): If True, do not raise exceptions for HTTP errors. Defaults to False. + **post_kwargs: Additional keyword arguments to pass to requests.post. + + Returns: + Response: The response from the server. + """ https_or_http_present(address) if as_json: - r = post(address, json=data, **post_kwargs) + r = post(address, json=data, **post_kwargs) # type: ignore[arg-type] else: - r = post(address, data=data, **post_kwargs) + r = post(address, data=data, **post_kwargs) # type: ignore[arg-type] if not ignore_errors: r.raise_for_status() return r -def http_get(address, data, as_json=True, ignore_errors=False, **post_kwargs): +def http_get( + address: str, + data: object, + as_json: bool = True, + ignore_errors: bool = False, + **post_kwargs: object, +) -> Response: + """Send an HTTP GET request. + + Args: + address (str): The URL to send the request to. + data (Any): The data to send in the request body. + as_json (bool): If True, send data as JSON. Defaults to True. + ignore_errors (bool): If True, do not raise exceptions for HTTP errors. Defaults to False. + **post_kwargs: Additional keyword arguments to pass to requests.get. + + Returns: + Response: The response from the server. + """ https_or_http_present(address) log = get_logger("utils.http_get") log.debug(f"GETTING {address} {data}") if as_json: - r = get(address, json=data, **post_kwargs) + r = get(address, json=data, **post_kwargs) # type: ignore[arg-type] else: - r = get(address, data=data, **post_kwargs) + r = get(address, data=data, **post_kwargs) # type: ignore[arg-type] log.debug(r.text) log.debug(r.status_code) @@ -320,32 +450,71 @@ def http_get(address, data, as_json=True, ignore_errors=False, **post_kwargs): return r -def http_patch(address, data, as_json=True, ignore_errors=False, **post_kwargs): +def http_patch( + address: str, + data: object, + as_json: bool = True, + ignore_errors: bool = False, + **post_kwargs: object, +) -> Response: + """Send an HTTP PATCH request. + + Args: + address (str): The URL to send the request to. + data (Any): The data to send in the request body. + as_json (bool): If True, send data as JSON. Defaults to True. + ignore_errors (bool): If True, do not raise exceptions for HTTP errors. Defaults to False. + **post_kwargs: Additional keyword arguments to pass to requests.patch. + + Returns: + Response: The response from the server. + """ https_or_http_present(address) if as_json: - r = patch(address, json=data, **post_kwargs) + r = patch(address, json=data, **post_kwargs) # type: ignore[arg-type] else: - r = patch(address, data=data, **post_kwargs) + r = patch(address, data=data, **post_kwargs) # type: ignore[arg-type] if not ignore_errors: r.raise_for_status() return r -def http_delete(address, data, as_json=True, ignore_errors=False, **post_kwargs): +def http_delete( + address: str, + data: object, + as_json: bool = True, + ignore_errors: bool = False, + **post_kwargs: object, +) -> None: + """Send an HTTP DELETE request. + + Args: + address (str): The URL to send the request to. + data (Any): The data to send in the request body. + as_json (bool): If True, send data as JSON. Defaults to True. + ignore_errors (bool): If True, do not raise exceptions for HTTP errors. Defaults to False. + **post_kwargs: Additional keyword arguments to pass to requests.delete. + """ https_or_http_present(address) if as_json: - r = delete(address, json=data, **post_kwargs) + r = delete(address, json=data, **post_kwargs) # type: ignore[arg-type] else: - r = delete(address, data=data, **post_kwargs) + r = delete(address, data=data, **post_kwargs) # type: ignore[arg-type] if not ignore_errors: r.raise_for_status() +class _ConnectivityService(Protocol): + def resolve(self, name: str, message_type: str) -> list[dict[str, object]]: ... + + class ControlType(Enum): + """Enumeration of control types for DUNE DAQ services.""" + Unknown = 0 gRPC = 1 REST_API = 2 @@ -353,6 +522,17 @@ class ControlType(Enum): def get_control_type_and_uri_from_cli(cli_args: list[str]) -> tuple[ControlType, str]: + """Extract control type and URI from CLI arguments. + + Args: + cli_args (list[str]): The CLI arguments to parse. + + Returns: + tuple[ControlType, str]: A tuple of (control_type, uri). + + Raises: + DruncSetupException: If protocol is not 'grpc://' or 'rest://'. + """ for arg in cli_args: if arg.startswith("rest://"): uri = arg.replace("rest://", "") @@ -366,19 +546,35 @@ def get_control_type_and_uri_from_cli(cli_args: list[str]) -> tuple[ControlType, def get_control_type_and_uri_from_connectivity_service( - connectivity_service, + connectivity_service: _ConnectivityService, name: str, timeout: int = 10, # seconds retry_wait: float = 0.1, # seconds progress_bar: bool = False, - title: str = None, + title: str | None = None, ) -> tuple[ControlType, str]: - uris = [] + """Get control type and URI from connectivity service. + + Args: + connectivity_service (object): The connectivity service instance. + name (str): The name of the service to resolve. + timeout (int): Maximum time to wait for resolution in seconds. Defaults to 10. + retry_wait (float): Time to wait between retries in seconds. Defaults to 0.1. + progress_bar (bool): Whether to display a progress bar. Defaults to False. + title (str | None): Title for the progress bar. Defaults to None. + + Returns: + tuple[ControlType, str]: A tuple of (control_type, uri). + + Raises: + ApplicationLookupUnsuccessful: If the URI cannot be resolved. + """ + uris: list[dict[str, object]] = [] logger = get_logger("utils.get_control_type_and_uri_from_connectivity_service") shared_console = get_shared_rich_console(logger) start = time.time() - elapsed = 0 + elapsed = 0.0 if progress_bar: with Progress( @@ -438,21 +634,25 @@ def get_control_type_and_uri_from_connectivity_service( f"Could not resolve the URI for '{name}_control' in the connectivity service, got response {uris}" ) - uri = uris[0]["uri"] + uri = cast(str, uris[0]["uri"]) return get_control_type_and_uri_from_cli([uri]) -def print_with_timestamp(message): +def print_with_timestamp(message: str) -> None: + """Print a message with a timestamp. + + Args: + message (str): The message to print. + """ now = datetime.now() now_str = now.isoformat() print(f"{now_str}: {message}") def format_name_for_cli(name: str) -> str: - """ - Format a command name or argument name to be CLI-friendly by replacing underscores - with hyphens and converting to lowercase. + """Format a command name or argument name to be CLI-friendly by replacing + underscores with hyphens and converting to lowercase. Args: name (str): The original command name. @@ -464,16 +664,16 @@ def format_name_for_cli(name: str) -> str: def resolve_target_ip(host: str) -> str | None: - """ - Intelligently resolves the host. + """Intelligently resolve a host to its IP address. + If host is 'localhost' or '127.0.0.1', it finds the actual LAN IP. Args: - host - the name of the host to reolve to LAN IP + host (str): The name of the host to resolve to LAN IP. Returns: - str - LAN IP of the host - None - if the host could not be resolved, None is returned + str: LAN IP of the host. + None: If the host could not be resolved, None is returned. """ log = get_logger("utils.resolve_target_ip") @@ -491,7 +691,7 @@ def resolve_target_ip(host: str) -> str | None: # blocked from sending data outside of the LAN. Use connect - this does # send any data, just establishes the connection. s.connect(("10.255.255.255", 1)) - return s.getsockname()[0] + return str(s.getsockname()[0]) except Exception: # Return the loopback address. log.warning(f"Failed to resolve the IP address of {host}") @@ -507,17 +707,15 @@ def resolve_target_ip(host: str) -> str | None: def is_port_available(host: str, port: int, timeout: int = 2) -> bool: - """ - Check if the given port number on a specified host is available. + """Check if the given port number on a specified host is available. Args: - host - the host name to check - port - the port number to check - timeout - timeout of attempting to establish the connection + host (str): The host name to check. + port (int): The port number to check. + timeout (int): Timeout of attempting to establish the connection. Defaults to 2. Returns: - true - the port is available - false - the port is not available + bool: True if the port is available, False otherwise. """ log = get_logger("utils.is_port_available") @@ -554,13 +752,12 @@ def is_port_available(host: str, port: int, timeout: int = 2) -> bool: def file_is_read_only(file_path: str) -> bool: - """ - Runs checks to see if the file path is read only. + """Check if a file is read-only. Args: - file_path - path of file to read + file_path (str): Path of file to check. Returns: - bool - true is file is read only, false otherwise + bool: True if the file is read-only, False otherwise. """ return not os.access(file_path, os.W_OK) From cedbf88bf3046843f4ea83959c5bf80d1476d9be Mon Sep 17 00:00:00 2001 From: James Paul Turner Date: Mon, 13 Jul 2026 18:20:07 +0100 Subject: [PATCH 15/20] Phased merge 9 --- .../drunc/integtest}/integ_test_utils.py | 0 src/drunc/process_manager/configuration.py | 100 ++++--- .../process_manager/interface/context.py | 23 +- .../interface/process_manager.py | 14 +- src/drunc/process_manager/interface/shell.py | 2 - .../process_manager/k8s_process_manager.py | 253 ++++++++++-------- src/drunc/process_manager/process_manager.py | 109 +------- .../process_manager/ssh_process_manager.py | 21 +- src/drunc/process_manager/utils.py | 41 ++- src/drunc/processes/process_metadata.py | 37 ++- ...ss_lifetime_manager_from_forked_process.py | 6 +- .../ssh_process_lifetime_manager_paramiko.py | 6 +- .../ssh_process_lifetime_manager_shell.py | 79 +++--- src/drunc/session_manager/configuration.py | 5 +- .../interface/session_manager.py | 4 +- src/drunc/unified_shell/context.py | 67 +++-- src/drunc/unified_shell/shell.py | 26 +- 17 files changed, 395 insertions(+), 398 deletions(-) rename {integtest => src/drunc/integtest}/integ_test_utils.py (100%) diff --git a/integtest/integ_test_utils.py b/src/drunc/integtest/integ_test_utils.py similarity index 100% rename from integtest/integ_test_utils.py rename to src/drunc/integtest/integ_test_utils.py diff --git a/src/drunc/process_manager/configuration.py b/src/drunc/process_manager/configuration.py index da9198722..5eedbaa7c 100644 --- a/src/drunc/process_manager/configuration.py +++ b/src/drunc/process_manager/configuration.py @@ -3,7 +3,7 @@ import sys from enum import Enum from importlib import resources -from typing import TYPE_CHECKING, Any, Dict, Union +from typing import TYPE_CHECKING, Any, Dict, Self, Union from urllib.parse import unquote, urlparse from jsonschema import ValidationError @@ -12,7 +12,6 @@ from opmonlib.publisher import OpMonPublisher from opmonlib.utils import parse_opmon_conf -from drunc.broadcast.server.configuration import KafkaBroadcastSenderConfData from drunc.exceptions import DruncCommandException from drunc.process_manager.exceptions import UnknownProcessManagerType from drunc.utils.configuration import ConfHandler @@ -38,93 +37,92 @@ class ProcessManagerTypes(Enum): SSH_PARAMIKO = 3 -class ProcessManagerConfData: - def __init__(self): +class ProcessManagerConfHandler(ConfHandler): + """Handler for process manager configuration.""" + + log_path: str = "./" + + def populate_from_dict(self, data: dict[str, object]) -> None: self.broadcaster = None self.authoriser = None - self.type = ProcessManagerTypes.Unknown + self.pm_type = ProcessManagerTypes.Unknown self.command_address = "" self.environment = {} self.settings = {} + self.opmon_conf = None self.opmon_uri = None self.opmon_publisher = None - -class ProcessManagerConfHandler(ConfHandler): - def __init__(self, log_path: str, *args, **kwargs): - super().__init__(*args, **kwargs) - self.log_path = log_path - self.log = get_logger("process_manager.conf_handler") - - def get_log_path(self): - return self.log_path - - def _parse_dict(self, data): - new_data = ProcessManagerConfData() - if data.get("broadcaster"): - new_data.broadcaster = KafkaBroadcastSenderConfData.from_dict( - data.get("broadcaster") - ) - new_data.environment = data.get("environment", {}) - new_data.settings = data.get("settings", {}) + self.environment = data.get("environment", {}) + self.settings = data.get("settings", {}) + self.opmon_conf = data.get("opmon_conf") + self.opmon_uri = data.get("opmon_uri") match data["type"].lower(): case "ssh": - new_data.type = ProcessManagerTypes.SSH_SHELL - new_data.kill_timeout = data.get("kill_timeout", 0.5) + self.pm_type = ProcessManagerTypes.SSH_SHELL + self.kill_timeout = data.get("kill_timeout", 0.5) case "ssh-paramiko": - new_data.type = ProcessManagerTypes.SSH_PARAMIKO - new_data.kill_timeout = data.get("kill_timeout", 0.5) + self.pm_type = ProcessManagerTypes.SSH_PARAMIKO + self.kill_timeout = data.get("kill_timeout", 0.5) case "k8s": - new_data.type = ProcessManagerTypes.K8s - new_data.image = data.get("image", "ghcr.io/dune-daq/alma9:latest") + self.pm_type = ProcessManagerTypes.K8s + self.image = data.get("image", "ghcr.io/dune-daq/alma9:latest") case _: raise UnknownProcessManagerType(data["type"]) - # opmon_publisher left as default None - self.opmon_conf = parse_opmon_conf( + @classmethod + def from_json( + cls, path: str, session_name: str | None = None, log_path: str = "./" + ) -> Self: + """Create handler from JSON file with optional log path.""" + instance = super().from_json(path, session_name) + instance.log_path = log_path + instance.log = get_logger("process_manager.conf_handler") + return instance + + def get_log_path(self): + return self.log_path + + def _post_process_oks(self) -> None: + """Post-process to handle OpMon configuration.""" + opmon_conf = parse_opmon_conf( log=self.log, - conf=data.get("opmon_conf", None), - uri=data.get("opmon_uri", None), - session=new_data.type.name, + conf=getattr(self, "opmon_conf", None), + uri=getattr(self, "opmon_uri", None), + session=getattr(self, "pm_type", ProcessManagerTypes.Unknown).name, application="process_manager", ) - if self.opmon_conf.path == "./info.json": - self.opmon_conf.path = ( - "./info." - + self.opmon_conf.session - + "." - + self.opmon_conf.application - + ".json" + if opmon_conf.path == "./info.json": + opmon_conf.path = ( + "./info." + opmon_conf.session + "." + opmon_conf.application + ".json" ) self.log.debug( - "Initializing process manager OpMon with configuration %s", self.opmon_conf + "Initializing process manager OpMon with configuration %s", opmon_conf ) try: - if self.opmon_conf.opmon_type == "stream": - new_data.opmon_publisher = KafkaOpMonPublisher(self.opmon_conf) + if opmon_conf.opmon_type == "stream": + self.opmon_publisher = KafkaOpMonPublisher(opmon_conf) self.log.debug( "KafkaOpMonPublisher initialized with configuration %s", - self.opmon_conf, + opmon_conf, ) else: - new_data.opmon_publisher = OpMonPublisher( - conf=self.opmon_conf, rich_handler=True + self.opmon_publisher = OpMonPublisher( + conf=opmon_conf, rich_handler=True ) self.log.debug( "%s OpMonPublisher initialized with configuration %s", - self.opmon_conf.opmon_type, - self.opmon_conf, + opmon_conf.opmon_type, + opmon_conf, ) except Exception as e: self.log.error("Failed to initialize OpMonPublisher: %s", e) raise DruncCommandException("Failed to initialize OpMonPublisher.") - return new_data - def get_commandline_parameters( config_filename: str, diff --git a/src/drunc/process_manager/interface/context.py b/src/drunc/process_manager/interface/context.py index a81396130..d6a30f592 100644 --- a/src/drunc/process_manager/interface/context.py +++ b/src/drunc/process_manager/interface/context.py @@ -1,24 +1,21 @@ -from collections.abc import Mapping +from collections.abc import MutableMapping from druncschema.token_pb2 import Token -from drunc.broadcast.client.broadcast_handler import BroadcastHandler -from drunc.broadcast.client.configuration import BroadcastClientConfHandler from drunc.process_manager.process_manager_driver import ProcessManagerDriver -from drunc.utils.configuration import ConfTypes from drunc.utils.shell_utils import ( ShellContext, create_dummy_token_from_uname, ) -from drunc.utils.utils import get_logger, resolve_localhost_to_hostname +from drunc.utils.utils import resolve_localhost_to_hostname -class ProcessManagerContext(ShellContext): # boilerplatefest +class ProcessManagerContext(ShellContext): def __init__(self, *args, **kwargs): self.status_receiver = None super(ProcessManagerContext, self).__init__(*args, **kwargs) - def reset(self, address: str = None): + def reset(self, address: str = "", **kwargs): self.address = resolve_localhost_to_hostname(address) super(ProcessManagerContext, self)._reset( name="process_manager_context", @@ -26,7 +23,7 @@ def reset(self, address: str = None): driver_args={}, ) - def create_drivers(self, **kwargs) -> Mapping[str, object]: + def create_drivers(self, **kwargs) -> MutableMapping[str, object]: if not self.address: return {} return { @@ -39,16 +36,6 @@ def create_drivers(self, **kwargs) -> Mapping[str, object]: def create_token(self, **kwargs) -> Token: return create_dummy_token_from_uname() - def start_listening(self, broadcaster_conf): - bcch = BroadcastClientConfHandler( - data=broadcaster_conf, - type=ConfTypes.ProtobufAny, - ) - self.status_receiver = BroadcastHandler(bcch) - get_logger("process_manager.shell").info( - f":ear: Listening to the Process Manager at {self.address}" - ) - def terminate(self): if self.status_receiver: self.status_receiver.stop() diff --git a/src/drunc/process_manager/interface/process_manager.py b/src/drunc/process_manager/interface/process_manager.py index 6e3d1b294..eb2fa4e13 100644 --- a/src/drunc/process_manager/interface/process_manager.py +++ b/src/drunc/process_manager/interface/process_manager.py @@ -22,7 +22,7 @@ ) from drunc.process_manager.process_manager import ProcessManager from drunc.process_manager.utils import get_log_path -from drunc.utils.configuration import parse_conf_url +from drunc.utils.configuration import ConfTypes, parse_conf_url from drunc.utils.grpc_utils import RichErrorServerInterceptor from drunc.utils.utils import ( get_logger, @@ -60,14 +60,16 @@ def run_pm( log.debug("Process manager configuration is valid.") conf_path, conf_type = parse_conf_url(pm_conf) + path_or_url = conf_path.split(":")[1] - pmch = ProcessManagerConfHandler( - log_path=log_path, type=conf_type, data=conf_path.split(":")[1] - ) + if conf_type == ConfTypes.JsonFileName: + pmch = ProcessManagerConfHandler.from_json(path=path_or_url, log_path=log_path) + else: + pmch = ProcessManagerConfHandler.from_pyobject(data=path_or_url) log_path = get_log_path( user=getpass.getuser(), - session_name=pmch.data.type.name, + session_name=getattr(pmch, "pm_type", pmch.type).name, application_name=appName, override_logs=override_logs, app_log_path=log_path, @@ -76,7 +78,7 @@ def run_pm( # Logger has been added to process_manager, so everything will be logged add_handler(log, HandlerType.File, True, path=log_path) - for key, value in pmch.data.environment.items(): + for key, value in pmch.environment.items(): os.environ[key] = value pm = ProcessManager.get(pmch, name="process_manager") diff --git a/src/drunc/process_manager/interface/shell.py b/src/drunc/process_manager/interface/shell.py index 5da7434c5..18f7962e2 100644 --- a/src/drunc/process_manager/interface/shell.py +++ b/src/drunc/process_manager/interface/shell.py @@ -71,8 +71,6 @@ def process_manager_shell(ctx, process_manager_address: str, log_level: str) -> process_manager_shell_log.info( f"Connected to {process_manager_address}, running '{desc.name}.{desc.session}' (name.session), starting listening..." ) - if desc.HasField("broadcast"): - ctx.obj.start_listening(desc.broadcast) def cleanup(): ctx.obj.terminate() diff --git a/src/drunc/process_manager/k8s_process_manager.py b/src/drunc/process_manager/k8s_process_manager.py index 152da18c6..c237b0725 100644 --- a/src/drunc/process_manager/k8s_process_manager.py +++ b/src/drunc/process_manager/k8s_process_manager.py @@ -9,10 +9,10 @@ import urllib.error import urllib.request import uuid +from dataclasses import dataclass from time import sleep, time # Local Application Imports -from druncschema.broadcast_pb2 import BroadcastType from druncschema.process_manager_pb2 import ( BootRequest, LogLines, @@ -42,6 +42,7 @@ ) from drunc.process_manager.process_manager import ProcessManager from drunc.process_manager.utils import ( + compute_role_from_boot_request, format_hostname, on_parent_exit, validate_k8s_session_name, @@ -49,6 +50,21 @@ from drunc.utils.utils import get_logger, resolve_localhost_to_hostname +@dataclass +class _LcsSessionState: + """ + Holds all Local Connection Server (LCS) state for a single session (k8s namespace). + + Using a dataclass instead of parallel dicts prevents the fields from drifting + out of sync and makes it obvious that they all describe the same LCS instance. + A session that has no LCS simply has no entry in K8sProcessManager._lcs_state. + """ + + podname: str | None = None + node_port: int | None = None + is_booted: bool = False + + class K8sPodWatcherThread(threading.Thread): def __init__(self, pm) -> None: """ @@ -208,7 +224,11 @@ def __init__(self, configuration: ProcessManagerConfHandler, **kwargs) -> None: self.uuids_pending_deletion = set() self.termination_complete_event = threading.Event() self.final_exit_codes = {} - self.local_connection_server_is_booted = False + # Per-session LCS state. Keyed by session (k8s namespace) name. + # Sessions that have no LCS simply have no entry here. + # Use _lcs_state_for(session) to read and _lcs_state.pop(session) to clean up. + self._lcs_state: dict[str, _LcsSessionState] = {} + self._lcs_state_lock = threading.Lock() # Host verification cache: {hostname: (is_valid, timestamp)} self._host_cache = {} @@ -217,7 +237,7 @@ def __init__(self, configuration: ProcessManagerConfHandler, **kwargs) -> None: # Get settings from configuration JSON file # Any comments following this one will relate to the parameters retrieved from # the configuration file if the comment starts as "CONFIGURATION -" - settings = getattr(self.configuration.data, "settings", {}) + settings = getattr(self.configuration, "settings", {}) # CONFIGURATION - label defaults labels = settings.get("labels", {}) @@ -229,10 +249,6 @@ def __init__(self, configuration: ProcessManagerConfHandler, **kwargs) -> None: # Readout app selector self.perf_selector = settings.get("readout_app_selector", "runp").lower() - # CONFIGURATION - connection server connection port numbers - self.connection_server_port = None - self.connection_server_node_port = None - # CONFIGURATION - per-pod service port number service = settings.get("service", {}) self.headless_discovery_port = service.get("headless_discovery_port", 80) @@ -336,8 +352,8 @@ def notify_termination( """ Callback for when a pod terminates. - Updates the final exit code, broadcasts a status update, and signals - the termination_complete_event when all pending deletions are confirmed. + Updates the final exit code and signals the termination_complete_event + when all pending deletions are confirmed. Args: proc_uuid: The UUID string of the terminated process. @@ -358,7 +374,20 @@ def notify_termination( # Publish this information self.log.info(end_str) - self.broadcast(end_str, BroadcastType.SUBPROCESS_STATUS_UPDATE) + + # If the terminated pod was the LCS for its session, remove all LCS + # state for this session. A partial reset (e.g. only is_booted=False) + # would leave stale podname/port/node_port fields that could mislead + # a future LCS boot — including one that lands on a different node. + # Popping the entry is consistent with the invariant: "no LCS = no entry". + with self._lcs_state_lock: + lcs = self._lcs_state.get(session) + if lcs is not None and lcs.podname == meta.name: + self.log.info( + f"LCS pod '{meta.name}' for session '{session}' has terminated; " + "removing LCS state so any future boot starts from a clean slate." + ) + self._lcs_state.pop(session, None) # Clear the list of processes being removed if proc_uuid in self.uuids_pending_deletion: @@ -505,6 +534,18 @@ def _is_root_controller(self, tree_labels: dict[str, str]) -> bool: """ return tree_labels.get(f"role.{self.drunc_label}") == "root-controller" + def _lcs_state_for(self, session: str) -> _LcsSessionState: + """ + Return the LCS state for *session*, creating a fresh entry if absent. + + Centralises dict access so callers never deal with missing-key logic. + Thread-safe: the check-then-create is performed under _lcs_state_lock. + """ + with self._lcs_state_lock: + if session not in self._lcs_state: + self._lcs_state[session] = _LcsSessionState() + return self._lcs_state[session] + def _is_host_cached(self, host: str) -> None | bool: """ Check if host is cached and not expired. @@ -765,20 +806,26 @@ def _create_headless_service(self, podname, session, pod_uid) -> None: if e.status != 409: self.log.error(f"Failed to create headless service for {podname}: {e}") - def _create_nodeport_service(self, podname, session, pod_uid) -> None: + def _create_nodeport_service( + self, + podname: str, + session: str, + pod_uid: str, + port: int, + node_port: int, + ) -> None: """ Create a NodePort Kubernetes Service for external access. - Builds and creates a NodePort Service with externalTrafficPolicy=Local, - mapping the connection_server_port to a fixed NodePort - (connection_server_node_port). The service is owned by the pod via an - OwnerReference. Raises a DruncK8sException if the NodePort is already - allocated or another API error occurs. + Builds and creates a NodePort Service with externalTrafficPolicy=Local. + The service is owned by the pod via an OwnerReference. Args: - podname - the name of the pod (also used as the service name) - session - the Kubernetes namespace (session) to create the service in - pod_uid - the UID of the owning pod for the OwnerReference + podname - the name of the pod (also used as the service name) + session - the Kubernetes namespace (session) to create the service in + pod_uid - the UID of the owning pod for the OwnerReference + port - the container port to expose + node_port - the fixed NodePort to allocate on cluster nodes Raises: DruncK8sException - if the NodePort is already in use or another API error occurs @@ -808,9 +855,9 @@ def _create_nodeport_service(self, podname, session, pod_uid) -> None: ports=[ client.V1ServicePort( protocol="TCP", - port=self.connection_server_port, - target_port=self.connection_server_port, - node_port=self.connection_server_node_port, + port=port, + target_port=port, + node_port=node_port, ) ], ), @@ -820,8 +867,8 @@ def _create_nodeport_service(self, podname, session, pod_uid) -> None: namespace=session, body=service_manifest ) self.log.info( - f'Created NodePort service "{session}.{podname}" on port {self.connection_server_port} ' - f"(NodePort: {self.connection_server_node_port} for external access)" + f'Created NodePort service "{session}.{podname}" on port {port} ' + f"(NodePort: {node_port} for external access)" ) except self._api_error_v1_api as e: @@ -836,9 +883,8 @@ def _create_nodeport_service(self, podname, session, pod_uid) -> None: is_port_conflict = True if is_port_conflict: - port = self.connection_server_node_port error_message = ( - f"NodePort {port} is already in use by another service. " + f"NodePort {node_port} is already in use by another service. " f"Cannot start '{podname}'." ) self.log.error(error_message) @@ -1049,50 +1095,15 @@ def _get_tree_labels( A dictionary of labels containing 'tree-id.{drunc_label}' and 'role.{drunc_label}' keys with their corresponding values. """ - role = "unknown" + role = compute_role_from_boot_request(boot_request) labels = {f"tree-id.{self.drunc_label}": tree_id} - - if not tree_id: - role = "unknown" - elif self._is_controller_executable(boot_request): - if tree_id == "0": - role = "root-controller" - elif tree_id.startswith("0."): - role = "segment-controller" - else: - role = "infrastructure-applications" # controller outside segment tree - - else: - if tree_id.startswith("0."): - role = "application" - else: - role = "infrastructure-applications" - labels[f"role.{self.drunc_label}"] = role self.log.info( f"Assigning labels for '{podname}': role={role}, tree-id={tree_id}" ) return labels - def _is_controller_executable(self, boot_request: BootRequest) -> bool: - """ - Check whether the boot request's main executable is a drunc-controller. - - Inspects all executable-and-arguments entries in the process description - for the 'drunc-controller' executable name. - - Args: - boot_request: The BootRequest to inspect. - - Returns: - True if any executable entry is 'drunc-controller', False otherwise. - """ - for e_and_a in boot_request.process_description.executable_and_arguments: - if e_and_a.exec == "drunc-controller": - return True - return False - def _build_container_env( self, boot_request: BootRequest, tree_labels: dict[str, str] ) -> list[client.V1EnvVar]: @@ -1210,7 +1221,7 @@ def _build_pod_main_container( main_container - the fully configured V1Container object """ - pod_image = self.configuration.data.image + pod_image = self.configuration.image exec_and_args_list = boot_request.process_description.executable_and_arguments # Build command to exec @@ -1255,7 +1266,7 @@ def _build_pod_main_container( resource_reqs = None is_perf_app = self.perf_selector in podname.lower() if is_perf_app: - settings = getattr(self.configuration.data, "settings", {}) + settings = getattr(self.configuration, "settings", {}) host_configs = settings.get("host_configs", {}) if not target_host or target_host not in host_configs: @@ -1318,7 +1329,6 @@ def _build_pod_main_container( self._is_local_connection_server(tree_labels, podname) and lcs_port is not None ): - self.connection_server_name = podname container_ports.append( client.V1ContainerPort(container_port=lcs_port, name="http-port") ) @@ -1394,32 +1404,36 @@ def _get_pod_host_aliases( host_aliases - a list containing a single V1HostAlias mapping localhost to the connection server IP, or None if not applicable """ - host_aliases = None + with self._lcs_state_lock: + lcs = self._lcs_state.get(session) if ( - not self._is_local_connection_server(tree_labels, podname) - and self.local_connection_server_is_booted + self._is_local_connection_server(tree_labels, podname) + or lcs is None + or not lcs.is_booted ): - connection_server_ip = None - retry_count = 0 - max_retries = 10 - while not connection_server_ip and retry_count < max_retries: - connection_server_ip = self._get_connection_server_cluster_ip(session) - if not connection_server_ip: - sleep(1) - retry_count += 1 - - if connection_server_ip: - host_aliases = [ - client.V1HostAlias(ip=connection_server_ip, hostnames=["localhost"]) - ] - self.log.info( - f"Pod '{podname}' will resolve localhost to connection server IP {connection_server_ip}" - ) - else: - self.log.warning( - f"Could not get connection server ClusterIP for pod '{podname}'" - ) - return host_aliases + return None + + connection_server_ip = None + retry_count = 0 + max_retries = 10 + while not connection_server_ip and retry_count < max_retries: + connection_server_ip = self._get_connection_server_cluster_ip(session) + if not connection_server_ip: + sleep(1) + retry_count += 1 + + if connection_server_ip: + self.log.info( + f"Pod '{podname}' will resolve localhost to connection server IP {connection_server_ip}" + ) + return [ + client.V1HostAlias(ip=connection_server_ip, hostnames=["localhost"]) + ] + + self.log.warning( + f"Could not get connection server ClusterIP for pod '{podname}'" + ) + return None def _determine_service_type( self, podname: str, boot_request: BootRequest, tree_labels: dict[str, str] @@ -1625,21 +1639,20 @@ def _create_associated_service( raise DruncK8sException( "LCS service creation failed: port was not extracted." ) - # This call uses class variables set in _create_pod - self._create_nodeport_service(podname, session, pod_uid) + self._create_nodeport_service( + podname, session, pod_uid, port=lcs_port, node_port=lcs_port + ) elif self._is_root_controller(tree_labels): self.log.info( f"'{podname}' is the root controller, checking for NodePort service." ) - # This call also relies on class variables, so we must set them - # here, just as the original logic did. port = self._extract_port_from_cmd(boot_request) if port: self.log.info(f"Extracted port {port} for '{podname}' NodePort.") - self.connection_server_port = port - self.connection_server_node_port = port - self._create_nodeport_service(podname, session, pod_uid) + self._create_nodeport_service( + podname, session, pod_uid, port=port, node_port=port + ) else: # This case should be caught by _determine_service_type, # but we handle it just in case. @@ -1672,12 +1685,13 @@ def _create_pod( """ try: lcs_port = None - # Early Port Extraction and Class Variable Setup for LCS + # Early Port Extraction and Per-Session LCS State Setup if self._is_local_connection_server(tree_labels, podname): lcs_port = self._extract_port_from_cmd(boot_request) if lcs_port: - self.connection_server_port = lcs_port - self.connection_server_node_port = lcs_port + lcs = self._lcs_state_for(session) + lcs.podname = podname + lcs.node_port = lcs_port else: raise DruncK8sException( f"Could not extract port for LCS '{podname}'." @@ -1760,8 +1774,8 @@ def _get_connection_server_cluster_ip(self, session: str) -> str: Get the ClusterIP of the connection server's Kubernetes Service. Reads the named service from the session namespace and returns its - clusterIP. Returns None if the service cannot be found or an API - error occurs. + clusterIP. Returns None if the session has no LCS, the service cannot + be found, or an API error occurs. Args: session - the Kubernetes namespace (session) containing the service @@ -1769,9 +1783,16 @@ def _get_connection_server_cluster_ip(self, session: str) -> str: Returns: cluster_ip - the ClusterIP string, or None on failure """ + with self._lcs_state_lock: + lcs = self._lcs_state.get(session) + if lcs is None or lcs.podname is None: + self.log.warning( + f"No local connection server registered for session '{session}'" + ) + return None try: service = self._core_v1_api.read_namespaced_service( - name=self.connection_server_name, namespace=session + name=lcs.podname, namespace=session ) return service.spec.cluster_ip except self._api_error_v1_api as e: @@ -2130,7 +2151,7 @@ def _wait_for_lcs_readiness(self, podname: str, session: str) -> None: Stage 1: waits for the pod to be Running and Ready in the Kubernetes API. Stage 2: waits for the NodePort to be externally reachable via HTTP. - Sets local_connection_server_is_booted to True on success. + Sets _lcs_state[session].is_booted to True on success. Args: podname - the name of the LCS pod to wait for @@ -2147,7 +2168,8 @@ def _wait_for_lcs_readiness(self, podname: str, session: str) -> None: node_name = self._wait_for_pod_api_ready(podname, session, total_timeout) # --- STAGE 2: Wait for NodePort to be externally reachable (using HTTP urllib) --- - url = f"http://{node_name}:{self.connection_server_node_port}" + lcs = self._lcs_state_for(session) + url = f"http://{node_name}:{lcs.node_port}" # Calculate remaining time for stage 2, preserving original logic elapsed_stage1 = time() - start_time @@ -2160,7 +2182,12 @@ def _wait_for_lcs_readiness(self, podname: str, session: str) -> None: self._wait_for_nodeport_http_ready(url, remaining_time) - self.local_connection_server_is_booted = True + with self._lcs_state_lock: + # Re-check the entry still exists: the watcher may have popped it + # if the pod died in the narrow window between stage 2 succeeding + # and this assignment. + if self._lcs_state.get(session) is lcs: + lcs.is_booted = True self.log.info(f"Connection server '{podname}' is fully ready.") def _wait_for_controller_readiness( @@ -2484,7 +2511,7 @@ def _kill_impl(self, query: ProcessQuery) -> ProcessInstanceList: return ProcessInstanceList(values=[]) self.log.info( - f"Starting staged termination for {len(targeted_uuids)} pod(s)..." + f"Starting staged termination for {len(targeted_uuids)} process(es)..." ) # Define the blocking kill_and_wait helper @@ -2496,7 +2523,7 @@ def kill_and_wait(uuids, grace_period=None) -> None: if grace_period == 0 else "Gracefully terminating" ) - self.log.info(f"{action} {len(uuids)} pod(s)...") + self.log.info(f"{action} {len(uuids)} process(es)...") self.termination_complete_event.clear() self.uuids_pending_deletion.update(uuids) @@ -2574,19 +2601,19 @@ def kill_and_wait(uuids, grace_period=None) -> None: for depth in sorted(by_depth.keys(), reverse=True): depth_uuids = by_depth[depth] self.log.info( - f"--- Termination Step: Shutting down role '{role}' at depth {depth} " + f"--- Termination of role '{role}' at depth {depth} " f"({len(depth_uuids)} pod(s)) ---" ) kill_and_wait(depth_uuids) # This call is blocking self.log.info( - f"--- Termination Step: Role '{role}' at depth {depth} complete ---" + f"--- Termination of role '{role}' at depth {depth} complete ---" ) else: self.log.info( - f"--- Termination Step: Shutting down role '{role}' ({len(uuids_in_step)} pod(s)) ---" + f"--- Termination of role '{role}' ({len(uuids_in_step)} pod(s)) ---" ) kill_and_wait(uuids_in_step) # This call is blocking - self.log.info(f"--- Termination Step: Role '{role}' complete ---") + self.log.info(f"--- Termination of role '{role}' complete ---") # Finalize and clean up final_ret = [] @@ -2616,6 +2643,8 @@ def kill_and_wait(uuids, grace_period=None) -> None: self.log.info(f'Session "{session}" is empty, deleting namespace.') self._core_v1_api.delete_namespace(session) self.managed_sessions.remove(session) + with self._lcs_state_lock: + self._lcs_state.pop(session, None) except self._api_error_v1_api as e: self.log.warning(f"Failed during namespace cleanup: {e}") @@ -2633,7 +2662,7 @@ def _terminate_impl(self) -> ProcessInstanceList: A ProcessInstanceList containing DEAD-status entries for all terminated processes, or an empty list if there were no processes to terminate. """ - self.log.info("Terminating all known K8s processes.") + self.log.info("Terminating") if not self.boot_request: self.log.info("No processes to terminate.") return ProcessInstanceList(values=[]) diff --git a/src/drunc/process_manager/process_manager.py b/src/drunc/process_manager/process_manager.py index 07afa8b86..524f6b800 100644 --- a/src/drunc/process_manager/process_manager.py +++ b/src/drunc/process_manager/process_manager.py @@ -6,7 +6,6 @@ from daqpytools.logging import LogHandlerConf, exceptions, setup_daq_ers_logger from druncschema.authoriser_pb2 import ActionType, SystemType -from druncschema.broadcast_pb2 import BroadcastType from druncschema.description_pb2 import CommandDescription, Description from druncschema.opmon.process_manager_pb2 import ProcessStatus from druncschema.process_manager_pb2 import ( @@ -28,9 +27,6 @@ from drunc.authoriser.configuration import DummyAuthoriserConfHandler from drunc.authoriser.decorators import authentified_and_authorised from drunc.authoriser.dummy_authoriser import DummyAuthoriser -from drunc.broadcast.server.broadcast_sender import BroadcastSender -from drunc.broadcast.server.configuration import BroadcastSenderConfHandler -from drunc.broadcast.server.decorators import broadcasted from drunc.exceptions import ( DruncCommandException, DruncNotImplementedException, @@ -39,7 +35,6 @@ ProcessManagerConfHandler, ProcessManagerTypes, ) -from drunc.utils.configuration import ConfTypes from drunc.utils.utils import get_logger, pid_info_str @@ -50,18 +45,14 @@ def __init__(self, txt): class ProcessManager(abc.ABC, ProcessManagerServicer): def __init__( - self, - configuration: ProcessManagerConfHandler, - name: str, - session: str = None, - **kwargs, + self, configuration: ProcessManagerConfHandler, name: str, session: str ): """C'tor. Note that this takes the ERS env variables from the json files defined in data/process_manager!""" super().__init__() self.log = get_logger( - f"process_manager.{configuration.get_data_type_name()}_process_manager", + f"process_manager.{configuration.pm_type.name}_process_manager", ) self.log.debug(pid_info_str()) self.log.debug("Initialized ProcessManager") @@ -81,16 +72,12 @@ def __init__( self.name = name self.session = session - self._create_broadcast_service(self.name, self.session) - - dach = DummyAuthoriserConfHandler( - data=self.configuration.get_data_authoriser(), type=ConfTypes.PyObject + dach = DummyAuthoriserConfHandler.from_pyobject( + data=self.configuration.authoriser ) - self.opmon_publisher = getattr( - self.configuration.get_data(), "opmon_publisher", None - ) - interval_s = getattr(self.configuration.get_data(), "interval_s", 10.0) + self.opmon_publisher = self.configuration.opmon_publisher + interval_s = self.configuration.opmon_conf["interval_s"] self.authoriser = DummyAuthoriser(dach, SystemType.PROCESS_MANAGER) self.process_store = {} # dict[str, sh.RunningCommand] # str = uuid @@ -154,8 +141,6 @@ def __init__( ), ] - self.broadcast(message="ready", btype=BroadcastType.SERVER_READY) - if self.opmon_publisher is not None: self.stop_event = threading.Event() self.thread = threading.Thread( @@ -168,21 +153,6 @@ def __init__( def get_log_path(self): return self.configuration.get_log_path() - def _create_broadcast_service(self, name, session): - bsch = BroadcastSenderConfHandler( - data=self.configuration.get_data_broadcaster(), type=ConfTypes.PyObject - ) - - self.broadcast_service = ( - BroadcastSender( - name=name, - session=session, - configuration=bsch, - ) - if bsch.data - else None - ) - def __del__(self): if hasattr(self, "opmon_publisher") and self.opmon_publisher is not None: self.stop_event.set() @@ -247,48 +217,10 @@ def find_by_uuid(pi_list, target_uuid: str): time.sleep(interval_s) - """ - A couple of simple pass-through functions to the broadcasting service - """ - - def broadcast(self, *args, **kwargs): - self.log.debug(f"{self.name} broadcasting") - return ( - self.broadcast_service.broadcast(*args, **kwargs) - if self.broadcast_service - else None - ) - - def can_broadcast(self, *args, **kwargs): - self.log.debug(f"Checking if {self.name} can broadcast") - return ( - self.broadcast_service.can_broadcast(*args, **kwargs) - if self.broadcast_service - else False - ) - - def describe_broadcast(self, *args, **kwargs): - self.log.debug(f"Describing {self.name} broadcast") - return ( - self.broadcast_service.describe_broadcast(*args, **kwargs) - if self.broadcast_service - else None - ) - - def interrupt_with_exception(self, *args, **kwargs): - self.log.debug(f"Interrupting {self.name} broadcast with exception") - return ( - self.broadcast_service._interrupt_with_exception(*args, **kwargs) - if self.broadcast_service - else None - ) - @abc.abstractmethod def _boot_impl(self, boot_request: BootRequest) -> ProcessInstanceList: raise NotImplementedError - # ORDER MATTERS! - @broadcasted # outer most wrapper 1st step @authentified_and_authorised( action=ActionType.CREATE, system=SystemType.PROCESS_MANAGER ) # 2nd step @@ -321,8 +253,6 @@ def boot( def _terminate_impl(self) -> ProcessInstanceList: raise NotImplementedError - # ORDER MATTERS! - @broadcasted # outer most wrapper 1st step @authentified_and_authorised( action=ActionType.DELETE, system=SystemType.PROCESS_MANAGER ) # 2nd step @@ -356,8 +286,6 @@ def terminate( def _restart_impl(self, query: ProcessQuery) -> ProcessInstanceList: raise NotImplementedError - # ORDER MATTERS! - @broadcasted # outer most wrapper 1st step @authentified_and_authorised( action=ActionType.DELETE, system=SystemType.PROCESS_MANAGER ) # 2nd step @@ -388,8 +316,6 @@ def restart( def _kill_impl(self, query: ProcessQuery) -> ProcessInstanceList: raise NotImplementedError - # ORDER MATTERS! - @broadcasted # outer most wrapper 1st step @authentified_and_authorised( action=ActionType.DELETE, system=SystemType.PROCESS_MANAGER ) # 2nd step @@ -420,8 +346,6 @@ def kill( def _ps_impl(self, query: ProcessQuery) -> ProcessInstanceList: raise NotImplementedError - # ORDER MATTERS! - @broadcasted # outer most wrapper 1st step @authentified_and_authorised( action=ActionType.READ, system=SystemType.PROCESS_MANAGER ) # 2nd step @@ -452,8 +376,6 @@ def ps( def _flush_impl(self, query: ProcessQuery) -> ProcessInstanceList: raise NotImplementedError - # ORDER MATTERS! - @broadcasted # outer most wrapper 1st step @authentified_and_authorised( action=ActionType.DELETE, system=SystemType.PROCESS_MANAGER ) # 2nd step @@ -493,8 +415,6 @@ def flush( return response - # ORDER MATTERS! - @broadcasted # outer most wrapper 1st step @authentified_and_authorised( action=ActionType.READ, system=SystemType.PROCESS_MANAGER ) # 2nd step @@ -511,17 +431,12 @@ def describe(self, request: Request, context: ServicerContext) -> Description: token=None, ) - if broadcast_description := self.describe_broadcast(): - response.broadcast.Pack(broadcast_description) - return response @abc.abstractmethod def _logs_impl(self, log_request: LogRequest) -> LogLines: raise NotImplementedError - # ORDER MATTERS! - @broadcasted # outer most wrapper 1st step @authentified_and_authorised( action=ActionType.READ, system=SystemType.PROCESS_MANAGER ) # 2nd step @@ -784,19 +699,19 @@ def _get_process_uid( def get(conf, **kwargs): log = get_logger("process_manager.get") - if conf.data.type == ProcessManagerTypes.SSH_SHELL: + if conf.pm_type == ProcessManagerTypes.SSH_SHELL: from drunc.process_manager.ssh_process_manager_shell import ( SSHProcessManagerShell, ) log.debug("Starting [green]SSH Shell process_manager[/green]") return SSHProcessManagerShell(conf, **kwargs) - elif conf.data.type == ProcessManagerTypes.K8s: + elif conf.pm_type == ProcessManagerTypes.K8s: from drunc.process_manager.k8s_process_manager import K8sProcessManager log.debug("Starting [green]K8s process_manager[/green]") return K8sProcessManager(conf, **kwargs) - elif conf.data.type == ProcessManagerTypes.SSH_PARAMIKO: + elif conf.pm_type == ProcessManagerTypes.SSH_PARAMIKO: from drunc.process_manager.ssh_process_manager_paramiko_client import ( SSHProcessManagerParamikoClient, ) @@ -804,7 +719,5 @@ def get(conf, **kwargs): log.debug("Starting [green]SSH Paramiko process_manager[/green]") return SSHProcessManagerParamikoClient(conf, **kwargs) else: - log.error(f"ProcessManager type {conf.get('type')} is unsupported!") - raise RuntimeError( - f"ProcessManager type {conf.get('type')} is unsupported!" - ) + log.error(f"ProcessManager type {conf.pm_type} is unsupported!") + raise RuntimeError(f"ProcessManager type {conf.pm_type} is unsupported!") diff --git a/src/drunc/process_manager/ssh_process_manager.py b/src/drunc/process_manager/ssh_process_manager.py index df7b16770..24b14356e 100644 --- a/src/drunc/process_manager/ssh_process_manager.py +++ b/src/drunc/process_manager/ssh_process_manager.py @@ -3,7 +3,6 @@ import uuid from typing import List, Optional -from druncschema.broadcast_pb2 import BroadcastType from druncschema.process_manager_pb2 import ( BootRequest, LogLines, @@ -37,13 +36,11 @@ def __init__( self.disable_localhost_host_key_check = False self.disable_host_key_check = False - if self.configuration.data.settings: - self.disable_localhost_host_key_check = ( - self.configuration.data.settings.get( - "disable_localhost_host_key_check", False - ) + if self.configuration.settings: + self.disable_localhost_host_key_check = self.configuration.settings.get( + "disable_localhost_host_key_check", False ) - self.disable_host_key_check = self.configuration.data.settings.get( + self.disable_host_key_check = self.configuration.settings.get( "disable_host_key_check", False ) @@ -102,7 +99,7 @@ def _build_process_instance( def _get_process_timeouts(self, uuids: List[str]) -> dict[str, float]: process_timeouts = {} for process_uuid in uuids: - process_timeouts[process_uuid] = self.configuration.data.kill_timeout + process_timeouts[process_uuid] = self.configuration.kill_timeout return process_timeouts def _on_ssh_process_exit( @@ -222,8 +219,6 @@ def _terminate_impl(self) -> ProcessInstanceList: self.log.info("Terminating") if self.boot_request: - self.log.info("Killing all the known processes before exiting") - # Build query to match all processes query = ProcessQuery(names=[".*"]) uuids = ProcessManager._match_processes_against_query( @@ -319,10 +314,8 @@ def _logs_impl(self, log_request: LogRequest) -> LogLines: ) def notify_join(self, name, session, user, exit_status: ExitStatus): - self.log.debug(f"{self.name} sending broadcast after ssh process exit") end_str = exit_status.get_process_manager_log_message(name, session, user) self.log.info(end_str) - self.broadcast(end_str, BroadcastType.SUBPROCESS_STATUS_UPDATE) def __boot(self, boot_request: BootRequest, uuid: str) -> ProcessInstance: """ @@ -530,7 +523,7 @@ def _restart_impl(self, query: ProcessQuery) -> ProcessInstanceList: self.add_process_to_expected_dead_processes(uuid) exit_status = self.ssh_lifetime_manager.kill_process( - uuid, self.configuration.data.kill_timeout + uuid, self.configuration.kill_timeout ) if exit_status is not None: self.archived_exit_statuses[uuid] = exit_status @@ -694,7 +687,7 @@ def _flush_impl(self, query: ProcessQuery) -> ProcessInstanceList: del self.boot_request[proc_uuid] # Clean data associated with the process from the lifetime manager self.ssh_lifetime_manager.kill_process( - proc_uuid, self.configuration.data.kill_timeout + proc_uuid, self.configuration.kill_timeout ) pi_return_code = ( diff --git a/src/drunc/process_manager/utils.py b/src/drunc/process_manager/utils.py index 859a74702..8793cc8d4 100644 --- a/src/drunc/process_manager/utils.py +++ b/src/drunc/process_manager/utils.py @@ -5,6 +5,7 @@ import click from druncschema.process_manager_pb2 import ( + BootRequest, ProcessInstance, ProcessInstanceList, ProcessQuery, @@ -18,10 +19,36 @@ ProcessManagerTypes, get_process_manager_configuration, ) -from drunc.utils.configuration import parse_conf_url +from drunc.processes.process_metadata import ProcessMetadata +from drunc.utils.configuration import ConfTypes, parse_conf_url from drunc.utils.utils import now_str +def compute_role_from_boot_request(boot_request: BootRequest) -> str: + """ + Determine the process role from a BootRequest. + + Extracts tree_id from the process metadata and checks + executable_and_arguments for the drunc-controller executable, + then delegates to ProcessMetadata.compute_role_from_tree_id. + + Args: + boot_request: The BootRequest describing the process to launch. + + Returns: + Role string: "root-controller", "segment-controller", "application", + "infrastructure-applications", or "unknown". + """ + tree_id = boot_request.process_description.metadata.tree_id + is_controller = any( + e.exec == "drunc-controller" + for e in boot_request.process_description.executable_and_arguments + ) + return ProcessMetadata.compute_role_from_tree_id( + tree_id, is_controller=is_controller + ) + + def generate_process_query( f, at_least_one: bool, all_processes_by_default: bool = False ): @@ -320,11 +347,15 @@ def get_pm_type_from_name(pm_name: str) -> ProcessManagerTypes: pm_conf_file = get_process_manager_configuration(pm_name) conf_path, conf_type = parse_conf_url(pm_conf_file) - pmch = ProcessManagerConfHandler( - log_path="./", type=conf_type, data=conf_path.split(":")[1] - ) + path_or_url = conf_path.split(":")[1] + + if conf_type == ConfTypes.JsonFileName: + pmch = ProcessManagerConfHandler.from_json(path=path_or_url) + else: + # OKS or other types - fallback to from_pyobject + pmch = ProcessManagerConfHandler.from_pyobject(data=path_or_url) - return pmch.data.type + return getattr(pmch, "pm_type", pmch.type) def format_hostname(hostname: str) -> str: diff --git a/src/drunc/processes/process_metadata.py b/src/drunc/processes/process_metadata.py index 8ee034b61..5a7f0ca6b 100644 --- a/src/drunc/processes/process_metadata.py +++ b/src/drunc/processes/process_metadata.py @@ -80,27 +80,36 @@ def from_json(cls, json_str: str) -> "ProcessMetadata": return cls.from_dict(json.loads(json_str)) @staticmethod - def compute_role_from_tree_id(tree_id: str) -> str: + def compute_role_from_tree_id(tree_id: str, is_controller: bool = False) -> str: """ - Determines the role of a process based on its tree_id. + Determines the role of a process based on its tree_id and executable type. + + - empty tree_id -> "unknown" + - is_controller + tree_id == "0" -> "root-controller" + - is_controller + tree_id starts with "0." -> "segment-controller" + - is_controller + otherwise -> "infrastructure-applications" + - not controller + tree_id starts with "0." -> "application" + - not controller + otherwise -> "infrastructure-applications" Args: - tree_id: Hierarchical identifier in format session.segment.application + tree_id: Dot-separated hierarchical identifier (e.g. "0", "0.1", "0.1.2.3"). + is_controller: True if the process executable is a drunc-controller. Returns: - Role string: "root-controller", "local-connection-server", - "segment-controller", "application", or "unknown" + Role string: "root-controller", "segment-controller", "application", + "infrastructure-applications", or "unknown" """ if not tree_id: return "unknown" - elif tree_id == "0": - return "root-controller" - elif tree_id == "1": - return "local-connection-server" - else: - depth = tree_id.count(".") - if depth == 1: + elif is_controller: + if tree_id == "0": + return "root-controller" + elif tree_id.startswith("0."): return "segment-controller" - elif depth == 2: + else: + return "infrastructure-applications" # controller outside segment tree + else: + if tree_id.startswith("0."): return "application" - return "unknown" + else: + return "infrastructure-applications" diff --git a/src/drunc/processes/ssh_process_lifetime_manager_from_forked_process.py b/src/drunc/processes/ssh_process_lifetime_manager_from_forked_process.py index 1a79636ab..6541ec57b 100644 --- a/src/drunc/processes/ssh_process_lifetime_manager_from_forked_process.py +++ b/src/drunc/processes/ssh_process_lifetime_manager_from_forked_process.py @@ -229,7 +229,11 @@ def __init__( The exception is reconstructed as a RuntimeError from the serialised message forwarded by the child process. """ - self.log = logger if logger is not None else get_logger(__name__) + self.log = ( + logger + if logger is not None + else get_logger("ssh_process_lifetime_manager_forked", rich_handler=True) + ) self._on_process_exit = on_process_exit # Queues for IPC between parent and child. diff --git a/src/drunc/processes/ssh_process_lifetime_manager_paramiko.py b/src/drunc/processes/ssh_process_lifetime_manager_paramiko.py index e918ee6e8..eaacbe27b 100644 --- a/src/drunc/processes/ssh_process_lifetime_manager_paramiko.py +++ b/src/drunc/processes/ssh_process_lifetime_manager_paramiko.py @@ -46,7 +46,11 @@ def __init__( """ self.disable_host_key_check = disable_host_key_check self.disable_localhost_host_key_check = disable_localhost_host_key_check - self.log = logger if logger else get_logger(__name__) + self.log = ( + logger + if logger + else get_logger("ssh_process_lifetime_manager_paramiko", rich_handler=True) + ) self.log.warning( "The paramiko-based SSH process manager is NOT actively maintatined. Consider using the shell-based SSH process manager instead." ) diff --git a/src/drunc/processes/ssh_process_lifetime_manager_shell.py b/src/drunc/processes/ssh_process_lifetime_manager_shell.py index bbdb661e3..7d25a3ff1 100644 --- a/src/drunc/processes/ssh_process_lifetime_manager_shell.py +++ b/src/drunc/processes/ssh_process_lifetime_manager_shell.py @@ -318,7 +318,11 @@ def __init__( """ self.disable_host_key_check = disable_host_key_check self.disable_localhost_host_key_check = disable_localhost_host_key_check - self.log = logger if logger else get_logger(__name__) + self.log = ( + logger + if logger + else get_logger("ssh_process_lifetime_manager_shell", rich_handler=True) + ) self._on_process_exit = on_process_exit # Create SSH command wrapper @@ -687,34 +691,29 @@ def kill_processes_by_role( uuids_to_kill = [] with self.lock: for uuid in candidate_uuids: - metadata = self.metadata.get(uuid, None) - if metadata and metadata.role == role: - uuids_to_kill.append(uuid) - if uuid not in process_timeouts: - process_timeouts[uuid] = ( - self.DEFAULT_TIMEOUT_FOR_KILLING_PROCESS - ) + metadata = self.metadata.get(uuid) + if metadata is None or metadata.role != role: + continue + + uuids_to_kill.append(uuid) + process_timeouts.setdefault( + uuid, + self.DEFAULT_TIMEOUT_FOR_KILLING_PROCESS, + ) if not uuids_to_kill: - self.log.debug(f"No processes found with role '{role}' in candidate list") return {} - self.log.info( - f"Killing {len(uuids_to_kill)} process(es) with role '{role}' " - f"from {len(candidate_uuids)} candidates" - ) + self.log.info(f"Killing {len(uuids_to_kill)} process(es) with role '{role}'") exit_statuses: Dict[str, Optional[ExitStatus]] = {} - # Terminate processes asynchronously using thread pool with ThreadPoolExecutor(max_workers=len(uuids_to_kill)) as executor: - # Submit kill tasks for all matching processes future_to_uuid = { executor.submit(self.kill_process, uuid, process_timeouts[uuid]): uuid for uuid in uuids_to_kill } - # Collect results as they complete for future in as_completed(future_to_uuid): uuid = future_to_uuid[future] try: @@ -758,12 +757,24 @@ def kill_processes( process_timeouts[uuid] = self.DEFAULT_TIMEOUT_FOR_KILLING_PROCESS all_exit_statuses: Dict[str, Optional[ExitStatus]] = {} - killed_uuids: set[str] = set() + killed_uuids = set() # Execute role-based shutdown in stages for role in PROCESS_SHUTDOWN_ORDERING: + with self.lock: + uuids_in_role = [ + uuid + for uuid in uuids + if (metadata := self.metadata.get(uuid)) is not None + and metadata.role == role + ] + + # Match k8s PM behavior: if role is absent, do not log/start/end a stage. + if not uuids_in_role: + continue + self.log.info( - f"--- Shutdown stage: Terminating role '{role}' from provided UUIDs ---" + f"--- Termination of role '{role}' ({len(uuids_in_role)} process(es)) ---" ) role_exit_statuses = self.kill_processes_by_role( role, uuids, process_timeouts=process_timeouts @@ -1002,27 +1013,29 @@ def read_process_metadata( ProcessMetadata instance if file exists and is valid, None otherwise """ try: - # Build user@host string for SSH connection user_host = f"{user}@{hostname}" - # Build SSH arguments including connection parameters - arguments = self._build_ssh_arguments(hostname, user_host) + # Metadata read is non-interactive and machine-readable. + arguments = self._build_ssh_arguments( + hostname, + user_host, + use_tty=False, + ) - # Remote command: wait for file to exist, then read it - # Polls every 50ms, times out after specified duration remote_command = ( - f"timeout {timeout} bash -c '" - f"while [ ! -f {metadata_file} ]; do sleep 0.05; done; " - f"cat {metadata_file}" + f"timeout {timeout} sh -c '" + f'metadata_file="{metadata_file}"; ' + f'while [ ! -s "$metadata_file" ]; do sleep 0.05; done; ' + f'cat "$metadata_file"' f"'" ) arguments.append(remote_command) - # Execute SSH command to wait for and read file (single round-trip) result = self.ssh(*arguments) json_content = str(result).strip() - # Parse JSON content and instantiate metadata object + self.log.debug(f"Metadata content for {uuid}: {json_content!r}") + metadata = ProcessMetadata.from_json(json_content) with self.lock: @@ -1033,7 +1046,7 @@ def read_process_metadata( return metadata except Exception as e: - self.log.debug(f"Failed to read metadata for {uuid}: {e}") + self.log.warning(f"Failed to read metadata for {uuid}: {e}") return None def _handle_external_client_sigquit( @@ -1127,7 +1140,13 @@ def _execute_bootrequest_via_ssh( metadata_file = SSHProcessLifetimeManagerShell.get_metadata_file_path(uuid) tree_id = boot_request.process_description.metadata.tree_id name = boot_request.process_description.metadata.name - role = ProcessMetadata.compute_role_from_tree_id(tree_id) + is_controller = any( + e_and_a.exec == "drunc-controller" + for e_and_a in boot_request.process_description.executable_and_arguments + ) + role = ProcessMetadata.compute_role_from_tree_id( + tree_id, is_controller=is_controller + ) remote_metadata_json = ( f"{{\"pid\": '$PID', " diff --git a/src/drunc/session_manager/configuration.py b/src/drunc/session_manager/configuration.py index 2211eb6ac..096ea553f 100644 --- a/src/drunc/session_manager/configuration.py +++ b/src/drunc/session_manager/configuration.py @@ -4,6 +4,7 @@ class SessionManagerConfHandler(ConfHandler): - """TODO: Change this exception to something more useful.""" + """Handler for session manager configuration.""" - pass + def populate_from_dict(self, data: dict[str, object]) -> None: + pass diff --git a/src/drunc/session_manager/interface/session_manager.py b/src/drunc/session_manager/interface/session_manager.py index afb6e68b4..0e7a79849 100644 --- a/src/drunc/session_manager/interface/session_manager.py +++ b/src/drunc/session_manager/interface/session_manager.py @@ -38,7 +38,7 @@ def serve(session_manager: SessionManager, address: str) -> None: @click.command() -def session_manager_cli()-> None: +def session_manager_cli() -> None: """CLI interface for the Drunc session manager. This command starts the session manager service, which allows clients to manage @@ -51,7 +51,7 @@ def session_manager_cli()-> None: logger = get_logger(app_name, rich_handler=True) # Load the configuration for the session manager. - config = SessionManagerConfHandler() + config = SessionManagerConfHandler.from_pyobject(data=None) logger.info(f"Using '{config}' as the SessionManager configuration.") # Load the session manager. diff --git a/src/drunc/unified_shell/context.py b/src/drunc/unified_shell/context.py index f5c9c759e..87963e605 100644 --- a/src/drunc/unified_shell/context.py +++ b/src/drunc/unified_shell/context.py @@ -1,8 +1,11 @@ from collections.abc import Mapping from enum import Enum +import grpc +from druncschema.process_manager_pb2 import ProcessQuery from druncschema.token_pb2 import Token +from drunc.utils.grpc_utils import ServerTimeout, ServerUnreachable from drunc.utils.shell_utils import ShellContext @@ -66,36 +69,52 @@ def set_controller_driver(self, address_controller, **kwargs) -> None: self._token, ) - # This will raise an exception if the driver already exists - # self.set_driver("controller", driver) - def create_token(self, **kwargs) -> Token: from drunc.utils.shell_utils import create_dummy_token_from_uname token = create_dummy_token_from_uname() return token - def start_listening_pm(self, broadcaster_conf) -> None: - from drunc.broadcast.client.broadcast_handler import BroadcastHandler - from drunc.broadcast.client.configuration import BroadcastClientConfHandler - from drunc.utils.configuration import ConfTypes - - bcch = BroadcastClientConfHandler( - type=ConfTypes.ProtobufAny, - data=broadcaster_conf, - ) - self.status_receiver_pm = BroadcastHandler(broadcast_configuration=bcch) - - def start_listening_controller(self, broadcaster_conf) -> None: - from drunc.broadcast.client.broadcast_handler import BroadcastHandler - from drunc.broadcast.client.configuration import BroadcastClientConfHandler - from drunc.utils.configuration import ConfTypes - - bcch = BroadcastClientConfHandler( - type=ConfTypes.ProtobufAny, - data=broadcaster_conf, - ) - self.status_receiver_controller = BroadcastHandler(broadcast_configuration=bcch) + def get_endpoint_display_host_overrides(self) -> dict[str, str]: + """ + Return a mapping of process name -> preferred display hostname for endpoint + rendering in the UI. + + These values are cosmetic only. The controller's advertised endpoint remains + the authoritative connect address. + + Returns: + dict[str, str]: Mapping from process name to preferred display hostname. + """ + # The PM driver may not be registered if the user connected directly to a + # controller without going through the process manager (e.g. standalone boot). + # In that case hostname overrides are unavailable and we fall back to + # get_hostname_smart in the endpoint rendering path. + pm_driver = self.get_driver("process_manager", quiet_fail=True) + if not pm_driver: + return {} + + if not self.session_name: + raise RuntimeError("session name must be set before querying process list") + query = ProcessQuery(names=[".*"], session=self.session_name) + try: + proc_list = pm_driver.ps(query) + except (ServerUnreachable, ServerTimeout, grpc.RpcError): + return {} + + overrides: dict[str, str] = {} + + for proc in proc_list.values: + metadata = proc.process_description.metadata + proc_name = getattr(metadata, "name", "") + host_name = getattr(metadata, "hostname", "") + + if not proc_name or not host_name: + continue + + overrides[proc_name] = host_name + + return overrides def terminate(self) -> None: if self.status_receiver_pm: diff --git a/src/drunc/unified_shell/shell.py b/src/drunc/unified_shell/shell.py index af8bedbec..eb8d67b5d 100644 --- a/src/drunc/unified_shell/shell.py +++ b/src/drunc/unified_shell/shell.py @@ -12,7 +12,6 @@ import click_shell import conffwk from daqpytools.logging import logging_log_levels -from druncschema.description_pb2 import Description from druncschema.process_manager_pb2 import ProcessQuery from drunc.connectivity_service.client import ConnectivityServiceClient @@ -62,7 +61,7 @@ from drunc.unified_shell.commands import boot, start_shell from drunc.unified_shell.context import UnifiedShellMode from drunc.unified_shell.shell_utils import generate_fsm_sequence_command -from drunc.utils.configuration import ConfTypes, OKSKey +from drunc.utils.configuration import OKSKey from drunc.utils.grpc_utils import ServerUnreachable from drunc.utils.utils import ( format_name_for_cli, @@ -277,9 +276,8 @@ def unified_shell( ctx.obj.reset(address_pm=process_manager_address) # Run a simple command (describe) to check the connection with the process manager - desc: Description | None = None try: - desc = ctx.obj.get_driver().describe() + ctx.obj.get_driver().describe() except Exception as e: ctx.obj.log.error( f"[red]Could not connect to the process manager at the address: [/red]" @@ -305,13 +303,6 @@ def unified_shell( sys.exit(1) - # Broadcasting configuration if requested - if desc.HasField("broadcast"): - ctx.obj.log.debug("Broadcasting") - ctx.obj.start_listening_pm( - broadcaster_conf=desc.broadcast, - ) - # Add the unified shell Click commands to the CLI ctx.obj.log.debug("Adding [green]unified_shell[/green] commands") ctx.command.add_command(boot, "boot") @@ -335,9 +326,8 @@ def unified_shell( # configuration and getting the FSM transitions from it. ctx.obj.log.debug("Defining the pseudo controller to get its FSM commands") controller_name = session_dal.segment.controller.id - controller_configuration = ControllerConfHandler( - type=ConfTypes.OKSFileName, - data=ctx.obj.configuration_file, + controller_configuration = ControllerConfHandler.from_oks( + url=ctx.obj.configuration_file, oks_key=OKSKey( schema_file="schema/confmodel/dunedaq.schema.xml", class_name="RCApplication", @@ -356,7 +346,7 @@ def unified_shell( # live with it. At least until controller.core uses file handler instead of stream get_logger("controller.core.FSM", log_level="CRITICAL") - fsmch = FSMConfHandler(data=controller_configuration.data.controller.fsm) + fsmch = FSMConfHandler.from_pyobject(data=controller_configuration.controller.fsm) ctx.obj.log.debug("Initializing the [green]StatefulNode[/green]") stateful_node = StatefulNode(fsm_configuration=fsmch, top_segment_controller=False) @@ -521,9 +511,9 @@ def cleanup(): ctx.obj.log.debug("Process manager terminated") ctx.obj.log.info("[green]unified_shell exited successfully[/green]") - logging.shutdown() # Shutdown logging - ctx.obj.terminate() # Terminate the broadcasters in the context - ctx.exit() # Close the click context + logging.shutdown() + ctx.obj.terminate() + ctx.exit() ctx.call_on_close(cleanup) From a4940138bfaf9d2ff326e76fe5ba2afde89f22b7 Mon Sep 17 00:00:00 2001 From: James Paul Turner Date: Tue, 14 Jul 2026 10:10:48 +0100 Subject: [PATCH 16/20] Phased merge 10 --- pyproject.toml | 15 ++++----------- 1 file changed, 4 insertions(+), 11 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 5b831d130..be11c1c90 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -34,7 +34,7 @@ dependencies = [ [project.optional-dependencies] prod = ["paramiko[gssapi]"] -dev = ["ruff", "pre-commit", "pytest", "pytest-cov", "grpcio-testing", "grpcio==1.75", "grpcio-tools==1.75", "grpcio-status==1.75", "pybind11-stubgen"] +dev = ["ruff", "pre-commit", "pytest", "pytest-cov", "grpcio-testing", "grpcio==1.75", "grpcio-tools==1.75", "grpcio-status==1.75"] test = ["pytest", "pytest-cov", "grpcio-testing", "grpcio==1.75", "grpcio-tools==1.75", "grpcio-status==1.75"] [project.scripts] @@ -107,18 +107,11 @@ disallow_any_unimported = true warn_unused_configs = true show_error_codes = true -# Tell mypy where to find stubs -mypy_path = "src:typings" - # These overrides are because the library stubs dont exist and not on typeshed [[tool.mypy.overrides]] -module = "google.rpc.*" +module = ["google.rpc.*"] ignore_missing_imports = true [[tool.mypy.overrides]] -module = "conffwk.*" -disallow_untyped_defs = false -disallow_incomplete_defs = false -check_untyped_defs = false -ignore_missing_imports = true -ignore_errors = true \ No newline at end of file +module = ["conffwk"] +ignore_missing_imports = true \ No newline at end of file From 1b7f9faba96f19ea9cd34d6540d435acbe800e93 Mon Sep 17 00:00:00 2001 From: PawelPlesniak Date: Thu, 16 Jul 2026 17:31:28 +0200 Subject: [PATCH 17/20] Merge complete --- src/drunc/fsm/configuration.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/src/drunc/fsm/configuration.py b/src/drunc/fsm/configuration.py index 946b49e3a..2a9e2703c 100644 --- a/src/drunc/fsm/configuration.py +++ b/src/drunc/fsm/configuration.py @@ -126,8 +126,8 @@ def _post_process_oks(self) -> None: tr.arguments += post_transitions.get_arguments() # Store the pre and post transition sequences for the transition - self._pre_transitions[tr] = pre_transitions - self._post_transitions[tr] = post_transitions + self.pre_transitions[tr] = pre_transitions + self.post_transitions[tr] = post_transitions # Add the transition to the list of transitions self.transitions += [tr] @@ -142,10 +142,10 @@ def get_actions(self) -> dict[str, FSMActionProtocol]: return self.actions def get_initial_state(self) -> str: - return self.data.initial_state + return self.initial_state def get_states(self) -> list[str]: - return self.data.states + return self.states def get_transitions(self) -> list[Transition]: return self.transitions From 52ef4bc99d564c5023b4a54550b8514448d0e0e1 Mon Sep 17 00:00:00 2001 From: PawelPlesniak Date: Tue, 21 Jul 2026 15:56:07 +0200 Subject: [PATCH 18/20] Corrected with conffwk and daqconf to have compile-time static classes instantiated from the base class only --- src/drunc/fsm/configuration.py | 14 +++++++------- src/drunc/fsm/core.py | 17 ++++++++++------- src/drunc/utils/grpc_utils.py | 7 +++++-- 3 files changed, 22 insertions(+), 16 deletions(-) diff --git a/src/drunc/fsm/configuration.py b/src/drunc/fsm/configuration.py index 2a9e2703c..e031033ec 100644 --- a/src/drunc/fsm/configuration.py +++ b/src/drunc/fsm/configuration.py @@ -1,8 +1,8 @@ from __future__ import annotations -from typing import List +from typing import List, cast -import conffwk +from conffwk.dal import FSMData, FSMxTransition from druncschema.controller_pb2 import FSMSequence from drunc.fsm._protocols import FSMActionProtocol @@ -16,13 +16,13 @@ class FSMConfHandler(ConfHandler): """Handler for FSM configuration.""" - data: conffwk.dal.FSMData + data: FSMData def _fill_pre_post_transition_sequence_oks( self, prefix: str, transition: Transition, - data: List[conffwk.dal.FSMxTransition] | None, + data: List[FSMxTransition] | None, ) -> PreOrPostTransitionSequence: """ Fill the pre or post transition sequence for a given transition. @@ -79,7 +79,7 @@ def _post_process_oks(self) -> None: Raises: None """ - raw = self._raw_data + raw = cast(FSMData, self._raw_data) # Define the data structures to store the FSM configuration self.log.debug("_post_process_oks configuration") @@ -95,12 +95,12 @@ def _post_process_oks(self) -> None: # Fill the actions dictionary with the FSMAction objects corresponding to the # action names defined in the configuration - for action in raw.actions: # type: 'conffwk.dal.FSMAction' + for action in raw.actions: # type: 'FSMAction' self.actions[action.id] = FSMActionFactory.get().get_action( action.id, action ) - for transition in raw.transitions: # type: 'conffwk.dal.FSMTransition' + for transition in raw.transitions: # type: 'FSMTransition' tr = Transition( name=transition.id, source=transition.source, diff --git a/src/drunc/fsm/core.py b/src/drunc/fsm/core.py index e9611ce7a..8d5b0e2e0 100644 --- a/src/drunc/fsm/core.py +++ b/src/drunc/fsm/core.py @@ -1,6 +1,6 @@ from __future__ import annotations -from typing import TYPE_CHECKING, Dict, Optional, Union +from typing import TYPE_CHECKING, Dict, Optional, Union, cast if TYPE_CHECKING: from drunc.fsm._protocols import ( @@ -111,20 +111,23 @@ def execute( transition_data = "{}" try: - input_data = json.loads(transition_data) + input_data = cast(dict[str, object], json.loads(transition_data)) except: raise fsme.TransitionDataOfIncorrectFormat(transition_data) for callback in self.sequence: try: - self.log.debug(f"data before callback: {input_data}") - self.log.debug( + self.log.critical(f"data before callback: {input_data}") + self.log.critical( f"executing the callback: {callback.method.__name__} from {callback.method.__module__}" ) - input_data = callback.method( - _input_data=input_data, _context=ctx, **transition_args + input_data = cast( + dict[str, object], + callback.method( + _input_data=input_data, _context=ctx, **transition_args + ), ) - self.log.debug(f"data after callback: {input_data}") + self.log.critical(f"data after callback: {input_data}") if input_data: ctx.runinfo.update(input_data) diff --git a/src/drunc/utils/grpc_utils.py b/src/drunc/utils/grpc_utils.py index a228de8a7..6741ebb67 100644 --- a/src/drunc/utils/grpc_utils.py +++ b/src/drunc/utils/grpc_utils.py @@ -3,7 +3,7 @@ from __future__ import annotations from dataclasses import dataclass -from typing import Callable, NoReturn, cast +from typing import Callable, NoReturn, TypeVar, cast import grpc from druncschema.generic_pb2 import PlainText @@ -76,7 +76,10 @@ def pack_to_any(data: Message) -> any_pb2.Any: return any -def unpack_any(data: any_pb2.Any, format: type[Message]) -> Message: +T = TypeVar("T") + + +def unpack_any(data: any_pb2.Any, format: type[T]) -> T: """Unpack an Any message into a specific protobuf format. Args: From a6bb890c56e7da0195650c06446116de47d049a6 Mon Sep 17 00:00:00 2001 From: PawelPlesniak Date: Tue, 21 Jul 2026 17:29:18 +0200 Subject: [PATCH 19/20] Removing un-necessary iterator type specifiers --- src/drunc/fsm/configuration.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/drunc/fsm/configuration.py b/src/drunc/fsm/configuration.py index e031033ec..6ef108be5 100644 --- a/src/drunc/fsm/configuration.py +++ b/src/drunc/fsm/configuration.py @@ -95,12 +95,12 @@ def _post_process_oks(self) -> None: # Fill the actions dictionary with the FSMAction objects corresponding to the # action names defined in the configuration - for action in raw.actions: # type: 'FSMAction' + for action in raw.actions: self.actions[action.id] = FSMActionFactory.get().get_action( action.id, action ) - for transition in raw.transitions: # type: 'FSMTransition' + for transition in raw.transitions: tr = Transition( name=transition.id, source=transition.source, From 22a200c868736b39fd3abcd86435f68ab68002ab Mon Sep 17 00:00:00 2001 From: PawelPlesniak Date: Tue, 28 Jul 2026 17:48:18 +0200 Subject: [PATCH 20/20] Adding requests types --- pyproject.toml | 110 ++++++++++++++++++++++++------------------------- 1 file changed, 54 insertions(+), 56 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index be11c1c90..97f26b32f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,57 +1,56 @@ [build-system] -requires = ["setuptools >= 61.0"] build-backend = "setuptools.build_meta" +requires = ["setuptools >= 61.0"] [project] -name = "drunc" description = "A flexible run control infrastructure for a distributed DAQ system" -version = "1.1.4" +name = "drunc" readme = "docs/README.md" requires-python = ">=3.10" +version = "1.1.4" dependencies = [ - "click", - "click-shell", - "protobuf", - "types-protobuf", - "grpcio", - "grpcio-status", - "grpcio-tools", - "types-grpcio", - "gunicorn", - "kafka-python", - "nest-asyncio", - "rich", - "requests", - "Flask", - "Flask-RESTful", - "sh", - "kubernetes", - "pytz", - "psutil", - "paramiko" + "click", + "click-shell", + "protobuf", + "types-protobuf", + "grpcio", + "grpcio-status", + "grpcio-tools", + "types-grpcio", + "gunicorn", + "kafka-python", + "nest-asyncio", + "rich", + "requests", + "Flask", + "Flask-RESTful", + "sh", + "kubernetes", + "pytz", + "psutil", + "paramiko", ] [project.optional-dependencies] +dev = ["ruff", "pre-commit", "pytest", "pytest-cov", "grpcio-testing", "grpcio==1.75", "grpcio-tools==1.75", "grpcio-status==1.75", "types-requests"] prod = ["paramiko[gssapi]"] -dev = ["ruff", "pre-commit", "pytest", "pytest-cov", "grpcio-testing", "grpcio==1.75", "grpcio-tools==1.75", "grpcio-status==1.75"] test = ["pytest", "pytest-cov", "grpcio-testing", "grpcio==1.75", "grpcio-tools==1.75", "grpcio-status==1.75"] [project.scripts] -fake_daq_application = "drunc.apps.fake_daqapp_rest:main" +application-registry-service = "drunc.apps.app_connectivity_server:main" +drunc-check-np0x-cluster = "drunc.apps.check_np0x_cluster:main" +drunc-check-np0x-hw = "drunc.apps.check_np0x_hw_status:main" drunc-controller = "drunc.apps.controller:main" drunc-controller-shell = "drunc.apps.controller_shell:main" +drunc-fsm-tests = "drunc.tests.fsm:main" drunc-process-manager = "drunc.apps.pm:main" drunc-process-manager-shell = "drunc.apps.pm_shell:main" drunc-session-manager = "drunc.apps.session_manager:main" -drunc-unified-shell = "drunc.apps.unified_shell:main" -drunc-fsm-tests = "drunc.tests.fsm:main" -application-registry-service = "drunc.apps.app_connectivity_server:main" -drunc-ssh-doctor = "drunc.apps.ssh_doctor:main" drunc-setup-ssh-config = "drunc.apps.ssh_configurator:main" -drunc-check-np0x-hw = "drunc.apps.check_np0x_hw_status:main" -drunc-check-np0x-cluster = "drunc.apps.check_np0x_cluster:main" - +drunc-ssh-doctor = "drunc.apps.ssh_doctor:main" +drunc-unified-shell = "drunc.apps.unified_shell:main" +fake_daq_application = "drunc.apps.fake_daqapp_rest:main" [tool.setuptools.packages.find] where = ["src"] @@ -61,57 +60,56 @@ where = ["src"] "drunc.data.process_manager.schema" = ["*.json"] [tool.pytest.ini_options] +addopts = "-v --tb=short --cov=drunc --cov=src/drunc" markers = [ - "grpc: marks tests for gRPC isolation (run with --test-grpc)", - "paramiko: marks tests for paramiko isolation (run with --test-paramiko)", + "grpc: marks tests for gRPC isolation (run with --test-grpc)", + "paramiko: marks tests for paramiko isolation (run with --test-paramiko)", ] -addopts = "-v --tb=short --cov=drunc --cov=src/drunc" testpaths = ["tests"] [tool.coverage.run] -source = ["drunc"] omit = ["tests/*"] +source = ["drunc"] # * See https://docs.astral.sh/ruff/rules/ for details on Ruff's linting options [tool.ruff.lint] -select = [ - "E", # pycodestyle errors - "F", # check for errors using PyFlakes - "I", # best practices for import calls - "UP", # suggestions for code modernization - "RUF", # build in Ruff warnings - "R", # refactoring suggestions -] ignore = [ - "E501", # Don't enforce line lengths within a linting context + "E501", # Don't enforce line lengths within a linting context +] +select = [ + "E", # pycodestyle errors + "F", # check for errors using PyFlakes + "I", # best practices for import calls + "UP", # suggestions for code modernization + "RUF", # build in Ruff warnings + "R", # refactoring suggestions ] - [tool.mypy] -disallow_untyped_defs = true -disallow_incomplete_defs = true check_untyped_defs = true +disallow_any_generics = true +disallow_incomplete_defs = true +disallow_subclassing_any = true +disallow_untyped_defs = true no_implicit_optional = true +strict_equality = true +warn_redundant_casts = true warn_return_any = true warn_unused_ignores = true -warn_redundant_casts = true -strict_equality = true -disallow_any_generics = true -disallow_subclassing_any = true +disallow_any_decorated = true disallow_any_explicit = true disallow_any_expr = false -disallow_any_decorated = true disallow_any_unimported = true -warn_unused_configs = true show_error_codes = true +warn_unused_configs = true # These overrides are because the library stubs dont exist and not on typeshed [[tool.mypy.overrides]] -module = ["google.rpc.*"] ignore_missing_imports = true +module = ["google.rpc.*"] [[tool.mypy.overrides]] +ignore_missing_imports = true module = ["conffwk"] -ignore_missing_imports = true \ No newline at end of file