"""Host contract of FixedPatchStepper.step(n): one stream wait per call (mocked driver). All n base intervals are enqueued in order on the backend stream and the host waits once before step returns. Launch order and buffer arguments must equal n step(1) calls; completion metadata, raw-exchange and window ledgers, and failure invalidation are published only after that wait. """ from contextlib import nullcontext from unittest.mock import Mock import unittest import numpy as np from CelerisLab.config import LBMConfig from tests.unit.test_raw_exchange import body from tests.unit.test_refinement_three_level_multipatch import host_module BODY_LEVEL, HOPS = 3, 8 # body on L3: 2**3 hops per base interval def backend(window=False): """Four-level host backend; every launch, memset and stream wait goes to one log.""" host = host_module() log, allocations = [], [] def allocate(size): allocations.append(Mock(size=size)) return allocations[-1] module = Mock() module.get_global.return_value = (Mock(), 20) module.get_function.side_effect = lambda name: Mock( side_effect=lambda *a, **k: log.append((name, a, k))) host.compiler.load_module.return_value = module host.cuda.mem_alloc.side_effect = allocate host.cuda.memcpy_dtoh.side_effect = lambda array, pointer: array.fill(0) host.cuda.memset_d8_async.side_effect = lambda *a: log.append(("memset", a, {})) run = host.FixedPatchStepper( LBMConfig(nx=48, ny=48, viscosity=.02), context=Mock(activate=nullcontext, sm_arch="sm_70"), refinements=[[(6, 6, 26, 26)], [(20, 20, 44, 44)], [(56, 56, 72, 72)]], level_blocks=[(4, 4)] * 4, level_objects=[(), (), (), (body(),)], observe_raw_exchange=True, observe_window_exchange=window) run._stream.synchronize.side_effect = lambda: log.append(("sync", (), {})) run.initialize(*[np.full(level.shape, 1 / 9, np.float32) for level in run._levels]) log.clear() return run, log, allocations def signature(log, allocations): """Event names with device-pointer arguments replaced by allocation index.""" index = {id(a): i for i, a in enumerate(allocations)} return [(name, tuple(index.get(id(a), "-") for a in args)) for name, args, _ in log] def names(log): return [name for name, _, _ in log] class StepSync(unittest.TestCase): def test_step1_waits_once_after_all_launches(self): run, log, _ = backend() run.step(1) self.assertEqual(names(log).count("sync"), 1) self.assertEqual(names(log)[-1], "sync") self.assertEqual(names(log)[0], "memset") self.assertEqual(names(log).count("BlockStep"), 1 + 2 + 4 + 8) self.assertEqual(run.step_count, 1) run.close() def test_step_n_equals_n_step1_launches_with_one_wait(self): single, single_log, single_alloc = backend() for _ in range(3): single.step(1) batched, batched_log, batched_alloc = backend() batched.step(3) expected = [e for e in signature(single_log, single_alloc) if e[0] != "sync"] got = signature(batched_log, batched_alloc) self.assertEqual(got[-1][0], "sync") self.assertEqual(names(batched_log).count("sync"), 1) self.assertEqual(got[:-1], expected) # same order, same ping-pong buffers self.assertEqual(names(batched_log).count("memset"), 3) self.assertEqual(batched.step_count, single.step_count) index = [{id(a): i for i, a in enumerate(alloc)} for alloc in (single_alloc, batched_alloc)] self.assertEqual(index[1][id(batched._levels[0].current)], index[0][id(single._levels[0].current)]) single.close() batched.close() def test_step0_launches_and_waits_for_nothing(self): run, log, _ = backend() run.step(2) log.clear() run.step(0) self.assertEqual(log, []) self.assertEqual(run.step_count, 2) value = run.read_raw_exchange(0, level=BODY_LEVEL) # step(0) keeps the ledger valid self.assertEqual((value["time_start"], value["time_end"]), (1, 2)) run.close() def test_completion_metadata_covers_all_intervals(self): run, _, _ = backend() with self.assertRaisesRegex(RuntimeError, "completed"): run.read_raw_exchange(0, level=BODY_LEVEL) run.step(3) value = run.read_raw_exchange(0, level=BODY_LEVEL) self.assertEqual((value["time_start"], value["time_end"], value["sample_count"]), (2, 3, HOPS)) # step(n) retains only the final base interval run.step(2) self.assertEqual(run.read_raw_exchange(0, level=BODY_LEVEL)["time_end"], 5) run.close() def test_window_accounting_across_batched_calls(self): windows = [] for calls in ((1, 3), (4,), (1, 1, 1, 1)): run, log, _ = backend(window=True) run.step(1) run.start_exchange_window() log.clear() for n in calls: run.step(n) self.assertEqual(names(log).count("BlockAccumulateExchange"), 2 * 4 * HOPS) self.assertEqual(names(log).count("sync"), len(calls)) out = run.read_exchange_window(0, level=BODY_LEVEL) windows.append((out["time_start"], out["time_end"], out["sample_count"])) run.close() self.assertEqual(windows, [(1., 5, 4 * HOPS)] * 3) def test_failed_wait_publishes_nothing(self): run, _, _ = backend(window=True) run.step(2) run.start_exchange_window() run._stream.synchronize.side_effect = RuntimeError("device failure") with self.assertRaisesRegex(RuntimeError, "device failure"): run.step(3) self.assertEqual(run.step_count, 2) # value at entry to the failing call self.assertFalse(run._exchange_valid) self.assertFalse(run._window_valid) for read in (lambda: run.read_raw_exchange(0, level=BODY_LEVEL), lambda: run.read_exchange_window(0, level=BODY_LEVEL), run.download, lambda: run.step(1)): with self.assertRaisesRegex(RuntimeError, "failed"): read() run._stream.synchronize.side_effect = None run.close() def test_launch_error_mid_batch_skips_wait_and_fails(self): run, log, _ = backend() run.step(1) log.clear() hop, calls = run._hop, [] def failing(index, *args, **kw): if index == 0: calls.append(1) if len(calls) == 2: raise RuntimeError("launch failure") return hop(index, *args, **kw) run._hop = failing with self.assertRaisesRegex(RuntimeError, "launch failure"): run.step(3) self.assertNotIn("sync", names(log)) self.assertEqual(run.step_count, 1) self.assertFalse(run._exchange_valid) with self.assertRaisesRegex(RuntimeError, "failed"): run.read_raw_exchange(0, level=BODY_LEVEL) run.close() if __name__ == "__main__": unittest.main()