def test_base_layer():
properties_to_estimate = [
create_dummy_property(Density),
create_dummy_property(Density),
]
dummy_options = RequestOptions()
batch = server.Batch()
batch.queued_properties = properties_to_estimate
batch.options = dummy_options
batch.force_field_id = ""
batch.options.calculation_schemas = {
"Density": {
"DummyCalculationLayer": CalculationLayerSchema()
}
}
with tempfile.TemporaryDirectory() as temporary_directory:
with temporarily_change_directory(temporary_directory):
# Create a simple calculation backend to test with.
test_backend = DaskLocalCluster()
test_backend.start()
# Create a simple storage backend to test with.
test_storage = LocalFileStorage()
layer_directory = "dummy_layer"
makedirs(layer_directory)
def dummy_callback(returned_request):
assert len(returned_request.estimated_properties) == 1
assert len(returned_request.exceptions) == 2
dummy_layer = DummyCalculationLayer()
dummy_layer.schedule_calculation(
test_backend,
test_storage,
layer_directory,
batch,
dummy_callback,
True,
)
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_submission():
with tempfile.TemporaryDirectory() as directory:
with temporarily_change_directory(directory):
with DaskLocalCluster() as calculation_backend:
# Spin up a server instance.
server = EvaluatorServer(
calculation_backend=calculation_backend,
working_directory=directory,
)
with server:
# Connect a client.
client = EvaluatorClient()
# Submit an empty data set.
force_field_path = "smirnoff99Frosst-1.1.0.offxml"
force_field_source = SmirnoffForceFieldSource.from_path(
force_field_path
)
request, error = client.request_estimate(
PhysicalPropertyDataSet(), force_field_source
)
assert error is None
assert isinstance(request, Request)
result, error = request.results(polling_interval=0.01)
assert error is None
assert isinstance(result, RequestResult)
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 test_server_spin_up():
with tempfile.TemporaryDirectory() as directory:
with temporarily_change_directory(directory):
with DaskLocalCluster() as calculation_backend:
server = EvaluatorServer(
calculation_backend=calculation_backend,
working_directory=directory,
)
with server:
sleep(0.5)
def test_same_component_batching():
thermodynamic_state = ThermodynamicState(temperature=1.0 * unit.kelvin,
pressure=1.0 * unit.atmosphere)
data_set = PhysicalPropertyDataSet()
data_set.add_properties(
Density(
thermodynamic_state=thermodynamic_state,
substance=Substance.from_components("O", "C"),
value=0.0 * unit.kilogram / unit.meter**3,
),
EnthalpyOfVaporization(
thermodynamic_state=thermodynamic_state,
substance=Substance.from_components("O", "C"),
value=0.0 * unit.kilojoule / unit.mole,
),
Density(
thermodynamic_state=thermodynamic_state,
substance=Substance.from_components("O", "CO"),
value=0.0 * unit.kilogram / unit.meter**3,
),
EnthalpyOfVaporization(
thermodynamic_state=thermodynamic_state,
substance=Substance.from_components("O", "CO"),
value=0.0 * unit.kilojoule / unit.mole,
),
)
options = RequestOptions()
submission = EvaluatorClient._Submission()
submission.dataset = data_set
submission.options = options
with DaskLocalCluster() as calculation_backend:
server = EvaluatorServer(calculation_backend)
batches = server._batch_by_same_component(submission, "")
assert len(batches) == 2
assert len(batches[0].queued_properties) == 2
assert len(batches[1].queued_properties) == 2
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_launch_batch():
# Set up a dummy data set
data_set = PhysicalPropertyDataSet()
data_set.add_properties(create_dummy_property(Density),
create_dummy_property(Density))
batch = Batch()
batch.force_field_id = ""
batch.options = RequestOptions()
batch.options.calculation_layers = ["QuickCalculationLayer"]
batch.options.calculation_schemas = {
"Density": {
"QuickCalculationLayer": CalculationLayerSchema()
}
}
batch.parameter_gradient_keys = []
batch.queued_properties = [*data_set]
batch.validate()
with tempfile.TemporaryDirectory() as directory:
with temporarily_change_directory(directory):
with DaskLocalCluster() as calculation_backend:
server = EvaluatorServer(
calculation_backend=calculation_backend,
working_directory=directory,
)
server._queued_batches[batch.id] = batch
server._launch_batch(batch)
while len(batch.queued_properties) > 0:
sleep(0.01)
assert len(batch.estimated_properties) == 1
assert len(batch.unsuccessful_properties) == 1
def setup_server(
backend_type=BackendType.LocalCPU,
max_number_of_workers=1,
conda_environment="evaluator",
worker_memory=4 * unit.gigabyte,
port=8000,
cuda_version="10.1",
):
"""A convenience function to sets up an estimation server which will can advantage
of different compute backends.
Parameters
----------
backend_type: BackendType
The type of compute backend to use.
max_number_of_workers: int
The maximum number of workers to adaptively insert into
the queuing system.
conda_environment: str
The name of the conda environment in which the evaluator
package is installed.
worker_memory: pint.Quantity
The maximum amount of memory to request per worker.
port: int
The port that the server should listen for estimation requests on.
cuda_version: str
The version of CUDA to use if running on a backend which supports
GPUs.
Returns
-------
EvaluatorServer
The server object.
"""
calculation_backend = None
if backend_type == BackendType.LocalCPU:
calculation_backend = DaskLocalCluster(
number_of_workers=max_number_of_workers)
elif backend_type == BackendType.LocalGPU:
calculation_backend = DaskLocalCluster(
number_of_workers=max_number_of_workers,
resources_per_worker=ComputeResources(
1, 1, ComputeResources.GPUToolkit.CUDA),
)
elif backend_type == BackendType.GPU:
queue_resources = QueueWorkerResources(
number_of_threads=1,
number_of_gpus=1,
preferred_gpu_toolkit=QueueWorkerResources.GPUToolkit.CUDA,
per_thread_memory_limit=worker_memory,
wallclock_time_limit="05:59",
)
worker_script_commands = [
f"conda activate {conda_environment}",
f"module load cuda/{cuda_version}",
]
calculation_backend = DaskLSFBackend(
minimum_number_of_workers=1,
maximum_number_of_workers=max_number_of_workers,
resources_per_worker=queue_resources,
queue_name="gpuqueue",
setup_script_commands=worker_script_commands,
adaptive_interval="1000ms",
)
elif backend_type == BackendType.CPU:
queue_resources = QueueWorkerResources(
number_of_threads=1,
per_thread_memory_limit=worker_memory,
wallclock_time_limit="01:30",
)
worker_script_commands = [f"conda activate {conda_environment}"]
calculation_backend = DaskLSFBackend(
minimum_number_of_workers=1,
maximum_number_of_workers=max_number_of_workers,
resources_per_worker=queue_resources,
queue_name="cpuqueue",
setup_script_commands=worker_script_commands,
adaptive_interval="1000ms",
)
calculation_backend.start()
# Spin up the server object.
return server.EvaluatorServer(calculation_backend=calculation_backend,
port=port)
# Currently the graph shouldn't merge with an
# addition
protocol_graph = ProtocolGraph()
protocol_graph.add_protocols(*protocols_a, *protocols_b)
dependants_graph = protocol_graph._build_dependants_graph(
protocol_graph.protocols, False, apply_reduction=False)
assert len(protocol_graph.protocols) == len(protocols_a) + len(protocols_b)
assert len(dependants_graph) == len(protocols_a) + len(protocols_b)
assert len(protocol_graph.root_protocols) == 2 * n_root_protocols
@pytest.mark.parametrize(
"calculation_backend, compute_resources",
[(DaskLocalCluster(), None), (None, ComputeResources())],
)
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:
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,
)
from evaluator.backends import ComputeResources
from evaluator.backends.dask import DaskLocalCluster
from evaluator.protocols.groups import ConditionalGroup
from evaluator.tests.test_workflow.utils import DummyInputOutputProtocol
from evaluator.thermodynamics import ThermodynamicState
from evaluator.workflow import Workflow, WorkflowResult, WorkflowSchema
from evaluator.workflow.schemas import ProtocolReplicator
from evaluator.workflow.utils import ProtocolPath, ReplicatorValue
@pytest.mark.parametrize(
"calculation_backend, compute_resources, exception",
[
(None, None, False),
(None, ComputeResources(number_of_threads=1), False),
(DaskLocalCluster(), None, False),
(DaskLocalCluster(), ComputeResources(number_of_threads=1), True),
],
)
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]