simpler.worker¶
Generated from the source. For the curated view — which **config keys
Worker accepts, CallConfig defaults, and the argument-order footguns — see
Python API.
simpler.worker ¶
Worker — unified factory for all hierarchy levels.
Callable identity is exposed as an opaque CallableHandle returned by
Worker.register(callable). L2 Worker.run and hierarchical
Orchestrator.submit_next_level / submit_sub consume handles, never raw
ChipCallable objects. L3+ Worker.run keeps the existing raw Python
orchestration-function entry point; that function captures handles and submits
them through the Orchestrator. L≥3 targets resolve the handle's stable SHA-256
digest to a private L2-side slot; later Python registrations are serialized
and sent through the mailbox control plane.
Usage::
# L2: one NPU chip
w = Worker(level=2, device_id=8, platform="a2a3", runtime="tensormap_and_ringbuffer")
w.init()
chip_handle = w.register(chip_callable) # L2 may register pre or post init()
w.run(chip_handle, chip_args, config)
w.close()
# L3: multiple chips + SubWorkers, auto-discovery in init()
w = Worker(level=3, device_ids=[8, 9], num_sub_workers=2,
platform="a2a3", runtime="tensormap_and_ringbuffer")
chip_handle = w.register(chip_callable) # ChipCallable, before init()
sub_handle = w.register(lambda args: postprocess()) # Python sub, before init()
w.init()
def my_orch(orch, args, cfg):
r = orch.submit_next_level(chip_handle, chip_args_ptr, cfg, worker=0)
orch.submit_sub(sub_handle, sub_args)
w.run(my_orch, my_args, my_config)
w.close()
# L4: recursive composition — L3 Workers as children
l3 = Worker(level=3, device_ids=[8, 9], num_sub_workers=1,
platform="a2a3", runtime="tensormap_and_ringbuffer")
w4 = Worker(level=4, num_sub_workers=1)
l3_handle = w4.register(my_l3_orch)
verify_handle = w4.register(lambda args: verify())
l3_worker_id = w4.add_worker(l3)
w4.init()
def my_l4_orch(orch, args, config):
orch.submit_next_level(l3_handle, chip_args, config, worker=l3_worker_id)
orch.submit_sub(verify_handle)
w4.run(my_l4_orch)
w4.close()
RemoteCallable
dataclass
¶
RemoteWorkerSpec
dataclass
¶
Describes a remote L3 worker to attach via Worker.add_remote_worker.
transport selects the data plane. The shipped daemon accepts
only the host_tcp profile today.
MpiL3GroupSpec
dataclass
¶
Describes L3 workers launched by one parent-owned mpirun.
command_port_base, health_port_base, session_listen_hosts,
connect_hosts, allow_wildcard_session_bind, ready_host, and
ready_port are accepted only for source compatibility with PR #1623.
MPI groups ignore them and use the local named mailbox plus MPI collectives.
Non-MPI RemoteWorkerSpec continues to use its TCP fields.
CleanupJournal ¶
Post-success resource cleanup journal shared by _teardown_ready_tree and _abort_hierarchical. Entries removed only after native free succeeds.
drive_kinds ¶
Drive every retained entry whose resource kind is in kinds.
InitCancelled ¶
Bases: RuntimeError
Raised from Worker.init() when a concurrent close() cancelled the
startup epoch before it committed READY.
An ordinary RuntimeError — the init-owner thread is a normal caller, so
the cancellation must be catchable by except Exception. Distinct from
_StartupCancelled, which is signal-delivered inside a forked child.
RunHandle ¶
Completion handle returned by :meth:Worker.submit.
A handle owns the run's Python keepalives and keeps its Worker alive until the native completion fence has fired and run-owned resources are cleaned up. Waiting is idempotent; every waiter observes the same terminal result.
done
property
¶
True once waiting on this run would not block.
This answers the native fence — the device is drained and every task
has reached its terminal state — plus this handle's own terminal flag.
It does not say the run's fence-owned cleanup has happened: the
CommDomain, L3-L2 and remote-slot teardown runs after the fence, in
whichever thread waits. Anything that has to be ordered behind that
teardown keys on _cleanup_published instead.
Reads False while another waiter is crossing the fence, because the native run identity can disappear in that interval.
wait ¶
Wait for completion, raising TimeoutError or the run's error.
result ¶
Alias for :meth:wait; successful runs have no return value.
Worker ¶
Unified worker for all hierarchy levels.
level=2: wraps the C++ ChipWorker (one NPU device). level=3: wraps the C++ Worker composite with ChipWorker×N + SubWorker×M, auto-created in init() from device_ids and num_sub_workers. level=4+: wraps the C++ Worker composite with Worker(level-1)×N as NEXT_LEVEL children + SubWorker×M. Children are added via add_worker() before init().
live_domains
property
¶
Read-only snapshot of currently-live dynamic CommDomain handles.
Useful for debugging. Mutating the returned dict has no effect; use
handle.release() or orch.release_domain(handle) to free.
aicpu_dlopen_count
property
¶
L2 only: number of distinct callable identities the AICPU has dlopened for.
Used by tests to assert that register + repeated run(handle)
calls do not retrigger the AICPU dlopen for an already-seen identity.
Returns 0 on non-L2 workers.
host_dlopen_count
property
¶
L2 only: number of host-side orch SO dlopens (hbg variants).
Mirrors aicpu_dlopen_count for the host_build_graph path. Returns
0 on non-L2 workers or device-orch variants (trb).
run_stream_set_create_count
property
¶
L2 only: number of AICore run streams the runner has created.
One AICPU + AICore pair serves every run for the runner's lifetime. The AICPU stream persists; the AICore stream is recreated when a new code upload makes it stale, and destroyed when an unproven completion retires it, so this advances per publication or unproven retirement rather than once per run or per pipeline slot. Returns 0 on non-L2 workers and on platforms whose runs use the persistent bootstrap stream pair (simulation, a5).
add_remote_worker ¶
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).
add_mpirun_worker_group ¶
Register L3 workers that will be launched by one parent-owned mpirun.
A single named mailbox is created for the group during init().
The returned ids remain exact NEXT_LEVEL targets, but all of them route
through that group mailbox and the MPI collective dispatcher.
remote_malloc ¶
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().
remote_free ¶
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.
remote_copy_to ¶
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.
remote_copy_from ¶
remote_copy_from(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.
remote_export ¶
remote_export(handle: RemoteBufferHandle, *, offset: int = 0, nbytes: int | None = None, access: str | int = 'readwrite', transport_profile: str = HOST_TCP_TRANSPORT_PROFILE) -> 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.
remote_import ¶
remote_import(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.
remote_release_import ¶
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.
register ¶
Register a callable for dispatch and return an opaque handle.
Integer execution slots remain private to the local target process. Submit APIs consume the returned handle and dispatch by its stable SHA-256 callable identity.
A post-init dynamic register re-validates eligibility against the
frozen topology (_eligible_target_need), same as init().
unregister ¶
Drop a CallableHandle from the registry and propagate cleanup.
Symmetric to Worker.register for the dynamic post-init path.
The target-local resources become reusable for the next
register call — the only practical way to keep a long-running worker under the
MAX_REGISTERED_CALLABLE_IDS ceiling when JIT or plugin code
churns through callables.
Failure semantics (docs section 8): unregister is best-effort. If any chip child reports an error, the parent warns and still pops the registry entry — orch_so_table_ on the AICPU side will be overwritten on target-local resource reuse, and refusing to release a known-bad entry would just exhaust the resource space faster.
Raises:
| Type | Description |
|---|---|
KeyError
|
handle was never registered. |
add_worker ¶
Add a lower-level Worker as a NEXT_LEVEL child. Must be called before init().
The child Worker must NOT be init'd — init happens inside the forked child process (so the child's own children are forked in the right process tree). Returns this child's stable NEXT_LEVEL worker id.
init ¶
Initialize the worker and bring its whole subtree to READY.
For an L3+ worker init is the single startup submission point: it
forks every local child (sub / chip / next-level), waits for the whole
subtree — recursively, for L4+ — to publish INIT_READY, activates any
remote L3 sessions, starts the C++ scheduler, and only then publishes
READY in one atomic commit. It returns with the tree ready to run, or
raises after a bounded rollback that reaps the children it forked
best-effort (a child wedged in native code past the deadline may be left
behind — see the deferred un-reaped-child / nested-shm items).
run / create_buffer / the remote register/memory APIs never
trigger startup.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
prewarm_config
|
CallConfig | None
|
Optional CallConfig. When given, its ring sizing
( |
None
|
_startup_deadline
|
float | None
|
Internal. Absolute |
None
|
malloc ¶
Allocate device memory on this L2 worker's own chip; returns a DEVICE_MALLOC Buffer.
Name a task arg with handle.tensor(shapes, dtype) and release with worker.free(handle). L3+
allocates child device memory with alloc_child_tensor(worker_id, ...) instead — a Worker is
the only allocator, the Orchestrator never allocates.
alloc_child_tensor ¶
Allocate device memory on next-level worker_id sized for shapes × dtype; returns a
DEVICE_MALLOC Buffer (successor of orch.malloc + child_memory).
Called from within an orchestration fn (capture the Worker in the closure). The pointer is
private to worker_id; name the arg with handle.tensor(shapes, dtype), dispatch it only to
that worker, and load host data with copy_to. Not auto-freed at end-of-task.
free ¶
Free a device Buffer allocated by malloc / alloc_child_tensor.
The operation lease is re-entrant, so an in-run orch.free that delegates here nests safely.
committed_device_memory ¶
Total device HBM (bytes) committed by chip worker worker_id's
MemoryAllocator (tensors + pooled arenas + runtime buffers; excludes
HCCL/VMM comm windows). Useful for downstream
runtimes to subtract simpler's own HBM from their cache budget.
Level 2 returns the in-process chip worker's committed bytes directly;
level 3 forwards a CTRL_COMMITTED_DEVICE_MEMORY query to the forked
chip child worker_id (sum across worker_ids for a multi-chip total).
device_memory_info ¶
Return the target device's ACL_HBM_MEM free/total byte snapshot.
Level 2 queries the in-process chip worker. Level 3 routes by logical
worker_id to the matching forked chip child. Simulator backends do
not synthesize device-wide memory and raise NotImplementedError.
copy_to ¶
copy_to(dst: Buffer, src, *, dst_offset: int = 0, src_offset: int = 0, nbytes: int | None = None) -> None
H2D: copy nbytes from src_offset in host src to dst_offset in device dst.
src is a host Buffer; the chip child resolves both handles through its
ImportRegistry and reads the host backing directly. At L2 the chip worker shares this
process, so a torch tensor or any writable buffer works too.
nbytes defaults to the rest of the host side after src_offset, so a plain
copy_to(dst, src) still transfers the whole host backing. An offset names a range of the
allocation dst already names; there is no way to name a sub-range with a handle built at
an interior address, because such a handle names no allocation at all.
copy_from ¶
copy_from(dst, src: Buffer, *, dst_offset: int = 0, src_offset: int = 0, nbytes: int | None = None) -> None
D2H: copy nbytes from src_offset in device src to dst_offset in host dst.
dst is a host Buffer; the chip child resolves both handles through its
ImportRegistry and writes the host backing directly. At L2 the chip worker shares this
process, so a torch tensor or any writable buffer works too.
nbytes defaults to the rest of the host side after dst_offset, so a plain
copy_from(dst, src) still transfers a whole host backing's worth.
create_buffer ¶
Allocate a shared Buffer owned by this Worker (P1-B).
The backing is a POSIX shm; the Buffer carries a typed canonical identity and a
self-describing descriptor, so a consumer can resolve it with no prior handshake: the
descriptor travels embedded in every Tensor built over this Buffer and the consumer
materializes it lazily on first receipt (map-once, keyed by canonical identity). At L3+ that
consumer is a forked child; at L2 (a leaf, no children) the Worker itself materializes the
tensor in-process on run. Build a tensor over buffer.shm.buf with the buffer protocol.
Not thread-safe against a concurrent run/create/free on the same Worker.
alloc_shared_tensor ¶
Allocate a runtime-managed intermediate buffer (the Tensor form of orch.alloc).
Called inside an orchestration fn. The backing comes from the orchestrator's HeapRing
(MAP_SHARED, visible to forked children) and is auto-reclaimed once every downstream consumer
has completed and the scope ends — no manual free. Returns a FORK_SHM Buffer whose
canonical identity is registered in the tensormap so a view of it (handle.tensor(shapes, dtype))
dependency-wires to this producer slot. Chip-A→chip-B intermediates: name it as an OUTPUT of the
producing task and an INPUT of the consumer.
make_tensor_arg ¶
make_tensor_arg(tensor, shapes: tuple[int, ...], dtype: int, *, strides: tuple[int, ...] | None = None)
Name a pre-fork host tensor as a Tensor over a memoized FORK_SHM handle.
The torch (or buffer-protocol) tensor MUST be allocated before init() so its VA is
fork-inherited by the children (the mainline "fork-inherited" contract). A share_memory_()
tensor is MAP_SHARED — read-write across the fork, so usable as an OUTPUT the parent reads back;
a plain tensor is COW read-only (input only). The handle is memoized by the tensor's storage
base, so every ref over the same storage shares one canonical identity and dependencies key on
it; the byte_offset this computes is what then separates two views that do not intersect.
At L2 (no fork) any host tensor works. dtype is the DataType int value.
release_buffer ¶
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 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. _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 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 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 converse does not hold, and is the one asymmetry here: a
Buffer.close() whose shm.close() raises has still unlinked the name (its finally
runs the unlink), so the backing can be nameless while descendants keep mappings this call
never told them to drop. A descendant materializing that identity afterwards gets the named
FileNotFoundError from ImportRegistry.materialize, not a silent bad mapping.
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.
submit ¶
Submit one task (L2) or one DAG (L3+) and return its completion handle.
Dispatch
- L2:
callableis aCallableHandlereturned byWorker.register(chip_callable). Routes to the private slot carried by the handle and returns a live completion handle. - L3+:
callableis a Python orch fn invoked with theOrchestratorhandle. Graph construction completes synchronously; device completion is reported by the returned handle.
args : TaskArgs (optional)
config: CallConfig (optional, default-constructed if None)
Graph construction remains serialized. How many runs may be admitted is the depth this worker's backends negotiated, not a constant: at the negotiated depth two, one active plus one prepared run are permitted and a third submission blocks before invoking its graph callback; where a backend publishes depth one, the second submission already blocks there. A5 tensor-map-and-ring-buffer publishes depth two, while its local endpoint retains one mailbox frame and serial device execution; A5 host-build-graph publishes no contract and stays at depth one. A caller whose first run only completes because a later callback runs would deadlock on a depth-one backend. Completion and cleanup stay attached to each handle.
run ¶
Execute one task or DAG synchronously as submit(...).wait().
Per-stage run timing (host wall, on-NPU device wall +
AICPU phase breakdown) is no longer returned — the platform emits it as
[STRACE] log markers from each L2 simpler_run, so the L3
dispatcher and its L2 children are observed uniformly. Parse the markers
with simpler_setup.tools.strace_timing (see
docs/dfx/host-trace.md).
close ¶
Release this worker's resources. Publicly terminal and retryable.
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-progressinit()on another thread cooperatively cancels it: the init epoch unwinds and this call proceeds to teardown. Cancellation is observed only at cooperative points, so an init blocked inside a native segment is not interrupted — this call raises after_CLOSE_CANCEL_UNWIND_TIMEOUT_Srather than blocking forever. Closing from the init-owner thread itself is rejected.- A concurrent
close()joins the in-flight attempt and observes its result; teardown never runs twice at once. - A later
close()retries journaled teardown debt. The journal keeps each resource until its native free succeeds and preserves the child pid/mailbox pair untilwaitpidproves the child is gone. - Native teardown runs on the
init()-owner thread, being device-bound.
attach_exception_note ¶
Mirror BaseException.add_note onto error.__notes__.
Interpreters without add_note still expose the list on __notes__.