Refresh final TrainerRank admission and propagate counter failures

src/art/trainer_rank/_impl.py

@@ -3803,6 +3803,31 @@ def _select_next_micro_batch(
         start: int,
         *,
         checkpoint: AdapterSelection = Unset,
+    ) -> _CandidateMicroBatch[ForwardInputsT]:
+        candidate = self._search_next_micro_batch(items, start, checkpoint=checkpoint)
+        # A later width check may observe less available memory.
+        # Do not return an earlier width's cached budget as the final admission.
+        check = self._memory_check_required(
+            candidate.check.estimated_required_bytes,
+            sync_across_dp=True,
+        )
+        if not check.fits:
+            self._snapshot_planning_telemetry(candidate.plan, check)
+            raise _memory_error(
+                context="forward_micro_batches",
+                message="selected microbatch exceeds freshly sampled available memory",
+                packed_tokens=candidate.plan.packed_tokens,
+                logical_tokens=candidate.plan.logical_tokens,
+                check=check,
+            )
+        return replace(candidate, check=check)
+
+    def _search_next_micro_batch(
+        self,
+        items: Sequence[ForwardInputsT],
+        start: int,
+        *,
+        checkpoint: AdapterSelection = Unset,
     ) -> _CandidateMicroBatch[ForwardInputsT]:
         dp_rank, dp_size = self._dp_rank_and_size()
         remaining, min_width, granularity = _wave_geometry(len(items), start, dp_size)
@@ -4005,12 +4030,27 @@ def candidate(width: int) -> _CandidateMicroBatch[ForwardInputsT]:
                 refusal_prefix = (
                     "smallest DP microbatch is predicted to exceed available memory"
                 )
-                found = self._find_admissible_forward(
-                    list(_flatten(local_inputs)),
-                    checkpoint=checkpoint,
-                    refusal_prefix=refusal_prefix,
-                )
-                agreed = self._all_ranks_true(not isinstance(found, _ForwardRefusal))
+                admission_error: BaseException | None = None
+                try:
+                    found = self._find_admissible_forward(
+                        list(_flatten(local_inputs)),
+                        checkpoint=checkpoint,
+                        refusal_prefix=refusal_prefix,
+                    )
+                except BaseException as exc:
+                    admission_error, found = exc, None
+                try:
+                    agreed = self._all_ranks_true(
+                        admission_error is None
+                        and not isinstance(found, _ForwardRefusal)
+                    )
+                except BaseException:
+                    if admission_error is not None:
+                        raise admission_error
+                    raise
+                if admission_error is not None:
+                    raise admission_error
+                assert found is not None
                 if isinstance(found, _ForwardRefusal):
                     self._snapshot_planning_telemetry(found.plan, found.check)
                     raise found.error("forward_micro_batches")
@@ -4020,7 +4060,8 @@ def candidate(width: int) -> _CandidateMicroBatch[ForwardInputsT]:
                         context="forward_micro_batches",
                         message=(
                             f"{refusal_prefix} on another DP rank, which was "
-                            "unable to find a feasible split for its share"
+                            "unable to complete admission or find a feasible split "
+                            "for its share"
                         ),
                         packed_tokens=first.plan.packed_tokens,
                         logical_tokens=first.plan.logical_tokens,
@@ -4952,18 +4993,36 @@ def _memory_check_required(
         *,
         sync_across_dp: bool = False,
     ) -> _MemoryCheck:
-        available = self._available_memory_bytes()
         if dist.is_available() and dist.is_initialized():
             group = None if sync_across_dp else self._forward_memory_group()
             values = torch.tensor(
-                [float(required), float(available)],
+                [float(required), 0.0],
                 device=self.device if self.device.type == "cuda" else "cpu",
                 dtype=torch.float64,
             )
             dist.all_reduce(values[0], op=dist.ReduceOp.MAX, group=group)
-            dist.all_reduce(values[1], op=dist.ReduceOp.MIN, group=group)
             required = int(values[0].item())
-            available = int(values[1].item())
+            error: BaseException | None = None
+            try:
+                available = self._available_memory_bytes()
+            except BaseException as exc:
+                error, available = exc, -1
+            try:
+                # A healthy communicator carries local failure to every peer
+                # in the existing MIN. This cannot repair a poisoned backend.
+                values[1] = available
+                dist.all_reduce(values[1], op=dist.ReduceOp.MIN, group=group)
+                available = int(values[1].item())
+            except BaseException:
+                if error is not None:
+                    raise error
+                raise
+            if error is not None:
+                raise error
+            if available < 0:
+                raise RuntimeError("Memory admission failed on another rank")
+        else:
+            available = self._available_memory_bytes()
         return _MemoryCheck(
             estimated_required_bytes=required,
             available_bytes=available,

tests/unit/test_trainer_rank_cuda_budget.py

@@ -1,5 +1,9 @@
 """Native allocator budget contracts; CPU counters, not CUDA qualification."""
 
+import asyncio
+from contextlib import nullcontext
+from types import SimpleNamespace
+from typing import cast
 from unittest.mock import Mock
 
 import pytest
@@ -178,3 +182,228 @@ def reduce(value, op, group):
         110,
         False,
     )
+
+
+@pytest.mark.parametrize("failed_locally", [False, True])
+@pytest.mark.parametrize(
+    "error_type", [RuntimeError, KeyboardInterrupt, SystemExit, asyncio.CancelledError]
+)
+def test_admission_failure_sentinel_stops_healthy_peers(
+    budget, monkeypatch, failed_locally, error_type
+):
+    rank, stats = budget
+    stats["active_bytes.all.current"] = 80
+    monkeypatch.setattr(_impl.dist, "is_available", lambda: True)
+    monkeypatch.setattr(_impl.dist, "is_initialized", lambda: True)
+    monkeypatch.setattr(rank, "_forward_memory_group", lambda: None)
+    tensor = torch.tensor
+    monkeypatch.setattr(
+        torch, "tensor", lambda values, **kwargs: tensor(values, dtype=kwargs["dtype"])
+    )
+    original = error_type("recoverable local budget read failure")
+    events = []
+
+    def read(_):
+        events.append("sample_failure")
+        raise original
+
+    monkeypatch.setattr(torch.cuda, "mem_get_info", read)
+    if not failed_locally:
+        monkeypatch.setattr(torch.cuda, "mem_get_info", lambda _: (200, 1000))
+
+    def reduce(value, op, group):
+        events.append((op, float(value.item())))
+        value.fill_(150 if op == _impl.dist.ReduceOp.MAX else -1)
+
+    monkeypatch.setattr(_impl.dist, "all_reduce", reduce)
+    with pytest.raises(error_type if failed_locally else RuntimeError) as caught:
+        for required in (0, 75):
+            rank._memory_check_required(required)
+    assert (
+        caught.value is original
+        if failed_locally
+        else str(caught.value) == "Memory admission failed on another rank"
+    )
+    assert events == [
+        (_impl.dist.ReduceOp.MAX, 0),
+        *(["sample_failure"] if failed_locally else []),
+        (_impl.dist.ReduceOp.MIN, -1 if failed_locally else 170),
+    ]
+
+
+@pytest.mark.parametrize("stage", ["assignment", "min", "result"])
+@pytest.mark.parametrize("failed_locally", [False, True])
+def test_secondary_collective_failure_preserves_local_primary(
+    budget, monkeypatch, stage, failed_locally
+):
+    rank, stats = budget
+    stats["active_bytes.all.current"] = 80
+    monkeypatch.setattr(_impl.dist, "is_available", lambda: True)
+    monkeypatch.setattr(_impl.dist, "is_initialized", lambda: True)
+    monkeypatch.setattr(rank, "_forward_memory_group", lambda: None)
+    original = RuntimeError("original local budget read error")
+    secondary = RuntimeError("secondary CUDA or communicator failure")
+
+    class Values:
+        def __init__(self, values, **kwargs):
+            self.values = values
+
+        def __getitem__(self, index):
+            def item():
+                if stage == "result" and index == 1:
+                    raise secondary
+                return self.values[index]
+
+            return SimpleNamespace(
+                item=item, fill_=lambda x: self.values.__setitem__(index, x)
+            )
+
+        def __setitem__(self, index, value):
+            if stage == "assignment":
+                raise secondary
+            self.values[index] = value
+
+    monkeypatch.setattr(torch, "tensor", Values)
+    monkeypatch.setattr(torch.cuda, "mem_get_info", Mock(side_effect=original))
+    if not failed_locally:
+        monkeypatch.setattr(torch.cuda, "mem_get_info", lambda _: (200, 1000))
+
+    def reduce(value, op, group):
+        if op == _impl.dist.ReduceOp.MAX:
+            value.fill_(150)
+        elif stage == "min":
+            raise secondary
+
+    monkeypatch.setattr(_impl.dist, "all_reduce", reduce)
+    with pytest.raises(RuntimeError) as caught:
+        rank._memory_check_required(0)
+    assert caught.value is (original if failed_locally else secondary)
+
+
+@pytest.mark.parametrize("distributed", [False, True])
+@pytest.mark.parametrize("available", [10, 210])
+def test_final_selection_uses_pure_fresh_budget_and_original_demand(
+    budget, monkeypatch, distributed, available
+):
+    rank, stats = budget
+    stats["active_bytes.all.current"] = 80
+    plan = cast(
+        _impl._FlatForwardPlan, SimpleNamespace(packed_tokens=64, logical_tokens=64)
+    )
+    stale = _impl._MemoryCheck(
+        estimated_required_bytes=192, available_bytes=256, fits=True
+    )
+    candidate = _impl._CandidateMicroBatch([], (), plan, stale, 64, 1, False)
+    monkeypatch.setattr(rank, "_search_next_micro_batch", lambda *a, **k: candidate)
+    monkeypatch.setattr(
+        rank,
+        "_estimate_required_memory_bytes_from_values",
+        Mock(side_effect=AssertionError("must keep original selected demand")),
+    )
+    monkeypatch.setattr(rank, "_snapshot_planning_telemetry", Mock())
+    monkeypatch.setattr(
+        torch.cuda, "mem_get_info", Mock(return_value=(available + 30, 1000))
+    )
+    monkeypatch.setattr(_impl.dist, "is_available", lambda: distributed)
+    monkeypatch.setattr(_impl.dist, "is_initialized", lambda: distributed)
+    tensor = torch.tensor
+    monkeypatch.setattr(
+        torch, "tensor", lambda values, **kwargs: tensor(values, dtype=kwargs["dtype"])
+    )
+    calls = []
+    monkeypatch.setattr(
+        _impl.dist,
+        "all_reduce",
+        lambda value, op, group: calls.append((op, value.item(), group)),
+    )
+    if available < 192:
+        with pytest.raises(_impl.TrainerRankMemoryError) as caught:
+            rank._select_next_micro_batch([], 0)
+        assert caught.value.predicted_peak_bytes == 192
+        assert caught.value.usable_limit_bytes == available
+    else:
+        selected = rank._select_next_micro_batch([], 0)
+        assert selected.check.estimated_required_bytes == 192
+        assert selected.check.available_bytes == available and selected.check.fits
+        assert selected.plan is plan
+    assert stale.available_bytes == 256
+    assert calls == (
+        [
+            (_impl.dist.ReduceOp.MAX, 192, None),
+            (_impl.dist.ReduceOp.MIN, available, None),
+        ]
+        if distributed
+        else []
+    )
+    torch.cuda.empty_cache.assert_not_called()
+    torch.cuda.mem_get_info.assert_called_once_with(rank.device)
+
+
+def test_available_sample_follows_required_collective(budget, monkeypatch):
+    rank, stats = budget
+    stats["active_bytes.all.current"] = 80
+    free = [100]
+    events = []
+
+    def read(_):
+        events.append("sample")
+        return free[0], 1000
+
+    monkeypatch.setattr(torch.cuda, "mem_get_info", read)
+    monkeypatch.setattr(_impl.dist, "is_available", lambda: True)
+    monkeypatch.setattr(_impl.dist, "is_initialized", lambda: True)
+    tensor = torch.tensor
+    monkeypatch.setattr(
+        torch, "tensor", lambda values, **kwargs: tensor(values, dtype=kwargs["dtype"])
+    )
+
+    def reduce(value, op, group):
+        events.append(op)
+        if op == _impl.dist.ReduceOp.MAX:
+            value.fill_(60)
+            free[0] = 50
+
+    monkeypatch.setattr(_impl.dist, "all_reduce", reduce)
+    check = rank._memory_check_required(0, sync_across_dp=True)
+    assert (check.estimated_required_bytes, check.available_bytes, check.fits) == (
+        60,
+        20,
+        False,
+    )
+    assert events == [_impl.dist.ReduceOp.MAX, "sample", _impl.dist.ReduceOp.MIN]
+    torch.cuda.empty_cache.assert_not_called()
+
+
+def test_execution_failure_retains_final_admission_without_resampling(
+    budget, monkeypatch
+):
+    rank, stats = budget
+    stats["active_bytes.all.current"] = 80
+    free = [100]
+    read = Mock(side_effect=lambda _: (free[0], 1000))
+    monkeypatch.setattr(torch.cuda, "mem_get_info", read)
+    admitted = rank._memory_check_required(50)
+    assert admitted.fits and admitted.available_bytes == 70
+    original = torch.cuda.OutOfMemoryError("synthetic later execution failure")
+
+    def execute(_):
+        free[0] = 10
+        raise original
+
+    monkeypatch.setattr(rank, "_execute_flat_plan", execute)
+    monkeypatch.setattr(rank, "_telemetry_signature", lambda _: {})
+    monkeypatch.setattr(rank, "_telemetry_plan_signature", lambda _: {})
+    monkeypatch.setattr(_impl, "_telemetry_phase", lambda *a, **k: nullcontext())
+    monkeypatch.setattr(torch.cuda, "synchronize", Mock())
+    monkeypatch.setattr(torch.cuda, "reset_peak_memory_stats", Mock())
+    with pytest.raises(_impl.TrainerRankMemoryError) as caught:
+        rank._run_flat_plan_with_memory_tracking(
+            SimpleNamespace(packed_tokens=1, logical_tokens=1),
+            check=admitted,
+            context="CPU final-admission witness",
+        )
+    assert caught.value.__cause__ is original
+    assert caught.value.usable_limit_bytes == admitted.available_bytes == 70
+    read.assert_called_once_with(rank.device)
+    assert rank._available_memory_bytes() == 0
+    torch.cuda.empty_cache.assert_not_called()

tests/unit/test_trainer_rank_split_peak.py

@@ -294,11 +294,21 @@ def profiled(**_):
             patch.setattr(rank, "_all_ranks_true", agree)
             patch.setattr(rank, "_all_ranks_have_memory_profile", profiled)
 
+            search = rank._search_next_micro_batch
+            search_finished = []
+
+            def searched(*args, **kwargs):
+                result = search(*args, **kwargs)
+                search_finished.append(True)
+                return result
+
+            patch.setattr(rank, "_search_next_micro_batch", searched)
+
             def reduce(value, op, group):
                 trace.append(("global" if group is None else "local", str(op)))
                 if group is None:
                     value.fill_(
-                        max(value.item(), 200)
+                        max(value.item(), 100 if search_finished else 200)
                         if op == tr.dist.ReduceOp.MAX
                         else min(value.item(), 100)
                     )
@@ -314,6 +324,10 @@ def reduce(value, op, group):
                 tr._SplitForwardPlan if dp_rank == 0 else tr._FlatForwardPlan,
             )
             traces.append([event for event in trace if event[0] == "global"])
+            assert traces[-1][-2:] == [
+                ("global", str(tr.dist.ReduceOp.MAX)),
+                ("global", str(tr.dist.ReduceOp.MIN)),
+            ]
     assert traces[0] == traces[1]