Skip to content

job

job

There are two types of nodes and one type of edge.

Nodes: - TrainingNode (takes a checkpoint or nothing, and continues training it) - EvaluationNode (takes a checkpoint, and labels it with evaluation)

Edges: - Apply

Examples: For [T]raining and [E]valuation nodes.

Imagine I'm doing a "standard" ML training recipe. This may look like:

[T] pretrain -> [T] midtrain -> [T] RL | | | v v v [E] blimp [E] smolchat [E] gsm8k [E] hellaswag

In fact each of these is one training step:

[T] -> [T] -> [T] -> [T] | v [E]

I could also be doing some topology transfer. For instance, you may wonder if you can thoughtbubblesify a normal transformer. So this may look like:

[T] load (contrib.Qwen) -> [T] midtrain (forking.Thoughtbubbles) | | v v [E] gsm8k [E] gsm8k

The edges are actually a bit of a simplification, in the sense that each "checkpoint" is actually an edge, some of which has evaluation.

BasicJob(spec: ExecutionSpec, base: Optional[Node] = None)

Bases: _BaseJob, Generic[C]

tick() -> None

Record the current node and advance to its logical successor.

state_restore(state: ValueRow) -> None abstractmethod

state_init() -> None abstractmethod

Initialize the job as the beginning of a branch.

Notes

Specifically, calling this function should allocate hardware memory, and start initialization etc., and start the job from scratch. A node has been created

state_reset() -> None

Reset state hook. We'll call this before every non-resume call.

run() -> None abstractmethod

Run the job, assuming all hosts have setup

node_name(spec: JobSpec) -> str staticmethod

Return a unique name for the node, based on the spec

setup(resume: bool = False) -> None

Prepare state and the run node once, without running or synchronizing.

The first successful setup fixes the branch/resume mode. Later calls preserve the existing state and node, including interactive edits. Resuming selects the saved node; advancing belongs to the running job.

CheckpointedJob(spec: ExecutionSpec, base: Optional[Node] = None)

Bases: BasicJob[C], Generic[C]

template: PyTree[Any] abstractmethod property

"Template PyTree with which to do loading/saving; could be abstract shapes.

initialize() -> None abstractmethod

Initialize and install a complete state.

surgery(partial: PyTree[bool]) -> PyTree[Any] abstractmethod

Initialize marked leaves and return None for all other leaves.

apply(state: PyTree[Any], metadata: ValueRow) -> None abstractmethod

"Load state to be consistent with the above.

Notes

Could probably just be self.state=state + step accounting. But usually involves a bit more wedding planning.

reset() -> None

Reset state hook. For instance for midtraining.

save(state: PyTree[Any], metadata: dict[str, float | int | str]) -> None

"Save the state into the current node.

state_restore(state: ValueRow) -> None

Restore and apply the checkpoint referenced by a stored node view.

LoggingJob(spec: ExecutionSpec, base: Optional[Node] = None)

Bases: BasicJob[C], Generic[C]

Job that reads and records scalar logs and PyTree payloads.

This class deliberately does not implement checkpointing. PyTree records use the simple :class:theseus.store.RecordStore path: they are accepted only on one host and inefficiently serialized as MessagePack by JAX process 0. Stateful jobs should also inherit :class:CheckpointedJob, which owns distributed Orbax checkpoints, synchronization, randomness, configuration, and job metadata.

get(node: Node | str) -> dict[str, Any]

Return a node's folded values and decoded MessagePack records.

Parameters:

Name Type Description Default
node Node | str

A node or its serialized identity.

 required

Returns:

Type Description
dict[str, Any]

The object-store values. When the node's blob directory contains
dict[str, Any]

MessagePack records, they are returned as payload in filename order.
dict[str, Any]

The original blob path remains available as blob.

log(payload: Mapping[str, Any]) -> None

Record scalar values on the job's current node.

Parameters:

Name Type Description Default
payload Mapping[str, Any]

Named scalar values. NumPy and JAX scalars are converted to their host Python representations.

 required

Raises:

Type Description
RuntimeError

If the job has not started or the value writer failed.
ValueError

If a value is not scalar.
TypeError

If a scalar has an unsupported type.

artifact(suffix: str, description: Mapping[str, Any], payload: PyTree[Any]) -> None

Store a PyTree artifact on the job's current node.

The artifact is written only by JAX process 0. Use log() for scalar logs and artifact() for analysis data and other payloads. Model checkpoints use save().

Parameters:

Name Type Description Default
suffix str

Unique filename stem within the node. Reusing it replaces the prior artifact.

 required
description Mapping[str, Any]

Scalar metadata used to discover the artifact.

 required
payload PyTree[Any]

PyTree serialized as MessagePack by the record store.

 required

Raises:

Type Description
RuntimeError

If the job has not started or a writer failed.
ValueError

If suffix is invalid or description values are not scalar.
TypeError

If a description scalar has an unsupported type.

CometLoggingJob(spec: ExecutionSpec, base: Optional[Node] = None)

Bases: LoggingJob[C], Generic[C]

Mirror scalar logs to Comet when logging/remote is enabled.

Include config() in the concrete job's schemas and call super().run() before its workload. One experiment follows a nonce across resumes; sequence numbers are metric steps. PyTree records remain in RecordStore.

run() -> None

Open remote logging on process 0 before the concrete workload.

log(payload: Mapping[str, Any]) -> None

Keep local scalar records and mirror them at the current sequence.

finish() -> None

Flush Comet once and always close the inherited job resources.

RestoreableJob(spec: ExecutionSpec, base: Optional[Node] = None)

Bases: CheckpointedJob[C], Generic[C]

Checkpointed job reconstructable from a stored node.

from_node(node: Node, spec: ExecutionSpec, runtime_cfg: Any | None = None, resume: bool = False) -> tuple[Self, Any] classmethod

Return the concrete job and configuration recorded for a node.

save(state: PyTree[Any], metadata: dict[str, float | int | str]) -> None

Save a checkpoint labeled with its registered job type.