store

`store` ¶

Append-only object storage for Theseus experiment nodes.

Scalar metadata is written asynchronously into ClickHouse, while blobs are exposed as directories whose relative locations are recorded in the same metadata stream. RecordStore adds deliberately simple, process-0-only MessagePack artifacts; it is not a checkpointing implementation. Stateful jobs should use CheckpointedJob for distributed state, configuration, and job metadata.

`SerializedQuery` ¶

Bases: TypedDict

Serializable state produced by a query builder.

`ObjectReader(cluster: Cluster)` ¶

Query ClickHouse metadata and resolve blobs against an object root.

Cluster.metadata is a persistent chDB directory/URI or an HTTP(S) ClickHouse DSN. Omission selects objects/values/v2 under the cluster root. ROOT/.id identifies the metadata table and must already exist when opening a reader. Embedded connections hold an exclusive filesystem lock until close; use a server for concurrent processes.

`local(root_dir: str | Path, metadata: str | None = None) -> Self` `classmethod` ¶

Read a local object root using the selected metadata connection.

`start() -> None` ¶

Open a connection, failing immediately if an embedded path is owned.

`close() -> None` ¶

Close the connection and release embedded database ownership.

`query(serialized: SerializedQuery | None = None) -> QueryBuilder` ¶

Return a fresh query builder for this reader.

`root() -> Path` ¶

Return the object-store root used to resolve blob paths.

`parents(node: Node) -> list[Node]` ¶

Return the distinct direct parents recorded for an exact node.

`ObjectStore(hardware: HardwareResult, execution_id: str | None = None, tag: str | None = None)` ¶

Bases: ObjectReader

Store blobs and scalar metadata associated with experiment DAG nodes.

Constructing a store immediately starts its asynchronous ClickHouse writer. Values become queryable after a size- or time-based flush, or after :meth:close drains the writer. Concurrent processes require a ClickHouse server; query conflict resolution gives lower process indices precedence. A root's stable .id is created atomically on its first store construction.

Parameters:

Name	Type	Description	Default
`hardware`	`HardwareResult`	Hardware allocation whose current process host determines the object-store root.	required

`local(root_dir: str | Path, metadata: str | None = None) -> Self` `classmethod` ¶

Open a local store using a persistent chDB path or ClickHouse DSN.

`close() -> None` ¶

Drain queued values and release the connection, including on failure.

`blob(node: Node, metadata: Mapping[str, float | int | str] | None = None) -> Iterator[Path]` ¶

Open the directory used to write a node's blob.

On normal context exit, process 0 creates nodespec.json if absent and every process records the blob's relative path with the supplied metadata. If the context exits with an exception, the partial directory is left in place but no metadata row is recorded.

Parameters:

Name	Type	Description	Default
`node`	`Node`	Node that owns the blob.	required
`metadata`	`Mapping[str, float \| int \| str] \| None`	Optional scalar metadata to publish with the blob location.	`None`

Yields:

Type	Description
`Path`	The directory into which the caller should write blob contents.

Raises:

Type	Description
`RuntimeError`	If the asynchronous value writer has previously failed.
`TypeError`	If metadata contains a non-scalar value.

`value(node: Node, kv: Mapping[str, float | int | str | None]) -> None` ¶

Queue a scalar metadata chunk associated with a node.

The input mapping is copied before internal node, process, parent, and write metadata is added. Multiple chunks may target the same node; query selection folds their non-conflicting keys together. Keys beginning with _x_ are reserved for the store and are not returned by queries.

Parameters:

Name	Type	Description	Default
`node`	`Node`	Exact node identity to associate with the values.	required
`kv`	`Mapping[str, float \| int \| str \| None]`	Scalar keys and values to enqueue.	required

Raises:

Type	Description
`RuntimeError`	If the asynchronous writer has failed.
`TypeError`	If `kv` contains a non-scalar value.

`start() -> None` ¶

Open the metadata connection and start the batched writer.

`RecordStore(hardware: HardwareResult, execution_id: str | None = None, tag: str | None = None)` ¶

Bases: ObjectStore

Write and decode lightweight PyTree artifacts attached to nodes.

Records are accepted only on one host (JAX process 0) and inefficiently serialized as MessagePack there. A bounded worker queue keeps that work off the caller's path, but this intentionally simple path is not suitable for checkpoints. Stateful jobs should use :class:theseus.job.CheckpointedJob, which owns distributed Orbax saves along with synchronization, randomness, configuration, and job metadata.

Parameters:

Name	Type	Description	Default
`hardware`	`HardwareResult`	Hardware allocation whose process-0 host writes record payloads.	required

`close() -> None` ¶

Flush records, then flush and close the underlying object store.

Raises:

Type	Description
`RuntimeError`	If the record or scalar writer failed.
`OSError`	If record publication fails.

`query(serialized: SerializedQuery | None = None) -> _RecordQueryBuilder` ¶

Return a fresh query builder that decodes MessagePack records.

Parameters:

Name	Type	Description	Default
`serialized`	`SerializedQuery \| None`	Optional state previously produced for an execution base.	`None`

Returns:

Type	Description
`_RecordQueryBuilder`	A record-aware query builder for this store.

`get_record(node: Node | str) -> dict[str, Any]` ¶

Return an exact node with any MessagePack records decoded.

Parameters:

Name	Type	Description	Default
`node`	`Node \| str`	A node or its serialized identity.	required

Returns:

Type	Description
`dict[str, Any]`	Folded metadata with decoded MessagePack records added as `payload`.
`dict[str, Any]`	The original blob path remains available as `blob`.

`artifact(node: Node, suffix: str, description: Mapping[str, float | int | str], payload: PyTree[Any]) -> None` ¶

Queue a PyTree artifact owned by an exact node.

Only one host writes artifacts: calls from JAX processes other than process 0 are no-ops. Reusing a suffix for the same node replaces that artifact.

Parameters:

Name	Type	Description	Default
`node`	`Node`	Exact node that owns the artifact.	required
`suffix`	`str`	Unique filename stem within the node's blob directory.	required
`description`	`Mapping[str, float \| int \| str]`	Scalar metadata stored with the artifact marker.	required
`payload`	`PyTree[Any]`	PyTree serialized as MessagePack on process 0.	required

Raises:

Type	Description
`TypeError`	If description metadata is not scalar.
`ValueError`	If `suffix` is not a plain, non-empty filename.
`RuntimeError`	If either asynchronous writer has failed.

`QueryBuilder(store: ObjectReader, serialized: SerializedQuery | None = None)` ¶

Build queries over the nodes recorded by an :class:ObjectStore.

Query methods are cumulative and return the builder, so filters can be chained before calling :meth:all or :meth:select. User values written in separate chunks are treated as values of the same node. If a key was written more than once, the query uses lower-process and last-write precedence.

Reserved _x_* fields cannot be passed to :meth:where, :meth:has, or :meth:sort. Use :meth:name, :meth:nonce, :meth:seq, :meth:job, :meth:node, :meth:checkpoint, :meth:artifact, or :meth:spec to query supported node metadata.

Example

Start with store.query(), chain filters such as .spec(group="pretrain").checkpoint(), then call .where("eval/score", ">", 0.8).all() to return matching nodes.

Parameters:

Name	Type	Description	Default
`store`	`ObjectReader`	Object store or reader containing the metadata table to query.	required
`serialized`	`SerializedQuery \| None`	Optional previously serialized query state.	`None`

`node(node: Node | str) -> Self` ¶

Restrict the query to a serialized node identity.

Parameters:

Name	Type	Description	Default
`node`	`Node \| str`	Exact node or a value produced by :meth:`theseus.base.Node.serialize`.	required

Returns:

Type	Description
`Self`	This builder, for continued chaining.

`name(name: str) -> Self` ¶

Restrict the query to an exact node name.

Parameters:

Name	Type	Description	Default
`name`	`str`	Fully qualified node name.	required

Returns:

Type	Description
`Self`	This builder, for continued chaining.

`nonce(nonce: str) -> Self` ¶

Restrict the query to an exact node nonce.

Parameters:

Name	Type	Description	Default
`nonce`	`str`	Branch nonce shared by nodes in the same lineage.	required

Returns:

Type	Description
`Self`	This builder, for continued chaining.

`seq(seq: int) -> Self` ¶

Restrict the query to an exact node sequence number.

Parameters:

Name	Type	Description	Default
`seq`	`int`	Sequence number within a node lineage.	required

Returns:

Type	Description
`Self`	This builder, for continued chaining.

`execution(execution_id: str) -> Self` ¶

Restrict the query to one execution.

`tag(tag: str) -> Self` ¶

Restrict the query to an exact execution tag.

`finished(is_finished: bool = True) -> Self` ¶

Restrict the query by completion of its execution/tag scope.

`latest() -> Self` ¶

Return only the most recently written matching node.

`job(job: str) -> Self` ¶

Restrict the query to nodes checkpointed by a registered job.

Parameters:

Name	Type	Description	Default
`job`	`str`	Exact registry key recorded in `_x_job` checkpoint metadata.	required

Returns:

Type	Description
`Self`	This builder, for continued chaining.

`checkpoint() -> Self` ¶

Restrict the query to nodes recorded as checkpoints.

Returns:

Type	Description
`Self`	This builder, for continued chaining.

`artifact() -> Self` ¶

Restrict the query to nodes containing a recorded PyTree artifact.

Artifact and checkpoint status are independent, so a node may match both :meth:artifact and :meth:checkpoint.

Returns:

Type	Description
`Self`	This builder, for continued chaining.

`spec(project: str | None = None, group: str | None = None, run: str | None = None) -> Self` ¶

Restrict the query by job-spec name components.

The node-name convention is delegated to :meth:theseus.job.BasicJob.node_name. A None component is left unrestricted, so any subset of project, group, and run may be supplied.

Parameters:

Name	Type	Description	Default
`project`	`str \| None`	Project component to match, or `None` for any project.	`None`
`group`	`str \| None`	Group component to match, or `None` for any group.	`None`
`run`	`str \| None`	Run-name component to match, or `None` for any run.	`None`

Returns:

Type	Description
`Self`	This builder, for continued chaining.

`where(key: str, operator: str, value: float | int | str) -> Self` ¶

Restrict the query using a comparison against a user value.

Parameters:

Name	Type	Description	Default
`key`	`str`	User metadata key to compare.	required
`operator`	`str`	One of `=`, `!=`, `<`, `<=`, `>`, or `>=`.	required
`value`	`float \| int \| str`	Scalar value to compare against the folded value of `key`.	required

Returns:

Type	Description
`Self`	This builder, for continued chaining.

Raises:

Type	Description
`ValueError`	If `key` is reserved or `operator` is unsupported.
`TypeError`	If `value` is not a supported scalar query value.

`has(key: str) -> Self` ¶

Restrict the query to nodes containing a user metadata key.

Parameters:

Name	Type	Description	Default
`key`	`str`	User metadata key whose folded value must be non-null.	required

Returns:

Type	Description
`Self`	This builder, for continued chaining.

Raises:

Type	Description
`ValueError`	If `key` is a reserved `_x_*` field.

`sort(key: str, ascending: bool = True) -> Self` ¶

Order matching nodes by a folded user metadata value.

Missing values are placed last. Name, nonce, and sequence are used as deterministic tie-breakers. Calling this method again replaces the prior sort key.

Parameters:

Name	Type	Description	Default
`key`	`str`	User metadata key whose folded value determines the ordering.	required
`ascending`	`bool`	Whether to place smaller values first.	`True`

Returns:

Type	Description
`Self`	This builder, for continued chaining.

`all() -> list[Node]` ¶

Return all matching node identities and clear the query.

Returns:

Type	Description
`list[Node]`	Matching nodes in the requested order. Returns an empty list when the
`list[Node]`	store has no finalized parts or a requested key has never been
`list[Node]`	recorded.

`select(return_nodes: bool = False, raw: bool = False, keys: Sequence[str] | None = None) -> list[ValueRow] | list[tuple[Node, ValueRow]]` ¶

select(return_nodes: Literal[False] = False, raw: bool = False, keys: Sequence[str] | None = None) -> list[ValueRow]

select(return_nodes: Literal[True], raw: bool = False, keys: Sequence[str] | None = None) -> list[tuple[Node, ValueRow]]

Return matching folded values and clear the accumulated query.

Values use lower-process and last-write precedence. Reserved storage metadata is omitted unless raw is true, while a recorded blob is exposed as blob and resolved against the current object-store root.

Parameters:

Name	Type	Description	Default
`return_nodes`	`bool`	Pair every value row with its exact node identity.	`False`
`raw`	`bool`	Include selected reserved `_x_*` storage metadata.	`False`
`keys`	`Sequence[str] \| None`	Optional columns to project instead of reading every value.	`None`

Returns:

Type	Description
`list[ValueRow] \| list[tuple[Node, ValueRow]]`	Matching values in the requested order, optionally paired with their
`list[ValueRow] \| list[tuple[Node, ValueRow]]`	nodes. Stores without matching finalized rows return an empty list.

store

store ¶

SerializedQuery ¶

ObjectReader(cluster: Cluster) ¶

local(root_dir: str | Path, metadata: str | None = None) -> Self classmethod ¶

start() -> None ¶

close() -> None ¶

query(serialized: SerializedQuery | None = None) -> QueryBuilder ¶

root() -> Path ¶

parents(node: Node) -> list[Node] ¶

ObjectStore(hardware: HardwareResult, execution_id: str | None = None, tag: str | None = None) ¶

local(root_dir: str | Path, metadata: str | None = None) -> Self classmethod ¶

close() -> None ¶

blob(node: Node, metadata: Mapping[str, float | int | str] | None = None) -> Iterator[Path] ¶

value(node: Node, kv: Mapping[str, float | int | str | None]) -> None ¶

start() -> None ¶

RecordStore(hardware: HardwareResult, execution_id: str | None = None, tag: str | None = None) ¶

close() -> None ¶

query(serialized: SerializedQuery | None = None) -> _RecordQueryBuilder ¶

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

artifact(node: Node, suffix: str, description: Mapping[str, float | int | str], payload: PyTree[Any]) -> None ¶

QueryBuilder(store: ObjectReader, serialized: SerializedQuery | None = None) ¶

node(node: Node | str) -> Self ¶

name(name: str) -> Self ¶

nonce(nonce: str) -> Self ¶

seq(seq: int) -> Self ¶

execution(execution_id: str) -> Self ¶

tag(tag: str) -> Self ¶

finished(is_finished: bool = True) -> Self ¶

latest() -> Self ¶

job(job: str) -> Self ¶

checkpoint() -> Self ¶

artifact() -> Self ¶

spec(project: str | None = None, group: str | None = None, run: str | None = None) -> Self ¶

where(key: str, operator: str, value: float | int | str) -> Self ¶

has(key: str) -> Self ¶

sort(key: str, ascending: bool = True) -> Self ¶

all() -> list[Node] ¶

select(return_nodes: bool = False, raw: bool = False, keys: Sequence[str] | None = None) -> list[ValueRow] | list[tuple[Node, ValueRow]] ¶

`store` ¶

`SerializedQuery` ¶

`ObjectReader(cluster: Cluster)` ¶

`local(root_dir: str | Path, metadata: str | None = None) -> Self` `classmethod` ¶

`start() -> None` ¶

`close() -> None` ¶

`query(serialized: SerializedQuery | None = None) -> QueryBuilder` ¶

`root() -> Path` ¶

`parents(node: Node) -> list[Node]` ¶

`ObjectStore(hardware: HardwareResult, execution_id: str | None = None, tag: str | None = None)` ¶

`local(root_dir: str | Path, metadata: str | None = None) -> Self` `classmethod` ¶

`close() -> None` ¶

`blob(node: Node, metadata: Mapping[str, float | int | str] | None = None) -> Iterator[Path]` ¶

`value(node: Node, kv: Mapping[str, float | int | str | None]) -> None` ¶

`start() -> None` ¶

`RecordStore(hardware: HardwareResult, execution_id: str | None = None, tag: str | None = None)` ¶

`close() -> None` ¶

`query(serialized: SerializedQuery | None = None) -> _RecordQueryBuilder` ¶

`get_record(node: Node | str) -> dict[str, Any]` ¶

`artifact(node: Node, suffix: str, description: Mapping[str, float | int | str], payload: PyTree[Any]) -> None` ¶

`QueryBuilder(store: ObjectReader, serialized: SerializedQuery | None = None)` ¶

`node(node: Node | str) -> Self` ¶

`name(name: str) -> Self` ¶

`nonce(nonce: str) -> Self` ¶

`seq(seq: int) -> Self` ¶

`execution(execution_id: str) -> Self` ¶

`tag(tag: str) -> Self` ¶

`finished(is_finished: bool = True) -> Self` ¶

`latest() -> Self` ¶

`job(job: str) -> Self` ¶

`checkpoint() -> Self` ¶

`artifact() -> Self` ¶

`spec(project: str | None = None, group: str | None = None, run: str | None = None) -> Self` ¶

`where(key: str, operator: str, value: float | int | str) -> Self` ¶

`has(key: str) -> Self` ¶

`sort(key: str, ascending: bool = True) -> Self` ¶

`all() -> list[Node]` ¶

`select(return_nodes: bool = False, raw: bool = False, keys: Sequence[str] | None = None) -> list[ValueRow] | list[tuple[Node, ValueRow]]` ¶