vertebrae.utils.memory
Memory estimation and admission helpers.
Classes
Resolved memory budget from a MemoryConfig. |
|
Estimated memory footprint for an embedding artifact. |
|
A matrix assembled either in memory or through a temporary local memmap. |
|
Opaque reference to one incrementally staged matrix row. |
|
Opaque reference to one incrementally staged metadata record. |
|
Stage compatible matrix rows without retaining extractor batch buffers. |
|
Stage per-row metadata without retaining an unbounded candidate list. |
|
Store ordered matrix-row references without an unbounded Python list. |
Functions
|
Resolve an effective memory budget using psutil. |
|
Estimate full embedding memory from a probe batch. |
|
Estimate memory required to hold a dense or sparse matrix. |
|
Conservatively estimate the retained size of nested metadata. |
|
Estimate final row-aligned metadata before allocating retained containers. |
|
Admit retained row metadata and the current output's transient peak. |
|
Estimate resident bytes from embedding metadata. |
|
Estimate the hypothetical dense footprint from embedding metadata. |
|
Estimate the actual dense or sparse representation consumed by scoring. |
|
Fail fast when a planned allocation exceeds the memory budget. |
|
Stack compatible matrix rows without forcing an over-budget dense allocation. |
|
Estimate the largest sample fraction that fits the memory budget. |
|
Estimate resident bytes for an existing scipy sparse matrix. |
|
Estimate CSR sparse bytes from non-zero count and row count. |
|
Estimate sparse matrix bytes from shape, dtype, and density. |
Module Contents
- class vertebrae.utils.memory.EmbeddingMemoryEstimate[source]
Estimated memory footprint for an embedding artifact.
- 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:
memory_config (vertebrae.config.MemoryConfig)
purpose (str)
- 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:
- 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:
- 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
- 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:
memory_config (vertebrae.config.MemoryConfig)
purpose (str)
matrix_stager (IncrementalMatrixStager)
- 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:
- 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:
reference (MetadataRowReference)
output_name (str)
- 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
- 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:
memory_config (vertebrae.config.MemoryConfig)
purpose (str)
matrix_stager (IncrementalMatrixStager)
- 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:
output_name (str)
position (int)
reference (MatrixRowReference)
- 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]
- 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:
- 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:
- 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_factoraccounts 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:
stager (IncrementalMatrixStager)
required_bytes (int)
purpose (str)
retained_bytes (int)
transient_bytes (int)
- 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:
- 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:
rows (list[Any])
memory_config (vertebrae.config.MemoryConfig)
purpose (str)
- Return type:
- vertebrae.utils.memory.largest_fitting_subsample_rate(required_bytes, memory_config)[source]
Estimate the largest sample fraction that fits the memory budget.
- Parameters:
required_bytes (int) – Estimated bytes for the full sample set.
memory_config (vertebrae.config.MemoryConfig) – Memory configuration.
- 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