Algorithms over pieces take MPI ranks that hold no pieces

.github/workflows/ci.yml

@@ -410,7 +410,7 @@ jobs:
   mpi:
     name: mpi
     runs-on: ubuntu-latest
-    timeout-minutes: 15
+    timeout-minutes: 30
     env:
       TACK_REQUIRE_MPI: "1"
       OMPI_MCA_rmaps_base_oversubscribe: "1"
@@ -450,6 +450,11 @@ jobs:
           uv run pytest -q -rs packages/tack-vis/tests/test_vtk_algorithms.py
           -k "mpi or world"
 
+      # Every algorithm over pieces with ranks that hold none: about 1.5 minutes a
+      # layout of pieces, four layouts.
+      - name: Run the algorithms over pieces with empty ranks
+        run: uv run pytest -q -rs packages/tack-vis/tests/test_empty_ranks.py
+
   # ── Packaging gate ─────────────────────────────────────────────────
   # tack-core's [cuda]/[metal]/[hip]/[level_zero] extras deliberately do not
   # include llvmlite — it is a CPU-only dependency. So a GPU install must

changelog.d/empty-ranks.md

@@ -0,0 +1,22 @@
+### Fixed
+
+- Over MPI, the algorithms over pieces take ranks that hold no pieces, as a reader
+  that gives some ranks nothing makes them. Each rank must make the same calls
+  that communicate, so a rank must not decide from its own pieces alone what the
+  others decide from theirs:
+  - `probe` of a `PartitionedDataSet` or an AMR hierarchy, `streamlines` and
+    `advect` over pieces, `voxelize` of pieces or a hierarchy, and `Advection`
+    (and the path filters) over pieces refused or failed on a rank without
+    pieces ("needs pieces on this rank to start from", "a rank with no pieces has
+    no piece to start its seeds from"), leaving the others waiting. Such a rank
+    now takes part through a piece of nothing (`_stand_in`): its points and seeds
+    go to the pieces that hold them, and the answers come back.
+  - `update_ghosts`, `extract_surface` of pieces, AMR `ghost_cells` and
+    `average_down`, `integrate` over pieces, and the VTK `Statistics`, `Entropy`
+    and `Integrate` algorithms chose their fields (or a shortcut) from this rank's
+    pieces; a rank without any now gets the others' choice.
+  - `merge_pieces` onto a root rank without pieces refused; the merge is made on
+    the first rank that has pieces and moved to the root.
+  - `test_empty_ranks.py` runs every algorithm over pieces, AMR hierarchies,
+    octrees, advection and the VTK algorithms over pieces on 3 ranks, with the
+    pieces on one, two or none of them, against the same pieces in one process.

packages/tack-vis/src/tack/dataset/algorithms/_common.py

@@ -356,3 +356,25 @@ def _csr_rows(offsets, connectivity, cells):
     at = (np.arange(starts[-1], dtype=np.int64) - np.repeat(starts[:-1], sizes)
           + np.repeat(offsets[cells], sizes))
     return connectivity[at], sizes
+
+
+def _stand_in(data, origin):
+    """For a rank that holds none of ``data``'s pieces (a ``PartitionedDataSet``), a
+    piece of nothing that lets it take part in an exchange over the pieces: no
+    point and no cell (so its block's box holds nothing), of the pieces' float
+    type, spatial unit and time, at ``origin``, on the default device; ``None`` on
+    a rank with pieces. Collective: every rank of ``data.world`` asks, in turn."""
+    import numpy as np
+
+    from tack.dataset.algorithms._extract import _vertices
+
+    pieces = list(data)
+    mine = None if not pieces else (pieces[0].dtype, pieces[0].unit, pieces[0].time,
+                                    pieces[0].time_unit)
+    given = next((m for m in data.world._allgather(mine) if m is not None), None)
+    if pieces:
+        return None
+    dtype, unit, time, time_unit = given if given is not None else (
+        tack.get_default_float(), None, None, None)
+    return _vertices(np.zeros((0, 3)), None, dtype=dtype, origin=origin, unit=unit,
+                     time=time, time_unit=time_unit)

packages/tack-vis/src/tack/dataset/algorithms/_extract.py

@@ -332,13 +332,16 @@ def _every(flags, stride):
         flags[i] = 1 if i % stride == 0 else 0
 
 
-def _vertices(points, like=None):
+def _vertices(points, like=None, dtype=None, **frame):
     """Vertex cells at ``points`` (an ``(n, 3)`` host array or an array of 3-vectors),
-    with ``like``'s id type (or wider if needed), origin, spatial unit and time."""
+    with ``like``'s id type (or wider if needed), origin, spatial unit and time;
+    without ``like``, of ``dtype`` (f32 by default) in ``frame`` (``origin``,
+    ``unit``, ...)."""
     from tack.dataset.topology import UnstructuredTopology
 
     n = int(points.shape[0])
     idt = ids.for_output(like, n) if like is not None else ids.choose(None, n)
+    if dtype is None:
+        dtype = like.dtype if like is not None else tack.f32
     return DataSet(UnstructuredTopology(*_vertex_cells(n, idt), num_points=n), points,
-                   dtype=like.dtype if like is not None else tack.f32,
-                   **(like._frame() if like is not None else {}))
+                   dtype=dtype, **{**(like._frame() if like is not None else {}), **frame})

packages/tack-vis/src/tack/dataset/algorithms/_particles.py

@@ -41,7 +41,7 @@
     ZERO_VELOCITY,
     _run,
 )
-from tack.dataset.algorithms._common import _time_unit
+from tack.dataset.algorithms._common import _stand_in, _time_unit
 from tack.dataset.algorithms._extract import _vertices
 from tack.dataset.dataset import DataSet, Field, is_field_key
 from tack.dataset.ghosts import Ghosts, visible
@@ -132,12 +132,14 @@ def advect_pieces(data, field, seeds, step_size, steps, integrator="rk4", unit="
 
     ``seeds`` are this rank's, an ``(n, 3)`` array; across MPI ranks each rank
     passes its own, numbered after the ranks before it."""
-    blocks, pieces, finished, _ = _rounds(data, field, seeds, step_size, steps, integrator,
-                                          record=False, unit=unit, neighbors=neighbors,
-                                          adaptive=adaptive)
+    blocks, pieces, finished, _, nothing, time_unit = _rounds(
+        data, field, seeds, step_size, steps, integrator, record=False, unit=unit,
+        neighbors=neighbors, adaptive=adaptive)
     out = []
     for g, piece in zip(blocks.local, pieces):
         done = finished[g]
+        if piece is nothing and not len(done):
+            continue                # a rank without pieces keeps its seeds no piece held
         order = np.argsort(done.seed, kind="stable")
         done = done.take(order)
         with blocks.device(g):
@@ -146,8 +148,7 @@ def advect_pieces(data, field, seeds, step_size, steps, integrator="rk4", unit="
             points = Values(vertices, "points")
             vertices.fields["steps"] = Field(points, tack.array(done.steps, tack.i32))
             vertices.fields["time"] = Field(points, tack.array(
-                done.time.astype(piece.dtype.numpy_dtype), piece.dtype),
-                _time_unit(piece, field))
+                done.time.astype(piece.dtype.numpy_dtype), piece.dtype), time_unit)
             vertices.fields["status"] = Field(points, tack.array(done.status, tack.i32))
             vertices.fields["seed"] = Field(points, tack.array(done.seed, tack.i64))
         out.append(vertices)
@@ -165,11 +166,13 @@ def streamlines_pieces(data, field, seeds, step_size, steps, integrator="rk4", u
     copy marks ``DUPLICATEPOINT``, so each path's points are owned once."""
     from tack.dataset.polyhedra import PolylineTopology
 
-    blocks, pieces, _, traced = _rounds(data, field, seeds, step_size, steps, integrator,
-                                        record=True, unit=unit, neighbors=neighbors,
-                                        adaptive=adaptive)
+    blocks, pieces, _, traced, nothing, time_unit = _rounds(
+        data, field, seeds, step_size, steps, integrator, record=True, unit=unit,
+        neighbors=neighbors, adaptive=adaptive)
     out = []
     for g, piece in zip(blocks.local, pieces):
+        if piece is nothing:
+            continue                # its seeds were traced on the pieces that held them
         stretches = sorted(traced[g], key=lambda s: (s[0], s[1]))
         positions, times, marks, offsets, members, seed, segment = [], [], [], [0], [], [], []
         count = 0
@@ -202,16 +205,18 @@ def streamlines_pieces(data, field, seeds, step_size, steps, integrator="rk4", u
                                                        tack.i32))
             lines.fields["time"] = Field(Values(lines, "points"), tack.array(
                 (np.concatenate(times) if times else np.zeros(0)).astype(dtype),
-                piece.dtype), _time_unit(piece, field))
+                piece.dtype), time_unit)
         out.append(lines)
     return PartitionedDataSet(out, world=blocks.world)
 
 
 def _rounds(data, field, seeds, step_size, steps, integrator, record, unit="time",
             neighbors=False, adaptive=None):
-    """The rounds: ``(blocks, pieces, finished, traced)``, the particles each piece
-    finished with, and with ``record`` each piece's stretches ``(seed, segment,
-    positions, times)``."""
+    """The rounds: ``(blocks, pieces, finished, traced, nothing, time_unit)``, the
+    particles each piece finished with, with ``record`` each piece's stretches
+    ``(seed, segment, positions, times)``, the piece of nothing standing in for a
+    rank without pieces (``_stand_in``, else ``None``) and the particles' unit of
+    time."""
     if int(steps) < 1:
         raise ValueError(f"steps must be at least 1, not {steps!r}")
     if not is_field_key(field):
@@ -221,13 +226,21 @@ def _rounds(data, field, seeds, step_size, steps, integrator, record, unit="time
     for piece in data:                       # each piece's work on its own device
         with piece.device:
             pieces.append(visible(piece, field)[0])
+    origin = common_origin(data, "advection over pieces")
+    # A rank with no pieces takes part through a piece of nothing: its block holds
+    # no point, and sends its seeds to the pieces that do.
+    nothing = _stand_in(data, origin)
+    if nothing is not None:
+        pieces = [nothing]
     blocks = tc.Blocks([piece.device for piece in pieces], data.world)
     locators = {}
     for g, piece in zip(blocks.local, pieces):
-        with piece.device:
-            locators[g] = cached_locator(piece)
+        if piece.num_points:
+            with piece.device:
+                locators[g] = cached_locator(piece)
     boxes = _boxes(blocks, pieces, locators)
-    pending = _seeded(blocks, seeds, common_origin(data, "advection over pieces"))
+    pending = _seeded(blocks, seeds, origin)
+    time_unit = _time_unit(data, field)
     finished = {g: [] for g in blocks.local}
     traced = {g: [] for g in blocks.local}
     for _ in range(int(steps) + 2):
@@ -250,7 +263,8 @@ def _rounds(data, field, seeds, step_size, steps, integrator, record, unit="time
             finished[g].append(done)
     else:
         raise RuntimeError("particles were still in flight after a round per step")
-    return blocks, pieces, {g: _Particles.joined(finished[g]) for g in blocks.local}, traced
+    return (blocks, pieces, {g: _Particles.joined(finished[g]) for g in blocks.local},
+            traced, nothing, time_unit)
 
 
 def _boxes(blocks, pieces, locators):
@@ -270,8 +284,6 @@ def _seeded(blocks, seeds, origin):
         host = np.asarray(seeds.numpy(), np.float64).reshape(-1, 3)
     else:
         host = np.asarray(seeds, np.float64).reshape(-1, 3) - np.asarray(origin)
-    if len(host) and not len(blocks.local):
-        raise ValueError("a rank with no pieces has no piece to start its seeds from")
     first = blocks.local.start if len(blocks.local) else None
     offsets, _ = blocks.exclusive_scan({g: len(host) if g == first else 0
                                         for g in blocks.local})

packages/tack-vis/src/tack/dataset/algorithms/amr_ghosts.py

@@ -876,9 +876,12 @@ def amr_ghost_cells(amr, layers, fields):
                                         tack.vec3(piece.dtype)) if need else None)
         arrays_by_kind.append(assembled)
         grown_boxes.append((i, level, k, piece, box, ext))
-    for name in names:
-        kind = _kind(pieces[0].fields[name])
-        fills[kind].run([arrays[name] for arrays in arrays_by_kind])
+    # Each field's kind, as a rank with boxes knows it: a rank without any runs the
+    # same fills, with nothing of its own to put in them.
+    mine = [_kind(pieces[0].fields[name]) for name in names] if pieces else None
+    kinds = next((k for k in amr.world._allgather(mine) if k is not None), [])
+    for i, kind in enumerate(kinds):
+        fills[kind].run([arrays[names[i]] for arrays in arrays_by_kind])
     if need:
         fills["points"].run(positions)
 
@@ -1069,7 +1072,8 @@ def average_down(amr, field):
     if not isinstance(amr, AMRDataSet):
         raise TypeError(f"average_down takes an AMRDataSet, not {type(amr).__name__}")
     boxes = [box for _, box in amr]
-    kinds = {_kind(box.fields[field]) for box in boxes}
+    # Every rank's boxes' say: a rank without boxes knows nothing of the field.
+    kinds = set().union(*amr.world._allgather({_kind(box.fields[field]) for box in boxes}))
     if len(kinds) != 1 or None in kinds:
         raise TypeError(f"average_down averages a point or cell field; {field!r} is not one")
     kind = kinds.pop()

packages/tack-vis/src/tack/dataset/algorithms/extract_surface.py

@@ -564,7 +564,9 @@ def _surface_of_pieces(pds, fields):
 
     world = pds.world
     pieces = [visible(piece) for piece in pds]
-    if all(p.topology.dimension == 2 for p in pieces):
+    # Every rank's pieces' say: a rank without pieces must not take the shortcut
+    # while the others match faces.
+    if all(world._allgather(all(p.topology.dimension == 2 for p in pieces))):
         # A surface's pieces are their own surfaces, as one surface is.
         return PartitionedDataSet([extract_surface(p, fields=fields) for p in pieces],
                                   world=world)

packages/tack-vis/src/tack/dataset/algorithms/ghost_cells.py

@@ -212,6 +212,8 @@ def update_ghosts(data, *fields):
                          "ghost_points made, which keeps the exchanges")
     pieces = list(data)
     blocks = plans["points"].blocks
+    if not blocks.count:
+        return                              # no pieces anywhere: nothing to refresh
     by_kind = {"points": {g: {} for g in blocks.local}, "cells": {g: {} for g in blocks.local}}
     for name in fields:
         for g, piece in zip(blocks.local, pieces):
@@ -227,7 +229,9 @@ def update_ghosts(data, *fields):
                                 f"a {type(field.values).__name__}, not a Tack array")
             by_kind[kind][g][field_label(name)] = field.values
     for kind, values in by_kind.items():
-        if any(values.values()):
+        # Every rank's pieces' say, a rank without pieces too: it takes part in the
+        # exchange the others run.
+        if blocks.all_reduce({g: int(bool(values[g])) for g in blocks.local}, "max"):
             plans[kind].run(values)
 
 

packages/tack-vis/src/tack/dataset/algorithms/integrate.py

@@ -379,23 +379,37 @@ def integrate(data, fields=None):
         from tack.dataset.algorithms.statistics import _over_pieces
 
         pieces = list(data)
-        names = _names(pieces[0], fields) if pieces else []
-        if any(isinstance(field, Field) for _, field in names):
+        # Refused alike on every rank, a rank without pieces too: from the arguments.
+        given = (list(fields.values()) if hasattr(fields, "items") else
+                 [] if fields is None else
+                 [fields] if is_field_key(fields) or isinstance(fields, Field) else list(fields))
+        if any(isinstance(field, Field) for field in given):
             raise TypeError("integrate over a PartitionedDataSet takes fields by name or "
                             "as objects whose on(piece) gives the piece's: a Field belongs "
                             "to one partition's topology")
+        names = _names(pieces[0], fields) if pieces else None
+        # The names and the measure's unit as a rank with pieces finds them: a rank
+        # without any gets the same answer (the sums, and the fields' units, are every
+        # rank's).
+        mine = None if names is None else ([name for name, _ in names],
+                                           _measure_unit(pieces[0]))
+        labels, measure_unit = next((m for m in data.world._allgather(mine) if m is not None),
+                                    ([], None))
+        if names is None:
+            names = [(name, name) for name in labels]
         parts = _over_pieces(data, lambda piece: _piece(piece, names))
         measure = math.fsum(m for m, _, _ in parts)
-        values = {name: np.sum([v[name] for _, v, _ in parts], axis=0) for name, _ in names}
+        values = {name: np.sum([v[name] for _, v, _ in parts], axis=0) for name in labels}
         field_units = parts[0][2] if parts else {}
-        first = pieces[0] if pieces else None
+        out_units = {name: units.derived((field_units[name], 1), (measure_unit, 1))
+                     for name in labels}
     else:
         names = _names(data, fields)
         measure, values, field_units = _piece(data, names)
-        first = data
-    out_units = {name: units.derived((field_units[name], 1), (_measure_unit(first), 1))
-                 if first is not None else None for name, _ in names}
-    return Integrals(measure=measure, measure_unit=_measure_unit(first) if first else None,
+        measure_unit = _measure_unit(data)
+        out_units = {name: units.derived((field_units[name], 1), (measure_unit, 1))
+                     for name, _ in names}
+    return Integrals(measure=measure, measure_unit=measure_unit,
                      values={name: (float(v) if np.ndim(v) == 0 else tuple(map(float, v)))
                              for name, v in values.items()},
                      units=out_units)