diff --git a/src/drunc/process_manager/configuration.py b/src/drunc/process_manager/configuration.py index da9198722..3b6a5a429 100644 --- a/src/drunc/process_manager/configuration.py +++ b/src/drunc/process_manager/configuration.py @@ -3,14 +3,16 @@ import sys from enum import Enum from importlib import resources -from typing import TYPE_CHECKING, Any, Dict, Union +from typing import Protocol, cast from urllib.parse import unquote, urlparse -from jsonschema import ValidationError +from jsonschema import ValidationError # type: ignore[import-untyped] from jsonschema import validate as js_validate -from kafkaopmon.OpMonPublisher import OpMonPublisher as KafkaOpMonPublisher -from opmonlib.publisher import OpMonPublisher -from opmonlib.utils import parse_opmon_conf +from kafkaopmon.OpMonPublisher import ( # type: ignore[import-untyped] + OpMonPublisher as KafkaOpMonPublisher, +) +from opmonlib.publisher import OpMonPublisher # type: ignore[import-untyped] +from opmonlib.utils import parse_opmon_conf # type: ignore[import-untyped] from drunc.broadcast.server.configuration import KafkaBroadcastSenderConfData from drunc.exceptions import DruncCommandException @@ -18,8 +20,39 @@ from drunc.utils.configuration import ConfHandler from drunc.utils.utils import get_logger -if TYPE_CHECKING: - import conffwk + +class _RunsOnInner(Protocol): + id: str + + +class _RunsOn(Protocol): + runs_on: _RunsOnInner + + +class _Service(Protocol): + id: str + port: int + protocol: str + + +class _SessionDal(Protocol): + id: str + controller_log_level: str + + +class _AppLike(Protocol): + id: str + runs_on: _RunsOn + exposes_service: list[_Service] + + def oksTypes(self) -> list[str]: ... + + +class _OpMonConf(Protocol): + path: str + session: str + application: str + opmon_type: str PROCESS_SHUTDOWN_ORDERING = [ @@ -39,55 +72,77 @@ class ProcessManagerTypes(Enum): class ProcessManagerConfData: - def __init__(self): - self.broadcaster = None - self.authoriser = None + def __init__(self) -> None: + self.broadcaster: KafkaBroadcastSenderConfData | None = None + self.authoriser: object | None = None self.type = ProcessManagerTypes.Unknown self.command_address = "" - self.environment = {} - self.settings = {} - self.opmon_uri = None - self.opmon_publisher = None + self.environment: dict[str, str] = {} + self.settings: dict[str, object] = {} + self.opmon_uri: str | None = None + self.opmon_publisher: object | None = None + self.kill_timeout: float = 0.5 + self.image: str = "" class ProcessManagerConfHandler(ConfHandler): - def __init__(self, log_path: str, *args, **kwargs): - super().__init__(*args, **kwargs) + @staticmethod + def _coerce_timeout(value: object, default: float = 0.5) -> float: + if isinstance(value, (int, float)): + return float(value) + if isinstance(value, str): + try: + return float(value) + except ValueError: + return default + return default + + def __init__(self, log_path: str, *args: object, **kwargs: object) -> None: + super().__init__(*args, **kwargs) # type: ignore[arg-type] self.log_path = log_path self.log = get_logger("process_manager.conf_handler") - def get_log_path(self): + def get_log_path(self) -> str: return self.log_path - def _parse_dict(self, data): + def _parse_dict(self, data: dict[str, object]) -> ProcessManagerConfData: new_data = ProcessManagerConfData() if data.get("broadcaster"): new_data.broadcaster = KafkaBroadcastSenderConfData.from_dict( - data.get("broadcaster") + cast(dict[str, object], data.get("broadcaster")) ) - new_data.environment = data.get("environment", {}) - new_data.settings = data.get("settings", {}) + new_data.environment = cast(dict[str, str], data.get("environment", {})) + new_data.settings = cast(dict[str, object], data.get("settings", {})) - match data["type"].lower(): + conf_type = str(data["type"]).lower() + + match conf_type: case "ssh": new_data.type = ProcessManagerTypes.SSH_SHELL - new_data.kill_timeout = data.get("kill_timeout", 0.5) + new_data.kill_timeout = self._coerce_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) + new_data.kill_timeout = self._coerce_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") + new_data.image = str(data.get("image", "ghcr.io/dune-daq/alma9:latest")) case _: - raise UnknownProcessManagerType(data["type"]) + raise UnknownProcessManagerType(str(data["type"])) # opmon_publisher left as default None - self.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, - application="process_manager", + self.opmon_conf = cast( + _OpMonConf, + parse_opmon_conf( + log=self.log, + conf=data.get("opmon_conf", None), + uri=data.get("opmon_uri", None), + session=new_data.type.name, + application="process_manager", + ), ) if self.opmon_conf.path == "./info.json": @@ -128,9 +183,9 @@ def _parse_dict(self, data): def get_commandline_parameters( config_filename: str, - session_dal: "conffwk.dal.Session", + session_dal: _SessionDal, session_name: str, - obj: Any, + obj: _AppLike, ) -> list[str]: """ Build CLI arguments used to launch a process manager application. @@ -205,7 +260,7 @@ def get_process_manager_configuration(process_manager_conf_filename: str) -> str return process_manager_conf_filename -def _load_pm_schema_from_package() -> Dict[str, Any]: +def _load_pm_schema_from_package() -> dict[str, object]: """Load JSON Schema from packaged file; raise if missing or unreadable.""" try: # Package path for schema JSON: drunc/data/process_manager/schema/process_manager.schema.json @@ -216,7 +271,7 @@ def _load_pm_schema_from_package() -> Dict[str, Any]: if not hasattr(schema_resource, "open"): raise FileNotFoundError("process_manager.schema.json resource not found") with schema_resource.open("r", encoding="utf-8") as f: - return json.load(f) + return cast(dict[str, object], json.load(f)) except Exception as e: logger = get_logger("process_manager.config_validation") logger.error(f"Failed to load packaged schema: {e}") @@ -225,7 +280,7 @@ def _load_pm_schema_from_package() -> Dict[str, Any]: ) -def _load_config_from_source(source: Union[str, Dict[str, Any]]) -> Dict[str, Any]: +def _load_config_from_source(source: str | dict[str, object]) -> dict[str, object]: """ Accepts: - file URLs (file:///...), @@ -237,7 +292,7 @@ def _load_config_from_source(source: Union[str, Dict[str, Any]]) -> Dict[str, An s = source.lstrip() # Raw JSON text if s.startswith("{") or s.startswith("["): - return json.loads(source) + return cast(dict[str, object], json.loads(source)) # URL? if "://" in source: @@ -245,7 +300,7 @@ def _load_config_from_source(source: Union[str, Dict[str, Any]]) -> Dict[str, An if u.scheme == "file": path = unquote(u.path) with open(path, "r", encoding="utf-8") as f: - return json.load(f) + return cast(dict[str, object], json.load(f)) raise FileNotFoundError(f"Unsupported URL scheme: {u.scheme}") # Name or filesystem path: resolve and then load @@ -253,7 +308,7 @@ def _load_config_from_source(source: Union[str, Dict[str, Any]]) -> Dict[str, An u = urlparse(resolved_url) path = unquote(u.path) with open(path, "r", encoding="utf-8") as f: - return json.load(f) + return cast(dict[str, object], json.load(f)) # Already a dict elif isinstance(source, dict): @@ -262,7 +317,7 @@ def _load_config_from_source(source: Union[str, Dict[str, Any]]) -> Dict[str, An raise TypeError("validate_config() expects dict, path, URL, or raw JSON text") -def validate_pm_config(config_or_source: Union[str, Dict[str, Any]]) -> bool: +def validate_pm_config(config_or_source: str | dict[str, object]) -> bool: try: pm_conf = _load_config_from_source(config_or_source) schema = _load_pm_schema_from_package() diff --git a/src/drunc/process_manager/exceptions.py b/src/drunc/process_manager/exceptions.py index 1e2f119d7..d716d1ac0 100644 --- a/src/drunc/process_manager/exceptions.py +++ b/src/drunc/process_manager/exceptions.py @@ -2,7 +2,7 @@ class UnknownProcessManagerType(DruncSetupException): - def __init__(self, pm_type): + def __init__(self, pm_type: str) -> None: super().__init__(f"'{pm_type}' is not handled/unknown") diff --git a/src/drunc/process_manager/interface/cli_argument.py b/src/drunc/process_manager/interface/cli_argument.py index 27e84872f..d71d1ba10 100644 --- a/src/drunc/process_manager/interface/cli_argument.py +++ b/src/drunc/process_manager/interface/cli_argument.py @@ -1,14 +1,25 @@ +from collections.abc import Callable +from typing import ParamSpec, TypeVar + import click +from click import Context +from click.core import Parameter -from drunc.process_manager.utils import generate_process_query +R = TypeVar("R") +P = ParamSpec("P") -def validate_conf_string(ctx, param, boot_configuration): +def validate_conf_string( + ctx: Context, + param: Parameter, + boot_configuration: str, +) -> str: + del ctx, param return boot_configuration -def add_query_options(at_least_one: bool, all_processes_by_default: bool = False): - def wrapper(f0): +def add_query_options() -> Callable[[Callable[P, R]], Callable[P, R]]: + def wrapper(f0: Callable[P, R]) -> Callable[P, R]: f1 = click.option( "-s", "--session", @@ -38,6 +49,6 @@ def wrapper(f0): multiple=True, help="Select the process of a particular UUIDs", )(f3) - return generate_process_query(f4, at_least_one, all_processes_by_default) + return f4 return wrapper diff --git a/src/drunc/process_manager/interface/commands.py b/src/drunc/process_manager/interface/commands.py index 129f3ea78..3d2a0fb1c 100644 --- a/src/drunc/process_manager/interface/commands.py +++ b/src/drunc/process_manager/interface/commands.py @@ -10,11 +10,22 @@ validate_conf_string, ) from drunc.process_manager.interface.context import ProcessManagerContext -from drunc.process_manager.utils import tabulate_process_instance_list +from drunc.process_manager.process_manager_driver import ProcessManagerDriver +from drunc.process_manager.utils import ( + build_process_query, + tabulate_process_instance_list, +) from drunc.utils.shell_utils import InterruptedCommand from drunc.utils.utils import get_logger +def _pm_driver(obj: ProcessManagerContext) -> ProcessManagerDriver: + driver = obj.get_driver("process_manager") + if not isinstance(driver, ProcessManagerDriver): + raise RuntimeError("Process manager driver is not initialized") + return driver + + @click.command("boot") @click.option( "-u", @@ -43,7 +54,8 @@ def boot( override_logs: bool, ) -> None: log = get_logger("process_manager.shell") - processes = obj.get_driver("process_manager").ps(ProcessQuery(user=user)) + pm_driver = _pm_driver(obj) + processes = pm_driver.ps(ProcessQuery(user=user)) if len(processes.values) > 0: click.confirm( @@ -55,7 +67,7 @@ def boot( f"Booting session {session_name} with boot configuration file {configuration_file} and id {configuration_id}, requested by user {user}" ) try: - results = obj.get_driver("process_manager").boot( + results = pm_driver.boot( conf_file=configuration_file, conf_id=configuration_id, user=user, @@ -63,6 +75,8 @@ def boot( log_level="INFO", ## Unused anyway!! override_logs=override_logs, ) + if results is None: + return for result in results: if not result: break @@ -75,7 +89,7 @@ def boot( log.exception(e) raise e - controller_address = obj.get_driver("process_manager").controller_address + controller_address = pm_driver.controller_address if controller_address: obj.print( Panel( @@ -129,11 +143,12 @@ def dummy_boot( session_name: str, ) -> None: log = get_logger("process_manager.shell") + pm_driver = _pm_driver(obj) log.debug( f"Running dummy_boot with {n_processes} processes for {sleep} seconds {n_sleeps} times, requested by user {user}" ) try: - results = obj.get_driver("process_manager").dummy_boot( + results = pm_driver.dummy_boot( user=user, session_name=session_name, n_processes=n_processes, @@ -162,7 +177,7 @@ def dummy_boot( def terminate(obj: ProcessManagerContext, width: int | None) -> None: log = get_logger("process_manager.shell") log.debug("Terminating") - result = obj.get_driver("process_manager").terminate() + result = _pm_driver(obj).terminate() if not result: return obj.print( @@ -179,7 +194,7 @@ def terminate(obj: ProcessManagerContext, width: int | None) -> None: default=None, help="Table width. Default is automatically calculated", ) -@add_query_options(at_least_one=True) +@add_query_options() @click.option( "--crash", is_flag=True, @@ -187,10 +202,21 @@ def terminate(obj: ProcessManagerContext, width: int | None) -> None: help="Simulate a crash: send SIGKILL without any cleanup, leaving the process manager in an unexpected-death state.", ) @click.pass_obj -def kill(obj: ProcessManagerContext, query: ProcessQuery, width: int | None) -> None: +def kill( + obj: ProcessManagerContext, + crash: bool, + session: str | None, + name: tuple[str, ...], + user: str | None, + uuid: tuple[str, ...], + width: int | None, +) -> None: log = get_logger("process_manager.shell") + query = build_process_query( + session, name, user, uuid, at_least_one=True, crash=crash + ) log.debug(f"Killing with query {query}") - result = obj.get_driver("process_manager").kill(query) + result = _pm_driver(obj).kill(query) if not result: return obj.print( @@ -206,16 +232,22 @@ def kill(obj: ProcessManagerContext, query: ProcessQuery, width: int | None) -> default=None, help="Table width. Default is automatically calculated", ) -@add_query_options(at_least_one=False, all_processes_by_default=True) +@add_query_options() @click.pass_obj def flush( obj: ProcessManagerContext, - query: ProcessQuery, + session: str | None, + name: tuple[str, ...], + user: str | None, + uuid: tuple[str, ...], width: int | None, ) -> None: log = get_logger("process_manager.shell") + query = build_process_query( + session, name, user, uuid, at_least_one=False, all_processes_by_default=True + ) log.debug(f"Flushing with query {query}") - result = obj.get_driver("process_manager").flush(query) + result = _pm_driver(obj).flush(query) if not result: return obj.print( @@ -224,7 +256,7 @@ def flush( @click.command("logs") -@add_query_options(at_least_one=True) +@add_query_options() @click.option( "--how-far", type=int, @@ -235,16 +267,23 @@ def flush( @click.option("--grep", type=str, default=None) @click.pass_obj def logs( - obj: ProcessManagerContext, how_far: int, grep: str, query: ProcessQuery + obj: ProcessManagerContext, + how_far: int, + grep: str, + session: str | None, + name: tuple[str, ...], + user: str | None, + uuid: tuple[str, ...], ) -> None: log = get_logger("process_manager.shell") + query = build_process_query(session, name, user, uuid, at_least_one=True) log.debug(f"Running logs with query {query}") log_req = LogRequest( how_far=how_far, query=query, ) - result = obj.get_driver("process_manager").logs(log_req) + result = _pm_driver(obj).logs(log_req) if result is None: return @@ -273,16 +312,23 @@ def logs( @click.command("restart") -@add_query_options(at_least_one=True) +@add_query_options() @click.pass_obj -def restart(obj: ProcessManagerContext, query: ProcessQuery) -> None: +def restart( + obj: ProcessManagerContext, + session: str | None, + name: tuple[str, ...], + user: str | None, + uuid: tuple[str, ...], +) -> None: log = get_logger("process_manager.shell") + query = build_process_query(session, name, user, uuid, at_least_one=True) log.debug(f"Restarting with query {query}") - obj.get_driver("process_manager").restart(query) + _pm_driver(obj).restart(query) @click.command("ps") -@add_query_options(at_least_one=False, all_processes_by_default=True) +@add_query_options() @click.option( "-l", "--long-format", @@ -301,13 +347,19 @@ def restart(obj: ProcessManagerContext, query: ProcessQuery) -> None: @click.pass_obj def ps( obj: ProcessManagerContext, - query: ProcessQuery, + session: str | None, + name: tuple[str, ...], + user: str | None, + uuid: tuple[str, ...], long_format: bool, width: int | None, ) -> None: log = get_logger("process_manager.shell") + query = build_process_query( + session, name, user, uuid, at_least_one=False, all_processes_by_default=True + ) log.debug(f"Running ps with query {query}") - results = obj.get_driver("process_manager").ps(query) + results = _pm_driver(obj).ps(query) if not results: return obj.print( diff --git a/src/drunc/process_manager/interface/context.py b/src/drunc/process_manager/interface/context.py index a81396130..ac2132ee7 100644 --- a/src/drunc/process_manager/interface/context.py +++ b/src/drunc/process_manager/interface/context.py @@ -14,19 +14,22 @@ class ProcessManagerContext(ShellContext): # boilerplatefest - def __init__(self, *args, **kwargs): - self.status_receiver = None + def __init__(self, *args: object, **kwargs: object) -> None: + self.status_receiver: BroadcastHandler | None = None super(ProcessManagerContext, self).__init__(*args, **kwargs) - def reset(self, address: str = None): - self.address = resolve_localhost_to_hostname(address) + def reset(self, **kwargs: object) -> None: + address = kwargs.get("address") + resolved_address = address if isinstance(address, str) else "" + self.address = resolve_localhost_to_hostname(resolved_address) super(ProcessManagerContext, self)._reset( name="process_manager_context", token_args={}, driver_args={}, ) - def create_drivers(self, **kwargs) -> Mapping[str, object]: + def create_drivers(self, **kwargs: object) -> Mapping[str, object]: + del kwargs if not self.address: return {} return { @@ -36,10 +39,11 @@ def create_drivers(self, **kwargs) -> Mapping[str, object]: ) } - def create_token(self, **kwargs) -> Token: + def create_token(self, **kwargs: object) -> Token: + del kwargs return create_dummy_token_from_uname() - def start_listening(self, broadcaster_conf): + def start_listening(self, broadcaster_conf: object) -> None: bcch = BroadcastClientConfHandler( data=broadcaster_conf, type=ConfTypes.ProtobufAny, @@ -49,6 +53,6 @@ def start_listening(self, broadcaster_conf): f":ear: Listening to the Process Manager at {self.address}" ) - def terminate(self): + def terminate(self) -> None: 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..54db61bd0 100644 --- a/src/drunc/process_manager/interface/process_manager.py +++ b/src/drunc/process_manager/interface/process_manager.py @@ -4,6 +4,8 @@ import signal import sys import types +from collections.abc import Callable +from typing import Protocol import click import grpc @@ -31,7 +33,16 @@ resolve_localhost_and_127_ip_to_network_ip, ) -_cleanup_coroutines = [] + +class _ReadyEvent(Protocol): + def set(self) -> None: ... + + +class _GeneratedPort(Protocol): + value: int + + +_cleanup_coroutines: list[Callable[[], None]] = [] def run_pm( @@ -39,10 +50,10 @@ def run_pm( pm_address: str, log_level: str, override_logs: bool, - log_path: str = None, - ready_event: bool = None, - signal_handler: bool = None, - generated_port: bool = None, + log_path: str | None = None, + ready_event: _ReadyEvent | None = None, + signal_handler: Callable[[], None] | None = None, + generated_port: _GeneratedPort | None = None, ) -> None: appName = "process_manager" log = get_logger(logger_name=appName, rich_handler=True) @@ -61,8 +72,10 @@ def run_pm( conf_path, conf_type = parse_conf_url(pm_conf) + pm_log_path = log_path or "" + pmch = ProcessManagerConfHandler( - log_path=log_path, type=conf_type, data=conf_path.split(":")[1] + log_path=pm_log_path, type=conf_type, data=conf_path.split(":")[1] ) log_path = get_log_path( @@ -70,7 +83,7 @@ def run_pm( session_name=pmch.data.type.name, application_name=appName, override_logs=override_logs, - app_log_path=log_path, + app_log_path=pm_log_path, ) # Logger has been added to process_manager, so everything will be logged @@ -131,7 +144,7 @@ def server_shutdown() -> None: server = None return - def handle_sigterm(signum: int, frame: types.FrameType) -> None: + def handle_sigterm(signum: int, frame: types.FrameType | None) -> None: """ Handle the SIGTERM signal to gracefully shut down the server. @@ -141,6 +154,7 @@ def handle_sigterm(signum: int, frame: types.FrameType) -> None: """ log.debug("SIGTERM received, shutting down server...") + del signum, frame server_shutdown() return @@ -184,7 +198,11 @@ def handle_sigterm(signum: int, frame: types.FrameType) -> None: help="Log path of process_manager logs.", ) def process_manager_cli( - pm_conf: str, pm_port: int, log_level: str, override_logs: bool, log_path: str + pm_conf: str, + pm_port: int, + log_level: str, + override_logs: bool, + log_path: str | None, ) -> None: get_root_logger(log_level) pm_conf = get_process_manager_configuration(pm_conf) diff --git a/src/drunc/process_manager/interface/shell.py b/src/drunc/process_manager/interface/shell.py index 5da7434c5..6d2617281 100644 --- a/src/drunc/process_manager/interface/shell.py +++ b/src/drunc/process_manager/interface/shell.py @@ -2,7 +2,7 @@ import os import click -import click_shell +import click_shell # type: ignore[import-untyped] from daqpytools.logging import HandlerType, add_handler, logging_log_levels from drunc.process_manager.interface.commands import ( @@ -39,7 +39,9 @@ ) @click.argument("process-manager-address", type=str, callback=validate_command_facility) @click.pass_context -def process_manager_shell(ctx, process_manager_address: str, log_level: str) -> None: +def process_manager_shell( # type: ignore[misc] + ctx: click.Context, process_manager_address: str, log_level: str +) -> None: get_root_logger(log_level) process_manager_log = get_logger( logger_name="process_manager", @@ -74,7 +76,7 @@ def process_manager_shell(ctx, process_manager_address: str, log_level: str) -> if desc.HasField("broadcast"): ctx.obj.start_listening(desc.broadcast) - def cleanup(): + def cleanup() -> None: ctx.obj.terminate() process_manager_log.warning( f"[green]{getpass.getuser()}[/green] disconnected from the process manager through a [green]drunc-process-manager-shell[/green]" @@ -82,13 +84,14 @@ def cleanup(): ctx.call_on_close(cleanup) - ctx.command.add_command(boot, "boot") - ctx.command.add_command(terminate, "terminate") - ctx.command.add_command(kill, "kill") - ctx.command.add_command(flush, "flush") - ctx.command.add_command(logs, "logs") - ctx.command.add_command(restart, "restart") - ctx.command.add_command(ps, "ps") - ctx.command.add_command(dummy_boot, "dummy_boot") + if isinstance(ctx.command, click.Group): + ctx.command.add_command(boot, "boot") + ctx.command.add_command(terminate, "terminate") + ctx.command.add_command(kill, "kill") + ctx.command.add_command(flush, "flush") + ctx.command.add_command(logs, "logs") + ctx.command.add_command(restart, "restart") + ctx.command.add_command(ps, "ps") + ctx.command.add_command(dummy_boot, "dummy_boot") process_manager_shell_log.info("Ready") diff --git a/src/drunc/process_manager/k8s_process_manager.py b/src/drunc/process_manager/k8s_process_manager.py index 04116cd61..b0fdb93b2 100644 --- a/src/drunc/process_manager/k8s_process_manager.py +++ b/src/drunc/process_manager/k8s_process_manager.py @@ -10,6 +10,7 @@ import urllib.request import uuid from time import sleep, time +from typing import cast # Local Application Imports from druncschema.broadcast_pb2 import BroadcastType @@ -26,9 +27,11 @@ ) # Third-Party Imports -from kubernetes import client, config, watch -from kubernetes.client.rest import ApiException -from kubernetes.config.config_exception import ConfigException +from kubernetes import client, config, watch # type: ignore[import-untyped] +from kubernetes.client.rest import ApiException # type: ignore[import-untyped] +from kubernetes.config.config_exception import ( # type: ignore[import-untyped] + ConfigException, +) from drunc.k8s_exceptions import ( DruncK8sException, @@ -51,7 +54,7 @@ class K8sPodWatcherThread(threading.Thread): - def __init__(self, pm) -> None: + def __init__(self, pm: "K8sProcessManager") -> None: """ Initialize the pod watcher thread that monitors and notifies on pod events. @@ -61,7 +64,7 @@ def __init__(self, pm) -> None: threading.Thread.__init__(self) self.pm = pm self.daemon = True - self.processed_uuids = set() + self.processed_uuids: set[str] = set() def run(self) -> None: """ @@ -157,7 +160,9 @@ def run(self) -> None: class K8sProcessManager(ProcessManager): - def __init__(self, configuration: ProcessManagerConfHandler, **kwargs) -> None: + def __init__( + self, configuration: ProcessManagerConfHandler, **kwargs: object + ) -> None: """ Manages processes as Kubernetes Pods. This ProcessManager interfaces with the Kubernetes API to start, stop, and monitor @@ -176,7 +181,13 @@ def __init__(self, configuration: ProcessManagerConfHandler, **kwargs) -> None: # Get the username for the session. This is needed as k8s does not pass the # username through to the pod self.session = getpass.getuser() - super().__init__(configuration=configuration, session=self.session, **kwargs) + name = str(kwargs.pop("name", "k8s-process-manager")) + super().__init__( + configuration=configuration, + name=name, + session=self.session, + **kwargs, + ) # Setup the loger self.log = get_logger("process_manager.k8s-process-manager") @@ -202,17 +213,17 @@ def __init__(self, configuration: ProcessManagerConfHandler, **kwargs) -> None: self._api_error_v1_api = client.rest.ApiException # Storage for process orchestrator parameters - self.managed_sessions = set() - self.watchers = [] + self.managed_sessions: set[str] = set() + self.watchers: list[K8sPodWatcherThread] = [] self._start_watcher() - self.sessions_pending_deletion = set() - self.uuids_pending_deletion = set() + self.sessions_pending_deletion: set[str] = set() + self.uuids_pending_deletion: set[str] = set() self.termination_complete_event = threading.Event() - self.final_exit_codes = {} + self.final_exit_codes: dict[str, int] = {} self.local_connection_server_is_booted = False # Host verification cache: {hostname: (is_valid, timestamp)} - self._host_cache = {} + self._host_cache: dict[str, tuple[bool, float]] = {} self._host_cache_lock = threading.Lock() # Get settings from configuration JSON file @@ -231,8 +242,8 @@ def __init__(self, configuration: ProcessManagerConfHandler, **kwargs) -> None: 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 + self.connection_server_port: int | None = None + self.connection_server_node_port: int | None = None # CONFIGURATION - per-pod service port number service = settings.get("service", {}) @@ -308,7 +319,7 @@ def _setup_signal_handlers(self) -> None: parent process dies unexpectedly. """ - def signal_handler(signum, frame): + def signal_handler(signum: int, frame: object) -> None: self.log.info(f"Received signal {signum}, cleaning up all pods...") try: self._terminate_impl() @@ -386,7 +397,7 @@ def is_alive(self, podname: str, session: str) -> bool: try: # Attempt to get the pod status, if you can the pod is alive pod_status = self._core_v1_api.read_namespaced_pod_status(podname, session) - return pod_status.status.phase == "Running" + return cast(bool, pod_status.status.phase == "Running") except self._api_error_v1_api as e: # Error 404 implies that if pod is not found, i.e. it is not alive if e.status == 404: @@ -647,7 +658,7 @@ def _create_namespace_and_wait_for_active(self, session: str) -> None: self._add_creator_label(session, "namespace") self.managed_sessions.add(session) - def _prepare_namespace(self, session) -> None: + def _prepare_namespace(self, session: str) -> None: """ If the namespace already exists and is in 'Terminating' state, waits for it to be fully deleted before recreating it. If the namespace exists and is active, @@ -711,7 +722,9 @@ def _prepare_namespace(self, session) -> None: else: raise DruncK8sException(f"Failed to check namespace '{session}': {e}") - def _create_headless_service(self, podname, session, pod_uid) -> None: + def _create_headless_service( + self, podname: str, session: str, pod_uid: str + ) -> None: """ Create a headless Kubernetes Service for inter-pod DNS discovery. @@ -766,7 +779,9 @@ 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 + ) -> None: """ Create a NodePort Kubernetes Service for external access. @@ -852,7 +867,7 @@ def _create_nodeport_service(self, podname, session, pod_uid) -> None: def _get_pod_volumes_and_mounts( self, boot_request: BootRequest - ) -> tuple[list[client.V1Volume], list[client.V1VolumeMount]]: + ) -> tuple[list[object], list[object]]: """ Prepares all pod volumes and container mounts, including static configs, performance hardware, and dynamic data/home mounts. @@ -1061,7 +1076,7 @@ def _get_tree_labels( def _build_container_env( self, boot_request: BootRequest, tree_labels: dict[str, str] - ) -> list[client.V1EnvVar]: + ) -> list[object]: """ Builds the list of environment variables for the container. @@ -1114,7 +1129,8 @@ def _build_container_env( # here to ensure that if it is used, it is set to a consistent and expected path # as other functions (e.g. os.expanduser) may not work as expected in the k8s # environment without it. - home_path: str = self.home_path_base + "/" + env_vars.get("USER") + home_path_base = str(self.home_path_base) + home_path: str = home_path_base + "/" + env_vars.get("USER", "") self.log.debug(f"Setting HOME environment variable to: {home_path}") if "HOME" in env_vars: self.log.warning( @@ -1151,9 +1167,9 @@ def _build_pod_main_container( podname: str, boot_request: BootRequest, lcs_port: int | None, - container_volume_mounts: list[client.V1VolumeMount], + container_volume_mounts: list[object], tree_labels: dict[str, str], - ) -> client.V1Container: + ) -> object: """ Build the primary pod container manifest from a boot request. @@ -1305,7 +1321,7 @@ def _build_pod_main_container( def _get_pod_node_selector( self, podname: str, restriction: ProcessRestriction - ) -> dict: + ) -> dict[str, str]: """ Build the Kubernetes node selector for a pod based on host restrictions. @@ -1342,7 +1358,7 @@ def _get_pod_node_selector( def _get_pod_host_aliases( self, podname: str, session: str, tree_labels: dict[str, str] - ) -> list[client.V1HostAlias] | None: + ) -> list[object] | None: """ Build host aliases to redirect localhost to the connection server ClusterIP. @@ -1447,13 +1463,13 @@ def _build_pod_manifest( self, podname: str, session: str, - main_container: client.V1Container, - node_selector: dict, - host_aliases: list[client.V1HostAlias] | None, - pod_volumes: list[client.V1Volume], + main_container: object, + node_selector: dict[str, str], + host_aliases: list[object] | None, + pod_volumes: list[object], extra_labels: dict[str, str] | None = None, use_host_network: bool = True, - ) -> client.V1Pod: + ) -> object: """ Assemble the final V1Pod manifest from its component parts. @@ -1509,7 +1525,7 @@ def _build_pod_manifest( ) def _execute_pod_creation_api( - self, session: str, podname: str, pod_manifest: client.V1Pod + self, session: str, podname: str, pod_manifest: object ) -> str: """ Attempts to create the pod via the API. If a 409 Conflict error occurs @@ -1536,7 +1552,7 @@ def _execute_pod_creation_api( session, pod_manifest ) self.log.info(f'Creating pod "{session}.{podname}"') - return created_pod.metadata.uid + return str(created_pod.metadata.uid) except self._api_error_v1_api as e: is_409_conflict = e.status == 409 @@ -1618,7 +1634,11 @@ def _create_associated_service( self._create_headless_service(podname, session, pod_uid) def _create_pod( - self, podname, session, boot_request: BootRequest, tree_labels: dict[str, str] + self, + podname: str, + session: str, + boot_request: BootRequest, + tree_labels: dict[str, str], ) -> None: """ Orchestrates the full pod creation pipeline: extracts the LCS port if @@ -1637,7 +1657,7 @@ def _create_pod( DruncK8sException - if pod or service creation fails for any reason """ try: - lcs_port = None + lcs_port: int | None = None # Early Port Extraction and Class Variable Setup for LCS if self._is_local_connection_server(tree_labels, podname): lcs_port = self._extract_port_from_cmd(boot_request) @@ -1721,7 +1741,7 @@ def _create_pod( f"Failed to create pod '{session}.{podname}': {e}" ) from e - def _get_connection_server_cluster_ip(self, session: str) -> str: + def _get_connection_server_cluster_ip(self, session: str) -> str | None: """ Get the ClusterIP of the connection server's Kubernetes Service. @@ -1739,12 +1759,12 @@ def _get_connection_server_cluster_ip(self, session: str) -> str: service = self._core_v1_api.read_namespaced_service( name=self.connection_server_name, namespace=session ) - return service.spec.cluster_ip + return cast(str | None, service.spec.cluster_ip) except self._api_error_v1_api as e: self.log.error(f"Failed to get connection server service IP: {e}") return None - def _extract_port_from_cmd(self, boot_request) -> int | None: + def _extract_port_from_cmd(self, boot_request: BootRequest) -> int | None: """ Parses the boot request's command arguments to find a port. @@ -1829,7 +1849,12 @@ def _extract_port_from_cmd(self, boot_request) -> int | None: return None - def _get_process_uid(self, query: ProcessQuery, order_by: str = None) -> list[str]: + def _get_process_uid( + self, + query: ProcessQuery, + in_boot_request: bool = False, + order_by: str = "random", + ) -> list[str]: """ Finds process UUIDs matching a query. @@ -1988,7 +2013,7 @@ def _wait_for_pod_api_ready( self.log.info( f"Stage 1: Pod '{podname}' is API Ready on node {node_name}." ) - return node_name # Success! + return cast(str, node_name) # Success! except self._api_error_v1_api as e: if e.status == 404: @@ -2395,7 +2420,12 @@ def _restart_impl(self, query: ProcessQuery) -> ProcessInstanceList: return ProcessInstanceList(values=ret) - def _kill_pod(self, podname, session, grace_period_seconds=None) -> None: + def _kill_pod( + self, + podname: str, + session: str, + grace_period_seconds: int | None = None, + ) -> None: """ Deletes a specific pod from a namespace. @@ -2454,7 +2484,9 @@ def _kill_impl(self, query: ProcessQuery) -> ProcessInstanceList: ) # Define the blocking kill_and_wait helper - def kill_and_wait(uuids, grace_period=None) -> None: + def kill_and_wait( + uuids: set[str] | list[str], grace_period: int | None = None + ) -> None: if not uuids: return action = ( @@ -2499,7 +2531,7 @@ def kill_and_wait(uuids, grace_period=None) -> None: self.log.error(f"Could not list pods for kill operation: {e}") # Map pods by their role label - pods_by_role = { + pods_by_role: dict[str, list[str]] = { "unknown": [], "application": [], "segment-controller": [], @@ -2514,13 +2546,14 @@ def kill_and_wait(uuids, grace_period=None) -> None: tree_id_label_key = f"tree-id.{self.drunc_label}" for pod in all_pods: - uuid = pod.metadata.labels.get(uuid_label_key) - if uuid and uuid in targeted_uuids: + pod_uuid = pod.metadata.labels.get(uuid_label_key) + if pod_uuid and pod_uuid in targeted_uuids: role = pod.metadata.labels.get(role_label_key, "unknown") - pods_by_role[role].append(uuid) - if role == "segment-controller": + role_key = role if role in pods_by_role else "unknown" + pods_by_role[role_key].append(pod_uuid) + if role_key == "segment-controller": tree_id = pod.metadata.labels.get(tree_id_label_key, "") - segment_controller_depths[uuid] = ( + segment_controller_depths[pod_uuid] = ( tree_id.count(".") if tree_id else 0 ) diff --git a/src/drunc/process_manager/oks_parser.py b/src/drunc/process_manager/oks_parser.py index f679e7f1a..4e2cfbff2 100644 --- a/src/drunc/process_manager/oks_parser.py +++ b/src/drunc/process_manager/oks_parser.py @@ -1,14 +1,83 @@ import os -from typing import TYPE_CHECKING, Any, Dict, List +from typing import Protocol, cast -import confmodel_dal +import confmodel_dal # type: ignore[import-untyped] from drunc.exceptions import DruncException, DruncSetupException from drunc.process_manager.configuration import get_commandline_parameters from drunc.utils.utils import get_logger -if TYPE_CHECKING: - import conffwk + +class _VariableLike(Protocol): + name: str + value: str + contains: list["_VariableLike"] + + def className(self) -> str: ... + + +class _RunsOnInner(Protocol): + id: str + + +class _RunsOn(Protocol): + runs_on: _RunsOnInner + + +class _ServiceLike(Protocol): + id: str + port: int + protocol: str + + +class _WriterParams(Protocol): + directory_path: str + + +class _Writer(Protocol): + data_store_params: _WriterParams + + +class _AppLike(Protocol): + id: str + application_name: str + commandline_parameters: list[str] + log_path: str + runs_on: _RunsOn + application_environment: list[_VariableLike] + exposes_service: list[_ServiceLike] + data_writers: list[_Writer] + tp_writer: _Writer + + def oksTypes(self) -> list[str]: ... + + +class _SegmentLike(Protocol): + id: str + controller: _AppLike + segments: list["_SegmentLike"] + applications: list[_AppLike] + + +class _SessionLike(Protocol): + id: str + environment: list[_VariableLike] + controller_log_level: str + segment: _SegmentLike + disabled: list["_HasId"] + infrastructure_applications: list[_AppLike] + + +class _DbLike(Protocol): + _obj: object + + +class _HasId(Protocol): + id: str + + +class _LoggerLike(Protocol): + def debug(self, msg: str) -> None: ... def get_full_db_path(db_path: str) -> str: @@ -32,7 +101,7 @@ def get_full_db_path(db_path: str) -> str: # Get the env var that points to the configuration files. If it doesn't exist, raise # an exception - search_path_str: str = os.environ.get("DUNEDAQ_DB_PATH", None) + search_path_str = os.environ.get("DUNEDAQ_DB_PATH") if not search_path_str: err_str = "DUNEDAQ_DB_PATH not set, exiting." raise DruncSetupException(err_str) @@ -81,7 +150,7 @@ def get_full_db_path(db_path: str) -> str: return resolved_path -def collect_variables(variables, env_dict: Dict[str, str]) -> None: +def collect_variables(variables: list[_VariableLike], env_dict: dict[str, str]) -> None: """!Process a dal::Variable object, placing key/value pairs in a dictionary @param variables A Variable/VariableSet object @@ -101,7 +170,7 @@ class EnvironmentVariableCannotBeSet(DruncException): def component_disabled_from_session_dal( - session_dal_obj: "conffwk.dal.Session", component_id: str + session_dal_obj: _SessionLike, component_id: str ) -> bool: """ Replaces the following without any db dependence @@ -142,13 +211,13 @@ def component_disabled_from_session_dal( def collect_apps( config_filename: str, session_name: str, - session_dal_obj: "conffwk.dal.Session", - segment_obj: "conffwk.dal.Segment", - env: Dict[str, str], - tree_prefix: List[int] = [ + session_dal_obj: _SessionLike, + segment_obj: _SegmentLike, + env: dict[str, str], + tree_prefix: list[int] = [ 0, ], -) -> List[Dict]: +) -> list[dict[str, object]]: """! Recustively collect (daq) application belonging to segment and its subsegments @param session_dal_obj The session the segment belongs to @@ -170,7 +239,7 @@ def collect_apps( collect_variables(session_dal_obj.environment, defenv) - apps = [] + apps: list[dict[str, object]] = [] # Add controller for this segment to list of apps controller = segment_obj.controller @@ -219,56 +288,58 @@ def collect_apps( except Exception as e: log.exception(e) raise e - for app in sub_apps: - apps.append(app) + for sub_app in sub_apps: + apps.append(sub_app) # Get all the enabled applications of this segment # Start app_index after sub-segment indices to avoid tree_id collisions app_index = len(segment_obj.segments) - for app in segment_obj.applications: - log.debug(f"Considering app {app.id}") - if "Resource" in app.oksTypes(): - enabled = not component_disabled_from_session_dal(session_dal_obj, app.id) - log.debug(f"{app.id} {enabled=}") + for segment_app in segment_obj.applications: + log.debug(f"Considering app {segment_app.id}") + if "Resource" in segment_app.oksTypes(): + enabled = not component_disabled_from_session_dal( + session_dal_obj, segment_app.id + ) + log.debug(f"{segment_app.id} {enabled=}") else: enabled = True - log.debug(f"{app.id} {enabled=}") + log.debug(f"{segment_app.id} {enabled=}") if not enabled: - log.debug(f"Ignoring disabled app {app.id}") + log.debug(f"Ignoring disabled app {segment_app.id}") continue app_env = defenv.copy() # Override with any app specific environment from Application - collect_variables(app.application_environment, app_env) - app_env["DUNEDAQ_APPLICATION_NAME"] = app.id + collect_variables(segment_app.application_environment, app_env) + app_env["DUNEDAQ_APPLICATION_NAME"] = segment_app.id app_tree_id_str = ".".join(map(str, tree_prefix + [app_index])) - host = app.runs_on.runs_on.id + host = segment_app.runs_on.runs_on.id args = get_commandline_parameters( config_filename=config_filename, session_dal=session_dal_obj, session_name=session_name, - obj=app, + obj=segment_app, ) - log.debug(f"Collecting app {app.id} with args {args}") + log.debug(f"Collecting app {segment_app.id} with args {args}") - data_path = get_writer_directory_path(app, log) + data_path = get_writer_directory_path(segment_app, log) if not data_path: - log.debug(f"No data path found for app {app.id}") + log.debug(f"No data path found for app {segment_app.id}") apps.append( { - "name": app.id, - "type": app.application_name, + "name": segment_app.id, + "type": segment_app.application_name, "args": args, "restriction": host, "host": host, "env": app_env, "tree_id": app_tree_id_str, - "log_path": app.log_path, + "log_path": segment_app.log_path, "data_path": data_path, } ) @@ -277,7 +348,7 @@ def collect_apps( return apps -def get_writer_directory_path(app, log) -> str | None: +def get_writer_directory_path(app: _AppLike, log: _LoggerLike) -> str | None: # Map known OKS types to their specific writer attribute APP_TYPE_TO_WRITER_ATTR = { "DFApplication": "data_writers", @@ -307,7 +378,7 @@ def get_writer_directory_path(app, log) -> str | None: writer = writers[0] params = getattr(writer, "data_store_params", None) if params and getattr(params, "directory_path", None): - directory_path = params.directory_path + directory_path = cast(str, params.directory_path) log.debug(f"data path for app {app.id}: {directory_path}") return directory_path @@ -315,10 +386,10 @@ def get_writer_directory_path(app, log) -> str | None: def collect_infra_apps( - session: "conffwk.dal.Session", - env: Dict[str, str], - tree_prefix: List[int], -) -> List[Dict[str, Any]]: + session: _SessionLike, + env: dict[str, str], + tree_prefix: list[int], +) -> list[dict[str, object]]: """! Collect infrastructure applications @param session The session @@ -337,7 +408,7 @@ def collect_infra_apps( collect_variables(session.environment, defenv) - apps = [] + apps: list[dict[str, object]] = [] for app_index, app in enumerate(session.infrastructure_applications): # Skip applications that do not define an application name @@ -370,9 +441,14 @@ def collect_infra_apps( # Search segment and all contained segments for apps controlled by # given controller. Return separate lists of apps and sub-controllers -def find_controlled_apps(db, session, mycontroller, segment): - apps = [] - controllers = [] +def find_controlled_apps( + db: _DbLike, + session: _SessionLike, + mycontroller: str, + segment: _SegmentLike, +) -> tuple[list[str], list[str]]: + apps: list[str] = [] + controllers: list[str] = [] if segment.controller.id == mycontroller: for app in segment.applications: apps.append(app.id) diff --git a/src/drunc/process_manager/process_manager.py b/src/drunc/process_manager/process_manager.py index 07afa8b86..483a276cf 100644 --- a/src/drunc/process_manager/process_manager.py +++ b/src/drunc/process_manager/process_manager.py @@ -3,6 +3,7 @@ import sys import threading import time +from typing import Protocol, cast from daqpytools.logging import LogHandlerConf, exceptions, setup_daq_ers_logger from druncschema.authoriser_pb2 import ActionType, SystemType @@ -44,20 +45,25 @@ class BadQuery(DruncCommandException): - def __init__(self, txt): + def __init__(self, txt: str) -> None: super(BadQuery, self).__init__(txt, code_pb2.INVALID_ARGUMENT) +class _OpMonPublisher(Protocol): + def publish(self, *, message: ProcessStatus) -> None: ... + + class ProcessManager(abc.ABC, ProcessManagerServicer): def __init__( self, configuration: ProcessManagerConfHandler, name: str, - session: str = None, - **kwargs, - ): + session: str | None = None, + **kwargs: object, + ) -> None: """C'tor. Note that this takes the ERS env variables from the - json files defined in data/process_manager!""" + json files defined in data/process_manager! + """ super().__init__() self.log = get_logger( @@ -87,18 +93,25 @@ def __init__( data=self.configuration.get_data_authoriser(), type=ConfTypes.PyObject ) - self.opmon_publisher = getattr( - self.configuration.get_data(), "opmon_publisher", None + self.opmon_publisher = cast( + _OpMonPublisher | None, + getattr(self.configuration.get_data(), "opmon_publisher", None), + ) + interval_raw = getattr(self.configuration.get_data(), "interval_s", 10.0) + try: + interval_s = float(interval_raw) + except (TypeError, ValueError): + interval_s = 10.0 + self.authoriser = DummyAuthoriser( + cast(SystemType, SystemType.PROCESS_MANAGER), dach ) - interval_s = getattr(self.configuration.get_data(), "interval_s", 10.0) - self.authoriser = DummyAuthoriser(dach, SystemType.PROCESS_MANAGER) - self.process_store = {} # dict[str, sh.RunningCommand] # str = uuid - self.boot_request = {} # dict[str, BootRequest] # str = uuid + self.process_store: dict[str, object] = {} # str = uuid + self.boot_request: dict[str, BootRequest] = {} # str = uuid # Define a list of applications that we expect to die, and a lock to read the memory self.dead_process_lock = threading.Lock() - self.expected_dead_applications = {} # dict[str, BootRequest] # str == uuid + self.expected_dead_applications: dict[str, BootRequest] = {} # str == uuid # TODO, probably need to think of a better way to do this? # Maybe I should "bind" the commands to their methods, and have something looping over this list to generate the gRPC functions @@ -165,10 +178,10 @@ def __init__( ) self.thread.start() - def get_log_path(self): + def get_log_path(self) -> str: return self.configuration.get_log_path() - def _create_broadcast_service(self, name, session): + def _create_broadcast_service(self, name: str, session: str | None) -> None: bsch = BroadcastSenderConfHandler( data=self.configuration.get_data_broadcaster(), type=ConfTypes.PyObject ) @@ -176,20 +189,25 @@ def _create_broadcast_service(self, name, session): self.broadcast_service = ( BroadcastSender( name=name, - session=session, + session=session or "", configuration=bsch, ) if bsch.data else None ) - def __del__(self): + def __del__(self) -> None: if hasattr(self, "opmon_publisher") and self.opmon_publisher is not None: self.stop_event.set() self.thread.join() - def publish(self, q: ProcessQuery, interval_s: float = 10.0): - def find_by_uuid(pi_list, target_uuid: str): + def publish(self, q: ProcessQuery, interval_s: float = 10.0) -> None: + if self.opmon_publisher is None: + return + + def find_by_uuid( + pi_list: ProcessInstanceList, target_uuid: str + ) -> ProcessInstance | None: """Identifies the process from a list by uuid""" for pi in pi_list.values: if pi.uuid.uuid == target_uuid: @@ -197,7 +215,7 @@ def find_by_uuid(pi_list, target_uuid: str): return None n_dead_prev = 0 - dead_processes_prev = set() + dead_processes_prev: set[str] = set() while not self.stop_event.is_set(): results = self._ps_impl(q) @@ -233,6 +251,8 @@ def find_by_uuid(pi_list, target_uuid: str): ) continue pi = find_by_uuid(results, diff) + if pi is None: + continue pi_return_code = ( pi.return_code if pi.HasField("return_code") else "NONE" ) @@ -251,7 +271,7 @@ def find_by_uuid(pi_list, target_uuid: str): A couple of simple pass-through functions to the broadcasting service """ - def broadcast(self, *args, **kwargs): + def broadcast(self, *args: object, **kwargs: object) -> object | None: self.log.debug(f"{self.name} broadcasting") return ( self.broadcast_service.broadcast(*args, **kwargs) @@ -259,7 +279,7 @@ def broadcast(self, *args, **kwargs): else None ) - def can_broadcast(self, *args, **kwargs): + def can_broadcast(self, *args: object, **kwargs: object) -> bool: self.log.debug(f"Checking if {self.name} can broadcast") return ( self.broadcast_service.can_broadcast(*args, **kwargs) @@ -267,7 +287,7 @@ def can_broadcast(self, *args, **kwargs): else False ) - def describe_broadcast(self, *args, **kwargs): + def describe_broadcast(self, *args: object, **kwargs: object) -> object | None: self.log.debug(f"Describing {self.name} broadcast") return ( self.broadcast_service.describe_broadcast(*args, **kwargs) @@ -275,7 +295,9 @@ def describe_broadcast(self, *args, **kwargs): else None ) - def interrupt_with_exception(self, *args, **kwargs): + def interrupt_with_exception( + self, *args: object, **kwargs: object + ) -> object | None: self.log.debug(f"Interrupting {self.name} broadcast with exception") return ( self.broadcast_service._interrupt_with_exception(*args, **kwargs) @@ -292,7 +314,7 @@ def _boot_impl(self, boot_request: BootRequest) -> ProcessInstanceList: @authentified_and_authorised( action=ActionType.CREATE, system=SystemType.PROCESS_MANAGER ) # 2nd step - def boot( + def boot( # type: ignore[misc] self, request: BootRequest, context: ServicerContext ) -> ProcessInstanceList: self.log.debug( @@ -326,7 +348,7 @@ def _terminate_impl(self) -> ProcessInstanceList: @authentified_and_authorised( action=ActionType.DELETE, system=SystemType.PROCESS_MANAGER ) # 2nd step - def terminate( + def terminate( # type: ignore[misc] self, request: Request, context: ServicerContext ) -> ProcessInstanceList: self.log.debug(f"{self.name} running terminate") @@ -361,7 +383,7 @@ def _restart_impl(self, query: ProcessQuery) -> ProcessInstanceList: @authentified_and_authorised( action=ActionType.DELETE, system=SystemType.PROCESS_MANAGER ) # 2nd step - def restart( + def restart( # type: ignore[misc] self, request: ProcessQuery, context: ServicerContext ) -> ProcessInstanceList: self.log.debug(f"{self.name} running restart") @@ -393,7 +415,7 @@ def _kill_impl(self, query: ProcessQuery) -> ProcessInstanceList: @authentified_and_authorised( action=ActionType.DELETE, system=SystemType.PROCESS_MANAGER ) # 2nd step - def kill( + def kill( # type: ignore[misc] self, request: ProcessQuery, context: ServicerContext ) -> ProcessInstanceList: self.log.debug(f"{self.name} running kill") @@ -425,7 +447,7 @@ def _ps_impl(self, query: ProcessQuery) -> ProcessInstanceList: @authentified_and_authorised( action=ActionType.READ, system=SystemType.PROCESS_MANAGER ) # 2nd step - def ps( + def ps( # type: ignore[misc] self, request: ProcessQuery, context: ServicerContext ) -> ProcessInstanceList: self.log.debug(f"{self.name} running ps") @@ -457,7 +479,7 @@ def _flush_impl(self, query: ProcessQuery) -> ProcessInstanceList: @authentified_and_authorised( action=ActionType.DELETE, system=SystemType.PROCESS_MANAGER ) # 2nd step - def flush( + def flush( # type: ignore[misc] self, request: ProcessQuery, context: ServicerContext ) -> ProcessInstanceList: """Remove dead processes from tracking so they no longer appear in ps. @@ -498,7 +520,9 @@ def flush( @authentified_and_authorised( action=ActionType.READ, system=SystemType.PROCESS_MANAGER ) # 2nd step - def describe(self, request: Request, context: ServicerContext) -> Description: + def describe( # type: ignore[misc] + self, request: Request, context: ServicerContext + ) -> Description: self.log.debug(f"{self.name} running describe") response = Description( @@ -525,7 +549,9 @@ def _logs_impl(self, log_request: LogRequest) -> LogLines: @authentified_and_authorised( action=ActionType.READ, system=SystemType.PROCESS_MANAGER ) # 2nd step - def logs(self, request: LogRequest, context: ServicerContext) -> LogLines: + def logs( # type: ignore[misc] + self, request: LogRequest, context: ServicerContext + ) -> LogLines: """Fetch logs for a process. Args: @@ -587,7 +613,7 @@ def _ensure_one_process( def _match_processes_against_query( query: ProcessQuery, available_uuids: list[str], - boot_request_dict: dict, + boot_request_dict: dict[str, BootRequest], order_by: str = "random", ) -> list[str]: """ @@ -621,7 +647,7 @@ def _match_processes_against_query( # relevent reading here: https://github.com/protocolbuffers/protobuf/blob/main/docs/field_presence.md # Filter processes based on query criteria - processes = [] + processes: list[str] = [] for uuid in available_uuids: accepted = False meta = boot_request_dict[uuid].process_description.metadata @@ -781,8 +807,10 @@ def _get_process_uid( return matched_processes @staticmethod - def get(conf, **kwargs): + def get(conf: ProcessManagerConfHandler, **kwargs: object) -> "ProcessManager": log = get_logger("process_manager.get") + name = str(kwargs.get("name", "process_manager")) + forwarded_kwargs = {k: v for k, v in kwargs.items() if k != "name"} if conf.data.type == ProcessManagerTypes.SSH_SHELL: from drunc.process_manager.ssh_process_manager_shell import ( @@ -790,21 +818,19 @@ def get(conf, **kwargs): ) log.debug("Starting [green]SSH Shell process_manager[/green]") - return SSHProcessManagerShell(conf, **kwargs) + return SSHProcessManagerShell(conf, name=name, **forwarded_kwargs) elif conf.data.type == ProcessManagerTypes.K8s: from drunc.process_manager.k8s_process_manager import K8sProcessManager log.debug("Starting [green]K8s process_manager[/green]") - return K8sProcessManager(conf, **kwargs) + return K8sProcessManager(conf, name=name, **forwarded_kwargs) elif conf.data.type == ProcessManagerTypes.SSH_PARAMIKO: from drunc.process_manager.ssh_process_manager_paramiko_client import ( SSHProcessManagerParamikoClient, ) log.debug("Starting [green]SSH Paramiko process_manager[/green]") - return SSHProcessManagerParamikoClient(conf, **kwargs) + return SSHProcessManagerParamikoClient(conf, name=name, **forwarded_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.data.type} is unsupported!") + raise RuntimeError(f"ProcessManager type {conf.data.type} is unsupported!") diff --git a/src/drunc/process_manager/process_manager_driver.py b/src/drunc/process_manager/process_manager_driver.py index 1f9e50aeb..eccc6fafd 100644 --- a/src/drunc/process_manager/process_manager_driver.py +++ b/src/drunc/process_manager/process_manager_driver.py @@ -6,13 +6,17 @@ import time from collections.abc import Iterator from time import sleep -from typing import Dict, List +from typing import NotRequired, Protocol, TypedDict, cast import conffwk import grpc -from daqconf.set_connectivity_service_port import set_connectivity_service_port -from daqconf.set_rc_controller_port import set_rc_controller_port -from daqconf.utils import find_free_port +from daqconf.set_connectivity_service_port import ( # type: ignore[import-untyped] + set_connectivity_service_port, +) +from daqconf.set_rc_controller_port import ( # type: ignore[import-untyped] + set_rc_controller_port, +) +from daqconf.utils import find_free_port # type: ignore[import-untyped] from druncschema.description_pb2 import Description from druncschema.process_manager_pb2 import ( BootRequest, @@ -51,8 +55,85 @@ ) +class BootApp(TypedDict): + name: str + type: str + args: list[str] + env: dict[str, str] + log_path: str | None + restriction: str + tree_id: str + data_path: NotRequired[str | None] + + +class _HasId(Protocol): + id: str + + +class _RunsOnNode(Protocol): + id: str + + +class _RunsOnRef(Protocol): + runs_on: _RunsOnNode + + +class _ServiceRef(Protocol): + id: str + port: int + protocol: str + + +class _ControllerRef(Protocol): + id: str + runs_on: _RunsOnRef + exposes_service: list[_ServiceRef] + + +class _SegmentRef(Protocol): + controller: _ControllerRef + + +class _ConnectivityEndpoint(Protocol): + port: int + + +class _ConnectivityServiceRef(Protocol): + host: str + service: _ConnectivityEndpoint + + +class SessionDalProto(Protocol): + id: str + segment: _SegmentRef + infrastructure_applications: list[_HasId] + log_path: str + rte_script: str | None + connectivity_service: _ConnectivityServiceRef | None + environment: object + + +class _DalClassNamed(Protocol): + def className(self) -> str: ... + + +class _DalVariable(_DalClassNamed, Protocol): + name: str + value: object + + +class _DalVariableSet(_DalClassNamed, Protocol): + contains: list[object] + + +class _ConfigurationProto(Protocol): + active_database: str + + def get_dal(self, class_name: str, uid: str) -> object: ... + + class ProcessManagerDriver: - controller_address = "" + controller_address: str | None = None def __init__(self, address: str, token: Token): self.log = get_logger("process_manager_driver", rich_handler=True) @@ -96,8 +177,8 @@ def boot( sleep_between_app_boot: ( int | float ) = 0, # This may be useful if you have are using SSHPM, and have SSHD's maxstartups setting set to a low value. - **kwargs, - ) -> Iterator[ProcessInstanceList] | None: + **kwargs: object, + ) -> Iterator[ProcessInstanceList]: self.log.info(f"Booting session [green]{session_name}[/green]") # Step 1 - consolidate configuration @@ -115,7 +196,7 @@ def boot( ) # Step 5 - track boot timings per host - last_boot_on_host_at = {} + last_boot_on_host_at: dict[str, float] = {} previous_host = None # Step 6: iterate over boot requests @@ -129,7 +210,7 @@ def boot( ): if not request: self.log.error("[red]No boot request was generated, ending boot.[/red]") - return None + return if request.process_description.metadata.name in [ app.id for app in session_dal.infrastructure_applications ]: @@ -181,29 +262,33 @@ def boot( def _collect_all_apps( self, oks_conf: str, - session_dal: "conffwk.dal.Session", + session_dal: SessionDalProto, session_name: str, - ) -> List[Dict]: + ) -> list[BootApp]: from drunc.process_manager.oks_parser import collect_apps, collect_infra_apps env = { "DUNEDAQ_SESSION": session_name, } - apps = collect_apps( + raw_apps = collect_apps( session_name=session_name, config_filename=oks_conf, - session_dal_obj=session_dal, - segment_obj=session_dal.segment, + session_dal_obj=session_dal, # type: ignore[arg-type] + segment_obj=session_dal.segment, # type: ignore[arg-type] env=env, tree_prefix=[ 0, ], ) + apps: list[BootApp] = [cast(BootApp, app) for app in raw_apps] # Next line gets the max of all the first number in the tree id, and adds 1 to it. next_tree_id = max([int(app["tree_id"].split(".")[0]) for app in apps]) + 1 - infra_apps = collect_infra_apps(session_dal, env, tree_prefix=[next_tree_id]) + infra_apps = cast( + list[BootApp], + collect_infra_apps(session_dal, env, tree_prefix=[next_tree_id]), # type: ignore[arg-type] + ) apps = infra_apps + apps @@ -212,8 +297,11 @@ def _collect_all_apps( return apps def _prepare_exec_and_args( - self, session_dal, exe: str, args: List[str] - ) -> List[ProcessDescription.ExecAndArgs]: + self, + session_dal: SessionDalProto, + exe: str, + args: list[str], + ) -> list[ProcessDescription.ExecAndArgs]: """ Prepare """ @@ -247,10 +335,10 @@ def _prepare_exec_and_args( def _build_boot_request( self, - app: Dict, + app: BootApp, user: str, session_name: str, - session_dal, + session_dal: SessionDalProto, session_log_path: str, override_logs: bool, pwd: str, @@ -262,8 +350,12 @@ def _build_boot_request( env = app["env"] app_log_path = app["log_path"] data_path = app.get("data_path") - env["DUNE_DAQ_BASE_RELEASE"] = os.getenv("DUNE_DAQ_BASE_RELEASE") - env["SPACK_RELEASES_DIR"] = os.getenv("SPACK_RELEASES_DIR") + dune_base_release = os.getenv("DUNE_DAQ_BASE_RELEASE") + if dune_base_release is not None: + env["DUNE_DAQ_BASE_RELEASE"] = dune_base_release + spack_releases_dir = os.getenv("SPACK_RELEASES_DIR") + if spack_releases_dir is not None: + env["SPACK_RELEASES_DIR"] = spack_releases_dir tree_id = app["tree_id"] self.log.debug(f"{name}:\n{json.dumps(app, indent=4)}") @@ -318,7 +410,7 @@ def _convert_oks_to_boot_request( self, oks_conf: str, user: str, - session_dal, + session_dal: SessionDalProto, session_name: str, override_logs: bool, ) -> Iterator[BootRequest]: @@ -344,11 +436,11 @@ def _convert_oks_to_boot_request( except DruncSetupException as e: log = get_logger("utils.boot_req_generator") log.error(f"[red]Caught exception in boot generator [/red]: {e}") - yield None + continue yield breq - def _consolidate_config(self, session_name, conf_file: str) -> str | None: - from daqconf.consolidate import consolidate_db + def _consolidate_config(self, session_name: str, conf_file: str) -> None: + from daqconf.consolidate import consolidate_db # type: ignore[import-untyped] self.log.debug(f"Validating {session_name} configuration") @@ -373,21 +465,26 @@ def _consolidate_config(self, session_name, conf_file: str) -> str | None: def update_connectivity_port_dal( self, - env_variables: list["conffwk.dal.Variable | conffwk.dal.VariableSet"], + env_variables: list[object], new_port: int, ) -> None: """Process a dal::Variable object, placing key/value pairs in a dictionary""" for item in env_variables: - if item.className() == "VariableSet": - self.update_connectivity_port_dal(item.contains, new_port) + class_named = cast(_DalClassNamed, item) + if class_named.className() == "VariableSet": + variable_set = cast(_DalVariableSet, item) + self.update_connectivity_port_dal(variable_set.contains, new_port) else: - if item.className() == "Variable": - if item.name == "CONNECTION_PORT": - item.value = new_port + if class_named.className() == "Variable": + variable = cast(_DalVariable, item) + if variable.name == "CONNECTION_PORT": + variable.value = new_port def check_port_conflicts( - self, db: conffwk.Configuration, session_dal: "conffwk.dal.Session" - ) -> tuple[conffwk.Configuration, "conffwk.dal.Session"]: + self, + db: _ConfigurationProto, + session_dal: SessionDalProto, + ) -> tuple[_ConfigurationProto, SessionDalProto]: """ Check that the ports allocated in the configuration file are available. If the file is editable, make the changes in the file itself. Otherwise, make the @@ -423,7 +520,7 @@ def check_port_conflicts( # Check that the address of the root controller is available, otherwise change # it to one that is available root_controller_host: str = session_dal.segment.controller.runs_on.runs_on.id - root_controller_service_list: int = [ + root_controller_service_list: list[_ServiceRef] = [ service for service in session_dal.segment.controller.exposes_service if "_control" in service.id @@ -449,7 +546,10 @@ def check_port_conflicts( # If a local connectivity service is being used, perform the same checks # Temporarily removed to allow integration tests to pass without restructuring # Note - if infrastructure applications outside of the connectivity service are spawned, this will need to be adjusted. - if session_dal.infrastructure_applications: # Check if the own application needs to be spawned, or if an externally managed one is in use (e.g. if using ehn1 connectivity service or integration tests.) + if ( + session_dal.infrastructure_applications + and session_dal.connectivity_service is not None + ): # Check if the own application needs to be spawned, or if an externally managed one is in use (e.g. if using ehn1 connectivity service or integration tests.) connectivity_service_host: str = session_dal.connectivity_service.host connectivity_service_port = session_dal.connectivity_service.service.port if not is_port_available( @@ -487,9 +587,13 @@ def check_port_conflicts( return db, session_dal else: # If the configuration file has been modified, instantiate a new DAL - updated_db = conffwk.Configuration("oksconflibs:" + configuration_file) - updated_session_dal = updated_db.get_dal( - class_name="Session", uid=configuration_id + updated_db = cast( + _ConfigurationProto, + conffwk.Configuration("oksconflibs:" + configuration_file), + ) + updated_session_dal = cast( + SessionDalProto, + updated_db.get_dal(class_name="Session", uid=configuration_id), ) self.log.info( "Configuration required updates and file is writable, re-instantiating DAL to reflect changes in the file." @@ -499,17 +603,23 @@ def check_port_conflicts( updated_session_dal, ) - def _initialise_session(self, conf_file: str, conf_id: str) -> tuple: + def _initialise_session( + self, conf_file: str, conf_id: str + ) -> tuple[_ConfigurationProto, SessionDalProto]: import conffwk # isort: skip - db = conffwk.Configuration(conf_file) - session_dal = db.get_dal(class_name="Session", uid=conf_id) + db = cast(_ConfigurationProto, conffwk.Configuration(conf_file)) + session_dal = cast( + SessionDalProto, db.get_dal(class_name="Session", uid=conf_id) + ) return db, session_dal def _connect_to_service( - self, session_dal: "conffwk.dal.Session", session_name: str - ) -> ConnectivityServiceClient | None: - if session_dal.connectivity_service: + self, + session_dal: SessionDalProto, + session_name: str, + ) -> tuple[ConnectivityServiceClient | None, str | None, int | None]: + if session_dal.connectivity_service is not None: connection_server = session_dal.connectivity_service.host connection_port = session_dal.connectivity_service.service.port @@ -528,12 +638,12 @@ def _connect_to_service( def _discover_controller( self, - session_dal: "conffwk.dal.Session", + session_dal: SessionDalProto, session_name: str, csc: ConnectivityServiceClient | None, - connection_server: str, - connection_port: int, - ): + connection_server: str | None, + connection_port: int | None, + ) -> None: """ Attempts to discover the controller address after booting applications. Tries dynamic lookup via connectivity service first, then falls back @@ -545,11 +655,13 @@ def _discover_controller( self.log.error(f"Could not determine controller name from OKS: {e}") top_controller_name = "Unknown-Controller" # Set a default - def get_controller_address(session_dal, session_name): + def get_controller_address( + session_dal: SessionDalProto, session_name: str + ) -> str | None: from drunc.process_manager.oks_parser import collect_variables - env = {} - collect_variables(session_dal.environment, env) + env: dict[str, str] = {} + collect_variables(session_dal.environment, env) # type: ignore[arg-type] # 1: Try dynamic lookup via Connectivity Service if csc: @@ -583,12 +695,13 @@ def get_controller_address(session_dal, session_name): f"Connectivity service lookup failed: Application '{top_controller_name}' not found." ) # Log the original failure details - self._log_controller_lookup_failure( - session_name, - top_controller_name, - connection_server, - connection_port, - ) + if connection_server is not None and connection_port is not None: + self._log_controller_lookup_failure( + session_name, + top_controller_name, + connection_server, + connection_port, + ) self.log.warning( "Falling back to static OKS configuration for address resolution." ) @@ -714,16 +827,17 @@ def get_controller_address(session_dal, session_name): ) return final_address - def keyboard_interrupt_on_sigint(signal, frame): + def keyboard_interrupt_on_sigint(_signal: object, _frame: object) -> None: self.log.warning("Interrupted") raise KeyboardInterrupt original_sigint_handler = signal.getsignal(signal.SIGINT) signal.signal(signal.SIGINT, keyboard_interrupt_on_sigint) try: - self.controller_address = get_controller_address(session_dal, session_name) + discovered_address = get_controller_address(session_dal, session_name) + self.controller_address = discovered_address except KeyboardInterrupt: - if session_dal.connectivity_service: + if session_dal.connectivity_service is not None: connection_server = session_dal.connectivity_service.host connection_port = session_dal.connectivity_service.service.port self._log_controller_interrupt( @@ -778,7 +892,9 @@ def dummy_boot( ) handle_grpc_error(e) - def _prepare_exec_and_args_dummy_boot(self, sleep: int, n_sleeps: int) -> list: + def _prepare_exec_and_args_dummy_boot( + self, sleep: int, n_sleeps: int + ) -> list[ProcessDescription.ExecAndArgs]: args = [ ProcessDescription.ExecAndArgs(exec="echo", args=["Starting dummy_boot."]) ] @@ -793,7 +909,12 @@ def _prepare_exec_and_args_dummy_boot(self, sleep: int, n_sleeps: int) -> list: return args def _build_boot_request_dummy_boot( - self, user: str, session_name: str, process: int, exec_args: list, pwd: str + self, + user: str, + session_name: str, + process: int, + exec_args: list[ProcessDescription.ExecAndArgs], + pwd: str, ) -> BootRequest: return BootRequest( token=copy_token(self.token), @@ -820,7 +941,9 @@ def terminate( request = Request(token=copy_token(self.token)) try: - response = self.stub.terminate(request, timeout=timeout) + response = cast( + ProcessInstanceList, self.stub.terminate(request, timeout=timeout) + ) except grpc.RpcError as e: try: error_details = extract_grpc_rich_error(e) @@ -840,7 +963,9 @@ def kill( request.token.CopyFrom(self.token) try: - response = self.stub.kill(request, timeout=timeout) + response = cast( + ProcessInstanceList, self.stub.kill(request, timeout=timeout) + ) except grpc.RpcError as e: try: error_details = extract_grpc_rich_error(e) @@ -858,14 +983,14 @@ def logs(self, request: LogRequest, timeout: int | float = 60) -> LogLines | Non request.token.CopyFrom(self.token) try: - response = self.stub.logs(request, timeout=timeout) + response = cast(LogLines, self.stub.logs(request, timeout=timeout)) # Check if the response indicates a BadQuery error if response.flag == ResponseFlag.NOT_EXECUTED_BAD_REQUEST_FORMAT: - lines = response.lines - if len(lines) == 1: - lines = lines[0] - self.log.warning(f"Bad query for logs: {lines}") + display_lines = ( + response.lines[0] if len(response.lines) == 1 else response.lines + ) + self.log.warning(f"Bad query for logs: {display_lines}") return None # Check for other error flags @@ -893,7 +1018,7 @@ def ps( request.token.CopyFrom(self.token) try: - response = self.stub.ps(request, timeout=timeout) + response = cast(ProcessInstanceList, self.stub.ps(request, timeout=timeout)) except grpc.RpcError as e: try: error_details = extract_grpc_rich_error(e) @@ -914,7 +1039,9 @@ def flush( request.token.CopyFrom(self.token) try: - response = self.stub.flush(request, timeout=timeout) + response = cast( + ProcessInstanceList, self.stub.flush(request, timeout=timeout) + ) except grpc.RpcError as e: try: error_details = extract_grpc_rich_error(e) @@ -935,7 +1062,9 @@ def restart( request.token.CopyFrom(self.token) try: - response = self.stub.restart(request, timeout=timeout) + response = cast( + ProcessInstanceList, self.stub.restart(request, timeout=timeout) + ) except grpc.RpcError as e: try: error_details = extract_grpc_rich_error(e) @@ -954,7 +1083,7 @@ def describe(self, timeout: int | float = 60) -> Description: request = Request(token=copy_token(self.token)) try: - response = self.stub.describe(request, timeout=timeout) + response = cast(Description, self.stub.describe(request, timeout=timeout)) except grpc.RpcError as e: try: error_details = extract_grpc_rich_error(e) @@ -972,8 +1101,12 @@ def describe(self, timeout: int | float = 60) -> Description: # ----- logging helpers ----- def _log_controller_lookup_failure( - self, session_name, top_controller_name, connection_server, connection_port - ): + self, + session_name: str, + top_controller_name: str, + connection_server: str, + connection_port: int, + ) -> None: # Logs detailed troubleshooting steps self.log.error( f""" @@ -999,8 +1132,8 @@ def _log_controller_lookup_failure( ) def _log_controller_interrupt( - self, top_controller_name, connection_server, connection_port - ): + self, top_controller_name: str, connection_server: str, connection_port: int + ) -> None: # Logs recovery instructions after user interrupts controller lookup self.log.warning( f"""This shell didn't connect to the {top_controller_name}. diff --git a/src/drunc/process_manager/ssh_process_manager.py b/src/drunc/process_manager/ssh_process_manager.py index e915ecd73..78073a212 100644 --- a/src/drunc/process_manager/ssh_process_manager.py +++ b/src/drunc/process_manager/ssh_process_manager.py @@ -1,7 +1,7 @@ import getpass import threading import uuid -from typing import List, Optional +from typing import cast from druncschema.broadcast_pb2 import BroadcastType from druncschema.process_manager_pb2 import ( @@ -18,6 +18,7 @@ from druncschema.request_response_pb2 import ResponseFlag from drunc.exceptions import DruncCommandException +from drunc.process_manager.configuration import ProcessManagerConfHandler from drunc.process_manager.process_manager import ProcessManager from drunc.processes.exit_status import ExitStatus from drunc.processes.ssh_process_lifetime_manager import ProcessLifetimeManager @@ -25,14 +26,22 @@ class SSHProcessManager(ProcessManager): def __init__( - self, configuration, LifetimeManagerClass: ProcessLifetimeManager, **kwargs - ): + self, + configuration: ProcessManagerConfHandler, + LifetimeManagerClass: type[ProcessLifetimeManager], + name: str = "process_manager", + **kwargs: object, + ) -> None: # Used to prevent races between process exit callbacks and ps/kill/flush queries self.boot_request_lock = threading.Lock() - self.ssh_lifetime_manager: Optional[ProcessLifetimeManager] = None self.session = getpass.getuser() # unfortunate - super().__init__(configuration=configuration, session=self.session, **kwargs) + super().__init__( + configuration=configuration, + name=name, + session=self.session, + **kwargs, + ) self.disable_localhost_host_key_check = False self.disable_host_key_check = False @@ -50,7 +59,7 @@ def __init__( # self.children_logs_depth = 1000 # self.children_logs = {} - self.ssh_lifetime_manager = LifetimeManagerClass( + self.ssh_lifetime_manager: ProcessLifetimeManager = LifetimeManagerClass( # type: ignore[call-arg] disable_host_key_check=self.disable_host_key_check, disable_localhost_host_key_check=self.disable_localhost_host_key_check, logger=self.log, @@ -62,8 +71,8 @@ def __init__( def _build_process_instance( self, uuid: str, - status_code, - return_code: int, + status_code: int, + return_code: int | None, ) -> ProcessInstance: """ Construct a ProcessInstance from boot request data and runtime state. @@ -93,14 +102,14 @@ def _build_process_instance( return ProcessInstance( process_description=pd, process_restriction=pr, - status_code=status_code, + status_code=cast(ProcessInstance.StatusCode.ValueType, status_code), return_code=return_code, uuid=pu, remote_pid="not available", ) - def _get_process_timeouts(self, uuids: List[str]) -> dict[str, float]: - process_timeouts = {} + def _get_process_timeouts(self, uuids: list[str]) -> dict[str, float]: + process_timeouts: dict[str, float] = {} for process_uuid in uuids: process_timeouts[process_uuid] = self.configuration.data.kill_timeout return process_timeouts @@ -108,8 +117,8 @@ def _get_process_timeouts(self, uuids: List[str]) -> dict[str, float]: def _on_ssh_process_exit( self, uuid: str, - exit_status: Optional[ExitStatus], - exception: Optional[Exception], + exit_status: ExitStatus | None, + exception: Exception | None, ) -> None: if uuid not in self.boot_request: return @@ -151,7 +160,7 @@ def _on_ssh_process_exit( exit_status=self.archived_exit_statuses[uuid], ) - def kill_processes(self, uuids: list) -> ProcessInstanceList: + def kill_processes(self, uuids: list[str]) -> ProcessInstanceList: """ Kill processes by their UUIDs. @@ -176,18 +185,21 @@ def kill_processes(self, uuids: list) -> ProcessInstanceList: self.archived_exit_statuses[proc_uuid] = exit_status # Build ProcessInstance objects from termination results - ret = [ - self._build_process_instance( - uuid=uuid, - status_code=ProcessInstance.StatusCode.DEAD, - return_code=( - exit_statuses[uuid].get_reported_exit_code() - if exit_statuses.get(uuid) is not None - else None - ), + ret: list[ProcessInstance] = [] + for proc_uuid in uuids: + exit_status = exit_statuses.get(proc_uuid) + return_code = ( + exit_status.get_reported_exit_code() + if exit_status is not None + else None + ) + ret.append( + self._build_process_instance( + uuid=proc_uuid, + status_code=ProcessInstance.StatusCode.DEAD, + return_code=return_code, + ) ) - for uuid in uuids - ] return ProcessInstanceList( name=self.name, @@ -196,18 +208,14 @@ def kill_processes(self, uuids: list) -> ProcessInstanceList: flag=ResponseFlag.EXECUTED_SUCCESSFULLY, ) - def _get_active_process_keys(self) -> list: + def _get_active_process_keys(self) -> list[str]: """ Retrieve a list of active process UUIDs managed by the SSH process manager. Returns: List of active process UUID strings """ - return ( - list(self.ssh_lifetime_manager.get_active_process_keys()) - if self.ssh_lifetime_manager is not None - else [] - ) + return list(self.ssh_lifetime_manager.get_active_process_keys()) def _terminate_impl(self) -> ProcessInstanceList: """ @@ -316,7 +324,9 @@ def _logs_impl(self, log_request: LogRequest) -> LogLines: flag=ResponseFlag.UNHANDLED_EXCEPTION_THROWN, ) - def notify_join(self, name, session, user, exit_status: ExitStatus): + def notify_join( + self, name: str, session: str, user: str, exit_status: ExitStatus + ) -> None: 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) @@ -465,7 +475,7 @@ def _ps_impl(self, query: ProcessQuery) -> ProcessInstanceList: if remote_pid_result.successful: pi.remote_pid = str(remote_pid_result.pid) else: - pi.remote_pid = remote_pid_result.reason + pi.remote_pid = remote_pid_result.reason or "not available" ret += [pi] continue @@ -488,7 +498,7 @@ def _ps_impl(self, query: ProcessQuery) -> ProcessInstanceList: if remote_pid_result.successful: pi.remote_pid = str(remote_pid_result.pid) else: - pi.remote_pid = remote_pid_result.reason + pi.remote_pid = remote_pid_result.reason or "not available" ret += [pi] return ProcessInstanceList( @@ -590,7 +600,7 @@ def _kill_impl(self, query: ProcessQuery) -> ProcessInstanceList: flag=ResponseFlag.EXECUTED_SUCCESSFULLY, ) - def _crash_processes(self, uuids: list) -> ProcessInstanceList: + def _crash_processes(self, uuids: list[str]) -> ProcessInstanceList: """ Simulate crashes for processes identified by their UUIDs. @@ -654,7 +664,7 @@ def _flush_impl(self, query: ProcessQuery) -> ProcessInstanceList: # Perform liveness checks outside the lock — these may involve SSH calls # and must not block the publish thread for extended periods. - dead_uuids = [] + dead_uuids: list[str] = [] for proc_uuid in candidate_uuids: if not self.ssh_lifetime_manager.is_process_alive(proc_uuid): dead_uuids.append(proc_uuid) @@ -663,7 +673,7 @@ def _flush_impl(self, query: ProcessQuery) -> ProcessInstanceList: f"Process {proc_uuid} is still running — skipping flush." ) - flushed = [] + flushed: list[ProcessInstance] = [] # Perform all mutations to boot_request under the lock so ps command always sees a # consistent boot_request diff --git a/src/drunc/process_manager/ssh_process_manager_paramiko_client.py b/src/drunc/process_manager/ssh_process_manager_paramiko_client.py index 0ffe00015..33a4032f9 100644 --- a/src/drunc/process_manager/ssh_process_manager_paramiko_client.py +++ b/src/drunc/process_manager/ssh_process_manager_paramiko_client.py @@ -1,13 +1,26 @@ +from typing import cast + +from drunc.process_manager.configuration import ProcessManagerConfHandler from drunc.process_manager.ssh_process_manager import SSHProcessManager +from drunc.processes.ssh_process_lifetime_manager import ProcessLifetimeManager from drunc.processes.ssh_process_lifetime_manager_paramiko import ( SSHProcessLifetimeManagerParamiko, ) class SSHProcessManagerParamikoClient(SSHProcessManager): - def __init__(self, configuration, **kwargs): + def __init__( + self, + configuration: ProcessManagerConfHandler, + name: str = "process_manager", + **kwargs: object, + ) -> None: super().__init__( configuration=configuration, - LifetimeManagerClass=SSHProcessLifetimeManagerParamiko, + LifetimeManagerClass=cast( + type[ProcessLifetimeManager], + SSHProcessLifetimeManagerParamiko, + ), + name=name, **kwargs, ) diff --git a/src/drunc/process_manager/ssh_process_manager_shell.py b/src/drunc/process_manager/ssh_process_manager_shell.py index d4e74a0ba..93f0ee8b6 100644 --- a/src/drunc/process_manager/ssh_process_manager_shell.py +++ b/src/drunc/process_manager/ssh_process_manager_shell.py @@ -1,13 +1,26 @@ +from typing import cast + +from drunc.process_manager.configuration import ProcessManagerConfHandler from drunc.process_manager.ssh_process_manager import SSHProcessManager +from drunc.processes.ssh_process_lifetime_manager import ProcessLifetimeManager from drunc.processes.ssh_process_lifetime_manager_from_forked_process import ( SSHProcessLifetimeManagerShellOnForkedProcess, ) class SSHProcessManagerShell(SSHProcessManager): - def __init__(self, configuration, **kwargs): + def __init__( + self, + configuration: ProcessManagerConfHandler, + name: str = "process_manager", + **kwargs: object, + ) -> None: super().__init__( configuration=configuration, - LifetimeManagerClass=SSHProcessLifetimeManagerShellOnForkedProcess, + LifetimeManagerClass=cast( + type[ProcessLifetimeManager], + SSHProcessLifetimeManagerShellOnForkedProcess, + ), + name=name, **kwargs, ) diff --git a/src/drunc/process_manager/utils.py b/src/drunc/process_manager/utils.py index b7d881d7e..7fa2cbbcb 100644 --- a/src/drunc/process_manager/utils.py +++ b/src/drunc/process_manager/utils.py @@ -1,7 +1,8 @@ import copy as cp import os import re -from functools import update_wrapper +from collections.abc import Callable +from typing import cast import click from druncschema.process_manager_pb2 import ( @@ -49,42 +50,39 @@ def compute_role_from_boot_request(boot_request: BootRequest) -> str: ) -def generate_process_query( - f, at_least_one: bool, all_processes_by_default: bool = False -): - @click.pass_context - def new_func(ctx, session, name, user, uuid, **kwargs): - is_trivial_query = bool( - (len(uuid) == 0) - and (session is None) - and (len(name) == 0) - and (user is None) - ) - - if is_trivial_query and at_least_one: - raise click.BadParameter( - "You need to provide at least a '--uuid', '--session', '--user' or '--name'!\nAll these values are presented with 'ps'.\nIf you want to kill everything, use 'ps' and 'kill'." - ) - - if all_processes_by_default and is_trivial_query: - name = [".*"] - - uuids = [ProcessUUID(uuid=uuid_) for uuid_ in uuid] +def build_process_query( + session: str | None, + name: tuple[str, ...], + user: str | None, + uuid: tuple[str, ...], + at_least_one: bool, + all_processes_by_default: bool = False, + crash: bool = False, +) -> ProcessQuery: + is_trivial_query = bool( + (len(uuid) == 0) and (session is None) and (len(name) == 0) and (user is None) + ) - query = ProcessQuery( - session=session, - names=name, - user=user, - uuids=uuids, - crash=kwargs.pop("crash", False), + if is_trivial_query and at_least_one: + raise click.BadParameter( + "You need to provide at least a '--uuid', '--session', '--user' or '--name'!\nAll these values are presented with 'ps'.\nIf you want to kill everything, use 'ps' and 'kill'." ) - # print(query) - return ctx.invoke(f, query=query, **kwargs) - return update_wrapper(new_func, f) + query_names = list(name) + if all_processes_by_default and is_trivial_query: + query_names = [".*"] + + uuids = [ProcessUUID(uuid=uuid_) for uuid_ in uuid] + return ProcessQuery( + session=session or "", + names=query_names, + user=user or "", + uuids=uuids, + crash=crash, + ) -def make_tree(values): +def make_tree(values: list[ProcessInstance]) -> list[str]: lines = [] for result in values: m = result.process_description.metadata @@ -95,19 +93,19 @@ def make_tree(values): return lines -def order_process_by_name(processes: list[ProcessInstance]): +def order_process_by_name(processes: list[ProcessInstance]) -> list[ProcessInstance]: """Given a list of processes, perform a tiered order by the name""" - by_session = {} + by_session: dict[str, list[ProcessInstance]] = {} for process in processes: m = process.process_description.metadata by_session.setdefault(m.session, []).append(process) - ordered = [] + ordered: list[ProcessInstance] = [] for session in sorted(by_session.keys()): session_processes = by_session[session] - node_by_id = {} - children = {} - roots = [] + node_by_id: dict[str, list[ProcessInstance]] = {} + children: dict[str, list[str]] = {} + roots: list[str] = [] for process in session_processes: tree_id = process.process_description.metadata.tree_id or "" @@ -129,11 +127,11 @@ def order_process_by_name(processes: list[ProcessInstance]): else: children.setdefault(parent_id, []).append(tree_id) - def sort_key(tree_id): + def sort_key(tree_id: str) -> tuple[str, str]: m = node_by_id[tree_id][0].process_description.metadata return (m.name, tree_id) - def walk(tree_id): + def walk(tree_id: str) -> None: ordered.extend(node_by_id[tree_id]) for child_id in sorted(children.get(tree_id, []), key=sort_key): walk(child_id) @@ -146,7 +144,7 @@ def walk(tree_id): def tabulate_process_instance_list( pil: ProcessInstanceList, title: str, long: bool = False, width: int | None = None -): +) -> Table: t = Table(title=title, width=width) t.add_column("session") t.add_column("friendly name") @@ -156,7 +154,7 @@ def tabulate_process_instance_list( t.add_column("alive") t.add_column("exit-code") - sorted_pil = order_process_by_name(pil.values) + sorted_pil = order_process_by_name(list(pil.values)) show_remote_pid = long and any( process.HasField("remote_pid") for process in sorted_pil @@ -200,7 +198,7 @@ def tabulate_process_instance_list( return t -def strip_env_for_rte(env): +def strip_env_for_rte(env: dict[str, str]) -> dict[str, str]: env_stripped = cp.deepcopy(env) for key in env.keys(): if key in [ @@ -217,7 +215,7 @@ def strip_env_for_rte(env): return env_stripped -def get_version(): +def get_version() -> str: version = os.getenv("DUNE_DAQ_BASE_RELEASE") if not version: raise RuntimeError( @@ -226,7 +224,7 @@ def get_version(): return version -def get_releases_dir(): +def get_releases_dir() -> str: releases_dir = os.getenv("SPACK_RELEASES_DIR") if not releases_dir: raise RuntimeError( @@ -235,7 +233,7 @@ def get_releases_dir(): return releases_dir -def release_or_dev(): +def release_or_dev() -> str: is_release = os.getenv("DBT_SETUP_RELEASE_SCRIPT_SOURCED") if is_release: return "rel" @@ -245,7 +243,7 @@ def release_or_dev(): return "rel" -def get_rte_script(): +def get_rte_script() -> str: script = "" if release_or_dev() == "rel": ver = get_version() @@ -254,6 +252,8 @@ def get_rte_script(): else: dbt_install_dir = os.getenv("DBT_INSTALL_DIR") + if not dbt_install_dir: + raise DruncSetupException("DBT_INSTALL_DIR is not set in the environment") script = os.path.join(dbt_install_dir, "daq_app_rte.sh") if not os.path.exists(script): @@ -266,9 +266,9 @@ def get_log_path( session_name: str, application_name: str, override_logs: bool, - app_log_path: str = None, - session_log_path: str = None, -): + app_log_path: str | None = None, + session_log_path: str | None = None, +) -> str: pwd = os.getcwd() if app_log_path == "./": app_log_path = pwd @@ -299,12 +299,12 @@ class PrCtlError(DruncException): pass -def on_parent_exit(signum): +def on_parent_exit(signum: int) -> Callable[[], None]: """Return a function to be run in a child process which will trigger SIGNAME to be sent when the parent process dies """ - def set_parent_exit_signal(): + def set_parent_exit_signal() -> None: from ctypes import cdll # http://linux.die.net/man/2/prctl @@ -351,7 +351,7 @@ def get_pm_type_from_name(pm_name: str) -> ProcessManagerTypes: log_path="./", type=conf_type, data=conf_path.split(":")[1] ) - return pmch.data.type + return cast(ProcessManagerTypes, pmch.data.type) def format_hostname(hostname: str) -> str: diff --git a/tests/process_manager/interface/test_commands.py b/tests/process_manager/interface/test_commands.py index 3571690e2..282250230 100644 --- a/tests/process_manager/interface/test_commands.py +++ b/tests/process_manager/interface/test_commands.py @@ -61,6 +61,15 @@ def mock_logger(): yield logger +@pytest.fixture(autouse=True) +def mock_pm_driver(): + with patch( + "drunc.process_manager.interface.commands._pm_driver", + side_effect=lambda obj: obj.get_driver("process_manager"), + ): + yield + + @pytest.fixture def mock_tabulate(): with patch(