From 5874013f5f95951348de2bd8c4294f3bb02b7ab5 Mon Sep 17 00:00:00 2001 From: Berk Geveci Date: Thu, 5 Nov 2026 13:14:58 -0500 Subject: [PATCH] Algorithms over pieces take MPI ranks that hold no pieces A reader that gives some ranks nothing makes ranks without pieces. Each rank must make the same calls that communicate, so none may decide from its own pieces alone what the others decide from theirs. - probe of pieces or an AMR hierarchy, streamlines and advect over pieces, voxelize, and Advection (with the path filters) refused or failed on a rank without pieces, leaving the others waiting. Such a rank now takes part through a piece of nothing (_common._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 is made on the first rank that has pieces and moved to the root. test_empty_ranks.py runs 36 cases (pieces, AMR hierarchies, octrees, advection, VTK algorithms) on 3 ranks with the pieces on one, two or none of them, against the same pieces in one process; CI's MPI job runs it. --- .github/workflows/ci.yml | 7 +- changelog.d/empty-ranks.md | 22 + .../src/tack/dataset/algorithms/_common.py | 22 + .../src/tack/dataset/algorithms/_extract.py | 11 +- .../src/tack/dataset/algorithms/_particles.py | 50 ++- .../src/tack/dataset/algorithms/amr_ghosts.py | 12 +- .../dataset/algorithms/extract_surface.py | 4 +- .../tack/dataset/algorithms/ghost_cells.py | 6 +- .../src/tack/dataset/algorithms/integrate.py | 30 +- .../tack/dataset/algorithms/merge_pieces.py | 118 ++++- .../src/tack/dataset/algorithms/probe.py | 38 +- .../src/tack/dataset/algorithms/voxelize.py | 11 +- .../tack-vis/src/tack/dataset/temporal.py | 11 +- .../src/tack/interop/vtk/algorithms.py | 24 +- packages/tack-vis/tests/_empty_ranks_mpi.py | 418 ++++++++++++++++++ packages/tack-vis/tests/test_empty_ranks.py | 25 ++ 16 files changed, 735 insertions(+), 74 deletions(-) create mode 100644 changelog.d/empty-ranks.md create mode 100644 packages/tack-vis/tests/_empty_ranks_mpi.py create mode 100644 packages/tack-vis/tests/test_empty_ranks.py diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 652fdda43..d550baa52 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 000000000..76b6af87e --- /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 52b29ac15..5c46cabcc 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 4f85cdc01..bfbafda63 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 5cd53b497..6a4d7e607 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 18cb01e6e..25df01621 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 92d3545bb..d5997af2d 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 5a53bddb4..50b28b64f 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 0c7d014da..a6d502bbe 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 268655e7f..d39d14282 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 c7e0c16e0..0efe48ef3 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 0a39b0308..0b4afbbb5 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 66c2bfef8..362e8f967 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 54f3dc60e..6a458d501 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 000000000..320d92163 --- /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 000000000..31ba0e2d3 --- /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:]