From 19db7d362fa14d5fe6ff119c26e1bd01adbed62b Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 27 Oct 2026 01:10:15 +0000 Subject: [PATCH] refinement: free-memory preflight for every level before any device work Restores the PeriodicBlockStepper safeguard dropped in 8e1feff, for all FixedPatchStepper levels: both G buffers, halo pairs and boundary columns of every level are compared with cuda.mem_get_info() before the stream, any module or any buffer exists, raising MemoryError('level buffers exceed free device memory'). Transfer tables are not counted (lower bound). Tests: the single-level MemoryError case is back beside the param-layout check (no stream, no allocation); host test_single_level covers depth 1 and depth 2 (no stream, build or allocation) and the wait failure after a good step (count stays 1). Driver-mocking fixtures report ample free memory. Co-Authored-By: Claude Opus 5.5 --- src/CelerisLab/lbm/refinement_stepper.py | 18 +++++++++++++++ .../test_refinement_three_level_multipatch.py | 1 + tests/unit/test_single_level.py | 23 +++++++++++++++++++ .../multigrid/test_block_stepper.py | 10 ++++++-- .../multigrid/test_refinement_blocked_base.py | 1 + .../multigrid/test_refinement_blocked_fine.py | 4 ++-- .../multigrid/test_refinement_followup.py | 4 ++-- .../multigrid/test_refinement_jets.py | 2 +- .../multigrid/test_refinement_three_level.py | 2 +- 9 files changed, 57 insertions(+), 8 deletions(-) diff --git a/src/CelerisLab/lbm/refinement_stepper.py b/src/CelerisLab/lbm/refinement_stepper.py index 5f5a4be..e884616 100644 --- a/src/CelerisLab/lbm/refinement_stepper.py +++ b/src/CelerisLab/lbm/refinement_stepper.py @@ -324,6 +324,9 @@ class FixedPatchStepper: self._block = (cfg.threads_per_block, 1, 1) with self.ctx.activate(): try: + free, _ = cuda.mem_get_info() + if self._planned_bytes() > free: # before any stream, module or buffer + raise MemoryError("level buffers exceed free device memory") self._stream = cuda.Stream() for level in range(depth): self._levels.append(self._build_level(level)) @@ -369,6 +372,21 @@ class FixedPatchStepper: pass # Preserve the original construction failure. raise + def _planned_bytes(self): + """Both G buffers, halo pairs and boundary columns of every level (transfer tables excluded).""" + total = 0 + for index, layout in enumerate(self._layouts): + cells = int(np.prod(layout.blocked_shape if layout is not None + else (9, *self.hierarchy.kinds[index].shape))) + total += 2 * 4 * cells + if layout is not None: + total += layout.halo_dst.nbytes + layout.halo_src.nbytes + b = self._level_boundaries[index] + if b is not None: + total += sum(np.asarray(v).nbytes for k, v in b.tables.items() + if k == "flags" or k in BLOCKED_BOUNDARY_DEVICE_COLUMNS) + return total + def _build_level(self, index): """Compile one level's step module and allocate its two G buffers.""" layout, config = self._layouts[index], self._level_cfgs[index] diff --git a/tests/unit/test_refinement_three_level_multipatch.py b/tests/unit/test_refinement_three_level_multipatch.py index 90ed5f3..5a6299d 100644 --- a/tests/unit/test_refinement_three_level_multipatch.py +++ b/tests/unit/test_refinement_three_level_multipatch.py @@ -28,6 +28,7 @@ def host_module(): from CelerisLab.body.jets import surface_link_velocity module.__dict__.update(cuda=Mock(),compiler=Mock(), surface_link_velocity=surface_link_velocity) + module.cuda.mem_get_info.return_value=(1<<40,1<<40) # free-memory preflight passes exec(compile(tree,str(path),'exec'),module.__dict__) return module diff --git a/tests/unit/test_single_level.py b/tests/unit/test_single_level.py index b8c85ae..4404e6d 100644 --- a/tests/unit/test_single_level.py +++ b/tests/unit/test_single_level.py @@ -100,6 +100,29 @@ class SingleLevel(unittest.TestCase): with self.assertRaisesRegex(RuntimeError, "failed"): run.download() + def test_free_memory_preflight_before_stream_and_buffers(self): + for cfg, options in ((LBMConfig(nx=32, ny=24, viscosity=.1), + dict(refinements=[], level_blocks=[(12, 16)])), + (LBMConfig(nx=48, ny=48, viscosity=.02), # every level counts + dict(refinements=[[(6, 6, 26, 26)]], level_blocks=[(4, 4)] * 2))): + with self.subTest(depth=len(options["refinements"]) + 1): + host = host_module() + host.cuda.mem_get_info.return_value = (1024, 1 << 30) + with self.assertRaisesRegex(MemoryError, "exceed free device memory"): + host.FixedPatchStepper(cfg, context=Mock(activate=nullcontext, sm_arch="sm_70"), + **options) + host.cuda.Stream.assert_not_called() + host.cuda.mem_alloc.assert_not_called() + host.compiler.build_module.assert_not_called() + + def test_wait_failure_after_a_good_step_keeps_its_count(self): + run, _, _ = backend() + run.step(1) + run._stream.synchronize.side_effect = RuntimeError("device fault") + with self.assertRaises(RuntimeError): + run.step(2) + self.assertEqual((run.step_count, run._status), (1, "failed")) + def test_final_wait_failure_fails_instance_and_keeps_step_count(self): run, _, _ = backend() run._stream.synchronize.side_effect = RuntimeError("device error") diff --git a/tests/validation/multigrid/test_block_stepper.py b/tests/validation/multigrid/test_block_stepper.py index d131497..1cc2c4a 100644 --- a/tests/validation/multigrid/test_block_stepper.py +++ b/tests/validation/multigrid/test_block_stepper.py @@ -159,11 +159,17 @@ class CpuContracts(unittest.TestCase): with self.assertRaisesRegex(ValueError, "launch limits"): single(context=self.forbidden_context()) - def test_param_layout_rejected_before_allocation(self): + def test_memory_and_param_layout_rejected_before_allocation(self): context = Mock(activate=nullcontext, sm_arch="sm_70") + with patch.object(cuda, "mem_get_info", return_value=(1024, 1 << 30)), \ + patch.object(cuda, "Stream") as stream, patch.object(cuda, "mem_alloc") as alloc: + with self.assertRaisesRegex(MemoryError, "exceed free device memory"): + single(context=context, objects=[circle()]) + stream.assert_not_called() + alloc.assert_not_called() module = Mock() module.get_global.return_value = (Mock(), 16) - with patch.object(cuda, "Stream"), patch.object(compiler, "generate_config"), \ + with patch.object(cuda, "Stream"), patch.object(cuda, "mem_get_info", return_value=(1 << 40, 1 << 40)), patch.object(compiler, "generate_config"), \ patch.object(compiler, "compile_kernel") as compile_kernel, \ patch.object(compiler, "cleanup_build_dir"), \ patch.object(compiler, "load_module", return_value=module), \ diff --git a/tests/validation/multigrid/test_refinement_blocked_base.py b/tests/validation/multigrid/test_refinement_blocked_base.py index 6591c62..cfe6232 100644 --- a/tests/validation/multigrid/test_refinement_blocked_base.py +++ b/tests/validation/multigrid/test_refinement_blocked_base.py @@ -73,6 +73,7 @@ def mocked_backend(mode="periodic", fail_alloc=None): stack.enter_context(patch.object(compiler, name)) stack.enter_context(patch.object(compiler, "load_module", return_value=module)) stack.enter_context(patch.object(cuda, "Stream", return_value=stream)) + stack.enter_context(patch.object(cuda, "mem_get_info", return_value=(1 << 40, 1 << 40))) stack.enter_context(patch.object(cuda, "mem_alloc", side_effect=allocate)) stack.enter_context(patch.object(cuda, "memcpy_htod", side_effect=lambda p, a: uploads.append((p, a)))) diff --git a/tests/validation/multigrid/test_refinement_blocked_fine.py b/tests/validation/multigrid/test_refinement_blocked_fine.py index f7f4e08..d82d40a 100644 --- a/tests/validation/multigrid/test_refinement_blocked_fine.py +++ b/tests/validation/multigrid/test_refinement_blocked_fine.py @@ -122,7 +122,7 @@ class BlockedFineCPU(unittest.TestCase): raise AssertionError("zero-byte device allocation") sizes.append(size) return Mock() - with patch.object(cuda, "Stream"), patch.object(compiler, "generate_config"), \ + with patch.object(cuda, "Stream"), patch.object(cuda, "mem_get_info", return_value=(1 << 40, 1 << 40)), patch.object(compiler, "generate_config"), \ patch.object(compiler, "compile_kernel"), patch.object(compiler, "cleanup_build_dir"), \ patch.object(compiler, "load_module", return_value=module), \ patch.object(cuda, "mem_alloc", side_effect=alloc), patch.object(cuda, "memcpy_htod"): @@ -172,7 +172,7 @@ class BlockedFineCPU(unittest.TestCase): return functions.setdefault(name, Mock(side_effect=lambda *a, **k: calls.append((name, a)))) module = Mock(get_function=Mock(side_effect=get_function)) module.get_global.return_value = (Mock(), 20) - with patch.object(cuda, "Stream"), patch.object(compiler, "generate_config"), \ + with patch.object(cuda, "Stream"), patch.object(cuda, "mem_get_info", return_value=(1 << 40, 1 << 40)), patch.object(compiler, "generate_config"), \ patch.object(compiler, "compile_kernel"), patch.object(compiler, "cleanup_build_dir"), \ patch.object(compiler, "load_module", return_value=module), \ patch.object(cuda, "mem_alloc", side_effect=lambda size: Mock()), \ diff --git a/tests/validation/multigrid/test_refinement_followup.py b/tests/validation/multigrid/test_refinement_followup.py index da7384e..6907ab1 100644 --- a/tests/validation/multigrid/test_refinement_followup.py +++ b/tests/validation/multigrid/test_refinement_followup.py @@ -166,7 +166,7 @@ class FaultGuardsCPU(unittest.TestCase): # A pair tau inconsistent with its level config is rejected twice over, both before # any device allocation or upload: first by the hierarchy chain check, and, if that # is bypassed, by the per-level omega/tau check in _build_level. - with patch.object(cuda, "Stream"), bad, patch.object(cuda, "mem_alloc") as alloc, \ + with patch.object(cuda, "Stream"), patch.object(cuda, "mem_get_info", return_value=(1 << 40, 1 << 40)), bad, patch.object(cuda, "mem_alloc") as alloc, \ patch.object(cuda, "memcpy_htod") as upload: with self.assertRaisesRegex(RuntimeError, "config chain is inconsistent"): PeriodicPatchStepper(cfg, context=context, refinements=[p.bounds]) @@ -177,7 +177,7 @@ class FaultGuardsCPU(unittest.TestCase): upload.assert_not_called() module = Mock() module.get_global.return_value = (Mock(), 16) - with patch.object(cuda, "Stream"), patch.object(compiler, "generate_config"), \ + with patch.object(cuda, "Stream"), patch.object(cuda, "mem_get_info", return_value=(1 << 40, 1 << 40)), patch.object(compiler, "generate_config"), \ patch.object(compiler, "compile_kernel"), patch.object(compiler, "cleanup_build_dir"), \ patch.object(compiler, "load_module", return_value=module), \ patch.object(cuda, "memcpy_htod") as upload: diff --git a/tests/validation/multigrid/test_refinement_jets.py b/tests/validation/multigrid/test_refinement_jets.py index 802bfec..de99afe 100644 --- a/tests/validation/multigrid/test_refinement_jets.py +++ b/tests/validation/multigrid/test_refinement_jets.py @@ -36,7 +36,7 @@ def mocked(objects, jets, fine_block): module = Mock(get_function=Mock(return_value=Mock())) module.get_global.return_value = (Mock(), 20) htod = Mock() - with patch.object(cuda, "Stream"), patch.object(compiler, "generate_config"), \ + with patch.object(cuda, "Stream"), patch.object(cuda, "mem_get_info", return_value=(1 << 40, 1 << 40)), patch.object(compiler, "generate_config"), \ patch.object(compiler, "compile_kernel"), patch.object(compiler, "cleanup_build_dir"), \ patch.object(compiler, "load_module", return_value=module), \ patch.object(cuda, "mem_alloc", side_effect=lambda size: Mock()), \ diff --git a/tests/validation/multigrid/test_refinement_three_level.py b/tests/validation/multigrid/test_refinement_three_level.py index b612b23..fc7a1d6 100644 --- a/tests/validation/multigrid/test_refinement_three_level.py +++ b/tests/validation/multigrid/test_refinement_three_level.py @@ -56,7 +56,7 @@ def mocked_child(): context = Mock(activate=nullcontext, sm_arch="sm_70") module = Mock(get_function=Mock(side_effect=lambda name: Mock(name=name))) module.get_global.return_value = (Mock(), 20) - with patch.object(cuda, "Stream"), patch.object(compiler, "generate_config"), \ + with patch.object(cuda, "Stream"), patch.object(cuda, "mem_get_info", return_value=(1 << 40, 1 << 40)), patch.object(compiler, "generate_config"), \ patch.object(compiler, "compile_kernel"), patch.object(compiler, "cleanup_build_dir"), \ patch.object(compiler, "load_module", return_value=module), \ patch.object(cuda, "mem_alloc", side_effect=lambda size: Mock()), \