mccode-plumber 0.6.0__py3-none-any.whl → 0.7.0__py3-none-any.whl

mccode_plumber/file_writer_control/CommandChannel.py

@@ -0,0 +1,236 @@
+import atexit
+import threading
+from datetime import datetime
+from queue import Queue
+from typing import Dict, List, Optional, Union
+
+from kafka import KafkaConsumer
+from kafka.errors import NoBrokersAvailable
+
+from file_writer_control.CommandStatus import CommandState, CommandStatus
+from file_writer_control.InThreadStatusTracker import (
+    DEAD_ENTITY_TIME_LIMIT,
+    InThreadStatusTracker,
+)
+from file_writer_control.JobStatus import JobStatus
+from file_writer_control.KafkaTopicUrl import KafkaTopicUrl
+from file_writer_control.WorkerStatus import WorkerStatus
+
+
+def thread_function(
+    host_port: str,
+    topic: str,
+    in_queue: Queue,
+    out_queue: Queue,
+    kafka_config: Dict[str, str] = {},
+):
+    """
+    Background thread for consuming Kafka messages.
+    :param host_port: The host + port of the Kafka broker that we are using.
+    :param topic: The Kafka topic that we are listening to.
+    :param in_queue: A queue for sending "exit" messages to the thread.
+    .. note:: The queue will exit upon the reception of the string "exit" on this queue.
+    :param out_queue: The queue to which status updates are published.
+    """
+    status_tracker = InThreadStatusTracker(out_queue)
+    while True:
+        try:
+            consumer = KafkaConsumer(
+                topic,
+                bootstrap_servers=host_port,
+                fetch_max_bytes=52428800 * 6,
+                max_partition_fetch_bytes=52428800 * 10,
+                consumer_timeout_ms=100,
+                **kafka_config
+            )  # Roughly 300MB
+            break
+        except NoBrokersAvailable:
+            pass  # Do not fail if the broker is not immediately available.
+        if not in_queue.empty():
+            new_msg = in_queue.get()
+            if new_msg == "exit":
+                return
+    while True:
+        for message in consumer:
+            status_tracker.process_message(message.value)
+        status_tracker.check_for_lost_connections()
+        if not in_queue.empty():
+            new_msg = in_queue.get()
+            if new_msg == "exit":
+                break
+    consumer.close(True)
+
+
+class CommandChannel(object):
+    """
+    A class that implements the functionality for receiving and interpreting messages that are published to the
+    Kafka command topic of a pool of file-writers.
+    .. note:: This class implements a thread that will continuously attempt to connect to a Kafka broker.
+    """
+
+    def __init__(self, command_topic_url: str, kafka_config: Dict[str, str] = {}):
+        """
+        Constructor.
+        :param command_topic_url: The url of the Kafka topic to where the file-writer status/command messages are published.
+        """
+        kafka_address = KafkaTopicUrl(command_topic_url)
+        self.status_queue = Queue()
+        self.to_thread_queue = Queue()
+        thread_kwargs = {
+            "host_port": kafka_address.host_port,
+            "topic": kafka_address.topic,
+            "in_queue": self.to_thread_queue,
+            "out_queue": self.status_queue,
+            "kafka_config": kafka_config,
+        }
+        self.map_of_workers: Dict[str, WorkerStatus] = {}
+        self.map_of_jobs: Dict[str, JobStatus] = {}
+        self.map_of_commands: Dict[str, CommandStatus] = {}
+        self.run_thread = True
+        self.thread = threading.Thread(
+            target=thread_function, daemon=True, kwargs=thread_kwargs
+        )
+        self.thread.start()
+
+        def do_exit():
+            self.stop_thread()
+
+        atexit.register(do_exit)
+
+    def add_job_id(self, job_id: str):
+        """
+        Add a job identifier to the list of known jobs before it has been encountered on the command topic.
+        :param job_id: The identifier of the new job.
+        """
+        if job_id not in self.map_of_jobs:
+            self.map_of_jobs[job_id] = JobStatus(job_id)
+
+    def add_command_id(self, job_id: str, command_id: str):
+        """
+        Add a command identifier to the list of known commands before it has been encountered on the command topic.
+        :param job_id: The job identifier of the new command.
+        :param command_id: The identifier of the new command.
+        """
+        if command_id not in self.map_of_commands:
+            self.map_of_commands[command_id] = CommandStatus(job_id, command_id)
+            self.map_of_commands[command_id].state = CommandState.WAITING_RESPONSE
+
+    def stop_thread(self):
+        """
+        Stop the thread that is continuously getting command topic messages in the background. Should only be called if
+        we are about to get rid of the current instance of CommandChannel.
+        """
+        self.to_thread_queue.put("exit")
+        try:
+            self.thread.join()
+        except RuntimeError:
+            pass  # Do not throw an exception if the thread has not yet been started.
+
+    def __del__(self):
+        self.stop_thread()
+
+    def update_workers(self, current_time: Optional[datetime] = None):
+        """
+        Update the list of known workers, jobs and commands. This is a non-blocking call but it might take some time
+        to execute if the queue of updates is long. This member function is called by many of the other member functions
+        in this class.
+        """
+        if current_time is None:
+            current_time = datetime.now()
+
+        def handle_worker_status(status_update):
+            if status_update.service_id not in self.map_of_workers:
+                self.map_of_workers[status_update.service_id] = status_update
+            self.map_of_workers[status_update.service_id].update_status(status_update)
+
+        def handle_job_status(status_update):
+            if status_update.job_id not in self.map_of_jobs:
+                self.map_of_jobs[status_update.job_id] = status_update
+            self.map_of_jobs[status_update.job_id].update_status(status_update)
+
+        def handle_command_status(status_update):
+            if status_update.command_id not in self.map_of_commands:
+                self.map_of_commands[status_update.command_id] = status_update
+            self.map_of_commands[status_update.command_id].update_status(status_update)
+
+        status_updater_map = {
+            WorkerStatus: handle_worker_status,
+            CommandStatus: handle_command_status,
+            JobStatus: handle_job_status,
+        }
+        while not self.status_queue.empty():
+            status_update = self.status_queue.get()
+            status_updater_map[type(status_update)](status_update)
+
+        for entity in (
+            list(self.map_of_workers.values())
+            + list(self.map_of_commands.values())
+            + list(self.map_of_jobs.values())
+        ):
+            entity.check_if_outdated(current_time)
+
+        def pruner(entities_dictionary):
+            for key in list(entities_dictionary.keys()):
+                if (
+                    entities_dictionary[key].last_update + DEAD_ENTITY_TIME_LIMIT
+                    < current_time
+                ):
+                    del entities_dictionary[key]
+
+        pruner(self.map_of_commands)
+        pruner(self.map_of_workers)
+        pruner(self.map_of_jobs)
+
+    def list_workers(self) -> List[WorkerStatus]:
+        """
+        :return: A list of the (known) workers with state and status information.
+        """
+        self.update_workers()
+        return list(self.map_of_workers.values())
+
+    def list_jobs(self) -> List[JobStatus]:
+        """
+        :return: A list of the (known) jobs with state and status information.
+        """
+        self.update_workers()
+        return list(self.map_of_jobs.values())
+
+    def list_commands(self) -> List[CommandStatus]:
+        """
+        :return: A list of the (known) commands and their outcomes.
+        """
+        self.update_workers()
+        return list(self.map_of_commands.values())
+
+    def get_job(self, job_id: str) -> Union[JobStatus, None]:
+        """
+        Get the status of a single job.
+        :param job_id: The job identifier of the job we are interested in.
+        :return: The job status or None if the job is not known.
+        """
+        self.update_workers()
+        if job_id in self.map_of_jobs:
+            return self.map_of_jobs[job_id]
+        return None
+
+    def get_worker(self, service_id: str) -> Union[WorkerStatus, None]:
+        """
+        Get the status of a single worker.
+        :param service_id: The service identifier of the worker we are interested in.
+        :return: The worker status or None if the service id is not known.
+        """
+        self.update_workers()
+        if service_id in self.map_of_workers:
+            return self.map_of_workers[service_id]
+        return None
+
+    def get_command(self, command_id: str) -> Union[CommandStatus, None]:
+        """
+        Get the status of a single command.
+        :param command_id: The command identifier of the command we are interested in.
+        :return: The command status/outcome or None if the command is not known.
+        """
+        self.update_workers()
+        if command_id in self.map_of_commands:
+            return self.map_of_commands[command_id]
+        return None

mccode_plumber/file_writer_control/CommandHandler.py

@@ -0,0 +1,58 @@
+from datetime import timedelta
+
+from file_writer_control.CommandChannel import CommandChannel
+from file_writer_control.CommandStatus import CommandState
+
+
+class CommandHandler:
+    """
+    A stand in for (more easily) checking the state of a command sent to a file-writer.
+    """
+
+    def __init__(self, command_channel: CommandChannel, command_id: str):
+        """
+        Constructor.
+        :param command_channel: The instance of a CommandChannel that this class uses for getting the command status from.
+        :param command_id: The (unique) command identifier.
+        """
+        self.command_id = command_id
+        self.command_channel = command_channel
+
+    def get_state(self) -> CommandState:
+        """
+        Get the command state.
+        :return: The state of the command if possible. CommandState.UNKNOWN if the the state can not be determined.
+        """
+        command = self.command_channel.get_command(self.command_id)
+        if command is None:
+            return CommandState.UNKNOWN
+        return command.state
+
+    def is_done(self) -> bool:
+        """
+        :return: True if the command completed successfully. False otherwise.
+        """
+        current_state = self.command_channel.get_command(self.command_id).state
+        if current_state == CommandState.ERROR:
+            raise RuntimeError(
+                f'Command failed with error message "{self.get_message()}".'
+            )
+        if current_state == CommandState.TIMEOUT_RESPONSE:
+            raise RuntimeError("Timed out while trying to send command.")
+        return current_state == CommandState.SUCCESS
+
+    def get_message(self) -> str:
+        """
+        :return: If there was an error executing the command, this member function will return the error string as
+        sent by the file-writer. Will return an empty string otherwise.
+        """
+        command = self.command_channel.get_command(self.command_id)
+        if command is None:
+            return ""
+        return command.message
+
+    def set_timeout(self, new_timeout: timedelta):
+        self.command_channel.get_command(self.command_id).timeout = new_timeout
+
+    def get_timeout(self):
+        return self.command_channel.get_command(self.command_id).timeout

mccode_plumber/file_writer_control/CommandStatus.py

@@ -0,0 +1,151 @@
+from datetime import datetime, timedelta
+from enum import Enum, auto
+from typing import Optional
+
+COMMAND_STATUS_TIMEOUT = timedelta(seconds=60)
+
+
+class CommandState(Enum):
+    """
+    The state of a command.
+    """
+
+    UNKNOWN = auto()
+    NO_COMMAND = auto()
+    WAITING_RESPONSE = auto()
+    TIMEOUT_RESPONSE = auto()
+    ERROR = auto()
+    SUCCESS = auto()
+
+
+class CommandStatus(object):
+    """
+    The status of a command.
+    """
+
+    def __init__(
+        self,
+        job_id: str,
+        command_id: str,
+        command_timeout: timedelta = COMMAND_STATUS_TIMEOUT,
+    ):
+        self._command_timeout = command_timeout
+        self._job_id = job_id
+        self._command_id = command_id
+        self._last_update = datetime.now()
+        self._state = CommandState.NO_COMMAND
+        self._message = ""
+        self._response_code = None
+
+    def __eq__(self, other_status: "CommandStatus"):
+        if not isinstance(other_status, CommandStatus):
+            raise NotImplementedError
+        return (
+            other_status.command_id == self.command_id
+            and other_status.job_id == self.job_id
+            and other_status.state == self.state
+            and other_status.response_code == self.response_code
+        )
+
+    def update_status(self, new_status: "CommandStatus"):
+        """
+        Updates the status/state of this instance of the CommandStatus class using another instance.
+        .. note:: The command identifier of both this instance and the other one must be identical.
+        :param new_status: The other instance of the CommandStatus class.
+        """
+        if new_status.command_id != self.command_id:
+            raise RuntimeError(
+                f"Command id of status update is not correct ({self.command_id} vs {new_status.command_id})"
+            )
+        self._state = new_status.state
+        self._response_code = new_status.response_code
+        if new_status.message:
+            self._message = new_status.message
+        self._last_update = new_status.last_update
+
+    def check_if_outdated(self, current_time: datetime):
+        """
+        Given the current time, state and the time of the last update: Have we lost the connection?
+        :param current_time: The current time
+        """
+        if (
+            self.state != CommandState.SUCCESS
+            and self.state != CommandState.ERROR
+            and self.state != CommandState.TIMEOUT_RESPONSE
+            and current_time - self.last_update > self._command_timeout
+        ):
+            self._state = CommandState.TIMEOUT_RESPONSE
+            self._last_update = current_time
+
+    @property
+    def response_code(self) -> Optional[int]:
+        """
+        A code that mirrors the result of a command if a response has been received. Is set to None otherwise.
+        """
+        return self._response_code
+
+    @response_code.setter
+    def response_code(self, new_code: int):
+        """
+        Set the current response code.
+        """
+        self._response_code = new_code
+        self._last_update = datetime.now()
+
+    @property
+    def job_id(self) -> str:
+        """
+        The job identifier under which this command is executed.
+        """
+        return self._job_id
+
+    @property
+    def command_id(self) -> str:
+        """
+        The unique command identifier of this command.
+        """
+        return self._command_id
+
+    @property
+    def message(self) -> str:
+        """
+        A status/error message as returned by the file-writer that responded to the command.
+        :return:
+        """
+        return self._message
+
+    @message.setter
+    def message(self, new_message: str):
+        if new_message:
+            self._message = new_message
+            self._last_update = datetime.now()
+
+    @property
+    def state(self) -> CommandState:
+        """
+        The current state of the command.
+        """
+        return self._state
+
+    @state.setter
+    def state(self, new_state: CommandState):
+        self._state = new_state
+        self._last_update = datetime.now()
+
+    @property
+    def last_update(self) -> datetime:
+        """
+        The local time stamp of the last update of the status of the command.
+        """
+        return self._last_update
+
+    @property
+    def timeout(self) -> timedelta:
+        """
+        Timeout for waiting for response to command.
+        """
+        return self._command_timeout
+
+    @timeout.setter
+    def timeout(self, new_timeout: timedelta):
+        self._command_timeout = new_timeout

mccode_plumber/file_writer_control/InThreadStatusTracker.py

@@ -0,0 +1,228 @@
+import json
+from datetime import datetime, timedelta
+from queue import Queue
+from typing import Dict
+
+from streaming_data_types import deserialise_6s4t as deserialise_stop_time
+from streaming_data_types import deserialise_answ as deserialise_answer
+from streaming_data_types import deserialise_pl72 as deserialise_start
+from streaming_data_types import deserialise_wrdn as deserialise_stopped
+from streaming_data_types import deserialise_x5f2 as deserialise_status
+from streaming_data_types.action_response_answ import FILE_IDENTIFIER as ANSW_IDENTIFIER
+from streaming_data_types.action_response_answ import Response
+from streaming_data_types.finished_writing_wrdn import (
+    FILE_IDENTIFIER as STOPPED_IDENTIFIER,
+)
+from streaming_data_types.finished_writing_wrdn import WritingFinished
+from streaming_data_types.run_start_pl72 import FILE_IDENTIFIER as START_IDENTIFIER
+from streaming_data_types.run_start_pl72 import RunStartInfo
+from streaming_data_types.run_stop_6s4t import FILE_IDENTIFIER as STOP_TIME_IDENTIFIER
+from streaming_data_types.run_stop_6s4t import RunStopInfo
+from streaming_data_types.status_x5f2 import FILE_IDENTIFIER as STAT_IDENTIFIER
+from streaming_data_types.status_x5f2 import StatusMessage
+from streaming_data_types.utils import get_schema
+
+from file_writer_control.CommandStatus import CommandState, CommandStatus
+from file_writer_control.JobStatus import JobState, JobStatus
+from file_writer_control.StateExtractor import (
+    extract_job_state_from_answer,
+    extract_state_from_command_answer,
+    extract_worker_state_from_status,
+)
+from file_writer_control.WorkerStatus import WorkerState, WorkerStatus
+
+DEAD_ENTITY_TIME_LIMIT = timedelta(hours=1)
+
+
+class InThreadStatusTracker:
+    """
+    Implements de-coding of flatbuffer messages and sends updates of worker, job and command state/status back to the
+    "main"-thread if there has been changes.
+    """
+
+    def __init__(self, status_queue: Queue):
+        """
+        Constructor.
+        :param status_queue: The output queue to which state/status updates are pushed.
+        """
+        self.queue = status_queue
+        self.known_workers: Dict[str, WorkerStatus] = {}
+        self.known_jobs: Dict[str, JobStatus] = {}
+        self.known_commands: Dict[str, CommandStatus] = {}
+
+    def process_message(self, message: bytes):
+        """
+        Process a binary message.
+        :param message: The binary message to be processed.
+        """
+        current_schema = get_schema(message).encode("utf-8")
+        update_time = datetime.now()
+        msg_process_map = {
+            ANSW_IDENTIFIER: lambda msg: self.process_answer(deserialise_answer(msg)),
+            STAT_IDENTIFIER: lambda msg: self.process_status(deserialise_status(msg)),
+            STOP_TIME_IDENTIFIER: lambda msg: self.process_set_stop_time(
+                deserialise_stop_time(msg)
+            ),
+            START_IDENTIFIER: lambda msg: self.process_start(deserialise_start(msg)),
+            STOPPED_IDENTIFIER: lambda msg: self.process_stopped(
+                deserialise_stopped(msg)
+            ),
+        }
+        if current_schema in msg_process_map:
+            msg_process_map[current_schema](message)
+
+        self.send_status_if_updated(update_time)
+
+    def send_status_if_updated(self, limit_time: datetime):
+        """
+        Sends status updates of workers, jobs and commands (to the status queue) if there has been any updates on or
+        after the limit_time.
+        :param limit_time: The cut-off time for deciding which updates should be sent to the status queue.
+        """
+        for entity in (
+            list(self.known_workers.values())
+            + list(self.known_jobs.values())
+            + list(self.known_commands.values())
+        ):
+            if entity.last_update >= limit_time:
+                self.queue.put(entity)
+
+    def check_for_worker_presence(self, service_id: str):
+        """
+        Check if a service_id is known and add it to a list of known ones if it is not.
+        :param service_id: The service identifier to look for.
+        """
+        if service_id not in self.known_workers:
+            self.known_workers[service_id] = WorkerStatus(service_id)
+
+    def check_for_job_presence(self, job_id: str):
+        """
+        Check if a job identifier is known and add it to a list of known ones if it is not.
+        :param job_id: The job identifier to look for.
+        """
+        if job_id not in self.known_jobs:
+            new_job = JobStatus(job_id)
+            self.known_jobs[job_id] = new_job
+
+    def check_for_command_presence(self, job_id: str, command_id: str):
+        """
+        Check if a command identifier is known and add it to a list of known ones if it is not.
+        :param job_id: The job identifier of the command that we are looking for. (Only used if we need to add the
+        command.)
+        :param command_id: The command identifier to look for.
+        """
+        if command_id not in self.known_commands:
+            new_command = CommandStatus(job_id, command_id)
+            self.known_commands[command_id] = new_command
+
+    def check_for_lost_connections(self):
+        """
+        Check workers, commands and jobs for the last update time and change the state of these if a timeout has been
+        reached.
+        """
+        now = datetime.now()
+        for entity in (
+            list(self.known_workers.values())
+            + list(self.known_jobs.values())
+            + list(self.known_commands.values())
+        ):
+            entity.check_if_outdated(now)
+        self.send_status_if_updated(now)
+
+    def prune_dead_entities(self, current_time: datetime):
+        """
+        Will remove old jobs, workers and commands that have not been updated recently.
+        :return:
+        """
+
+        def pruner(entities_dictionary):
+            for key in list(entities_dictionary.keys()):
+                if (
+                    entities_dictionary[key].last_update + DEAD_ENTITY_TIME_LIMIT
+                    < current_time
+                ):
+                    del entities_dictionary[key]
+
+        pruner(self.known_workers)
+        pruner(self.known_commands)
+        pruner(self.known_jobs)
+
+    def process_answer(self, answer: Response):
+        """
+        Update workers, jobs and commands based on information in a response message.
+        :param answer: The response/answer message to use for status updates.
+        """
+        self.check_for_worker_presence(answer.service_id)
+        self.check_for_job_presence(answer.job_id)
+        self.check_for_command_presence(answer.job_id, answer.command_id)
+        new_job_state = extract_job_state_from_answer(answer)
+        if new_job_state is not None:
+            self.known_jobs[answer.job_id].state = new_job_state
+        current_command = self.known_commands[answer.command_id]
+        current_command.state = extract_state_from_command_answer(answer)
+        current_command.message = answer.message
+        current_command.response_code = Response.status_code
+        self.known_jobs[answer.job_id].message = answer.message
+        try:
+            self.known_jobs[answer.job_id].service_id = answer.service_id
+        except RuntimeError:
+            pass
+
+    def process_status(self, status_update: StatusMessage):
+        """
+        Update workers and jobs based on information in a status message.
+        :param status_update: The status message to use for updates.
+        """
+        self.check_for_worker_presence(status_update.service_id)
+        current_state = extract_worker_state_from_status(status_update)
+        self.known_workers[status_update.service_id].state = current_state
+        if current_state == WorkerState.WRITING:
+            try:
+                json_data = json.loads(status_update.status_json)
+                job_id = json_data["job_id"]
+            except (KeyError, json.JSONDecodeError):
+                raise RuntimeError("Unable to extract JSON data from status message.")
+            file_name = json_data["file_being_written"]
+            self.check_for_job_presence(job_id)
+            self.known_jobs[job_id].state = JobState.WRITING
+            self.known_jobs[job_id].file_name = file_name
+            self.known_jobs[job_id].metadata = json_data
+            # For some jobs, we will only know the service-id when a worker starts working on a job.
+            # Thus we need the following statement to update the (known) service-id of a job.
+            try:
+                self.known_jobs[job_id].service_id = status_update.service_id
+            except RuntimeError:
+                pass  # Expected error (i.e. the job is not known), do nothing
+
+    def process_set_stop_time(self, stop_time: RunStopInfo):
+        """
+        Update commands and jobs based on information in a "set stop time" message.
+        :param stop_time: The "stop" message to use for updates.
+        """
+        self.check_for_command_presence(stop_time.job_id, stop_time.command_id)
+        self.known_commands[stop_time.command_id].state = CommandState.WAITING_RESPONSE
+
+    def process_start(self, start: RunStartInfo):
+        """
+        Update commands and jobs based on information in a "start" message.
+        :param start: The "start" message to use for updates.
+        """
+        self.check_for_job_presence(start.job_id)
+        self.check_for_command_presence(start.job_id, start.job_id)
+        self.known_commands[start.job_id].state = CommandState.WAITING_RESPONSE
+
+    def process_stopped(self, stopped: WritingFinished):
+        """
+        Update workers and jobs based on information in a "has stopped" message.
+        :param stopped: The "stopped" message to use for updates.
+        """
+        self.check_for_job_presence(stopped.job_id)
+        self.check_for_worker_presence(stopped.service_id)
+        current_job = self.known_jobs[stopped.job_id]
+        if stopped.error_encountered:
+            current_job.state = JobState.ERROR
+        else:
+            current_job.state = JobState.DONE
+            current_job.metadata = json.loads(stopped.metadata)
+        current_job.message = stopped.message
+        self.known_workers[stopped.service_id].state = WorkerState.IDLE