feat: results library phase 2 — completeness over a declared grid + enrich

autofit/__init__.py

@@ -39,6 +39,7 @@
 from .aggregator.summary.aggregate_fits import AggregateFITS
 from autofit.aggregator.fit_interface import Fit
 from .aggregator.search_output import SearchOutput
+from .results.library import ResultsLibrary
 from .mapper import prior
 from .mapper.model import AbstractModel
 from .mapper.model import ModelInstance

autofit/aggregator/aggregator.py

@@ -711,6 +711,42 @@ def value_of(search_output, attribute):
             errors=self.errors,
         )
 
+    def enrich(
+        self,
+        fn,
+        name: str,
+        overwrite: bool = False,
+        shard=None,
+    ):
+        """
+        Add a product computed from each saved result to the result itself, written to
+        ``files/<name>`` atomically, without rerunning any fit
+        (``autofit.results.enrich``).
+
+        Parameters
+        ----------
+        fn
+            ``fn(search_output)`` returns the product: a ``dict`` / ``list`` (JSON),
+            ``str``, ``bytes``, a path to a file, or ``None`` to skip that result.
+        name
+            The product's name (``files/<name>``; ``.json`` is added for JSON content
+            when it has no suffix).
+        overwrite
+            Rewrite the product where a result already holds it.
+        shard
+            ``"i/N"`` or ``(i, N)``: only shard ``i`` of ``N`` disjoint shards, so the
+            pass can be split between Slurm array tasks or processes.
+
+        Returns
+        -------
+        An ``EnrichReport``: what was written, already there, declined or failed.
+        """
+        from autofit.results.enrich import enrich
+
+        return enrich(
+            self.search_outputs, fn, name=name, overwrite=overwrite, shard=shard
+        )
+
     def values(self, name: str, parser=lambda o: o) -> Iterator:
         """
         Get values from outputs with a given name.

autofit/non_linear/paths/abstract.py

@@ -18,7 +18,12 @@
 from autofit.non_linear.samples.summary import SamplesSummary
 
 from autofit.text import text_util
-from autofit.tools.util import compress_type_for, open_, zip_directory
+from autofit.tools.util import (
+    compress_type_for,
+    open_,
+    write_zip_member,
+    zip_directory,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -61,47 +66,13 @@ def _matches_archived(info: zipfile.ZipInfo, file_path) -> bool:
 def _replace_zip_member(zip_path, arcname: str, file_path):
     """
     Replace the member ``arcname`` of the zip at ``zip_path`` with the file at
-    ``file_path``.
-
-    ``zipfile`` cannot overwrite a member in place, so every other member is
-    streamed into a new archive alongside the replacement. The new archive is
-    written next to the original (same directory, so the final ``os.replace``
-    is atomic on one filesystem) and only swapped in once complete: an error
-    part-way leaves the original archive untouched.
+    ``file_path``, atomically (``write_zip_member``), keeping its compression.
 
     Rebuilding the archive from the output directory instead (as the search
     itself does at completion) is not an option here, because with
     ``remove_files`` the directory holds only the file just written.
     """
-    zip_path = Path(zip_path)
-    temporary_path = zip_path.with_name(f"{zip_path.name}.replace.tmp")
-
-    try:
-        with zipfile.ZipFile(zip_path, "r") as source:
-            replaced = source.getinfo(arcname)
-
-            with zipfile.ZipFile(temporary_path, "w") as destination:
-                for info in source.infolist():
-                    if info.filename == arcname:
-                        continue
-
-                    with source.open(info) as member:
-                        with destination.open(info, "w") as target:
-                            shutil.copyfileobj(member, target)
-
-                destination.write(
-                    file_path,
-                    arcname,
-                    compress_type=replaced.compress_type,
-                )
-
-        os.replace(temporary_path, zip_path)
-    except BaseException:
-        try:
-            os.remove(temporary_path)
-        except FileNotFoundError:
-            pass
-        raise
+    write_zip_member(zip_path, arcname, file_path=file_path)
 
 
 _zip_complete_cache = {}

autofit/non_linear/search/abstract_search.py

@@ -524,7 +524,7 @@ def result_via_completed_fit(
 
         Optional tasks update the results on hard-disk depending on the `general.yaml` config file's `output`
         section: `force_visualize_overwrite=True` performs the visualization again, and `force_pickle_overwrite=True`
-        recreates the output files of the model-fit.
+        recreates the output files (to add a product to finished results, use `Aggregator.enrich`).
 
         """
         status_notify(self, "skipped")

autofit/results/__init__.py

@@ -5,10 +5,16 @@
 - ``ResultsIndex``: an incremental index of every result under one or more roots, one
   JSON record per result (``ResultRecord``), each classified into a ``ResultState``;
 - ``ResultsCache``: a local SQLite cache of an index for queries;
-- ``python -m autofit.results {index,status,count}``: the command line.
+- ``ResultsLibrary``: what a set of fits is expected to hold, and ``check`` for what is
+  missing (``CompletenessReport``);
+- ``enrich``: add a product to finished results (``Aggregator.enrich``);
+- ``python -m autofit.results {index,status,count,check,enrich}``: the command line.
 """
 from autofit.non_linear.status.states import ResultState
 from autofit.results.cache import ResultsCache
 from autofit.results.classify import classify
 from autofit.results.index import DEFAULT_LATEST_KEY, ResultsIndex, find_results
 from autofit.results.record import ResultRecord
+from autofit.results.completeness import CompletenessReport
+from autofit.results.enrich import EnrichReport, enrich
+from autofit.results.library import ExpectedResult, LibraryException, ResultsLibrary

autofit/results/__main__.py

@@ -4,16 +4,27 @@
     python -m autofit.results index output/ [more_roots/ ...] [--index-dir DIR]
     python -m autofit.results status output/ [--index-dir DIR] [--state corrupt_zip]
     python -m autofit.results count output/ [--index-dir DIR] [--state completed]
+    python -m autofit.results check library.yaml [--no-update] [--output-dir DIR]
+    python -m autofit.results enrich output/ --fn module:function --name NAME \
+        [--index-dir DIR] [--shard i/N] [--overwrite]
 
 ``index`` brings the index up to date (only new or changed results are opened);
 ``status`` prints one line per result state, or with ``--state`` the results in that
 state; ``count`` prints a number. ``status`` and ``count`` read the existing index and
 never open a result; run ``index`` first.
+
+``check`` compares a declared library (``library.yaml``, see
+``autofit.results.library``) with the index and writes ``<name>.missing.csv``; it
+exits 1 when anything expected is not done. ``enrich`` adds the product
+``function(search_output)`` returns to every completed result as ``files/<NAME>``
+(``autofit.results.enrich``); ``--shard i/N`` takes one of N disjoint parts.
 """
 import argparse
+import importlib
 import json
 import sys
 import time
+from pathlib import Path
 
 from autofit.non_linear.status.states import ResultState
 from autofit.results.index import ResultsIndex
@@ -105,6 +116,62 @@ def command_count(arguments) -> int:
     return 0
 
 
+def command_check(arguments) -> int:
+    from autofit.results.library import ResultsLibrary
+
+    # 1 means "something expected is not done", so every error is 2.
+    try:
+        library = ResultsLibrary.from_yaml(arguments.library)
+        report = library.check(update=not arguments.no_update)
+        path = report.write_missing_csv(
+            arguments.output_dir or Path(arguments.library).parent
+        )
+    except Exception as exception:
+        print(f"check failed: {type(exception).__name__}: {exception}", file=sys.stderr)
+        return 2
+    print(report.summary())
+    print(f"gaps: {path}")
+    return 0 if report.is_complete else 1
+
+
+def _function(spec: str):
+    module, _, name = spec.partition(":")
+    if not module or not name:
+        raise ValueError(f"--fn is module:function, not {spec!r}.")
+    function = importlib.import_module(module)
+    for part in name.split("."):
+        function = getattr(function, part)
+    return function
+
+
+def command_enrich(arguments) -> int:
+    from autofit.aggregator.aggregator import Aggregator
+
+    sys.path.insert(0, str(Path.cwd()))
+    try:
+        function = _function(arguments.fn)
+    except (ImportError, AttributeError, ValueError) as exception:
+        print(f"Cannot load {arguments.fn}: {exception}", file=sys.stderr)
+        return 2
+    try:
+        aggregator = Aggregator.from_index(
+            arguments.roots, index_dir=arguments.index_dir
+        )
+        report = aggregator.enrich(
+            function,
+            name=arguments.name,
+            overwrite=arguments.overwrite,
+            shard=arguments.shard,
+        )
+    except Exception as exception:
+        print(f"enrich failed: {type(exception).__name__}: {exception}", file=sys.stderr)
+        return 2
+    print(report)
+    for output_path, error in report.errors:
+        print(f"  error {output_path}: {error}", file=sys.stderr)
+    return 1 if report.errors else 0
+
+
 def main(argv=None) -> int:
     parser = argparse.ArgumentParser(
         prog="python -m autofit.results",
@@ -148,6 +215,34 @@ def add(name, function, help_):
         if name == "status":
             command.add_argument("--json", action="store_true")
 
+    check = commands.add_parser(
+        "check", help="Completeness of a declared library; writes <name>.missing.csv."
+    )
+    check.add_argument("library", help="The library.yaml.")
+    check.add_argument(
+        "--no-update",
+        action="store_true",
+        help="Use the index as it stands instead of updating it first.",
+    )
+    check.add_argument(
+        "--output-dir",
+        default=None,
+        help="Where <name>.missing.csv goes (default: beside library.yaml).",
+    )
+    check.set_defaults(function=command_check)
+
+    enrich = add("enrich", command_enrich, "Add a product to every completed result.")
+    enrich.add_argument(
+        "--fn", required=True, help="module:function, called with each search output."
+    )
+    enrich.add_argument("--name", required=True, help="Written to files/<name>.")
+    enrich.add_argument("--shard", default=None, help="i/N: shard i of N.")
+    enrich.add_argument(
+        "--overwrite",
+        action="store_true",
+        help="Rewrite the product where a result already holds it.",
+    )
+
     arguments = parser.parse_args(argv)
     return arguments.function(arguments)