From 7746f23f6677b384c8ab75902501ba841de433bc Mon Sep 17 00:00:00 2001 From: Marvin Poul Date: Thu, 8 Oct 2026 16:43:07 -0400 Subject: [PATCH 1/2] fix: close a stage's LAMMPS session when the stage raises Every stage method (run_averaging, run_integration, reversible_scaling, temperature_scaling, pressure_scaling, scan_temperature_range) creates its own LAMMPS session and closes it as its last statement, so an exception anywhere in the stage leaves the session open. In library mode that keeps the pylammpsmpi/executorlib workers alive until the process exits. `BaseRunner` gains `__enter__`/`__exit__`, and `Phase` gains a context that closes the session opened by the stage it wraps. The routines now run each stage call as `with job: job.stage(...)`. Stage bodies are untouched apart from creating the session via `self._open_lammps()`, which remembers it. On success the stage has already closed the session and the second close is a no-op on both backends (`LibraryRunner` guards on `_closed`; the executable runner's buffer is already empty and the restart-file removal is guarded). If `close()` itself fails while an exception is propagating, the close error is logged and the original exception wins. Co-Authored-By: Claude Opus 5.5 (1M context) --- calphy/alchemy.py | 6 ++-- calphy/dcci.py | 3 +- calphy/liquid.py | 4 +-- calphy/phase.py | 30 +++++++++++++++--- calphy/routines.py | 36 ++++++++++++++-------- calphy/runner.py | 28 +++++++++++++++++ calphy/solid.py | 4 +-- tests/test_runner.py | 73 ++++++++++++++++++++++++++++++++++++++++++++ 8 files changed, 159 insertions(+), 25 deletions(-) diff --git a/calphy/alchemy.py b/calphy/alchemy.py index eb29495..3fab12c 100644 --- a/calphy/alchemy.py +++ b/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 _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") diff --git a/calphy/dcci.py b/calphy/dcci.py index f5e91b4..3983679 100644 --- a/calphy/dcci.py +++ b/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) diff --git a/calphy/liquid.py b/calphy/liquid.py index 4f323df..eedc8e3 100644 --- a/calphy/liquid.py +++ b/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") diff --git a/calphy/phase.py b/calphy/phase.py index af04d49..0490a38 100644 --- a/calphy/phase.py +++ b/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() diff --git a/calphy/routines.py b/calphy/routines.py index 34a87b5..7d86162 100644 --- a/calphy/routines.py +++ b/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)) diff --git a/calphy/runner.py b/calphy/runner.py index 7bd4922..ba12dab 100644 --- a/calphy/runner.py +++ b/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 ``.log.lammps`` in the sim folder.""" diff --git a/calphy/solid.py b/calphy/solid.py index 8b02156..f0a9d62 100644 --- a/calphy/solid.py +++ b/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) diff --git a/tests/test_runner.py b/tests/test_runner.py index 7914629..7b3d88b 100644 --- a/tests/test_runner.py +++ b/tests/test_runner.py @@ -633,3 +633,76 @@ def test_imports_without_pylammpsmpi(): text=True, cwd=REPO_ROOT) assert proc.returncode == 0, proc.stdout + proc.stderr assert "OK" in proc.stdout + + +# --------------------------------------------------------------------------- # +# BaseRunner as a context manager +# --------------------------------------------------------------------------- # +class _RecordingRunner(R.BaseRunner): + """Minimal BaseRunner subclass that records close() calls.""" + + def __init__(self, directory=".", close_error=None): + super().__init__(directory) + self.closes = 0 + self.close_error = close_error + + def close(self): + self.closes += 1 + if self.close_error is not None: + raise self.close_error + + +def test_context_manager_yields_the_runner_and_closes(): + runner = _RecordingRunner() + with runner as lmp: + assert lmp is runner + assert runner.closes == 0 + assert runner.closes == 1 + + +def test_context_manager_closes_on_exception_and_does_not_suppress(): + runner = _RecordingRunner() + with pytest.raises(ValueError, match="physics"): + with runner: + raise ValueError("physics went wrong") + assert runner.closes == 1 + + +def test_close_failure_propagates_when_no_exception_is_in_flight(): + runner = _RecordingRunner(close_error=RuntimeError("close failed")) + with pytest.raises(RuntimeError, match="close failed"): + with runner: + pass + + +def test_close_failure_does_not_mask_the_original_exception(): + """A dead LAMMPS worker can make close() fail; the physics error must win.""" + runner = _RecordingRunner(close_error=RuntimeError("worker already dead")) + with pytest.raises(ValueError, match="physics"): + with runner: + raise ValueError("physics went wrong") + assert runner.closes == 1 + + +def test_phase_context_closes_the_session_of_a_failed_stage(monkeypatch): + """``with job: job.stage()`` closes the session the stage opened if it raises.""" + import calphy.phase as cph + + runner = _RecordingRunner() + monkeypatch.setattr(cph.ph, "create_object", lambda calc, simfolder: runner) + job = cph.Phase.__new__(cph.Phase) + job.calc, job.simfolder = None, "." + + def stage(): + job._open_lammps() + raise ValueError("physics went wrong") + + with pytest.raises(ValueError, match="physics"): + with job: + stage() + assert runner.closes == 1 + + # a context with no stage in it is a no-op + with job: + pass + assert runner.closes == 1 From a0f5bc422f421158eca80285d5876c56602e782f Mon Sep 17 00:00:00 2001 From: Marvin Poul Date: Thu, 8 Oct 2026 17:11:25 -0400 Subject: [PATCH 2/2] fix: close the runner if session setup raises `Phase._open_lammps()` only learns about the runner once `ph.create_object()` returns, so a failure while setting the session up left it open: `create_object` emits the init commands after constructing the runner, and `LibraryRunner.__init__` activates ML-IAP after the executor already exists. Close the runner in both places before re-raising, through its `__exit__` so a failing close cannot mask the setup error. Co-Authored-By: Claude Opus 5.5 (1M context) --- calphy/helpers.py | 18 ++++++++++++++-- calphy/library_runner.py | 9 +++++++- tests/test_library_runner.py | 41 ++++++++++++++++++++++++++++++++++++ 3 files changed, 65 insertions(+), 3 deletions(-) diff --git a/calphy/helpers.py b/calphy/helpers.py index bf111d2..91b6977 100644 --- a/calphy/helpers.py +++ b/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): diff --git a/calphy/library_runner.py b/calphy/library_runner.py index 1d3ee26..414b25b 100644 --- a/calphy/library_runner.py +++ b/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. diff --git a/tests/test_library_runner.py b/tests/test_library_runner.py index 3de8a8c..c74e6dd 100644 --- a/tests/test_library_runner.py +++ b/tests/test_library_runner.py @@ -254,6 +254,47 @@ def test_create_object_library_mode(tmp_path, fake_pylammpsmpi, make_calc): assert len(lmp.logical_commands) == 4 +@pytest.fixture +def recorded_libraries(fake_pylammpsmpi, monkeypatch): + """Every FakeLammpsLibrary created, so a test can check it got closed.""" + created = [] + + class Recording(FakeLammpsLibrary): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + created.append(self) + + monkeypatch.setattr(fake_pylammpsmpi, "LammpsLibrary", Recording) + return created + + +def test_failed_mliap_activation_closes_the_session( + tmp_path, recorded_libraries, fake_lammps_mliap, monkeypatch +): + def boom(self): + raise RuntimeError("activation failed") + + monkeypatch.setattr(_FakeConcurrent, "activate_mliappy", boom) + with pytest.raises(RuntimeError, match="activation failed"): + make_runner(tmp_path) + assert [lib.closed for lib in recorded_libraries] == [True] + + +def test_create_object_closes_the_session_if_init_commands_fail( + tmp_path, recorded_libraries, make_calc, monkeypatch +): + import calphy.helpers as ph + + def boom(lmp, calc): + raise ValueError("bad init command") + + monkeypatch.setattr(ph, "emit_init_commands", boom) + calc = make_calc("B1", execution_mode="library") + with pytest.raises(ValueError, match="bad init command"): + ph.create_object(calc, str(tmp_path)) + assert [lib.closed for lib in recorded_libraries] == [True] + + def test_execution_mode_validation(make_calc): with pytest.raises(Exception, match="execution_mode"): make_calc("B1", execution_mode="banana")