Backport fix from PR #212

src/somd2/runner/_base.py

@@ -636,6 +636,11 @@ def __init__(self, system, config):
         # used to skip duplicate rows on restart.
         self._last_ec_time = {}
 
+        # Per-window energy-components rows collected since the last checkpoint,
+        # and the number of rows at which they are written out regardless.
+        self._ec_rows = {}
+        self._max_ec_rows = 10000
+
         # Per-window cache of the integrator's integration force groups bitmask.
         self._integration_groups = {}
 
@@ -2500,6 +2505,9 @@ def _checkpoint(
                             df.iloc[-self._energy_per_block :],
                         )
 
+            if not is_post_equilibration:
+                self._flush_energy_components(index)
+
         except Exception as e:
             return index, e
 
@@ -2644,8 +2652,11 @@ def _backup_checkpoint(self, index):
 
     def _save_energy_components(self, index, context, time_ns):
         """
-        Internal function to save the energy components for each force group to a
-        Parquet file.
+        Internal function to record the energy components for each force group.
+        Rows are buffered and written to the Parquet file by
+        _flush_energy_components() at checkpoint time, so that the file is only
+        rewritten once per checkpoint and stays consistent with the other
+        checkpoint files.
 
         Parameters
         ----------
@@ -2660,10 +2671,7 @@ def _save_energy_components(self, index, context, time_ns):
             The current simulation time in nanoseconds.
         """
 
-        import json as _json
         import openmm
-        import pandas as _pd
-        import pyarrow as _pa
         import pyarrow.parquet as _pq_local
 
         filepath = self._filenames[index]["energy_components"]
@@ -2701,8 +2709,37 @@ def _save_energy_components(self, index, context, time_ns):
                 openmm.unit.kilocalories_per_mole
             )
 
-        row = {"time": round(time_ns, 6)} | energies
-        df = _pd.DataFrame([row])
+        rows = self._ec_rows.setdefault(index, [])
+        rows.append({"time": round(time_ns, 6)} | energies)
+        self._last_ec_time[index] = time_ns
+
+        # Bound the memory used when checkpoints are rare or disabled.
+        if len(rows) >= self._max_ec_rows:
+            self._flush_energy_components(index)
+
+    def _flush_energy_components(self, index):
+        """
+        Write the energy components buffered by _save_energy_components() to
+        the Parquet file for a window.
+
+        Parameters
+        ----------
+
+        index : int
+            The index of the window or replica.
+        """
+
+        import json as _json
+        import pandas as _pd
+        import pyarrow as _pa
+        import pyarrow.parquet as _pq_local
+
+        rows = self._ec_rows.pop(index, [])
+        if not rows:
+            return
+
+        filepath = self._filenames[index]["energy_components"]
+        df = _pd.DataFrame(rows)
 
         path = _Path(filepath)
         if path.exists() and path.stat().st_size > 0:
@@ -2719,8 +2756,6 @@ def _save_energy_components(self, index, context, time_ns):
             )
             _pq_local.write_table(table, filepath)
 
-        self._last_ec_time[index] = time_ns
-
     def _restore_backup_files(self):
         """
         Restore backup files in the working directory.

tests/runner/test_energy_components.py

@@ -0,0 +1,131 @@
+from pathlib import Path
+import tempfile
+
+import pyarrow.parquet as pq
+import pytest
+
+from somd2.runner import Runner, RepexRunner
+from somd2.config import Config
+
+from tests.conftest import has_cuda
+
+
+def _config(tmpdir, runtime, platform, repex, restart=False):
+    return Config(
+        runtime=runtime,
+        restart=restart,
+        output_directory=tmpdir,
+        energy_frequency="4fs",
+        checkpoint_frequency="12fs",
+        frame_frequency="12fs",
+        platform=platform,
+        max_threads=1,
+        num_lambda=2,
+        replica_exchange=repex,
+        save_energy_components=True,
+    )
+
+
+def _times(tmpdir, lam):
+    table = pq.read_table(Path(tmpdir) / f"energy_components_{lam}.parquet")
+    return table.column("time").to_pylist()
+
+
+@pytest.mark.parametrize(
+    "runner_class, platform",
+    [
+        (Runner, "CPU"),
+        pytest.param(
+            RepexRunner,
+            "cuda",
+            marks=pytest.mark.skipif(not has_cuda, reason="CUDA not available."),
+        ),
+    ],
+)
+def test_energy_components_rows(ethane_methanol, runner_class, platform):
+    """
+    Validate that the energy components file holds one row per energy save,
+    including across a restart.
+    """
+    repex = runner_class is RepexRunner
+    with tempfile.TemporaryDirectory() as tmpdir:
+        runner = runner_class(ethane_methanol, _config(tmpdir, "12fs", platform, repex))
+        runner.run()
+
+        for lam in ("0.00000", "1.00000"):
+            times = _times(tmpdir, lam)
+            assert len(times) == 3
+            assert times == sorted(set(times))
+
+        runner = runner_class(
+            ethane_methanol, _config(tmpdir, "24fs", platform, repex, restart=True)
+        )
+        runner.run()
+
+        for lam in ("0.00000", "1.00000"):
+            times = _times(tmpdir, lam)
+            assert len(times) == 6
+            assert times == sorted(set(times))
+
+
+@pytest.mark.skipif(not has_cuda, reason="CUDA not available.")
+def test_energy_components_written_at_checkpoint(ethane_methanol):
+    """
+    Validate that energy components are buffered between checkpoints rather
+    than written on every energy save.
+    """
+    with tempfile.TemporaryDirectory() as tmpdir:
+        runner = RepexRunner(ethane_methanol, _config(tmpdir, "24fs", "cuda", True))
+
+        saves = []
+        flushes = []
+        save = runner._save_energy_components
+        flush = runner._flush_energy_components
+
+        def counting_save(index, context, time_ns):
+            saves.append(index)
+            return save(index, context, time_ns)
+
+        def counting_flush(index):
+            flushes.append(index)
+            return flush(index)
+
+        runner._save_energy_components = counting_save
+        runner._flush_energy_components = counting_flush
+
+        runner.run()
+
+        # Six energy saves and two checkpoints per replica.
+        assert saves.count(0) == 6
+        assert flushes.count(0) == 2
+        assert len(_times(tmpdir, "0.00000")) == 6
+
+
+@pytest.mark.skipif(not has_cuda, reason="CUDA not available.")
+def test_energy_components_buffer_limit(ethane_methanol):
+    """
+    Validate that the buffer is written out when it reaches its size limit,
+    so that memory is bounded when checkpoints are rare.
+    """
+    with tempfile.TemporaryDirectory() as tmpdir:
+        runner = RepexRunner(ethane_methanol, _config(tmpdir, "24fs", "cuda", True))
+        runner._max_ec_rows = 2
+
+        flushes = []
+        flush = runner._flush_energy_components
+
+        def counting_flush(index):
+            flushes.append(index)
+            return flush(index)
+
+        runner._flush_energy_components = counting_flush
+
+        runner.run()
+
+        # The limit is reached after the second and fifth saves, since the
+        # checkpoint after the third empties the buffer, plus one flush at each
+        # of the two checkpoints.
+        assert flushes.count(0) == 4
+        times = _times(tmpdir, "0.00000")
+        assert len(times) == 6
+        assert times == sorted(set(times))