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

[Rust TreeCore] Harden runtime and CI parity (#37303)

Browse Source
This commit is contained in:
Jialin Ouyang
2026-09-09 10:40:22 +08:00
committed by GitHub
parent e54ff1efb9
commit 7a464a7014
22 changed files with 2390 additions and 1137 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.github/actions/download-rust-ext/action.yml
+1 -1
View File
@@ -53,7 +53,7 @@ runs:
path: |
python/sglang/srt/rust_extensions/_*.so
python/sglang/srt/mem_cache/rust_tree_core/mem_cache*.so
key: ${{ inputs.cache_key_prefix }}-${{ hashFiles('rust/**', 'python/setup.py', 'python/pyproject.toml', 'python/sglang/srt/rust_extensions/torch_build.py', '.github/workflows/_pr-test-rust-ext-build.yml', 'scripts/ci/utils/stage_rust_ext_modules.sh') }}
key: ${{ inputs.cache_key_prefix }}-${{ hashFiles('rust/**', 'proto/**', 'python/setup.py', 'python/pyproject.toml', 'python/sglang/srt/rust_extensions/torch_build.py', '.github/workflows/_pr-test-rust-ext-build.yml', 'scripts/ci/utils/stage_rust_ext_modules.sh') }}
# Job-wide, but only setup.py reads it, and only while building.
# Whether the modules suit this interpreter is not decided here:
.github/workflows/_pr-test-rust-ext-build.yml
+6 -3
View File
@@ -88,10 +88,12 @@ jobs:
ref: ${{ inputs.git_ref || github.sha }}
# Just what the cache key hashes, plus the action and script this job
# runs: the workspace is cold here and the rest of the tree is mostly
# docs. Both jobs must hash the same rust/** set, which this preserves.
# docs. Both jobs must hash the same Rust extension inputs, which this
# preserves.
# Cone mode off is what allows naming a single file.
sparse-checkout: |
rust
proto
python/setup.py
python/pyproject.toml
python/sglang/srt/rust_extensions/torch_build.py
@@ -115,7 +117,7 @@ jobs:
path: |
python/sglang/srt/rust_extensions/_*.so
python/sglang/srt/mem_cache/rust_tree_core/mem_cache*.so
key: ${{ inputs.cache_key_prefix }}-${{ hashFiles('rust/**', 'python/setup.py', 'python/pyproject.toml', 'python/sglang/srt/rust_extensions/torch_build.py', '.github/workflows/_pr-test-rust-ext-build.yml', 'scripts/ci/utils/stage_rust_ext_modules.sh') }}
key: ${{ inputs.cache_key_prefix }}-${{ hashFiles('rust/**', 'proto/**', 'python/setup.py', 'python/pyproject.toml', 'python/sglang/srt/rust_extensions/torch_build.py', '.github/workflows/_pr-test-rust-ext-build.yml', 'scripts/ci/utils/stage_rust_ext_modules.sh') }}
# On a miss: different hash = rust/setup.py moved; no entries = evicted.
- name: Report cache lookup
@@ -324,6 +326,7 @@ jobs:
# both have to check out the same set for hashFiles to agree.
sparse-checkout: |
rust
proto
python/setup.py
python/pyproject.toml
python/sglang/srt/rust_extensions/torch_build.py
@@ -359,7 +362,7 @@ jobs:
path: |
python/sglang/srt/rust_extensions/_*.so
python/sglang/srt/mem_cache/rust_tree_core/mem_cache*.so
key: ${{ inputs.cache_key_prefix }}-${{ hashFiles('rust/**', 'python/setup.py', 'python/pyproject.toml', 'python/sglang/srt/rust_extensions/torch_build.py', '.github/workflows/_pr-test-rust-ext-build.yml', 'scripts/ci/utils/stage_rust_ext_modules.sh') }}
key: ${{ inputs.cache_key_prefix }}-${{ hashFiles('rust/**', 'proto/**', 'python/setup.py', 'python/pyproject.toml', 'python/sglang/srt/rust_extensions/torch_build.py', '.github/workflows/_pr-test-rust-ext-build.yml', 'scripts/ci/utils/stage_rust_ext_modules.sh') }}
- name: Upload extension modules
uses: actions/upload-artifact@v4
.github/workflows/seed-rust-ext-cache.yml
+1
View File
@@ -9,6 +9,7 @@ on:
branches: [main]
paths:
- 'rust/**'
- 'proto/**'
- 'python/setup.py'
- 'python/pyproject.toml'
- 'python/sglang/srt/rust_extensions/torch_build.py'
python/sglang/srt/rust_extensions/loader.py
+22 -7
View File
@@ -55,6 +55,7 @@ class _CrateSpec:
manifest: Path
workspace: Path
features: tuple[str, ...]
source_inputs: tuple[Path, ...]
@dataclass(frozen=True)
@@ -80,7 +81,8 @@ def load_rust_extension(
The crate is discovered from the workspace under ``rust/``: the one whose
Cargo manifest declares ``[package.metadata.sglang] python-module`` equal
to ``python_module`` (the same metadata setup.py uses for wheel builds), so
new crates need no registration here.
new crates need no registration here. Crates may declare ``source-inputs``
relative to their manifest for build inputs outside the Rust workspace.
``auto`` prefers a module bundled in an installed wheel. In a source tree,
it ignores unverified in-package artifacts and uses the fingerprinted cache
@@ -152,9 +154,12 @@ def load_rust_extension(
features=features,
build_environment=build_environment,
)
if _source_digest(crate.workspace) != context.source_digest:
if (
_source_digest(crate.workspace, crate.source_inputs)
!= context.source_digest
):
raise RuntimeError(
f"Rust sources under {crate.workspace} changed during the build; "
f"Rust extension sources for {crate.package} changed during the build; "
"the result was not cached"
)
_stage_atomically(artifact, extension_path)
@@ -216,6 +221,10 @@ def _discover_crate(workspace: Path, python_module: str) -> _CrateSpec:
manifest=manifest,
workspace=crate_workspace,
features=tuple(sglang_metadata.get("features", ())),
source_inputs=tuple(
(manifest.parent / path).resolve()
for path in sglang_metadata.get("source-inputs", ())
),
)
)
@@ -245,7 +254,7 @@ def _build_context(
features = crate.features
if extension_module is None:
extension_module = crate.python_module
source_digest = _source_digest(crate.workspace)
source_digest = _source_digest(crate.workspace, crate.source_inputs)
toolchain = {
"cargo": _command_version(
"cargo", "--version", "--verbose", cwd=crate.workspace
@@ -289,10 +298,16 @@ def _build_context(
)
def _source_digest(workspace: Path) -> str:
def _source_digest(workspace: Path, source_inputs: tuple[Path, ...] = ()) -> str:
digest = hashlib.sha256()
for path in _source_files(workspace):
relative_path = path.relative_to(workspace).as_posix().encode()
paths = set(_source_files(workspace))
for source_input in source_inputs:
if source_input.is_dir():
paths.update(_source_files(source_input))
else:
paths.add(source_input)
for path in sorted(paths, key=lambda item: os.path.relpath(item, workspace)):
relative_path = Path(os.path.relpath(path, workspace)).as_posix().encode()
digest.update(len(relative_path).to_bytes(8, "big"))
digest.update(relative_path)
if path.is_symlink():
rust/sglang-grpc/Cargo.toml
+3
View File
@@ -9,6 +9,9 @@ license.workspace = true
# of the main sglang wheel at the given import path.
[package.metadata.sglang]
python-module = "sglang.srt.rust_extensions._grpc"
# build.rs compiles schemas from outside the Rust workspace; include them in
# the source-loader artifact fingerprint.
source-inputs = ["../../proto"]
# Always build optimized, even for an editable install.
debug = false
rust/sglang-radix-tree/src/components/full.rs
+8 -4
View File
@@ -49,6 +49,8 @@ impl<K: ChildKeyType> TreeComponent<K> for FullComponent {
&self,
tree_core: &UnifiedTreeCore<K>,
mut result: MatchResult,
last_device_node_idx: NodeIdx_,
best_match_node_idx: NodeIdx_,
params: &MatchPrefixParams<'_, K>,
value_chunks: &[Tensor],
best_value_len: usize,
@@ -56,9 +58,8 @@ impl<K: ChildKeyType> TreeComponent<K> for FullComponent {
// Compute Full KV host hit length: walk from last_host_node up to
// last_device_node, summing host_value lengths of evicted nodes.
let mut kv_host_hit = 0;
let mut node_idx = tree_core.arena.resolve(result.best_match_node_id);
let last_device_idx = tree_core.arena.resolve(result.last_device_node_id);
while node_idx != last_device_idx {
let mut node_idx = best_match_node_idx;
while node_idx != last_device_node_idx {
let node = tree_core.arena.node(node_idx);
let parent = node.try_parent().unwrap_or_else(|| {
panic!(
@@ -469,7 +470,10 @@ impl<K: ChildKeyType> TreeComponent<K> for FullComponent {
{
let mut offset = 0i64;
for &loaded_id in transfer.nodes_to_load.iter().flatten() {
let loaded_idx = tree_core.arena.resolve(loaded_id);
let loaded_idx = tree_core
.arena
.resolve(loaded_id)
.expect("load-back transfers must reference live nodes");
let loaded = tree_core.arena.node_mut(loaded_idx);
let n_len = loaded.host_value_len(FULL) as i64;
loaded
rust/sglang-radix-tree/src/components/mamba.rs
+9 -4
View File
@@ -127,6 +127,8 @@ impl<K: ChildKeyType> TreeComponent<K> for MambaComponent {
&self,
tree_core: &UnifiedTreeCore<K>,
mut result: MatchResult,
_last_device_node_idx: NodeIdx_,
best_match_node_idx: NodeIdx_,
_params: &MatchPrefixParams<'_, K>,
_value_chunks: &[Tensor],
_best_value_len: usize,
@@ -143,9 +145,7 @@ impl<K: ChildKeyType> TreeComponent<K> for MambaComponent {
// HiCache: if mamba was evicted from device but has host backup,
// ensure mamba_host_hit_length >= 1 so load_back is triggered.
let last_node = tree_core
.arena
.node(tree_core.arena.resolve(result.best_match_node_id));
let last_node = tree_core.arena.node(best_match_node_idx);
if !last_node.has_device_value(MAMBA) && last_node.has_host_value(MAMBA) {
result.mamba_host_hit_length = result.mamba_host_hit_length.max(1);
}
@@ -643,7 +643,12 @@ impl<K: ChildKeyType> TreeComponent<K> for MambaComponent {
let target_node_id = insert_result
.as_deref()
.and_then(|result| result.inserted_host_node)
.map(|id| tree_core.arena.resolve(id));
.map(|id| {
tree_core
.arena
.resolve(id)
.expect("prefetch insert results must reference live nodes")
});
let attach_target = match (host_indices, target_node_id) {
(Some(_), Some(target))
if loaded && !tree_core.arena.has_host_value(target, MAMBA) =>
rust/sglang-radix-tree/src/components/mod.rs
+2
View File
@@ -130,6 +130,8 @@ pub trait TreeComponent<K: ChildKeyType> {
&self,
tree_core: &UnifiedTreeCore<K>,
result: MatchResult,
_last_device_node_idx: NodeIdx_,
_best_match_node_idx: NodeIdx_,
params: &MatchPrefixParams<'_, K>,
value_chunks: &[Tensor],
best_value_len: usize,
rust/sglang-radix-tree/src/components/swa.rs
+22 -5
View File
@@ -225,7 +225,12 @@ impl SwaComponent {
});
let target = insert_result
.and_then(|result| result.inserted_host_node)
.map(|id| tree_core.arena.resolve(id));
.map(|id| {
tree_core
.arena
.resolve(id)
.expect("prefetch insert results must reference live nodes")
});
let (Some(target), Some(host_indices)) = (target, transfer.host_indices.as_ref()) else {
if let Some(host_indices) = &transfer.host_indices {
@@ -233,6 +238,15 @@ impl SwaComponent {
}
return;
};
// Cache-mode graft commit only (buffer fills never reach here):
// a hit-shrunk window mid-tree is missing its head, so drop it.
// Root anchors are complete windows of their own.
if node_id != tree_core.arena.root()
&& window_require_pages < self.sliding_window_size.div_ceil(page_size)
{
self.release_swa_host_(host_indices.shallow_clone(), cache_actions);
return;
}
if window_require_pages == 0 || loaded_pages < window_require_pages {
self.release_swa_host_(host_indices.shallow_clone(), cache_actions);
return;
@@ -338,6 +352,8 @@ impl<K: ChildKeyType> TreeComponent<K> for SwaComponent {
&self,
tree_core: &UnifiedTreeCore<K>,
mut result: MatchResult,
_last_device_node_idx: NodeIdx_,
best_match_node_idx: NodeIdx_,
params: &MatchPrefixParams<'_, K>,
value_chunks: &[Tensor],
best_value_len: usize,
@@ -347,9 +363,7 @@ impl<K: ChildKeyType> TreeComponent<K> for SwaComponent {
// toward the SWA host hit.
let mut n_swa = 0;
let mut swa_host_hit = 0;
let mut node = tree_core
.arena
.node(tree_core.arena.resolve(result.best_match_node_id));
let mut node = tree_core.arena.node(best_match_node_idx);
while !node.is_root() && n_swa < self.sliding_window_size {
if node.has_device_value(SWA) {
n_swa += node.device_value_len(SWA);
@@ -965,7 +979,10 @@ impl<K: ChildKeyType> TreeComponent<K> for SwaComponent {
let mut swa_chunks: Vec<Tensor> = Vec::new();
let mut offset = 0i64;
for &loaded_id in transfer.nodes_to_load.iter().flatten() {
let loaded_idx = tree_core.arena.resolve(loaded_id);
let loaded_idx = tree_core
.arena
.resolve(loaded_id)
.expect("load-back transfers must reference live nodes");
let n_tokens = tree_core.arena.host_value_len(loaded_idx, SWA) as i64;
let swa_chunk = device_indices.narrow(0, offset, n_tokens).copy();
tree_core.set_component_device_value_(
rust/sglang-radix-tree/src/node.rs
+18 -13
View File
@@ -688,14 +688,21 @@ pub struct ValueState {
// Tree-core runtime errors.
/// A public node handle does not name a live arena node.
#[derive(Clone, Copy, Debug, Eq, PartialEq, thiserror::Error)]
#[error("node {node_id} is not allocated")]
pub struct NodeAccessError {
pub node_id: NodeId,
}
/// Errors surfaced from the tree-core runtime API when a caller violates a documented
/// contract (freeing an unallocated node, allocating under a freed parent).
#[allow(clippy::enum_variant_names)]
#[derive(Debug, thiserror::Error)]
pub enum TreeCoreRuntimeError {
/// A public NodeId no longer names a live arena node.
#[error("node {node_id} is not allocated")]
NodeNotAllocated { node_id: NodeId },
#[error(transparent)]
NodeAccess(#[from] NodeAccessError),
/// `begin_insert`/`insert` called while a resumable insert is suspended.
#[error("concurrent insert walks")]
ConcurrentInsertWalk,
@@ -736,6 +743,10 @@ pub enum TreeCoreRuntimeError {
/// A host insert below a non-root anchor must remain in that anchor's namespace.
#[error("insert_host namespace does not match non-root anchor {node_id}")]
InsertHostNamespaceMismatch { node_id: NodeId },
/// An inspection-only invariant check failed without mutating the tree.
#[cfg(any(test, feature = "inspection"))]
#[error("{0}")]
InspectionAssertion(String),
}
// Unigram and bigram child keys.
@@ -1058,18 +1069,12 @@ impl<K: ChildKeyType> NodeArena<K> {
self.root = self.alloc_root();
}
/// The live slot for an external handle; panics on a freed or unknown id.
#[track_caller]
pub fn resolve(&self, id: NodeId) -> NodeIdx_ {
*self
.id_map
/// The live slot for an external handle.
pub fn resolve(&self, id: NodeId) -> Result<NodeIdx_, NodeAccessError> {
self.id_map
.get(&id)
.unwrap_or_else(|| panic!("node {id} is not allocated"))
}
/// The live slot for an external handle, or None if freed/unknown.
pub fn try_resolve(&self, id: NodeId) -> Option<NodeIdx_> {
self.id_map.get(&id).copied()
.copied()
.ok_or(NodeAccessError { node_id: id })
}
/// Mint the next external handle for the slot and index it.
rust/sglang-radix-tree/src/python_bindings.rs
+256 -154
View File
@@ -12,7 +12,7 @@ use tch::{Device, Kind, Tensor};
use crate::components::{ComponentType, FULL, MAMBA, SWA};
use crate::node::ChildKeyType;
use crate::node::{KeyNamespaceRef, NodeId, TreeCoreRuntimeError};
use crate::node::{KeyNamespaceRef, NodeAccessError, NodeId, TreeCoreRuntimeError};
use crate::unified_tree_core::KvCacheEvent;
use crate::unified_tree_core::{
BufferBackupSnapshot, BufferBackupState, CacheAction, CacheInitParams, CacheTransferPhase,
@@ -65,17 +65,21 @@ fn parse_evict_layer(target: u8) -> PyResult<EvictLayer> {
}
}
fn node_access_error(error: NodeAccessError) -> PyErr {
PyKeyError::new_err(error.node_id)
}
/// Convert an expected tree-core contract failure without unwinding through PyO3.
fn tree_core_runtime_error(error: TreeCoreRuntimeError) -> PyErr {
match error {
TreeCoreRuntimeError::NodeNotAllocated { node_id } => PyKeyError::new_err(node_id),
TreeCoreRuntimeError::NodeAccess(error) => node_access_error(error),
error => PyRuntimeError::new_err(error.to_string()),
}
}
fn tree_core_assertion_error(error: TreeCoreRuntimeError) -> PyErr {
match error {
TreeCoreRuntimeError::NodeNotAllocated { node_id } => PyKeyError::new_err(node_id),
TreeCoreRuntimeError::NodeAccess(error) => node_access_error(error),
error => PyAssertionError::new_err(error.to_string()),
}
}
@@ -1082,10 +1086,12 @@ impl<K: ChildKeyType + Send + Sync> TreeCoreBinding<K> {
.into_iter()
.map(parse_component_type)
.collect::<PyResult<Vec<_>>>()?;
let result = py.allow_threads(|| {
self.core()
.inc_lock_ref_with_skip(node_id, &skip_lock_components)
});
let result = py
.allow_threads(|| {
self.core()
.inc_lock_ref_with_skip(node_id, &skip_lock_components)
})
.map_err(node_access_error)?;
Ok(IncLockRefResultBinding::from_result(result))
}
@@ -1098,7 +1104,8 @@ impl<K: ChildKeyType + Send + Sync> TreeCoreBinding<K> {
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));
py.allow_threads(|| self.core().dec_lock_ref(node_id, params.as_ref(), skip_swa))
.map_err(node_access_error)?;
Ok(())
}
@@ -1116,18 +1123,20 @@ impl<K: ChildKeyType + Send + Sync> TreeCoreBinding<K> {
.into_iter()
.map(|(ct, node_ids)| Ok((parse_component_type(ct)?, node_ids)))
.collect::<PyResult<HashMap<_, _>>>()?;
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(
node_id,
swa_uuid_for_lock,
Some(&skip_lock_node_ids),
&mut device_frees,
&mut host_frees,
);
(device_frees, host_frees)
});
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(
node_id,
swa_uuid_for_lock,
Some(&skip_lock_node_ids),
&mut device_frees,
&mut host_frees,
)?;
Ok((device_frees, host_frees))
})
.map_err(node_access_error)?;
Ok((frees_to_py(py, device_frees)?, frees_to_py(py, host_frees)?))
}
@@ -1157,7 +1166,8 @@ impl<K: ChildKeyType + Send + Sync> TreeCoreBinding<K> {
py.allow_threads(|| {
self.core()
.set_component_device_value(node_id, component_type, value)
});
})
.map_err(node_access_error)?;
Ok(())
}
@@ -1169,11 +1179,13 @@ impl<K: ChildKeyType + Send + Sync> TreeCoreBinding<K> {
component_type: u8,
) -> PyResult<Option<PyTensor>> {
let component_type = parse_component_type(component_type)?;
let value = py.allow_threads(|| {
self.core()
.get_component_device_value(node_id, component_type)
.map(|tensor| tensor.shallow_clone())
});
let value = py
.allow_threads(|| {
self.core()
.get_component_device_value(node_id, component_type)
.map(|value| value.map(|tensor| tensor.shallow_clone()))
})
.map_err(node_access_error)?;
Ok(value.map(PyTensor))
}
@@ -1221,11 +1233,13 @@ impl<K: ChildKeyType + Send + Sync> TreeCoreBinding<K> {
py: Python<'_>,
node_id: NodeId,
) -> PyResult<EvictDeviceLeafResultBinding> {
let (backup, result) = py.allow_threads(move || {
let mut core = self.core();
let is_write_back = core.is_write_back;
core.evict_device_leaf(node_id, is_write_back)
});
let (backup, result) = py
.allow_threads(move || {
let mut core = self.core();
let is_write_back = core.is_write_back;
core.evict_device_leaf(node_id, is_write_back)
})
.map_err(node_access_error)?;
Ok(EvictDeviceLeafResultBinding {
backup_kv: backup
.map(|backup| cache_action_to_py(py, CacheAction::BackupKV(backup)))
@@ -1264,11 +1278,14 @@ impl<K: ChildKeyType + Send + Sync> TreeCoreBinding<K> {
py: Python<'_>,
from_node_id: NodeId,
until_node_id: NodeId,
) -> PyTensor {
PyTensor(py.allow_threads(|| {
self.core()
.collect_full_device_indices(from_node_id, until_node_id)
}))
) -> PyResult<PyTensor> {
let value = py
.allow_threads(|| {
self.core()
.collect_full_device_indices(from_node_id, until_node_id)
})
.map_err(node_access_error)?;
Ok(PyTensor(value))
}
/// Every FULL device value in the tree, concatenated.
@@ -1337,8 +1354,9 @@ impl<K: ChildKeyType + Send + Sync> TreeCoreBinding<K> {
}
/// Whether the node's FULL device value has been evicted.
fn is_full_device_evicted(&self, py: Python<'_>, node_id: NodeId) -> bool {
fn is_full_device_evicted(&self, py: Python<'_>, node_id: NodeId) -> PyResult<bool> {
py.allow_threads(|| self.core().is_full_device_evicted(node_id))
.map_err(node_access_error)
}
/// Mark the host tier (HiCache) as wired.
@@ -1382,7 +1400,7 @@ impl<K: ChildKeyType + Send + Sync> TreeCoreBinding<K> {
}
let result = py
.allow_threads(move || {
self.core().try_insert_host_in_namespace(
self.core().insert_host_in_namespace(
node_id,
KeyNamespaceRef::new(extra_key.as_deref(), cache_salt.as_deref()),
key,
@@ -1400,8 +1418,9 @@ impl<K: ChildKeyType + Send + Sync> TreeCoreBinding<K> {
py: Python<'_>,
node_id: NodeId,
) -> PyResult<(PyTensor, Py<PyDict>)> {
let (device_value, comp_xfers) =
py.allow_threads(|| self.core().build_backup_spec(node_id));
let (device_value, comp_xfers) = py
.allow_threads(|| self.core().build_backup_spec(node_id))
.map_err(node_access_error)?;
Ok((PyTensor(device_value), comp_xfers_to_py(py, comp_xfers)?))
}
@@ -1412,10 +1431,12 @@ impl<K: ChildKeyType + Send + Sync> TreeCoreBinding<K> {
node_id: NodeId,
pass_prefix_keys: bool,
) -> PyResult<Option<StorageBackupSpecBinding>> {
let spec = py.allow_threads(|| {
self.core()
.build_storage_backup_spec(node_id, pass_prefix_keys)
});
let spec = py
.allow_threads(|| {
self.core()
.build_storage_backup_spec(node_id, pass_prefix_keys)
})
.map_err(node_access_error)?;
let Some(spec) = spec else {
return Ok(None);
};
@@ -1450,7 +1471,7 @@ impl<K: ChildKeyType + Send + Sync> TreeCoreBinding<K> {
let host_indices = host_indices.map(|t| t.0);
let transfers = py
.allow_threads(|| {
self.core().try_build_hicache_transfers(
self.core().build_hicache_transfers(
component_type,
node_id,
phase,
@@ -1477,37 +1498,37 @@ impl<K: ChildKeyType + Send + Sync> TreeCoreBinding<K> {
py: Python<'_>,
node_id: NodeId,
) -> PyResult<(Option<String>, Option<String>)> {
py.allow_threads(|| self.core().try_prefetch_anchor_info(node_id))
.map_err(tree_core_runtime_error)
py.allow_threads(|| self.core().prefetch_anchor_info(node_id))
.map_err(node_access_error)
}
/// Whether the node's Full KV is present on host.
fn node_backuped(&self, py: Python<'_>, node_id: NodeId) -> PyResult<bool> {
py.allow_threads(|| self.core().try_node_backuped(node_id))
.map_err(tree_core_runtime_error)
py.allow_threads(|| self.core().node_backuped(node_id))
.map_err(node_access_error)
}
/// Whether the node is a (default or named) root.
fn is_root(&self, py: Python<'_>, node_id: NodeId) -> PyResult<bool> {
py.allow_threads(|| self.core().try_is_root(node_id))
.map_err(tree_core_runtime_error)
py.allow_threads(|| self.core().is_root(node_id))
.map_err(node_access_error)
}
/// The node's last page hash, or None when it was never hashed.
fn get_last_hash_value(&self, py: Python<'_>, node_id: NodeId) -> PyResult<Option<String>> {
py.allow_threads(|| self.core().try_get_last_hash_value(node_id))
.map_err(tree_core_runtime_error)
py.allow_threads(|| self.core().get_last_hash_value(node_id))
.map_err(node_access_error)
}
/// The hash chain of the node's ancestors, in root-to-parent order.
fn get_prefix_hash_values(&self, py: Python<'_>, node_id: NodeId) -> PyResult<Vec<String>> {
py.allow_threads(|| self.core().try_get_prefix_hash_values(node_id))
.map_err(tree_core_runtime_error)
py.allow_threads(|| self.core().get_prefix_hash_values(node_id))
.map_err(node_access_error)
}
fn get_hash_values(&self, py: Python<'_>, node_id: NodeId) -> PyResult<Vec<String>> {
py.allow_threads(|| self.core().try_get_hash_values(node_id))
.map_err(tree_core_runtime_error)
py.allow_threads(|| self.core().get_hash_values(node_id))
.map_err(node_access_error)
}
fn snapshot_buffer_backup(
@@ -1546,8 +1567,8 @@ impl<K: ChildKeyType + Send + Sync> TreeCoreBinding<K> {
}
fn dfs_weight_order(&self, py: Python<'_>, node_ids: Vec<NodeId>) -> PyResult<Vec<usize>> {
py.allow_threads(|| self.core().try_dfs_weight_order(&node_ids))
.map_err(tree_core_runtime_error)
py.allow_threads(|| self.core().dfs_weight_order(&node_ids))
.map_err(node_access_error)
}
/// Commit each component's HiCache transfers; returns the new cache actions.
@@ -1582,22 +1603,24 @@ impl<K: ChildKeyType + Send + Sync> TreeCoreBinding<K> {
})
})
.transpose()?;
let (cache_actions, mamba_exist) = py.allow_threads(move || {
let mut cache_actions = Vec::new();
let mut insert_result = insert_result;
self.core().commit_hicache_transfers(
node_id,
phase,
comp_xfers,
&mut cache_actions,
insert_result.as_mut(),
pool_storage_result.as_ref(),
);
(
cache_actions,
insert_result.map(|result| result.mamba_exist),
)
});
let (cache_actions, mamba_exist) = py
.allow_threads(move || {
let mut cache_actions = Vec::new();
let mut insert_result = insert_result;
self.core().commit_hicache_transfers(
node_id,
phase,
comp_xfers,
&mut cache_actions,
insert_result.as_mut(),
pool_storage_result.as_ref(),
)?;
Ok((
cache_actions,
insert_result.map(|result| result.mamba_exist),
))
})
.map_err(node_access_error)?;
Ok((cache_actions_to_py(py, cache_actions)?, mamba_exist))
}
@@ -1611,7 +1634,8 @@ impl<K: ChildKeyType + Send + Sync> TreeCoreBinding<K> {
) -> PyResult<()> {
let comp_xfers = comp_xfers_from_args(comp_xfers)?;
let host_indices = host_indices.0;
py.allow_threads(move || self.core().commit_backup(node_id, host_indices, comp_xfers));
py.allow_threads(move || self.core().commit_backup(node_id, host_indices, comp_xfers))
.map_err(node_access_error)?;
Ok(())
}
@@ -1626,7 +1650,7 @@ impl<K: ChildKeyType + Send + Sync> TreeCoreBinding<K> {
mamba_pool_idx: mamba_pool_idx.map(|t| t.0),
};
let (kv_xfer, comp_xfers) = py
.allow_threads(move || self.core().try_build_load_back_spec(node_id, Some(&req)))
.allow_threads(move || self.core().build_load_back_spec(node_id, Some(&req)))
.map_err(tree_core_assertion_error)?;
Ok((
transfer_to_py(py, kv_xfer)?,
@@ -1646,17 +1670,19 @@ impl<K: ChildKeyType + Send + Sync> TreeCoreBinding<K> {
let kv_xfer = transfer_from_args(kv_xfer)?;
let comp_xfers = comp_xfers_from_args(comp_xfers)?;
let device_indices = device_indices.0;
let actions = py.allow_threads(move || {
self.core()
.commit_load_back(node_id, device_indices, kv_xfer, comp_xfers)
});
let actions = py
.allow_threads(move || {
self.core()
.commit_load_back(node_id, device_indices, kv_xfer, comp_xfers)
})
.map_err(node_access_error)?;
cache_actions_to_py(py, actions)
}
/// Release a node's device KV once its host copy exists.
fn demote(&self, py: Python<'_>, node_id: NodeId) -> PyResult<DemoteResultBinding> {
let result = py
.allow_threads(move || self.core().try_demote(node_id))
.allow_threads(move || self.core().demote(node_id))
.map_err(tree_core_assertion_error)?;
Ok(DemoteResultBinding {
tracker: tracker_to_py(result.tracker),
@@ -1688,7 +1714,9 @@ impl<K: ChildKeyType + Send + Sync> TreeCoreBinding<K> {
py: Python<'_>,
tail_node_id: NodeId,
) -> PyResult<HostEvictionResultBinding> {
let result = py.allow_threads(move || self.core().evict_excess_path_states(tail_node_id));
let result = py
.allow_threads(move || self.core().evict_excess_path_states(tail_node_id))
.map_err(node_access_error)?;
Ok(HostEvictionResultBinding {
tracker: tracker_to_py(result.tracker),
new_device_frees: frees_to_py(py, result.device_frees)?,
@@ -1697,9 +1725,15 @@ impl<K: ChildKeyType + Send + Sync> TreeCoreBinding<K> {
}
/// Bump the reference count on a node's host-side component locks.
fn inc_host_lock_ref(&self, py: Python<'_>, node_id: NodeId) -> IncLockRefResultBinding {
let result = py.allow_threads(|| self.core().inc_host_lock_ref(node_id));
IncLockRefResultBinding {
fn inc_host_lock_ref(
&self,
py: Python<'_>,
node_id: NodeId,
) -> PyResult<IncLockRefResultBinding> {
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,
@@ -1708,7 +1742,7 @@ impl<K: ChildKeyType + Send + Sync> TreeCoreBinding<K> {
.into_iter()
.map(|(ct, node_ids)| (component_type_to_u8(ct), node_ids))
.collect(),
}
})
}
/// Decrease the reference count on a node's host-side component locks.
@@ -1719,7 +1753,8 @@ impl<K: ChildKeyType + Send + Sync> TreeCoreBinding<K> {
params: Option<&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()));
py.allow_threads(|| self.core().dec_host_lock_ref(node_id, params.as_ref()))
.map_err(node_access_error)?;
Ok(())
}
@@ -1811,7 +1846,9 @@ impl<K: ChildKeyType + Send + Sync> TreeCoreBinding<K> {
py: Python<'_>,
node_id: NodeId,
) -> PyResult<DropSubtreeResultBinding> {
let (dropped, result) = py.allow_threads(move || self.core().drop_subtree_no_host(node_id));
let (dropped, result) = py
.allow_threads(move || self.core().drop_subtree_no_host(node_id))
.map_err(node_access_error)?;
Ok(DropSubtreeResultBinding {
dropped,
tracker: tracker_to_py(result.tracker),
@@ -1826,18 +1863,26 @@ impl<K: ChildKeyType + Send + Sync> TreeCoreBinding<K> {
py: Python<'_>,
node_ids: Vec<NodeId>,
ack_id: NodeId,
) -> Vec<NodeId> {
py.allow_threads(|| self.core().mark_write_through_pending(node_ids, ack_id))
) -> PyResult<Vec<NodeId>> {
py.allow_threads(move || self.core().mark_write_through_pending(node_ids, ack_id))
.map_err(node_access_error)
}
/// Clear the write-through-pending mark on the acked nodes.
fn finish_write_through(&self, py: Python<'_>, node_ids: Vec<NodeId>, ack_id: NodeId) {
py.allow_threads(|| self.core().finish_write_through(node_ids, ack_id));
fn finish_write_through(
&self,
py: Python<'_>,
node_ids: Vec<NodeId>,
ack_id: NodeId,
) -> PyResult<()> {
py.allow_threads(move || self.core().finish_write_through(node_ids, ack_id))
.map_err(node_access_error)
}
/// Clear the in-flight H->D marks on the anchor's root path at ack time.
fn finish_load_back(&self, py: Python<'_>, anchor_node_id: NodeId) {
py.allow_threads(|| self.core().finish_load_back(anchor_node_id));
fn finish_load_back(&self, py: Python<'_>, anchor_node_id: NodeId) -> PyResult<()> {
py.allow_threads(|| self.core().finish_load_back(anchor_node_id))
.map_err(node_access_error)
}
/// Order-sensitive digest of reclaimed coexisting host values.
@@ -1853,7 +1898,8 @@ impl<K: ChildKeyType + Send + Sync> TreeCoreBinding<K> {
component_type: u8,
) -> PyResult<bool> {
let ct = parse_component_type(component_type)?;
Ok(py.allow_threads(|| self.core().component_has_host_value_only(node_id, ct)))
py.allow_threads(|| self.core().component_has_host_value_only(node_id, ct))
.map_err(node_access_error)
}
}
@@ -1863,24 +1909,33 @@ impl<K: ChildKeyType + Send + Sync> TreeCoreBinding<K> {
py.allow_threads(|| self.core().inspect_contains_node(node_id))
}
fn inspect_get_parent_node_id(&self, py: Python<'_>, node_id: NodeId) -> Option<NodeId> {
fn inspect_get_parent_node_id(
&self,
py: Python<'_>,
node_id: NodeId,
) -> PyResult<Option<NodeId>> {
py.allow_threads(|| self.core().inspect_get_parent_node_id(node_id))
.map_err(node_access_error)
}
fn inspect_get_child_node_ids(&self, py: Python<'_>, node_id: NodeId) -> Vec<NodeId> {
fn inspect_get_child_node_ids(&self, py: Python<'_>, node_id: NodeId) -> PyResult<Vec<NodeId>> {
py.allow_threads(|| self.core().inspect_get_child_node_ids(node_id))
.map_err(node_access_error)
}
fn inspect_get_node_key_length(&self, py: Python<'_>, node_id: NodeId) -> usize {
fn inspect_get_node_key_length(&self, py: Python<'_>, node_id: NodeId) -> PyResult<usize> {
py.allow_threads(|| self.core().inspect_get_node_key_length(node_id))
.map_err(node_access_error)
}
fn inspect_get_node_token_ids(&self, py: Python<'_>, node_id: NodeId) -> Vec<i64> {
fn inspect_get_node_token_ids(&self, py: Python<'_>, node_id: NodeId) -> PyResult<Vec<i64>> {
py.allow_threads(|| self.core().inspect_get_node_token_ids(node_id))
.map_err(node_access_error)
}
fn inspect_is_node_key_bigram(&self, py: Python<'_>, node_id: NodeId) -> bool {
fn inspect_is_node_key_bigram(&self, py: Python<'_>, node_id: NodeId) -> PyResult<bool> {
py.allow_threads(|| self.core().inspect_is_node_key_bigram(node_id))
.map_err(node_access_error)
}
fn inspect_get_component_host_value(
@@ -1890,12 +1945,12 @@ impl<K: ChildKeyType + Send + Sync> TreeCoreBinding<K> {
component_type: u8,
) -> PyResult<Option<PyTensor>> {
let component_type = parse_component_type(component_type)?;
Ok(py
.allow_threads(|| {
self.core()
.inspect_get_component_host_value(node_id, component_type)
})
.map(PyTensor))
py.allow_threads(|| {
self.core()
.inspect_get_component_host_value(node_id, component_type)
})
.map(|value| value.map(PyTensor))
.map_err(node_access_error)
}
fn inspect_get_component_device_lock_ref(
@@ -1905,22 +1960,25 @@ impl<K: ChildKeyType + Send + Sync> TreeCoreBinding<K> {
component_type: u8,
) -> PyResult<u32> {
let component_type = parse_component_type(component_type)?;
Ok(py.allow_threads(|| {
py.allow_threads(|| {
self.core()
.inspect_get_component_device_lock_ref(node_id, component_type)
}))
})
.map_err(node_access_error)
}
fn inspect_get_node_hit_count(&self, py: Python<'_>, node_id: NodeId) -> i64 {
fn inspect_get_node_hit_count(&self, py: Python<'_>, node_id: NodeId) -> PyResult<i64> {
py.allow_threads(|| self.core().inspect_get_node_hit_count(node_id))
.map_err(node_access_error)
}
fn inspect_get_write_through_pending_id(
&self,
py: Python<'_>,
node_id: NodeId,
) -> Option<usize> {
) -> PyResult<Option<usize>> {
py.allow_threads(|| self.core().inspect_get_write_through_pending_id(node_id))
.map_err(node_access_error)
}
fn inspect_is_node_in_device_lru(
@@ -1930,10 +1988,11 @@ impl<K: ChildKeyType + Send + Sync> TreeCoreBinding<K> {
component_type: u8,
) -> PyResult<bool> {
let component_type = parse_component_type(component_type)?;
Ok(py.allow_threads(|| {
py.allow_threads(|| {
self.core()
.inspect_is_node_in_device_lru(node_id, component_type)
}))
})
.map_err(node_access_error)
}
fn inspect_is_node_in_host_lru(
@@ -1943,10 +2002,11 @@ impl<K: ChildKeyType + Send + Sync> TreeCoreBinding<K> {
component_type: u8,
) -> PyResult<bool> {
let component_type = parse_component_type(component_type)?;
Ok(py.allow_threads(|| {
py.allow_threads(|| {
self.core()
.inspect_is_node_in_host_lru(node_id, component_type)
}))
})
.map_err(node_access_error)
}
fn inspect_get_component_device_lru_node_ids(
@@ -1969,8 +2029,9 @@ impl<K: ChildKeyType + Send + Sync> TreeCoreBinding<K> {
py.allow_threads(|| self.core().inspect_is_host_evictable_leaf(node_id))
}
fn inspect_is_device_leaf(&self, py: Python<'_>, node_id: NodeId) -> bool {
fn inspect_is_device_leaf(&self, py: Python<'_>, node_id: NodeId) -> PyResult<bool> {
py.allow_threads(|| self.core().inspect_is_device_leaf(node_id))
.map_err(node_access_error)
}
fn inspect_get_all_node_ids(&self, py: Python<'_>) -> Vec<NodeId> {
@@ -1991,11 +2052,12 @@ impl<K: ChildKeyType + Send + Sync> TreeCoreBinding<K> {
py: Python<'_>,
node_id: NodeId,
hash_values: Option<Vec<String>>,
) {
) -> PyResult<()> {
py.allow_threads(move || {
self.core()
.inspect_set_node_hash_values(node_id, hash_values)
});
})
.map_err(node_access_error)
}
fn inspect_set_component_device_value_raw(
@@ -2010,8 +2072,8 @@ impl<K: ChildKeyType + Send + Sync> TreeCoreBinding<K> {
py.allow_threads(move || {
self.core()
.inspect_set_component_device_value_raw(node_id, component_type, value)
});
Ok(())
})
.map_err(node_access_error)
}
fn inspect_set_component_host_value_raw(
@@ -2026,8 +2088,8 @@ impl<K: ChildKeyType + Send + Sync> TreeCoreBinding<K> {
py.allow_threads(move || {
self.core()
.inspect_set_component_host_value_raw(node_id, component_type, value)
});
Ok(())
})
.map_err(node_access_error)
}
fn inspect_set_component_device_lock_ref(
@@ -2041,8 +2103,8 @@ impl<K: ChildKeyType + Send + Sync> TreeCoreBinding<K> {
py.allow_threads(|| {
self.core()
.inspect_set_component_device_lock_ref(node_id, component_type, lock_ref)
});
Ok(())
})
.map_err(node_access_error)
}
fn inspect_remove_node_from_device_lru(
@@ -2055,8 +2117,8 @@ impl<K: ChildKeyType + Send + Sync> TreeCoreBinding<K> {
py.allow_threads(|| {
self.core()
.inspect_remove_node_from_device_lru(node_id, component_type)
});
Ok(())
})
.map_err(node_access_error)
}
fn inspect_insert_node_into_host_lru(
@@ -2069,8 +2131,8 @@ impl<K: ChildKeyType + Send + Sync> TreeCoreBinding<K> {
py.allow_threads(|| {
self.core()
.inspect_insert_node_into_host_lru(node_id, component_type)
});
Ok(())
})
.map_err(node_access_error)
}
fn inspect_set_component_evictable_size(
@@ -2101,8 +2163,9 @@ impl<K: ChildKeyType + Send + Sync> TreeCoreBinding<K> {
Ok(())
}
fn inspect_update_duplicate_tracking(&self, py: Python<'_>, node_id: NodeId) {
py.allow_threads(|| self.core().inspect_update_duplicate_tracking(node_id));
fn inspect_update_duplicate_tracking(&self, py: Python<'_>, node_id: NodeId) -> PyResult<()> {
py.allow_threads(|| self.core().inspect_update_duplicate_tracking(node_id))
.map_err(node_access_error)
}
fn inspect_advance_insert_walk_once(&self, py: Python<'_>) -> PyResult<()> {
@@ -2119,10 +2182,12 @@ impl<K: ChildKeyType + Send + Sync> TreeCoreBinding<K> {
) -> PyResult<HostEvictionResultBinding> {
let component_type = parse_component_type(component_type)?;
let target = parse_evict_layer(target)?;
let result = py.allow_threads(|| {
self.core()
.inspect_evict_component(node_id, component_type, target)
});
let result = py
.allow_threads(|| {
self.core()
.inspect_evict_component(node_id, component_type, target)
})
.map_err(node_access_error)?;
HostEvictionResultBinding::from_eviction_step(py, result)
}
@@ -2139,7 +2204,7 @@ impl<K: ChildKeyType + Send + Sync> TreeCoreBinding<K> {
self.core()
.inspect_validate_cascade_evict(node_id, component_type, target)
})
.map_err(PyAssertionError::new_err)
.map_err(tree_core_assertion_error)
}
fn inspect_cleanup_tombstone_ancestors(
@@ -2147,7 +2212,9 @@ impl<K: ChildKeyType + Send + Sync> TreeCoreBinding<K> {
py: Python<'_>,
node_id: NodeId,
) -> PyResult<HostEvictionResultBinding> {
let result = py.allow_threads(|| self.core().inspect_cleanup_tombstone_ancestors(node_id));
let result = py
.allow_threads(|| self.core().inspect_cleanup_tombstone_ancestors(node_id))
.map_err(node_access_error)?;
HostEvictionResultBinding::from_eviction_step(py, result)
}
@@ -2205,6 +2272,7 @@ impl<K: ChildKeyType + Send + Sync> TreeCoreBinding<K> {
best_value_len,
)
});
let result = result.map_err(node_access_error)?;
MatchResultBinding::from_match_result(py, result)
}
@@ -2213,11 +2281,12 @@ impl<K: ChildKeyType + Send + Sync> TreeCoreBinding<K> {
py: Python<'_>,
node_id: NodeId,
write_back: bool,
) -> Vec<NodeId> {
) -> PyResult<Vec<NodeId>> {
py.allow_threads(|| {
self.core()
.inspect_build_backup_node_ids(node_id, write_back)
})
.map_err(node_access_error)
}
}
@@ -2422,7 +2491,7 @@ macro_rules! tree_core_binding {
py: Python<'_>,
from_node_id: NodeId,
until_node_id: NodeId,
) -> PyTensor {
) -> PyResult<PyTensor> {
self.inner
.collect_full_device_indices(py, from_node_id, until_node_id)
}
@@ -2484,7 +2553,7 @@ macro_rules! tree_core_binding {
}
/// Whether the node's FULL device value has been evicted.
fn is_full_device_evicted(&self, py: Python<'_>, node_id: NodeId) -> bool {
fn is_full_device_evicted(&self, py: Python<'_>, node_id: NodeId) -> PyResult<bool> {
self.inner.is_full_device_evicted(py, node_id)
}
@@ -2739,7 +2808,11 @@ macro_rules! tree_core_binding {
}
/// Bump the reference count on a node's host-side component locks.
fn inc_host_lock_ref(&self, py: Python<'_>, node_id: NodeId) -> IncLockRefResultBinding {
fn inc_host_lock_ref(
&self,
py: Python<'_>,
node_id: NodeId,
) -> PyResult<IncLockRefResultBinding> {
self.inner.inc_host_lock_ref(py, node_id)
}
@@ -2810,17 +2883,22 @@ macro_rules! tree_core_binding {
py: Python<'_>,
node_ids: Vec<NodeId>,
ack_id: NodeId,
) -> Vec<NodeId> {
) -> PyResult<Vec<NodeId>> {
self.inner.mark_write_through_pending(py, node_ids, ack_id)
}
/// Clear the write-through-pending mark on the acked nodes.
fn finish_write_through(&self, py: Python<'_>, node_ids: Vec<NodeId>, ack_id: NodeId) {
fn finish_write_through(
&self,
py: Python<'_>,
node_ids: Vec<NodeId>,
ack_id: NodeId,
) -> PyResult<()> {
self.inner.finish_write_through(py, node_ids, ack_id)
}
/// Clear the in-flight H->D marks on the anchor's root path at ack time.
fn finish_load_back(&self, py: Python<'_>, anchor_node_id: NodeId) {
fn finish_load_back(&self, py: Python<'_>, anchor_node_id: NodeId) -> PyResult<()> {
self.inner.finish_load_back(py, anchor_node_id)
}
@@ -2853,7 +2931,7 @@ macro_rules! tree_core_binding {
&self,
py: Python<'_>,
node_id: NodeId,
) -> Option<NodeId> {
) -> PyResult<Option<NodeId>> {
self.inner.inspect_get_parent_node_id(py, node_id)
}
@@ -2862,22 +2940,34 @@ macro_rules! tree_core_binding {
&self,
py: Python<'_>,
node_id: NodeId,
) -> Vec<NodeId> {
) -> PyResult<Vec<NodeId>> {
self.inner.inspect_get_child_node_ids(py, node_id)
}
#[cfg(feature = "inspection")]
fn inspect_get_node_key_length(&self, py: Python<'_>, node_id: NodeId) -> usize {
fn inspect_get_node_key_length(
&self,
py: Python<'_>,
node_id: NodeId,
) -> PyResult<usize> {
self.inner.inspect_get_node_key_length(py, node_id)
}
#[cfg(feature = "inspection")]
fn inspect_get_node_token_ids(&self, py: Python<'_>, node_id: NodeId) -> Vec<i64> {
fn inspect_get_node_token_ids(
&self,
py: Python<'_>,
node_id: NodeId,
) -> PyResult<Vec<i64>> {
self.inner.inspect_get_node_token_ids(py, node_id)
}
#[cfg(feature = "inspection")]
fn inspect_is_node_key_bigram(&self, py: Python<'_>, node_id: NodeId) -> bool {
fn inspect_is_node_key_bigram(
&self,
py: Python<'_>,
node_id: NodeId,
) -> PyResult<bool> {
self.inner.inspect_is_node_key_bigram(py, node_id)
}
@@ -2904,7 +2994,11 @@ macro_rules! tree_core_binding {
}
#[cfg(feature = "inspection")]
fn inspect_get_node_hit_count(&self, py: Python<'_>, node_id: NodeId) -> i64 {
fn inspect_get_node_hit_count(
&self,
py: Python<'_>,
node_id: NodeId,
) -> PyResult<i64> {
self.inner.inspect_get_node_hit_count(py, node_id)
}
@@ -2913,7 +3007,7 @@ macro_rules! tree_core_binding {
&self,
py: Python<'_>,
node_id: NodeId,
) -> Option<usize> {
) -> PyResult<Option<usize>> {
self.inner
.inspect_get_write_through_pending_id(py, node_id)
}
@@ -2969,7 +3063,11 @@ macro_rules! tree_core_binding {
}
#[cfg(feature = "inspection")]
fn inspect_is_device_leaf(&self, py: Python<'_>, node_id: NodeId) -> bool {
fn inspect_is_device_leaf(
&self,
py: Python<'_>,
node_id: NodeId,
) -> PyResult<bool> {
self.inner.inspect_is_device_leaf(py, node_id)
}
@@ -2995,7 +3093,7 @@ macro_rules! tree_core_binding {
py: Python<'_>,
node_id: NodeId,
hash_values: Option<Vec<String>>,
) {
) -> PyResult<()> {
self.inner
.inspect_set_node_hash_values(py, node_id, hash_values)
}
@@ -3095,7 +3193,11 @@ macro_rules! tree_core_binding {
}
#[cfg(feature = "inspection")]
fn inspect_update_duplicate_tracking(&self, py: Python<'_>, node_id: NodeId) {
fn inspect_update_duplicate_tracking(
&self,
py: Python<'_>,
node_id: NodeId,
) -> PyResult<()> {
self.inner.inspect_update_duplicate_tracking(py, node_id)
}
@@ -3171,7 +3273,7 @@ macro_rules! tree_core_binding {
py: Python<'_>,
node_id: NodeId,
write_back: bool,
) -> Vec<NodeId> {
) -> PyResult<Vec<NodeId>> {
self.inner
.inspect_build_backup_node_ids(py, node_id, write_back)
}
rust/sglang-radix-tree/src/tests/components/base.rs
+2
View File
@@ -119,6 +119,8 @@ fn finalize_match_result_default_returns_result_unchanged() {
let out = DefaultComponentForTest.finalize_match_result_in_tree_core(
&tc,
result,
tc.arena.root(),
tc.arena.root(),
&MatchPrefixParams {
key: &Vec::new(),
namespace: Default::default(),
rust/sglang-radix-tree/src/tests/components/full.rs
+197 -65
View File
@@ -1,5 +1,6 @@
use super::*;
use crate::components::FULL;
use crate::node::NodeAccessError;
use crate::test_utils::accumulate_step;
use crate::unified_tree_core::CacheInitParams;
@@ -392,14 +393,16 @@ fn host_drive_reclaims_coexisting_host_values_while_sparing_the_device_leaf() {
let leaf_handle = tc
.match_prefix(&match_params(&vec![1, 2, 3]))
.best_match_node_id;
let leaf = tc.arena.resolve(leaf_handle);
let leaf = tc.arena.resolve(leaf_handle).expect("live test node");
let parent = tc.arena.node(leaf).parent();
tc.commit_backup(
tc.arena.node(parent).id,
Tensor::from_slice(&[20i64, 21]),
HashMap::new(),
);
tc.commit_backup(leaf_handle, Tensor::from_slice(&[22i64]), HashMap::new());
)
.expect("live test node");
tc.commit_backup(leaf_handle, Tensor::from_slice(&[22i64]), HashMap::new())
.expect("live test node");
assert!(tc.evictable_host_leaves.is_empty());
let (mut tr, mut df, mut hf) = (tracker(), frees(), frees());
@@ -424,8 +427,10 @@ fn host_drive_spares_coexisting_host_values_under_an_in_flight_transfer() {
let handle = tc
.match_prefix(&match_params(&vec![1, 2]))
.best_match_node_id;
tc.commit_backup(handle, Tensor::from_slice(&[20i64, 21]), HashMap::new());
tc.mark_write_through_pending(vec![handle], /* ack_id = */ handle);
tc.commit_backup(handle, Tensor::from_slice(&[20i64, 21]), HashMap::new())
.expect("live test node");
tc.mark_write_through_pending(vec![handle], /* ack_id = */ handle)
.expect("live test node");
let (mut tr, mut df, mut hf) = (tracker(), frees(), frees());
accumulate_step(
@@ -435,9 +440,14 @@ fn host_drive_spares_coexisting_host_values_under_an_in_flight_transfer() {
&mut hf,
);
assert_eq!(tr[&FULL], 0);
assert!(tc.arena.node(tc.arena.resolve(handle)).has_host_value(FULL));
assert!(
tc.arena
.node(tc.arena.resolve(handle).expect("live test node"))
.has_host_value(FULL)
);
tc.finish_write_through(vec![handle], handle);
tc.finish_write_through(vec![handle], handle)
.expect("live test node");
accumulate_step(
tc.drive_host_eviction(FULL, /* num_tokens = */ 2),
&mut tr,
@@ -445,7 +455,11 @@ fn host_drive_spares_coexisting_host_values_under_an_in_flight_transfer() {
&mut hf,
);
assert_eq!(tr[&FULL], 2);
assert!(!tc.arena.node(tc.arena.resolve(handle)).has_host_value(FULL));
assert!(
!tc.arena
.node(tc.arena.resolve(handle).expect("live test node"))
.has_host_value(FULL)
);
tc.sanity_check(&[], &[]);
}
@@ -610,7 +624,9 @@ fn lock_chain(tc: &mut UnifiedTreeCore<Vec<i64>>) -> (NodeIdx_, NodeIdx_) {
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);
let result = tc
.inc_lock_ref(tc.arena.node(n2).id)
.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);
@@ -625,8 +641,11 @@ 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);
let result = tc.inc_lock_ref(tc.arena.node(n2).id);
tc.inc_lock_ref(tc.arena.node(n2).id)
.expect("live test node");
let result = tc
.inc_lock_ref(tc.arena.node(n2).id)
.expect("live test node");
assert_eq!(result.delta, Some(0));
assert_eq!(tc.arena.device_lock_ref(n1, FULL), 2);
assert_eq!(tc.arena.device_lock_ref(n2, FULL), 2);
@@ -640,8 +659,11 @@ 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);
let result = tc.inc_lock_ref(tc.arena.node(n2).id);
tc.inc_lock_ref(tc.arena.node(n1).id)
.expect("live test node");
let result = tc
.inc_lock_ref(tc.arena.node(n2).id)
.expect("live test node");
assert_eq!(result.delta, Some(3));
assert_eq!(tc.arena.device_lock_ref(n1, FULL), 2);
assert_eq!(tc.arena.device_lock_ref(n2, FULL), 1);
@@ -686,7 +708,9 @@ fn inc_lock_ref_collects_the_evicted_bottom_segment() {
.set_device_value(n1, FULL, Tensor::from_slice(&[0i64, 1]));
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);
let result = tc
.inc_lock_ref(tc.arena.node(n3).id)
.expect("live test node");
assert_eq!(result.delta, Some(2));
assert_eq!(
result.skip_lock_node_ids[&FULL],
@@ -703,7 +727,9 @@ fn inc_lock_ref_collects_the_evicted_bottom_segment() {
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);
let result = tc
.inc_lock_ref(tc.arena.node(root).id)
.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.
@@ -712,7 +738,8 @@ fn lock_round_trips_on_a_root_anchor_are_noops() {
tc.arena.node(root).id,
/* params = */ None,
/* skip_swa = */ false,
);
)
.expect("live test node");
assert_eq!(tc.arena.device_lock_ref(root, FULL), 1);
}
@@ -732,7 +759,9 @@ fn lock_walks_stop_at_the_root_of_a_salted_chain() {
tc.arena
.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);
let result = tc
.inc_lock_ref(tc.arena.node(n1).id)
.expect("live test node");
assert_eq!(result.delta, Some(2));
assert_eq!(tc.arena.device_lock_ref(n1, FULL), 1);
// The root keeps its construction-time lock untouched.
@@ -742,7 +771,8 @@ fn lock_walks_stop_at_the_root_of_a_salted_chain() {
tc.arena.node(n1).id,
/* params = */ None,
/* skip_swa = */ false,
);
)
.expect("live test node");
assert_eq!(tc.arena.device_lock_ref(n1, FULL), 0);
assert_eq!(tc.arena.device_lock_ref(lora, FULL), 1);
}
@@ -767,7 +797,9 @@ fn lock_walks_treat_a_present_but_empty_value_as_device_on() {
ValueSlotIdx::device(FULL),
Tensor::from_slice(&empty),
);
let result = tc.inc_lock_ref(tc.arena.node(n1).id);
let result = tc
.inc_lock_ref(tc.arena.node(n1).id)
.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));
@@ -778,7 +810,8 @@ fn lock_walks_treat_a_present_but_empty_value_as_device_on() {
tc.arena.node(n1).id,
/* params = */ None,
/* skip_swa = */ false,
);
)
.expect("live test node");
assert_eq!(tc.arena.device_lock_ref(n1, FULL), 0);
let state = tc.component_state(FULL);
assert_eq!(state.evictable_size, 0);
@@ -789,12 +822,14 @@ 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)
.expect("live test node");
tc.dec_lock_ref(
tc.arena.node(n2).id,
/* params = */ None,
/* 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);
let state = tc.component_state(FULL);
@@ -839,7 +874,9 @@ fn dec_lock_ref_replays_the_skip_set() {
tc.arena
.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);
let result = tc
.inc_lock_ref(tc.arena.node(n3).id)
.expect("live test node");
let params = DecLockRefParams {
skip_lock_node_ids: result.skip_lock_node_ids,
..Default::default()
@@ -849,7 +886,8 @@ fn dec_lock_ref_replays_the_skip_set() {
tc.arena.node(n3).id,
Some(&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);
@@ -896,7 +934,9 @@ fn temp_lock_skips_the_evicted_anchor_and_mirrors_on_release() {
.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.
let temp_lock = tc.inc_lock_ref(tc.arena.node(anchor).id);
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);
@@ -907,7 +947,9 @@ fn temp_lock_skips_the_evicted_anchor_and_mirrors_on_release() {
// 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);
let second_lock = tc
.inc_lock_ref(tc.arena.node(anchor).id)
.expect("live test node");
assert_eq!(tc.arena.device_lock_ref(anchor, FULL), 1);
assert_eq!(tc.arena.device_lock_ref(y, FULL), 2);
assert_eq!(tc.arena.device_lock_ref(a, FULL), 2);
@@ -920,7 +962,8 @@ fn temp_lock_skips_the_evicted_anchor_and_mirrors_on_release() {
tc.arena.node(anchor).id,
Some(&temp_params),
/* skip_swa = */ false,
);
)
.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);
@@ -932,7 +975,8 @@ fn temp_lock_skips_the_evicted_anchor_and_mirrors_on_release() {
tc.arena.node(anchor).id,
Some(&second_params),
/* skip_swa = */ false,
);
)
.expect("live test node");
assert_eq!(tc.arena.device_lock_ref(anchor, FULL), 0);
assert_eq!(tc.arena.device_lock_ref(y, FULL), 0);
assert_eq!(tc.arena.device_lock_ref(a, FULL), 0);
@@ -965,24 +1009,28 @@ 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)
.expect("live test node");
tc.dec_lock_ref(
tc.arena.node(n2).id,
/* params = */ None,
/* skip_swa = */ false,
);
)
.expect("live test node");
}
#[test]
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)
.expect("live test node");
tc.dec_lock_ref(
tc.arena.node(n2).id,
/* params = */ None,
/* skip_swa = */ true,
);
)
.expect("live test node");
assert_eq!(tc.arena.device_lock_ref(n2, FULL), 0);
}
@@ -991,13 +1039,16 @@ 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);
tc.inc_lock_ref(tc.arena.node(n2).id)
.expect("live test node");
tc.inc_lock_ref(tc.arena.node(n2).id)
.expect("live test node");
tc.dec_lock_ref(
tc.arena.node(n2).id,
/* params = */ None,
/* skip_swa = */ false,
);
)
.expect("live test node");
let state = tc.component_state(FULL);
assert_eq!(state.evictable_size, 0);
assert_eq!(state.protected_size, 5);
@@ -1006,7 +1057,8 @@ fn nested_locks_release_pairwise() {
tc.arena.node(n2).id,
/* params = */ None,
/* skip_swa = */ false,
);
)
.expect("live test node");
let state = tc.component_state(FULL);
assert_eq!(state.evictable_size, 5);
assert_eq!(state.protected_size, 0);
@@ -1022,7 +1074,8 @@ fn dec_lock_ref_panics_on_an_unlocked_node() {
tc.arena.node(n2).id,
/* params = */ None,
/* skip_swa = */ false,
);
)
.expect("live test node");
}
#[test]
@@ -1052,7 +1105,8 @@ 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)
.expect("live test node");
}
#[test]
@@ -1071,7 +1125,8 @@ 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)
.expect("live test node");
}
#[test]
@@ -1087,7 +1142,8 @@ fn dec_lock_ref_panics_on_protected_underflow() {
tc.arena.node(n2).id,
/* params = */ None,
/* skip_swa = */ false,
);
)
.expect("live test node");
}
fn write_back_core() -> UnifiedTreeCore<Vec<i64>> {
@@ -1123,7 +1179,9 @@ fn inc_host_lock_ref_pins_the_backuped_anchor() {
let mut tc = core();
let node = host_lock_anchor(&mut tc);
tc.component_state_mut(FULL).evictable_size = 7;
let result = tc.inc_host_lock_ref(tc.arena.node(node).id);
let result = tc
.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);
@@ -1139,8 +1197,10 @@ fn inc_host_lock_ref_pins_the_backuped_anchor() {
fn inc_host_lock_ref_again_only_bumps_the_counter() {
let mut tc = core();
let node = host_lock_anchor(&mut tc);
tc.inc_host_lock_ref(tc.arena.node(node).id);
tc.inc_host_lock_ref(tc.arena.node(node).id);
tc.inc_host_lock_ref(tc.arena.node(node).id)
.expect("live test node");
tc.inc_host_lock_ref(tc.arena.node(node).id)
.expect("live test node");
assert_eq!(tc.arena.host_lock_ref(node, FULL), 2);
assert!(!tc.evictable_host_leaves.contains(node));
}
@@ -1172,7 +1232,8 @@ fn inc_host_lock_ref_pins_only_the_anchor_not_its_ancestors() {
.set_host_value(n1, FULL, Tensor::from_slice(&[0i64]));
tc.arena
.set_host_value(n2, FULL, Tensor::from_slice(&[0i64]));
tc.inc_host_lock_ref(tc.arena.node(n2).id);
tc.inc_host_lock_ref(tc.arena.node(n2).id)
.expect("live test node");
assert_eq!(tc.arena.host_lock_ref(n2, FULL), 1);
assert_eq!(tc.arena.host_lock_ref(n1, FULL), 0);
}
@@ -1181,7 +1242,8 @@ fn inc_host_lock_ref_pins_only_the_anchor_not_its_ancestors() {
fn inc_host_lock_ref_skips_an_anchor_without_a_host_value() {
let mut tc = core();
let (_n1, n2) = lock_chain(&mut tc);
tc.inc_host_lock_ref(tc.arena.node(n2).id);
tc.inc_host_lock_ref(tc.arena.node(n2).id)
.expect("live test node");
assert_eq!(tc.arena.host_lock_ref(n2, FULL), 0);
}
@@ -1189,10 +1251,13 @@ fn inc_host_lock_ref_skips_an_anchor_without_a_host_value() {
fn host_lock_round_trips_on_a_root_anchor_are_noops() {
let mut tc = core();
let root = tc.arena.root();
let result = tc.inc_host_lock_ref(tc.arena.node(root).id);
let result = tc
.inc_host_lock_ref(tc.arena.node(root).id)
.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, /* params = */ None)
.expect("live test node");
assert_eq!(tc.arena.host_lock_ref(root, FULL), 0);
}
@@ -1200,7 +1265,8 @@ fn host_lock_round_trips_on_a_root_anchor_are_noops() {
fn inc_host_lock_ref_under_write_back_pins_a_device_only_anchor() {
let mut tc = write_back_core();
let (_n1, n2) = lock_chain(&mut tc);
tc.inc_host_lock_ref(tc.arena.node(n2).id);
tc.inc_host_lock_ref(tc.arena.node(n2).id)
.expect("live test node");
assert_eq!(tc.arena.host_lock_ref(n2, FULL), 1);
// The write-back host lock is a pure counter: no size shifts.
let state = tc.component_state(FULL);
@@ -1213,8 +1279,10 @@ fn dec_host_lock_ref_unpins_and_restores_the_h_leaf_set() {
let mut tc = core();
let node = host_lock_anchor(&mut tc);
tc.component_state_mut(FULL).evictable_size = 7;
tc.inc_host_lock_ref(tc.arena.node(node).id);
tc.dec_host_lock_ref(tc.arena.node(node).id, /* params = */ None);
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)
.expect("live test node");
assert_eq!(tc.arena.host_lock_ref(node, FULL), 0);
assert!(tc.evictable_host_leaves.contains(node));
let state = tc.component_state(FULL);
@@ -1226,7 +1294,8 @@ 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, /* params = */ None)
.expect("live test node");
assert_eq!(tc.arena.host_lock_ref(node, FULL), 0);
}
@@ -1235,9 +1304,11 @@ fn dec_host_lock_ref_keeps_the_counter_when_the_host_value_is_gone() {
// A host-evicted anchor keeps its pin count under write-through.
let mut tc = core();
let node = host_lock_anchor(&mut tc);
tc.inc_host_lock_ref(tc.arena.node(node).id);
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, /* params = */ None)
.expect("live test node");
assert_eq!(tc.arena.host_lock_ref(node, FULL), 1);
}
@@ -1245,8 +1316,10 @@ fn dec_host_lock_ref_keeps_the_counter_when_the_host_value_is_gone() {
fn host_lock_round_trip_under_write_back_is_a_pure_counter() {
let mut tc = write_back_core();
let (_n1, n2) = lock_chain(&mut tc);
tc.inc_host_lock_ref(tc.arena.node(n2).id);
tc.dec_host_lock_ref(tc.arena.node(n2).id, /* params = */ None);
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)
.expect("live test node");
assert_eq!(tc.arena.host_lock_ref(n2, FULL), 0);
let state = tc.component_state(FULL);
assert_eq!(state.evictable_size, 5);
@@ -1270,7 +1343,8 @@ fn acquire_host_arm_updates_the_h_leaf_set_without_the_dispatcher() {
fn release_host_arm_updates_the_h_leaf_set_without_the_dispatcher() {
let mut tc = core();
let node = host_lock_anchor(&mut tc);
tc.inc_host_lock_ref(tc.arena.node(node).id);
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,
);
@@ -1281,12 +1355,16 @@ fn release_host_arm_updates_the_h_leaf_set_without_the_dispatcher() {
fn nested_host_locks_release_pairwise() {
let mut tc = core();
let node = host_lock_anchor(&mut tc);
tc.inc_host_lock_ref(tc.arena.node(node).id);
tc.inc_host_lock_ref(tc.arena.node(node).id);
tc.dec_host_lock_ref(tc.arena.node(node).id, /* params = */ None);
tc.inc_host_lock_ref(tc.arena.node(node).id)
.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)
.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, /* params = */ None)
.expect("live test node");
assert_eq!(tc.arena.host_lock_ref(node, FULL), 0);
assert!(tc.evictable_host_leaves.contains(node));
}
@@ -1578,9 +1656,19 @@ fn host_hit_chain() -> (UnifiedTreeCore<Vec<i64>>, NodeIdx_, NodeIdx_) {
}
fn finalize(tc: &UnifiedTreeCore<Vec<i64>>, result: MatchResult) -> MatchResult {
let last_device_node_idx = tc
.arena
.resolve(result.last_device_node_id)
.expect("live test device node");
let best_match_node_idx = tc
.arena
.resolve(result.best_match_node_id)
.expect("live test best-match node");
FullComponent.finalize_match_result_in_tree_core(
tc,
result,
last_device_node_idx,
best_match_node_idx,
&MatchPrefixParams {
key: &Vec::new(),
namespace: Default::default(),
@@ -1791,18 +1879,62 @@ fn match_validator_panics_on_missing_node() {
}
#[test]
#[should_panic(expected = "is not allocated")]
fn finalize_panics_on_missing_best_match_node() {
fn inspect_finalize_rejects_missing_match_nodes() {
let tc = core();
let root = tc.arena.root();
finalize(
&tc,
let root_id = tc.arena.node(root).id;
let params = MatchPrefixParams {
key: &Vec::new(),
namespace: Default::default(),
};
let result = tc.inspect_finalize_component_match_result(
FULL,
MatchResult {
last_device_node_id: tc.arena.node(root).id,
last_device_node_id: root_id,
last_host_node_id: root_id,
best_match_node_id: 999,
host_hit_length: 0,
..tc.empty_match_result()
},
&params,
&[],
0,
);
assert!(matches!(result, Err(NodeAccessError { node_id: 999 })));
let result = tc.inspect_finalize_component_match_result(
FULL,
MatchResult {
last_device_node_id: 998,
last_host_node_id: root_id,
best_match_node_id: root_id,
host_hit_length: 0,
..tc.empty_match_result()
},
&params,
&[],
0,
);
assert!(matches!(result, Err(NodeAccessError { node_id: 998 })));
let result = tc.inspect_finalize_component_match_result(
FULL,
MatchResult {
last_device_node_id: root_id,
last_host_node_id: 997,
best_match_node_id: root_id,
host_hit_length: 0,
..tc.empty_match_result()
},
&params,
&[],
0,
);
assert!(matches!(result, Err(NodeAccessError { node_id: 997 })));
assert!(
tc.inspect_finalize_component_match_result(FULL, tc.empty_match_result(), &params, &[], 0,)
.is_ok()
);
}
rust/sglang-radix-tree/src/tests/components/mamba.rs
+125 -51
View File
@@ -44,7 +44,8 @@ fn hybrid_lock_core() -> (UnifiedTreeCore<Vec<i64>>, NodeIdx_, NodeIdx_) {
for (node, full_slot, swa_slot, mamba_slot) in [(parent, 10, 20, 30), (leaf, 11, 21, 31)] {
tc.arena
.set_device_value(node, FULL, Tensor::from_slice(&[full_slot]));
tc.set_component_device_value(tc.arena.node(node).id, SWA, Tensor::from_slice(&[swa_slot]));
tc.set_component_device_value(tc.arena.node(node).id, SWA, Tensor::from_slice(&[swa_slot]))
.expect("live test node");
set_mamba_device(&mut tc, node, mamba_slot);
tc.update_evictable_leaf_sets_(node);
}
@@ -229,7 +230,8 @@ fn device_value_round_trips_through_the_component() {
let [a] = chain::<1>(&mut tc);
let mamba = mamba_component();
assert!(tc.arena.try_device_value(a, MAMBA).is_none());
tc.set_component_device_value(tc.arena.node(a).id, MAMBA, Tensor::from_slice(&[42i64]));
tc.set_component_device_value(tc.arena.node(a).id, MAMBA, Tensor::from_slice(&[42i64]))
.expect("live test node");
assert!(
tc.arena
.try_device_value(a, MAMBA)
@@ -314,7 +316,9 @@ fn skip_aware_lock_records_only_the_mamba_target() {
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]);
let result = tc
.inc_lock_ref_with_skip(leaf_handle, &[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));
@@ -333,7 +337,8 @@ fn skip_aware_lock_records_only_the_mamba_target() {
..Default::default()
}),
/* 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);
}
@@ -342,8 +347,10 @@ fn skip_aware_lock_records_only_the_mamba_target() {
fn swa_only_release_honors_a_skipped_mamba_target() {
let (mut tc, _parent, leaf) = hybrid_lock_core();
let leaf_handle = tc.arena.node(leaf).id;
let owner = tc.inc_lock_ref(leaf_handle);
let skipped = tc.inc_lock_ref_with_skip(leaf_handle, &[MAMBA]);
let owner = tc.inc_lock_ref(leaf_handle).expect("live test node");
let skipped = tc
.inc_lock_ref_with_skip(leaf_handle, &[MAMBA])
.expect("live test node");
assert_eq!(tc.arena.node(leaf).device_lock_ref(MAMBA), 1);
let mut device_frees = HashMap::new();
@@ -354,7 +361,8 @@ fn swa_only_release_honors_a_skipped_mamba_target() {
Some(&skipped.skip_lock_node_ids),
&mut device_frees,
&mut host_frees,
);
)
.expect("live test node");
assert!(device_frees.is_empty());
assert!(host_frees.is_empty());
@@ -370,7 +378,8 @@ fn swa_only_release_honors_a_skipped_mamba_target() {
leaf_handle,
Some(&skipped_params),
/* skip_swa = */ true,
);
)
.expect("live test node");
let owner_params = DecLockRefParams {
swa_uuid_for_lock: owner.swa_uuid_for_lock,
skip_lock_node_ids: owner.skip_lock_node_ids,
@@ -380,7 +389,8 @@ fn swa_only_release_honors_a_skipped_mamba_target() {
leaf_handle,
Some(&owner_params),
/* skip_swa = */ false,
);
)
.expect("live test node");
assert_eq!(tc.protected_size_(MAMBA), 0);
}
@@ -483,14 +493,14 @@ fn insert_attaches_the_donated_slot_to_the_new_leaf() {
.best_match_node_id;
assert!(
tc.arena
.node(tc.arena.resolve(leaf))
.node(tc.arena.resolve(leaf).expect("live test node"))
.try_device_value(MAMBA)
.unwrap()
.equal(&Tensor::from_slice(&[7i64]))
);
assert!(
tc.device_lru_list(MAMBA)
.in_list(Some(tc.arena.resolve(leaf)))
.in_list(Some(tc.arena.resolve(leaf).expect("live test node")))
);
assert_eq!(tc.evictable_size_(MAMBA), 1);
}
@@ -508,7 +518,7 @@ fn reinsert_keeps_the_existing_slot_and_flags_the_caller() {
// The original slot stays; the caller frees the unused donated one.
assert!(
tc.arena
.node(tc.arena.resolve(leaf))
.node(tc.arena.resolve(leaf).expect("live test node"))
.try_device_value(MAMBA)
.unwrap()
.equal(&Tensor::from_slice(&[7i64]))
@@ -522,7 +532,8 @@ fn reinsert_full_backed_target_schedules_mamba_only_backup() {
let key = vec![1, 2];
tc.insert(&insert_params_mamba(&key, &[10, 11], Some(7)));
let leaf = tc.match_prefix(&match_params(&key)).best_match_node_id;
tc.commit_backup(leaf, Tensor::from_slice(&[100i64, 101]), HashMap::new());
tc.commit_backup(leaf, Tensor::from_slice(&[100i64, 101]), HashMap::new())
.expect("live test node");
let result = tc.insert(&insert_params_mamba(&key, &[20, 21], Some(8)));
let backups = result
@@ -536,7 +547,7 @@ fn reinsert_full_backed_target_schedules_mamba_only_backup() {
assert_eq!(backups.len(), 1);
assert_eq!(backups[0].node_ids, vec![leaf]);
let (full_device_indices, comp_xfers) = tc.build_backup_spec(leaf);
let (full_device_indices, comp_xfers) = tc.build_backup_spec(leaf).expect("live test node");
assert_eq!(full_device_indices.numel(), 0);
let mamba_xfers = &comp_xfers[&MAMBA];
assert_eq!(mamba_xfers.len(), 1);
@@ -548,7 +559,8 @@ fn reinsert_full_backed_target_schedules_mamba_only_backup() {
.equal(&Tensor::from_slice(&[7i64]))
);
tc.mark_write_through_pending(vec![leaf], /* ack_id = */ leaf);
tc.mark_write_through_pending(vec![leaf], /* ack_id = */ leaf)
.expect("live test node");
let pending = tc.insert(&insert_params_mamba(&key, &[30, 31], Some(9)));
assert!(
!pending
@@ -779,11 +791,25 @@ fn device_walk_advances_one_allocator_mutation_per_call() {
// The internal node is a complete step so its free can be reused before
// the walk hands out another victim.
assert_eq!(first, None);
assert!(!tc.arena.node(tc.arena.resolve(a)).has_device_value(MAMBA));
assert!(tc.arena.node(tc.arena.resolve(b)).has_device_value(MAMBA));
assert!(tc.arena.has_device_value(tc.arena.resolve(a), FULL));
assert!(
!tc.arena
.node(tc.arena.resolve(a).expect("live test node"))
.has_device_value(MAMBA)
);
assert!(
tc.arena
.node(tc.arena.resolve(b).expect("live test node"))
.has_device_value(MAMBA)
);
assert!(
tc.arena
.has_device_value(tc.arena.resolve(a).expect("live test node"), FULL)
);
assert_eq!(tracker[&MAMBA], 1);
assert!(!tc.device_lru_list(MAMBA).in_list(Some(tc.arena.resolve(a))));
assert!(
!tc.device_lru_list(MAMBA)
.in_list(Some(tc.arena.resolve(a).expect("live test node")))
);
let (second, step) = tc.evict_device_next_node(MAMBA, &tracker);
accumulate_step(step, &mut tracker, &mut device_frees, &mut host_frees);
@@ -801,7 +827,7 @@ fn device_walk_skips_locked_nodes() {
let b = tc
.match_prefix(&match_params(&vec![1, 2]))
.best_match_node_id;
let a_idx = tc.arena.resolve(a);
let a_idx = tc.arena.resolve(a).expect("live test node");
mamba_component().acquire_component_lock(
&mut tc,
a_idx,
@@ -816,7 +842,11 @@ fn device_walk_skips_locked_nodes() {
accumulate_step(step, &mut tracker, &mut device_frees, &mut host_frees);
// The locked internal node stays; the cursor starts on the leaf.
assert_eq!(next, Some(b));
assert!(tc.arena.node(tc.arena.resolve(a)).has_device_value(MAMBA));
assert!(
tc.arena
.node(tc.arena.resolve(a).expect("live test node"))
.has_device_value(MAMBA)
);
assert_eq!(tracker[&MAMBA], 0);
tc.evict_device_end(MAMBA);
}
@@ -877,12 +907,16 @@ fn host_eviction_takes_a_host_leaf_atomically() {
Tensor::from_slice(&[100i64]),
vec!["h0".to_string()],
)
.expect("live test node")
.inserted_host_node
.unwrap();
let leaf_idx = tc.arena.resolve(leaf);
let leaf_idx = tc.arena.resolve(leaf).expect("live test node");
set_mamba_host(&mut tc, leaf_idx, 8);
tc.host_lru_list_mut(MAMBA).insert_mru(leaf_idx);
assert!(tc.evictable_host_leaves.contains(tc.arena.resolve(leaf)));
assert!(
tc.evictable_host_leaves
.contains(tc.arena.resolve(leaf).expect("live test node"))
);
let mut tracker = HashMap::from([(MAMBA, 0)]);
let mut device_frees = HashMap::new();
let mut host_frees = HashMap::new();
@@ -897,7 +931,7 @@ fn host_eviction_takes_a_host_leaf_atomically() {
assert_eq!(tracker[&MAMBA], 1);
assert!(host_frees[&MAMBA][0].equal(&Tensor::from_slice(&[8i64])));
assert!(host_frees[&FULL][0].equal(&Tensor::from_slice(&[100i64])));
assert!(tc.arena.try_resolve(leaf).is_none());
assert!(tc.arena.resolve(leaf).is_err());
assert!(!tc.host_lru_list(MAMBA).in_list(Some(leaf_idx)));
}
@@ -996,6 +1030,7 @@ fn backup_host_build_carries_the_device_slot() {
0,
None,
)
.expect("live test node")
.unwrap();
assert_eq!(transfers.len(), 1);
assert_eq!(transfers[0].name, PoolName::Mamba);
@@ -1026,6 +1061,7 @@ fn backup_host_build_carries_the_device_slot() {
0,
None,
)
.expect("live test node")
.is_none()
);
}
@@ -1045,6 +1081,7 @@ fn load_back_build_restores_the_host_only_node() {
0,
None,
)
.expect("live test node")
.unwrap();
assert_eq!(transfers.len(), 1);
assert!(
@@ -1073,6 +1110,7 @@ fn load_back_build_skips_device_backed_and_bare_nodes() {
0,
None,
)
.expect("live test node")
.is_none()
);
}
@@ -1127,7 +1165,8 @@ fn backup_host_commit_stores_the_host_slot_once() {
&mut cache_actions,
None,
None,
);
)
.expect("live test node");
assert!(
tc.arena
.node(a)
@@ -1150,7 +1189,8 @@ fn backup_host_commit_stores_the_host_slot_once() {
&mut cache_actions,
None,
None,
);
)
.expect("live test node");
assert!(
tc.arena
.node(a)
@@ -1183,7 +1223,8 @@ fn load_back_commit_moves_the_node_onto_the_device_tier() {
&mut cache_actions,
None,
None,
);
)
.expect("live test node");
assert!(
tc.arena
.node(a)
@@ -1204,15 +1245,17 @@ fn mamba_device_eviction_skips_a_load_back_pinned_node() {
let [n] = chain::<1>(&mut tc);
set_full_host(&mut tc, n, 10);
set_mamba_host(&mut tc, n, 20);
let (kv_xfer, mut comp_xfers) =
tc.build_load_back_spec(tc.arena.node(n).id, /* req = */ None);
let (kv_xfer, mut comp_xfers) = tc
.build_load_back_spec(tc.arena.node(n).id, /* req = */ None)
.expect("live test node");
comp_xfers.get_mut(&MAMBA).unwrap()[0].device_indices = Some(Tensor::from_slice(&[40i64]));
tc.commit_load_back(
tc.arena.node(n).id,
Tensor::from_slice(&[30i64]),
kv_xfer,
comp_xfers,
);
)
.expect("live test node");
tc.evict_device_start(MAMBA, /* request_cnt = */ 1);
let (next, _) = tc.evict_device_next_node(MAMBA, &HashMap::new());
@@ -1220,7 +1263,8 @@ fn mamba_device_eviction_skips_a_load_back_pinned_node() {
tc.evict_device_end(MAMBA);
assert!(tc.arena.has_device_value(n, MAMBA));
tc.finish_load_back(tc.arena.node(n).id);
tc.finish_load_back(tc.arena.node(n).id)
.expect("live test node");
tc.evict_device_start(MAMBA, /* request_cnt = */ 1);
let (next, _) = tc.evict_device_next_node(MAMBA, &HashMap::new());
assert_eq!(next, Some(tc.arena.node(n).id));
@@ -1236,21 +1280,25 @@ fn mamba_host_eviction_skips_a_load_back_pinned_node() {
set_full_host(&mut tc, b, 11);
set_mamba_host(&mut tc, a, 20);
tc.host_lru_list_mut(MAMBA).insert_mru(a);
let (kv_xfer, comp_xfers) = tc.build_load_back_spec(tc.arena.node(b).id, /* req = */ None);
let (kv_xfer, comp_xfers) = tc
.build_load_back_spec(tc.arena.node(b).id, /* req = */ None)
.expect("live test node");
assert!(comp_xfers.is_empty());
tc.commit_load_back(
tc.arena.node(b).id,
Tensor::from_slice(&[30i64, 31]),
kv_xfer,
comp_xfers,
);
)
.expect("live test node");
let result = tc.drive_host_eviction(MAMBA, /* num_tokens = */ 1);
assert_eq!(result.tracker[&MAMBA], 0);
assert!(result.host_frees.is_empty());
assert!(tc.arena.has_host_value(a, MAMBA));
tc.finish_load_back(tc.arena.node(b).id);
tc.finish_load_back(tc.arena.node(b).id)
.expect("live test node");
let result = tc.drive_host_eviction(MAMBA, /* num_tokens = */ 1);
assert_eq!(result.tracker[&MAMBA], 1);
assert_eq!(result.host_frees[&MAMBA].len(), 1);
@@ -1278,7 +1326,8 @@ fn backup_storage_commit_is_a_noop() {
&mut cache_actions,
None,
None,
);
)
.expect("live test node");
assert!(tc.arena.node(a).has_host_value(MAMBA));
assert!(cache_actions.is_empty());
}
@@ -1298,6 +1347,7 @@ fn backup_storage_build_keys_the_trailing_hash() {
0,
None,
)
.expect("live test node")
.is_none()
);
set_mamba_host(&mut tc, a, 8);
@@ -1312,6 +1362,7 @@ fn backup_storage_build_keys_the_trailing_hash() {
0,
None,
)
.expect("live test node")
.is_none()
);
tc.arena.node_mut(a).hash_value = Some(vec!["h0".to_string(), "h1".to_string()]);
@@ -1325,6 +1376,7 @@ fn backup_storage_build_keys_the_trailing_hash() {
0,
None,
)
.expect("live test node")
.unwrap();
assert_eq!(transfers.len(), 1);
assert_eq!(transfers[0].keys, Some(vec!["h1".to_string()]));
@@ -1351,6 +1403,7 @@ fn prefetch_build_wraps_the_host_buffer_with_a_placeholder_key() {
0,
None,
)
.expect("live test node")
.unwrap();
assert_eq!(transfers.len(), 1);
assert_eq!(transfers[0].keys, Some(vec!["__placeholder__".to_string()]));
@@ -1370,6 +1423,7 @@ fn prefetch_commit_attaches_the_loaded_slot_to_the_inserted_node() {
Tensor::from_slice(&[100i64]),
vec!["h0".to_string()],
)
.expect("live test node")
.inserted_host_node
.unwrap();
let mut insert_result = InsertResult {
@@ -1395,17 +1449,18 @@ fn prefetch_commit_attaches_the_loaded_slot_to_the_inserted_node() {
kv_hit_pages: 1,
extra_pool_hit_pages: HashMap::from([(PoolName::Mamba, 1)]),
}),
);
)
.expect("live test node");
assert!(
tc.arena
.node(tc.arena.resolve(target))
.node(tc.arena.resolve(target).expect("live test node"))
.try_host_value(MAMBA)
.unwrap()
.equal(&Tensor::from_slice(&[50i64]))
);
assert!(
tc.host_lru_list(MAMBA)
.in_list(Some(tc.arena.resolve(target)))
.in_list(Some(tc.arena.resolve(target).expect("live test node")))
);
assert!(!insert_result.mamba_exist);
assert!(cache_actions.is_empty());
@@ -1424,6 +1479,7 @@ fn prefetch_commit_frees_the_buffer_when_it_cannot_attach() {
Tensor::from_slice(&[100i64]),
vec!["h0".to_string()],
)
.expect("live test node")
.inserted_host_node
.unwrap();
// Not loaded: the buffer frees and the caller keeps its slot flag.
@@ -1450,10 +1506,11 @@ fn prefetch_commit_frees_the_buffer_when_it_cannot_attach() {
kv_hit_pages: 1,
extra_pool_hit_pages: HashMap::new(),
}),
);
)
.expect("live test node");
assert!(
!tc.arena
.node(tc.arena.resolve(target))
.node(tc.arena.resolve(target).expect("live test node"))
.has_host_value(MAMBA)
);
assert!(insert_result.mamba_exist);
@@ -1468,7 +1525,7 @@ fn prefetch_commit_frees_the_buffer_when_it_cannot_attach() {
assert!(host_indices[0].equal(&Tensor::from_slice(&[50i64])));
// An already-hosted target frees the buffer too.
let target_idx = tc.arena.resolve(target);
let target_idx = tc.arena.resolve(target).expect("live test node");
set_mamba_host(&mut tc, target_idx, 8);
let mut insert_result = InsertResult {
total_len: 1,
@@ -1493,12 +1550,13 @@ fn prefetch_commit_frees_the_buffer_when_it_cannot_attach() {
kv_hit_pages: 1,
extra_pool_hit_pages: HashMap::from([(PoolName::Mamba, 1)]),
}),
);
)
.expect("live test node");
assert!(insert_result.mamba_exist);
assert_eq!(cache_actions.len(), 1);
assert!(
tc.arena
.node(tc.arena.resolve(target))
.node(tc.arena.resolve(target).expect("live test node"))
.try_host_value(MAMBA)
.unwrap()
.equal(&Tensor::from_slice(&[8i64]))
@@ -1528,7 +1586,9 @@ fn evict_excess_path_states_removes_the_shallowest_states_beyond_the_cap() {
set_mamba_device(&mut tc, a, 7);
set_mamba_device(&mut tc, b, 8);
set_mamba_device(&mut tc, c, 9);
let mut result = tc.evict_excess_path_states(tc.arena.node(c).id);
let mut result = tc
.evict_excess_path_states(tc.arena.node(c).id)
.expect("live test node");
let freed = result
.device_frees
.remove(&MAMBA)
@@ -1560,7 +1620,9 @@ fn evict_excess_path_states_preserves_forks_locked_nodes_and_the_tail() {
tc.arena
.node_mut(b)
.set_lock_ref_(ValueSlotIdx::device(MAMBA), 1);
let result = tc.evict_excess_path_states(tc.arena.node(c).id);
let result = tc
.evict_excess_path_states(tc.arena.node(c).id)
.expect("live test node");
assert!(result.device_frees.is_empty());
assert!(result.host_frees.is_empty());
assert!(tc.arena.node(a).try_device_value(MAMBA).is_some());
@@ -1574,7 +1636,9 @@ fn evict_excess_path_states_without_a_cap_is_a_no_op() {
let [a, b] = chain::<2>(&mut tc);
set_mamba_device(&mut tc, a, 7);
set_mamba_device(&mut tc, b, 8);
let result = tc.evict_excess_path_states(tc.arena.node(b).id);
let result = tc
.evict_excess_path_states(tc.arena.node(b).id)
.expect("live test node");
assert!(result.device_frees.is_empty());
assert!(result.host_frees.is_empty());
assert!(tc.arena.node(a).try_device_value(MAMBA).is_some());
@@ -1616,7 +1680,8 @@ fn swa_evict_on_a_full_locked_leaf_sweeps_mamba_and_spares_full() {
let [a] = chain::<1>(&mut tc);
tc.arena
.set_device_value(a, FULL, Tensor::from_slice(&[10i64]));
tc.set_component_device_value(tc.arena.node(a).id, SWA, Tensor::from_slice(&[20i64]));
tc.set_component_device_value(tc.arena.node(a).id, SWA, Tensor::from_slice(&[20i64]))
.expect("live test node");
set_mamba_device(&mut tc, a, 7);
// The held Full lock keeps the leaf out of the D-leaf set.
tc.arena
@@ -1721,11 +1786,20 @@ fn branching_from_a_host_full_hit_is_reusable_after_insert() {
b,
Tensor::from_slice(&[100i64, 101, 102, 103]),
HashMap::new(),
);
tc.demote(b);
)
.expect("live test node");
tc.demote(b).expect("valid demote state");
// The demote's cascade swept b's mamba slot: b is Full-host-only, no mamba.
assert!(!tc.arena.node(tc.arena.resolve(b)).has_device_value(MAMBA));
assert!(!tc.arena.node(tc.arena.resolve(b)).has_host_value(MAMBA));
assert!(
!tc.arena
.node(tc.arena.resolve(b).expect("live test node"))
.has_device_value(MAMBA)
);
assert!(
!tc.arena
.node(tc.arena.resolve(b).expect("live test node"))
.has_host_value(MAMBA)
);
let result = tc.match_prefix(&match_params(&vec![1, 2, 3, 4, 5, 6, 7]));
assert_eq!(result.best_match_node_id, a);
assert_eq!(result.last_device_node_id, a);
rust/sglang-radix-tree/src/tests/components/swa.rs
+154 -86
View File
@@ -172,7 +172,8 @@ fn store_swa_device(tc: &mut UnifiedTreeCore<Vec<i64>>, node: NodeIdx_) {
tc.arena.node(node).id,
SWA,
Tensor::from_slice(&vec![0i64; len]),
);
)
.expect("live test node");
}
#[test]
@@ -393,6 +394,7 @@ fn finalize(
best: NodeIdx_,
prior_swa_host_hit: usize,
) -> MatchResult {
let last_device_node_idx = tc.arena.root();
swa.finalize_match_result_in_tree_core(
tc,
MatchResult {
@@ -400,6 +402,8 @@ fn finalize(
swa_host_hit_length: prior_swa_host_hit,
..tc.empty_match_result()
},
last_device_node_idx,
best,
&MatchPrefixParams {
key: &Vec::new(),
namespace: Default::default(),
@@ -704,7 +708,8 @@ fn insert_overlap_with_live_swa_frees_the_whole_duplicate() {
tc.arena.node(leaf).id,
SWA,
Tensor::from_slice(&[50i64, 51, 52]),
);
)
.expect("live test node");
let result = tc.insert(&insert_params_swa(&vec![1, 2, 3], &[20, 21, 22], 0, 0));
assert_eq!(result.prefix_len, 3);
let [CacheAction::FreeDeviceKV(freed)] = result.cache_actions.as_slice() else {
@@ -957,7 +962,8 @@ fn insert_overlap_boundary_at_the_node_start_recovers_the_whole_node() {
tc.arena.node(a).id,
SWA,
Tensor::from_slice(&[50i64, 51, 52]),
);
)
.expect("live test node");
// The boundary lands exactly on b's start: full recovery, no split.
let result = tc.insert(&insert_params_swa(
&vec![1, 2, 3, 4, 5],
@@ -1285,7 +1291,8 @@ fn reinsert_with_live_swa_skips_recovery() {
tc.arena.node(leaf).id,
SWA,
Tensor::from_slice(&[50i64, 51, 52]),
);
)
.expect("live test node");
evict_full(&mut tc, leaf, /* remaining_size = */ 0);
let result = tc.insert(&insert_params_swa(&vec![1, 2, 3], &[20, 21, 22], 0, 0));
// The SWA value is already live: no rebuild is emitted.
@@ -1353,7 +1360,8 @@ fn walk_split_redistributes_the_live_swa_value() {
tc.arena.node(node).id,
SWA,
Tensor::from_slice(&[50i64, 51, 52, 53]),
);
)
.expect("live test node");
let result = tc.insert(&insert_params_swa(&vec![1, 2, 9], &[20, 21, 29], 0, 0));
assert_eq!(result.prefix_len, 2);
let parent = child_of(&tc, root, &[1]);
@@ -1424,7 +1432,8 @@ fn redistribute_on_node_split_keeps_device_valued_sides_off_the_host_lru() {
tc.arena.node(node).id,
SWA,
Tensor::from_slice(&[50i64, 51]),
);
)
.expect("live test node");
tc.arena
.set_host_value(node, SWA, Tensor::from_slice(&[70i64, 71]));
tc.arena
@@ -1614,7 +1623,9 @@ fn inc_lock_ref_runs_full_and_swa_walks_together() {
store_swa_device(&mut tc, a);
store_swa_device(&mut tc, b);
store_swa_device(&mut tc, c);
let result = tc.inc_lock_ref(tc.arena.node(c).id);
let result = tc
.inc_lock_ref(tc.arena.node(c).id)
.expect("live test node");
// FULL sees a valueless path (skip segment only); SWA locks its window.
assert_eq!(result.delta, Some(0));
assert_eq!(result.skip_lock_node_ids[&FULL].len(), 3);
@@ -1633,7 +1644,9 @@ fn inc_host_lock_ref_runs_full_and_swa_host_arms_together() {
}
tc.arena
.set_host_value(c, FULL, Tensor::from_slice(&[0i64]));
let result = tc.inc_host_lock_ref(tc.arena.node(c).id);
let result = tc
.inc_host_lock_ref(tc.arena.node(c).id)
.expect("live test node");
// FULL pins only the anchor; SWA walks its host window up to b.
assert_eq!(tc.arena.host_lock_ref(c, FULL), 1);
assert_eq!(tc.arena.host_lock_ref(b, FULL), 0);
@@ -1647,7 +1660,8 @@ fn inc_host_lock_ref_runs_full_and_swa_host_arms_together() {
skip_lock_node_ids: result.skip_lock_node_ids,
..Default::default()
};
tc.dec_host_lock_ref(tc.arena.node(c).id, Some(&params));
tc.dec_host_lock_ref(tc.arena.node(c).id, Some(&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);
assert_eq!(tc.arena.host_lock_ref(b, SWA), 0);
@@ -1663,8 +1677,11 @@ fn dec_host_lock_ref_with_the_inner_uuid_leaves_an_outer_window_pinned() {
set_swa_host(&mut tc, node);
}
// Overlapping host windows: {c, b} stamps its uuid at b, {b, a} at a.
let inner = tc.inc_host_lock_ref(tc.arena.node(c).id);
tc.inc_host_lock_ref(tc.arena.node(b).id);
let inner = tc
.inc_host_lock_ref(tc.arena.node(c).id)
.expect("live test node");
tc.inc_host_lock_ref(tc.arena.node(b).id)
.expect("live test node");
assert!(inner.swa_uuid_for_host_lock.is_some());
assert_eq!(tc.arena.host_lock_ref(c, SWA), 1);
assert_eq!(tc.arena.host_lock_ref(b, SWA), 2);
@@ -1676,7 +1693,8 @@ fn dec_host_lock_ref_with_the_inner_uuid_leaves_an_outer_window_pinned() {
skip_lock_node_ids: inner.skip_lock_node_ids,
..Default::default()
};
tc.dec_host_lock_ref(tc.arena.node(c).id, Some(&params));
tc.dec_host_lock_ref(tc.arena.node(c).id, Some(&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);
assert_eq!(tc.arena.host_lock_ref(a, SWA), 1);
@@ -1888,7 +1906,8 @@ fn evict_component_device_frees_the_full_indices_and_tombstones_swa() {
tc.arena.node(node).id,
SWA,
Tensor::from_slice(&[50i64, 51]),
);
)
.expect("live test node");
let mut device_frees = HashMap::new();
let mut host_frees = HashMap::new();
let (freed, host_freed) = swa_component(4).evict_component(
@@ -1921,7 +1940,8 @@ fn evict_component_device_parks_a_remaining_host_value() {
tc.arena.node(node).id,
SWA,
Tensor::from_slice(&[50i64, 51]),
);
)
.expect("live test node");
set_swa_host(&mut tc, node);
let mut device_frees = HashMap::new();
let mut host_frees = HashMap::new();
@@ -2177,7 +2197,9 @@ fn inc_then_dec_lock_ref_roundtrips_with_dec_params() {
store_swa_device(&mut tc, a);
store_swa_device(&mut tc, b);
store_swa_device(&mut tc, c);
let result = tc.inc_lock_ref(tc.arena.node(c).id);
let result = tc
.inc_lock_ref(tc.arena.node(c).id)
.expect("live test node");
let params = DecLockRefParams {
swa_uuid_for_lock: result.swa_uuid_for_lock,
swa_uuid_for_host_lock: result.swa_uuid_for_host_lock,
@@ -2187,7 +2209,8 @@ fn inc_then_dec_lock_ref_roundtrips_with_dec_params() {
tc.arena.node(c).id,
Some(&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);
@@ -2206,7 +2229,9 @@ 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);
let result = tc
.inc_lock_ref(tc.arena.node(c).id)
.expect("live test node");
let mut device_frees = HashMap::new();
let mut host_frees = HashMap::new();
tc.dec_swa_lock_only(
@@ -2214,7 +2239,8 @@ fn dec_swa_lock_only_releases_swa_while_full_stays_locked() {
result.swa_uuid_for_lock,
&mut device_frees,
&mut host_frees,
);
)
.expect("live test node");
// SWA is early-released; the FULL locks on the path stay.
assert_eq!(tc.arena.device_lock_ref(c, SWA), 0);
assert_eq!(tc.arena.device_lock_ref(b, SWA), 0);
@@ -2250,7 +2276,8 @@ fn dec_swa_lock_only_evicts_a_fully_unlocked_device_leaf() {
result.swa_uuid_for_lock,
&mut device_frees,
&mut host_frees,
);
)
.expect("live test node");
// The fully unlocked leaf c is device-evicted on release; b keeps its
// SWA value because its child still holds FULL KV.
assert!(!tc.arena.has_device_value(c, SWA));
@@ -2272,7 +2299,8 @@ fn dec_swa_lock_only_is_a_noop_without_the_swa_component() {
None,
&mut device_frees,
&mut host_frees,
);
)
.expect("live test node");
assert!(device_frees.is_empty());
}
@@ -2410,7 +2438,8 @@ fn dec_swa_lock_only_releases_the_window_exactly_once() {
first.swa_uuid_for_lock,
&mut device_frees,
&mut host_frees,
);
)
.expect("live test node");
// The second window still holds the lock: refs drop to 1, sizes stay.
assert_eq!(tc.arena.device_lock_ref(c, SWA), 1);
assert_eq!(tc.arena.device_lock_ref(b, SWA), 1);
@@ -2421,7 +2450,8 @@ fn dec_swa_lock_only_releases_the_window_exactly_once() {
first.swa_uuid_for_lock,
&mut device_frees,
&mut host_frees,
);
)
.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);
@@ -2456,7 +2486,8 @@ fn dec_swa_lock_only_leaves_out_of_window_swa_locks_alone() {
result.swa_uuid_for_lock,
&mut device_frees,
&mut host_frees,
);
)
.expect("live test node");
// Only the SWA window is released; a's out-of-window lock survives.
assert_eq!(tc.arena.device_lock_ref(a, SWA), 1);
assert_eq!(tc.arena.device_lock_ref(b, SWA), 0);
@@ -2835,11 +2866,13 @@ fn try_device_value_and_evictable_size_read_the_swa_slots() {
store_swa_device(&mut tc, a);
assert!(
tc.get_component_device_value(tc.arena.node(a).id, SWA)
.expect("live test node")
.unwrap()
.equal(&Tensor::from_slice(&[0i64]))
);
assert!(
tc.get_component_device_value(tc.arena.node(b).id, SWA)
.expect("live test node")
.is_none()
);
assert_eq!(tc.evictable_size_(SWA), 1);
@@ -3620,12 +3653,13 @@ fn backup_spec_reads_the_swa_value_recovered_by_an_earlier_action() {
.set_device_value(a, FULL, Tensor::from_slice(&[9i64]));
let a_id = tc.arena.node(a).id;
// The tombstone carries no SWA transfer into the backup spec.
let (_, xfers) = tc.build_backup_spec(a_id);
let (_, xfers) = tc.build_backup_spec(a_id).expect("live test node");
assert!(xfers.is_empty());
// The cache resolves the recover/rebuild action, then rebuilds the spec:
// the deferred read now captures the freshly stored SWA value.
tc.set_component_device_value(a_id, SWA, Tensor::from_slice(&[50i64]));
let (_, xfers) = tc.build_backup_spec(a_id);
tc.set_component_device_value(a_id, SWA, Tensor::from_slice(&[50i64]))
.expect("live test node");
let (_, xfers) = tc.build_backup_spec(a_id).expect("live test node");
let swa_xfer = &xfers[&SWA][0];
assert!(
swa_xfer
@@ -3762,7 +3796,7 @@ fn fallible_load_back_boundaries_reject_a_bare_window_node() {
let node_id = tc.arena.node(a).id;
assert!(matches!(
tc.try_build_hicache_transfers(
tc.build_hicache_transfers(
SWA,
node_id,
CacheTransferPhase::LoadBack,
@@ -3775,7 +3809,7 @@ fn fallible_load_back_boundaries_reject_a_bare_window_node() {
if missing == node_id
));
assert!(matches!(
tc.try_build_load_back_spec(node_id, /* req = */ None),
tc.build_load_back_spec(node_id, /* req = */ None),
Err(TreeCoreRuntimeError::SwaLoadBackMissingValue { node_id: missing })
if missing == node_id
));
@@ -3974,7 +4008,8 @@ fn commit_hicache_transfers_routes_to_the_component() {
&mut cache_actions,
/* insert_result = */ None,
/* pool_storage_result = */ None,
);
)
.expect("live test node");
assert_eq!(cache_actions.len(), 1);
}
@@ -4109,9 +4144,9 @@ fn prefetch_commit_without_a_target_releases_the_whole_buffer() {
}
#[test]
fn prefetch_commit_releases_the_out_of_path_prefix() {
fn prefetch_commit_releases_a_shortened_window_under_a_non_root_anchor() {
// root -> a -> b -> c, one token each; anchor b, target c: the loaded
// window spans two tokens but the leaf->anchor path covers only c's one.
// window is missing its head and cannot be reused as a complete SWA window.
let mut tc = swa_core(/* window = */ 4, /* page_size = */ 1);
let [a, b, c] = chain::<3>(&mut tc);
let mut cache_actions = Vec::new();
@@ -4137,21 +4172,15 @@ fn prefetch_commit_releases_the_out_of_path_prefix() {
Some(&mut insert_result),
Some(&storage_result),
);
// c (on path) fills with the buffer tail; the out-of-path prefix releases.
assert!(
tc.arena
.node(c)
.host_value(SWA)
.equal(&Tensor::from_slice(&[31i64]))
);
assert!(!tc.arena.node(c).has_host_value(SWA));
assert!(!tc.arena.node(b).has_host_value(SWA));
assert!(!tc.arena.node(a).has_host_value(SWA));
assert!(tc.host_lru_list(SWA).in_list(Some(c)));
assert!(!tc.host_lru_list(SWA).in_list(Some(c)));
assert_eq!(cache_actions.len(), 1);
let CacheAction::FreeComponentHostSlot { host_indices, .. } = &cache_actions[0] else {
panic!("expected a host free");
};
assert!(host_indices[0].equal(&Tensor::from_slice(&[30i64])));
assert!(host_indices[0].equal(&Tensor::from_slice(&[30i64, 31])));
}
#[test]
@@ -4317,8 +4346,9 @@ fn build_load_back_spec_includes_the_swa_transfers() {
set_full_host(&mut tc, n);
tc.arena
.set_host_value(n, SWA, Tensor::from_slice(&[30i64]));
let (kv_xfer, mut comp_xfers) =
tc.build_load_back_spec(tc.arena.node(n).id, /* req = */ None);
let (kv_xfer, mut comp_xfers) = tc
.build_load_back_spec(tc.arena.node(n).id, /* req = */ None)
.expect("live test node");
assert_eq!(kv_xfer.nodes_to_load, Some(vec![tc.arena.node(n).id]));
let swa_xfers = comp_xfers.get_mut(&SWA).unwrap();
assert_eq!(swa_xfers.len(), 1);
@@ -4332,12 +4362,14 @@ fn build_load_back_spec_includes_the_swa_transfers() {
assert_eq!(swa_xfers[0].nodes_to_load, Some(vec![tc.arena.node(n).id]));
// The orchestrator fills each transfer's device side from the pool load.
swa_xfers[0].device_indices = Some(Tensor::from_slice(&[60i64]));
let actions = tc.commit_load_back(
tc.arena.node(n).id,
Tensor::from_slice(&[50i64]),
kv_xfer,
comp_xfers,
);
let actions = tc
.commit_load_back(
tc.arena.node(n).id,
Tensor::from_slice(&[50i64]),
kv_xfer,
comp_xfers,
)
.expect("live test node");
assert!(
tc.arena
.device_value(n, FULL)
@@ -4383,7 +4415,8 @@ fn auxiliary_load_does_not_reuse_a_full_pending_pin() {
..Default::default()
},
HashMap::new(),
);
)
.expect("live test node");
assert_eq!(tc.arena.node(shared).load_back_pending_id, Some(shared_id));
let swa_xfer = PoolTransfer {
@@ -4403,7 +4436,8 @@ fn auxiliary_load_does_not_reuse_a_full_pending_pin() {
..Default::default()
},
HashMap::from([(SWA, vec![swa_xfer])]),
);
)
.expect("live test node");
assert_eq!(tc.arena.node(shared).load_back_pending_id, Some(shared_id));
assert_eq!(
@@ -4426,22 +4460,25 @@ fn swa_device_eviction_skips_a_load_back_pinned_node() {
let [n] = chain::<1>(&mut tc);
set_full_host(&mut tc, n);
set_swa_host(&mut tc, n);
let (kv_xfer, mut comp_xfers) =
tc.build_load_back_spec(tc.arena.node(n).id, /* req = */ None);
let (kv_xfer, mut comp_xfers) = tc
.build_load_back_spec(tc.arena.node(n).id, /* req = */ None)
.expect("live test node");
comp_xfers.get_mut(&SWA).unwrap()[0].device_indices = Some(Tensor::from_slice(&[60i64]));
tc.commit_load_back(
tc.arena.node(n).id,
Tensor::from_slice(&[50i64]),
kv_xfer,
comp_xfers,
);
)
.expect("live test node");
// The pin alone keeps the in-flight SWA slice out of every eviction branch.
tc.evict_device_start(SWA, 4);
let (next, _) = tc.evict_device_next_node(SWA, &HashMap::new());
assert_eq!(next, None);
tc.evict_device_end(SWA);
assert!(tc.arena.has_device_value(n, SWA));
tc.finish_load_back(tc.arena.node(n).id);
tc.finish_load_back(tc.arena.node(n).id)
.expect("live test node");
tc.evict_device_start(SWA, 4);
let (next, _) = tc.evict_device_next_node(SWA, &HashMap::new());
assert_eq!(next, Some(tc.arena.node(n).id));
@@ -4465,8 +4502,9 @@ fn swa_host_eviction_skips_a_load_back_pinned_node() {
set_full_host(&mut tc, b);
set_swa_host(&mut tc, b);
tc.host_lru_list_mut(SWA).insert_mru(a);
let (kv_xfer, mut comp_xfers) =
tc.build_load_back_spec(tc.arena.node(b).id, /* req = */ None);
let (kv_xfer, mut comp_xfers) = tc
.build_load_back_spec(tc.arena.node(b).id, /* req = */ None)
.expect("live test node");
assert_eq!(
comp_xfers.get(&SWA).unwrap()[0].nodes_to_load,
Some(vec![tc.arena.node(b).id])
@@ -4477,14 +4515,16 @@ fn swa_host_eviction_skips_a_load_back_pinned_node() {
Tensor::from_slice(&[50i64, 51]),
kv_xfer,
comp_xfers,
);
)
.expect("live test node");
let result = tc.drive_host_eviction(SWA, /* num_tokens = */ 1);
assert_eq!(result.tracker[&SWA], 0);
assert!(result.host_frees.is_empty());
assert!(tc.arena.has_host_value(a, SWA));
tc.finish_load_back(tc.arena.node(b).id);
tc.finish_load_back(tc.arena.node(b).id)
.expect("live test node");
let result = tc.drive_host_eviction(SWA, /* num_tokens = */ 1);
assert_eq!(result.tracker[&SWA], 1);
assert_eq!(result.host_frees[&SWA].len(), 1);
@@ -4511,21 +4551,28 @@ fn build_load_back_spec_degrades_to_empty_on_a_foreign_pin() {
set_full_host(&mut tc, b);
set_swa_host(&mut tc, b);
// Anchor `a` models a Full-only load whose SWA slice remains host-only.
let (kv_xfer, _comp_xfers) =
tc.build_load_back_spec(tc.arena.node(a).id, /* req = */ None);
let (kv_xfer, _comp_xfers) = tc
.build_load_back_spec(tc.arena.node(a).id, /* req = */ None)
.expect("live test node");
tc.commit_load_back(
tc.arena.node(a).id,
Tensor::from_slice(&[50i64]),
kv_xfer,
HashMap::new(),
);
)
.expect("live test node");
// Anchor `b` must reject its SWA window because `a` has a foreign pin.
let (kv_xfer, comp_xfers) = tc.build_load_back_spec(tc.arena.node(b).id, /* req = */ None);
let (kv_xfer, comp_xfers) = tc
.build_load_back_spec(tc.arena.node(b).id, /* req = */ None)
.expect("live test node");
assert_eq!(kv_xfer.host_indices.unwrap().numel(), 0);
assert_eq!(kv_xfer.nodes_to_load, Some(vec![]));
assert!(comp_xfers.is_empty());
tc.finish_load_back(tc.arena.node(a).id);
let (kv_xfer, comp_xfers) = tc.build_load_back_spec(tc.arena.node(b).id, /* req = */ None);
tc.finish_load_back(tc.arena.node(a).id)
.expect("live test node");
let (kv_xfer, comp_xfers) = tc
.build_load_back_spec(tc.arena.node(b).id, /* req = */ None)
.expect("live test node");
assert_eq!(kv_xfer.nodes_to_load, Some(vec![tc.arena.node(b).id]));
assert_eq!(
comp_xfers.get(&SWA).unwrap()[0].nodes_to_load,
@@ -4551,7 +4598,8 @@ fn host_drive_reclaims_swa_coexisting_host_values_when_the_host_lru_is_empty() {
let leaf_idx = child_of(&tc, parent_idx, &[3]);
let (parent, leaf) = (tc.arena.node(parent_idx).id, tc.arena.node(leaf_idx).id);
for (handle, slots) in [(parent, vec![30i64, 31]), (leaf, vec![32i64])] {
tc.set_component_device_value(handle, SWA, Tensor::from_slice(&slots));
tc.set_component_device_value(handle, SWA, Tensor::from_slice(&slots))
.expect("live test node");
}
for (handle, host) in [(parent, vec![20i64, 21]), (leaf, vec![22i64])] {
let swa_xfer = PoolTransfer {
@@ -4563,7 +4611,8 @@ fn host_drive_reclaims_swa_coexisting_host_values_when_the_host_lru_is_empty() {
handle,
Tensor::from_slice(&host),
HashMap::from([(SWA, vec![swa_xfer])]),
);
)
.expect("live test node");
}
assert_eq!(tc.host_lru_list(SWA).len(), 0);
@@ -4667,12 +4716,14 @@ fn write_through_offloads_a_boundary_split_leaf() {
tc.arena.node(parent).id,
Tensor::from_slice(&[100i64, 101]),
HashMap::new(),
);
)
.expect("live test node");
tc.commit_backup(
tc.arena.node(leaf).id,
Tensor::from_slice(&[102i64, 103]),
HashMap::new(),
);
)
.expect("live test node");
let mut tracker = swa_tracker();
let (mut df, mut hf) = (HashMap::new(), HashMap::new());
tc.evict_device_start(FULL, /* request_cnt = */ 100);
@@ -4680,7 +4731,9 @@ fn write_through_offloads_a_boundary_split_leaf() {
let (next, step) = tc.evict_device_next_node(FULL, &tracker);
accumulate_step(step, &mut tracker, &mut df, &mut hf);
let Some(next) = next else { break };
let (backup, step) = tc.evict_device_leaf(next, /* is_write_back = */ false);
let (backup, step) = tc
.evict_device_leaf(next, /* is_write_back = */ false)
.expect("live test node");
assert!(backup.is_none());
accumulate_step(step, &mut tracker, &mut df, &mut hf);
}
@@ -4727,7 +4780,8 @@ fn deep_swa_tree_survives_backup_evict_and_load_back_rounds() {
source_value,
} = action
{
tc.set_component_device_value(*node_id, SWA, source_value.copy());
tc.set_component_device_value(*node_id, SWA, source_value.copy())
.expect("live test node");
}
}
tc.sanity_check(&[], &[]);
@@ -4744,7 +4798,8 @@ fn deep_swa_tree_survives_backup_evict_and_load_back_rounds() {
tc.arena.node(node).id,
Tensor::from_slice(&vec![0i64; len]),
HashMap::new(),
);
)
.expect("live test node");
}
tc.sanity_check(&[], &[]);
// Stepwise eviction rounds: half the Full budget, then the whole SWA budget.
@@ -4757,7 +4812,9 @@ fn deep_swa_tree_survives_backup_evict_and_load_back_rounds() {
let (leaf, step) = tc.evict_device_next_node(FULL, &tracker);
accumulate_step(step, &mut tracker, &mut df, &mut hf);
let Some(leaf) = leaf else { break };
let (_, step) = tc.evict_device_leaf(leaf, /* is_write_back = */ false);
let (_, step) = tc
.evict_device_leaf(leaf, /* is_write_back = */ false)
.expect("live test node");
accumulate_step(step, &mut tracker, &mut df, &mut hf);
}
tc.evict_device_end(FULL);
@@ -4770,7 +4827,9 @@ fn deep_swa_tree_survives_backup_evict_and_load_back_rounds() {
let (leaf, step) = tc.evict_device_next_node(SWA, &tracker);
accumulate_step(step, &mut tracker, &mut df, &mut hf);
let Some(leaf) = leaf else { break };
let (_, step) = tc.evict_device_leaf(leaf, /* is_write_back = */ false);
let (_, step) = tc
.evict_device_leaf(leaf, /* is_write_back = */ false)
.expect("live test node");
accumulate_step(step, &mut tracker, &mut df, &mut hf);
}
tc.evict_device_end(SWA);
@@ -4781,24 +4840,31 @@ fn deep_swa_tree_survives_backup_evict_and_load_back_rounds() {
// commit-then-lock sequence.
for key in [vec![1i64, 2, 3, 4, 5, 6], vec![1i64, 2]] {
let anchor = tc.match_prefix(&match_params(&key)).best_match_node_id;
if !tc.is_root(anchor) && tc.is_full_device_evicted(anchor) {
let (kv_xfer, comp_xfers) = tc.build_load_back_spec(anchor, /* req = */ None);
if !tc.is_root(anchor).expect("live test node")
&& tc.is_full_device_evicted(anchor).expect("live test node")
{
let (kv_xfer, comp_xfers) = tc
.build_load_back_spec(anchor, /* req = */ None)
.expect("live test node");
let loaded = kv_xfer.host_indices.as_ref().unwrap().numel();
let actions = tc.commit_load_back(
anchor,
Tensor::from_slice(&vec![0i64; loaded]),
kv_xfer,
comp_xfers,
);
let actions = tc
.commit_load_back(
anchor,
Tensor::from_slice(&vec![0i64; loaded]),
kv_xfer,
comp_xfers,
)
.expect("live test node");
assert!(actions.is_empty());
let lock = tc.inc_lock_ref(anchor);
let lock = tc.inc_lock_ref(anchor).expect("live test node");
let params = DecLockRefParams {
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,
};
tc.dec_lock_ref(anchor, Some(&params), /* skip_swa = */ false);
tc.finish_load_back(anchor);
tc.dec_lock_ref(anchor, Some(&params), /* skip_swa = */ false)
.expect("live test node");
tc.finish_load_back(anchor).expect("live test node");
}
tc.sanity_check(&[], &[]);
}
@@ -4836,7 +4902,8 @@ fn recovered_swa_span_evicts_before_the_window_leaf() {
);
};
assert_eq!(*node_id, tc.arena.node(leaf).id);
tc.set_component_device_value(*node_id, SWA, source_value.copy());
tc.set_component_device_value(*node_id, SWA, source_value.copy())
.expect("live test node");
assert!(!tc.arena.has_device_value(prefix, SWA));
// The fully-in-window re-insert recovers the prefix at its walk barrier.
@@ -4866,7 +4933,8 @@ fn recovered_swa_span_evicts_before_the_window_leaf() {
100i64, 101, 102, 103, 104, 105, 106, 107
])));
assert_eq!(*node_id, tc.arena.node(prefix).id);
tc.set_component_device_value(*node_id, SWA, source_value.copy());
tc.set_component_device_value(*node_id, SWA, source_value.copy())
.expect("live test node");
let done = tc.resume_insert();
assert_eq!(
done.result.expect("the resumed walk completes").prefix_len,
rust/sglang-radix-tree/src/tests/node.rs
+13 -8
View File
@@ -5,7 +5,7 @@ use tch::Tensor;
use super::*;
use crate::components::{FULL, MAMBA, SWA};
use crate::node::TreeCoreRuntimeError;
use crate::node::{NodeAccessError, TreeCoreRuntimeError};
static COUNTED_KEY_CONSTRUCTIONS: AtomicUsize = AtomicUsize::new(0);
@@ -1264,7 +1264,7 @@ fn alloc_stamps_self_id_and_a_fresh_access_tick() -> Result<(), TreeCoreRuntimeE
)?;
assert_eq!(arena.node(a).id, 1);
assert_eq!(arena.node(b).id, 2);
assert_eq!(arena.resolve(arena.node(a).id), a);
assert_eq!(arena.resolve(arena.node(a).id).expect("live test node"), a);
// Construction stamps strictly increasing ticks: root, then a, then b;
// both stamps share the node's single construction tick.
let root_tick = arena.node(root).last_access_counter;
@@ -1711,11 +1711,13 @@ fn failed_alloc_child_mints_no_id_and_keeps_the_freelist() -> Result<(), TreeCor
}
#[test]
#[should_panic(expected = "is not allocated")]
fn resolve_panics_on_a_never_minted_handle() {
fn resolve_returns_err_for_a_never_minted_handle() {
let arena = arena();
arena.root();
arena.resolve(1_000_000);
assert!(matches!(
arena.resolve(1_000_000),
Err(NodeAccessError { node_id: 1_000_000 })
));
}
#[test]
@@ -1815,7 +1817,7 @@ fn id_map_stays_consistent_across_free_and_realloc() -> Result<(), TreeCoreRunti
)?;
let b_id = arena.node(b).id;
arena.free_leaf(b)?;
assert!(arena.try_resolve(b_id).is_none());
assert!(arena.resolve(b_id).is_err());
// The freed slot is recycled with a fresh handle; the old one stays dead.
let c = arena.alloc_child(
root,
@@ -1825,10 +1827,13 @@ fn id_map_stays_consistent_across_free_and_realloc() -> Result<(), TreeCoreRunti
)?;
assert_eq!(c, b);
assert_ne!(arena.node(c).id, b_id);
assert!(arena.try_resolve(b_id).is_none());
assert!(arena.resolve(b_id).is_err());
// Every live slot resolves back from its own handle.
for idx in arena.live_ids().collect::<Vec<_>>() {
assert_eq!(arena.resolve(arena.node(idx).id), idx);
assert_eq!(
arena.resolve(arena.node(idx).id).expect("live test node"),
idx
);
}
let _ = a;
Ok(())
rust/sglang-radix-tree/src/tests/unified_tree_core.rs
+861 -413
View File
File diff suppressed because it is too large Load Diff
rust/sglang-radix-tree/src/unified_tree_core.rs
+315 -316
View File
File diff suppressed because it is too large Load Diff
test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_full.py
+2
View File
@@ -7,6 +7,7 @@ from sglang.test.test_utils import (
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
DEFAULT_URL_FOR_TEST,
CustomTestCase,
is_in_amd_ci,
popen_launch_server,
terminate_and_kill_process_tree,
unified_radix_tree_server_env,
@@ -49,6 +50,7 @@ class TestUnifiedFullRadixCache(UnifiedRadixTreeTestMixin, CustomTestCase):
terminate_and_kill_process_tree(cls.process, wait_timeout=60)
@unittest.skipIf(is_in_amd_ci(), "Rust TreeCore is not packaged in AMD CI")
class TestRustUnifiedFullRadixCache(TestUnifiedFullRadixCache):
tree_core_backend = "rust"
test/registered/rust/test_rust_extension.py
+38 -1
View File
@@ -162,6 +162,7 @@ crate-type = ["cdylib"]
self.assertEqual(crate.library, "demo_extension")
self.assertEqual(crate.python_module, "demo._core")
self.assertEqual(crate.features, ("python",))
self.assertEqual(crate.source_inputs, ())
with self.assertRaisesRegex(
ModuleNotFoundError, r"declared modules: \['demo\._core'\]"
@@ -208,6 +209,36 @@ crate-type = ["cdylib"]
inspection.target_fingerprint,
)
def test_fingerprint_covers_declared_external_source_inputs(self):
with TemporaryDirectory() as directory:
root = Path(directory)
workspace = self._workspace(root)
proto = root / "proto/demo.proto"
proto.parent.mkdir()
proto.write_text("message Demo {}\n", encoding="utf-8")
manifest = workspace / "demo/Cargo.toml"
manifest.write_text(
manifest.read_text(encoding="utf-8").replace(
'features = ["python"]',
'features = ["python"]\nsource-inputs = ["../../proto"]',
),
encoding="utf-8",
)
crate = rust_extension._discover_crate(workspace, "demo._core")
self.assertEqual(crate.source_inputs, (proto.parent.resolve(),))
with mock.patch.object(
rust_extension,
"_command_version",
side_effect=lambda command, *args, **kwargs: f"{command} 1.0",
):
first = rust_extension._build_context(crate)
proto.write_text("message Changed {}\n", encoding="utf-8")
changed = rust_extension._build_context(crate)
self.assertNotEqual(first.fingerprint, changed.fingerprint)
self.assertEqual(first.target_fingerprint, changed.target_fingerprint)
def test_auto_builds_once_then_uses_cache(self):
with TemporaryDirectory() as directory:
root = Path(directory)
@@ -527,30 +558,35 @@ crate-type = ["cdylib"]
self.assertNotIn(module_name, sys.modules)
def test_checked_in_crates_are_discovered_from_wheel_metadata(self):
for python_module, package, library, features in (
grpc_proto = (rust_extension._RUST_WORKSPACE.parent / "proto").resolve()
for python_module, package, library, features, source_inputs in (
(
"sglang.srt.rust_extensions._server",
"sglang-server",
"sglang_server",
(),
(),
),
(
"sglang.srt.rust_extensions._grpc",
"sglang-grpc",
"sglang_grpc_core",
(),
(grpc_proto,),
),
(
"sglang.srt.rust_extensions._multimodal",
"sglang-mm",
"sglang_mm_core",
("python", "parallel"),
(),
),
(
"sglang.srt.mem_cache.rust_tree_core.mem_cache",
"sglang-radix-tree",
"mem_cache",
("python-extension",),
(),
),
):
crate = rust_extension._discover_crate(
@@ -559,6 +595,7 @@ crate-type = ["cdylib"]
self.assertEqual(crate.package, package)
self.assertEqual(crate.library, library)
self.assertEqual(crate.features, features)
self.assertEqual(crate.source_inputs, source_inputs)
if __name__ == "__main__":
test/registered/unit/mem_cache/test_rust_tree_core_integration.py
+280 -6
View File
@@ -29,6 +29,7 @@ from sglang.srt.mem_cache.base_prefix_cache import (
InsertParams,
InsertResult,
MatchPrefixParams,
MatchResult,
)
from sglang.srt.mem_cache.cache_init_params import CacheInitParams
from sglang.srt.mem_cache.hicache_storage import (
@@ -173,6 +174,51 @@ def test_stale_handle_reads_raise_key_error_without_poisoning_the_core():
assert core.is_root(live_root)
def test_stale_match_finalizer_handles_raise_key_error_without_poisoning_the_core():
from rust_unified_tree_core_inspector import RustUnifiedTreeCoreInspector
core = RustUnifiedTreeCoreInspector(
CacheInitParams(
disable=False,
req_to_token_pool=None,
token_to_kv_pool_allocator=None,
page_size=1,
tree_components=(ComponentType.FULL,),
)
)
stale_root = core.root_node_handle()
core.reset()
live_root = core.root_node_handle()
result = MatchResult(
device_indices=torch.empty(0, dtype=torch.int64),
last_device_node=live_root,
last_host_node=live_root,
best_match_node=live_root,
)
params = MatchPrefixParams(key=_key([]))
for field in ("last_device_node", "last_host_node", "best_match_node"):
with pytest.raises(KeyError) as exc_info:
core.finalize_component_match_result(
ComponentType.FULL,
result._replace(**{field: stale_root}),
params,
value_chunks=[],
best_value_len=0,
)
assert exc_info.value.args == (stale_root,), field
assert core.is_root(live_root), field
finalized = core.finalize_component_match_result(
ComponentType.FULL,
result,
params,
value_chunks=[],
best_value_len=0,
)
assert finalized.best_match_node == live_root
def test_stale_handle_operations_raise_key_error_without_poisoning_the_core():
from sglang.srt.mem_cache.unified_cache.components import CacheTransferPhase
@@ -180,15 +226,122 @@ def test_stale_handle_operations_raise_key_error_without_poisoning_the_core():
stale_root = core.root_node_handle()
core.reset()
live_root = core.root_node_handle()
empty = torch.empty(0, dtype=torch.int64)
operations = (
lambda: core.demote(stale_root),
lambda: core.build_hicache_transfers(
operations = {
"inc_lock_ref": lambda: core.inc_lock_ref(stale_root),
"dec_lock_ref": lambda: core.dec_lock_ref(stale_root),
"dec_swa_lock_only": lambda: core.dec_swa_lock_only(stale_root, None),
"evict_device_leaf": lambda: core.evict_device_leaf(stale_root, False),
"drop_subtree_no_host": lambda: core.drop_subtree_no_host(stale_root),
"demote": lambda: core.demote(stale_root),
"is_full_device_evicted": lambda: core.is_full_device_evicted(stale_root),
"collect_full_device_indices/from": lambda: core.collect_full_device_indices(
stale_root, live_root
),
"collect_full_device_indices/until": lambda: core.collect_full_device_indices(
live_root, stale_root
),
"insert_host": lambda: core.insert_host(
stale_root, _key([1]), empty, ["0" * 64]
),
"build_backup_spec": lambda: core.build_backup_spec(stale_root),
"build_storage_backup_spec": lambda: core.build_storage_backup_spec(
stale_root, False
),
"build_hicache_transfers": lambda: core.build_hicache_transfers(
ComponentType.FULL, stale_root, CacheTransferPhase.BACKUP_STORAGE
),
lambda: core.build_load_back_spec(stale_root),
lambda: core.get_hash_values(stale_root),
lambda: core.dfs_weight_order([stale_root]),
"commit_backup": lambda: core.commit_backup(stale_root, empty, {}),
"commit_hicache_transfers": lambda: core.commit_hicache_transfers(
stale_root,
CacheTransferPhase.BACKUP_HOST,
{},
cache_actions=[],
),
"commit_load_back": lambda: core.commit_load_back(
stale_root, empty, PoolTransfer(name=PoolName.KV), {}
),
"build_load_back_spec": lambda: core.build_load_back_spec(stale_root),
"evict_excess_path_states": lambda: core.evict_excess_path_states(
stale_root, {}, {}
),
"inc_host_lock_ref": lambda: core.inc_host_lock_ref(stale_root),
"dec_host_lock_ref": lambda: core.dec_host_lock_ref(stale_root),
"mark_write_through_pending": lambda: core.mark_write_through_pending(
[stale_root], stale_root
),
"finish_write_through": lambda: core.finish_write_through(
[stale_root], stale_root
),
"finish_load_back": lambda: core.finish_load_back(stale_root),
"get_component_device_value": lambda: core.get_component_device_value(
stale_root, ComponentType.FULL
),
"component_has_host_value_only": lambda: core.component_has_host_value_only(
stale_root, ComponentType.FULL
),
"get_hash_values": lambda: core.get_hash_values(stale_root),
"dfs_weight_order": lambda: core.dfs_weight_order([stale_root]),
}
for name, operation in operations.items():
with pytest.raises(KeyError) as exc_info:
operation()
assert exc_info.value.args == (stale_root,), name
assert core.is_root(live_root), name
def test_stale_handles_nested_in_transfer_results_do_not_poison_the_core():
from sglang.srt.mem_cache.unified_cache.components import CacheTransferPhase
core = _tree_core()
stale_root = core.root_node_handle()
core.reset()
live_root = core.root_node_handle()
stale_transfer = PoolTransfer(name=PoolName.KV, nodes_to_load=[stale_root])
operations = (
lambda: core.commit_hicache_transfers(
live_root,
CacheTransferPhase.LOAD_BACK,
{ComponentType.FULL: [stale_transfer]},
cache_actions=[],
),
lambda: core.commit_hicache_transfers(
live_root,
CacheTransferPhase.PREFETCH,
{},
cache_actions=[],
insert_result=InsertResult(prefix_len=0, inserted_host_node=stale_root),
),
lambda: core.commit_load_back(
live_root,
torch.empty(0, dtype=torch.int64),
stale_transfer,
{},
),
)
for operation in operations:
with pytest.raises(KeyError) as exc_info:
operation()
assert exc_info.value.args == (stale_root,)
assert core.is_root(live_root)
def test_stale_handle_component_access_does_not_poison_the_core():
core = _tree_core(
tree_components=(ComponentType.FULL, ComponentType.SWA),
sliding_window_size=8,
)
stale_root = core.root_node_handle()
core.reset()
live_root = core.root_node_handle()
operations = (
lambda: core.set_component_device_value(
stale_root, ComponentType.SWA, torch.empty(0, dtype=torch.int64)
),
lambda: core.get_component_device_value(stale_root, ComponentType.SWA),
)
for operation in operations:
with pytest.raises(KeyError) as exc_info:
@@ -1261,6 +1414,10 @@ def test_buffer_backup_snapshot_round_trips_and_detects_a_split():
)
assert core.validate_buffer_backup(leaf, len(snapshot.key)) is None
core.reset()
assert core.snapshot_buffer_backup(leaf, pass_prefix_keys=True) is None
assert core.validate_buffer_backup(leaf, len(snapshot.key)) is None
def test_buffer_backup_snapshot_preserves_bigram_keys():
core = _tree_core(is_eagle=True)
@@ -1996,5 +2153,122 @@ def test_bigram_insert_value_shorter_than_the_bigram_count_raises():
)
def test_stale_inspection_handles_raise_key_error_or_report_absence():
from rust_unified_tree_core_inspector import RustUnifiedTreeCoreInspector
from sglang.srt.mem_cache.unified_cache.components import EvictLayer
core = RustUnifiedTreeCoreInspector(
CacheInitParams(
disable=False,
req_to_token_pool=None,
token_to_kv_pool_allocator=None,
page_size=1,
tree_components=(ComponentType.FULL,),
)
)
stale_root = core.root_node_handle()
core.reset()
live_root = core.root_node_handle()
operations = {
"get_parent_node_id": lambda: core.get_parent_node_id(stale_root),
"get_child_node_ids": lambda: core.get_child_node_ids(stale_root),
"get_node_key_length": lambda: core.get_node_key_length(stale_root),
"get_node_token_ids": lambda: core.get_node_token_ids(stale_root),
"is_node_key_bigram": lambda: core.is_node_key_bigram(stale_root),
"get_component_host_value": lambda: core.get_component_host_value(
stale_root, ComponentType.FULL
),
"get_component_device_lock_ref": lambda: core.get_component_device_lock_ref(
stale_root, ComponentType.FULL
),
"get_node_hit_count": lambda: core.get_node_hit_count(stale_root),
"get_write_through_pending_id": lambda: core.get_write_through_pending_id(
stale_root
),
"is_node_in_device_lru": lambda: core.is_node_in_device_lru(
stale_root, ComponentType.FULL
),
"is_node_in_host_lru": lambda: core.is_node_in_host_lru(
stale_root, ComponentType.FULL
),
"is_device_leaf": lambda: core.is_device_leaf(stale_root),
"set_node_hash_values": lambda: core.set_node_hash_values(stale_root, None),
"set_component_device_value_raw": lambda: core.set_component_device_value_raw(
stale_root, ComponentType.FULL, None
),
"set_component_host_value_raw": lambda: core.set_component_host_value_raw(
stale_root, ComponentType.FULL, None
),
"set_component_device_lock_ref": lambda: core.set_component_device_lock_ref(
stale_root, ComponentType.FULL, 0
),
"remove_node_from_device_lru": lambda: core.remove_node_from_device_lru(
stale_root, ComponentType.FULL
),
"insert_node_into_host_lru": lambda: core.insert_node_into_host_lru(
stale_root, ComponentType.FULL
),
"update_duplicate_tracking": lambda: core.update_duplicate_tracking(stale_root),
"evict_component": lambda: core.evict_component(
stale_root, ComponentType.FULL, EvictLayer.DEVICE
),
"validate_cascade_evict": lambda: core.validate_cascade_evict(
stale_root, ComponentType.FULL, EvictLayer.DEVICE
),
"cleanup_tombstone_ancestors": lambda: core.cleanup_tombstone_ancestors(
stale_root
),
"build_backup_node_ids": lambda: core.build_backup_node_ids(stale_root),
}
for name, operation in operations.items():
with pytest.raises(KeyError) as exc_info:
operation()
assert exc_info.value.args == (stale_root,), name
assert core.is_root(live_root), name
disabled_component_operations = {
"get_component_host_value": lambda: core.get_component_host_value(
stale_root, ComponentType.SWA
),
"get_component_device_lock_ref": lambda: core.get_component_device_lock_ref(
stale_root, ComponentType.SWA
),
"set_component_device_value_raw": lambda: core.set_component_device_value_raw(
stale_root, ComponentType.SWA, None
),
"set_component_host_value_raw": lambda: core.set_component_host_value_raw(
stale_root, ComponentType.SWA, None
),
"set_component_device_lock_ref": lambda: core.set_component_device_lock_ref(
stale_root, ComponentType.SWA, 0
),
"remove_node_from_device_lru": lambda: core.remove_node_from_device_lru(
stale_root, ComponentType.SWA
),
"insert_node_into_host_lru": lambda: core.insert_node_into_host_lru(
stale_root, ComponentType.SWA
),
"evict_component": lambda: core.evict_component(
stale_root, ComponentType.SWA, EvictLayer.DEVICE
),
"validate_cascade_evict": lambda: core.validate_cascade_evict(
stale_root, ComponentType.SWA, EvictLayer.DEVICE
),
}
for name, operation in disabled_component_operations.items():
with pytest.raises(KeyError) as exc_info:
operation()
assert exc_info.value.args == (stale_root,), name
assert core.is_root(live_root), name
assert not core.contains_node(stale_root)
assert not core.is_device_evictable_leaf(stale_root)
assert not core.is_host_evictable_leaf(stale_root)
assert not core.is_node_in_device_lru(stale_root, ComponentType.SWA)
assert not core.is_node_in_host_lru(stale_root, ComponentType.SWA)
if __name__ == "__main__":
sys.exit(pytest.main([__file__]))
test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py
+55
View File
@@ -8065,6 +8065,61 @@ class _InsertWalkSuite(CustomTestCase):
)
@unittest.skipUnless(torch.cuda.is_available(), "cache fixtures need CUDA")
class TestUnifiedTreeCoreSWAPrefetchBackends(_InsertWalkSuite):
cfg = CacheConfig(
components=(ComponentType.FULL, ComponentType.SWA), sliding_window_size=4
)
def test_mid_tree_shortened_swa_prefetch_is_released(self):
cache, allocator, req_to_token_pool = build_fixture(self.cfg)
prefix = [1, 2]
self._insert(cache, allocator, req_to_token_pool, prefix)
anchor = cache.match_prefix(
MatchPrefixParams(key=RadixKey(array("q", prefix)))
).last_device_node
cache.tree_core.is_write_back = True
suffix = [3, 4]
insert_result = cache.tree_core.insert_host(
anchor,
RadixKey(array("q", suffix)),
torch.tensor([100, 101], dtype=torch.int64),
["h3", "h4"],
)
self.assertIsNotNone(insert_result.inserted_host_node)
swa_host_indices = torch.tensor([30, 31], dtype=torch.int64)
actions = []
cache.tree_core.commit_hicache_transfers(
anchor,
CacheTransferPhase.PREFETCH,
{
ComponentType.SWA: [
PoolTransfer(
name=PoolName.SWA,
host_indices=swa_host_indices,
)
]
},
cache_actions=actions,
insert_result=insert_result,
pool_storage_result=PoolTransferResult(
kv_hit_pages=len(suffix),
extra_pool_hit_pages={PoolName.SWA: len(suffix)},
),
)
self.assertIsNone(
_host_value(cache, insert_result.inserted_host_node, ComponentType.SWA)
)
self.assertEqual(len(actions), 1)
self.assertIsInstance(actions[0], FreeComponentHostSlot)
self.assertEqual(actions[0].component_type, ComponentType.SWA)
self.assertEqual(len(actions[0].host_indices), 1)
self.assertTrue(torch.equal(actions[0].host_indices[0], swa_host_indices))
@unittest.skipUnless(torch.cuda.is_available(), "cache fixtures need CUDA")
class TestResumableInsertWalk(_InsertWalkSuite):
cfg = CacheConfig()