vertebrae.utils.memory

Memory estimation and admission helpers.

Classes

MemoryBudget

Resolved memory budget from a MemoryConfig.

EmbeddingMemoryEstimate

Estimated memory footprint for an embedding artifact.

MatrixAssembly

A matrix assembled either in memory or through a temporary local memmap.

MatrixRowReference

Opaque reference to one incrementally staged matrix row.

MetadataRowReference

Opaque reference to one incrementally staged metadata record.

IncrementalMatrixStager

Stage compatible matrix rows without retaining extractor batch buffers.

IncrementalMetadataStager

Stage per-row metadata without retaining an unbounded candidate list.

IncrementalMatrixReferenceStager

Store ordered matrix-row references without an unbounded Python list.

Functions

resolve_memory_budget(config)

Resolve an effective memory budget using psutil.

estimate_embedding_from_probe(probe_embeddings, ...[, ...])

Estimate full embedding memory from a probe batch.

estimate_matrix_resident_bytes(matrix)

Estimate memory required to hold a dense or sparse matrix.

estimate_object_resident_bytes(value[, _seen])

Conservatively estimate the retained size of nested metadata.

estimate_final_row_metadata_bytes(rows, *, ...)

Estimate final row-aligned metadata before allocating retained containers.

admit_final_metadata(stager, required_bytes, *, purpose)

Admit retained row metadata and the current output's transient peak.

estimate_metadata_resident_bytes(metadata)

Estimate resident bytes from embedding metadata.

estimate_metadata_dense_scoring_bytes(metadata)

Estimate the hypothetical dense footprint from embedding metadata.

estimate_metadata_scoring_input_bytes(metadata, *[, ...])

Estimate the actual dense or sparse representation consumed by scoring.

assert_within_memory(required_bytes, memory_config, ...)

Fail fast when a planned allocation exceeds the memory budget.

assemble_matrix_rows(rows, memory_config, *, purpose)

Stack compatible matrix rows without forcing an over-budget dense allocation.

largest_fitting_subsample_rate(required_bytes, ...)

Estimate the largest sample fraction that fits the memory budget.

sparse_matrix_nbytes(matrix)

Estimate resident bytes for an existing scipy sparse matrix.

sparse_nbytes_from_nnz(nnz, n_rows, dtype)

Estimate CSR sparse bytes from non-zero count and row count.

estimate_sparse_bytes(n_samples, n_features, dtype, ...)

Estimate sparse matrix bytes from shape, dtype, and density.

Module Contents

class vertebrae.utils.memory.MemoryBudget[source]

Resolved memory budget from a MemoryConfig.

total_bytes[source]

Total system memory.

available_bytes[source]

Currently available system memory.

reserve_system_bytes[source]

Bytes reserved for system use.

max_memory_bytes[source]

Maximum bytes vertebrae should use for planned work.

class vertebrae.utils.memory.EmbeddingMemoryEstimate[source]

Estimated memory footprint for an embedding artifact.

n_samples[source]

Number of embedding rows.

embedding_dim[source]

Number of embedding columns.

dtype[source]

Embedding dtype string.

resident_bytes[source]

Estimated bytes to hold the embedding artifact in memory.

scoring_input_bytes[source]

Estimated bytes for the representation consumed by scoring.

dense_scoring_bytes[source]

Hypothetical dense footprint of the scoring matrix.

batch_embedding_bytes[source]

Estimated bytes for one embedding batch.

strategy[source]

Planned strategy: “in_memory” or “stream_to_disk”.

to_dict()[source]

Serialize the estimate for result metadata.

Returns:

JSON-compatible memory estimate.

Return type:

dict[str, Any]

class vertebrae.utils.memory.MatrixAssembly[source]

A matrix assembled either in memory or through a temporary local memmap.

class vertebrae.utils.memory.MatrixRowReference[source]

Opaque reference to one incrementally staged matrix row.

class vertebrae.utils.memory.MetadataRowReference[source]

Opaque reference to one incrementally staged metadata record.

class vertebrae.utils.memory.IncrementalMatrixStager(memory_config, *, purpose)[source]

Stage compatible matrix rows without retaining extractor batch buffers.

When disk spill is enabled, every row is written immediately to an append-only per-output staging file. This keeps memory bounded even when a streaming extractor emits several named outputs. When spill is disabled, rows remain in memory only while their aggregate footprint fits the configured budget; otherwise appending fails before another row is retained.

Parameters:
append(output_name, row)[source]

Append one dense or sparse matrix row and return its lightweight reference.

Parameters:
  • output_name (str)

  • row (Any)

Return type:

MatrixRowReference

assemble(output_name, references, *, purpose, force_disk=False)[source]

Assemble selected rows for one output and release that output’s staging data.

Parameters:
  • output_name (str)

  • references (Iterable[MatrixRowReference])

  • purpose (str)

  • force_disk (bool)

Return type:

MatrixAssembly

reserve_metadata(required_bytes, *, purpose)[source]

Reserve retained metadata bytes against the shared no-spill budget.

Parameters:
  • required_bytes (int)

  • purpose (str)

Return type:

None

release_metadata(released_bytes)[source]

Release metadata bytes previously reserved by a paired stager.

Parameters:

released_bytes (int)

Return type:

None

discard_output(output_name)[source]

Release all staged candidate rows for one named output.

Parameters:

output_name (str)

Return type:

None

close()[source]

Remove all temporary candidate staging files and release row references.

Return type:

None

class vertebrae.utils.memory.IncrementalMetadataStager(memory_config, *, purpose, matrix_stager)[source]

Stage per-row metadata without retaining an unbounded candidate list.

Disk-spill mode writes each record directly to an append-only pickle stream and stores only ordering and selection indexes in SQLite. Iteration loads one record at a time. No-spill mode retains records only after reserving their estimated resident size against the paired matrix stager’s memory budget.

Parameters:
property resident_bytes: int[source]

Estimated retained metadata bytes in no-spill mode.

Return type:

int

append(output_name, row, *, priority_key='', group_key='', row_key=0, column_key=0)[source]

Append one metadata record with optional deterministic ordering keys.

Parameters:
  • output_name (str)

  • row (Dict[str, Any])

  • priority_key (str)

  • group_key (str)

  • row_key (int)

  • column_key (int)

Return type:

MetadataRowReference

iter_rows(output_name, *, order='ordinal', selected_only=False)[source]

Iterate records one at a time in insertion, priority, or final order.

Parameters:
  • output_name (str)

  • order (str)

  • selected_only (bool)

Return type:

Iterator[Tuple[MetadataRowReference, Dict[str, Any]]]

count_rows(output_name, *, selected_only=False)[source]

Return the number of staged records for an output.

Parameters:
  • output_name (str)

  • selected_only (bool)

Return type:

int

mark_selected(reference, output_name)[source]

Mark one staged record as retained by deterministic filtering.

Parameters:
Return type:

None

counter_value(output_name, namespace, key)[source]

Read a bounded or disk-backed selection counter.

Parameters:
  • output_name (str)

  • namespace (str)

  • key (str)

Return type:

int

increment_counter(output_name, namespace, key, amount=1)[source]

Increment a bounded or disk-backed selection counter.

Parameters:
  • output_name (str)

  • namespace (str)

  • key (str)

  • amount (int)

Return type:

int

add_member(output_name, namespace, group_key, member_key)[source]

Add one exact membership value and report whether it was new.

Parameters:
  • output_name (str)

  • namespace (str)

  • group_key (str)

  • member_key (str)

Return type:

bool

has_member(output_name, namespace, group_key, member_key)[source]

Return whether an exact membership value has been observed.

Parameters:
  • output_name (str)

  • namespace (str)

  • group_key (str)

  • member_key (str)

Return type:

bool

member_count(output_name, namespace, group_key='')[source]

Count exact members in one state namespace and group.

Parameters:
  • output_name (str)

  • namespace (str)

  • group_key (str)

Return type:

int

discard_output(output_name)[source]

Release all records and selection state for one output.

Parameters:

output_name (str)

Return type:

None

close()[source]

Release retained records and remove all temporary spill files.

Return type:

None

class vertebrae.utils.memory.IncrementalMatrixReferenceStager(memory_config, *, purpose, matrix_stager)[source]

Store ordered matrix-row references without an unbounded Python list.

Spill-enabled runs keep reference payloads and ordering indexes in the metadata stager’s disk-backed SQLite/pickle storage. No-spill runs charge reference records, exact-position membership, and deterministic sorting to the paired matrix stager’s memory budget.

Parameters:
property strategy: str[source]

Return disk or memory for reference bookkeeping.

Return type:

str

property resident_bytes: int[source]

Return charged resident reference bytes in no-spill mode.

Return type:

int

count_rows(output_name)[source]

Return the number of staged references for one output.

Parameters:

output_name (str)

Return type:

int

append(output_name, position, reference)[source]

Stage one unique matrix reference at its final sample position.

Parameters:
Return type:

None

iter_references(output_name, *, expected_rows)[source]

Yield references in exact sample order after validating full coverage.

Parameters:
  • output_name (str)

  • expected_rows (int)

Return type:

Iterator[MatrixRowReference]

assemble(output_name, *, expected_rows, purpose, force_disk=False)[source]

Assemble one output in sample order and release its reference state.

Parameters:
  • output_name (str)

  • expected_rows (int)

  • purpose (str)

  • force_disk (bool)

Return type:

MatrixAssembly

close()[source]

Release all in-memory or disk-backed reference state.

Return type:

None

vertebrae.utils.memory.resolve_memory_budget(config)[source]

Resolve an effective memory budget using psutil.

Parameters:

config (vertebrae.config.MemoryConfig) – Memory configuration.

Returns:

Resolved memory budget.

Return type:

MemoryBudget

vertebrae.utils.memory.estimate_embedding_from_probe(probe_embeddings, n_samples, batch_size, memory_config, scoring_row_multiplier=1.0)[source]

Estimate full embedding memory from a probe batch.

Parameters:
  • probe_embeddings (Any) – Dense or sparse probe embedding batch.

  • n_samples (int) – Full dataset sample count.

  • batch_size (int) – Planned embedding batch size.

  • memory_config (vertebrae.config.MemoryConfig) – Memory configuration.

  • scoring_row_multiplier (float) – Expected row expansion performed by the scorer.

Returns:

Estimated embedding footprint and strategy.

Return type:

EmbeddingMemoryEstimate

vertebrae.utils.memory.estimate_matrix_resident_bytes(matrix)[source]

Estimate memory required to hold a dense or sparse matrix.

Parameters:

matrix (Any) – Dense or sparse matrix.

Returns:

Estimated resident bytes.

Return type:

int

vertebrae.utils.memory.estimate_object_resident_bytes(value, _seen=None)[source]

Conservatively estimate the retained size of nested metadata.

Parameters:
  • value (Any)

  • _seen (Optional[set[int]])

Return type:

int

vertebrae.utils.memory.estimate_final_row_metadata_bytes(rows, *, expected_rows, expansion_factor, purpose)[source]

Estimate final row-aligned metadata before allocating retained containers.

Staged rows are intentionally read one at a time. expansion_factor accounts for the distinct labels, IDs, groups, annotations, and provenance containers built from each staged row by a materializer.

Parameters:
  • rows (Iterable[Mapping[str, Any]])

  • expected_rows (int)

  • expansion_factor (float)

  • purpose (str)

Return type:

int

vertebrae.utils.memory.admit_final_metadata(stager, required_bytes, *, purpose, retained_bytes=0, transient_bytes=0)[source]

Admit retained row metadata and the current output’s transient peak.

Parameters:
Return type:

int

vertebrae.utils.memory.estimate_metadata_resident_bytes(metadata)[source]

Estimate resident bytes from embedding metadata.

Parameters:

metadata (dict[str, Any]) – Embedding metadata dictionary.

Returns:

Estimated bytes, or None if metadata is incomplete.

Return type:

Optional[int]

vertebrae.utils.memory.estimate_metadata_dense_scoring_bytes(metadata)[source]

Estimate the hypothetical dense footprint from embedding metadata.

Parameters:

metadata (dict[str, Any]) – Embedding metadata dictionary.

Returns:

Estimated dense footprint, or None if metadata is incomplete.

Return type:

Optional[int]

vertebrae.utils.memory.estimate_metadata_scoring_input_bytes(metadata, *, scoring_row_multiplier=1.0)[source]

Estimate the actual dense or sparse representation consumed by scoring.

Parameters:
  • metadata (dict[str, Any])

  • scoring_row_multiplier (float)

Return type:

Optional[int]

vertebrae.utils.memory.assert_within_memory(required_bytes, memory_config, purpose)[source]

Fail fast when a planned allocation exceeds the memory budget.

Parameters:
  • required_bytes (int) – Estimated required bytes.

  • memory_config (vertebrae.config.MemoryConfig) – Memory configuration.

  • purpose (str) – Human-readable description of the planned work.

Returns:

Resolved memory budget.

Raises:

ValueError – If required_bytes exceeds the configured budget.

Return type:

MemoryBudget

vertebrae.utils.memory.assemble_matrix_rows(rows, memory_config, *, purpose)[source]

Stack compatible matrix rows without forcing an over-budget dense allocation.

Dense and sparse results that exceed the resolved memory budget are assembled into temporary disk-backed arrays when spill is enabled. Staging files are removed automatically once the returned matrix is no longer referenced.

Parameters:
Return type:

MatrixAssembly

vertebrae.utils.memory.largest_fitting_subsample_rate(required_bytes, memory_config)[source]

Estimate the largest sample fraction that fits the memory budget.

Parameters:
Returns:

Fraction in (0, 1] that should fit the configured budget.

Return type:

float

vertebrae.utils.memory.sparse_matrix_nbytes(matrix)[source]

Estimate resident bytes for an existing scipy sparse matrix.

Parameters:

matrix (Any)

Return type:

int

vertebrae.utils.memory.sparse_nbytes_from_nnz(nnz, n_rows, dtype)[source]

Estimate CSR sparse bytes from non-zero count and row count.

Parameters:
  • nnz (int)

  • n_rows (int)

  • dtype (numpy.dtype)

Return type:

int

vertebrae.utils.memory.estimate_sparse_bytes(n_samples, n_features, dtype, density)[source]

Estimate sparse matrix bytes from shape, dtype, and density.

Parameters:
  • n_samples (int)

  • n_features (int)

  • dtype (numpy.dtype)

  • density (float)

Return type:

int