Refactor: L3 SceneTestCase — HostCallable as Python DAG + ChipCallable + SubCallable

Extend SceneTestCase to support L3 distributed tests (multi-chip + SubWorker).

L3 CALLABLE mirrors L2 structurally:
- orchestration: Python DAG function (vs C++ binary for L2)
- callables: list of ChipCallable entries (compiled) + SubCallable entries
  (registered), distinguished by field presence

Key additions:
- CallableNamespace: dot-access container for orch functions to access
  compiled ChipCallables and registered SubCallable IDs by name, with
  keep() for lifetime management of transient objects past drain()
- _compile_chip_callable_from_spec: extracted from compile_chip_callable,
  reused by both L2 and L3 compilation paths
- _run_and_validate_l3: builds CallableNamespace, wraps orch in Task,
  compares all tensors against golden
- conftest.py L3 branch: registers SubCallable entries from CALLABLE,
  reads device_count/num_sub_workers from case config dict
- Rewrites test_l3_dependency.py and test_l3_group.py as SceneTestCase
  subclasses with module-level orch + sub functions

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: ChaoWao

conftest.py

@@ -137,15 +137,24 @@ def st_worker(request, st_platform, device_pool):
         device_pool.release(ids)
 
     elif level == 3:
-        max_devices = max((c.get("device_count", 1) for c in cls.CASES), default=1)
-        max_subs = max((c.get("num_sub_workers", 0) for c in cls.CASES), default=0)
+        max_devices = max((c.get("config", {}).get("device_count", 1) for c in cls.CASES), default=1)
+        max_subs = max((c.get("config", {}).get("num_sub_workers", 0) for c in cls.CASES), default=0)
         ids = device_pool.allocate(max_devices)
         if not ids:
             pytest.fail(f"need {max_devices} devices")
 
         from simpler.worker import Worker  # noqa: PLC0415
 
         w = Worker(level=3, device_ids=ids, num_sub_workers=max_subs, platform=st_platform, runtime=runtime)
+
+        # Register SubCallable entries from cls.CALLABLE
+        sub_ids = {}
+        for entry in cls.CALLABLE.get("callables", []):
+            if "callable" in entry:
+                cid = w.register(entry["callable"])
+                sub_ids[entry["name"]] = cid
+        cls._st_sub_ids = sub_ids
+
         w.init()
         yield w
         w.close()

python/bindings/task_interface.cpp

@@ -222,6 +222,14 @@ NB_MODULE(_task_interface, m) {
                 return reinterpret_cast<uint64_t>(&self);
             },
             "Return the memory address of the underlying C++ object."
+        )
+
+        .def_static(
+            "sizeof",
+            []() -> size_t {
+                return sizeof(ChipStorageTaskArgs);
+            },
+            "Return sizeof(ChipStorageTaskArgs) in bytes."
         );
 
     // --- TensorArgType enum ---

python/simpler/worker.py

@@ -272,26 +272,14 @@ def _init_level3(self) -> None:
         device_ids = self._config.get("device_ids", [])
         n_sub = self._config.get("num_sub_workers", 0)
 
-        # 1. Allocate mailboxes
+        # 1. Allocate sub-worker mailboxes
         for _ in range(n_sub):
             shm = SharedMemory(create=True, size=DIST_SUB_MAILBOX_SIZE)
             assert shm.buf is not None
             struct.pack_into("i", shm.buf, _OFF_STATE, _IDLE)
             self._shms.append(shm)
 
-        # 2. Fork SubWorker processes (MUST be before any C++ threads)
-        registry = self._callable_registry
-        for i in range(n_sub):
-            pid = os.fork()
-            if pid == 0:
-                buf = self._shms[i].buf
-                assert buf is not None
-                _sub_worker_loop(buf, registry)
-                os._exit(0)
-            else:
-                self._pids.append(pid)
-
-        # 3. Fork ChipWorker processes (only if device_ids provided)
+        # 2. Prepare chip-worker config (but do NOT fork yet — deferred to _start_level3)
         if device_ids:
             from runtime_builder import RuntimeBuilder  # noqa: PLC0415
 
@@ -302,39 +290,71 @@ def _init_level3(self) -> None:
             builder = RuntimeBuilder(platform)
             binaries = builder.get_binaries(runtime, build=False)
 
-            # Determine args_size (sizeof ChipStorageTaskArgs)
-            _objs = [_CSA() for _ in range(5)]
-            _ptrs = [o.__ptr__() for o in _objs]
-            args_size = min(abs(_ptrs[i + 1] - _ptrs[i]) for i in range(len(_ptrs) - 1))
-            del _objs, _ptrs
-
-            host_lib_path = str(binaries.host_path)
-            aicpu_path = str(binaries.aicpu_path)
-            aicore_path = str(binaries.aicore_path)
-            sim_ctx_path = str(binaries.sim_context_path) if getattr(binaries, "sim_context_path", None) else ""
+            self._l3_args_size = _CSA.sizeof()
+            self._l3_host_lib_path = str(binaries.host_path)
+            self._l3_aicpu_path = str(binaries.aicpu_path)
+            self._l3_aicore_path = str(binaries.aicore_path)
+            self._l3_sim_ctx_path = (
+                str(binaries.sim_context_path) if getattr(binaries, "sim_context_path", None) else ""
+            )
 
-            for dev_id in device_ids:
+            # Allocate chip mailboxes (shared memory, no fork yet)
+            for _ in device_ids:
                 shm = SharedMemory(create=True, size=DIST_CHIP_MAILBOX_SIZE)
                 assert shm.buf is not None
                 struct.pack_into("i", shm.buf, _CHIP_OFF_STATE, _IDLE)
                 self._chip_shms.append(shm)
 
+        self._l3_started = False
+
+    def _start_level3(self) -> None:
+        """Fork child processes and start C++ scheduler. Called on first run()."""
+        if self._l3_started:
+            return
+        self._l3_started = True
+
+        device_ids = self._config.get("device_ids", [])
+        n_sub = self._config.get("num_sub_workers", 0)
+
+        # Fork SubWorker processes (MUST be before any C++ threads)
+        registry = self._callable_registry
+        for i in range(n_sub):
+            pid = os.fork()
+            if pid == 0:
+                buf = self._shms[i].buf
+                assert buf is not None
+                _sub_worker_loop(buf, registry)
+                os._exit(0)
+            else:
+                self._pids.append(pid)
+
+        # Fork ChipWorker processes
+        if device_ids:
+            for idx, dev_id in enumerate(device_ids):
                 pid = os.fork()
                 if pid == 0:
-                    buf = shm.buf
+                    buf = self._chip_shms[idx].buf
                     assert buf is not None
-                    _chip_process_loop(buf, host_lib_path, dev_id, aicpu_path, aicore_path, sim_ctx_path, args_size)
+                    _chip_process_loop(
+                        buf,
+                        self._l3_host_lib_path,
+                        dev_id,
+                        self._l3_aicpu_path,
+                        self._l3_aicore_path,
+                        self._l3_sim_ctx_path,
+                        self._l3_args_size,
+                    )
                     os._exit(0)
                 else:
                     self._chip_pids.append(pid)
 
-        # 4. Create DistWorker and wire chip processes + sub workers
+        # Create DistWorker and wire chip processes + sub workers
         dw = DistWorker(3)
         self._dist_worker = dw
 
         if device_ids:
             for shm in self._chip_shms:
-                cp = DistChipProcess(_mailbox_addr(shm), args_size)
+                cp = DistChipProcess(_mailbox_addr(shm), self._l3_args_size)
                 self._dist_chip_procs.append(cp)
                 dw.add_chip_process(cp)
 
@@ -343,7 +363,7 @@ def _init_level3(self) -> None:
             self._dist_sub_workers.append(sw)
             dw.add_sub_worker(sw)
 
-        # 6. Start Scheduler + WorkerThreads (C++ threads start here, after fork)
+        # Start Scheduler + WorkerThreads (C++ threads start here, after fork)
         dw.init()
 
     # ------------------------------------------------------------------
@@ -377,6 +397,7 @@ def run(self, task_or_payload, args=None, **kwargs) -> None:
                 # run(callable, args, **kwargs)
                 self._chip_worker.run(task_or_payload, args, **kwargs)
         else:
+            self._start_level3()
             assert self._dist_worker is not None
             task = task_or_payload
             task.orch(self, task.args)

simpler_setup/__init__.py

@@ -13,9 +13,10 @@
 from .platform_info import parse_platform
 from .pto_isa import ensure_pto_isa_root
 from .runtime_builder import RuntimeBuilder
-from .scene_test import Scalar, SceneTestCase, TaskArgsBuilder, Tensor, scene_test
+from .scene_test import CallableNamespace, Scalar, SceneTestCase, TaskArgsBuilder, Tensor, scene_test
 
 __all__ = [
+    "CallableNamespace",
     "KernelCompiler",
     "RuntimeBuilder",
     "Scalar",