dew.artifacts

What an objective’s evaluation produces, as typed values.

Objectives return these values from scoring and preview hooks. A metric consumes scoring artifacts; a tracker renders previews by type. Array leaves cross jit, while optional captions and decoded text remain host metadata.

Name	Summary
T
ImageGrid	Images in [-1, 1], [N, H, W, C], with the text each was conditioned on where there was any.
VideoGrid	Clips in [-1, 1], [N, T, H, W, C].
TextSamples	Generated token rows, with optional decoded preview text and prompt.
Representations	Encoder outputs [N, D] and the labels of the records they came from, for a probe to score.
TokenScores	Teacher-forced per-token losses [N, L] and the weight of each target, 1 where it counts and 0 where it is padding or a document’s first token.
Artifact
Artifacts
uint8_pixels	[-1, 1] pixels, as ImageGrid and VideoGrid hold them, as uint8 in [0, 255].
host	An artifact whose arrays are host-local numpy.
GATHER_BYTES	The bytes of global leaves collective_host gathers in one computation.
collective_host	Materialize an evaluation tree on every rank with transfer consensus.
broadcast_from_process_zero	A JSON-encodable value broadcast from rank zero to every rank.
agreed	Run operation on every rank, then agree on the outcome before going on.
PeerFailure	Another rank failed at a phase agreement this rank passed.
AGREEMENT_PATIENCE_SECONDS	How long ranks that reached an agreement wait on the host for the rest.
agree_process_phase	Propagate host errors, then count ranks declaring availability.
FAILURE_DIRECTORY
FAILURE_KEY	The coordination-service key a failing process writes its error under.
FAILURE_GRACE_SECONDS	How long a published failure may go unheard before every process ends.
this_process	This process as host:pid, named without the backend, which may be the thing that failed or be opening in another thread.
publish_failure	Write this process’s failure where every process’s watch sees it.
withdraw_failure	Remove the published failure once the pool has heard it at an agreement.
end_pool_on_failure	End this process when a failure goes unheard, or when it fails itself.

T

attribute source

T = TypeVar('T')

ImageGrid

dataclass source

class ImageGrid(captions: tuple[str, ...] = struct.field(pytree_node=False, default=()))

Images in [-1, 1], [N, H, W, C], with the text each was conditioned on where there was any.

VideoGrid

dataclass source

class VideoGrid(captions: tuple[str, ...] = struct.field(pytree_node=False, default=()))

Clips in [-1, 1], [N, T, H, W, C].

TextSamples

dataclass source

class TextSamples(
    prompt: str = struct.field(pytree_node=False, default=''),
    texts: tuple[str, ...] = struct.field(pytree_node=False, default=()),
)

Generated token rows, with optional decoded preview text and prompt.

Representations

dataclass source

class Representations()

Encoder outputs [N, D] and the labels of the records they came from, for a probe to score.

TokenScores

dataclass source

class TokenScores()

Teacher-forced per-token losses [N, L] and the weight of each target, 1 where it counts and 0 where it is padding or a document’s first token. A perplexity is exp of the weighted mean over a whole pass, so a batch with no counted target weighs nothing.

Artifact

attribute source

Artifact = ImageGrid | VideoGrid | TextSamples | Representations | TokenScores

Artifacts

attribute source

Artifacts = Artifact | tuple[Artifact, ...]

uint8_pixels

function source

def uint8_pixels(images: ArrayLike) -> NDArray[np.uint8]

[-1, 1] pixels, as ImageGrid and VideoGrid hold them, as uint8 in [0, 255].

Each value goes to its nearest level, (x + 1) * 127.5 in float32 rounded half to even (np.rint), then clipped, because a sample can leave the range. A metric scores and a tracker previews these bytes, so the two see the same image.

host

function source

def host(value: T) -> T

An artifact whose arrays are host-local numpy.

Scoring and drawing happen on the host: a metric reads the arrays with numpy, a tracker draws them. On one process that is a device transfer. On a pool the arrays are shards of a global array, which numpy cannot read at all, and the gather that completes them is a collective, so every process has to make the same call. That is why the trainer brings an artifact home once for the whole pool before any metric or tracker, which run on one process, sees it.

GATHER_BYTES

attribute source

GATHER_BYTES = 256 * 2 ** 20

The bytes of global leaves collective_host gathers in one computation.

A group’s leaves are replicated by one computation and agreed on once, so a tree of many small leaves pays one round trip a group rather than one a leaf, and a device holds at most one group’s values, or its one larger leaf, beside the tree’s shards. With held_by="first" a rank that holds nothing still holds one replicated group on its devices while that group’s computation runs, as process_allgather replicated each leaf.

collective_host

function source

def collective_host(
    value: T,
    *,
    phase: str,
    held_by: Literal['every', 'first'] = 'every',
) -> T | None

Materialize an evaluation tree on every rank with transfer consensus.

All ranks must call this, even for entirely local trees. Every array leaf is waited on before any data gather, so a computation that failed on a rank reports at the preflight rather than inside a gather collective; a wait raises what a host copy would, and moves nothing. Ranks then agree the ordered global gather plan and each transfer outcome. Local-only trees may differ, as with root-only decoded previews. A device failure inside an in-flight collective still needs runtime termination.

Global leaves are gathered in groups (GATHER_BYTES), one computation and one agreement a group, and a rank that fails in a group reports at that group’s agreement, before any rank starts the next. With held_by “first”, only process 0 of the pool, jax.process_index() == 0, copies the gathered tree to its host and returns it; the others take part in every computation and agreement and return None.

broadcast_from_process_zero

function source

def broadcast_from_process_zero(value)

A JSON-encodable value broadcast from rank zero to every rank.

agreed

function source

def agreed(phase: str, operation: Callable[[], T]) -> T

Run operation on every rank, then agree on the outcome before going on.

A rank that fails reports its error at the agreement point instead of raising alone, so its peers hear about it there rather than hanging at the next collective. The peers raise PeerFailure; the failing rank re-raises its own error. On one process this is a plain call.

PeerFailure

class source

class PeerFailure(RuntimeError)

Another rank failed at a phase agreement this rank passed.

AGREEMENT_PATIENCE_SECONDS

attribute source

AGREEMENT_PATIENCE_SECONDS = 24 * 3600

How long ranks that reached an agreement wait on the host for the rest.

Host work before an agreement, such as process 0 uploading the final checkpoint, can take hours; the pool bounds device executions, not this.

agree_process_phase

function source

def agree_process_phase(
    error: BaseException | None,
    *,
    phase: str,
    available: bool = True,
) -> int

Propagate host errors, then count ranks declaring availability.

All live ranks must reach this boundary. It cannot rescue a blocked device collective: a rank that failed between steps meets peers still inside the next step’s collectives, which wait for it for ever on a GPU. So a failing rank first publishes its error (publish_failure), which ends the pool within FAILURE_GRACE_SECONDS unless the agreement completes and withdraws it. Errors take priority over unavailable input.

The ranks meet on the host, at a coordination-service barrier, before the device collectives that carry the outcome. A rank still busy on its host keeps the others waiting there rather than inside an execution, which the pool’s bound (dew.training.runtime.EXECUTION_TIMEOUT) would end.

FAILURE_DIRECTORY

attribute source

FAILURE_DIRECTORY = 'dew/failure/'

FAILURE_KEY

attribute source

FAILURE_KEY = FAILURE_DIRECTORY + 'published'

The coordination-service key a failing process writes its error under.

FAILURE_GRACE_SECONDS

attribute source

FAILURE_GRACE_SECONDS = 60.0

How long a published failure may go unheard before every process ends.

this_process

function source

def this_process() -> str

This process as host:pid, named without the backend, which may be the thing that failed or be opening in another thread.

publish_failure

function source

def publish_failure(error: BaseException, where: str) -> bool

Write this process’s failure where every process’s watch sees it.

Returns False when another failure is already there: the first stands.

withdraw_failure

function source

def withdraw_failure() -> None

Remove the published failure once the pool has heard it at an agreement.

end_pool_on_failure

function source

def end_pool_on_failure(grace: float = FAILURE_GRACE_SECONDS) -> None

End this process when a failure goes unheard, or when it fails itself.

A process of a pool that raises past its program, or meets a peer’s failure it cannot hear, would otherwise hang: its peers wait for it in a collective no GPU backend times out, and jax.distributed’s shutdown barrier holds an exiting process up to shutdown_timeout_seconds (300 s) for them. So an uncaught exception prints, publishes and leaves at once, and a watch thread ends the process grace seconds after any published failure that no agreement withdrew. dew launch, srun and a pod’s scheduler then see the failure and stop the rest.

It needs only the pool’s coordination service, not the backend, so it goes in before the backend opens: a process whose devices fail to open after the pool has formed would otherwise wait in that barrier for peers that wait for its devices.