minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 10 Packages Projects Releases Wiki Activity

[HiCache] Replace skip_lock_node_ids with a segment lock protocol (#36848)

Browse Source
This commit is contained in:
Zhiqiang Xie
2026-09-09 14:55:51 -07:00
committed by GitHub
parent a84ffd1326
commit beaf3d9252
41 changed files with 2116 additions and 1146 deletions
Download Patch File Download Diff File Expand all files Collapse all files
rust/sglang-radix-tree/src/components/full.rs
+24 -34
View File
@@ -2,7 +2,7 @@
//! the rest from the `TreeComponent` defaults.
use std::cmp::Reverse;
use std::collections::{BinaryHeap, HashMap, HashSet};
use std::collections::{BinaryHeap, HashMap};
use tch::{Kind, Tensor};
@@ -277,8 +277,6 @@ impl<K: ChildKeyType> TreeComponent<K> for FullComponent {
mut result: IncLockRefResult,
lock_host: bool,
) -> IncLockRefResult {
let ct = FULL;
// Only the last host node needs to be protected.
if lock_host {
let node = tree_core.arena.node_mut(node_id);
@@ -291,20 +289,17 @@ impl<K: ChildKeyType> TreeComponent<K> for FullComponent {
return result;
}
// Skip the bottom evicted segment, recording it for the matching release.
let on_boundary = |node: &Node<K>| node.is_root() || node.has_device_value(FULL);
// The bottom device-evicted segment is locked too (no ledger move —
// nothing is on device); a load-back that materializes a value under
// lock credits protected directly.
let mut cur = node_id;
let mut node = tree_core.arena.node(cur);
if !on_boundary(node) {
let skip_lock_node_ids = result.skip_lock_node_ids.entry(ct).or_default();
loop {
skip_lock_node_ids.insert(node.id);
cur = node.parent();
node = tree_core.arena.node(cur);
if on_boundary(node) {
break;
}
loop {
let node = tree_core.arena.node_mut(cur);
if node.is_root() || node.has_device_value(FULL) {
break;
}
node.inc_device_lock_ref(FULL);
cur = node.parent();
}
// Lock the device-on segment up to the root.
@@ -341,11 +336,9 @@ impl<K: ChildKeyType> TreeComponent<K> for FullComponent {
&self,
tree_core: &mut UnifiedTreeCore<K>,
node_id: NodeIdx_,
params: Option<&DecLockRefParams>,
_params: &DecLockRefParams,
lock_host: bool,
) {
let ct = FULL;
if lock_host {
let node = tree_core.arena.node_mut(node_id);
if node.host_lock_ref(FULL) == 0 {
@@ -360,10 +353,6 @@ impl<K: ChildKeyType> TreeComponent<K> for FullComponent {
return;
}
let empty = HashSet::new();
let skip_lock_node_ids = params
.and_then(|p| p.skip_lock_node_ids.get(&ct))
.unwrap_or(&empty);
let mut cur = node_id;
loop {
let node = tree_core.arena.node_mut(cur);
@@ -371,20 +360,12 @@ impl<K: ChildKeyType> TreeComponent<K> for FullComponent {
break;
}
let parent = node.parent();
if skip_lock_node_ids.contains(&node.id) {
cur = parent;
continue;
}
assert!(
node.has_device_value(FULL),
"release_component_lock: node {cur} has no FULL device value"
);
let old_lock_ref = node.device_lock_ref(FULL);
assert!(
old_lock_ref > 0,
"release_component_lock: node {cur} is not locked"
"FULL segment release hit lock_ref=0 on node {cur}"
);
let newly_unlocked_len = if old_lock_ref == 1 {
let newly_unlocked_len = if old_lock_ref == 1 && node.has_device_value(FULL) {
Some(node.device_value_len(FULL))
} else {
None
@@ -393,6 +374,8 @@ impl<K: ChildKeyType> TreeComponent<K> for FullComponent {
if let Some(key_len) = newly_unlocked_len {
tree_core.dec_protected_size(FULL, key_len);
tree_core.inc_evictable_size(FULL, key_len);
}
if old_lock_ref == 1 {
tree_core.update_evictable_leaf_sets_(cur);
}
cur = parent;
@@ -478,9 +461,16 @@ impl<K: ChildKeyType> TreeComponent<K> for FullComponent {
let n_len = loaded.host_value_len(FULL) as i64;
loaded
.set_device_value(FULL, device_indices.narrow(0, offset, n_len).copy());
let locked = loaded.device_lock_ref(FULL) > 0;
offset += n_len;
// Full uses leaf sets, not LRU.
tree_core.inc_evictable_size(FULL, n_len as usize);
// Full uses leaf sets, not LRU. A value materialized
// under lock is protected; the last release moves it
// to evictable.
if locked {
tree_core.inc_protected_size(FULL, n_len as usize);
} else {
tree_core.inc_evictable_size(FULL, n_len as usize);
}
tree_core.update_evictable_leaf_sets_(loaded_idx);
}
}
rust/sglang-radix-tree/src/components/mamba.rs
+35 -55
View File
@@ -178,16 +178,7 @@ impl<K: ChildKeyType> TreeComponent<K> for MambaComponent {
return;
}
if !tree_core.arena.has_device_value(node_id, MAMBA) {
// Tombstone refill: the node moves from the host LRU to the device LRU.
tree_core
.arena
.set_device_value(node_id, MAMBA, mamba_value.shallow_clone());
let host_lru = tree_core.host_lru_list_mut(MAMBA);
if host_lru.in_list(Some(node_id)) {
host_lru.remove_node(node_id);
}
tree_core.device_lru_list_mut(MAMBA).insert_mru(node_id);
tree_core.inc_evictable_size(MAMBA, slot_len);
tree_core.set_component_device_value_(node_id, MAMBA, mamba_value.shallow_clone());
let tick = tree_core.arena.get_and_bump_access_counter();
tree_core.arena.node_mut(node_id).last_access_counter = tick;
self.emit_excess_path_states_eviction_(tree_core.arena.node(node_id).id, cache_actions);
@@ -417,24 +408,19 @@ impl<K: ChildKeyType> TreeComponent<K> for MambaComponent {
&self,
tree_core: &mut UnifiedTreeCore<K>,
node_id: NodeIdx_,
mut result: IncLockRefResult,
result: IncLockRefResult,
lock_host: bool,
) -> IncLockRefResult {
let node = tree_core.arena.node(node_id);
if node.is_root() {
return result;
}
// A node in skip_lock_node_ids was a tombstone when this lock was acquired.
if !Self::has_value(node, lock_host) {
result
.skip_lock_node_ids
.entry(MAMBA)
.or_default()
.insert(node.id);
return result;
}
// Tombstones are counted too; ledger/LRU track only data-bearing
// nodes (a value materialized under lock is credited to protected
// at the materialization site).
let has_value = Self::has_value(node, lock_host);
if lock_host {
if node.host_lock_ref(MAMBA) == 0 {
if node.host_lock_ref(MAMBA) == 0 && has_value {
let host_lru = tree_core.host_lru_list_mut(MAMBA);
if host_lru.in_list(Some(node_id)) {
host_lru.remove_node(node_id);
@@ -443,7 +429,7 @@ impl<K: ChildKeyType> TreeComponent<K> for MambaComponent {
tree_core.arena.inc_host_lock_ref(node_id, MAMBA);
} else {
let value_len = node.device_value_len(MAMBA);
if node.device_lock_ref(MAMBA) == 0 {
if node.device_lock_ref(MAMBA) == 0 && has_value {
tree_core.dec_evictable_size(MAMBA, value_len);
tree_core.inc_protected_size(MAMBA, value_len);
}
@@ -457,43 +443,44 @@ impl<K: ChildKeyType> TreeComponent<K> for MambaComponent {
&self,
tree_core: &mut UnifiedTreeCore<K>,
node_id: NodeIdx_,
params: Option<&DecLockRefParams>,
_params: &DecLockRefParams,
lock_host: bool,
) {
if tree_core.arena.node(node_id).is_root() {
return;
}
if let Some(params) = params
&& params
.skip_lock_node_ids
.get(&MAMBA)
.is_some_and(|ids| ids.contains(&tree_core.arena.node(node_id).id))
{
return;
}
if lock_host {
let node = tree_core.arena.node_mut(node_id);
assert!(
node.host_lock_ref(MAMBA) > 0,
"Mamba release hit host_lock_ref=0 on node {node_id}"
);
node.dec_host_lock_ref(MAMBA);
if node.host_lock_ref(MAMBA) == 0
&& !node.has_device_value(MAMBA)
&& node.has_host_value(MAMBA)
{
let host_lru = tree_core.host_lru_list_mut(MAMBA);
if !host_lru.in_list(Some(node_id)) {
host_lru.insert_mru(node_id);
if node.host_lock_ref(MAMBA) == 0 {
if !node.has_device_value(MAMBA) && node.has_host_value(MAMBA) {
let host_lru = tree_core.host_lru_list_mut(MAMBA);
if !host_lru.in_list(Some(node_id)) {
host_lru.insert_mru(node_id);
}
}
tree_core.update_evictable_leaf_sets_(node_id);
}
return;
}
let node = tree_core.arena.node(node_id);
let device_lock_ref = node.device_lock_ref(MAMBA);
if device_lock_ref > 0 {
if device_lock_ref == 1 {
let value_len = node.device_value_len(MAMBA);
tree_core.inc_evictable_size(MAMBA, value_len);
tree_core.dec_protected_size(MAMBA, value_len);
}
tree_core.arena.dec_device_lock_ref(node_id, MAMBA);
assert!(
device_lock_ref > 0,
"Mamba release hit lock_ref=0 on node {node_id}"
);
if device_lock_ref == 1 && node.has_device_value(MAMBA) {
let value_len = node.device_value_len(MAMBA);
tree_core.inc_evictable_size(MAMBA, value_len);
tree_core.dec_protected_size(MAMBA, value_len);
}
tree_core.arena.dec_device_lock_ref(node_id, MAMBA);
if device_lock_ref == 1 {
tree_core.update_evictable_leaf_sets_(node_id);
}
}
@@ -613,16 +600,9 @@ impl<K: ChildKeyType> TreeComponent<K> for MambaComponent {
return;
};
if let Some(device_indices) = &transfer.device_indices {
let node = tree_core.arena.node_mut(node_id);
node.set_device_value(MAMBA, device_indices.copy());
let count = node.device_value_len(MAMBA);
// Move from host LRU to device LRU
let host_lru = tree_core.host_lru_list_mut(MAMBA);
if host_lru.in_list(Some(node_id)) {
host_lru.remove_node(node_id);
}
tree_core.device_lru_list_mut(MAMBA).insert_mru(node_id);
tree_core.inc_evictable_size(MAMBA, count);
// The materialization primitive owns the ledger/LRU moves,
// including crediting protected when restored under lock.
tree_core.set_component_device_value_(node_id, MAMBA, device_indices.copy());
}
}
// The python elif chain has no BACKUP_STORAGE arm.
rust/sglang-radix-tree/src/components/mod.rs
+44 -1
View File
@@ -344,7 +344,7 @@ pub trait TreeComponent<K: ChildKeyType> {
&self,
tree_core: &mut UnifiedTreeCore<K>,
node_id: NodeIdx_,
params: Option<&DecLockRefParams>,
params: &DecLockRefParams,
lock_host: bool,
);
@@ -466,6 +466,49 @@ pub const BASE_COMPONENT_TYPE: ComponentType = ComponentType::Full;
/// Slots per tier — the arrays are sized to this, not the enabled subset.
pub const NUM_COMPONENT_TYPES: usize = ComponentType::Mamba as usize + 1;
/// A set of component types (bitmask over `ComponentType::idx`), e.g. the
/// components an `inc_lock_ref` left untaken.
#[derive(Copy, Clone, Default, PartialEq, Eq, Debug)]
pub struct ComponentSet(u8);
impl ComponentSet {
pub const EMPTY: ComponentSet = ComponentSet(0);
/// The set holding exactly one component.
pub const fn of(component_type: ComponentType) -> ComponentSet {
ComponentSet(1 << component_type.idx())
}
pub fn insert(&mut self, component_type: ComponentType) {
self.0 |= 1 << component_type.idx();
}
pub const fn contains(self, component_type: ComponentType) -> bool {
self.0 & (1 << component_type.idx()) != 0
}
pub const fn is_empty(self) -> bool {
self.0 == 0
}
/// The members, in component-index order.
pub fn iter(self) -> impl Iterator<Item = ComponentType> {
(0..NUM_COMPONENT_TYPES)
.filter(move |idx| self.0 & (1 << idx) != 0)
.map(ComponentType::from_idx)
}
}
impl FromIterator<ComponentType> for ComponentSet {
fn from_iter<I: IntoIterator<Item = ComponentType>>(iter: I) -> Self {
let mut set = ComponentSet::EMPTY;
for component_type in iter {
set.insert(component_type);
}
set
}
}
impl ComponentType {
/// Index into a per-component array.
pub const fn idx(self) -> usize {
rust/sglang-radix-tree/src/components/swa.rs
+87 -88
View File
@@ -3,7 +3,7 @@
//! SWA values arrive pool-resolved; the full->SWA index translation happens at
//! the cache boundary.
use std::collections::{HashMap, HashSet};
use std::collections::HashMap;
use tch::{Kind, Tensor};
@@ -34,10 +34,15 @@ impl SwaComponent {
impl SwaComponent {
/// Build the driver from the tree's init params.
pub fn new(params: &CacheInitParams) -> Self {
let sliding_window_size = params
.swa_sliding_window_size
.expect("the SWA component requires swa_sliding_window_size");
assert!(
sliding_window_size > 0,
"swa_sliding_window_size must be positive"
);
SwaComponent {
sliding_window_size: params
.swa_sliding_window_size
.expect("the SWA component requires swa_sliding_window_size"),
sliding_window_size,
}
}
@@ -408,11 +413,9 @@ impl<K: ChildKeyType> TreeComponent<K> for SwaComponent {
}
let swa_evicted_seqlen = params.swa_evicted_seqlen;
assert_eq!(
node.device_lock_ref(SWA),
0,
"tombstone Swa lock_ref should be 0, node {node_id}"
);
// A locked tombstone is legal (segment locks count every node); the
// full-value swap below is safe because full lock_ref >= swa
// lock_ref, so a locked-SWA node always takes the Recover branch.
assert_eq!(
swa_evicted_seqlen % tree_core.page_size,
0,
@@ -495,11 +498,6 @@ impl<K: ChildKeyType> TreeComponent<K> for SwaComponent {
if node.has_device_value(SWA) {
return;
}
assert_eq!(
node.device_lock_ref(SWA),
0,
"tombstone Swa lock_ref should be 0 on unevict, node {node_id}"
);
let swa_evicted_seqlen = params.swa_evicted_seqlen;
assert_eq!(
swa_evicted_seqlen % tree_core.page_size,
@@ -587,26 +585,33 @@ impl<K: ChildKeyType> TreeComponent<K> for SwaComponent {
let (new_parent, child) = tree_core.arena.node_pair_mut(new_parent_id, child_id);
let split_len = new_parent.key.atom_len() as i64;
new_parent.copy_device_lock_ref(SWA, child);
new_parent.copy_host_lock_ref(SWA, child);
if child.has_device_value(SWA) {
Node::redistribute_child_device_value(new_parent, child, SWA, split_len);
}
if child.has_host_value(SWA) {
Node::redistribute_child_host_value(new_parent, child, SWA, split_len);
// Device-tombstoned sides park in the host LRU.
let parent_is_tombstone = !new_parent.has_device_value(SWA);
let child_is_tombstone = !child.has_device_value(SWA);
// Device-tombstoned sides park in the host LRU. Host-locked
// halves stay out of it: in-flight IO holds them, and host
// acquire removed the node at 0->1.
let parent_parks =
!new_parent.has_device_value(SWA) && new_parent.host_lock_ref(SWA) == 0;
let child_parks = !child.has_device_value(SWA) && child.host_lock_ref(SWA) == 0;
let host_lru = tree_core.host_lru_list_mut(SWA);
if parent_is_tombstone {
if parent_parks {
host_lru.insert_mru(new_parent_id);
}
if child_is_tombstone && !host_lru.in_list(Some(child_id)) {
if child_parks && !host_lru.in_list(Some(child_id)) {
host_lru.insert_mru(child_id);
}
}
// parent inherits the swa_uuid from child for swa lock ref
// The window-boundary uuids mark the node's older edge, which the
// split moves to the parent — both tiers migrate with it.
let swa_uuid = tree_core.arena.node_mut(child_id).swa_uuid.take();
tree_core.arena.node_mut(new_parent_id).swa_uuid = swa_uuid;
let swa_host_uuid = tree_core.arena.node_mut(child_id).swa_host_uuid.take();
tree_core.arena.node_mut(new_parent_id).swa_host_uuid = swa_host_uuid;
}
fn evict_component(
@@ -1033,46 +1038,44 @@ impl<K: ChildKeyType> TreeComponent<K> for SwaComponent {
mut result: IncLockRefResult,
lock_host: bool,
) -> IncLockRefResult {
let ct = SWA;
// Lock the contiguous segment covering the trailing window.
//
// Every node in [node, boundary] is counted, tombstones included, so
// the paired release decrements the same contiguous segment with no
// carried skip state. Coverage is position-based (key length); the
// boundary node is always uuid-stamped, so a release without a uuid
// means the segment reached the root. Ledger/LRU transitions track
// only data-bearing nodes; a value materialized later under lock is
// credited to protected by set_component_device_value.
let sliding_window_size = self.sliding_window_size;
let mut swa_lock_size = 0;
let mut covered = 0;
let mut swa_uuid = None;
// Tombstoned nodes (cd.value is None) have no SWA chunk to protect
// skip them and keep walking up. This path is hit when HiCache
// backs up a FULL present internal node whose SWA was already evicted.
let mut cur = node_id;
loop {
let node = tree_core.arena.node_mut(cur);
if node.is_root() || swa_lock_size >= sliding_window_size {
if node.is_root() || covered >= sliding_window_size {
break;
}
let parent = node.parent();
if !Self::has_value(node, lock_host) {
result
.skip_lock_node_ids
.entry(ct)
.or_default()
.insert(node.id);
cur = parent;
continue;
}
let key_len = node.key.atom_len();
let has_value = Self::has_value(node, lock_host);
let value_len = Self::value_len(node, lock_host);
let newly_locked = Self::lock_ref(node, lock_host) == 0;
Self::inc_lock_ref(node, lock_host);
swa_lock_size += Self::value_len(node, lock_host);
if newly_locked {
if newly_locked && has_value {
if lock_host {
let host_lru = tree_core.host_lru_list_mut(SWA);
if host_lru.in_list(Some(cur)) {
host_lru.remove_node(cur);
}
} else {
tree_core.dec_evictable_size(SWA, key_len);
tree_core.inc_protected_size(SWA, key_len);
tree_core.dec_evictable_size(SWA, value_len);
tree_core.inc_protected_size(SWA, value_len);
}
}
if swa_lock_size >= sliding_window_size {
covered += key_len;
if covered >= sliding_window_size {
swa_uuid = Some(Self::ensure_swa_uuid(tree_core, cur, lock_host));
}
cur = parent;
@@ -1090,23 +1093,15 @@ impl<K: ChildKeyType> TreeComponent<K> for SwaComponent {
&self,
tree_core: &mut UnifiedTreeCore<K>,
node_id: NodeIdx_,
params: Option<&DecLockRefParams>,
params: &DecLockRefParams,
lock_host: bool,
) {
let ct = SWA;
let swa_uuid_for_lock = params.and_then(|p| {
if lock_host {
p.swa_uuid_for_host_lock
} else {
p.swa_uuid_for_lock
}
});
let empty = HashSet::new();
let skip_lock_node_ids = params
.and_then(|p| p.skip_lock_node_ids.get(&ct))
.unwrap_or(&empty);
let swa_uuid_for_lock = if lock_host {
params.swa_uuid_for_host_lock
} else {
params.swa_uuid_for_lock
};
// A node in skip_lock_node_ids was a tombstone when this lock was acquired.
let mut cur = node_id;
loop {
let node = tree_core.arena.node_mut(cur);
@@ -1114,30 +1109,36 @@ impl<K: ChildKeyType> TreeComponent<K> for SwaComponent {
break;
}
let parent = node.parent();
if skip_lock_node_ids.contains(&node.id) {
cur = parent;
continue;
}
let lock_ref = Self::lock_ref(node, lock_host);
if lock_ref == 0 {
cur = parent;
continue;
}
if lock_ref == 1 {
// Acquire counted every segment node and splits copy refs, so a
// zero here means the release does not mirror its acquire.
assert!(
lock_ref > 0,
"SWA segment release hit {}lock_ref=0 on node {cur}",
if lock_host { "host_" } else { "" }
);
let has_value = Self::has_value(node, lock_host);
let value_len = Self::value_len(node, lock_host);
if lock_ref == 1 && has_value {
if lock_host {
if !node.has_device_value(SWA) && node.has_host_value(SWA) {
if !node.has_device_value(SWA) {
let host_lru = tree_core.host_lru_list_mut(SWA);
if !host_lru.in_list(Some(cur)) {
host_lru.insert_mru(cur);
}
}
} else {
let key_len = node.device_value_len(SWA);
tree_core.inc_evictable_size(SWA, key_len);
tree_core.dec_protected_size(SWA, key_len);
tree_core.inc_evictable_size(SWA, value_len);
tree_core.dec_protected_size(SWA, value_len);
}
}
Self::dec_lock_ref(tree_core.arena.node_mut(cur), lock_host);
if lock_ref == 1 {
// This may have been the last lock holding the node out of
// the evictable-leaf sets; refresh it here rather than rely
// on the Full walk running after this one.
tree_core.update_evictable_leaf_sets_(cur);
}
if swa_uuid_for_lock.is_some()
&& Self::swa_uuid(tree_core.arena.node(cur), lock_host) == swa_uuid_for_lock
{
@@ -1147,15 +1148,14 @@ impl<K: ChildKeyType> TreeComponent<K> for SwaComponent {
}
}
/// Early-release the SWA lock along [node, swa_uuid_for_lock] while
/// leaving Full and Mamba locks intact.
/// Early-release the SWA lock along [node, swa_uuid_for_lock]; this
/// method touches only SWA state. The wrapping `dec_swa_lock_only` also
/// drops strictly-lower-priority co-located locks (e.g. Mamba) per the
/// receipt; the Full lock stays so the request's prefix is protected.
///
/// Called when a request's decode position has advanced past the sliding
/// window — the SWA portion of the tree lock is no longer needed but the
/// Full lock must stay so the request's prefix is protected.
///
/// Caller (UnifiedRadixCache.dec_swa_lock_only) must ensure this is
/// invoked at most once per (node, swa_uuid_for_lock) pair.
/// window. The caller must invoke this at most once per
/// (node, swa_uuid_for_lock) pair.
fn release_window_lock(
&self,
tree_core: &mut UnifiedTreeCore<K>,
@@ -1172,21 +1172,20 @@ impl<K: ChildKeyType> TreeComponent<K> for SwaComponent {
break;
}
let parent = node.parent();
// Acquire skips tombstoned nodes; release must skip them too. Same
// for nodes with lock_ref == 0 — acquire never credited them.
if !node.has_device_value(SWA) || node.device_lock_ref(SWA) == 0 {
if swa_uuid_for_lock.is_some() && node.swa_uuid == swa_uuid_for_lock {
break;
}
cur = parent;
continue;
}
assert!(
node.device_lock_ref(SWA) > 0,
"SWA window release hit lock_ref=0 on node {cur}"
);
let has_value = node.has_device_value(SWA);
let value_len = node.device_value_len(SWA);
node.dec_device_lock_ref(SWA);
if node.device_lock_ref(SWA) == 0 {
let key_len = node.key.atom_len();
tree_core.dec_protected_size(SWA, key_len);
tree_core.inc_evictable_size(SWA, key_len);
let now_unlocked = node.device_lock_ref(SWA) == 0;
if now_unlocked {
tree_core.update_evictable_leaf_sets_(cur);
}
if now_unlocked && has_value {
tree_core.dec_protected_size(SWA, value_len);
tree_core.inc_evictable_size(SWA, value_len);
if tree_core.is_evictable_device_leaf_(tree_core.arena.node(cur)) {
tree_core.evict_component_and_detach_lru_(
cur,
rust/sglang-radix-tree/src/node.rs
+6
View File
@@ -283,6 +283,12 @@ impl<K: ChildKeyType> Node<K> {
self.set_lock_ref_(slot, src_node.lock_ref_(slot));
}
/// Copy the component's host lock refcount from `src_node`.
pub fn copy_host_lock_ref(&mut self, component_type: ComponentType, src_node: &Node<K>) {
let slot = ValueSlotIdx::host(component_type);
self.set_lock_ref_(slot, src_node.lock_ref_(slot));
}
/// Split the component's device value between a new parent and the child.
pub fn redistribute_child_device_value(
parent_node: &mut Node<K>,
rust/sglang-radix-tree/src/python_bindings.rs
+55 -73
View File
@@ -1,7 +1,7 @@
//! Python bindings: the `mem_cache` extension module and its TreeCore adapter.
use std::borrow::Cow;
use std::collections::{HashMap, HashSet};
use std::collections::HashMap;
use std::sync::Mutex;
use pyo3::buffer::PyBuffer;
@@ -10,7 +10,7 @@ use pyo3::prelude::*;
use pyo3::types::{PyBytes, PyDict, PyList};
use tch::{Device, Kind, Tensor};
use crate::components::{ComponentType, FULL, MAMBA, SWA};
use crate::components::{ComponentSet, ComponentType, FULL, MAMBA, SWA};
use crate::node::ChildKeyType;
use crate::node::{KeyNamespaceRef, NodeAccessError, NodeId, TreeCoreRuntimeError};
use crate::unified_tree_core::KvCacheEvent;
@@ -648,24 +648,27 @@ impl InsertResultBinding {
#[pyclass(get_all, set_all)]
#[derive(Clone, Default)]
pub struct DecLockRefParamsBinding {
pub node_id: Option<NodeId>,
pub swa_uuid_for_lock: Option<i64>,
pub swa_uuid_for_host_lock: Option<i64>,
pub skip_lock_node_ids: HashMap<u8, HashSet<NodeId>>,
pub skipped_lock_components: Vec<u8>,
}
#[pymethods]
impl DecLockRefParamsBinding {
#[new]
#[pyo3(signature = (swa_uuid_for_lock = None, swa_uuid_for_host_lock = None, skip_lock_node_ids = None))]
#[pyo3(signature = (node_id = None, swa_uuid_for_lock = None, swa_uuid_for_host_lock = None, skipped_lock_components = Vec::new()))]
fn new(
node_id: Option<NodeId>,
swa_uuid_for_lock: Option<i64>,
swa_uuid_for_host_lock: Option<i64>,
skip_lock_node_ids: Option<HashMap<u8, HashSet<NodeId>>>,
skipped_lock_components: Vec<u8>,
) -> Self {
DecLockRefParamsBinding {
node_id,
swa_uuid_for_lock,
swa_uuid_for_host_lock,
skip_lock_node_ids: skip_lock_node_ids.unwrap_or_default(),
skipped_lock_components,
}
}
}
@@ -674,44 +677,49 @@ impl DecLockRefParamsBinding {
/// Convert into the tree core's dec-lock params.
fn to_dec_lock_ref_params(&self) -> PyResult<DecLockRefParams> {
Ok(DecLockRefParams {
node_id: self.node_id,
swa_uuid_for_lock: self.swa_uuid_for_lock,
swa_uuid_for_host_lock: self.swa_uuid_for_host_lock,
skip_lock_node_ids: self
.skip_lock_node_ids
.iter()
.map(|(ct, node_ids)| {
Ok::<_, PyErr>((parse_component_type(*ct)?, node_ids.clone()))
})
.collect::<PyResult<_>>()?,
skipped_lock_components: component_set_from_py(&self.skipped_lock_components)?,
})
}
}
/// Python-visible inc_lock_ref result; hand skip_lock_node_ids back to the
/// matching dec_lock_ref.
/// Python-visible inc_lock_ref result; the receipt (anchor node, boundary
/// uuids, skipped components) is handed back to the matching dec_lock_ref.
#[pyclass(get_all)]
pub struct IncLockRefResultBinding {
delta: Option<usize>,
node_id: Option<NodeId>,
swa_uuid_for_lock: Option<i64>,
swa_uuid_for_host_lock: Option<i64>,
skip_lock_node_ids: HashMap<u8, HashSet<NodeId>>,
skipped_lock_components: Vec<u8>,
}
impl IncLockRefResultBinding {
fn from_result(result: crate::unified_tree_core::IncLockRefResult) -> Self {
Self {
delta: result.delta,
node_id: result.node_id,
swa_uuid_for_lock: result.swa_uuid_for_lock,
swa_uuid_for_host_lock: result.swa_uuid_for_host_lock,
skip_lock_node_ids: result
.skip_lock_node_ids
.into_iter()
.map(|(ct, node_ids)| (component_type_to_u8(ct), node_ids))
skipped_lock_components: result
.skipped_lock_components
.iter()
.map(|ct| ct.idx() as u8)
.collect(),
}
}
}
/// Parse Python component-type ids into a component set.
fn component_set_from_py(component_types: &[u8]) -> PyResult<ComponentSet> {
component_types
.iter()
.map(|ct| parse_component_type(*ct))
.collect()
}
/// Convert a Python component-keyed tracker into the core's counts.
fn tracker_from_py(tracker: HashMap<u8, usize>) -> PyResult<HashMap<ComponentType, usize>> {
tracker
@@ -1074,23 +1082,17 @@ impl<K: ChildKeyType + Send + Sync> TreeCoreBinding<K> {
cache_actions_to_py(py, actions)
}
/// Bump the reference count on a node's component locks.
/// Bump the reference count on a node's component locks; the listed
/// components are left untaken and recorded in the receipt.
fn inc_lock_ref(
&self,
py: Python<'_>,
node_id: NodeId,
skip_lock_components: Option<Vec<u8>>,
skip_lock_components: Vec<u8>,
) -> PyResult<IncLockRefResultBinding> {
let skip_lock_components = skip_lock_components
.unwrap_or_default()
.into_iter()
.map(parse_component_type)
.collect::<PyResult<Vec<_>>>()?;
let skip = component_set_from_py(&skip_lock_components)?;
let result = py
.allow_threads(|| {
self.core()
.inc_lock_ref_with_skip(node_id, &skip_lock_components)
})
.allow_threads(|| self.core().inc_lock_ref(node_id, skip))
.map_err(node_access_error)?;
Ok(IncLockRefResultBinding::from_result(result))
}
@@ -1100,11 +1102,11 @@ impl<K: ChildKeyType + Send + Sync> TreeCoreBinding<K> {
&self,
py: Python<'_>,
node_id: NodeId,
params: Option<&DecLockRefParamsBinding>,
params: &DecLockRefParamsBinding,
skip_swa: bool,
) -> PyResult<()> {
let params = params.map(|p| p.to_dec_lock_ref_params()).transpose()?;
py.allow_threads(|| self.core().dec_lock_ref(node_id, params.as_ref(), skip_swa))
let params = params.to_dec_lock_ref_params()?;
py.allow_threads(|| self.core().dec_lock_ref(node_id, &params, skip_swa))
.map_err(node_access_error)?;
Ok(())
}
@@ -1115,22 +1117,16 @@ impl<K: ChildKeyType + Send + Sync> TreeCoreBinding<K> {
&self,
py: Python<'_>,
node_id: NodeId,
swa_uuid_for_lock: Option<i64>,
skip_lock_node_ids: Option<HashMap<u8, HashSet<NodeId>>>,
params: &DecLockRefParamsBinding,
) -> PyResult<(Py<PyDict>, Py<PyDict>)> {
let skip_lock_node_ids = skip_lock_node_ids
.unwrap_or_default()
.into_iter()
.map(|(ct, node_ids)| Ok((parse_component_type(ct)?, node_ids)))
.collect::<PyResult<HashMap<_, _>>>()?;
let params = params.to_dec_lock_ref_params()?;
let (device_frees, host_frees) = py
.allow_threads(|| {
let mut device_frees = HashMap::new();
let mut host_frees = HashMap::new();
self.core().dec_swa_lock_only_with_skip(
self.core().dec_swa_lock_only(
node_id,
swa_uuid_for_lock,
Some(&skip_lock_node_ids),
&params,
&mut device_frees,
&mut host_frees,
)?;
@@ -1733,16 +1729,7 @@ impl<K: ChildKeyType + Send + Sync> TreeCoreBinding<K> {
let result = py
.allow_threads(|| self.core().inc_host_lock_ref(node_id))
.map_err(node_access_error)?;
Ok(IncLockRefResultBinding {
delta: result.delta,
swa_uuid_for_lock: result.swa_uuid_for_lock,
swa_uuid_for_host_lock: result.swa_uuid_for_host_lock,
skip_lock_node_ids: result
.skip_lock_node_ids
.into_iter()
.map(|(ct, node_ids)| (component_type_to_u8(ct), node_ids))
.collect(),
})
Ok(IncLockRefResultBinding::from_result(result))
}
/// Decrease the reference count on a node's host-side component locks.
@@ -1750,10 +1737,10 @@ impl<K: ChildKeyType + Send + Sync> TreeCoreBinding<K> {
&self,
py: Python<'_>,
node_id: NodeId,
params: Option<&DecLockRefParamsBinding>,
params: &DecLockRefParamsBinding,
) -> PyResult<()> {
let params = params.map(|p| p.to_dec_lock_ref_params()).transpose()?;
py.allow_threads(|| self.core().dec_host_lock_ref(node_id, params.as_ref()))
let params = params.to_dec_lock_ref_params()?;
py.allow_threads(|| self.core().dec_host_lock_ref(node_id, &params))
.map_err(node_access_error)?;
Ok(())
}
@@ -2371,23 +2358,24 @@ macro_rules! tree_core_binding {
}
/// Bump the reference count on a node's component locks.
#[pyo3(signature = (node_id, skip_lock_components = None))]
#[pyo3(signature = (node_id, skip_lock_components = Vec::new()))]
fn inc_lock_ref(
&self,
py: Python<'_>,
node_id: NodeId,
skip_lock_components: Option<Vec<u8>>,
skip_lock_components: Vec<u8>,
) -> PyResult<IncLockRefResultBinding> {
self.inner.inc_lock_ref(py, node_id, skip_lock_components)
}
/// Decrease the reference count on a node's component locks.
#[pyo3(signature = (node_id, params = None, skip_swa = false))]
/// Decrease the reference count on a node's component locks. The
/// receipt is required: a release must replay its acquire's evidence.
#[pyo3(signature = (node_id, params, skip_swa = false))]
fn dec_lock_ref(
&self,
py: Python<'_>,
node_id: NodeId,
params: Option<&DecLockRefParamsBinding>,
params: &DecLockRefParamsBinding,
skip_swa: bool,
) -> PyResult<()> {
self.inner.dec_lock_ref(py, node_id, params, skip_swa)
@@ -2395,20 +2383,14 @@ macro_rules! tree_core_binding {
/// Early-release the SWA portion of a request's tree lock; returns this
/// release's per-component (device_frees, host_frees).
#[pyo3(signature = (node_id, swa_uuid_for_lock = None, skip_lock_node_ids = None))]
#[pyo3(signature = (node_id, params))]
fn dec_swa_lock_only(
&self,
py: Python<'_>,
node_id: NodeId,
swa_uuid_for_lock: Option<i64>,
skip_lock_node_ids: Option<HashMap<u8, HashSet<NodeId>>>,
params: &DecLockRefParamsBinding,
) -> PyResult<(Py<PyDict>, Py<PyDict>)> {
self.inner.dec_swa_lock_only(
py,
node_id,
swa_uuid_for_lock,
skip_lock_node_ids,
)
self.inner.dec_swa_lock_only(py, node_id, params)
}
/// Store a component's device value on a node (the SWA rebuild write-back).
@@ -2817,12 +2799,12 @@ macro_rules! tree_core_binding {
}
/// Decrease the reference count on a node's host-side component locks.
#[pyo3(signature = (node_id, params = None))]
/// The receipt is required, as for dec_lock_ref.
fn dec_host_lock_ref(
&self,
py: Python<'_>,
node_id: NodeId,
params: Option<&DecLockRefParamsBinding>,
params: &DecLockRefParamsBinding,
) -> PyResult<()> {
self.inner.dec_host_lock_ref(py, node_id, params)
}
rust/sglang-radix-tree/src/tests/components/base.rs
+1 -1
View File
@@ -70,7 +70,7 @@ impl TreeComponent<Vec<i64>> for DefaultComponentForTest {
&self,
tree_core: &mut UnifiedTreeCore<Vec<i64>>,
node_id: NodeIdx_,
params: Option<&DecLockRefParams>,
params: &DecLockRefParams,
lock_host: bool,
) {
unimplemented!()
rust/sglang-radix-tree/src/tests/components/full.rs
+120 -80
View File
@@ -1,5 +1,5 @@
use super::*;
use crate::components::FULL;
use crate::components::{ComponentSet, FULL};
use crate::node::NodeAccessError;
use crate::test_utils::accumulate_step;
use crate::unified_tree_core::CacheInitParams;
@@ -625,10 +625,9 @@ fn inc_lock_ref_locks_the_device_path() {
let mut tc = core();
let (n1, n2) = lock_chain(&mut tc);
let result = tc
.inc_lock_ref(tc.arena.node(n2).id)
.inc_lock_ref(tc.arena.node(n2).id, ComponentSet::EMPTY)
.expect("live test node");
assert_eq!(result.delta, Some(5));
assert!(result.skip_lock_node_ids.is_empty());
assert_eq!(tc.arena.device_lock_ref(n1, FULL), 1);
assert_eq!(tc.arena.device_lock_ref(n2, FULL), 1);
let state = tc.component_state(FULL);
@@ -641,10 +640,10 @@ fn inc_lock_ref_locks_the_device_path() {
fn inc_lock_ref_again_only_bumps_the_refs() {
let mut tc = core();
let (n1, n2) = lock_chain(&mut tc);
tc.inc_lock_ref(tc.arena.node(n2).id)
tc.inc_lock_ref(tc.arena.node(n2).id, ComponentSet::EMPTY)
.expect("live test node");
let result = tc
.inc_lock_ref(tc.arena.node(n2).id)
.inc_lock_ref(tc.arena.node(n2).id, ComponentSet::EMPTY)
.expect("live test node");
assert_eq!(result.delta, Some(0));
assert_eq!(tc.arena.device_lock_ref(n1, FULL), 2);
@@ -659,10 +658,10 @@ fn inc_lock_ref_counts_only_newly_locked_nodes() {
// n1 is already locked via its own path; locking n2 moves only n2's tokens.
let mut tc = core();
let (n1, n2) = lock_chain(&mut tc);
tc.inc_lock_ref(tc.arena.node(n1).id)
tc.inc_lock_ref(tc.arena.node(n1).id, ComponentSet::EMPTY)
.expect("live test node");
let result = tc
.inc_lock_ref(tc.arena.node(n2).id)
.inc_lock_ref(tc.arena.node(n2).id, ComponentSet::EMPTY)
.expect("live test node");
assert_eq!(result.delta, Some(3));
assert_eq!(tc.arena.device_lock_ref(n1, FULL), 2);
@@ -673,8 +672,9 @@ fn inc_lock_ref_counts_only_newly_locked_nodes() {
}
#[test]
fn inc_lock_ref_collects_the_evicted_bottom_segment() {
// n2 and n3 are evicted (no device value): the walk records both and locks only n1.
fn inc_lock_ref_counts_the_evicted_bottom_segment() {
// n2 and n3 are evicted (no device value): counted in the segment with no
// ledger move; only n1's tokens turn protected.
let mut tc = core();
let root = tc.arena.root();
let n1 = tc
@@ -709,16 +709,12 @@ fn inc_lock_ref_collects_the_evicted_bottom_segment() {
tc.component_state_mut(FULL).evictable_size = 2;
tc.evictable_device_leaves.add(n1);
let result = tc
.inc_lock_ref(tc.arena.node(n3).id)
.inc_lock_ref(tc.arena.node(n3).id, ComponentSet::EMPTY)
.expect("live test node");
assert_eq!(result.delta, Some(2));
assert_eq!(
result.skip_lock_node_ids[&FULL],
HashSet::from([tc.arena.node(n2).id, tc.arena.node(n3).id])
);
assert_eq!(tc.arena.device_lock_ref(n1, FULL), 1);
assert_eq!(tc.arena.device_lock_ref(n2, FULL), 0);
assert_eq!(tc.arena.device_lock_ref(n3, FULL), 0);
assert_eq!(tc.arena.device_lock_ref(n2, FULL), 1);
assert_eq!(tc.arena.device_lock_ref(n3, FULL), 1);
// The locked ancestor leaves the D-leaf set.
assert!(!tc.evictable_device_leaves.contains(n1));
}
@@ -728,15 +724,18 @@ fn lock_round_trips_on_a_root_anchor_are_noops() {
let mut tc = core();
let root = tc.arena.root();
let result = tc
.inc_lock_ref(tc.arena.node(root).id)
.inc_lock_ref(tc.arena.node(root).id, ComponentSet::EMPTY)
.expect("live test node");
assert_eq!(result.delta, Some(0));
assert!(result.skip_lock_node_ids.is_empty());
// The protected root keeps its construction-time lock through the pair.
assert_eq!(tc.arena.device_lock_ref(root, FULL), 1);
tc.dec_lock_ref(
tc.arena.node(root).id,
/* params = */ None,
/* params = */
&DecLockRefParams {
skipped_lock_components: ComponentSet::EMPTY,
..Default::default()
},
/* skip_swa = */ false,
)
.expect("live test node");
@@ -760,7 +759,7 @@ fn lock_walks_stop_at_the_root_of_a_salted_chain() {
.set_device_value(n1, FULL, Tensor::from_slice(&[0i64, 1]));
tc.component_state_mut(FULL).evictable_size = 2;
let result = tc
.inc_lock_ref(tc.arena.node(n1).id)
.inc_lock_ref(tc.arena.node(n1).id, ComponentSet::EMPTY)
.expect("live test node");
assert_eq!(result.delta, Some(2));
assert_eq!(tc.arena.device_lock_ref(n1, FULL), 1);
@@ -769,7 +768,11 @@ fn lock_walks_stop_at_the_root_of_a_salted_chain() {
// The release walk stops at the same boundary.
tc.dec_lock_ref(
tc.arena.node(n1).id,
/* params = */ None,
/* params = */
&DecLockRefParams {
skipped_lock_components: ComponentSet::EMPTY,
..Default::default()
},
/* skip_swa = */ false,
)
.expect("live test node");
@@ -798,17 +801,20 @@ fn lock_walks_treat_a_present_but_empty_value_as_device_on() {
Tensor::from_slice(&empty),
);
let result = tc
.inc_lock_ref(tc.arena.node(n1).id)
.inc_lock_ref(tc.arena.node(n1).id, ComponentSet::EMPTY)
.expect("live test node");
// A present-but-empty value is device-on (Python `value is not None`):
// locked, zero tokens moved.
assert_eq!(result.delta, Some(0));
assert_eq!(tc.arena.device_lock_ref(n1, FULL), 1);
assert!(result.skip_lock_node_ids.is_empty());
// The release side moves the same zero tokens back.
tc.dec_lock_ref(
tc.arena.node(n1).id,
/* params = */ None,
/* params = */
&DecLockRefParams {
skipped_lock_components: ComponentSet::EMPTY,
..Default::default()
},
/* skip_swa = */ false,
)
.expect("live test node");
@@ -822,11 +828,15 @@ fn lock_walks_treat_a_present_but_empty_value_as_device_on() {
fn dec_lock_ref_unlocks_and_restores_sizes() {
let mut tc = core();
let (n1, n2) = lock_chain(&mut tc);
tc.inc_lock_ref(tc.arena.node(n2).id)
tc.inc_lock_ref(tc.arena.node(n2).id, ComponentSet::EMPTY)
.expect("live test node");
tc.dec_lock_ref(
tc.arena.node(n2).id,
/* params = */ None,
/* params = */
&DecLockRefParams {
skipped_lock_components: ComponentSet::EMPTY,
..Default::default()
},
/* skip_swa = */ false,
)
.expect("live test node");
@@ -875,19 +885,15 @@ fn dec_lock_ref_replays_the_skip_set() {
.set_device_value(n1, FULL, Tensor::from_slice(&[0i64, 1]));
tc.component_state_mut(FULL).evictable_size = 2;
let result = tc
.inc_lock_ref(tc.arena.node(n3).id)
.inc_lock_ref(tc.arena.node(n3).id, ComponentSet::EMPTY)
.expect("live test node");
let params = DecLockRefParams {
skip_lock_node_ids: result.skip_lock_node_ids,
skipped_lock_components: result.skipped_lock_components,
..Default::default()
};
// The still-evicted n2 and n3 are skipped instead of tripping the lock asserts.
tc.dec_lock_ref(
tc.arena.node(n3).id,
Some(&params),
/* skip_swa = */ false,
)
.expect("live test node");
tc.dec_lock_ref(tc.arena.node(n3).id, &params, /* skip_swa = */ false)
.expect("live test node");
assert_eq!(tc.arena.device_lock_ref(n1, FULL), 0);
assert_eq!(tc.arena.device_lock_ref(n2, FULL), 0);
assert_eq!(tc.arena.device_lock_ref(n3, FULL), 0);
@@ -897,7 +903,7 @@ fn dec_lock_ref_replays_the_skip_set() {
}
#[test]
fn temp_lock_skips_the_evicted_anchor_and_mirrors_on_release() {
fn temp_lock_counts_the_evicted_anchor_and_mirrors_on_release() {
// Chain root -> a -> y -> anchor with FULL device values; the anchor is evicted.
let mut tc = core();
let root = tc.arena.root();
@@ -933,34 +939,32 @@ fn temp_lock_skips_the_evicted_anchor_and_mirrors_on_release() {
tc.arena
.set_device_value(y, FULL, Tensor::from_slice(&[0i64]));
tc.component_state_mut(FULL).evictable_size = 3;
// The temp lock records the evicted anchor and locks only its ancestors.
// The temp lock counts the evicted anchor too (no ledger move).
let temp_lock = tc
.inc_lock_ref(tc.arena.node(anchor).id)
.expect("live test node");
assert_eq!(tc.arena.device_lock_ref(anchor, FULL), 0);
assert_eq!(tc.arena.device_lock_ref(y, FULL), 1);
assert_eq!(tc.arena.device_lock_ref(a, FULL), 1);
assert_eq!(
temp_lock.skip_lock_node_ids[&FULL],
HashSet::from([tc.arena.node(anchor).id])
);
// A load-back restores the anchor; the second acquire covers it.
tc.arena
.set_device_value(anchor, FULL, Tensor::from_slice(&[0i64]));
let second_lock = tc
.inc_lock_ref(tc.arena.node(anchor).id)
.inc_lock_ref(tc.arena.node(anchor).id, ComponentSet::EMPTY)
.expect("live test node");
assert_eq!(tc.arena.device_lock_ref(anchor, FULL), 1);
assert_eq!(tc.arena.device_lock_ref(y, FULL), 1);
assert_eq!(tc.arena.device_lock_ref(a, FULL), 1);
// A load-back restores the anchor under the held lock: credited to
// protected, as the production commit does. The second acquire stacks.
tc.arena
.set_device_value(anchor, FULL, Tensor::from_slice(&[0i64]));
tc.inc_protected_size(FULL, 1);
let second_lock = tc
.inc_lock_ref(tc.arena.node(anchor).id, ComponentSet::EMPTY)
.expect("live test node");
assert_eq!(tc.arena.device_lock_ref(anchor, FULL), 2);
assert_eq!(tc.arena.device_lock_ref(y, FULL), 2);
assert_eq!(tc.arena.device_lock_ref(a, FULL), 2);
// Releasing the temp lock mirrors its skip set: the anchor keeps its lock.
// Each release takes back exactly its own refs.
let temp_params = DecLockRefParams {
skip_lock_node_ids: temp_lock.skip_lock_node_ids,
skipped_lock_components: temp_lock.skipped_lock_components,
..Default::default()
};
tc.dec_lock_ref(
tc.arena.node(anchor).id,
Some(&temp_params),
&temp_params,
/* skip_swa = */ false,
)
.expect("live test node");
@@ -968,12 +972,12 @@ fn temp_lock_skips_the_evicted_anchor_and_mirrors_on_release() {
assert_eq!(tc.arena.device_lock_ref(y, FULL), 1);
assert_eq!(tc.arena.device_lock_ref(a, FULL), 1);
let second_params = DecLockRefParams {
skip_lock_node_ids: second_lock.skip_lock_node_ids,
skipped_lock_components: second_lock.skipped_lock_components,
..Default::default()
};
tc.dec_lock_ref(
tc.arena.node(anchor).id,
Some(&second_params),
&second_params,
/* skip_swa = */ false,
)
.expect("live test node");
@@ -983,9 +987,9 @@ fn temp_lock_skips_the_evicted_anchor_and_mirrors_on_release() {
}
#[test]
#[should_panic(expected = "has no FULL device value")]
fn dec_lock_ref_panics_without_replaying_the_skip_set() {
// Dropping the acquire's skip set makes the release walk hit the tombstone.
#[should_panic(expected = "FULL segment release hit lock_ref=0")]
fn dec_lock_ref_panics_on_double_release() {
// The second, unpaired release hits the already-unlocked segment.
let mut tc = core();
let root = tc.arena.root();
let n1 = tc
@@ -1009,11 +1013,25 @@ fn dec_lock_ref_panics_without_replaying_the_skip_set() {
tc.arena
.set_device_value(n1, FULL, Tensor::from_slice(&[0i64, 1]));
tc.component_state_mut(FULL).evictable_size = 2;
tc.inc_lock_ref(tc.arena.node(n2).id)
tc.inc_lock_ref(tc.arena.node(n2).id, ComponentSet::EMPTY)
.expect("live test node");
tc.dec_lock_ref(
tc.arena.node(n2).id,
/* params = */ None,
/* params = */
&DecLockRefParams {
skipped_lock_components: ComponentSet::EMPTY,
..Default::default()
},
/* skip_swa = */ false,
)
.expect("live test node");
tc.dec_lock_ref(
tc.arena.node(n2).id,
/* params = */
&DecLockRefParams {
skipped_lock_components: ComponentSet::EMPTY,
..Default::default()
},
/* skip_swa = */ false,
)
.expect("live test node");
@@ -1023,11 +1041,15 @@ fn dec_lock_ref_panics_without_replaying_the_skip_set() {
fn dec_lock_ref_with_skip_swa_still_releases_full() {
let mut tc = core();
let (_n1, n2) = lock_chain(&mut tc);
tc.inc_lock_ref(tc.arena.node(n2).id)
tc.inc_lock_ref(tc.arena.node(n2).id, ComponentSet::EMPTY)
.expect("live test node");
tc.dec_lock_ref(
tc.arena.node(n2).id,
/* params = */ None,
/* params = */
&DecLockRefParams {
skipped_lock_components: ComponentSet::EMPTY,
..Default::default()
},
/* skip_swa = */ true,
)
.expect("live test node");
@@ -1039,13 +1061,17 @@ fn nested_locks_release_pairwise() {
// Two acquires then two releases: sizes move only on the outermost pair.
let mut tc = core();
let (_n1, n2) = lock_chain(&mut tc);
tc.inc_lock_ref(tc.arena.node(n2).id)
tc.inc_lock_ref(tc.arena.node(n2).id, ComponentSet::EMPTY)
.expect("live test node");
tc.inc_lock_ref(tc.arena.node(n2).id)
tc.inc_lock_ref(tc.arena.node(n2).id, ComponentSet::EMPTY)
.expect("live test node");
tc.dec_lock_ref(
tc.arena.node(n2).id,
/* params = */ None,
/* params = */
&DecLockRefParams {
skipped_lock_components: ComponentSet::EMPTY,
..Default::default()
},
/* skip_swa = */ false,
)
.expect("live test node");
@@ -1055,7 +1081,11 @@ fn nested_locks_release_pairwise() {
assert!(!tc.evictable_device_leaves.contains(n2));
tc.dec_lock_ref(
tc.arena.node(n2).id,
/* params = */ None,
/* params = */
&DecLockRefParams {
skipped_lock_components: ComponentSet::EMPTY,
..Default::default()
},
/* skip_swa = */ false,
)
.expect("live test node");
@@ -1066,13 +1096,17 @@ fn nested_locks_release_pairwise() {
}
#[test]
#[should_panic(expected = "is not locked")]
#[should_panic(expected = "FULL segment release hit lock_ref=0")]
fn dec_lock_ref_panics_on_an_unlocked_node() {
let mut tc = core();
let (_n1, n2) = lock_chain(&mut tc);
tc.dec_lock_ref(
tc.arena.node(n2).id,
/* params = */ None,
/* params = */
&DecLockRefParams {
skipped_lock_components: ComponentSet::EMPTY,
..Default::default()
},
/* skip_swa = */ false,
)
.expect("live test node");
@@ -1105,7 +1139,7 @@ fn inc_lock_ref_panics_on_an_evicted_ancestor() {
tc.arena
.set_device_value(n2, FULL, Tensor::from_slice(&[0i64, 1, 2]));
tc.component_state_mut(FULL).evictable_size = 3;
tc.inc_lock_ref(tc.arena.node(n2).id)
tc.inc_lock_ref(tc.arena.node(n2).id, ComponentSet::EMPTY)
.expect("live test node");
}
@@ -1125,7 +1159,7 @@ fn inc_lock_ref_panics_when_evictable_size_is_unaccounted() {
.unwrap();
tc.arena
.set_device_value(n1, FULL, Tensor::from_slice(&[0i64]));
tc.inc_lock_ref(tc.arena.node(n1).id)
tc.inc_lock_ref(tc.arena.node(n1).id, ComponentSet::EMPTY)
.expect("live test node");
}
@@ -1140,7 +1174,11 @@ fn dec_lock_ref_panics_on_protected_underflow() {
.set_lock_ref_(ValueSlotIdx::device(FULL), 1);
tc.dec_lock_ref(
tc.arena.node(n2).id,
/* params = */ None,
/* params = */
&DecLockRefParams {
skipped_lock_components: ComponentSet::EMPTY,
..Default::default()
},
/* skip_swa = */ false,
)
.expect("live test node");
@@ -1183,7 +1221,6 @@ fn inc_host_lock_ref_pins_the_backuped_anchor() {
.inc_host_lock_ref(tc.arena.node(node).id)
.expect("live test node");
assert_eq!(result.delta, None);
assert!(result.skip_lock_node_ids.is_empty());
assert_eq!(tc.arena.host_lock_ref(node, FULL), 1);
// The pinned anchor leaves the H-leaf set; the device tier is untouched.
assert!(!tc.evictable_host_leaves.contains(node));
@@ -1256,7 +1293,7 @@ fn host_lock_round_trips_on_a_root_anchor_are_noops() {
.expect("live test node");
assert_eq!(result.delta, None);
assert_eq!(tc.arena.host_lock_ref(root, FULL), 0);
tc.dec_host_lock_ref(tc.arena.node(root).id, /* params = */ None)
tc.dec_host_lock_ref(tc.arena.node(root).id, &DecLockRefParams::default())
.expect("live test node");
assert_eq!(tc.arena.host_lock_ref(root, FULL), 0);
}
@@ -1281,7 +1318,7 @@ fn dec_host_lock_ref_unpins_and_restores_the_h_leaf_set() {
tc.component_state_mut(FULL).evictable_size = 7;
tc.inc_host_lock_ref(tc.arena.node(node).id)
.expect("live test node");
tc.dec_host_lock_ref(tc.arena.node(node).id, /* params = */ None)
tc.dec_host_lock_ref(tc.arena.node(node).id, &DecLockRefParams::default())
.expect("live test node");
assert_eq!(tc.arena.host_lock_ref(node, FULL), 0);
assert!(tc.evictable_host_leaves.contains(node));
@@ -1294,7 +1331,7 @@ fn dec_host_lock_ref_unpins_and_restores_the_h_leaf_set() {
fn dec_host_lock_ref_on_an_unlocked_anchor_is_a_noop() {
let mut tc = core();
let node = host_lock_anchor(&mut tc);
tc.dec_host_lock_ref(tc.arena.node(node).id, /* params = */ None)
tc.dec_host_lock_ref(tc.arena.node(node).id, &DecLockRefParams::default())
.expect("live test node");
assert_eq!(tc.arena.host_lock_ref(node, FULL), 0);
}
@@ -1307,7 +1344,7 @@ fn dec_host_lock_ref_keeps_the_counter_when_the_host_value_is_gone() {
tc.inc_host_lock_ref(tc.arena.node(node).id)
.expect("live test node");
let _ = tc.arena.take_host_value(node, FULL);
tc.dec_host_lock_ref(tc.arena.node(node).id, /* params = */ None)
tc.dec_host_lock_ref(tc.arena.node(node).id, &DecLockRefParams::default())
.expect("live test node");
assert_eq!(tc.arena.host_lock_ref(node, FULL), 1);
}
@@ -1318,7 +1355,7 @@ fn host_lock_round_trip_under_write_back_is_a_pure_counter() {
let (_n1, n2) = lock_chain(&mut tc);
tc.inc_host_lock_ref(tc.arena.node(n2).id)
.expect("live test node");
tc.dec_host_lock_ref(tc.arena.node(n2).id, /* params = */ None)
tc.dec_host_lock_ref(tc.arena.node(n2).id, &DecLockRefParams::default())
.expect("live test node");
assert_eq!(tc.arena.host_lock_ref(n2, FULL), 0);
let state = tc.component_state(FULL);
@@ -1346,7 +1383,10 @@ fn release_host_arm_updates_the_h_leaf_set_without_the_dispatcher() {
tc.inc_host_lock_ref(tc.arena.node(node).id)
.expect("live test node");
FullComponent.release_component_lock(
&mut tc, node, /* params = */ None, /* lock_host = */ true,
&mut tc,
node,
&DecLockRefParams::default(),
/* lock_host = */ true,
);
assert!(tc.evictable_host_leaves.contains(node));
}
@@ -1359,11 +1399,11 @@ fn nested_host_locks_release_pairwise() {
.expect("live test node");
tc.inc_host_lock_ref(tc.arena.node(node).id)
.expect("live test node");
tc.dec_host_lock_ref(tc.arena.node(node).id, /* params = */ None)
tc.dec_host_lock_ref(tc.arena.node(node).id, &DecLockRefParams::default())
.expect("live test node");
assert_eq!(tc.arena.host_lock_ref(node, FULL), 1);
assert!(!tc.evictable_host_leaves.contains(node));
tc.dec_host_lock_ref(tc.arena.node(node).id, /* params = */ None)
tc.dec_host_lock_ref(tc.arena.node(node).id, &DecLockRefParams::default())
.expect("live test node");
assert_eq!(tc.arena.host_lock_ref(node, FULL), 0);
assert!(tc.evictable_host_leaves.contains(node));
rust/sglang-radix-tree/src/tests/components/mamba.rs
+89 -63
View File
@@ -1,5 +1,5 @@
use super::*;
use crate::components::{FULL, MAMBA, SWA};
use crate::components::{ComponentSet, FULL, MAMBA, SWA};
use crate::test_utils::{accumulate_step, action_kinds};
use crate::unified_lru_list::UnifiedLRUList;
@@ -292,7 +292,6 @@ fn device_lock_moves_the_slot_between_evictable_and_protected_once() {
IncLockRefResult::default(),
/* lock_host = */ false,
);
assert!(result.skip_lock_node_ids.is_empty());
assert_eq!(tc.evictable_size_(MAMBA), 0);
assert_eq!(tc.protected_size_(MAMBA), 1);
mamba.acquire_component_lock(
@@ -303,25 +302,34 @@ fn device_lock_moves_the_slot_between_evictable_and_protected_once() {
);
assert_eq!(tc.protected_size_(MAMBA), 1);
assert_eq!(tc.arena.node(a).device_lock_ref(MAMBA), 2);
mamba.release_component_lock(&mut tc, a, None, /* lock_host = */ false);
mamba.release_component_lock(
&mut tc,
a,
&DecLockRefParams::default(),
/* lock_host = */ false,
);
assert_eq!(tc.protected_size_(MAMBA), 1);
mamba.release_component_lock(&mut tc, a, None, /* lock_host = */ false);
mamba.release_component_lock(
&mut tc,
a,
&DecLockRefParams::default(),
/* lock_host = */ false,
);
assert_eq!(tc.evictable_size_(MAMBA), 1);
assert_eq!(tc.protected_size_(MAMBA), 0);
assert_eq!(tc.arena.node(a).device_lock_ref(MAMBA), 0);
}
#[test]
fn skip_aware_lock_records_only_the_mamba_target() {
fn lock_without_mamba_records_the_receipt_and_leaves_mamba_evictable() {
let (mut tc, parent, leaf) = hybrid_lock_core();
let leaf_handle = tc.arena.node(leaf).id;
let result = tc
.inc_lock_ref_with_skip(leaf_handle, &[MAMBA])
.inc_lock_ref(leaf_handle, ComponentSet::of(MAMBA))
.expect("live test node");
assert_eq!(result.skip_lock_node_ids[&MAMBA].len(), 1);
assert!(result.skip_lock_node_ids[&MAMBA].contains(&leaf_handle));
assert!(result.skipped_lock_components.contains(MAMBA));
assert_eq!(tc.arena.node(parent).device_lock_ref(MAMBA), 0);
assert_eq!(tc.arena.node(leaf).device_lock_ref(MAMBA), 0);
assert_eq!(tc.evictable_size_(MAMBA), 2);
@@ -329,36 +337,45 @@ fn skip_aware_lock_records_only_the_mamba_target() {
assert_eq!(tc.arena.node(parent).device_lock_ref(FULL), 1);
assert_eq!(tc.arena.node(leaf).device_lock_ref(FULL), 1);
tc.dec_lock_ref(
leaf_handle,
Some(&DecLockRefParams {
swa_uuid_for_lock: result.swa_uuid_for_lock,
skip_lock_node_ids: result.skip_lock_node_ids,
..Default::default()
}),
/* skip_swa = */ false,
)
.expect("live test node");
// The receipt replays exactly what was taken: FULL only.
let params = DecLockRefParams {
swa_uuid_for_lock: result.swa_uuid_for_lock,
skipped_lock_components: result.skipped_lock_components,
..Default::default()
};
tc.dec_lock_ref(leaf_handle, &params, /* skip_swa = */ false)
.expect("live test node");
assert_eq!(tc.arena.node(parent).device_lock_ref(FULL), 0);
assert_eq!(tc.arena.node(leaf).device_lock_ref(FULL), 0);
assert_eq!(tc.evictable_size_(MAMBA), 2);
assert_eq!(tc.protected_size_(MAMBA), 0);
}
#[test]
fn swa_only_release_honors_a_skipped_mamba_target() {
fn swa_only_release_spares_another_holders_mamba_lock() {
let (mut tc, _parent, leaf) = hybrid_lock_core();
let leaf_handle = tc.arena.node(leaf).id;
let owner = tc.inc_lock_ref(leaf_handle).expect("live test node");
let skipped = tc
.inc_lock_ref_with_skip(leaf_handle, &[MAMBA])
let owner = tc
.inc_lock_ref(leaf_handle, ComponentSet::EMPTY)
.expect("live test node");
let holder = tc
.inc_lock_ref(leaf_handle, ComponentSet::of(MAMBA))
.expect("live test node");
assert!(!owner.skipped_lock_components.contains(MAMBA));
assert!(holder.skipped_lock_components.contains(MAMBA));
assert_eq!(tc.arena.node(leaf).device_lock_ref(MAMBA), 1);
// The holder's early SWA release must not drop the owner's mamba lock.
let holder_params = DecLockRefParams {
swa_uuid_for_lock: holder.swa_uuid_for_lock,
skipped_lock_components: holder.skipped_lock_components,
..Default::default()
};
let mut device_frees = HashMap::new();
let mut host_frees = HashMap::new();
tc.dec_swa_lock_only_with_skip(
tc.dec_swa_lock_only(
leaf_handle,
skipped.swa_uuid_for_lock,
Some(&skipped.skip_lock_node_ids),
&holder_params,
&mut device_frees,
&mut host_frees,
)
@@ -369,33 +386,22 @@ fn swa_only_release_honors_a_skipped_mamba_target() {
assert_eq!(tc.arena.node(leaf).device_lock_ref(MAMBA), 1);
assert_eq!(tc.protected_size_(MAMBA), 1);
let skipped_params = DecLockRefParams {
swa_uuid_for_lock: skipped.swa_uuid_for_lock,
skip_lock_node_ids: skipped.skip_lock_node_ids,
..Default::default()
};
tc.dec_lock_ref(
leaf_handle,
Some(&skipped_params),
/* skip_swa = */ true,
)
.expect("live test node");
tc.dec_lock_ref(leaf_handle, &holder_params, /* skip_swa = */ true)
.expect("live test node");
assert_eq!(tc.arena.node(leaf).device_lock_ref(MAMBA), 1);
let owner_params = DecLockRefParams {
swa_uuid_for_lock: owner.swa_uuid_for_lock,
skip_lock_node_ids: owner.skip_lock_node_ids,
skipped_lock_components: owner.skipped_lock_components,
..Default::default()
};
tc.dec_lock_ref(
leaf_handle,
Some(&owner_params),
/* skip_swa = */ false,
)
.expect("live test node");
tc.dec_lock_ref(leaf_handle, &owner_params, /* skip_swa = */ false)
.expect("live test node");
assert_eq!(tc.arena.node(leaf).device_lock_ref(MAMBA), 0);
assert_eq!(tc.protected_size_(MAMBA), 0);
}
#[test]
fn tombstone_lock_is_recorded_and_replayed_at_release() {
fn tombstone_lock_is_counted_with_no_ledger_move() {
let mut tc = mamba_core(/* page_size = */ 1);
let [a] = chain::<1>(&mut tc);
let mamba = mamba_component();
@@ -405,14 +411,15 @@ fn tombstone_lock_is_recorded_and_replayed_at_release() {
IncLockRefResult::default(),
/* lock_host = */ false,
);
assert!(result.skip_lock_node_ids[&MAMBA].contains(&tc.arena.node(a).id));
assert_eq!(tc.arena.node(a).device_lock_ref(MAMBA), 0);
// The replayed skip set keeps the release from touching the node.
assert_eq!(tc.arena.node(a).device_lock_ref(MAMBA), 1);
assert_eq!(tc.evictable_size_(MAMBA), 0);
assert_eq!(tc.protected_size_(MAMBA), 0);
// The paired release decrements the counted tombstone, ledger untouched.
let params = DecLockRefParams {
skip_lock_node_ids: result.skip_lock_node_ids.clone(),
skipped_lock_components: result.skipped_lock_components,
..DecLockRefParams::default()
};
mamba.release_component_lock(&mut tc, a, Some(&params), /* lock_host = */ false);
mamba.release_component_lock(&mut tc, a, &params, /* lock_host = */ false);
assert_eq!(tc.arena.node(a).device_lock_ref(MAMBA), 0);
assert_eq!(tc.evictable_size_(MAMBA), 0);
}
@@ -428,8 +435,12 @@ fn root_locks_are_noops() {
IncLockRefResult::default(),
/* lock_host = */ false,
);
assert!(result.skip_lock_node_ids.is_empty());
mamba.release_component_lock(&mut tc, root, None, /* lock_host = */ false);
mamba.release_component_lock(
&mut tc,
root,
&DecLockRefParams::default(),
/* lock_host = */ false,
);
assert_eq!(tc.evictable_size_(MAMBA), 0);
}
@@ -448,7 +459,12 @@ fn host_lock_detaches_and_reattaches_the_host_lru() {
);
assert!(!tc.host_lru_list(MAMBA).in_list(Some(a)));
assert_eq!(tc.arena.node(a).host_lock_ref(MAMBA), 1);
mamba.release_component_lock(&mut tc, a, None, /* lock_host = */ true);
mamba.release_component_lock(
&mut tc,
a,
&DecLockRefParams::default(),
/* lock_host = */ true,
);
assert!(tc.host_lru_list(MAMBA).in_list(Some(a)));
assert_eq!(tc.arena.node(a).host_lock_ref(MAMBA), 0);
}
@@ -466,7 +482,12 @@ fn host_unlock_skips_the_lru_for_device_backed_nodes() {
IncLockRefResult::default(),
/* lock_host = */ true,
);
mamba.release_component_lock(&mut tc, a, None, /* lock_host = */ true);
mamba.release_component_lock(
&mut tc,
a,
&DecLockRefParams::default(),
/* lock_host = */ true,
);
assert!(!tc.host_lru_list(MAMBA).in_list(Some(a)));
}
@@ -1820,7 +1841,7 @@ fn branching_from_a_host_full_hit_is_reusable_after_insert() {
}
#[test]
fn skip_set_release_after_a_restore_and_relock_keeps_the_new_lock() {
fn release_after_a_restore_and_relock_keeps_the_other_lock() {
let mut tc = mamba_core(/* page_size = */ 1);
let [a] = chain::<1>(&mut tc);
let mamba = mamba_component();
@@ -1830,28 +1851,33 @@ fn skip_set_release_after_a_restore_and_relock_keeps_the_new_lock() {
IncLockRefResult::default(),
/* lock_host = */ false,
);
assert!(first.skip_lock_node_ids[&MAMBA].contains(&tc.arena.node(a).id));
// The tombstone is restored and a second request locks it before the
// first release replays its skip set.
set_mamba_device(&mut tc, a, 7);
assert_eq!(tc.arena.node(a).device_lock_ref(MAMBA), 1);
// The tombstone is restored under the held lock (credited to protected)
// and a second request stacks its own lock on it.
tc.set_component_device_value_(a, MAMBA, Tensor::from_slice(&[7i64]));
let _ = mamba.acquire_component_lock(
&mut tc,
a,
IncLockRefResult::default(),
/* lock_host = */ false,
);
assert_eq!(tc.arena.node(a).device_lock_ref(MAMBA), 1);
assert_eq!(tc.arena.node(a).device_lock_ref(MAMBA), 2);
assert_eq!(tc.evictable_size_(MAMBA), 0);
assert_eq!(tc.protected_size_(MAMBA), 1);
let params = DecLockRefParams {
skip_lock_node_ids: first.skip_lock_node_ids.clone(),
skipped_lock_components: first.skipped_lock_components,
..DecLockRefParams::default()
};
mamba.release_component_lock(&mut tc, a, Some(&params), /* lock_host = */ false);
// The replayed skip keeps the restored node's fresh lock intact.
mamba.release_component_lock(&mut tc, a, &params, /* lock_host = */ false);
// The first release takes back exactly its own ref.
assert_eq!(tc.arena.node(a).device_lock_ref(MAMBA), 1);
assert_eq!(tc.protected_size_(MAMBA), 1);
mamba.release_component_lock(&mut tc, a, None, /* lock_host = */ false);
mamba.release_component_lock(
&mut tc,
a,
&DecLockRefParams::default(),
/* lock_host = */ false,
);
assert_eq!(tc.arena.node(a).device_lock_ref(MAMBA), 0);
assert_eq!(tc.evictable_size_(MAMBA), 1);
assert_eq!(tc.protected_size_(MAMBA), 0);
rust/sglang-radix-tree/src/tests/components/swa.rs
+273 -100
View File
@@ -1,5 +1,5 @@
use super::*;
use crate::components::{FULL, MAMBA, SWA};
use crate::components::{ComponentSet, FULL, MAMBA, SWA};
use crate::test_utils::{accumulate_step, action_kinds};
use crate::unified_tree_core::CacheInitParams;
@@ -766,18 +766,32 @@ fn insert_overlap_recovers_a_tombstone_inside_the_window() {
}
#[test]
#[should_panic(expected = "tombstone Swa lock_ref should be 0, node")]
fn insert_overlap_panics_on_a_locked_swa_tombstone() {
fn insert_overlap_recovers_a_locked_swa_tombstone() {
let mut tc = swa_core(/* window = */ 8, /* page_size = */ 1);
tc.insert(&insert_params_swa(&vec![1, 2, 3], &[10, 11, 12], 0, 0));
let root = tc.arena.root();
let leaf = child_of(&tc, root, &[1]);
// The rebuild is deferred, so the leaf is still an SWA tombstone; a raw
// lock on it breaks the tombstones-are-unlocked contract.
// A segment lock may hold an SWA tombstone; the co-held FULL lock (the
// full >= swa protocol invariant) forces the Recover branch, so the
// locked full stays on the node.
tc.arena
.node_mut(leaf)
.set_lock_ref_(ValueSlotIdx::device(SWA), 1);
tc.insert(&insert_params_swa(&vec![1, 2, 3], &[20, 21, 22], 0, 0));
tc.arena
.node_mut(leaf)
.set_lock_ref_(ValueSlotIdx::device(FULL), 1);
let result = tc.insert(&insert_params_swa(&vec![1, 2, 3], &[20, 21, 22], 0, 0));
assert!(
tc.arena
.device_value(leaf, FULL)
.equal(&Tensor::from_slice(&[10i64, 11, 12]))
);
assert!(
result
.cache_actions
.iter()
.any(|action| matches!(action, CacheAction::RecoverSwaWithLockedFull { .. }))
);
}
#[test]
@@ -1576,7 +1590,7 @@ fn acquire_lock_reuses_the_stamped_uuid_and_shifts_sizes_once() {
}
#[test]
fn acquire_lock_skips_tombstones_and_records_them() {
fn acquire_lock_counts_tombstones_toward_the_window() {
let mut tc = swa_core(/* window = */ 2, /* page_size = */ 1);
let [a, b, c] = chain(&mut tc);
store_swa_device(&mut tc, a);
@@ -1587,14 +1601,13 @@ fn acquire_lock_skips_tombstones_and_records_them() {
IncLockRefResult::default(),
/* lock_host = */ false,
);
// The valueless b is recorded and skipped; the window fills at a.
// The valueless b is counted too (no ledger move); position-based
// coverage fills the window at b, so a stays outside the segment.
assert_eq!(tc.arena.device_lock_ref(c, SWA), 1);
assert_eq!(tc.arena.device_lock_ref(b, SWA), 0);
assert_eq!(tc.arena.device_lock_ref(a, SWA), 1);
assert_eq!(result.skip_lock_node_ids[&SWA].len(), 1);
assert!(result.skip_lock_node_ids[&SWA].contains(&tc.arena.node(b).id));
assert_eq!(tc.arena.device_lock_ref(b, SWA), 1);
assert_eq!(tc.arena.device_lock_ref(a, SWA), 0);
assert!(result.swa_uuid_for_lock.is_some());
assert_eq!(node_swa_uuid(&tc, a), result.swa_uuid_for_lock);
assert_eq!(node_swa_uuid(&tc, b), result.swa_uuid_for_lock);
}
#[test]
@@ -1624,11 +1637,14 @@ fn inc_lock_ref_runs_full_and_swa_walks_together() {
store_swa_device(&mut tc, b);
store_swa_device(&mut tc, c);
let result = tc
.inc_lock_ref(tc.arena.node(c).id)
.inc_lock_ref(tc.arena.node(c).id, ComponentSet::EMPTY)
.expect("live test node");
// FULL sees a valueless path (skip segment only); SWA locks its window.
// FULL counts its valueless bottom segment (no ledger move); SWA locks
// its window.
assert_eq!(result.delta, Some(0));
assert_eq!(result.skip_lock_node_ids[&FULL].len(), 3);
assert_eq!(tc.arena.device_lock_ref(c, FULL), 1);
assert_eq!(tc.arena.device_lock_ref(b, FULL), 1);
assert_eq!(tc.arena.device_lock_ref(a, FULL), 1);
assert!(result.swa_uuid_for_lock.is_some());
assert_eq!(tc.arena.device_lock_ref(c, SWA), 1);
assert_eq!(tc.arena.device_lock_ref(b, SWA), 1);
@@ -1657,10 +1673,10 @@ fn inc_host_lock_ref_runs_full_and_swa_host_arms_together() {
// The release replays the acquire's uuid and unwinds both arms.
let params = DecLockRefParams {
swa_uuid_for_host_lock: result.swa_uuid_for_host_lock,
skip_lock_node_ids: result.skip_lock_node_ids,
skipped_lock_components: result.skipped_lock_components,
..Default::default()
};
tc.dec_host_lock_ref(tc.arena.node(c).id, Some(&params))
tc.dec_host_lock_ref(tc.arena.node(c).id, &params)
.expect("live test node");
assert_eq!(tc.arena.host_lock_ref(c, FULL), 0);
assert_eq!(tc.arena.host_lock_ref(c, SWA), 0);
@@ -1690,10 +1706,10 @@ fn dec_host_lock_ref_with_the_inner_uuid_leaves_an_outer_window_pinned() {
// window's lock above the boundary survives.
let params = DecLockRefParams {
swa_uuid_for_host_lock: inner.swa_uuid_for_host_lock,
skip_lock_node_ids: inner.skip_lock_node_ids,
skipped_lock_components: inner.skipped_lock_components,
..Default::default()
};
tc.dec_host_lock_ref(tc.arena.node(c).id, Some(&params))
tc.dec_host_lock_ref(tc.arena.node(c).id, &params)
.expect("live test node");
assert_eq!(tc.arena.host_lock_ref(c, SWA), 0);
assert_eq!(tc.arena.host_lock_ref(b, SWA), 1);
@@ -1758,7 +1774,7 @@ fn acquire_host_lock_reuses_the_stamped_uuid_and_skips_unlisted_nodes() {
}
#[test]
fn acquire_host_lock_skips_host_tombstones_and_records_them() {
fn acquire_host_lock_counts_host_tombstones_toward_the_window() {
let mut tc = swa_core(/* window = */ 2, /* page_size = */ 1);
let [a, b, c] = chain(&mut tc);
set_swa_host(&mut tc, a);
@@ -1769,12 +1785,12 @@ fn acquire_host_lock_skips_host_tombstones_and_records_them() {
IncLockRefResult::default(),
/* lock_host = */ true,
);
// The host-valueless b is counted too; position-based coverage fills
// the window at b, so a stays outside the segment.
assert_eq!(tc.arena.host_lock_ref(c, SWA), 1);
assert_eq!(tc.arena.host_lock_ref(b, SWA), 0);
assert_eq!(tc.arena.host_lock_ref(a, SWA), 1);
assert_eq!(result.skip_lock_node_ids[&SWA].len(), 1);
assert!(result.skip_lock_node_ids[&SWA].contains(&tc.arena.node(b).id));
assert_eq!(node_swa_host_uuid(&tc, a), result.swa_uuid_for_host_lock);
assert_eq!(tc.arena.host_lock_ref(b, SWA), 1);
assert_eq!(tc.arena.host_lock_ref(a, SWA), 0);
assert_eq!(node_swa_host_uuid(&tc, b), result.swa_uuid_for_host_lock);
assert!(result.swa_uuid_for_host_lock.is_some());
}
@@ -1996,11 +2012,12 @@ fn release_lock_returns_the_window_to_evictable() {
/* lock_host = */ false,
);
let params = DecLockRefParams {
node_id: None,
swa_uuid_for_lock: result.swa_uuid_for_lock,
swa_uuid_for_host_lock: result.swa_uuid_for_host_lock,
skip_lock_node_ids: result.skip_lock_node_ids,
skipped_lock_components: result.skipped_lock_components,
};
swa.release_component_lock(&mut tc, c, Some(&params), /* lock_host = */ false);
swa.release_component_lock(&mut tc, c, &params, /* lock_host = */ false);
assert_eq!(tc.arena.device_lock_ref(c, SWA), 0);
assert_eq!(tc.arena.device_lock_ref(b, SWA), 0);
assert_eq!(tc.arena.device_lock_ref(a, SWA), 0);
@@ -2029,11 +2046,12 @@ fn release_lock_keeps_sizes_while_other_locks_remain() {
/* lock_host = */ false,
);
let params = DecLockRefParams {
node_id: None,
swa_uuid_for_lock: first.swa_uuid_for_lock,
swa_uuid_for_host_lock: first.swa_uuid_for_host_lock,
skip_lock_node_ids: first.skip_lock_node_ids,
skipped_lock_components: first.skipped_lock_components,
};
swa.release_component_lock(&mut tc, c, Some(&params), /* lock_host = */ false);
swa.release_component_lock(&mut tc, c, &params, /* lock_host = */ false);
assert_eq!(tc.arena.device_lock_ref(c, SWA), 1);
assert_eq!(tc.arena.device_lock_ref(b, SWA), 1);
assert_eq!(tc.swa_evictable_size(), 1);
@@ -2056,11 +2074,12 @@ fn release_lock_replays_the_tombstone_skips() {
// b gained a device value AFTER the acquire recorded it as a tombstone.
store_swa_device(&mut tc, b);
let params = DecLockRefParams {
node_id: None,
swa_uuid_for_lock: result.swa_uuid_for_lock,
swa_uuid_for_host_lock: result.swa_uuid_for_host_lock,
skip_lock_node_ids: result.skip_lock_node_ids,
skipped_lock_components: result.skipped_lock_components,
};
swa.release_component_lock(&mut tc, c, Some(&params), /* lock_host = */ false);
swa.release_component_lock(&mut tc, c, &params, /* lock_host = */ false);
assert_eq!(tc.arena.device_lock_ref(c, SWA), 0);
assert_eq!(tc.arena.device_lock_ref(b, SWA), 0);
assert_eq!(tc.arena.device_lock_ref(a, SWA), 0);
@@ -2087,11 +2106,12 @@ fn release_lock_stops_at_the_window_uuid() {
.node_mut(a)
.set_lock_ref_(ValueSlotIdx::device(SWA), 1);
let params = DecLockRefParams {
node_id: None,
swa_uuid_for_lock: result.swa_uuid_for_lock,
swa_uuid_for_host_lock: result.swa_uuid_for_host_lock,
skip_lock_node_ids: result.skip_lock_node_ids,
skipped_lock_components: result.skipped_lock_components,
};
swa.release_component_lock(&mut tc, c, Some(&params), /* lock_host = */ false);
swa.release_component_lock(&mut tc, c, &params, /* lock_host = */ false);
assert_eq!(tc.arena.device_lock_ref(c, SWA), 0);
assert_eq!(tc.arena.device_lock_ref(b, SWA), 0);
assert_eq!(tc.arena.device_lock_ref(a, SWA), 1);
@@ -2120,11 +2140,12 @@ fn release_host_lock_stops_at_the_host_uuid_boundary() {
/* lock_host = */ true,
);
let params = DecLockRefParams {
node_id: None,
swa_uuid_for_lock: result.swa_uuid_for_lock,
swa_uuid_for_host_lock: result.swa_uuid_for_host_lock,
skip_lock_node_ids: result.skip_lock_node_ids,
skipped_lock_components: result.skipped_lock_components,
};
swa.release_component_lock(&mut tc, c, Some(&params), /* lock_host = */ true);
swa.release_component_lock(&mut tc, c, &params, /* lock_host = */ true);
assert_eq!(tc.arena.host_lock_ref(c, SWA), 0);
assert_eq!(tc.arena.host_lock_ref(b, SWA), 0);
assert_eq!(tc.arena.host_lock_ref(a, SWA), 1);
@@ -2132,34 +2153,27 @@ fn release_host_lock_stops_at_the_host_uuid_boundary() {
}
#[test]
fn release_lock_without_params_passes_over_an_unlocked_middle_node() {
#[should_panic(expected = "SWA segment release hit lock_ref=0")]
fn release_lock_without_the_boundary_uuid_dies_at_the_segment_edge() {
let mut tc = swa_core(/* window = */ 1, /* page_size = */ 1);
let [a, b, c] = chain(&mut tc);
store_swa_device(&mut tc, a);
store_swa_device(&mut tc, b);
let [_a, _b, c] = chain(&mut tc);
store_swa_device(&mut tc, c);
let swa = swa_component(1);
// The 1-atom window locks only the acquired node: c and a, never b.
// The 1-atom window locks only the acquired node c.
let _ = swa.acquire_component_lock(
&mut tc,
c,
IncLockRefResult::default(),
/* lock_host = */ false,
);
let _ = swa.acquire_component_lock(
// A receipt-less release overshoots the boundary into unlocked territory
// and dies there instead of silently stealing whatever it crosses.
swa.release_component_lock(
&mut tc,
a,
IncLockRefResult::default(),
c,
&DecLockRefParams::default(),
/* lock_host = */ false,
);
swa.release_component_lock(
&mut tc, c, /* params = */ None, /* lock_host = */ false,
);
assert_eq!(tc.arena.device_lock_ref(c, SWA), 0);
assert_eq!(tc.arena.device_lock_ref(b, SWA), 0);
assert_eq!(tc.arena.device_lock_ref(a, SWA), 0);
assert_eq!(tc.swa_evictable_size(), 3);
assert_eq!(tc.swa_protected_size(), 0);
}
#[test]
@@ -2178,11 +2192,12 @@ fn release_host_lock_reparks_tombstoned_host_nodes() {
/* lock_host = */ true,
);
let params = DecLockRefParams {
node_id: None,
swa_uuid_for_lock: result.swa_uuid_for_lock,
swa_uuid_for_host_lock: result.swa_uuid_for_host_lock,
skip_lock_node_ids: result.skip_lock_node_ids,
skipped_lock_components: result.skipped_lock_components,
};
swa.release_component_lock(&mut tc, c, Some(&params), /* lock_host = */ true);
swa.release_component_lock(&mut tc, c, &params, /* lock_host = */ true);
assert_eq!(tc.arena.host_lock_ref(c, SWA), 0);
assert_eq!(tc.arena.host_lock_ref(b, SWA), 0);
assert!(tc.host_lru_list(SWA).in_list(Some(c)));
@@ -2198,19 +2213,16 @@ fn inc_then_dec_lock_ref_roundtrips_with_dec_params() {
store_swa_device(&mut tc, b);
store_swa_device(&mut tc, c);
let result = tc
.inc_lock_ref(tc.arena.node(c).id)
.inc_lock_ref(tc.arena.node(c).id, ComponentSet::EMPTY)
.expect("live test node");
let params = DecLockRefParams {
node_id: None,
swa_uuid_for_lock: result.swa_uuid_for_lock,
swa_uuid_for_host_lock: result.swa_uuid_for_host_lock,
skip_lock_node_ids: result.skip_lock_node_ids,
skipped_lock_components: result.skipped_lock_components,
};
tc.dec_lock_ref(
tc.arena.node(c).id,
Some(&params),
/* skip_swa = */ false,
)
.expect("live test node");
tc.dec_lock_ref(tc.arena.node(c).id, &params, /* skip_swa = */ false)
.expect("live test node");
assert_eq!(tc.arena.device_lock_ref(c, SWA), 0);
assert_eq!(tc.arena.device_lock_ref(b, SWA), 0);
assert_eq!(tc.swa_evictable_size(), 3);
@@ -2230,13 +2242,17 @@ fn dec_swa_lock_only_releases_swa_while_full_stays_locked() {
// Fund FULL's evictable counter for its lock walk (raw slot sets skip it).
tc.component_state_mut(FULL).evictable_size = 3;
let result = tc
.inc_lock_ref(tc.arena.node(c).id)
.inc_lock_ref(tc.arena.node(c).id, ComponentSet::EMPTY)
.expect("live test node");
let mut device_frees = HashMap::new();
let mut host_frees = HashMap::new();
tc.dec_swa_lock_only(
tc.arena.node(c).id,
result.swa_uuid_for_lock,
&DecLockRefParams {
swa_uuid_for_lock: result.swa_uuid_for_lock,
skipped_lock_components: ComponentSet::EMPTY,
..Default::default()
},
&mut device_frees,
&mut host_frees,
)
@@ -2273,7 +2289,11 @@ fn dec_swa_lock_only_evicts_a_fully_unlocked_device_leaf() {
let mut host_frees = HashMap::new();
tc.dec_swa_lock_only(
tc.arena.node(c).id,
result.swa_uuid_for_lock,
&DecLockRefParams {
swa_uuid_for_lock: result.swa_uuid_for_lock,
skipped_lock_components: ComponentSet::EMPTY,
..Default::default()
},
&mut device_frees,
&mut host_frees,
)
@@ -2296,7 +2316,11 @@ fn dec_swa_lock_only_is_a_noop_without_the_swa_component() {
let mut host_frees = HashMap::new();
tc.dec_swa_lock_only(
tc.arena.node(root).id,
None,
&DecLockRefParams {
swa_uuid_for_lock: None,
skipped_lock_components: ComponentSet::EMPTY,
..Default::default()
},
&mut device_frees,
&mut host_frees,
)
@@ -2374,11 +2398,12 @@ fn release_lock_skip_set_leaves_a_relocked_tombstone_credited() {
/* lock_host = */ false,
);
let params = DecLockRefParams {
node_id: None,
swa_uuid_for_lock: first.swa_uuid_for_lock,
swa_uuid_for_host_lock: first.swa_uuid_for_host_lock,
skip_lock_node_ids: first.skip_lock_node_ids,
skipped_lock_components: first.skipped_lock_components,
};
swa.release_component_lock(&mut tc, c, Some(&params), /* lock_host = */ false);
swa.release_component_lock(&mut tc, c, &params, /* lock_host = */ false);
assert_eq!(tc.arena.device_lock_ref(c, SWA), 1);
assert_eq!(tc.arena.device_lock_ref(b, SWA), 1);
assert_eq!(tc.arena.device_lock_ref(a, SWA), 0);
@@ -2387,28 +2412,27 @@ fn release_lock_skip_set_leaves_a_relocked_tombstone_credited() {
}
#[test]
fn release_lock_passes_over_uncredited_nodes_without_params() {
#[should_panic(expected = "SWA segment release hit lock_ref=0")]
fn double_release_with_one_receipt_dies_loud() {
let mut tc = swa_core(/* window = */ 2, /* page_size = */ 1);
let [a, b, c] = chain(&mut tc);
store_swa_device(&mut tc, a);
store_swa_device(&mut tc, b);
store_swa_device(&mut tc, c);
let swa = swa_component(2);
let _ = swa.acquire_component_lock(
let result = swa.acquire_component_lock(
&mut tc,
c,
IncLockRefResult::default(),
/* lock_host = */ false,
);
// No params: the walk crosses the never-credited a up to the root.
swa.release_component_lock(
&mut tc, c, /* params = */ None, /* lock_host = */ false,
);
assert_eq!(tc.arena.device_lock_ref(c, SWA), 0);
assert_eq!(tc.arena.device_lock_ref(b, SWA), 0);
assert_eq!(tc.arena.device_lock_ref(a, SWA), 0);
assert_eq!(tc.swa_evictable_size(), 3);
assert_eq!(tc.swa_protected_size(), 0);
let params = DecLockRefParams {
swa_uuid_for_lock: result.swa_uuid_for_lock,
..Default::default()
};
swa.release_component_lock(&mut tc, c, &params, /* lock_host = */ false);
// Consuming the same receipt twice dies at the first unlocked node.
swa.release_component_lock(&mut tc, c, &params, /* lock_host = */ false);
}
#[test]
@@ -2435,7 +2459,11 @@ fn dec_swa_lock_only_releases_the_window_exactly_once() {
let mut host_frees = HashMap::new();
tc.dec_swa_lock_only(
tc.arena.node(c).id,
first.swa_uuid_for_lock,
&DecLockRefParams {
swa_uuid_for_lock: first.swa_uuid_for_lock,
skipped_lock_components: ComponentSet::EMPTY,
..Default::default()
},
&mut device_frees,
&mut host_frees,
)
@@ -2447,7 +2475,11 @@ fn dec_swa_lock_only_releases_the_window_exactly_once() {
assert_eq!(tc.swa_protected_size(), 2);
tc.dec_swa_lock_only(
tc.arena.node(c).id,
first.swa_uuid_for_lock,
&DecLockRefParams {
swa_uuid_for_lock: first.swa_uuid_for_lock,
skipped_lock_components: ComponentSet::EMPTY,
..Default::default()
},
&mut device_frees,
&mut host_frees,
)
@@ -2483,7 +2515,11 @@ fn dec_swa_lock_only_leaves_out_of_window_swa_locks_alone() {
let mut host_frees = HashMap::new();
tc.dec_swa_lock_only(
tc.arena.node(c).id,
result.swa_uuid_for_lock,
&DecLockRefParams {
swa_uuid_for_lock: result.swa_uuid_for_lock,
skipped_lock_components: ComponentSet::EMPTY,
..Default::default()
},
&mut device_frees,
&mut host_frees,
)
@@ -2496,7 +2532,8 @@ fn dec_swa_lock_only_leaves_out_of_window_swa_locks_alone() {
}
#[test]
fn release_window_lock_passes_over_an_unlocked_valued_node() {
#[should_panic(expected = "SWA window release hit lock_ref=0")]
fn release_window_lock_without_the_uuid_dies_past_the_boundary() {
let mut tc = swa_core(/* window = */ 2, /* page_size = */ 1);
let [a, b, c] = chain(&mut tc);
store_swa_device(&mut tc, a);
@@ -2511,11 +2548,9 @@ fn release_window_lock_passes_over_an_unlocked_valued_node() {
);
let mut device_frees = HashMap::new();
let mut host_frees = HashMap::new();
// No uuid bound: the walk crosses the valued-but-unlocked a to the root.
// Without the boundary uuid the walk crosses the segment edge into the
// unlocked a and dies there instead of stealing.
swa.release_window_lock(&mut tc, c, None, &mut device_frees, &mut host_frees);
assert_eq!(tc.arena.device_lock_ref(c, SWA), 0);
assert_eq!(tc.arena.device_lock_ref(b, SWA), 0);
assert_eq!(tc.arena.device_lock_ref(a, SWA), 0);
}
#[test]
@@ -2533,7 +2568,8 @@ fn release_window_lock_passes_over_a_mid_chain_tombstone_without_a_uuid() {
);
let mut device_frees = HashMap::new();
let mut host_frees = HashMap::new();
// No uuid bound: the walk crosses the mid-chain tombstone b and releases a.
// Walked-to-root acquire (window > chain): the uuid-less release counts
// back through the mid-chain tombstone b and releases a.
swa.release_window_lock(&mut tc, c, None, &mut device_frees, &mut host_frees);
assert_eq!(tc.arena.device_lock_ref(c, SWA), 0);
assert_eq!(tc.arena.device_lock_ref(b, SWA), 0);
@@ -2558,11 +2594,12 @@ fn release_host_lock_does_not_repark_a_node_whose_host_value_was_taken() {
// device value either, so the release has nothing to park.
let _ = tc.arena.take_host_value(a, SWA);
let params = DecLockRefParams {
node_id: None,
swa_uuid_for_lock: result.swa_uuid_for_lock,
swa_uuid_for_host_lock: result.swa_uuid_for_host_lock,
skip_lock_node_ids: result.skip_lock_node_ids,
skipped_lock_components: result.skipped_lock_components,
};
swa.release_component_lock(&mut tc, a, Some(&params), /* lock_host = */ true);
swa.release_component_lock(&mut tc, a, &params, /* lock_host = */ true);
assert_eq!(tc.arena.host_lock_ref(a, SWA), 0);
assert!(!tc.host_lru_list(SWA).in_list(Some(a)));
}
@@ -2584,11 +2621,12 @@ fn release_host_lock_skips_reparking_device_valued_nodes() {
/* lock_host = */ true,
);
let params = DecLockRefParams {
node_id: None,
swa_uuid_for_lock: result.swa_uuid_for_lock,
swa_uuid_for_host_lock: result.swa_uuid_for_host_lock,
skip_lock_node_ids: result.skip_lock_node_ids,
skipped_lock_components: result.skipped_lock_components,
};
swa.release_component_lock(&mut tc, c, Some(&params), /* lock_host = */ true);
swa.release_component_lock(&mut tc, c, &params, /* lock_host = */ true);
// Device-valued nodes never re-park in the host LRU on host release.
assert!(!tc.host_lru_list(SWA).in_list(Some(c)));
assert!(!tc.host_lru_list(SWA).in_list(Some(b)));
@@ -2612,11 +2650,12 @@ fn release_host_lock_leaves_an_already_listed_node_listed() {
// Something re-listed b while the lock was held (e.g. a split re-park).
tc.host_lru_list_mut(SWA).insert_mru(b);
let params = DecLockRefParams {
node_id: None,
swa_uuid_for_lock: result.swa_uuid_for_lock,
swa_uuid_for_host_lock: result.swa_uuid_for_host_lock,
skip_lock_node_ids: result.skip_lock_node_ids,
skipped_lock_components: result.skipped_lock_components,
};
swa.release_component_lock(&mut tc, c, Some(&params), /* lock_host = */ true);
swa.release_component_lock(&mut tc, c, &params, /* lock_host = */ true);
assert!(tc.host_lru_list(SWA).in_list(Some(b)));
assert!(tc.host_lru_list(SWA).in_list(Some(c)));
let _ = a;
@@ -2890,6 +2929,29 @@ fn redistribute_on_node_split_moves_the_swa_uuid_to_the_parent() {
assert_eq!(node_swa_uuid(&tc, node), None);
}
#[test]
fn redistribute_on_node_split_preserves_host_lock_state() {
let mut tc = swa_core(/* window = */ 4, /* page_size = */ 1);
tc.insert(&insert_params_swa(&vec![1, 2], &[10, 11], 0, 0));
let root = tc.arena.root();
let node = child_of(&tc, root, &[1]);
// A host-locked tombstone mid-IO: both halves must stay pinned and out
// of the host LRU, and the host boundary uuid moves to the parent.
tc.arena
.set_host_value(node, SWA, Tensor::from_slice(&[70i64, 71]));
tc.arena
.node_mut(node)
.set_lock_ref_(ValueSlotIdx::host(SWA), 1);
tc.arena.node_mut(node).swa_host_uuid = Some(9);
let (parent, _) = tc.split_node_(node, /* split_len = */ 1);
assert_eq!(tc.arena.host_lock_ref(parent, SWA), 1);
assert_eq!(tc.arena.host_lock_ref(node, SWA), 1);
assert_eq!(node_swa_host_uuid(&tc, parent), Some(9));
assert_eq!(node_swa_host_uuid(&tc, node), None);
assert!(!tc.host_lru_list(SWA).in_list(Some(parent)));
assert!(!tc.host_lru_list(SWA).in_list(Some(node)));
}
#[test]
fn finalize_window_arithmetic_at_page_boundaries() {
let mut tc = swa_core(/* window = */ 4, /* page_size = */ 2);
@@ -3277,17 +3339,24 @@ fn reinsert_rejects_a_page_misaligned_boundary() {
}
#[test]
#[should_panic(expected = "tombstone Swa lock_ref should be 0 on unevict")]
fn reinsert_rejects_a_locked_tombstone() {
fn reinsert_rebuilds_a_locked_tombstone() {
let mut tc = swa_core(/* window = */ 8, /* page_size = */ 1);
tc.insert(&insert_params_swa(&vec![1, 2], &[10, 11], 0, 0));
let root = tc.arena.root();
let node = child_of(&tc, root, &[1]);
evict_full(&mut tc, node, /* remaining_size = */ 0);
// Segment locks count evicted nodes, so a locked tombstone is legal and
// the re-insert rebuilds its SWA from the fresh KV.
tc.arena
.node_mut(node)
.set_lock_ref_(ValueSlotIdx::device(SWA), 1);
tc.insert(&insert_params_swa(&vec![1, 2], &[20, 21], 0, 0));
let result = tc.insert(&insert_params_swa(&vec![1, 2], &[20, 21], 0, 0));
assert!(
result
.cache_actions
.iter()
.any(|action| matches!(action, CacheAction::SwaRebuild { .. }))
);
}
fn set_full_host(tc: &mut UnifiedTreeCore<Vec<i64>>, node: NodeIdx_) {
@@ -4856,13 +4925,16 @@ fn deep_swa_tree_survives_backup_evict_and_load_back_rounds() {
)
.expect("live test node");
assert!(actions.is_empty());
let lock = tc.inc_lock_ref(anchor).expect("live test node");
let lock = tc
.inc_lock_ref(anchor, ComponentSet::EMPTY)
.expect("live test node");
let params = DecLockRefParams {
node_id: None,
swa_uuid_for_lock: lock.swa_uuid_for_lock,
swa_uuid_for_host_lock: lock.swa_uuid_for_host_lock,
skip_lock_node_ids: lock.skip_lock_node_ids,
skipped_lock_components: lock.skipped_lock_components,
};
tc.dec_lock_ref(anchor, Some(&params), /* skip_swa = */ false)
tc.dec_lock_ref(anchor, &params, /* skip_swa = */ false)
.expect("live test node");
tc.finish_load_back(anchor).expect("live test node");
}
@@ -4960,3 +5032,104 @@ fn recovered_swa_span_evicts_before_the_window_leaf() {
assert!(tc.arena.has_device_value(leaf, FULL));
tc.sanity_check(&[], &[]);
}
#[test]
fn aux_release_refreshes_the_leaf_set_whatever_the_release_order() {
let mut tc = swa_core(/* window = */ 1, /* page_size = */ 1);
tc.insert(&insert_params_swa(&vec![1], &[10], 0, 0));
let leaf = child_of(&tc, tc.arena.root(), &[1]);
store_swa_device(&mut tc, leaf);
assert!(tc.evictable_device_leaves.contains(leaf));
let result = tc
.inc_lock_ref(tc.arena.node(leaf).id, ComponentSet::EMPTY)
.expect("live test node");
assert!(!tc.evictable_device_leaves.contains(leaf));
let params = DecLockRefParams {
node_id: result.node_id,
swa_uuid_for_lock: result.swa_uuid_for_lock,
swa_uuid_for_host_lock: None,
skipped_lock_components: ComponentSet::EMPTY,
};
// Full first: its walk still sees the SWA lock, so the leaf stays out.
crate::components::FullComponent
.release_component_lock(&mut tc, leaf, &params, /* lock_host = */ false);
assert!(!tc.evictable_device_leaves.contains(leaf));
// The SWA release drops the last lock and must readmit the leaf itself.
swa_component(1).release_component_lock(&mut tc, leaf, &params, /* lock_host = */ false);
assert!(tc.evictable_device_leaves.contains(leaf));
}
#[test]
#[should_panic(expected = "lock receipt anchored on node")]
fn dec_lock_ref_rejects_a_receipt_from_another_node() {
let mut tc = swa_core(/* window = */ 2, /* page_size = */ 1);
let [a, _b, c] = chain(&mut tc);
store_swa_device(&mut tc, c);
let result = tc
.inc_lock_ref(tc.arena.node(c).id, ComponentSet::EMPTY)
.expect("live test node");
let params = DecLockRefParams {
node_id: result.node_id,
swa_uuid_for_lock: result.swa_uuid_for_lock,
swa_uuid_for_host_lock: None,
skipped_lock_components: result.skipped_lock_components,
};
// Same receipt, wrong anchor: the walk would otherwise release a's
// segment, which this holder never locked.
tc.dec_lock_ref(tc.arena.node(a).id, &params, /* skip_swa = */ false)
.expect("live test node");
}
#[test]
#[should_panic(expected = "lock receipt anchored on node")]
fn dec_host_lock_ref_rejects_a_receipt_from_another_node() {
let mut tc = swa_core(/* window = */ 2, /* page_size = */ 1);
let [a, _b, c] = chain(&mut tc);
set_swa_host(&mut tc, c);
tc.host_lru_list_mut(SWA).insert_mru(c);
let result = tc
.inc_host_lock_ref(tc.arena.node(c).id)
.expect("live test node");
let params = DecLockRefParams {
node_id: result.node_id,
swa_uuid_for_lock: None,
swa_uuid_for_host_lock: result.swa_uuid_for_host_lock,
skipped_lock_components: ComponentSet::EMPTY,
};
tc.dec_host_lock_ref(tc.arena.node(a).id, &params)
.expect("live test node");
}
#[test]
fn receipt_anchor_follows_the_locked_node_through_a_split() {
let mut tc = swa_core(/* window = */ 1, /* page_size = */ 1);
tc.insert(&insert_params_swa(&vec![1, 2], &[10, 11], 0, 0));
let leaf = child_of(&tc, tc.arena.root(), &[1]);
store_swa_device(&mut tc, leaf);
let leaf_id = tc.arena.node(leaf).id;
let result = tc
.inc_lock_ref(leaf_id, ComponentSet::EMPTY)
.expect("live test node");
assert_eq!(result.node_id, Some(leaf_id));
// Diverge inside the node: the split keeps the id on the deeper half.
tc.insert(&insert_params_swa(&vec![1, 3], &[12, 13], 0, 0));
let params = DecLockRefParams {
node_id: result.node_id,
swa_uuid_for_lock: result.swa_uuid_for_lock,
swa_uuid_for_host_lock: None,
skipped_lock_components: ComponentSet::EMPTY,
};
tc.dec_lock_ref(leaf_id, &params, /* skip_swa = */ false)
.expect("live test node");
assert_eq!(
tc.arena
.device_lock_ref(tc.arena.resolve(leaf_id).expect("live test node"), SWA),
0
);
}
#[test]
#[should_panic(expected = "swa_sliding_window_size must be positive")]
fn new_panics_on_a_zero_sliding_window_size() {
SwaComponent::new(&swa_params_with_window(0));
}
rust/sglang-radix-tree/src/tests/unified_tree_core.rs
+85 -35
View File
@@ -3,7 +3,7 @@ use std::sync::Mutex;
use tch::Tensor;
use super::*;
use crate::components::{FULL, MAMBA, SWA};
use crate::components::{ComponentSet, FULL, MAMBA, SWA};
use crate::node::{NodeAccessError, ValueSlotIdx};
use crate::test_utils::{accumulate_step, action_kinds};
@@ -91,7 +91,7 @@ impl TreeComponent<Vec<i64>> for RecordingComponentForTest {
&self,
_tree_core: &mut UnifiedTreeCore<Vec<i64>>,
_node_id: NodeIdx_,
_params: Option<&DecLockRefParams>,
_params: &DecLockRefParams,
_lock_host: bool,
) {
unimplemented!()
@@ -213,7 +213,7 @@ impl TreeComponent<Vec<i64>> for CountingComponentForTest {
&self,
_tree_core: &mut UnifiedTreeCore<Vec<i64>>,
_node_id: NodeIdx_,
_params: Option<&DecLockRefParams>,
_params: &DecLockRefParams,
_lock_host: bool,
) {
unimplemented!()
@@ -293,11 +293,11 @@ impl TreeComponent<Vec<i64>> for LowPriorityComponentForTest {
&self,
_tree_core: &mut UnifiedTreeCore<Vec<i64>>,
_node_id: NodeIdx_,
params: Option<&DecLockRefParams>,
params: &DecLockRefParams,
lock_host: bool,
) {
assert!(!lock_host);
assert!(params.is_some_and(|p| p.swa_uuid_for_lock.is_some()));
assert!(params.swa_uuid_for_lock.is_some());
panic!("low-priority release dispatched");
}
}
@@ -368,7 +368,7 @@ impl TreeComponent<Vec<i64>> for SwaComponentForTest {
&self,
_tree_core: &mut UnifiedTreeCore<Vec<i64>>,
_node_id: NodeIdx_,
_params: Option<&DecLockRefParams>,
_params: &DecLockRefParams,
_lock_host: bool,
) {
unimplemented!()
@@ -481,7 +481,7 @@ impl TreeComponent<Vec<i64>> for SwaEvictionComponentForTest {
&self,
_tree_core: &mut UnifiedTreeCore<Vec<i64>>,
_node_id: NodeIdx_,
_params: Option<&DecLockRefParams>,
_params: &DecLockRefParams,
_lock_host: bool,
) {
unimplemented!()
@@ -503,7 +503,7 @@ fn locked_anchor_for_dispatch(tc: &mut UnifiedTreeCore<Vec<i64>>) -> NodeIdx_ {
tc.arena
.set_device_value(n1, FULL, Tensor::from_slice(&[0i64, 1]));
tc.component_state_mut(FULL).evictable_size = 2;
tc.inc_lock_ref(tc.arena.node(n1).id)
tc.inc_lock_ref(tc.arena.node(n1).id, ComponentSet::EMPTY)
.expect("live test node");
n1
}
@@ -516,7 +516,11 @@ fn dec_lock_ref_skip_swa_skips_the_swa_component() {
// The skipped Swa driver is never dispatched, so its stub cannot panic.
tc.dec_lock_ref(
tc.arena.node(n1).id,
/* params = */ None,
/* params = */
&DecLockRefParams {
skipped_lock_components: ComponentSet::EMPTY,
..Default::default()
},
/* skip_swa = */ true,
)
.expect("live test node");
@@ -541,7 +545,7 @@ fn inc_lock_ref_reaches_every_component() {
tc.arena
.set_device_value(n1, FULL, Tensor::from_slice(&[0i64, 1]));
tc.component_state_mut(FULL).evictable_size = 2;
let _ = tc.inc_lock_ref(tc.arena.node(n1).id);
let _ = tc.inc_lock_ref(tc.arena.node(n1).id, ComponentSet::EMPTY);
}
#[test]
@@ -552,7 +556,11 @@ fn dec_lock_ref_without_skip_swa_reaches_every_component() {
tc.register_component_(Arc::new(SwaComponentForTest));
let _ = tc.dec_lock_ref(
tc.arena.node(n1).id,
/* params = */ None,
/* params = */
&DecLockRefParams {
skipped_lock_components: ComponentSet::EMPTY,
..Default::default()
},
/* skip_swa = */ false,
);
}
@@ -630,7 +638,11 @@ fn dec_swa_lock_only_dispatches_lower_priority_releases() {
let mut host_frees = HashMap::new();
let _ = tc.dec_swa_lock_only(
tc.arena.node(root).id,
Some(7),
&DecLockRefParams {
swa_uuid_for_lock: Some(7),
skipped_lock_components: ComponentSet::EMPTY,
..Default::default()
},
&mut device_frees,
&mut host_frees,
);
@@ -671,7 +683,11 @@ fn dec_swa_lock_only_returns_device_frees_in_the_device_dict() {
let mut host_frees = HashMap::new();
tc.dec_swa_lock_only(
tc.arena.node(a).id,
result.swa_uuid_for_lock,
&DecLockRefParams {
swa_uuid_for_lock: result.swa_uuid_for_lock,
skipped_lock_components: ComponentSet::EMPTY,
..Default::default()
},
&mut device_frees,
&mut host_frees,
)
@@ -3755,11 +3771,15 @@ fn commit_load_back_reattaches_device_slices_and_restores_the_match() {
assert_eq!(tc.full_evictable_size(), 4);
// The orchestrator re-locks the loaded path right after commit; that lock walk
// also re-evaluates the parent's transient D-leaf membership.
tc.inc_lock_ref(tc.arena.node(child).id)
tc.inc_lock_ref(tc.arena.node(child).id, ComponentSet::EMPTY)
.expect("live test node");
tc.dec_lock_ref(
tc.arena.node(child).id,
/* params = */ None,
/* params = */
&DecLockRefParams {
skipped_lock_components: ComponentSet::EMPTY,
..Default::default()
},
/* skip_swa = */ false,
)
.expect("live test node");
@@ -5261,7 +5281,7 @@ fn stale_handle_returns_err_after_its_node_is_freed() {
tc.evict_device_leaf(leaf, /* is_write_back = */ false)
.expect("live test node");
assert!(matches!(
tc.inc_lock_ref(leaf),
tc.inc_lock_ref(leaf, ComponentSet::EMPTY),
Err(NodeAccessError { node_id }) if node_id == leaf
));
}
@@ -6164,7 +6184,7 @@ fn reset_restores_a_fresh_tree() {
..insert_params(&vec![7, 8], &[20, 21])
});
let matched = tc.match_prefix(&match_params(&vec![1, 2, 3]));
tc.inc_lock_ref(matched.best_match_node_id)
tc.inc_lock_ref(matched.best_match_node_id, ComponentSet::EMPTY)
.expect("live match node");
assert_eq!(tc.protected_size(), 3);
// Seed aux LRU, host LRU, and host-leaf state so the reset must clear each.
@@ -6216,7 +6236,7 @@ fn size_accessors_mirror_the_full_component_state() {
assert_eq!(tc.protected_size(), 0);
assert_eq!(tc.component_evictable_size(FULL), 3);
let matched = tc.match_prefix(&match_params(&vec![1, 2, 3]));
tc.inc_lock_ref(matched.best_match_node_id)
tc.inc_lock_ref(matched.best_match_node_id, ComponentSet::EMPTY)
.expect("live match node");
assert_eq!(tc.protected_size(), 3);
assert_eq!(tc.full_protected_size(), 3);
@@ -6306,7 +6326,7 @@ fn walk_for_kv_canary_chains_slots_across_namespaces() {
fn walk_for_kv_canary_unlocked_only_skips_locked_nodes_but_keeps_the_chain() {
let mut tc = core();
let (a, _b) = matched_chain(&mut tc);
tc.inc_lock_ref(tc.arena.node(a).id)
tc.inc_lock_ref(tc.arena.node(a).id, ComponentSet::EMPTY)
.expect("live test node");
assert_eq!(
sorted_canary_rows(tc.walk_for_kv_canary(true, false)),
@@ -6659,13 +6679,18 @@ fn sanity_check_passes_on_a_healthy_tree() {
let leaf = tc
.match_prefix(&match_params(&vec![1, 2, 9]))
.best_match_node_id;
tc.inc_lock_ref(leaf).expect("live test node");
tc.inc_lock_ref(leaf, ComponentSet::EMPTY)
.expect("live test node");
tc.sanity_check(&[(1, leaf)], &[(2, leaf)]);
tc.dec_lock_ref(
tc.arena
.node(tc.arena.resolve(leaf).expect("live test node"))
.id,
/* params = */ None,
/* params = */
&DecLockRefParams {
skipped_lock_components: ComponentSet::EMPTY,
..Default::default()
},
/* skip_swa = */ false,
)
.expect("live test node");
@@ -6774,16 +6799,28 @@ fn sanity_check_detects_an_evicted_parent_prefix() {
}
#[test]
#[should_panic(expected = "evicted but lock_ref")]
fn sanity_check_detects_a_locked_tombstone() {
fn sanity_check_accepts_a_locked_tombstone() {
// Segment locks count evicted nodes, so a device-locked tombstone is a
// legal state the checker must not flag.
let mut tc = sane_tree();
// write_back spares the tombstone's ancestors the backup-chain rule.
tc.is_write_back = true;
let leaf = tc
.match_prefix(&match_params(&vec![1, 2, 9]))
.best_match_node_id;
tc.inc_lock_ref(leaf).expect("live test node");
let _ = tc
.arena
.take_device_value(tc.arena.resolve(leaf).expect("live test node"), FULL);
let leaf_idx = tc.arena.resolve(leaf).expect("live test node");
// Tombstone the leaf consistently first (host copy, ledger, leaf sets),
// then lock through it: the bottom segment counts the tombstone.
tc.arena
.set_host_value(leaf_idx, FULL, Tensor::from_slice(&[9i64]));
let taken = tc.arena.take_device_value(leaf_idx, FULL);
tc.dec_evictable_size(FULL, taken.size()[0] as usize);
tc.update_evictable_leaf_sets_(leaf_idx);
let parent_idx = tc.arena.node(leaf_idx).parent();
tc.update_evictable_leaf_sets_(parent_idx);
tc.inc_lock_ref(leaf, ComponentSet::EMPTY)
.expect("live test node");
assert_eq!(tc.arena.device_lock_ref(leaf_idx, FULL), 1);
tc.sanity_check(&[], &[]);
}
@@ -8049,13 +8086,16 @@ fn run_random_op_sequence(mut tc: UnifiedTreeCore<Vec<i64>>, page: usize, mamba:
2 => {
// Balanced lock round trip on whatever the key matches.
let anchor = tc.match_prefix(&match_params(&key)).best_match_node_id;
let lock = tc.inc_lock_ref(anchor).expect("live match anchor");
let lock = tc
.inc_lock_ref(anchor, ComponentSet::EMPTY)
.expect("live match anchor");
let params = DecLockRefParams {
node_id: None,
swa_uuid_for_lock: lock.swa_uuid_for_lock,
swa_uuid_for_host_lock: lock.swa_uuid_for_host_lock,
skip_lock_node_ids: lock.skip_lock_node_ids,
skipped_lock_components: lock.skipped_lock_components,
};
tc.dec_lock_ref(anchor, Some(&params), /* skip_swa = */ false)
tc.dec_lock_ref(anchor, &params, /* skip_swa = */ false)
.expect("live match anchor");
}
_ => {
@@ -8063,7 +8103,9 @@ fn run_random_op_sequence(mut tc: UnifiedTreeCore<Vec<i64>>, page: usize, mamba:
let matched = tc.match_prefix(&match_params(&key));
let anchor = matched.best_match_node_id;
let matched_len = matched.device_indices.numel() as usize;
let lock = tc.inc_lock_ref(anchor).expect("live match anchor");
let lock = tc
.inc_lock_ref(anchor, ComponentSet::EMPTY)
.expect("live match anchor");
tc.insert(&sequence_insert_params(
&key,
matched_len,
@@ -8072,11 +8114,12 @@ fn run_random_op_sequence(mut tc: UnifiedTreeCore<Vec<i64>>, page: usize, mamba:
mamba,
));
let params = DecLockRefParams {
node_id: None,
swa_uuid_for_lock: lock.swa_uuid_for_lock,
swa_uuid_for_host_lock: lock.swa_uuid_for_host_lock,
skip_lock_node_ids: lock.skip_lock_node_ids,
skipped_lock_components: lock.skipped_lock_components,
};
tc.dec_lock_ref(anchor, Some(&params), /* skip_swa = */ false)
tc.dec_lock_ref(anchor, &params, /* skip_swa = */ false)
.expect("live match anchor");
}
}
@@ -8224,10 +8267,17 @@ fn a_zero_length_match_anchors_at_the_root() {
.best_match_node_id;
assert_eq!(anchor, tc.root_node_handle(Some("salted")));
// The root handle stays valid across a full namespace eviction.
tc.inc_lock_ref(anchor).expect("live root");
tc.inc_lock_ref(anchor, ComponentSet::EMPTY)
.expect("live root");
drain_full_device(&mut tc);
tc.dec_lock_ref(
anchor, /* params = */ None, /* skip_swa = */ false,
anchor,
/* params = */
&DecLockRefParams {
skipped_lock_components: ComponentSet::EMPTY,
..Default::default()
},
/* skip_swa = */ false,
)
.expect("live root");
assert!(tc.arena.resolve(anchor).is_ok());
rust/sglang-radix-tree/src/unified_tree_core.rs
+146 -96
View File
@@ -8,7 +8,9 @@ use std::sync::Arc;
use sha2::{Digest, Sha256};
use tch::{Device, Kind, Tensor};
use crate::components::{self, FullComponent, MambaComponent, SwaComponent, TreeComponent};
use crate::components::{
self, ComponentSet, FullComponent, MambaComponent, SwaComponent, TreeComponent,
};
use crate::components::{
BASE_COMPONENT_TYPE, ComponentType, FULL, MAMBA, NUM_COMPONENT_TYPES, SWA,
};
@@ -34,28 +36,41 @@ fn next_coexist_reclaim_digest(current: i64, node_id: NodeId, component_idx: usi
// ---- interface types ----
/// Result of `inc_lock_ref`, handed back to the matching `dec_lock_ref`.
///
/// The receipt a release needs is per-component lock evidence: the SWA
/// segment boundary uuid (None means the segment reached the root) and
/// whether the single-node Mamba lock was taken (the decode hold opts
/// out). Locks count every node in their contiguous segment, so no
/// per-node skip state exists. Receipt fields default to nothing-acquired;
/// `inc_lock_ref` stamps what it actually took.
#[derive(Default)]
pub struct IncLockRefResult {
/// Tokens newly protected (moved out of evictable) by this lock.
pub delta: Option<usize>,
/// The node the lock was taken on; a release replays the receipt there only.
pub node_id: Option<NodeId>,
/// SWA lock-window uuid minted/reused by the device lock walk.
pub swa_uuid_for_lock: Option<i64>,
/// SWA lock-window uuid minted/reused by the host lock walk.
pub swa_uuid_for_host_lock: Option<i64>,
/// Per-component nodes that were tombstones at acquire time; replayed at
/// release so the unlock skips them.
pub skip_lock_node_ids: HashMap<ComponentType, HashSet<NodeId>>,
/// Components the acquire left untaken; the release skips them too.
pub skipped_lock_components: ComponentSet,
}
/// Params for `dec_lock_ref`.
/// Params for `dec_lock_ref`. Receipt fields default to nothing-acquired so
/// a lost receipt under-releases (a leak sanity checks report) instead of
/// releasing a lock another holder owns.
#[derive(Default)]
pub struct DecLockRefParams {
/// The node the matching acquire locked; None only for receipts that did
/// not come from this core (a mispaired anchor is a protocol violation).
pub node_id: Option<NodeId>,
/// SWA lock-window uuid the device unlock stops at, from the matching acquire.
pub swa_uuid_for_lock: Option<i64>,
/// SWA lock-window uuid the host unlock stops at, from the matching acquire.
pub swa_uuid_for_host_lock: Option<i64>,
/// Per-component nodes the unlock walk skips (from the matching acquire).
pub skip_lock_node_ids: HashMap<ComponentType, HashSet<NodeId>>,
/// Components the matching acquire left untaken.
pub skipped_lock_components: ComponentSet,
}
/// Result of `dec_lock_ref`.
@@ -792,58 +807,90 @@ impl<K: ChildKeyType> UnifiedTreeCore<K> {
self.swa_uuid_counter
}
/// Bump the reference count on a node's component locks.
pub fn inc_lock_ref(&mut self, node_id: NodeId) -> Result<IncLockRefResult, NodeAccessError> {
self.inc_lock_ref_with_skip(node_id, &[])
}
/// Bump component locks, leaving explicitly skipped target components evictable.
pub fn inc_lock_ref_with_skip(
/// Bump the reference count on a node's component locks. Components in
/// `skip_lock_components` are left untaken; the receipt records the anchor
/// node and the skipped set so the paired release mirrors them.
pub fn inc_lock_ref(
&mut self,
node_id: NodeId,
skip_lock_components: &[ComponentType],
skip_lock_components: ComponentSet,
) -> Result<IncLockRefResult, NodeAccessError> {
let node_id = self.arena.resolve(node_id)?;
let node = self.arena.node(node_id);
let node_handle = node.id;
let is_root = node.is_root();
let mut result = IncLockRefResult::default();
let node_idx = self.arena.resolve(node_id)?;
let mut result = IncLockRefResult {
node_id: Some(self.arena.node(node_idx).id),
skipped_lock_components: skip_lock_components,
..Default::default()
};
for i in 0..self.components.len() {
let component_type = self.components[i].component_type();
if skip_lock_components.contains(&component_type) {
if !is_root {
result
.skip_lock_node_ids
.entry(component_type)
.or_default()
.insert(node_handle);
}
let component = Arc::clone(&self.components[i]);
if skip_lock_components.contains(component.component_type()) {
continue;
}
let component = Arc::clone(&self.components[i]);
result = component
.acquire_component_lock(self, node_id, result, /* lock_host = */ false);
.acquire_component_lock(self, node_idx, result, /* lock_host = */ false);
}
self.update_evictable_leaf_sets_(node_id);
self.update_evictable_leaf_sets_(node_idx);
Ok(result)
}
/// Decrease the reference count on a node's component locks.
/// A receipt releases only the node its acquire returned; a mispaired
/// node would silently release (or steal) another holder's segment.
fn assert_receipt_anchor_(&self, node_idx: NodeIdx_, params: &DecLockRefParams) {
if let Some(anchor) = params.node_id {
let node_handle = self.arena.node(node_idx).id;
assert!(
anchor == node_handle,
"lock receipt anchored on node {anchor} released on node {node_handle}"
);
}
}
/// Release each component this receipt acquired. Auxiliaries go first so
/// Full, whose walk refreshes leaf membership on every node it unlocks,
/// sees their final refs; the auxiliary walks also refresh the nodes they
/// unlock, so the order is not load-bearing for the sets.
fn release_components_(
&mut self,
node_idx: NodeIdx_,
params: &DecLockRefParams,
lock_host: bool,
skip_swa_and_below: bool,
) {
let swa_priority = if skip_swa_and_below {
self.try_component_by_type_(SWA)
.map(|swa| swa.eviction_priority(/* is_leaf = */ false))
} else {
None
};
for i in (0..self.components.len()).rev() {
let component = Arc::clone(&self.components[i]);
let ct = component.component_type();
if params.skipped_lock_components.contains(ct) {
continue;
}
if let Some(swa_priority) = swa_priority
&& (ct == SWA || component.eviction_priority(/* is_leaf = */ false) < swa_priority)
{
continue;
}
component.release_component_lock(self, node_idx, params, lock_host);
}
}
/// Decrease the reference count on a node's component locks. The receipt
/// is required: a release must replay its acquire's evidence. After an SWA
/// early release (`dec_swa_lock_only`), `skip_swa` leaves SWA and the
/// lower-priority components it already dropped alone.
pub fn dec_lock_ref(
&mut self,
node_id: NodeId,
params: Option<&DecLockRefParams>,
params: &DecLockRefParams,
skip_swa: bool,
) -> Result<DecLockRefResult, NodeAccessError> {
let node_id = self.arena.resolve(node_id)?;
for i in 0..self.components.len() {
if skip_swa && self.components[i].component_type() == SWA {
continue;
}
let component = Arc::clone(&self.components[i]);
component.release_component_lock(self, node_id, params, /* lock_host = */ false);
}
self.update_evictable_leaf_sets_(node_id);
let node_idx = self.arena.resolve(node_id)?;
self.assert_receipt_anchor_(node_idx, params);
self.release_components_(node_idx, params, /* lock_host = */ false, skip_swa);
self.update_evictable_leaf_sets_(node_idx);
// TODO: delta is not aggregated from components; no caller uses it yet.
Ok(DecLockRefResult::default())
}
@@ -853,50 +900,37 @@ impl<K: ChildKeyType> UnifiedTreeCore<K> {
pub fn dec_swa_lock_only(
&mut self,
node_id: NodeId,
swa_uuid_for_lock: Option<i64>,
params: &DecLockRefParams,
device_frees: &mut HashMap<ComponentType, Vec<Tensor>>,
host_frees: &mut HashMap<ComponentType, Vec<Tensor>>,
) -> Result<(), NodeAccessError> {
self.dec_swa_lock_only_with_skip(
node_id,
swa_uuid_for_lock,
/* skip_lock_node_ids = */ None,
device_frees,
host_frees,
)
}
/// Skip-aware variant used when an acquire deliberately omitted a component.
pub fn dec_swa_lock_only_with_skip(
&mut self,
node_id: NodeId,
swa_uuid_for_lock: Option<i64>,
skip_lock_node_ids: Option<&HashMap<ComponentType, HashSet<NodeId>>>,
device_frees: &mut HashMap<ComponentType, Vec<Tensor>>,
host_frees: &mut HashMap<ComponentType, Vec<Tensor>>,
) -> Result<(), NodeAccessError> {
let node_id = self.arena.resolve(node_id)?;
let node_idx = self.arena.resolve(node_id)?;
self.assert_receipt_anchor_(node_idx, params);
let Some(swa) = self.try_component_by_type_(SWA) else {
return Ok(());
};
swa.release_window_lock(self, node_id, swa_uuid_for_lock, device_frees, host_frees);
swa.release_window_lock(
self,
node_idx,
params.swa_uuid_for_lock,
device_frees,
host_frees,
);
// Drop strictly-lower-priority locks (e.g. Mamba) co-located on the node.
// Drop strictly-lower-priority locks co-located on the node, skipping
// any the paired inc never took.
let swa_priority = swa.eviction_priority(/* is_leaf = */ false);
let dec_params = DecLockRefParams {
swa_uuid_for_lock,
skip_lock_node_ids: skip_lock_node_ids.cloned().unwrap_or_default(),
..Default::default()
};
for i in 0..self.components.len() {
for i in (0..self.components.len()).rev() {
let component = Arc::clone(&self.components[i]);
if params
.skipped_lock_components
.contains(component.component_type())
{
continue;
}
if component.eviction_priority(/* is_leaf = */ false) < swa_priority {
component.release_component_lock(
self,
node_id,
Some(&dec_params),
/* lock_host = */ false,
);
component
.release_component_lock(self, node_idx, params, /* lock_host = */ false);
}
}
Ok(())
@@ -925,29 +959,31 @@ impl<K: ChildKeyType> UnifiedTreeCore<K> {
&mut self,
node_id: NodeId,
) -> Result<IncLockRefResult, NodeAccessError> {
let node_id = self.arena.resolve(node_id)?;
let mut result = IncLockRefResult::default();
let node_idx = self.arena.resolve(node_id)?;
let mut result = IncLockRefResult {
node_id: Some(self.arena.node(node_idx).id),
..Default::default()
};
for i in 0..self.components.len() {
let component = Arc::clone(&self.components[i]);
result = component
.acquire_component_lock(self, node_id, result, /* lock_host = */ true);
.acquire_component_lock(self, node_idx, result, /* lock_host = */ true);
}
self.update_evictable_leaf_sets_(node_id);
self.update_evictable_leaf_sets_(node_idx);
Ok(result)
}
/// Decrease the reference count on a node's host-side component locks.
/// The receipt is required, as for `dec_lock_ref`.
pub fn dec_host_lock_ref(
&mut self,
node_id: NodeId,
params: Option<&DecLockRefParams>,
params: &DecLockRefParams,
) -> Result<DecLockRefResult, NodeAccessError> {
let node_id = self.arena.resolve(node_id)?;
for i in 0..self.components.len() {
let component = Arc::clone(&self.components[i]);
component.release_component_lock(self, node_id, params, /* lock_host = */ true);
}
self.update_evictable_leaf_sets_(node_id);
let node_idx = self.arena.resolve(node_id)?;
self.assert_receipt_anchor_(node_idx, params);
self.release_components_(node_idx, params, /* lock_host = */ true, false);
self.update_evictable_leaf_sets_(node_idx);
Ok(DecLockRefResult::default())
}
@@ -1877,7 +1913,14 @@ impl<K: ChildKeyType> UnifiedTreeCore<K> {
pub fn unevict_node_on_insert_(&mut self, node_id: NodeIdx_, fresh_value: &Tensor) {
self.arena
.set_device_value(node_id, FULL, fresh_value.copy());
self.inc_evictable_size(FULL, fresh_value.size()[0] as usize);
let tokens = fresh_value.size()[0] as usize;
// A value materialized under lock is protected; the last release
// moves it to evictable.
if self.arena.device_lock_ref(node_id, FULL) > 0 {
self.inc_protected_size(FULL, tokens);
} else {
self.inc_evictable_size(FULL, tokens);
}
self.update_evictable_leaf_sets_(node_id);
self.update_full_coexisting_host_tracking_(node_id);
if let Some(parent_id) = self.arena.node(node_id).try_parent() {
@@ -2644,6 +2687,11 @@ impl<K: ChildKeyType> UnifiedTreeCore<K> {
if node.is_host_locked() {
return false;
}
// Segment locks count evicted nodes too: a device-locked candidate is
// a live segment's anchor, and evict_host_leaf_ would delete it.
if node.is_device_locked() {
return false;
}
if !node.children.is_empty() {
return false;
}
@@ -3702,7 +3750,13 @@ impl<K: ChildKeyType> UnifiedTreeCore<K> {
host_lru.remove_node(node_id);
}
self.device_lru_list_mut(component_type).insert_mru(node_id);
self.inc_evictable_size(component_type, tokens);
// A value materialized under lock is protected; the last release
// moves it to evictable.
if self.arena.device_lock_ref(node_id, component_type) > 0 {
self.inc_protected_size(component_type, tokens);
} else {
self.inc_evictable_size(component_type, tokens);
}
}
/// The component's device value on the node, or None if evicted.
@@ -3931,12 +3985,8 @@ impl<K: ChildKeyType> UnifiedTreeCore<K> {
device_state.lock_ref
));
}
if device_state.value.is_none() && device_state.lock_ref > 0 {
errors.push(format!(
"node {node_id} {ct:?} evicted but lock_ref={}",
device_state.lock_ref
));
}
// Locked tombstones are legal: segment locks count every
// node in [start, boundary], data-bearing or not.
}
// Collect expected leaf qualification (single pass)