diff --git a/.github/workflows/prek.yml b/.github/workflows/prek.yml index 6c8a29d72..1f8f27e54 100644 --- a/.github/workflows/prek.yml +++ b/.github/workflows/prek.yml @@ -226,6 +226,8 @@ jobs: tests/unit/test_trainer_rank_validation.py \ tests/unit/test_trainer_rank_weird_shapes.py \ tests/unit/test_trainer_rank_split.py \ + tests/unit/test_trainer_rank_cache_recovery.py::test_dense_cp_exact_demand_fits_after_recovery \ + tests/unit/test_trainer_rank_cache_recovery.py::test_dense_cp_exact_demand_refuses_after_recovery \ tests/acceptance/trainer_rank_planner \ tests/integration/megatron/test_sft_packing.py::test_sft_packing_preserves_training_targets \ tests/integration/megatron/model_support/test_dispatcher_graph_retention.py \ @@ -240,6 +242,8 @@ jobs: - name: Run unit tests run: | uv run --no-sync pytest --nbval --current-env --tb=short tests/unit \ + --deselect=tests/unit/test_trainer_rank_cache_recovery.py::test_dense_cp_exact_demand_fits_after_recovery \ + --deselect=tests/unit/test_trainer_rank_cache_recovery.py::test_dense_cp_exact_demand_refuses_after_recovery \ --ignore=tests/unit/test_megatron_reference_logprobs.py \ --ignore=tests/unit/test_moe_routing_replay.py \ --ignore=tests/unit/test_moe_routing_real_path.py \ diff --git a/src/art/trainer_rank/_impl.py b/src/art/trainer_rank/_impl.py index 6e2fc5a94..4ee194dfb 100644 --- a/src/art/trainer_rank/_impl.py +++ b/src/art/trainer_rank/_impl.py @@ -26,6 +26,7 @@ import struct import threading import time +import traceback from types import TracebackType from typing import ( TYPE_CHECKING, @@ -566,6 +567,18 @@ class TrainerRankSlotStateError(RuntimeError): pass +@dataclass +class _CacheRecoveryState: + # Local, lifetime measurements; reductions never replace these ledgers. + cost: float = 0.0 + work: float = 0.0 + high: float = 0.0 + first_consumed: bool = False + invalid: bool = False + owner: object | None = None + lock: Any = dataclass_field(default_factory=threading.Lock) + + @dataclass(frozen=True) class _MemoryCheck: estimated_required_bytes: int @@ -1043,9 +1056,9 @@ def _memory_error( @dataclass(frozen=True) class _ForwardRefusal: - """Why no admissible plan was found. ``plan`` is the unsplit call.""" + """Why admission failed, with the last relevant local plan and check.""" - plan: _FlatForwardPlan + plan: _AnyForwardPlan check: _MemoryCheck message: str @@ -1506,6 +1519,7 @@ def memory_field(name: str, default: Any = None) -> Any: self._hybridep_graph_tracking = False self._hybridep_buffer_id: int | None = None self._hybridep_rows_high_water = 0 + self._cache_recovery_state = _CacheRecoveryState() self._memory_profiles: dict[_MemorySignature, _MemoryProfile] = {} self._split_memory_floors: dict[bytes, int] = {} self._split_memory_floor_status = "not_observed" @@ -2511,35 +2525,43 @@ def _execute_admitted_plan( def _execute_split_plan_with_memory_tracking( self, plan: _SplitForwardPlan, *, check: _MemoryCheck, context: str ) -> tuple[list[AnyForwardOutput], int | None, int]: - baseline, peak = None, 0 - merged: list[AnyForwardOutput | None] = [None] * plan.request_count - for ordinal, (subforward, indices) in enumerate( - zip(plan.subforwards, plan.request_indices, strict=True) - ): - try: - outputs, child_baseline = self._run_flat_plan_with_memory_tracking( - subforward, check=check, context=context - ) - if child_baseline is not None: - if baseline is None: - baseline = child_baseline - peak = max(peak, int(torch.cuda.max_memory_allocated(self.device))) - except TrainerRankMemoryError as error: - # Model execution already began, so no replanning is possible - # and the caller must not mistake this for an up-front refusal. - raise TrainerRankPartialExecutionError( - f"{context}: subforward {ordinal + 1} of " - f"{plan.subforward_count} failed during execution " - f"({ordinal} of {plan.subforward_count} completed). {error}", - predicted_peak_bytes=error.predicted_peak_bytes, - usable_limit_bytes=error.usable_limit_bytes, - suggestion=error.suggestion, - ) from error - for index, output in zip(indices, outputs, strict=True): - merged[index] = output - if any(output is None for output in merged): - raise AssertionError("split execution did not cover every request") - return cast(list[AnyForwardOutput], merged), baseline, peak + state = self._recovery_state() + work_before = state.work + try: + baseline, peak = None, 0 + merged: list[AnyForwardOutput | None] = [None] * plan.request_count + for ordinal, (subforward, indices) in enumerate( + zip(plan.subforwards, plan.request_indices, strict=True) + ): + try: + outputs, child_baseline = self._run_flat_plan_with_memory_tracking( + subforward, check=check, context=context + ) + if child_baseline is not None: + if baseline is None: + baseline = child_baseline + peak = max( + peak, int(torch.cuda.max_memory_allocated(self.device)) + ) + except TrainerRankMemoryError as error: + # Model execution already began, so no replanning is possible + # and the caller must not mistake this for an up-front refusal. + raise TrainerRankPartialExecutionError( + f"{context}: subforward {ordinal + 1} of " + f"{plan.subforward_count} failed during execution " + f"({ordinal} of {plan.subforward_count} completed). {error}", + predicted_peak_bytes=error.predicted_peak_bytes, + usable_limit_bytes=error.usable_limit_bytes, + suggestion=error.suggestion, + ) from error + for index, output in zip(indices, outputs, strict=True): + merged[index] = output + if any(output is None for output in merged): + raise AssertionError("split execution did not cover every request") + return cast(list[AnyForwardOutput], merged), baseline, peak + except BaseException: + state.work = work_before + raise def _plan_admissible_forward( self, @@ -2548,19 +2570,22 @@ def _plan_admissible_forward( checkpoint: AdapterSelection, context: str, ) -> tuple[_AnyForwardPlan, _MemoryCheck]: - """Plan one forward (splitting if needed), recording telemetry, or raise.""" - - found = self._find_admissible_forward( - requests, - checkpoint=checkpoint, - refusal_prefix="forward is predicted to exceed available memory", + # DP-local recovery may retry asymmetrically; checkpoint setup is WORLD. + self._ensure_checkpoint_slots_for(requests, checkpoint=checkpoint) + result = self._recover_admission( + lambda: self._find_admissible_forward( + requests, + checkpoint=checkpoint, + refusal_prefix="forward is predicted to exceed available memory", + ensure_slots=False, + ), + lambda value: value, + lambda value, check: (value[0], check), + context=context, + sync_across_dp=False, ) - if isinstance(found, _ForwardRefusal): - self._snapshot_planning_telemetry(found.plan, found.check) - raise found.error(context) - plan, check = found - self._snapshot_planning_telemetry(plan, check) - return plan, check + self._snapshot_planning_telemetry(*result) + return result def _find_admissible_forward( self, @@ -2568,6 +2593,7 @@ def _find_admissible_forward( *, checkpoint: AdapterSelection, refusal_prefix: str, + ensure_slots: bool = True, ) -> tuple[_AnyForwardPlan, _MemoryCheck] | _ForwardRefusal: """Find an admissible plan: unsplit first, then the bounded split ladder. @@ -2579,12 +2605,14 @@ def _find_admissible_forward( refusal worded as "unable to find a feasible split": the search is bounded, so this is not a claim that none exists. - Checkpoint slots are ensured exactly once, up front; everything after - plans with ``ensure_slots=False`` so the number of collectives this - rank performs does not depend on its (DP-local) inputs. + Checkpoint slots are ensured up front unless the caller already did + so before entering a DP-local retry loop. Everything after plans with + ``ensure_slots=False`` so the number of collectives this rank performs + does not depend on its (DP-local) inputs. """ - self._ensure_checkpoint_slots_for(requests, checkpoint=checkpoint) + if ensure_slots: + self._ensure_checkpoint_slots_for(requests, checkpoint=checkpoint) plan = self._plan_flat_forward( requests, checkpoint=checkpoint, ensure_slots=False ) @@ -3882,23 +3910,13 @@ def _select_next_micro_batch( *, 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, + return self._recover_admission( + lambda: self._search_next_micro_batch(items, start, checkpoint=checkpoint), + lambda value: (value.plan, value.check), + lambda value, check: replace(value, check=check), + context="forward_micro_batches", 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, @@ -3906,7 +3924,7 @@ def _search_next_micro_batch( start: int, *, checkpoint: AdapterSelection = Unset, - ) -> _CandidateMicroBatch[ForwardInputsT]: + ) -> _CandidateMicroBatch[ForwardInputsT] | _ForwardRefusal: dp_rank, dp_size = self._dp_rank_and_size() remaining, min_width, granularity = _wave_geometry(len(items), start, dp_size) if min_width <= 0: @@ -4133,9 +4151,12 @@ def candidate(width: int) -> _CandidateMicroBatch[ForwardInputsT]: except BaseException as exc: admission_error, found = exc, None try: - agreed = self._all_ranks_true( - admission_error is None - and not isinstance(found, _ForwardRefusal) + outcome = self._admission_outcome( + 0 + if admission_error is not None + else 1 + if isinstance(found, _ForwardRefusal) + else 2 ) except BaseException: if admission_error is not None: @@ -4144,21 +4165,16 @@ def candidate(width: int) -> _CandidateMicroBatch[ForwardInputsT]: if admission_error is not None: raise admission_error assert found is not None + if outcome == 0: + raise RuntimeError("Memory admission failed on another DP rank") if isinstance(found, _ForwardRefusal): - self._snapshot_planning_telemetry(found.plan, found.check) - raise found.error("forward_micro_batches") - if not agreed: - self._snapshot_planning_telemetry(first.plan, first.check) - raise _memory_error( - context="forward_micro_batches", - message=( - f"{refusal_prefix} on another DP rank, which was " - "unable to complete admission or find a feasible split " - "for its share" - ), - packed_tokens=first.plan.packed_tokens, - logical_tokens=first.plan.logical_tokens, - check=first.check, + return found + if outcome == 1: + return _ForwardRefusal( + found[0], + found[1], + f"{refusal_prefix} on another DP rank, which was unable " + "to find a feasible split for its share", ) split_plan, split_check = found return _CandidateMicroBatch( @@ -4753,6 +4769,7 @@ def _run_flat_plan_with_memory_tracking( torch.cuda.reset_peak_memory_stats(self.device) else: baseline = None + started = self._recovery_clock() if baseline is not None else None try: with _telemetry_phase( "forward", @@ -4774,10 +4791,17 @@ def _run_flat_plan_with_memory_tracking( # can enter collectives that successful peers never reach. check=check, ) from exc + finished = self._recovery_clock() if baseline is not None else None + seconds = None if started is None or finished is None else finished - started if baseline is not None: self._update_peak_memory_profile( plan, baseline, int(torch.cuda.memory_allocated(self.device)) ) + if seconds is not None and plan.packed_tokens > 0: + try: + self._record_recovery_work(context, seconds) + except Exception: + self._recovery_state().invalid = True return outputs, baseline def _update_peak_memory_profile( @@ -5115,6 +5139,337 @@ def _memory_check( ) return self._memory_check_required(required, sync_across_dp=sync_across_dp) + def _admission_outcome(self, local: int) -> int: + """Existing world fallback MIN: error=0, refusal=1, fit=2.""" + if not (dist.is_available() and dist.is_initialized()): + return local + value = torch.tensor( + local, + device=self.device if self.device.type == "cuda" else "cpu", + dtype=torch.int32, + ) + dist.all_reduce(value, op=dist.ReduceOp.MIN) + return int(value.item()) + + def _recover_admission( + self, + search: Callable[[], Any], + describe: Callable[[Any], tuple[_AnyForwardPlan, _MemoryCheck]], + update: Callable[[Any, _MemoryCheck], Any], + *, + context: str, + sync_across_dp: bool, + ) -> Any: + """Pure search, at most one smaller-plan refresh, then one recovery.""" + original: TrainerRankMemoryError | None = None + refused: _ForwardRefusal | None = None + + def finish(value: Any) -> Any: + nonlocal refused + if isinstance(value, _ForwardRefusal): + refused = value + return None + plan, check = describe(value) + if sync_across_dp: + check = self._memory_check_required( + check.estimated_required_bytes, sync_across_dp=True + ) + if check.fits: + return update(value, check) + refused = _ForwardRefusal( + plan, check, "selected plan exceeds freshly sampled available memory" + ) + return None + + value = search() + result = finish(value) + if result is not None: + return result + assert refused is not None + original = refused.error(context) + state = self._recovery_state() + started = self._recovery_clock() + owner = object() + with state.lock: + if state.owner is None: + state.owner = owner + primary: BaseException | None = None + try: + if not isinstance(value, _ForwardRefusal): + # A formerly fitting width is not proof that the minimum cannot fit. + value = search() + result = finish(value) + if result is not None: + return result + if not isinstance(value, _ForwardRefusal): + # Counters moved again: do not loop or reclaim for a large width. + assert refused is not None + self._snapshot_planning_telemetry(refused.plan, refused.check) + raise refused.error(context) from original + assert refused is not None + if self._try_cache_recovery( + refused.check, + sync_across_dp=sync_across_dp, + owner=owner, + started=started, + ): + value = search() + result = finish(value) + if result is not None: + return result + assert refused is not None + self._snapshot_planning_telemetry(refused.plan, refused.check) + latest = refused.error(context) + raise latest from original + except BaseException as exc: + primary = exc + raise + finally: + # Includes control, release, resample and repeated search; no idle. + try: + finished = self._recovery_clock() + elapsed = ( + None if started is None or finished is None else finished - started + ) + with state.lock: + if ( + elapsed is None + or not math.isfinite(elapsed) + or elapsed <= 0 + or not math.isfinite(state.cost + elapsed) + ): + state.invalid = True + else: + state.cost += elapsed + state.high = max(state.high, elapsed) + except Exception: + state.invalid = True + except BaseException as exc: + state.invalid = True + if primary is None: + primary = exc + raise + finally: + try: + with state.lock: + if state.owner is owner: + state.owner = None + except Exception: + state.invalid = True + except BaseException: + state.invalid = True + if primary is None: + raise + + def _recovery_clock(self) -> float | None: + try: + value = time.perf_counter() + if type(value) is float and math.isfinite(value): + return value + except Exception: + pass + self._recovery_state().invalid = True + return None + + def _record_recovery_work(self, context: str, seconds: float) -> None: + if context not in ("forward_micro_batches", "dp_rank_forward"): + return + state = self._recovery_state() + try: + with state.lock: + if ( + not math.isfinite(seconds) + or seconds < 0 + or not math.isfinite(state.work + seconds) + ): + state.invalid = True + else: + state.work += seconds + except Exception: + state.invalid = True + + def _recovery_state(self) -> _CacheRecoveryState: + state = getattr(self, "_cache_recovery_state", None) + if state is None: + state = self._cache_recovery_state = _CacheRecoveryState() + return state + + def _recovery_reduce( + self, + values: list[float], + *, + op: Literal["MAX", "MIN", "SUM"], + sync_across_dp: bool, + ) -> list[float]: + if not (dist.is_available() and dist.is_initialized()): + return values + tensor = torch.tensor( + values, + device=self.device if self.device.type == "cuda" else "cpu", + dtype=torch.float64, + ) + dist.all_reduce( + tensor, + op=getattr(dist.ReduceOp, op), + group=None if sync_across_dp else self._forward_memory_group(), + ) + return [float(value) for value in tensor.tolist()] + + @staticmethod + def _memory_error_with_reduction_note( + error: BaseException, exchange_error: BaseException | None + ) -> BaseException: + # Raise outside the exchange handler to preserve the local error's chain. + # A secondary poisoned-communicator failure is diagnostic, not the primary. + if exchange_error is not None and exchange_error is not error: + try: + BaseException.add_note( + error, + "Secondary memory reduction failure:\n" + + "".join(traceback.format_exception(exchange_error)), + ) + except BaseException: + # Exception rendering must never replace the original failure. + pass + return error + + def _try_cache_recovery( + self, + check: _MemoryCheck, + *, + sync_across_dp: bool, + owner: object, + started: float | None, + ) -> bool: + state = self._recovery_state() + now = self._recovery_clock() + elapsed = None if now is None or started is None else now - started + with state.lock: + invalid = ( + state.invalid + or state.owner is not owner + or elapsed is None + or not math.isfinite(elapsed) + or elapsed < 0 + ) + projected = state.cost if elapsed is None else state.cost + elapsed + invalid |= any( + not math.isfinite(value) or value < 0 + for value in ( + projected, + state.work, + state.high, + projected + state.high, + ) + ) + local_cost = [0.0, 0.0] if invalid else [projected, state.high] + # SUM costs deliberately overcharges parallel ranks; unlike MAX of + # lifetime costs, it cannot miss episodes with different slow ranks. + costs = self._recovery_reduce( + local_cost, op="SUM", sync_across_dp=sync_across_dp + ) + invalid |= any(not math.isfinite(value) for value in (*costs, sum(costs))) + values = self._recovery_reduce( + [ + float(check.estimated_required_bytes), + 0.0 if invalid else state.work, + float(state.first_consumed), + float(invalid), + ], + op="MAX", + sync_across_dp=sync_across_dp, + ) + required = int(values[0]) + state.first_consumed = bool(values[2]) + state.invalid |= bool(values[3]) + if state.invalid: + return False + error: BaseException | None = None + needed = False + cap_blocks = False + try: + available = self._available_memory_bytes() + if ( + available < required + and self.device.type == "cuda" + and torch.cuda.is_available() + and torch.cuda.get_allocator_backend() == "native" + ): + free, total = torch.cuda.mem_get_info(self.device) + needed = int(free) < required + int(total * _MEMORY_RESERVE_FRACTION) + if os.environ.get(_TEST_HOOKS_ENV) == "1": + limit = os.environ.get(_TEST_MEMORY_LIMIT_ENV) + if limit: + cap_blocks = required > max( + 0, + int(limit) - int(torch.cuda.memory_allocated(self.device)), + ) + except BaseException as exc: + error, available = exc, -1 + exchange_error: BaseException | None = None + try: + sampled = self._recovery_reduce( + [ + float(available), + -float(needed), + float(not cap_blocks), + float(not state.invalid), + ], + op="MIN", + sync_across_dp=sync_across_dp, + ) + except BaseException as exc: + if error is None: + raise + exchange_error = exc + if error is not None: + raise self._memory_error_with_reduction_note(error, exchange_error) + if sampled[0] < 0: + raise RuntimeError("Memory recovery sampling failed on another rank") + state.invalid |= not bool(sampled[3]) + if required <= sampled[0]: + return True + if sampled[1] == 0 or sampled[2] == 0 or state.invalid: + return False + if state.first_consumed and not ( + costs[1] > 0 and sum(costs) <= 0.05 * values[1] + ): + return False + # Every peer enters this block. Only locally deficient native ranks call + # the process-wide allocator operation. Test-only caps cannot trigger it. + state.first_consumed = True + attempted = False + try: + if needed: + # The control exchange can change allocator state. Check the + # physical condition again immediately before the sole call. + free, total = torch.cuda.mem_get_info(self.device) + if int(free) < required + int(total * _MEMORY_RESERVE_FRACTION): + attempted = True + torch.cuda.empty_cache() + available = self._available_memory_bytes() + except BaseException as exc: + error, available = exc, -1 + exchange_error: BaseException | None = None + try: + sampled = self._recovery_reduce( + [float(available), -float(attempted)], + op="MIN", + sync_across_dp=sync_across_dp, + ) + except BaseException as exc: + if error is None: + raise + exchange_error = exc + else: + state.first_consumed |= bool(sampled[1]) + if error is not None: + raise self._memory_error_with_reduction_note(error, exchange_error) + if sampled[0] < 0: + raise RuntimeError("Memory recovery failed on another rank") + # Rebuild pure search caches even when this fresh sample decreased. + return True + def _memory_check_required( self, required: int, @@ -5135,18 +5490,19 @@ def _memory_check_required( available = self._available_memory_bytes() except BaseException as exc: error, available = exc, -1 + exchange_error: BaseException | None = None 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 + except BaseException as exc: + if error is None: + raise + exchange_error = exc if error is not None: - raise error + raise self._memory_error_with_reduction_note(error, exchange_error) if available < 0: raise RuntimeError("Memory admission failed on another rank") else: @@ -5341,23 +5697,10 @@ def _available_memory_bytes(self) -> int: return 1 << 60 free, total = torch.cuda.mem_get_info(self.device) if torch.cuda.get_allocator_backend() == "native": - stats = torch.cuda.memory_stats(self.device) - allocated = stats.get("allocated_bytes.all.current") - active = stats.get("active_bytes.all.current") - reserved = stats.get("reserved_bytes.all.current") - if ( - type(allocated) is int - and type(active) is int - and type(reserved) is int - and 0 <= allocated <= active <= reserved - ): - # Pending frees remain active until ordinary event collection. - # This excludes them, not split/private-pool incompatibilities. - reusable_reserved = reserved - active - else: - # Incomplete native counters cannot establish reusable cache. - allocated = int(torch.cuda.memory_allocated(self.device)) - reusable_reserved = 0 + # Cached bytes are not physical free memory. This sample does not + # reserve memory for execution or the caller's later backward. + allocated = int(torch.cuda.memory_allocated(self.device)) + reusable_reserved = 0 else: # Preserve the previous, unqualified policy for other backends. allocated = int(torch.cuda.memory_allocated(self.device)) diff --git a/tests/unit/test_trainer_rank_active_memory.py b/tests/unit/test_trainer_rank_active_memory.py index d9c1d710a..25c1e8c89 100644 --- a/tests/unit/test_trainer_rank_active_memory.py +++ b/tests/unit/test_trainer_rank_active_memory.py @@ -187,16 +187,16 @@ def test_warm_admission_rechecks_current_residency(monkeypatch): ) monkeypatch.delenv("ART_TRAINER_RANK_TEST_HOOKS", raising=False) # Fresh memory accounting observes newly resident state without discarding - # a valid incremental profile; cached free blocks remain reusable. + # a valid incremental profile; cache alone is not physical availability. for allocated, reserved, fits in [ (10_000, 10_000, True), (90_000, 90_000, False), - (10_000, 90_000, True), + (10_000, 90_000, False), ]: state.update(allocated=allocated, reserved=reserved) check = rank._memory_check(plan) assert check.estimated_required_bytes == required - assert check.available_bytes == total - allocated - 3000 + assert check.available_bytes == total - reserved - 3000 assert check.fits == fits assert rank._memory_profiles[plan.signature] == profile diff --git a/tests/unit/test_trainer_rank_cache_recovery.py b/tests/unit/test_trainer_rank_cache_recovery.py new file mode 100644 index 000000000..3b6b08f8c --- /dev/null +++ b/tests/unit/test_trainer_rank_cache_recovery.py @@ -0,0 +1,899 @@ +"""Actual admission methods with scalar allocator/search/clock facades, no CUDA.""" + +from contextlib import nullcontext +import math +import os +import types +from typing import cast +import unittest +from unittest.mock import patch + +from art.trainer_rank import _impl + +Refusal = _impl.TrainerRankMemoryError +Partial = _impl.TrainerRankPartialExecutionError +OOM = _impl.torch.cuda.OutOfMemoryError + + +class Scalar: + def __init__(self, t, i): + self.t, self.i = t, i + + def item(self): + return self.t.data[self.i] + + +class Tensor: + def __init__(self, data, **kw): + self.data = data if isinstance(data, list) else [data] + + def __getitem__(self, i): + return Scalar(self, i) + + def __setitem__(self, i, v): + self.data[i] = v + + def tolist(self): + return self.data + + def item(self): + return self.data[0] + + +class CUDA: + OutOfMemoryError = OOM + + def synchronize(self, device): + self.events.append("sync") + + def reset_peak_memory_stats(self, device): + self.events.append("reset") + + def max_memory_allocated(self, device): + return self.allocated + + def __init__(self): + self.free = 40 + self.total = 1000 + self.allocated = 0 + self.backend = "native" + self.events = [] + self.failure = None + + def is_available(self): + return True + + def get_allocator_backend(self): + return self.backend + + def mem_get_info(self, device): + self.events.append("sample") + return self.free, self.total + + def memory_allocated(self, device): + return self.allocated + + def memory_reserved(self, device): + return self.allocated + + def empty_cache(self): + self.events.append("release") + if self.failure is not None: + raise self.failure + self.free = 200 + + +class Clock: + def __init__(self): + self.value = 0.0 + self.next = None + + def perf_counter(self): + if self.next is not None: + return self.next + self.value += 0.01 + return self.value + + +def plan(required=80): + return types.SimpleNamespace(packed_tokens=required, logical_tokens=required) + + +def fail(ns, required=80): + return ns["_ForwardRefusal"]( + plan(required), + ns["_MemoryCheck"](required, 10, False), + "smallest actual refusal", + ) + + +def success(ns, required=80): + return (plan(required), ns["_MemoryCheck"](required, 170, True)) + + +def run(q, items, *, sync=True): + calls = [] + + def search(): + calls.append(1) + value = items[min(len(calls) - 1, len(items) - 1)] + if isinstance(value, BaseException): + raise value + return value + + try: + value = q._recover_admission( + search, + lambda v: v, + lambda v, c: (v[0], c), + context="forward_micro_batches" if sync else "dp_rank_forward", + sync_across_dp=sync, + ) + except BaseException as e: + return None, e, len(calls) + return value, None, len(calls) + + +class TestRecovery(unittest.TestCase): + def make(self): + cuda = CUDA() + clock = Clock() + for name, value in dict( + torch=types.SimpleNamespace( + cuda=cuda, tensor=Tensor, float64="f64", int32="i32" + ), + time=clock, + dist=types.SimpleNamespace( + is_available=lambda: False, is_initialized=lambda: False + ), + _telemetry_phase=lambda *a, **kw: nullcontext(), + _TEST_HOOKS_ENV="CONTROL_ART_HOOK", + _TEST_MEMORY_LIMIT_ENV="CONTROL_ART_LIMIT", + ).items(): + patcher = patch.object(_impl, name, value) + patcher.start() + self.addCleanup(patcher.stop) + q = object.__new__(_impl.TrainerRank) + q.device = types.SimpleNamespace(type="cuda") + q._update_peak_memory_profile = lambda *a: None + q._execute_flat_plan = lambda p: [object() for _ in range(p.request_count)] + q._telemetry_signature = lambda p: {} + q._telemetry_plan_signature = lambda p: {} + q._snapshot_planning_telemetry = lambda *a: None + q._forward_memory_group = lambda: None + return ( + q, + cuda, + clock, + dict( + _MemoryCheck=_impl._MemoryCheck, _ForwardRefusal=_impl._ForwardRefusal + ), + ) + + def test_first_once_and_fresh_return(self): + q, c, k, n = self.make() + v, e, calls = run(q, [fail(n), success(n)]) + self.assertIsNone(e) + self.assertEqual(calls, 2) + self.assertEqual(c.events.count("release"), 1) + self.assertEqual(v[1].available_bytes, 170) + self.assertTrue(q._recovery_state().first_consumed) + + def test_small_fit_without_release(self): + q, c, k, n = self.make() + c.free = 200 + v, e, count = run(q, [success(n)]) + self.assertIsNone(e) + self.assertEqual(count, 1) + self.assertNotIn("release", c.events) + + def test_one_bounded_smaller_refresh(self): + q, c, k, n = self.make() + v, e, count = run(q, [success(n, 192), success(n, 5)]) + self.assertIsNone(e) + self.assertEqual(count, 2) + self.assertEqual(v[1].available_bytes, 10) + self.assertNotIn("release", c.events) + self.assertGreater(q._recovery_state().cost, 0) + self.assertFalse(q._recovery_state().first_consumed) + + def test_moving_final_counters_do_not_reclaim(self): + q, c, k, n = self.make() + v, e, count = run(q, [success(n, 192), success(n, 192)]) + self.assertIsInstance(e, Refusal) + self.assertEqual(count, 2) + self.assertNotIn("release", c.events) + + def test_final_refresh_then_exhaustion_then_one_recovery(self): + q, c, k, n = self.make() + v, e, count = run(q, [success(n, 192), fail(n), success(n)]) + self.assertIsNone(e) + self.assertEqual(count, 3) + self.assertEqual(c.events.count("release"), 1) + + def test_search_runtime_error_never_recovered(self): + q, c, k, n = self.make() + original = RuntimeError("peer runtime error") + v, e, count = run(q, [original]) + self.assertIs(e, original) + self.assertEqual(count, 1) + self.assertNotIn("release", c.events) + + def test_failed_release_original_and_consumed(self): + for typ in (RuntimeError, KeyboardInterrupt, SystemExit): + q, c, k, n = self.make() + original = typ("release error") + c.failure = original + v, e, count = run(q, [fail(n), success(n)]) + self.assertIs(e, original) + self.assertEqual(count, 1) + self.assertTrue(q._recovery_state().first_consumed) + self.assertIsNone(q._recovery_state().owner) + + def test_second_failed_fit_no_second_release(self): + q, c, k, n = self.make() + v, e, count = run(q, [fail(n), fail(n, 200)]) + self.assertIsInstance(e, Refusal) + self.assertEqual(count, 2) + self.assertEqual(c.events.count("release"), 1) + + def test_lower_after_release_refuses(self): + q, c, k, n = self.make() + + def release(): + c.events.append("release") + c.free = 31 + + c.empty_cache = release + v, e, count = run(q, [fail(n), success(n, 80)]) + self.assertIsInstance(e, Refusal) + self.assertEqual(e.usable_limit_bytes, 1) + self.assertEqual(c.events.count("release"), 1) + + def test_quota_stops_with_persistent_first_debt(self): + q, c, k, n = self.make() + run(q, [fail(n), success(n)]) + c.free = 40 + v, e, count = run(q, [fail(n), success(n)]) + self.assertIsInstance(e, Refusal) + self.assertEqual(count, 1) + self.assertEqual(c.events.count("release"), 1) + self.assertGreater(q._recovery_state().cost, 0) + + def test_completed_forward_earns_next_trial(self): + q, c, k, n = self.make() + run(q, [fail(n), success(n)]) + q._record_recovery_work("dp_rank_forward", 2.0) + c.free = 40 + v, e, count = run(q, [fail(n), success(n)]) + self.assertIsNone(e) + self.assertEqual(c.events.count("release"), 2) + + def test_no_first_trial_per_entrypoint(self): + q, c, k, n = self.make() + run(q, [fail(n), success(n)]) + c.free = 40 + v, e, count = run(q, [fail(n), success(n)], sync=False) + self.assertIsInstance(e, Refusal) + self.assertEqual(c.events.count("release"), 1) + + def test_invalid_clock_does_not_release(self): + for x in (math.nan, math.inf): + q, c, k, n = self.make() + k.next = x + v, e, count = run(q, [fail(n), success(n)]) + self.assertIsInstance(e, Refusal) + self.assertNotIn("release", c.events) + self.assertTrue(q._recovery_state().invalid) + + def test_test_cap_never_triggers_release_even_physical_low(self): + for free in (40, 200): + q, c, k, n = self.make() + c.free = free + os.environ["CONTROL_ART_HOOK"] = "1" + os.environ["CONTROL_ART_LIMIT"] = "20" + try: + v, e, count = run(q, [fail(n), success(n)]) + finally: + os.environ.pop("CONTROL_ART_HOOK") + os.environ.pop("CONTROL_ART_LIMIT") + self.assertIsInstance(e, Refusal) + self.assertNotIn("release", c.events) + + def test_sufficient_space_after_refusal_retries_without_release(self): + q, c, k, n = self.make() + c.free = 200 + v, e, count = run(q, [fail(n), success(n)]) + self.assertIsNone(e) + self.assertEqual(count, 2) + self.assertNotIn("release", c.events) + + def test_original_over_timing_secondary(self): + q, c, k, n = self.make() + original = RuntimeError("release original") + c.failure = original + old = k.perf_counter + calls = [] + + def timer(): + calls.append(1) + if len(calls) > 2: + raise ValueError("secondary timer") + return old() + + k.perf_counter = timer + v, e, count = run(q, [fail(n), success(n)]) + self.assertIs(e, original) + self.assertTrue(q._recovery_state().invalid) + self.assertIsNone(q._recovery_state().owner) + + def test_no_work_for_unowned_context(self): + q, c, k, n = self.make() + q._record_recovery_work("arbitrary caller time", 100.0) + self.assertEqual(q._recovery_state().work, 0) + + def test_foreign_owner_is_retained(self): + q, c, k, n = self.make() + foreign = object() + q._recovery_state().owner = foreign + v, e, count = run(q, [fail(n), success(n)]) + self.assertIsInstance(e, Refusal) + self.assertIs(q._recovery_state().owner, foreign) + self.assertTrue(q._recovery_state().invalid) + self.assertNotIn("release", c.events) + + def test_cap_refusal_in_later_trial_does_not_use_work(self): + q, c, k, n = self.make() + run(q, [fail(n), success(n)]) + q._record_recovery_work("dp_rank_forward", 100) + c.free = 40 + os.environ["CONTROL_ART_HOOK"] = "1" + os.environ["CONTROL_ART_LIMIT"] = "20" + try: + v, e, count = run(q, [fail(n), success(n)]) + finally: + os.environ.pop("CONTROL_ART_HOOK") + os.environ.pop("CONTROL_ART_LIMIT") + self.assertIsInstance(e, Refusal) + self.assertEqual(c.events.count("release"), 1) + + def test_fresh_pre_release_fit_skips_call_consumes_trial(self): + q, c, k, n = self.make() + reads = [] + old = c.mem_get_info + + def sample(d): + reads.append(1) + if len(reads) >= 3: + c.free = 200 + return old(d) + + c.mem_get_info = sample + v, e, count = run(q, [fail(n), success(n)]) + self.assertIsNone(e) + self.assertNotIn("release", c.events) + self.assertTrue(q._recovery_state().first_consumed) + + def test_overflowing_work_disables_recovery(self): + q, c, k, n = self.make() + q._record_recovery_work("dp_rank_forward", 1e308) + q._record_recovery_work("dp_rank_forward", 1e308) + self.assertTrue(q._recovery_state().invalid) + v, e, count = run(q, [fail(n), success(n)]) + self.assertNotIn("release", c.events) + + def test_finite_inputs_overflow_derived_quota(self): + q, c, k, n = self.make() + state = q._recovery_state() + state.first_consumed = True + state.high = 1e308 + state.cost = 1e308 + state.work = 1e308 + v, e, count = run(q, [fail(n), success(n)]) + self.assertTrue(state.invalid) + self.assertNotIn("release", c.events) + + def test_forward_clock_failures_preserve_success(self): + for failing_call in (1, 2): + q, c, k, n = self.make() + old = k.perf_counter + calls = [] + + def timer(): + calls.append(1) + if len(calls) == failing_call: + raise ValueError("instrumentation only") + return old() + + k.perf_counter = timer + p = plan() + p.request_count = 1 + outputs, baseline = q._run_flat_plan_with_memory_tracking( + p, check=n["_MemoryCheck"](80, 170, True), context="dp_rank_forward" + ) + self.assertEqual(len(outputs), 1) + self.assertTrue(q._recovery_state().invalid) + self.assertEqual(q._recovery_state().work, 0) + + def test_forward_recording_failure_preserves_success(self): + q, c, k, n = self.make() + p = plan() + p.request_count = 1 + + def record(*a): + raise ValueError("recording only") + + q._record_recovery_work = record + outputs, baseline = q._run_flat_plan_with_memory_tracking( + p, check=n["_MemoryCheck"](80, 170, True), context="dp_rank_forward" + ) + self.assertEqual(len(outputs), 1) + self.assertTrue(q._recovery_state().invalid) + + def test_execution_oom_keeps_admission_and_cause(self): + q, c, k, n = self.make() + p = plan() + p.request_count = 1 + check = n["_MemoryCheck"](80, 170, True) + original = OOM("execution") + + def execute(p): + raise original + + q._execute_flat_plan = execute + try: + q._run_flat_plan_with_memory_tracking( + p, check=check, context="dp_rank_forward" + ) + except Refusal as e: + self.assertIs(e.__cause__, original) + self.assertEqual(e.usable_limit_bytes, check.available_bytes) + self.assertEqual(e.predicted_peak_bytes, check.estimated_required_bytes) + else: + self.fail("expected exact original wrapped OOM") + self.assertNotIn("release", c.events) + self.assertEqual(q._recovery_state().work, 0) + + def test_split_children_credit_once_and_rollback(self): + for fail_at in (None, 2): + q, c, k, n = self.make() + p = plan() + p.request_count = 1 + count = [] + + def execute(p): + count.append(1) + if len(count) == fail_at: + raise OOM("second child") + return [object()] + + q._execute_flat_plan = execute + split = types.SimpleNamespace( + subforwards=(p, p), + request_indices=((0,), (1,)), + request_count=2, + subforward_count=2, + ) + state = q._recovery_state() + state.work = 1.0 + try: + outputs, baseline, peak = q._execute_split_plan_with_memory_tracking( + split, + check=n["_MemoryCheck"](80, 170, True), + context="dp_rank_forward", + ) + except Partial: + self.assertEqual(fail_at, 2) + self.assertEqual(state.work, 1.0) + else: + self.assertIsNone(fail_at) + self.assertAlmostEqual(state.work, 1.02) + + def test_split_mapping_error_rolls_back(self): + q, c, k, n = self.make() + p = plan() + p.request_count = 2 + split = types.SimpleNamespace( + subforwards=(p,), + request_indices=((0,),), + request_count=1, + subforward_count=1, + ) + state = q._recovery_state() + state.work = 1.0 + with self.assertRaises(ValueError): + q._execute_split_plan_with_memory_tracking( + split, check=n["_MemoryCheck"](80, 170, True), context="dp_rank_forward" + ) + self.assertEqual(state.work, 1.0) + + def test_cross_entrypoint_progress_with_work_and_no_new_first_trial(self): + q, c, k, n = self.make() + run(q, [fail(n), success(n)], sync=True) + q._record_recovery_work("dp_rank_forward", 2.0) + c.free = 40 + v, e, count = run(q, [fail(n), success(n)], sync=False) + self.assertIsNone(e) + self.assertEqual(c.events.count("release"), 2) + + def test_normal_dp_local_has_no_added_check(self): + q, c, k, n = self.make() + value = success(n) + q._find_admissible_forward = lambda *a, **kw: value + q._ensure_checkpoint_slots_for = lambda *a, **kw: None + + def forbidden(*a, **kw): + raise AssertionError("unnecessary added admission collective") + + q._memory_check_required = forbidden + result = q._plan_admissible_forward( + [], checkpoint=None, context="dp_rank_forward" + ) + self.assertIs(result[1], value[1]) + self.assertNotIn("release", c.events) + + def test_final_bookkeeping_failure_preserves_success(self): + q, c, k, n = self.make() + calls = [] + + def search(): + calls.append(1) + if len(calls) == 1: + return fail(n) + q._recovery_state().high = object() + return success(n) + + result = q._recover_admission( + search, + lambda v: v, + lambda v, c: (v[0], c), + context="forward_micro_batches", + sync_across_dp=True, + ) + self.assertTrue(result[1].fits) + self.assertTrue(q._recovery_state().invalid) + self.assertIsNone(q._recovery_state().owner) + + def test_owner_cleanup_keeps_original_runtime_error(self): + for secondary_type in (KeyboardInterrupt, SystemExit): + q, c, k, n = self.make() + state = q._recovery_state() + original = RuntimeError("original release failure") + secondary = secondary_type("owner cleanup cancellation") + c.failure = original + + class Lock: + calls = 0 + + def __enter__(self): + self.calls += 1 + if self.calls == 4: + raise secondary + + def __exit__(self, *args): + return False + + state.lock = Lock() + _, error, _ = run(q, [fail(n), success(n)]) + self.assertIs(error, original) + self.assertEqual(state.lock.calls, 4) + self.assertTrue(state.invalid) + + def test_first_bookkeeping_cancellation_keeps_identity(self): + q, c, k, n = self.make() + state = q._recovery_state() + original = KeyboardInterrupt("first bookkeeping cancellation") + secondary = SystemExit("second owner cleanup cancellation") + + class Lock: + calls = 0 + + def __enter__(self): + self.calls += 1 + if self.calls == 3: + raise original + if self.calls == 4: + raise secondary + + def __exit__(self, *args): + return False + + state.lock = Lock() + _, error, _ = run(q, [fail(n), success(n)]) + self.assertIs(error, original) + self.assertEqual(state.lock.calls, 4) + self.assertTrue(state.invalid) + + def test_ordinary_owner_cleanup_failure_keeps_success(self): + q, c, k, n = self.make() + state = q._recovery_state() + + class Lock: + calls = 0 + + def __enter__(self): + self.calls += 1 + if self.calls == 4: + raise RuntimeError("ordinary owner cleanup diagnostic") + + def __exit__(self, *args): + return False + + state.lock = Lock() + value, error, _ = run(q, [fail(n), success(n)]) + self.assertIsNone(error) + self.assertTrue(value[1].fits) + self.assertTrue(state.invalid) + + def test_sampling_and_reduction_errors_keep_both_diagnostics(self): + for phase in ("pre_sample", "release", "post_sample"): + for mode in ("local", "transport", "both"): + with self.subTest(phase=phase, mode=mode): + q, c, k, n = self.make() + local = ValueError("local sample or release failed") + cause, context = TypeError("original cause"), LookupError("context") + local.__cause__, local.__context__ = cause, context + local.__suppress_context__ = True + local.add_note("existing note") + transport = RuntimeError("memory reduction failed") + transport.__cause__ = OSError("transport cause") + samples, reductions = [], [] + + def available(): + samples.append(1) + if mode != "transport" and ( + phase == "pre_sample" + and len(samples) == 1 + or phase == "post_sample" + and len(samples) == 2 + ): + raise local + return 10 + + def reduce(values, *, op, sync_across_dp): + reductions.append((op, list(values), sync_across_dp)) + if mode != "local" and len(reductions) == ( + 3 if phase == "pre_sample" else 4 + ): + raise transport + return values + + q._available_memory_bytes = available + q._recovery_reduce = reduce + if phase == "release" and mode != "transport": + c.failure = local + _, error, calls = run(q, [fail(n), success(n)]) + self.assertIs(error, transport if mode == "transport" else local) + if mode != "transport": + self.assertIs(local.__cause__, cause) + self.assertIs(local.__context__, context) + self.assertTrue(local.__suppress_context__) + self.assertEqual(local.__notes__[0], "existing note") + self.assertEqual( + len(local.__notes__), 2 if mode == "both" else 1 + ) + if mode == "both": + self.assertIn( + "OSError: transport cause", local.__notes__[1] + ) + self.assertIn( + "RuntimeError: memory reduction failed", + local.__notes__[1], + ) + self.assertIn("raise transport", local.__notes__[1]) + self.assertEqual(calls, 1) + self.assertEqual( + [x[0] for x in reductions], + ["SUM", "MAX", "MIN"] + + ([] if phase == "pre_sample" else ["MIN"]), + ) + state = q._recovery_state() + self.assertIsNone(state.owner) + self.assertEqual(state.work, 0) + self.assertEqual(state.cost, state.high) + self.assertGreater(state.cost, 0) + self.assertEqual(state.first_consumed, phase != "pre_sample") + + def test_admission_sampling_and_reduction_error_keep_original_chain(self): + q, c, k, n = self.make() + local = ValueError("admission sample") + context, cause = LookupError("context"), TypeError("cause") + local.__context__, local.__cause__ = context, cause + transport = RuntimeError("admission reduction") + calls = [] + + def available(): + raise local + + def all_reduce(value, *, op, group): + calls.append(op) + if op == "MIN": + raise transport + + q._available_memory_bytes = available + with patch.object( + _impl, + "dist", + types.SimpleNamespace( + is_available=lambda: True, + is_initialized=lambda: True, + ReduceOp=types.SimpleNamespace(MAX="MAX", MIN="MIN"), + all_reduce=all_reduce, + ), + ): + with self.assertRaises(ValueError) as captured: + q._memory_check_required(80, sync_across_dp=True) + self.assertIs(captured.exception, local) + self.assertIs(local.__context__, context) + self.assertIs(local.__cause__, cause) + self.assertIn("RuntimeError: admission reduction", local.__notes__[0]) + self.assertEqual(calls, ["MAX", "MIN"]) + + def test_secondary_note_failure_never_replaces_primary(self): + primary, secondary = ValueError("primary"), RuntimeError("secondary") + context = LookupError("original context") + primary.__context__ = context + helper = _impl.TrainerRank._memory_error_with_reduction_note + with patch.object( + _impl.traceback, "format_exception", side_effect=SystemExit("renderer") + ): + self.assertIs(helper(primary, secondary), primary) + self.assertIs(primary.__context__, context) + primary.__dict__["__notes__"] = 42 + self.assertIs(helper(primary, secondary), primary) + self.assertEqual(primary.__notes__, 42) + self.assertIs(primary.__context__, context) + + +def _check_dense_cp_exact_demand_recovery(monkeypatch, *, fits_after_release): + import pytest + from test_trainer_rank_recompute_memory import _hybrid_rank + import torch + + from art.megatron.model_support.handlers.qwen3_5 import Qwen35DenseHandler + from art.megatron.routing_replay import ParallelTopology + from art.trainer_rank import ForwardInput, TrainerRankMemoryError + + # Declared CPU topology and inert model fixture. The real CP/GDN layout + # planner computes retained tokens; its result is never stubbed. + rank = _hybrid_rank(monkeypatch, 2) + monkeypatch.delattr(rank, "_topology_key") + topology = ParallelTopology(tp=2, ep=1, cp=2, sp=True) + monkeypatch.setattr(rank, "_topology", lambda: topology) + monkeypatch.setattr(rank, "_dp_rank_and_size", lambda: (0, topology.dp)) + rank.runtime.provider.tensor_model_parallel_size = 2 + rank.runtime.provider.context_parallel_size = 2 + # _hybrid_rank supplies an inert namespace, not TrainingRuntime's property. + runtime = cast(types.SimpleNamespace, rank.runtime) + assert isinstance(runtime, types.SimpleNamespace) + runtime.model_support_handler = Qwen35DenseHandler() + assert rank._topology_key() == (1, 2, 2, 1) + assert rank._dp_rank_and_size() == (0, 1) + assert rank.device.type == "cpu" and not rank._geometry.moe_experts + requests = [ + ForwardInput(input_tokens=torch.arange(4096), hidden_states=True), + ForwardInput( + input_tokens=torch.cat( + (torch.arange(1228), torch.arange(1228, 4096) + 10000) + ), + hidden_states=True, + ), + ] + budget, search_number = 0, 0 + plan_fields: tuple[int, int, int] | None = None + exact_rows, estimates, searches, releases, errors = [], [], [], [], [] + monkeypatch.setattr(rank, "_available_memory_bytes", lambda: budget) + monkeypatch.setattr( + rank, "_execute_flat_plan", lambda *a, **kw: pytest.fail("No model execution") + ) + monkeypatch.setattr( + torch.cuda, "empty_cache", lambda *a, **kw: pytest.fail("No CUDA release") + ) + estimate = rank._estimate_required_memory_bytes_from_values + check = rank._memory_check + search = rank._search_next_micro_batch + cheap = rank._estimate_flat_forward + error_factory = _impl._ForwardRefusal.error + + def observed_estimate(**kwargs): + value = estimate(**kwargs) + if plan_fields is not None: + packed, retained, segments = plan_fields + assert kwargs["retained_tokens"] == retained + assert kwargs["gdn_segments"] == segments + exact_rows.append((search_number, packed, retained, segments, value)) + return value + + def observed_check(plan, **kwargs): + nonlocal plan_fields + previous = plan_fields + plan_fields = ( + plan.packed_tokens, + rank._plan_retained_tokens(plan), + plan.grad_segment_count, + ) + assert plan_fields[1] > 0 and plan_fields[2] > 0 + try: + return check(plan, **kwargs) + finally: + plan_fields = previous + + def observed_cheap(*args, **kwargs): + value = cheap(*args, **kwargs) + estimates.append((search_number, value is None)) + assert value is None, "Dense CP must reach exact-plan pricing" + return value + + def observed_search(*args, **kwargs): + nonlocal search_number + search_number += 1 + value = search(*args, **kwargs) + searches.append(value) + return value + + def observed_error(refused, context): + error = error_factory(refused, context) + errors.append(error) + return error + + def release_facade(refused_check, **kwargs): + nonlocal budget + # Only availability changes. The actual outer recovery/search remain; + # this CPU test does not exercise allocator policy or CUDA release. + assert kwargs["sync_across_dp"] is True + assert not releases and refused_check.estimated_required_bytes > 1 + # A split refusal can carry only an optimistic lower bound. Restore + # the first actual flat-plan demand observed before this release. + exact = next(row[1:] for row in exact_rows if row[0] == search_number) + assert refused_check.estimated_required_bytes <= exact[-1] + releases.append(exact) + budget = exact[-1] if fits_after_release else 1 + return True + + monkeypatch.setattr( + rank, "_estimate_required_memory_bytes_from_values", observed_estimate + ) + monkeypatch.setattr(rank, "_memory_check", observed_check) + monkeypatch.setattr(rank, "_estimate_flat_forward", observed_cheap) + monkeypatch.setattr(rank, "_search_next_micro_batch", observed_search) + monkeypatch.setattr(_impl._ForwardRefusal, "error", observed_error) + monkeypatch.setattr(rank, "_try_cache_recovery", release_facade) + if fits_after_release: + selected = rank._select_next_micro_batch([requests], 0) + assert selected.check.fits + assert selected.check.estimated_required_bytes == releases[0][-1] + assert selected.check.available_bytes == budget and len(errors) == 1 + else: + with pytest.raises(TrainerRankMemoryError) as caught: + rank._select_next_micro_batch([requests], 0) + assert len(errors) == 2 and caught.value is errors[1] + assert caught.value.__cause__ is errors[0] and errors[0] is not errors[1] + assert search_number == len(searches) == 2 and len(releases) == 1 + assert isinstance(searches[0], _impl._ForwardRefusal) + assert {number for number, _ in estimates} == {1, 2} + assert all(unavailable for _, unavailable in estimates) + before = {row[1:] for row in exact_rows if row[0] == 1} + after = {row[1:] for row in exact_rows if row[0] == 2} + assert before and after and after <= before + if fits_after_release: + plan = selected.plan + assert isinstance(plan, _impl._FlatForwardPlan) + assert ( + plan.packed_tokens, + rank._plan_retained_tokens(plan), + plan.grad_segment_count, + selected.check.estimated_required_bytes, + ) == releases[0] + assert releases[0] in before + else: + assert before == after + assert rank._recovery_state().owner is None + assert not torch.cuda.is_initialized() + + +def test_dense_cp_exact_demand_fits_after_recovery(monkeypatch): + _check_dense_cp_exact_demand_recovery(monkeypatch, fits_after_release=True) + + +def test_dense_cp_exact_demand_refuses_after_recovery(monkeypatch): + _check_dense_cp_exact_demand_recovery(monkeypatch, fits_after_release=False) diff --git a/tests/unit/test_trainer_rank_cuda_budget.py b/tests/unit/test_trainer_rank_cuda_budget.py index 858ec6102..f373e8349 100644 --- a/tests/unit/test_trainer_rank_cuda_budget.py +++ b/tests/unit/test_trainer_rank_cuda_budget.py @@ -41,36 +41,37 @@ def budget(monkeypatch): return rank, stats -@pytest.mark.parametrize("active,available", [(10, 140), (30, 120), (80, 70)]) -def test_native_pending_is_not_reusable_credit(budget, active, available): +@pytest.mark.parametrize("active", [10, 30, 80]) +def test_native_cache_and_pending_are_not_physical_credit(budget, active): rank, stats = budget stats["active_bytes.all.current"] = active - assert rank._available_memory_bytes() == available - torch.cuda.memory_stats.assert_called_once_with(rank.device) - torch.cuda.memory_allocated.assert_not_called() - torch.cuda.memory_reserved.assert_not_called() + assert rank._available_memory_bytes() == 70 + cast(Mock, torch.cuda.memory_stats).assert_not_called() + cast(Mock, torch.cuda.memory_allocated).assert_called_once_with(rank.device) + cast(Mock, torch.cuda.memory_reserved).assert_not_called() -def test_pending_credit_changes_admission_before_execution(budget): +def test_reclaimable_cache_does_not_make_progress_without_physical_free(budget): rank, stats = budget check = rank._memory_check_required(130) - assert check.available_bytes == 120 + assert check.available_bytes == 70 assert not check.fits # Normal allocator collection can later make the same bytes inactive. # The budget itself neither polls events nor forces collection. stats["active_bytes.all.current"] = 10 - assert rank._memory_check_required(130).fits + assert not rank._memory_check_required(130).fits + assert rank._memory_check_required(70).fits -def test_split_and_private_credit_remain_explicit_residuals(budget): +def test_split_and_private_cache_do_not_inflate_physical_sample(budget): rank, stats = budget stats["active_bytes.all.current"] = 10 stats["inactive_split_bytes.all.current"] = 70 - assert rank._available_memory_bytes() == 140 + assert rank._available_memory_bytes() == 70 stats["inactive_split_bytes.all.current"] = 0 # A whole inactive retained private pool has this same scalar geometry. - # This partial correction does not establish pool compatibility. - assert rank._available_memory_bytes() == 140 + # Neither layout can become physical reserve through accounting. + assert rank._available_memory_bytes() == 70 @pytest.mark.parametrize( @@ -96,7 +97,7 @@ def test_incomplete_native_counters_grant_no_cache_credit(budget, field, value): else: stats[key] = value assert rank._available_memory_bytes() == 70 - torch.cuda.memory_allocated.assert_called_once_with(rank.device) + cast(Mock, torch.cuda.memory_allocated).assert_called_once_with(rank.device) @pytest.mark.parametrize("backend", ["cudaMallocAsync", "unrecognized"]) @@ -104,9 +105,9 @@ def test_other_backends_retain_legacy_unqualified_credit(budget, monkeypatch, ba rank, _ = budget monkeypatch.setattr(torch.cuda, "get_allocator_backend", lambda: backend) assert rank._available_memory_bytes() == 140 - torch.cuda.memory_stats.assert_not_called() - torch.cuda.memory_allocated.assert_called_once_with(rank.device) - torch.cuda.memory_reserved.assert_called_once_with(rank.device) + cast(Mock, torch.cuda.memory_stats).assert_not_called() + cast(Mock, torch.cuda.memory_allocated).assert_called_once_with(rank.device) + cast(Mock, torch.cuda.memory_reserved).assert_called_once_with(rank.device) @pytest.mark.parametrize("missing", [False, True]) @@ -120,17 +121,25 @@ def test_existing_test_limit_stays_relative_to_allocated(budget, monkeypatch, mi monkeypatch.setenv(_impl._TEST_MEMORY_LIMIT_ENV, "5") assert rank._available_memory_bytes() == 0 monkeypatch.setenv("ART_TRAINER_RANK_TEST_HOOKS", "0") - assert rank._available_memory_bytes() == (70 if missing else 120) + assert rank._available_memory_bytes() == 70 + monkeypatch.setenv("ART_TRAINER_RANK_TEST_HOOKS", "1") + monkeypatch.setenv(_impl._TEST_MEMORY_LIMIT_ENV, "10000") + assert rank._available_memory_bytes() == 70 @pytest.mark.parametrize( - "api", ["mem_get_info", "get_allocator_backend", "memory_stats"] + "api", ["mem_get_info", "get_allocator_backend", "memory_allocated"] +) +@pytest.mark.parametrize( + "error_type", [RuntimeError, KeyboardInterrupt, SystemExit, asyncio.CancelledError] ) -def test_actual_api_error_identity_is_not_swallowed(budget, monkeypatch, api): +def test_actual_api_error_identity_is_not_swallowed( + budget, monkeypatch, api, error_type +): rank, _ = budget - error = RuntimeError("native API failed") + error = error_type("native API failed") monkeypatch.setattr(torch.cuda, api, Mock(side_effect=error)) - with pytest.raises(RuntimeError) as caught: + with pytest.raises(error_type) as caught: rank._available_memory_bytes() assert caught.value is error @@ -152,10 +161,14 @@ def test_cpu_budget_avoids_all_new_cuda_api_calls(budget, monkeypatch): torch.cuda, "get_allocator_backend", Mock(side_effect=AssertionError) ) assert rank._available_memory_bytes() == 1 << 60 - torch.cuda.memory_stats.assert_not_called() + cast(Mock, torch.cuda.memory_stats).assert_not_called() -def test_required_max_then_available_min_collectives_unchanged(budget, monkeypatch): +@pytest.mark.parametrize("sync_across_dp", [False, True]) +@pytest.mark.parametrize("required", [0, 130]) +def test_required_max_then_available_min_collectives_unchanged( + budget, monkeypatch, sync_across_dp, required +): rank, _ = budget monkeypatch.setattr(_impl.dist, "is_available", lambda: True) monkeypatch.setattr(_impl.dist, "is_initialized", lambda: True) @@ -169,21 +182,57 @@ def test_required_max_then_available_min_collectives_unchanged(budget, monkeypat def reduce(value, op, group): calls.append((float(value.item()), op, group)) - value.fill_(150 if op == _impl.dist.ReduceOp.MAX else 110) + value.fill_(150 if op == _impl.dist.ReduceOp.MAX else 60) monkeypatch.setattr(_impl.dist, "all_reduce", reduce) - check = rank._memory_check_required(130) + check = rank._memory_check_required(required, sync_across_dp=sync_across_dp) + expected_group = None if sync_across_dp else group assert calls == [ - (130, _impl.dist.ReduceOp.MAX, group), - (120, _impl.dist.ReduceOp.MIN, group), + (required, _impl.dist.ReduceOp.MAX, expected_group), + (70, _impl.dist.ReduceOp.MIN, expected_group), ] assert (check.estimated_required_bytes, check.available_bytes, check.fits) == ( 150, - 110, + 60, False, ) +def test_oom_preserves_admission_and_original_cause_after_free_changes( + budget, monkeypatch +): + rank, _ = budget + 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("later allocation") + + 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()) + plan = SimpleNamespace(packed_tokens=1, logical_tokens=1) + with pytest.raises(_impl.TrainerRankMemoryError) as caught: + rank._run_flat_plan_with_memory_tracking( + plan, check=admitted, context="physical-free CPU witness" + ) + assert caught.value.__cause__ is original + assert caught.value.usable_limit_bytes == admitted.available_bytes == 70 + read.assert_called_once_with(rank.device) + # A separate later sample is zero. It cannot rewrite the admitted check or + # imply that the previous physical sample reserved bytes through execution. + assert rank._available_memory_bytes() == 0 + assert admitted.available_bytes == 70 + + @pytest.mark.parametrize("failed_locally", [False, True]) @pytest.mark.parametrize( "error_type", [RuntimeError, KeyboardInterrupt, SystemExit, asyncio.CancelledError] @@ -332,11 +381,14 @@ def test_final_selection_uses_pure_fresh_budget_and_original_demand( (_impl.dist.ReduceOp.MAX, 192, None), (_impl.dist.ReduceOp.MIN, available, None), ] + * (2 if available < 192 else 1) if distributed else [] ) - torch.cuda.empty_cache.assert_not_called() - torch.cuda.mem_get_info.assert_called_once_with(rank.device) + cast(Mock, torch.cuda.empty_cache).assert_not_called() + assert cast(Mock, torch.cuda.mem_get_info).call_count == ( + 2 if available < 192 else 1 + ) def test_available_sample_follows_required_collective(budget, monkeypatch): @@ -371,7 +423,7 @@ def reduce(value, op, group): False, ) assert events == [_impl.dist.ReduceOp.MAX, "sample", _impl.dist.ReduceOp.MIN] - torch.cuda.empty_cache.assert_not_called() + cast(Mock, torch.cuda.empty_cache).assert_not_called() def test_execution_failure_retains_final_admission_without_resampling( @@ -406,4 +458,4 @@ def execute(_): 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() + cast(Mock, torch.cuda.empty_cache).assert_not_called() diff --git a/tests/unit/test_trainer_rank_recovery_slots.py b/tests/unit/test_trainer_rank_recovery_slots.py new file mode 100644 index 000000000..abcbbd889 --- /dev/null +++ b/tests/unit/test_trainer_rank_recovery_slots.py @@ -0,0 +1,98 @@ +"""Checkpoint setup precedes the DP-local recovery loop exactly once.""" + +from types import SimpleNamespace +from typing import cast + +import pytest + +from art.trainer_rank import TrainerRank, _impl + + +@pytest.mark.parametrize("search_count", (1, 2, 3)) +@pytest.mark.parametrize("empty", (False, True)) +def test_dp_recovery_ensures_before_all_searches(search_count, empty): + rank = TrainerRank.__new__(TrainerRank) + rank.device = _impl.torch.device("cpu") + requests = [] if empty else [object()] + checkpoint = object() + events = [] + plan = cast( + _impl._AnyForwardPlan, SimpleNamespace(packed_tokens=1, logical_tokens=1) + ) + bad = _impl._MemoryCheck(80, 10, False) + fit = (plan, _impl._MemoryCheck(80, 200, True)) + refused = _impl._ForwardRefusal(plan, bad, "too large") + results = {1: [fit], 2: [refused, fit], 3: [(plan, bad), refused, fit]}[ + search_count + ] + + def ensure(actual, **kwargs): + assert actual is requests and kwargs == {"checkpoint": checkpoint} + events.append("ensure") + + def search(actual, **kwargs): + assert actual is requests + assert kwargs == dict( + checkpoint=checkpoint, + refusal_prefix="forward is predicted to exceed available memory", + ensure_slots=False, + ) + events.append("search") + return results.pop(0) + + rank._ensure_checkpoint_slots_for = ensure + rank._find_admissible_forward = search + rank._snapshot_planning_telemetry = lambda *args: None + rank._try_cache_recovery = lambda *args, **kwargs: True + result = rank._plan_admissible_forward( + requests, checkpoint=checkpoint, context="dp_rank_forward" + ) + assert result == fit and not results + assert events == ["ensure"] + ["search"] * search_count + + +@pytest.mark.parametrize("error_type", (ValueError, KeyboardInterrupt)) +def test_checkpoint_error_precedes_search_and_preserves_identity(error_type): + rank = TrainerRank.__new__(TrainerRank) + error = error_type("checkpoint setup failed") + error.__cause__, error.__context__ = LookupError("cause"), KeyError("context") + error.__suppress_context__ = True + cause, context = error.__cause__, error.__context__ + events = [] + + def ensure(*args, **kwargs): + events.append("ensure") + raise error + + def forbidden(*args, **kwargs): + raise AssertionError("search/recovery must not begin") + + rank._ensure_checkpoint_slots_for = ensure + rank._recover_admission = forbidden + rank._find_admissible_forward = forbidden + with pytest.raises(error_type) as captured: + rank._plan_admissible_forward([], checkpoint=None, context="dp_rank_forward") + assert captured.value is error + assert error.__cause__ is cause and error.__context__ is context + assert error.__suppress_context__ and events == ["ensure"] + + +@pytest.mark.parametrize("ensure_slots", (None, False, True)) +def test_direct_search_keeps_default_setup(ensure_slots): + rank = TrainerRank.__new__(TrainerRank) + events = [] + plan, check = object(), _impl._MemoryCheck(80, 200, True) + rank._ensure_checkpoint_slots_for = lambda *a, **kw: events.append("ensure") + + def materialize(*args, **kwargs): + assert kwargs == dict(checkpoint=None, ensure_slots=False) + events.append("plan") + return plan + + rank._plan_flat_forward = materialize + rank._memory_check = lambda actual: check if actual is plan else None + options = {} if ensure_slots is None else {"ensure_slots": ensure_slots} + assert rank._find_admissible_forward( + [], checkpoint=None, refusal_prefix="refused", **options + ) == (plan, check) + assert events == (["plan"] if ensure_slots is False else ["ensure", "plan"]) diff --git a/tests/unit/test_trainer_rank_recovery_slots_distributed.py b/tests/unit/test_trainer_rank_recovery_slots_distributed.py new file mode 100644 index 000000000..3326663cb --- /dev/null +++ b/tests/unit/test_trainer_rank_recovery_slots_distributed.py @@ -0,0 +1,164 @@ +"""Two real Gloo DP peers, native checkpoint gather, CPU admission facades.""" + +import json +import os +from pathlib import Path +import subprocess +import sys + +import pytest + +HEAD = Path(__file__).resolve().parents[2] / "src" + + +@pytest.mark.parametrize("mode", ("fit", "both", "asymmetric")) +def test_native_checkpoint_gather_after_recovery(tmp_path, mode): + selected = HEAD + children = [] + logs = [] + try: + for rank in range(2): + log = (tmp_path / f"rank-{rank}.log").open("w") + logs.append(log) + env = os.environ | { + "PYTHONPATH": str(selected), + "PYTHONDONTWRITEBYTECODE": "1", + } + child = subprocess.Popen( + [sys.executable, __file__, "worker", str(rank), mode, str(tmp_path)], + env=env, + stdout=log, + stderr=subprocess.STDOUT, + ) + children.append(child) + for index, child in enumerate(children): + assert child.wait(timeout=25) == 0, ( + tmp_path / f"rank-{index}.log" + ).read_text() + finally: + for child in children: + if child.poll() is None: + child.terminate() + for child in children: + try: + child.wait(timeout=3) + except subprocess.TimeoutExpired: + child.kill() + child.wait(timeout=3) + for log in logs: + log.close() + rows = [ + json.loads((tmp_path / f"rank-{rank}.json").read_text()) for rank in range(2) + ] + assert all(row["error"] is None for row in rows), rows + assert all(row["barrier_error"] is None for row in rows), rows + assert all(row["ensures"] == 1 for row in rows), rows + + +def worker(index, mode, directory): + from datetime import timedelta + import threading + from types import SimpleNamespace + + import torch + import torch.distributed as dist + + from art.trainer_rank import TrainerRank, _impl + + torch.set_num_threads(1) + dist.init_process_group( + "gloo", + rank=index, + world_size=2, + init_method=f"file://{directory}/rendezvous", + timeout=timedelta(seconds=3), + ) + rank = TrainerRank.__new__(TrainerRank) + rank.device = torch.device("cpu") + rank._checkpoint_mutation_lock = threading.RLock() + rank._checkpoint_prefetch_lock = threading.Lock() + rank._checkpoint_slots = {} + rank._checkpoint_group_lock = threading.Lock() + # Native _ensure_checkpoint_slots uses these all-rank groups unchanged. + rank._checkpoint_process_group = dist.new_group( + backend="gloo", timeout=timedelta(seconds=3) + ) + rank._checkpoint_finalize_process_group = dist.new_group( + backend="gloo", timeout=timedelta(seconds=3) + ) + groups = [ + dist.new_group([r], backend="gloo", timeout=timedelta(seconds=3)) + for r in range(2) + ] + rank._forward_memory_group = lambda: groups[index] + plan = SimpleNamespace( + packed_tokens=1, + logical_tokens=1, + active_logical_tokens=1, + grad_segment_count=0, + output_bytes=0, + signature=_impl._MemorySignature( + topology=(2, 1, 1, 1), + planner_coefficients=(0, None), + slot_group_count=1, + request_mix=("hidden_states",), + grad_enabled=False, + grad_modes=(False,), + ), + ) + rank._plan_flat_forward = lambda *args, **kwargs: plan + rank._estimate_required_memory_bytes_from_values = lambda **kwargs: 80 + rank._snapshot_planning_telemetry = lambda *args: None + reads = [] + ensures = [] + + def available(): + reads.append(1) + deficient = mode == "both" or (mode == "asymmetric" and index == 0) + return 10 if deficient and len(reads) <= 2 else 200 + + rank._available_memory_bytes = available + native = rank._ensure_checkpoint_slots + + def ensure(values): + ensures.append(1) + return native(values) + + rank._ensure_checkpoint_slots = ensure + request = SimpleNamespace( + target_tokens=None, + logits=False, + top_k=None, + hidden_states=True, + checkpoint=None, + ) + error = barrier_error = None + try: + rank._plan_admissible_forward( + [request], checkpoint=None, context="dp_rank_forward" + ) + except BaseException as exc: + error = {"type": type(exc).__name__, "message": str(exc)} + try: + dist.barrier() + except BaseException as exc: + barrier_error = {"type": type(exc).__name__, "message": str(exc)} + (directory / f"rank-{index}.json").write_text( + json.dumps( + { + "source": _impl.__file__, + "rank": index, + "mode": mode, + "ensures": len(ensures), + "samples": len(reads), + "error": error, + "barrier_error": barrier_error, + }, + indent=2, + ) + ) + dist.destroy_process_group() + + +if __name__ == "__main__": + worker(int(sys.argv[2]), sys.argv[3], Path(sys.argv[4])) diff --git a/tests/unit/test_trainer_rank_split.py b/tests/unit/test_trainer_rank_split.py index 7b8bb812f..359700def 100644 --- a/tests/unit/test_trainer_rank_split.py +++ b/tests/unit/test_trainer_rank_split.py @@ -232,6 +232,8 @@ def plan(requests, **kwargs): # bounded ladder (2, 4, 8 subforwards), whose failed rungs are rejected # with cheap bounds — the planner runs only for the unsplit attempts. _packed_budget(monkeypatch, rank, 9) + # Recovery must observe the same controlled budget as ordinary admission. + monkeypatch.setattr(rank, "_available_memory_bytes", lambda: 9) with pytest.raises(TrainerRankMemoryError) as exc_info: rank.dp_rank_forward(inputs) diff --git a/tests/unit/test_trainer_rank_split_peak.py b/tests/unit/test_trainer_rank_split_peak.py index 54473e12e..dbfb21d54 100644 --- a/tests/unit/test_trainer_rank_split_peak.py +++ b/tests/unit/test_trainer_rank_split_peak.py @@ -211,6 +211,12 @@ def execute(plan): def test_completed_iterator_preserves_caller_peak_for_next_admission(monkeypatch): rank, requests, counters = _counter_split(monkeypatch) + # This fixture's 10,000-byte admission budget is synthetic, not a physical + # deficit. Keep its learned-floor refusal independent of cache recovery. + monkeypatch.setattr(torch.cuda, "get_allocator_backend", lambda: "native") + monkeypatch.setattr(torch.cuda, "mem_get_info", lambda _: (1_000_000, 1_000_000)) + releases = [] + monkeypatch.setattr(torch.cuda, "empty_cache", lambda: releases.append(True)) iterator = rank.forward_micro_batches([requests], yield_empty=True) batch = next(iterator) assert batch.stats.subforward_count == counters["executed"] == 2 @@ -225,6 +231,8 @@ def test_completed_iterator_preserves_caller_peak_for_next_admission(monkeypatch next(rank.forward_micro_batches([requests], yield_empty=True)) assert counters["executed"] == 2 + assert releases == [] + @pytest.mark.parametrize("termination", ["throw", "close"]) def test_incomplete_caller_does_not_learn_split_peak(monkeypatch, termination): @@ -306,7 +314,7 @@ def searched(*args, **kwargs): patch.setattr(rank, "_search_next_micro_batch", searched) - def reduce(value, op, group): + def reduce(value, op, group=None): trace.append(("global" if group is None else "local", str(op))) if group is None: value.fill_(