Add NoBatch mode (#602)

* add nobatch

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* add test

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: lilyminium

openff/evaluator/_tests/test_server.py

@@ -2,23 +2,26 @@
 Units tests for the openff.evaluator.server module.
 """
 
+import os
 import tempfile
 from time import sleep
 
+import pytest
 from openff.units import unit
 
 from openff.evaluator._tests.utils import create_dummy_property
 from openff.evaluator.backends.dask import DaskLocalCluster
 from openff.evaluator.client import EvaluatorClient, RequestOptions
 from openff.evaluator.datasets import PhysicalPropertyDataSet
+from openff.evaluator.forcefield import SmirnoffForceFieldSource
 from openff.evaluator.layers import (
     CalculationLayer,
     CalculationLayerResult,
     CalculationLayerSchema,
     calculation_layer,
 )
 from openff.evaluator.properties import Density, EnthalpyOfVaporization
-from openff.evaluator.server.server import Batch, EvaluatorServer
+from openff.evaluator.server.server import Batch, BatchMode, EvaluatorServer
 from openff.evaluator.substances import Substance
 from openff.evaluator.thermodynamics import ThermodynamicState
 from openff.evaluator.utils.utils import temporarily_change_directory
@@ -106,7 +109,8 @@ def test_launch_batch():
                 assert len(batch.unsuccessful_properties) == 1
 
 
-def test_same_component_batching():
+@pytest.fixture
+def c_o_dataset():
     thermodynamic_state = ThermodynamicState(
         temperature=1.0 * unit.kelvin, pressure=1.0 * unit.atmosphere
     )
@@ -134,61 +138,66 @@ def test_same_component_batching():
             value=0.0 * unit.kilojoule / unit.mole,
         ),
     )
+    return data_set
 
+
+@pytest.fixture
+def dataset_submission(c_o_dataset):
     options = RequestOptions()
 
     submission = EvaluatorClient._Submission()
-    submission.dataset = data_set
+    submission.dataset = c_o_dataset
     submission.options = options
+    submission.force_field_source = SmirnoffForceFieldSource.from_path(
+        "openff-2.1.0.offxml"
+    )
+
+    return submission
 
+
+def test_same_component_batching(dataset_submission, tmp_path):
+    os.chdir(tmp_path)
     with DaskLocalCluster() as calculation_backend:
         server = EvaluatorServer(calculation_backend)
-        batches = server._batch_by_same_component(submission, "")
+        batches = server._batch_by_same_component(dataset_submission, "")
 
     assert len(batches) == 2
 
     assert len(batches[0].queued_properties) == 2
     assert len(batches[1].queued_properties) == 2
 
 
-def test_shared_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,
-        ),
-    )
+def test_shared_component_batching(dataset_submission, tmp_path):
+    os.chdir(tmp_path)
+    with DaskLocalCluster() as calculation_backend:
+        server = EvaluatorServer(calculation_backend)
+        batches = server._batch_by_shared_component(dataset_submission, "")
 
-    options = RequestOptions()
+    assert len(batches) == 1
+    assert len(batches[0].queued_properties) == 4
 
-    submission = EvaluatorClient._Submission()
-    submission.dataset = data_set
-    submission.options = options
 
+def test_nobatching(dataset_submission, tmp_path):
+    os.chdir(tmp_path)
     with DaskLocalCluster() as calculation_backend:
         server = EvaluatorServer(calculation_backend)
-        batches = server._batch_by_shared_component(submission, "")
+        batches = server._no_batch(dataset_submission, "")
 
     assert len(batches) == 1
     assert len(batches[0].queued_properties) == 4
+    assert batches[0].id == "batch_0000"
+
+
+def test_prepare_batches_nobatch(dataset_submission, tmp_path):
+    os.chdir(tmp_path)
+    dataset_submission.options.batch_mode = BatchMode.NoBatch
+    with DaskLocalCluster() as calculation_backend:
+        server = EvaluatorServer(calculation_backend)
+        server._batch_ids_per_client_id["request_id"] = []
+        batches = server._prepare_batches(dataset_submission, "request_id")
+
+        assert len(batches) == 1
+        assert len(batches[0].queued_properties) == 4
+        assert batches[0].id == "batch_0000"
+        assert server._queued_batches["batch_0000"] is batches[0]
+        assert server._batch_ids_per_client_id["request_id"] == ["batch_0000"]

openff/evaluator/client/client.py

@@ -83,13 +83,16 @@ class BatchMode(Enum):
       common component will be batched together. E.g.The densities of 80:20 and 20:80
       mixtures of ethanol and water, and the pure densities of ethanol and water would
       be batched together.
+    * NoBatch: No batching will be performed. Each property will be estimated in a
+      single, sequentially-increasing batch.
 
     Properties will only be marked as estimated by the server when all properties in a
     single batch are completed.
     """
 
     SameComponents = "SameComponents"
     SharedComponents = "SharedComponents"
+    NoBatch = "NoBatch"
 
 
 class Request(AttributeClass):

openff/evaluator/server/server.py

@@ -280,6 +280,29 @@ def _query_request_status(self, client_request_id):
 
         return request_results, None
 
+    def _no_batch(self, submission, force_field_id):
+        """Returns a single Batch."""
+
+        reserved_batch_ids = {
+            *self._queued_batches.keys(),
+            *self._finished_batches.keys(),
+        }
+        batch = Batch()
+        batch.force_field_id = force_field_id
+        batch.enable_data_caching = self._enable_data_caching
+        batch.queued_properties = list(submission.dataset.properties)
+        batch.options = RequestOptions.parse_json(submission.options.json())
+        batch.parameter_gradient_keys = copy.deepcopy(
+            submission.parameter_gradient_keys
+        )
+
+        n_batchs = len(reserved_batch_ids)
+        batch.id = f"batch_{n_batchs:04d}"
+        while batch.id in reserved_batch_ids:
+            n_batchs += 1
+            batch.id = f"batch_{n_batchs:04d}"
+        return [batch]
+
     def _batch_by_same_component(self, submission, force_field_id):
         """Batches a set of requested properties based on which substance they were
         measured for. Properties which were measured for substances containing the
@@ -447,6 +470,8 @@ def _prepare_batches(self, submission, request_id):
             batches = self._batch_by_same_component(submission, force_field_id)
         elif batch_mode == BatchMode.SharedComponents:
             batches = self._batch_by_shared_component(submission, force_field_id)
+        elif batch_mode == BatchMode.NoBatch:
+            batches = self._no_batch(submission, force_field_id)
         else:
             raise NotImplementedError()