CelerisLab/tests/unit/test_step_sync.py

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