[Rust TreeCore] Support external cache linker (#37306)

This commit is contained in:
Jialin Ouyang
2026-09-10 19:22:24 +08:00
committed by GitHub
parent 334e94d8ac
commit 908226fea2
18 changed files with 1549 additions and 90 deletions
@@ -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>,
+22
View File
@@ -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]
+141 -3
View File
@@ -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,