CelerisLab/tests/unit/test_build_contracts.py

"""CPU-only regressions for configuration and isolated compilation inputs."""
import hashlib
import json
import os
from pathlib import Path
import re
import subprocess
import sys
import tempfile
from concurrent.futures import ThreadPoolExecutor
from threading import Barrier
import unittest
from unittest.mock import patch

from CelerisLab.config import LBMConfig, load_lbm_config
from CelerisLab.cuda import compiler


class ConfigContracts(unittest.TestCase):
    def test_explicit_missing_config_never_falls_back(self):
        with tempfile.TemporaryDirectory() as d:
            with self.assertRaises(FileNotFoundError):
                load_lbm_config(str(Path(d) / "missing.json"))
        with self.assertRaises(FileNotFoundError):
            load_lbm_config("")

    def test_float_literals_preserve_selected_physics(self):
        cfg = LBMConfig(rho=1.04, omega_min=.01234, outlet_blend_alpha=.712345)
        macros = cfg.to_macros()
        self.assertEqual(float(macros["RHO"]), cfg.rho)
        self.assertEqual(float(macros["OMEGA_COLLISION_MIN"].rstrip("f")), cfg.omega_min)
        self.assertEqual(float(macros["OUTLET_BLEND_ALPHA"].rstrip("f")), cfg.outlet_blend_alpha)
        cfg.rho = .04
        self.assertEqual(float(cfg.to_macros()["RHO"]), .04)

    def test_invalid_inputs_raise_even_without_assertions(self):
        for kwargs in ({"nx": 0}, {"threads_per_block": 0}, {"threads_per_block": 2048},
                       {"rho": -1}, {"rho": float("nan")}, {"viscosity": float("inf")},
                       {"inlet_profile": "missing"}, {"outlet_mode": "missing"}):
            with self.subTest(kwargs=kwargs), self.assertRaises(ValueError):
                LBMConfig(**kwargs).validate()

    def test_d3_schema_remains_distinct_from_runtime_capability(self):
        cfg = LBMConfig(lattice_model="D3Q19", nz=8)
        cfg.validate()
        self.assertEqual((cfg.dim, cfg.nq), (3, 19))
        with self.assertRaisesRegex(NotImplementedError, "dimension-complete"):
            cfg.validate_runtime()

    def test_runtime_grid_requires_interior_and_ghost_cells(self):
        for kwargs in ({"nx": 1}, {"nx": 2}, {"ny": 1}, {"ny": 2}, {"ny": 3}):
            with self.subTest(kwargs=kwargs):
                cfg = LBMConfig(**kwargs)
                cfg.validate()
                with self.assertRaisesRegex(ValueError, "wall-adjacent rows"):
                    cfg.validate_runtime()
        LBMConfig(nx=3, ny=4).validate_runtime()

    def test_allocation_signature_tracks_layout_not_collision(self):
        cfg = LBMConfig()
        signature = cfg.allocation_signature()
        cfg.collision = "MRT"
        cfg.threads_per_block = 128
        self.assertEqual(signature, cfg.allocation_signature())
        for name, value in (("nx", 1024), ("store_precision", "FP16S"),
                            ("streaming", "esopull"), ("ddf_shifting", True)):
            changed = LBMConfig(**{name: value})
            self.assertNotEqual(signature, changed.allocation_signature())

    def test_cpu_imports_do_not_require_pycuda(self):
        # Refuse the dependency even on development machines with it installed.
        code = '''
import sys
sys.modules["pycuda"] = None
from CelerisLab.config import LBMConfig
from CelerisLab.lbm.equilibrium import compute_feq_d2q9
from CelerisLab.body.geometry.circle import CircleGeometry
from CelerisLab.cuda import compiler
assert "pycuda.driver" not in sys.modules
assert abs(sum(compute_feq_d2q9(1.,0.,0.)) - 1.) < 1e-6
'''
        subprocess.run([sys.executable, "-c", code], check=True)


class BuildContracts(unittest.TestCase):
    def test_block_root_modes_are_rendered_and_recorded(self):
        with tempfile.TemporaryDirectory() as root:
            for mode in (0, 1, 2):
                with self.subTest(mode=mode):
                    build = compiler.generate_config(LBMConfig(), build_root=root, block_root_bc=mode)
                    try:
                        method = Path(build, "config/config_method.h").read_text()
                        self.assertIn(f"#define BLOCK_ROOT_BC {mode}\n", method)
                        metadata = json.loads(Path(build, ".celeris-build").read_text())
                        self.assertEqual(metadata["block_root_bc"], mode)
                        self.assertNotIn("block_root_bc", metadata["config"])
                    finally:
                        compiler.cleanup_build_dir(build)
            build = compiler.generate_config(LBMConfig(), build_root=root)
            try:
                self.assertEqual(json.loads(Path(build, ".celeris-build").read_text())["block_root_bc"], 0)
            finally:
                compiler.cleanup_build_dir(build)

    def test_invalid_block_root_mode_fails_before_snapshot(self):
        for mode in (-1, 3, True, False, 1.0, "1", None):
            with self.subTest(mode=mode), patch.object(compiler.tempfile, "mkdtemp") as allocate:
                with self.assertRaisesRegex(ValueError, "block_root_bc"):
                    compiler.generate_config(LBMConfig(), block_root_bc=mode)
                allocate.assert_not_called()

    @staticmethod
    def fake_compile(command, check):
        source = Path(command[-3])
        grid = (source.parent / "config/config_grid.h").read_text()
        objects = (source.parent / "config/config_objects.h").read_text()
        data = {"nx": int(re.search(r"#define NX\s+(\d+)", grid)[1]),
                "objects": int(re.search(r"#define N_OBJS\s+(\d+)", objects)[1])}
        Path(command[-1]).write_text(json.dumps(data))

    def test_parallel_builds_never_share_headers_or_outputs(self):
        sources = list(Path(compiler.kernel_path("config")).glob("*.h"))
        before = {p: hashlib.sha256(p.read_bytes()).hexdigest() for p in sources}
        barrier = Barrier(2)
        with tempfile.TemporaryDirectory() as root:
            def work(nx, count):
                build = compiler.generate_config(LBMConfig(nx=nx), count, build_root=root)
                barrier.wait()
                output = compiler.compile_kernel(build_dir=build)
                data = json.loads(Path(output).read_text())
                self.assertEqual(data, {"nx": nx, "objects": count})
                metadata = json.loads(Path(output + ".json").read_text())
                self.assertEqual(metadata["config"]["nx"], nx)
                compiler.cleanup_build_dir(build)
                return output
            with patch.object(compiler.subprocess, "run", self.fake_compile):
                with ThreadPoolExecutor(max_workers=2) as pool:
                    a, b = pool.submit(work, 64, 1), pool.submit(work, 256, 3)
                    self.assertNotEqual(a.result(), b.result())
        self.assertEqual(before, {p: hashlib.sha256(p.read_bytes()).hexdigest() for p in sources})

    def test_explicit_build_survives_interleaved_generation(self):
        with tempfile.TemporaryDirectory() as root:
            a = compiler.generate_config(LBMConfig(nx=64), 1, build_root=root)
            b = compiler.generate_config(LBMConfig(nx=128), 2, build_root=root)
            with patch.object(compiler.subprocess, "run", self.fake_compile):
                output = compiler.compile_kernel(build_dir=a)
            self.assertEqual(json.loads(Path(output).read_text())["nx"], 64)
            self.assertTrue(Path(b, "config/config_grid.h").exists())
            with self.assertRaises(ValueError):
                compiler.compile_kernel(build_dir=a, output="../unsafe.ptx")
            with self.assertRaises(ValueError):
                compiler.cleanup_build_dir(compiler.kernel_path(""))
            compiler.cleanup_build_dir(a)
            compiler.cleanup_build_dir(a)
            compiler.cleanup_build_dir(b)
            with self.assertRaises(RuntimeError):
                compiler.compile_kernel(build_dir=a)  # removed snapshot is never compiled

    def test_failed_compile_cannot_implicitly_load_previous_artifact(self):
        with tempfile.TemporaryDirectory() as root:
            build = compiler.generate_config(LBMConfig(nx=64), build_root=root)
            with patch.object(compiler.subprocess, "run", self.fake_compile):
                previous = compiler.compile_kernel(build_dir=build)
            old_bytes = Path(previous).read_bytes()
            failure = subprocess.CalledProcessError(1, "nvcc")
            with patch.object(compiler.subprocess, "run", side_effect=failure):
                with self.assertRaises(subprocess.CalledProcessError):
                    compiler.compile_kernel(build_dir=build)
            self.assertEqual(Path(previous).read_bytes(), old_bytes)
            self.assertEqual(sorted(Path(build).glob("uniform-*.ptx")), [])  # no partial output
            compiler.cleanup_build_dir(build)


if __name__ == "__main__":
    unittest.main()