def test_exec_dag_sink(self):
'''
Tests that the last function in a DAG executes correctly and stores the
result in the KVS.
'''
def func(_, x):
return x * x
fname = 'square'
arg = 2
dag = create_linear_dag([func], [fname], self.kvs_client, 'dag')
schedule, triggers = self._create_fn_schedule(dag, arg, fname, [fname])
exec_dag_function(self.pusher_cache, self.kvs_client, triggers, func,
schedule, self.user_library, {}, {})
# Assert that there have been 0 messages sent.
self.assertEqual(len(self.socket.outbox), 0)
# Retrieve the result, ensure it is a LWWPairLattice, then deserialize
# it.
result = self.kvs_client.get(schedule.id)[schedule.id]
self.assertEqual(type(result), LWWPairLattice)
result = serializer.load_lattice(result)
# Check that the output is equal to a local function execution.
self.assertEqual(result, func('', arg))
def test_create_gpu_dag_no_resources(self):
# Create a simple two-function DAG and add it to the inbound socket.
dag_name = 'dag'
dag = create_linear_dag([None], ['fn'], self.kvs_client, dag_name)
dag.functions[0].gpu = True
self.socket.inbox.append(dag.SerializeToString())
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), 0)
# Check that no additional messages were sent.
self.assertEqual(len(self.policy.unpinned_cpu_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)
def test_dag_exec_multi_input_true(self):
def incr(_, x):
x + 1
iname = 'incr'
def square(_, x):
return x * x
sname = 'square'
arg = 2
# Create a DAG and a trigger for the first function in the DAG.
dag = create_linear_dag([incr, square], [iname, sname],
self.kvs_client, 'dag')
ref = dag.functions[1] if dag.functions[1].name == sname else \
dag.functions[0]
ref.type = MULTIEXEC
ref.invalid_results.append(serializer.dump(5))
schedule, triggers = self._create_fn_schedule(dag, arg, sname,
[iname, sname])
result = exec_dag_function(self.pusher_cache, self.kvs_client,
[triggers], square, [schedule],
self.user_library, {}, {}, [], False)
self.assertTrue(result)
data = self.kvs_client.get(schedule.id)
self.assertEqual(serializer.load_lattice(data[schedule.id]),
square(None, arg))
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 test_exec_causal_dag_sink(self):
'''
Tests that the last function in a causal DAG executes correctly and
stores the result in the KVS. Also checks to make sure that causal
metadata is properly created.
'''
def func(_, x):
return x * x
fname = 'square'
arg = 2
dag = create_linear_dag([func], [fname], self.kvs_client, 'dag',
MultiKeyCausalLattice)
schedule, triggers = self._create_fn_schedule(dag, arg, fname, [fname],
MULTI)
schedule.output_key = 'output_key'
schedule.client_id = '12'
# We know that there is only one trigger. We populate dependencies
# explicitly in this trigger message to make sure that they are
# reflected in the final result.
kv = triggers['BEGIN'].dependencies.add()
kv.key = 'dependency'
DEFAULT_VC.serialize(kv.vector_clock)
exec_dag_function(self.pusher_cache, self.kvs_client, triggers, func,
schedule, self.user_library, {}, {})
# Assert that there have been 0 messages sent.
self.assertEqual(len(self.socket.outbox), 0)
# Retrieve the result and check its value and its metadata.
result = self.kvs_client.get(schedule.output_key)[schedule.output_key]
self.assertEqual(type(result), MultiKeyCausalLattice)
# Check that the vector clock of the output corresponds ot the client
# ID.
self.assertEqual(result.vector_clock,
VectorClock({schedule.client_id: 1}, True))
# Check that the dependencies of the output match those specified in
# the trigger.
self.assertEqual(len(result.dependencies.reveal()), 1)
self.assertTrue(kv.key in result.dependencies.reveal())
self.assertEqual(result.dependencies.reveal()[kv.key], DEFAULT_VC)
# Check that the output is equal to a local function execution.
result = serializer.load_lattice(result)[0]
self.assertEqual(result, func('', arg))
def test_exec_causal_dag_non_sink(self):
'''
Creates and executes a non-sink function in a causal-mode DAG. This
should be exactly the same as the non-causal version of the test,
except we ensure that the causal metadata is empty, because we don't
have any KVS accesses.
'''
# Create two functions intended to be used in sequence.
def incr(_, x):
x + 1
iname = 'incr'
def square(_, x):
return x * x
sname = 'square'
arg = 1
# Create a DAG and a trigger for the first function in the DAG.
dag = create_linear_dag([incr, square], [iname, sname],
self.kvs_client, 'dag', MultiKeyCausalLattice)
schedule, triggers = self._create_fn_schedule(dag, arg, iname,
[iname, sname], MULTI)
exec_dag_function(self.pusher_cache, self.kvs_client, triggers, incr,
schedule, self.user_library, {}, {})
# Assert that there has been a message sent.
self.assertEqual(len(self.pusher_cache.socket.outbox), 1)
# Extract that message and check its contents.
trigger = DagTrigger()
trigger.ParseFromString(self.pusher_cache.socket.outbox[0])
self.assertEqual(trigger.id, schedule.id)
self.assertEqual(trigger.target_function, sname)
self.assertEqual(trigger.source, iname)
self.assertEqual(len(trigger.arguments.values), 1)
self.assertEqual(len(trigger.version_locations), 0)
self.assertEqual(len(trigger.dependencies), 0)
val = serializer.load(trigger.arguments.values[0])
self.assertEqual(val, incr('', arg))
def test_exec_dag_non_sink(self):
'''
Executes a non-sink function in a DAG and ensures that the correct
downstream trigger was sent with a correct execution of the function.
'''
# Create two functions intended to be used in sequence.
def incr(_, x):
x + 1
iname = 'incr'
def square(_, x):
return x * x
sname = 'square'
arg = 1
# Create a DAG and a trigger for the first function in the DAG.
dag = create_linear_dag([incr, square], [iname, sname],
self.kvs_client, 'dag')
schedule, triggers = self._create_fn_schedule(dag, arg, iname,
[iname, sname])
exec_dag_function(self.pusher_cache, self.kvs_client, triggers, incr,
schedule, self.user_library, {}, {})
# Assert that there has been a message sent.
self.assertEqual(len(self.pusher_cache.socket.outbox), 1)
# Extract that message and check its contents.
trigger = DagTrigger()
trigger.ParseFromString(self.pusher_cache.socket.outbox[0])
self.assertEqual(trigger.id, schedule.id)
self.assertEqual(trigger.target_function, sname)
self.assertEqual(trigger.source, iname)
self.assertEqual(len(trigger.arguments.values), 1)
val = serializer.load(trigger.arguments.values[0])
self.assertEqual(val, incr('', arg))
def test_create_dag_already_exists(self):
'''
This test attempts to create a DAG that already exists and makes sure
that the server correctly rejects the request.
'''
# 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 this to the existing server metadata.
dags = {dag.name: (dag, {source})}
# Add relevant metadata to the policy engine.
address_set = {(self.ip, 1), (self.ip, 2)}
self.policy.unpinned_executors.update(address_set)
# Attempt to create the DAG.
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, DAG_ALREADY_EXISTS)
# Check that no additional metadata was created or sent.
self.assertEqual(len(self.pusher_cache.socket.outbox), 0)
self.assertEqual(len(self.policy.unpinned_executors), 2)
self.assertEqual(len(self.policy.function_locations), 0)
self.assertEqual(len(self.policy.pending_dags), 0)
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)
def test_create_dag(self):
'''
This test creates a new DAG, checking that the correct pin messages are
sent to executors and that it is persisted in the KVS correctly. It
also checks that the server metadata was updated as expected.
'''
# 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.
address_set = {(self.ip, 1), (self.ip, 2)}
self.policy.unpinned_executors.update(address_set)
# Prepopulate the pin_accept socket with sufficient success messages.
self.pin_socket.inbox.append(sutils.ok_resp)
self.pin_socket.inbox.append(sutils.ok_resp)
# Call the DAG creation method.
dags = {}
call_frequency = {}
create_dag(self.socket, self.pusher_cache, self.kvs_client, dags,
self.policy, call_frequency)
# Test that the correct metadata was created.
self.assertTrue(dag_name in dags)
created, dag_source = dags[dag_name]
self.assertEqual(created, dag)
self.assertEqual(len(dag_source), 1)
self.assertEqual(list(dag_source)[0], source)
self.assertTrue(source in call_frequency)
self.assertTrue(sink in call_frequency)
self.assertEqual(call_frequency[source], 0)
self.assertEqual(call_frequency[sink], 0)
# Test that the DAG is stored in the KVS correctly.
result = self.kvs_client.get(dag_name)[dag_name]
created = Dag()
created.ParseFromString(result.reveal())
self.assertEqual(created, dag)
# Test that the correct response was returned to the user.
self.assertTrue(len(self.socket.outbox), 1)
response = GenericResponse()
response.ParseFromString(self.socket.outbox.pop())
self.assertTrue(response.success)
# Test that the correct pin messages were sent.
self.assertEqual(len(self.pusher_cache.socket.outbox), 2)
messages = self.pusher_cache.socket.outbox
function_set = {source, sink}
for message in messages:
self.assertTrue(':' in message)
ip, fname = message.split(':')
self.assertEqual(ip, self.ip)
self.assertTrue(fname in function_set)
function_set.discard(fname)
self.assertEqual(len(function_set), 0)
for address in address_set:
self.assertTrue(
get_pin_address(*address) in self.pusher_cache.addresses)
# Test that the policy engine has the correct metadata stored.
self.assertEqual(len(self.policy.unpinned_executors), 0)
self.assertEqual(len(self.policy.pending_dags), 0)
self.assertTrue(source in self.policy.function_locations)
self.assertTrue(sink in self.policy.function_locations)
self.assertEqual(len(self.policy.function_locations[source]), 1)
self.assertEqual(len(self.policy.function_locations[sink]), 1)
def test_exec_causal_dag_non_sink_with_ref(self):
'''
Creates and executes a non-sink function in a causal-mode DAG. This
version accesses a KVS key, so we ensure that data is appropriately
cached and the metadata is passed downstream.
'''
# Create two functions intended to be used in sequence.
def incr(_, x):
x + 1
iname = 'incr'
def square(_, x):
return x * x
sname = 'square'
# Put tthe argument into the KVS.
arg_name = 'arg'
arg_value = 1
arg = serializer.dump_lattice(arg_value, MultiKeyCausalLattice)
self.kvs_client.put(arg_name, arg)
# Create a DAG and a trigger for the first function in the DAG.
dag = create_linear_dag([incr, square], [iname, sname],
self.kvs_client, 'dag', MultiKeyCausalLattice)
schedule, triggers = self._create_fn_schedule(
dag, CloudburstReference(arg_name, True), iname, [iname, sname],
MULTI)
exec_dag_function(self.pusher_cache, self.kvs_client, triggers, incr,
schedule, self.user_library, {}, {})
# Assert that there has been a message sent.
self.assertEqual(len(self.pusher_cache.socket.outbox), 1)
# Extract that message and check its contents.
trigger = DagTrigger()
trigger.ParseFromString(self.pusher_cache.socket.outbox[0])
self.assertEqual(trigger.id, schedule.id)
self.assertEqual(trigger.target_function, sname)
self.assertEqual(trigger.source, iname)
self.assertEqual(len(trigger.arguments.values), 1)
# Check the metadata of the key that is cached here after execution.
locs = trigger.version_locations
self.assertEqual(len(locs), 1)
self.assertTrue(self.ip in locs.keys())
self.assertEqual(len(locs[self.ip].keys), 1)
kv = locs[self.ip].keys[0]
self.assertEqual(kv.key, arg_name)
self.assertEqual(VectorClock(dict(kv.vector_clock), True),
arg.vector_clock)
# Check the metatada of the causal dependency passed downstream.
self.assertEqual(len(trigger.dependencies), 1)
kv = trigger.dependencies[0]
self.assertEqual(kv.key, arg_name)
self.assertEqual(VectorClock(dict(kv.vector_clock), True),
arg.vector_clock)
val = serializer.load(trigger.arguments.values[0])
self.assertEqual(val, incr('', arg_value))
def test_create_gpu_dag(self):
# Create a simple two-function DAG and add it to the inbound socket.
dag_name = 'dag'
fn = 'fn'
dag = create_linear_dag([None], [fn], self.kvs_client, dag_name)
dag.functions[0].gpu = True
self.socket.inbox.append(dag.SerializeToString())
dags = {}
call_frequency = {}
address_set = {(self.ip, 1)}
self.policy.unpinned_gpu_executors.update(address_set)
self.pin_socket.inbox.append(sutils.ok_resp)
create_dag(self.socket, self.pusher_cache, self.kvs_client, dags,
self.policy, call_frequency)
# Test that the correct metadata was created.
self.assertTrue(dag_name in dags)
created, dag_source = dags[dag_name]
self.assertEqual(created, dag)
self.assertEqual(len(dag_source), 1)
self.assertEqual(list(dag_source)[0], fn)
self.assertTrue(fn in call_frequency)
self.assertEqual(call_frequency[fn], 0)
# Test that the DAG is stored in the KVS correctly.
result = self.kvs_client.get(dag_name)[dag_name]
created = Dag()
created.ParseFromString(result.reveal())
self.assertEqual(created, dag)
# Test that the correct response was returned to the user.
self.assertTrue(len(self.socket.outbox), 1)
response = GenericResponse()
response.ParseFromString(self.socket.outbox.pop())
self.assertTrue(response.success)
# Test that the correct pin messages were sent.
self.assertEqual(len(self.pusher_cache.socket.outbox), 1)
messages = self.pusher_cache.socket.outbox
function_set = {fn}
for message in messages:
pin_msg = PinFunction()
pin_msg.ParseFromString(message)
self.assertEqual(pin_msg.response_address, self.ip)
self.assertTrue(pin_msg.name in function_set)
function_set.discard(pin_msg.name)
self.assertEqual(len(function_set), 0)
for address in address_set:
self.assertTrue(
get_pin_address(*address) in self.pusher_cache.addresses)
# Test that the policy engine has the correct metadata stored.
self.assertEqual(len(self.policy.unpinned_cpu_executors), 0)
self.assertEqual(len(self.policy.pending_dags), 0)
self.assertTrue(fn in self.policy.function_locations)
self.assertEqual(len(self.policy.function_locations[fn]), 1)
