Add: docstrings for the remote-buffer API and Worker lifecycle

python/simpler/orchestrator.py

@@ -453,9 +453,15 @@ def create_l3_l2_queue(self, *, worker_id: int, depth: int, input_arena_bytes: i
     #         orch.submit_next_level(c, ..., worker=0)  # outer-scope ring
 
     def scope_begin(self) -> None:
+        """Open a nested scope explicitly.
+
+        Prefer the ``scope()`` context manager, which pairs the end for you. Every
+        ``scope_begin()`` must be matched by a ``scope_end()``.
+        """
         self._o.scope_begin()
 
     def scope_end(self) -> None:
+        """Close the scope opened by the matching ``scope_begin()``."""
         self._o.scope_end()
 
     @contextlib.contextmanager

python/simpler/task_interface.py

@@ -163,6 +163,13 @@ def _assert_bindings_match_source_tree() -> None:
 
 
 class RemoteAddressSpace(IntEnum):
+    """How a remote buffer's bytes are reached.
+
+    ``HOST_INLINE`` carries the payload in the message itself rather than
+    naming remote memory. ``REMOTE_WINDOW`` and ``UB_LDST`` are protocol
+    placeholders: the shipped transport is simulation-backed.
+    """
+
     HOST_INLINE = 1
     REMOTE_DEVICE = 2
     REMOTE_WINDOW = 3
@@ -177,6 +184,22 @@ class RemoteAddressSpace(IntEnum):
 
 
 class RemoteBufferHandle:
+    """A reference to memory on a remote L3 worker.
+
+    Returned by ``Worker.remote_malloc`` (an *owner* handle) or by
+    ``Worker.remote_import`` (an *imported* handle, told apart by
+    ``is_imported``). The two are not interchangeable: owner handles are freed
+    with ``remote_free``, imported ones with ``remote_release_import``.
+
+    ``RemoteTensorRef.host_inline`` produces a third form, with
+    ``address_space`` of ``HOST_INLINE``: it carries its bytes in the message
+    and names no remote allocation, so neither release call applies —
+    ``remote_free`` rejects it outright and it is never ``is_imported``.
+
+    Construct only through ``Worker`` or ``RemoteTensorRef.host_inline``; the
+    constructor is token-guarded.
+    """
+
     __slots__ = (
         "_worker_id",
         "_owner_worker_id",
@@ -328,34 +351,42 @@ def _from_imported_mapping(  # noqa: PLR0913
 
     @property
     def worker_id(self) -> int:
+        """Worker holding this reference — the importer, for an imported handle."""
         return self._worker_id
 
     @property
     def owner_worker_id(self) -> int:
+        """Worker that owns the underlying allocation."""
         return self._owner_worker_id
 
     @property
     def import_id(self) -> int:
+        """Nonzero on an imported handle; ``0`` on an owner handle."""
         return self._import_id
 
     @property
     def address_space(self) -> RemoteAddressSpace:
+        """How these bytes are reached; see ``RemoteAddressSpace``."""
         return self._address_space
 
     @property
     def nbytes(self) -> int:
+        """Size of the allocation in bytes, or of the payload for ``HOST_INLINE``."""
         return self._nbytes
 
     @property
     def released(self) -> bool:
+        """Whether the handle has been freed or released."""
         return self._released
 
     @property
     def access_flags(self) -> int:
+        """Permitted access as a read/write bitmask; an export may only narrow it."""
         return self._access_flags
 
     @property
     def is_imported(self) -> bool:
+        """Whether this came from ``remote_import`` rather than ``remote_malloc``."""
         return self._import_id != 0
 
     def _mark_released(self) -> None:
@@ -524,26 +555,32 @@ def __setattr__(self, name: str, value: Any) -> None:
 
     @property
     def owner_worker_id(self) -> int:
+        """Worker that owns the exported allocation."""
         return self._owner_worker_id
 
     @property
     def address_space(self) -> RemoteAddressSpace:
+        """How the exported bytes are reached."""
         return self._address_space
 
     @property
     def offset(self) -> int:
+        """Start of the exported range within the owner buffer."""
         return self._offset
 
     @property
     def nbytes(self) -> int:
+        """Length of the exported range in bytes."""
         return self._nbytes
 
     @property
     def access_flags(self) -> int:
+        """Access granted here; a subset of the owner handle's flags."""
         return self._access_flags
 
     @property
     def transport_profile(self) -> str:
+        """Transport this export was minted for."""
         return self._transport_profile
 
     def __repr__(self) -> str:
@@ -585,6 +622,8 @@ class _RemoteTaskArgsSidecar:
 
 @dataclass(frozen=True)
 class RemoteTensorRef:
+    """A tensor argument that lives on, or travels to, a remote worker."""
+
     handle: RemoteBufferHandle
     offset: int = 0
     shape: tuple[int, ...] = ()
@@ -620,6 +659,11 @@ def __post_init__(self) -> None:
 
     @classmethod
     def host_inline(cls, payload: bytes, *, shape: tuple[int, ...], dtype: DataType) -> RemoteTensorRef:
+        """Build a reference whose payload travels inline, naming no remote memory.
+
+        ``payload`` length must equal the byte size implied by ``shape`` and
+        ``dtype``, and shape entries must be non-negative.
+        """
         data = bytes(payload)
         shape_tuple = tuple(int(x) for x in shape)
         if any(x < 0 for x in shape_tuple):
@@ -847,6 +891,10 @@ class CommBufferSpec:
 
 @dataclass
 class ChipDomainContext:
+    """Per-domain view handed to a chip worker: its rank within the domain and
+    the local slice of the symmetric window.
+    """
+
     name: str
     domain_rank: int
     domain_size: int
@@ -1346,8 +1394,10 @@ def comm_destroy_all(self) -> None:
 
     @property
     def device_id(self):
+        """ACL device ordinal this worker is bound to."""
         return self._impl.device_id
 
     @property
     def initialized(self):
+        """Whether the underlying native worker has completed init."""
         return self._impl.initialized

python/simpler/worker.py

@@ -377,15 +377,23 @@ def __post_init__(self) -> None:
 
     @property
     def module(self) -> str:
+        """Module half of the ``module:qualname`` target."""
         return self.target.split(":", 1)[0]
 
     @property
     def qualname(self) -> str:
+        """Qualified-name half of the ``module:qualname`` target."""
         return self.target.split(":", 1)[1]
 
 
 @dataclass(frozen=True)
 class RemoteWorkerSpec:
+    """Describes a remote L3 worker to attach via ``Worker.add_remote_worker``.
+
+    ``transport`` selects the data plane and is simulation-backed today; the
+    daemon rejects any other value.
+    """
+
     # endpoint is "host:port"; host must be a numeric IP (or "localhost").
     # Hostnames are rejected at add_remote_worker time — getaddrinfo resolution is
     # unbounded and uncancellable and would risk pinning startup on a hung DNS.
@@ -1903,9 +1911,10 @@ class _CloseAttempt:
 
     Teardown is single-shot and terminal: once it *runs* (``_teardown_attempted``
     latches True), an un-reclaimed resource LEAKS — a later close() never re-drives
-    a half-torn tree. The one retry path is a *drain-timeout*, which leaves
-    teardown UN-attempted and the tree intact; a later close() may then drive
-    drain+teardown once the in-flight operation finishes.
+    a half-torn tree. Teardown stays UN-attempted, with the tree intact and a
+    later close() free to drive drain+teardown, only where the drain itself did
+    not complete: in-flight leases outlived the cleanup budget, or an async
+    interruption left an accepted run fence undrained.
     """
 
     __slots__ = ("done", "error", "incomplete")
@@ -2396,6 +2405,16 @@ def _allocate_next_level_worker_id(self) -> int:
         return worker_id
 
     def add_remote_worker(self, spec: RemoteWorkerSpec) -> int:
+        """Register a remote L3 worker and return its NEXT_LEVEL worker id.
+
+        Must be called before ``init()`` — the topology freezes there — and only
+        on a ``level >= 4`` parent. ``spec.endpoint`` is validated here rather
+        than at activation, so a bad address fails before any process is forked;
+        its host must be a numeric IPv4 address or ``localhost``. IPv6 is not
+        reachable here: the endpoint is parsed as a single ``host:port`` pair, so
+        a literal carrying more than one colon is rejected before the numeric
+        check runs (see ``RemoteWorkerSpec``).
+        """
         # Hold the lifecycle lock across the state check and the topology
         # mutation so a concurrent init() cannot freeze the topology snapshot
         # between them.
@@ -2737,6 +2756,11 @@ def _send_remote_release_import_fields(self, fields: Any) -> None:
         )
 
     def remote_malloc(self, *, worker: int, nbytes: int) -> RemoteBufferHandle:
+        """Allocate ``nbytes`` on a started remote worker and return an owner handle.
+
+        ``nbytes`` must be positive. The target remote worker must already be
+        started, so this is callable only after ``init()``.
+        """
         worker_id = int(worker)
         size = int(nbytes)
         if size <= 0:
@@ -2757,6 +2781,14 @@ def remote_malloc(self, *, worker: int, nbytes: int) -> RemoteBufferHandle:
             )
 
     def remote_free(self, handle: RemoteBufferHandle) -> None:
+        """Free an owner remote allocation.
+
+        Idempotent: freeing an already-released handle is a no-op. Rejects
+        imported handles (use ``remote_release_import``) and ``HOST_INLINE``
+        handles, which are not remote allocations. If the buffer is still
+        referenced by a live task slot or by an outstanding import, the free is
+        recorded and deferred until those references drop rather than issued now.
+        """
         if not isinstance(handle, RemoteBufferHandle):
             raise TypeError("expected a RemoteBufferHandle returned by Worker.remote_malloc/import")
         if handle.address_space == RemoteAddressSpace.HOST_INLINE:
@@ -2778,6 +2810,11 @@ def remote_free(self, handle: RemoteBufferHandle) -> None:
             handle._mark_released()
 
     def remote_copy_to(self, handle: RemoteBufferHandle, host_ptr: Any, nbytes: int, *, offset: int = 0) -> None:
+        """Copy ``nbytes`` from host memory into an owner remote buffer.
+
+        Requires an owner handle, not an imported one. ``offset + nbytes`` must
+        fall within ``handle.nbytes``.
+        """
         with self._operation_lease("remote_copy_to"):
             self._require_live_remote_buffer(handle)
             if handle.is_imported:
@@ -2800,6 +2837,11 @@ def remote_copy_to(self, handle: RemoteBufferHandle, host_ptr: Any, nbytes: int,
             )
 
     def remote_copy_from(self, handle: RemoteBufferHandle, host_ptr: Any, nbytes: int, *, offset: int = 0) -> None:
+        """Copy ``nbytes`` out of an owner remote buffer into host memory.
+
+        Requires an owner handle, not an imported one. ``offset + nbytes`` must
+        fall within ``handle.nbytes``.
+        """
         with self._operation_lease("remote_copy_from"):
             self._require_live_remote_buffer(handle)
             if handle.is_imported:
@@ -2830,6 +2872,12 @@ def remote_export(
         access: str | int = "readwrite",
         transport_profile: str = "sim",
     ) -> RemoteBufferExport:
+        """Export a range of an owner buffer so another worker can import it.
+
+        ``nbytes=None`` exports from ``offset`` to the end of the buffer. The
+        requested ``access`` must be a subset of the handle's own access flags —
+        an export can narrow permissions but never widen them.
+        """
         with self._operation_lease("remote_export"):
             return self._remote_export_locked(
                 handle, offset=offset, nbytes=nbytes, access=access, transport_profile=transport_profile
@@ -2889,6 +2937,11 @@ def _remote_export_locked(
     def remote_import(
         self, exported: RemoteBufferExport, *, worker: int, access: str | int | None = None
     ) -> RemoteBufferHandle:
+        """Import an exported buffer on ``worker`` and return an imported handle.
+
+        ``access`` defaults to the export's own flags. Rejects an export minted
+        by a different ``Worker`` and one whose owner buffer has been freed.
+        """
         # Argument validation (type / forged / stale) is independent of lifecycle
         # and runs before admission; the lease guards the actual transport.
         if not isinstance(exported, RemoteBufferExport):
@@ -2957,6 +3010,12 @@ def _remote_import_locked(
             raise
 
     def remote_release_import(self, handle: RemoteBufferHandle) -> None:
+        """Release an imported remote handle.
+
+        Idempotent, and rejects owner handles (use ``remote_free``). Deferred
+        while a live task slot still references it. Releasing the last import of
+        a buffer whose owner already called ``remote_free`` completes that free.
+        """
         if not isinstance(handle, RemoteBufferHandle):
             raise TypeError("expected a RemoteBufferHandle returned by Worker.remote_import")
         if not handle.is_imported:
@@ -6214,6 +6273,26 @@ def _describe_live_resources(self) -> str:
         return ", ".join(parts) if parts else "(none)"
 
     def close(self) -> None:  # noqa: PLR0912, PLR0915 -- lifecycle linearization: reentrancy / init-guard / join / owner / claim / drain / teardown
+        """Release this worker's resources. Terminal and single-shot.
+
+        A permanent commitment, not a reversible attempt: CLOSED is published
+        atomically and never reverts to READY, and the leased live-tree APIs are
+        rejected from then on. Put the call in a ``finally`` — a worker that is
+        never closed keeps its device held.
+
+        - Reentrant ``close()`` from inside a leased operation is rejected.
+        - ``close()`` during an in-progress ``init()`` fails fast; this worker
+          does not cancel initialization. Wait for READY or FAILED.
+        - A concurrent ``close()`` joins the in-flight attempt and observes its
+          result; teardown never runs twice at once.
+        - Teardown is single-shot. Once it runs, a later ``close()`` re-raises
+          the same result rather than re-driving a half-torn tree. A retry is
+          permitted only where the drain did not complete and so left teardown
+          un-attempted and the tree intact: either in-flight leases outlived
+          the cleanup budget, or an asynchronous interruption left an accepted
+          run fence undrained.
+        - Native teardown runs on the ``init()``-owner thread, being device-bound.
+        """
         # close() is a permanent commitment against a resource, not a reversible
         # attempt: it publishes CLOSED atomically (the sole public admission
         # fence — the leased live-tree APIs are rejected once CLOSED) and NEVER
@@ -6225,9 +6304,11 @@ def close(self) -> None:  # noqa: PLR0912, PLR0915 -- lifecycle linearization: r
         #     result; the same worker's teardown never runs twice at once;
         #   - teardown is single-shot and TERMINAL: once it runs, an un-reclaimed
         #     resource leaks and a later close() re-raises the same result — it
-        #     never re-drives a half-torn tree. Only a drain-timeout (teardown
-        #     un-attempted, tree intact) lets a later close() retry once the
-        #     in-flight op finishes; a tree with a live op is never torn down;
+        #     never re-drives a half-torn tree. A later close() may retry only
+        #     where the drain left teardown un-attempted and the tree intact —
+        #     leases outliving the budget, or an async interruption leaving an
+        #     accepted run fence undrained; a tree with a live op is never torn
+        #     down;
         #   - native teardown runs only on the init-owner thread (device-bound).
         # `attempt` is None until the claim installs it. The pre-claim checks
         # raise/return before that, so the finally skips completion for them. From
@@ -6300,7 +6381,8 @@ def close(self) -> None:  # noqa: PLR0912, PLR0915 -- lifecycle linearization: r
                 # rejects new leases; a tree with a live op is never torn down.
                 # If an op outlives the budget, teardown stays UN-attempted and
                 # the tree intact so a later close() can retry once it drains —
-                # the one retryable close() path.
+                # one of the two paths that keep close() retryable (the other is
+                # an async interruption leaving an accepted run fence undrained).
                 if self._active_ops > 0:
                     drain_deadline = time.monotonic() + _ROLLBACK_GRACEFUL_TIMEOUT_S
                     while self._active_ops > 0: