From 349bd0eb2f8c3721f6b0207f77fe443304679bb6 Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Fri, 30 Oct 2026 09:34:57 +0000 Subject: [PATCH] =?UTF-8?q?feat:=20read=20zipped=20results=20in=20place=20?= =?UTF-8?q?=E2=80=94=20ZipMember,=20from=5Fzip,=20open=5Fmember,=20value(k?= =?UTF-8?q?ind=3D),=20Aggregator.from=5Findex=20(results-library=20p1,=20P?= =?UTF-8?q?R=201b)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - ZipMember: a pathlib-like path into a result zip that holds no open file (each open reopens the zip), so thousands of zipped outputs never hit the open-file limit; member names cached per (size, mtime). - FileOutput / SearchOutput read through path.open(), so a SearchOutput can be backed by a zip (SearchOutput.from_zip) with nothing extracted; FITS members load from memory. - SearchOutput.open_member(path) public; value(name, kind=...) chooses between files sharing a name (tracer.json vs tracer.fits) and warns once when a bare name is ambiguous (order unchanged: json, pickle, csv, fits). - Aggregator.from_index(roots, index_dir, states=('completed',)) aggregates over the results index, superseded results left out by default. Refs #1739 Co-Authored-By: Claude Opus 5.5 --- autofit/aggregator/aggregator.py | 60 +++++++++ autofit/aggregator/file_output.py | 16 ++- autofit/aggregator/search_output.py | 108 ++++++++++++---- autofit/aggregator/zip_member.py | 161 ++++++++++++++++++++++++ test_autofit/results/test_lazy_reads.py | 157 +++++++++++++++++++++++ 5 files changed, 476 insertions(+), 26 deletions(-) create mode 100644 autofit/aggregator/zip_member.py create mode 100644 test_autofit/results/test_lazy_reads.py diff --git a/autofit/aggregator/aggregator.py b/autofit/aggregator/aggregator.py index f6f7212cd..4d4578de4 100755 --- a/autofit/aggregator/aggregator.py +++ b/autofit/aggregator/aggregator.py @@ -475,6 +475,66 @@ def scan_all(scan_directory): errors=errors, ) + @classmethod + def from_index( + cls, + roots: Union[str, os.PathLike, List[Union[str, os.PathLike]]], + index_dir: Optional[Union[str, os.PathLike]] = None, + states: Tuple[str, ...] = ("completed",), + reference: Optional[dict] = None, + update: bool = True, + ) -> "Aggregator": + """ + Aggregate the results a results index (``autofit.results.ResultsIndex``) lists, + reading zipped results in place: nothing is extracted to disk. + + Parameters + ---------- + roots + The directory (or directories) holding the results. + index_dir + Where the index is kept (default ``/.autofit_index``; required with + several roots or a read-only tree). + states + The result states to include. The default, ``("completed",)``, leaves out + superseded results (an older sibling of a newer completed fit); add + ``"superseded"`` to keep them. + reference + A dictionary mapping paths to types to be used when loading models from disk. + update + Bring the index up to date first (only new or changed results are opened). + ``False`` uses the index as it stands. + """ + from autofit.results import ResultsIndex + + index = ResultsIndex(roots, index_dir=index_dir) + if update: + index.update() + + states = {str(state) for state in states} + search_outputs = [] + for record in index.records(): + if record.state not in states: + continue + output_path = Path(record.output_path) + if record.has_zip and record.state != "corrupt_zip": + search_outputs.append( + SearchOutput.from_zip(f"{output_path}.zip", reference=reference) + ) + elif record.has_dir: + search_outputs.append(SearchOutput(output_path, reference=reference)) + + logger.info( + f"Aggregator.from_index: {len(search_outputs)} search outputs in states " + f"{sorted(states)}." + ) + if not search_outputs: + logger.warning( + f"Aggregator.from_index found no search outputs in states " + f"{sorted(states)} (index: {index.index_dir}, counts: {index.counts()})." + ) + return cls(search_outputs, []) + def add_directory( self, directory: Union[str, Path], diff --git a/autofit/aggregator/file_output.py b/autofit/aggregator/file_output.py index e9b4ffb74..ecb4fc5aa 100644 --- a/autofit/aggregator/file_output.py +++ b/autofit/aggregator/file_output.py @@ -1,3 +1,4 @@ +import io import json from abc import abstractmethod, ABC from pathlib import Path @@ -33,7 +34,8 @@ def __init__(self, name: str, path: Path): The name of the output file. This is the path to the file from the root output directory separated by '.' and without a suffix. e.g. output/directory/info.json -> directory.info path - The path to the output file + The path to the output file: a ``Path``, or a ``ZipMember`` when the + output is read from inside a zip. """ self.name = name self.path = path @@ -56,7 +58,8 @@ def value(self) -> np.ndarray: """ The array stored in the csv file """ - return np.loadtxt(self.path, delimiter=",") + with self.path.open() as f: + return np.loadtxt(f, delimiter=",") class JSONOutput(FileOutput): @@ -65,7 +68,7 @@ def dict(self) -> Union[dict, list]: """ The dictionary stored in the json file """ - with open(self.path) as f: + with self.path.open() as f: return json.load(f) @property @@ -88,7 +91,7 @@ def value(self): """ The object stored in the pickle file """ - with open(self.path, "rb") as f: + with self.path.open("rb") as f: return dill.load(f) @@ -100,4 +103,7 @@ def value(self): """ from astropy.io import fits - return fits.open(self.path) + if isinstance(self.path, Path): + return fits.open(self.path) + # A member of a zip (``ZipMember``): read into memory, nothing extracted. + return fits.open(io.BytesIO(self.path.read_bytes())) diff --git a/autofit/aggregator/search_output.py b/autofit/aggregator/search_output.py index 43af08b0b..7d143cbb8 100644 --- a/autofit/aggregator/search_output.py +++ b/autofit/aggregator/search_output.py @@ -24,6 +24,7 @@ JSONOutput, FileOutput, ) +from autofit.aggregator.zip_member import ZipMember from autofit.mapper.identifier import Identifier from autofit.non_linear.samples.sample import samples_from_iterator from autonerves.dictable import from_dict @@ -53,6 +54,13 @@ def _create_file_handle(*args, **kwargs): dill._dill._create_filehandle = _create_file_handle +_KINDS = (".json", ".pickle", ".csv", ".fits") +""" +The file kinds ``value`` loads, in the order a bare name resolves them. +""" + +_warned_ambiguous = set() + class TemporaryExtraction: """ @@ -117,7 +125,7 @@ def completed_at(self) -> Optional[str]: """ When the search completed (ISO 8601), or ``None`` if it has not. - Taken from the status record beside the output (``.status.json``) + Taken from the status record beside the output (``.status.json``) when it says ``completed``, else from the ``.completed`` member of the zip beside the output, else from the ``.completed`` file's modification time. The zip's own modification time is never used: copying or re-zipping resets it. @@ -125,21 +133,28 @@ def completed_at(self) -> Optional[str]: ``.completed`` time is the extraction time; index the tree (``python -m autofit.results index``) to compare such outputs. """ - directory = Path(self.directory) + if isinstance(self.directory, ZipMember): + zip_path = self.directory.zip_path + output_path = zip_path.with_name(zip_path.name[: -len(".zip")]) + else: + output_path = Path(self.directory) + zip_path = Path(f"{output_path}.zip") try: - with open(f"{directory}.status.json") as f: + with open(f"{output_path}.status.json") as f: status = json.load(f) if status.get("state") == "completed" and status.get("updated"): return status["updated"] except (OSError, ValueError, AttributeError): pass try: - with zipfile.ZipFile(f"{directory}.zip") as f: + with zipfile.ZipFile(zip_path) as f: return datetime(*f.getinfo(".completed").date_time).isoformat() except (OSError, KeyError, zipfile.BadZipFile): pass + if isinstance(self.directory, ZipMember): + return None try: - mtime = (directory / ".completed").stat().st_mtime + mtime = (output_path / ".completed").stat().st_mtime except OSError: return None return datetime.fromtimestamp(mtime, timezone.utc).isoformat( @@ -241,9 +256,24 @@ def image(self, name: str) -> Image.Image: ------- The loaded image """ - return Image.open(self.directory / "image" / f"{name}.png") + return Image.open((self.directory / "image" / f"{name}.png").open("rb")) - def value(self, name: str): + def open_member(self, path: str, mode: str = "rb"): + """ + Open a file of this output by its path relative to the output (e.g. + ``"files/tracer.fits"`` or ``"image/fit.png"``), whether the output is a folder + or a zip read in place. + + Parameters + ---------- + path + The file's path inside the output, with ``/`` separators. + mode + ``"rb"`` (the default) or ``"r"``. + """ + return (self.directory / path).open(mode) + + def value(self, name: str, kind: Optional[str] = None): """ Load the value of some object in the files directory for the search. @@ -258,19 +288,47 @@ def value(self, name: str): ---------- name The name of the file to load without a file suffix. + kind + Which file to load when several share ``name``: ``"json"``, ``"pickle"``, + ``"csv"`` or ``"fits"``. Without it a JSON file wins, then pickle, csv and + fits, and a warning names the files ``name`` could mean (e.g. + ``tracer.json`` and ``tracer.fits``), so pass ``kind`` to be explicit. Returns ------- The loaded object """ - for item in self.jsons: - if item.name == name: - return item.value_using_reference(self._reference) - for item in self.pickles + self.arrays + self.fits: - if item.name == name: - return item.value + if kind is not None: + suffix = f".{kind.lstrip('.')}" + if suffix not in _KINDS: + raise ValueError( + f"kind must be one of {sorted(k[1:] for k in _KINDS)}, not {kind!r}." + ) + for item in self._outputs_by_suffix[suffix]: + if item.name == name: + if suffix == ".json": + return item.value_using_reference(self._reference) + return item.value + return None - return None + matches = [ + item + for suffix in _KINDS + for item in self._outputs_by_suffix[suffix] + if item.name == name + ] + if not matches: + return None + if len(matches) > 1 and name not in _warned_ambiguous: + _warned_ambiguous.add(name) + logging.warning( + f"'{name}' names several files ({', '.join(m.path.name for m in matches)}); " + f"the {matches[0].path.suffix} file is loaded. Pass kind= to choose." + ) + item = matches[0] + if isinstance(item, JSONOutput): + return item.value_using_reference(self._reference) + return item.value class SearchOutput(AbstractSearchOutput, fit_interface.Fit): @@ -308,6 +366,14 @@ def __init__( self.directory = directory + @classmethod + def from_zip(cls, zip_path, reference: Optional[dict] = None) -> "SearchOutput": + """ + The output archived in the zip at ``zip_path``, read in place: every file is + read straight out of the zip when it is needed and nothing is extracted. + """ + return cls(ZipMember(zip_path), reference=reference) + @property def samples_summary(self) -> SamplesSummary: """ @@ -421,7 +487,7 @@ def _load_samples(self, name=None, model=None): try: info_json = JSONOutput("info", directory / "samples_info.json").dict - with open(directory / "samples.csv") as f: + with (directory / "samples.csv").open() as f: sample_list = samples_from_iterator(csv.reader(f)) if model is None: @@ -469,14 +535,14 @@ def children(self): """ A list of child analyses loaded from the analyses directory """ - return list(map(SearchOutput, Path(self.directory).glob("analyses/*"))) + return list(map(SearchOutput, self.directory.glob("analyses/*"))) @property def model_results(self) -> str: """ Reads the model.results file """ - with open(self.directory / "model.results") as f: + with (self.directory / "model.results").open() as f: return f.read() @property @@ -484,7 +550,7 @@ def mask(self): """ A pickled mask object """ - with open(self.files_path / "mask.pickle", "rb") as f: + with (self.files_path / "mask.pickle").open("rb") as f: return dill.load(f) @property @@ -503,11 +569,11 @@ def search(self): """ if self.__search is None: try: - with open(self.files_path / "search.json") as f: + with (self.files_path / "search.json").open() as f: self.__search = from_dict(json.load(f)) except (FileNotFoundError, ModuleNotFoundError): try: - with open(self.files_path / "search.pickle", "rb") as f: + with (self.files_path / "search.pickle").open("rb") as f: self.__search = pickle.load(f) except (FileNotFoundError, ModuleNotFoundError): logging.warning("Could not load search") @@ -548,7 +614,7 @@ def non_linear_search(self) -> Optional[str]: """ search_path = self.directory / "files" / "search.json" try: - with open(search_path) as f: + with search_path.open() as f: class_path = json.load(f)["class_path"] except (FileNotFoundError, KeyError, json.JSONDecodeError): pass diff --git a/autofit/aggregator/zip_member.py b/autofit/aggregator/zip_member.py new file mode 100644 index 000000000..fb597fcd5 --- /dev/null +++ b/autofit/aggregator/zip_member.py @@ -0,0 +1,161 @@ +""" +``ZipMember``: a path to a member (file or folder) inside a result zip, read in place +with nothing extracted to disk. + +It supports the subset of ``pathlib.Path`` that ``SearchOutput`` and ``FileOutput`` +use (``/``, ``exists``, ``is_file``, ``is_dir``, ``iterdir``, ``glob``, ``rglob``, +``open``, ``read_text``, ``read_bytes``, ``relative_to``, ``name``, ``suffix``, +``stem``), so a search output can be backed by a zip instead of a folder. + +Unlike ``zipfile.Path``, a ``ZipMember`` holds no open file: every ``open`` opens the +zip, opens the member and closes the zip again (``zipfile`` keeps the underlying file +open until the returned stream is closed). Aggregating thousands of zipped results +therefore never runs into the open-file limit. The member names of each zip are read +once and cached while the zip's size and modification time are unchanged. +""" +import fnmatch +import io +import os +import zipfile +from pathlib import Path, PurePosixPath +from typing import Dict, Iterator, Tuple + +_names_cache: Dict[str, Tuple[Tuple[int, int], Tuple[str, ...]]] = {} + + +def zip_names(zip_path) -> Tuple[str, ...]: + """ + The member names of the zip at ``zip_path`` (files only, no folder entries), + cached per path while its size and modification time are unchanged. + """ + zip_path = str(zip_path) + stat = os.stat(zip_path) + key = (stat.st_size, stat.st_mtime_ns) + cached = _names_cache.get(zip_path) + if cached is not None and cached[0] == key: + return cached[1] + with zipfile.ZipFile(zip_path) as f: + names = tuple(name for name in f.namelist() if not name.endswith("/")) + _names_cache[zip_path] = (key, names) + return names + + +class ZipMember: + def __init__(self, zip_path, at: str = ""): + """ + The member ``at`` (a POSIX path inside the zip, ``""`` for the zip's root) of + the zip at ``zip_path``. + """ + self.zip_path = Path(zip_path) + self.at = str(PurePosixPath(at)) if at not in ("", ".") else "" + + @property + def _names(self) -> Tuple[str, ...]: + return zip_names(self.zip_path) + + def _child(self, at: str) -> "ZipMember": + return ZipMember(self.zip_path, at) + + def __truediv__(self, other) -> "ZipMember": + other = str(other) + return self._child(f"{self.at}/{other}" if self.at else other) + + def __str__(self): + return f"{self.zip_path}/{self.at}" if self.at else str(self.zip_path) + + def __repr__(self): + return f"" + + def __eq__(self, other): + return ( + isinstance(other, ZipMember) + and self.zip_path == other.zip_path + and self.at == other.at + ) + + def __hash__(self): + return hash((self.zip_path, self.at)) + + @property + def name(self) -> str: + return PurePosixPath(self.at).name + + @property + def suffix(self) -> str: + return PurePosixPath(self.at).suffix + + @property + def stem(self) -> str: + return PurePosixPath(self.at).stem + + @property + def _prefix(self) -> str: + return f"{self.at}/" if self.at else "" + + def is_file(self) -> bool: + return bool(self.at) and self.at in self._names + + def is_dir(self) -> bool: + prefix = self._prefix + return any(name.startswith(prefix) for name in self._names) + + def exists(self) -> bool: + return self.is_file() or self.is_dir() + + def _descendants(self) -> Iterator[str]: + """ + Every file and folder below this member, as paths relative to it. + """ + prefix = self._prefix + seen = set() + for name in self._names: + if not name.startswith(prefix) or name == prefix: + continue + relative = PurePosixPath(name[len(prefix) :]) + parents = list(relative.parents)[:-1] + for path in reversed(parents): + if path not in seen: + seen.add(path) + yield str(path) + yield str(relative) + + def iterdir(self) -> Iterator["ZipMember"]: + for relative in self._descendants(): + if "/" not in relative: + yield self / relative + + def glob(self, pattern: str) -> Iterator["ZipMember"]: + depth = pattern.count("/") + for relative in self._descendants(): + if relative.count("/") == depth and fnmatch.fnmatchcase(relative, pattern): + yield self / relative + + def rglob(self, pattern: str) -> Iterator["ZipMember"]: + for relative in self._descendants(): + if fnmatch.fnmatchcase(PurePosixPath(relative).name, pattern): + yield self / relative + + def relative_to(self, other: "ZipMember") -> PurePosixPath: + return PurePosixPath(self.at).relative_to(PurePosixPath(other.at or ".")) + + def open(self, mode: str = "r", encoding: str = "utf-8"): + """ + Open the member for reading (``"r"`` text, ``"rb"`` bytes). Raises + ``FileNotFoundError`` if it is not a file in the zip. + """ + if mode not in ("r", "rb"): + raise ValueError(f"ZipMember is read-only; mode {mode!r} is not supported.") + if not self.is_file(): + raise FileNotFoundError(str(self)) + with zipfile.ZipFile(self.zip_path) as f: + stream = f.open(self.at) + if mode == "rb": + return stream + return io.TextIOWrapper(stream, encoding=encoding) + + def read_bytes(self) -> bytes: + with self.open("rb") as f: + return f.read() + + def read_text(self, encoding: str = "utf-8") -> str: + return self.read_bytes().decode(encoding) diff --git a/test_autofit/results/test_lazy_reads.py b/test_autofit/results/test_lazy_reads.py new file mode 100644 index 000000000..b41e4d7d1 --- /dev/null +++ b/test_autofit/results/test_lazy_reads.py @@ -0,0 +1,157 @@ +""" +Reading zipped results in place: ``ZipMember``, ``SearchOutput.from_zip``, +``open_member``, ``value(kind=...)`` and ``Aggregator.from_index`` with nothing +extracted to disk. +""" +import logging +import zipfile +from pathlib import Path + +import numpy as np +import pytest + +import autofit as af +from autofit.aggregator import search_output as search_output_module +from autofit.aggregator.aggregator import Aggregator +from autofit.aggregator.search_output import SearchOutput +from autofit.aggregator.zip_member import ZipMember +from autonerves import conf + +CONFIG = Path(__file__).parents[1] / "config" + +pytestmark = pytest.mark.filterwarnings("ignore::FutureWarning") + +FIXTURE = Path(__file__).parents[1] / "aggregator" / "search_output" + + +@pytest.fixture(name="output", autouse=True) +def make_output(tmp_path): + conf.instance.push(new_path=str(CONFIG), output_path=str(tmp_path / "output")) + return tmp_path / "output" + + +@pytest.fixture(name="zipped_fixture") +def make_zipped_fixture(tmp_path): + """ + The aggregator's directory fixture (json, pickle, csv and fits files, analyses) + archived as a zip. + """ + zip_path = tmp_path / "search_output.zip" + with zipfile.ZipFile(zip_path, "w") as f: + for path in FIXTURE.rglob("*"): + if path.is_file(): + f.write(path, path.relative_to(FIXTURE).as_posix()) + return zip_path + + +def _fit(name, unique_tag, total_draws=5): + search = af.Drawer(name=name, unique_tag=unique_tag, total_draws=total_draws) + search.fit( + model=af.Model(af.ex.Gaussian), + analysis=af.ex.Analysis( + data=np.ones(20), noise_map=np.ones(20) * 0.1, use_jax=False + ), + ) + return search + + +def test_zip_member_paths(zipped_fixture): + root = ZipMember(zipped_fixture) + files = root / "files" + + assert files.is_dir() and not files.is_file() + assert (files / "model.json").is_file() + assert not (files / "absent.json").exists() + assert (files / "model.json").suffix == ".json" + assert (files / "model.json").stem == "model" + assert (files / "model.json").relative_to(files).as_posix() == "model.json" + assert {p.name for p in root.iterdir()} >= {"files", "analyses", ".completed"} + assert sorted(p.name for p in root.glob("analyses/*")) == [ + "analysis_0", + "analysis_1", + ] + assert any(p.name == "psf.fits" for p in root.rglob("*.fits")) + with pytest.raises(FileNotFoundError): + (files / "absent.json").open() + with pytest.raises(ValueError): + (files / "model.json").open("w") + + +def test_zipped_output_reads_like_the_folder(zipped_fixture): + folder = SearchOutput(FIXTURE) + zipped = SearchOutput.from_zip(zipped_fixture) + + assert zipped.is_complete + assert zipped.id == folder.id + assert zipped.model.prior_count == folder.model.prior_count + assert zipped.search.__class__ is folder.search.__class__ + assert len(zipped.children) == len(folder.children) == 2 + assert sorted(o.name for o in zipped.jsons) == sorted(o.name for o in folder.jsons) + assert zipped.value("psf", kind="fits")[0].data.shape == ( + folder.value("psf", kind="fits")[0].data.shape + ) + assert zipped.value("info") == folder.value("info") + + +def test_open_member(zipped_fixture): + zipped = SearchOutput.from_zip(zipped_fixture) + with zipped.open_member("files/search.json", "r") as f: + assert "class_path" in f.read() + with SearchOutput(FIXTURE).open_member("files/search.json") as f: + assert b"class_path" in f.read() + + +def test_value_kind_resolves_a_name_shared_by_json_and_fits(tmp_path, caplog): + directory = tmp_path / "out" + (directory / "files").mkdir(parents=True) + (directory / "files" / "tracer.json").write_text('{"a": 1}') + from astropy.io import fits + + fits.HDUList([fits.PrimaryHDU(np.ones((2, 2)))]).writeto( + directory / "files" / "tracer.fits" + ) + output = SearchOutput(directory) + search_output_module._warned_ambiguous.discard("tracer") + + with caplog.at_level(logging.WARNING): + assert output.value("tracer") == {"a": 1} + assert "Pass kind= to choose" in caplog.text + + assert output.value("tracer", kind="fits")[0].data.shape == (2, 2) + assert output.value("tracer", kind="json") == {"a": 1} + assert output.value("tracer", kind="pickle") is None + with pytest.raises(ValueError): + output.value("tracer", kind="png") + + +def test_from_index_extracts_nothing(output, tmp_path, monkeypatch): + _fit("search", "dataset_0", total_draws=5) + _fit("search", "dataset_0", total_draws=6) + _fit("search", "dataset_1") + + def refuse(*args, **kwargs): + raise AssertionError("from_index must not extract") + + monkeypatch.setattr(zipfile.ZipFile, "extractall", refuse) + monkeypatch.setattr(zipfile.ZipFile, "extract", refuse) + before = sorted(output.rglob("*")) + + aggregator = Aggregator.from_index(output, index_dir=tmp_path / "index") + + # The older dataset_0 fit is superseded and left out by default. + assert len(aggregator) == 2 + assert all(isinstance(o.directory, ZipMember) for o in aggregator) + assert sorted(o.unique_tag for o in aggregator) == ["dataset_0", "dataset_1"] + assert all(o.samples_summary is not None for o in aggregator) + assert all(o.completed_at is not None for o in aggregator) + assert len(list(aggregator.values("samples_summary"))) == 2 + assert sorted(output.rglob("*")) == before + + everything = Aggregator.from_index( + output, + index_dir=tmp_path / "index", + states=("completed", "superseded"), + update=False, + ) + assert len(everything) == 3 + assert len(everything.latest_per()) == 2