diff --git a/autofit/non_linear/search/nest/nautilus/search.py b/autofit/non_linear/search/nest/nautilus/search.py index b971ff531..87219158f 100644 --- a/autofit/non_linear/search/nest/nautilus/search.py +++ b/autofit/non_linear/search/nest/nautilus/search.py @@ -5,7 +5,7 @@ import logging import os import sys -from contextlib import nullcontext +from contextlib import contextmanager, nullcontext from pathlib import Path from typing import Dict, Optional, Tuple, TYPE_CHECKING @@ -44,6 +44,84 @@ def _resumed_by_nautilus(path): ) +NAUTILUS_N_LIVE_MESSAGE = ( + "Nautilus cannot fit this {prior_count}-parameter model with {argument}={value}: " + "nautilus fits an ellipsoid to the live points at its first bound, which needs " + "more points than dimensions (nautilus/bounds/basic.py, Ellipsoid.compute), so " + "the run would start and then fail. Set {argument} to at least {floor} " + "(n_live > n_dim, n_points_min >= n_dim + 1)." +) +""" +The one message raised when a Nautilus live-point setting is below the model's floor. +""" + +NAUTILUS_HDF5_LIBVER = ("v108", "latest") +""" +The HDF5 file-format bounds nautilus's checkpoint is opened with (see +``_large_attribute_hdf5``). +""" + + +class _LargeAttributeH5py: + """ + A stand-in for the ``h5py`` module inside ``nautilus.sampler`` whose ``File`` + defaults to ``libver=NAUTILUS_HDF5_LIBVER`` when the caller passes none; every + other attribute is the real ``h5py``'s. + """ + + def __init__(self, h5py): + self._h5py = h5py + + def File(self, *args, **kwargs): + kwargs.setdefault("libver", NAUTILUS_HDF5_LIBVER) + return self._h5py.File(*args, **kwargs) + + def __getattr__(self, name): + return getattr(self._h5py, name) + + +@contextmanager +def _large_attribute_hdf5(): + """ + Let nautilus write its checkpoint for models with 91 or more parameters. + + nautilus-sampler 1.0.5 stores each ellipsoid bound's ``n_dim x n_dim`` matrices + (``A``, ``B``, ``B_inv``) as HDF5 attributes (``Ellipsoid.write``, + ``nautilus/bounds/basic.py`` :404-407) in a file it creates with + ``h5py.File(filepath, 'x')`` (``nautilus/sampler.py`` :1270), i.e. HDF5's + earliest file format, whose 64 KiB object-header message limit an + ``n_dim**2`` float64 array exceeds from ``n_dim = 91``. The first checkpoint + write after an ellipsoid bound then raises ``OSError: Unable to synchronously + create attribute (object header message is too large)``. + + For the duration of the ``with`` block, ``nautilus.sampler.h5py`` is a proxy + whose ``File`` defaults to ``libver=('v108', 'latest')``, the HDF5 1.8 format + whose dense attribute storage holds attributes of any size. The checkpoint + round-trips and resumes; a file written this way needs HDF5 >= 1.8 to read, + which every supported h5py has. The real module is restored on exit. + + Remove this once nautilus stores the bound matrices as datasets (or opens the + file with a newer libver) upstream. + """ + try: + import nautilus.sampler as nautilus_sampler + except ImportError: + yield + return + + h5py = getattr(nautilus_sampler, "h5py", None) + + if h5py is None or isinstance(h5py, _LargeAttributeH5py): + yield + return + + nautilus_sampler.h5py = _LargeAttributeH5py(h5py) + try: + yield + finally: + nautilus_sampler.h5py = h5py + + # How often `_LikelihoodWorkerPool.map` wakes up to check that the pool's # workers are the ones it was built with. Read at call time (not bound at @@ -241,6 +319,7 @@ def __init__( force_x1_cpu: bool = False, session: Optional[sa.orm.Session] = None, use_jax_vmap: Optional[bool] = None, + checkpoint: bool = True, **kwargs, ): """ @@ -264,7 +343,15 @@ def __init__( The name of a unique tag for this model-fit, which will be given a unique entry in the sqlite database and also acts as the folder after the path prefix and before the search name. n_live - Number of live points used for sampling. + Number of live points used for sampling. Must exceed the number of free + parameters: nautilus fits an ellipsoid to the live points at its first + bound, which needs more points than dimensions. The search raises a + ``SearchException`` before any likelihood call when ``n_live`` is at or + below the model's ``prior_count``. + n_points_min + The minimum number of points each ellipsoid is fitted to (nautilus + defaults it to ``n_dim + 50``). When set, it must be at least + ``prior_count + 1``, checked like ``n_live``. n_batch Number of likelihood evaluations performed at each step. n_like_max @@ -291,6 +378,16 @@ def __init__( If True, force single-CPU mode even when number_of_cores > 1. session An SQLalchemy session instance so the results of the model-fit are written to an SQLite database. + checkpoint + If True (the default), nautilus writes its resume state, ``checkpoint.hdf5``, + in ``search_internal/`` and an interrupted fit resumes from it. If False, no + filepath is passed to nautilus, so nothing is written and an interrupted fit + starts again from scratch: an explicit resume-less mode, for example when + the checkpoint's disk I/O is unwanted. It does not change the samples, so it + is not an identifier field. Models with 91 or more parameters checkpoint + normally: the search writes the file in the HDF5 1.8 format, because + nautilus 1.0.5's default format cannot store their bounds + (``_large_attribute_hdf5``). """ super().__init__( @@ -325,6 +422,7 @@ def __init__( self.verbose = verbose self.force_x1_cpu = force_x1_cpu + self.checkpoint = checkpoint if use_jax_vmap is not None: from autofit.non_linear.search.abstract_search import ( warn_deprecated_jax_knob, @@ -355,6 +453,47 @@ def fitness_overrides(self, analysis) -> dict: } return {} + def check_model(self, model: AbstractPriorModel): + """ + Refuse a model nautilus cannot bound, before any likelihood call. + + nautilus needs more live points than dimensions (``n_live > prior_count``) and, + when ``n_points_min`` is set, ``n_points_min >= prior_count + 1``; neither + nautilus nor its ``Sampler`` checks this up front, so the run otherwise + starts and fails at its first ellipsoid bound. + + Raises + ------ + exc.SearchException + With ``NAUTILUS_N_LIVE_MESSAGE``, naming the setting and its floor. + """ + super().check_model(model=model) + + if model is None: + return + + prior_count = model.prior_count + + if self.n_live <= prior_count: + raise exc.SearchException( + NAUTILUS_N_LIVE_MESSAGE.format( + prior_count=prior_count, + argument="n_live", + value=self.n_live, + floor=prior_count + 1, + ) + ) + + if self.n_points_min is not None and self.n_points_min < prior_count + 1: + raise exc.SearchException( + NAUTILUS_N_LIVE_MESSAGE.format( + prior_count=prior_count, + argument="n_points_min", + value=self.n_points_min, + floor=prior_count + 1, + ) + ) + def run(self, ctx): """ Run Nautilus (the ``run(ctx)`` hook; see ``docs/design/run_ctx.md``) and return @@ -369,8 +508,11 @@ def run(self, ctx): ctx The fit's ``FitContext``. """ + # `fit` runs the same guard; repeated here for callers that drive `run(ctx)` + # directly. + self.check_model(model=ctx.model) - if ctx.paths.supports_checkpointing: + if ctx.paths.supports_checkpointing and self.checkpoint_file is not None: checkpoint_exists = Path(self.checkpoint_file).exists() else: checkpoint_exists = False @@ -421,7 +563,11 @@ def checkpoint_file(self): The path to the file used for checkpointing. If autofit is not outputting results to hard-disk (e.g. paths is `NullPaths`), this function is bypassed. + It is None when the search was built with ``checkpoint=False``, so nautilus + writes no checkpoint at all. """ + if not getattr(self, "checkpoint", True): + return None try: return self.paths.search_internal_path / NAUTILUS_FILENAME except TypeError: @@ -455,31 +601,32 @@ def fit_x1_cpu(self, fitness, model, analysis=None, update=None): "Running search where parallelization is disabled." ) - search_internal = self.sampler_cls( - prior=PriorVectorized(model=model), - likelihood=fitness.call_wrap, - n_dim=model.prior_count, - filepath=self.checkpoint_file, - pool=None, - vectorized=fitness.batched, - n_live=self.n_live, - n_update=self.n_update, - enlarge_per_dim=self.enlarge_per_dim, - n_points_min=self.n_points_min, - split_threshold=self.split_threshold, - n_networks=self.n_networks, - n_batch=self.n_batch, - n_like_new_bound=self.n_like_new_bound, - seed=self.backend_seed, - ) + with _large_attribute_hdf5(): + search_internal = self.sampler_cls( + prior=PriorVectorized(model=model), + likelihood=fitness.call_wrap, + n_dim=model.prior_count, + filepath=self.checkpoint_file, + pool=None, + vectorized=fitness.batched, + n_live=self.n_live, + n_update=self.n_update, + enlarge_per_dim=self.enlarge_per_dim, + n_points_min=self.n_points_min, + split_threshold=self.split_threshold, + n_networks=self.n_networks, + n_batch=self.n_batch, + n_like_new_bound=self.n_like_new_bound, + seed=self.backend_seed, + ) - return self.call_search( - search_internal=search_internal, - model=model, - analysis=analysis, - fitness=fitness, - update=update, - ) + return self.call_search( + search_internal=search_internal, + model=model, + analysis=analysis, + fitness=fitness, + update=update, + ) def fit_multiprocessing(self, fitness, model, analysis=None, pools=None, update=None): """ @@ -550,7 +697,7 @@ def fit_multiprocessing(self, fitness, model, analysis=None, pools=None, update= ) ) - with pool_context as pool: + with pool_context as pool, _large_attribute_hdf5(): search_internal = self.sampler_cls( prior=PriorVectorized(model=model), likelihood=fitness.call_wrap, diff --git a/docs/cookbooks/search.md b/docs/cookbooks/search.md index 8be5b8743..16effa08d 100644 --- a/docs/cookbooks/search.md +++ b/docs/cookbooks/search.md @@ -388,6 +388,7 @@ search = af.Nautilus( n_shell=1, n_eff=500, discard_exploration=False, + checkpoint=True, # False: write no checkpoint.hdf5 (an interrupted fit restarts) ) ``` diff --git a/docs/design/checkpointing.md b/docs/design/checkpointing.md index 783a283aa..a0f7f0753 100644 --- a/docs/design/checkpointing.md +++ b/docs/design/checkpointing.md @@ -53,6 +53,44 @@ dynesty's `save_sampler` writes `savestate.save.tmp` and renames it; NSS's can leave them unreadable. `NativeFileCheckpointer.save` and `finalize` are no-ops: autofit never writes or rewrites a native file. +## Nautilus at 91 or more parameters + +nautilus-sampler 1.0.5 cannot write its own checkpoint for a model with 91 or more +free parameters. `Ellipsoid.write` (`nautilus/bounds/basic.py`, lines 404-407) stores +each ellipsoid bound's `n_dim x n_dim` float64 matrices (`A`, `B`, `B_inv`) as HDF5 +attributes, and `Sampler.write` creates the file with `h5py.File(filepath, 'x')` +(`nautilus/sampler.py`, line 1270), which uses HDF5's earliest file format. That +format caps an object-header message at 64 KiB, so from `n_dim = 91` +(`91**2 * 8` bytes > 64 KiB) the first write after an ellipsoid bound raises +`OSError: Unable to synchronously create attribute (object header message is too large)` +and the run dies, although the sampler itself is fine. + +The workaround is `_large_attribute_hdf5()` in `autofit/non_linear/search/nest/nautilus/search.py`. +Around sampler construction and the sampling loop, on both execution paths +(`fit_x1_cpu` and `fit_multiprocessing`), it replaces `nautilus.sampler.h5py` with a +proxy whose `File` defaults to `libver=('v108', 'latest')`. That is the HDF5 1.8 format, +which stores attributes of any size densely. Every other attribute is the real `h5py`'s, +and the module is restored when the block exits. The checkpoint then round-trips and +resumes at any `n_dim`. A file written this way needs HDF5 >= 1.8 to read, which every +supported h5py has. The workaround is to be removed once nautilus fixes this upstream; +`test__d100_ellipsoid_write_fails_in_nautilus_default_format` is a strict xfail that +starts XPASSing (and so fails) when that happens. + +Two related settings on `af.Nautilus`: + +- **Live-point floor.** nautilus's ellipsoid fit needs more points than dimensions + (`Ellipsoid.compute`, `nautilus/bounds/basic.py`, lines 298-300), and nothing checks + this before the first bound. `Nautilus.check_model` (run by `fit`, and again at the + top of `run(ctx)`) raises `SearchException` (`NAUTILUS_N_LIVE_MESSAGE`) before any + likelihood call unless `n_live > prior_count` and, when set, + `n_points_min >= prior_count + 1`. +- **`checkpoint=False`.** This passes no filepath to nautilus, so no `checkpoint.hdf5` + is written and an interrupted fit starts again from scratch. It is an explicit + resume-less mode rather than a fix. It does not change the samples, so it is not an + identifier field. + +Upstream issue: . + ## Corruption A corrupt file is never silently replaced by a different source of samples: diff --git a/test_autofit/non_linear/search/conformance_roster.py b/test_autofit/non_linear/search/conformance_roster.py index ee4d9ae74..d02f69666 100644 --- a/test_autofit/non_linear/search/conformance_roster.py +++ b/test_autofit/non_linear/search/conformance_roster.py @@ -482,6 +482,7 @@ class Golden: | { "session", "number_of_cores", + "checkpoint", "discard_exploration", "enlarge_per_dim", "f_live", diff --git a/test_autofit/non_linear/search/nest/test_nautilus.py b/test_autofit/non_linear/search/nest/test_nautilus.py index 84829ca96..5d5954d18 100644 --- a/test_autofit/non_linear/search/nest/test_nautilus.py +++ b/test_autofit/non_linear/search/nest/test_nautilus.py @@ -427,3 +427,267 @@ def test__likelihood_pool_healthy_map_completes(monkeypatch): finally: pool.terminate() pool.join() + + +# --- D >= 91: the live-point guard and the HDF5 checkpoint format (#1697) --- + + +class _VectorGaussianAnalysis(af.Analysis): + """ + An isotropic Gaussian over every free parameter, counting likelihood calls. + """ + + def __init__(self, n_dim): + super().__init__() + self.names = [f"x{index}" for index in range(n_dim)] + self.calls = 0 + + def log_likelihood_function(self, instance): + self.calls += 1 + vector = np.array([getattr(instance, name) for name in self.names]) + return -0.5 * float(np.sum(((vector - 0.5) / 0.1) ** 2)) + + +def _vector_model(n_dim): + return af.Collection( + **{f"x{index}": af.UniformPrior(lower_limit=0.0, upper_limit=1.0) for index in range(n_dim)} + ) + + +def _ellipsoid_d100(): + from nautilus.bounds.basic import Ellipsoid + + rng = np.random.default_rng(1) + return Ellipsoid.compute(rng.random((300, 100))) + + +@requires_nautilus +def test__d100_ellipsoid_round_trips_through_the_proxied_file(tmp_path): + """ + A D = 100 ellipsoid's ``n_dim x n_dim`` attributes write and read back when the + file is opened through the proxy's ``File`` (HDF5 1.8 format). + """ + import h5py + from nautilus.bounds.basic import Ellipsoid + + from autofit.non_linear.search.nest.nautilus.search import _LargeAttributeH5py + + bound = _ellipsoid_d100() + filepath = tmp_path / "checkpoint.hdf5" + + with _LargeAttributeH5py(h5py).File(filepath, "x") as fstream: + bound.write(fstream.create_group("bound")) + + with h5py.File(filepath, "r+") as fstream: + read = Ellipsoid.read(fstream["bound"]) + + assert np.array_equal(read.B, bound.B) + assert np.array_equal(read.B_inv, bound.B_inv) + + +@requires_nautilus +@pytest.mark.xfail( + strict=True, + raises=OSError, + reason="nautilus 1.0.5 writes D x D attributes in HDF5's earliest format; " + "XPASSes (and fails this strict xfail) once nautilus fixes it upstream, " + "the cue to remove _large_attribute_hdf5", +) +def test__d100_ellipsoid_write_fails_in_nautilus_default_format(tmp_path): + """ + The control: nautilus's own ``h5py.File(filepath, 'x')`` cannot hold the bound. + """ + import h5py + + bound = _ellipsoid_d100() + + with h5py.File(tmp_path / "checkpoint.hdf5", "x") as fstream: + bound.write(fstream.create_group("bound")) + + +def test__large_attribute_hdf5_restores_nautilus_h5py(): + pytest.importorskip("nautilus") + import nautilus.sampler as nautilus_sampler + + from autofit.non_linear.search.nest.nautilus import search as nautilus_search + + original = nautilus_sampler.h5py + + with pytest.raises(RuntimeError): + with nautilus_search._large_attribute_hdf5(): + assert isinstance(nautilus_sampler.h5py, nautilus_search._LargeAttributeH5py) + raise RuntimeError + + assert nautilus_sampler.h5py is original + + +class _StopAtFirstUpdate(Exception): + pass + + +def _fast_ellipsoid_bound_cls(): + """ + A stand-in for nautilus's ``NautilusBound`` that keeps its D x D ellipsoid + matrices but none of its cost. + + With ~100 live points at D = 100, nautilus's real bound rejection-samples an + enlarged ellipsoid against the unit cube (``Union.sample`` / + ``NeuralBound.sample``) with a near-zero acceptance rate, which takes minutes, so + the test cannot afford it. This bound is a real ``Ellipsoid`` fitted to the live + set (written through ``Ellipsoid.write``: the ``A``, ``B``, ``B_inv`` attributes + that overflow nautilus's default HDF5 format), sampled from the unit cube and + declared smaller than it so the sampler accepts it and writes the checkpoint. + """ + from nautilus.bounds.basic import Ellipsoid + + class FastEllipsoidBound(Ellipsoid): + @classmethod + def compute(cls, points, log_l, log_l_min, log_v_target, enlarge_per_dim=1.1, rng=None, **kwargs): + bound = super().compute( + points[log_l >= log_l_min], enlarge_per_dim=enlarge_per_dim, rng=rng + ) + return bound + + @property + def log_v(self): + return -1.0 + + def sample(self, n_points=100, return_points=True, pool=None): + points = self.rng.random((n_points, self.n_dim)) + return points if return_points else None + + def update(self, group): + pass + + return FastEllipsoidBound + + +@requires_nautilus +def test__d100_fit_writes_a_checkpoint_that_nautilus_resumes(tmp_path, monkeypatch): + """ + A D = 100 fit reaches its first checkpoint write after an ellipsoid bound + (nautilus-sampler 1.0.5 raised ``object header message is too large`` there), + and a fresh nautilus ``Sampler`` on the same file resumes it. + + The fit is stopped at its first full update, as an interrupted run would be: + a completed fit deletes its resume state. nautilus's bound class is replaced by + ``_fast_ellipsoid_bound_cls`` (same ellipsoid serialisation, no rejection + sampling) to keep the test to seconds; everything else is the search's real + ``fit`` -> ``run(ctx)`` -> ``fit_x1_cpu`` path and nautilus's real write and + resume. + """ + from nautilus import Sampler + + from autofit.conf import instance + + monkeypatch.delenv("PYAUTO_TEST_MODE", raising=False) + monkeypatch.setattr(instance, "output_path", str(tmp_path)) + + n_dim = 100 + model = _vector_model(n_dim) + analysis = _VectorGaussianAnalysis(n_dim) + + search = af.Nautilus( + name="nautilus_d100_checkpoint", + n_live=101, + n_networks=1, + n_batch=101, + n_like_max=404, + seed=1, + verbose=False, + ) + + import nautilus.sampler as nautilus_sampler + + monkeypatch.setattr(nautilus_sampler, "NautilusBound", _fast_ellipsoid_bound_cls()) + + stopped = {} + + def stop(search_internal=None, **kwargs): + stopped["n_like"] = search_internal.n_like + stopped["bounds"] = len(search_internal.bounds) + raise _StopAtFirstUpdate + + monkeypatch.setattr(search, "perform_update", stop) + + with pytest.raises(_StopAtFirstUpdate): + search.fit(model=model, analysis=analysis) + + checkpoint_file = search.checkpoint_file + + assert str(checkpoint_file).startswith(str(tmp_path)) + assert checkpoint_file.exists() + assert stopped["bounds"] >= 2 + + resumed = Sampler( + prior=lambda x: x, + likelihood=lambda x: 0.0, + n_dim=n_dim, + n_live=101, + filepath=checkpoint_file, + resume=True, + ) + + assert resumed.n_like == stopped["n_like"] + assert len(resumed.bounds) >= 2 + + +def test__checkpoint_false_passes_no_filepath(): + search = af.Nautilus(name="nautilus_no_checkpoint", checkpoint=False) + + assert search.checkpoint is False + assert search.checkpoint_file is None + assert "checkpoint" not in search.__identifier_fields__ + + checkpointed = af.Nautilus(name="nautilus_no_checkpoint") + + assert checkpointed.checkpoint_file is not None + + +def _guard_fit(search, n_dim, monkeypatch): + monkeypatch.delenv("PYAUTO_TEST_MODE", raising=False) + analysis = _VectorGaussianAnalysis(n_dim) + + def backend_must_not_run(*args, **kwargs): + raise AssertionError("the backend ran before the live-point guard") + + monkeypatch.setattr(search, "_fit", backend_must_not_run) + + with pytest.raises(af.exc.SearchException) as error: + search.fit(model=_vector_model(n_dim), analysis=analysis) + + assert analysis.calls == 0 + return str(error.value) + + +def test__n_live_at_prior_count_raises_before_any_likelihood_call(monkeypatch): + message = _guard_fit(af.Nautilus(n_live=100), 100, monkeypatch) + + assert message.startswith( + "Nautilus cannot fit this 100-parameter model with n_live=100" + ) + assert "at least 101" in message + + +def test__n_points_min_below_floor_raises(monkeypatch): + message = _guard_fit(af.Nautilus(n_live=200, n_points_min=50), 100, monkeypatch) + + assert "n_points_min=50" in message + assert "at least 101" in message + + +def test__live_point_floor_passes_the_gate(): + model = _vector_model(100) + + af.Nautilus(n_live=101).check_model(model) + af.Nautilus(n_live=101, n_points_min=101).check_model(model) + + with pytest.raises(af.exc.SearchException): + af.Nautilus(n_live=101, n_points_min=100).check_model(model) + + +def test__run_ctx_repeats_the_guard(): + from types import SimpleNamespace + + with pytest.raises(af.exc.SearchException, match="n_live=10"): + af.Nautilus(n_live=10).run(SimpleNamespace(model=_vector_model(20)))