Add: post-init eligibility re-check in Worker.register()

python/simpler/worker.py

@@ -5141,13 +5141,8 @@ def register(self, target, *, workers: list[int] | None = None) -> CallableHandl
         Submit APIs consume the returned handle and dispatch by its stable
         SHA-256 callable identity.
 
-        Target eligibility (a callable's kind having a resolving child) is
-        checked only at init(), over the pre-init registrations
-        (``_validate_eligible_targets``). A post-init dynamic register does NOT
-        re-validate against the frozen topology, so registering e.g. a
-        ChipCallable on a chipless worker yields a handle that never dispatches.
-        Unifying the two paths is a follow-up (needs a device-free chip-child
-        test harness).
+        A post-init dynamic register re-validates eligibility against the
+        frozen topology (``_eligible_target_need``), same as init().
         """
         if isinstance(target, RemoteCallable) and self.level < 4:
             raise TypeError("Worker.register(RemoteCallable): remote L3 dispatch requires a level >= 4 parent")
@@ -5171,6 +5166,11 @@ def register(self, target, *, workers: list[int] | None = None) -> CallableHandl
             handle = self._register_into_snapshot_or_wait(reg)
             if handle is not None:
                 return handle
+            need = self._eligible_target_need(reg.target_namespace, reg.eligible_worker_ids)
+            if need is not None:
+                raise ValueError(
+                    f"Worker.register(): {reg.target_namespace} callable has no eligible dispatch target (needs {need})"
+                )
             # Post-start broadcast touches the live tree; hold a lease so close()
             # drains it before teardown (re-checks READY, closing the
             # gate-then-teardown race).
@@ -5180,6 +5180,11 @@ def register(self, target, *, workers: list[int] | None = None) -> CallableHandl
             handle = self._register_into_snapshot_or_wait(reg)
             if handle is not None:
                 return handle
+            need = self._eligible_target_need(reg.target_namespace, reg.eligible_worker_ids)
+            if need is not None:
+                raise ValueError(
+                    f"Worker.register(): {reg.target_namespace} callable has no eligible dispatch target (needs {need})"
+                )
             # Post-start publication touches the live tree; hold a lease across
             # the whole transaction, publication included, so close() drains it
             # before teardown (re-checks READY, closing the gate-then-teardown
@@ -6009,17 +6014,17 @@ def _eligible_target_need(self, namespace: str | None, eligible_worker_ids) -> s
           - ``REMOTE_TASK_DISPATCHER`` only onto its named remote worker(s).
         An L2 worker (or any non-dispatch namespace) is always eligible.
 
-        Used only by ``_validate_eligible_targets`` at init (the *startup*
-        eligibility gate). The post-init dynamic ``register`` path does NOT yet
-        apply this rule — see that method for the deferred inconsistency.
+        Applied at init by ``_validate_eligible_targets`` and on every
+        post-init dynamic ``register`` path.
         """
         if self.level < 3:
             return None
         if namespace == "LOCAL_PYTHON":
             has_python_child = self._config.get("num_sub_workers", 0) > 0 or bool(self._next_level_workers)
             return None if has_python_child else "a SUB or next-level child"
         if namespace == "LOCAL_CHIP":
-            return None if bool(self._config.get("device_ids")) else "a chip device (device_ids)"
+            has_chip_child = bool(self._config.get("device_ids")) or bool(self._next_level_workers)
+            return None if has_chip_child else "a chip device (device_ids)"
         if namespace == "REMOTE_TASK_DISPATCHER":
             has_remote_workers = set(self._remote_worker_ids)
             ok = bool(has_remote_workers) and set(eligible_worker_ids) <= has_remote_workers

tests/ut/py/test_worker/test_admission_fence.py

@@ -40,6 +40,8 @@
 from _task_interface import ChipCallable  # pyright: ignore[reportMissingImports]
 from simpler.worker import Worker
 
+from ._harness import SIM_PLATFORM, SIM_RUNTIME, install_fake_chip
+
 # Comfortably above every wait these tests actually take, well under any hang.
 _TEST_WALL_BUDGET_S = 30.0
 # How long close() must be observed NOT to return while a fenced operation is
@@ -53,9 +55,17 @@ def _chip_callable(func_name: str = "admission_fence") -> ChipCallable:
 
 
 @pytest.fixture
-def ready_worker():
-    """A READY L3 worker with one SUB child and no chips — no NPU required."""
-    w = Worker(level=3, num_sub_workers=1, startup_timeout_s=_TEST_WALL_BUDGET_S)
+def ready_worker(monkeypatch):
+    """A READY L3 worker with one SUB child and a fake chip — no NPU required."""
+    install_fake_chip(monkeypatch)
+    w = Worker(
+        level=3,
+        device_ids=[0],
+        platform=SIM_PLATFORM,
+        runtime=SIM_RUNTIME,
+        num_sub_workers=1,
+        startup_timeout_s=_TEST_WALL_BUDGET_S,
+    )
     w.init()
     try:
         yield w

tests/ut/py/test_worker/test_host_worker.py

@@ -1196,14 +1196,9 @@ def test_l2_rejects_python_callable(self):
 
     def test_close_releases_registered_callables(self):
         # close() must drop every Worker-held reference to registered callables.
-        # A ChipCallable is a nanobind instance; if close() leaves it in one of
-        # the registries, a closed Worker kept alive past interpreter exit (e.g.
-        # a failing test's traceback pinning the frame's `worker` local) keeps
-        # the instance live, which blocks nanobind's module unload and prints a
-        # leak dump at shutdown.
-        hw = Worker(level=3, num_sub_workers=0)
+        hw = Worker(level=3, num_sub_workers=1)
         hw.init()
-        handle = hw.register(_unique_chip_callable(1))
+        handle = hw.register(lambda args: None)
         assert _slot_for(hw, handle) in hw._callable_registry
         assert hw._identity_registry and hw._live_handles
         hw.close()
@@ -1245,7 +1240,7 @@ def test_prepare_python_fn_after_start_no_python_children_raises(self):
         hw.init()
         try:
             hw.run(lambda orch, args, cfg: None)
-            with pytest.raises(RuntimeError, match="no Python-capable child"):
+            with pytest.raises(ValueError, match=r"\(needs a SUB or next-level child\)"):
                 hw.register(lambda args: None)
         finally:
             hw.close()
@@ -1428,23 +1423,20 @@ def test_prepare_chip_callable_surfaces_chip_child_register_failure(self, monkey
             assert hw._callable_registry == {}
 
     def test_prepare_chip_callable_at_cid_overflow_raises(self):
-        # cid budget is enforced under the new dynamic-prepare path too:
-        # pre-fill registry with lambdas pre-init, init, then attempt one
-        # post-init ChipCallable prepare and observe the existing
-        # MAX_REGISTERED_CALLABLE_IDS RuntimeError. A sub child gives the
-        # pre-registered python callables an eligible dispatch target.
+        # cid budget is enforced under the new dynamic-prepare path.
+        # Pre-fill registry with lambdas pre-init, init, then attempt one
+        # post-init ChipCallable prepare. A sub child gives the pre-registered
+        # python callables an eligible dispatch target.
         #
-        # This worker is chipless, so the ChipCallable is over the cid budget
-        # AND ineligible under the startup dispatch-target rule. Only the cid
-        # budget is asserted; the relative order of the two checks is unpinned
-        # here (see the b2 xfail in
-        # test_startup_readiness.py::TestEligibleTargetPrecheck).
+        # This worker is chipless, so the post-init ChipCallable fails the
+        # eligibility re-check before reaching the cid budget — eligibility
+        # fires first by design (more actionable than "out of slots").
         hw = Worker(level=3, num_sub_workers=1)
         try:
             for i in range(MAX_REGISTERED_CALLABLE_IDS):
                 hw.register(_unique_py_callable(i))
             hw.init()
-            with pytest.raises(RuntimeError, match="MAX_REGISTERED_CALLABLE_IDS"):
+            with pytest.raises(ValueError, match=r"\(needs a chip device \(device_ids\)\)"):
                 hw.register(chip_callable())
         finally:
             hw.close()
@@ -1478,7 +1470,7 @@ def test_unregister_chip_callable_after_init_succeeds(self, monkeypatch):
             assert _slot_for(hw, handle_b) == slot_a, "smallest-unused-cid policy should reuse the freed slot"
 
     def test_prepare_chip_callable_broadcast_runs_without_registry_lock(self):
-        hw = Worker(level=3, num_sub_workers=0)
+        hw = Worker(level=3, num_sub_workers=0, device_ids=[0])
         hw._lifecycle = worker_mod._Lifecycle.READY
         callable_obj = ChipCallable.build(signature=[], func_name="x", binary=b"\x00", children=[])
         observed = {}
@@ -1500,7 +1492,7 @@ def fake_post_init_register(target, digest, *, is_new):
     def test_register_child_chip_broadcast_runs_without_registry_lock(self):
         from simpler.worker import _build_callable_registration  # noqa: PLC0415
 
-        hw = Worker(level=3, num_sub_workers=0)
+        hw = Worker(level=3, num_sub_workers=0, device_ids=[0])
         hw._lifecycle = worker_mod._Lifecycle.READY
         callable_obj = ChipCallable.build(signature=[], func_name="x", binary=b"\x00", children=[])
         digest = _build_callable_registration(hw, callable_obj).digest
@@ -2211,7 +2203,7 @@ def broadcast_register_all(self, blob_ptr, blob_size, digest):
                 calls.append(("binary_register", blob_size, digest))
                 return [_FakeControlResult("NEXT_LEVEL", 0, True)]
 
-        hw = Worker(level=3, num_sub_workers=1)
+        hw = Worker(level=3, num_sub_workers=1, device_ids=[0])
         hw._lifecycle = worker_mod._Lifecycle.READY
         hw._worker = FakeWorker()
         callable_obj = ChipCallable.build(signature=[], func_name="x", binary=b"\x00", children=[])
@@ -2246,7 +2238,7 @@ def control_digest_only(self, worker_type, worker_id, sub_cmd, digest, timeout_s
                 calls.append(("cleanup_one", worker_type, worker_id, sub_cmd, digest))
                 return _FakeControlResult("NEXT_LEVEL", worker_id, True)
 
-        hw = Worker(level=3, num_sub_workers=1)
+        hw = Worker(level=3, num_sub_workers=1, device_ids=[0])
         hw._lifecycle = worker_mod._Lifecycle.READY
         hw._worker = FakeWorker()
         callable_obj = ChipCallable.build(signature=[], func_name="x", binary=b"\x00", children=[])
@@ -2277,7 +2269,7 @@ def broadcast_unregister_all(self, digest):
                 calls.append(("cleanup", digest))
                 return ["cleanup failed"]
 
-        hw = Worker(level=3, num_sub_workers=1)
+        hw = Worker(level=3, num_sub_workers=1, device_ids=[0])
         hw._lifecycle = worker_mod._Lifecycle.READY
         hw._worker = FakeWorker()
         callable_obj = ChipCallable.build(signature=[], func_name="x", binary=b"\x00", children=[])
@@ -2296,13 +2288,13 @@ def test_unregister_middle_cid_reuses_hole(self):
         # len(registry). The bug it guards against: fill slots 0/1/2,
         # unregister slot 1, next register would silently overwrite the
         # existing cid=2 under a `len(registry)` policy.
-        hw = Worker(level=3, num_sub_workers=0)
+        hw = Worker(level=3, num_sub_workers=1)
         hw.init()
         try:
-            cb0 = _unique_chip_callable(0)
-            cb1 = _unique_chip_callable(1)
-            cb2 = _unique_chip_callable(2)
-            cb3 = _unique_chip_callable(3)
+            cb0 = _unique_py_callable(0)
+            cb1 = _unique_py_callable(1)
+            cb2 = _unique_py_callable(2)
+            cb3 = _unique_py_callable(3)
             handle0 = hw.register(cb0)
             handle1 = hw.register(cb1)
             handle2 = hw.register(cb2)
@@ -2316,7 +2308,7 @@ def test_unregister_middle_cid_reuses_hole(self):
             # cid=2 entry must still be the original callable, not silently overwritten.
             assert hw._callable_registry[slot2] is cb2
             # Next register fills cid=3 since 0..2 are all occupied.
-            next_handle = hw.register(_unique_chip_callable(4))
+            next_handle = hw.register(_unique_py_callable(4))
             assert _slot_for(hw, next_handle) == 3
         finally:
             hw.close()

tests/ut/py/test_worker/test_startup_readiness.py

@@ -1501,10 +1501,6 @@ def test_chipless_l3_with_chip_callable_rejected_at_init(self):
         finally:
             w.close()
 
-    @pytest.mark.xfail(
-        strict=True,
-        reason="P0.2-b2: post-init register does not re-apply the startup eligibility rule",
-    )
     def test_post_init_chip_callable_on_chipless_l3_rejected(self):
         # The same LOCAL_CHIP rule, one epoch later: an L3 that came up without
         # a chip child cannot resolve a ChipCallable handed to it post-init
@@ -1514,7 +1510,7 @@ def test_post_init_chip_callable_on_chipless_l3_rejected(self):
         try:
             with _hard_timeout(_TEST_WALL_BUDGET_S):
                 w.init()
-                with pytest.raises(RuntimeError, match=r"\(needs a chip device \(device_ids\)\)"):
+                with pytest.raises(ValueError, match=r"\(needs a chip device \(device_ids\)\)"):
                     w.register(chip_callable())
         finally:
             w.close()