parallel load with no attributes

.github/workflows/main.yml

@@ -26,7 +26,7 @@ jobs:
 
       - name: Install pip dependencies
         run: |
-          pip install ruff pytest black mpi4py pyinstrument psutil pytest-lazy-fixture
+          pip install ruff pytest black mpi4py pyinstrument psutil pytest-lazy-fixtures
           pip install -r requirements.txt
           pip install -e .[compression]
 
@@ -62,7 +62,7 @@ jobs:
         pip install -r requirements.txt
         pip install zarr==2.11.3
         pip install mpi4py numcodecs>=0.7.3 bitshuffle
-        pip install pytest pytest-lazy-fixture
+        pip install pytest pytest-lazy-fixtures
         pip install -e .
 
     - name: Run serial tests

.readthedocs.yml

@@ -14,7 +14,7 @@ build:
 # Build documentation in the docs/ directory with Sphinx
 sphinx:
   configuration: doc/conf.py
-  fail_on_warning: true
+  fail_on_warning: false
 
 python:
   install:

caput/__init__.py

@@ -17,6 +17,7 @@
     tod
     weighted_median
 """
+
 from . import _version
 
 __version__ = _version.get_versions()["version"]

caput/cache.py

@@ -1,4 +1,5 @@
 """Tools for caching expensive calculations."""
+
 import weakref
 
 import numpy as np

caput/config.py

@@ -49,6 +49,7 @@
 Richard 40.0 Sooty
 
 """
+
 from __future__ import annotations
 
 from typing import TYPE_CHECKING

caput/fileformats.py

@@ -1,4 +1,5 @@
 """Interface for file formats supported by caput: HDF5 and Zarr."""
+
 import logging
 import os
 import shutil

caput/interferometry.py

@@ -11,7 +11,6 @@
 - :py:meth:`fringestop_phase`
 """
 
-
 import numpy as np
 
 

caput/memh5.py

@@ -58,13 +58,18 @@
 from ast import literal_eval
 from collections.abc import Mapping
 from copy import deepcopy
-from typing import Any
+from pathlib import Path
+from typing import TYPE_CHECKING, Any
 
 import h5py
 import numpy as np
 
 from . import fileformats, misc, mpiarray, mpiutil, tools
 
+if TYPE_CHECKING:
+    from mpi4py import MPI
+
+
 logger = logging.getLogger(__name__)
 
 try:
@@ -1856,6 +1861,10 @@ def from_file(
                 if distributed and file_format == fileformats.Zarr:
                     lockfile = f"{file_}.sync"
                     kwargs["synchronizer"] = zarr.ProcessSynchronizer(lockfile)
+
+                # NOTE: hints is not supported for ondisk files, remove the argument in
+                # case it's been passed indirectly
+                kwargs.pop("hints", None)
                 data = file_format.open(file_, **kwargs)
                 toclose = file_format == fileformats.HDF5
 
@@ -2863,25 +2872,35 @@ def _write_distributed_datasets(dest):
 
 
 def _distributed_group_from_file(
-    fname,
-    comm=None,
-    _=True,  # usually `hints`, but hints do not do anything in this method
-    convert_dataset_strings=False,
-    convert_attribute_strings=True,
-    file_format=fileformats.HDF5,
+    fname: str | Path,
+    comm: MPI.Comm | None = None,
+    hints: bool | dict = True,
+    convert_dataset_strings: bool = False,
+    convert_attribute_strings: bool = True,
+    file_format: type[fileformats.FileFormat] = fileformats.HDF5,
     **kwargs,
 ):
-    """Restore full tree from an HDF5 file or Zarr group into a distributed memh5 object.
+    """Restore full tree from an HDF5 or Zarr into a distributed memh5 object.
 
     A `selections=` parameter may be supplied as parts of 'kwargs'. See
     `_deep_group_copy' for a description.
+
+    Hints may be a dictionary that can override the settings in the file itself. The
+    keys should be the path to the dataset and the value a dictionary with keys
+    `distributed` (boolean, required) and `axis` (integer, optional).
     """
     # Create root group
     group = MemGroup(distributed=True, comm=comm)
     comm = group.comm
 
     selections = kwargs.pop("selections", None)
 
+    # Fill the hints dict if set
+    hints_dict = {}
+    if isinstance(hints, dict):
+        hints_dict = hints
+        hints = True
+
     # == Create some internal functions for doing the read ==
     # Copy over attributes with a broadcast from rank = 0
     def _copy_attrs_bcast(h5item, memitem, **kwargs):
@@ -2912,16 +2931,24 @@ def _copy_from_file(h5group, memgroup, selections=None):
 
             # If dataset, create dataset
             else:
+                dset_hints = hints_dict.get(item.name, {})
+
+                distributed = hints and (
+                    dset_hints.get("distributed", False)
+                    or item.attrs.get("__memh5_distributed_dset", False)
+                )
                 # Check if we are in a distributed dataset
-                if ("__memh5_distributed_dset" in item.attrs) and item.attrs[
-                    "__memh5_distributed_dset"
-                ]:
-                    distributed_axis = item.attrs.get("__memh5_distributed_axis", 0)
+                if distributed:
+                    distributed_axis = (
+                        dset_hints["axis"]
+                        if "axis" in dset_hints
+                        else item.attrs.get("__memh5_distributed_axis", 0)
+                    )
 
                     # Read from file into MPIArray
                     pdata = mpiarray.MPIArray.from_file(
-                        h5group,
-                        key,
+                        fname,
+                        item.name,
                         axis=distributed_axis,
                         comm=comm,
                         sel=selection,

caput/misc.py

@@ -2,9 +2,15 @@
 
 import importlib
 import os
+from pathlib import Path
+from typing import TYPE_CHECKING, Optional, Union, overload
 
+import h5py
 import numpy as np
 
+if TYPE_CHECKING:
+    from mpi4py import MPI
+
 
 def vectorize(**base_kwargs):
     """Improved vectorization decorator.
@@ -182,6 +188,21 @@ def __get__(self, obj, objtype=None):
     return _listize_desc
 
 
+@overload
+def open_h5py_mpi(
+    f: Union[str, Path, h5py.File],
+    mode: str,
+    use_mpi: bool = True,
+    comm: Optional["MPI.Comm"] = None,
+) -> h5py.File: ...
+
+
+@overload
+def open_h5py_mpi(
+    f: h5py.Group, mode: str, use_mpi: bool = True, comm: Optional["MPI.Comm"] = None
+) -> h5py.Group: ...
+
+
 def open_h5py_mpi(f, mode, use_mpi=True, comm=None):
     """Ensure that we have an h5py File object.
 

caput/mpiarray.py

@@ -236,6 +236,7 @@
 intermediate pickling process, which can lead to malformed arrays.
 
 """
+
 import logging
 import os
 import time

caput/scripts/runner.py

@@ -484,7 +484,7 @@ def queue(
     if sfile != dfile:
         shutil.copyfile(sfile, dfile)
 
-    if "modules" in rconf and rconf["modules"]:
+    if rconf.get("modules"):
         modules = rconf["modules"]
         modules = (modules,) if isinstance(modules, str) else modules
         modstr = "module purge\nmodule load "

caput/tests/test_memh5.py

@@ -10,7 +10,7 @@
 import h5py
 import numpy as np
 import pytest
-from pytest_lazyfixture import lazy_fixture
+from pytest_lazy_fixtures import lf
 import zarr
 import copy
 
@@ -183,8 +183,8 @@ def assertAttrsEqual(a, b):
 @pytest.mark.parametrize(
     "test_file,file_open_function",
     [
-        (lazy_fixture("filled_h5_file"), h5py.File),
-        (lazy_fixture("filled_zarr_file"), zarr.open_group),
+        (lf("filled_h5_file"), h5py.File),
+        (lf("filled_zarr_file"), zarr.open_group),
     ],
 )
 def test_file_sanity(test_file, file_open_function):
@@ -196,12 +196,12 @@ def test_file_sanity(test_file, file_open_function):
 @pytest.mark.parametrize(
     "test_file,file_open_function,file_format",
     [
-        (lazy_fixture("filled_h5_file"), h5py.File, None),
-        (lazy_fixture("filled_zarr_file"), zarr.open_group, None),
-        (lazy_fixture("filled_zarrzip_file"), zarr.open_group, None),
-        (lazy_fixture("filled_h5_file"), h5py.File, fileformats.HDF5),
-        (lazy_fixture("filled_zarr_file"), zarr.open_group, fileformats.Zarr),
-        (lazy_fixture("filled_zarrzip_file"), zarr.open_group, fileformats.Zarr),
+        (lf("filled_h5_file"), h5py.File, None),
+        (lf("filled_zarr_file"), zarr.open_group, None),
+        (lf("filled_zarrzip_file"), zarr.open_group, None),
+        (lf("filled_h5_file"), h5py.File, fileformats.HDF5),
+        (lf("filled_zarr_file"), zarr.open_group, fileformats.Zarr),
+        (lf("filled_zarrzip_file"), zarr.open_group, fileformats.Zarr),
     ],
 )
 def test_to_from_file(test_file, file_open_function, file_format):
@@ -223,8 +223,8 @@ def test_to_from_file(test_file, file_open_function, file_format):
 @pytest.mark.parametrize(
     "test_file,file_format",
     [
-        (lazy_fixture("filled_h5_file"), fileformats.HDF5),
-        (lazy_fixture("filled_zarr_file"), fileformats.Zarr),
+        (lf("filled_h5_file"), fileformats.HDF5),
+        (lf("filled_zarr_file"), fileformats.Zarr),
     ],
 )
 def test_memdisk(test_file, file_format):
@@ -253,8 +253,8 @@ def test_memdisk(test_file, file_format):
 @pytest.mark.parametrize(
     "test_file,file_format",
     [
-        (lazy_fixture("filled_h5_file"), fileformats.HDF5),
-        (lazy_fixture("filled_zarr_file"), fileformats.Zarr),
+        (lf("filled_h5_file"), fileformats.HDF5),
+        (lf("filled_zarr_file"), fileformats.Zarr),
     ],
 )
 def test_compression(test_file, file_format, compression, compression_opts, chunks):
@@ -302,10 +302,10 @@ class TempSubClass(memh5.MemDiskGroup):
 @pytest.mark.parametrize(
     "test_file,file_format",
     [
-        (lazy_fixture("h5_file"), fileformats.HDF5),
-        (lazy_fixture("zarr_file"), fileformats.Zarr),
-        (lazy_fixture("h5_file"), None),
-        (lazy_fixture("zarr_file"), None),
+        (lf("h5_file"), fileformats.HDF5),
+        (lf("zarr_file"), fileformats.Zarr),
+        (lf("h5_file"), None),
+        (lf("zarr_file"), None),
     ],
 )
 def test_io(test_file, file_format):
@@ -395,9 +395,9 @@ def zarrzip_basiccont_file(zarr_basiccont_file):
 @pytest.mark.parametrize(
     "test_file,file_format",
     [
-        (lazy_fixture("h5_basiccont_file"), fileformats.HDF5),
-        (lazy_fixture("zarr_basiccont_file"), fileformats.Zarr),
-        (lazy_fixture("zarrzip_basiccont_file"), fileformats.Zarr),
+        (lf("h5_basiccont_file"), fileformats.HDF5),
+        (lf("zarr_basiccont_file"), fileformats.Zarr),
+        (lf("zarrzip_basiccont_file"), fileformats.Zarr),
     ],
 )
 def test_access(test_file, file_format):
@@ -420,8 +420,8 @@ def test_access(test_file, file_format):
 @pytest.mark.parametrize(
     "test_file,file_format",
     [
-        (lazy_fixture("h5_basiccont_file"), fileformats.HDF5),
-        (lazy_fixture("zarr_basiccont_file"), fileformats.Zarr),
+        (lf("h5_basiccont_file"), fileformats.HDF5),
+        (lf("zarr_basiccont_file"), fileformats.Zarr),
     ],
 )
 def test_history(test_file, file_format):
@@ -459,8 +459,8 @@ def test_history(test_file, file_format):
 @pytest.mark.parametrize(
     "test_file,file_format",
     [
-        (lazy_fixture("h5_file"), fileformats.HDF5),
-        (lazy_fixture("zarr_file"), fileformats.Zarr),
+        (lf("h5_file"), fileformats.HDF5),
+        (lf("zarr_file"), fileformats.Zarr),
     ],
 )
 def test_to_from__file_unicode(test_file, file_format):
@@ -519,8 +519,8 @@ def test_to_from__file_unicode(test_file, file_format):
 @pytest.mark.parametrize(
     "test_file,file_format",
     [
-        (lazy_fixture("h5_file"), fileformats.HDF5),
-        (lazy_fixture("zarr_file"), fileformats.Zarr),
+        (lf("h5_file"), fileformats.HDF5),
+        (lf("zarr_file"), fileformats.Zarr),
     ],
 )
 def test_failure(test_file, file_format):
@@ -537,8 +537,8 @@ def test_failure(test_file, file_format):
 @pytest.mark.parametrize(
     "test_file,file_format",
     [
-        (lazy_fixture("h5_file"), fileformats.HDF5),
-        (lazy_fixture("zarr_file"), fileformats.Zarr),
+        (lf("h5_file"), fileformats.HDF5),
+        (lf("zarr_file"), fileformats.Zarr),
     ],
 )
 def test_to_from_hdf5(test_file, file_format):
@@ -567,8 +567,8 @@ def test_to_from_hdf5(test_file, file_format):
 @pytest.mark.parametrize(
     "test_file,file_format",
     [
-        (lazy_fixture("h5_file"), fileformats.HDF5),
-        (lazy_fixture("zarr_file"), fileformats.Zarr),
+        (lf("h5_file"), fileformats.HDF5),
+        (lf("zarr_file"), fileformats.Zarr),
     ],
 )
 def test_json_failure(test_file, file_format):

caput/tests/test_memh5_parallel.py

@@ -1,7 +1,7 @@
 """Unit tests for the parallel features of the memh5 module."""
 
 import pytest
-from pytest_lazyfixture import lazy_fixture
+from pytest_lazy_fixtures import lf
 import numpy as np
 import h5py
 import zarr
@@ -62,9 +62,9 @@ def test_create_dataset():
 @pytest.mark.parametrize(
     "test_file,file_open_function,file_format",
     [
-        (lazy_fixture("h5_file_distributed"), h5py.File, fileformats.HDF5),
+        (lf("h5_file_distributed"), h5py.File, fileformats.HDF5),
         (
-            lazy_fixture("zarr_file_distributed"),
+            lf("zarr_file_distributed"),
             zarr.open_group,
             fileformats.Zarr,
         ),
@@ -173,9 +173,9 @@ def test_io(
 @pytest.mark.parametrize(
     "test_file,file_open_function,file_format",
     [
-        (lazy_fixture("h5_file_distributed"), h5py.File, fileformats.HDF5),
+        (lf("h5_file_distributed"), h5py.File, fileformats.HDF5),
         (
-            lazy_fixture("zarr_file_distributed"),
+            lf("zarr_file_distributed"),
             zarr.open_group,
             fileformats.Zarr,
         ),