Skip to content

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

Import-path descriptor for a parent-facing remote L3 callable.

module property

module: str

Module half of the module:qualname target.

qualname property

qualname: str

Qualified-name half of the module:qualname target.

RemoteWorkerSpec dataclass

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.

HostBuffer dataclass

Handle for a worker-allocated, born-shared host buffer (zero-copy).

Returned by Worker.create_host_buffer. buffer is a memoryview over shared memory already attached into every local L3 child; wrap it with torch.frombuffer / np.frombuffer to get a real tensor whose writes land directly in the child-visible pages — no per-run copy. token / data_ptr / nbytes identify the mapping; pass this handle back to free_host_buffer to release it.

eq=False keeps object-identity equality/hash so the (unhashable) memoryview field never blocks using the handle as a dict key or set member.

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

done: bool

True once this run is terminal and its cleanup has been published.

Reads False while another waiter is still publishing cleanup, so a run whose native fence has already fired can report False until that waiter finishes.

wait

wait(timeout: float | None = None) -> None

Wait for completion, raising TimeoutError or the run's error.

result

result(timeout: float | None = None) -> None

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

live_domains: dict[str, CommDomainHandle]

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

aicpu_dlopen_count: int

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

host_dlopen_count: int

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

run_stream_set_create_count: int

L2 only: number of AICore run streams the runner has created.

AICPU streams belong to pipeline slots for the worker's lifetime, while each run creates and retires its own AICore stream, so this advances once per run. Returns 0 on non-L2 workers and on platforms whose runs use the persistent bootstrap stream pair (simulation, a5).

add_remote_worker

add_remote_worker(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).

remote_malloc

remote_malloc(*, 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().

remote_free

remote_free(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.

remote_copy_to

remote_copy_to(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.

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 = '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.

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

remote_release_import(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.

register

register(target, *, workers: list[int] | None = None) -> CallableHandle

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.

Target eligibility (a callable's kind having a resolving child) is checked only at init(), over the pre-init registrations (_validate_eligible_targets). A post-init dynamic register does NOT re-validate against the frozen topology, so registering e.g. a ChipCallable on a chipless worker yields a handle that never dispatches. Unifying the two paths is a follow-up (needs a device-free chip-child test harness).

unregister

unregister(handle_or_slot) -> None

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_worker(worker: Worker) -> int

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

init(prewarm_config: CallConfig | None = None, *, _startup_deadline: float | None = None) -> None

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_host_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 (runtime_env.ring_task_window / ring_heap / ring_dep_pool) is built + cached so the first run with the same sizing skips the (~800ms) cold prebuilt runtime-arena build. An L2 worker prewarms here; an L3+ worker prewarms each chip child during hierarchy startup, before it publishes INIT_READY. A no-op for runtimes without a prebuilt arena (host_build_graph). None (default) disables prewarm.

None
_startup_deadline float | None

Internal. Absolute time.monotonic() deadline inherited from a parent's startup epoch so a recursive descendant consumes the parent's remaining budget instead of restarting the timeout. None starts a fresh epoch.

None

malloc

malloc(size: int, worker_id: int = 0) -> int

Allocate memory on next-level chip worker worker_id. Returns a pointer.

free

free(ptr: int, worker_id: int = 0) -> None

Free memory allocated by malloc().

copy_to

copy_to(dst: int, src: int, size: int, worker_id: int = 0) -> None

Copy size bytes from host src to chip worker dst.

copy_from

copy_from(dst: int, src: int, size: int, worker_id: int = 0) -> None

Copy size bytes from chip worker src to host dst.

create_host_buffer

create_host_buffer(nbytes: int) -> HostBuffer

Allocate a born-shared host buffer, attached into every local L3 child, that a later run() reads/writes with no per-run copy.

Local L3 children are forked during init(); host memory allocated afterwards is not in their address space. This hands you memory that is born in a shm already attached into every child, so there is nothing to copy: the child reads and writes the same physical pages the parent sees.

Returns a :class:HostBuffer whose buffer is a memoryview over that shm. Build a tensor over it with the buffer protocol, framework of your choice, and pass it to run() as usual::

buf = worker.create_host_buffer(n * 4)
t = torch.frombuffer(buf.buffer, dtype=torch.float32, count=n)
t.uniform_(0, 1)                       # in place → lands in the shm
worker.run(orch(chip, t, out), args=None, config=CallConfig())
worker.free_host_buffer(buf)           # drop the tensor first

simpler stays framework-free: torch/numpy appear only on the user's side (frombuffer). Blocks until every local L3 child has attached the buffer; not thread-safe against a concurrent run / create / free on the same Worker — drive them from one thread, as the L3 worker is otherwise.

free_host_buffer

free_host_buffer(handle: HostBuffer) -> None

Release a born-shared buffer created by create_host_buffer.

Unmaps it from every local L3 child and frees the parent shm. Drop every tensor / memoryview you built over handle.buffer first: a live view keeps the shm's pages exported, so close() cannot release them and the buffer only warns (and is reclaimed once the last view is gone).

Best-effort and idempotent: a stale handle whose token no longer matches (e.g. freed twice) is a silent no-op.

submit

submit(callable, args=None, config=None) -> RunHandle

Submit one task (L2) or one DAG (L3+) and return its completion handle.

Dispatch
  • L2: callable is a CallableHandle returned by Worker.register(chip_callable). Routes to the private slot carried by the handle. The current L2 backend remains blocking, so the returned handle is already complete.
  • L3+: callable is a Python orch fn invoked with the Orchestrator handle. 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. A later submission waits until prior dispatches are accepted; completion and cleanup stay attached to each returned handle.

run

run(callable, args=None, config=None) -> None

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

close() -> None

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.