From 97df0e071fdae9006096dc1004813c4fd1dc7ce0 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Mon, 14 Sep 2026 21:18:28 +0000 Subject: [PATCH] Exclude pending native allocator bytes from reusable credit --- src/art/trainer_rank/_impl.py | 26 ++- tests/unit/test_trainer_rank_active_memory.py | 10 + tests/unit/test_trainer_rank_cuda_budget.py | 180 ++++++++++++++++++ 3 files changed, 213 insertions(+), 3 deletions(-) create mode 100644 tests/unit/test_trainer_rank_cuda_budget.py diff --git a/src/art/trainer_rank/_impl.py b/src/art/trainer_rank/_impl.py index a183cc117..073fb61e5 100644 --- a/src/art/trainer_rank/_impl.py +++ b/src/art/trainer_rank/_impl.py @@ -4855,9 +4855,29 @@ def _available_memory_bytes(self) -> int: if not (torch.cuda.is_available() and self.device.type == "cuda"): return 1 << 60 free, total = torch.cuda.mem_get_info(self.device) - allocated = int(torch.cuda.memory_allocated(self.device)) - reserved = int(torch.cuda.memory_reserved(self.device)) - reusable_reserved = max(0, reserved - allocated) + 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 + else: + # Preserve the previous, unqualified policy for other backends. + allocated = int(torch.cuda.memory_allocated(self.device)) + reserved = int(torch.cuda.memory_reserved(self.device)) + reusable_reserved = max(0, reserved - allocated) reserve = int(total * _MEMORY_RESERVE_FRACTION) available = max(0, int(free) + reusable_reserved - reserve) if os.environ.get(_TEST_HOOKS_ENV) == "1": diff --git a/tests/unit/test_trainer_rank_active_memory.py b/tests/unit/test_trainer_rank_active_memory.py index a81905fe1..5964ed8a0 100644 --- a/tests/unit/test_trainer_rank_active_memory.py +++ b/tests/unit/test_trainer_rank_active_memory.py @@ -167,6 +167,16 @@ def test_warm_admission_rechecks_current_residency(monkeypatch): ) monkeypatch.setattr(torch.cuda, "memory_allocated", lambda _: state["allocated"]) monkeypatch.setattr(torch.cuda, "memory_reserved", lambda _: state["reserved"]) + monkeypatch.setattr(torch.cuda, "get_allocator_backend", lambda: "native") + monkeypatch.setattr( + torch.cuda, + "memory_stats", + lambda _: { + "allocated_bytes.all.current": state["allocated"], + "active_bytes.all.current": state["allocated"], + "reserved_bytes.all.current": state["reserved"], + }, + ) 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. diff --git a/tests/unit/test_trainer_rank_cuda_budget.py b/tests/unit/test_trainer_rank_cuda_budget.py new file mode 100644 index 000000000..3575fe5e2 --- /dev/null +++ b/tests/unit/test_trainer_rank_cuda_budget.py @@ -0,0 +1,180 @@ +"""Native allocator budget contracts; CPU counters, not CUDA qualification.""" + +from unittest.mock import Mock + +import pytest +import torch + +from art.trainer_rank import TrainerRank, _impl + + +@pytest.fixture +def budget(monkeypatch): + rank = object.__new__(TrainerRank) + rank.device = torch.device("cuda") + stats = { + "allocated_bytes.all.current": 10, + "active_bytes.all.current": 30, + "reserved_bytes.all.current": 80, + "inactive_split_bytes.all.current": 20, + } + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(torch.cuda, "mem_get_info", lambda _: (100, 1000)) + monkeypatch.setattr(torch.cuda, "get_allocator_backend", lambda: "native") + monkeypatch.setattr(torch.cuda, "memory_stats", Mock(side_effect=lambda _: stats)) + monkeypatch.setattr(torch.cuda, "memory_allocated", Mock(return_value=10)) + monkeypatch.setattr(torch.cuda, "memory_reserved", Mock(return_value=80)) + monkeypatch.setattr(_impl.dist, "is_available", lambda: False) + for name in ( + "empty_cache", + "synchronize", + "memory_snapshot", + "reset_peak_memory_stats", + ): + monkeypatch.setattr(torch.cuda, name, Mock(side_effect=AssertionError(name))) + monkeypatch.delenv("ART_TRAINER_RANK_TEST_HOOKS", raising=False) + monkeypatch.delenv("ART_TRAINER_RANK_TEST_MEMORY_LIMIT_BYTES", raising=False) + 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): + 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() + + +def test_pending_credit_changes_admission_before_execution(budget): + rank, stats = budget + check = rank._memory_check_required(130) + assert check.available_bytes == 120 + 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 + + +def test_split_and_private_credit_remain_explicit_residuals(budget): + rank, stats = budget + stats["active_bytes.all.current"] = 10 + stats["inactive_split_bytes.all.current"] = 70 + assert rank._available_memory_bytes() == 140 + 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 + + +@pytest.mark.parametrize( + "field,value", + [ + ("allocated", None), + ("active", None), + ("reserved", None), + ("allocated", True), + ("active", True), + ("reserved", 80.0), + ("allocated", -1), + ("active", 9), + ("active", 81), + ("reserved", -1), + ], +) +def test_incomplete_native_counters_grant_no_cache_credit(budget, field, value): + rank, stats = budget + key = field + "_bytes.all.current" + if value is None: + stats.pop(key) + else: + stats[key] = value + assert rank._available_memory_bytes() == 70 + torch.cuda.memory_allocated.assert_called_once_with(rank.device) + + +@pytest.mark.parametrize("backend", ["cudaMallocAsync", "unrecognized"]) +def test_other_backends_retain_legacy_unqualified_credit(budget, monkeypatch, backend): + 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) + + +@pytest.mark.parametrize("missing", [False, True]) +def test_existing_test_limit_stays_relative_to_allocated(budget, monkeypatch, missing): + rank, stats = budget + if missing: + stats.pop("active_bytes.all.current") + monkeypatch.setenv("ART_TRAINER_RANK_TEST_HOOKS", "1") + monkeypatch.setenv(_impl._TEST_MEMORY_LIMIT_ENV, "50") + assert rank._available_memory_bytes() == 40 + 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) + + +@pytest.mark.parametrize( + "api", ["mem_get_info", "get_allocator_backend", "memory_stats"] +) +def test_actual_api_error_identity_is_not_swallowed(budget, monkeypatch, api): + rank, _ = budget + error = RuntimeError("native API failed") + monkeypatch.setattr(torch.cuda, api, Mock(side_effect=error)) + with pytest.raises(RuntimeError) as caught: + rank._available_memory_bytes() + assert caught.value is error + + +def test_missing_counter_fallback_api_error_is_preserved(budget, monkeypatch): + rank, stats = budget + stats.pop("active_bytes.all.current") + error = RuntimeError("allocated read failed") + monkeypatch.setattr(torch.cuda, "memory_allocated", Mock(side_effect=error)) + with pytest.raises(RuntimeError) as caught: + rank._available_memory_bytes() + assert caught.value is error + + +def test_cpu_budget_avoids_all_new_cuda_api_calls(budget, monkeypatch): + rank, _ = budget + rank.device = torch.device("cpu") + monkeypatch.setattr( + torch.cuda, "get_allocator_backend", Mock(side_effect=AssertionError) + ) + assert rank._available_memory_bytes() == 1 << 60 + torch.cuda.memory_stats.assert_not_called() + + +def test_required_max_then_available_min_collectives_unchanged(budget, monkeypatch): + rank, _ = budget + monkeypatch.setattr(_impl.dist, "is_available", lambda: True) + monkeypatch.setattr(_impl.dist, "is_initialized", lambda: True) + group = object() + monkeypatch.setattr(rank, "_forward_memory_group", lambda: group) + tensor = torch.tensor + monkeypatch.setattr( + torch, "tensor", lambda values, **kwargs: tensor(values, dtype=kwargs["dtype"]) + ) + calls = [] + + def reduce(value, op, group): + calls.append((float(value.item()), op, group)) + value.fill_(150 if op == _impl.dist.ReduceOp.MAX else 110) + + monkeypatch.setattr(_impl.dist, "all_reduce", reduce) + check = rank._memory_check_required(130) + assert calls == [ + (130, _impl.dist.ReduceOp.MAX, group), + (120, _impl.dist.ReduceOp.MIN, group), + ] + assert (check.estimated_required_bytes, check.available_bytes, check.fits) == ( + 150, + 110, + False, + )