[Graph] Add tests for member-ndarray yield_on= / graph_do_while + schema-v3 round-trip

Follows up on the member-ndarray support commit with a wider test surface:

- Behavioural yield/resume for `yield_on=params.flag` (dataclass mirror of the existing data_oriented test).
- Error paths for `yield_on=self.nonexistent_attr`, `qd.graph_do_while(self.nonexistent_attr)`, and
  `yield_on=self.scalar_attr` (non-ndarray) -- pins the user-facing diagnostic for the attribute forms.
- Nested `qd.graph_do_while(self.outer)` containing `qd.graph_do_while(self.inner)` -- exercises the level-table
  machinery with `@qd.data_oriented` member ndarrays end-to-end.
- Direct ``CacheValue`` round-trip unit test for schema v3 (`cachevalue-v3-ast-resolved-ids`): covers the new
  3-tuple `graph_do_while_levels` + `checkpoint_yield_on_args` / `checkpoint_yield_on_cpp_arg_ids` /
  `checkpoint_user_labels_by_cp_id` fields the loader/storer now plumb through.
- Cross-process fast-cache restore test for a `@qd.kernel(graph=True, checkpoints=True, fastcache=True)` kernel
  with `yield_on=self.flag` -- without the schema-v3 restore the launch path's
  `forward_yield_on_table_to_ctx` would be a no-op and yield/resume would silently break on fast-cached
  checkpoint kernels.

Also fixes an existing `test_src_hasher_store_validate` assertion that indexed `src_hasher.load(...)` by
position; the loader now returns a `CacheValue` (or `None`) so the test is updated to use attribute access plus
asserts the new default-empty fields are present on the round-tripped object.

Author: hughperkins

tests/python/quadrants/lang/fast_caching/test_src_hasher.py

@@ -109,19 +109,17 @@ def get_fileinfos(functions: list[Callable]) -> list[_wrap_inspect.FunctionSourc
 
     kernel_cache_key = "I'm a kernel cache key"
 
-    load_res = src_hasher.load(fast_cache_key)
-    assert load_res[0] is None and load_res[1] is None
+    assert src_hasher.load(fast_cache_key) is None
 
     some_used_vars = {"fee", "fi", "fo"}
     src_hasher.store(kernel_cache_key, fast_cache_key, fileinfos, some_used_vars)
 
     def assert_loaded(cache_key: str) -> None:
         res = src_hasher.load(cache_key)
-        assert res[0] is not None and res[1] is not None
+        assert res is not None and res.frontend_cache_key == kernel_cache_key
 
     def assert_not_loaded(cache_key: str) -> None:
-        res = src_hasher.load(cache_key)
-        assert res[0] is None and res[1] is None
+        assert src_hasher.load(cache_key) is None
 
     assert_loaded(fast_cache_key)
 
@@ -136,7 +134,66 @@ def assert_not_loaded(cache_key: str) -> None:
 
     assert_not_loaded("abcdefg")
 
-    assert src_hasher.load(fast_cache_key)[0] == some_used_vars
+    loaded = src_hasher.load(fast_cache_key)
+    assert loaded is not None
+    assert loaded.used_py_dataclass_parameters == some_used_vars
+    # The new schema-v3 AST-resolved fields default to empty for kernels with no graph_do_while / checkpoint
+    # metadata, exercising the BaseModel default path on round-trip.
+    assert loaded.graph_do_while_levels is None
+    assert loaded.checkpoint_yield_on_args == []
+    assert loaded.checkpoint_yield_on_cpp_arg_ids == []
+    assert loaded.checkpoint_user_labels_by_cp_id == []
+
+
+@test_utils.test()
+def test_src_hasher_store_validate_round_trips_schema_v3_metadata(
+    monkeypatch: pytest.MonkeyPatch, tmp_path: pathlib.Path, temporary_module
+) -> None:
+    """Schema v3 (`cachevalue-v3-ast-resolved-ids`) added AST-resolved arg-id fields to the persisted ``CacheValue``
+    so the launch path can forward them after a fast-cache restore (which skips AST transformation). This test
+    pins the round-trip for the new fields -- ``graph_do_while_levels`` as 3-tuples carrying ``cond_cpp_arg_id``,
+    plus ``checkpoint_yield_on_args`` / ``checkpoint_yield_on_cpp_arg_ids`` / ``checkpoint_user_labels_by_cp_id``.
+    Without this, a schema bug (wrong tuple arity, dropped field, mis-typed BaseModel default) would only surface
+    via a hard-to-debug functional regression in a fast-cached checkpoint / graph_do_while kernel."""
+    test_files_path = pathlib.Path("tests/python/quadrants/lang/fast_caching/test_files")
+
+    offline_cache_path = tmp_path / "cache"
+    temp_import_path = tmp_path / "temp_import"
+    temp_import_path.mkdir(exist_ok=True)
+
+    qd_init_same_arch(offline_cache_file_path=str(offline_cache_path))
+
+    monkeypatch.syspath_prepend(temp_import_path)
+    shutil.copy2(test_files_path / "child_diff_base.py", temp_import_path / "child_diff_schema_v3.py")
+    mod = temporary_module("child_diff_schema_v3")
+    info, _src = _wrap_inspect.get_source_info_and_src(mod.f1.fn)
+    fileinfos = [info]
+    fast_cache_key = src_hasher.create_cache_key(False, info, [], [])
+    assert fast_cache_key is not None
+
+    gdw_levels = [("self.outer", -1, 4), ("self.inner", 0, 5)]
+    cp_yield_args = ["self.flag", None, "params.flag"]
+    cp_yield_cpp_ids = [3, -1, 7]
+    cp_user_labels = [10, None, 20]
+
+    src_hasher.store(
+        "kernel_cache_key_v3",
+        fast_cache_key,
+        fileinfos,
+        {"used_var"},
+        graph_do_while_levels=gdw_levels,
+        checkpoint_yield_on_args=cp_yield_args,
+        checkpoint_yield_on_cpp_arg_ids=cp_yield_cpp_ids,
+        checkpoint_user_labels_by_cp_id=cp_user_labels,
+    )
+
+    loaded = src_hasher.load(fast_cache_key)
+    assert loaded is not None
+    assert loaded.frontend_cache_key == "kernel_cache_key_v3"
+    assert loaded.graph_do_while_levels == [("self.outer", -1, 4), ("self.inner", 0, 5)]
+    assert loaded.checkpoint_yield_on_args == cp_yield_args
+    assert loaded.checkpoint_yield_on_cpp_arg_ids == cp_yield_cpp_ids
+    assert loaded.checkpoint_user_labels_by_cp_id == cp_user_labels
 
 
 # Should be enough to run these on cpu I think, and anything involving

tests/python/test_checkpoint.py

@@ -19,16 +19,24 @@
 (IF conditional node count) are guarded behind ``_is_checkpoint_if_path_native``.
 """
 
+import os
+import pathlib
+import subprocess
+import sys
 from enum import IntEnum
 
 import numpy as np
+import pydantic
 import pytest
 
 import quadrants as qd
 from quadrants.lang import impl
 
 from tests import test_utils
 
+TEST_RAN = "test ran"
+RET_SUCCESS = 42
+
 
 def _on_cuda():
     return impl.current_cfg().arch == qd.cuda
@@ -927,6 +935,99 @@ def step(params: Params):
     assert len(cpp_ids) == 1 and cpp_ids[0] >= 0
 
 
+@test_utils.test()
+def test_checkpoint_yield_on_dataclass_member_yields_and_resumes():
+    """Behavioural round-trip for `yield_on=params.flag` -- mirror of the `self.flag` test below, using a
+    `@dataclasses.dataclass` kernel parameter instead of a `@qd.data_oriented` owner. The dataclass-member access
+    is pre-rewritten to a flattened parameter, so verifying the full yield/resume contract end-to-end is the only
+    way to confirm the right ndarray is wired up at launch."""
+    import dataclasses  # pylint: disable=import-outside-toplevel
+
+    if not _supports_checkpoint_yield_resume():
+        pytest.skip("backend does not implement checkpoint yield/resume")
+    N = 4
+
+    @dataclasses.dataclass
+    class Params:
+        x: qd.types.NDArray[qd.i32, 1]
+        flag: qd.types.NDArray[qd.i32, 0]
+
+    @qd.kernel(graph=True, checkpoints=True)
+    def step(params: Params):
+        with qd.checkpoint(7, yield_on=params.flag):
+            for i in range(params.x.shape[0]):
+                params.x[i] = params.x[i] + 1
+                params.flag[()] = 1
+        with qd.checkpoint(8, yield_on=params.flag):
+            for i in range(params.x.shape[0]):
+                params.x[i] = params.x[i] + 10
+
+    params = Params(
+        x=qd.ndarray(qd.i32, shape=(N,)),
+        flag=qd.ndarray(qd.i32, shape=()),
+    )
+    params.x.from_numpy(np.zeros(N, dtype=np.int32))
+    params.flag.from_numpy(np.array(0, dtype=np.int32))
+    status = step(params)
+    assert status.yielded
+    assert status.checkpoint == 7
+    np.testing.assert_array_equal(params.x.to_numpy(), np.ones(N, dtype=np.int32))
+    params.flag.from_numpy(np.array(0, dtype=np.int32))
+    status = step.resume(from_checkpoint=8)
+    assert not status.yielded
+    np.testing.assert_array_equal(params.x.to_numpy(), np.full(N, 11, dtype=np.int32))
+
+
+@test_utils.test()
+def test_checkpoint_yield_on_member_nonexistent_attribute_raises():
+    """`yield_on=self.nonexistent_attr` (attribute does not exist on the `@qd.data_oriented` owner) must raise a
+    user-facing `QuadrantsSyntaxError` at the `with` site -- the AST-time resolver wraps the underlying attribute
+    lookup failure in the same `does not resolve to an ndarray kernel parameter` diagnostic as the bare-name
+    nonexistent case, so users see one consistent error pattern."""
+    N = 4
+
+    @qd.data_oriented
+    class Sim:
+        def __init__(self):
+            self.x = qd.ndarray(qd.i32, shape=(N,))
+            self.flag = qd.ndarray(qd.i32, shape=())
+
+        @qd.kernel(graph=True, checkpoints=True)
+        def step(self):
+            with qd.checkpoint(0, yield_on=self.nonexistent_flag):  # type: ignore[attr-defined]
+                for i in range(self.x.shape[0]):
+                    self.x[i] = self.x[i] + 1
+
+    sim = Sim()
+    with pytest.raises(qd.QuadrantsSyntaxError, match="does not resolve to an ndarray kernel parameter"):
+        sim.step()
+
+
+@test_utils.test()
+def test_checkpoint_yield_on_member_non_ndarray_attribute_raises():
+    """`yield_on=self.scalar` where `self.scalar` is a Python int (not an ndarray) must raise the same
+    `does not resolve to an ndarray kernel parameter` diagnostic -- the AST-time resolver builds the expression
+    but rejects it because the resulting Expr is not an `ExternalTensorExpression`. Pinning this so future
+    refactors of the resolver can't silently accept non-ndarray attributes and crash later in the launcher."""
+    N = 4
+
+    @qd.data_oriented
+    class Sim:
+        def __init__(self):
+            self.x = qd.ndarray(qd.i32, shape=(N,))
+            self.scalar = 7
+
+        @qd.kernel(graph=True, checkpoints=True)
+        def step(self):
+            with qd.checkpoint(0, yield_on=self.scalar):
+                for i in range(self.x.shape[0]):
+                    self.x[i] = self.x[i] + 1
+
+    sim = Sim()
+    with pytest.raises(qd.QuadrantsSyntaxError, match="does not resolve to an ndarray kernel parameter"):
+        sim.step()
+
+
 @test_utils.test()
 def test_checkpoint_yield_on_data_oriented_member_yields_and_resumes():
     """Behavioural round-trip for `yield_on=self.flag`: setting the member flag from inside the kernel yields, and
@@ -967,6 +1068,102 @@ def step(self):
     np.testing.assert_array_equal(sim.x.to_numpy(), np.full(N, 11, dtype=np.int32))
 
 
+# Module-level kernel for the fastcache-restoration test below. Lives outside any test so the child subprocess can
+# import the test module and reach it without re-creating the (closure-captured) outer scope. The kernel has to be
+# annotated with `fastcache=True` (=> implies `pure`) and lifted out of any decorator-bound owner so it qualifies
+# for the src_ll_cache path. We model the data_oriented owner as the `_FastcacheYieldOnSelfCheckpoint` class below.
+
+
+@qd.data_oriented
+class _FastcacheYieldOnSelfCheckpoint:
+    def __init__(self, n: int):
+        self.x = qd.ndarray(qd.i32, shape=(n,))
+        self.flag = qd.ndarray(qd.i32, shape=())
+
+    @qd.kernel(graph=True, checkpoints=True, fastcache=True)
+    def step(self):
+        with qd.checkpoint(0, yield_on=self.flag):
+            for i in range(self.x.shape[0]):
+                self.x[i] = self.x[i] + 1
+
+
+class _FastcacheCheckpointArgs(pydantic.BaseModel):
+    arch: str
+    offline_cache_file_path: str
+    expect_loaded_from_fastcache: bool
+
+
+def _fastcache_checkpoint_child(args: list[str]) -> None:
+    args_obj = _FastcacheCheckpointArgs.model_validate_json(args[0])
+    qd.init(
+        arch=getattr(qd, args_obj.arch),
+        offline_cache=True,
+        offline_cache_file_path=args_obj.offline_cache_file_path,
+        src_ll_cache=True,
+    )
+
+    N = 8
+    sim = _FastcacheYieldOnSelfCheckpoint(N)
+    sim.x.from_numpy(np.zeros(N, dtype=np.int32))
+    sim.flag.from_numpy(np.array(0, dtype=np.int32))
+    sim.step()
+    np.testing.assert_array_equal(sim.x.to_numpy(), np.ones(N, dtype=np.int32))
+
+    primal = type(sim).step._primal
+    # The schema-v3 fast-cache restore path must repopulate `checkpoint_yield_on_args` and
+    # `checkpoint_yield_on_cpp_arg_ids` from the cached `CacheValue` (since AST transformation is skipped on a
+    # cache hit). A regression here would surface as an empty `_forward_yield_on_table_to_ctx` call, silently
+    # breaking yield/resume on fast-cached checkpoint kernels.
+    labels = primal.checkpoint_yield_on_args
+    cpp_ids = primal.checkpoint_yield_on_cpp_arg_ids
+    assert (
+        labels and len(labels) == 1 and labels[0] is not None and "flag" in labels[0]
+    ), f"checkpoint_yield_on_args should round-trip with one slot containing 'flag', got {labels!r}"
+    assert (
+        len(cpp_ids) == 1 and cpp_ids[0] >= 0
+    ), f"checkpoint_yield_on_cpp_arg_ids should round-trip with one valid id, got {cpp_ids!r}"
+    assert primal.checkpoint_user_labels_by_cp_id == [
+        0
+    ], f"checkpoint_user_labels_by_cp_id should round-trip as [0], got {primal.checkpoint_user_labels_by_cp_id!r}"
+    assert primal.src_ll_cache_observations.cache_loaded == args_obj.expect_loaded_from_fastcache, (
+        f"cache_loaded={primal.src_ll_cache_observations.cache_loaded!r} but expected "
+        f"{args_obj.expect_loaded_from_fastcache!r}"
+    )
+
+    print(TEST_RAN)
+    sys.exit(RET_SUCCESS)
+
+
+@test_utils.test()
+def test_checkpoint_fastcache_restores_self_member_yield_on(tmp_path: pathlib.Path):
+    """After a fast-cache restore in a fresh process, a `@qd.kernel(graph=True, checkpoints=True, fastcache=True)`
+    kernel with `yield_on=self.flag` must repopulate `checkpoint_yield_on_args` /
+    `checkpoint_yield_on_cpp_arg_ids` / `checkpoint_user_labels_by_cp_id` from the persisted ``CacheValue`` --
+    not from the AST transformer, which is skipped on a cache hit. Without the schema-v3 round-trip the launch
+    path's `forward_yield_on_table_to_ctx` would be a no-op and yield/resume would silently break for fast-cached
+    checkpoint kernels."""
+    assert qd.lang is not None
+    arch = qd.lang.impl.current_cfg().arch.name
+    env = dict(os.environ)
+    env["PYTHONPATH"] = "."
+
+    for expect_loaded in [False, True]:
+        args_obj = _FastcacheCheckpointArgs(
+            arch=arch,
+            offline_cache_file_path=str(tmp_path / "cache"),
+            expect_loaded_from_fastcache=expect_loaded,
+        )
+        cmd_line = [sys.executable, __file__, _fastcache_checkpoint_child.__name__, args_obj.model_dump_json()]
+        proc = subprocess.run(cmd_line, capture_output=True, text=True, env=env)
+        if proc.returncode != RET_SUCCESS:
+            print(" ".join(cmd_line))
+            print(proc.stdout)
+            print("-" * 100)
+            print(proc.stderr)
+        assert TEST_RAN in proc.stdout
+        assert proc.returncode == RET_SUCCESS
+
+
 # ----------------------------------------------------------------------------------------------------------------------
 # CUDA-native introspection (slice 1c).
 # ----------------------------------------------------------------------------------------------------------------------
@@ -1001,3 +1198,11 @@ def k(x: qd.types.ndarray(qd.i32, ndim=1), flag: qd.types.ndarray(qd.i32, ndim=0
         assert (
             _num_checkpoints_on_last_call() == 3
         ), f"expected 3 IF conditional nodes (2 implicit + 1 explicit), got {_num_checkpoints_on_last_call()}"
+
+
+# Subprocess dispatch for fast-cache restoration tests above (mirrors the pattern in `test_graph_do_while.py`). The
+# parent test invokes us via `subprocess.run([sys.executable, __file__, <child_fn_name>, <json_args>])` so the
+# child runs in a fresh interpreter with a clean `qd.init` -- the only way to exercise the cross-process fast-cache
+# load path that ``Kernel._try_load_fastcache`` takes after a previous run has populated the on-disk cache.
+if __name__ == "__main__":
+    globals()[sys.argv[1]](sys.argv[2:])