diff --git a/python/simpler/buffer.py b/python/simpler/buffer.py index 46fa362a5e..a03c833199 100644 --- a/python/simpler/buffer.py +++ b/python/simpler/buffer.py @@ -685,6 +685,19 @@ def resolve(self, identity: CanonicalIdentity) -> ImportedBuffer: raise KeyError(f"ImportRegistry: no buffer registered for {identity}") return imported + def unregister(self, identity: CanonicalIdentity) -> None: + """Drop ``identity``'s mapping if this endpoint made one; a no-op otherwise. + + Consumer-side only — unlinking belongs to the owning Worker, so this never destroys a + backing, only this endpoint's own view of it. The owner broadcasts this on + ``release_buffer()`` so a long-lived endpoint does not keep every backing it ever saw + mapped for its entire lifetime; best-effort by design, since an endpoint that never + materialized ``identity`` has nothing to drop. + """ + imported = self._by_identity.pop(identity, None) + if imported is not None and imported.shm is not None: + imported.shm.close() + def close(self) -> None: """Close every mapping this endpoint made. Consumer-side only — unlinking belongs to the owning Worker, so this never destroys a backing.""" diff --git a/python/simpler/worker.py b/python/simpler/worker.py index fc3648d0f1..8d9c6c1691 100644 --- a/python/simpler/worker.py +++ b/python/simpler/worker.py @@ -103,6 +103,7 @@ def my_l4_orch(orch, args, config): from . import _log as _simpler_log from .buffer import ( + OWNER_INSTANCE_ID_BYTES, AccessMode, AddressSpace, BackendKind, @@ -495,6 +496,13 @@ def _shm_name(token: str, suffix: str): _CTRL_PY_REGISTER = 10 _CTRL_PY_UNREGISTER = 11 _CTRL_PY_IMPORT_REGISTER = 12 +# Best-effort import-cache invalidation: the owner released a Buffer, this tells every consumer +# (chip or SUB child, or a nested NEXT_LEVEL Worker forwarding it further down its own tree) to +# drop its own materialized mapping for the identity if it has one. Reached via the generic +# broadcast_control_all path (Worker._broadcast_py_control), the same mechanism _CTRL_PY_REGISTER +# et al. already use — the digest slot carries a CanonicalIdentity's three meaningful fields +# packed by _pack_identity_wire, not a callable hash. +_CTRL_IMPORT_RELEASE = 13 # Operation names a child puts in its error message when a control command # fails, so the parent's re-raised text names the operation and not just a # numeric sub-command. Absent entries fall back to the raw number. @@ -504,8 +512,10 @@ def _shm_name(token: str, suffix: str): _CTRL_PY_REGISTER: "py_register", _CTRL_PY_IMPORT_REGISTER: "py_register", _CTRL_PY_UNREGISTER: "py_unregister", + _CTRL_IMPORT_RELEASE: "import_release", } + _CTRL_WORKER_CHIP_REGION_CREATE = 16 _CTRL_WORKER_CHIP_REGION_RELEASE = 17 _CTRL_COMMITTED_DEVICE_MEMORY = 18 @@ -1651,6 +1661,26 @@ def _read_control_digest(buf) -> bytes: return bytes(buf[_OFF_CONTROL_CALLABLE_HASH : _OFF_CONTROL_CALLABLE_HASH + CALLABLE_HASH_DIGEST_BYTES]) +_IDENTITY_WIRE = struct.Struct(f"<{OWNER_INSTANCE_ID_BYTES}sQI") +assert _IDENTITY_WIRE.size <= CALLABLE_HASH_DIGEST_BYTES + + +def _pack_identity_wire(identity: CanonicalIdentity) -> bytes: + """CanonicalIdentity's three meaningful fields, packed to fit the digest-sized control slot + _CTRL_IMPORT_RELEASE reuses. Deliberately not the identity's own bytes — see + ``CanonicalIdentity``'s binding docstring on why it exposes no ``pack()``: this reconstructs + the wire form field-by-field, the same way ``remote_l3_protocol.py`` encodes one for the + cross-machine wire, rather than dumping the object's raw memory. + """ + body = _IDENTITY_WIRE.pack(identity.owner_instance_id, identity.buffer_id, identity.generation) + return body.ljust(CALLABLE_HASH_DIGEST_BYTES, b"\x00") + + +def _unpack_identity_wire(buf: bytes) -> CanonicalIdentity: + owner_instance_id, buffer_id, generation = _IDENTITY_WIRE.unpack_from(buf, 0) + return CanonicalIdentity(owner_instance_id, buffer_id, generation) + + def _read_task_digest(buf) -> bytes: return bytes(buf[_OFF_TASK_CALLABLE_HASH : _OFF_TASK_CALLABLE_HASH + CALLABLE_HASH_DIGEST_BYTES]) @@ -1922,14 +1952,17 @@ def handle_task(task_buf) -> tuple[int, str]: def handle_control(sub_cmd: int) -> tuple[int, str]: try: - _handle_py_callable_control( - buf, - registry, - identity_table, - identity_refs, - sub_cmd, - context="sub_worker", - ) + if sub_cmd == _CTRL_IMPORT_RELEASE: + import_registry.unregister(_unpack_identity_wire(_read_control_digest(buf))) + else: + _handle_py_callable_control( + buf, + registry, + identity_table, + identity_refs, + sub_cmd, + context="sub_worker", + ) except Exception as e: # noqa: BLE001 return 1, _format_exc("sub_worker control", e) return 0, "" @@ -2705,6 +2738,8 @@ def handle_control( # noqa: PLR0912, PLR0915 -- one branch per control sub-comm _handle_ctrl_worker_chip_region_release(buf, worker_chip_region_store) elif sub_cmd == _CTRL_COMMITTED_DEVICE_MEMORY: struct.pack_into("Q", buf, _CTRL_OFF_RESULT, cw.committed_device_memory) + elif sub_cmd == _CTRL_IMPORT_RELEASE: + import_registry.unregister(_unpack_identity_wire(_read_control_digest(buf))) elif sub_cmd == CTRL_GLOBAL_DOMAIN_PREPARE: _handle_ctrl_global_domain_prepare(cw, buf, global_domain_store) elif sub_cmd == CTRL_GLOBAL_DOMAIN_IMPORT: @@ -3263,6 +3298,8 @@ def handle_control(sub_cmd: int) -> tuple[int, str]: digest = _read_control_digest(buf) inner_worker._unregister_child_digest(digest=digest) _remove_local_identity(registry, identity_table, identity_refs, digest) + elif sub_cmd == _CTRL_IMPORT_RELEASE: + inner_worker._release_import_recursive(_unpack_identity_wire(_read_control_digest(buf))) elif sub_cmd in (_CTRL_PY_REGISTER, _CTRL_PY_IMPORT_REGISTER, _CTRL_PY_UNREGISTER): _handle_py_callable_control( buf, @@ -6694,6 +6731,48 @@ def _broadcast_unregister(self, digest: bytes) -> None: ) sys.stderr.flush() + def _broadcast_import_release(self, identity: CanonicalIdentity) -> None: + """Broadcast _CTRL_IMPORT_RELEASE to every NEXT_LEVEL child (chip or nested Worker) and + every SUB child, via the generic control channel _CTRL_PY_REGISTER et al. already use. + + A no-op before ``init()`` (``self._worker`` unset) — a Worker that never started has no + forked children to reach regardless of what ``_python_worker_types()`` would report. + Otherwise unconditional on both target lists: a Worker with no children of one kind sends + to an empty list, which the underlying broadcast already no-ops on — no need to duplicate + ``_python_worker_types()``'s gating here (that helper only counts nested-Worker NEXT_LEVEL + children, not chip children, so it would wrongly skip the chip-only case). + + Best-effort, mirroring ``_broadcast_unregister``: a child that never materialized + ``identity`` has nothing to drop, and a slow or dead child must not block or fail + ``release_buffer()`` — the Buffer is already closed on the owner side by the time this + runs, so this call is pure cleanup, not a correctness gate. + """ + if self._worker is None: + return + errors = self._broadcast_py_control( + [WorkerType.NEXT_LEVEL, WorkerType.SUB], + _CTRL_IMPORT_RELEASE, + digest=_pack_identity_wire(identity), + strict=False, + ) + if errors: + sys.stderr.write( + f"Worker.release_buffer(identity={identity}): {len(errors)} children reported errors " + f"(continuing best-effort). First error: {errors[0]}\n" + ) + sys.stderr.flush() + + def _release_import_recursive(self, identity: CanonicalIdentity) -> None: + """Drop ``identity`` from this Worker's own same-process import cache, then forward one + more hop down this Worker's own children — same shape as ``_unregister_child_digest``'s + recursive forward for callable cleanup, since a NEXT_LEVEL child may itself have + materialized ``identity`` further down its own tree (chip/SUB leaves, or its own + NEXT_LEVEL children in turn). + """ + if self._chip_import_registry is not None: + self._chip_import_registry.unregister(identity) + self._broadcast_import_release(identity) + def add_worker(self, worker: Worker) -> int: """Add a lower-level Worker as a NEXT_LEVEL child. Must be called before init(). @@ -9720,7 +9799,8 @@ def _close_chip_import_registry(self) -> None: self._chip_import_registry = None def release_buffer(self, buffer: Buffer) -> None: - """Close + unlink one owner Buffer and drop its registry entry. + """Close + unlink one owner Buffer, drop its registry entry, and tell every descendant to + drop its own cached import for the identity. Rejects outright if any currently in-flight L3+ run (not yet past ``_cleanup_published``) sent this identity as a NEXT_LEVEL Tensor arg, or any in-flight L2 direct-chip run sent it — @@ -9735,7 +9815,9 @@ def release_buffer(self, buffer: Buffer) -> None: two checks run sequentially rather than under one shared lock. Neither is checked once ``buffer`` is already closed, matching ``Buffer.close()``'s own idempotency. The entry survives a failed close, so ``_release_all_buffers`` still reports the leak at close() rather - than losing it here. The slot is dropped only when it still holds *this* buffer: a buffer_id + than losing it here — the import-cache broadcast only fires once close() has actually + succeeded, so a failed release never tells a descendant to drop a mapping the owner still + considers live. The slot is dropped only when it still holds *this* buffer: a buffer_id minted elsewhere can collide with a registry key, and evicting the live entry it names would strand that backing.""" if not buffer.closed: @@ -9750,6 +9832,7 @@ def release_buffer(self, buffer: Buffer) -> None: f"release_buffer: {buffer.identity} is still referenced by an in-flight L2 run" ) buffer.close() + self._release_import_recursive(buffer.identity) with self._registry_lock: buffer_id = int(buffer.identity.buffer_id) if self._buffers.get(buffer_id) is buffer: @@ -10032,16 +10115,17 @@ def _submit_l2_locked(self, callable_id: int, args, cfg: CallConfig) -> RunHandl """ assert self._chip_worker is not None touched = self._identities_in_args(args) if args is not None else set() - chip_args = self._materialize_l2_args(args) - # Publish touched_identities BEFORE dispatch, not after: release_buffer() reads this + # Publish touched_identities BEFORE materializing, not after: release_buffer() reads this # dict to decide whether a Buffer is still in flight, so if it were only written after - # _submit_chip_run_direct() returns, a release racing the window between dispatch and - # publication would see no entry for a run that is already running on the chip. + # _materialize_l2_args() (which populates self._chip_import_registry, the very cache + # release_buffer() pops), a release racing that window would see no entry for a run that + # has already cached the mapping it is about to pop out from under it. with self._registry_lock: self._chip_run_seq += 1 run_id = self._chip_run_seq self._chip_run_touched_identities[run_id] = touched try: + chip_args = self._materialize_l2_args(args) chip_run = self._chip_worker._impl._submit_chip_run_direct(callable_id, chip_args, cfg) except BaseException: with self._registry_lock: diff --git a/tests/ut/py/test_buffer.py b/tests/ut/py/test_buffer.py index 1ef3c8300c..8135454dce 100644 --- a/tests/ut/py/test_buffer.py +++ b/tests/ut/py/test_buffer.py @@ -183,6 +183,28 @@ def test_resolve_unregistered_raises(): reg.resolve(_identity()) +def test_unregister_drops_a_materialized_mapping(): + oid = mint_owner_instance_id() + buffer = create_host_shared_buffer(nbytes=64, owner_instance_id=oid, buffer_id=1) + reg = ImportRegistry() + try: + reg.materialize(buffer.to_descriptor()) + reg.unregister(buffer.identity) + with pytest.raises(KeyError): + reg.resolve(buffer.identity) + # Re-materializing after unregister must re-open the shm fresh, not fail or hit a stale + # cache entry — the same identity mapped, closed, and mapped again. + reg.materialize(buffer.to_descriptor()) + finally: + reg.close() + buffer.close() + + +def test_unregister_is_a_no_op_for_an_identity_never_materialized(): + reg = ImportRegistry() + reg.unregister(_identity()) # must not raise + + def test_tensor_full_view_is_contiguous(): oid = mint_owner_instance_id() h = create_host_shared_buffer(nbytes=1024, owner_instance_id=oid, buffer_id=1) diff --git a/tests/ut/py/test_remote_l3_lifecycle.py b/tests/ut/py/test_remote_l3_lifecycle.py index 24289c3af0..4fa0aec37c 100644 --- a/tests/ut/py/test_remote_l3_lifecycle.py +++ b/tests/ut/py/test_remote_l3_lifecycle.py @@ -959,6 +959,8 @@ def _bare_l3_worker(): w._accepted_run_handles = set() w._submit_mu = threading.Lock() w._chip_run_touched_identities = {} + w._chip_import_registry = None + w._worker = None return w diff --git a/tests/ut/py/test_worker/test_create_buffer.py b/tests/ut/py/test_worker/test_create_buffer.py index 3c360ff603..e3c7806897 100644 --- a/tests/ut/py/test_worker/test_create_buffer.py +++ b/tests/ut/py/test_worker/test_create_buffer.py @@ -39,6 +39,8 @@ def _bare_worker(level: int, *, chip: int = 0, sub: int = 0, next_level: int = 0 w._accepted_run_handles = set() w._submit_mu = threading.Lock() w._chip_run_touched_identities = {} + w._chip_import_registry = None + w._worker = None return w diff --git a/tests/ut/py/test_worker/test_release_buffer.py b/tests/ut/py/test_worker/test_release_buffer.py index f08392874d..4d7862d053 100644 --- a/tests/ut/py/test_worker/test_release_buffer.py +++ b/tests/ut/py/test_worker/test_release_buffer.py @@ -31,6 +31,8 @@ from simpler.task_interface import CallConfig, TaskArgs from simpler.worker import RunHandle, Worker, _RunResources +from ._harness import fake_chip_l3 + _F32 = 0 # DataType.FLOAT32 value _OID = mint_owner_instance_id() @@ -272,6 +274,37 @@ def _submit_chip_run_direct(self, callable_id, chip_args, cfg): submit_thread.join(timeout=5) assert not buf.closed + def test_publishes_touched_identities_before_materializing_not_after(self, monkeypatch): + # _submit_l2_locked must publish _chip_run_touched_identities BEFORE calling + # _materialize_l2_args, not just before dispatch -- _materialize_l2_args is what + # populates self._chip_import_registry, the very cache release_buffer()'s broadcast + # pops. A release racing the window between materialize and publication would see no + # entry, pass its check, and pop a mapping a dispatch already in progress just cached. + # Block materialize mid-call and confirm release_buffer already rejects at that point. + w, buf = _l2_with_registered_buffer() + entered_materialize = threading.Event() + proceed = threading.Event() + real_materialize = Worker._materialize_l2_args + + def _blocking_materialize(self, args): + entered_materialize.set() + proceed.wait(timeout=5) + return real_materialize(self, args) + + monkeypatch.setattr(Worker, "_materialize_l2_args", _blocking_materialize) + + submit_thread = threading.Thread(target=lambda: w._submit_l2_locked(3, _buffer_args(buf), CallConfig())) + submit_thread.start() + assert entered_materialize.wait(timeout=5) + + try: + with pytest.raises(RuntimeError, match="in-flight L2 run"): + w.release_buffer(buf) + finally: + proceed.set() + submit_thread.join(timeout=5) + assert not buf.closed + class TestReleaseBufferL2InFlight: def test_rejects_while_an_l2_run_is_in_flight(self): @@ -290,3 +323,16 @@ def test_succeeds_once_the_l2_run_is_gone(self): w.release_buffer(buf) assert buf.closed assert int(buf.identity.buffer_id) not in w._buffers + + +class TestReleaseBufferImportBroadcast: + def test_broadcasts_to_a_real_forked_chip_child_without_error(self, monkeypatch, capsys): + # Exercises the real wire round-trip against a live forked chip child (device-free via the + # fake-chip-L3 harness, but _run_chip_main_loop and its ImportRegistry are production code, + # not faked): the new _CTRL_IMPORT_RELEASE sub_cmd must reach the child's handle_control and + # return cleanly, proving the sub_cmd numbering and CanonicalIdentity wire packing agree on + # both ends rather than just in the mocked unit tests above. + with fake_chip_l3(monkeypatch, device_ids=(0,)) as w: + buf = w.create_buffer(64) + w.release_buffer(buf) + assert "reported errors" not in capsys.readouterr().err