def test_conditional_protocol_group():
with tempfile.TemporaryDirectory() as directory:
initial_value = 2 * unit.kelvin
value_protocol_a = DummyInputOutputProtocol("protocol_a")
value_protocol_a.input_value = initial_value
add_values = AddValues("add_values")
add_values.values = [
ProtocolPath("output_value", value_protocol_a.id),
ProtocolPath("output_value", value_protocol_a.id),
]
condition = ConditionalGroup.Condition()
condition.left_hand_value = ProtocolPath("result", add_values.id)
condition.right_hand_value = ProtocolPath("output_value",
value_protocol_a.id)
condition.type = ConditionalGroup.Condition.Type.GreaterThan
protocol_group = ConditionalGroup("protocol_group")
protocol_group.conditions.append(condition)
protocol_group.add_protocols(value_protocol_a, add_values)
protocol_group.execute(directory, ComputeResources())
assert (protocol_group.get_value(ProtocolPath(
"result", add_values.id)) == 4 * unit.kelvin)
def test_conditional_group_self_reference():
"""Tests that protocols within a conditional group
can access the outputs of its parent, such as the
current iteration of the group."""
max_iterations = 10
criteria = random.randint(1, max_iterations - 1)
group = ConditionalGroup("conditional_group")
group.max_iterations = max_iterations
protocol = DummyInputOutputProtocol("protocol_a")
protocol.input_value = ProtocolPath("current_iteration", group.id)
condition_1 = ConditionalGroup.Condition()
condition_1.left_hand_value = ProtocolPath("output_value", group.id,
protocol.id)
condition_1.right_hand_value = criteria
condition_1.type = ConditionalGroup.Condition.Type.GreaterThan
condition_2 = ConditionalGroup.Condition()
condition_2.left_hand_value = ProtocolPath("current_iteration", group.id)
condition_2.right_hand_value = criteria
condition_2.type = ConditionalGroup.Condition.Type.GreaterThan
group.add_protocols(protocol)
group.add_condition(condition_1)
group.add_condition(condition_2)
with tempfile.TemporaryDirectory() as directory:
group.execute(directory, ComputeResources())
assert protocol.output_value == criteria + 1
def test_conditional_protocol_group_fail():
with tempfile.TemporaryDirectory() as directory:
initial_value = 2 * unit.kelvin
value_protocol_a = DummyInputOutputProtocol("protocol_a")
value_protocol_a.input_value = initial_value
add_values = AddValues("add_values")
add_values.values = [
ProtocolPath("output_value", value_protocol_a.id),
ProtocolPath("output_value", value_protocol_a.id),
]
condition = ConditionalGroup.Condition()
condition.left_hand_value = ProtocolPath("result", add_values.id)
condition.right_hand_value = ProtocolPath("output_value",
value_protocol_a.id)
condition.type = ConditionalGroup.Condition.Type.LessThan
protocol_group = ConditionalGroup("protocol_group")
protocol_group.conditions.append(condition)
protocol_group.max_iterations = 10
protocol_group.add_protocols(value_protocol_a, add_values)
with pytest.raises(RuntimeError):
protocol_group.execute(directory, ComputeResources())
def test_from_schema():
protocol_a = DummyInputOutputProtocol("protocol_a")
protocol_a.input_value = 1 * unit.kelvin
schema = WorkflowSchema()
schema.protocol_schemas = [protocol_a.schema]
workflow = Workflow.from_schema(schema, {}, unique_id="")
assert workflow is not None
rebuilt_schema = workflow.schema
rebuilt_schema.gradients_sources = UNDEFINED
rebuilt_schema.outputs_to_store = UNDEFINED
assert rebuilt_schema.json(format=True) == schema.json(format=True)
def test_protocol_graph_execution(calculation_backend, compute_resources):
if calculation_backend is not None:
calculation_backend.start()
protocol_a = DummyInputOutputProtocol("protocol_a")
protocol_a.input_value = 1
protocol_b = DummyInputOutputProtocol("protocol_b")
protocol_b.input_value = ProtocolPath("output_value", protocol_a.id)
protocol_graph = ProtocolGraph()
protocol_graph.add_protocols(protocol_a, protocol_b)
with tempfile.TemporaryDirectory() as directory:
results = protocol_graph.execute(directory, calculation_backend,
compute_resources)
final_result = results[protocol_b.id]
if calculation_backend is not None:
final_result = final_result.result()
with open(final_result[1]) as file:
results_b = json.load(file, cls=TypedJSONDecoder)
assert results_b[".output_value"] == protocol_a.input_value
if compute_resources is not None:
assert protocol_b.output_value == protocol_a.input_value
if calculation_backend is not None:
calculation_backend.stop()
def test_index_replicated_protocol():
replicator = ProtocolReplicator("replicator")
replicator.template_values = ["a", "b", "c", "d"]
replicated_protocol = DummyInputOutputProtocol(
f"protocol_{replicator.placeholder_id}")
replicated_protocol.input_value = ReplicatorValue(replicator.id)
schema = WorkflowSchema()
schema.protocol_replicators = [replicator]
schema.protocol_schemas = [replicated_protocol.schema]
for index in range(len(replicator.template_values)):
indexing_protocol = DummyInputOutputProtocol(
f"indexing_protocol_{index}")
indexing_protocol.input_value = ProtocolPath("output_value",
f"protocol_{index}")
schema.protocol_schemas.append(indexing_protocol.schema)
schema.validate()
workflow = Workflow({})
workflow.schema = schema
def test_nested_input():
dict_protocol = DummyInputOutputProtocol("dict_protocol")
dict_protocol.input_value = {"a": ThermodynamicState(1.0 * unit.kelvin)}
quantity_protocol = DummyInputOutputProtocol("quantity_protocol")
quantity_protocol.input_value = ProtocolPath("output_value[a].temperature",
dict_protocol.id)
schema = WorkflowSchema()
schema.protocol_schemas = [dict_protocol.schema, quantity_protocol.schema]
schema.validate()
workflow = Workflow({})
workflow.schema = schema
workflow_graph = workflow.to_graph()
with tempfile.TemporaryDirectory() as temporary_directory:
with DaskLocalCluster() as calculation_backend:
results_futures = workflow_graph.execute(temporary_directory,
calculation_backend)
assert len(results_futures) == 1
result = results_futures[0].result()
assert isinstance(result, WorkflowResult)
def test_simple_workflow_graph(calculation_backend, compute_resources,
exception):
expected_value = (1 * unit.kelvin).plus_minus(0.1 * unit.kelvin)
protocol_a = DummyInputOutputProtocol("protocol_a")
protocol_a.input_value = expected_value
protocol_b = DummyInputOutputProtocol("protocol_b")
protocol_b.input_value = ProtocolPath("output_value", protocol_a.id)
schema = WorkflowSchema()
schema.protocol_schemas = [protocol_a.schema, protocol_b.schema]
schema.final_value_source = ProtocolPath("output_value", protocol_b.id)
schema.validate()
workflow = Workflow({})
workflow.schema = schema
workflow_graph = workflow.to_graph()
with tempfile.TemporaryDirectory() as directory:
if calculation_backend is not None:
with DaskLocalCluster() as calculation_backend:
if exception:
with pytest.raises(AssertionError):
workflow_graph.execute(directory, calculation_backend,
compute_resources)
return
else:
results_futures = workflow_graph.execute(
directory, calculation_backend, compute_resources)
assert len(results_futures) == 1
result = results_futures[0].result()
else:
result = workflow_graph.execute(directory, calculation_backend,
compute_resources)[0]
if exception:
with pytest.raises(AssertionError):
workflow_graph.execute(directory, calculation_backend,
compute_resources)
return
assert isinstance(result, WorkflowResult)
assert result.value.value == expected_value.value
def build_easy_graph():
protocol_a = DummyInputOutputProtocol("protocol_a")
protocol_a.input_value = 1
protocol_b = DummyInputOutputProtocol("protocol_b")
protocol_b.input_value = 1
return [protocol_a], [protocol_b]
def build_protocols(prefix):
# .-------------------.
# | / i - j -|- b
# a - | g - h - | |
# | \ k - l -|- c
# .-------------------.
protocol_a = DummyInputOutputProtocol(prefix + "protocol_a")
protocol_a.input_value = 1
fork_protocols = build_fork(prefix)
fork_protocols[0].input_value = ProtocolPath("output_value",
protocol_a.id)
protocol_group = ProtocolGroup(prefix + "protocol_group")
protocol_group.add_protocols(*fork_protocols)
protocol_b = DummyInputOutputProtocol(prefix + "protocol_b")
protocol_b.input_value = ProtocolPath("output_value",
protocol_group.id, "protocol_j")
protocol_c = DummyInputOutputProtocol(prefix + "protocol_c")
protocol_c.input_value = ProtocolPath("output_value",
protocol_group.id, "protocol_l")
return [protocol_a, protocol_group, protocol_b, protocol_c]
def test_protocol_group_execution():
protocol_a = DummyInputOutputProtocol("protocol_a")
protocol_a.input_value = 1
protocol_b = DummyInputOutputProtocol("protocol_b")
protocol_b.input_value = ProtocolPath("output_value", protocol_a.id)
protocol_group = ProtocolGroup("protocol_group")
protocol_group.add_protocols(protocol_a, protocol_b)
with tempfile.TemporaryDirectory() as directory:
protocol_group.execute(directory, ComputeResources())
value_path = ProtocolPath("output_value", protocol_group.id, protocol_b.id)
final_value = protocol_group.get_value(value_path)
assert final_value == protocol_a.input_value
def test_workflow_with_groups():
expected_value = (1 * unit.kelvin).plus_minus(0.1 * unit.kelvin)
protocol_a = DummyInputOutputProtocol("protocol_a")
protocol_a.input_value = expected_value
protocol_b = DummyInputOutputProtocol("protocol_b")
protocol_b.input_value = ProtocolPath("output_value", protocol_a.id)
conditional_group = ConditionalGroup("conditional_group")
conditional_group.add_protocols(protocol_a, protocol_b)
condition = ConditionalGroup.Condition()
condition.right_hand_value = 2 * unit.kelvin
condition.type = ConditionalGroup.Condition.Type.LessThan
condition.left_hand_value = ProtocolPath("output_value.value",
conditional_group.id,
protocol_b.id)
conditional_group.add_condition(condition)
schema = WorkflowSchema()
schema.protocol_schemas = [conditional_group.schema]
schema.final_value_source = ProtocolPath("output_value",
conditional_group.id,
protocol_b.id)
schema.validate()
workflow = Workflow({})
workflow.schema = schema
workflow_graph = workflow.to_graph()
with tempfile.TemporaryDirectory() as directory:
with DaskLocalCluster() as calculation_backend:
results_futures = workflow_graph.execute(directory,
calculation_backend)
assert len(results_futures) == 1
result = results_futures[0].result()
assert isinstance(result, WorkflowResult)
assert result.value.value == expected_value.value
def test_nested_protocol_paths():
value_protocol_a = DummyInputOutputProtocol("protocol_a")
value_protocol_a.input_value = (1 * unit.kelvin).plus_minus(0.1 *
unit.kelvin)
assert (value_protocol_a.get_value(ProtocolPath("input_value.value")) ==
value_protocol_a.input_value.value)
value_protocol_b = DummyInputOutputProtocol("protocol_b")
value_protocol_b.input_value = (2 * unit.kelvin).plus_minus(0.05 *
unit.kelvin)
value_protocol_c = DummyInputOutputProtocol("protocol_c")
value_protocol_c.input_value = (4 * unit.kelvin).plus_minus(0.01 *
unit.kelvin)
add_values_protocol = AddValues("add_values")
add_values_protocol.values = [
ProtocolPath("output_value", value_protocol_a.id),
ProtocolPath("output_value", value_protocol_b.id),
ProtocolPath("output_value", value_protocol_b.id),
5,
]
with pytest.raises(ValueError):
add_values_protocol.get_value(ProtocolPath("valus[string]"))
with pytest.raises(ValueError):
add_values_protocol.get_value(ProtocolPath("values[string]"))
input_values = add_values_protocol.get_value_references(
ProtocolPath("values"))
assert isinstance(input_values, dict) and len(input_values) == 3
for index, value_reference in enumerate(input_values):
input_value = add_values_protocol.get_value(value_reference)
assert input_value.full_path == add_values_protocol.values[
index].full_path
add_values_protocol.set_value(value_reference, index)
assert set(add_values_protocol.values) == {0, 1, 2, 5}
dummy_dict_protocol = DummyInputOutputProtocol("dict_protocol")
dummy_dict_protocol.input_value = {
"value_a": ProtocolPath("output_value", value_protocol_a.id),
"value_b": ProtocolPath("output_value", value_protocol_b.id),
}
input_values = dummy_dict_protocol.get_value_references(
ProtocolPath("input_value"))
assert isinstance(input_values, dict) and len(input_values) == 2
for index, value_reference in enumerate(input_values):
input_value = dummy_dict_protocol.get_value(value_reference)
dummy_dict_keys = list(dummy_dict_protocol.input_value.keys())
assert (input_value.full_path == dummy_dict_protocol.input_value[
dummy_dict_keys[index]].full_path)
dummy_dict_protocol.set_value(value_reference, index)
add_values_protocol_2 = AddValues("add_values")
add_values_protocol_2.values = [
[ProtocolPath("output_value", value_protocol_a.id)],
[
ProtocolPath("output_value", value_protocol_b.id),
ProtocolPath("output_value", value_protocol_b.id),
],
]
with pytest.raises(ValueError):
add_values_protocol_2.get_value(ProtocolPath("valus[string]"))
with pytest.raises(ValueError):
add_values_protocol.get_value(ProtocolPath("values[string]"))
pass
def build_fork(prefix):
# / i - j
# g - h - |
# \ k - l
protocol_g = DummyInputOutputProtocol(prefix + "protocol_g")
protocol_g.input_value = 3
protocol_h = DummyInputOutputProtocol(prefix + "protocol_h")
protocol_h.input_value = ProtocolPath("output_value", protocol_g.id)
protocol_i = DummyInputOutputProtocol(prefix + "protocol_i")
protocol_i.input_value = ProtocolPath("output_value", protocol_h.id)
protocol_j = DummyInputOutputProtocol(prefix + "protocol_j")
protocol_j.input_value = ProtocolPath("output_value", protocol_i.id)
protocol_k = DummyInputOutputProtocol(prefix + "protocol_k")
protocol_k.input_value = ProtocolPath("output_value", protocol_h.id)
protocol_l = DummyInputOutputProtocol(prefix + "protocol_l")
protocol_l.input_value = ProtocolPath("output_value", protocol_k.id)
return [
protocol_g, protocol_h, protocol_i, protocol_j, protocol_k, protocol_l
]
def build_merge(prefix):
# a - b \
# | - e - f
# c - d /
protocol_a = DummyInputOutputProtocol(prefix + "protocol_a")
protocol_a.input_value = 1
protocol_b = DummyInputOutputProtocol(prefix + "protocol_b")
protocol_b.input_value = ProtocolPath("output_value", protocol_a.id)
protocol_c = DummyInputOutputProtocol(prefix + "protocol_c")
protocol_c.input_value = 2
protocol_d = DummyInputOutputProtocol(prefix + "protocol_d")
protocol_d.input_value = ProtocolPath("output_value", protocol_c.id)
protocol_e = DummyInputOutputProtocol(prefix + "protocol_e")
protocol_e.input_value = [
ProtocolPath("output_value", protocol_b.id),
ProtocolPath("output_value", protocol_d.id),
]
protocol_f = DummyInputOutputProtocol(prefix + "protocol_f")
protocol_f.input_value = ProtocolPath("output_value", protocol_e.id)
return [
protocol_a,
protocol_b,
protocol_c,
protocol_d,
protocol_e,
protocol_f,
]
def test_group_replicators():
dummy_schema = WorkflowSchema()
replicator_id = "replicator"
dummy_replicated_protocol = DummyInputOutputProtocol(
f"dummy_$({replicator_id})")
dummy_replicated_protocol.input_value = ReplicatorValue(replicator_id)
dummy_group = ProtocolGroup("dummy_group")
dummy_group.add_protocols(dummy_replicated_protocol)
dummy_protocol_single_value = DummyInputOutputProtocol(
f"dummy_single_$({replicator_id})")
dummy_protocol_single_value.input_value = ProtocolPath(
"output_value", dummy_group.id, dummy_replicated_protocol.id)
dummy_protocol_list_value = AddValues("dummy_list")
dummy_protocol_list_value.values = ProtocolPath(
"output_value", dummy_group.id, dummy_replicated_protocol.id)
dummy_schema.protocol_schemas = [
dummy_group.schema,
dummy_protocol_single_value.schema,
dummy_protocol_list_value.schema,
]
replicator = ProtocolReplicator(replicator_id)
replicator.template_values = [
(1.0 * unit.kelvin).plus_minus(1.0 * unit.kelvin),
(2.0 * unit.kelvin).plus_minus(2.0 * unit.kelvin),
]
dummy_schema.protocol_replicators = [replicator]
dummy_schema.validate()
dummy_property = create_dummy_property(Density)
dummy_metadata = Workflow.generate_default_metadata(
dummy_property, "smirnoff99Frosst-1.1.0.offxml", [])
dummy_workflow = Workflow(dummy_metadata, "")
dummy_workflow.schema = dummy_schema
assert len(dummy_workflow.protocols) == 4
assert (dummy_workflow.protocols[dummy_group.id].protocols["dummy_0"].
input_value.value == replicator.template_values[0].value)
assert (dummy_workflow.protocols[dummy_group.id].protocols["dummy_1"].
input_value.value == replicator.template_values[1].value)
assert dummy_workflow.protocols[
"dummy_single_0"].input_value == ProtocolPath("output_value",
dummy_group.id,
"dummy_0")
assert dummy_workflow.protocols[
"dummy_single_1"].input_value == ProtocolPath("output_value",
dummy_group.id,
"dummy_1")
assert len(dummy_workflow.protocols["dummy_list"].values) == 2
assert dummy_workflow.protocols["dummy_list"].values[0] == ProtocolPath(
"output_value", dummy_group.id, "dummy_0")
assert dummy_workflow.protocols["dummy_list"].values[1] == ProtocolPath(
"output_value", dummy_group.id, "dummy_1")
def test_protocol_group_resume():
"""A test that protocol groups can recover after being killed
(e.g. by a worker being killed due to hitting a wallclock limit)
"""
compute_resources = ComputeResources()
# Fake a protocol group which executes the first
# two protocols and then 'gets killed'.
protocol_a = DummyInputOutputProtocol("protocol_a")
protocol_a.input_value = 1
protocol_b = DummyInputOutputProtocol("protocol_b")
protocol_b.input_value = ProtocolPath("output_value", protocol_a.id)
protocol_group_a = ProtocolGroup("group_a")
protocol_group_a.add_protocols(protocol_a, protocol_b)
protocol_graph = ProtocolGraph()
protocol_graph.add_protocols(protocol_group_a)
protocol_graph.execute("graph_a", compute_resources=compute_resources)
# Remove the output file so it appears the the protocol group had not
# completed.
os.unlink(
os.path.join("graph_a", protocol_group_a.id,
f"{protocol_group_a.id}_output.json"))
# Build the 'full' group with the last two protocols which
# 'had not been exited' after the group was 'killed'
protocol_a = DummyInputOutputProtocol("protocol_a")
protocol_a.input_value = 1
protocol_b = DummyInputOutputProtocol("protocol_b")
protocol_b.input_value = ProtocolPath("output_value", protocol_a.id)
protocol_c = DummyInputOutputProtocol("protocol_c")
protocol_c.input_value = ProtocolPath("output_value", protocol_b.id)
protocol_d = DummyInputOutputProtocol("protocol_d")
protocol_d.input_value = ProtocolPath("output_value", protocol_c.id)
protocol_group_a = ProtocolGroup("group_a")
protocol_group_a.add_protocols(protocol_a, protocol_b, protocol_c,
protocol_d)
protocol_graph = ProtocolGraph()
protocol_graph.add_protocols(protocol_group_a)
protocol_graph.execute("graph_a", compute_resources=compute_resources)
assert all(x != UNDEFINED for x in protocol_group_a.outputs.values())
def test_workflow_layer():
"""Test the `WorkflowLayer` calculation layer. As the `SimulationLayer`
is the simplest implementation of the abstract layer, we settle for
testing this."""
properties_to_estimate = [
create_dummy_property(Density),
create_dummy_property(Density),
]
# Create a very simple workflow which just returns some placeholder
# value.
estimated_value = (1 * unit.kelvin).plus_minus(0.1 * unit.kelvin)
protocol_a = DummyInputOutputProtocol("protocol_a")
protocol_a.input_value = estimated_value
schema = WorkflowSchema()
schema.protocol_schemas = [protocol_a.schema]
schema.final_value_source = ProtocolPath("output_value", protocol_a.id)
layer_schema = SimulationSchema()
layer_schema.workflow_schema = schema
options = RequestOptions()
options.add_schema("SimulationLayer", "Density", layer_schema)
batch = server.Batch()
batch.queued_properties = properties_to_estimate
batch.options = options
with tempfile.TemporaryDirectory() as directory:
with temporarily_change_directory(directory):
# Create a directory for the layer.
layer_directory = "simulation_layer"
os.makedirs(layer_directory)
# Set-up a simple storage backend and add a force field to it.
force_field = SmirnoffForceFieldSource.from_path(
"smirnoff99Frosst-1.1.0.offxml")
storage_backend = LocalFileStorage()
batch.force_field_id = storage_backend.store_force_field(
force_field)
# Create a simple calculation backend to test with.
with DaskLocalCluster() as calculation_backend:
def dummy_callback(returned_request):
assert len(returned_request.estimated_properties) == 2
assert len(returned_request.exceptions) == 0
simulation_layer = SimulationLayer()
simulation_layer.schedule_calculation(
calculation_backend,
storage_backend,
layer_directory,
batch,
dummy_callback,
True,
)