From b920a303a49d83465780dc44c2a21ed37ba1c4fb Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 30 Oct 2026 07:58:40 +0000 Subject: [PATCH] simulation: keep the owner context active through moving diagnostics readback moving_diagnostics is context_bound, so stream synchronization and the DTOH copy run inside one owner activation. Tests: owner stays active through sync and download; launch order moving -> static curved -> OneStep with a real static link count and the buffer swap between steps. Co-Authored-By: Claude Opus 5.5 --- src/CelerisLab/simulation.py | 1 + tests/unit/test_moving_bodies.py | 46 +++++++++++++++++++++++++++++--- 2 files changed, 44 insertions(+), 3 deletions(-) diff --git a/src/CelerisLab/simulation.py b/src/CelerisLab/simulation.py index 1bb90d4..3916a83 100644 --- a/src/CelerisLab/simulation.py +++ b/src/CelerisLab/simulation.py @@ -226,6 +226,7 @@ class Simulation: from .body.motion import MovingTables return MovingTables.upload([(obj.motion, obj.obj_id) for obj in moving]) + @context_bound def moving_diagnostics(self) -> np.ndarray: """Per translating circle (in body order): cme_x, cme_y, wen_x, wen_y, link_dm, n_links accumulated since obs were last zeroed (diagnostics only).""" diff --git a/tests/unit/test_moving_bodies.py b/tests/unit/test_moving_bodies.py index c757707..a95d1ca 100644 --- a/tests/unit/test_moving_bodies.py +++ b/tests/unit/test_moving_bodies.py @@ -261,21 +261,61 @@ def test_moving_tables_layout_and_attachment(monkeypatch): assert sim.stepper.moving is None and sim.stepper.blocked is not None +class CountingOwner(Owner): + """Owner double that counts nested activations like CudaContext.""" + def __init__(self): + super().__init__() + self.active = 0 + + def activate(self): + from contextlib import contextmanager + + @contextmanager + def scope(): + self.active += 1 + try: + yield self + finally: + self.active -= 1 + return scope() + + +def test_public_diagnostics_keep_owner_active_through_sync_and_download(monkeypatch): + sim = make_simulation(monkeypatch) + sim.ctx, sim._failed = CountingOwner(), None + seen = [] + sim._stream = SimpleNamespace(synchronize=lambda: seen.append(("sync", sim.ctx.active))) + expected = np.arange(6, dtype=np.float32).reshape(1, 6) + + def download(): + seen.append(("download", sim.ctx.active)) + return expected + sim.stepper = SimpleNamespace(moving=SimpleNamespace(download_diagnostics=download)) + assert sim.moving_diagnostics() is expected + assert [name for name, _ in seen] == ["sync", "download"] + assert all(active > 0 for _, active in seen) and sim.ctx.active == 0 + + def test_stepper_launches_moving_links_before_static_curved(): from CelerisLab.lbm.stepper import LBMStepper launches = [] module = SimpleNamespace(get_function=lambda name: lambda *a, **k: launches.append((name, a, k))) cfg = SimpleNamespace(streaming="double_buffer", threads_per_block=128, is_d2q9=True, nx=64, ny=32, nz=1) - field = SimpleNamespace(ddf_gpu="ddf", temp_gpu="temp", flag_gpu="flag", n_curved=0, - invalidate_host_ddf_cache=lambda: None) + curved = SimpleNamespace(**{f"{name}_gpu": name for name in ( + "fluid_idx", "dir", "q", "rx", "ry", "rz", "fallback_class", "body_id")}, + surface_velocity=None, surface_velocity_gpu=None, link_exchange_gpu=None) + field = SimpleNamespace(ddf_gpu="ddf", temp_gpu="temp", flag_gpu="flag", n_curved=5, + curved=curved, invalidate_host_ddf_cache=lambda: None) stepper = LBMStepper(field, module, cfg) stepper._launch_force_region = stepper._launch_sensor = lambda *a, **k: None stepper.moving = SimpleNamespace(band_gpu="band", geom_gpu="geom", diag_gpu="diag", slot_debug_gpu=None, n_bodies=1, n_slots=300) stepper.step(2, action_gpu="act", obs_gpu="obs") names = [name for name, _, _ in launches] - assert names == ["MovingCircleLinks", "OneStep"] * 2 and stepper.step_count == 2 + assert names == ["MovingCircleLinks", "CurvedBoundaryKernel", "OneStep"] * 2 + assert stepper.step_count == 2 + assert launches[1][1][0] == "ddf" and launches[3][1][0] == "temp" # swap between steps args, kw = launches[0][1], launches[0][2] assert args[:4] == ("ddf", "flag", "band", "geom") and int(args[5]) == 300 assert kw["grid"] == (3, 1, 1) and kw["block"] == (128, 1, 1)