Runs¶

mujou.run ¶

run(
    root: Node | Network,
    *,
    timer: Timer | None = None,
    inputs: Mapping[P, Message] | None = None,
    state: Mapping[str, Any] | None = None,
    workers: int | None = None,
) -> Run

Run root, a node or a network, in the background; return once every node has started.

The nodes are at their paths from the root: those of a network at "l1", "body.l1", ..., and a node given as the root at "". To run several side by side, put them in a Group. The host writes and reads the ports themselves (run.write({node.prev: x})).

Every node's on_start (and, with state, the node's load_state_dict) has returned, and the run's timer runs when this returns. A node runs once: run new instances to run again (the attributes of the nodes of a run stay readable after it). A run refused before any node started (by any error but NetworkError) has changed nothing, and the same instances run once the cause is fixed.

The run's nodes share a pool of workers threads, which it creates for itself; a node with dedicated_thread=True runs on a thread of its own instead. A worker runs one hook of one node at a time, the updates that are due earliest first. A hook that blocks holds its worker, so the other nodes' dt grows meanwhile.

PARAMETER	DESCRIPTION
`root`	The node or network to run. TYPE: `Node \| Network`
`timer`	The time axis of the run: every node's `dt`, every message's time and every pause of the run is on it. `None` is `mujou.time.global_timer`, which every run without a timer of its own shares. TYPE: `Timer \| None` DEFAULT: `None`
`inputs`	The initial values of unconnected ports that receive, as `Run.write` takes them (and copies them, to the same places), written when the run starts. Written after `state`, they replace its values. TYPE: `Mapping[P, Message] \| None` DEFAULT: `None`
`state`	A tree from `Run.state_dict` to restore. The slots, the nodes' states (through each node's `load_state_dict`, after its `on_start`) and their last update times continue from the time the tree was taken, on a new `Timer` of the run's own that starts at that time and time scale (so `global_timer` and the other runs do not move, and a pause of `global_timer` does not pause it). Build the tree to run as the one the tree came from: the same node paths, ports and peers. TYPE: `Mapping[str, Any] \| None` DEFAULT: `None`
`workers`	The number of worker threads of the run (>= 1). `None` is as many as the CPUs the process may use (the logical cores, or fewer when the CPU affinity or a quota limits them), or half of them (at least 1) while the GIL is enabled at run time (`sys._is_gil_enabled()`), never more than the nodes on the workers. With no such node, the run has no worker. TYPE: `int \| None` DEFAULT: `None`

Example, with the node Error of the example of Run::

node = Error()
with mujou.run(node, inputs={node.prev: torch.ones(2)}) as run:
    run.timer.sleep(0.1)
    error = run.read(node.next)

RAISES	DESCRIPTION
`RuntimeError`	A node of the tree has already run (`"Leaky has already run; create a new instance (its attributes stay readable)"`); or this is called from a node's hook.
`(TypeError, ValueError)`	The tree is invalid (as `Node` and `Network` say); or `inputs` or the `initial` value of a port do not match their specs (as `Run.write` raises).
`ValueError`	`workers` is less than 1; `timer` and `state` are both given (a restored run runs on a timer of its own); or `state` does not have the timer, nodes, ports and peers of the tree.
`NetworkError`	A node raised while starting; the run has stopped.

mujou.Run ¶

Run(wiring: Wiring, runtime: Runtime, timer: Timer, clock: ReadOnlyTimer)

A run of a tree of nodes, as mujou.run returns it: what the host uses to talk to the running nodes.

The nodes run in the background, on the run's worker threads (or threads of their own); the host's thread only writes inputs, reads outputs and takes the state:

write inputs, read / read_with_time / subscribe / wait_for outputs, take a state_dict.
shutdown() to stop the run, or join() to wait until it stops.
After a normal end: take the final state_dict, and read the nodes' attributes (a module they learned, say).

A with statement shuts the run down at its end. These methods are for the host: calling them from a node's hook raises RuntimeError.

If a node raises, the whole run stops. If that happens while mujou.run starts the nodes, mujou.run raises NetworkError; otherwise every later call of write, read, read_with_time, subscribe (and its iterator), wait_for, state_dict, shutdown, join and ManualRun.step raises it. Its __cause__ is the node's exception. The state at the exception is not kept. A with left by another exception lets that exception through instead.

A run keeps running when the host drops it, until shutdown or, at the latest, the interpreter's exit: once the non-daemon threads have ended, the exit shuts down every run still running and waits for their on_stop. A thread that iterates over subscribe keeps the exit waiting until its run ends, so shut the run down before the program ends (with does).

mujou.run and mujou.testing.run make it.

Example::

from typing import override

import torch

import mujou
from mujou import Node, Port, TensorSpec

VECTOR = TensorSpec((2,))


class Error(Node):  # sends its input minus its teacher
    def __init__(self) -> None:
        super().__init__()
        self.prev = Port(recv=VECTOR)
        self.teacher = Port(recv=VECTOR)
        self.next = Port(send=VECTOR)

    @override
    def on_update(self, dt: float) -> None:
        super().on_update(dt)
        self.next.send(self.prev.recv() - self.teacher.recv())


node = Error()
x, y = torch.ones(2), torch.zeros(2)
with mujou.run(node, inputs={node.prev: x, node.teacher: y}) as run:
    run.write({node.prev: 2 * x, node.teacher: x})  # switch together
    run.timer.sleep(0.1)  # present the input for 0.1 s of the run's time
    error, t = run.read_with_time(node.next)
state = run.state_dict()  # the state at the end, to save

timer `property` ¶

timer: Timer

The run's timer: the timer given to mujou.run, mujou.time.global_timer when none was given, or, for a run restored from a state, a new Timer that started at the time of the state.

Every node's dt and every time of the run is on it. The host waits, pauses and scales the time of the run with it.

write ¶

write(values: Mapping[P, Message]) -> Timestamp

Write messages to unconnected ports that receive, and return the time they are stamped with.

The values are detached and copied. A dense tensor on the CUDA device of the port's node is copied on that device, on this thread's current stream of it, which write waits for; any other tensor is copied to host memory (so a node on another device or on the CPU never waits for what this thread queues on a GPU). The nodes receive the values on their own devices. All of them switch together and bear the same time: at any one moment, a node sees either all of them or none. (A node that reads two of them one after the other may still see a later write in between; compare the times to tell.) Any such port of the run can be written, exposed by the root or not.

The time is the run's timer at the write, as a node's recv_with_time and the "in" slots of state_dict show it; read_with_time shows the times of sends, not of writes. An empty values changes nothing and returns the time all the same.

RAISES	DESCRIPTION
`RuntimeError`	The run has stopped normally (no node reads what is written then).
`ValueError`	A port is not an unconnected `Port` of this run that receives; or a shape or a dtype differs from the spec.
`TypeError`	A value is not of its field's kind (a tensor, an `int`, ...), or a dict message does not have the spec's keys.
`OverflowError`	An `int` is out of the range of its `IntSpec`.
`NetworkError`	The run stopped because a node raised.

read ¶

read(port: Port[S, Any] | MultiPort[S, Any]) -> S

Return the latest message port sent.

Any port of the run that sends can be read, so a connected port shows an inner node. Only send (to every peer, not MultiPort.send_to) updates the value read here; before the first one it is the port's initial (zeros by default) or the value restored from state. The tensors are new copies on the CPU. read_with_time also returns the time it was sent at.

RAISES	DESCRIPTION
`ValueError`	`port` is not a port of this run that sends.
`NetworkError`	The run stopped because a node raised.

read_with_time ¶

read_with_time(port: Port[S, Any] | MultiPort[S, Any]) -> tuple[S, Timestamp | None]

Return the latest message port sent, as read does, and the time it was sent at: None before the first send (the port's initial, or the restored value of a port that had not sent).

The time is on the run's timer. A Timestamp does not say which timer it is on: compare it only with times of the same timer (run.timer).

RAISES	DESCRIPTION
`ValueError`	As `read` raises.
`NetworkError`	The run stopped because a node raised.

subscribe(
    port: Port[S, Any] | MultiPort[S, Any], *, timeout: float = math.inf
) -> Iterator[tuple[S, Timestamp | None]]

Iterate over the messages port sends, with their times.

The first step yields the latest message at the time of this call without waiting (before the first send, the port's initial or the restored value, with the time read_with_time gives it). Each later step waits until a newer message is sent and yields the latest; messages sent in between are skipped. Only the consumer waits: the sending node never does. The iterator ends when the run shuts down normally. Without a timeout, a thread iterating here ends only with the run. In a ManualRun nothing is sent unless step is called, so a wait here without a timeout on the thread that steps never ends.

To wait for one message that meets a condition, use wait_for.

PARAMETER	DESCRIPTION
`port`	The port to follow. TYPE: `Port[S, Any] \| MultiPort[S, Any]`
`timeout`	Real (wall-clock) seconds, counted from this call, that the whole iteration may take, the time spent in the loop's body included. `math.inf` iterates as long as the run runs. TYPE: `float` DEFAULT: `inf`

RAISES	DESCRIPTION
`ValueError`	As `read` raises.
`TimeoutError`	From the iterator, at a step after the first that begins once `timeout` has passed, or that would wait past it. The iterator ends; the run keeps running.
`NetworkError`	The run stopped because a node raised (also from the iterator).

wait_for ¶

wait_for(
    port: Port[S, Any] | MultiPort[S, Any],
    until: Callable[[S], bool],
    *,
    timeout: float = math.inf,
) -> tuple[S, Timestamp | None]

Wait until port sends a message that meets until, and return it with its time (None if the message that meets it is the one before the first send, as read_with_time gives it).

Checks the latest message first, then each newer one as subscribe yields them (messages sent in between are skipped). Only the host waits: the nodes never do. In a ManualRun nothing is sent unless step is called, so a wait without a timeout on the thread that steps never ends. Here with the node Error of the example of Run::

node = Error()
x = torch.ones(2)
with mujou.run(node, inputs={node.prev: x}) as run:
    error, t = run.wait_for(
        node.next, lambda e: torch.equal(e, x), timeout=5.0
    )

PARAMETER	DESCRIPTION
`port`	The port to follow. TYPE: `Port[S, Any] \| MultiPort[S, Any]`
`until`	Called on each message on the host's thread; `True` ends the wait. TYPE: `Callable[[S], bool]`
`timeout`	Real (wall-clock) seconds, counted from this call, to wait at most. `math.inf` waits as long as the run runs. TYPE: `float` DEFAULT: `inf`

RAISES	DESCRIPTION
`ValueError`	As `read` raises.
`TimeoutError`	`timeout` passed first. The run keeps running.
`RuntimeError`	The run shut down normally first.
`NetworkError`	The run stopped because a node raised.

state_dict ¶

state_dict(*, timeout: float = math.inf) -> dict[str, Any]

Copy the state of the run at one time, while it runs or after it ended normally.

While it runs, this pauses the run's timer and nodes, copies the state and resumes, and returns once the run has resumed. It waits at most timeout for every node to pause. After a normal end, this copies the state as the nodes left it, calling each node's state_dict on this thread. The tree's timer holds the time at which the run was asked to stop, as if the run had paused then: no update or message of the run is later, though its time scale is the timer's when the state is taken. A run restored from it goes on from there, without the time that passed after the end.

The tree is a plain dict of copies on the CPU (later updates of the nodes do not change it): save it with torch.save, load it with torch.load (weights_only=True reads it), and pass it to mujou.run(..., state=tree) to restore it::

{
    "timer": {"time_ns": int, "time_scale": float},
    "nodes": {
        "body.l1": {
            "state": ...,  # the node's state_dict()
            "last_update_ns": int,
            "ports": {
                "prev": {"in": E},
                "next": {"out": E, "to": {"body.l2.prev": E}},
            },
        },
        ...,
    },
}

A node is keyed by its path from the root ("" for a node that is the root itself), and each of its ports by its attribute. A port that sends has "out", its latest message (what read returns), and "to", its latest message to each peer (MultiPort.send_to writes only one), by the path of the peer's port. An unconnected Port that receives has "in", what the host wrote last. Each E is {"value": message, "time_ns": int | None}: time_ns is None while nothing was written. As every time of the tree, a Timestamp in a message (a TimestampSpec) is held as its nanoseconds, an int.

PARAMETER	DESCRIPTION
`timeout`	Real (wall-clock) seconds to wait for every node to pause (the end of its current `on_update`). `math.inf` waits as long as it takes. TYPE: `float` DEFAULT: `inf`

RAISES	DESCRIPTION
`RuntimeError`	This is called from a node's hook.
`TimeoutError`	A node did not pause within `timeout`; no state is taken and the run keeps running.
`NetworkError`	The run stopped because a node raised (its state is not kept).
`Exception`	After a normal end, what a node's `state_dict` raised; the run stays ended normally.

shutdown ¶

shutdown(*, timeout: float = math.inf) -> None

Stop every node and return once every on_stop has returned.

Does nothing after a normal end.

PARAMETER	DESCRIPTION
`timeout`	Real (wall-clock) seconds to wait for the end. `math.inf` waits as long as it takes. TYPE: `float` DEFAULT: `inf`

RAISES	DESCRIPTION
`TimeoutError`	The run did not stop within `timeout`. The stop stays asked: the nodes take no more updates and the run goes on stopping in the background (an `on_stop` may still be running). `join` waits for the end.
`ValueError`	`timeout` is negative or NaN; nothing changes.
`NetworkError`	The run stopped because a node raised.

join ¶

join(*, timeout: float = math.inf) -> None

Wait until the run stops, without asking it to stop.

It returns once another thread called shutdown and every on_stop has returned; if a node raises, it raises instead. If nothing stops the run and timeout is math.inf, it never returns.

PARAMETER	DESCRIPTION
`timeout`	Real (wall-clock) seconds to wait for the end. `math.inf` waits as long as it takes. TYPE: `float` DEFAULT: `inf`

RAISES	DESCRIPTION
`TimeoutError`	The run did not stop within `timeout`. Nothing changes: a running run keeps running.
`NetworkError`	The run stopped because a node raised.

mujou.NetworkError ¶

Bases: RuntimeError

The network stopped because a node raised.

__cause__ is the first exception a node raised. Once a run has stopped this way, the host's calls on it (write, read, state_dict, shutdown, ...) raise NetworkError too.