Merge pull request #5021 from epretti/safe-save

Always safely overwrite checkpoint and state files

wrappers/python/openmm/app/checkpointreporter.py

@@ -32,10 +32,6 @@ from __future__ import absolute_import
 __author__ = "Robert McGibbon"
 __version__ = "1.0"
 
-import openmm as mm
-import os
-import os.path
-
 __all__ = ['CheckpointReporter']
 
 
@@ -126,29 +122,15 @@ class CheckpointReporter(object):
         state : State
             The current state of the simulation
         """
-        if isinstance(self._file, str):
-            # Do a safe save.
+        isFileObj = not isinstance(self._file, str)
+        if isFileObj:
+            self._file.seek(0)
 
-            tempFilename1 = self._file+".backup1"
-            tempFilename2 = self._file+".backup2"
         if self._writeState:
-                simulation.saveState(tempFilename1)
-            else:
-                simulation.saveCheckpoint(tempFilename1)
-            exists = os.path.exists(self._file)
-            if exists:
-                os.rename(self._file, tempFilename2)
-            os.rename(tempFilename1, self._file)
-            if exists:
-                os.remove(tempFilename2)
+            simulation.saveState(self._file)
         else:
-            # Replace the contents of the file.
+            simulation.saveCheckpoint(self._file)
 
-            self._file.seek(0)
-            if self._writeState:
-                state = simulation.context.getState(positions=True, velocities=True, parameters=True, integratorParameters=True)
-                self._file.write(mm.XmlSerializer.serialize(state))
-            else:
-                self._file.write(simulation.context.createCheckpoint())
+        if isFileObj:
             self._file.truncate()
             self._file.flush()

wrappers/python/openmm/app/internal/safesave.py

+"""
+safesave.py: Helper module to ensure atomic overwrite/backup of existing files.
+
+This is part of the OpenMM molecular simulation toolkit originating from
+Simbios, the NIH National Center for Physics-Based Simulation of
+Biological Structures at Stanford, funded under the NIH Roadmap for
+Medical Research, grant U54 GM072970. See https://simtk.org.
+
+Portions copyright (c) 2025 Stanford University and the Authors.
+Authors: Evan Pretti
+Contributors:
+
+Permission is hereby granted, free of charge, to any person obtaining a
+copy of this software and associated documentation files (the "Software"),
+to deal in the Software without restriction, including without limitation
+the rights to use, copy, modify, merge, publish, distribute, sublicense,
+and/or sell copies of the Software, and to permit persons to whom the
+Software is furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in
+all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
+THE AUTHORS, CONTRIBUTORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM,
+DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR
+OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE
+USE OR OTHER DEALINGS IN THE SOFTWARE.
+"""
+
+import itertools
+import os
+
+def _getTempFilename(prefix):
+    """
+    Returns the name of a temporary file starting with a given prefix that is
+    guaranteed not to exist already.  Upon successful return of this function,
+    an empty file with the name returned will have been created.
+
+    Parameters
+    ----------
+    prefix : str
+        The prefix of the temporary file name to create.  If a path with
+        multiple components, the directory in which to create the file must
+        exist.
+
+    Returns
+    -------
+    str
+        The temporary file name created.
+    """
+
+    for index in itertools.count():
+        name = f'{prefix}.{index}.tmp'
+        try:
+            with open(name, 'x'):
+                return name
+        except FileExistsError:
+            pass
+
+def save(data, filename):
+    """
+    Saves data to a specified file.  If the file exists, it will be overwritten
+    atomically, or if this is not possible, a backup copy of the existing data
+    will be created during overwriting and deleted once it is successful.
+
+    Parameters
+    ----------
+    data : bytes or str
+        The data to write.  If bytes, the file will be opened in binary mode; if
+        str, in text mode.
+    filename : str
+        The filename to write to.
+    """
+
+    if isinstance(data, bytes):
+        mode = 'wb'
+    elif isinstance(data, str):
+        mode = 'w'
+    else:
+        raise ValueError('Expected bytes or str')
+
+    tempFilename1 = _getTempFilename(filename)
+    with open(tempFilename1, mode) as file:
+        file.write(data)
+
+    try:
+        # If the target file already exists, rename() should overwrite
+        # atomically on POSIX and fail with a FileExistsError on Windows.
+        os.rename(tempFilename1, filename)
+    except FileExistsError:
+        # Make a backup copy since replace() on Windows may not be atomic.
+        tempFilename2 = _getTempFilename(filename)
+        os.replace(filename, tempFilename2)
+        os.replace(tempFilename1, filename)
+        os.remove(tempFilename2)

wrappers/python/openmm/app/simulation.py

@@ -34,6 +34,7 @@ __version__ = "1.0"
 
 import openmm as mm
 import openmm.unit as unit
+from openmm.app.internal import safesave
 import sys
 from datetime import datetime, timedelta
 try:
@@ -300,8 +301,7 @@ class Simulation(object):
             filename
         """
         if isinstance(file, str):
-            with open(file, 'wb') as f:
-                f.write(self.context.createCheckpoint())
+            safesave.save(self.context.createCheckpoint(), file)
         else:
             file.write(self.context.createCheckpoint())
 
@@ -341,8 +341,7 @@ class Simulation(object):
         state = self.context.getState(positions=True, velocities=True, parameters=True, integratorParameters=True)
         xml = mm.XmlSerializer.serialize(state)
         if isinstance(file, str):
-            with open(file, 'w') as f:
-                f.write(xml)
+            safesave.save(xml, file)
         else:
             file.write(xml)
 

wrappers/python/tests/TestCheckpointReporter.py

 import os
 import unittest
 import tempfile
+from io import BytesIO, StringIO
 from openmm import app
 import openmm as mm
 from openmm import unit
@@ -41,5 +42,32 @@ class TestCheckpointReporter(unittest.TestCase):
                 newPositions = self.simulation.context.getState(getPositions=True).getPositions()
                 self.assertSequenceEqual(positions, newPositions)
 
+    def testFileObj(self):
+        """Test writing to a file object.  This should truncate so that only the most recent frame is present in the output."""
+
+        # Test checkpoint saving.
+
+        checkpointBuffer = BytesIO()
+        self.simulation.reporters.clear()
+        self.simulation.reporters.append(app.CheckpointReporter(checkpointBuffer, 1, writeState=False))
+        self.simulation.step(5)
+        checkpointData = checkpointBuffer.getvalue()
+
+        checkpointBuffer = BytesIO()
+        self.simulation.saveCheckpoint(checkpointBuffer)
+        self.assertSequenceEqual(checkpointData, checkpointBuffer.getvalue())
+
+        # Test state saving.
+
+        stateBuffer = StringIO()
+        self.simulation.reporters.clear()
+        self.simulation.reporters.append(app.CheckpointReporter(stateBuffer, 1, writeState=True))
+        self.simulation.step(5)
+        stateData = stateBuffer.getvalue()
+
+        stateBuffer = StringIO()
+        self.simulation.saveState(stateBuffer)
+        self.assertSequenceEqual(stateData, stateBuffer.getvalue())
+
 if __name__ == '__main__':
     unittest.main()

wrappers/python/tests/TestSimulation.py

 import unittest
 import tempfile
 from datetime import datetime, timedelta
+from io import BytesIO, StringIO
 from openmm import *
 from openmm.app import *
 from openmm.unit import *
@@ -98,6 +99,77 @@ class TestSimulation(unittest.TestCase):
         self.assertEqual(initialState.getPositions(), state.getPositions())
         self.assertEqual(initialState.getVelocities(), state.getVelocities())
 
+    def testSafeSave(self):
+        """Test that the safe saving feature works as expected."""
+        pdb = PDBFile('systems/alanine-dipeptide-implicit.pdb')
+        ff = ForceField('amber99sb.xml', 'tip3p.xml')
+        system = ff.createSystem(pdb.topology)
+        integrator = VerletIntegrator(0.001*picoseconds)
+
+        # Create a Simulation.
+
+        simulation = Simulation(pdb.topology, system, integrator, Platform.getPlatform('Reference'))
+        simulation.context.setPositions(pdb.positions)
+        simulation.context.setVelocitiesToTemperature(300*kelvin)
+
+        # Get reference checkpoint and state data.
+
+        checkpointBuffer = BytesIO()
+        simulation.saveCheckpoint(checkpointBuffer)
+        checkpointData = checkpointBuffer.getvalue()
+
+        stateBuffer = StringIO()
+        simulation.saveState(stateBuffer)
+        stateData = stateBuffer.getvalue()
+
+        # Try a safe save of a checkpoint.
+
+        with tempfile.TemporaryDirectory() as directory:
+            tempPath = os.path.join(directory, 'testSafeSaveCheckpoint.dat')
+
+            # Make a file that should get overwritten by the safe save, and some that shouldn't.
+
+            with open(tempPath, 'w') as testFile:
+                testFile.write('Test')
+            with open(f'{tempPath}.0.tmp', 'w') as testFile:
+                testFile.write('Test0')
+            with open(f'{tempPath}.1.tmp', 'w') as testFile:
+                testFile.write('Test1')
+
+            # Perform and verify the safe save and that the contents of the test files were not overwritten.
+
+            simulation.saveCheckpoint(tempPath)
+            with open(tempPath, 'rb') as checkpointFile:
+                self.assertSequenceEqual(checkpointData, checkpointFile.read())
+            with open(f'{tempPath}.0.tmp', 'r') as testFile:
+                self.assertSequenceEqual('Test0', testFile.read())
+            with open(f'{tempPath}.1.tmp', 'r') as testFile:
+                self.assertSequenceEqual('Test1', testFile.read())
+
+        # Try a safe save of a state.
+
+        with tempfile.TemporaryDirectory() as directory:
+            tempPath = os.path.join(directory, 'testSafeSaveState.dat')
+
+            # Make a file that should get overwritten by the safe save, and some that shouldn't.
+
+            with open(tempPath, 'w') as testFile:
+                testFile.write('Test')
+            with open(f'{tempPath}.0.tmp', 'w') as testFile:
+                testFile.write('Test0')
+            with open(f'{tempPath}.1.tmp', 'w') as testFile:
+                testFile.write('Test1')
+
+            # Perform and verify the safe save and that the contents of the test files were not overwritten.
+
+            simulation.saveState(tempPath)
+            with open(tempPath, 'r') as stateFile:
+                self.assertSequenceEqual(stateData, stateFile.read())
+            with open(f'{tempPath}.0.tmp', 'r') as testFile:
+                self.assertSequenceEqual('Test0', testFile.read())
+            with open(f'{tempPath}.1.tmp', 'r') as testFile:
+                self.assertSequenceEqual('Test1', testFile.read())
+
     def testStep(self):
         """Test the step() method."""
         pdb = PDBFile('systems/alanine-dipeptide-implicit.pdb')