diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 652fdda4..d550baa5 100644 --- a/.github/workflows/ci.yml +++ b/.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 diff --git a/changelog.d/empty-ranks.md b/changelog.d/empty-ranks.md new file mode 100644 index 00000000..76b6af87 --- /dev/null +++ b/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. diff --git a/packages/tack-vis/src/tack/dataset/algorithms/_common.py b/packages/tack-vis/src/tack/dataset/algorithms/_common.py index 52b29ac1..5c46cabc 100644 --- a/packages/tack-vis/src/tack/dataset/algorithms/_common.py +++ b/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) diff --git a/packages/tack-vis/src/tack/dataset/algorithms/_extract.py b/packages/tack-vis/src/tack/dataset/algorithms/_extract.py index 4f85cdc0..bfbafda6 100644 --- a/packages/tack-vis/src/tack/dataset/algorithms/_extract.py +++ b/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}) diff --git a/packages/tack-vis/src/tack/dataset/algorithms/_particles.py b/packages/tack-vis/src/tack/dataset/algorithms/_particles.py index 5cd53b49..6a4d7e60 100644 --- a/packages/tack-vis/src/tack/dataset/algorithms/_particles.py +++ b/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}) diff --git a/packages/tack-vis/src/tack/dataset/algorithms/amr_ghosts.py b/packages/tack-vis/src/tack/dataset/algorithms/amr_ghosts.py index 18cb01e6..25df0162 100644 --- a/packages/tack-vis/src/tack/dataset/algorithms/amr_ghosts.py +++ b/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() diff --git a/packages/tack-vis/src/tack/dataset/algorithms/extract_surface.py b/packages/tack-vis/src/tack/dataset/algorithms/extract_surface.py index 92d3545b..d5997af2 100644 --- a/packages/tack-vis/src/tack/dataset/algorithms/extract_surface.py +++ b/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) diff --git a/packages/tack-vis/src/tack/dataset/algorithms/ghost_cells.py b/packages/tack-vis/src/tack/dataset/algorithms/ghost_cells.py index 5a53bddb..50b28b64 100644 --- a/packages/tack-vis/src/tack/dataset/algorithms/ghost_cells.py +++ b/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) diff --git a/packages/tack-vis/src/tack/dataset/algorithms/integrate.py b/packages/tack-vis/src/tack/dataset/algorithms/integrate.py index 0c7d014d..a6d502bb 100644 --- a/packages/tack-vis/src/tack/dataset/algorithms/integrate.py +++ b/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) diff --git a/packages/tack-vis/src/tack/dataset/algorithms/merge_pieces.py b/packages/tack-vis/src/tack/dataset/algorithms/merge_pieces.py index 268655e7..d39d1428 100644 --- a/packages/tack-vis/src/tack/dataset/algorithms/merge_pieces.py +++ b/packages/tack-vis/src/tack/dataset/algorithms/merge_pieces.py @@ -39,7 +39,8 @@ def merge_pieces(data, root=0, fields=None, device=None, tolerance=0.0): Across MPI ranks every rank passes its pieces and the result is on the ``root`` rank (``None`` elsewhere), on ``device`` (by default the device - of the root rank's first piece).""" + of the root rank's first piece). A root with no pieces gets the merge made + on the first rank that has some, moved to it (on its default device).""" if not isinstance(data, PartitionedDataSet): raise TypeError(f"merge_pieces takes a PartitionedDataSet, not {type(data).__name__}") from tack.dataset.algorithms.ghost_cells import _carried, _kind @@ -54,21 +55,112 @@ def merge_pieces(data, root=0, fields=None, device=None, tolerance=0.0): raise ValueError(f"root is a rank, 0 to {world.size - 1}, not {root}") firsts = blocks.all_gather({g: g == blocks.local.start for g in blocks.local}) starts = [g for g, first in enumerate(firsts) if first] - target = next((g for g in starts if blocks.rank_of(g) == root), None) - if target is None: - raise ValueError(f"merge_pieces gathers onto a piece of rank {root}, which has none") - on_root = target in blocks.local + if not starts: + raise ValueError("merge_pieces has no pieces on any rank to merge") + # Onto the root's first piece; a root with no pieces gets the merge made on + # the first rank that has some (``_to_root``). + target = next((g for g in starts if blocks.rank_of(g) == root), starts[0]) + on_target = target in blocks.local out_device = (tack.device(device) if device is not None else - blocks.device(target)) if on_root else None + blocks.device(target)) if on_target else None structured = blocks.all_gather({g: isinstance(p.topology, StructuredTopology) for g, p in zip(blocks.local, pieces)}) - if all(structured): - tiled = _box(blocks, pieces) - if tiled is not None: - return _merged_structured(blocks, pieces, names, _kind, target, out_device, - tiled) - return _merged_unstructured(blocks, pieces, names, _kind, target, out_device, - tolerance) + tiled = _box(blocks, pieces) if all(structured) else None + if tiled is not None: + merged = _merged_structured(blocks, pieces, names, _kind, target, out_device, tiled) + else: + merged = _merged_unstructured(blocks, pieces, names, _kind, target, out_device, + tolerance) + source = blocks.rank_of(target) + if source == root: + return merged + return _to_root(world, merged, source, root, device) + + +def _to_root(world, merged, source, root, device): + """The merged dataset, made on rank ``source``, moved to rank ``root`` (on + ``device``, by default the default device there): its description in a small + gather, its arrays in one exchange. ``None`` on every other rank.""" + carrier = tc.Blocks([tack.device(device) if device is not None + else tack.get_default_device()], world) # one block a rank + description, payload = _described(merged) if world.rank == source else (None, {}) + description = world._allgather(description)[source] + outgoing = {world.rank: ({root: payload} if world.rank == source else {})} + incoming = carrier.exchange(outgoing) + if world.rank != root: + return None + with carrier.device(root): + return _rebuilt(description, incoming[root].get(source, {})) + + +def _described(data): + """A merged dataset as ``(description, arrays)``: everything but the arrays in + small Python values, the arrays (explicit, by name) to send.""" + payload = {} + + def put(name, values): + payload[name] = arrays.materialize(values) + return name + + topology = data.topology + if isinstance(topology, StructuredTopology): + shape = ("structured", topology.point_dims, topology.offset, topology.id_dtype) + else: + shape = ("unstructured", topology.num_points, topology.id_dtype, + put("\0types", topology.types), put("\0offsets", topology.offsets), + put("\0connectivity", topology.connectivity)) + coordinates = data.geometry.values + if isinstance(coordinates, UniformCoordinates): + geometry = ("uniform", coordinates.dims, tuple(coordinates._origin), + tuple(coordinates._spacing), np.asarray(coordinates._matrix).tolist()) + elif isinstance(coordinates, CartesianProduct): + geometry = ("rectilinear", put("\0x", coordinates.xs), put("\0y", coordinates.ys), + put("\0z", coordinates.zs)) + else: + geometry = ("points", put("\0positions", coordinates)) + fields = [] + for name, field in data.fields.items(): + space = type(field.space).__name__ + on = getattr(field.space, "on", None) + fields.append((name, space, on, field.unit, put(f"field {len(fields)}", field.values))) + ghosts = data.ghosts + marks = (None if ghosts is None or ghosts.cells is None else put("\0cell marks", ghosts.cells), + None if ghosts is None or ghosts.points is None + else put("\0point marks", ghosts.points)) + ids = data.global_ids + numbers = (None if ids is None or ids.cells is None else put("\0cell ids", ids.cells), + None if ids is None or ids.points is None else put("\0point ids", ids.points)) + return (shape, geometry, data.dtype, data._frame(), fields, marks, numbers), payload + + +def _rebuilt(description, received): + """``_described``'s dataset again, from its description and the arrays received.""" + shape, geometry, dtype, frame, fields, marks, numbers = description + if shape[0] == "structured": + topology = StructuredTopology(shape[1], id_dtype=shape[3], offset=shape[2]) + else: + topology = UnstructuredTopology(received[shape[3]], received[shape[4]], + received[shape[5]], num_points=shape[1], + id_dtype=shape[2]) + if geometry[0] == "uniform": + values = UniformCoordinates(geometry[1], geometry[2], geometry[3], geometry[4], + dtype=dtype) + elif geometry[0] == "rectilinear": + values = CartesianProduct(*(received[name] for name in geometry[1:]), dtype=dtype) + else: + values = received[geometry[1]] + cell_marks, point_marks = (received[m] if m is not None else None for m in marks) + cell_ids, point_ids = (received[n] if n is not None else None for n in numbers) + out = DataSet(topology, values, dtype=dtype, + ghosts=Ghosts(cells=cell_marks, points=point_marks) + if cell_marks is not None or point_marks is not None else None, + global_ids=GlobalIds(cells=cell_ids, points=point_ids) + if cell_ids is not None or point_ids is not None else None, **frame) + spaces = {"H1": lambda on: H1(out), "Constant": lambda on: Constant(out), + "Values": lambda on: Values(out, on)} + for name, space, on, unit, key in fields: + out.fields[name] = Field(spaces[space](on), received[key], unit) + return out def _box(blocks, pieces): diff --git a/packages/tack-vis/src/tack/dataset/algorithms/probe.py b/packages/tack-vis/src/tack/dataset/algorithms/probe.py index c7e0c16e..0efe48ef 100644 --- a/packages/tack-vis/src/tack/dataset/algorithms/probe.py +++ b/packages/tack-vis/src/tack/dataset/algorithms/probe.py @@ -8,7 +8,7 @@ import tack from tack import arrays from tack.dataset import ids, units -from tack.dataset.algorithms._common import _field +from tack.dataset.algorithms._common import _field, _stand_in from tack.dataset.algorithms._extract import _vertices from tack.dataset.buckets import bucket_groups from tack.dataset.by_id import field_by_id @@ -281,14 +281,28 @@ def _probe_pieces(data, where, fields, levels=None): for piece in data: # each piece's work on its own device with piece.device: pieces.append(visible(piece)) + origin = common_origin(data, "probe of a PartitionedDataSet") + # What each piece gives back (its fields' names, element types and units, from + # a probe of no points) and the pieces' geometry, from a rank that has pieces: + # every rank needs them, one without pieces too, which probes through the + # others' and receives their answers. + mine = None + if pieces: + with pieces[0].device: + empty = probe(pieces[0], np.zeros((0, 3)), fields) + mine = [(n, f.values.dtype, f.unit) for n, f in empty.fields.items()] + given = next((m for m in data.world._allgather(mine) if m is not None), None) + # A rank with no pieces takes part through a piece of nothing: its block holds + # no point, and sends its targets to the pieces that do. + nothing = _stand_in(data, origin) + if nothing is not None: + pieces = [nothing] + if levels is not None: + levels = [0] blocks = tc.Blocks([piece.device for piece in pieces], data.world) - if not len(blocks.local): - raise ValueError("probe of a PartitionedDataSet needs pieces on this rank to " - "start from") # This rank's targets, their points end to end, from its first block, relative # to the pieces' one origin. - origin = common_origin(data, "probe of a PartitionedDataSet") spatial = units.common("probe of a PartitionedDataSet's pieces' geometry", [piece.unit for piece in pieces]) many = isinstance(where, PartitionedDataSet) @@ -336,12 +350,12 @@ def holds(g, x): routing = tc.Routing(blocks, positions, boxes, holds) arrived = routing.send() - # What each piece gives back, and its element types, from a probe of no - # points (so a rank whose points went nowhere still knows them). - with pieces[0].device: - empty = probe(pieces[0], np.zeros((0, 3)), fields) - names = [n for n in empty.fields if n not in ("valid",)] + ["valid"] - like = {n: empty.fields[n].values.dtype for n in names} + # What each piece gives back (so a rank whose points went nowhere still knows + # it); with no pieces anywhere, only ``valid``. + found = {n: (dtype, unit) for n, dtype, unit in (given or [])} + found.setdefault("valid", (tack.i32, None)) + names = [n for n in found if n != "valid"] + ["valid"] + like = {n: found[n][0] for n in names} answers = {} for g, piece in zip(blocks.local, pieces): x, _ = arrived[g] @@ -359,7 +373,7 @@ def holds(g, x): for name in names: values = gather(replied[name], rows).to_device(target.device) space = Values(target, "points") if name == "valid" else H1(target) - kept[name] = Field(space, values, empty.fields[name].unit) + kept[name] = Field(space, values, found[name][1]) start += n out.append(keep_marks(target, DataSet(target.topology, target.geometry, fields=kept, sets=dict(target.sets)))) diff --git a/packages/tack-vis/src/tack/dataset/algorithms/voxelize.py b/packages/tack-vis/src/tack/dataset/algorithms/voxelize.py index 0a39b030..0b4afbbb 100644 --- a/packages/tack-vis/src/tack/dataset/algorithms/voxelize.py +++ b/packages/tack-vis/src/tack/dataset/algorithms/voxelize.py @@ -63,17 +63,20 @@ def _means(values, valid, v0, per, out, n): def _like(data): """A dataset standing for ``data``'s frame and type: itself, a first piece or box, - an octree's roots.""" + an octree's roots; on a rank with no pieces or boxes, a piece of nothing in the + others' frame (``_stand_in``, so collective for pieces and hierarchies).""" + from tack.dataset.algorithms._common import _stand_in from tack.dataset.amr import AMRDataSet from tack.dataset.octree import OctreeDataSet - from tack.dataset.partitioned import PartitionedDataSet + from tack.dataset.partitioned import PartitionedDataSet, common_origin if isinstance(data, OctreeDataSet): return data.roots if isinstance(data, AMRDataSet): - return next(box for _, box in data) + data = data._as_pieces() if isinstance(data, PartitionedDataSet): - return next(iter(data)) + nothing = _stand_in(data, common_origin(data, "voxelize")) + return nothing if nothing is not None else next(iter(data)) return data diff --git a/packages/tack-vis/src/tack/dataset/temporal.py b/packages/tack-vis/src/tack/dataset/temporal.py index 66c2bfef..362e8f96 100644 --- a/packages/tack-vis/src/tack/dataset/temporal.py +++ b/packages/tack-vis/src/tack/dataset/temporal.py @@ -368,8 +368,15 @@ def _kept_dataset(data, field, copy): def _frame_of(data): - """A dataset's, or the first piece's, ``_frame()`` and device.""" - first = next(iter(data)) if isinstance(data, PartitionedDataSet) else data + """A dataset's, or the first piece's, ``_frame()`` and device; over pieces, on a rank + with none, the others' frame on the default device (``_stand_in``: collective).""" + if not isinstance(data, PartitionedDataSet): + return data._frame(), data.device + from tack.dataset.algorithms._common import _stand_in + from tack.dataset.partitioned import common_origin + + nothing = _stand_in(data, common_origin(data, "advection over pieces")) + first = nothing if nothing is not None else next(iter(data)) return first._frame(), first.device diff --git a/packages/tack-vis/src/tack/interop/vtk/algorithms.py b/packages/tack-vis/src/tack/interop/vtk/algorithms.py index 54f3dc60..6a458d50 100644 --- a/packages/tack-vis/src/tack/interop/vtk/algorithms.py +++ b/packages/tack-vis/src/tack/interop/vtk/algorithms.py @@ -2507,13 +2507,20 @@ def RequestData(self, request, inInfo, outInfo): def _selected(self, pieces): """``(label, where, key, component)`` for each value to reduce: the chosen fields (``fields``, or ``SetInputArrayToProcess``'s, or every field), a vector field's - components one by one (``component`` None for a scalar).""" + components one by one (``component`` None for a scalar). As a rank with pieces + finds them: a rank without any reduces the same values, with nothing of its + own (collective).""" + mine = self._selected_here(pieces) + return next((s for s in pieces.world._allgather(mine) if s is not None), []) + + def _selected_here(self, pieces): + """``_selected`` from this rank's first piece; ``None`` without one.""" import tack.dataset as td from tack.arrays import width_of first = next(iter(pieces), None) if first is None: - return [] + return None chosen = self.fields if chosen is None: chosen = [a for a in map(self._input_array, range(self._arrays_set())) @@ -2670,8 +2677,16 @@ def reduce(self, pieces): selected = {} for name, where, key, _ in self._selected(pieces): selected.setdefault((name, where), key) + # The first piece's origin over every rank, and the widest cells: alike on a + # rank without pieces, which takes part with nothing. first = next(iter(pieces), None) - origin = np.asarray(first.origin if first is not None else (0.0, 0.0, 0.0), float) + found = pieces.world._allgather(( + None if first is None else tuple(first.origin), + max((piece.topology.dimension for piece in pieces), default=0))) + origin = np.asarray(next((o for o, _ in found if o is not None), (0.0, 0.0, 0.0)), + float) + dimension = max(d for _, d in found) or 3 + names = [f"{where}:{name}" for name, where in selected] + ["__position__"] out = [] for piece in pieces: fields = {f"{where}:{name}": piece.fields[key] @@ -2689,7 +2704,7 @@ def reduce(self, pieces): ghosts=piece.ghosts, global_ids=piece.global_ids, **piece._frame())) integrals = td.integrate(td.PartitionedDataSet(out, world=pieces.world), - fields=list(fields) if out else None) + fields=names) measure = integrals.measure centre = origin + (np.asarray(integrals.values["__position__"], float) / measure if measure else 0.0) @@ -2706,7 +2721,6 @@ def reduce(self, pieces): array.SetName(name) (grid.GetPointData() if where == "points" else grid.GetCellData()).AddArray(array) _label(array, integrals.units.get(f"{where}:{name}")) - dimension = max((piece.topology.dimension for piece in pieces), default=3) total = vtkDoubleArray() total.InsertNextValue(measure) total.SetName({1: "Length", 2: "Area"}.get(dimension, "Volume")) diff --git a/packages/tack-vis/tests/_empty_ranks_mpi.py b/packages/tack-vis/tests/_empty_ranks_mpi.py new file mode 100644 index 00000000..320d9216 --- /dev/null +++ b/packages/tack-vis/tests/_empty_ranks_mpi.py @@ -0,0 +1,418 @@ +# SPDX-FileCopyrightText: Copyright (c) Kitware Inc. +# SPDX-License-Identifier: BSD-3-Clause + +"""``test_empty_ranks``' check under MPI: ``mpiexec -n 3 python _empty_ranks_mpi.py +LAYOUT [CASE ...]``. Every algorithm over pieces runs with ranks that hold none: +four pieces all on rank 0 (``one``), on ranks 1 and 2 (``first-empty``), on ranks 0 +and 1 (``last-empty``), or nowhere (``none``). Each case's answer, combined over the +ranks, must be what the same four pieces give in one process; with no pieces +anywhere it must be empty, or the same error on every rank. A case that hangs is +stopped by a deadline (``faulthandler``), which names where. Rank 0 prints ``ok`` +or the failures.""" + +import faulthandler +import os +import sys +import traceback + +import numpy as np +from mpi4py import MPI + +import tack +import tack.dataset as td +from tack.communication import World + +LAYOUTS = {"one": [[0, 1, 2, 3], [], []], "first-empty": [[], [0, 1], [2, 3]], + "last-empty": [[0, 1], [2, 3], []], "none": [[], [], []]} + + +# ── The pieces ──────────────────────────────────────────────────────── + +def _wavelet_pieces(): + return list(td.split(td.wavelet((0, 8, 0, 8, 0, 8)), (2, 2, 1))) + + +def structured(): + return _wavelet_pieces() + + +def tetrahedra(): + return [td.tetrahedralize(piece) for piece in _wavelet_pieces()] + + +def polyhedra(): + return [td.as_polyhedra(piece) for piece in _wavelet_pieces()] + + +def flowing(): + """The tetrahedra with a velocity ``v``.""" + return _with_velocity(tetrahedra(), 0.0) + + +def cloud(): + rng = np.random.default_rng(7) + points = rng.random((2000, 3)) * 8.0 + out = [] + for k in range(4): + mine = points[(points[:, 0] >= 2 * k) & (points[:, 0] < 2 * k + 2)] + piece = td.point_cloud(mine) + piece.fields["h"] = td.Field(td.H1(piece), tack.array(np.full(len(mine), 0.6, + np.float32))) + piece.fields["mass"] = td.Field(td.H1(piece), tack.array(np.ones(len(mine), + np.float32))) + out.append(piece) + return out + + +def pds(pieces, world, layout): + mine = LAYOUTS[layout][world.rank] + return td.PartitionedDataSet([pieces[k] for k in mine], world=world) + + +# ── The cases: (pieces, run(data) -> this rank's part, combine(parts) -> answer) ── + +def _cells(data): + return sum(p.num_cells for p in data) + + +def _owned_cells(data): + return sum(int((p.ghosts.cells.numpy() == 0).sum()) if p.ghosts is not None + else p.num_cells for p in data) + + +def _total(parts): + return sum(parts) + + +def _union(parts): + return sorted(set().union(*parts)) + + +def _first(parts): + found = [p for p in parts if p is not None] + return found[0] if found else None + + +def _same_everywhere(parts): + assert all(repr(p) == repr(parts[0]) for p in parts), parts + return parts[0] + + +SEEDS = np.array([[0.5, 0.5, 0.5], [6.5, 6.5, 4.0], [3.0, 7.0, 7.5]]) +WHERE = np.array([[0.25, 0.25, 0.25], [4.0, 4.0, 4.0], [7.9, 7.9, 7.9], [20.0, 0.0, 0.0]]) +GRID = td.uniform_grid((5, 5, 5), spacing=(2.0, 2.0, 2.0)) + + +def _ids(data): + return set().union(*[set(p.global_ids.points.numpy().tolist()) for p in data]) \ + if len(data) else set() + + +CASES = { + "ghost_cells/tetrahedra": (tetrahedra, lambda d: _owned_cells(td.ghost_cells(d)), _total), + "ghost_cells/structured": (structured, lambda d: _owned_cells(td.ghost_cells(d)), _total), + "ghost_cells/polyhedra": (polyhedra, lambda d: _owned_cells(td.ghost_cells(d)), _total), + "update_ghosts": (tetrahedra, lambda d: _owned_cells(td.update_ghosts(td.ghost_cells(d), + "RTData") + or td.ghost_cells(d)), _total), + "generate_global_ids/tetrahedra": (tetrahedra, lambda d: _ids(td.generate_global_ids(d)), + _union), + "generate_global_ids/structured": (structured, lambda d: _ids(td.generate_global_ids(d)), + _union), + "merge_pieces/tetrahedra": (tetrahedra, lambda d: (lambda m: None if m is None + else m.num_points)(td.merge_pieces(d)), + _first), + "merge_pieces/structured": (structured, lambda d: (lambda m: None if m is None + else m.num_points)(td.merge_pieces(d)), + _first), + "extract_surface": (tetrahedra, lambda d: _cells(td.extract_surface(d)), _total), + "statistics": (tetrahedra, lambda d: round(td.statistics(td.generate_global_ids(d), + "RTData").mean, 6), + _same_everywhere), + "entropy": (tetrahedra, lambda d: round(float(td.entropy(td.generate_global_ids(d), + "RTData")), 6), + _same_everywhere), + "integrate": (tetrahedra, lambda d: round(td.integrate(td.generate_global_ids(d)).measure, + 9), _same_everywhere), + "connected_components": (tetrahedra, + lambda d: _union_regions(td.connected_components(d)), _union), + "connected_components/radius": (cloud, lambda d: _union_regions( + td.connected_components(d, radius=0.3)), _union), + "probe": (tetrahedra, lambda d: _valid(td.probe(d, WHERE)), _same_everywhere), + "streamlines/seeds on every rank": ( # each rank releases its share of them + flowing, lambda d: _paths(td.streamlines(d, "v", SEEDS[d.world.rank::d.world.size], + 0.1, 40)), _total), + "advect/seeds on rank 0": ( + flowing, lambda d: _paths(td.advect(d, "v", SEEDS if _rank() == 0 + else np.zeros((0, 3)), 0.1, 40)), _total), + "ghost_points": (cloud, lambda d: _cells(td.ghost_points(d, radius=0.5)), _total), + "voxel_downsample": (cloud, lambda d: _cells(td.voxel_downsample(d, 1.0)), _total), + "outliers": (cloud, lambda d: sum(int(p.fields["outlier"].values.numpy().sum()) + for p in td.outliers(d, radius=0.3, min_neighbors=2)), + _total), + "deposit": (cloud, lambda d: round(float(td.deposit(d, GRID).values.numpy().sum()), 3), + _same_everywhere), + "voxelize": (tetrahedra, lambda d: round(float(td.voxelize(d, GRID, samples=1) + .fields["fraction"].values.numpy().sum()), + 6), _same_everywhere), +} + + +def _rank(): + return MPI.COMM_WORLD.rank + + +# ── Hierarchies, octrees, steps and VTK, spread by the same layouts ────────── + +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) + + +class Spread: + """What a case spreads itself: ``spread(world, layout)`` gives this rank's share, + and ``alone()`` the whole in one process.""" + + def __init__(self, spread, alone): + self.spread, self.alone = spread, alone + + +def _buckets(n, world, layout): + """Which of ``n`` items this rank holds: item i in bucket i % 4, the buckets + spread by ``layout``.""" + mine = set(LAYOUTS[layout][world.rank]) + return [i for i in range(n) if i % 4 in mine] + + +def _amr(world, layout): + import test_amr + + every = [list(level) for level in test_amr.hierarchy().levels] + items = [(lv, k) for lv, boxes in enumerate(every) for k in range(len(boxes))] + chosen = {items[i] for i in _buckets(len(items), world, layout)} + mine = [[test_amr._unmarked(box) for k, box in enumerate(boxes) if (lv, k) in chosen] + for lv, boxes in enumerate(every)] + return td.AMRDataSet(mine, [2, 2], world=world) + + +def _amr_alone(): + import test_amr + + return test_amr.hierarchy() + + +def _octree(world, layout): + import test_octrees + + roots, level, centre, rho, temperature = test_octrees.ramses_like(3, trees=(3, 2, 2)) + whole = td.OctreeDataSet.from_cells(roots, level, centre, fields={"T": temperature}) + trees = _buckets(roots.num_cells, world, layout) + mine = np.isin(whole._tree, trees) + return td.OctreeDataSet(roots, whole._tree[mine], whole._level[mine], whole._key[mine], + fields={"T": whole.fields["T"].numpy()[mine]}, world=world) + + +def _octree_alone(): + import test_octrees + + roots, level, centre, rho, temperature = test_octrees.ramses_like(3, trees=(3, 2, 2)) + return td.OctreeDataSet.from_cells(roots, level, centre, fields={"T": temperature}) + + +def _with_velocity(pieces, t): + out = [] + for piece in pieces: + piece.time = t + v = np.zeros((piece.num_points, 3), np.float32) + v[:, 0], v[:, 1] = 1.0, 0.3 * t + piece.fields["v"] = td.Field(td.H1(piece), tack.array(v, tack.vec3())) + out.append(piece) + return out + + +def _advected(data_at): + """Particles from SEEDS (rank 0's) carried over three steps: how many, and where + (summed).""" + advection = td.Advection("v", seeds=SEEDS if _rank() == 0 else np.zeros((0, 3)), + step_size=0.1) + for t in (0.0, 0.5, 1.0): + cloud = advection.update(data_at(t)) + clouds = list(cloud) if isinstance(cloud, td.PartitionedDataSet) else [cloud] + return sum(c.num_points for c in clouds) + + +def _vtk_pieces(world, layout): + from vtkmodules.vtkCommonDataModel import vtkPartitionedDataSet + + from tack.interop.vtk import dataset_to_vtk + + out = vtkPartitionedDataSet() + for k in LAYOUTS[layout][world.rank]: + out.append(dataset_to_vtk(tetrahedra()[k])) + return out + + +def _vtk_alone(): + from vtkmodules.vtkCommonDataModel import vtkPartitionedDataSet + + from tack.interop.vtk import dataset_to_vtk + + out = vtkPartitionedDataSet() + for piece in tetrahedra(): + out.append(dataset_to_vtk(piece)) + return out + + +def _vtk(algorithm_of, read): + """A VTK algorithm (made by ``algorithm_of(world)``) over the VTK pieces.""" + def run(vtk_data, world): + algorithm = algorithm_of(world) + out = algorithm(vtk_data) + if algorithm.error is not None: + raise algorithm.error + return read(out) + return run + + +def _vtk_probe(vtk_data, world): + """VTK's Probe of WHERE (the same points on every rank) over the pieces, as its + source: how many points a piece holds.""" + import vtkmodules.vtkCommonDataModel as dm + + from tack.interop.vtk.algorithms import Probe + + where = dm.vtkPolyData() + where.points = __import__("vtkmodules.vtkCommonCore", fromlist=["x"]).vtkPoints( + data=WHERE.copy()) + probe = Probe(world=world) + probe.SetInputDataObject(0, where) + probe.SetSourceData(vtk_data) + probe.Update() + if probe.error is not None: + raise probe.error + return int(np.asarray(probe.GetOutputDataObject(0).point_data["valid"]).view(np.int8).sum()) + + +def _leaf_cells(out): + if hasattr(out, "GetNumberOfPartitions"): + return sum(out.GetPartition(i).GetNumberOfCells() + for i in range(out.GetNumberOfPartitions())) + return out.GetNumberOfCells() + + +SPREAD_CASES = { + "amr/statistics": (Spread(_amr, _amr_alone), + lambda a, w: td.statistics(a, "c").n, _same_everywhere), + "amr/probe": (Spread(_amr, _amr_alone), + lambda a, w: round(float(td.probe(a, WHERE * 0.9).fields["c"].values + .numpy().sum()), 4), _same_everywhere), + "amr/ghost_cells": (Spread(_amr, _amr_alone), + lambda a, w: sum(b.num_cells for _, b in td.ghost_cells(a)), _total), + "amr/average_down": (Spread(_amr, _amr_alone), + lambda a, w: (td.average_down(a, "c"), td.statistics(a, "c").n)[1], + _same_everywhere), + "octree/statistics": (Spread(_octree, _octree_alone), + lambda o, w: round(td.statistics(o, "T").mean, 6), _same_everywhere), + "octree/probe": (Spread(_octree, _octree_alone), + lambda o, w: int(td.probe(o, WHERE * 0.1).fields["valid"].values.numpy() + .sum()), _same_everywhere), + "advection": (Spread(lambda w, layout: (w, layout), lambda: None), + lambda wl, w: _advected( + (lambda t: td.PartitionedDataSet( + _with_velocity([tetrahedra()[k] for k in LAYOUTS[wl[1]][w.rank]], t), + world=w)) if wl is not None else + (lambda t: td.PartitionedDataSet(_with_velocity(tetrahedra(), t), + world=World()))), _total), + "vtk/GenerateGlobalIds": (Spread(_vtk_pieces, _vtk_alone), _vtk( + lambda w: __import__("tack.interop.vtk.algorithms", fromlist=["x"]) + .GenerateGlobalIds(world=w), _leaf_cells), _total), + "vtk/Statistics": (Spread(_vtk_pieces, _vtk_alone), _vtk( + lambda w: __import__("tack.interop.vtk.algorithms", fromlist=["x"]) + .Statistics(fields=["RTData"], world=w), + lambda t: t.GetColumnByName("n").GetValue(0)), _same_everywhere), + "vtk/MergePieces": (Spread(_vtk_pieces, _vtk_alone), _vtk( + lambda w: __import__("tack.interop.vtk.algorithms", fromlist=["x"]) + .MergePieces(world=w), lambda m: m.GetNumberOfPoints() or None), _first), + "vtk/Integrate": (Spread(_vtk_pieces, _vtk_alone), _vtk( + lambda w: __import__("tack.interop.vtk.algorithms", fromlist=["x"]) + .Integrate(fields=["RTData"], world=w), + lambda t: round(float(np.asarray(t.cell_data["Volume"])[0]), 9)), _same_everywhere), + "vtk/GhostCells": (Spread(_vtk_pieces, _vtk_alone), _vtk( + lambda w: __import__("tack.interop.vtk.algorithms", fromlist=["x"]) + .GhostCells(world=w), _leaf_cells), _total), + "vtk/Probe": (Spread(_vtk_pieces, _vtk_alone), _vtk_probe, _same_everywhere), + "vtk/ExtractSurface": (Spread(_vtk_pieces, _vtk_alone), _vtk( + lambda w: __import__("tack.interop.vtk.algorithms", fromlist=["x"]) + .ExtractSurface(world=w), _leaf_cells), _total), +} + + +def _union_regions(data): + return {int(r) for p in data for r in np.unique(p.fields["region"].values.numpy()) + if r >= 0} + + +def _valid(result): + if isinstance(result, td.PartitionedDataSet): + return sum(int(p.fields["valid"].values.numpy().sum()) for p in result) + return int(result.fields["valid"].values.numpy().sum()) + + +def _paths(result): + return sum(p.num_cells for p in result) if isinstance(result, td.PartitionedDataSet) \ + else result.num_cells + + +# ── Running ─────────────────────────────────────────────────────────── + +def main(): + comm = MPI.COMM_WORLD + layout, names = sys.argv[1], sys.argv[2:] or [*CASES, *SPREAD_CASES] + try: + import vtkmodules.vtkCommonDataModel # noqa: F401 + except ImportError: # the VTK cases need VTK (interop extra) + names = [n for n in names if not n.startswith("vtk/")] + tack.set_default_device("cpu") + world = World(comm) + failures = [] + for name in names: + spread = name in SPREAD_CASES + make, run, combine = SPREAD_CASES[name] if spread else CASES[name] + traces = os.environ.get("TACK_EMPTY_RANKS_TRACES") + out = open(os.path.join(traces, f"rank{comm.rank}.txt"), "w") if traces else sys.stderr + faulthandler.dump_traceback_later( + int(os.environ.get("TACK_EMPTY_RANKS_DEADLINE", "120")) + 5 * comm.rank, + exit=True, file=out) + try: + if spread: + part = run(make.spread(world, layout), world) + else: + part = run(pds(make(), world, layout)) + error = None + except Exception as e: + part, error = None, f"{type(e).__name__}: {e}" + if os.environ.get("TACK_EMPTY_RANKS_TRACE"): + traceback.print_exc() + faulthandler.cancel_dump_traceback_later() + parts, errors = comm.allgather(part), comm.allgather(error) + # The same error on every rank with no pieces anywhere is a refusal, and the + # ranks go on together; any other error ends the run on every rank alike, + # since they may disagree from there. + refused = layout == "none" and any(errors) and len(set(errors)) == 1 + if any(errors) and not refused: + if comm.rank == 0: + failures.append(f"{name}: {errors}") + break + if comm.rank == 0 and not refused: + got = combine(parts) + if layout == "none": + continue + alone = World() + expected = combine([run(make.alone(), alone) if spread else + run(td.PartitionedDataSet(make(), world=alone))]) + if repr(got) != repr(expected): + failures.append(f"{name}: {got!r} against {expected!r} in one process") + if comm.rank == 0: + print("ok" if not failures else "FAILED\n" + "\n".join(failures), flush=True) + + +if __name__ == "__main__": + main() diff --git a/packages/tack-vis/tests/test_empty_ranks.py b/packages/tack-vis/tests/test_empty_ranks.py new file mode 100644 index 00000000..31ba0e2d --- /dev/null +++ b/packages/tack-vis/tests/test_empty_ranks.py @@ -0,0 +1,25 @@ +# SPDX-FileCopyrightText: Copyright (c) Kitware Inc. +# SPDX-License-Identifier: BSD-3-Clause + +"""Every algorithm over pieces with MPI ranks that hold none of them. + +A rank with no pieces is common: four files read on eight ranks, a reader +that gives rank 0 nothing, a hierarchy whose fine boxes are on a few ranks. +Each algorithm that communicates must take part from such a rank, deciding +nothing from its own pieces alone that the others decide differently, or the +others wait for it forever. ``_empty_ranks_mpi.py`` runs each on 3 ranks with +four pieces (boxes, trees) all on rank 0, on ranks 1 and 2, on ranks 0 and 1, +and on none, and holds the answers to the same pieces in one process. +""" + +import pathlib + +import pytest + +SCRIPT = pathlib.Path(__file__).with_name("_empty_ranks_mpi.py") + + +@pytest.mark.parametrize("layout", ["one", "first-empty", "last-empty", "none"]) +def test_algorithms_over_pieces_take_ranks_without_pieces(mpiexec, layout): + run = mpiexec(3, SCRIPT, layout) + assert run.returncode == 0 and "ok" in run.stdout.split(), run.stdout + run.stderr[-3000:]