Skip to content

Experiment DAG and Storage

Coming Soon

Heyoooo so I ran out of time writing docs as I have to do actual machine learning. I'll eventually catch this up but as of right now here's my friend gpt-6-astra who will do the talking.


An experiment produces several kinds of things at different times: training metrics, checkpoints, evaluation scores, and perhaps a plot made days later. theseus gives them a common address: a node in an experiment's lineage. The node identifies the point being described. Parquet records describe that point, and blob files hold the larger objects associated with it.

The DAG is this recorded lineage. It is distinct from the execution recipe that tells the dispatcher which jobs to run. Dispatch can create branches and resume existing nodes, but the stored graph describes what the jobs recorded.

A node is an identity, not a row

A node's identity consists of three fields:

Field Meaning
name The job's project.group.run name, defaulting to general.default.<run> when project and group are omitted.
nonce A short generated identifier distinguishing lineages that use the same name.
seq The integer position on that lineage's logical clock, starting at zero.

Node.serialize() encodes those three fields into a URL-safe base64 string. It is an address, not a content hash. Repeating a run name usually creates a new nonce; two runs with identical settings are still separate runs.

A Node also carries an optional parent, expressed as another node's serialized identity. The parent is not included in the child's identity. Deserializing an address therefore recovers the name, nonce, and sequence, but does not fetch its parents. ObjectReader.parents(node) reads the recorded direct parent edges.

Several writes may share exactly the same identity. A loss report, an evaluation score, and a checkpoint marker can arrive in separate records and still describe one logical node. Likewise, several JAX processes may record that node. They do not create separate logical nodes merely by writing separate Parquet rows.

Ticking creates a successor

BasicJob.tick() records the current node if the job has not already written it, then advances the live node to its successor. The successor keeps the same name and nonce, increments seq, and records the old identity as its parent.

flowchart LR
    A["run A · seq 0"] --> B["run A · seq 1"]
    B --> C["run A · seq 2"]
    B --> D["run B · seq 0"]
    D --> E["run B · seq 1"]

The horizontal chains represent ticking. The edge from A's sequence 1 to B's sequence 0 represents a new branch from that point. The new branch can have a new run name, and gets a new nonce even if it reuses the old name.

Ticking mutates the existing live Node object through Node.update() instead of replacing it. Dataloaders, profilers, and borrowed analyses can hold a reference to that object and observe the same clock as the job.

A tick is not a filesystem flush, optimizer update, or checkpoint save. It only advances the logical position after ensuring the old position has a record queued. Job code decides which work belongs to each position. The standard trainer performs its update, reports metrics, runs any scheduled evaluation or analysis, and requests checkpointing before ticking.

The optimizer's state.step and the node's seq are consequently different concepts. They need not be numerically equal: the first sequence is zero, a branch starts a new sequence, and jobs other than trainers use the same node abstraction. Use recorded token counts or optimizer metadata when that is the quantity you intend to compare.

Why the dataloader shares the clock

The training loader uses the node sequence to choose a deterministic batch plan. Repeated batch() calls at the same node return the cached batch instead of consuming more examples. This lets a debugger or analysis revisit the inputs associated with the current training position. Advancing the node allows the loader to advance its batch.

Restore initially selects the saved node. An analysis can inspect that position without ticking. When the standard trainer actually resumes training at its base node, it ticks once before consuming the next training batch, then continues from the successor of the saved optimizer step.

Replay depends on keeping the relevant dataset, mixture, seed, and batching configuration compatible. The identity alone does not preserve the entire input pipeline. See Dataloader I/O for the batch-plan and replay invariants, and interactive inspection for using that behavior.

Branching, resuming, and attaching work later

A job constructed with a base node restores state from that node. Its setup mode determines the destination identity:

Mode Destination Typical use
No base New lineage at sequence zero A fresh experiment.
Branch from a base New lineage at sequence zero, with the base as parent Try a different training recipe or keep a separate analysis lineage.
Resume a base The selected base identity itself Continue training, or attach analysis to an existing checkpoint.

Setup is idempotent: after successful setup, another setup call preserves the prepared state and node. Normal job invocation synchronizes the node across hosts so their writes refer to the same identity.

Historical analysis uses resume semantics deliberately. It restores a checkpoint, replays the appropriate batch, and adds its artifact to that exact node without advancing the clock. The plot's later write time and its historical node position are both meaningful; attaching the plot does not pretend that the computation occurred during the original training run.

Each individual record has at most one parent field, but records for one node can contain different parent values. parents() returns their distinct union rather than applying scalar conflict resolution to this graph relationship. The normal tick/branch lifecycle creates forward lineage edges. The low-level store does not perform general cycle detection, validate arbitrary parent references, or merge model parameters when it sees multiple parents.

The physical layout

With the default cluster directories, the object store looks like this:

<root>/objects/
    values/
        part-<uuid>.parquet
        part-<uuid>.parquet
    blobs/
        <serialized-node>/
            nodespec.json
            config.yaml
            job.json
            rng.npy
            checkpoint/          # distributed checkpoint state
            attention.msgpack    # an optional analysis artifact

The blob contents depend on what was saved; not every node has a blob or a checkpoint. _x_blob stores a relative path under objects/blobs, which readers resolve against their object-store root. Analysis PDFs/JSON also have their normal result files under the configured results directory.

ObjectReader only reads. ObjectStore adds the scalar writer and blob-directory interface. RecordStore adds lightweight MessagePack payloads for artifacts. Model checkpoints use CheckpointedJob and the Orbax-backed checkpoint manager, not the MessagePack artifact path.

This separation also allows a reader to use a local mirror of Parquet parts while resolving blob paths against the original object store. Sharing or mirroring metadata does not implicitly copy checkpoint contents. See JuiceFS Integration for a shared-root deployment.

Publishing Parquet parts

A call to value(node, fields) copies scalar fields into a queue and adds the node identity, parent, process index, execution/tag information, and write order. Keys beginning with _x_ are reserved metadata. Large arrays and figures belong in checkpoints or artifacts rather than scalar columns.

Each store has a background writer. It batches up to 4,096 queued rows or flushes on a 30-second deadline. Closing the store drains the queue. Each batch's schema is the union of its keys, with nulls where an individual row lacks a field; later parts can introduce new metric columns.

The writer claims a unique hidden temporary filename, writes a complete Parquet table, then renames it to part-<uuid>.parquet. Readers query the final *.parquet files, so an unfinished temporary file is not queryable. A successful queue operation is therefore not an immediate durable write or a guarantee that a concurrent reader can already see it.

Closing also compacts parts owned by that store instance. It unifies their schemas and copies the rows into a replacement part before deleting the source parts. It does not rewrite another writer's parts or fold logical nodes during compaction. Interruption after publication can leave duplicate rows; deleting inputs only after publication avoids removing the only copy first.

This is file publication, not a database transaction covering every part and blob. A crash can lose queued records. Writer failures are surfaced on later writes or close; the current retry loop covers temporary-file claim permission errors with waits of 1, 2, 4, and 8 seconds. It does not make permanent permission failures harmless or retry every storage operation.

ClickHouse is the query engine

theseus calls embedded ClickHouse through chdb to query the Parquet files and return Arrow tables. This storage path does not require a separate ClickHouse server or insert rows into a server-managed table. The durable scalar data is in the Parquet directory.

Queries infer a union schema across parts and group by (_x_name, _x_nonce, _x_seq). For every requested field, an argMaxIf aggregate chooses the highest-priority non-null value. Nonconflicting fields accumulate; a later sparse record does not erase keys it omits.

For example, these physical records all target the same node:

Process Write order train/loss eval/accuracy
0 Earlier 2.4
1 Later 9.0 0.6
0 Latest 2.1

The logical result contains train/loss=2.1 and eval/accuracy=0.6. Process zero wins the loss conflict; process one's accuracy still contributes because process zero did not supply that field.

Priority is lexicographic: lower process index first, then the explicit write order when present, with file time, filename, and row position as fallbacks and tie-breakers. The writer makes its timestamp-derived order monotonic within the store instance. This is a defined conflict policy, not an average across replicas or a globally synchronized distributed clock. Distributed metric reduction must happen before logging if an average is what you need.

The wide-schema query uses COLUMNS(...) APPLY(...) to express the same aggregate for many fields without repeating its full SQL text for every field. It still names the selected columns and still folds them; a raw SELECT * would return physical fragments instead of logical nodes. The parser's max_query_size is set before parsing the generated SELECT, and expression-derived output names are mapped back to the public field names.

Filtering and selecting logical nodes

Identity and execution-scope filters restrict the physical input. Value predicates such as .where("train/loss", "<=", 2.5) are applied to the grouped result, so they use the same resolved field value that selection returns. Checkpoint and artifact filters test whether the group contains their respective markers.

.all() returns node identities. .select() returns folded values; .select(return_nodes=True) pairs each row with its identity. keys=[...] limits the returned fields, while keys=None returns all available user fields. Reserved _x_ fields require raw=True; blob is exposed as a resolved path. If a requested field is absent from the entire available schema, the current query returns no rows. A missing value within one node is omitted from that node's returned mapping.

latest() means most recently written among the matching groups, with sequence and identity tie-breakers. It does not simply mean the largest training sequence. Adding an artifact to an old checkpoint can make that node the most recently written one. Sort explicitly by sequence when sequence order is your intent. Query execution clears accumulated filters and predicates; use a fresh builder for a separate question.

Interpolation belongs to the view

Folding combines records with the same identity. It does not fill missing measurements from neighboring sequences, follow a parent to inherit metrics, or interpolate between checkpoints. A store query only returns recorded values.

The UI can interpolate when positioning a cursor or a checkpoint marker on a plotted series. Its interpolation is linear between neighboring points and clamps to endpoint values outside their range. For a custom x-axis, it can first map a checkpoint's sequence to an x-coordinate, then find a y-coordinate on the curve. Smoothing and plot downsampling are also presentation operations.

Those displayed coordinates are not new stored measurements. Selecting a node for restoration still selects an actual recorded identity and its actual blob; an interpolated point is not an interpolated model checkpoint.

Checkpoint and artifact completion are separate

A checkpoint blob contains model state, saved configuration, job metadata, and random-state information. Its scalar record carries _x_checkpoint. Orbax saves are asynchronous: the manager waits for a previous save before starting another, and closing waits for outstanding saves. Publishing a blob's metadata does not by itself prove every asynchronous checkpoint file is already complete.

Artifacts use a separate process-zero worker to serialize MessagePack, replace the payload file, and publish a blob/record marker. Reusing an artifact suffix at the same node replaces that payload. Scalar history is append-oriented, but blob files are not immutable versions of every previous write.

The useful boundary is the logical node: it ties state, observations, and lineage together while allowing their storage and computation to proceed independently. For the concrete APIs, see Analysis System, Evaluation System, and the training tutorial. The implementation lives in theseus/base/dag.py, theseus/job.py, theseus/store.py, and theseus/checkpoint.py.