def test_delete_dag(self):
'''
We attempt to delete a DAG that has already been created and check to
ensure that the correct unpin messages are sent to executors and that
the metadata is updated appropriately.
'''
# Create a simple two fucntion DAG and add it to the system metadata.
source = 'source'
sink = 'sink'
dag_name = 'dag'
dag = create_linear_dag([None, None], [source, sink], self.kvs_client,
dag_name)
dags = {}
call_frequency = {}
dags[dag.name] = (dag, {source})
call_frequency[source] = 100
call_frequency[sink] = 100
# Add the correct metadata to the policy engine.
source_location = (self.ip, 1)
sink_location = (self.ip, 2)
self.policy.function_locations[source] = {source_location}
self.policy.function_locations[sink] = {sink_location}
self.socket.inbox.append(dag.name)
# Attempt to delete the DAG.
delete_dag(self.socket, dags, self.policy, call_frequency)
# Check that the correct unpin messages were sent.
messages = self.pusher_cache.socket.outbox
self.assertEqual(len(messages), 2)
self.assertEqual(messages[0], source)
self.assertEqual(messages[1], sink)
addresses = self.pusher_cache.addresses
self.assertEqual(len(addresses), 2)
self.assertEqual(addresses[0], get_unpin_address(*source_location))
self.assertEqual(addresses[1], get_unpin_address(*sink_location))
# Check that the server metadata was updated correctly.
self.assertEqual(len(dags), 0)
self.assertEqual(len(call_frequency), 0)
# Check that the correct message was sent to the user.
self.assertEqual(len(self.socket.outbox), 1)
response = GenericResponse()
response.ParseFromString(self.socket.outbox.pop())
self.assertTrue(response.success)
def discard_dag(self, dag, pending=False):
pinned_locations = []
if pending:
if dag.name in self.pending_dags:
# If the DAG was pending, we can simply look at the sequestered
# pending metadata.
pinned_locations = list(self.pending_dags[dag.name])
del self.pending_dags[dag.name]
else:
# If the DAG was not pinned, we construct a set of all the
# locations where functions were pinned for this DAG.
for function_ref in dag.functions:
for location in self.function_locations[function_ref.name]:
pinned_locations.append((function_ref.name, location))
# For each location, we fire-and-forget an unpin message.
for function_name, location in pinned_locations:
ip, tid = location
sckt = self.pusher_cache.get(get_unpin_address(ip, tid))
sckt.send_string(function_name)
def test_create_dag_insufficient_resources(self):
'''
This test attempts to create a DAG even though there are not enough
free executors in the system. It checks that a pin message is attempted
to be sent, we run out of resources, and then the request is rejected.
We check that the metadata is properly restored back to its original
state.
'''
# Create a simple two-function DAG and add it to the inbound socket.
source = 'source'
sink = 'sink'
dag_name = 'dag'
dag = create_linear_dag([None, None], [source, sink], self.kvs_client,
dag_name)
self.socket.inbox.append(dag.SerializeToString())
# Add relevant metadata to the policy engine, but set the number of
# executors to fewer than needed.
address_set = {(self.ip, 1)}
self.policy.unpinned_executors.update(address_set)
# Prepopulate the pin_accept socket with sufficient success messages.
self.pin_socket.inbox.append(sutils.ok_resp)
# Attempt to create the DAG.
dags = {}
call_frequency = {}
create_dag(self.socket, self.pusher_cache, self.kvs_client, dags,
self.policy, call_frequency)
# Check that an error was returned to the user.
self.assertEqual(len(self.socket.outbox), 1)
response = GenericResponse()
response.ParseFromString(self.socket.outbox[0])
self.assertFalse(response.success)
self.assertEqual(response.error, NO_RESOURCES)
# Test that the correct pin messages were sent.
self.assertEqual(len(self.pusher_cache.socket.outbox), 2)
messages = self.pusher_cache.socket.outbox
# Checks for the pin message.
self.assertTrue(':' in messages[0])
ip, fname = messages[0].split(':')
self.assertEqual(ip, self.ip)
self.assertEqual(source, fname)
# Checks for the unpin message.
self.assertEqual(messages[1], source)
address = random.sample(address_set, 1)[0]
addresses = self.pusher_cache.addresses
self.assertEqual(get_pin_address(*address), addresses[0])
self.assertEqual(get_unpin_address(*address), addresses[1])
# Check that no additional messages were sent.
self.assertEqual(len(self.policy.unpinned_executors), 0)
self.assertEqual(len(self.policy.function_locations), 0)
self.assertEqual(len(self.policy.pending_dags), 0)
# Check that no additional metadata was created or sent.
self.assertEqual(len(call_frequency), 0)
self.assertEqual(len(dags), 0)
