diff --git a/autofit/__init__.py b/autofit/__init__.py index add6579db..4cd27e10b 100644 --- a/autofit/__init__.py +++ b/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 diff --git a/autofit/aggregator/aggregator.py b/autofit/aggregator/aggregator.py index 4d4578de4..d170bbcef 100755 --- a/autofit/aggregator/aggregator.py +++ b/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/`` 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/``; ``.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. diff --git a/autofit/non_linear/paths/abstract.py b/autofit/non_linear/paths/abstract.py index a280d1574..0a75f0ad6 100644 --- a/autofit/non_linear/paths/abstract.py +++ b/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 = {} diff --git a/autofit/non_linear/search/abstract_search.py b/autofit/non_linear/search/abstract_search.py index 6865857fa..6a9b4947e 100644 --- a/autofit/non_linear/search/abstract_search.py +++ b/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") diff --git a/autofit/results/__init__.py b/autofit/results/__init__.py index dd5612f0c..f3340368a 100644 --- a/autofit/results/__init__.py +++ b/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 diff --git a/autofit/results/__main__.py b/autofit/results/__main__.py index 7eecd0e32..9cdf81758 100644 --- a/autofit/results/__main__.py +++ b/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 ``.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/`` +(``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 .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 .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/.") + 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) diff --git a/autofit/results/completeness.py b/autofit/results/completeness.py new file mode 100644 index 000000000..5c91c2e42 --- /dev/null +++ b/autofit/results/completeness.py @@ -0,0 +1,332 @@ +""" +Completeness of a ``ResultsLibrary``: every expected result is ``done``, ``missing``, +``incomplete``, ``failed`` or ``duplicate``, each with a reason, read from the results +index — a zip is opened only to list its members when a path product is required. + +- ``done``: one completed result holding every required product; +- ``missing``: no result at all (``no_result``); +- ``incomplete``: a result that is not complete yet (``not_completed:``: + running, interrupted, never started), a completed one lacking a required product + (``missing_product``, one gap per product), or one that sits where another object's + result belongs (``location_mismatch``); +- ``failed``: the result failed or its zip / checkpoint is corrupt + (``not_completed:``), or where it should be is not known (a campaign manifest + line missing a field, ``manifest_line: ...``); +- ``duplicate``: several completed results match (``duplicate:``) — the library + cannot say which is the one. + +Superseded results (an older sibling of a newer completed fit) never count, except at +an exact location: there the result asked about is the one at that path. An exact +(campaign) expectation also needs a zip (``no_zip``) holding its path products as +members, as ``autofit.campaign.collector.check_success`` does. + +``CompletenessReport.write_missing_csv`` writes one row per gap to +``.missing.csv``: an expected result that is missing, failed or incomplete as a +whole has one row with an empty ``product``; one lacking products has a row per product. +""" +import csv +import logging +import os +from collections import Counter, defaultdict +from dataclasses import dataclass, field +from pathlib import Path, PurePosixPath +from typing import Dict, Iterable, List, Optional, Sequence + +from autofit.non_linear.status.states import Product, ResultState +from autofit.results.library import ExpectedResult +from autofit.results.record import ResultRecord +from autofit.tools.util import open_atomic + +logger = logging.getLogger(__name__) + +NAMED_PRODUCTS = { + "samples", + "samples_summary", + "latent", + "image", + "dataset_identity", + "search_internal", +} + +STATUSES = ("done", "missing", "incomplete", "failed", "duplicate") + +MISSING_COLUMNS = ("dataset", "search", "product", "status", "reason", "output_path") + +_FAILED_STATES = { + str(ResultState.FAILED), + str(ResultState.CORRUPT_ZIP), + str(ResultState.CORRUPT_CKPT), +} + +# Which unfinished result speaks for an expected result when there are several: the +# one with the most to say about why it is not done. +_PRECEDENCE = [ + str(state) + for state in ( + ResultState.FAILED, + ResultState.CORRUPT_ZIP, + ResultState.CORRUPT_CKPT, + ResultState.INTERRUPTED_VALID_CKPT, + ResultState.RUNNING, + ResultState.NEVER_STARTED, + ) +] + + +@dataclass +class ResultCheck: + expected: ExpectedResult + status: str + reason: Optional[str] = None + record: Optional[ResultRecord] = None + missing_products: List[str] = field(default_factory=list) + + @property + def output_path(self) -> Optional[str]: + return self.record.output_path if self.record is not None else None + + def gaps(self) -> List[Dict[str, str]]: + """ + The rows this result adds to ``.missing.csv`` (none when ``done``). + """ + if self.status == "done": + return [] + row = { + "dataset": self.expected.dataset, + "search": self.expected.search, + "status": self.status, + "output_path": self.output_path or "", + } + if self.missing_products: + return [ + {**row, "product": product, "reason": f"missing_product:{product}"} + for product in self.missing_products + ] + return [{**row, "product": "", "reason": self.reason or ""}] + + +@dataclass +class CompletenessReport: + name: str + results: List[ResultCheck] + unexpected: List[ResultRecord] = field(default_factory=list) + """Completed results of a declared search that no expected result claims, e.g. a + fit of an object outside the manifest or a tree indexed by mistake.""" + + def counts(self) -> Dict[str, int]: + """ + Expected results per status, plus ``expected`` and ``unexpected``. + """ + statuses = Counter(result.status for result in self.results) + return { + "expected": len(self.results), + **{status: statuses.get(status, 0) for status in STATUSES}, + "unexpected": len(self.unexpected), + } + + def product_counts(self) -> Dict[str, Dict[str, int]]: + """ + For each required product: how many expected results require it and how many + of the completed ones lack it. + """ + counts = defaultdict(lambda: {"required": 0, "missing": 0}) + for result in self.results: + for product in result.expected.products: + counts[product]["required"] += 1 + for product in result.missing_products: + counts[product]["missing"] += 1 + return dict(counts) + + @property + def is_complete(self) -> bool: + return all(result.status == "done" for result in self.results) + + def gaps(self) -> List[Dict[str, str]]: + return [row for result in self.results for row in result.gaps()] + + def write_missing_csv(self, directory: os.PathLike = ".") -> Path: + """ + Write every gap to ``/.missing.csv`` (header only when the + library is complete), atomically: a reader never sees half a file. + """ + path = Path(directory) / f"{self.name}.missing.csv" + with open_atomic(path, "w") as f: + writer = csv.DictWriter(f, fieldnames=MISSING_COLUMNS, lineterminator="\n") + writer.writeheader() + for row in self.gaps(): + writer.writerow(row) + return path + + def summary(self) -> str: + counts = self.counts() + lines = [ + f"{self.name}: {counts['done']} of {counts['expected']} expected results done" + ] + for status in STATUSES[1:]: + if counts[status]: + lines.append(f" {status:12s} {counts[status]}") + for product, product_counts in sorted(self.product_counts().items()): + if product_counts["missing"]: + lines.append( + f" missing {product}: {product_counts['missing']} of " + f"{product_counts['required']}" + ) + if counts["unexpected"]: + lines.append( + f" {'unexpected':12s} {counts['unexpected']} completed results no " + f"expected result claims" + ) + return "\n".join(lines) + + +def _resolved(path: str) -> str: + return str(Path(path).resolve()) + + +def _name_of(record: ResultRecord) -> Optional[str]: + # A fit that failed before writing ``files/search.json`` leaves only its status + # record: its name is then where it sits, ``//``. + return record.name or Path(record.output_path).parent.name + + +def _dataset_keys(record: ResultRecord) -> List[str]: + keys = [record.unique_tag, record.object_id, record.path_prefix] + if record.path_prefix: + keys.append(PurePosixPath(record.path_prefix).name) + elif record.name is None: + keys.append(Path(record.output_path).parent.parent.name) + return [key for key in keys if key] + + +def has_product(record: ResultRecord, product: str, zip_only: bool = False) -> bool: + """ + Whether the result ``record`` describes holds ``product`` (a named product or a + path inside the result; a trailing ``/`` names a folder). With ``zip_only`` a path + product must be a member of the zip (named products come from the index record). + """ + if product in NAMED_PRODUCTS: + return bool(record.products.get(product)) + + from autofit.aggregator.zip_member import zip_names + + output_path = Path(record.output_path) + folder = product.endswith("/") + member = product.rstrip("/") + if record.has_dir and not zip_only: + path = output_path / member + if path.is_dir() if folder else path.is_file(): + return True + if record.has_zip: + try: + names = zip_names(f"{output_path}{Product.ZIP}") + except OSError as exception: + logger.warning(f"Could not list {output_path}{Product.ZIP}: {exception}") + return False + if folder: + return any(name.startswith(f"{member}/") for name in names) + return member in names + return False + + +def _check_one( + expected: ExpectedResult, candidates: Sequence[ResultRecord] +) -> ResultCheck: + if expected.error is not None: + return ResultCheck(expected, "failed", expected.error) + if not expected.exact: + # Only the latest of several fits of one search to one object counts; at an + # exact location the result there is the one asked about, superseded or not. + candidates = [ + record + for record in candidates + if record.state != str(ResultState.SUPERSEDED) + ] + if not candidates: + return ResultCheck(expected, "missing", "no_result") + + finished = (str(ResultState.COMPLETED), str(ResultState.SUPERSEDED)) + completed = [record for record in candidates if record.state in finished] + if len(completed) > 1: + latest = max(completed, key=lambda r: (r.completed_at or "", r.output_path)) + return ResultCheck(expected, "duplicate", f"duplicate:{len(completed)}", latest) + + if not completed: + record = min( + candidates, + key=lambda r: ( + _PRECEDENCE.index(r.state) if r.state in _PRECEDENCE else len(_PRECEDENCE), + r.output_path, + ), + ) + reason = f"not_completed:{record.state}" + if record.reason: + reason = f"{reason} ({record.reason})" + status = "failed" if record.state in _FAILED_STATES else "incomplete" + return ResultCheck(expected, status, reason, record) + + record = completed[0] + if expected.exact and not record.has_zip: + return ResultCheck(expected, "incomplete", "no_zip", record) + if record.location_mismatch and not expected.exact: + return ResultCheck(expected, "incomplete", "location_mismatch", record) + missing = [ + product + for product in expected.products + if not has_product(record, product, zip_only=expected.exact) + ] + if missing: + return ResultCheck( + expected, + "incomplete", + "missing_product", + record, + missing_products=missing, + ) + return ResultCheck(expected, "done", None, record) + + +def check(library, records: Iterable[ResultRecord]) -> CompletenessReport: + """ + The ``CompletenessReport`` of ``library`` against the index ``records``. + """ + records = list(records) + by_path = {} + by_parent = defaultdict(list) + by_name = defaultdict(list) + for record in records: + path = _resolved(record.output_path) + by_path[path] = record + by_parent[str(Path(path).parent)].append(record) + by_name[_name_of(record)].append(record) + + claimed = set() + results = [] + for expected in library.expected: + if expected.error is not None: + candidates = [] + elif expected.location is not None: + location = _resolved(expected.location) + if location in by_path: + candidates = [by_path[location]] + elif expected.exact: + candidates = [] + else: + candidates = by_parent.get(location, []) + else: + candidates = [ + record + for record in by_name.get(expected.search, []) + if expected.dataset in _dataset_keys(record) + ] + claimed.update(record.output_path for record in candidates) + results.append(_check_one(expected, candidates)) + + searches = {expected.search for expected in library.expected} + unexpected = [ + record + for record in records + if record.state == str(ResultState.COMPLETED) + and record.output_path not in claimed + and record.name in searches + ] + return CompletenessReport(library.name, results, unexpected) diff --git a/autofit/results/enrich.py b/autofit/results/enrich.py new file mode 100644 index 000000000..4a33866af --- /dev/null +++ b/autofit/results/enrich.py @@ -0,0 +1,213 @@ +""" +``enrich``: add a new product to results that already exist, computed from each saved +result, without rerunning a fit. + +``fn(search_output)`` returns the product's content, written to ``files/``: + +- a ``dict`` or ``list``: JSON (``files/.json`` when ``name`` has no suffix); +- ``str`` or ``bytes``: the content as is; +- an ``os.PathLike``: the file at that path; +- ``None``: nothing (``fn`` declines this result). + +A zipped result gets the member atomically (``write_zip_member``): every other member +keeps its content and a write killed part-way leaves the zip as it was. A result kept +as a folder gets the file written atomically. A result with both (the default, without +``remove_files``) gets both, the zip first; a product counts as present only when every +copy holds it, so rerunning a stopped pass finishes it. + +Each write rewrites the whole zip, so two passes must not write to the same results at +the same time (two products at once, or a pass while those searches run): one of the +members would be lost. Run one pass after another, or give each its own shard. + +With ``overwrite=False`` (the default) a result that already holds the product is left +alone, so an enrich pass that was stopped can simply be run again. ``shard`` (``"i/N"`` +or ``(i, N)``) takes one of ``N`` disjoint parts of the results +(``autofit.results.shard``), so the pass can be shared between Slurm array tasks or a +process pool. A result whose ``fn`` or write fails is recorded in ``errors`` and the pass +goes on. + +``enrich`` replaces running a search again with ``force_pickle_overwrite=True`` (and a +reload job on the cluster) to add a product such as ``coolest.json`` or ``wcs.json``. +""" +import json +import logging +import os +from dataclasses import dataclass, field +from pathlib import Path +from typing import Callable, List, Optional, Tuple + +from autofit.non_linear.status.states import Product +from autofit.results.shard import ShardLike, in_shard, parse_shard +from autofit.tools.util import NumpyEncoder, open_atomic, write_zip_member + +logger = logging.getLogger(__name__) + + +@dataclass +class EnrichReport: + name: str + """The member written, ``files/``.""" + written: List[str] = field(default_factory=list) + """The output paths that got the product.""" + exists: List[str] = field(default_factory=list) + """The output paths that already held it (``overwrite=False``).""" + declined: List[str] = field(default_factory=list) + """The output paths ``fn`` returned ``None`` for.""" + errors: List[Tuple[str, str]] = field(default_factory=list) + """``(output_path, error)`` of every result whose ``fn`` or write failed.""" + skipped_shard: int = 0 + """Results in another shard.""" + + def counts(self): + return { + "written": len(self.written), + "exists": len(self.exists), + "declined": len(self.declined), + "errors": len(self.errors), + "other_shards": self.skipped_shard, + } + + def __str__(self): + counts = self.counts() + return f"enrich {self.name}: " + ", ".join( + f"{count} {key}" for key, count in counts.items() + ) + + +def member_name(name: str) -> str: + """ + The member ``files/`` a product named ``name`` is written to. + """ + name = str(name).strip("/") + if not name or name.startswith(".") or ".." in Path(name).parts: + raise ValueError(f"{name!r} is not a product name.") + return f"files/{name}" + + +def _content(value, name: str) -> Tuple[Optional[str], Optional[bytes], Optional[Path]]: + """ + ``(arcname, data, file_path)`` for the value ``fn`` returned. + """ + arcname = member_name(name) + if isinstance(value, (dict, list)): + if not Path(arcname).suffix: + arcname = f"{arcname}.json" + data = json.dumps(value, indent=4, cls=NumpyEncoder).encode("utf-8") + return arcname, data, None + if isinstance(value, str): + return arcname, value.encode("utf-8"), None + if isinstance(value, (bytes, bytearray)): + return arcname, bytes(value), None + if isinstance(value, os.PathLike): + return arcname, None, Path(value) + raise TypeError( + f"fn returned a {type(value).__name__}; return a dict, list, str, bytes, a " + f"path or None." + ) + + +def _locations(search_output) -> Tuple[Optional[Path], Optional[Path]]: + """ + ``(zip_path, folder)`` a search output is stored in (either may be ``None``). + """ + from autofit.aggregator.zip_member import ZipMember + + directory = search_output.directory + if isinstance(directory, ZipMember): + # Without ``remove_files`` the folder stays beside the zip, and a rerun of the + # search re-zips the folder: the product must go into both. + zip_path = directory.zip_path + folder = zip_path.with_name(zip_path.name[: -len(str(Product.ZIP))]) + return zip_path, (folder if folder.is_dir() else None) + folder = Path(directory) + zip_path = Path(f"{folder}{Product.ZIP}") + return (zip_path if zip_path.exists() else None), folder + + +def _holds(zip_path: Optional[Path], folder: Optional[Path], arcname: str) -> bool: + """ + Whether every copy of the result (zip and folder, where they exist) holds + ``arcname``: a pass stopped between the two writes is finished by the next one. + """ + from autofit.aggregator.zip_member import zip_names + + if zip_path is not None and arcname not in zip_names(zip_path): + return False + if folder is not None and not (folder / arcname).is_file(): + return False + return zip_path is not None or folder is not None + + +def _write(zip_path: Optional[Path], folder: Optional[Path], arcname, data, file_path): + # The zip first: it is what a ``remove_files`` resume restores from, and a folder + # copy without one is finished by rerunning the pass (``_holds``). + if zip_path is not None: + write_zip_member(zip_path, arcname, data=data, file_path=file_path) + if folder is not None and folder.is_dir(): + with open_atomic(folder / arcname, "wb") as f: + if data is not None: + f.write(data) + else: + f.write(Path(file_path).read_bytes()) + + +def enrich( + search_outputs, + fn: Callable, + name: str, + overwrite: bool = False, + shard: ShardLike = None, +) -> EnrichReport: + """ + Add the product ``fn`` computes to each of ``search_outputs`` (see the module + docstring). + + Parameters + ---------- + fn + ``fn(search_output)`` returns the product's content, or ``None`` to skip. + name + The product's name; it is written to ``files/``. + overwrite + Rewrite the product where a result already holds it. + shard + ``"i/N"`` or ``(i, N)``: only the results in shard ``i`` of ``N``. + """ + shard = parse_shard(shard) + report = EnrichReport(name=member_name(name)) + for search_output in search_outputs: + if getattr(search_output, "_temporary_directory", None) is not None: + raise ValueError( + "enrich cannot write to results extracted to a temporary directory " + "(Aggregator.from_directory(unzip_temporary=True)); build the " + "aggregator with Aggregator.from_index or from_directory instead." + ) + zip_path, folder = _locations(search_output) + output_path = str(folder if folder is not None else zip_path) + if output_path.endswith(str(Product.ZIP)): + output_path = output_path[: -len(str(Product.ZIP))] + if not in_shard(output_path, shard): + report.skipped_shard += 1 + continue + try: + arcname = member_name(name) + if not overwrite and ( + _holds(zip_path, folder, arcname) + or _holds(zip_path, folder, f"{arcname}.json") + ): + report.exists.append(output_path) + continue + value = fn(search_output) + if value is None: + report.declined.append(output_path) + continue + arcname, data, file_path = _content(value, name) + report.name = arcname + _write(zip_path, folder, arcname, data, file_path) + except Exception as exception: + logger.warning(f"enrich {name}: {output_path}: {exception}") + report.errors.append((output_path, f"{type(exception).__name__}: {exception}")) + continue + report.written.append(output_path) + logger.info(str(report)) + return report diff --git a/autofit/results/library.py b/autofit/results/library.py new file mode 100644 index 000000000..6fd70e07a --- /dev/null +++ b/autofit/results/library.py @@ -0,0 +1,365 @@ +""" +``ResultsLibrary``: what a set of fits is *expected* to have produced — datasets × +searches, and the products each result must hold — so ``check`` can say what is +missing rather than only counting what is there. + +A library is declared in one of three ways: + +- by API, ``ResultsLibrary.grid(datasets, searches, roots=...)`` or a list of + ``ExpectedResult``; +- in ``library.yaml`` (``ResultsLibrary.from_yaml``): + + .. code-block:: yaml + + name: dr1_v1 # names .missing.csv (default: file stem) + roots: [output] # where the results are, relative to this file + index_dir: .autofit_index # optional, as for ``ResultsIndex`` + datasets_file: manifest.txt # or ``datasets: [a, b, ...]`` + searches: + source_lp: + products: [image, latent, files/model.fits] + location: "{dataset}/source_lp" # optional, relative to the first root + mass_total: [image, samples_summary] + +- or from an hpc-campaign file (``from_campaign: campaign.yaml`` in ``library.yaml``, + or ``ResultsLibrary.from_campaign``): every manifest line is a dataset, every stage + with a ``success.zip_complete`` is a search found at that path, and its + ``success.members`` are its required products. The campaign file stays the single + declaration of what was run; a library never keeps a second manifest. + +A ``dataset`` without a ``location`` is matched to results by the search ``name`` and +the dataset as the result's ``unique_tag``, ``object_id`` (``dataset_identity.json``), +``path_prefix`` or last ``path_prefix`` segment. + +A product is either a name the index already records — ``samples``, +``samples_summary``, ``latent``, ``image``, ``dataset_identity``, ``search_internal`` +— or a path inside the result (``files/model.fits``; ``image/`` for a folder), checked +against the zip's member list without reading any member. +""" +import os +from dataclasses import dataclass, field +from pathlib import Path +from typing import Dict, Iterable, List, Mapping, Optional, Sequence, Tuple, Union + +import yaml + +from autofit.results.index import DEFAULT_LATEST_KEY, ResultsIndex + +PathLike = Union[str, os.PathLike] + +LIBRARY_FIELDS = { + "name", + "roots", + "index_dir", + "datasets", + "datasets_file", + "searches", + "from_campaign", + "stages", + "products", + "note", +} + + +class LibraryException(Exception): + pass + + +@dataclass(frozen=True) +class ExpectedResult: + dataset: str + """The dataset (object) the result is a fit of.""" + search: str + """The search's ``name`` (for a campaign, the stage).""" + products: Tuple[str, ...] = () + """What the result must hold beyond being complete.""" + location: Optional[str] = None + """Where the result is: its ``output_path``, or the folder holding its identifier + folders (``/``). ``None`` matches by dataset and search name.""" + exact: bool = False + """``location`` is the ``output_path`` itself and the result must be a zip there + holding the products as members — how an hpc-campaign's ``success`` check reads + it (``autofit.campaign.collector.check_success``).""" + error: Optional[str] = None + """Why where the result should be is not known (e.g. a manifest line missing a + field); such a result is reported ``failed`` with this reason.""" + + +def _datasets_from_file(path: Path) -> List[str]: + datasets = [] + for line in path.read_text().splitlines(): + line = line.strip() + if line and not line.startswith("#"): + datasets.append(line.split()[0]) + return datasets + + +def _products(value) -> Tuple[str, ...]: + if value is None: + return () + if isinstance(value, str): + return (value,) + return tuple(str(product) for product in value) + + +def _products_and_location(value) -> Tuple[Tuple[str, ...], Optional[str]]: + if value is None or isinstance(value, str): + return _products(value), None + if isinstance(value, (list, tuple)): + return tuple(str(product) for product in value), None + if isinstance(value, Mapping): + unknown = set(value) - {"products", "location"} + if unknown: + raise LibraryException( + f"A search takes 'products' and 'location', not {sorted(unknown)}." + ) + return _products(value.get("products")), value.get("location") + raise LibraryException( + f"A search is a list of products or a mapping, not {value!r}." + ) + + +class ResultsLibrary: + def __init__( + self, + expected: Iterable[ExpectedResult], + roots: Union[PathLike, Sequence[PathLike]], + index_dir: Optional[PathLike] = None, + name: str = "library", + latest_key: Tuple[str, ...] = DEFAULT_LATEST_KEY, + ): + """ + The results ``expected`` under ``roots``. + + Parameters + ---------- + expected + Every result the library should hold. + roots + Where the results are (as for ``ResultsIndex``). + index_dir + Where the index is kept (as for ``ResultsIndex``). + name + Names the gaps file, ``.missing.csv``. + """ + self.expected = list(expected) + if isinstance(roots, (str, os.PathLike)): + roots = [roots] + self.roots = [Path(root) for root in roots] + if len(self.roots) != 1 and index_dir is None: + raise LibraryException("A library with several roots needs an index_dir.") + self.index_dir = index_dir + self.name = name + self.latest_key = tuple(latest_key) + + @classmethod + def grid( + cls, + datasets: Iterable[str], + searches: Mapping[str, Union[Sequence[str], Mapping, None]], + roots: Union[PathLike, Sequence[PathLike]], + **kwargs, + ) -> "ResultsLibrary": + """ + Every dataset × every search. + + Parameters + ---------- + searches + Search name to its required products, or to ``{products: [...], + location: "..."}`` where ``location`` is a template with ``{dataset}``, + relative to the first root. + """ + if isinstance(roots, (str, os.PathLike)): + roots = [roots] + first_root = Path(roots[0]) + expected = [] + for dataset in datasets: + for search, value in searches.items(): + products, location = _products_and_location(value) + if location is not None: + location = str(first_root / location.format(dataset=dataset)) + expected.append( + ExpectedResult( + dataset=str(dataset), + search=str(search), + products=products, + location=location, + ) + ) + return cls(expected, roots=roots, **kwargs) + + @classmethod + def from_campaign( + cls, + campaign, + stages: Optional[Sequence[str]] = None, + products: Optional[Mapping[str, Sequence[str]]] = None, + roots: Optional[Union[PathLike, Sequence[PathLike]]] = None, + **kwargs, + ) -> "ResultsLibrary": + """ + The library an hpc-campaign declares: each manifest line × each stage. + + Parameters + ---------- + campaign + An ``autofit.campaign.Campaign`` or the path of its ``campaign.yaml``. + stages + The stages to expect (default: every stage with a ``success.zip_complete``). + products + Required products added per stage to its ``success.members``. + roots + Where to index (default: the campaign's ``status_root``). + """ + from autofit.campaign.spec import Campaign + + if not isinstance(campaign, Campaign): + campaign = Campaign.load(campaign) + products = dict(products or {}) + + if stages is None: + stages = [ + name + for name, stage in campaign.stages.items() + if stage.success.zip_complete + ] + for name in stages: + if name not in campaign.stages: + raise LibraryException(f"Campaign {campaign.name} has no stage {name}.") + if not campaign.stages[name].success.zip_complete: + raise LibraryException( + f"Stage {name} declares no success.zip_complete, so where its " + f"results are is not known." + ) + unknown = set(products) - set(stages) + if unknown: + raise LibraryException(f"products names unknown stages {sorted(unknown)}.") + + expected = [] + for index in campaign.indices(): + fields = campaign.fields(index) + for name in stages: + success = campaign.stages[name].success + required = tuple(success.members) + _products(products.get(name)) + try: + location = campaign.path(success.zip_complete.format(**fields)) + except (KeyError, IndexError) as exception: + expected.append( + ExpectedResult( + dataset=fields["line"], + search=name, + products=required, + exact=True, + error=f"manifest_line:{index} has no field {exception}", + ) + ) + continue + if location.suffix == ".zip": + location = location.with_suffix("") + expected.append( + ExpectedResult( + dataset=fields["line"], + search=name, + products=required, + location=str(location), + exact=True, + ) + ) + + kwargs.setdefault("name", str(campaign.name)) + return cls( + expected, + roots=roots if roots is not None else campaign.path(campaign.status_root), + **kwargs, + ) + + @classmethod + def from_yaml(cls, path: PathLike) -> "ResultsLibrary": + """ + The library declared in ``library.yaml`` (see the module docstring). Relative + paths are relative to the file. + """ + path = Path(path) + try: + raw = yaml.safe_load(path.read_text()) + except (OSError, yaml.YAMLError) as exception: + raise LibraryException(f"Library file {path} cannot be read: {exception}") + if not isinstance(raw, dict): + raise LibraryException(f"Library file {path} must be a mapping.") + unknown = set(raw) - LIBRARY_FIELDS + if unknown: + raise LibraryException(f"Library file {path}: unknown keys {sorted(unknown)}.") + + def relative(value): + value = Path(value) + return value if value.is_absolute() else path.parent / value + + roots = raw.get("roots") + if isinstance(roots, (str, os.PathLike)): + roots = [roots] + roots = [relative(root) for root in roots] if roots else None + kwargs = {"name": raw.get("name") or path.stem} + if raw.get("index_dir") is not None: + kwargs["index_dir"] = relative(raw["index_dir"]) + + if raw.get("from_campaign"): + clashes = {"datasets", "datasets_file", "searches"} & set(raw) + if clashes: + raise LibraryException( + f"Library file {path}: from_campaign takes its datasets and " + f"searches from the campaign; drop {sorted(clashes)}." + ) + return cls.from_campaign( + relative(raw["from_campaign"]), + stages=raw.get("stages"), + products=raw.get("products"), + roots=roots, + **kwargs, + ) + + for key in ("stages", "products"): + if key in raw: + raise LibraryException( + f"Library file {path}: '{key}' only applies with from_campaign." + ) + if roots is None: + raise LibraryException(f"Library file {path} names no roots.") + if ("datasets" in raw) == ("datasets_file" in raw): + raise LibraryException( + f"Library file {path} needs exactly one of datasets and datasets_file." + ) + datasets = ( + [str(dataset) for dataset in raw["datasets"]] + if "datasets" in raw + else _datasets_from_file(relative(raw["datasets_file"])) + ) + searches = raw.get("searches") + if not isinstance(searches, dict) or not searches: + raise LibraryException(f"Library file {path} declares no searches.") + return cls.grid(datasets, searches, roots=roots, **kwargs) + + def index(self) -> ResultsIndex: + return ResultsIndex( + self.roots, index_dir=self.index_dir, latest_key=self.latest_key + ) + + def check(self, update: bool = True): + """ + Compare the library with the results index. + + Parameters + ---------- + update + Bring the index up to date first (only new or changed results are opened). + + Returns + ------- + A ``CompletenessReport``. + """ + from autofit.results.completeness import check + + index = self.index() + if update: + index.update() + return check(self, index.records()) diff --git a/autofit/results/shard.py b/autofit/results/shard.py new file mode 100644 index 000000000..3f11610b0 --- /dev/null +++ b/autofit/results/shard.py @@ -0,0 +1,60 @@ +""" +Split a set of results into ``N`` disjoint shards, so one pass over thousands of +results (``Aggregator.enrich``, an export) can be shared between Slurm array tasks, a +laptop process pool or anything else, with no coordination between the workers. + +A result's shard depends only on its ``output_path`` (without ``.zip``), never on the +order results are found in, so every worker that sees the tree at the same path puts a +result in the same shard and the shards are disjoint. Run the workers of one pass on +one mount of the tree. +""" +import hashlib +import os +from typing import Optional, Tuple, Union + +ShardLike = Union[None, str, Tuple[int, int]] + + +def parse_shard(shard: ShardLike) -> Optional[Tuple[int, int]]: + """ + ``(i, n)`` from ``"i/N"`` or ``(i, N)`` (``None`` for no sharding). + + Raises + ------ + ValueError + Unless ``0 <= i < N``. + """ + if shard is None: + return None + if isinstance(shard, str): + try: + i, n = (int(part) for part in shard.split("/")) + except ValueError: + raise ValueError(f"A shard is 'i/N' (e.g. '0/4'), not {shard!r}.") + else: + i, n = (int(part) for part in shard) + if n < 1 or not 0 <= i < n: + raise ValueError(f"Shard {i}/{n} is not one of 0/{n} .. {n - 1}/{n}.") + return i, n + + +def shard_of(output_path: Union[str, os.PathLike], n: int) -> int: + """ + Which of ``n`` shards the result at ``output_path`` (with or without ``.zip``) + belongs to. + """ + path = str(output_path) + if path.endswith(".zip"): + path = path[: -len(".zip")] + return int(hashlib.sha1(path.encode("utf-8")).hexdigest(), 16) % n + + +def in_shard(output_path: Union[str, os.PathLike], shard: ShardLike) -> bool: + """ + Whether the result at ``output_path`` is in ``shard`` (always, for ``None``). + """ + shard = parse_shard(shard) + if shard is None: + return True + i, n = shard + return shard_of(output_path, n) == i diff --git a/autofit/tools/util.py b/autofit/tools/util.py index 0bff1d1f7..794aec10b 100644 --- a/autofit/tools/util.py +++ b/autofit/tools/util.py @@ -1,12 +1,13 @@ import inspect import json import os +import shutil import sys import zipfile from contextlib import contextmanager from functools import wraps from pathlib import Path -from typing import Union +from typing import Optional, Union import numpy as np @@ -123,6 +124,65 @@ def zip_directory(source_directory, output=None): os.replace(tmp, output) +def write_zip_member( + zip_path, + arcname: str, + data: Optional[bytes] = None, + file_path=None, + compress_type: Optional[int] = None, +): + """ + Add the member ``arcname`` to the zip at ``zip_path``, or replace it if the zip + already holds one, from ``data`` or from the file at ``file_path``. + + ``zipfile`` can neither overwrite a member in place nor append without risk: a + process killed during an append leaves a zip whose central directory is gone, + which no later read can open. So every other member is streamed into a new + archive written next to the original (same directory, so the final + ``os.replace`` is atomic on one filesystem), and the new archive is only swapped + in once complete. Every other member keeps its name, compression and content; a + failure part-way leaves the original zip untouched. + + Parameters + ---------- + compress_type + How the member is stored; by default the replaced member's compression, or + for a new member ``compress_type_for(arcname)``. + """ + if (data is None) == (file_path is None): + raise ValueError("Pass exactly one of data or file_path.") + + zip_path = Path(zip_path) + tmp = zip_path.with_name(f"{zip_path.name}.{os.getpid()}.tmp") + + try: + with zipfile.ZipFile(zip_path, "r") as source: + if compress_type is None: + try: + compress_type = source.getinfo(arcname).compress_type + except KeyError: + compress_type = compress_type_for(arcname) + + with zipfile.ZipFile(tmp, "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) + + if file_path is not None: + destination.write(file_path, arcname, compress_type=compress_type) + else: + destination.writestr(arcname, data, compress_type=compress_type) + + os.replace(tmp, zip_path) + except BaseException: + tmp.unlink(missing_ok=True) + raise + + def open_(filename, *flags): directory = Path(filename) os.makedirs(directory.parent, exist_ok=True) diff --git a/test_autofit/results/test_completeness.py b/test_autofit/results/test_completeness.py new file mode 100644 index 000000000..9b88b55be --- /dev/null +++ b/test_autofit/results/test_completeness.py @@ -0,0 +1,281 @@ +import csv +import textwrap +from pathlib import Path + +import pytest + +from autofit.results import ( + LibraryException, + ResultsLibrary, +) +from autofit.results.__main__ import main +from autofit.tools.util import write_zip_member + +from .tree import write_folder, write_status, write_zip + +PRODUCTS = ["image", "files/model.fits"] + + +def _with_model(output_path): + write_zip_member(f"{output_path}.zip", "files/model.fits", data=b"SIMPLE") + return output_path + + +@pytest.fixture(name="holes") +def make_holes(results_root): + """ + One result per way a library can have a hole. + """ + _with_model(write_zip(results_root, prefix="d0")) # done + # d1: no result at all + write_zip(results_root, prefix="d2", completed=False) # zip without .completed + write_zip(results_root, prefix="d3") # completed, no files/model.fits + write_folder(results_root, prefix="d4", completed=False, checkpoint="valid") + _with_model(write_zip(results_root, prefix="d5", seed=1, date_time=(2026, 9, 1, 0, 0, 0))) + _with_model(write_zip(results_root, prefix="d5", seed=2, date_time=(2026, 10, 1, 0, 0, 0))) + failed = results_root / "d6" / "search" / ("f" * 32) + failed.parent.mkdir(parents=True) + write_status(failed, "failed", failure_class="oom") # failed before search.json + _with_model(write_zip(results_root, prefix="a/d7", unique_tag="d7", seed=3)) + _with_model(write_zip(results_root, prefix="b/d7", unique_tag="d7", seed=4)) + _with_model(write_zip(results_root, prefix="zz")) # not declared + return results_root + + +def _by_dataset(report): + return {result.expected.dataset: result for result in report.results} + + +def test_every_hole_has_its_reason(holes): + library = ResultsLibrary.grid( + [f"d{i}" for i in range(8)], {"search": PRODUCTS}, roots=holes + ) + report = library.check() + results = _by_dataset(report) + + assert results["d0"].status == "done" + assert (results["d1"].status, results["d1"].reason) == ("missing", "no_result") + assert results["d2"].status == "failed" + assert results["d2"].reason.startswith("not_completed:corrupt_zip") + assert results["d3"].status == "incomplete" + assert results["d3"].missing_products == ["files/model.fits"] + assert results["d4"].status == "incomplete" + assert results["d4"].reason.startswith("not_completed:interrupted_valid_ckpt") + assert results["d5"].status == "done" # the superseded sibling does not count + assert results["d6"].status == "failed" + assert results["d6"].reason.startswith("not_completed:failed") + assert (results["d7"].status, results["d7"].reason) == ("duplicate", "duplicate:2") + + assert report.counts() == { + "expected": 8, + "done": 2, + "missing": 1, + "incomplete": 2, + "failed": 2, + "duplicate": 1, + "unexpected": 1, + } + assert report.product_counts()["files/model.fits"] == {"required": 8, "missing": 1} + assert [record.path_prefix for record in report.unexpected] == ["zz"] + assert not report.is_complete + + +def test_missing_csv(holes, tmp_path): + library = ResultsLibrary.grid( + [f"d{i}" for i in range(8)], {"search": PRODUCTS}, roots=holes, name="lib" + ) + path = library.check().write_missing_csv(tmp_path) + + assert path == tmp_path / "lib.missing.csv" + with open(path, newline="") as f: + rows = list(csv.DictReader(f)) + gaps = {(row["dataset"], row["product"]): row for row in rows} + assert set(gaps) == { + ("d1", ""), + ("d2", ""), + ("d3", "files/model.fits"), + ("d4", ""), + ("d6", ""), + ("d7", ""), + } + assert gaps[("d3", "files/model.fits")]["reason"] == "missing_product:files/model.fits" + assert gaps[("d3", "files/model.fits")]["output_path"].endswith(".zip") is False + assert gaps[("d1", "")]["output_path"] == "" + assert not list(tmp_path.glob("*.tmp")) + + +def test_named_and_folder_products(results_root): + write_zip(results_root, prefix="d0") + write_folder(results_root, prefix="d1") + library = ResultsLibrary.grid( + ["d0", "d1"], {"search": ["image", "image/", "samples", "latent"]}, roots=results_root + ) + results = _by_dataset(library.check()) + + assert results["d0"].missing_products == ["latent"] + assert results["d1"].missing_products == ["latent"] + + +def test_location(results_root): + write_zip(results_root, prefix="d0") + write_zip(results_root, prefix="moved", unique_tag="d1") + library = ResultsLibrary.grid( + ["d0", "d1"], + {"search": {"products": ["image"], "location": "{dataset}/search"}}, + roots=results_root, + ) + results = _by_dataset(library.check()) + + assert results["d0"].status == "done" + # Found by key it would be done; at its declared location there is nothing. + assert results["d1"].reason == "no_result" + + +def test_from_yaml(holes, tmp_path): + (tmp_path / "manifest.txt").write_text("# objects\nd0 extra\n\nd3\n") + path = tmp_path / "lib.yaml" + path.write_text( + textwrap.dedent( + f""" + roots: [{holes}] + datasets_file: manifest.txt + searches: + search: [image, files/model.fits] + """ + ) + ) + library = ResultsLibrary.from_yaml(path) + report = library.check() + + assert library.name == "lib" + assert [e.dataset for e in library.expected] == ["d0", "d3"] + assert report.counts()["done"] == 1 + + +@pytest.mark.parametrize( + "text", + [ + "roots: [o]\nsearches: {s: []}\n", + "roots: [o]\ndatasets: [a]\ndatasets_file: m.txt\nsearches: {s: []}\n", + "roots: [o]\ndatasets: [a]\n", + "datasets: [a]\nsearches: {s: []}\n", + "roots: [o]\ndatasets: [a]\nsearches: {s: []}\nstages: [x]\n", + "roots: [o]\ndatasets: [a]\nsearches: {s: {product: [x]}}\n", + "roots: [o]\ndatasets: [a]\nsearches: {s: []}\ntypo: 1\n", + "from_campaign: c.yaml\ndatasets: [a]\n", + ], +) +def test_from_yaml_refuses(tmp_path, text): + path = tmp_path / "lib.yaml" + path.write_text(text) + with pytest.raises(LibraryException): + ResultsLibrary.from_yaml(path) + + +def _campaign_zip(output, tile, completed=True, members=True): + """ + A result zip exactly where the test campaign's ``success.zip_complete`` names it: + ``output/dr1//vis_lp.zip``. + """ + path = write_zip(output, prefix="dr1", name=tile, completed=completed) + target = output / "dr1" / tile / "vis_lp" + Path(f"{path}.zip").rename(f"{target}.zip") + return target + + +def test_from_campaign(tmp_path): + from autofit.campaign.collector import check_success + from autofit.campaign.spec import Campaign + from test_autofit.campaign.conftest import write_campaign + + project = tmp_path / "project" + campaign_file = write_campaign(project, n=4) + output = project / "output" + _campaign_zip(output, "Tile000") + _campaign_zip(output, "Tile001", completed=False) + # Under /: not where the campaign's success check looks. + write_zip(output, prefix="dr1/Tile003", name="vis_lp") + + library = ResultsLibrary.from_campaign(campaign_file, products={"vis_lp": "image"}) + report = library.check() + results = _by_dataset(report) + + assert library.name == "demo" + assert [e.search for e in library.expected] == ["vis_lp"] * 4 + assert library.expected[0].products == ("files/samples_summary.json", "image") + assert results["dr1 Tile000"].status == "done" + assert results["dr1 Tile001"].status == "failed" + assert results["dr1 Tile002"].reason == "no_result" + assert results["dr1 Tile003"].reason == "no_result" + + # The library and the campaign collector agree on every task. + campaign = Campaign.load(campaign_file) + stage = campaign.stages["vis_lp"] + for index in campaign.indices(): + done = results[campaign.fields(index)["line"]].status == "done" + assert done == bool(check_success(campaign, stage, index)) + + yaml_path = project / "lib.yaml" + yaml_path.write_text("from_campaign: campaign.yaml\nproducts: {vis_lp: [image]}\n") + assert ResultsLibrary.from_yaml(yaml_path).check().counts() == report.counts() + + with pytest.raises(LibraryException): + ResultsLibrary.from_campaign(campaign_file, stages=["vis_pix"]) + + +def test_from_campaign_short_manifest_line(tmp_path): + from test_autofit.campaign.conftest import write_campaign + + project = tmp_path / "project" + campaign_file = write_campaign(project, n=2) + (project / "manifest.txt").write_text("dr1 Tile000\ndr1\n") + + results = _by_dataset(ResultsLibrary.from_campaign(campaign_file).check()) + + assert results["dr1"].status == "failed" + assert results["dr1"].reason.startswith("manifest_line:1 has no field") + assert results["dr1 Tile000"].reason == "no_result" + + +def test_exact_location_counts_a_superseded_result(results_root): + from autofit.results import ExpectedResult + + original = write_zip(results_root, prefix="d0") + write_zip(results_root / "zbackup", prefix="d0") # same key, wins the tie + library = ResultsLibrary( + [ExpectedResult("d0", "search", (), str(original), exact=True)], + roots=results_root, + ) + assert library.check().results[0].status == "done" + + +def test_string_product_is_one_product(results_root): + write_zip(results_root, prefix="d0") + library = ResultsLibrary.grid( + ["d0"], {"search": {"products": "image"}}, roots=results_root + ) + assert library.expected[0].products == ("image",) + assert library.check().results[0].status == "done" + + +def test_several_roots_need_an_index_dir(tmp_path): + with pytest.raises(LibraryException): + ResultsLibrary([], roots=[tmp_path / "a", tmp_path / "b"]) + + +def test_cli_check(holes, tmp_path, capsys): + path = tmp_path / "lib.yaml" + path.write_text(f"roots: [{holes}]\ndatasets: [d0, d3]\nsearches: {{search: [files/model.fits]}}\n") + + assert main(["check", str(path)]) == 1 + out = capsys.readouterr().out + assert "1 of 2 expected results done" in out + assert (tmp_path / "lib.missing.csv").exists() + + path.write_text(f"roots: [{holes}]\ndatasets: [d0]\nsearches: {{search: [image]}}\n") + assert main(["check", str(path), "--no-update", "--output-dir", str(tmp_path / "o")]) == 0 + assert (tmp_path / "o" / "lib.missing.csv").read_text().count("\n") == 1 + + # An error is 2, never 1 ("incomplete"). + path.write_text(f"roots: [{holes}, {tmp_path}]\ndatasets: [d0]\nsearches: {{search: []}}\n") + assert main(["check", str(path)]) == 2 diff --git a/test_autofit/results/test_enrich.py b/test_autofit/results/test_enrich.py new file mode 100644 index 000000000..42928bb2b --- /dev/null +++ b/test_autofit/results/test_enrich.py @@ -0,0 +1,216 @@ +import hashlib +import json +import zipfile +from pathlib import Path + +import pytest + +from autofit.aggregator.aggregator import Aggregator +from autofit.results.__main__ import main +from autofit.results.shard import in_shard, parse_shard, shard_of + +from autofit.tools.util import zip_directory + +from .tree import write_folder, write_zip + + +def _digests(zip_path): + with zipfile.ZipFile(zip_path) as f: + return { + info.filename: (hashlib.sha256(f.read(info)).hexdigest(), info.compress_type) + for info in f.infolist() + } + + +@pytest.fixture(name="library") +def make_library(results_root): + return [write_zip(results_root, prefix=f"d{i:03d}", seed=i) for i in range(100)] + + +def _extra(search_output): + return {"identifier": search_output.directory.zip_path.stem} + + +def test_enrich_adds_member_and_keeps_the_rest(results_root, library): + before = {path: _digests(f"{path}.zip") for path in library} + + report = Aggregator.from_index(results_root).enrich(_extra, name="extra") + + assert report.counts()["written"] == 100 + assert report.name == "files/extra.json" + for path in library: + after = _digests(f"{path}.zip") + assert {k: v for k, v in after.items() if k != "files/extra.json"} == before[path] + with zipfile.ZipFile(f"{path}.zip") as f: + assert json.loads(f.read("files/extra.json")) == {"identifier": path.name} + assert not list(results_root.rglob("*.tmp")) + + +def test_overwrite(results_root, library): + aggregator = Aggregator.from_index(results_root) + aggregator.enrich(lambda o: "one", name="note.txt") + + report = aggregator.enrich(lambda o: "two", name="note.txt") + assert report.counts()["exists"] == 100 + + report = aggregator.enrich(lambda o: "two", name="note.txt", overwrite=True) + assert report.counts()["written"] == 100 + with zipfile.ZipFile(f"{library[0]}.zip") as f: + assert f.read("files/note.txt") == b"two" + assert f.namelist().count("files/note.txt") == 1 + + +def test_shards_partition(results_root, library): + aggregator = Aggregator.from_index(results_root) + written = [] + for i in range(4): + report = aggregator.enrich(_extra, name="extra", shard=f"{i}/4") + assert report.skipped_shard == 100 - len(report.written) + written.extend(report.written) + assert sorted(written) == sorted(str(path) for path in library) + + +def test_shard_helpers(): + assert parse_shard("1/4") == (1, 4) + assert parse_shard((0, 1)) == (0, 1) + assert parse_shard(None) is None + for bad in ("4/4", "-1/2", "x", "1/0"): + with pytest.raises(ValueError): + parse_shard(bad) + assert shard_of("/x/d0/search/abc", 7) == shard_of("/x/d0/search/abc.zip", 7) + # Objects whose identifiers coincide still spread over the shards. + shards = {shard_of(f"/x/d{i}/slam/abc", 4) for i in range(40)} + assert shards == {0, 1, 2, 3} + assert in_shard("/x", None) + + +def test_failure_keeps_zip_and_goes_on(results_root, library): + original = Path(f"{library[0]}.zip").read_bytes() + + def fn(search_output): + if search_output.directory.zip_path.stem == library[0].name: + raise RuntimeError("boom") + return {"ok": True} + + report = Aggregator.from_index(results_root).enrich(fn, name="extra") + + assert report.counts()["written"] == 99 + assert report.errors == [(str(library[0]), "RuntimeError: boom")] + assert Path(f"{library[0]}.zip").read_bytes() == original + + +def test_killed_write_leaves_zip_intact(results_root, monkeypatch): + path = write_zip(results_root, prefix="d0") + original = Path(f"{path}.zip").read_bytes() + aggregator = Aggregator.from_index(results_root) + + import autofit.tools.util as util + + def killed(*args, **kwargs): + raise KeyboardInterrupt + + monkeypatch.setattr(util.shutil, "copyfileobj", killed) + with pytest.raises(KeyboardInterrupt): + aggregator.enrich(_extra, name="extra") + + assert Path(f"{path}.zip").read_bytes() == original + assert not list(results_root.rglob("*.tmp")) + + +def test_declined_and_bad_values(results_root, library): + aggregator = Aggregator.from_index(results_root) + assert aggregator.enrich(lambda o: None, name="x").counts()["declined"] == 100 + report = aggregator.enrich(lambda o: 1.5, name="x") + assert report.counts()["errors"] == 100 + assert "TypeError" in report.errors[0][1] + with pytest.raises(ValueError): + aggregator.enrich(_extra, name="../escape") + + +def test_file_path_value(results_root, library, tmp_path): + image = tmp_path / "model.fits" + image.write_bytes(b"SIMPLE") + Aggregator.from_index(results_root).enrich(lambda o: image, name="model.fits") + with zipfile.ZipFile(f"{library[0]}.zip") as f: + info = f.getinfo("files/model.fits") + assert f.read(info) == b"SIMPLE" + assert info.compress_type == zipfile.ZIP_STORED + + +def test_folder_result(results_root): + path = write_folder(results_root, prefix="d0") + report = Aggregator.from_directory(results_root).enrich(lambda o: [1, 2], name="extra") + assert report.counts()["written"] == 1 + assert json.loads((path / "files" / "extra.json").read_text()) == [1, 2] + + +def test_temporary_extraction_refused(results_root, library): + aggregator = Aggregator.from_directory(results_root, unzip_temporary=True) + with pytest.raises(ValueError, match="temporary"): + aggregator.enrich(_extra, name="extra") + + +def test_cli_enrich(results_root, library, tmp_path, monkeypatch): + (tmp_path / "products.py").write_text( + "def note(search_output):\n return 'hello'\n" + ) + monkeypatch.chdir(tmp_path) + monkeypatch.syspath_prepend(str(tmp_path)) + assert main(["enrich", str(results_root), "--fn", "products:note", "--name", "note.txt", "--shard", "0/2"]) == 0 + assert main(["enrich", str(results_root), "--fn", "products:note", "--name", "note.txt", "--shard", "1/2"]) == 0 + for path in library: + with zipfile.ZipFile(f"{path}.zip") as f: + assert f.read("files/note.txt") == b"hello" + assert main(["enrich", str(results_root), "--fn", "products:nope", "--name", "x"]) == 2 + + +def _folder_and_zip(root, prefix="d0"): + path = write_folder(root, prefix=prefix) + zip_directory(path, f"{path}.zip") + return path + + +def test_folder_and_zip_both_written(results_root): + path = _folder_and_zip(results_root) + + report = Aggregator.from_index(results_root).enrich(_extra, name="extra") + + assert report.counts()["written"] == 1 + assert (path / "files" / "extra.json").is_file() + # A rerun of the search re-zips the folder: the product survives it. + zip_directory(path, f"{path}.zip") + with zipfile.ZipFile(f"{path}.zip") as f: + assert "files/extra.json" in f.namelist() + + +def test_stopped_between_copies_is_finished_by_a_rerun(results_root, monkeypatch): + path = _folder_and_zip(results_root) + aggregator = Aggregator.from_index(results_root) + + import importlib + + enrich_module = importlib.import_module("autofit.results.enrich") + + real_open_atomic = enrich_module.open_atomic + + def full_disk(*args, **kwargs): + raise OSError("disk full") + + monkeypatch.setattr(enrich_module, "open_atomic", full_disk) + assert aggregator.enrich(_extra, name="extra").counts()["errors"] == 1 + monkeypatch.setattr(enrich_module, "open_atomic", real_open_atomic) + + report = aggregator.enrich(_extra, name="extra") + assert report.counts() == { + "written": 1, + "exists": 0, + "declined": 0, + "errors": 0, + "other_shards": 0, + } + assert (path / "files" / "extra.json").is_file() + assert aggregator.enrich(_extra, name="extra").counts()["exists"] == 1 + + +def test_cli_enrich_errors_exit_2(results_root, library): + assert main(["enrich", str(results_root), "--fn", "json:dumps", "--name", "x", "--shard", "9/2"]) == 2