Fix: close 4 of #1729's 5 deferred buffer-registry cleanup items

python/simpler/buffer.py

@@ -442,9 +442,11 @@ class ImportedBuffer:
 @dataclass
 class MappedArg:
     """A Python compute (sub-worker) task arg: a ``Tensor`` materialized into this process, exposing a
-    writable ``buffer`` at the view origin plus the view geometry. The callable computes with e.g.
+    ``buffer`` at the view origin plus the view geometry. The callable computes with e.g.
     ``torch.frombuffer(arg.buffer, dtype=<from arg.dtype>, count=prod(arg.shapes))`` — reads/writes
-    land in the shared backing the owner sees.
+    land in the shared backing the owner sees, except when the descriptor's ``access`` is
+    ``AccessMode.READ``, where ``buffer`` is a read-only view (a ``FORK_COW`` backing's writes are
+    invisible to the owner and are the reason ``access`` is forced to ``READ`` for it).
     """
 
     imported: ImportedBuffer
@@ -455,7 +457,8 @@ class MappedArg:
 
     @property
     def buffer(self) -> memoryview:
-        """A memoryview over the mapped backing at this view's origin (``byte_offset``)."""
+        """A memoryview over the mapped backing at this view's origin (``byte_offset``); read-only
+        when the descriptor's ``access`` is ``AccessMode.READ``."""
         ib = self.imported
         if ib.shm is not None:
             base = ib.shm.buf
@@ -464,7 +467,10 @@ def buffer(self) -> memoryview:
             # FORK_SHM / FORK_COW: no shm object — the base is a host VA inherited across the
             # fork, so wrap that range directly.
             base = memoryview((ctypes.c_char * ib.nbytes).from_address(ib.base))
-        return base[self.byte_offset :]
+        view = base[self.byte_offset :]
+        if ib.descriptor is not None and ib.descriptor.access == AccessMode.READ:
+            return view.toreadonly()
+        return view
 
 
 class MappedArgs(Sequence):
@@ -617,23 +623,25 @@ def materialize(self, desc: BufferDescriptor) -> ImportedBuffer:
         self._by_identity[key] = imported
         return imported
 
-    def materialization_map(self) -> dict[CanonicalIdentity, tuple[int, int]]:
-        """Snapshot for ``materialize_tensor_blob`` / ``materialize_task_args``:
-        identity -> (local base, address_space)."""
-        return {key: (ib.base, int(ib.address_space)) for key, ib in self._by_identity.items()}
-
     def materialize_blob(self, blob_ptr: int, capacity: int) -> dict[CanonicalIdentity, tuple[int, int]]:
-        """Materialize every embedded descriptor in a task-args blob and return the resolved map."""
+        """Materialize every embedded descriptor in a task-args blob and return the resolved map:
+        identity -> (local base, address_space), scoped to this call's own tensors."""
         args = read_args_from_blob(blob_ptr, capacity)
+        resolved: dict[CanonicalIdentity, tuple[int, int]] = {}
         for i in range(args.tensor_count()):
-            self.materialize(args.tensor(i).buffer)
-        return self.materialization_map()
+            desc = args.tensor(i).buffer
+            imported = self.materialize(desc)
+            resolved[desc.identity] = (imported.base, int(imported.address_space))
+        return resolved
 
     def materialize_args(self, args) -> dict[CanonicalIdentity, tuple[int, int]]:
         """The same, for a ``TaskArgs`` already held in this process (the L2-leaf path)."""
+        resolved: dict[CanonicalIdentity, tuple[int, int]] = {}
         for i in range(args.tensor_count()):
-            self.materialize(args.tensor(i).buffer)
-        return self.materialization_map()
+            desc = args.tensor(i).buffer
+            imported = self.materialize(desc)
+            resolved[desc.identity] = (imported.base, int(imported.address_space))
+        return resolved
 
     def mapped_args_from_blob(self, blob_ptr: int, capacity: int) -> MappedArgs:
         """Materialize a task-args blob into a Python compute callable's args: every tensor becomes a
@@ -663,13 +671,6 @@ def resolve(self, identity: CanonicalIdentity) -> ImportedBuffer:
             raise KeyError(f"ImportRegistry: no buffer registered for {identity}")
         return imported
 
-    def unregister(self, identity: CanonicalIdentity) -> None:
-        """Drop one import and close its mapping. The owner still holds the backing; only this
-        endpoint's view of it goes away. A no-op for an identity that was never materialized."""
-        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."""

tests/ut/py/test_buffer.py

@@ -16,6 +16,8 @@
 registry and the Buffer constructors that are genuinely defined there.
 """
 
+import ctypes
+
 import pytest
 from _task_interface import OWNER_INSTANCE_ID_BYTES, DataType
 from simpler.buffer import (
@@ -26,6 +28,7 @@
     CanonicalIdentity,
     ImportContext,
     ImportRegistry,
+    MappedArg,
     Tensor,
     create_host_shared_buffer,
     intern_worker_path,
@@ -203,6 +206,30 @@ def test_device_malloc_wrap_materialize():
     assert imp.shm is None
 
 
+def test_materialize_args_scopes_the_returned_map_to_this_calls_tensors():
+    # materialize_args must not hand back an endpoint's entire materialize-once history — only the
+    # identities the tensors in THIS call touched. A dispatch that reuses a registry across many
+    # tasks (the real chip/L2-leaf usage) would otherwise pay O(every identity ever seen) per call.
+    import simpler.task_interface as ti  # noqa: PLC0415
+
+    oid = mint_owner_instance_id()
+    reg = ImportRegistry(ImportContext(is_host_endpoint=False, owning_chip_instance_id=oid))
+    h1 = wrap_device_malloc(0xDEAD0000, 4096, oid, buffer_id=1)
+    h2 = wrap_device_malloc(0xBEEF0000, 4096, oid, buffer_id=2)
+
+    args1 = ti.TaskArgs()
+    args1.add_tensor(h1.tensor(shapes=(16,), dtype=DataType.FLOAT32))
+    resolved1 = reg.materialize_args(args1)
+    assert set(resolved1) == {h1.identity}
+
+    args2 = ti.TaskArgs()
+    args2.add_tensor(h2.tensor(shapes=(16,), dtype=DataType.FLOAT32))
+    resolved2 = reg.materialize_args(args2)
+    # h1 is still live in the registry's own history (map-once), but this call's args never
+    # referenced it, so it must not appear in this call's returned map.
+    assert set(resolved2) == {h2.identity}
+
+
 def test_materialize_remote_sidecar_rejected():
     desc = BufferDescriptor(
         identity=_identity(),
@@ -429,3 +456,23 @@ def test_fork_backend_is_stated_not_inferred_from_access():
     )
     with pytest.raises(ValueError, match="FORK_COW"):
         bad.to_descriptor()
+
+
+def test_mapped_arg_buffer_is_read_only_for_a_read_access_descriptor():
+    # FORK_COW's whole contract is that a write is invisible to the owner (copy-on-write splits the
+    # page privately), so `MappedArg.buffer` must not hand back a writable view for it. `.cast("B")`
+    # is what makes the write attempt itself meaningful: the raw `<c` format memoryview this property
+    # returns for a ctypes-backed mapping does not support slice assignment at all (readonly or not),
+    # so asserting through the native format would prove nothing about the readonly flag specifically.
+    data = bytearray(16)
+    addr = ctypes.addressof((ctypes.c_char * 16).from_buffer(data))
+    oid = mint_owner_instance_id()
+    buf = wrap_fork_inherited(addr, 16, oid, buffer_id=1)  # default: access=READ, backend=FORK_COW
+    reg = ImportRegistry()
+    imported = reg.materialize(buf.to_descriptor())
+    arg = MappedArg(imported, byte_offset=0, shapes=(16,), strides=(1,), dtype=DataType.UINT8)
+
+    view = arg.buffer
+    assert view.readonly
+    with pytest.raises(TypeError):
+        view.cast("B")[0:4] = b"\x01\x02\x03\x04"