def _get_func_list(self, prefix=None):
msg = prefix if prefix else ''
self.list_sock.send_string(msg)
flist = StringSet()
flist.ParseFromString(self.list_sock.recv())
return flist.keys
def update(self):
# Periodically clean up the running counts map to drop any times older
# than 5 seconds.
for executor in self.running_counts:
new_set = set()
for ts in self.running_counts[executor]:
if time.time() - ts < 5:
new_set.add(ts)
self.running_counts[executor] = new_set
# Clean up any backoff messages that were added more than 5 seconds ago
# -- this should be enough to drain a queue.
remove_set = set()
for executor in self.backoff:
if time.time() - self.backoff[executor] > 5:
remove_set.add(executor)
for executor in remove_set:
del self.backoff[executor]
executors = set(map(lambda status: status.ip,
self.thread_statuses.values()))
# Update the sets of keys that are being cached at each IP address.
self.key_locations.clear()
for ip in executors:
key = get_cache_ip_key(ip)
# This is of type LWWPairLattice, which has a StringSet protobuf
# packed into it; we want the keys in that StringSet protobuf.
lattice = self.kvs_client.get(key)[key]
if lattice is None:
# We will only get None if this executor is still joining; if
# so, we just ignore this for now and move on.
continue
st = StringSet()
st.ParseFromString(lattice.reveal())
for key in st.keys:
if key not in self.key_locations:
self.key_locations[key] = []
self.key_locations[key].append(ip)
def get_ip_set(request_ip, socket_cache, exec_threads=True):
sckt = socket_cache.get(request_ip)
# we can send an empty request because the response is always thes same
sckt.send(b'')
ips = StringSet()
ips.ParseFromString(sckt.recv())
result = set()
if exec_threads:
for ip in ips.keys:
for i in range(NUM_EXEC_THREADS):
result.add((ip, i))
return result
else:
return set(ips.keys)
def get_ip_set(management_request_socket, exec_threads=True):
# we can send an empty request because the response is always the same
management_request_socket.send(b'')
try:
ips = StringSet()
ips.ParseFromString(management_request_socket.recv())
result = set()
if exec_threads:
for ip in ips.keys:
for i in range(NUM_EXEC_THREADS):
result.add((ip, i))
return result
else:
return set(ips.keys)
except zmq.ZMQError as e:
if e.errno == zmq.EAGAIN:
return None
else:
raise e