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
|
import argparse
import logging
import time
import uuid
import zmq
from zmq.utils.monitor import recv_monitor_message
from parsl.addresses import get_all_addresses, tcp_url
logger = logging.getLogger(__name__)
def probe_addresses(addresses, task_port, timeout=120):
"""
Parameters
----------
addresses: [string]
List of addresses as strings
task_port: int
Task port on the interchange
timeout: int
Timeout in seconds
Returns
-------
None or string address
"""
context = zmq.Context()
addr_map = {}
for addr in addresses:
socket = context.socket(zmq.DEALER)
socket.setsockopt(zmq.LINGER, 0)
socket.setsockopt(zmq.IPV6, True)
url = tcp_url(addr, task_port)
logger.debug("Trying to connect back on {}".format(url))
socket.connect(url)
addr_map[addr] = {'sock': socket,
'mon_sock': socket.get_monitor_socket(events=zmq.EVENT_CONNECTED)}
start_t = time.time()
first_connected = None
while time.time() < start_t + timeout and not first_connected:
for addr in addr_map:
try:
recv_monitor_message(addr_map[addr]['mon_sock'], zmq.NOBLOCK)
first_connected = addr
logger.info("Connected to interchange on {}".format(first_connected))
break
except zmq.Again:
pass
time.sleep(0.01)
for addr in addr_map:
addr_map[addr]['sock'].close()
return first_connected
class TestWorker:
def __init__(self, addresses, port):
uid = str(uuid.uuid4())
self.context = zmq.Context()
self.task_incoming = self.context.socket(zmq.DEALER)
self.task_incoming.setsockopt(zmq.IDENTITY, uid.encode('utf-8'))
# Linger is set to 0, so that the manager can exit even when there might be
# messages in the pipe
self.task_incoming.setsockopt(zmq.LINGER, 0)
address = probe_addresses(addresses, port)
print("Viable address :", address)
self.task_incoming.connect(tcp_url(address, port))
def heartbeat(self):
""" Send heartbeat to the incoming task queue
"""
HEARTBEAT_CODE = (2 ** 32) - 1
heartbeat = (HEARTBEAT_CODE).to_bytes(4, "little")
r = self.task_incoming.send(heartbeat)
print("Return from heartbeat: {}".format(r))
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("-p", "--port", required=True,
help="Port to connect to")
args = parser.parse_args()
addresses = get_all_addresses()
worker = TestWorker(addresses, args.port)
worker.heartbeat()
|
