File: zmq_router.py

package info (click to toggle)
python-parsl 2025.11.10%2Bds-1
  • links: PTS, VCS
  • area: main
  • in suites: forky
  • size: 12,124 kB
  • sloc: python: 24,375; makefile: 352; sh: 252; ansic: 45
file content (192 lines) | stat: -rw-r--r-- 7,826 bytes parent folder | download | duplicates (2)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
from __future__ import annotations

import logging
import multiprocessing.queues as mpq
import os
import queue
import time
from multiprocessing.context import SpawnProcess as SpawnProcessType
from multiprocessing.queues import Queue as QueueType
from multiprocessing.synchronize import Event as EventType
from typing import Tuple

import typeguard
import zmq

from parsl.addresses import tcp_url
from parsl.log_utils import set_file_logger
from parsl.monitoring.errors import MonitoringRouterStartError
from parsl.monitoring.radios.multiprocessing import MultiprocessingQueueRadioSender
from parsl.monitoring.types import TaggedMonitoringMessage
from parsl.multiprocessing import (
    SpawnEvent,
    SpawnProcess,
    SpawnQueue,
    join_terminate_close_proc,
)
from parsl.process_loggers import wrap_with_logs
from parsl.utils import setproctitle

logger = logging.getLogger(__name__)


class MonitoringRouter:

    def __init__(self,
                 *,
                 address: str,
                 port_range: Tuple[int, int] = (55050, 56000),

                 run_dir: str = ".",
                 logging_level: int = logging.INFO,
                 resource_msgs: mpq.Queue,
                 exit_event: EventType,
                 ):
        """ Initializes a monitoring configuration class.

        Parameters
        ----------
        address : str
             The ip address at which the workers will be able to reach the Hub.
        port_range : tuple(int, int)
             The MonitoringHub picks ports at random from the range which will be used by Hub.
             Default: (55050, 56000)
        run_dir : str
             Parsl log directory paths. Logs and temp files go here. Default: '.'
        logging_level : int
             Logging level as defined in the logging module. Default: logging.INFO
        resource_msgs : multiprocessing.Queue
            A multiprocessing queue to receive messages to be routed onwards to the database process
        exit_event : Event
            An event that the main Parsl process will set to signal that the monitoring router should shut down.
        """
        os.makedirs(run_dir, exist_ok=True)
        set_file_logger(f"{run_dir}/monitoring_zmq_router.log",
                        level=logging_level)
        logger.debug("Monitoring router starting")

        self.address = address

        self.loop_freq = 10.0  # milliseconds

        self._context = zmq.Context()
        self.zmq_receiver_channel = self._context.socket(zmq.DEALER)
        self.zmq_receiver_channel.setsockopt(zmq.LINGER, 0)
        self.zmq_receiver_channel.set_hwm(0)
        self.zmq_receiver_channel.RCVTIMEO = int(self.loop_freq)  # in milliseconds
        logger.debug("address: {}. port_range {}".format(address, port_range))
        self.zmq_receiver_port = self.zmq_receiver_channel.bind_to_random_port(tcp_url(address),
                                                                               min_port=port_range[0],
                                                                               max_port=port_range[1])

        self.target_radio = MultiprocessingQueueRadioSender(resource_msgs)
        self.exit_event = exit_event

    @wrap_with_logs
    def start(self) -> None:
        logger.info("Starting ZMQ listener")
        try:
            while not self.exit_event.is_set():
                try:
                    dfk_loop_start = time.time()
                    while time.time() - dfk_loop_start < 1.0:  # TODO make configurable
                        # note that nothing checks that msg really is of the annotated type
                        msg: TaggedMonitoringMessage
                        msg = self.zmq_receiver_channel.recv_pyobj()

                        assert isinstance(msg, tuple), "ZMQ Receiver expects only tuples, got {}".format(msg)
                        assert len(msg) >= 1, "ZMQ Receiver expects tuples of length at least 1, got {}".format(msg)
                        assert len(msg) == 2, "ZMQ Receiver expects message tuples of exactly length 2, got {}".format(msg)

                        self.target_radio.send(msg)
                except zmq.Again:
                    pass
                except Exception:
                    # This will catch malformed messages. What happens if the
                    # channel is broken in such a way that it always raises
                    # an exception? Looping on this would maybe be the wrong
                    # thing to do.
                    logger.warning("Failure processing a ZMQ message", exc_info=True)

            logger.info("ZMQ listener finishing normally")
        finally:
            logger.info("ZMQ listener finished")


@wrap_with_logs
@typeguard.typechecked
def zmq_router_starter(*,
                       comm_q: mpq.Queue,
                       resource_msgs: mpq.Queue,
                       exit_event: EventType,

                       address: str,
                       port_range: Tuple[int, int],

                       run_dir: str,
                       logging_level: int) -> None:
    setproctitle("parsl: monitoring zmq router")
    try:
        router = MonitoringRouter(address=address,
                                  port_range=port_range,
                                  run_dir=run_dir,
                                  logging_level=logging_level,
                                  resource_msgs=resource_msgs,
                                  exit_event=exit_event)
    except Exception as e:
        logger.error("MonitoringRouter construction failed.", exc_info=True)
        comm_q.put(f"Monitoring router construction failed: {e}")
    else:
        comm_q.put(router.zmq_receiver_port)
        router.start()


class ZMQRadioReceiver():
    def __init__(self, *, process: SpawnProcessType, exit_event: EventType, port: int) -> None:
        self.process = process
        self.exit_event = exit_event
        self.port = port

    def close(self) -> None:
        self.exit_event.set()
        join_terminate_close_proc(self.process)


def start_zmq_receiver(*,
                       monitoring_messages: QueueType,
                       loopback_address: str,
                       port_range: Tuple[int, int],
                       logdir: str,
                       worker_debug: bool) -> ZMQRadioReceiver:
    comm_q = SpawnQueue(maxsize=10)

    router_exit_event = SpawnEvent()

    router_proc = SpawnProcess(target=zmq_router_starter,
                               kwargs={"comm_q": comm_q,
                                       "resource_msgs": monitoring_messages,
                                       "exit_event": router_exit_event,
                                       "address": loopback_address,
                                       "port_range": port_range,
                                       "run_dir": logdir,
                                       "logging_level": logging.DEBUG if worker_debug else logging.INFO,
                                       },
                               name="Monitoring-ZMQ-Router-Process",
                               daemon=True,
                               )
    router_proc.start()

    try:
        logger.debug("Waiting for router process to report port")
        comm_q_result = comm_q.get(block=True, timeout=120)
        comm_q.close()
        comm_q.join_thread()
    except queue.Empty:
        logger.error("Monitoring ZMQ Router has not reported port in 120s")
        raise MonitoringRouterStartError()

    if isinstance(comm_q_result, str):
        logger.error("MonitoringRouter sent an error message: %s", comm_q_result)
        raise RuntimeError(f"MonitoringRouter failed to start: {comm_q_result}")

    return ZMQRadioReceiver(process=router_proc, exit_event=router_exit_event, port=comm_q_result)