Close a stage's LAMMPS session when the stage raises

calphy/alchemy.py

@@ -175,7 +175,7 @@ def _run_integration_coupling(self, iteration=1):
         """Coupling-mode integration: ramp ONLY the last component 0->1
         (forward) and 1->0 (backward) over the fixed base.
         W = int <U_last>_lam dlam;  dF = (W_f + W_b)/2 per iteration."""
-        lmp = ph.create_object(self.calc, self.simfolder)
+        lmp = self._open_lammps()
         conf = os.path.join(self.simfolder, "conf.equilibration.data")
         # style only before read_data (pair_coeff needs the box)
         ns = self.calc._pair_style_with_options
@@ -253,7 +253,7 @@ def run_averaging(self):
         Fix lattice option is not implemented at present.
         At the end of the run, the averaged box dimensions are calculated.
         """
-        lmp = ph.create_object(self.calc, self.simfolder)
+        lmp = self._open_lammps()
 
         if self.calc.alchemy_coupling:
             # equilibrate on the BASE system (lam = 0): full base, ramped
@@ -345,7 +345,7 @@ def run_integration(self, iteration=1):
         initial, final = self._end_states()
 
         # create lammps object
-        lmp = ph.create_object(self.calc, self.simfolder)
+        lmp = self._open_lammps()
 
         # Adiabatic switching parameters.
         lmp.command("variable        li       equal   1.0")

calphy/dcci.py

@@ -345,7 +345,8 @@ def equilibrate_cells(self):
         """
         def equilibrate(phase):
             self.logger.info("Equilibrating the %s cell at %.2f K, %.2f bar" % (phase, self.t0, self.p_start))
-            self.jobs[phase].run_averaging()
+            with self.jobs[phase]:
+                self.jobs[phase].run_averaging()
             self.logger.info("%s cell: %.4f A^3/atom" % (phase, self.jobs[phase].volatom))
 
         self._for_cells(equilibrate)

calphy/helpers.py

@@ -24,6 +24,7 @@
 """
 
 import os
+import sys
 import shutil
 import warnings
 import logging
@@ -96,7 +97,7 @@ def create_object(calc, directory):
             cmdargs=calc.md.cmdargs,
             directory=directory,
         )
-        return emit_init_commands(lmp, calc)
+        return _emit_init_commands_or_close(lmp, calc)
 
     binary = resolve_lammps_executable(calc.lammps_executable)
     mpi_command = (
@@ -113,7 +114,20 @@ def create_object(calc, directory):
         directory=directory,
         dry_run=False,
     )
-    return emit_init_commands(lmp, calc)
+    return _emit_init_commands_or_close(lmp, calc)
+
+
+def _emit_init_commands_or_close(lmp, calc):
+    """Run :func:`emit_init_commands`, closing ``lmp`` if it raises.
+
+    The caller only gets hold of the runner once :func:`create_object`
+    returns, so a failure here would otherwise leave the session (in library
+    mode, the executor behind it) open."""
+    try:
+        return emit_init_commands(lmp, calc)
+    except BaseException:
+        lmp.__exit__(*sys.exc_info())
+        raise
 
 
 def emit_init_commands(lmp, calc):

calphy/library_runner.py

@@ -37,6 +37,7 @@
 other part of calphy) never requires it.
 """
 import os
+import sys
 import logging
 
 from calphy.runner import BaseRunner, _normalize_cmdargs
@@ -81,7 +82,13 @@ def __init__(self, *, cores, cmdargs, directory):
         self.lmp = LammpsLibrary(
             cores=cores, working_directory=directory, cmdargs=cmdargs
         )
-        self._activate_mliap()
+        # the executor is live from here on; nobody else holds a reference to
+        # this runner yet, so close it ourselves if the rest of setup fails
+        try:
+            self._activate_mliap()
+        except BaseException:
+            self.__exit__(*sys.exc_info())
+            raise
 
     def _activate_mliap(self):
         """Register the mliappy coupling in the live session when available.

calphy/liquid.py

@@ -162,7 +162,7 @@ def run_averaging(self):
         At the end of the run, the averaged box dimensions are calculated.
         """
         # create lammps object
-        lmp = ph.create_object(self.calc, self.simfolder)
+        lmp = self._open_lammps()
 
         lmp = ph.set_pair_style(lmp, self.calc)
 
@@ -237,7 +237,7 @@ def run_integration(self, iteration=1):
         Run the integration routine where the initial and final systems are connected using
         the lambda parameter. See algorithm 4 in publication.
         """
-        lmp = ph.create_object(self.calc, self.simfolder)
+        lmp = self._open_lammps()
 
         # Adiabatic switching parameters.
         lmp.command("variable        li       equal   1.0")

calphy/phase.py

@@ -1497,7 +1497,7 @@ def _reversible_scaling_forward(self, iteration: int = 1) -> None:
             iteration, t0, tf, li, lf, pi, pf,
         )
 
-        lmp = ph.create_object(self.calc, self.simfolder)
+        lmp = self._open_lammps()
 
         lmp.command("echo              log")
         lmp.command("variable          li equal %f" % li)
@@ -1705,7 +1705,7 @@ def _reversible_scaling_backward(self, iteration: int = 1) -> None:
             iteration, tf, t0, lf, li, pf, pi,
         )
 
-        lmp = ph.create_object(self.calc, self.simfolder)
+        lmp = self._open_lammps()
 
         lmp.command("echo              log")
         lmp.command("variable          li equal %f" % li)
@@ -1983,7 +1983,7 @@ def scan_temperature_range(self):
         )
 
         # ── Build the LAMMPS object and load the equilibrated configuration ──
-        lmp = ph.create_object(self.calc, self.simfolder)
+        lmp = self._open_lammps()
 
         lmp.command("echo              log")
         lmp = ph.set_pair_style(lmp, self.calc)
@@ -2173,7 +2173,7 @@ def temperature_scaling(self, iteration=1):
         p0 = self.calc._pressure
 
         # create lammps object
-        lmp = ph.create_object(self.calc, self.simfolder)
+        lmp = self._open_lammps()
 
         lmp.command("echo              log")
         lmp.command("variable          li equal %f" % li)
@@ -2323,7 +2323,7 @@ def pressure_scaling(self, iteration=1):
         pf = self.calc._pressure_stop
 
         # create lammps object
-        lmp = ph.create_object(self.calc, self.simfolder)
+        lmp = self._open_lammps()
 
         lmp.command("echo              log")
         lmp.command("variable          li equal %f" % li)
@@ -2482,5 +2482,25 @@ def clean_up(self):
         with open(os.path.join(self.simfolder, "metadata.yaml"), "w") as fout:
             yaml.safe_dump(metadata, fout)
 
+    def _open_lammps(self):
+        """Create a stage's LAMMPS session and remember it, so that a
+        surrounding ``with job:`` can close it if the stage raises."""
+        self._lmp = ph.create_object(self.calc, self.simfolder)
+        return self._lmp
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc, tb):
+        """Close the session opened by the last stage, exception or not.
+
+        On success the stage has already closed it itself; closing again is a
+        no-op.  The exception is never suppressed.
+        """
+        lmp, self._lmp = getattr(self, "_lmp", None), None
+        if lmp is not None:
+            lmp.__exit__(exc_type, exc, tb)
+        return False
+
     def lammps_close(self, lmp):
         lmp.close()

calphy/routines.py

@@ -271,7 +271,8 @@ def run_jobs(self):
         self.logger.info("Starting solid reversible scaling run")
         for i in range(self.soljob.calc.n_iterations):
             try:
-                self.soljob.reversible_scaling(iteration=(i + 1))
+                with self.soljob:
+                    self.soljob.reversible_scaling(iteration=(i + 1))
             except MeltedError:
                 self.logger.info("Solid system melted during reversible scaling run")
                 return 2
@@ -290,7 +291,8 @@ def run_jobs(self):
         self.logger.info("Starting liquid reversible scaling calculation")
         for i in range(self.lqdjob.calc.n_iterations):
             try:
-                self.lqdjob.reversible_scaling(iteration=(i + 1))
+                with self.lqdjob:
+                    self.lqdjob.reversible_scaling(iteration=(i + 1))
             except SolidifiedError:
                 self.logger.info("Liquid froze during reversible scaling calculation")
                 return 3
@@ -527,14 +529,16 @@ def routine_fe(job):
     Perform an FE calculation routine
     """
     ts = time.time()
-    job.run_averaging()
+    with job:
+        job.run_averaging()
     te = time.time() - ts
     job.logger.info("Averaging routine finished in %f s" % te)
 
     # now run integration loops
     for i in range(job.calc.n_iterations):
         ts = time.time()
-        job.run_integration(iteration=(i + 1))
+        with job:
+            job.run_integration(iteration=(i + 1))
         te = time.time() - ts
         job.logger.info("Integration cycle %d finished in %f s" % (i + 1, te))
 
@@ -556,14 +560,16 @@ def routine_ts(job):
     # in the single-phase region.  Disabled by default (mode='none').
     if job.calc.phase_transition_detection.mode != "none":
         ts = time.time()
-        job.scan_temperature_range()
+        with job:
+            job.scan_temperature_range()
         te = time.time() - ts
         job.logger.info("Pre-flight temperature-range scan finished in %f s" % te)
 
     # now do rev scale steps
     for i in range(job.calc.n_iterations):
         ts = time.time()
-        job.reversible_scaling(iteration=(i + 1))
+        with job:
+            job.reversible_scaling(iteration=(i + 1))
         te = time.time() - ts
         job.logger.info("TS integration cycle %d finished in %f s" % (i + 1, te))
 
@@ -581,7 +587,8 @@ def routine_tscale(job):
     # now do rev scale steps
     for i in range(job.calc.n_iterations):
         ts = time.time()
-        job.temperature_scaling(iteration=(i + 1))
+        with job:
+            job.temperature_scaling(iteration=(i + 1))
         te = time.time() - ts
         job.logger.info("Temperature scaling cycle %d finished in %f s" % (i + 1, te))
 
@@ -599,7 +606,8 @@ def routine_pscale(job):
     # now do rev scale steps
     for i in range(job.calc.n_iterations):
         ts = time.time()
-        job.pressure_scaling(iteration=(i + 1))
+        with job:
+            job.pressure_scaling(iteration=(i + 1))
         te = time.time() - ts
         job.logger.info("Pressure scaling cycle %d finished in %f s" % (i + 1, te))
 
@@ -613,14 +621,16 @@ def routine_alchemy(job):
     Perform an FE calculation routine
     """
     ts = time.time()
-    job.run_averaging()
+    with job:
+        job.run_averaging()
     te = time.time() - ts
     job.logger.info("Averaging routine finished in %f s" % te)
 
     # now run integration loops
     for i in range(job.calc.n_iterations):
         ts = time.time()
-        job.run_integration(iteration=(i + 1))
+        with job:
+            job.run_integration(iteration=(i + 1))
         te = time.time() - ts
         job.logger.info("Alchemy integration cycle %d finished in %f s" % (i + 1, te))
 
@@ -742,14 +752,16 @@ def routine_composition_scaling(job):
 
     # now start cycle
     ts = time.time()
-    job.run_averaging()
+    with job:
+        job.run_averaging()
     te = time.time() - ts
     job.logger.info("Averaging routine finished in %f s" % te)
 
     # now run integration loops
     for i in range(job.calc.n_iterations):
         ts = time.time()
-        job.run_integration(iteration=(i + 1))
+        with job:
+            job.run_integration(iteration=(i + 1))
         te = time.time() - ts
         job.logger.info("Alchemy integration cycle %d finished in %f s" % (i + 1, te))
 

calphy/runner.py

@@ -366,6 +366,34 @@ def close(self):
         """End the session and release backend resources."""
         raise NotImplementedError
 
+    def __enter__(self):
+        """Enter a runner scope; the runner itself is the ``as`` target."""
+        return self
+
+    def __exit__(self, exc_type, exc, tb):
+        """Close the session on the way out, exception or not.
+
+        This is what makes a stage's LAMMPS session -- and, in library mode,
+        the executorlib executor behind it -- go away deterministically even
+        when an exception escapes the driver.  The driver's exception is never
+        suppressed.
+
+        If :meth:`close` itself fails *while an exception is propagating*, the
+        close error is logged and swallowed rather than replacing the original
+        one: the interesting failure is the physics error that got us here, not
+        the fallout of shutting a half-dead session down.
+        """
+        try:
+            self.close()
+        except Exception:
+            if exc_type is None:
+                raise
+            logger.exception(
+                "closing the LAMMPS session failed while handling %s; the "
+                "original exception is propagating", exc_type.__name__,
+            )
+        return False
+
     def rotate_logs(self, stage_name):
         """Collect the LAMMPS log output produced since the last rotation into
         ``<stage_name>.log.lammps`` in the sim folder."""

calphy/solid.py

@@ -192,7 +192,7 @@ def run_averaging(self):
         At the end of the run, the averaged box dimensions are calculated.
         """
 
-        lmp = ph.create_object(self.calc, self.simfolder)
+        lmp = self._open_lammps()
 
         # set up potential
         lmp = ph.set_pair_style(lmp, self.calc)
@@ -268,7 +268,7 @@ def run_integration(self, iteration=1):
         Run the integration routine where the initial and final systems are connected using
         the lambda parameter. See algorithm 4 in publication.
         """
-        lmp = ph.create_object(self.calc, self.simfolder)
+        lmp = self._open_lammps()
 
         # set up potential
         lmp = ph.set_pair_style(lmp, self.calc)