def test_succesful_pin(self): ''' This test executes a pin operation that is supposed to be successful, and it checks to make sure that that the correct metadata for execution and reporting is generated. ''' # Create a new function in the KVS. fname = 'incr' def func(_, x): return x + 1 create_function(func, self.kvs_client, fname) # Create a pin message and put it into the socket. msg = PinFunction(name=fname, response_address=self.ip) self.socket.inbox.append(msg.SerializeToString()) # Execute the pin operation. pin(self.socket, self.pusher_cache, self.kvs_client, self.status, self.pinned_functions, self.runtimes, self.exec_counts, self.user_library, False, False) # Check that the correct messages were sent and the correct metadata # created. self.assertEqual(len(self.pusher_cache.socket.outbox), 1) response = GenericResponse() response.ParseFromString(self.pusher_cache.socket.outbox[0]) self.assertTrue(response.success) self.assertEqual(func('', 1), self.pinned_functions[fname]('', 1)) self.assertTrue(fname in self.pinned_functions) self.assertTrue(fname in self.runtimes) self.assertTrue(fname in self.exec_counts) self.assertTrue(fname in self.status.functions)
def pin_function(self, dag_name, function_ref): # If there are no functions left to choose from, then we return None, # indicating that we ran out of resources to use. if len(self.unpinned_executors) == 0: return False if dag_name not in self.pending_dags: self.pending_dags[dag_name] = [] # Make a copy of the set of executors, so that we don't modify the # system's metadata. candidates = set(self.unpinned_executors) # Construct a PinFunction message to be sent to executors. pin_msg = PinFunction() pin_msg.name = function_ref.name pin_msg.response_address = self.ip serialized = pin_msg.SerializeToString() while True: # Pick a random executor from the set of candidates and attempt to # pin this function there. node, tid = sys_random.sample(candidates, 1)[0] sckt = self.pusher_cache.get(get_pin_address(node, tid)) sckt.send(serialized) response = GenericResponse() try: response.ParseFromString(self.pin_accept_socket.recv()) except zmq.ZMQError: logging.error('Pin operation to %s:%d timed out. Retrying.' % (node, tid)) continue # Do not use this executor either way: If it rejected, it has # something else pinned, and if it accepted, it has pinned what we # just asked it to pin. # In local model allow executors to have multiple functions pinned if not self.local: self.unpinned_executors.discard((node, tid)) candidates.discard((node, tid)) if response.success: # The pin operation succeeded, so we return the node and thread # ID to the caller. self.pending_dags[dag_name].append((function_ref.name, (node, tid))) return True else: # The pin operation was rejected, remove node and try again. logging.error('Node %s:%d rejected pin for %s. Retrying.' % (node, tid, function_ref.name)) continue
def test_occupied_pin(self): ''' This test attempts to pin a function onto a node where another function is already pinned. We currently only allow one pinned node per machine, so this operation should fail. ''' # Create a new function in the KVS. fname = 'incr' def func(_, x): return x + 1 create_function(func, self.kvs_client, fname) # Create a pin message and put it into the socket. msg = PinFunction(name=fname, response_address=self.ip) self.socket.inbox.append(msg.SerializeToString()) # Add an already pinned_function, so that we reject the request. self.pinned_functions['square'] = lambda _, x: x * x self.runtimes['square'] = [] self.exec_counts['square'] = [] self.status.functions.append('square') # Execute the pin operation. pin(self.socket, self.pusher_cache, self.kvs_client, self.status, self.pinned_functions, self.runtimes, self.exec_counts, self.user_library, False, False) # Check that the correct messages were sent and the correct metadata # created. self.assertEqual(len(self.pusher_cache.socket.outbox), 1) response = GenericResponse() response.ParseFromString(self.pusher_cache.socket.outbox[0]) self.assertFalse(response.success) # Make sure that none of the metadata was corrupted with this failed # pin attempt self.assertTrue(fname not in self.pinned_functions) self.assertTrue(fname not in self.runtimes) self.assertTrue(fname not in self.exec_counts) self.assertTrue(fname not in self.status.functions)
def pin_function(self, dag_name, function_ref, colocated): # If there are no functions left to choose from, then we return None, # indicating that we ran out of resources to use. if function_ref.gpu and len(self.unpinned_gpu_executors) == 0: return False elif not function_ref.gpu and len(self.unpinned_cpu_executors) == 0: return False if dag_name not in self.pending_dags: self.pending_dags[dag_name] = [] # Make a copy of the set of executors, so that we don't modify the # system's metadata. if function_ref.gpu: candidates = set(self.unpinned_gpu_executors) elif len(colocated) == 0: # If this is not a GPU function, just look at all of the unpinned # executors. candidates = set(self.unpinned_cpu_executors) else: candidates = set() already_pinned = set() for fn, thread in self.pending_dags[dag_name]: if fn in colocated: already_pinned.add((fn, thread)) candidate_nodes = set() if len(already_pinned) > 0: for fn, thread in already_pinned: candidate_nodes.add(thread[0]) # The node's IP for node, tid in self.unpinned_cpu_executors: if node in candidate_nodes: candidates.add((node, tid)) else: # If this is the first colocate to be pinned, try to assign to # an empty node. nodes = {} for node, tid in self.unpinned_cpu_executors: if node not in nodes: nodes[node] = 0 nodes[node] += 1 for node in nodes: if nodes[node] == NUM_EXECUTOR_THREADS: for i in range(NUM_EXECUTOR_THREADS): candidates.add((node, i)) if len(candidates) == 0: # There no valid executors to colocate on. return self.pin_function(dag_name, function_ref, []) # Construct a PinFunction message to be sent to executors. pin_msg = PinFunction() pin_msg.name = function_ref.name pin_msg.batching = function_ref.batching pin_msg.response_address = self.ip serialized = pin_msg.SerializeToString() while True: # Pick a random executor from the set of candidates and attempt to # pin this function there. node, tid = sys_random.sample(candidates, 1)[0] sckt = self.pusher_cache.get(get_pin_address(node, tid)) sckt.send(serialized) response = GenericResponse() try: response.ParseFromString(self.pin_accept_socket.recv()) except zmq.ZMQError: logging.error('Pin operation to %s:%d timed out. Retrying.' % (node, tid)) continue # Do not use this executor either way: If it rejected, it has # something else pinned, and if it accepted, it has pinned what we # just asked it to pin. In local mode, however we allow executors # to have multiple functions pinned. if not self.local: if function_ref.gpu: self.unpinned_gpu_executors.discard((node, tid)) candidates.discard((node, tid)) else: self.unpinned_cpu_executors.discard((node, tid)) candidates.discard((node, tid)) if response.success: # The pin operation succeeded, so we return the node and thread # ID to the caller. self.pending_dags[dag_name].append((function_ref.name, (node, tid))) return True else: # The pin operation was rejected, remove node and try again. logging.error('Node %s:%d rejected pin for %s. Retrying.' % (node, tid, function_ref.name)) continue if len(candidates) == 0 and len(colocated) > 0: # Try again without colocation. return self.pin_function(self, dag_name, function_ref, [])