mirror of
https://github.com/ruvnet/RuView
synced 2026-07-26 18:01:48 +00:00
d803bfe2b1
git-subtree-dir: vendor/ruvector git-subtree-split: b64c21726f2bb37286d9ee36a7869fef60cc6900
2284 lines
74 KiB
Rust
2284 lines
74 KiB
Rust
//! Block-based storage engine for temporal tensor compression (ADR-018).
|
|
//!
|
|
//! Provides tiered quantized storage with CRC32 integrity checking,
|
|
//! access-pattern tracking, and eviction support. Each block of tensor
|
|
//! data is quantized at the bit width appropriate for its storage tier
|
|
//! and tracked with rich metadata for tier-promotion/demotion decisions.
|
|
//!
|
|
//! # Storage Tiers
|
|
//!
|
|
//! | Tier | Bits | Description |
|
|
//! |-------|------|-------------------------------------|
|
|
//! | Tier0 | 0 | Evicted: metadata only, no payload |
|
|
//! | Tier1 | 8 | Hot: full fidelity quantization |
|
|
//! | Tier2 | 7 | Warm: moderate compression |
|
|
//! | Tier3 | 3 | Cold: aggressive compression |
|
|
//!
|
|
//! # Example
|
|
//!
|
|
//! ```rust
|
|
//! use ruvector_temporal_tensor::store::{BlockKey, Tier, TieredStore, ReconstructPolicy};
|
|
//!
|
|
//! let mut store = TieredStore::new(4096);
|
|
//! let key = BlockKey { tensor_id: 1, block_index: 0 };
|
|
//! let data = vec![1.0f32; 64];
|
|
//!
|
|
//! store.put(key, &data, Tier::Tier1, 0).unwrap();
|
|
//! assert_eq!(store.block_count(), 1);
|
|
//!
|
|
//! let mut out = vec![0.0f32; 64];
|
|
//! let n = store.get(key, &mut out, 1).unwrap();
|
|
//! assert_eq!(n, 64);
|
|
//! ```
|
|
|
|
use std::collections::HashMap;
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Core types
|
|
// ---------------------------------------------------------------------------
|
|
|
|
/// Unique identifier for a tensor block.
|
|
///
|
|
/// Composed of the owning tensor's 128-bit ID and a block index within
|
|
/// that tensor, allowing fine-grained block-level storage and retrieval.
|
|
#[derive(Clone, Copy, PartialEq, Eq, Hash, Debug)]
|
|
pub struct BlockKey {
|
|
pub tensor_id: u128,
|
|
pub block_index: u32,
|
|
}
|
|
|
|
/// Storage tier for a block.
|
|
///
|
|
/// Tiers form a hierarchy from hot (high fidelity, fast access) to evicted
|
|
/// (metadata-only, zero payload bytes).
|
|
#[derive(Clone, Copy, PartialEq, Eq, Debug, Hash)]
|
|
#[repr(u8)]
|
|
pub enum Tier {
|
|
/// Evicted: compressed to zero bits, only metadata remains.
|
|
Tier0 = 0,
|
|
/// Hot: 8-bit quantization, full fidelity.
|
|
Tier1 = 1,
|
|
/// Warm: 7-bit quantization.
|
|
Tier2 = 2,
|
|
/// Cold: 3-bit quantization.
|
|
Tier3 = 3,
|
|
}
|
|
|
|
/// Data type of the original tensor.
|
|
#[derive(Clone, Copy, PartialEq, Eq, Debug)]
|
|
#[repr(u8)]
|
|
pub enum DType {
|
|
F32 = 0,
|
|
F16 = 1,
|
|
BF16 = 2,
|
|
}
|
|
|
|
/// Reconstruction policy for evicted (Tier0) blocks.
|
|
#[derive(Clone, Copy, PartialEq, Eq, Debug)]
|
|
#[repr(u8)]
|
|
pub enum ReconstructPolicy {
|
|
/// No reconstruction possible. Reads fail or return zeros.
|
|
None = 0,
|
|
/// Reconstruct from base + delta chain.
|
|
Delta = 1,
|
|
/// Reconstruct from stored low-rank factors.
|
|
Factor = 2,
|
|
}
|
|
|
|
/// Complete metadata for a single block.
|
|
#[derive(Clone, Debug)]
|
|
pub struct BlockMeta {
|
|
pub key: BlockKey,
|
|
pub dtype: DType,
|
|
pub tier: Tier,
|
|
/// Quantization bit width (8, 7, 5, or 3).
|
|
pub bits: u8,
|
|
/// Quantization scale: `max(|v|) / qmax`.
|
|
pub scale: f32,
|
|
/// Quantization zero point (0 for symmetric).
|
|
pub zero_point: i16,
|
|
/// Tick at which this block was created.
|
|
pub created_at: u64,
|
|
/// Tick of the most recent access.
|
|
pub last_access_at: u64,
|
|
/// Cumulative access count.
|
|
pub access_count: u32,
|
|
/// Exponential moving average of access rate.
|
|
pub ema_rate: f32,
|
|
/// Sliding-window bitset for the last 64 ticks.
|
|
pub window: u64,
|
|
/// CRC32 checksum of quantized payload concatenated with scale bytes.
|
|
pub checksum: u32,
|
|
/// How to reconstruct if evicted.
|
|
pub reconstruct: ReconstructPolicy,
|
|
/// Number of ticks spent in the current tier.
|
|
pub tier_age: u32,
|
|
/// Optional parent tensor ID for delta-chain lineage.
|
|
pub lineage_parent: Option<u128>,
|
|
/// Size of this block's quantized payload in bytes.
|
|
pub block_bytes: u32,
|
|
}
|
|
|
|
/// Errors produced by the storage engine.
|
|
#[derive(Clone, Debug, PartialEq, Eq)]
|
|
pub enum StoreError {
|
|
/// The block has been evicted to Tier0 and cannot be read directly.
|
|
TensorEvicted,
|
|
/// No block exists for the given key.
|
|
BlockNotFound,
|
|
/// CRC32 verification failed after read.
|
|
ChecksumMismatch,
|
|
/// An underlying I/O operation failed.
|
|
IOError,
|
|
/// The memory budget has been exhausted.
|
|
BudgetExhausted,
|
|
/// The block data is malformed or invalid.
|
|
InvalidBlock,
|
|
/// A delta reconstruction chain exceeded the maximum depth.
|
|
DeltaChainTooLong,
|
|
/// Reconstruction of an evicted block failed.
|
|
ReconstructionFailed,
|
|
/// The provided data is malformed or could not be parsed.
|
|
InvalidData,
|
|
/// The delta chain is at maximum length and cannot accept more deltas.
|
|
ChainFull,
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Traits
|
|
// ---------------------------------------------------------------------------
|
|
|
|
/// Clock abstraction for deterministic time in tests and production.
|
|
pub trait Clock {
|
|
/// Returns the current tick count.
|
|
fn now_ticks(&self) -> u64;
|
|
}
|
|
|
|
/// Block I/O for reading and writing raw quantized data.
|
|
pub trait BlockIO {
|
|
/// Read quantized bytes for `key` from the given `tier` into `dst`.
|
|
/// Returns the number of bytes written to `dst`.
|
|
fn read_block(&self, tier: Tier, key: BlockKey, dst: &mut [u8]) -> Result<usize, StoreError>;
|
|
|
|
/// Write raw quantized bytes `src` for `key` into the given `tier`.
|
|
fn write_block(&mut self, tier: Tier, key: BlockKey, src: &[u8]) -> Result<(), StoreError>;
|
|
|
|
/// Delete the raw data for `key` from the given `tier`.
|
|
fn delete_block(&mut self, tier: Tier, key: BlockKey) -> Result<(), StoreError>;
|
|
}
|
|
|
|
/// Metadata log for append-only persistence of block metadata.
|
|
pub trait MetaLog {
|
|
/// Append (or upsert) a metadata record.
|
|
fn append(&mut self, rec: &BlockMeta) -> Result<(), StoreError>;
|
|
|
|
/// Look up metadata by key.
|
|
fn get(&self, key: BlockKey) -> Option<&BlockMeta>;
|
|
|
|
/// Iterate over all metadata records.
|
|
fn iter(&self) -> Box<dyn Iterator<Item = &BlockMeta> + '_>;
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// CRC32
|
|
// ---------------------------------------------------------------------------
|
|
|
|
/// Compute CRC32 using the standard reflected polynomial (0xEDB88320).
|
|
///
|
|
/// This is the same algorithm used by zlib/gzip/PNG. No lookup table is
|
|
/// used to keep the binary small; the byte-at-a-time loop is sufficient
|
|
/// for the block sizes involved.
|
|
pub fn crc32(data: &[u8]) -> u32 {
|
|
let mut crc: u32 = 0xFFFF_FFFF;
|
|
for &byte in data {
|
|
crc ^= byte as u32;
|
|
for _ in 0..8 {
|
|
if crc & 1 != 0 {
|
|
crc = (crc >> 1) ^ 0xEDB8_8320;
|
|
} else {
|
|
crc >>= 1;
|
|
}
|
|
}
|
|
}
|
|
!crc
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Helpers
|
|
// ---------------------------------------------------------------------------
|
|
|
|
/// Return the default bit width for a storage tier.
|
|
fn bits_for_tier(tier: Tier) -> u8 {
|
|
match tier {
|
|
Tier::Tier0 => 0,
|
|
Tier::Tier1 => 8,
|
|
Tier::Tier2 => 7,
|
|
Tier::Tier3 => 3,
|
|
}
|
|
}
|
|
|
|
/// Compute the maximum representable signed magnitude for a given bit width.
|
|
///
|
|
/// `qmax = 2^(bits-1) - 1`. Returns 0 for invalid widths.
|
|
#[inline]
|
|
fn qmax(bits: u8) -> i32 {
|
|
if bits == 0 || bits > 8 {
|
|
return 0;
|
|
}
|
|
(1i32 << (bits - 1)) - 1
|
|
}
|
|
|
|
/// Internal representation of a stored quantized block.
|
|
struct BlockData {
|
|
/// Number of original f32 elements (needed for exact dequantization).
|
|
element_count: u32,
|
|
/// Packed quantized bytes.
|
|
packed: Vec<u8>,
|
|
}
|
|
|
|
/// Quantize an f32 slice using symmetric quantization at the given bit width.
|
|
///
|
|
/// Returns the packed byte vector and the computed scale factor.
|
|
fn quantize_block(data: &[f32], bits: u8) -> (Vec<u8>, f32) {
|
|
let qm = qmax(bits);
|
|
if qm == 0 || data.is_empty() {
|
|
return (Vec::new(), 0.0);
|
|
}
|
|
let qm_f = qm as f32;
|
|
|
|
// Find the maximum finite absolute value.
|
|
let max_abs = data
|
|
.iter()
|
|
.filter(|v| v.is_finite())
|
|
.fold(0.0f32, |acc, v| acc.max(v.abs()));
|
|
|
|
let scale = if max_abs == 0.0 { 0.0 } else { max_abs / qm_f };
|
|
let inv_scale = if scale == 0.0 { 0.0 } else { 1.0 / scale };
|
|
|
|
let bits_u32 = bits as u32;
|
|
let needed = (data.len() * bits as usize).div_ceil(8);
|
|
let mut packed = Vec::with_capacity(needed);
|
|
|
|
let mut acc: u64 = 0;
|
|
let mut acc_bits: u32 = 0;
|
|
|
|
for &v in data {
|
|
let q = if v.is_finite() {
|
|
(v * inv_scale).round() as i32
|
|
} else {
|
|
0
|
|
}
|
|
.clamp(-qm, qm);
|
|
|
|
let u = (q + qm) as u32;
|
|
acc |= (u as u64) << acc_bits;
|
|
acc_bits += bits_u32;
|
|
|
|
while acc_bits >= 8 {
|
|
packed.push((acc & 0xFF) as u8);
|
|
acc >>= 8;
|
|
acc_bits -= 8;
|
|
}
|
|
}
|
|
|
|
if acc_bits > 0 {
|
|
packed.push((acc & 0xFF) as u8);
|
|
}
|
|
|
|
(packed, scale)
|
|
}
|
|
|
|
/// Dequantize packed bytes back to f32 using the given scale and bit width.
|
|
///
|
|
/// Writes up to `count` values into `out` and returns how many were written.
|
|
fn dequantize_block(packed: &[u8], scale: f32, bits: u8, count: usize, out: &mut [f32]) -> usize {
|
|
let qm = qmax(bits);
|
|
if qm == 0 || packed.is_empty() {
|
|
return 0;
|
|
}
|
|
|
|
let bits_u32 = bits as u32;
|
|
let mask = (1u64 << bits_u32) - 1;
|
|
let limit = count.min(out.len());
|
|
|
|
let mut acc: u64 = 0;
|
|
let mut acc_bits: u32 = 0;
|
|
let mut byte_idx: usize = 0;
|
|
let mut written: usize = 0;
|
|
|
|
while written < limit {
|
|
while acc_bits < bits_u32 && byte_idx < packed.len() {
|
|
acc |= (packed[byte_idx] as u64) << acc_bits;
|
|
acc_bits += 8;
|
|
byte_idx += 1;
|
|
}
|
|
if acc_bits < bits_u32 {
|
|
break;
|
|
}
|
|
|
|
let u = (acc & mask) as i32;
|
|
acc >>= bits_u32;
|
|
acc_bits -= bits_u32;
|
|
|
|
out[written] = (u - qm) as f32 * scale;
|
|
written += 1;
|
|
}
|
|
|
|
written
|
|
}
|
|
|
|
/// Compute the CRC32 checksum over quantized payload concatenated with scale.
|
|
fn block_checksum(packed: &[u8], scale: f32) -> u32 {
|
|
let scale_bytes = scale.to_le_bytes();
|
|
let total = packed.len() + scale_bytes.len();
|
|
let mut buf = Vec::with_capacity(total);
|
|
buf.extend_from_slice(packed);
|
|
buf.extend_from_slice(&scale_bytes);
|
|
crc32(&buf)
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// TickResult
|
|
// ---------------------------------------------------------------------------
|
|
|
|
/// Summary of actions taken during a budgeted maintenance tick.
|
|
#[derive(Debug, Default)]
|
|
pub struct TickResult {
|
|
/// Number of blocks promoted to a hotter tier.
|
|
pub upgrades: u32,
|
|
/// Number of blocks demoted to a colder tier.
|
|
pub downgrades: u32,
|
|
/// Number of blocks evicted to Tier0.
|
|
pub evictions: u32,
|
|
/// Total bytes freed by evictions and downgrades.
|
|
pub bytes_freed: usize,
|
|
/// Number of budget operations consumed.
|
|
pub ops_used: u32,
|
|
/// Total migration candidates identified before budget limits.
|
|
pub candidates_found: u32,
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Type adapters: store types <-> tiering types
|
|
// ---------------------------------------------------------------------------
|
|
|
|
/// Convert a store [`Tier`] to a [`crate::tiering::Tier`].
|
|
fn to_tiering_tier(tier: Tier) -> crate::tiering::Tier {
|
|
match tier {
|
|
Tier::Tier0 => crate::tiering::Tier::Tier0,
|
|
Tier::Tier1 => crate::tiering::Tier::Tier1,
|
|
Tier::Tier2 => crate::tiering::Tier::Tier2,
|
|
Tier::Tier3 => crate::tiering::Tier::Tier3,
|
|
}
|
|
}
|
|
|
|
/// Convert a [`crate::tiering::Tier`] to a store [`Tier`].
|
|
fn from_tiering_tier(tier: crate::tiering::Tier) -> Tier {
|
|
match tier {
|
|
crate::tiering::Tier::Tier0 => Tier::Tier0,
|
|
crate::tiering::Tier::Tier1 => Tier::Tier1,
|
|
crate::tiering::Tier::Tier2 => Tier::Tier2,
|
|
crate::tiering::Tier::Tier3 => Tier::Tier3,
|
|
}
|
|
}
|
|
|
|
/// Build a [`crate::tiering::BlockMeta`] from a store [`BlockMeta`] at time `now`.
|
|
fn to_tiering_meta(meta: &BlockMeta, now: u64) -> crate::tiering::BlockMeta {
|
|
crate::tiering::BlockMeta {
|
|
ema_rate: meta.ema_rate,
|
|
access_window: meta.window,
|
|
last_access: meta.last_access_at,
|
|
access_count: meta.access_count as u64,
|
|
current_tier: to_tiering_tier(meta.tier),
|
|
tier_since: now.saturating_sub(meta.tier_age as u64),
|
|
}
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// TieredStore
|
|
// ---------------------------------------------------------------------------
|
|
|
|
/// In-memory tiered storage engine for quantized tensor blocks.
|
|
///
|
|
/// Provides put/get with automatic quantization and dequantization,
|
|
/// per-block metadata tracking, access-pattern statistics, and
|
|
/// eviction to Tier0.
|
|
pub struct TieredStore {
|
|
/// Nominal block size hint (bytes). Stored for reference; actual block
|
|
/// sizes are determined by the data passed to [`put`].
|
|
block_bytes: usize,
|
|
|
|
/// Block metadata index keyed by [`BlockKey`].
|
|
index: HashMap<BlockKey, BlockMeta>,
|
|
|
|
/// Tier1 (hot, 8-bit) quantized data.
|
|
tier1_data: HashMap<BlockKey, BlockData>,
|
|
/// Tier2 (warm, 7-bit) quantized data.
|
|
tier2_data: HashMap<BlockKey, BlockData>,
|
|
/// Tier3 (cold, 3-bit) quantized data.
|
|
tier3_data: HashMap<BlockKey, BlockData>,
|
|
|
|
/// Keys present in each tier, for candidate-selection scans.
|
|
tier1_keys: Vec<BlockKey>,
|
|
tier2_keys: Vec<BlockKey>,
|
|
tier3_keys: Vec<BlockKey>,
|
|
|
|
/// Witness log for auditing tiering decisions.
|
|
witness_log: crate::metrics::WitnessLog,
|
|
|
|
/// Optional coherence checker for read-after-write validation.
|
|
coherence: Option<crate::coherence::CoherenceCheck>,
|
|
/// Epoch tracker for staleness detection.
|
|
epoch_tracker: crate::coherence::EpochTracker,
|
|
/// Metrics time-series for trend analysis.
|
|
metrics_series: crate::metrics::MetricsSeries,
|
|
}
|
|
|
|
/// Smoothing constant for the exponential moving average of access rate.
|
|
const EMA_ALPHA: f32 = 0.1;
|
|
|
|
impl TieredStore {
|
|
/// Create a new store with the given nominal block size (in bytes).
|
|
pub fn new(block_bytes: usize) -> Self {
|
|
Self {
|
|
block_bytes,
|
|
index: HashMap::new(),
|
|
tier1_data: HashMap::new(),
|
|
tier2_data: HashMap::new(),
|
|
tier3_data: HashMap::new(),
|
|
tier1_keys: Vec::new(),
|
|
tier2_keys: Vec::new(),
|
|
tier3_keys: Vec::new(),
|
|
witness_log: crate::metrics::WitnessLog::new(10_000),
|
|
coherence: None,
|
|
epoch_tracker: crate::coherence::EpochTracker::new(),
|
|
metrics_series: crate::metrics::MetricsSeries::new(256),
|
|
}
|
|
}
|
|
|
|
/// Nominal block size hint (bytes) configured at construction.
|
|
#[inline]
|
|
pub fn block_bytes(&self) -> usize {
|
|
self.block_bytes
|
|
}
|
|
|
|
/// Access the witness log.
|
|
pub fn witness_log(&self) -> &crate::metrics::WitnessLog {
|
|
&self.witness_log
|
|
}
|
|
|
|
/// Access the witness log mutably.
|
|
pub fn witness_log_mut(&mut self) -> &mut crate::metrics::WitnessLog {
|
|
&mut self.witness_log
|
|
}
|
|
|
|
/// Enable coherence checking with the given configuration.
|
|
///
|
|
/// When enabled, every `put()` records a write epoch in the epoch tracker.
|
|
/// Callers can use [`coherence_check()`](Self::coherence_check) to validate
|
|
/// read-after-write consistency.
|
|
pub fn enable_coherence(&mut self, check: crate::coherence::CoherenceCheck) {
|
|
self.coherence = Some(check);
|
|
}
|
|
|
|
/// Disable coherence checking.
|
|
pub fn disable_coherence(&mut self) {
|
|
self.coherence = None;
|
|
}
|
|
|
|
/// Access the epoch tracker.
|
|
pub fn epoch_tracker(&self) -> &crate::coherence::EpochTracker {
|
|
&self.epoch_tracker
|
|
}
|
|
|
|
/// Access the epoch tracker mutably.
|
|
pub fn epoch_tracker_mut(&mut self) -> &mut crate::coherence::EpochTracker {
|
|
&mut self.epoch_tracker
|
|
}
|
|
|
|
/// Access the metrics time-series.
|
|
pub fn metrics_series(&self) -> &crate::metrics::MetricsSeries {
|
|
&self.metrics_series
|
|
}
|
|
|
|
/// Access the metrics time-series mutably.
|
|
pub fn metrics_series_mut(&mut self) -> &mut crate::metrics::MetricsSeries {
|
|
&mut self.metrics_series
|
|
}
|
|
|
|
/// Perform a coherence check on a recently written block.
|
|
///
|
|
/// Returns `None` if coherence checking is not enabled.
|
|
/// Returns `Some(Err(...))` if the block doesn't exist or is evicted.
|
|
/// Returns `Some(Ok(result))` with the coherence result.
|
|
pub fn coherence_check(
|
|
&mut self,
|
|
key: BlockKey,
|
|
original_data: &[f32],
|
|
now: u64,
|
|
) -> Option<Result<crate::coherence::CoherenceResult, StoreError>> {
|
|
let check = self.coherence.clone()?;
|
|
Some(check.check_coherence(self, key, original_data, now))
|
|
}
|
|
|
|
/// Compute current aggregate metrics.
|
|
pub fn metrics(&self) -> crate::metrics::StoreMetrics {
|
|
let mut m = crate::metrics::StoreMetrics::new();
|
|
m.total_blocks = self.index.len() as u64;
|
|
m.tier0_blocks = self
|
|
.index
|
|
.values()
|
|
.filter(|b| b.tier == Tier::Tier0)
|
|
.count() as u64;
|
|
m.tier1_blocks = self.tier1_keys.len() as u64;
|
|
m.tier2_blocks = self.tier2_keys.len() as u64;
|
|
m.tier3_blocks = self.tier3_keys.len() as u64;
|
|
m.tier1_bytes = self
|
|
.tier1_data
|
|
.values()
|
|
.map(|d| d.packed.len() as u64)
|
|
.sum();
|
|
m.tier2_bytes = self
|
|
.tier2_data
|
|
.values()
|
|
.map(|d| d.packed.len() as u64)
|
|
.sum();
|
|
m.tier3_bytes = self
|
|
.tier3_data
|
|
.values()
|
|
.map(|d| d.packed.len() as u64)
|
|
.sum();
|
|
m.total_evictions = self.witness_log.count_evictions() as u64;
|
|
m.tier_flips_last_minute = self.witness_log.tier_flip_rate(60, self.index.len() as u64);
|
|
m
|
|
}
|
|
|
|
/// Quantize `data` at the bit width for `tier` and store the block.
|
|
///
|
|
/// If a block with the same key already exists, it is replaced (the old
|
|
/// data is removed from whatever tier it resided in).
|
|
///
|
|
/// Returns [`StoreError::InvalidBlock`] if `tier` is [`Tier::Tier0`]
|
|
/// (you cannot directly write to the evicted tier).
|
|
pub fn put(
|
|
&mut self,
|
|
key: BlockKey,
|
|
data: &[f32],
|
|
tier: Tier,
|
|
now: u64,
|
|
) -> Result<(), StoreError> {
|
|
if tier == Tier::Tier0 {
|
|
return Err(StoreError::InvalidBlock);
|
|
}
|
|
|
|
let bits = bits_for_tier(tier);
|
|
let (packed, scale) = quantize_block(data, bits);
|
|
let checksum = block_checksum(&packed, scale);
|
|
|
|
// If the key already exists, remove old data first.
|
|
if let Some(old_meta) = self.index.get(&key) {
|
|
let old_tier = old_meta.tier;
|
|
self.remove_data(old_tier, key);
|
|
self.remove_from_bucket(old_tier, key);
|
|
}
|
|
|
|
let byte_count = packed.len() as u32;
|
|
let block = BlockData {
|
|
element_count: data.len() as u32,
|
|
packed,
|
|
};
|
|
|
|
match tier {
|
|
Tier::Tier1 => {
|
|
self.tier1_data.insert(key, block);
|
|
}
|
|
Tier::Tier2 => {
|
|
self.tier2_data.insert(key, block);
|
|
}
|
|
Tier::Tier3 => {
|
|
self.tier3_data.insert(key, block);
|
|
}
|
|
Tier::Tier0 => unreachable!(),
|
|
}
|
|
self.add_to_bucket(tier, key);
|
|
|
|
let meta = BlockMeta {
|
|
key,
|
|
dtype: DType::F32,
|
|
tier,
|
|
bits,
|
|
scale,
|
|
zero_point: 0,
|
|
created_at: now,
|
|
last_access_at: now,
|
|
access_count: 1,
|
|
ema_rate: 0.0,
|
|
window: 1,
|
|
checksum,
|
|
reconstruct: ReconstructPolicy::None,
|
|
tier_age: 0,
|
|
lineage_parent: None,
|
|
block_bytes: byte_count,
|
|
};
|
|
self.index.insert(key, meta);
|
|
|
|
// Record witness event for the write.
|
|
self.witness_log.record(
|
|
now,
|
|
crate::metrics::WitnessEvent::Access {
|
|
key,
|
|
score: 0.0,
|
|
tier,
|
|
},
|
|
);
|
|
|
|
// Record write epoch for staleness detection.
|
|
self.epoch_tracker.record_write(key);
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// Dequantize the block identified by `key` into `out`.
|
|
///
|
|
/// `now` is the current tick counter, used to update access statistics
|
|
/// and record a witness event.
|
|
///
|
|
/// Returns the number of f32 elements written to `out`.
|
|
///
|
|
/// # Errors
|
|
///
|
|
/// - [`StoreError::TensorEvicted`] if the block resides in Tier0.
|
|
/// - [`StoreError::BlockNotFound`] if no block exists for `key`.
|
|
/// - [`StoreError::ChecksumMismatch`] if the stored checksum does not
|
|
/// match a freshly computed checksum of the payload.
|
|
pub fn get(&mut self, key: BlockKey, out: &mut [f32], now: u64) -> Result<usize, StoreError> {
|
|
let meta = self.index.get(&key).ok_or(StoreError::BlockNotFound)?;
|
|
|
|
if meta.tier == Tier::Tier0 {
|
|
return Err(StoreError::TensorEvicted);
|
|
}
|
|
|
|
let tier = meta.tier;
|
|
let scale = meta.scale;
|
|
let bits = meta.bits;
|
|
let checksum = meta.checksum;
|
|
|
|
let block = self
|
|
.data_map(tier)
|
|
.and_then(|m| m.get(&key))
|
|
.ok_or(StoreError::BlockNotFound)?;
|
|
|
|
// Verify integrity.
|
|
let actual_crc = block_checksum(&block.packed, scale);
|
|
if actual_crc != checksum {
|
|
return Err(StoreError::ChecksumMismatch);
|
|
}
|
|
|
|
let n = dequantize_block(
|
|
&block.packed,
|
|
scale,
|
|
bits,
|
|
block.element_count as usize,
|
|
out,
|
|
);
|
|
|
|
// Update access statistics.
|
|
self.touch(key, now);
|
|
|
|
// Record witness event.
|
|
self.witness_log.record(
|
|
now,
|
|
crate::metrics::WitnessEvent::Access {
|
|
key,
|
|
score: 0.0, // score not computed during basic get
|
|
tier,
|
|
},
|
|
);
|
|
|
|
Ok(n)
|
|
}
|
|
|
|
/// Update access statistics for `key` at tick `now`.
|
|
///
|
|
/// Increments `access_count`, refreshes `last_access_at`, updates the
|
|
/// sliding-window bitset, and recalculates the EMA access rate.
|
|
/// Does nothing if the key is not present.
|
|
pub fn touch(&mut self, key: BlockKey, now: u64) {
|
|
if let Some(meta) = self.index.get_mut(&key) {
|
|
let delta = now.saturating_sub(meta.last_access_at);
|
|
|
|
// Update sliding-window bitset.
|
|
if delta >= 64 {
|
|
meta.window = 1;
|
|
} else if delta > 0 {
|
|
meta.window = (meta.window << delta) | 1;
|
|
}
|
|
// delta == 0: same tick, window unchanged but count still bumps.
|
|
|
|
// Update EMA access rate.
|
|
if delta > 0 {
|
|
let instant_rate = 1.0 / delta as f32;
|
|
meta.ema_rate = EMA_ALPHA * instant_rate + (1.0 - EMA_ALPHA) * meta.ema_rate;
|
|
}
|
|
|
|
meta.last_access_at = now;
|
|
meta.access_count = meta.access_count.saturating_add(1);
|
|
}
|
|
}
|
|
|
|
/// Return a reference to the metadata for `key`, if it exists.
|
|
pub fn meta(&self, key: BlockKey) -> Option<&BlockMeta> {
|
|
self.index.get(&key)
|
|
}
|
|
|
|
/// Total number of blocks tracked (including Tier0 evicted blocks).
|
|
pub fn block_count(&self) -> usize {
|
|
self.index.len()
|
|
}
|
|
|
|
/// Number of blocks currently in the given tier.
|
|
pub fn tier_count(&self, tier: Tier) -> usize {
|
|
match tier {
|
|
Tier::Tier0 => self
|
|
.index
|
|
.values()
|
|
.filter(|m| m.tier == Tier::Tier0)
|
|
.count(),
|
|
Tier::Tier1 => self.tier1_keys.len(),
|
|
Tier::Tier2 => self.tier2_keys.len(),
|
|
Tier::Tier3 => self.tier3_keys.len(),
|
|
}
|
|
}
|
|
|
|
/// Total bytes of quantized data stored across all active tiers.
|
|
pub fn total_bytes(&self) -> usize {
|
|
let sum = |map: &HashMap<BlockKey, BlockData>| -> usize {
|
|
map.values().map(|b| b.packed.len()).sum()
|
|
};
|
|
sum(&self.tier1_data) + sum(&self.tier2_data) + sum(&self.tier3_data)
|
|
}
|
|
|
|
/// Slice of block keys currently residing in the given tier.
|
|
///
|
|
/// Returns an empty slice for [`Tier::Tier0`].
|
|
pub fn blocks_in_tier(&self, tier: Tier) -> &[BlockKey] {
|
|
match tier {
|
|
Tier::Tier0 => &[],
|
|
Tier::Tier1 => &self.tier1_keys,
|
|
Tier::Tier2 => &self.tier2_keys,
|
|
Tier::Tier3 => &self.tier3_keys,
|
|
}
|
|
}
|
|
|
|
/// Evict a block to Tier0, removing its quantized payload.
|
|
///
|
|
/// The block's metadata is preserved with the specified
|
|
/// [`ReconstructPolicy`] so that higher-level code can decide how
|
|
/// (or whether) to reconstruct the data on future reads.
|
|
///
|
|
/// Returns [`StoreError::BlockNotFound`] if the key does not exist.
|
|
pub fn evict(&mut self, key: BlockKey, policy: ReconstructPolicy) -> Result<(), StoreError> {
|
|
let meta = self.index.get_mut(&key).ok_or(StoreError::BlockNotFound)?;
|
|
let old_tier = meta.tier;
|
|
|
|
if old_tier == Tier::Tier0 {
|
|
// Already evicted; just update the policy.
|
|
meta.reconstruct = policy;
|
|
return Ok(());
|
|
}
|
|
|
|
let bytes_freed = meta.block_bytes as usize;
|
|
let evict_ts = meta.last_access_at;
|
|
|
|
// Mutate metadata before touching the data maps (avoids a second
|
|
// lookup since we already have the mutable reference).
|
|
meta.tier = Tier::Tier0;
|
|
meta.reconstruct = policy;
|
|
meta.tier_age = 0;
|
|
meta.block_bytes = 0;
|
|
meta.bits = 0;
|
|
|
|
// Drop the mutable borrow so we can call helper methods.
|
|
self.remove_data(old_tier, key);
|
|
self.remove_from_bucket(old_tier, key);
|
|
|
|
// Record witness event for the eviction.
|
|
self.witness_log.record(
|
|
evict_ts,
|
|
crate::metrics::WitnessEvent::Eviction {
|
|
key,
|
|
score: 0.0,
|
|
bytes_freed,
|
|
},
|
|
);
|
|
|
|
Ok(())
|
|
}
|
|
|
|
// -- private helpers ----------------------------------------------------
|
|
|
|
/// Return a reference to the data map for the given tier.
|
|
fn data_map(&self, tier: Tier) -> Option<&HashMap<BlockKey, BlockData>> {
|
|
match tier {
|
|
Tier::Tier0 => None,
|
|
Tier::Tier1 => Some(&self.tier1_data),
|
|
Tier::Tier2 => Some(&self.tier2_data),
|
|
Tier::Tier3 => Some(&self.tier3_data),
|
|
}
|
|
}
|
|
|
|
/// Remove raw data for `key` from the given tier's map.
|
|
fn remove_data(&mut self, tier: Tier, key: BlockKey) {
|
|
match tier {
|
|
Tier::Tier1 => {
|
|
self.tier1_data.remove(&key);
|
|
}
|
|
Tier::Tier2 => {
|
|
self.tier2_data.remove(&key);
|
|
}
|
|
Tier::Tier3 => {
|
|
self.tier3_data.remove(&key);
|
|
}
|
|
Tier::Tier0 => {}
|
|
}
|
|
}
|
|
|
|
/// Remove `key` from the tier's candidate-selection bucket.
|
|
fn remove_from_bucket(&mut self, tier: Tier, key: BlockKey) {
|
|
let bucket = match tier {
|
|
Tier::Tier1 => &mut self.tier1_keys,
|
|
Tier::Tier2 => &mut self.tier2_keys,
|
|
Tier::Tier3 => &mut self.tier3_keys,
|
|
Tier::Tier0 => return,
|
|
};
|
|
if let Some(pos) = bucket.iter().position(|k| *k == key) {
|
|
bucket.swap_remove(pos);
|
|
}
|
|
}
|
|
|
|
/// Add `key` to the tier's candidate-selection bucket.
|
|
fn add_to_bucket(&mut self, tier: Tier, key: BlockKey) {
|
|
match tier {
|
|
Tier::Tier1 => self.tier1_keys.push(key),
|
|
Tier::Tier2 => self.tier2_keys.push(key),
|
|
Tier::Tier3 => self.tier3_keys.push(key),
|
|
Tier::Tier0 => {}
|
|
}
|
|
}
|
|
|
|
// -- tiering-aware methods -----------------------------------------------
|
|
|
|
/// Run a budgeted maintenance tick.
|
|
///
|
|
/// Evaluates all blocks, selects migration candidates, and executes
|
|
/// tier transitions within the given byte and operation budgets.
|
|
/// Returns a summary of actions taken.
|
|
pub fn tick(
|
|
&mut self,
|
|
config: &crate::tiering::TierConfig,
|
|
now: u64,
|
|
budget_bytes: usize,
|
|
budget_ops: u32,
|
|
) -> TickResult {
|
|
let mut result = TickResult::default();
|
|
|
|
// Step 1: Collect all blocks and convert to tiering types.
|
|
// Use sequential indices as tiering::BlockKey values to avoid collisions.
|
|
let store_keys: Vec<BlockKey> = self.index.keys().copied().collect();
|
|
if store_keys.is_empty() {
|
|
return result;
|
|
}
|
|
|
|
let tiering_blocks: Vec<(crate::tiering::BlockKey, crate::tiering::BlockMeta)> = store_keys
|
|
.iter()
|
|
.enumerate()
|
|
.map(|(idx, key)| {
|
|
let meta = &self.index[key];
|
|
(
|
|
crate::tiering::BlockKey(idx as u64),
|
|
to_tiering_meta(meta, now),
|
|
)
|
|
})
|
|
.collect();
|
|
|
|
let blocks_ref: Vec<(crate::tiering::BlockKey, &crate::tiering::BlockMeta)> =
|
|
tiering_blocks.iter().map(|(k, m)| (*k, m)).collect();
|
|
|
|
// Step 2: Select migration candidates (upgrades first by highest score,
|
|
// then downgrades by lowest score).
|
|
let candidates = crate::tiering::select_candidates(config, now, &blocks_ref);
|
|
result.candidates_found = candidates.len() as u32;
|
|
|
|
// Step 3: Process candidates within budget.
|
|
let mut remaining_bytes = budget_bytes;
|
|
let mut remaining_ops = budget_ops;
|
|
let mut migrated = std::collections::HashSet::new();
|
|
|
|
for candidate in &candidates {
|
|
if remaining_ops == 0 {
|
|
break;
|
|
}
|
|
|
|
let store_key = store_keys[candidate.key.0 as usize];
|
|
let target_tier = from_tiering_tier(candidate.target_tier);
|
|
let current_tier = from_tiering_tier(candidate.current_tier);
|
|
|
|
let old_bytes = self
|
|
.index
|
|
.get(&store_key)
|
|
.map(|m| m.block_bytes as usize)
|
|
.unwrap_or(0);
|
|
|
|
// Check byte budget.
|
|
if old_bytes > remaining_bytes {
|
|
continue;
|
|
}
|
|
|
|
if target_tier == Tier::Tier0 {
|
|
// Eviction.
|
|
if self.evict(store_key, ReconstructPolicy::None).is_ok() {
|
|
result.evictions += 1;
|
|
result.bytes_freed += old_bytes;
|
|
remaining_ops -= 1;
|
|
result.ops_used += 1;
|
|
remaining_bytes = remaining_bytes.saturating_sub(old_bytes);
|
|
migrated.insert(store_key);
|
|
}
|
|
} else {
|
|
// Tier migration.
|
|
let warm_bytes: usize = self.tier2_data.values().map(|b| b.packed.len()).sum();
|
|
let target_bits =
|
|
crate::tiering::bits_for_tier(config, to_tiering_tier(target_tier), warm_bytes);
|
|
|
|
let old_tier_u8 = current_tier as u8;
|
|
let new_tier_u8 = target_tier as u8;
|
|
|
|
if self
|
|
.migrate_block(store_key, target_tier, target_bits)
|
|
.is_ok()
|
|
{
|
|
let new_bytes = self
|
|
.index
|
|
.get(&store_key)
|
|
.map(|m| m.block_bytes as usize)
|
|
.unwrap_or(0);
|
|
|
|
if new_tier_u8 < old_tier_u8 {
|
|
// Upgrade (hotter tier).
|
|
result.upgrades += 1;
|
|
} else {
|
|
// Downgrade (colder tier).
|
|
result.downgrades += 1;
|
|
result.bytes_freed += old_bytes.saturating_sub(new_bytes);
|
|
}
|
|
|
|
// Record witness event for the tier change.
|
|
let reason = if new_tier_u8 < old_tier_u8 {
|
|
crate::metrics::TierChangeReason::ScoreUpgrade
|
|
} else {
|
|
crate::metrics::TierChangeReason::ScoreDowngrade
|
|
};
|
|
self.witness_log.record(
|
|
now,
|
|
crate::metrics::WitnessEvent::TierChange {
|
|
key: store_key,
|
|
from_tier: current_tier,
|
|
to_tier: target_tier,
|
|
score: candidate.score,
|
|
reason,
|
|
},
|
|
);
|
|
|
|
remaining_ops -= 1;
|
|
result.ops_used += 1;
|
|
remaining_bytes = remaining_bytes.saturating_sub(old_bytes);
|
|
migrated.insert(store_key);
|
|
}
|
|
}
|
|
}
|
|
|
|
// Step 4: For blocks not migrated, increment tier_age and call tick_decay.
|
|
for key in &store_keys {
|
|
if migrated.contains(key) {
|
|
continue;
|
|
}
|
|
if let Some(meta) = self.index.get_mut(key) {
|
|
meta.tier_age = meta.tier_age.saturating_add(1);
|
|
// Apply tick_decay via the tiering module.
|
|
let mut tm = crate::tiering::BlockMeta {
|
|
ema_rate: meta.ema_rate,
|
|
access_window: meta.window,
|
|
last_access: meta.last_access_at,
|
|
access_count: meta.access_count as u64,
|
|
current_tier: to_tiering_tier(meta.tier),
|
|
tier_since: now.saturating_sub(meta.tier_age as u64),
|
|
};
|
|
crate::tiering::tick_decay(config, &mut tm);
|
|
meta.ema_rate = tm.ema_rate;
|
|
meta.window = tm.access_window;
|
|
}
|
|
}
|
|
|
|
// Record a maintenance witness event.
|
|
self.witness_log.record(
|
|
now,
|
|
crate::metrics::WitnessEvent::Maintenance {
|
|
upgrades: result.upgrades,
|
|
downgrades: result.downgrades,
|
|
evictions: result.evictions,
|
|
bytes_freed: result.bytes_freed,
|
|
budget_remaining_bytes: remaining_bytes.min(u32::MAX as usize) as u32,
|
|
budget_remaining_ops: remaining_ops,
|
|
},
|
|
);
|
|
|
|
// Auto-record a metrics snapshot for trend analysis.
|
|
let snapshot_metrics = self.metrics();
|
|
self.metrics_series.record(now, snapshot_metrics);
|
|
|
|
result
|
|
}
|
|
|
|
/// Migrate a single block from one tier to another.
|
|
///
|
|
/// Re-quantizes the data at the target tier's bit width. The block's
|
|
/// metadata is updated with the new tier, bits, scale, checksum, and
|
|
/// `tier_age` is reset to 0.
|
|
fn migrate_block(
|
|
&mut self,
|
|
key: BlockKey,
|
|
target_tier: Tier,
|
|
target_bits: u8,
|
|
) -> Result<(), StoreError> {
|
|
// Read current metadata (copy fields to release the borrow).
|
|
let meta = self.index.get(&key).ok_or(StoreError::BlockNotFound)?;
|
|
let old_tier = meta.tier;
|
|
let old_bits = meta.bits;
|
|
let old_scale = meta.scale;
|
|
|
|
if old_tier == Tier::Tier0 {
|
|
return Err(StoreError::TensorEvicted);
|
|
}
|
|
if target_tier == Tier::Tier0 {
|
|
return Err(StoreError::InvalidBlock);
|
|
}
|
|
|
|
// Dequantize the old data to f32 within a limited scope so the
|
|
// immutable borrow on self (through data_map) is released before
|
|
// we need mutable access.
|
|
let (element_count, f32_data) = {
|
|
let block = self
|
|
.data_map(old_tier)
|
|
.and_then(|m| m.get(&key))
|
|
.ok_or(StoreError::BlockNotFound)?;
|
|
let ec = block.element_count;
|
|
let mut data = vec![0.0f32; ec as usize];
|
|
dequantize_block(&block.packed, old_scale, old_bits, ec as usize, &mut data);
|
|
(ec, data)
|
|
};
|
|
|
|
// Re-quantize at the target bit width.
|
|
let (packed, scale) = quantize_block(&f32_data, target_bits);
|
|
let checksum = block_checksum(&packed, scale);
|
|
let byte_count = packed.len() as u32;
|
|
let new_block = BlockData {
|
|
element_count,
|
|
packed,
|
|
};
|
|
|
|
// Remove from old tier.
|
|
self.remove_data(old_tier, key);
|
|
self.remove_from_bucket(old_tier, key);
|
|
|
|
// Insert into target tier.
|
|
match target_tier {
|
|
Tier::Tier1 => {
|
|
self.tier1_data.insert(key, new_block);
|
|
}
|
|
Tier::Tier2 => {
|
|
self.tier2_data.insert(key, new_block);
|
|
}
|
|
Tier::Tier3 => {
|
|
self.tier3_data.insert(key, new_block);
|
|
}
|
|
Tier::Tier0 => unreachable!(),
|
|
}
|
|
self.add_to_bucket(target_tier, key);
|
|
|
|
// Update metadata.
|
|
let meta = self.index.get_mut(&key).unwrap();
|
|
meta.tier = target_tier;
|
|
meta.bits = target_bits;
|
|
meta.scale = scale;
|
|
meta.checksum = checksum;
|
|
meta.tier_age = 0;
|
|
meta.block_bytes = byte_count;
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// Compute the current score for a block using the enhanced tiering
|
|
/// algorithm (EMA + popcount + recency).
|
|
///
|
|
/// Returns `None` if the block does not exist.
|
|
pub fn score_block(
|
|
&self,
|
|
key: BlockKey,
|
|
config: &crate::tiering::TierConfig,
|
|
now: u64,
|
|
) -> Option<f32> {
|
|
let meta = self.index.get(&key)?;
|
|
let tm = to_tiering_meta(meta, now);
|
|
Some(crate::tiering::compute_score(config, now, &tm))
|
|
}
|
|
|
|
/// Record an access event using the enhanced tiering algorithm.
|
|
///
|
|
/// Updates `ema_rate`, `access_window`, `last_access_at`, and
|
|
/// `access_count` using the configurable alpha from [`TierConfig`].
|
|
/// Does nothing if the key is not present.
|
|
pub fn touch_block(&mut self, key: BlockKey, config: &crate::tiering::TierConfig, now: u64) {
|
|
if let Some(meta) = self.index.get_mut(&key) {
|
|
let mut tm = crate::tiering::BlockMeta {
|
|
ema_rate: meta.ema_rate,
|
|
access_window: meta.window,
|
|
last_access: meta.last_access_at,
|
|
access_count: meta.access_count as u64,
|
|
current_tier: to_tiering_tier(meta.tier),
|
|
tier_since: now.saturating_sub(meta.tier_age as u64),
|
|
};
|
|
crate::tiering::touch(config, now, &mut tm);
|
|
meta.ema_rate = tm.ema_rate;
|
|
meta.window = tm.access_window;
|
|
meta.last_access_at = tm.last_access;
|
|
meta.access_count = tm.access_count.min(u32::MAX as u64) as u32;
|
|
}
|
|
}
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Trait implementations for TieredStore
|
|
// ---------------------------------------------------------------------------
|
|
|
|
impl BlockIO for TieredStore {
|
|
fn read_block(&self, tier: Tier, key: BlockKey, dst: &mut [u8]) -> Result<usize, StoreError> {
|
|
let map = self.data_map(tier).ok_or(StoreError::BlockNotFound)?;
|
|
let block = map.get(&key).ok_or(StoreError::BlockNotFound)?;
|
|
let n = block.packed.len().min(dst.len());
|
|
dst[..n].copy_from_slice(&block.packed[..n]);
|
|
Ok(n)
|
|
}
|
|
|
|
fn write_block(&mut self, tier: Tier, key: BlockKey, src: &[u8]) -> Result<(), StoreError> {
|
|
if tier == Tier::Tier0 {
|
|
return Err(StoreError::InvalidBlock);
|
|
}
|
|
let block = BlockData {
|
|
element_count: 0, // raw write; element count unknown
|
|
packed: src.to_vec(),
|
|
};
|
|
match tier {
|
|
Tier::Tier1 => {
|
|
self.tier1_data.insert(key, block);
|
|
}
|
|
Tier::Tier2 => {
|
|
self.tier2_data.insert(key, block);
|
|
}
|
|
Tier::Tier3 => {
|
|
self.tier3_data.insert(key, block);
|
|
}
|
|
Tier::Tier0 => unreachable!(),
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
fn delete_block(&mut self, tier: Tier, key: BlockKey) -> Result<(), StoreError> {
|
|
let removed = match tier {
|
|
Tier::Tier1 => self.tier1_data.remove(&key).is_some(),
|
|
Tier::Tier2 => self.tier2_data.remove(&key).is_some(),
|
|
Tier::Tier3 => self.tier3_data.remove(&key).is_some(),
|
|
Tier::Tier0 => false,
|
|
};
|
|
if removed {
|
|
Ok(())
|
|
} else {
|
|
Err(StoreError::BlockNotFound)
|
|
}
|
|
}
|
|
}
|
|
|
|
impl MetaLog for TieredStore {
|
|
fn append(&mut self, rec: &BlockMeta) -> Result<(), StoreError> {
|
|
self.index.insert(rec.key, rec.clone());
|
|
Ok(())
|
|
}
|
|
|
|
fn get(&self, key: BlockKey) -> Option<&BlockMeta> {
|
|
self.index.get(&key)
|
|
}
|
|
|
|
fn iter(&self) -> Box<dyn Iterator<Item = &BlockMeta> + '_> {
|
|
Box::new(self.index.values())
|
|
}
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Tests
|
|
// ---------------------------------------------------------------------------
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use std::collections::hash_map::DefaultHasher;
|
|
use std::hash::{Hash, Hasher};
|
|
|
|
fn make_key(tid: u128, idx: u32) -> BlockKey {
|
|
BlockKey {
|
|
tensor_id: tid,
|
|
block_index: idx,
|
|
}
|
|
}
|
|
|
|
// -- CRC32 -------------------------------------------------------------
|
|
|
|
#[test]
|
|
fn test_crc32_known_vector() {
|
|
// The CRC32 of the ASCII string "123456789" is 0xCBF43926.
|
|
let data = b"123456789";
|
|
assert_eq!(crc32(data), 0xCBF4_3926);
|
|
}
|
|
|
|
#[test]
|
|
fn test_crc32_empty() {
|
|
assert_eq!(crc32(&[]), 0x0000_0000);
|
|
}
|
|
|
|
#[test]
|
|
fn test_crc32_single_byte() {
|
|
// CRC32 of [0x00] is 0xD202EF8D.
|
|
assert_eq!(crc32(&[0x00]), 0xD202_EF8D);
|
|
}
|
|
|
|
// -- BlockKey hashing --------------------------------------------------
|
|
|
|
#[test]
|
|
fn test_block_key_equality() {
|
|
let a = make_key(1, 0);
|
|
let b = make_key(1, 0);
|
|
let c = make_key(1, 1);
|
|
assert_eq!(a, b);
|
|
assert_ne!(a, c);
|
|
}
|
|
|
|
#[test]
|
|
fn test_block_key_hash_differs() {
|
|
fn hash_of(k: &BlockKey) -> u64 {
|
|
let mut h = DefaultHasher::new();
|
|
k.hash(&mut h);
|
|
h.finish()
|
|
}
|
|
let a = make_key(1, 0);
|
|
let b = make_key(2, 0);
|
|
let c = make_key(1, 1);
|
|
// Different keys should (almost certainly) hash differently.
|
|
assert_ne!(hash_of(&a), hash_of(&b));
|
|
assert_ne!(hash_of(&a), hash_of(&c));
|
|
}
|
|
|
|
#[test]
|
|
fn test_block_key_hash_stable() {
|
|
fn hash_of(k: &BlockKey) -> u64 {
|
|
let mut h = DefaultHasher::new();
|
|
k.hash(&mut h);
|
|
h.finish()
|
|
}
|
|
let a = make_key(42, 7);
|
|
let b = make_key(42, 7);
|
|
assert_eq!(hash_of(&a), hash_of(&b));
|
|
}
|
|
|
|
// -- qmax helper -------------------------------------------------------
|
|
|
|
#[test]
|
|
fn test_qmax_values() {
|
|
assert_eq!(qmax(8), 127);
|
|
assert_eq!(qmax(7), 63);
|
|
assert_eq!(qmax(5), 15);
|
|
assert_eq!(qmax(3), 3);
|
|
assert_eq!(qmax(1), 0);
|
|
assert_eq!(qmax(0), 0);
|
|
assert_eq!(qmax(9), 0);
|
|
}
|
|
|
|
// -- Quantization roundtrip --------------------------------------------
|
|
|
|
#[test]
|
|
fn test_quantize_roundtrip_8bit() {
|
|
let data: Vec<f32> = (0..128).map(|i| (i as f32 - 64.0) * 0.1).collect();
|
|
let (packed, scale) = quantize_block(&data, 8);
|
|
let mut out = vec![0.0f32; 128];
|
|
let n = dequantize_block(&packed, scale, 8, 128, &mut out);
|
|
assert_eq!(n, 128);
|
|
for (i, (&orig, &dec)) in data.iter().zip(out.iter()).enumerate() {
|
|
let err = (orig - dec).abs();
|
|
let tol = if orig.abs() > 0.01 {
|
|
orig.abs() * 0.02
|
|
} else {
|
|
0.1
|
|
};
|
|
assert!(err < tol, "i={i} orig={orig} dec={dec} err={err}");
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_quantize_roundtrip_3bit() {
|
|
let data: Vec<f32> = (0..64).map(|i| (i as f32 - 32.0) * 0.5).collect();
|
|
let (packed, scale) = quantize_block(&data, 3);
|
|
let mut out = vec![0.0f32; 64];
|
|
let n = dequantize_block(&packed, scale, 3, 64, &mut out);
|
|
assert_eq!(n, 64);
|
|
let max_val = data.iter().map(|v| v.abs()).fold(0.0f32, f32::max);
|
|
for (&orig, &dec) in data.iter().zip(out.iter()) {
|
|
let err = (orig - dec).abs();
|
|
assert!(err < max_val * 0.35, "orig={orig} dec={dec} err={err}");
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_quantize_zeros() {
|
|
let data = vec![0.0f32; 64];
|
|
let (packed, scale) = quantize_block(&data, 8);
|
|
assert_eq!(scale, 0.0);
|
|
let mut out = vec![1.0f32; 64];
|
|
let n = dequantize_block(&packed, scale, 8, 64, &mut out);
|
|
assert_eq!(n, 64);
|
|
for &v in &out {
|
|
assert_eq!(v, 0.0);
|
|
}
|
|
}
|
|
|
|
// -- TieredStore put/get -----------------------------------------------
|
|
|
|
#[test]
|
|
fn test_store_put_get_roundtrip() {
|
|
let mut store = TieredStore::new(4096);
|
|
let key = make_key(1, 0);
|
|
let data: Vec<f32> = (0..64).map(|i| i as f32 * 0.25).collect();
|
|
|
|
store.put(key, &data, Tier::Tier1, 0).unwrap();
|
|
|
|
let mut out = vec![0.0f32; 64];
|
|
let n = TieredStore::get(&mut store, key, &mut out, 1).unwrap();
|
|
assert_eq!(n, 64);
|
|
|
|
for (i, (&orig, &dec)) in data.iter().zip(out.iter()).enumerate() {
|
|
let err = (orig - dec).abs();
|
|
let tol = if orig.abs() > 0.01 {
|
|
orig.abs() * 0.02
|
|
} else {
|
|
0.15
|
|
};
|
|
assert!(err < tol, "i={i} orig={orig} dec={dec} err={err}");
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_store_put_tier3_roundtrip() {
|
|
let mut store = TieredStore::new(4096);
|
|
let key = make_key(10, 5);
|
|
let data: Vec<f32> = (0..32).map(|i| (i as f32 - 16.0) * 0.5).collect();
|
|
|
|
store.put(key, &data, Tier::Tier3, 100).unwrap();
|
|
|
|
let meta = store.meta(key).unwrap();
|
|
assert_eq!(meta.tier, Tier::Tier3);
|
|
assert_eq!(meta.bits, 3);
|
|
assert_eq!(meta.created_at, 100);
|
|
|
|
let mut out = vec![0.0f32; 32];
|
|
let n = TieredStore::get(&mut store, key, &mut out, 101).unwrap();
|
|
assert_eq!(n, 32);
|
|
|
|
let max_val = data.iter().map(|v| v.abs()).fold(0.0f32, f32::max);
|
|
for (&orig, &dec) in data.iter().zip(out.iter()) {
|
|
let err = (orig - dec).abs();
|
|
assert!(err < max_val * 0.35, "orig={orig} dec={dec} err={err}");
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_store_get_not_found() {
|
|
let mut store = TieredStore::new(4096);
|
|
let key = make_key(99, 0);
|
|
let mut out = vec![0.0f32; 8];
|
|
assert_eq!(
|
|
TieredStore::get(&mut store, key, &mut out, 0),
|
|
Err(StoreError::BlockNotFound)
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn test_store_put_tier0_rejected() {
|
|
let mut store = TieredStore::new(4096);
|
|
let key = make_key(1, 0);
|
|
let data = vec![1.0f32; 8];
|
|
assert_eq!(
|
|
store.put(key, &data, Tier::Tier0, 0),
|
|
Err(StoreError::InvalidBlock)
|
|
);
|
|
}
|
|
|
|
// -- Eviction ----------------------------------------------------------
|
|
|
|
#[test]
|
|
fn test_eviction() {
|
|
let mut store = TieredStore::new(4096);
|
|
let key = make_key(1, 0);
|
|
let data = vec![1.0f32; 64];
|
|
|
|
store.put(key, &data, Tier::Tier1, 0).unwrap();
|
|
assert_eq!(store.tier_count(Tier::Tier1), 1);
|
|
assert!(store.total_bytes() > 0);
|
|
|
|
store.evict(key, ReconstructPolicy::Delta).unwrap();
|
|
|
|
let meta = store.meta(key).unwrap();
|
|
assert_eq!(meta.tier, Tier::Tier0);
|
|
assert_eq!(meta.reconstruct, ReconstructPolicy::Delta);
|
|
assert_eq!(meta.block_bytes, 0);
|
|
assert_eq!(meta.bits, 0);
|
|
assert_eq!(meta.tier_age, 0);
|
|
|
|
// Data is gone; read should fail with TensorEvicted.
|
|
let mut out = vec![0.0f32; 64];
|
|
assert_eq!(
|
|
TieredStore::get(&mut store, key, &mut out, 1),
|
|
Err(StoreError::TensorEvicted)
|
|
);
|
|
|
|
// Tier1 should be empty; Tier0 count should be 1.
|
|
assert_eq!(store.tier_count(Tier::Tier1), 0);
|
|
assert_eq!(store.tier_count(Tier::Tier0), 1);
|
|
|
|
// Block still exists in the index (metadata preserved).
|
|
assert_eq!(store.block_count(), 1);
|
|
}
|
|
|
|
#[test]
|
|
fn test_eviction_not_found() {
|
|
let mut store = TieredStore::new(4096);
|
|
let key = make_key(1, 0);
|
|
assert_eq!(
|
|
store.evict(key, ReconstructPolicy::None),
|
|
Err(StoreError::BlockNotFound),
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn test_eviction_idempotent() {
|
|
let mut store = TieredStore::new(4096);
|
|
let key = make_key(1, 0);
|
|
store.put(key, &[1.0; 16], Tier::Tier2, 0).unwrap();
|
|
|
|
store.evict(key, ReconstructPolicy::None).unwrap();
|
|
// Evicting again should succeed and update the policy.
|
|
store.evict(key, ReconstructPolicy::Factor).unwrap();
|
|
|
|
let meta = store.meta(key).unwrap();
|
|
assert_eq!(meta.reconstruct, ReconstructPolicy::Factor);
|
|
}
|
|
|
|
// -- Tier counts -------------------------------------------------------
|
|
|
|
#[test]
|
|
fn test_tier_counts() {
|
|
let mut store = TieredStore::new(4096);
|
|
let data = vec![1.0f32; 16];
|
|
|
|
store.put(make_key(1, 0), &data, Tier::Tier1, 0).unwrap();
|
|
store.put(make_key(2, 0), &data, Tier::Tier1, 0).unwrap();
|
|
store.put(make_key(3, 0), &data, Tier::Tier2, 0).unwrap();
|
|
store.put(make_key(4, 0), &data, Tier::Tier3, 0).unwrap();
|
|
store.put(make_key(5, 0), &data, Tier::Tier3, 0).unwrap();
|
|
store.put(make_key(6, 0), &data, Tier::Tier3, 0).unwrap();
|
|
|
|
assert_eq!(store.block_count(), 6);
|
|
assert_eq!(store.tier_count(Tier::Tier0), 0);
|
|
assert_eq!(store.tier_count(Tier::Tier1), 2);
|
|
assert_eq!(store.tier_count(Tier::Tier2), 1);
|
|
assert_eq!(store.tier_count(Tier::Tier3), 3);
|
|
|
|
assert_eq!(store.blocks_in_tier(Tier::Tier1).len(), 2);
|
|
assert_eq!(store.blocks_in_tier(Tier::Tier0).len(), 0);
|
|
}
|
|
|
|
// -- Total bytes -------------------------------------------------------
|
|
|
|
#[test]
|
|
fn test_total_bytes() {
|
|
let mut store = TieredStore::new(4096);
|
|
assert_eq!(store.total_bytes(), 0);
|
|
|
|
let data = vec![1.0f32; 64];
|
|
store.put(make_key(1, 0), &data, Tier::Tier1, 0).unwrap();
|
|
let bytes_after_one = store.total_bytes();
|
|
assert!(bytes_after_one > 0);
|
|
|
|
store.put(make_key(2, 0), &data, Tier::Tier2, 0).unwrap();
|
|
assert!(store.total_bytes() > bytes_after_one);
|
|
}
|
|
|
|
#[test]
|
|
fn test_total_bytes_decreases_on_evict() {
|
|
let mut store = TieredStore::new(4096);
|
|
let data = vec![1.0f32; 64];
|
|
let key = make_key(1, 0);
|
|
|
|
store.put(key, &data, Tier::Tier1, 0).unwrap();
|
|
let before = store.total_bytes();
|
|
|
|
store.evict(key, ReconstructPolicy::None).unwrap();
|
|
assert_eq!(store.total_bytes(), before - before); // back to 0
|
|
}
|
|
|
|
// -- Touch / access stats ----------------------------------------------
|
|
|
|
#[test]
|
|
fn test_touch_updates_stats() {
|
|
let mut store = TieredStore::new(4096);
|
|
let key = make_key(1, 0);
|
|
store.put(key, &[1.0; 16], Tier::Tier1, 0).unwrap();
|
|
|
|
// Initial state after put.
|
|
let meta = store.meta(key).unwrap();
|
|
assert_eq!(meta.access_count, 1);
|
|
assert_eq!(meta.last_access_at, 0);
|
|
assert_eq!(meta.window, 1);
|
|
|
|
// Touch at tick 5.
|
|
store.touch(key, 5);
|
|
let meta = store.meta(key).unwrap();
|
|
assert_eq!(meta.access_count, 2);
|
|
assert_eq!(meta.last_access_at, 5);
|
|
// Window should have shifted left by 5 and gained bit 0.
|
|
assert_eq!(meta.window, (1u64 << 5) | 1);
|
|
assert!(meta.ema_rate > 0.0);
|
|
|
|
// Touch at tick 5 again (same tick).
|
|
store.touch(key, 5);
|
|
let meta = store.meta(key).unwrap();
|
|
assert_eq!(meta.access_count, 3);
|
|
// Window unchanged on same-tick touch.
|
|
assert_eq!(meta.window, (1u64 << 5) | 1);
|
|
}
|
|
|
|
#[test]
|
|
fn test_touch_window_overflow() {
|
|
let mut store = TieredStore::new(4096);
|
|
let key = make_key(1, 0);
|
|
store.put(key, &[1.0; 16], Tier::Tier1, 0).unwrap();
|
|
|
|
// Touch after more than 64 ticks clears the window entirely.
|
|
store.touch(key, 100);
|
|
let meta = store.meta(key).unwrap();
|
|
assert_eq!(meta.window, 1);
|
|
assert_eq!(meta.last_access_at, 100);
|
|
}
|
|
|
|
#[test]
|
|
fn test_touch_nonexistent_noop() {
|
|
let mut store = TieredStore::new(4096);
|
|
// Should not panic.
|
|
store.touch(make_key(42, 0), 10);
|
|
}
|
|
|
|
// -- Overwrite ---------------------------------------------------------
|
|
|
|
#[test]
|
|
fn test_put_overwrite() {
|
|
let mut store = TieredStore::new(4096);
|
|
let key = make_key(1, 0);
|
|
|
|
store.put(key, &[1.0; 16], Tier::Tier1, 0).unwrap();
|
|
assert_eq!(store.tier_count(Tier::Tier1), 1);
|
|
|
|
// Overwrite into a different tier.
|
|
store.put(key, &[2.0; 16], Tier::Tier3, 10).unwrap();
|
|
assert_eq!(store.block_count(), 1);
|
|
assert_eq!(store.tier_count(Tier::Tier1), 0);
|
|
assert_eq!(store.tier_count(Tier::Tier3), 1);
|
|
|
|
let meta = store.meta(key).unwrap();
|
|
assert_eq!(meta.tier, Tier::Tier3);
|
|
assert_eq!(meta.created_at, 10);
|
|
}
|
|
|
|
// -- Checksum ----------------------------------------------------------
|
|
|
|
#[test]
|
|
fn test_checksum_stored_correctly() {
|
|
let mut store = TieredStore::new(4096);
|
|
let key = make_key(1, 0);
|
|
let data: Vec<f32> = (0..32).map(|i| i as f32).collect();
|
|
|
|
store.put(key, &data, Tier::Tier1, 0).unwrap();
|
|
|
|
let meta = store.meta(key).unwrap();
|
|
assert_ne!(meta.checksum, 0);
|
|
|
|
// Manually verify the checksum matches.
|
|
let (packed, scale) = quantize_block(&data, 8);
|
|
let expected = block_checksum(&packed, scale);
|
|
assert_eq!(meta.checksum, expected);
|
|
}
|
|
|
|
// -- BlockIO trait ------------------------------------------------------
|
|
|
|
#[test]
|
|
fn test_block_io_write_read() {
|
|
let mut store = TieredStore::new(4096);
|
|
let key = make_key(1, 0);
|
|
let raw = vec![0xAA, 0xBB, 0xCC, 0xDD];
|
|
|
|
store.write_block(Tier::Tier1, key, &raw).unwrap();
|
|
|
|
let mut dst = vec![0u8; 8];
|
|
let n = store.read_block(Tier::Tier1, key, &mut dst).unwrap();
|
|
assert_eq!(n, 4);
|
|
assert_eq!(&dst[..4], &raw);
|
|
}
|
|
|
|
#[test]
|
|
fn test_block_io_delete() {
|
|
let mut store = TieredStore::new(4096);
|
|
let key = make_key(1, 0);
|
|
store.write_block(Tier::Tier2, key, &[1, 2, 3]).unwrap();
|
|
|
|
store.delete_block(Tier::Tier2, key).unwrap();
|
|
|
|
let mut dst = vec![0u8; 4];
|
|
assert_eq!(
|
|
store.read_block(Tier::Tier2, key, &mut dst),
|
|
Err(StoreError::BlockNotFound),
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn test_block_io_write_tier0_rejected() {
|
|
let mut store = TieredStore::new(4096);
|
|
let key = make_key(1, 0);
|
|
assert_eq!(
|
|
store.write_block(Tier::Tier0, key, &[1]),
|
|
Err(StoreError::InvalidBlock),
|
|
);
|
|
}
|
|
|
|
// -- MetaLog trait ------------------------------------------------------
|
|
|
|
#[test]
|
|
fn test_meta_log_append_get() {
|
|
let mut store = TieredStore::new(4096);
|
|
let key = make_key(1, 0);
|
|
let meta = BlockMeta {
|
|
key,
|
|
dtype: DType::F32,
|
|
tier: Tier::Tier1,
|
|
bits: 8,
|
|
scale: 0.5,
|
|
zero_point: 0,
|
|
created_at: 42,
|
|
last_access_at: 42,
|
|
access_count: 1,
|
|
ema_rate: 0.0,
|
|
window: 1,
|
|
checksum: 0,
|
|
reconstruct: ReconstructPolicy::None,
|
|
tier_age: 0,
|
|
lineage_parent: None,
|
|
block_bytes: 64,
|
|
};
|
|
|
|
MetaLog::append(&mut store, &meta).unwrap();
|
|
let retrieved = MetaLog::get(&store, key).unwrap();
|
|
assert_eq!(retrieved.key, key);
|
|
assert_eq!(retrieved.created_at, 42);
|
|
}
|
|
|
|
#[test]
|
|
fn test_meta_log_iter() {
|
|
let mut store = TieredStore::new(4096);
|
|
let data = vec![1.0f32; 8];
|
|
|
|
store.put(make_key(1, 0), &data, Tier::Tier1, 0).unwrap();
|
|
store.put(make_key(2, 0), &data, Tier::Tier2, 0).unwrap();
|
|
store.put(make_key(3, 0), &data, Tier::Tier3, 0).unwrap();
|
|
|
|
let entries: Vec<_> = MetaLog::iter(&store).collect();
|
|
assert_eq!(entries.len(), 3);
|
|
}
|
|
|
|
// -- bits_for_tier -----------------------------------------------------
|
|
|
|
#[test]
|
|
fn test_bits_for_tier() {
|
|
assert_eq!(bits_for_tier(Tier::Tier0), 0);
|
|
assert_eq!(bits_for_tier(Tier::Tier1), 8);
|
|
assert_eq!(bits_for_tier(Tier::Tier2), 7);
|
|
assert_eq!(bits_for_tier(Tier::Tier3), 3);
|
|
}
|
|
|
|
// -- Tier enum ---------------------------------------------------------
|
|
|
|
#[test]
|
|
fn test_tier_repr() {
|
|
assert_eq!(Tier::Tier0 as u8, 0);
|
|
assert_eq!(Tier::Tier1 as u8, 1);
|
|
assert_eq!(Tier::Tier2 as u8, 2);
|
|
assert_eq!(Tier::Tier3 as u8, 3);
|
|
}
|
|
|
|
#[test]
|
|
fn test_dtype_repr() {
|
|
assert_eq!(DType::F32 as u8, 0);
|
|
assert_eq!(DType::F16 as u8, 1);
|
|
assert_eq!(DType::BF16 as u8, 2);
|
|
}
|
|
|
|
#[test]
|
|
fn test_reconstruct_policy_repr() {
|
|
assert_eq!(ReconstructPolicy::None as u8, 0);
|
|
assert_eq!(ReconstructPolicy::Delta as u8, 1);
|
|
assert_eq!(ReconstructPolicy::Factor as u8, 2);
|
|
}
|
|
|
|
// -- Integration: multi-block workflow ---------------------------------
|
|
|
|
#[test]
|
|
fn test_multi_block_workflow() {
|
|
let mut store = TieredStore::new(4096);
|
|
|
|
// Insert 10 blocks across tiers.
|
|
for i in 0..10u32 {
|
|
let key = make_key(1, i);
|
|
let data: Vec<f32> = (0..32).map(|j| (i * 32 + j) as f32 * 0.1).collect();
|
|
let tier = match i % 3 {
|
|
0 => Tier::Tier1,
|
|
1 => Tier::Tier2,
|
|
_ => Tier::Tier3,
|
|
};
|
|
store.put(key, &data, tier, i as u64).unwrap();
|
|
}
|
|
|
|
assert_eq!(store.block_count(), 10);
|
|
assert_eq!(store.tier_count(Tier::Tier1), 4); // 0,3,6,9
|
|
assert_eq!(store.tier_count(Tier::Tier2), 3); // 1,4,7
|
|
assert_eq!(store.tier_count(Tier::Tier3), 3); // 2,5,8
|
|
|
|
// Touch some blocks.
|
|
store.touch(make_key(1, 0), 20);
|
|
store.touch(make_key(1, 5), 25);
|
|
|
|
// Evict a cold block.
|
|
store
|
|
.evict(make_key(1, 8), ReconstructPolicy::Delta)
|
|
.unwrap();
|
|
assert_eq!(store.tier_count(Tier::Tier3), 2);
|
|
assert_eq!(store.tier_count(Tier::Tier0), 1);
|
|
assert_eq!(store.block_count(), 10); // metadata preserved
|
|
|
|
// Read back a hot block.
|
|
let mut out = vec![0.0f32; 32];
|
|
let n = TieredStore::get(&mut store, make_key(1, 0), &mut out, 30).unwrap();
|
|
assert_eq!(n, 32);
|
|
}
|
|
|
|
// -- tick / score / touch_block -----------------------------------------
|
|
|
|
#[test]
|
|
fn test_tick_empty_store() {
|
|
let mut store = TieredStore::new(4096);
|
|
let config = crate::tiering::TierConfig::default();
|
|
let result = store.tick(&config, 100, 1_000_000, 100);
|
|
assert_eq!(result.upgrades, 0);
|
|
assert_eq!(result.downgrades, 0);
|
|
assert_eq!(result.evictions, 0);
|
|
assert_eq!(result.bytes_freed, 0);
|
|
assert_eq!(result.ops_used, 0);
|
|
assert_eq!(result.candidates_found, 0);
|
|
}
|
|
|
|
#[test]
|
|
fn test_tick_migrates_cold_to_hot() {
|
|
let mut store = TieredStore::new(4096);
|
|
let key = make_key(1, 0);
|
|
let data: Vec<f32> = (0..64).map(|i| i as f32 * 0.1).collect();
|
|
|
|
// Put block in Tier3 (cold).
|
|
store.put(key, &data, Tier::Tier3, 0).unwrap();
|
|
assert_eq!(store.tier_count(Tier::Tier3), 1);
|
|
|
|
// Simulate a highly-accessed block by directly setting metadata
|
|
// fields so that the tiering score exceeds t1 + hysteresis.
|
|
if let Some(meta) = store.index.get_mut(&key) {
|
|
meta.ema_rate = 1.0;
|
|
meta.window = u64::MAX; // all 64 bits set
|
|
meta.last_access_at = 100;
|
|
meta.access_count = 100;
|
|
meta.tier_age = 10; // past default min_residency (5)
|
|
}
|
|
|
|
let config = crate::tiering::TierConfig::default();
|
|
let result = store.tick(&config, 100, 1_000_000, 100);
|
|
|
|
assert!(
|
|
result.upgrades > 0,
|
|
"expected at least one upgrade, got {}",
|
|
result.upgrades
|
|
);
|
|
assert_eq!(result.downgrades, 0);
|
|
assert!(result.candidates_found > 0);
|
|
|
|
let meta = store.meta(key).unwrap();
|
|
assert_eq!(
|
|
meta.tier,
|
|
Tier::Tier1,
|
|
"block should be in Tier1 after upgrade"
|
|
);
|
|
assert_eq!(meta.bits, 8, "Tier1 should use 8-bit quantization");
|
|
assert_eq!(meta.tier_age, 0, "tier_age should reset after migration");
|
|
|
|
// The block should still be readable.
|
|
let mut out = vec![0.0f32; 64];
|
|
let n = TieredStore::get(&mut store, key, &mut out, 101).unwrap();
|
|
assert_eq!(n, 64);
|
|
}
|
|
|
|
#[test]
|
|
fn test_tick_respects_budget_ops() {
|
|
let mut store = TieredStore::new(4096);
|
|
let data: Vec<f32> = (0..64).map(|i| i as f32 * 0.1).collect();
|
|
|
|
// Create 5 blocks in Tier3, all hot enough to warrant migration.
|
|
for i in 0..5u32 {
|
|
let key = make_key(i as u128 + 1, 0);
|
|
store.put(key, &data, Tier::Tier3, 0).unwrap();
|
|
if let Some(meta) = store.index.get_mut(&key) {
|
|
meta.ema_rate = 1.0;
|
|
meta.window = u64::MAX;
|
|
meta.last_access_at = 100;
|
|
meta.access_count = 100;
|
|
meta.tier_age = 10;
|
|
}
|
|
}
|
|
|
|
let config = crate::tiering::TierConfig::default();
|
|
// Budget only 2 ops.
|
|
let result = store.tick(&config, 100, 1_000_000, 2);
|
|
|
|
assert_eq!(result.ops_used, 2, "should use exactly 2 ops");
|
|
assert_eq!(result.upgrades, 2, "should upgrade only 2 blocks");
|
|
assert!(result.candidates_found >= 5, "should find all 5 candidates");
|
|
}
|
|
|
|
#[test]
|
|
fn test_touch_block_updates_ema_and_window() {
|
|
let mut store = TieredStore::new(4096);
|
|
let key = make_key(1, 0);
|
|
store.put(key, &[1.0; 16], Tier::Tier1, 0).unwrap();
|
|
|
|
let config = crate::tiering::TierConfig::default();
|
|
|
|
// Initial state: ema_rate is 0 after put.
|
|
let meta = store.meta(key).unwrap();
|
|
assert_eq!(meta.ema_rate, 0.0);
|
|
|
|
// Touch at tick 5.
|
|
store.touch_block(key, &config, 5);
|
|
let meta = store.meta(key).unwrap();
|
|
|
|
// tiering::touch sets ema_rate = alpha + (1 - alpha) * old_ema
|
|
// = 0.3 + 0.7 * 0.0 = 0.3
|
|
assert!(
|
|
(meta.ema_rate - config.alpha).abs() < 1e-6,
|
|
"ema_rate={}, expected={}",
|
|
meta.ema_rate,
|
|
config.alpha,
|
|
);
|
|
assert_eq!(meta.last_access_at, 5);
|
|
// Window should have bit 0 set after touch.
|
|
assert_ne!(meta.window & 1, 0, "bit 0 should be set");
|
|
// Elapsed = 5 ticks from 0, so window = (initial << 5) | 1.
|
|
// Initial window from put is 1, so: (1 << 5) | 1 = 0b100001.
|
|
assert_eq!(meta.window, (1u64 << 5) | 1);
|
|
}
|
|
|
|
#[test]
|
|
fn test_score_block_none_for_missing() {
|
|
let store = TieredStore::new(4096);
|
|
let config = crate::tiering::TierConfig::default();
|
|
let result = store.score_block(make_key(99, 0), &config, 100);
|
|
assert_eq!(result, None);
|
|
}
|
|
|
|
// -----------------------------------------------------------------------
|
|
// Coherence integration
|
|
// -----------------------------------------------------------------------
|
|
|
|
#[test]
|
|
fn test_epoch_tracker_wired_into_put() {
|
|
let mut store = TieredStore::new(4096);
|
|
let key = BlockKey {
|
|
tensor_id: 1,
|
|
block_index: 0,
|
|
};
|
|
let data = vec![1.0f32; 64];
|
|
|
|
assert_eq!(store.epoch_tracker().check_epoch(key), None);
|
|
|
|
store.put(key, &data, Tier::Tier1, 0).unwrap();
|
|
assert!(store.epoch_tracker().check_epoch(key).is_some());
|
|
|
|
let epoch1 = store.epoch_tracker().check_epoch(key).unwrap();
|
|
store.put(key, &data, Tier::Tier1, 1).unwrap();
|
|
let epoch2 = store.epoch_tracker().check_epoch(key).unwrap();
|
|
assert!(epoch2 > epoch1, "epoch should increment on overwrite");
|
|
}
|
|
|
|
#[test]
|
|
fn test_coherence_disabled_by_default() {
|
|
let mut store = TieredStore::new(4096);
|
|
let key = BlockKey {
|
|
tensor_id: 1,
|
|
block_index: 0,
|
|
};
|
|
let data = vec![1.0f32; 64];
|
|
store.put(key, &data, Tier::Tier1, 0).unwrap();
|
|
|
|
assert!(store.coherence_check(key, &data, 1).is_none());
|
|
}
|
|
|
|
#[test]
|
|
fn test_coherence_enabled_passes() {
|
|
let mut store = TieredStore::new(4096);
|
|
store.enable_coherence(crate::coherence::CoherenceCheck::default());
|
|
|
|
let key = BlockKey {
|
|
tensor_id: 1,
|
|
block_index: 0,
|
|
};
|
|
let data: Vec<f32> = (0..64).map(|i| (i as f32 + 1.0) * 0.25).collect();
|
|
store.put(key, &data, Tier::Tier1, 0).unwrap();
|
|
|
|
let result = store.coherence_check(key, &data, 1).unwrap().unwrap();
|
|
assert!(
|
|
result.passed,
|
|
"Tier1 coherence should pass; err={}",
|
|
result.max_error
|
|
);
|
|
}
|
|
|
|
// -----------------------------------------------------------------------
|
|
// MetricsSeries integration
|
|
// -----------------------------------------------------------------------
|
|
|
|
#[test]
|
|
fn test_metrics_series_wired_into_tick() {
|
|
use crate::tiering::TierConfig;
|
|
|
|
let mut store = TieredStore::new(4096);
|
|
let config = TierConfig::default();
|
|
|
|
// Put a few blocks.
|
|
for i in 0..5u128 {
|
|
let key = BlockKey {
|
|
tensor_id: i,
|
|
block_index: 0,
|
|
};
|
|
store.put(key, &vec![1.0f32; 64], Tier::Tier1, 0).unwrap();
|
|
}
|
|
|
|
assert!(store.metrics_series().is_empty());
|
|
|
|
// Run a tick -- should auto-record a metrics snapshot.
|
|
store.tick(&config, 100, 1_000_000, 100);
|
|
assert_eq!(store.metrics_series().len(), 1);
|
|
|
|
// Run another tick.
|
|
store.tick(&config, 200, 1_000_000, 100);
|
|
assert_eq!(store.metrics_series().len(), 2);
|
|
|
|
// Latest snapshot should reflect current state.
|
|
let (ts, m) = store.metrics_series().latest().unwrap();
|
|
assert_eq!(*ts, 200);
|
|
assert_eq!(m.total_blocks, 5);
|
|
}
|
|
|
|
// -----------------------------------------------------------------------
|
|
// Benchmarks
|
|
// -----------------------------------------------------------------------
|
|
//
|
|
// Run with: cargo test bench_ -- --nocapture
|
|
// These use std::time::Instant and std::hint::black_box for stable timing.
|
|
|
|
#[test]
|
|
fn bench_batch_scoring_10k() {
|
|
use crate::tiering::{
|
|
compute_score, compute_scores_batch, BlockMeta as TBlockMeta, Tier as TTier, TierConfig,
|
|
};
|
|
use std::time::Instant;
|
|
|
|
let cfg = TierConfig::default();
|
|
let metas: Vec<TBlockMeta> = (0..10_000)
|
|
.map(|i| TBlockMeta {
|
|
ema_rate: (i as f32) * 0.0001,
|
|
access_window: 0x5555_5555_5555_5555,
|
|
last_access: 50 + (i as u64 % 100),
|
|
access_count: i as u64,
|
|
current_tier: TTier::Tier1,
|
|
tier_since: 0,
|
|
})
|
|
.collect();
|
|
|
|
let iters = 1000;
|
|
|
|
// Individual scoring
|
|
let start = Instant::now();
|
|
for _ in 0..iters {
|
|
for m in &metas {
|
|
std::hint::black_box(compute_score(&cfg, 100, m));
|
|
}
|
|
}
|
|
let individual = start.elapsed();
|
|
|
|
// Batch scoring
|
|
let start = Instant::now();
|
|
for _ in 0..iters {
|
|
std::hint::black_box(compute_scores_batch(&cfg, 100, &metas));
|
|
}
|
|
let batch = start.elapsed();
|
|
|
|
eprintln!(
|
|
"Individual scoring 10k x {iters}: {:?} ({:.0} ns/block)",
|
|
individual,
|
|
individual.as_nanos() as f64 / (iters * 10_000) as f64
|
|
);
|
|
eprintln!(
|
|
"Batch scoring 10k x {iters}: {:?} ({:.0} ns/block)",
|
|
batch,
|
|
batch.as_nanos() as f64 / (iters * 10_000) as f64
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn bench_dequant_5bit_4096() {
|
|
use std::time::Instant;
|
|
|
|
let data: Vec<f32> = (0..4096).map(|i| (i as f32 - 2048.0) * 0.01).collect();
|
|
let (packed, scale) = quantize_block(&data, 5);
|
|
let mut out = vec![0.0f32; 4096];
|
|
|
|
let iters = 10_000;
|
|
let start = Instant::now();
|
|
for _ in 0..iters {
|
|
std::hint::black_box(dequantize_block(&packed, scale, 5, 4096, &mut out));
|
|
}
|
|
let elapsed = start.elapsed();
|
|
|
|
let total_bytes = 4096u64 * 4 * iters as u64;
|
|
let gbs = total_bytes as f64 / elapsed.as_secs_f64() / 1e9;
|
|
eprintln!(
|
|
"Dequant 5-bit 4096 x {iters}: {:?} ({:.2} GB/s output throughput)",
|
|
elapsed, gbs
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn bench_dequant_7bit_4096() {
|
|
use std::time::Instant;
|
|
|
|
let data: Vec<f32> = (0..4096).map(|i| (i as f32 - 2048.0) * 0.01).collect();
|
|
let (packed, scale) = quantize_block(&data, 7);
|
|
let mut out = vec![0.0f32; 4096];
|
|
|
|
let iters = 10_000;
|
|
let start = Instant::now();
|
|
for _ in 0..iters {
|
|
std::hint::black_box(dequantize_block(&packed, scale, 7, 4096, &mut out));
|
|
}
|
|
let elapsed = start.elapsed();
|
|
|
|
let total_bytes = 4096u64 * 4 * iters as u64;
|
|
let gbs = total_bytes as f64 / elapsed.as_secs_f64() / 1e9;
|
|
eprintln!(
|
|
"Dequant 7-bit 4096 x {iters}: {:?} ({:.2} GB/s output throughput)",
|
|
elapsed, gbs
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn bench_quant_5bit_4096() {
|
|
use std::time::Instant;
|
|
|
|
let data: Vec<f32> = (0..4096).map(|i| (i as f32 - 2048.0) * 0.01).collect();
|
|
|
|
let iters = 10_000;
|
|
let start = Instant::now();
|
|
for _ in 0..iters {
|
|
std::hint::black_box(quantize_block(&data, 5));
|
|
}
|
|
let elapsed = start.elapsed();
|
|
|
|
let total_bytes = 4096u64 * 4 * iters as u64;
|
|
let gbs = total_bytes as f64 / elapsed.as_secs_f64() / 1e9;
|
|
eprintln!(
|
|
"Quant 5-bit 4096 x {iters}: {:?} ({:.2} GB/s input throughput)",
|
|
elapsed, gbs
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn bench_svd_adaptive_64x64() {
|
|
use crate::delta::FactorSet;
|
|
use std::time::Instant;
|
|
|
|
let (rows, cols) = (64, 64);
|
|
let data: Vec<f32> = (0..rows * cols)
|
|
.map(|i| (i as f32 * 0.37).sin() + (i as f32 * 0.73).cos())
|
|
.collect();
|
|
|
|
let iters = 100;
|
|
let start = Instant::now();
|
|
for _ in 0..iters {
|
|
std::hint::black_box(FactorSet::from_data_adaptive(&data, rows, cols, 16, 0.05));
|
|
}
|
|
let elapsed = start.elapsed();
|
|
|
|
eprintln!(
|
|
"SVD adaptive 64x64 (max_rank=16, target=0.05) x {iters}: {:?} ({:.2} ms/iter)",
|
|
elapsed,
|
|
elapsed.as_secs_f64() * 1000.0 / iters as f64
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn bench_format_report() {
|
|
use crate::metrics::StoreMetrics;
|
|
use std::time::Instant;
|
|
|
|
let m = StoreMetrics {
|
|
total_blocks: 10_000,
|
|
tier0_blocks: 500,
|
|
tier1_blocks: 4000,
|
|
tier2_blocks: 3500,
|
|
tier3_blocks: 2000,
|
|
tier1_bytes: 4_000_000,
|
|
tier2_bytes: 2_500_000,
|
|
tier3_bytes: 750_000,
|
|
total_reads: 1_000_000,
|
|
total_writes: 500_000,
|
|
total_evictions: 5000,
|
|
total_upgrades: 12_000,
|
|
total_downgrades: 8000,
|
|
total_reconstructions: 200,
|
|
total_checksum_failures: 0,
|
|
total_compactions: 150,
|
|
tier_flips_last_minute: 0.023,
|
|
avg_score_tier1: 0.85,
|
|
avg_score_tier2: 0.45,
|
|
avg_score_tier3: 0.12,
|
|
};
|
|
|
|
let iters = 10_000;
|
|
let start = Instant::now();
|
|
for _ in 0..iters {
|
|
std::hint::black_box(m.format_report());
|
|
}
|
|
let elapsed = start.elapsed();
|
|
|
|
eprintln!(
|
|
"format_report x {iters}: {:?} ({:.0} ns/call)",
|
|
elapsed,
|
|
elapsed.as_nanos() as f64 / iters as f64
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn bench_format_json() {
|
|
use crate::metrics::StoreMetrics;
|
|
use std::time::Instant;
|
|
|
|
let m = StoreMetrics {
|
|
total_blocks: 10_000,
|
|
tier0_blocks: 500,
|
|
tier1_blocks: 4000,
|
|
tier2_blocks: 3500,
|
|
tier3_blocks: 2000,
|
|
tier1_bytes: 4_000_000,
|
|
tier2_bytes: 2_500_000,
|
|
tier3_bytes: 750_000,
|
|
total_reads: 1_000_000,
|
|
total_writes: 500_000,
|
|
total_evictions: 5000,
|
|
total_upgrades: 12_000,
|
|
total_downgrades: 8000,
|
|
total_reconstructions: 200,
|
|
total_checksum_failures: 0,
|
|
total_compactions: 150,
|
|
tier_flips_last_minute: 0.023,
|
|
avg_score_tier1: 0.85,
|
|
avg_score_tier2: 0.45,
|
|
avg_score_tier3: 0.12,
|
|
};
|
|
|
|
let iters = 10_000;
|
|
let start = Instant::now();
|
|
for _ in 0..iters {
|
|
std::hint::black_box(m.format_json());
|
|
}
|
|
let elapsed = start.elapsed();
|
|
|
|
eprintln!(
|
|
"format_json x {iters}: {:?} ({:.0} ns/call)",
|
|
elapsed,
|
|
elapsed.as_nanos() as f64 / iters as f64
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn bench_metrics_series_trend_100() {
|
|
use crate::metrics::{MetricsSeries, StoreMetrics};
|
|
use std::time::Instant;
|
|
|
|
let mut series = MetricsSeries::new(256);
|
|
for i in 0..100u64 {
|
|
series.record(
|
|
i,
|
|
StoreMetrics {
|
|
total_blocks: 1000 + i,
|
|
tier1_blocks: 400 + i % 50,
|
|
tier2_blocks: 350,
|
|
tier3_blocks: 250,
|
|
tier1_bytes: 400_000 + i * 100,
|
|
tier2_bytes: 250_000,
|
|
tier3_bytes: 75_000,
|
|
total_evictions: i * 3,
|
|
..Default::default()
|
|
},
|
|
);
|
|
}
|
|
|
|
let iters = 10_000;
|
|
let start = Instant::now();
|
|
for _ in 0..iters {
|
|
std::hint::black_box(series.trend());
|
|
}
|
|
let elapsed = start.elapsed();
|
|
|
|
eprintln!(
|
|
"MetricsSeries trend (100 snapshots) x {iters}: {:?} ({:.0} ns/call)",
|
|
elapsed,
|
|
elapsed.as_nanos() as f64 / iters as f64
|
|
);
|
|
}
|
|
}
|