diff --git a/python/simpler/orchestrator.py b/python/simpler/orchestrator.py index 789918ec2..8b457870e 100644 --- a/python/simpler/orchestrator.py +++ b/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 diff --git a/python/simpler/task_interface.py b/python/simpler/task_interface.py index b18142329..4a817b5d5 100644 --- a/python/simpler/task_interface.py +++ b/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 diff --git a/python/simpler/worker.py b/python/simpler/worker.py index 2c3fa7bb6..14f6fe4aa 100644 --- a/python/simpler/worker.py +++ b/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: