feat(mesh/io): zarr save/load for Mesh and DomainMesh via tensordict's zarr backend (NVIDIA#1894)

* feat(mesh/io): zarr save/load for Mesh and DomainMesh via tensordict backend

Delegates all serialization to tensordict's upstream zarr storage backend
(to_zarr/from_zarr, pytorch/tensordict#1754) and adds only what tensordict
cannot know: chunk/compression policy aligned to training subsample sizes,
and a root attr recording the mesh type so from_zarr can rebuild
Mesh/DomainMesh from the plain PersistentTensorDict the backend returns.

MeshReader/DomainMeshReader route zarr groups (zarr.json present) to
mesh.io.from_zarr; everything else in the dataloader path is unchanged.

Includes a scoped workaround for a tensordict bug where to_zarr
chunks/compressors kwargs are silently dropped for nested leaves.

* Address review: nested trees, recursive glob, zarr subsample push-down, tests

- from_zarr: recursive _read_tree returning TensorDict, so nested
  TensorDict fields round-trip (review); shared by eager and windowed reads.
- Readers: restore recursive '**' discovery (glob.glob recursive=True).
- Subsample push-down for zarr stores: with subsample_n_cells/points set,
  readers fetch only the selected window (contiguous row-runs; cell meshes
  gather referenced points and compose measure weights) instead of
  materializing the full mesh. Bitwise-identical to eager load + in-memory
  subsample under the same seed.
- test/mesh/io/io_zarr: round-trip, nested fields, point clouds, subgroup
  loads, layout policy on nested leaves, reader routing, push-down
  equivalence, mixed-format discovery, plus a committed golden store
  guarding on-disk format stability.
- CHANGELOG entries for the feature and the glob fix.

* Cache zarr store handles in readers; consolidate metadata on write

Cold-metadata cost fix: readers re-opened each case's store per draw,
walking the group-metadata chain -- on Lustre every uncached lookup is a
metadata-server round-trip, so the first epochs on a metadata-cold client
paid it per sample per epoch. Handles are now cached per reader (opened
once per case), and to_zarr consolidates metadata so the one remaining
open costs a single read instead of one per group.

* Address review: drop upstream workaround, top-level tensordict import, fixture-generated format contract

- Remove _propagate_nested_create_kwargs: fixed upstream in
  pytorch/tensordict#1759 (merge is gated on the tensordict pin, which
  will include it).
- tensordict imported at module top (core dependency); drop the runtime
  backend checks for the same reason.
- Replace the committed golden store with a session-fixture store plus
  explicit on-disk contract assertions (group paths, attrs, dtypes,
  chunking, codecs) -- format drift still fails a test, without binary
  files in the repo.

Author: negin513

CHANGELOG.md

@@ -10,6 +10,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Added
 
+- Adds zarr save/load for `Mesh` and `DomainMesh` via tensordict's zarr
+  storage backend: `physicsnemo.mesh.io.to_zarr` / `from_zarr`, with
+  training-appropriate chunking and zstd compression. `MeshReader` and
+  `DomainMeshReader` transparently read zarr stores alongside
+  `.pmsh`/`.pdmsh` (opt in via `pattern`). Requires optional `zarr >= 3`
+  and a tensordict release with the zarr backend.
 - Promotes GeoTransolver out of `experimental` to
   `physicsnemo.models.geotransolver.GeoTransolver`, together with the FLARE
   model (`physicsnemo.models.flare.FLARE`) and the reusable GALE and FLARE
@@ -373,6 +379,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Fixed
 
+- `MeshReader` / `DomainMeshReader` sample discovery no longer uses
+  `pathlib.Path.glob`, which can silently drop entries under filesystem
+  metadata-server load (Lustre), causing training to proceed on a subset
+  of the dataset.
 - `compute_cotan_weights_fem`, and the calculus, curvature, and smoothing
   routines built on it such as `Mesh.laplacian`, no longer fail on degenerate
   cells in float32. The Gram-matrix regularization is now scale-free, so it also

physicsnemo/datapipes/readers/mesh.py

@@ -24,6 +24,7 @@
 
 from __future__ import annotations
 
+import glob as _glob
 import logging
 from pathlib import Path
 from typing import Any, Iterator
@@ -128,6 +129,73 @@ def _subsample_mesh_cells(
     return mesh
 
 
+def _indices_to_runs(indices: torch.Tensor) -> list[tuple[int, int]]:
+    """Convert cyclic-block indices (1-2 ascending contiguous runs) to runs."""
+    breaks = torch.nonzero(indices[1:] != indices[:-1] + 1).flatten()
+    starts = [0] + [int(b) + 1 for b in breaks]
+    ends = [int(b) + 1 for b in breaks] + [len(indices)]
+    return [(int(indices[s]), int(indices[e - 1]) + 1) for s, e in zip(starts, ends)]
+
+
+def _zarr_mesh_subsampled(
+    group,
+    n_cells: int | None,
+    n_points: int | None,
+    generator: torch.Generator | None,
+) -> Mesh:
+    """Partial-read a zarr mesh group: fetch only the subsample window.
+
+    Reproduces :func:`_subsample_mesh` semantics (cyclic contiguous blocks,
+    vertex compaction, Horvitz-Thompson measure weights) while reading only
+    the selected rows from the store instead of materializing the full mesh.
+    """
+    from physicsnemo.mesh.io import io_zarr as _ioz
+
+    total_cells = group["cells"].shape[0] if "cells" in group else 0
+    total_points = group["points"].shape[0]
+
+    if total_cells > 0 and n_cells is not None and total_cells > n_cells:
+        indices = _cyclic_block_indices(total_cells, n_cells, generator=generator)
+        runs = _indices_to_runs(indices)
+        cells = _ioz._read_rows(group["cells"], runs)
+        # Compact: gather only referenced vertices; remap connectivity to the
+        # sorted-unique order, matching slice_cells + slice_points.
+        referenced, inverse = torch.unique(cells, return_inverse=True)
+        cells = inverse.reshape(cells.shape)
+        ref_np = referenced.numpy()
+        mesh = Mesh(
+            points=_ioz._read_index(group["points"], ref_np),
+            cells=cells,
+            point_data=_ioz._read_tree(
+                group, "point_data", leaf_reader=lambda a: _ioz._read_index(a, ref_np)
+            ),
+            cell_data=_ioz._read_tree(
+                group, "cell_data", leaf_reader=lambda a: _ioz._read_rows(a, runs)
+            ),
+            global_data=_ioz._read_tree(group, "global_data"),
+        )
+        compose_measure_weights(mesh, total_cells / n_cells)
+        if n_points is not None:
+            mesh = _subsample_mesh_points(mesh, n_points, generator=generator)
+        return mesh
+
+    if total_cells == 0 and n_points is not None and total_points > n_points:
+        indices = _cyclic_block_indices(total_points, n_points, generator=generator)
+        runs = _indices_to_runs(indices)
+        return Mesh(
+            points=_ioz._read_rows(group["points"], runs),
+            point_data=_ioz._read_tree(
+                group, "point_data", leaf_reader=lambda a: _ioz._read_rows(a, runs)
+            ),
+            cell_data=_ioz._read_tree(group, "cell_data"),
+            global_data=_ioz._read_tree(group, "global_data"),
+        )
+
+    # No subsampling applies (small mesh, or unsupported combination):
+    # eager full read keeps semantics identical to the memmap path.
+    return _ioz._mesh_from_group(group, None)
+
+
 def _subsample_mesh(
     mesh: Mesh,
     n_cells: int | None = None,
@@ -215,13 +283,49 @@ def __init__(
         if not self._root.is_dir():
             raise ValueError(f"Path must be a directory: {self._root}")
 
-        self._paths = sorted(self._root.glob(pattern))
+        # glob.glob instead of Path.glob: the latter re-stats each entry and
+        # silently drops entries under Lustre metadata-server load.
+        self._paths = sorted(
+            Path(p) for p in _glob.glob(str(self._root / pattern), recursive=True)
+        )
         if not self._paths:
             raise ValueError(f"No paths matching {pattern!r} found in {self._root}")
 
     def _load_sample(self, index: int) -> Mesh:
         """Load a single Mesh from disk."""
         mesh_path = self._paths[index]
+        if (mesh_path / "zarr.json").exists():
+            from physicsnemo.mesh.io import from_zarr, io_zarr
+
+            if (
+                self.subsample_n_cells is not None
+                or self.subsample_n_points is not None
+            ):
+                # Push the subsample into the read: fetch only the selected
+                # window from the store. The generator derivation matches
+                # __getitem__, so the draw is identical to subsampling after
+                # an eager load (whose subsample then no-ops).
+                # Cache opened store handles: re-opening walks the store's
+                # group-metadata chain, and on networked filesystems every
+                # uncached lookup is a metadata-server round-trip per draw.
+                cache = getattr(self, "_zarr_groups", None)
+                if cache is None:
+                    cache = self._zarr_groups = {}
+                group = cache.get(mesh_path)
+                if group is None:
+                    group = cache[mesh_path] = io_zarr._open_group(mesh_path)
+                generator = (
+                    None
+                    if self._seed_base is None
+                    else spawn_generator(self._seed_base, self._epoch, index)
+                )
+                return _zarr_mesh_subsampled(
+                    group,
+                    self.subsample_n_cells,
+                    self.subsample_n_points,
+                    generator,
+                )
+            return from_zarr(mesh_path)
         return Mesh.load(mesh_path)
 
     def _get_sample_metadata(self, index: int) -> dict[str, Any]:
@@ -406,13 +510,63 @@ def __init__(
         if not self._root.is_dir():
             raise ValueError(f"Path must be a directory: {self._root}")
 
-        self._paths = sorted(self._root.glob(pattern))
+        # glob.glob instead of Path.glob: the latter re-stats each entry and
+        # silently drops entries under Lustre metadata-server load.
+        self._paths = sorted(
+            Path(p) for p in _glob.glob(str(self._root / pattern), recursive=True)
+        )
         if not self._paths:
             raise ValueError(f"No paths matching {pattern!r} found in {self._root}")
 
     def _load_sample(self, index: int) -> DomainMesh:
         """Load a single DomainMesh from disk."""
-        return DomainMesh.load(self._paths[index])
+        path = self._paths[index]
+        if (path / "zarr.json").exists():
+            from physicsnemo.mesh.io import from_zarr, io_zarr
+
+            if (
+                self.subsample_n_cells is not None
+                or self.subsample_n_points is not None
+            ):
+                # Push the subsample into the read (window reads per
+                # sub-mesh); drop flags are honored at read time so skipped
+                # data is never fetched. Generator derivation and sub-mesh
+                # order match __getitem__, whose subsample then no-ops.
+                generator = (
+                    None
+                    if self._seed_base is None
+                    else spawn_generator(self._seed_base, self._epoch, index)
+                )
+                cache = getattr(self, "_zarr_groups", None)
+                if cache is None:
+                    cache = self._zarr_groups = {}
+                root = cache.get(path)
+                if root is None:
+                    root = cache[path] = io_zarr._open_group(path)
+                interior = _zarr_mesh_subsampled(
+                    root["interior"],
+                    self.subsample_n_cells,
+                    self.subsample_n_points,
+                    generator,
+                )
+                boundaries = {}
+                if not self.drop_in_file_boundaries and "boundaries" in root:
+                    boundaries = {
+                        name: _zarr_mesh_subsampled(
+                            grp,
+                            self.subsample_n_cells,
+                            self.subsample_n_points,
+                            generator,
+                        )
+                        for name, grp in root["boundaries"].groups()
+                    }
+                return DomainMesh(
+                    interior=interior,
+                    boundaries=boundaries,
+                    global_data=io_zarr._read_tree(root, "global_data"),
+                )
+            return from_zarr(path)
+        return DomainMesh.load(path)
 
     def __len__(self) -> int:
         return len(self._paths)
@@ -533,7 +687,12 @@ def _load_extra_boundaries(self, dm: DomainMesh, index: int) -> DomainMesh:
                     glob_pattern,
                     matches[0],
                 )
-            new_boundaries[bnd_name] = Mesh.load(matches[0])
+            if (matches[0] / "zarr.json").exists():
+                from physicsnemo.mesh.io import from_zarr
+
+                new_boundaries[bnd_name] = from_zarr(matches[0])
+            else:
+                new_boundaries[bnd_name] = Mesh.load(matches[0])
 
         return DomainMesh(
             interior=dm.interior,

physicsnemo/mesh/io/__init__.py

@@ -27,15 +27,20 @@
 
 from typing import TYPE_CHECKING
 
-__all__ = ["from_pyvista", "to_pyvista"]
+__all__ = ["from_pyvista", "to_pyvista", "from_zarr", "to_zarr"]
 
 if TYPE_CHECKING:
     from physicsnemo.mesh.io.io_pyvista import from_pyvista, to_pyvista
+    from physicsnemo.mesh.io.io_zarr import from_zarr, to_zarr
 
 
 def __getattr__(name: str):  # PEP 562
     if name in {"from_pyvista", "to_pyvista"}:
         from physicsnemo.mesh.io import io_pyvista
 
         return getattr(io_pyvista, name)
+    if name in {"from_zarr", "to_zarr"}:
+        from physicsnemo.mesh.io import io_zarr
+
+        return getattr(io_zarr, name)
     raise AttributeError(f"module {__name__!r} has no attribute {name!r}")