diff --git a/autofit/non_linear/fitness.py b/autofit/non_linear/fitness.py index 796af8bce..2604d2569 100644 --- a/autofit/non_linear/fitness.py +++ b/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 diff --git a/autofit/non_linear/paths/directory.py b/autofit/non_linear/paths/directory.py index ae153a4d5..e4ede6874 100644 --- a/autofit/non_linear/paths/directory.py +++ b/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): diff --git a/autofit/non_linear/search/mle/multi_start_gradient/search.py b/autofit/non_linear/search/mle/multi_start_gradient/search.py index 30af2edc4..eb2ad4059 100644 --- a/autofit/non_linear/search/mle/multi_start_gradient/search.py +++ b/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)) diff --git a/autofit/tools/util.py b/autofit/tools/util.py index ebcdfb3e7..cc2c6a555 100644 --- a/autofit/tools/util.py +++ b/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: diff --git a/test_autofit/tools/test_atomic_write.py b/test_autofit/tools/test_atomic_write.py new file mode 100644 index 000000000..cb9bc5abf --- /dev/null +++ b/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))