From 1f5e36ee83a5eadbe725bcc1a246f4a0a72d3acc Mon Sep 17 00:00:00 2001 From: atulcthakur Date: Tue, 28 Jul 2026 21:12:22 -0700 Subject: [PATCH] fix: scope DD communication to domain groups Signed-off-by: atulcthakur --- .../distributed/_core/gather_primitives.py | 49 ++-- nvalchemi/distributed/_core/particle_halo.py | 45 ++-- .../distributed/_dynamics_coordinator.py | 6 +- nvalchemi/distributed/sharded_batch.py | 2 +- nvalchemi/distributed/strategy.py | 4 + .../_core/test_subgroup_communication.py | 246 ++++++++++++++++++ .../test_2d_subgroup_collectives.py | 136 ++++++++++ 7 files changed, 452 insertions(+), 36 deletions(-) create mode 100644 test/distributed/_core/test_subgroup_communication.py create mode 100644 test/distributed/test_2d_subgroup_collectives.py diff --git a/nvalchemi/distributed/_core/gather_primitives.py b/nvalchemi/distributed/_core/gather_primitives.py index e010f7c2..f2abe15c 100644 --- a/nvalchemi/distributed/_core/gather_primitives.py +++ b/nvalchemi/distributed/_core/gather_primitives.py @@ -49,16 +49,16 @@ def mesh_group(mesh: Any) -> Any: - """Return the default ``ProcessGroup`` for *mesh*. + """Return the concrete ``ProcessGroup`` for *mesh*. - Accepts either a real ``DeviceMesh`` or a test-harness mock; returns - ``None`` for both "no distribution configured" and "mesh present but not - group-capable". + Accepts a concrete ``ProcessGroup`` unchanged, or resolves one from a real + ``DeviceMesh`` or test-harness mock. Returns ``None`` for both "no + distribution configured" and "mesh present but not group-capable". Parameters ---------- - mesh : DeviceMesh or object or None - The device mesh to resolve a group from. + mesh : ProcessGroup, DeviceMesh, object, or None + The process group to preserve or device mesh to resolve. Returns ------- @@ -68,6 +68,8 @@ def mesh_group(mesh: Any) -> Any: """ if mesh is None: return None + if isinstance(mesh, dist.ProcessGroup): + return mesh get_group = getattr(mesh, "get_group", None) if get_group is None: return None @@ -100,15 +102,18 @@ def _funcol_group_arg(mesh: Any) -> Any: Prefers the ``(DeviceMesh, 0)`` spec — the form Dynamo special-cases (and the one physicsnemo's own compiled collectives use). Real distributed inference always carries a real ``DeviceMesh``, so the compiled path takes - this branch. Falls back to the resolved ``ProcessGroup`` only for eager - test harnesses that pass a lightweight (non-``DeviceMesh``) mesh; the - ``isinstance`` is a compile-time constant, so under ``torch.compile`` only - the traceable branch survives. + this branch. A concrete ``ProcessGroup`` is returned unchanged. Otherwise, + falls back to the resolved ``ProcessGroup`` for eager test harnesses that + pass a lightweight (non-``DeviceMesh``) mesh; the ``isinstance`` checks are + compile-time constants, so under ``torch.compile`` only the traceable branch + survives. """ from torch.distributed.device_mesh import DeviceMesh # noqa: PLC0415 if isinstance(mesh, DeviceMesh): return (mesh, 0) + if isinstance(mesh, dist.ProcessGroup): + return mesh return funcol_group(mesh) @@ -343,18 +348,24 @@ def _neighbor_p2p_v_1d( if send_counts[r] > 0: recv[r_off[r] : r_off[r + 1]].copy_(send[s_off[r] : s_off[r + 1]]) continue + peer = dist.get_global_rank(group, r) if group is not None else r if send_counts[r] > 0: ops.append( dist.P2POp( dist.isend, send[s_off[r] : s_off[r + 1]].contiguous(), - r, + peer, group=group, ) ) if recv_counts[r] > 0: ops.append( - dist.P2POp(dist.irecv, recv[r_off[r] : r_off[r + 1]], r, group=group) + dist.P2POp( + dist.irecv, + recv[r_off[r] : r_off[r + 1]], + peer, + group=group, + ) ) if ops: for work in dist.batch_isend_irecv(ops): @@ -421,8 +432,11 @@ def _neighbor_p2p_fixed( if r == rank: continue sl = slice(r * m, (r + 1) * m) - ops.append(dist.P2POp(dist.isend, send_rows[sl].contiguous(), r, group=group)) - ops.append(dist.P2POp(dist.irecv, recv[sl], r, group=group)) + peer = dist.get_global_rank(group, r) if group is not None else r + ops.append( + dist.P2POp(dist.isend, send_rows[sl].contiguous(), peer, group=group) + ) + ops.append(dist.P2POp(dist.irecv, recv[sl], peer, group=group)) if ops: for work in dist.batch_isend_irecv(ops): work.wait() @@ -467,7 +481,7 @@ def halo_exchange_fixed( is_nccl = False if is_nccl: return _neighbor_p2p_fixed(send_rows, world_size, neighbors, group) - return funcol_all_to_all_fixed(send_rows, world_size, None) + return funcol_all_to_all_fixed(send_rows, world_size, group) def funcol_all_to_all_v_rows( @@ -581,10 +595,11 @@ def _isend_irecv_v_1d( if r == rank: recv_slice.copy_(send_slice) else: + peer = dist.get_global_rank(group, r) if group is not None else r if send_slice.numel() > 0: - ops.append(dist.isend(send_slice, dst=r, group=group)) + ops.append(dist.isend(send_slice, dst=peer, group=group)) if recv_slice.numel() > 0: - ops.append(dist.irecv(recv_slice, src=r, group=group)) + ops.append(dist.irecv(recv_slice, src=peer, group=group)) for op in ops: op.wait() diff --git a/nvalchemi/distributed/_core/particle_halo.py b/nvalchemi/distributed/_core/particle_halo.py index f3c90efe..a133c7ea 100644 --- a/nvalchemi/distributed/_core/particle_halo.py +++ b/nvalchemi/distributed/_core/particle_halo.py @@ -55,6 +55,19 @@ logger = logging.getLogger(__name__) +# ``torch.library.custom_op`` schemas cannot carry a Python ``ProcessGroup``. +# Publish the active halo group immediately before each model forward so the +# eager custom-op bodies can resolve their collectives to the same domain group +# as the surrounding DD strategy. ``None`` preserves the default-group behavior +# for direct/single-mesh callers. +_HALO_PROCESS_GROUP: Any = None + + +def set_halo_process_group(group: Any) -> None: + """Publish the domain process group used by compiled halo custom ops.""" + global _HALO_PROCESS_GROUP + _HALO_PROCESS_GROUP = group + # ====================================================================== # Ghost identification @@ -725,19 +738,21 @@ def _halo_a2a_v_default_group( rank: int, world_size: int, ) -> torch.Tensor: - """:func:`_funcol_indexed_all_to_all_v_rows` over the DEFAULT process group. + """:func:`_funcol_indexed_all_to_all_v_rows` over the active halo group. Used inside the halo-correction custom op, which runs eagerly at runtime and - cannot take a ``DeviceMesh`` arg. Valid for single-domain-mesh-dim topology - (the domain group IS the world group); ``funcol`` mesh=None resolves the - default group. + cannot take a ``DeviceMesh`` or ``ProcessGroup`` arg. The surrounding halo + strategy publishes its domain group before the model forward; direct callers + that publish no group retain the default-world behavior. """ send_rows = torch.cat( [tensor.index_select(0, indices[j]) for j in range(world_size)], dim=0 ) send_counts = [int(sizes[rank][j]) for j in range(world_size)] recv_counts = [int(sizes[i][rank]) for i in range(world_size)] - return funcol_all_to_all_v_rows(send_rows, send_counts, recv_counts, None) + return funcol_all_to_all_v_rows( + send_rows, send_counts, recv_counts, _HALO_PROCESS_GROUP + ) def _halo_scatter_correct_dense( @@ -749,7 +764,7 @@ def _halo_scatter_correct_dense( world_size: int, ) -> torch.Tensor: """``halo_forward_exchange(halo_reverse_exchange(padded))`` as pure - compute + collective (no autograd.Function, default group).""" + compute + collective (no autograd.Function, active halo group).""" # reverse: fold borrowed halo rows back into their owners. halo = padded[n_owned:].contiguous() rev_indices: list[torch.Tensor] = [] @@ -892,10 +907,10 @@ def halo_forward_op( this rank's halo (ghost) region. Compile-safe counterpart of :func:`halo_forward_exchange`; runs eagerly at - runtime (default group), while the trace sees only the registered fake. Its - adjoint (backward) is :func:`halo_scatter_correct_op`. The marker arrays ride - as ``int[]`` (not tensors) so inductor lowering does not see a real-Tensor - constant alongside the fake ``owned`` input. + runtime on the active halo group, while the trace sees only the registered + fake. Its adjoint (backward) is :func:`halo_scatter_correct_op`. The marker + arrays ride as ``int[]`` (not tensors) so inductor lowering does not see a + real-Tensor constant alongside the fake ``owned`` input. Parameters ---------- @@ -1199,10 +1214,10 @@ def halo_forward_static_op( world_size: int, ) -> torch.Tensor: """Fixed-shape owned->[owned|ghost] refresh. Runs eagerly at runtime - (default group); the trace sees only the fake. Owned rows pass through; + on the active halo group; the trace sees only the fake. Owned rows pass through; ghost rows are gathered from neighbors via a uniform-split all_to_all.""" send_rows = padded_in.index_select(0, send_index) - recv = halo_exchange_fixed(send_rows, world_size, None) + recv = halo_exchange_fixed(send_rows, world_size, _HALO_PROCESS_GROUP) recv = recv * _row_mask_like(recv_real, recv).to(recv.dtype) ghost_acc = torch.zeros_like(padded_in).index_add(0, recv_dest, recv) rowidx = torch.arange(padded_in.shape[0], device=padded_in.device) @@ -1232,7 +1247,7 @@ def _hfs_backward(ctx, grad_out): # type: ignore[no-untyped-def] grad_recv = grad_ghost.index_select(0, recv_dest) * _row_mask_like( recv_real, grad_out ).to(grad_out.dtype) - grad_send = halo_exchange_fixed(grad_recv, ws, None) + grad_send = halo_exchange_fixed(grad_recv, ws, _HALO_PROCESS_GROUP) grad_in = grad_out * (~ghostmask).to(grad_out.dtype) grad_in = grad_in.index_add(0, send_index, grad_send) return grad_in, None, None, None, None, None @@ -1262,13 +1277,13 @@ def halo_scatter_correct_static_op( # reverse: ghost rows (recv-slot order) -> owning rank -> index_add into owners. ghost_rows = padded_in.index_select(0, recv_dest) * recv_real_f - back = halo_exchange_fixed(ghost_rows, world_size, None) + back = halo_exchange_fixed(ghost_rows, world_size, _HALO_PROCESS_GROUP) owned_only = (padded_in * (~ghostmask).to(padded_in.dtype)).to(acc_dt) owned_acc = owned_only.index_add(0, send_index, back.to(acc_dt)).to(padded_in.dtype) # forward: re-broadcast corrected owners to ghosts. send_rows = owned_acc.index_select(0, send_index) - recv = halo_exchange_fixed(send_rows, world_size, None) * recv_real_f + recv = halo_exchange_fixed(send_rows, world_size, _HALO_PROCESS_GROUP) * recv_real_f ghost_acc = torch.zeros_like(padded_in).index_add(0, recv_dest, recv) return torch.where(ghostmask, ghost_acc, owned_acc) diff --git a/nvalchemi/distributed/_dynamics_coordinator.py b/nvalchemi/distributed/_dynamics_coordinator.py index bd55f717..28a874c4 100644 --- a/nvalchemi/distributed/_dynamics_coordinator.py +++ b/nvalchemi/distributed/_dynamics_coordinator.py @@ -256,8 +256,8 @@ def _restore(self) -> None: def broadcast_state(self, batch: Batch) -> None: """Broadcast the replicated controller + cell state (the integrator's - declared ``__dd_replicated__``) from rank 0 so floating-point divergence - cannot accumulate over a long run. + declared ``__dd_replicated__``) from group-local rank 0 so floating-point + divergence cannot accumulate over a long run. With global KE/DOF and identical config the controller evolves identically on every rank, so this is anti-drift insurance, not a @@ -266,7 +266,7 @@ def broadcast_state(self, batch: Batch) -> None: if not self.active: return group = self._strategy.process_group - src = 0 + src = dist.get_global_rank(group, 0) if group is not None else 0 state = getattr(self._dyn, "_state", None) fields = getattr(self._dyn, "__dd_replicated__", ()) for name in fields: diff --git a/nvalchemi/distributed/sharded_batch.py b/nvalchemi/distributed/sharded_batch.py index d2293f11..42eccae4 100644 --- a/nvalchemi/distributed/sharded_batch.py +++ b/nvalchemi/distributed/sharded_batch.py @@ -270,7 +270,7 @@ def rank_assignment(self) -> torch.Tensor: # Single all_gather into a flat (world_size,) tensor, one sync. n_owned_t = torch.tensor([self.n_owned], dtype=torch.int64, device=device) sizes_t = torch.empty(world_size, dtype=torch.int64, device=device) - dist.all_gather_into_tensor(sizes_t, n_owned_t) + dist.all_gather_into_tensor(sizes_t, n_owned_t, group=mesh_group(self.mesh)) # Build the block-constant assignment via repeat_interleave — no # per-rank Python loop or slicing. diff --git a/nvalchemi/distributed/strategy.py b/nvalchemi/distributed/strategy.py index 0c231977..a0893973 100644 --- a/nvalchemi/distributed/strategy.py +++ b/nvalchemi/distributed/strategy.py @@ -52,6 +52,7 @@ mesh_group, set_halo_neighbor_ranks, ) +from nvalchemi.distributed._core.particle_halo import set_halo_process_group if TYPE_CHECKING: from nvalchemi.data.batch import Batch @@ -307,6 +308,9 @@ def run_forward( """Run the model forward on this rank's halo-padded (owned + ghost) shard.""" dist_model._dist_ctx.cap_atoms = self.caps_atoms dist_model._dist_ctx.strategy = self + # Compiled halo custom ops cannot carry a Python ProcessGroup through + # their dispatcher schema, so publish this strategy's exact domain group. + set_halo_process_group(self.process_group) return _halo_run_forward(dist_model, state, wired_fields) def on_cell_change(self, state: ShardState, cell: torch.Tensor | None) -> None: diff --git a/test/distributed/_core/test_subgroup_communication.py b/test/distributed/_core/test_subgroup_communication.py new file mode 100644 index 00000000..bd27d884 --- /dev/null +++ b/test/distributed/_core/test_subgroup_communication.py @@ -0,0 +1,246 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Communication primitives over domain subgroups of a two-dimensional mesh.""" + +from __future__ import annotations + +import sys +from pathlib import Path + +import pytest +import torch + +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) +from _gloo_harness import run_gloo # noqa: E402 + + +def _domain_group(rank: int): # type: ignore[no-untyped-def] + """Return this rank's domain row in a (pipeline=2, domain=2) layout.""" + import torch.distributed as dist + + first_stage = dist.new_group([0, 1]) + second_stage = dist.new_group([2, 3]) + return first_stage if rank < 2 else second_stage + + +def _expected_received(rank: int) -> list[float]: + """Expected source-ordered values within the rank's two-member domain row.""" + stage_first_rank = (rank // 2) * 2 + domain_rank = rank % 2 + return [ + float(source * 10 + domain_rank) + for source in (stage_first_rank, stage_first_rank + 1) + ] + + +def _check_p2p_helper(rank: int, world_size: int, queue, helper_name: str) -> None: # type: ignore[no-untyped-def] + import torch.distributed as dist + + from nvalchemi.distributed._core.gather_primitives import ( + _isend_irecv_v_1d, + _neighbor_p2p_fixed, + _neighbor_p2p_v_1d, + ) + + assert world_size == 4 + group = _domain_group(rank) + domain_world_size = dist.get_world_size(group) + domain_rank = dist.get_rank(group) + + # Slot d contains the value this global rank is sending to domain-local rank d. + send = torch.tensor( + [float(rank * 10 + destination) for destination in range(domain_world_size)] + ) + + received: list[float] | None = None + error: str | None = None + try: + if helper_name == "neighbor_variable": + result = _neighbor_p2p_v_1d( + send, + [1] * domain_world_size, + [1] * domain_world_size, + group, + ) + elif helper_name == "neighbor_fixed": + result = _neighbor_p2p_fixed( + send, + domain_world_size, + list(range(domain_world_size)), + group, + ) + else: + result = torch.empty_like(send) + _isend_irecv_v_1d( + send, + [1] * domain_world_size, + result, + [1] * domain_world_size, + group, + ) + received = result.tolist() + except Exception as exc: # noqa: BLE001 + error = f"{type(exc).__name__}: {exc}" + + # Keep the first domain row alive until the second row has also completed + # (or reported its rank-translation error). + dist.barrier() + queue.put( + ( + rank, + domain_rank, + received, + _expected_received(rank), + error, + ) + ) + + +@pytest.mark.parametrize( + "helper_name", + ["neighbor_variable", "neighbor_fixed", "gloo_fallback"], +) +def test_p2p_helpers_use_global_peers_in_domain_subgroups(helper_name: str) -> None: + """A domain-local peer index must be translated before calling P2P APIs.""" + results = run_gloo( + world_size=4, + fn=_check_p2p_helper, + args=(helper_name,), + timeout_sec=30.0, + ) + assert len(results) == 4, results + for rank, domain_rank, received, expected, error in results: + assert error is None, ( + f"global rank {rank} (domain rank {domain_rank}) failed: {error}" + ) + assert received == expected, ( + f"global rank {rank} (domain rank {domain_rank}) " + f"received {received}, expected {expected}" + ) + + +def _check_fixed_halo_group(rank: int, world_size: int, queue) -> None: # type: ignore[no-untyped-def] + import torch.distributed as dist + + from nvalchemi.distributed._core.gather_primitives import ( + halo_exchange_fixed, + mesh_group, + ) + + assert world_size == 4 + group = _domain_group(rank) + assert mesh_group(group) is group + domain_world_size = dist.get_world_size(group) + domain_rank = dist.get_rank(group) + send_rows = torch.tensor( + [float(rank * 10 + destination) for destination in range(domain_world_size)] + ) + + received: list[float] | None = None + error: str | None = None + try: + result = halo_exchange_fixed(send_rows, domain_world_size, group) + received = result.tolist() + except Exception as exc: # noqa: BLE001 + error = f"{type(exc).__name__}: {exc}" + + dist.barrier() + queue.put( + ( + rank, + domain_rank, + received, + _expected_received(rank), + error, + ) + ) + + +def test_halo_exchange_fixed_uses_domain_subgroup() -> None: + """The fixed halo dispatcher must not replace its subgroup with WORLD.""" + results = run_gloo( + world_size=4, + fn=_check_fixed_halo_group, + timeout_sec=30.0, + ) + assert len(results) == 4, results + for rank, domain_rank, received, expected, error in results: + assert error is None, ( + f"global rank {rank} (domain rank {domain_rank}) failed: {error}" + ) + assert received == expected, ( + f"global rank {rank} (domain rank {domain_rank}) " + f"received {received}, expected {expected}" + ) + + +def _check_static_halo_op_group(rank: int, world_size: int, queue) -> None: # type: ignore[no-untyped-def] + import torch.distributed as dist + + from nvalchemi.distributed._core.particle_halo import ( + halo_forward_static_op, + set_halo_process_group, + ) + + assert world_size == 4 + group = _domain_group(rank) + domain_rank = dist.get_rank(group) + set_halo_process_group(group) + + # One owned row, one real ghost row from the peer, and one dead padding row. + # Each peer block has a fixed capacity of one row. The self-source block is + # masked as padding; the other source block lands in the ghost row. + padded = torch.tensor([float(rank), -1.0, -1.0], requires_grad=True) + send_index = torch.tensor([0, 0], dtype=torch.int64) + if domain_rank == 0: + recv_dest = torch.tensor([2, 1], dtype=torch.int64) + recv_real = torch.tensor([False, True]) + else: + recv_dest = torch.tensor([1, 2], dtype=torch.int64) + recv_real = torch.tensor([True, False]) + n_owned = torch.tensor(1, dtype=torch.int64) + + out = halo_forward_static_op( + padded, + send_index, + recv_dest, + recv_real, + n_owned, + dist.get_world_size(group), + ) + out.sum().backward() + + peer = rank + 1 if domain_rank == 0 else rank - 1 + queue.put( + (rank, out.tolist(), padded.grad.tolist(), [float(rank), float(peer), 0.0]) + ) + + +def test_static_halo_op_uses_published_domain_subgroup() -> None: + """The real compiled-path custom op must stay inside its pipeline stage.""" + results = run_gloo( + world_size=4, + fn=_check_static_halo_op_group, + timeout_sec=30.0, + ) + assert len(results) == 4, results + for rank, received, gradient, expected in results: + assert received == expected, ( + f"global rank {rank} received {received}, expected {expected}" + ) + assert gradient == [2.0, 0.0, 0.0], ( + f"global rank {rank} produced input gradient {gradient}" + ) diff --git a/test/distributed/test_2d_subgroup_collectives.py b/test/distributed/test_2d_subgroup_collectives.py new file mode 100644 index 00000000..322cab9d --- /dev/null +++ b/test/distributed/test_2d_subgroup_collectives.py @@ -0,0 +1,136 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Subgroup correctness for pipeline × domain (2-D) meshes. + +With a ``(pipeline=2, domain=2)`` mesh, the domain groups are global ranks +``{0, 1}`` and ``{2, 3}``. Both groups number their members locally as +``{0, 1}``, so these tests catch collectives that accidentally use a local rank +as a global rank or fall back to the four-rank world group. +""" + +from __future__ import annotations + +import sys +from pathlib import Path +from types import SimpleNamespace +from typing import Any + +import pytest +import torch +import torch.distributed as dist + +sys.path.insert(0, str(Path(__file__).resolve().parent)) +from _gloo_harness import run_gloo # noqa: E402 + +pytestmark = pytest.mark.skipif( + not dist.is_gloo_available(), reason="gloo backend required" +) + + +class _LocalRows: + """Small ShardTensor stand-in for properties that only call ``to_local``.""" + + def __init__(self, rows: torch.Tensor) -> None: + self._rows = rows + + def to_local(self) -> torch.Tensor: + return self._rows + + +def _domain_submesh() -> Any: + """Build this rank's two-rank domain row from a 2 × 2 device mesh.""" + from torch.distributed import init_device_mesh + + mesh_2d = init_device_mesh("cpu", (2, 2), mesh_dim_names=("pipeline", "domain")) + return mesh_2d["domain"] + + +def _broadcast_state_worker(rank: int, world_size: int, queue: Any) -> None: + del world_size + + from nvalchemi.distributed._dynamics_coordinator import ( + DynamicsDistributionCoordinator, + ) + + domain = _domain_submesh() + domain_rank = domain.get_local_rank() + state = SimpleNamespace(controller=torch.tensor([rank], dtype=torch.int64)) + dynamics = SimpleNamespace( + __dd_thermo_kind__="nhc", + __dd_replicated__=("controller",), + _state=state, + ) + strategy = SimpleNamespace(process_group=domain.get_group()) + + coordinator = DynamicsDistributionCoordinator(dynamics, strategy) + coordinator.broadcast_state(SimpleNamespace()) + + # Every domain row must receive its own local-rank-0 value: + # group {0, 1} -> 0 and group {2, 3} -> 2. + expected_lead = rank - domain_rank + torch.testing.assert_close( + state.controller, torch.tensor([expected_lead], dtype=torch.int64) + ) + queue.put((rank, int(state.controller.item()))) + + +def test_controller_state_broadcast_uses_each_domain_group_lead() -> None: + """The second domain row must broadcast from global rank 2, not rank 0.""" + results = run_gloo(world_size=4, fn=_broadcast_state_worker, timeout_sec=30.0) + assert sorted(results) == [(0, 0), (1, 0), (2, 2), (3, 2)] + + +def _rank_assignment_worker(rank: int, world_size: int, queue: Any) -> None: + del world_size + + from nvalchemi.distributed.sharded_batch import ShardedBatch + + domain = _domain_submesh() + domain_rank = domain.get_local_rank() + pipeline_rank = rank // 2 + + # Make the two domain rows deliberately different: + # group {0, 1}: local counts [1, 2] + # group {2, 3}: local counts [3, 4] + group_counts = [1, 2] if pipeline_rank == 0 else [3, 4] + local_n = group_counts[domain_rank] + n_global = sum(group_counts) + + sharded = ShardedBatch( + mesh=domain, + atom_fields={"positions": _LocalRows(torch.zeros(local_n, 3))}, + cell=torch.eye(3).unsqueeze(0), + pbc=torch.zeros(1, 3, dtype=torch.bool), + n_global=n_global, + ) + + got = sharded.rank_assignment + expected = torch.tensor( + [0] * group_counts[0] + [1] * group_counts[1], dtype=torch.int64 + ) + torch.testing.assert_close(got, expected) + queue.put((rank, got.tolist())) + + +def test_rank_assignment_gathers_counts_within_domain_group() -> None: + """Ownership counts from the other pipeline row must not be gathered.""" + results = run_gloo(world_size=4, fn=_rank_assignment_worker, timeout_sec=30.0) + assert sorted(results) == [ + (0, [0, 1, 1]), + (1, [0, 1, 1]), + (2, [0, 0, 0, 1, 1, 1, 1]), + (3, [0, 0, 0, 1, 1, 1, 1]), + ]