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 <noreply@anthropic.com>

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]

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
 

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")

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), \

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))))

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()), \

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:

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()), \

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()), \