Fix: harden the Buffer release path on the owner and consumer sides

python/simpler/buffer.py

@@ -147,6 +147,10 @@ class Buffer:
     # guard and free/copy key on (owner_worker_id, base).
     owner_worker_id: int = 0
     closed: bool = False
+    # Owner-side only: whether the backing's name has actually been removed. Distinct from `closed`,
+    # which is the derivation gate — a close() whose unlink raised leaves this false so a retry
+    # attempts the unlink again.
+    unlinked: bool = False
 
     def to_descriptor(self) -> BufferDescriptor:
         """The wire descriptor for this backing — what a consumer needs to resolve it."""
@@ -190,17 +194,26 @@ def close(self) -> None:
         """Release the backing. The owner unlinks it, so a later consumer map fails rather than
         resolving a name whose bytes are gone. Idempotent; a released Buffer's ``tensor()``/
         ``to_descriptor()`` are refused rather than building a view over memory that may already be
-        gone."""
-        if self.closed:
-            return
+        gone — that refusal holds from the first call on, whether or not the release itself
+        succeeded, since a partly-released backing is no safer to hand out than a fully released one.
+
+        Each OS action succeeds at most once and is retried until it does. ``shm.close()`` runs
+        first and never gates the unlink: the named backing outlives this process, so it is the leak
+        worth removing even when the local unmap raised. An unlink that raises leaves ``shm`` in
+        place, so a second ``close()`` attempts it again — that retry is what
+        ``Worker._release_all_buffers`` leaves the registry entry behind for.
+        """
         self.closed = True
         shm = self.shm
-        self.shm = None
-        if shm is not None:
-            try:
-                shm.close()
-            finally:
+        if shm is None:
+            return
+        try:
+            shm.close()
+        finally:
+            if not self.unlinked:
                 shm.unlink()
+                self.unlinked = True
+        self.shm = None
 
 
 def create_host_shared_buffer(
@@ -578,7 +591,7 @@ class ImportRegistry:
     still describes the backing that was mapped, so one identity can never come to mean two things.
     """
 
-    def __init__(self, context: ImportContext | None = None) -> None:
+    def __init__(self, context: ImportContext) -> None:
         self._by_identity: dict[CanonicalIdentity, ImportedBuffer] = {}
         self._context = context
 
@@ -616,7 +629,7 @@ def materialize(self, desc: BufferDescriptor) -> ImportedBuffer:
                 )
             return cached
         if desc.address_space == AddressSpace.DEVICE:
-            if self._context is None or self._context.is_host_endpoint:
+            if self._context.is_host_endpoint:
                 raise ValueError(
                     f"ImportRegistry: [{RegionAccessReasonCode.UNSUPPORTED_ENDPOINT_RELATION.value}] "
                     f"refusing to materialize a DEVICE backing ({desc.identity}) on a host endpoint"
@@ -643,7 +656,17 @@ def materialize(self, desc: BufferDescriptor) -> ImportedBuffer:
             base = int.from_bytes(desc.body, "little")
             imported = ImportedBuffer(desc.identity, base, desc.nbytes, desc.address_space, None, desc)
         elif desc.backend_kind == BackendKind.POSIX_SHM:
-            shm = SharedMemory(name=desc.body.decode("utf-8"))
+            name = desc.body.decode("utf-8")
+            try:
+                shm = SharedMemory(name=name)
+            except FileNotFoundError as exc:
+                # The owner unlinks on release, so a missing name is the expected shape of "this
+                # identity was released", not a corrupt descriptor. Naming both the identity and
+                # that reading separates it from a genuinely bad name at a glance.
+                raise FileNotFoundError(
+                    f"ImportRegistry: shm object {name!r} for {desc.identity} does not exist — its "
+                    f"owner has released the buffer, or it was never created"
+                ) from exc
             # `validate_tensor` admits a view when byte_offset + extent <= nbytes, so a backing
             # smaller than the nbytes its own descriptor advertises turns every one of those checks
             # into a comparison against a number no memory stands behind. The object's real size is
@@ -708,15 +731,36 @@ def unregister(self, identity: CanonicalIdentity) -> None:
         ``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.
+
+        The entry is dropped only once its mapping is really gone. ``close()`` on a mapping whose
+        consumer still holds a derived ``memoryview`` raises ``BufferError``, and an entry dropped
+        before that point is a mapping nothing can reach to retry — so the raise leaves the entry
+        in place for this registry's own ``close()`` to attempt again.
         """
-        imported = self._by_identity.pop(identity, None)
-        if imported is not None and imported.shm is not None:
+        imported = self._by_identity.get(identity)
+        if imported is None:
+            return
+        if imported.shm is not None:
             imported.shm.close()
+        del self._by_identity[identity]
 
     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."""
-        for imported in self._by_identity.values():
+        owning Worker, so this never destroys a backing.
+
+        Every mapping is attempted, and only the ones that closed are dropped: one endpoint holding
+        an exported view must not strand every mapping behind it in the iteration order. The first
+        error is raised once the sweep is done, so a caller still learns the endpoint leaked rather
+        than seeing a silent success.
+        """
+        errors: list[BaseException] = []
+        for identity, imported in list(self._by_identity.items()):
             if imported.shm is not None:
-                imported.shm.close()
-        self._by_identity.clear()
+                try:
+                    imported.shm.close()
+                except BaseException as exc:  # noqa: BLE001
+                    errors.append(exc)
+                    continue
+            del self._by_identity[identity]
+        if errors:
+            raise errors[0]

python/simpler/orchestrator.py

@@ -483,6 +483,8 @@ def submit_sub(self, callable_handle: Any, args: TaskArgs | None = None):
         )
         _reject_remote_sidecar_args(args, kind="orch.submit_sub")
         _reject_device_args(args, kind="orch.submit_sub")
+        if self._worker is not None:
+            self._worker._record_touched_identities(args)
         _admit_task_submission(self._worker)
         self._o.submit_sub(digest, kind, target_namespace, args)
 
@@ -497,6 +499,12 @@ def submit_sub_group(self, callable_handle: Any, args_list: list):
         for args in args_list:
             _reject_remote_sidecar_args(args, kind="orch.submit_sub_group")
             _reject_device_args(args, kind="orch.submit_sub_group")
+        # Second pass, as in submit_next_level_group: a member rejected above means no member is
+        # dispatched, and identities recorded for a group that never went out would refuse a
+        # release of buffers no task ever received.
+        if self._worker is not None:
+            for args in args_list:
+                self._worker._record_touched_identities(args)
         _admit_task_submission(self._worker)
         self._o.submit_sub_group(digest, kind, target_namespace, args_list)
 

python/simpler/worker.py

@@ -7309,14 +7309,20 @@ def _broadcast_import_release(self, identity: CanonicalIdentity) -> None:
             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
+        """Drop ``identity`` from this Worker's own same-process caches, 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).
+
+        Two caches name a released identity here, not one: the import cache holds a mapping of it,
+        and ``_reexport_by_source`` holds the forwarding handle built from it. A retained re-export
+        outlives its backing, and its ``to_descriptor()`` keeps answering — so a later forward of
+        the same identity would hand a child a descriptor for a name the owner has unlinked.
         """
         if self._chip_import_registry is not None:
             self._chip_import_registry.unregister(identity)
+        self._reexport_by_source.pop(identity, None)
         self._broadcast_import_release(identity)
 
     def add_worker(self, worker: Worker) -> int:
@@ -10082,9 +10088,11 @@ def _record_touched_identities(self, args: Any) -> None:
         """Add every tensor arg's identity in ``args`` to the current run's touched set.
 
         A no-op when no run is being built (``_building_run_resources is None``) — tracking is
-        opportunistic, only meaningful inside ``submit_next_level``/``submit_next_level_group``'s
-        run context, per the ``current_resources = self._building_run_resources; if ... is not
-        None`` idiom used elsewhere for the same "attach to the open run, if any" shape.
+        opportunistic, only meaningful inside the run context of the four orchestrator dispatch
+        entry points that carry Tensor args to another process (``submit_next_level``,
+        ``submit_next_level_group``, ``submit_sub``, ``submit_sub_group``), per the
+        ``current_resources = self._building_run_resources; if ... is not None`` idiom used
+        elsewhere for the same "attach to the open run, if any" shape.
         """
         resources = self._building_run_resources
         if resources is None:
@@ -10497,13 +10505,23 @@ def release_buffer(self, buffer: Buffer) -> None:
         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 —
-        a Buffer never goes away while a dispatched task still names it. The L3+ check takes
+        sent this identity as a NEXT_LEVEL or SUB Tensor arg, or any in-flight L2 direct-chip run
+        sent it — a Buffer never goes away while a dispatched task still names it. All three
+        dispatch paths retain: a SUB task maps the identity into a sub-worker process just as a
+        NEXT_LEVEL task maps it into a child, so unlinking the backing under either one faults the
+        consumer on a segment that no longer has a name. The L3+ check takes
         ``_submit_mu`` first: a handle is visible in ``_accepted_run_handles`` before its
         orchestration callback (where ``touched_identities`` gets populated) has run, and that
         callback is what ``_submit_mu`` already serializes graph construction against, so taking it
         here means the check only ever runs between callbacks, never mid-callback with a
-        not-yet-complete touched set. The L2 check is independent (a separate run-id namespace with
+        not-yet-complete touched set. ``_abandoned_run_handles`` is scanned in the same block and
+        without the ``_cleanup_published`` test: ``_publish_abandoned_run`` sets that flag and drops
+        the handle from the accepted set while the run itself stays retained until native teardown
+        drains it, so an abandoned run is exactly the case where the flag stops describing whether
+        the device is done with the backing. Such a buffer therefore stops being releasable through
+        this API for the Worker's remaining life, which strands nothing: ``close()`` reclaims it via
+        ``_release_all_buffers`` calling ``Buffer.close()`` directly. The L2 check is independent (a
+        separate run-id namespace with
         no callback to serialize against — ``_chip_run_touched_identities`` is written atomically
         alongside ``_chip_runs`` under ``_registry_lock`` instead, see ``_submit_l2_locked``), so the
         two checks run sequentially rather than under one shared lock. Neither is checked once
@@ -10523,6 +10541,13 @@ def release_buffer(self, buffer: Buffer) -> None:
                 for handle in self._accepted_run_handles:
                     if not handle._cleanup_published and buffer.identity in handle._resources.touched_identities:
                         raise RuntimeError(f"release_buffer: {buffer.identity} is still referenced by an in-flight run")
+                for handle in self._abandoned_run_handles:
+                    resources = handle._resources
+                    if resources is not None and buffer.identity in resources.touched_identities:
+                        raise RuntimeError(
+                            f"release_buffer: {buffer.identity} is still referenced by an abandoned run "
+                            f"whose native teardown has not completed"
+                        )
             with self._registry_lock:
                 for touched in self._chip_run_touched_identities.values():
                     if buffer.identity in touched:
@@ -11675,6 +11700,7 @@ def _step(fn) -> None:
             ("remote", "pending remote frees", self._flush_pending_remote_frees),
             ("buffer", "all owner Buffers", self._release_all_buffers),
             ("buffer", "fork-inherited tensor buffers", self._fork_tensor_handles.clear),
+            ("buffer", "re-exported forwarding handles", self._reexport_by_source.clear),
             ("buffer", "chip import registry", self._close_chip_import_registry),
         ):
             self._cleanup_journal.add_once(kind, identity, cleanup)