Files
CelerisLab/tests/unit/test_block_step_sync.py
T
ClaudeandClaude Opus 5.5 1cbdc479cf test_block_step_sync: locate block_stepper without importing it
The module-level import of CelerisLab.lbm.block_stepper loaded
pycuda.driver into the shared test process, unlike the other CPU host
tests (which load mocked copies via AST). find_spec gives the source path
without the import. Test content and assertions unchanged.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
2026-10-21 01:39:52 +00:00

105 lines
4.2 KiB
Python

"""Host contract of PeriodicBlockStepper.step(n): one stream wait per call (mocked driver)."""
import ast
from contextlib import nullcontext
import importlib.util
from pathlib import Path
import types
from unittest.mock import Mock
import unittest
import numpy as np
from CelerisLab.config import LBMConfig
from CelerisLab.lbm.blocks import BlockLayout2D
def host_module():
"""block_stepper with pycuda/compiler replaced by mocks; no GPU context."""
# Locate the source without importing it: importing block_stepper loads pycuda.driver,
# which CPU-only import-boundary tests (test_build_contracts) must never see.
path = Path(importlib.util.find_spec("CelerisLab.lbm.block_stepper").origin)
tree = ast.parse(path.read_text())
tree.body = [node for node in tree.body if not (
isinstance(node, ast.Import) and any(a.name == "pycuda.driver" for a in node.names)
or isinstance(node, ast.ImportFrom) and node.module == "cuda")]
module = types.ModuleType("CelerisLab.lbm._test_host_block_stepper")
module.__package__ = "CelerisLab.lbm"
module.__dict__.update(cuda=Mock(), compiler=Mock())
exec(compile(tree, str(path), "exec"), module.__dict__)
return module
def backend():
host, log = host_module(), []
host.cuda.mem_get_info.return_value = (1 << 40, 1 << 40)
host.cuda.mem_alloc.side_effect = lambda size: Mock(size=size) # distinct buffers
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[:2])))
host.compiler.load_module.return_value = module
cfg = LBMConfig(nx=16, ny=16, viscosity=.1)
run = host.PeriodicBlockStepper(cfg, BlockLayout2D.build(16, 16, 8, 8),
context=Mock(activate=nullcontext, sm_arch="sm_70"))
run._stream.synchronize.side_effect = lambda: log.append(("sync", ()))
run.initialize(np.full((9, 16, 16), 1 / 9))
log.clear()
return run, log
class BlockStepSync(unittest.TestCase):
def test_step_n_waits_once_and_matches_n_single_steps(self):
many, log_many = backend()
many.step(3)
single, log_single = backend()
for _ in range(3):
single.step(1)
names = [name for name, _ in log_many]
self.assertEqual(names.count("sync"), 1)
self.assertEqual(names[-1], "sync")
self.assertEqual(names.count("BlockStep"), 3)
# Same launch order and buffer roles as three step(1) calls.
strip = lambda log: [(n, tuple(id(x) for x in a)) for n, a in log if n != "sync"]
roles = lambda run, log: [(n, tuple({id(run._current): "c", id(run._spare): "s"}.get(x, "-")
for x in ids)) for n, ids in strip(log)]
self.assertEqual([n for n, _ in strip(log_many)], [n for n, _ in strip(log_single)])
self.assertEqual(roles(many, log_many), roles(single, log_single))
self.assertIsNot(many._current, many._spare)
self.assertEqual((many.step_count, single.step_count), (3, 3))
def test_step_zero_launches_and_waits_for_nothing(self):
run, log = backend()
run.step(0)
self.assertEqual((log, run.step_count), ([], 0))
def test_failure_keeps_entry_step_count(self):
run, log = 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_launch_failure_mid_batch_fails_closed_without_waiting(self):
run, log = backend()
calls = []
step = run._step
def failing(*args, **kwargs):
calls.append(1)
if len(calls) == 2:
raise RuntimeError("launch failure")
return step(*args, **kwargs)
run._step = failing
with self.assertRaises(RuntimeError):
run.step(3)
self.assertNotIn("sync", [name for name, _ in log])
self.assertEqual((run.step_count, run._status), (0, "failed"))
with self.assertRaises(RuntimeError):
run.download()
if __name__ == "__main__":
unittest.main()