fix: stop a crashed run poisoning every later run of the same search

autofit/non_linear/fitness.py

@@ -633,6 +633,34 @@ def check_log_likelihood(self, fitness):
             samples_summary = self.paths.load_samples_summary()
         except FileNotFoundError:
             return
+        except ValueError as e:
+            # A CORRUPT previous summary means the same thing as an ABSENT one
+            # for this check: there is no trustworthy old likelihood to compare
+            # against. Returning early is what the FileNotFoundError branch
+            # above already does for the no-previous-run case.
+            #
+            # `ValueError` is the catch because `json.JSONDecodeError`
+            # subclasses it -- which is exactly why this was missed by both the
+            # `FileNotFoundError` above and the `(FileNotFoundError, TypeError,
+            # KeyError)` guard on the multi-start resume path. A half-written
+            # file therefore aborted the whole run from inside an OPTIONAL
+            # sanity check -- and it stayed aborted on every subsequent run of
+            # the same search name, since nothing rewrites the file until a run
+            # gets far enough to finish.
+            #
+            # Warned rather than passed over in silence: an unreadable file is
+            # a real event, unlike a missing one, and the user is the only one
+            # who can decide whether the old results mattered.
+            logger.warning(
+                f"Could not read the previous samples summary while resuming "
+                f"({type(e).__name__}: {e}). It is missing or corrupt, most "
+                f"likely because an earlier run of this search was interrupted "
+                f"while writing its output. The likelihood-function sanity "
+                f"check is being SKIPPED for this run, and results are being "
+                f"recomputed. Delete the search's output directory if you want "
+                f"a guaranteed-clean start."
+            )
+            return
 
         try:
             max_log_likelihood_sample = samples_summary.max_log_likelihood_sample

autofit/non_linear/paths/directory.py

@@ -13,7 +13,7 @@
 from autonerves.dictable import to_dict, from_dict
 from autonerves.output import conditional_output, should_output
 from autofit.text import formatter
-from autofit.tools.util import open_, NumpyEncoder
+from autofit.tools.util import open_, open_atomic, NumpyEncoder
 from autofit.non_linear.samples.samples import Samples
 
 from .abstract import AbstractPaths, _test_mode_segment
@@ -76,11 +76,20 @@ def save_json(self, name, object_dict: Union[dict, list], prefix: str = ""):
         prefix
             A prefix to add to the path which is the name of the folder the file is saved in.
         """
-        # ``NumpyEncoder``: a stray ``np.float32`` (or any NumPy scalar that is
-        # not a ``float64``) would otherwise raise ``TypeError`` here, at the
-        # very end of a successful fit, throwing the whole run away at its
-        # output step. See the encoder's docstring for why float64 hid this.
-        with open_(self._path_for_json(name, prefix), "w+") as f:
+        # Two guards, and they answer different halves of the same incident.
+        #
+        # ``NumpyEncoder`` stops a stray ``np.float32`` (or any NumPy scalar
+        # that is not a ``float64``) raising ``TypeError`` here at the very end
+        # of a successful fit -- see the encoder's docstring for why float64
+        # hid this for so long.
+        #
+        # ``open_atomic`` handles the write failing for ANY reason: a plain
+        # "w+" truncates before it writes, so a failure partway leaves a
+        # HALF-WRITTEN file where a valid one was, and the next run of this
+        # search reads that while resuming and dies on it. Written to a sibling
+        # temp file and ``os.replace``d into place, a failed write leaves the
+        # previous file whole instead.
+        with open_atomic(self._path_for_json(name, prefix)) as f:
             json.dump(object_dict, f, indent=4, cls=NumpyEncoder)
 
     def load_json(self, name, prefix: str = ""):
@@ -196,7 +205,11 @@ def save_search_internal(self, obj):
         """
         filename = self.search_internal_path / "search_internal.dill"
 
-        with open_(filename, "wb") as f:
+        # Atomic for the same reason as ``save_json``, and it matters more
+        # here: ``search_internal`` is what a resumed run restores its step
+        # count and counters from, so a truncated dill does not merely fail to
+        # load -- it is the file the resume path most depends on.
+        with open_atomic(filename, "wb") as f:
             dill.dump(obj, f)
 
     def load_search_internal(self):

autofit/non_linear/search/mle/multi_start_gradient/search.py

@@ -1,4 +1,5 @@
 import inspect
+import pickle
 from typing import Optional
 
 import numpy as np
@@ -717,7 +718,32 @@ def batched_value_and_grad(params):
                 "Resuming MultiStartGradient search (previous samples found)."
             )
 
-        except (FileNotFoundError, TypeError, KeyError):
+        # ``EOFError``/``UnpicklingError``/``ValueError`` join the original
+        # three so that a CORRUPT ``search_internal`` starts a fresh run rather
+        # than killing this one. A dill truncated by an interrupted write is
+        # the same situation as an absent one -- there is no state to resume
+        # from -- but it raised out of this guard instead of falling into the
+        # fresh-start branch below, and it kept doing so on every rerun of the
+        # same search name, because nothing rewrites the file until a run
+        # finishes. ``save_search_internal`` is atomic now, so this is the
+        # belt to that braces: it also covers files left by older versions,
+        # killed processes and full disks.
+        except (
+            FileNotFoundError,
+            TypeError,
+            KeyError,
+            EOFError,
+            ValueError,
+            pickle.UnpicklingError,
+        ) as e:
+
+            if not isinstance(e, FileNotFoundError):
+                self.logger.warning(
+                    f"Could not restore the previous MultiStartGradient state "
+                    f"({type(e).__name__}: {e}); starting a fresh run. This "
+                    f"usually means an earlier run was interrupted while "
+                    f"writing its output."
+                )
 
             if not self.silence:
                 self.logger.info(self._compile_message(batched=False))

autofit/tools/util.py

@@ -90,6 +90,42 @@ def open_(filename, *flags):
     return open(filename, *flags)
 
 
+@contextmanager
+def open_atomic(filename, mode: str = "w+"):
+    """
+    Write ``filename`` so that it is either fully replaced or left untouched.
+
+    A plain ``open(path, "w+")`` **truncates first and writes second**. If the
+    write then fails partway -- a serialisation ``TypeError``, a full disk, a
+    killed process -- what is left on disk is a half-written file where a valid
+    one used to be. That truncated file is not inert: the next run of the same
+    search reads it while resuming and fails on it, so one crash propagates
+    into every subsequent run of that search until the output directory is
+    deleted by hand.
+
+    Writing to a temporary file in the *same directory* and then ``os.replace``
+    -ing it into place removes the window. ``os.replace`` is atomic on POSIX
+    and on Windows, and staying in one directory keeps it a rename within a
+    single filesystem, which is where that guarantee holds. On failure the
+    temporary file is removed and the original is still whole.
+    """
+    path = Path(filename)
+    os.makedirs(path.parent, exist_ok=True)
+
+    tmp = path.with_name(f"{path.name}.{os.getpid()}.tmp")
+
+    try:
+        with open(tmp, mode) as f:
+            yield f
+    except BaseException:
+        # BaseException, not Exception: a KeyboardInterrupt mid-write leaves
+        # exactly the same debris and must not leak a .tmp file either.
+        tmp.unlink(missing_ok=True)
+        raise
+
+    os.replace(tmp, path)
+
+
 @contextmanager
 def suppress_stdout():
     with open(os.devnull, "w") as devnull:

test_autofit/tools/test_atomic_write.py

@@ -0,0 +1,176 @@
+import json
+
+import numpy as np
+import pytest
+
+import autofit as af
+from autofit.tools.util import open_atomic
+
+
+class TestOpenAtomic:
+    def test__successful_write_replaces_the_file(self, tmp_path):
+        path = tmp_path / "f.json"
+        path.write_text('{"old": 1}')
+
+        with open_atomic(path) as f:
+            json.dump({"new": 2}, f)
+
+        assert json.loads(path.read_text()) == {"new": 2}
+
+    def test__failed_write_leaves_the_original_intact(self, tmp_path):
+        """
+        The whole point. A plain ``open(path, "w+")`` truncates first, so a
+        failure partway through destroys the previous file.
+        """
+        path = tmp_path / "f.json"
+        path.write_text('{"old": 1}')
+
+        with pytest.raises(TypeError):
+            with open_atomic(path) as f:
+                f.write('{"partial": ')
+                raise TypeError("Object of type float32 is not JSON serializable")
+
+        assert json.loads(path.read_text()) == {"old": 1}
+
+    def test__failed_write_leaves_no_temp_file_behind(self, tmp_path):
+        path = tmp_path / "f.json"
+        path.write_text('{"old": 1}')
+
+        with pytest.raises(TypeError):
+            with open_atomic(path) as f:
+                f.write("junk")
+                raise TypeError
+
+        assert [p.name for p in tmp_path.iterdir()] == ["f.json"]
+
+    def test__keyboard_interrupt_also_cleans_up(self, tmp_path):
+        """
+        ``BaseException``, not ``Exception``: an interrupt mid-write leaves the
+        same debris and must not leak a .tmp file either.
+        """
+        path = tmp_path / "f.json"
+        path.write_text('{"old": 1}')
+
+        with pytest.raises(KeyboardInterrupt):
+            with open_atomic(path) as f:
+                f.write("junk")
+                raise KeyboardInterrupt
+
+        assert json.loads(path.read_text()) == {"old": 1}
+        assert [p.name for p in tmp_path.iterdir()] == ["f.json"]
+
+    def test__creates_a_file_that_did_not_exist(self, tmp_path):
+        path = tmp_path / "sub" / "f.json"
+
+        with open_atomic(path) as f:
+            json.dump({"a": 1}, f)
+
+        assert json.loads(path.read_text()) == {"a": 1}
+
+    def test__binary_mode(self, tmp_path):
+        """``save_search_internal`` writes dill, so binary must work too."""
+        path = tmp_path / "f.bin"
+
+        with open_atomic(path, "wb") as f:
+            f.write(b"\x00\x01")
+
+        assert path.read_bytes() == b"\x00\x01"
+
+
+class TestSaveJsonIsAtomic:
+    def test__a_failed_save_json_does_not_destroy_the_previous_file(
+        self, output_directory
+    ):
+        """
+        The field sequence: a run writes a good summary, a later run dies while
+        rewriting it, and the half-written file poisons every run after that.
+        """
+        paths = af.DirectoryPaths(name="atomic_save_json")
+        paths._identifier = "id"
+
+        paths.save_json(name="counters", object_dict={"clipped": 4})
+        assert paths.load_json(name="counters") == {"clipped": 4}
+
+        class Unserialisable:
+            pass
+
+        with pytest.raises(TypeError):
+            paths.save_json(
+                name="counters", object_dict={"clipped": Unserialisable()}
+            )
+
+        # Still readable, and still the OLD value -- not truncated.
+        assert paths.load_json(name="counters") == {"clipped": 4}
+
+
+class TestCorruptOutputDoesNotPoisonTheNextRun:
+    def test__truncated_summary_lets_the_next_run_proceed(
+        self, output_directory
+    ):
+        """
+        End-to-end, because this bug is only visible across two runs.
+
+        Before the fix the second run died with an opaque
+        ``JSONDecodeError: Expecting value: line 1 column 13`` raised from
+        inside an OPTIONAL likelihood sanity check, and kept dying on every
+        rerun of the same name.
+        """
+        from autofit import example
+
+        xvalues = np.arange(60)
+        truth = example.Gaussian(centre=30.0, normalization=25.0, sigma=8.0)
+        data = np.asarray(truth.model_data_from(xvalues=xvalues))
+        noise_map = np.full(60, 1.0)
+
+        def analysis():
+            instance = example.Analysis(data=data, noise_map=noise_map)
+            instance._use_jax = False
+            return instance
+
+        def model():
+            built = af.Model(example.Gaussian)
+            built.centre = af.UniformPrior(lower_limit=5.0, upper_limit=35.0)
+            built.normalization = af.UniformPrior(
+                lower_limit=10.0, upper_limit=40.0
+            )
+            built.sigma = af.UniformPrior(lower_limit=1.0, upper_limit=20.0)
+            return built
+
+        name = "corrupt_resume"
+
+        search = af.LBFGS(name=name, maxiter=10)
+        search.fit(model=model(), analysis=analysis())
+
+        out = search.paths.output_path
+
+        # The test config prunes output files after a run, so write the
+        # corrupt state directly rather than depending on what survived. This
+        # is the same half-written file a truncating "w+" leaves behind.
+        summary = out / "files" / "samples_summary.json"
+        summary.parent.mkdir(parents=True, exist_ok=True)
+        summary.write_text('{"partial": ')
+
+        # Not `.completed` by rglob -- the marker's location is the paths
+        # object's business, and an interrupted run never wrote it anyway.
+        search.paths._has_completed_path.unlink(missing_ok=True)
+
+        resumed = af.LBFGS(name=name, maxiter=10)
+        assert not resumed.paths.is_complete, (
+            "the second run must take the resume path, not be short-circuited "
+            "as already-complete -- otherwise this asserts nothing"
+        )
+
+        # The regression: before the fix this raised JSONDecodeError out of an
+        # optional sanity check, and did so on every rerun of the same name.
+        result = resumed.fit(model=model(), analysis=analysis())
+
+        assert result is not None
+
+    def test__json_decode_error_is_a_value_error(self):
+        """
+        The crux of why this was missed. ``JSONDecodeError`` is neither a
+        ``FileNotFoundError``, a ``TypeError`` nor a ``KeyError``, so it fell
+        through every guard on the resume path.
+        """
+        assert issubclass(json.JSONDecodeError, ValueError)
+        assert not issubclass(json.JSONDecodeError, (TypeError, KeyError))