[Rust TreeCore] Support external cache linker (#37306)
This commit is contained in:
@@ -427,6 +427,25 @@ impl<K: ChildKeyType> TreeComponent<K> for FullComponent {
|
||||
})
|
||||
}
|
||||
|
||||
fn build_external_linker_offload_transfer(
|
||||
&self,
|
||||
tree_core: &UnifiedTreeCore<K>,
|
||||
node_id: NodeIdx_,
|
||||
) -> Option<PoolTransfer> {
|
||||
let node = tree_core.arena.node(node_id);
|
||||
let keys = node
|
||||
.hash_value
|
||||
.as_ref()
|
||||
.filter(|hashes| !hashes.is_empty())?;
|
||||
let device_indices = node.try_device_value(FULL)?;
|
||||
Some(PoolTransfer {
|
||||
name: PoolName::Kv,
|
||||
device_indices: Some(device_indices.to_kind(Kind::Int64)),
|
||||
keys: Some(keys.clone()),
|
||||
..Default::default()
|
||||
})
|
||||
}
|
||||
|
||||
fn commit_hicache_transfer(
|
||||
&self,
|
||||
tree_core: &mut UnifiedTreeCore<K>,
|
||||
|
||||
@@ -392,6 +392,15 @@ pub trait TreeComponent<K: ChildKeyType> {
|
||||
unimplemented!("TreeComponent.build_hicache_transfers")
|
||||
}
|
||||
|
||||
/// Build this component's direct device-to-external-store transfer for a node.
|
||||
fn build_external_linker_offload_transfer(
|
||||
&self,
|
||||
_tree_core: &UnifiedTreeCore<K>,
|
||||
_node_id: NodeIdx_,
|
||||
) -> Option<PoolTransfer> {
|
||||
None
|
||||
}
|
||||
|
||||
/// Post-transfer bookkeeping: store host indices, update LRU, etc.
|
||||
fn commit_hicache_transfer(
|
||||
&self,
|
||||
|
||||
@@ -951,6 +951,35 @@ impl<K: ChildKeyType> TreeComponent<K> for SwaComponent {
|
||||
})
|
||||
}
|
||||
|
||||
fn build_external_linker_offload_transfer(
|
||||
&self,
|
||||
tree_core: &UnifiedTreeCore<K>,
|
||||
node_id: NodeIdx_,
|
||||
) -> Option<PoolTransfer> {
|
||||
let node = tree_core.arena.node(node_id);
|
||||
let hashes = node
|
||||
.hash_value
|
||||
.as_ref()
|
||||
.filter(|hashes| !hashes.is_empty())?;
|
||||
let value = node.try_device_value(SWA)?;
|
||||
let num_pages = value.size()[0] as usize / tree_core.page_size;
|
||||
if num_pages == 0 {
|
||||
return None;
|
||||
}
|
||||
let num_tokens = num_pages * tree_core.page_size;
|
||||
Some(PoolTransfer {
|
||||
name: PoolName::Swa,
|
||||
device_indices: Some(
|
||||
value
|
||||
.narrow(0, value.size()[0] - num_tokens as i64, num_tokens as i64)
|
||||
.to_kind(Kind::Int64),
|
||||
),
|
||||
keys: Some(hashes[hashes.len().saturating_sub(num_pages)..].to_vec()),
|
||||
hit_policy: PoolHitPolicy::TrailingPages,
|
||||
..Default::default()
|
||||
})
|
||||
}
|
||||
|
||||
fn commit_hicache_transfer(
|
||||
&self,
|
||||
tree_core: &mut UnifiedTreeCore<K>,
|
||||
|
||||
@@ -175,6 +175,8 @@ pub struct Node<K: ChildKeyType> {
|
||||
/// Per-page hash chain; None when the node was never hashed.
|
||||
/// TODO: Store raw digests and hex-encode only at the Python or storage boundary.
|
||||
pub hash_value: Option<Vec<String>>,
|
||||
/// Whether this node is available through the direct external-cache linker.
|
||||
pub external_cache_stored: bool,
|
||||
/// The in-flight write-through backup's ack id.
|
||||
pub write_through_pending_id: Option<usize>,
|
||||
/// Load-back anchor currently reading this node's host slots.
|
||||
@@ -396,6 +398,7 @@ impl<K: ChildKeyType> Node<K> {
|
||||
swa_uuid: None,
|
||||
swa_host_uuid: None,
|
||||
hash_value: Some(Vec::new()),
|
||||
external_cache_stored: false,
|
||||
write_through_pending_id: None,
|
||||
load_back_pending_id: None,
|
||||
last_access_counter: 0,
|
||||
@@ -418,6 +421,7 @@ impl<K: ChildKeyType> Node<K> {
|
||||
swa_uuid: None,
|
||||
swa_host_uuid: None,
|
||||
hash_value: None,
|
||||
external_cache_stored: false,
|
||||
write_through_pending_id: None,
|
||||
load_back_pending_id: None,
|
||||
last_access_counter: 0,
|
||||
@@ -753,6 +757,24 @@ pub enum TreeCoreRuntimeError {
|
||||
#[cfg(any(test, feature = "inspection"))]
|
||||
#[error("{0}")]
|
||||
InspectionAssertion(String),
|
||||
/// Direct external-cache linking does not support this tree component.
|
||||
#[error("external cache linker does not support component {component_type:?}")]
|
||||
ExternalCacheLinkerUnsupportedComponent { component_type: ComponentType },
|
||||
/// The existing device anchor must be on the restored endpoint's root path.
|
||||
#[error("node {until_node_id} is not an ancestor of node {from_node_id}")]
|
||||
ExternalCachePathNotAncestor {
|
||||
from_node_id: NodeId,
|
||||
until_node_id: NodeId,
|
||||
},
|
||||
/// External offload lifecycle calls must observe valid state transitions.
|
||||
#[error(
|
||||
"invalid external offload state for node {node_id}: stored={stored}, pending={pending_id:?}"
|
||||
)]
|
||||
InvalidExternalCacheOffloadState {
|
||||
node_id: NodeId,
|
||||
stored: bool,
|
||||
pending_id: Option<NodeId>,
|
||||
},
|
||||
}
|
||||
|
||||
// Unigram and bigram child keys.
|
||||
|
||||
@@ -1775,6 +1775,17 @@ impl<K: ChildKeyType + Send + Sync> TreeCoreBinding<K> {
|
||||
py.allow_threads(|| self.core().enable_storage)
|
||||
}
|
||||
|
||||
/// Enable or disable the direct external-cache linker.
|
||||
fn set_enable_external_cache_linker(&self, py: Python<'_>, value: bool) -> PyResult<()> {
|
||||
py.allow_threads(|| self.core().set_enable_external_cache_linker(value))
|
||||
.map_err(tree_core_assertion_error)
|
||||
}
|
||||
|
||||
/// Whether the direct external-cache linker is wired.
|
||||
fn enable_external_cache_linker(&self, py: Python<'_>) -> bool {
|
||||
py.allow_threads(|| self.core().enable_external_cache_linker)
|
||||
}
|
||||
|
||||
/// Queue the all-cleared placement event.
|
||||
fn record_all_cleared_event(&self, py: Python<'_>) {
|
||||
py.allow_threads(|| self.core().record_all_cleared_event());
|
||||
@@ -1872,6 +1883,64 @@ impl<K: ChildKeyType + Send + Sync> TreeCoreBinding<K> {
|
||||
.map_err(node_access_error)
|
||||
}
|
||||
|
||||
/// Build transfers for a node with no stored or pending external copy.
|
||||
fn build_external_linker_offload_transfers(
|
||||
&self,
|
||||
py: Python<'_>,
|
||||
node_id: NodeId,
|
||||
) -> PyResult<Option<Vec<Py<PyAny>>>> {
|
||||
let transfers = py
|
||||
.allow_threads(|| self.core().build_external_linker_offload_transfers(node_id))
|
||||
.map_err(node_access_error)?;
|
||||
transfers
|
||||
.map(|transfers| {
|
||||
transfers
|
||||
.into_iter()
|
||||
.map(|transfer| transfer_to_py(py, transfer))
|
||||
.collect::<PyResult<Vec<_>>>()
|
||||
})
|
||||
.transpose()
|
||||
}
|
||||
|
||||
/// Mark an externally restored path, excluding its existing anchor.
|
||||
fn mark_external_cache_stored_path(
|
||||
&self,
|
||||
py: Python<'_>,
|
||||
from_node_id: NodeId,
|
||||
until_node_id: NodeId,
|
||||
) -> PyResult<()> {
|
||||
py.allow_threads(|| {
|
||||
self.core()
|
||||
.mark_external_cache_stored_path(from_node_id, until_node_id)
|
||||
})
|
||||
.map_err(tree_core_runtime_error)
|
||||
}
|
||||
|
||||
/// Publish an accepted external offload as pending.
|
||||
fn mark_external_linker_offload_pending(
|
||||
&self,
|
||||
py: Python<'_>,
|
||||
node_id: NodeId,
|
||||
) -> PyResult<()> {
|
||||
py.allow_threads(|| self.core().mark_external_linker_offload_pending(node_id))
|
||||
.map_err(tree_core_assertion_error)
|
||||
}
|
||||
|
||||
/// Finalize external-store state for an offload and its split fragments.
|
||||
fn finish_external_linker_offload(
|
||||
&self,
|
||||
py: Python<'_>,
|
||||
node_ids: Vec<NodeId>,
|
||||
ack_id: NodeId,
|
||||
success: bool,
|
||||
) -> PyResult<()> {
|
||||
py.allow_threads(|| {
|
||||
self.core()
|
||||
.finish_external_linker_offload(&node_ids, ack_id, success)
|
||||
})
|
||||
.map_err(tree_core_assertion_error)
|
||||
}
|
||||
|
||||
/// Order-sensitive digest of reclaimed coexisting host values.
|
||||
fn write_back_coexist_reclaim_digest(&self, py: Python<'_>) -> i64 {
|
||||
py.allow_threads(|| self.core().write_back_coexist_reclaim_digest)
|
||||
@@ -1968,6 +2037,11 @@ impl<K: ChildKeyType + Send + Sync> TreeCoreBinding<K> {
|
||||
.map_err(node_access_error)
|
||||
}
|
||||
|
||||
fn inspect_is_external_cache_stored(&self, py: Python<'_>, node_id: NodeId) -> PyResult<bool> {
|
||||
py.allow_threads(|| self.core().inspect_is_external_cache_stored(node_id))
|
||||
.map_err(node_access_error)
|
||||
}
|
||||
|
||||
fn inspect_is_node_in_device_lru(
|
||||
&self,
|
||||
py: Python<'_>,
|
||||
@@ -2839,6 +2913,20 @@ macro_rules! tree_core_binding {
|
||||
self.inner.enable_storage(py)
|
||||
}
|
||||
|
||||
/// Enable or disable the direct external-cache linker.
|
||||
fn set_enable_external_cache_linker(
|
||||
&self,
|
||||
py: Python<'_>,
|
||||
value: bool,
|
||||
) -> PyResult<()> {
|
||||
self.inner.set_enable_external_cache_linker(py, value)
|
||||
}
|
||||
|
||||
/// Whether the direct external-cache linker is wired.
|
||||
fn enable_external_cache_linker(&self, py: Python<'_>) -> bool {
|
||||
self.inner.enable_external_cache_linker(py)
|
||||
}
|
||||
|
||||
/// Queue the all-cleared placement event.
|
||||
fn record_all_cleared_event(&self, py: Python<'_>) {
|
||||
self.inner.record_all_cleared_event(py)
|
||||
@@ -2884,6 +2972,53 @@ macro_rules! tree_core_binding {
|
||||
self.inner.finish_load_back(py, anchor_node_id)
|
||||
}
|
||||
|
||||
/// Build transfers for a node with no stored or pending external copy.
|
||||
fn build_external_linker_offload_transfers(
|
||||
&self,
|
||||
py: Python<'_>,
|
||||
node_id: NodeId,
|
||||
) -> PyResult<Option<Vec<Py<PyAny>>>> {
|
||||
self.inner
|
||||
.build_external_linker_offload_transfers(py, node_id)
|
||||
}
|
||||
|
||||
/// Mark an externally restored path, excluding its existing anchor.
|
||||
fn mark_external_cache_stored_path(
|
||||
&self,
|
||||
py: Python<'_>,
|
||||
from_node_id: NodeId,
|
||||
until_node_id: NodeId,
|
||||
) -> PyResult<()> {
|
||||
self.inner.mark_external_cache_stored_path(
|
||||
py,
|
||||
from_node_id,
|
||||
until_node_id,
|
||||
)
|
||||
}
|
||||
|
||||
/// Publish an accepted external offload as pending.
|
||||
fn mark_external_linker_offload_pending(
|
||||
&self,
|
||||
py: Python<'_>,
|
||||
node_id: NodeId,
|
||||
) -> PyResult<()> {
|
||||
self.inner
|
||||
.mark_external_linker_offload_pending(py, node_id)
|
||||
}
|
||||
|
||||
/// Finalize external-store state for an offload and its split fragments.
|
||||
fn finish_external_linker_offload(
|
||||
&self,
|
||||
py: Python<'_>,
|
||||
node_ids: Vec<NodeId>,
|
||||
ack_id: NodeId,
|
||||
success: bool,
|
||||
) -> PyResult<()> {
|
||||
self.inner.finish_external_linker_offload(
|
||||
py, node_ids, ack_id, success,
|
||||
)
|
||||
}
|
||||
|
||||
/// Order-sensitive digest of reclaimed coexisting host values.
|
||||
#[pyo3(name = "write_back_duplicate_reclaim_digest")]
|
||||
fn write_back_coexist_reclaim_digest(&self, py: Python<'_>) -> i64 {
|
||||
@@ -2994,6 +3129,15 @@ macro_rules! tree_core_binding {
|
||||
.inspect_get_write_through_pending_id(py, node_id)
|
||||
}
|
||||
|
||||
#[cfg(feature = "inspection")]
|
||||
fn inspect_is_external_cache_stored(
|
||||
&self,
|
||||
py: Python<'_>,
|
||||
node_id: NodeId,
|
||||
) -> PyResult<bool> {
|
||||
self.inner.inspect_is_external_cache_stored(py, node_id)
|
||||
}
|
||||
|
||||
#[cfg(feature = "inspection")]
|
||||
fn inspect_is_node_in_device_lru(
|
||||
&self,
|
||||
|
||||
@@ -2201,6 +2201,248 @@ fn insert_threshold_crossing_emits_the_backup_kv_action() {
|
||||
assert_eq!(backups, vec![vec![leaf]]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn external_linker_hashes_new_nodes_and_triggers_offload_action() {
|
||||
let params = CacheInitParams {
|
||||
page_size: 2,
|
||||
write_through_threshold: 1,
|
||||
..Default::default()
|
||||
};
|
||||
let mut tc: UnifiedTreeCore<Vec<i64>> = UnifiedTreeCore::new(params, vec![FULL]);
|
||||
tc.set_enable_external_cache_linker(true).unwrap();
|
||||
|
||||
let result = tc.insert(&insert_params(&vec![1, 2], &[10, 11]));
|
||||
let leaf = result.last_device_node_id.unwrap();
|
||||
assert!(result.cache_actions.iter().any(|action| {
|
||||
matches!(action, CacheAction::BackupKV(backup) if backup.node_ids == vec![leaf])
|
||||
}));
|
||||
|
||||
let transfers = tc
|
||||
.build_external_linker_offload_transfers(leaf)
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
assert_eq!(transfers.len(), 1);
|
||||
assert_eq!(transfers[0].name, PoolName::Kv);
|
||||
assert_eq!(transfers[0].hit_policy, PoolHitPolicy::AllPages);
|
||||
assert!(
|
||||
transfers[0]
|
||||
.device_indices
|
||||
.as_ref()
|
||||
.unwrap()
|
||||
.equal(&Tensor::from_slice(&[10i64, 11]))
|
||||
);
|
||||
assert_eq!(
|
||||
transfers[0].keys,
|
||||
tc.arena
|
||||
.node(tc.arena.resolve(leaf).expect("live test node"))
|
||||
.hash_value
|
||||
.clone()
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn external_linker_swa_offload_uses_complete_trailing_pages() {
|
||||
let params = CacheInitParams {
|
||||
page_size: 2,
|
||||
swa_sliding_window_size: Some(4),
|
||||
..Default::default()
|
||||
};
|
||||
let mut tc: UnifiedTreeCore<Vec<i64>> = UnifiedTreeCore::new(params, vec![FULL, SWA]);
|
||||
tc.set_enable_external_cache_linker(true).unwrap();
|
||||
let root = tc.arena.root();
|
||||
let node = tc.add_new_node_(
|
||||
root,
|
||||
vec![1, 2, 3, 4, 5, 6],
|
||||
&Tensor::from_slice(&[10i64, 11, 12, 13, 14, 15]),
|
||||
0,
|
||||
None,
|
||||
);
|
||||
tc.arena.node_mut(node).values[SWA.idx()].value =
|
||||
Some(Tensor::from_slice(&[20i64, 21, 22, 23, 24]));
|
||||
|
||||
let transfers = tc
|
||||
.build_external_linker_offload_transfers(tc.arena.node(node).id)
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
assert_eq!(transfers.len(), 2);
|
||||
assert_eq!(transfers[0].name, PoolName::Kv);
|
||||
assert_eq!(transfers[1].name, PoolName::Swa);
|
||||
assert_eq!(transfers[1].hit_policy, PoolHitPolicy::TrailingPages);
|
||||
assert!(
|
||||
transfers[1]
|
||||
.device_indices
|
||||
.as_ref()
|
||||
.unwrap()
|
||||
.equal(&Tensor::from_slice(&[21i64, 22, 23, 24]))
|
||||
);
|
||||
assert_eq!(
|
||||
transfers[1].keys.as_ref().unwrap(),
|
||||
&tc.arena.node(node).hash_value.as_ref().unwrap()[1..]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn external_linker_rejects_mamba_trees() {
|
||||
let params = CacheInitParams {
|
||||
mamba_cache_chunk_size: Some(1),
|
||||
..Default::default()
|
||||
};
|
||||
let mut tc: UnifiedTreeCore<Vec<i64>> = UnifiedTreeCore::new(params, vec![FULL, MAMBA]);
|
||||
let error = tc.set_enable_external_cache_linker(true).unwrap_err();
|
||||
assert!(matches!(
|
||||
&error,
|
||||
TreeCoreRuntimeError::ExternalCacheLinkerUnsupportedComponent {
|
||||
component_type
|
||||
} if *component_type == MAMBA
|
||||
));
|
||||
assert!(error.to_string().contains("Mamba"));
|
||||
assert!(!tc.enable_external_cache_linker);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn external_linker_state_follows_load_offload_and_split_lifecycle() {
|
||||
let mut tc = core();
|
||||
tc.set_enable_external_cache_linker(true).unwrap();
|
||||
tc.insert(&insert_params(&vec![1, 2], &[10, 11]));
|
||||
tc.insert(&insert_params(&vec![1, 2, 3, 4], &[10, 11, 12, 13]));
|
||||
let anchor = tc
|
||||
.match_prefix(&match_params(&vec![1, 2]))
|
||||
.best_match_node_id;
|
||||
let leaf = tc
|
||||
.match_prefix(&match_params(&vec![1, 2, 3, 4]))
|
||||
.best_match_node_id;
|
||||
|
||||
tc.mark_external_cache_stored_path(leaf, anchor).unwrap();
|
||||
assert!(
|
||||
tc.arena
|
||||
.node(tc.arena.resolve(leaf).expect("live test node"))
|
||||
.external_cache_stored
|
||||
);
|
||||
assert!(
|
||||
!tc.arena
|
||||
.node(tc.arena.resolve(anchor).expect("live test node"))
|
||||
.external_cache_stored
|
||||
);
|
||||
assert!(
|
||||
tc.build_external_linker_offload_transfers(leaf)
|
||||
.unwrap()
|
||||
.is_none()
|
||||
);
|
||||
assert!(matches!(
|
||||
tc.mark_external_linker_offload_pending(leaf),
|
||||
Err(TreeCoreRuntimeError::InvalidExternalCacheOffloadState {
|
||||
node_id,
|
||||
stored: true,
|
||||
pending_id: None,
|
||||
}) if node_id == leaf
|
||||
));
|
||||
|
||||
let leaf_idx = tc.arena.resolve(leaf).expect("live test node");
|
||||
tc.arena.node_mut(leaf_idx).external_cache_stored = false;
|
||||
tc.mark_external_linker_offload_pending(leaf).unwrap();
|
||||
assert!(
|
||||
tc.build_external_linker_offload_transfers(leaf)
|
||||
.unwrap()
|
||||
.is_none()
|
||||
);
|
||||
assert!(matches!(
|
||||
tc.mark_external_linker_offload_pending(leaf),
|
||||
Err(TreeCoreRuntimeError::InvalidExternalCacheOffloadState {
|
||||
node_id,
|
||||
stored: false,
|
||||
pending_id: Some(pending_id),
|
||||
}) if node_id == leaf && pending_id == leaf
|
||||
));
|
||||
let (new_parent, action) = tc.split_node_(leaf_idx, 1);
|
||||
let new_parent_handle = tc.arena.node(new_parent).id;
|
||||
assert!(action.is_some());
|
||||
assert!(!tc.arena.node(new_parent).external_cache_stored);
|
||||
assert!(!tc.arena.node(leaf_idx).external_cache_stored);
|
||||
|
||||
tc.insert(&insert_params(&vec![9], &[19]));
|
||||
let independent = tc.match_prefix(&match_params(&vec![9])).best_match_node_id;
|
||||
tc.mark_external_linker_offload_pending(independent)
|
||||
.unwrap();
|
||||
assert!(matches!(
|
||||
tc.finish_external_linker_offload(&[independent, new_parent_handle], independent, false),
|
||||
Err(TreeCoreRuntimeError::InvalidExternalCacheOffloadState {
|
||||
node_id,
|
||||
stored: false,
|
||||
pending_id: Some(pending_id),
|
||||
}) if node_id == new_parent_handle && pending_id == leaf
|
||||
));
|
||||
assert_eq!(
|
||||
tc.arena
|
||||
.node(tc.arena.resolve(independent).expect("live test node"))
|
||||
.write_through_pending_id,
|
||||
Some(independent)
|
||||
);
|
||||
for node_id in [new_parent_handle, leaf] {
|
||||
let node = tc
|
||||
.arena
|
||||
.node(tc.arena.resolve(node_id).expect("live test node"));
|
||||
assert_eq!(node.write_through_pending_id, Some(leaf));
|
||||
assert!(!node.external_cache_stored);
|
||||
}
|
||||
|
||||
tc.finish_external_linker_offload(&[independent], independent, false)
|
||||
.unwrap();
|
||||
tc.finish_external_linker_offload(&[new_parent_handle, leaf], leaf, false)
|
||||
.unwrap();
|
||||
for node_id in [new_parent_handle, leaf] {
|
||||
let node = tc
|
||||
.arena
|
||||
.node(tc.arena.resolve(node_id).expect("live test node"));
|
||||
assert_eq!(node.write_through_pending_id, None);
|
||||
assert!(!node.external_cache_stored);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn failed_external_offload_preserves_independently_confirmed_state() {
|
||||
let mut tc = core();
|
||||
tc.set_enable_external_cache_linker(true).unwrap();
|
||||
tc.insert(&insert_params(&vec![1], &[10]));
|
||||
tc.insert(&insert_params(&vec![1, 2], &[10, 11]));
|
||||
let anchor = tc.match_prefix(&match_params(&vec![1])).best_match_node_id;
|
||||
let leaf = tc
|
||||
.match_prefix(&match_params(&vec![1, 2]))
|
||||
.best_match_node_id;
|
||||
|
||||
tc.mark_external_linker_offload_pending(leaf).unwrap();
|
||||
tc.mark_external_cache_stored_path(leaf, anchor).unwrap();
|
||||
tc.finish_external_linker_offload(&[leaf], leaf, false)
|
||||
.unwrap();
|
||||
|
||||
let leaf = tc
|
||||
.arena
|
||||
.node(tc.arena.resolve(leaf).expect("live test node"));
|
||||
assert_eq!(leaf.write_through_pending_id, None);
|
||||
assert!(leaf.external_cache_stored);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn external_linker_path_validation_is_atomic() {
|
||||
let mut tc = core();
|
||||
tc.insert(&insert_params(&vec![1], &[10]));
|
||||
tc.insert(&insert_params(&vec![2], &[20]));
|
||||
let left = tc.match_prefix(&match_params(&vec![1])).best_match_node_id;
|
||||
let right = tc.match_prefix(&match_params(&vec![2])).best_match_node_id;
|
||||
|
||||
assert!(matches!(
|
||||
tc.mark_external_cache_stored_path(left, right),
|
||||
Err(TreeCoreRuntimeError::ExternalCachePathNotAncestor {
|
||||
from_node_id,
|
||||
until_node_id,
|
||||
}) if from_node_id == left && until_node_id == right
|
||||
));
|
||||
assert!(
|
||||
!tc.arena
|
||||
.node(tc.arena.resolve(left).expect("live test node"))
|
||||
.external_cache_stored
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn mark_write_through_pending_stamps_the_supplied_ack() {
|
||||
let mut tc = core();
|
||||
@@ -2325,6 +2567,38 @@ fn backup_kv_action_chains_unbacked_ancestors_first() {
|
||||
assert_eq!(action.node_ids, vec![c]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn backup_kv_action_stops_at_an_externally_stored_or_pending_ancestor() {
|
||||
let mut tc = core();
|
||||
tc.set_enable_external_cache_linker(true).unwrap();
|
||||
tc.insert(&insert_params(&vec![1], &[10]));
|
||||
tc.insert(&insert_params(&vec![1, 2], &[10, 11]));
|
||||
tc.insert(&insert_params(&vec![1, 2, 3], &[10, 11, 12]));
|
||||
let a = tc.match_prefix(&match_params(&vec![1])).best_match_node_id;
|
||||
let b = tc
|
||||
.match_prefix(&match_params(&vec![1, 2]))
|
||||
.best_match_node_id;
|
||||
let c = tc
|
||||
.match_prefix(&match_params(&vec![1, 2, 3]))
|
||||
.best_match_node_id;
|
||||
let a_idx = tc.arena.resolve(a).expect("live test node");
|
||||
tc.arena.node_mut(a_idx).external_cache_stored = true;
|
||||
|
||||
let action = tc.build_backup_kv_action_(
|
||||
tc.arena.node(tc.arena.resolve(c).expect("live test node")),
|
||||
/* write_back = */ false,
|
||||
);
|
||||
assert_eq!(action.node_ids, vec![b, c]);
|
||||
|
||||
tc.arena.node_mut(a_idx).external_cache_stored = false;
|
||||
tc.mark_external_linker_offload_pending(a).unwrap();
|
||||
let action = tc.build_backup_kv_action_(
|
||||
tc.arena.node(tc.arena.resolve(c).expect("live test node")),
|
||||
/* write_back = */ false,
|
||||
);
|
||||
assert_eq!(action.node_ids, vec![b, c]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn split_of_a_pending_node_transfers_the_ack_and_emits_the_replace_action() {
|
||||
let mut tc = core();
|
||||
@@ -3958,6 +4232,35 @@ fn fallible_node_boundaries_reject_stale_handles() {
|
||||
Err(TreeCoreRuntimeError::NodeAccess(NodeAccessError { node_id }))
|
||||
if node_id == stale_root
|
||||
));
|
||||
assert!(matches!(
|
||||
tc.build_external_linker_offload_transfers(stale_root),
|
||||
Err(NodeAccessError { node_id }) if node_id == stale_root
|
||||
));
|
||||
assert!(matches!(
|
||||
tc.mark_external_cache_stored_path(stale_root, live_root),
|
||||
Err(TreeCoreRuntimeError::NodeAccess(NodeAccessError { node_id }))
|
||||
if node_id == stale_root
|
||||
));
|
||||
assert!(matches!(
|
||||
tc.mark_external_cache_stored_path(live_root, stale_root),
|
||||
Err(TreeCoreRuntimeError::NodeAccess(NodeAccessError { node_id }))
|
||||
if node_id == stale_root
|
||||
));
|
||||
assert!(matches!(
|
||||
tc.mark_external_linker_offload_pending(stale_root),
|
||||
Err(TreeCoreRuntimeError::NodeAccess(NodeAccessError { node_id }))
|
||||
if node_id == stale_root
|
||||
));
|
||||
assert!(matches!(
|
||||
tc.finish_external_linker_offload(&[live_root, stale_root], live_root, true),
|
||||
Err(TreeCoreRuntimeError::NodeAccess(NodeAccessError { node_id }))
|
||||
if node_id == stale_root
|
||||
));
|
||||
assert!(
|
||||
!tc.arena
|
||||
.node(tc.arena.resolve(live_root).expect("live root"))
|
||||
.external_cache_stored
|
||||
);
|
||||
assert!(matches!(
|
||||
tc.get_hash_values(stale_root),
|
||||
Err(NodeAccessError { node_id }) if node_id == stale_root
|
||||
@@ -7964,6 +8267,10 @@ fn inspection_rejects_stale_handles_without_panicking() {
|
||||
assert_eq!(tc.inspect_get_parent_node_id(stale_root), Err(expected));
|
||||
assert_eq!(tc.inspect_get_child_node_ids(stale_root), Err(expected));
|
||||
assert_eq!(tc.inspect_get_node_key_length(stale_root), Err(expected));
|
||||
assert_eq!(
|
||||
tc.inspect_is_external_cache_stored(stale_root),
|
||||
Err(expected)
|
||||
);
|
||||
assert_eq!(
|
||||
tc.inspect_set_node_hash_values(stale_root, None),
|
||||
Err(expected)
|
||||
@@ -7993,6 +8300,7 @@ fn inspection_rejects_stale_handles_without_panicking() {
|
||||
assert!(!tc.inspect_is_host_evictable_leaf(stale_root));
|
||||
|
||||
assert_eq!(tc.inspect_get_parent_node_id(live_root), Ok(None));
|
||||
assert_eq!(tc.inspect_is_external_cache_stored(live_root), Ok(false));
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
||||
@@ -534,6 +534,8 @@ pub struct UnifiedTreeCore<K: ChildKeyType> {
|
||||
pub(crate) enable_hicache: bool,
|
||||
/// Whether the storage tier (L3) is wired; gates page-hash computation.
|
||||
pub(crate) enable_storage: bool,
|
||||
/// Whether a direct device-to-external-cache linker is wired.
|
||||
pub(crate) enable_external_cache_linker: bool,
|
||||
/// Whether the cache wired a host SWA pool (HiCache).
|
||||
pub(crate) has_swa_host_pool: bool,
|
||||
/// Whether tree mutations emit BlockStored/BlockRemoved events.
|
||||
@@ -717,6 +719,7 @@ impl<K: ChildKeyType> UnifiedTreeCore<K> {
|
||||
is_write_back: params.is_write_back,
|
||||
enable_hicache: params.enable_hicache,
|
||||
enable_storage: false,
|
||||
enable_external_cache_linker: false,
|
||||
has_swa_host_pool: params.has_swa_host_pool,
|
||||
enable_kv_cache_events: params.enable_kv_cache_events,
|
||||
kv_event_queue: Vec::new(),
|
||||
@@ -1319,6 +1322,12 @@ impl<K: ChildKeyType> UnifiedTreeCore<K> {
|
||||
return false;
|
||||
}
|
||||
node.hit_count += 1;
|
||||
|
||||
if self.enable_external_cache_linker {
|
||||
return Self::needs_external_linker_offload_(node)
|
||||
&& node.hit_count >= self.write_through_threshold;
|
||||
}
|
||||
|
||||
self.enable_hicache && !node.backuped() && node.hit_count >= self.write_through_threshold
|
||||
}
|
||||
|
||||
@@ -1763,6 +1772,7 @@ impl<K: ChildKeyType> UnifiedTreeCore<K> {
|
||||
let child = self.arena.node(child_id);
|
||||
let parent_id = child.parent();
|
||||
let child_namespace = child.namespace.clone();
|
||||
let child_external_cache_stored = child.external_cache_stored;
|
||||
let (key_head, key_tail) = child.key.split_at(split_len);
|
||||
// key_head keeps the original key's first page, which keys the parent's child map.
|
||||
let parent_map_key = key_head.child_key(page_size);
|
||||
@@ -1778,6 +1788,7 @@ impl<K: ChildKeyType> UnifiedTreeCore<K> {
|
||||
(child_namespace.clone(), key_tail.child_key(page_size)),
|
||||
child_id,
|
||||
);
|
||||
self.arena.node_mut(new_node_id).external_cache_stored = child_external_cache_stored;
|
||||
|
||||
// The child's aux LRU cells detach while it is re-linked.
|
||||
self.for_each_component_lru_(
|
||||
@@ -1897,7 +1908,7 @@ impl<K: ChildKeyType> UnifiedTreeCore<K> {
|
||||
"add_new_node_: parent {parent_id} already has a child on the new node's page"
|
||||
);
|
||||
self.inc_evictable_size(FULL, value.size()[0] as usize);
|
||||
if self.enable_storage {
|
||||
if self.enable_storage || self.enable_external_cache_linker {
|
||||
let hash_values = self.arena.compute_node_hash_values(new_node_id, page_size);
|
||||
self.arena.node_mut(new_node_id).hash_value = Some(hash_values);
|
||||
}
|
||||
@@ -2708,6 +2719,22 @@ impl<K: ChildKeyType> UnifiedTreeCore<K> {
|
||||
self.enable_storage = value;
|
||||
}
|
||||
|
||||
/// Enable or disable the direct external-cache linker.
|
||||
pub fn set_enable_external_cache_linker(
|
||||
&mut self,
|
||||
value: bool,
|
||||
) -> Result<(), TreeCoreRuntimeError> {
|
||||
if value && self.components_by_type[MAMBA.idx()].is_some() {
|
||||
return Err(
|
||||
TreeCoreRuntimeError::ExternalCacheLinkerUnsupportedComponent {
|
||||
component_type: MAMBA,
|
||||
},
|
||||
);
|
||||
}
|
||||
self.enable_external_cache_linker = value;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
// ==== KV cache placement events ====
|
||||
|
||||
/// Append an event, coalescing it with a compatible queue tail.
|
||||
@@ -3463,14 +3490,19 @@ impl<K: ChildKeyType> UnifiedTreeCore<K> {
|
||||
Ok(order)
|
||||
}
|
||||
|
||||
/// Build the backup action for a node and its unbacked ancestors.
|
||||
/// Build the backup action for a node and its not-yet-persisted ancestors.
|
||||
pub fn build_backup_kv_action_(&self, node: &Node<K>, write_back: bool) -> BackupKV {
|
||||
let mut chain = vec![node.id];
|
||||
if !write_back {
|
||||
let mut ancestor = node.try_parent();
|
||||
while let Some(ancestor_idx) = ancestor {
|
||||
let ancestor_node = self.arena.node(ancestor_idx);
|
||||
if ancestor_node.is_root() || ancestor_node.backuped() {
|
||||
if ancestor_node.is_root()
|
||||
|| ancestor_node.backuped()
|
||||
|| ancestor_node.external_cache_stored
|
||||
|| (self.enable_external_cache_linker
|
||||
&& ancestor_node.write_through_pending_id.is_some())
|
||||
{
|
||||
break;
|
||||
}
|
||||
chain.push(ancestor_node.id);
|
||||
@@ -3696,6 +3728,103 @@ impl<K: ChildKeyType> UnifiedTreeCore<K> {
|
||||
depth
|
||||
}
|
||||
|
||||
/// Build transfers for a node with no stored or pending external copy.
|
||||
pub fn build_external_linker_offload_transfers(
|
||||
&self,
|
||||
node_id: NodeId,
|
||||
) -> Result<Option<Vec<PoolTransfer>>, NodeAccessError> {
|
||||
let node_id = self.arena.resolve(node_id)?;
|
||||
if !Self::needs_external_linker_offload_(self.arena.node(node_id)) {
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
let transfers = self
|
||||
.components
|
||||
.iter()
|
||||
.filter_map(|component| component.build_external_linker_offload_transfer(self, node_id))
|
||||
.collect();
|
||||
Ok(Some(transfers))
|
||||
}
|
||||
|
||||
fn needs_external_linker_offload_(node: &Node<K>) -> bool {
|
||||
!node.external_cache_stored && node.write_through_pending_id.is_none()
|
||||
}
|
||||
|
||||
/// Mark the path from `from_node_id` to, but excluding, `until_node_id` as
|
||||
/// available in the external cache.
|
||||
pub fn mark_external_cache_stored_path(
|
||||
&mut self,
|
||||
from_node_id: NodeId,
|
||||
until_node_id: NodeId,
|
||||
) -> Result<(), TreeCoreRuntimeError> {
|
||||
let from = self.arena.resolve(from_node_id)?;
|
||||
let until = self.arena.resolve(until_node_id)?;
|
||||
let mut path = Vec::new();
|
||||
let mut current = from;
|
||||
while current != until {
|
||||
let node = self.arena.node(current);
|
||||
let Some(parent) = node.try_parent() else {
|
||||
return Err(TreeCoreRuntimeError::ExternalCachePathNotAncestor {
|
||||
from_node_id,
|
||||
until_node_id,
|
||||
});
|
||||
};
|
||||
path.push(current);
|
||||
current = parent;
|
||||
}
|
||||
for node_id in path {
|
||||
self.arena.node_mut(node_id).external_cache_stored = true;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Publish an accepted external offload as pending.
|
||||
pub fn mark_external_linker_offload_pending(
|
||||
&mut self,
|
||||
node_id: NodeId,
|
||||
) -> Result<(), TreeCoreRuntimeError> {
|
||||
let node_idx = self.arena.resolve(node_id)?;
|
||||
let node = self.arena.node(node_idx);
|
||||
if !Self::needs_external_linker_offload_(node) {
|
||||
return Err(TreeCoreRuntimeError::InvalidExternalCacheOffloadState {
|
||||
node_id,
|
||||
stored: node.external_cache_stored,
|
||||
pending_id: node.write_through_pending_id,
|
||||
});
|
||||
}
|
||||
self.arena.node_mut(node_idx).write_through_pending_id = Some(node_id);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Finalize external-store state for an offload and its split fragments.
|
||||
pub fn finish_external_linker_offload(
|
||||
&mut self,
|
||||
node_ids: &[NodeId],
|
||||
ack_id: NodeId,
|
||||
success: bool,
|
||||
) -> Result<(), TreeCoreRuntimeError> {
|
||||
let node_indices = node_ids
|
||||
.iter()
|
||||
.map(|&node_id| self.arena.resolve(node_id))
|
||||
.collect::<Result<Vec<_>, _>>()?;
|
||||
for (&node_id, &node_idx) in node_ids.iter().zip(&node_indices) {
|
||||
let node = self.arena.node(node_idx);
|
||||
if node.write_through_pending_id != Some(ack_id) {
|
||||
return Err(TreeCoreRuntimeError::InvalidExternalCacheOffloadState {
|
||||
node_id,
|
||||
stored: node.external_cache_stored,
|
||||
pending_id: node.write_through_pending_id,
|
||||
});
|
||||
}
|
||||
}
|
||||
for node_id in node_indices {
|
||||
let node = self.arena.node_mut(node_id);
|
||||
node.write_through_pending_id = None;
|
||||
node.external_cache_stored |= success;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Clear the write-through-pending mark (when it matches ack_id) and record the
|
||||
/// host store event for each acked node.
|
||||
pub fn finish_write_through(
|
||||
@@ -4357,6 +4486,15 @@ impl<K: ChildKeyType> UnifiedTreeCore<K> {
|
||||
Ok(self.arena.node(node_id).write_through_pending_id)
|
||||
}
|
||||
|
||||
/// Whether a node is known to be stored in the external cache.
|
||||
pub fn inspect_is_external_cache_stored(
|
||||
&self,
|
||||
node_id: NodeId,
|
||||
) -> Result<bool, NodeAccessError> {
|
||||
let node_id = self.arena.resolve(node_id)?;
|
||||
Ok(self.arena.node(node_id).external_cache_stored)
|
||||
}
|
||||
|
||||
/// Whether a node is in a component's device LRU.
|
||||
pub fn inspect_is_node_in_device_lru(
|
||||
&self,
|
||||
|
||||
Reference in New Issue
Block a user