From bbd672a8ab7fd6c2e2cda4931f761b742d941d64 Mon Sep 17 00:00:00 2001 From: Jialin Ouyang Date: Mon, 31 Aug 2026 16:34:29 -0700 Subject: [PATCH 1/6] [Rust] Harden extension build cache inputs --- .github/actions/download-rust-ext/action.yml | 2 +- .github/workflows/_pr-test-rust-ext-build.yml | 8 ++-- .github/workflows/seed-rust-ext-cache.yml | 1 + python/sglang/srt/rust_extensions/loader.py | 29 ++++++++++---- rust/sglang-grpc/Cargo.toml | 3 ++ test/registered/rust/test_rust_extension.py | 39 ++++++++++++++++++- 6 files changed, 70 insertions(+), 12 deletions(-) diff --git a/.github/actions/download-rust-ext/action.yml b/.github/actions/download-rust-ext/action.yml index ee5183102d2c..7809793cefcb 100644 --- a/.github/actions/download-rust-ext/action.yml +++ b/.github/actions/download-rust-ext/action.yml @@ -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: diff --git a/.github/workflows/_pr-test-rust-ext-build.yml b/.github/workflows/_pr-test-rust-ext-build.yml index 227f6e6a045e..7c4ddfc4eec5 100644 --- a/.github/workflows/_pr-test-rust-ext-build.yml +++ b/.github/workflows/_pr-test-rust-ext-build.yml @@ -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 @@ -359,7 +361,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 diff --git a/.github/workflows/seed-rust-ext-cache.yml b/.github/workflows/seed-rust-ext-cache.yml index 7256b7a4a623..6740b16df8d7 100644 --- a/.github/workflows/seed-rust-ext-cache.yml +++ b/.github/workflows/seed-rust-ext-cache.yml @@ -9,6 +9,7 @@ on: branches: [main] paths: - 'rust/**' + - 'proto/**' - 'python/setup.py' - 'python/pyproject.toml' - 'python/sglang/srt/rust_extensions/torch_build.py' diff --git a/python/sglang/srt/rust_extensions/loader.py b/python/sglang/srt/rust_extensions/loader.py index bc76bb1af71e..2175f0e77462 100644 --- a/python/sglang/srt/rust_extensions/loader.py +++ b/python/sglang/srt/rust_extensions/loader.py @@ -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(): diff --git a/rust/sglang-grpc/Cargo.toml b/rust/sglang-grpc/Cargo.toml index d26625c77958..ae66f6b17d7f 100644 --- a/rust/sglang-grpc/Cargo.toml +++ b/rust/sglang-grpc/Cargo.toml @@ -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 diff --git a/test/registered/rust/test_rust_extension.py b/test/registered/rust/test_rust_extension.py index b80f45806cb9..53208137913d 100644 --- a/test/registered/rust/test_rust_extension.py +++ b/test/registered/rust/test_rust_extension.py @@ -162,6 +162,7 @@ def test_discovery_reads_crate_manifest_metadata(self): 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 @@ def test_fingerprint_is_content_based_and_covers_build_inputs(self): 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 @@ def test_failed_import_does_not_poison_sys_modules(self): 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 @@ def test_checked_in_crates_are_discovered_from_wheel_metadata(self): 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__": From a0395b30e19f0fb95b4efdf39c4c9f2c6a11e4ed Mon Sep 17 00:00:00 2001 From: Jialin Ouyang Date: Mon, 31 Aug 2026 04:08:56 -0700 Subject: [PATCH 2/6] [CI] Keep Rust TreeCore E2E off AMD --- .../unified_radix_tree/test_unified_radix_cache_kl_full.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_full.py b/test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_full.py index ff50de352866..b7588ac6fd2b 100644 --- a/test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_full.py +++ b/test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_full.py @@ -7,6 +7,7 @@ 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 @@ def tearDownClass(cls): 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" From d43934816a4d5688d18becce1455dd42baf92a62 Mon Sep 17 00:00:00 2001 From: Jialin Ouyang Date: Mon, 31 Aug 2026 16:35:07 -0700 Subject: [PATCH 3/6] [Rust TreeCore] Reject shortened SWA prefetches --- rust/sglang-radix-tree/src/components/swa.rs | 9 +++ .../src/tests/components/swa.rs | 16 ++---- .../test_unified_radix_cache_unittest.py | 55 +++++++++++++++++++ 3 files changed, 69 insertions(+), 11 deletions(-) diff --git a/rust/sglang-radix-tree/src/components/swa.rs b/rust/sglang-radix-tree/src/components/swa.rs index 27906e8a31d0..776a193d7f7c 100644 --- a/rust/sglang-radix-tree/src/components/swa.rs +++ b/rust/sglang-radix-tree/src/components/swa.rs @@ -233,6 +233,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; diff --git a/rust/sglang-radix-tree/src/tests/components/swa.rs b/rust/sglang-radix-tree/src/tests/components/swa.rs index 452858a9a0c8..7b80f61a620a 100644 --- a/rust/sglang-radix-tree/src/tests/components/swa.rs +++ b/rust/sglang-radix-tree/src/tests/components/swa.rs @@ -4109,9 +4109,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 +4137,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] diff --git a/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py b/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py index df4e258f10c5..b15775829a1e 100644 --- a/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py +++ b/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py @@ -8063,6 +8063,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() From f058a65af0af134c655dbf4c09517ae28e209ae6 Mon Sep 17 00:00:00 2001 From: Jialin Ouyang Date: Tue, 1 Sep 2026 12:28:28 -0700 Subject: [PATCH 4/6] [Rust TreeCore] Propagate stale node access errors --- rust/sglang-radix-tree/src/components/full.rs | 12 +- .../sglang-radix-tree/src/components/mamba.rs | 13 +- rust/sglang-radix-tree/src/components/mod.rs | 2 + rust/sglang-radix-tree/src/components/swa.rs | 18 +- rust/sglang-radix-tree/src/node.rs | 31 +- rust/sglang-radix-tree/src/python_bindings.rs | 410 ++++-- .../src/tests/components/base.rs | 2 + .../src/tests/components/full.rs | 262 +++- .../src/tests/components/mamba.rs | 176 ++- .../src/tests/components/swa.rs | 224 ++- rust/sglang-radix-tree/src/tests/node.rs | 21 +- .../src/tests/unified_tree_core.rs | 1267 +++++++++++------ .../src/unified_tree_core.rs | 629 ++++---- .../test_rust_tree_core_integration.py | 286 +++- 14 files changed, 2243 insertions(+), 1110 deletions(-) diff --git a/rust/sglang-radix-tree/src/components/full.rs b/rust/sglang-radix-tree/src/components/full.rs index be58d7158fc8..27f16b5b52f2 100644 --- a/rust/sglang-radix-tree/src/components/full.rs +++ b/rust/sglang-radix-tree/src/components/full.rs @@ -49,6 +49,8 @@ impl TreeComponent for FullComponent { &self, tree_core: &UnifiedTreeCore, 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 TreeComponent 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 TreeComponent 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 diff --git a/rust/sglang-radix-tree/src/components/mamba.rs b/rust/sglang-radix-tree/src/components/mamba.rs index d4f3c07499cf..8dde068d7e79 100644 --- a/rust/sglang-radix-tree/src/components/mamba.rs +++ b/rust/sglang-radix-tree/src/components/mamba.rs @@ -127,6 +127,8 @@ impl TreeComponent for MambaComponent { &self, tree_core: &UnifiedTreeCore, 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 TreeComponent 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 TreeComponent 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) => diff --git a/rust/sglang-radix-tree/src/components/mod.rs b/rust/sglang-radix-tree/src/components/mod.rs index b768b091726a..72add6b139f7 100644 --- a/rust/sglang-radix-tree/src/components/mod.rs +++ b/rust/sglang-radix-tree/src/components/mod.rs @@ -130,6 +130,8 @@ pub trait TreeComponent { &self, tree_core: &UnifiedTreeCore, result: MatchResult, + _last_device_node_idx: NodeIdx_, + _best_match_node_idx: NodeIdx_, params: &MatchPrefixParams<'_, K>, value_chunks: &[Tensor], best_value_len: usize, diff --git a/rust/sglang-radix-tree/src/components/swa.rs b/rust/sglang-radix-tree/src/components/swa.rs index 776a193d7f7c..e531999f683e 100644 --- a/rust/sglang-radix-tree/src/components/swa.rs +++ b/rust/sglang-radix-tree/src/components/swa.rs @@ -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 { @@ -347,6 +352,8 @@ impl TreeComponent for SwaComponent { &self, tree_core: &UnifiedTreeCore, mut result: MatchResult, + _last_device_node_idx: NodeIdx_, + best_match_node_idx: NodeIdx_, params: &MatchPrefixParams<'_, K>, value_chunks: &[Tensor], best_value_len: usize, @@ -356,9 +363,7 @@ impl TreeComponent 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); @@ -974,7 +979,10 @@ impl TreeComponent for SwaComponent { let mut swa_chunks: Vec = 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_( diff --git a/rust/sglang-radix-tree/src/node.rs b/rust/sglang-radix-tree/src/node.rs index 1e3758edb231..8d02421ead4a 100644 --- a/rust/sglang-radix-tree/src/node.rs +++ b/rust/sglang-radix-tree/src/node.rs @@ -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 NodeArena { 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 { + 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 { - self.id_map.get(&id).copied() + .copied() + .ok_or(NodeAccessError { node_id: id }) } /// Mint the next external handle for the slot and index it. diff --git a/rust/sglang-radix-tree/src/python_bindings.rs b/rust/sglang-radix-tree/src/python_bindings.rs index af470a613c75..a1ad8f413a61 100644 --- a/rust/sglang-radix-tree/src/python_bindings.rs +++ b/rust/sglang-radix-tree/src/python_bindings.rs @@ -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 { } } +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 TreeCoreBinding { .into_iter() .map(parse_component_type) .collect::>>()?; - 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 TreeCoreBinding { 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 TreeCoreBinding { .into_iter() .map(|(ct, node_ids)| Ok((parse_component_type(ct)?, node_ids))) .collect::>>()?; - 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 TreeCoreBinding { 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 TreeCoreBinding { component_type: u8, ) -> PyResult> { 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 TreeCoreBinding { py: Python<'_>, node_id: NodeId, ) -> PyResult { - 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 TreeCoreBinding { 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 { + 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 TreeCoreBinding { } /// 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 { 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 TreeCoreBinding { } 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 TreeCoreBinding { py: Python<'_>, node_id: NodeId, ) -> PyResult<(PyTensor, Py)> { - 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 TreeCoreBinding { node_id: NodeId, pass_prefix_keys: bool, ) -> PyResult> { - 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 TreeCoreBinding { 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 TreeCoreBinding { py: Python<'_>, node_id: NodeId, ) -> PyResult<(Option, Option)> { - 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 { - 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 { - 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> { - 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> { - 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> { - 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 TreeCoreBinding { } fn dfs_weight_order(&self, py: Python<'_>, node_ids: Vec) -> PyResult> { - 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 TreeCoreBinding { }) }) .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 TreeCoreBinding { ) -> 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 TreeCoreBinding { 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 TreeCoreBinding { 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 { 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 TreeCoreBinding { py: Python<'_>, tail_node_id: NodeId, ) -> PyResult { - 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 TreeCoreBinding { } /// 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 { + 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 TreeCoreBinding { .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 TreeCoreBinding { 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 TreeCoreBinding { py: Python<'_>, node_id: NodeId, ) -> PyResult { - 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 TreeCoreBinding { py: Python<'_>, node_ids: Vec, ack_id: NodeId, - ) -> Vec { - py.allow_threads(|| self.core().mark_write_through_pending(node_ids, ack_id)) + ) -> PyResult> { + 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, 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, + 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 TreeCoreBinding { component_type: u8, ) -> PyResult { 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 TreeCoreBinding { py.allow_threads(|| self.core().inspect_contains_node(node_id)) } - fn inspect_get_parent_node_id(&self, py: Python<'_>, node_id: NodeId) -> Option { + fn inspect_get_parent_node_id( + &self, + py: Python<'_>, + node_id: NodeId, + ) -> PyResult> { 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 { + fn inspect_get_child_node_ids(&self, py: Python<'_>, node_id: NodeId) -> PyResult> { 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 { 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 { + fn inspect_get_node_token_ids(&self, py: Python<'_>, node_id: NodeId) -> PyResult> { 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 { 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 TreeCoreBinding { component_type: u8, ) -> PyResult> { 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 TreeCoreBinding { component_type: u8, ) -> PyResult { 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 { 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 { + ) -> PyResult> { 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 TreeCoreBinding { component_type: u8, ) -> PyResult { 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 TreeCoreBinding { component_type: u8, ) -> PyResult { 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 TreeCoreBinding { 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 { 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 { @@ -1991,11 +2052,12 @@ impl TreeCoreBinding { py: Python<'_>, node_id: NodeId, hash_values: Option>, - ) { + ) -> 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 TreeCoreBinding { 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 TreeCoreBinding { 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 TreeCoreBinding { 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 TreeCoreBinding { 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 TreeCoreBinding { 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 TreeCoreBinding { 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 TreeCoreBinding { ) -> PyResult { 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 TreeCoreBinding { 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 TreeCoreBinding { py: Python<'_>, node_id: NodeId, ) -> PyResult { - 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 TreeCoreBinding { best_value_len, ) }); + let result = result.map_err(node_access_error)?; MatchResultBinding::from_match_result(py, result) } @@ -2213,11 +2281,12 @@ impl TreeCoreBinding { py: Python<'_>, node_id: NodeId, write_back: bool, - ) -> Vec { + ) -> PyResult> { 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 { 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 { 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 { self.inner.inc_host_lock_ref(py, node_id) } @@ -2810,17 +2883,22 @@ macro_rules! tree_core_binding { py: Python<'_>, node_ids: Vec, ack_id: NodeId, - ) -> Vec { + ) -> PyResult> { 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, ack_id: NodeId) { + fn finish_write_through( + &self, + py: Python<'_>, + node_ids: Vec, + 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 { + ) -> PyResult> { 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 { + ) -> PyResult> { 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 { 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 { + fn inspect_get_node_token_ids( + &self, + py: Python<'_>, + node_id: NodeId, + ) -> PyResult> { 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 { 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 { 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 { + ) -> PyResult> { 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 { 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>, - ) { + ) -> 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 { + ) -> PyResult> { self.inner .inspect_build_backup_node_ids(py, node_id, write_back) } diff --git a/rust/sglang-radix-tree/src/tests/components/base.rs b/rust/sglang-radix-tree/src/tests/components/base.rs index b0700bf6c956..6d9bdfa89be1 100644 --- a/rust/sglang-radix-tree/src/tests/components/base.rs +++ b/rust/sglang-radix-tree/src/tests/components/base.rs @@ -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(), diff --git a/rust/sglang-radix-tree/src/tests/components/full.rs b/rust/sglang-radix-tree/src/tests/components/full.rs index e6017c0f7786..0c224b136c06 100644 --- a/rust/sglang-radix-tree/src/tests/components/full.rs +++ b/rust/sglang-radix-tree/src/tests/components/full.rs @@ -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>) -> (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(¶ms), /* 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> { @@ -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>, NodeIdx_, NodeIdx_) { } fn finalize(tc: &UnifiedTreeCore>, 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() }, + ¶ms, + &[], + 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() + }, + ¶ms, + &[], + 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() + }, + ¶ms, + &[], + 0, + ); + assert!(matches!(result, Err(NodeAccessError { node_id: 997 }))); + + assert!( + tc.inspect_finalize_component_match_result(FULL, tc.empty_match_result(), ¶ms, &[], 0,) + .is_ok() ); } diff --git a/rust/sglang-radix-tree/src/tests/components/mamba.rs b/rust/sglang-radix-tree/src/tests/components/mamba.rs index 2fc8d8a9e765..ef7239877d03 100644 --- a/rust/sglang-radix-tree/src/tests/components/mamba.rs +++ b/rust/sglang-radix-tree/src/tests/components/mamba.rs @@ -44,7 +44,8 @@ fn hybrid_lock_core() -> (UnifiedTreeCore>, 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); diff --git a/rust/sglang-radix-tree/src/tests/components/swa.rs b/rust/sglang-radix-tree/src/tests/components/swa.rs index 7b80f61a620a..7a5ae3cdbb25 100644 --- a/rust/sglang-radix-tree/src/tests/components/swa.rs +++ b/rust/sglang-radix-tree/src/tests/components/swa.rs @@ -172,7 +172,8 @@ fn store_swa_device(tc: &mut UnifiedTreeCore>, 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(¶ms)); + tc.dec_host_lock_ref(tc.arena.node(c).id, Some(¶ms)) + .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(¶ms)); + tc.dec_host_lock_ref(tc.arena.node(c).id, Some(¶ms)) + .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(¶ms), /* 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); } @@ -4311,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); @@ -4326,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) @@ -4377,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 { @@ -4397,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!( @@ -4420,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)); @@ -4459,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]) @@ -4471,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); @@ -4505,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, @@ -4545,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 { @@ -4557,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); @@ -4661,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); @@ -4674,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); } @@ -4721,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(&[], &[]); @@ -4738,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. @@ -4751,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); @@ -4764,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); @@ -4775,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(¶ms), /* skip_swa = */ false); - tc.finish_load_back(anchor); + tc.dec_lock_ref(anchor, Some(¶ms), /* skip_swa = */ false) + .expect("live test node"); + tc.finish_load_back(anchor).expect("live test node"); } tc.sanity_check(&[], &[]); } @@ -4830,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. @@ -4860,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, diff --git a/rust/sglang-radix-tree/src/tests/node.rs b/rust/sglang-radix-tree/src/tests/node.rs index 1fd6dd39ce80..b834cf113001 100644 --- a/rust/sglang-radix-tree/src/tests/node.rs +++ b/rust/sglang-radix-tree/src/tests/node.rs @@ -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::>() { - 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(()) diff --git a/rust/sglang-radix-tree/src/tests/unified_tree_core.rs b/rust/sglang-radix-tree/src/tests/unified_tree_core.rs index 105f700148f7..f2d0f62d03ad 100644 --- a/rust/sglang-radix-tree/src/tests/unified_tree_core.rs +++ b/rust/sglang-radix-tree/src/tests/unified_tree_core.rs @@ -4,7 +4,7 @@ use tch::Tensor; use super::*; use crate::components::{FULL, MAMBA, SWA}; -use crate::node::ValueSlotIdx; +use crate::node::{NodeAccessError, ValueSlotIdx}; use crate::test_utils::{accumulate_step, action_kinds}; fn core() -> UnifiedTreeCore> { @@ -503,7 +503,8 @@ fn locked_anchor_for_dispatch(tc: &mut UnifiedTreeCore>) -> NodeIdx_ { tc.arena .set_device_value(n1, FULL, Tensor::from_slice(&[0i64, 1])); tc.component_state_mut(FULL).evictable_size = 2; - tc.inc_lock_ref(tc.arena.node(n1).id); + tc.inc_lock_ref(tc.arena.node(n1).id) + .expect("live test node"); n1 } @@ -517,7 +518,8 @@ fn dec_lock_ref_skip_swa_skips_the_swa_component() { tc.arena.node(n1).id, /* params = */ None, /* skip_swa = */ true, - ); + ) + .expect("live test node"); assert_eq!(tc.arena.device_lock_ref(n1, FULL), 0); } @@ -539,7 +541,7 @@ fn inc_lock_ref_reaches_every_component() { tc.arena .set_device_value(n1, FULL, Tensor::from_slice(&[0i64, 1])); tc.component_state_mut(FULL).evictable_size = 2; - tc.inc_lock_ref(tc.arena.node(n1).id); + let _ = tc.inc_lock_ref(tc.arena.node(n1).id); } #[test] @@ -548,7 +550,7 @@ fn dec_lock_ref_without_skip_swa_reaches_every_component() { let mut tc = core(); let n1 = locked_anchor_for_dispatch(&mut tc); tc.register_component_(Arc::new(SwaComponentForTest)); - tc.dec_lock_ref( + let _ = tc.dec_lock_ref( tc.arena.node(n1).id, /* params = */ None, /* skip_swa = */ false, @@ -573,7 +575,8 @@ fn set_component_device_value_sizes_by_the_value_length() { tc.arena.node(node).id, SWA, Tensor::from_slice(&[7i64, 8, 9]), - ); + ) + .expect("live test node"); assert!( tc.arena .device_value(node, SWA) @@ -597,9 +600,10 @@ fn set_component_device_value_rejects_an_occupied_slot() { /* extra_key = */ None, ) .unwrap(); - tc.set_component_device_value(tc.arena.node(node).id, SWA, Tensor::from_slice(&[7i64])); + tc.set_component_device_value(tc.arena.node(node).id, SWA, Tensor::from_slice(&[7i64])) + .expect("live test node"); tc.device_lru_list_mut(SWA).remove_node(node); - tc.set_component_device_value(tc.arena.node(node).id, SWA, Tensor::from_slice(&[8i64])); + let _ = tc.set_component_device_value(tc.arena.node(node).id, SWA, Tensor::from_slice(&[8i64])); } #[test] @@ -607,7 +611,7 @@ fn set_component_device_value_rejects_an_occupied_slot() { fn set_component_device_value_rejects_a_disabled_component() { let mut tc = core(); let root = tc.arena.root(); - tc.set_component_device_value(tc.arena.node(root).id, SWA, Tensor::from_slice(&[1i64])); + let _ = tc.set_component_device_value(tc.arena.node(root).id, SWA, Tensor::from_slice(&[1i64])); } #[test] @@ -624,7 +628,7 @@ fn dec_swa_lock_only_dispatches_lower_priority_releases() { let root = tc.arena.root(); let mut device_frees = HashMap::new(); let mut host_frees = HashMap::new(); - tc.dec_swa_lock_only( + let _ = tc.dec_swa_lock_only( tc.arena.node(root).id, Some(7), &mut device_frees, @@ -651,7 +655,8 @@ fn dec_swa_lock_only_returns_device_frees_in_the_device_dict() { .unwrap(); tc.arena .set_device_value(a, FULL, Tensor::from_slice(&[9i64])); - tc.set_component_device_value(tc.arena.node(a).id, SWA, Tensor::from_slice(&[7i64])); + tc.set_component_device_value(tc.arena.node(a).id, SWA, Tensor::from_slice(&[7i64])) + .expect("live test node"); let swa = SwaComponent::new(&CacheInitParams { swa_sliding_window_size: Some(2), ..Default::default() @@ -669,7 +674,8 @@ fn dec_swa_lock_only_returns_device_frees_in_the_device_dict() { result.swa_uuid_for_lock, &mut device_frees, &mut host_frees, - ); + ) + .expect("live test node"); // The fully unlocked D-leaf's SWA value is device-evicted on release; // the freed span is reported as the node's Full indices. assert!(!tc.arena.has_device_value(a, SWA)); @@ -866,8 +872,18 @@ fn new_node_ids_are_distinct_live_slots() { /* extra_key = */ None, ); assert_ne!(a, b); - assert_eq!(tc.arena.resolve(tc.arena.node(a).id), a); - assert_eq!(tc.arena.resolve(tc.arena.node(b).id), b); + assert_eq!( + tc.arena + .resolve(tc.arena.node(a).id) + .expect("live test node"), + a + ); + assert_eq!( + tc.arena + .resolve(tc.arena.node(b).id) + .expect("live test node"), + b + ); } // Chain root -> c with a 3-atom key and FULL device value, seeded as a D-leaf. @@ -1246,9 +1262,17 @@ fn match_prefix_splits_on_a_partial_match() { let prefix_node = result.best_match_node_id; assert_ne!(prefix_node, tc.arena.node(a).id); assert!(result.device_indices.equal(&Tensor::from_slice(&[10i64]))); - assert_eq!(tc.arena.node(tc.arena.resolve(prefix_node)).key, vec![1]); + assert_eq!( + tc.arena + .node(tc.arena.resolve(prefix_node).expect("live test node")) + .key, + vec![1] + ); assert_eq!(tc.arena.node(a).key, vec![2]); - assert_eq!(tc.arena.node(a).parent(), tc.arena.resolve(prefix_node)); + assert_eq!( + tc.arena.node(a).parent(), + tc.arena.resolve(prefix_node).expect("live test node") + ); } #[test] @@ -1762,7 +1786,8 @@ fn repeated_deep_swa_matches_keep_the_tree_sane() { tc.arena.node(node).id, SWA, Tensor::from_slice(&vec![0i64; len]), - ); + ) + .expect("live test node"); } for _ in 0..3 { let result = tc.match_prefix(&match_params(&vec![1, 2, 3, 4])); @@ -1866,7 +1891,7 @@ fn insert_reports_new_and_unevicted_full_ranges() { let leaf = tc .match_prefix(&match_params(&vec![1, 2, 3])) .best_match_node_id; - let leaf = tc.arena.resolve(leaf); + let leaf = tc.arena.resolve(leaf).expect("live test node"); let _ = tc.arena.take_device_value(leaf, FULL); tc.component_state_mut(FULL).evictable_size = 0; tc.evictable_device_leaves.discard(leaf); @@ -2032,20 +2057,26 @@ fn insert_unevicts_a_tombstoned_node() { let a = tc .match_prefix(&match_params(&vec![1, 2])) .best_match_node_id; - let _ = tc.arena.take_device_value(tc.arena.resolve(a), FULL); + let _ = tc + .arena + .take_device_value(tc.arena.resolve(a).expect("live test node"), FULL); tc.component_state_mut(FULL).evictable_size = 0; - tc.evictable_device_leaves.discard(tc.arena.resolve(a)); + tc.evictable_device_leaves + .discard(tc.arena.resolve(a).expect("live test node")); let result = tc.insert(&insert_params(&vec![1, 2], &[20, 21])); assert_eq!(result.prefix_len, 2); // The fresh KV revives the node; nothing is duplicate. assert!(result.cache_actions.is_empty()); assert!( tc.arena - .device_value(tc.arena.resolve(a), FULL) + .device_value(tc.arena.resolve(a).expect("live test node"), FULL) .equal(&Tensor::from_slice(&[20i64, 21])) ); assert_eq!(tc.evictable_size_(FULL), 2); - assert!(tc.evictable_device_leaves.contains(tc.arena.resolve(a))); + assert!( + tc.evictable_device_leaves + .contains(tc.arena.resolve(a).expect("live test node")) + ); } #[test] @@ -2059,7 +2090,12 @@ fn insert_priority_floor_applies_along_the_path() { priority: 5, ..insert_params(&vec![1, 2], &[20, 21]) }); - assert_eq!(tc.arena.node(tc.arena.resolve(a)).priority, 5); + assert_eq!( + tc.arena + .node(tc.arena.resolve(a).expect("live test node")) + .priority, + 5 + ); } #[test] @@ -2069,12 +2105,20 @@ fn insert_chunked_skips_the_hit_count() { let a = tc .match_prefix(&match_params(&vec![1, 2])) .best_match_node_id; - let hits_before = tc.arena.node(tc.arena.resolve(a)).hit_count; + let hits_before = tc + .arena + .node(tc.arena.resolve(a).expect("live test node")) + .hit_count; tc.insert(&InsertParams { chunked: true, ..insert_params(&vec![1, 2], &[20, 21]) }); - assert_eq!(tc.arena.node(tc.arena.resolve(a)).hit_count, hits_before); + assert_eq!( + tc.arena + .node(tc.arena.resolve(a).expect("live test node")) + .hit_count, + hits_before + ); } #[test] @@ -2084,10 +2128,15 @@ fn insert_extension_bumps_the_traversed_node_hit_count_once() { let a = tc .match_prefix(&match_params(&vec![1, 2, 3])) .best_match_node_id; - let hits_before = tc.arena.node(tc.arena.resolve(a)).hit_count; + let hits_before = tc + .arena + .node(tc.arena.resolve(a).expect("live test node")) + .hit_count; tc.insert(&insert_params(&vec![1, 2, 3, 4, 5], &[20, 21, 22, 13, 14])); assert_eq!( - tc.arena.node(tc.arena.resolve(a)).hit_count, + tc.arena + .node(tc.arena.resolve(a).expect("live test node")) + .hit_count, hits_before + 1 ); } @@ -2099,11 +2148,16 @@ fn insert_full_overlap_bumps_the_hit_count_once() { let a = tc .match_prefix(&match_params(&vec![1, 2, 3])) .best_match_node_id; - let hits_before = tc.arena.node(tc.arena.resolve(a)).hit_count; + let hits_before = tc + .arena + .node(tc.arena.resolve(a).expect("live test node")) + .hit_count; // The walk already counted the full overlap; the target is no new leaf. tc.insert(&insert_params(&vec![1, 2, 3], &[20, 21, 22])); assert_eq!( - tc.arena.node(tc.arena.resolve(a)).hit_count, + tc.arena + .node(tc.arena.resolve(a).expect("live test node")) + .hit_count, hits_before + 1 ); } @@ -2136,10 +2190,11 @@ fn mark_write_through_pending_stamps_the_supplied_ack() { let mut tc = core(); tc.insert(&insert_params(&vec![1], &[10])); let leaf = tc.match_prefix(&match_params(&vec![1])).best_match_node_id; - tc.mark_write_through_pending(vec![leaf], /* ack_id = */ leaf); + tc.mark_write_through_pending(vec![leaf], /* ack_id = */ leaf) + .expect("live test node"); assert_eq!( tc.arena - .node(tc.arena.resolve(leaf)) + .node(tc.arena.resolve(leaf).expect("live test node")) .write_through_pending_id, Some(leaf) ); @@ -2155,7 +2210,9 @@ fn mark_write_through_pending_stamps_one_ack_on_every_published_node() { .match_prefix(&match_params(&vec![1, 2])) .best_match_node_id; - let published = tc.mark_write_through_pending(vec![parent, leaf], /* ack_id = */ leaf); + let published = tc + .mark_write_through_pending(vec![parent, leaf], /* ack_id = */ leaf) + .expect("live test nodes"); assert_eq!(published, vec![parent, leaf]); for node_id in [parent, leaf] { @@ -2188,7 +2245,9 @@ fn mark_write_through_pending_returns_the_published_nodes_ancestors_first() { .best_match_node_id; // The caller merges per-component transfers, whose order is not tree order. - let published = tc.mark_write_through_pending(vec![leaf, parent], /* ack_id = */ leaf); + let published = tc + .mark_write_through_pending(vec![leaf, parent], /* ack_id = */ leaf) + .expect("live test nodes"); assert_eq!(published, vec![parent, leaf]); } @@ -2198,18 +2257,21 @@ fn finish_write_through_clears_only_the_matching_ack() { let mut tc = core(); tc.insert(&insert_params(&vec![1], &[10])); let leaf = tc.match_prefix(&match_params(&vec![1])).best_match_node_id; - tc.mark_write_through_pending(vec![leaf], /* ack_id = */ leaf); - tc.finish_write_through(vec![leaf], /* ack_id = */ 999_999); + tc.mark_write_through_pending(vec![leaf], /* ack_id = */ leaf) + .expect("live test node"); + tc.finish_write_through(vec![leaf], /* ack_id = */ 999_999) + .expect("live test node"); assert_eq!( tc.arena - .node(tc.arena.resolve(leaf)) + .node(tc.arena.resolve(leaf).expect("live test node")) .write_through_pending_id, Some(leaf) ); - tc.finish_write_through(vec![leaf], /* ack_id = */ leaf); + tc.finish_write_through(vec![leaf], /* ack_id = */ leaf) + .expect("live test node"); assert_eq!( tc.arena - .node(tc.arena.resolve(leaf)) + .node(tc.arena.resolve(leaf).expect("live test node")) .write_through_pending_id, None ); @@ -2229,15 +2291,18 @@ fn backup_kv_action_chains_unbacked_ancestors_first() { .match_prefix(&match_params(&vec![1, 2, 3])) .best_match_node_id; // a is backuped: the chain stops there and orders ancestors first. - tc.arena - .set_host_value(tc.arena.resolve(a), FULL, Tensor::from_slice(&[20i64])); + tc.arena.set_host_value( + tc.arena.resolve(a).expect("live test node"), + FULL, + Tensor::from_slice(&[20i64]), + ); let action = tc.build_backup_kv_action_( - tc.arena.node(tc.arena.resolve(c)), + tc.arena.node(tc.arena.resolve(c).expect("live test node")), /* write_back = */ false, ); assert_eq!(action.node_ids, vec![b, c]); let action = tc.build_backup_kv_action_( - tc.arena.node(tc.arena.resolve(c)), + tc.arena.node(tc.arena.resolve(c).expect("live test node")), /* write_back = */ true, ); assert_eq!(action.node_ids, vec![c]); @@ -2250,12 +2315,16 @@ fn split_of_a_pending_node_transfers_the_ack_and_emits_the_replace_action() { let node = tc .match_prefix(&match_params(&vec![1, 2, 3])) .best_match_node_id; - tc.mark_write_through_pending(vec![node], /* ack_id = */ node); - let (new_node, action) = tc.split_node_(tc.arena.resolve(node), /* split_len = */ 1); + tc.mark_write_through_pending(vec![node], /* ack_id = */ node) + .expect("live test node"); + let (new_node, action) = tc.split_node_( + tc.arena.resolve(node).expect("live test node"), + /* split_len = */ 1, + ); assert_eq!(tc.arena.node(new_node).write_through_pending_id, Some(node)); assert_eq!( tc.arena - .node(tc.arena.resolve(node)) + .node(tc.arena.resolve(node).expect("live test node")) .write_through_pending_id, Some(node) ); @@ -2288,7 +2357,12 @@ fn insert_does_not_hash_without_storage() { let leaf = tc .match_prefix(&match_params(&vec![1, 2])) .best_match_node_id; - assert_eq!(tc.arena.node(tc.arena.resolve(leaf)).hash_value, None); + assert_eq!( + tc.arena + .node(tc.arena.resolve(leaf).expect("live test node")) + .hash_value, + None + ); } #[test] @@ -2308,7 +2382,9 @@ fn insert_hashes_pages_chained_from_the_parent_when_storage_is_on() { .match_prefix(&match_params(&vec![1, 2])) .best_match_node_id; assert_eq!( - tc.arena.node(tc.arena.resolve(parent)).hash_value, + tc.arena + .node(tc.arena.resolve(parent).expect("live test node")) + .hash_value, Some(vec![ "34fb5c825de7ca4aea6e712f19d439c1da0c92c37b423936c5f618545ca4fa1f".to_string() ]) @@ -2317,21 +2393,25 @@ fn insert_hashes_pages_chained_from_the_parent_when_storage_is_on() { .match_prefix(&match_params(&vec![1, 2, 7, 8])) .best_match_node_id; assert_eq!( - tc.arena.node(tc.arena.resolve(child)).hash_value, + tc.arena + .node(tc.arena.resolve(child).expect("live test node")) + .hash_value, Some(vec![ "0bfa9b9c6fd727c7410b6d42b753439911022d34cc6ef99ac43ed7724aa48a75".to_string() ]) ); // The prefix walk concatenates the chain in root-to-node order. assert_eq!( - tc.arena.prefix_hash_values(Some(tc.arena.resolve(child))), + tc.arena + .prefix_hash_values(Some(tc.arena.resolve(child).expect("live test node"))), vec![ "34fb5c825de7ca4aea6e712f19d439c1da0c92c37b423936c5f618545ca4fa1f".to_string(), "0bfa9b9c6fd727c7410b6d42b753439911022d34cc6ef99ac43ed7724aa48a75".to_string(), ] ); assert_eq!( - tc.arena.prefix_hash_values(Some(tc.arena.resolve(parent))), + tc.arena + .prefix_hash_values(Some(tc.arena.resolve(parent).expect("live test node"))), vec!["34fb5c825de7ca4aea6e712f19d439c1da0c92c37b423936c5f618545ca4fa1f".to_string()] ); } @@ -2378,7 +2458,9 @@ fn insert_coalesces_parent_linked_block_stores() { .match_prefix(&match_params(&vec![1, 2, 7, 8])) .best_match_node_id; assert_eq!( - tc.arena.node(tc.arena.resolve(leaf)).hash_value, + tc.arena + .node(tc.arena.resolve(leaf).expect("live test node")) + .hash_value, Some(hashes) ); assert!(tc.salted_event_hashes.is_empty()); @@ -2399,7 +2481,7 @@ fn salted_event_hashes_are_sparse_and_removed_with_the_node() { let leaf = tc .match_prefix(&match_params_in_namespace(&key, None, Some("tenant-a"))) .best_match_node_id; - let leaf_idx = tc.arena.resolve(leaf); + let leaf_idx = tc.arena.resolve(leaf).expect("live test node"); assert_eq!(tc.salted_event_hashes[&leaf].len(), 2); assert_eq!( tc.arena.node(leaf_idx).hash_value, @@ -2411,7 +2493,9 @@ fn salted_event_hashes_are_sparse_and_removed_with_the_node() { tc.evict_device_start(FULL, key.len()); let (candidate, step) = tc.evict_device_next_node(FULL, &tracker); accumulate_step(step, &mut tracker, &mut device_frees, &mut host_frees); - let (_, step) = tc.evict_device_leaf(candidate.unwrap(), false); + let (_, step) = tc + .evict_device_leaf(candidate.unwrap(), false) + .expect("live test node"); accumulate_step(step, &mut tracker, &mut device_frees, &mut host_frees); tc.evict_device_end(FULL); tc.take_events(); @@ -2464,7 +2548,10 @@ fn salted_event_hashes_survive_node_split() { Some("tenant-a"), )) .best_match_node_id; - let split_parent_idx = tc.arena.node(tc.arena.resolve(split_child)).parent(); + let split_parent_idx = tc + .arena + .node(tc.arena.resolve(split_child).expect("live test node")) + .parent(); let split_parent = tc.arena.node(split_parent_idx).id; assert_eq!(tc.salted_event_hashes[&split_parent], original_hashes[..1]); assert_eq!(tc.salted_event_hashes[&split_child], original_hashes[1..]); @@ -2596,7 +2683,9 @@ fn eviction_emits_block_removed_with_all_page_hashes() { let (node, step) = tc.evict_device_next_node(FULL, &tracker); accumulate_step(step, &mut tracker, &mut device_frees, &mut host_frees); let Some(node) = node else { break }; - let (_, step) = tc.evict_device_leaf(node, /* is_write_back = */ false); + let (_, step) = tc + .evict_device_leaf(node, /* is_write_back = */ false) + .expect("live test node"); accumulate_step(step, &mut tracker, &mut device_frees, &mut host_frees); } tc.evict_device_end(FULL); @@ -2658,7 +2747,8 @@ fn finish_write_through_emits_cpu_stored_events() { tc.insert(&insert_params(&vec![1], &[10])); let leaf = tc.match_prefix(&match_params(&vec![1])).best_match_node_id; let _ = tc.take_events(); - tc.finish_write_through(vec![leaf], leaf); + tc.finish_write_through(vec![leaf], leaf) + .expect("live test node"); let hashes = crate::node::get_hash_str::>(&[1], None, 1); assert_eq!( tc.take_events(), @@ -2678,10 +2768,11 @@ fn demoted_events_leaf(tc: &mut UnifiedTreeCore>) -> NodeIdx_ { tc.set_hicache_enabled(); tc.insert(&insert_params(&vec![1], &[10])); let leaf = tc.match_prefix(&match_params(&vec![1])).best_match_node_id; - tc.commit_backup(leaf, Tensor::from_slice(&[100i64]), HashMap::new()); - tc.demote(leaf); + tc.commit_backup(leaf, Tensor::from_slice(&[100i64]), HashMap::new()) + .expect("live test node"); + tc.demote(leaf).expect("valid demote"); let _ = tc.take_events(); - tc.arena.resolve(leaf) + tc.arena.resolve(leaf).expect("live test node") } #[test] @@ -2703,13 +2794,16 @@ fn host_eviction_emits_a_cpu_block_removed() { fn load_back_commit_emits_gpu_stored_events() { let mut tc = events_core(1); let leaf = demoted_events_leaf(&mut tc); - let (kv_xfer, comp_xfers) = tc.build_load_back_spec(tc.arena.node(leaf).id, None); + let (kv_xfer, comp_xfers) = tc + .build_load_back_spec(tc.arena.node(leaf).id, None) + .expect("live test node"); tc.commit_load_back( tc.arena.node(leaf).id, Tensor::from_slice(&[50i64]), kv_xfer, comp_xfers, - ); + ) + .expect("live test node"); let hashes = crate::node::get_hash_str::>(&[1], None, 1); assert_eq!( tc.take_events(), @@ -2761,7 +2855,7 @@ fn drop_subtree_emits_removals_for_host_descendants_then_the_leaf() { let child = tc .arena .alloc_child( - tc.arena.resolve(leaf), + tc.arena.resolve(leaf).expect("live test node"), /* key = */ vec![3, 4], /* priority = */ 0, /* extra_key = */ None, @@ -2770,9 +2864,9 @@ fn drop_subtree_emits_removals_for_host_descendants_then_the_leaf() { tc.arena .set_host_value(child, FULL, Tensor::from_slice(&[20i64, 21])); tc.update_evictable_leaf_sets_(child); - tc.update_evictable_leaf_sets_(tc.arena.resolve(leaf)); + tc.update_evictable_leaf_sets_(tc.arena.resolve(leaf).expect("live test node")); let _ = tc.take_events(); - let (dropped, _step) = tc.drop_subtree_no_host(leaf); + let (dropped, _step) = tc.drop_subtree_no_host(leaf).expect("live test node"); assert!(dropped); // The leaf hashed lazily at its insert store event; the host-only // child hashes lazily at removal, chaining from the leaf. @@ -2834,14 +2928,18 @@ fn split_insert_stores_only_the_new_block_chained_to_the_split_parent() { .match_prefix(&match_params(&vec![1, 2])) .best_match_node_id; assert_eq!( - tc.arena.node(tc.arena.resolve(parent)).hash_value, + tc.arena + .node(tc.arena.resolve(parent).expect("live test node")) + .hash_value, Some(vec![base_hashes[0].clone()]) ); let child = tc .match_prefix(&match_params(&vec![1, 2, 3, 4])) .best_match_node_id; assert_eq!( - tc.arena.node(tc.arena.resolve(child)).hash_value, + tc.arena + .node(tc.arena.resolve(child).expect("live test node")) + .hash_value, Some(vec![base_hashes[1].clone()]) ); } @@ -2854,7 +2952,8 @@ fn finish_write_through_after_a_split_publishes_both_fragments() { let leaf = tc .match_prefix(&match_params(&vec![1, 2, 3, 4])) .best_match_node_id; - 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 _ = tc.take_events(); let result = tc.insert(&insert_params(&vec![1, 2, 5, 6], &[20, 21, 22, 23])); let new_node_id = result @@ -2884,9 +2983,12 @@ fn finish_write_through_after_a_split_publishes_both_fragments() { new_node_id, Tensor::from_slice(&[100i64, 101]), HashMap::new(), - ); - tc.commit_backup(leaf, Tensor::from_slice(&[102i64, 103]), HashMap::new()); - tc.finish_write_through(vec![new_node_id, leaf], /* ack_id = */ leaf); + ) + .expect("live test node"); + tc.commit_backup(leaf, Tensor::from_slice(&[102i64, 103]), HashMap::new()) + .expect("live test node"); + tc.finish_write_through(vec![new_node_id, leaf], /* ack_id = */ leaf) + .expect("live test nodes"); let hashes = crate::node::get_hash_str::>(&[1, 2, 3, 4], None, 2); assert_eq!( tc.take_events(), @@ -2905,13 +3007,13 @@ fn finish_write_through_after_a_split_publishes_both_fragments() { // The matching ack cleared the pending mark on both fragments. assert_eq!( tc.arena - .node(tc.arena.resolve(new_node_id)) + .node(tc.arena.resolve(new_node_id).expect("live test node")) .write_through_pending_id, None ); assert_eq!( tc.arena - .node(tc.arena.resolve(leaf)) + .node(tc.arena.resolve(leaf).expect("live test node")) .write_through_pending_id, None ); @@ -2928,7 +3030,10 @@ fn prefetch_anchor_info_maps_the_namespace() { let plain = tc .match_prefix(&match_params(&vec![1, 2])) .best_match_node_id; - assert_eq!(tc.prefetch_anchor_info(plain), (None, None)); + assert_eq!( + tc.prefetch_anchor_info(plain).expect("live test node"), + (None, None) + ); let salted = tc .match_prefix(&MatchPrefixParams { key: &vec![7, 8], @@ -2936,13 +3041,14 @@ fn prefetch_anchor_info_maps_the_namespace() { }) .best_match_node_id; assert_eq!( - tc.prefetch_anchor_info(salted), + tc.prefetch_anchor_info(salted).expect("live test node"), (Some("chat".to_string()), Some("tenant-a".to_string())) ); // A root anchor carries no namespace: the single root serves them all. let root = tc.arena.root(); assert_eq!( - tc.prefetch_anchor_info(tc.arena.node(root).id), + tc.prefetch_anchor_info(tc.arena.node(root).id) + .expect("live test node"), (None, None) ); // A node minted by a split inherits the namespace. @@ -2958,7 +3064,7 @@ fn prefetch_anchor_info_maps_the_namespace() { .best_match_node_id; assert_ne!(split_mid, salted); assert_eq!( - tc.prefetch_anchor_info(split_mid), + tc.prefetch_anchor_info(split_mid).expect("live test node"), (Some("chat".to_string()), Some("tenant-a".to_string())) ); } @@ -3036,26 +3142,34 @@ fn prefetch_node_accessors_cover_gate_and_hash_chain() { .match_prefix(&match_params(&vec![1, 2, 7, 8])) .best_match_node_id; - assert!(!tc.node_backuped(child)); - assert!(!tc.is_root(child)); + assert!(!tc.node_backuped(child).expect("live test node")); + assert!(!tc.is_root(child).expect("live test node")); assert_eq!( - tc.get_last_hash_value(child).as_deref(), + tc.get_last_hash_value(child) + .expect("live test node") + .as_deref(), Some("0bfa9b9c6fd727c7410b6d42b753439911022d34cc6ef99ac43ed7724aa48a75") ); assert_eq!( - tc.get_prefix_hash_values(child), + tc.get_prefix_hash_values(child).expect("live test node"), vec!["34fb5c825de7ca4aea6e712f19d439c1da0c92c37b423936c5f618545ca4fa1f".to_string()] ); - tc.commit_backup(child, Tensor::from_slice(&[102i64, 103]), HashMap::new()); - assert!(tc.node_backuped(child)); + tc.commit_backup(child, Tensor::from_slice(&[102i64, 103]), HashMap::new()) + .expect("live test node"); + assert!(tc.node_backuped(child).expect("live test node")); // Roots have no hashes of their own. let root = tc.arena.root(); - assert!(tc.is_root(tc.arena.node(root).id)); - assert_eq!(tc.get_last_hash_value(tc.arena.node(root).id), None); + assert!(tc.is_root(tc.arena.node(root).id).expect("live test node")); assert_eq!( - tc.get_prefix_hash_values(tc.arena.node(root).id), + tc.get_last_hash_value(tc.arena.node(root).id) + .expect("live test node"), + None + ); + assert_eq!( + tc.get_prefix_hash_values(tc.arena.node(root).id) + .expect("live test node"), Vec::::new() ); } @@ -3068,8 +3182,14 @@ fn storage_backup_spec_is_none_for_an_unbackuped_node() { .match_prefix(&match_params(&vec![1, 2])) .best_match_node_id; assert!( - tc.build_storage_backup_spec(tc.arena.node(tc.arena.resolve(leaf)).id, true) - .is_none() + tc.build_storage_backup_spec( + tc.arena + .node(tc.arena.resolve(leaf).expect("live test node")) + .id, + true + ) + .expect("live test node") + .is_none() ); } @@ -3092,12 +3212,20 @@ fn storage_backup_spec_gathers_the_chained_node() { let child = tc .match_prefix(&match_params(&vec![1, 2, 7, 8])) .best_match_node_id; - tc.commit_backup(parent, Tensor::from_slice(&[100i64, 101]), HashMap::new()); - tc.commit_backup(child, Tensor::from_slice(&[102i64, 103]), HashMap::new()); + tc.commit_backup(parent, Tensor::from_slice(&[100i64, 101]), HashMap::new()) + .expect("live test node"); + tc.commit_backup(child, Tensor::from_slice(&[102i64, 103]), HashMap::new()) + .expect("live test node"); let spec = tc - .build_storage_backup_spec(tc.arena.node(tc.arena.resolve(child)).id, true) - .unwrap(); + .build_storage_backup_spec( + tc.arena + .node(tc.arena.resolve(child).expect("live test node")) + .id, + true, + ) + .expect("live test node") + .expect("backuped node"); assert!(spec.host_value.equal(&Tensor::from_slice(&[102i64, 103]))); assert_eq!(spec.token_ids, vec![7, 8]); assert_eq!( @@ -3115,8 +3243,14 @@ fn storage_backup_spec_gathers_the_chained_node() { assert!(spec.comp_xfers.is_empty()); let spec = tc - .build_storage_backup_spec(tc.arena.node(tc.arena.resolve(child)).id, false) - .unwrap(); + .build_storage_backup_spec( + tc.arena + .node(tc.arena.resolve(child).expect("live test node")) + .id, + false, + ) + .expect("live test node") + .expect("backuped node"); assert_eq!(spec.prefix_keys, None); } @@ -3134,18 +3268,22 @@ fn prefix_hash_walk_stops_below_an_unhashed_ancestor() { fn insert_host_attaches_a_host_only_leaf_under_the_root() { let mut tc = core(); let root = tc.arena.root(); - let result = tc.insert_host( - tc.arena.node(root).id, - /* extra_key = */ None, - vec![1, 2], - Tensor::from_slice(&[100i64, 101]), - vec!["h0".to_string(), "h1".to_string()], - ); + let result = tc + .insert_host( + tc.arena.node(root).id, + /* extra_key = */ None, + vec![1, 2], + Tensor::from_slice(&[100i64, 101]), + vec!["h0".to_string(), "h1".to_string()], + ) + .expect("live test node"); assert_eq!(result.prefix_len, 0); assert_eq!(result.total_len, 2); assert!(!result.host_insert_dropped); let new_node = result.inserted_host_node.unwrap(); - let node = tc.arena.node(tc.arena.resolve(new_node)); + let node = tc + .arena + .node(tc.arena.resolve(new_node).expect("live test node")); assert!(node.evicted() && node.backuped()); assert!( node.host_value(FULL) @@ -3157,7 +3295,7 @@ fn insert_host_attaches_a_host_only_leaf_under_the_root() { ); assert!( tc.evictable_host_leaves - .contains(tc.arena.resolve(new_node)) + .contains(tc.arena.resolve(new_node).expect("live test node")) ); tc.sanity_check(&[], &[]); } @@ -3168,22 +3306,26 @@ fn insert_host_allows_a_suffix_under_an_unbacked_write_back_parent() { tc.is_write_back = true; tc.insert(&insert_params(&vec![1, 2], &[10, 11])); let root = tc.arena.root(); - let result = tc.insert_host( - tc.arena.node(root).id, - /* extra_key = */ None, - vec![1, 2, 3, 4], - Tensor::from_slice(&[100i64, 101, 102, 103]), - vec!["h0", "h1", "h2", "h3"] - .into_iter() - .map(String::from) - .collect(), - ); + let result = tc + .insert_host( + tc.arena.node(root).id, + /* extra_key = */ None, + vec![1, 2, 3, 4], + Tensor::from_slice(&[100i64, 101, 102, 103]), + vec!["h0", "h1", "h2", "h3"] + .into_iter() + .map(String::from) + .collect(), + ) + .expect("live test node"); assert_eq!(result.prefix_len, 2); assert_eq!(result.total_len, 4); assert!(!result.host_insert_dropped); - let new_node = tc - .arena - .node(tc.arena.resolve(result.inserted_host_node.unwrap())); + let new_node = tc.arena.node( + tc.arena + .resolve(result.inserted_host_node.unwrap()) + .expect("live test node"), + ); assert!( new_node .host_value(FULL) @@ -3201,16 +3343,18 @@ fn insert_host_drops_a_suffix_under_an_unbacked_write_through_parent() { tc.insert(&insert_params(&vec![1, 2], &[10, 11])); let root = tc.arena.root(); let nodes_before = tc.arena.len(); - let result = tc.insert_host( - tc.arena.node(root).id, - /* extra_key = */ None, - vec![1, 2, 3, 4], - Tensor::from_slice(&[100i64, 101, 102, 103]), - vec!["h0", "h1", "h2", "h3"] - .into_iter() - .map(String::from) - .collect(), - ); + let result = tc + .insert_host( + tc.arena.node(root).id, + /* extra_key = */ None, + vec![1, 2, 3, 4], + Tensor::from_slice(&[100i64, 101, 102, 103]), + vec!["h0", "h1", "h2", "h3"] + .into_iter() + .map(String::from) + .collect(), + ) + .expect("live test node"); assert_eq!(result.prefix_len, 2); assert_eq!(result.total_len, 4); @@ -3227,15 +3371,18 @@ fn insert_host_drop_preserves_split_actions_and_lengths() { let leaf = tc .match_prefix(&match_params(&vec![1, 2, 3])) .best_match_node_id; - 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 root = tc.arena.root(); - let result = tc.insert_host( - tc.arena.node(root).id, - /* extra_key = */ None, - vec![1, 9], - Tensor::from_slice(&[100i64, 101]), - vec!["h0".to_string(), "h1".to_string()], - ); + let result = tc + .insert_host( + tc.arena.node(root).id, + /* extra_key = */ None, + vec![1, 9], + Tensor::from_slice(&[100i64, 101]), + vec!["h0".to_string(), "h1".to_string()], + ) + .expect("live test node"); assert_eq!(result.prefix_len, 1); assert_eq!(result.total_len, 2); @@ -3260,18 +3407,23 @@ fn insert_host_splits_a_host_chain_and_divides_the_hash() { .into_iter() .map(String::from) .collect(), - ); - let result = tc.insert_host( - tc.arena.node(root).id, - /* extra_key = */ None, - vec![1, 9], - Tensor::from_slice(&[200i64, 201]), - vec!["g0".to_string(), "g1".to_string()], - ); + ) + .expect("live test node"); + let result = tc + .insert_host( + tc.arena.node(root).id, + /* extra_key = */ None, + vec![1, 9], + Tensor::from_slice(&[200i64, 201]), + vec!["g0".to_string(), "g1".to_string()], + ) + .expect("live test node"); assert_eq!(result.prefix_len, 1); - let new_node = tc - .arena - .node(tc.arena.resolve(result.inserted_host_node.unwrap())); + let new_node = tc.arena.node( + tc.arena + .resolve(result.inserted_host_node.unwrap()) + .expect("live test node"), + ); assert!( new_node .host_value(FULL) @@ -3304,18 +3456,23 @@ fn insert_host_hash_slices_by_pages_not_atoms() { vec![1, 2], Tensor::from_slice(&[100i64, 101]), vec!["h0".to_string()], - ); - let result = tc.insert_host( - tc.arena.node(root).id, - /* extra_key = */ None, - vec![1, 2, 3, 4], - Tensor::from_slice(&[200i64, 201, 202, 203]), - vec!["g0".to_string(), "g1".to_string()], - ); + ) + .expect("live test node"); + let result = tc + .insert_host( + tc.arena.node(root).id, + /* extra_key = */ None, + vec![1, 2, 3, 4], + Tensor::from_slice(&[200i64, 201, 202, 203]), + vec!["g0".to_string(), "g1".to_string()], + ) + .expect("live test node"); assert_eq!(result.prefix_len, 2); - let new_node = tc - .arena - .node(tc.arena.resolve(result.inserted_host_node.unwrap())); + let new_node = tc.arena.node( + tc.arena + .resolve(result.inserted_host_node.unwrap()) + .expect("live test node"), + ); // Two matched atoms are ONE page: only g0 is consumed. assert_eq!(new_node.hash_value, Some(vec!["g1".to_string()])); assert!( @@ -3334,29 +3491,33 @@ fn insert_host_full_match_reports_only_a_backuped_node() { .best_match_node_id; let root = tc.arena.root(); // The device-only match reports no host node. - let result = tc.insert_host( - tc.arena.node(root).id, - /* extra_key = */ None, - vec![1, 2], - Tensor::from_slice(&[100i64, 101]), - vec!["h0".to_string(), "h1".to_string()], - ); + let result = tc + .insert_host( + tc.arena.node(root).id, + /* extra_key = */ None, + vec![1, 2], + Tensor::from_slice(&[100i64, 101]), + vec!["h0".to_string(), "h1".to_string()], + ) + .expect("live test node"); assert_eq!(result.prefix_len, 2); assert_eq!(result.inserted_host_node, None); assert!(!result.host_insert_dropped); // Once backuped, the same insert reports the node. tc.arena.set_host_value( - tc.arena.resolve(leaf), + tc.arena.resolve(leaf).expect("live test node"), FULL, Tensor::from_slice(&[20i64, 21]), ); - let result = tc.insert_host( - tc.arena.node(root).id, - /* extra_key = */ None, - vec![1, 2], - Tensor::from_slice(&[100i64, 101]), - vec!["h0".to_string(), "h1".to_string()], - ); + let result = tc + .insert_host( + tc.arena.node(root).id, + /* extra_key = */ None, + vec![1, 2], + Tensor::from_slice(&[100i64, 101]), + vec!["h0".to_string(), "h1".to_string()], + ) + .expect("live test node"); assert_eq!(result.inserted_host_node, Some(leaf)); assert!(!result.host_insert_dropped); } @@ -3379,20 +3540,23 @@ fn insert_host_panics_on_a_colliding_page() { vec![1, 9], Tensor::from_slice(&[100i64, 101]), vec!["h0".to_string(), "h1".to_string()], - ); + ) + .expect("live test node"); } #[test] fn insert_host_empty_key_is_a_noop() { let mut tc = core(); let root = tc.arena.root(); - let result = tc.insert_host( - tc.arena.node(root).id, - /* extra_key = */ None, - vec![], - Tensor::from_slice(&[0i64; 0]), - vec![], - ); + let result = tc + .insert_host( + tc.arena.node(root).id, + /* extra_key = */ None, + vec![], + Tensor::from_slice(&[0i64; 0]), + vec![], + ) + .expect("live test node"); assert_eq!(result.prefix_len, 0); assert!(result.mamba_exist); assert_eq!(result.inserted_host_node, None); @@ -3407,8 +3571,11 @@ fn commit_backup_attaches_the_host_value() { let leaf = tc .match_prefix(&match_params(&vec![1, 2])) .best_match_node_id; - tc.commit_backup(leaf, Tensor::from_slice(&[100i64, 101]), HashMap::new()); - let node = tc.arena.node(tc.arena.resolve(leaf)); + tc.commit_backup(leaf, Tensor::from_slice(&[100i64, 101]), HashMap::new()) + .expect("live test node"); + let node = tc + .arena + .node(tc.arena.resolve(leaf).expect("live test node")); assert!(node.backuped()); assert!( node.host_value(FULL) @@ -3423,7 +3590,7 @@ fn build_backup_spec_reads_the_device_value() { let leaf = tc .match_prefix(&match_params(&vec![1, 2])) .best_match_node_id; - let (device_value, comp_xfers) = tc.build_backup_spec(leaf); + let (device_value, comp_xfers) = tc.build_backup_spec(leaf).expect("live test node"); assert!(device_value.equal(&Tensor::from_slice(&[10i64, 11]))); assert!(comp_xfers.is_empty()); } @@ -3436,10 +3603,15 @@ fn build_backup_spec_skips_full_kv_for_an_already_backuped_node() { .match_prefix(&match_params(&vec![1, 2])) .best_match_node_id; - tc.commit_backup(leaf, Tensor::from_slice(&[100i64, 101]), HashMap::new()); - assert!(tc.arena.node(tc.arena.resolve(leaf)).backuped()); + tc.commit_backup(leaf, Tensor::from_slice(&[100i64, 101]), HashMap::new()) + .expect("live test node"); + assert!( + tc.arena + .node(tc.arena.resolve(leaf).expect("live test node")) + .backuped() + ); - let (device_value, comp_xfers) = tc.build_backup_spec(leaf); + let (device_value, comp_xfers) = tc.build_backup_spec(leaf).expect("live test node"); assert_eq!(device_value.numel(), 0); assert!(comp_xfers.is_empty()); } @@ -3452,10 +3624,14 @@ fn commit_backup_preserves_full_kv_when_host_indices_are_empty() { .match_prefix(&match_params(&vec![1, 2])) .best_match_node_id; - tc.commit_backup(leaf, Tensor::from_slice(&[100i64, 101]), HashMap::new()); - tc.commit_backup(leaf, Tensor::from_slice(&[] as &[i64]), HashMap::new()); + tc.commit_backup(leaf, Tensor::from_slice(&[100i64, 101]), HashMap::new()) + .expect("live test node"); + tc.commit_backup(leaf, Tensor::from_slice(&[] as &[i64]), HashMap::new()) + .expect("live test node"); - let node = tc.arena.node(tc.arena.resolve(leaf)); + let node = tc + .arena + .node(tc.arena.resolve(leaf).expect("live test node")); assert!( node.host_value(FULL) .equal(&Tensor::from_slice(&[100i64, 101])) @@ -3472,14 +3648,20 @@ fn backuped_chain(tc: &mut UnifiedTreeCore>) -> (NodeIdx_, NodeIdx_) { let child = tc .match_prefix(&match_params(&vec![1, 2, 3, 4])) .best_match_node_id; - tc.commit_backup(parent, Tensor::from_slice(&[20i64, 21]), HashMap::new()); - tc.commit_backup(child, Tensor::from_slice(&[22i64, 23]), HashMap::new()); - (tc.arena.resolve(parent), tc.arena.resolve(child)) + tc.commit_backup(parent, Tensor::from_slice(&[20i64, 21]), HashMap::new()) + .expect("live test node"); + tc.commit_backup(child, Tensor::from_slice(&[22i64, 23]), HashMap::new()) + .expect("live test node"); + ( + tc.arena.resolve(parent).expect("live test node"), + tc.arena.resolve(child).expect("live test node"), + ) } // Demote `node_id` (device release of a backuped node), discarding the frees. fn demote_node(tc: &mut UnifiedTreeCore>, node_id: NodeIdx_) { - tc.demote(tc.arena.node(node_id).id); + tc.demote(tc.arena.node(node_id).id) + .expect("backuped live test node"); } #[test] @@ -3499,8 +3681,9 @@ fn build_load_back_spec_collects_the_evicted_chain_ancestors_first() { let (parent, child) = backuped_chain(&mut tc); demote_node(&mut tc, child); demote_node(&mut tc, parent); - let (kv_xfer, comp_xfers) = - tc.build_load_back_spec(tc.arena.node(child).id, /* req = */ None); + let (kv_xfer, comp_xfers) = tc + .build_load_back_spec(tc.arena.node(child).id, /* req = */ None) + .expect("live test node"); assert_eq!(kv_xfer.name, PoolName::Kv); assert!( kv_xfer @@ -3520,8 +3703,9 @@ fn build_load_back_spec_collects_the_evicted_chain_ancestors_first() { fn build_load_back_spec_returns_an_empty_transfer_for_a_device_backed_node() { let mut tc = core(); let (_parent, child) = backuped_chain(&mut tc); - let (kv_xfer, comp_xfers) = - tc.build_load_back_spec(tc.arena.node(child).id, /* req = */ None); + let (kv_xfer, comp_xfers) = tc + .build_load_back_spec(tc.arena.node(child).id, /* req = */ None) + .expect("live test node"); let host_indices = kv_xfer.host_indices.unwrap(); assert_eq!(host_indices.numel(), 0); assert_eq!(host_indices.kind(), Kind::Int64); @@ -3539,14 +3723,17 @@ fn commit_load_back_reattaches_device_slices_and_restores_the_match() { // after demotion. Remove them so this test observes the ack-time refresh. tc.full_coexisting_host_nodes.discard(parent); tc.full_coexisting_host_nodes.discard(child); - let (kv_xfer, comp_xfers) = - tc.build_load_back_spec(tc.arena.node(child).id, /* req = */ None); - let actions = tc.commit_load_back( - tc.arena.node(child).id, - Tensor::from_slice(&[50i64, 51, 52, 53]), - kv_xfer, - comp_xfers, - ); + let (kv_xfer, comp_xfers) = tc + .build_load_back_spec(tc.arena.node(child).id, /* req = */ None) + .expect("live test node"); + let actions = tc + .commit_load_back( + tc.arena.node(child).id, + Tensor::from_slice(&[50i64, 51, 52, 53]), + kv_xfer, + comp_xfers, + ) + .expect("live transfer nodes"); assert!(actions.is_empty()); assert!( tc.arena @@ -3567,13 +3754,16 @@ fn commit_load_back_reattaches_device_slices_and_restores_the_match() { assert_eq!(tc.full_evictable_size(), 4); // The orchestrator re-locks the loaded path right after commit; that lock walk // also re-evaluates the parent's transient D-leaf membership. - tc.inc_lock_ref(tc.arena.node(child).id); + tc.inc_lock_ref(tc.arena.node(child).id) + .expect("live test node"); tc.dec_lock_ref( tc.arena.node(child).id, /* params = */ None, /* skip_swa = */ false, - ); - tc.finish_load_back(tc.arena.node(child).id); + ) + .expect("live test node"); + tc.finish_load_back(tc.arena.node(child).id) + .expect("live test node"); assert!(tc.full_coexisting_host_nodes.contains(parent)); assert!(tc.full_coexisting_host_nodes.contains(child)); tc.sanity_check(&[], &[]); @@ -3599,14 +3789,16 @@ fn device_eviction_and_demote_skip_a_load_back_pinned_chain() { demote_node(&mut tc, parent); tc.full_coexisting_host_nodes.discard(parent); tc.full_coexisting_host_nodes.discard(child); - let (kv_xfer, comp_xfers) = - tc.build_load_back_spec(tc.arena.node(child).id, /* req = */ None); + let (kv_xfer, comp_xfers) = tc + .build_load_back_spec(tc.arena.node(child).id, /* req = */ None) + .expect("live test node"); tc.commit_load_back( tc.arena.node(child).id, Tensor::from_slice(&[50i64, 51, 52, 53]), kv_xfer, comp_xfers, - ); + ) + .expect("live transfer nodes"); let anchor_id = tc.arena.node(child).id; assert_eq!(tc.arena.node(parent).load_back_pending_id, Some(anchor_id)); assert_eq!(tc.arena.node(child).load_back_pending_id, Some(anchor_id)); @@ -3617,9 +3809,10 @@ fn device_eviction_and_demote_skip_a_load_back_pinned_chain() { let (next, _) = tc.evict_device_next_node(FULL, &HashMap::new()); assert_eq!(next, None); tc.evict_device_end(FULL); - tc.demote(tc.arena.node(child).id); + tc.demote(tc.arena.node(child).id).expect("live test node"); assert!(tc.arena.has_device_value(child, FULL)); - tc.finish_load_back(tc.arena.node(child).id); + tc.finish_load_back(tc.arena.node(child).id) + .expect("live test node"); assert!(!tc.arena.node(parent).is_load_back_pending()); assert!(!tc.arena.node(child).is_load_back_pending()); assert!(tc.full_coexisting_host_nodes.contains(parent)); @@ -3629,7 +3822,8 @@ fn device_eviction_and_demote_skip_a_load_back_pinned_chain() { let (next, _) = tc.evict_device_next_node(FULL, &HashMap::new()); assert_eq!(next, Some(tc.arena.node(child).id)); tc.evict_device_end(FULL); - tc.demote(tc.arena.node(child).id); + tc.demote(tc.arena.node(child).id) + .expect("backuped live test node"); assert!(!tc.arena.has_device_value(child, FULL)); tc.sanity_check(&[], &[]); } @@ -3639,37 +3833,44 @@ fn component_has_host_value_only_tracks_the_demote_and_load_back_cycle() { let mut tc = core(); tc.insert(&insert_params(&vec![1], &[10])); let leaf = tc.match_prefix(&match_params(&vec![1])).best_match_node_id; - assert!(!tc.component_has_host_value_only(leaf, FULL)); - tc.commit_backup(leaf, Tensor::from_slice(&[20i64]), HashMap::new()); + assert!( + !tc.component_has_host_value_only(leaf, FULL) + .expect("live test node") + ); + tc.commit_backup(leaf, Tensor::from_slice(&[20i64]), HashMap::new()) + .expect("live test node"); // Device value still present: backuped but not host-only. - assert!(!tc.component_has_host_value_only(leaf, FULL)); - let leaf_idx = tc.arena.resolve(leaf); + assert!( + !tc.component_has_host_value_only(leaf, FULL) + .expect("live test node") + ); + let leaf_idx = tc.arena.resolve(leaf).expect("live test node"); demote_node(&mut tc, leaf_idx); - assert!(tc.component_has_host_value_only(leaf, FULL)); - let (kv_xfer, comp_xfers) = tc.build_load_back_spec(leaf, /* req = */ None); - tc.commit_load_back(leaf, Tensor::from_slice(&[30i64]), kv_xfer, comp_xfers); - assert!(!tc.component_has_host_value_only(leaf, FULL)); - tc.finish_load_back(leaf); + assert!( + tc.component_has_host_value_only(leaf, FULL) + .expect("live test node") + ); + let (kv_xfer, comp_xfers) = tc + .build_load_back_spec(leaf, /* req = */ None) + .expect("live test node"); + tc.commit_load_back(leaf, Tensor::from_slice(&[30i64]), kv_xfer, comp_xfers) + .expect("live transfer nodes"); + assert!( + !tc.component_has_host_value_only(leaf, FULL) + .expect("live test node") + ); + tc.finish_load_back(leaf).expect("live test node"); tc.sanity_check(&[], &[]); } #[test] -#[should_panic(expected = "!node.evicted() && node.backuped()")] -fn demote_panics_on_an_unbackuped_node() { - let mut tc = core(); - tc.insert(&insert_params(&vec![1], &[10])); - let leaf = tc.match_prefix(&match_params(&vec![1])).best_match_node_id; - tc.demote(leaf); -} - -#[test] -fn try_demote_rejects_unbackuped_and_evicted_nodes() { +fn demote_rejects_unbackuped_and_evicted_nodes() { let mut tc = core(); tc.insert(&insert_params(&vec![1], &[10])); let leaf = tc.match_prefix(&match_params(&vec![1])).best_match_node_id; assert!(matches!( - tc.try_demote(leaf), + tc.demote(leaf), Err(TreeCoreRuntimeError::InvalidDemoteState { node_id, evicted: false, @@ -3677,10 +3878,11 @@ fn try_demote_rejects_unbackuped_and_evicted_nodes() { }) if node_id == leaf )); - tc.commit_backup(leaf, Tensor::from_slice(&[20i64]), HashMap::new()); - tc.demote(leaf); + tc.commit_backup(leaf, Tensor::from_slice(&[20i64]), HashMap::new()) + .expect("live test node"); + tc.demote(leaf).expect("backuped live test node"); assert!(matches!( - tc.try_demote(leaf), + tc.demote(leaf), Err(TreeCoreRuntimeError::InvalidDemoteState { node_id, evicted: true, @@ -3694,13 +3896,31 @@ fn fallible_node_boundaries_reject_stale_handles() { let mut tc = core(); let stale_root = tc.root_node_handle(/* extra_key = */ None); tc.reset(); + let live_root = tc.root_node_handle(/* extra_key = */ None); + + assert!(matches!( + tc.validate_node_handles(&[live_root, stale_root]), + Err(NodeAccessError { node_id }) if node_id == stale_root + )); + assert!(matches!( + tc.insert_host_in_namespace( + stale_root, + KeyNamespaceRef::default(), + vec![1], + Tensor::from_slice(&[10i64]), + vec!["hash".to_string()], + ), + Err(TreeCoreRuntimeError::NodeAccess(NodeAccessError { node_id })) + if node_id == stale_root + )); assert!(matches!( - tc.try_demote(stale_root), - Err(TreeCoreRuntimeError::NodeNotAllocated { node_id }) if node_id == stale_root + tc.demote(stale_root), + Err(TreeCoreRuntimeError::NodeAccess(NodeAccessError { node_id })) + if node_id == stale_root )); assert!(matches!( - tc.try_build_hicache_transfers( + tc.build_hicache_transfers( FULL, stale_root, CacheTransferPhase::BackupStorage, @@ -3709,23 +3929,24 @@ fn fallible_node_boundaries_reject_stale_handles() { /* prefetch_tokens = */ 0, /* last_hash = */ None, ), - Err(TreeCoreRuntimeError::NodeNotAllocated { node_id }) if node_id == stale_root + Err(TreeCoreRuntimeError::NodeAccess(NodeAccessError { node_id })) + if node_id == stale_root )); assert!(matches!( - tc.try_build_load_back_spec(stale_root, /* req = */ None), - Err(TreeCoreRuntimeError::NodeNotAllocated { node_id }) if node_id == stale_root + tc.build_load_back_spec(stale_root, /* req = */ None), + Err(TreeCoreRuntimeError::NodeAccess(NodeAccessError { node_id })) + if node_id == stale_root )); assert!(matches!( - tc.try_get_hash_values(stale_root), - Err(TreeCoreRuntimeError::NodeNotAllocated { node_id }) if node_id == stale_root + tc.get_hash_values(stale_root), + Err(NodeAccessError { node_id }) if node_id == stale_root )); assert!(matches!( - tc.try_dfs_weight_order(&[stale_root]), - Err(TreeCoreRuntimeError::NodeNotAllocated { node_id }) if node_id == stale_root + tc.dfs_weight_order(&[stale_root]), + Err(NodeAccessError { node_id }) if node_id == stale_root )); - let live_root = tc.root_node_handle(/* extra_key = */ None); - assert!(tc.is_root(live_root)); + assert!(tc.is_root(live_root).expect("live root")); } #[test] @@ -3740,8 +3961,9 @@ fn match_prefix_with_hicache_splits_a_host_only_backuped_node() { leaf, Tensor::from_slice(&[100i64, 101, 102, 103]), HashMap::new(), - ); - tc.demote(leaf); + ) + .expect("live test node"); + tc.demote(leaf).expect("backuped live test node"); // The partial match splits the host-only node; the host prefix stays usable. let result = tc.match_prefix(&match_params(&vec![1, 2, 9])); assert_eq!(result.device_indices.numel(), 0); @@ -3749,7 +3971,10 @@ fn match_prefix_with_hicache_splits_a_host_only_backuped_node() { assert_eq!(result.last_device_node_id, tc.arena.node(root).id); assert_eq!(result.host_hit_length, 2); assert_eq!(result.best_match_node_id, result.last_host_node_id); - let parent = tc.arena.resolve(result.best_match_node_id); + let parent = tc + .arena + .resolve(result.best_match_node_id) + .expect("live test node"); let child = tc.arena.node(parent).children[&(KeyNamespace::default(), vec![3])]; { let parent_node = tc.arena.node(parent); @@ -3786,15 +4011,18 @@ fn mixed_backup_evict_insert_keeps_the_leaf_sets_disjoint() { let first = tc .match_prefix(&match_params(&vec![1, 2])) .best_match_node_id; - tc.commit_backup(first, Tensor::from_slice(&[100i64, 101]), HashMap::new()); + tc.commit_backup(first, Tensor::from_slice(&[100i64, 101]), HashMap::new()) + .expect("live test node"); let second = tc .match_prefix(&match_params(&vec![101, 102])) .best_match_node_id; - tc.commit_backup(second, Tensor::from_slice(&[102i64, 103]), HashMap::new()); + tc.commit_backup(second, Tensor::from_slice(&[102i64, 103]), HashMap::new()) + .expect("live test node"); let third = tc .match_prefix(&match_params(&vec![201, 202])) .best_match_node_id; - tc.commit_backup(third, Tensor::from_slice(&[104i64, 105]), HashMap::new()); + tc.commit_backup(third, Tensor::from_slice(&[104i64, 105]), HashMap::new()) + .expect("live test node"); let mut tracker = HashMap::from([(FULL, 0)]); let (mut df, mut hf) = (HashMap::new(), HashMap::new()); tc.evict_device_start(FULL, /* request_cnt = */ 4); @@ -3802,7 +4030,9 @@ fn mixed_backup_evict_insert_keeps_the_leaf_sets_disjoint() { 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); @@ -3851,7 +4081,7 @@ fn unbacked_leaf_with_host_child(tc: &mut UnifiedTreeCore>) -> (NodeIdx let child = tc .arena .alloc_child( - tc.arena.resolve(leaf), + tc.arena.resolve(leaf).expect("live test node"), /* key = */ vec![3, 4], /* priority = */ 0, /* extra_key = */ None, @@ -3860,8 +4090,8 @@ fn unbacked_leaf_with_host_child(tc: &mut UnifiedTreeCore>) -> (NodeIdx tc.arena .set_host_value(child, FULL, Tensor::from_slice(&[20i64, 21])); tc.update_evictable_leaf_sets_(child); - tc.update_evictable_leaf_sets_(tc.arena.resolve(leaf)); - (tc.arena.resolve(leaf), child) + tc.update_evictable_leaf_sets_(tc.arena.resolve(leaf).expect("live test node")); + (tc.arena.resolve(leaf).expect("live test node"), child) } #[test] @@ -3870,7 +4100,9 @@ fn drop_subtree_no_host_frees_the_leaf_and_its_host_descendants() { let (leaf, _child) = unbacked_leaf_with_host_child(&mut tc); let mut tracker = HashMap::from([(FULL, 0)]); let (mut df, mut hf) = (HashMap::new(), HashMap::new()); - let (dropped, step) = tc.drop_subtree_no_host(tc.arena.node(leaf).id); + let (dropped, step) = tc + .drop_subtree_no_host(tc.arena.node(leaf).id) + .expect("live test node"); accumulate_step(step, &mut tracker, &mut df, &mut hf); assert!(dropped); // Under EvictLayer::All only device tokens enter the tracker; host @@ -3905,7 +4137,9 @@ fn drop_subtree_no_host_removes_a_deeper_host_chain_child_first() { tc.update_evictable_leaf_sets_(child); let mut tracker = HashMap::from([(FULL, 0)]); let (mut df, mut hf) = (HashMap::new(), HashMap::new()); - let (dropped, step) = tc.drop_subtree_no_host(tc.arena.node(leaf).id); + let (dropped, step) = tc + .drop_subtree_no_host(tc.arena.node(leaf).id) + .expect("live test node"); accumulate_step(step, &mut tracker, &mut df, &mut hf); assert!(dropped); assert_eq!(tracker[&FULL], 2); @@ -3923,7 +4157,9 @@ fn drop_subtree_no_host_bails_on_a_locked_descendant() { .set_lock_ref_(ValueSlotIdx::host(FULL), 1); let mut tracker = HashMap::from([(FULL, 0)]); let (mut df, mut hf) = (HashMap::new(), HashMap::new()); - let (dropped, step) = tc.drop_subtree_no_host(tc.arena.node(leaf).id); + let (dropped, step) = tc + .drop_subtree_no_host(tc.arena.node(leaf).id) + .expect("live test node"); accumulate_step(step, &mut tracker, &mut df, &mut hf); assert!(!dropped); assert_eq!(tracker[&FULL], 0); @@ -3938,7 +4174,9 @@ fn drop_subtree_no_host_bails_on_a_host_locked_root() { tc.arena .node_mut(leaf) .set_lock_ref_(ValueSlotIdx::host(FULL), 1); - let (dropped, _step) = tc.drop_subtree_no_host(tc.arena.node(leaf).id); + let (dropped, _step) = tc + .drop_subtree_no_host(tc.arena.node(leaf).id) + .expect("live test node"); assert!(!dropped); assert_eq!(tc.arena.len(), 3); } @@ -3952,7 +4190,7 @@ fn drop_subtree_no_host_panics_on_a_non_device_leaf() { let parent = tc .match_prefix(&match_params(&vec![1, 2])) .best_match_node_id; - tc.drop_subtree_no_host(parent); + let _ = tc.drop_subtree_no_host(parent); } #[test] @@ -3963,8 +4201,9 @@ fn drop_subtree_no_host_panics_on_a_backuped_leaf() { let leaf = tc .match_prefix(&match_params(&vec![1, 2])) .best_match_node_id; - tc.commit_backup(leaf, Tensor::from_slice(&[20i64, 21]), HashMap::new()); - tc.drop_subtree_no_host(leaf); + tc.commit_backup(leaf, Tensor::from_slice(&[20i64, 21]), HashMap::new()) + .expect("live test node"); + let _ = tc.drop_subtree_no_host(leaf); } #[test] @@ -3976,20 +4215,26 @@ fn write_back_eviction_frees_the_device_value_exactly_once() { let leaf = tc.match_prefix(&match_params(&vec![1])).best_match_node_id; let mut tracker = HashMap::from([(FULL, 0)]); let (mut df, mut hf) = (HashMap::new(), HashMap::new()); - let (action, step) = tc.evict_device_leaf(leaf, true); + let (action, step) = tc.evict_device_leaf(leaf, true).expect("live test node"); accumulate_step(step, &mut tracker, &mut df, &mut hf); assert_eq!(action.unwrap().node_ids, vec![leaf]); assert!(df.is_empty() && hf.is_empty()); - tc.commit_backup(leaf, Tensor::from_slice(&[20i64]), HashMap::new()); - let (action, step) = tc.evict_device_leaf(leaf, true); + tc.commit_backup(leaf, Tensor::from_slice(&[20i64]), HashMap::new()) + .expect("live test node"); + let (action, step) = tc.evict_device_leaf(leaf, true).expect("live test node"); accumulate_step(step, &mut tracker, &mut df, &mut hf); assert!(action.is_none()); assert_eq!(tracker[&FULL], 1); assert_eq!(df[&FULL].len(), 1); assert!(df[&FULL][0].equal(&Tensor::from_slice(&[10i64]))); assert!( - tc.arena.node(tc.arena.resolve(leaf)).evicted() - && tc.arena.node(tc.arena.resolve(leaf)).backuped() + tc.arena + .node(tc.arena.resolve(leaf).expect("live test node")) + .evicted() + && tc + .arena + .node(tc.arena.resolve(leaf).expect("live test node")) + .backuped() ); tc.sanity_check(&[], &[]); } @@ -4060,10 +4305,15 @@ fn dfs_weight_order_groups_the_heaviest_subtree_first() { .last_device_node_id; assert_eq!( - tc.dfs_weight_order(&[leaf_b, leaf_a2, leaf_a1, leaf_a1, branch_a]), + tc.dfs_weight_order(&[leaf_b, leaf_a2, leaf_a1, leaf_a1, branch_a]) + .expect("live test nodes"), vec![2, 3, 1, 4, 0] ); - assert_eq!(tc.dfs_weight_order(&[leaf_b, leaf_a2]), vec![1, 0]); + assert_eq!( + tc.dfs_weight_order(&[leaf_b, leaf_a2]) + .expect("live test nodes"), + vec![1, 0] + ); } #[test] @@ -4071,12 +4321,14 @@ fn get_hash_values_reads_the_nodes_own_hashes() { let mut tc = core(); let (a, _b) = matched_chain(&mut tc); assert_eq!( - tc.get_hash_values(tc.arena.node(a).id), + tc.get_hash_values(tc.arena.node(a).id) + .expect("live test node"), Vec::::new() ); tc.arena.node_mut(a).hash_value = Some(vec!["h0".to_string(), "h1".to_string()]); assert_eq!( - tc.get_hash_values(tc.arena.node(a).id), + tc.get_hash_values(tc.arena.node(a).id) + .expect("live test node"), vec!["h0".to_string(), "h1".to_string()] ); } @@ -4150,10 +4402,15 @@ fn insert_page_size_two_drops_the_unaligned_tail() { let leaf = tc .match_prefix(&match_params(&vec![1, 2])) .best_match_node_id; - assert_eq!(tc.arena.node(tc.arena.resolve(leaf)).key, vec![1, 2]); + assert_eq!( + tc.arena + .node(tc.arena.resolve(leaf).expect("live test node")) + .key, + vec![1, 2] + ); assert!( tc.arena - .device_value(tc.arena.resolve(leaf), FULL) + .device_value(tc.arena.resolve(leaf).expect("live test node"), FULL) .equal(&Tensor::from_slice(&[10i64, 11])) ); tc.sanity_check(&[], &[]); @@ -4177,7 +4434,12 @@ fn insert_page_size_two_splits_mid_page_divergence_at_the_page_boundary() { let prefix = tc .match_prefix(&match_params(&vec![1, 2])) .best_match_node_id; - assert_eq!(tc.arena.node(tc.arena.resolve(prefix)).key, vec![1, 2]); + assert_eq!( + tc.arena + .node(tc.arena.resolve(prefix).expect("live test node")) + .key, + vec![1, 2] + ); let matched = tc.match_prefix(&match_params(&vec![1, 2, 3, 4])); assert!( matched @@ -4206,7 +4468,11 @@ fn match_prefix_page_size_two_splits_at_a_page_boundary() { ); assert_eq!( tc.arena - .node(tc.arena.resolve(result.best_match_node_id)) + .node( + tc.arena + .resolve(result.best_match_node_id) + .expect("live test node") + ) .key, vec![1, 2] ); @@ -4243,7 +4509,9 @@ fn evict_walk_page_size_two_empties_the_tree() { 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); evicted += 1; } @@ -4869,7 +5137,9 @@ fn evict_walk_and_driver_empty_the_tree_end_to_end() { 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 (backup, step) = tc.evict_device_leaf(leaf, /* is_write_back = */ false); + let (backup, step) = tc + .evict_device_leaf(leaf, /* is_write_back = */ false) + .expect("live test node"); accumulate_step(step, &mut tracker, &mut df, &mut hf); assert!(backup.is_none()); evicted += 1; @@ -4898,7 +5168,7 @@ fn evict_driver_readmits_the_parent_into_the_walk() { 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, false); + let (_, step) = tc.evict_device_leaf(leaf, false).expect("live test node"); accumulate_step(step, &mut tracker, &mut df, &mut hf); evicted += 1; } @@ -4954,7 +5224,8 @@ fn evict_driver_deletes_through_a_tombstone_parent() { /* priority = */ 0, /* extra_key = */ None, ); - tc.evict_device_leaf(tc.arena.node(b).id, false); + tc.evict_device_leaf(tc.arena.node(b).id, false) + .expect("live test node"); assert_eq!(tc.arena.len(), 1); } @@ -4964,11 +5235,12 @@ fn insert_after_eviction_reuses_the_slot_but_never_the_handle() { tc.insert(&insert_params(&vec![1, 2, 3], &[10, 11, 12])); tc.insert(&insert_params(&vec![4], &[13])); let leaf = tc.match_prefix(&match_params(&vec![4])).best_match_node_id; - let leaf_idx = tc.arena.resolve(leaf); - tc.evict_device_leaf(leaf, /* is_write_back = */ false); + let leaf_idx = tc.arena.resolve(leaf).expect("live test node"); + tc.evict_device_leaf(leaf, /* is_write_back = */ false) + .expect("live test node"); assert_eq!(tc.arena.len(), 2); // The freed handle no longer resolves. - assert!(tc.arena.try_resolve(leaf).is_none()); + assert!(tc.arena.resolve(leaf).is_err()); // The splitting insert allocates its prefix node into the freed slot. tc.insert(&insert_params(&vec![1, 2, 9], &[20, 21, 29])); assert_eq!(tc.arena.len(), 4); @@ -4976,28 +5248,33 @@ fn insert_after_eviction_reuses_the_slot_but_never_the_handle() { .match_prefix(&match_params(&vec![1, 2])) .best_match_node_id; // Slot recycled, but the stale handle can never alias the new node. - assert_eq!(tc.arena.resolve(prefix), leaf_idx); + assert_eq!(tc.arena.resolve(prefix).expect("live test node"), leaf_idx); assert_ne!(prefix, leaf); } #[test] -#[should_panic(expected = "is not allocated")] -fn stale_handle_panics_after_its_node_is_freed() { +fn stale_handle_returns_err_after_its_node_is_freed() { let mut tc = core(); tc.insert(&insert_params(&vec![4], &[13])); let leaf = tc.match_prefix(&match_params(&vec![4])).best_match_node_id; - tc.evict_device_leaf(leaf, /* is_write_back = */ false); - tc.inc_lock_ref(leaf); + tc.evict_device_leaf(leaf, /* is_write_back = */ false) + .expect("live test node"); + assert!(matches!( + tc.inc_lock_ref(leaf), + Err(NodeAccessError { node_id }) if node_id == leaf + )); } #[test] -#[should_panic(expected = "is not allocated")] -fn pre_reset_handle_panics_after_reset() { +fn pre_reset_handle_returns_err_after_reset() { let mut tc = core(); tc.insert(&insert_params(&vec![4], &[13])); let leaf = tc.match_prefix(&match_params(&vec![4])).best_match_node_id; tc.reset(); - tc.arena.resolve(leaf); + assert!(matches!( + tc.arena.resolve(leaf), + Err(NodeAccessError { node_id }) if node_id == leaf + )); } #[test] @@ -5010,7 +5287,7 @@ fn evict_driver_rejects_a_non_leaf() { let prefix = tc .match_prefix(&match_params(&vec![1, 2])) .best_match_node_id; - tc.evict_device_leaf(prefix, false); + let _ = tc.evict_device_leaf(prefix, false); } #[test] @@ -5020,11 +5297,17 @@ fn evict_driver_write_back_returns_the_backup_action_for_an_unbacked_leaf() { let leaf = tc.match_prefix(&match_params(&vec![1])).best_match_node_id; let mut tracker = HashMap::from([(FULL, 0)]); let (mut df, mut hf) = (HashMap::new(), HashMap::new()); - let (action, step) = tc.evict_device_leaf(leaf, /* is_write_back = */ true); + let (action, step) = tc + .evict_device_leaf(leaf, /* is_write_back = */ true) + .expect("live test node"); accumulate_step(step, &mut tracker, &mut df, &mut hf); // Write-back carries only the leaf itself; nothing is freed yet. assert_eq!(action.unwrap().node_ids, vec![leaf]); - assert!(!tc.arena.node(tc.arena.resolve(leaf)).evicted()); + assert!( + !tc.arena + .node(tc.arena.resolve(leaf).expect("live test node")) + .evicted() + ); assert_eq!(tracker[&FULL], 0); assert!(df.is_empty() && hf.is_empty()); } @@ -5034,21 +5317,32 @@ fn evict_driver_demotes_a_backuped_leaf_to_host_only() { let mut tc = core(); tc.insert(&insert_params(&vec![1], &[10])); let leaf = tc.match_prefix(&match_params(&vec![1])).best_match_node_id; - tc.arena - .set_host_value(tc.arena.resolve(leaf), FULL, Tensor::from_slice(&[20i64])); + tc.arena.set_host_value( + tc.arena.resolve(leaf).expect("live test node"), + FULL, + Tensor::from_slice(&[20i64]), + ); let mut tracker = HashMap::from([(FULL, 0)]); let (mut df, mut hf) = (HashMap::new(), HashMap::new()); - let (action, step) = tc.evict_device_leaf(leaf, false); + let (action, step) = tc.evict_device_leaf(leaf, false).expect("live test node"); accumulate_step(step, &mut tracker, &mut df, &mut hf); assert!(action.is_none()); // The node stays in the tree, now host-only. - let node = tc.arena.node(tc.arena.resolve(leaf)); + let node = tc + .arena + .node(tc.arena.resolve(leaf).expect("live test node")); assert!(node.evicted() && node.backuped()); assert_eq!(tracker[&FULL], 1); assert_eq!(df[&FULL].len(), 1); assert!(hf.is_empty()); - assert!(!tc.evictable_device_leaves.contains(tc.arena.resolve(leaf))); - assert!(tc.evictable_host_leaves.contains(tc.arena.resolve(leaf))); + assert!( + !tc.evictable_device_leaves + .contains(tc.arena.resolve(leaf).expect("live test node")) + ); + assert!( + tc.evictable_host_leaves + .contains(tc.arena.resolve(leaf).expect("live test node")) + ); tc.sanity_check(&[], &[]); } @@ -5061,10 +5355,14 @@ fn evict_device_leaf_step_counts_are_independent_of_prior_evictions() { .match_prefix(&match_params(&vec![1, 2, 3])) .best_match_node_id; let second = tc.match_prefix(&match_params(&vec![4])).best_match_node_id; - let (_, step) = tc.evict_device_leaf(first, /* is_write_back = */ false); + let (_, step) = tc + .evict_device_leaf(first, /* is_write_back = */ false) + .expect("live test node"); assert_eq!(step.tracker[&FULL], 3); // The second step reports only its own leaf, not a running total. - let (_, step) = tc.evict_device_leaf(second, /* is_write_back = */ false); + let (_, step) = tc + .evict_device_leaf(second, /* is_write_back = */ false) + .expect("live test node"); assert_eq!(step.tracker[&FULL], 1); assert_eq!(step.device_frees[&FULL].len(), 1); } @@ -5865,7 +6163,8 @@ fn reset_restores_a_fresh_tree() { ..insert_params(&vec![7, 8], &[20, 21]) }); let matched = tc.match_prefix(&match_params(&vec![1, 2, 3])); - tc.inc_lock_ref(matched.best_match_node_id); + tc.inc_lock_ref(matched.best_match_node_id) + .expect("live match node"); assert_eq!(tc.protected_size(), 3); // Seed aux LRU, host LRU, and host-leaf state so the reset must clear each. let root = tc.arena.root(); @@ -5916,7 +6215,8 @@ fn size_accessors_mirror_the_full_component_state() { assert_eq!(tc.protected_size(), 0); assert_eq!(tc.component_evictable_size(FULL), 3); let matched = tc.match_prefix(&match_params(&vec![1, 2, 3])); - tc.inc_lock_ref(matched.best_match_node_id); + tc.inc_lock_ref(matched.best_match_node_id) + .expect("live match node"); assert_eq!(tc.protected_size(), 3); assert_eq!(tc.full_protected_size(), 3); assert_eq!(tc.evictable_size(), 0); @@ -6005,7 +6305,8 @@ fn walk_for_kv_canary_chains_slots_across_namespaces() { fn walk_for_kv_canary_unlocked_only_skips_locked_nodes_but_keeps_the_chain() { let mut tc = core(); let (a, _b) = matched_chain(&mut tc); - tc.inc_lock_ref(tc.arena.node(a).id); + tc.inc_lock_ref(tc.arena.node(a).id) + .expect("live test node"); assert_eq!( sorted_canary_rows(tc.walk_for_kv_canary(true, false)), vec![(12, 2, 11)] @@ -6081,11 +6382,18 @@ fn get_component_device_value_reads_the_full_value() { .best_match_node_id; assert!( tc.get_component_device_value(leaf, FULL) - .unwrap() + .expect("live test node") + .expect("device value") .equal(&Tensor::from_slice(&[10i64, 11, 12])) ); - let _ = tc.arena.take_device_value(tc.arena.resolve(leaf), FULL); - assert!(tc.get_component_device_value(leaf, FULL).is_none()); + let _ = tc + .arena + .take_device_value(tc.arena.resolve(leaf).expect("live test node"), FULL); + assert!( + tc.get_component_device_value(leaf, FULL) + .expect("live test node") + .is_none() + ); } #[test] @@ -6093,7 +6401,7 @@ fn get_component_device_value_reads_the_full_value() { fn get_component_device_value_panics_on_an_unregistered_component() { let tc = core(); let root = tc.arena.root(); - tc.get_component_device_value(tc.arena.node(root).id, SWA); + let _ = tc.get_component_device_value(tc.arena.node(root).id, SWA); } #[test] @@ -6103,6 +6411,7 @@ fn get_component_device_value_reads_the_registered_components_slot() { let (a, _b) = matched_chain(&mut tc); assert!( tc.get_component_device_value(tc.arena.node(a).id, SWA) + .expect("live test node") .is_none() ); tc.arena @@ -6110,7 +6419,8 @@ fn get_component_device_value_reads_the_registered_components_slot() { assert_eq!( Vec::::try_from( tc.get_component_device_value(tc.arena.node(a).id, SWA) - .unwrap() + .expect("live test node") + .expect("device value") ) .unwrap(), vec![5, 6] @@ -6150,9 +6460,11 @@ fn is_full_device_evicted_flips_when_the_value_tombstones() { let leaf = tc .match_prefix(&match_params(&vec![1, 2, 3])) .best_match_node_id; - assert!(!tc.is_full_device_evicted(leaf)); - let _ = tc.arena.take_device_value(tc.arena.resolve(leaf), FULL); - assert!(tc.is_full_device_evicted(leaf)); + assert!(!tc.is_full_device_evicted(leaf).expect("live test node")); + let _ = tc + .arena + .take_device_value(tc.arena.resolve(leaf).expect("live test node"), FULL); + assert!(tc.is_full_device_evicted(leaf).expect("live test node")); } #[test] @@ -6170,7 +6482,8 @@ fn set_component_device_value_stores_and_restamps_the_lru() { ) .unwrap(); assert!(!tc.arena.has_device_value(a, SWA)); - tc.set_component_device_value(tc.arena.node(a).id, SWA, Tensor::from_slice(&[5i64, 6])); + tc.set_component_device_value(tc.arena.node(a).id, SWA, Tensor::from_slice(&[5i64, 6])) + .expect("live test node"); assert!( tc.arena .device_value(a, SWA) @@ -6197,7 +6510,8 @@ fn set_component_device_value_migrates_the_node_off_the_host_lru() { ) .unwrap(); tc.host_lru_list_mut(SWA).insert_mru(a); - tc.set_component_device_value(tc.arena.node(a).id, SWA, Tensor::from_slice(&[5i64])); + tc.set_component_device_value(tc.arena.node(a).id, SWA, Tensor::from_slice(&[5i64])) + .expect("live test node"); assert!(!tc.host_lru_list(SWA).in_list(Some(a))); assert_eq!(tc.host_lru_list(SWA).len(), 0); assert_eq!(tc.device_lru_list(SWA).len(), 1); @@ -6208,7 +6522,7 @@ fn set_component_device_value_migrates_the_node_off_the_host_lru() { fn set_component_device_value_rejects_the_base_component() { let mut tc = core(); let root = tc.arena.root(); - tc.set_component_device_value( + let _ = tc.set_component_device_value( tc.arena.node(root).id, BASE_COMPONENT_TYPE, Tensor::from_slice(&[1i64]), @@ -6223,18 +6537,24 @@ fn collect_full_device_indices_concatenates_in_root_order() { let leaf = tc .match_prefix(&match_params(&vec![1, 2, 3, 4, 5])) .best_match_node_id; - let parent = tc.arena.node(tc.arena.resolve(leaf)).parent(); + let parent = tc + .arena + .node(tc.arena.resolve(leaf).expect("live test node")) + .parent(); let root = tc.arena.root(); assert!( tc.collect_full_device_indices(leaf, tc.arena.node(root).id) + .expect("live test nodes") .equal(&Tensor::from_slice(&[10i64, 11, 12, 13, 14])) ); assert!( tc.collect_full_device_indices(leaf, tc.arena.node(parent).id) + .expect("live test nodes") .equal(&Tensor::from_slice(&[13i64, 14])) ); assert_eq!( tc.collect_full_device_indices(tc.arena.node(root).id, tc.arena.node(root).id) + .expect("live test nodes") .numel(), 0 ); @@ -6249,7 +6569,10 @@ fn collect_full_device_indices_panics_on_an_evicted_path() { let leaf = tc .match_prefix(&match_params(&vec![1, 2, 3, 4, 5])) .best_match_node_id; - let parent = tc.arena.node(tc.arena.resolve(leaf)).parent(); + let parent = tc + .arena + .node(tc.arena.resolve(leaf).expect("live test node")) + .parent(); let _ = tc.arena.take_device_value(parent, FULL); let root = tc.arena.root(); let _ = tc.collect_full_device_indices(leaf, tc.arena.node(root).id); @@ -6295,7 +6618,9 @@ fn pretty_format_renders_every_namespace_and_component() { let leaf = tc .match_prefix(&match_params(&vec![1, 2, 3, 4, 5])) .best_match_node_id; - let _ = tc.arena.take_device_value(tc.arena.resolve(leaf), FULL); + let _ = tc + .arena + .take_device_value(tc.arena.resolve(leaf).expect("live test node"), FULL); tc.register_component_(Arc::new(SwaComponentForTest)); tc.arena.node_mut(NodeIdx_(1)).values[SWA.idx()].value = Some(Tensor::from_slice(&[0i64])); // Sibling render order follows HashMap iteration, so pin the line set. @@ -6333,13 +6658,16 @@ fn sanity_check_passes_on_a_healthy_tree() { let leaf = tc .match_prefix(&match_params(&vec![1, 2, 9])) .best_match_node_id; - tc.inc_lock_ref(leaf); + tc.inc_lock_ref(leaf).expect("live test node"); tc.sanity_check(&[(1, leaf)], &[(2, leaf)]); tc.dec_lock_ref( - tc.arena.node(tc.arena.resolve(leaf)).id, + tc.arena + .node(tc.arena.resolve(leaf).expect("live test node")) + .id, /* params = */ None, /* skip_swa = */ false, - ); + ) + .expect("live test node"); tc.sanity_check(&[], &[]); } @@ -6353,7 +6681,9 @@ fn sanity_check_passes_after_the_eviction_walk() { 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); @@ -6370,7 +6700,8 @@ fn sanity_check_detects_a_missing_device_leaf() { let leaf = tc .match_prefix(&match_params(&vec![1, 2, 9])) .best_match_node_id; - tc.evictable_device_leaves.discard(tc.arena.resolve(leaf)); + tc.evictable_device_leaves + .discard(tc.arena.resolve(leaf).expect("live test node")); tc.sanity_check(&[], &[]); } @@ -6381,7 +6712,10 @@ fn sanity_check_detects_an_extra_device_leaf() { let leaf = tc .match_prefix(&match_params(&vec![1, 2, 9])) .best_match_node_id; - let parent = tc.arena.node(tc.arena.resolve(leaf)).parent(); + let parent = tc + .arena + .node(tc.arena.resolve(leaf).expect("live test node")) + .parent(); tc.evictable_device_leaves.add(parent); tc.sanity_check(&[], &[]); } @@ -6401,7 +6735,9 @@ fn sanity_check_detects_a_dead_node() { let leaf = tc .match_prefix(&match_params(&vec![1, 2, 9])) .best_match_node_id; - let _ = tc.arena.take_device_value(tc.arena.resolve(leaf), FULL); + let _ = tc + .arena + .take_device_value(tc.arena.resolve(leaf).expect("live test node"), FULL); tc.sanity_check(&[], &[]); } @@ -6411,7 +6747,7 @@ fn try_sanity_check_returns_a_dead_node_error() { let leaf = tc .match_prefix(&match_params(&vec![1, 2, 9])) .best_match_node_id; - let leaf_idx = tc.arena.resolve(leaf); + let leaf_idx = tc.arena.resolve(leaf).expect("live test node"); let _ = tc.arena.take_device_value(leaf_idx, FULL); let error = tc.try_sanity_check(&[], &[]).unwrap_err(); @@ -6428,7 +6764,10 @@ fn sanity_check_detects_an_evicted_parent_prefix() { let leaf = tc .match_prefix(&match_params(&vec![1, 2, 9])) .best_match_node_id; - let parent = tc.arena.node(tc.arena.resolve(leaf)).parent(); + let parent = tc + .arena + .node(tc.arena.resolve(leaf).expect("live test node")) + .parent(); let _ = tc.arena.take_device_value(parent, FULL); tc.sanity_check(&[], &[]); } @@ -6440,8 +6779,10 @@ fn sanity_check_detects_a_locked_tombstone() { let leaf = tc .match_prefix(&match_params(&vec![1, 2, 9])) .best_match_node_id; - tc.inc_lock_ref(leaf); - let _ = tc.arena.take_device_value(tc.arena.resolve(leaf), FULL); + tc.inc_lock_ref(leaf).expect("live test node"); + let _ = tc + .arena + .take_device_value(tc.arena.resolve(leaf).expect("live test node"), FULL); tc.sanity_check(&[], &[]); } @@ -6461,8 +6802,10 @@ fn sanity_check_detects_an_aux_lru_mismatch() { .match_prefix(&match_params(&vec![1, 2, 9])) .best_match_node_id; tc.register_component_(Arc::new(SwaComponentForTest)); - tc.arena.node_mut(tc.arena.resolve(leaf)).values[SWA.idx()].value = - Some(Tensor::from_slice(&[0i64, 0, 0])); + tc.arena + .node_mut(tc.arena.resolve(leaf).expect("live test node")) + .values[SWA.idx()] + .value = Some(Tensor::from_slice(&[0i64, 0, 0])); tc.sanity_check(&[], &[]); } @@ -6549,7 +6892,7 @@ fn sanity_check_detects_a_broken_parent_pointer() { let leaf = tc .match_prefix(&match_params(&vec![1, 2, 9])) .best_match_node_id; - let leaf_idx = tc.arena.resolve(leaf); + let leaf_idx = tc.arena.resolve(leaf).expect("live test node"); tc.arena.node_mut(leaf_idx).parent = Some(NodeIdx_(0)); tc.sanity_check(&[], &[]); } @@ -6562,9 +6905,13 @@ fn sanity_check_detects_aux_device_without_full() { .match_prefix(&match_params(&vec![1, 2, 9])) .best_match_node_id; tc.register_component_(Arc::new(SwaComponentForTest)); - tc.arena.node_mut(tc.arena.resolve(leaf)).values[SWA.idx()].value = - Some(Tensor::from_slice(&[0i64])); - let _ = tc.arena.take_device_value(tc.arena.resolve(leaf), FULL); + tc.arena + .node_mut(tc.arena.resolve(leaf).expect("live test node")) + .values[SWA.idx()] + .value = Some(Tensor::from_slice(&[0i64])); + let _ = tc + .arena + .take_device_value(tc.arena.resolve(leaf).expect("live test node"), FULL); tc.sanity_check(&[], &[]); } @@ -6577,7 +6924,7 @@ fn sanity_check_detects_aux_host_without_full_host() { .best_match_node_id; tc.register_component_(Arc::new(SwaComponentForTest)); tc.arena - .node_mut(tc.arena.resolve(leaf)) + .node_mut(tc.arena.resolve(leaf).expect("live test node")) .state_mut_(ValueSlotIdx::host(SWA)) .value = Some(Tensor::from_slice(&[0i64])); tc.sanity_check(&[], &[]); @@ -6591,7 +6938,7 @@ fn sanity_check_detects_an_unbacked_parent_prefix() { .match_prefix(&match_params(&vec![1, 2, 9])) .best_match_node_id; tc.arena - .node_mut(tc.arena.resolve(leaf)) + .node_mut(tc.arena.resolve(leaf).expect("live test node")) .state_mut_(ValueSlotIdx::host(FULL)) .value = Some(Tensor::from_slice(&[30i64])); tc.sanity_check(&[], &[]); @@ -6611,11 +6958,11 @@ fn sanity_check_accepts_a_write_back_child_backed_up_before_its_parent() { .match_prefix(&match_params(&vec![1, 2, 3, 4, 5])) .best_match_node_id; tc.arena - .node_mut(tc.arena.resolve(leaf)) + .node_mut(tc.arena.resolve(leaf).expect("live test node")) .state_mut_(ValueSlotIdx::host(FULL)) .value = Some(Tensor::from_slice(&[13i64, 14])); // Register the host value set directly by the test. - tc.update_full_coexisting_host_tracking_(tc.arena.resolve(leaf)); + tc.update_full_coexisting_host_tracking_(tc.arena.resolve(leaf).expect("live test node")); tc.sanity_check(&[], &[]); } @@ -6627,7 +6974,10 @@ fn sanity_check_detects_an_aux_lock_above_full() { .match_prefix(&match_params(&vec![1, 2, 9])) .best_match_node_id; tc.register_component_(Arc::new(SwaComponentForTest)); - tc.arena.node_mut(tc.arena.resolve(leaf)).values[SWA.idx()].lock_ref = 5; + tc.arena + .node_mut(tc.arena.resolve(leaf).expect("live test node")) + .values[SWA.idx()] + .lock_ref = 5; tc.sanity_check(&[], &[]); } @@ -6638,10 +6988,15 @@ fn sanity_check_detects_a_missing_host_leaf() { let leaf = tc .match_prefix(&match_params(&vec![1, 2, 9])) .best_match_node_id; - let parent = tc.arena.node(tc.arena.resolve(leaf)).parent(); - let _ = tc.arena.take_device_value(tc.arena.resolve(leaf), FULL); + let parent = tc + .arena + .node(tc.arena.resolve(leaf).expect("live test node")) + .parent(); + let _ = tc + .arena + .take_device_value(tc.arena.resolve(leaf).expect("live test node"), FULL); tc.arena - .node_mut(tc.arena.resolve(leaf)) + .node_mut(tc.arena.resolve(leaf).expect("live test node")) .state_mut_(ValueSlotIdx::host(FULL)) .value = Some(Tensor::from_slice(&[30i64])); tc.arena @@ -6658,7 +7013,8 @@ fn sanity_check_detects_an_extra_host_leaf() { let leaf = tc .match_prefix(&match_params(&vec![1, 2, 9])) .best_match_node_id; - tc.evictable_host_leaves.add(tc.arena.resolve(leaf)); + tc.evictable_host_leaves + .add(tc.arena.resolve(leaf).expect("live test node")); tc.sanity_check(&[], &[]); } @@ -6669,7 +7025,8 @@ fn sanity_check_detects_a_leaf_in_both_sets() { let leaf = tc .match_prefix(&match_params(&vec![1, 2, 9])) .best_match_node_id; - tc.evictable_host_leaves.add(tc.arena.resolve(leaf)); + tc.evictable_host_leaves + .add(tc.arena.resolve(leaf).expect("live test node")); tc.sanity_check(&[], &[]); } @@ -6705,7 +7062,7 @@ fn sanity_check_detects_an_aux_host_lru_mismatch() { .match_prefix(&match_params(&vec![1, 2, 9])) .best_match_node_id; tc.register_component_(Arc::new(SwaComponentForTest)); - let leaf_idx2 = tc.arena.resolve(leaf); + let leaf_idx2 = tc.arena.resolve(leaf).expect("live test node"); tc.host_lru_list_mut(SWA).insert_mru(leaf_idx2); tc.sanity_check(&[], &[]); } @@ -6718,9 +7075,11 @@ fn sanity_check_detects_an_aux_node_in_both_lrus() { .match_prefix(&match_params(&vec![1, 2, 9])) .best_match_node_id; tc.register_component_(Arc::new(SwaComponentForTest)); - tc.arena.node_mut(tc.arena.resolve(leaf)).values[SWA.idx()].value = - Some(Tensor::from_slice(&[0i64, 0, 0])); - let leaf_idx = tc.arena.resolve(leaf); + tc.arena + .node_mut(tc.arena.resolve(leaf).expect("live test node")) + .values[SWA.idx()] + .value = Some(Tensor::from_slice(&[0i64, 0, 0])); + let leaf_idx = tc.arena.resolve(leaf).expect("live test node"); tc.device_lru_list_mut(SWA).insert_mru(leaf_idx); tc.host_lru_list_mut(SWA).insert_mru(leaf_idx); tc.sanity_check(&[], &[]); @@ -6833,21 +7192,18 @@ fn refresh_dispatches_fire_per_walk_phase_in_a_namespace() { .best_match_node_id; let refreshes = recorder.refreshes.lock().unwrap(); assert!(!refreshes.is_empty()); - assert!( - refreshes - .iter() - .any(|&(phase, node)| phase == LRURefreshPhase::Walkdown - && node == tc.arena.resolve(leaf)) - ); - assert!(refreshes.iter().any( - |&(phase, node)| phase == LRURefreshPhase::InsertEnd && node == tc.arena.resolve(leaf) - )); - assert!( - refreshes - .iter() - .any(|&(phase, node)| phase == LRURefreshPhase::MatchEnd - && node == tc.arena.resolve(leaf)) - ); + assert!(refreshes.iter().any(|&(phase, node)| { + phase == LRURefreshPhase::Walkdown + && node == tc.arena.resolve(leaf).expect("live test node") + })); + assert!(refreshes.iter().any(|&(phase, node)| { + phase == LRURefreshPhase::InsertEnd + && node == tc.arena.resolve(leaf).expect("live test node") + })); + assert!(refreshes.iter().any(|&(phase, node)| { + phase == LRURefreshPhase::MatchEnd + && node == tc.arena.resolve(leaf).expect("live test node") + })); } #[test] @@ -6874,7 +7230,7 @@ fn sanity_check_detects_a_reverse_map_mismatch() { let leaf = tc .match_prefix(&match_params(&vec![1, 2, 9])) .best_match_node_id; - let leaf_idx3 = tc.arena.resolve(leaf); + let leaf_idx3 = tc.arena.resolve(leaf).expect("live test node"); let parent = tc.arena.node(leaf_idx3).parent(); let key = tc.arena.node(leaf_idx3).key.child_key(1); let parent_node = tc.arena.node_mut(parent); @@ -6892,8 +7248,10 @@ fn sanity_check_detects_a_value_length_mismatch() { let leaf = tc .match_prefix(&match_params(&vec![1, 2, 9])) .best_match_node_id; - tc.arena.node_mut(tc.arena.resolve(leaf)).values[FULL.idx()].value = - Some(Tensor::from_slice(&[7i64, 8, 9, 10])); + tc.arena + .node_mut(tc.arena.resolve(leaf).expect("live test node")) + .values[FULL.idx()] + .value = Some(Tensor::from_slice(&[7i64, 8, 9, 10])); tc.sanity_check(&[], &[]); } @@ -6904,13 +7262,16 @@ fn sanity_check_detects_a_host_value_length_mismatch() { let leaf = tc .match_prefix(&match_params(&vec![1, 2, 9])) .best_match_node_id; - let parent = tc.arena.node(tc.arena.resolve(leaf)).parent(); + let parent = tc + .arena + .node(tc.arena.resolve(leaf).expect("live test node")) + .parent(); tc.arena .node_mut(parent) .state_mut_(ValueSlotIdx::host(FULL)) .value = Some(Tensor::from_slice(&[10i64, 11])); tc.arena - .node_mut(tc.arena.resolve(leaf)) + .node_mut(tc.arena.resolve(leaf).expect("live test node")) .state_mut_(ValueSlotIdx::host(FULL)) .value = Some(Tensor::from_slice(&[7i64, 8, 9, 10])); tc.sanity_check(&[], &[]); @@ -6923,7 +7284,9 @@ fn sanity_check_detects_an_empty_key() { let leaf = tc .match_prefix(&match_params(&vec![1, 2, 9])) .best_match_node_id; - tc.arena.node_mut(tc.arena.resolve(leaf)).key = vec![]; + tc.arena + .node_mut(tc.arena.resolve(leaf).expect("live test node")) + .key = vec![]; tc.sanity_check(&[], &[]); } @@ -6935,7 +7298,9 @@ fn sanity_check_detects_an_unaligned_key() { let leaf = tc .match_prefix(&match_params(&vec![1, 2, 3, 4])) .best_match_node_id; - tc.arena.node_mut(tc.arena.resolve(leaf)).key = vec![1]; + tc.arena + .node_mut(tc.arena.resolve(leaf).expect("live test node")) + .key = vec![1]; tc.sanity_check(&[], &[]); } @@ -6945,9 +7310,12 @@ fn cyclic_child_map_tree() -> UnifiedTreeCore> { let leaf = tc .match_prefix(&match_params(&vec![1, 2, 9])) .best_match_node_id; - let parent = tc.arena.node(tc.arena.resolve(leaf)).parent(); + let parent = tc + .arena + .node(tc.arena.resolve(leaf).expect("live test node")) + .parent(); tc.arena - .node_mut(tc.arena.resolve(leaf)) + .node_mut(tc.arena.resolve(leaf).expect("live test node")) .children .insert((KeyNamespace::default(), vec![50]), parent); tc @@ -6974,7 +7342,10 @@ fn sanity_check_detects_a_host_locked_value_missing_from_the_lru() { let leaf = tc .match_prefix(&match_params(&vec![1, 2, 9])) .best_match_node_id; - let parent = tc.arena.node(tc.arena.resolve(leaf)).parent(); + let parent = tc + .arena + .node(tc.arena.resolve(leaf).expect("live test node")) + .parent(); tc.register_component_(Arc::new(SwaComponentForTest)); // The arena was built Full-only; give the root the stub's lock too. tc.arena.node_mut(tc.arena.root()).values[SWA.idx()].lock_ref = 1; @@ -6982,7 +7353,9 @@ fn sanity_check_detects_a_host_locked_value_missing_from_the_lru() { .node_mut(parent) .state_mut_(ValueSlotIdx::host(FULL)) .value = Some(Tensor::from_slice(&[10i64, 11])); - let leaf_node = tc.arena.node_mut(tc.arena.resolve(leaf)); + let leaf_node = tc + .arena + .node_mut(tc.arena.resolve(leaf).expect("live test node")); leaf_node.state_mut_(ValueSlotIdx::host(FULL)).value = Some(Tensor::from_slice(&[30i64])); leaf_node.state_mut_(ValueSlotIdx::host(SWA)).value = Some(Tensor::from_slice(&[30i64])); leaf_node.state_mut_(ValueSlotIdx::host(SWA)).lock_ref = 1; @@ -6994,7 +7367,10 @@ fn host_only_aux_leaf(tc: &mut UnifiedTreeCore>) -> NodeIdx_ { let leaf = tc .match_prefix(&match_params(&vec![1, 2, 9])) .best_match_node_id; - let parent = tc.arena.node(tc.arena.resolve(leaf)).parent(); + let parent = tc + .arena + .node(tc.arena.resolve(leaf).expect("live test node")) + .parent(); tc.register_component_(Arc::new(SwaComponentForTest)); // The arena was built Full-only; give the root the stub's lock too. tc.arena.node_mut(tc.arena.root()).values[SWA.idx()].lock_ref = 1; @@ -7002,13 +7378,15 @@ fn host_only_aux_leaf(tc: &mut UnifiedTreeCore>) -> NodeIdx_ { .node_mut(parent) .state_mut_(ValueSlotIdx::host(FULL)) .value = Some(Tensor::from_slice(&[10i64, 11])); - let leaf_node = tc.arena.node_mut(tc.arena.resolve(leaf)); + let leaf_node = tc + .arena + .node_mut(tc.arena.resolve(leaf).expect("live test node")); leaf_node.state_mut_(ValueSlotIdx::host(FULL)).value = Some(Tensor::from_slice(&[30i64])); leaf_node.state_mut_(ValueSlotIdx::host(SWA)).value = Some(Tensor::from_slice(&[30i64])); // Register the host values set directly by the test. tc.update_full_coexisting_host_tracking_(parent); - tc.update_full_coexisting_host_tracking_(tc.arena.resolve(leaf)); - tc.arena.resolve(leaf) + tc.update_full_coexisting_host_tracking_(tc.arena.resolve(leaf).expect("live test node")); + tc.arena.resolve(leaf).expect("live test node") } #[test] @@ -7058,15 +7436,18 @@ fn insert_unevicts_a_tombstoned_deep_node() { let leaf = tc .match_prefix(&match_params(&vec![1, 2, 3, 4, 5])) .best_match_node_id; - let _ = tc.arena.take_device_value(tc.arena.resolve(leaf), FULL); + let _ = tc + .arena + .take_device_value(tc.arena.resolve(leaf).expect("live test node"), FULL); tc.component_state_mut(FULL).evictable_size = 3; - tc.evictable_device_leaves.discard(tc.arena.resolve(leaf)); + tc.evictable_device_leaves + .discard(tc.arena.resolve(leaf).expect("live test node")); let result = tc.insert(&insert_params(&vec![1, 2, 3, 4, 5], &[30, 31, 32, 33, 34])); assert_eq!(result.prefix_len, 5); // The revived leaf takes its own span of the fresh KV, not the key head. assert!( tc.arena - .device_value(tc.arena.resolve(leaf), FULL) + .device_value(tc.arena.resolve(leaf).expect("live test node"), FULL) .equal(&Tensor::from_slice(&[33i64, 34])) ); let [CacheAction::FreeDeviceKV(freed)] = result.cache_actions.as_slice() else { @@ -7170,7 +7551,7 @@ fn suspended_walk_core() -> (UnifiedTreeCore>, NodeIdx_, InsertStepResu .match_prefix(&match_params(&vec![1, 2, 3])) .best_match_node_id; let step = tc.begin_insert(&insert_params(&vec![1, 2, 3, 4, 5], &[20, 21, 22, 13, 14])); - let a_idx = tc.arena.resolve(a); + let a_idx = tc.arena.resolve(a).expect("live test node"); (tc, a_idx, step) } @@ -7254,12 +7635,13 @@ fn resume_insert_completes_after_an_on_path_host_leaf_is_evicted() { let h_leaf = tc .match_prefix(&match_params(&vec![1, 2, 3, 4, 5, 6, 7, 8])) .best_match_node_id; - let h_leaf_idx = tc.arena.resolve(h_leaf); + let h_leaf_idx = tc.arena.resolve(h_leaf).expect("live test node"); tc.commit_backup( h_leaf, Tensor::from_slice(&[104i64, 105, 106, 107]), HashMap::new(), - ); + ) + .expect("live test node"); demote_node(&mut tc, h_leaf_idx); let step = tc.begin_insert(&insert_params( &vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12], @@ -7273,20 +7655,26 @@ fn resume_insert_completes_after_an_on_path_host_leaf_is_evicted() { // The barrier's backup host-evicts the on-path H-leaf before committing. let mut tracker = HashMap::from([(FULL, 0)]); let (mut df, mut hf) = (HashMap::new(), HashMap::new()); - tc.evict_host_leaf_(tc.arena.resolve(h_leaf), &mut tracker, &mut df, &mut hf); + tc.evict_host_leaf_( + tc.arena.resolve(h_leaf).expect("live test node"), + &mut tracker, + &mut df, + &mut hf, + ); assert_eq!(tracker[&FULL], 4); tc.commit_backup( top, Tensor::from_slice(&[100i64, 101, 102, 103]), HashMap::new(), - ); + ) + .expect("live test node"); let done = tc.resume_insert(); let result = done.result.expect("the resumed walk completes"); assert_eq!(result.prefix_len, 4); assert!(!tc.has_ongoing_insert()); - assert!(tc.arena.try_resolve(h_leaf).is_none()); + assert!(tc.arena.resolve(h_leaf).is_err()); // The recreated suffix is top's single child, spanning the whole gap. - let top_idx = tc.arena.resolve(top); + let top_idx = tc.arena.resolve(top).expect("live test node"); assert_eq!(tc.arena.node(top_idx).children.len(), 1); let suffix = *tc.arena.node(top_idx).children.values().next().unwrap(); assert_eq!(tc.arena.node(suffix).key, vec![5, 6, 7, 8, 9, 10, 11, 12]); @@ -7315,7 +7703,8 @@ fn aborted_barrier_crossing_refires_on_the_next_insert() { tc.arena.node(a).id, Tensor::from_slice(&[100i64, 101, 102]), HashMap::new(), - ); + ) + .expect("live test node"); let done = tc.resume_insert(); assert_eq!( done.result.expect("the resumed walk completes").prefix_len, @@ -7341,12 +7730,13 @@ fn one_insert_walk_fires_two_crossings_around_a_backuped_middle() { let middle = tc .match_prefix(&match_params(&vec![1, 2, 3, 4, 5, 6, 7, 8])) .best_match_node_id; - let middle_idx = tc.arena.resolve(middle); + let middle_idx = tc.arena.resolve(middle).expect("live test node"); tc.commit_backup( middle, Tensor::from_slice(&[104i64, 105, 106, 107]), HashMap::new(), - ); + ) + .expect("live test node"); demote_node(&mut tc, middle_idx); // The device insert restores the middle and adds the unbacked deep leaf. @@ -7374,7 +7764,11 @@ fn one_insert_walk_fires_two_crossings_around_a_backuped_middle() { }) .collect(); assert_eq!(backups, vec![vec![top], vec![deep]]); - assert!(tc.arena.node(tc.arena.resolve(middle)).backuped()); + assert!( + tc.arena + .node(tc.arena.resolve(middle).expect("live test node")) + .backuped() + ); } #[test] @@ -7502,11 +7896,14 @@ fn reset_invalidates_every_prior_handle() { .best_match_node_id; tc.reset(); // Handles are never re-minted, so pre-reset ones miss instead of aliasing. - assert!(tc.arena.try_resolve(old_root).is_none()); - assert!(tc.arena.try_resolve(old_leaf).is_none()); + assert!(tc.arena.resolve(old_root).is_err()); + assert!(tc.arena.resolve(old_leaf).is_err()); let new_root = tc.root_node_handle(/* extra_key = */ None); assert_ne!(new_root, old_root); - assert_eq!(tc.arena.resolve(new_root), tc.arena.root()); + assert_eq!( + tc.arena.resolve(new_root).expect("live test node"), + tc.arena.root() + ); tc.insert(&insert_params(&vec![1, 2], &[10, 11])); assert_eq!( tc.match_prefix(&match_params(&vec![1, 2])) @@ -7516,12 +7913,56 @@ fn reset_invalidates_every_prior_handle() { ); } +#[test] +fn inspection_rejects_stale_handles_without_panicking() { + let mut tc = core(); + let stale_root = tc.root_node_handle(/* extra_key = */ None); + tc.reset(); + let live_root = tc.root_node_handle(/* extra_key = */ None); + let expected = NodeAccessError { + node_id: stale_root, + }; + + assert_eq!(tc.inspect_get_parent_node_id(stale_root), Err(expected)); + assert_eq!(tc.inspect_get_child_node_ids(stale_root), Err(expected)); + assert_eq!(tc.inspect_get_node_key_length(stale_root), Err(expected)); + assert_eq!( + tc.inspect_set_node_hash_values(stale_root, None), + Err(expected) + ); + assert_eq!( + tc.inspect_build_backup_node_ids(stale_root, /* write_back = */ false), + Err(expected) + ); + assert!(matches!( + tc.inspect_get_component_host_value(stale_root, SWA), + Err(error) if error == expected + )); + assert_eq!(tc.inspect_is_node_in_device_lru(stale_root, SWA), Ok(false)); + assert_eq!(tc.inspect_is_node_in_host_lru(stale_root, SWA), Ok(false)); + assert!(matches!( + tc.inspect_evict_component(stale_root, SWA, EvictLayer::Device), + Err(error) if error == expected + )); + assert!(matches!( + tc.inspect_validate_cascade_evict(stale_root, SWA, EvictLayer::Device), + Err(TreeCoreRuntimeError::NodeAccess(error)) if error == expected + )); + + // Presence-style probes intentionally treat stale handles as absent. + assert!(!tc.inspect_contains_node(stale_root)); + assert!(!tc.inspect_is_device_evictable_leaf(stale_root)); + assert!(!tc.inspect_is_host_evictable_leaf(stale_root)); + + assert_eq!(tc.inspect_get_parent_node_id(live_root), Ok(None)); +} + #[test] #[should_panic(expected = "is not enabled")] fn component_has_host_value_only_panics_on_a_disabled_component() { let tc = core(); let root = tc.root_node_handle(/* extra_key = */ None); - tc.component_has_host_value_only(root, SWA); + let _ = tc.component_has_host_value_only(root, SWA); } #[test] @@ -7607,20 +8048,21 @@ fn run_random_op_sequence(mut tc: UnifiedTreeCore>, page: usize, mamba: 2 => { // Balanced lock round trip on whatever the key matches. let anchor = tc.match_prefix(&match_params(&key)).best_match_node_id; - let lock = tc.inc_lock_ref(anchor); + let lock = tc.inc_lock_ref(anchor).expect("live match anchor"); 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(¶ms), /* skip_swa = */ false); + tc.dec_lock_ref(anchor, Some(¶ms), /* skip_swa = */ false) + .expect("live match anchor"); } _ => { // Insert-while-locked churn, the cache_finished_req shape. let matched = tc.match_prefix(&match_params(&key)); let anchor = matched.best_match_node_id; let matched_len = matched.device_indices.numel() as usize; - let lock = tc.inc_lock_ref(anchor); + let lock = tc.inc_lock_ref(anchor).expect("live match anchor"); tc.insert(&sequence_insert_params( &key, matched_len, @@ -7633,7 +8075,8 @@ fn run_random_op_sequence(mut tc: UnifiedTreeCore>, page: usize, mamba: 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(¶ms), /* skip_swa = */ false); + tc.dec_lock_ref(anchor, Some(¶ms), /* skip_swa = */ false) + .expect("live match anchor"); } } if step % 8 == 7 { @@ -7652,8 +8095,9 @@ fn run_random_op_sequence(mut tc: UnifiedTreeCore>, page: usize, mamba: &mut host_frees, ); let Some(leaf) = next else { break }; - let (_, evict_result) = - tc.evict_device_leaf(leaf, /* is_write_back = */ false); + let (_, evict_result) = tc + .evict_device_leaf(leaf, /* is_write_back = */ false) + .expect("live eviction candidate"); accumulate_step( evict_result, &mut tracker, @@ -7712,7 +8156,9 @@ fn drain_full_device(tc: &mut UnifiedTreeCore>) { 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); @@ -7734,7 +8180,7 @@ fn an_emptied_namespace_leaves_nothing_behind() { drain_full_device(&mut tc); // The namespace's nodes evict like any others; its edge map drops with them. assert!(!tc.arena.namespace_exists(Some("salted"))); - assert!(tc.arena.try_resolve(top).is_none()); + assert!(tc.arena.resolve(top).is_err()); assert_eq!(tc.arena.len(), 1); tc.sanity_check(&[], &[]); // A later insert respins the namespace from scratch. @@ -7777,11 +8223,12 @@ fn a_zero_length_match_anchors_at_the_root() { .best_match_node_id; assert_eq!(anchor, tc.root_node_handle(Some("salted"))); // The root handle stays valid across a full namespace eviction. - tc.inc_lock_ref(anchor); + tc.inc_lock_ref(anchor).expect("live root"); drain_full_device(&mut tc); tc.dec_lock_ref( anchor, /* params = */ None, /* skip_swa = */ false, - ); - assert!(tc.arena.try_resolve(anchor).is_some()); + ) + .expect("live root"); + assert!(tc.arena.resolve(anchor).is_ok()); tc.sanity_check(&[], &[]); } diff --git a/rust/sglang-radix-tree/src/unified_tree_core.rs b/rust/sglang-radix-tree/src/unified_tree_core.rs index 1e8426c1ddbe..322f64c6f545 100644 --- a/rust/sglang-radix-tree/src/unified_tree_core.rs +++ b/rust/sglang-radix-tree/src/unified_tree_core.rs @@ -16,7 +16,9 @@ use crate::node::EvictableNodeSet; use crate::node::Node; use crate::node::NodeArena; use crate::node::{ChildKeyType, HashDigest, KeyNamespace, KeyNamespaceRef}; -use crate::node::{NUM_VALUE_SLOTS, NodeId, NodeIdx_, TreeCoreRuntimeError, ValueSlotIdx}; +use crate::node::{ + NUM_VALUE_SLOTS, NodeAccessError, NodeId, NodeIdx_, TreeCoreRuntimeError, ValueSlotIdx, +}; use crate::unified_lru_list::UnifiedLRUList; use crate::unified_lru_list::{EvictionStrategy, PriorityKey, get_eviction_strategy}; @@ -791,7 +793,7 @@ impl UnifiedTreeCore { } /// Bump the reference count on a node's component locks. - pub fn inc_lock_ref(&mut self, node_id: NodeId) -> IncLockRefResult { + pub fn inc_lock_ref(&mut self, node_id: NodeId) -> Result { self.inc_lock_ref_with_skip(node_id, &[]) } @@ -800,8 +802,8 @@ impl UnifiedTreeCore { &mut self, node_id: NodeId, skip_lock_components: &[ComponentType], - ) -> IncLockRefResult { - let node_id = self.arena.resolve(node_id); + ) -> Result { + let node_id = self.arena.resolve(node_id)?; let node = self.arena.node(node_id); let node_handle = node.id; let is_root = node.is_root(); @@ -823,7 +825,7 @@ impl UnifiedTreeCore { .acquire_component_lock(self, node_id, result, /* lock_host = */ false); } self.update_evictable_leaf_sets_(node_id); - result + Ok(result) } /// Decrease the reference count on a node's component locks. @@ -832,8 +834,8 @@ impl UnifiedTreeCore { node_id: NodeId, params: Option<&DecLockRefParams>, skip_swa: bool, - ) -> DecLockRefResult { - let node_id = self.arena.resolve(node_id); + ) -> Result { + let node_id = self.arena.resolve(node_id)?; for i in 0..self.components.len() { if skip_swa && self.components[i].component_type() == SWA { continue; @@ -843,7 +845,7 @@ impl UnifiedTreeCore { } self.update_evictable_leaf_sets_(node_id); // TODO: delta is not aggregated from components; no caller uses it yet. - DecLockRefResult::default() + Ok(DecLockRefResult::default()) } /// Early-release the SWA portion of a request's tree lock, plus any @@ -854,14 +856,14 @@ impl UnifiedTreeCore { swa_uuid_for_lock: Option, device_frees: &mut HashMap>, host_frees: &mut HashMap>, - ) { + ) -> Result<(), NodeAccessError> { self.dec_swa_lock_only_with_skip( node_id, swa_uuid_for_lock, /* skip_lock_node_ids = */ None, device_frees, host_frees, - ); + ) } /// Skip-aware variant used when an acquire deliberately omitted a component. @@ -872,10 +874,10 @@ impl UnifiedTreeCore { skip_lock_node_ids: Option<&HashMap>>, device_frees: &mut HashMap>, host_frees: &mut HashMap>, - ) { - let node_id = self.arena.resolve(node_id); + ) -> Result<(), NodeAccessError> { + let node_id = self.arena.resolve(node_id)?; let Some(swa) = self.try_component_by_type_(SWA) else { - return; + return Ok(()); }; swa.release_window_lock(self, node_id, swa_uuid_for_lock, device_frees, host_frees); @@ -897,12 +899,16 @@ impl UnifiedTreeCore { ); } } + Ok(()) } /// Evict shallow Mamba device checkpoints beyond the per-path cap on the /// tail's root path; the mamba component drives the walk. - pub fn evict_excess_path_states(&mut self, tail_node_id: NodeId) -> EvictionStepResult { - let tail_node_id = self.arena.resolve(tail_node_id); + pub fn evict_excess_path_states( + &mut self, + tail_node_id: NodeId, + ) -> Result { + let tail_node_id = self.arena.resolve(tail_node_id)?; let mut result = EvictionStepResult::default(); let component = self.component_by_type_(MAMBA); component.evict_excess_path_states( @@ -911,12 +917,15 @@ impl UnifiedTreeCore { &mut result.device_frees, &mut result.host_frees, ); - result + Ok(result) } /// Bump the reference count on a node's host-side component locks. - pub fn inc_host_lock_ref(&mut self, node_id: NodeId) -> IncLockRefResult { - let node_id = self.arena.resolve(node_id); + pub fn inc_host_lock_ref( + &mut self, + node_id: NodeId, + ) -> Result { + let node_id = self.arena.resolve(node_id)?; let mut result = IncLockRefResult::default(); for i in 0..self.components.len() { let component = Arc::clone(&self.components[i]); @@ -924,7 +933,7 @@ impl UnifiedTreeCore { .acquire_component_lock(self, node_id, result, /* lock_host = */ true); } self.update_evictable_leaf_sets_(node_id); - result + Ok(result) } /// Decrease the reference count on a node's host-side component locks. @@ -932,14 +941,14 @@ impl UnifiedTreeCore { &mut self, node_id: NodeId, params: Option<&DecLockRefParams>, - ) -> DecLockRefResult { - let node_id = self.arena.resolve(node_id); + ) -> Result { + let node_id = self.arena.resolve(node_id)?; for i in 0..self.components.len() { let component = Arc::clone(&self.components[i]); component.release_component_lock(self, node_id, params, /* lock_host = */ true); } self.update_evictable_leaf_sets_(node_id); - DecLockRefResult::default() + Ok(DecLockRefResult::default()) } /// Match a key against the tree; returns device indices + boundary NodeIds. @@ -1189,6 +1198,8 @@ impl UnifiedTreeCore { result = component.finalize_match_result_in_tree_core( self, result, + best_match_device_node_id, + best_match_node_id, params, &value, best_match_device_value_len, @@ -1216,9 +1227,9 @@ impl UnifiedTreeCore { } /// Whether the node's FULL device value has been evicted. - pub fn is_full_device_evicted(&self, node_id: NodeId) -> bool { - let node_id = self.arena.resolve(node_id); - self.arena.node(node_id).evicted() + pub fn is_full_device_evicted(&self, node_id: NodeId) -> Result { + let node_id = self.arena.resolve(node_id)?; + Ok(self.arena.node(node_id).evicted()) } /// Concatenate FULL device values from ``from_node`` up to (exclusive) @@ -1227,9 +1238,9 @@ impl UnifiedTreeCore { &self, from_node_id: NodeId, until_node_id: NodeId, - ) -> Tensor { - let from_node_id = self.arena.resolve(from_node_id); - let until_node_id = self.arena.resolve(until_node_id); + ) -> Result { + let from_node_id = self.arena.resolve(from_node_id)?; + let until_node_id = self.arena.resolve(until_node_id)?; let mut prefix_chunks: Vec = Vec::new(); let mut node_id = from_node_id; while node_id != until_node_id { @@ -1238,10 +1249,10 @@ impl UnifiedTreeCore { node_id = node.parent(); } if prefix_chunks.is_empty() { - return self.empty_device_indices.shallow_clone(); + return Ok(self.empty_device_indices.shallow_clone()); } prefix_chunks.reverse(); - Tensor::cat(&prefix_chunks, 0) + Ok(Tensor::cat(&prefix_chunks, 0)) } /// Refresh a node's access tick and component LRU positions. @@ -2030,8 +2041,8 @@ impl UnifiedTreeCore { &mut self, node_id: NodeId, is_write_back: bool, - ) -> (Option, EvictionStepResult) { - let node_id = self.arena.resolve(node_id); + ) -> Result<(Option, EvictionStepResult), NodeAccessError> { + let node_id = self.arena.resolve(node_id)?; let mut result = EvictionStepResult::default(); { let node = self.arena.node(node_id); @@ -2047,12 +2058,12 @@ impl UnifiedTreeCore { &mut result.device_frees, &mut result.host_frees, ); - return (None, result); + return Ok((None, result)); } if is_write_back { let backup = self .build_backup_kv_action_(self.arena.node(node_id), /* write_back = */ true); - return (Some(backup), result); + return Ok((Some(backup), result)); } // Write-through: node has no backup, delete entirely. @@ -2062,15 +2073,18 @@ impl UnifiedTreeCore { &mut result.device_frees, &mut result.host_frees, ); - (None, result) + Ok((None, result)) } /// Write-back fallback when a D-leaf's D->H backup fails under host /// memory pressure: drop the subtree rooted at the unbacked leaf so /// device eviction keeps making progress instead of leaving its KV /// unevictable until host space frees up. - pub fn drop_subtree_no_host(&mut self, node_id: NodeId) -> (bool, EvictionStepResult) { - let node_id = self.arena.resolve(node_id); + pub fn drop_subtree_no_host( + &mut self, + node_id: NodeId, + ) -> Result<(bool, EvictionStepResult), NodeAccessError> { + let node_id = self.arena.resolve(node_id)?; let mut result = EvictionStepResult::default(); { let node = self.arena.node(node_id); @@ -2082,7 +2096,7 @@ impl UnifiedTreeCore { // no host state and no in-flight DMA reading its device slots. assert!(!node.backuped() && node.write_through_pending_id.is_none()); if node.is_host_locked() { - return (false, result); + return Ok((false, result)); } } let mut descendants: Vec = Vec::new(); @@ -2096,7 +2110,7 @@ impl UnifiedTreeCore { while let Some(cur_id) = stack.pop() { let cur = self.arena.node(cur_id); if cur.is_device_locked() || cur.is_host_locked() { - return (false, result); + return Ok((false, result)); } descendants.push(cur_id); stack.extend(cur.children.values().copied()); @@ -2127,7 +2141,7 @@ impl UnifiedTreeCore { &mut result.device_frees, &mut result.host_frees, ); - (true, result) + Ok((true, result)) } /// Free every component layer on the node and detach it from the LRU @@ -2289,28 +2303,8 @@ impl UnifiedTreeCore { /// Release a node's device KV once its host copy exists; the node stays in the /// tree, now host-only. - pub fn demote(&mut self, node_id: NodeId) -> EvictionStepResult { - let node_id = self.arena.resolve(node_id); - let mut result = EvictionStepResult::default(); - // Skip a deferred demote when a load-back now pins the device indices. - if self.arena.node(node_id).is_load_back_pending() { - return result; - } - self.demote_( - node_id, - &mut result.tracker, - &mut result.device_frees, - &mut result.host_frees, - ); - result - } - - /// Fallible variant of [`Self::demote`]. - pub fn try_demote( - &mut self, - node_id: NodeId, - ) -> Result { - let node_id = self.try_resolve_node_handle_(node_id)?; + pub fn demote(&mut self, node_id: NodeId) -> Result { + let node_id = self.arena.resolve(node_id)?; let mut result = EvictionStepResult::default(); // Skip a deferred demote when a load-back now pins the device indices. if self.arena.node(node_id).is_load_back_pending() { @@ -2898,7 +2892,7 @@ impl UnifiedTreeCore { key: K, host_value: Tensor, hash_value: Vec, - ) -> InsertResult { + ) -> Result { self.insert_host_in_namespace( node_id, KeyNamespaceRef::new(extra_key, /* cache_salt = */ None), @@ -2915,21 +2909,9 @@ impl UnifiedTreeCore { key: K, host_value: Tensor, hash_value: Vec, - ) -> InsertResult { - self.try_insert_host_in_namespace(node_id, namespace, key, host_value, hash_value) - .unwrap_or_else(|error| panic!("{error}")) - } - - pub fn try_insert_host_in_namespace( - &mut self, - node_id: NodeId, - namespace: KeyNamespaceRef<'_>, - key: K, - host_value: Tensor, - hash_value: Vec, ) -> Result { let total_len = key.atom_len(); - let mut node_id = self.arena.resolve(node_id); + let mut node_id = self.arena.resolve(node_id)?; let anchor = self.arena.node(node_id); if !anchor.is_root() && anchor.namespace.as_ref() != namespace { return Err(TreeCoreRuntimeError::InsertHostNamespaceMismatch { node_id: anchor.id }); @@ -3047,8 +3029,9 @@ impl UnifiedTreeCore { pub fn build_backup_spec( &self, node_id: NodeId, - ) -> (Tensor, HashMap>) { - self.build_backup_spec_(self.arena.node(self.arena.resolve(node_id))) + ) -> Result<(Tensor, HashMap>), NodeAccessError> { + let node_id = self.arena.resolve(node_id)?; + Ok(self.build_backup_spec_(self.arena.node(node_id))) } /// Gather device value backup spec. @@ -3096,11 +3079,11 @@ impl UnifiedTreeCore { &self, node_id: NodeId, pass_prefix_keys: bool, - ) -> Option { - let node_id = self.arena.resolve(node_id); + ) -> Result, NodeAccessError> { + let node_id = self.arena.resolve(node_id)?; let node = self.arena.node(node_id); if !node.backuped() { - return None; + return Ok(None); } let prefix_keys = pass_prefix_keys.then(|| self.arena.prefix_hash_values(node.parent)); let mut comp_xfers: HashMap> = HashMap::new(); @@ -3127,13 +3110,13 @@ impl UnifiedTreeCore { comp_xfers.insert(component_type, transfers); } } - Some(StorageBackupSpec { + Ok(Some(StorageBackupSpec { host_value: node.host_value(FULL).shallow_clone(), token_ids: K::raw_token_ids(node.key.as_ref()).into_owned(), hash_value: node.hash_value.clone(), prefix_keys, comp_xfers, - }) + })) } /// Route a build_hicache_transfers call to the component for the given type. @@ -3146,30 +3129,8 @@ impl UnifiedTreeCore { token_ids: Option<&[i64]>, prefetch_tokens: usize, last_hash: Option<&str>, - ) -> Option> { - self.try_build_hicache_transfers( - component_type, - node_id, - phase, - host_indices, - token_ids, - prefetch_tokens, - last_hash, - ) - .unwrap() - } - - pub fn try_build_hicache_transfers( - &self, - component_type: ComponentType, - node_id: NodeId, - phase: CacheTransferPhase, - host_indices: Option, - token_ids: Option<&[i64]>, - prefetch_tokens: usize, - last_hash: Option<&str>, ) -> Result>, TreeCoreRuntimeError> { - let node_id = self.try_resolve_node_handle_(node_id)?; + let node_id = self.arena.resolve(node_id)?; self.component_by_type_(component_type) .build_hicache_transfers( self, @@ -3188,18 +3149,10 @@ impl UnifiedTreeCore { &self, node_id: NodeId, req: Option<&Req>, - ) -> (PoolTransfer, HashMap>) { - self.try_build_load_back_spec(node_id, req).unwrap() - } - - pub fn try_build_load_back_spec( - &self, - node_id: NodeId, - req: Option<&Req>, ) -> Result<(PoolTransfer, HashMap>), TreeCoreRuntimeError> { let anchor_id = node_id; - let node_id = self.try_resolve_node_handle_(node_id)?; + let node_id = self.arena.resolve(node_id)?; // Component hooks take primitives, not Req: extract its fields here. let mamba_pool_idx = req.and_then(|r| r.mamba_pool_idx.as_ref()); let mut kv_transfers = self @@ -3239,23 +3192,21 @@ impl UnifiedTreeCore { } } // Reject transfers that would claim a node pinned by another load-back anchor. - let any_foreign_pin = kv_xfer - .nodes_to_load - .iter() - .chain( - comp_xfers - .values() - .flatten() - .filter_map(|xfer| xfer.nodes_to_load.as_ref()), - ) - .flatten() - .any(|&pinned_id| { - let pinned_idx = self.arena.resolve(pinned_id); - self.arena - .node(pinned_idx) - .load_back_pending_id - .is_some_and(|id| id != anchor_id) - }); + let mut any_foreign_pin = false; + for &pinned_id in kv_xfer.nodes_to_load.iter().flatten().chain( + comp_xfers + .values() + .flatten() + .filter_map(|xfer| xfer.nodes_to_load.as_ref()) + .flatten(), + ) { + let pinned_idx = self.arena.resolve(pinned_id)?; + any_foreign_pin |= self + .arena + .node(pinned_idx) + .load_back_pending_id + .is_some_and(|id| id != anchor_id); + } if any_foreign_pin { let empty_kv = PoolTransfer { name: PoolName::Kv, @@ -3268,23 +3219,32 @@ impl UnifiedTreeCore { Ok((kv_xfer, comp_xfers)) } - fn try_resolve_node_handle_(&self, node_id: NodeId) -> Result { - self.arena - .try_resolve(node_id) - .ok_or(TreeCoreRuntimeError::NodeNotAllocated { node_id }) + /// Validate that every external node handle names a live node. + pub(crate) fn validate_node_handles(&self, node_ids: &[NodeId]) -> Result<(), NodeAccessError> { + for &node_id in node_ids { + self.arena.resolve(node_id)?; + } + Ok(()) } - /// The anchor node's caller-defined key and cache salt. - pub fn prefetch_anchor_info(&self, node_id: NodeId) -> (Option, Option) { - self.try_prefetch_anchor_info(node_id) - .unwrap_or_else(|error| panic!("{error}")) + fn validate_pool_transfer_handles<'a>( + &self, + transfers: impl IntoIterator, + ) -> Result<(), NodeAccessError> { + for transfer in transfers { + if let Some(node_ids) = &transfer.nodes_to_load { + self.validate_node_handles(node_ids)?; + } + } + Ok(()) } - pub fn try_prefetch_anchor_info( + /// The anchor node's caller-defined key and cache salt. + pub fn prefetch_anchor_info( &self, node_id: NodeId, - ) -> Result<(Option, Option), TreeCoreRuntimeError> { - let node_id = self.try_resolve_node_handle_(node_id)?; + ) -> Result<(Option, Option), NodeAccessError> { + let node_id = self.arena.resolve(node_id)?; Ok(( self.arena.node_extra_key(node_id).map(str::to_string), self.arena.node_cache_salt(node_id).map(str::to_string), @@ -3292,38 +3252,20 @@ impl UnifiedTreeCore { } /// Whether the node's Full KV is present on host. - pub fn node_backuped(&self, node_id: NodeId) -> bool { - self.try_node_backuped(node_id) - .unwrap_or_else(|error| panic!("{error}")) - } - - pub fn try_node_backuped(&self, node_id: NodeId) -> Result { - let node_id = self.try_resolve_node_handle_(node_id)?; + pub fn node_backuped(&self, node_id: NodeId) -> Result { + let node_id = self.arena.resolve(node_id)?; Ok(self.arena.node(node_id).backuped()) } /// Whether the node is a (default or named) root. - pub fn is_root(&self, node_id: NodeId) -> bool { - self.try_is_root(node_id) - .unwrap_or_else(|error| panic!("{error}")) - } - - pub fn try_is_root(&self, node_id: NodeId) -> Result { - let node_id = self.try_resolve_node_handle_(node_id)?; + pub fn is_root(&self, node_id: NodeId) -> Result { + let node_id = self.arena.resolve(node_id)?; Ok(self.arena.node(node_id).is_root()) } /// The node's last page hash, or None when it was never hashed. - pub fn get_last_hash_value(&self, node_id: NodeId) -> Option { - self.try_get_last_hash_value(node_id) - .unwrap_or_else(|error| panic!("{error}")) - } - - pub fn try_get_last_hash_value( - &self, - node_id: NodeId, - ) -> Result, TreeCoreRuntimeError> { - let node_id = self.try_resolve_node_handle_(node_id)?; + pub fn get_last_hash_value(&self, node_id: NodeId) -> Result, NodeAccessError> { + let node_id = self.arena.resolve(node_id)?; Ok(self .arena .node(node_id) @@ -3332,32 +3274,16 @@ impl UnifiedTreeCore { } /// The hash chain of the node's ancestors, in root-to-parent order. - pub fn get_prefix_hash_values(&self, node_id: NodeId) -> Vec { - self.try_get_prefix_hash_values(node_id) - .unwrap_or_else(|error| panic!("{error}")) - } - - pub fn try_get_prefix_hash_values( - &self, - node_id: NodeId, - ) -> Result, TreeCoreRuntimeError> { - let node_id = self.try_resolve_node_handle_(node_id)?; + pub fn get_prefix_hash_values(&self, node_id: NodeId) -> Result, NodeAccessError> { + let node_id = self.arena.resolve(node_id)?; Ok(self .arena .prefix_hash_values(self.arena.node(node_id).parent)) } /// The hash values owned by this node, excluding its ancestors. - pub fn get_hash_values(&self, node_id: NodeId) -> Vec { - self.try_get_hash_values(node_id) - .unwrap_or_else(|error| panic!("{error}")) - } - - pub fn try_get_hash_values( - &self, - node_id: NodeId, - ) -> Result, TreeCoreRuntimeError> { - let node_id = self.try_resolve_node_handle_(node_id)?; + pub fn get_hash_values(&self, node_id: NodeId) -> Result, NodeAccessError> { + let node_id = self.arena.resolve(node_id)?; Ok(self .arena .node(node_id) @@ -3371,7 +3297,7 @@ impl UnifiedTreeCore { node_id: NodeId, pass_prefix_keys: bool, ) -> Option { - let node_id = self.arena.try_resolve(node_id)?; + let node_id = self.arena.resolve(node_id).ok()?; let node = self.arena.node(node_id); if node.is_root() || !node.has_device_value(FULL) { return None; @@ -3401,7 +3327,7 @@ impl UnifiedTreeCore { node_id: NodeId, expected_key_length: usize, ) -> Option { - let node_id = self.arena.try_resolve(node_id)?; + let node_id = self.arena.resolve(node_id).ok()?; let node = self.arena.node(node_id); if !node.has_device_value(FULL) || node.key.atom_len() != expected_key_length { return None; @@ -3436,18 +3362,10 @@ impl UnifiedTreeCore { } /// Return input indices in depth-first, subtree-weight order. - pub fn dfs_weight_order(&self, node_ids: &[NodeId]) -> Vec { - self.try_dfs_weight_order(node_ids) - .unwrap_or_else(|error| panic!("{error}")) - } - - pub fn try_dfs_weight_order( - &self, - node_ids: &[NodeId], - ) -> Result, TreeCoreRuntimeError> { + pub fn dfs_weight_order(&self, node_ids: &[NodeId]) -> Result, NodeAccessError> { let mut node_to_indices: HashMap> = HashMap::new(); for (index, &node_id) in node_ids.iter().enumerate() { - let node_id = self.try_resolve_node_handle_(node_id)?; + let node_id = self.arena.resolve(node_id)?; node_to_indices.entry(node_id).or_default().push(index); } @@ -3525,8 +3443,17 @@ impl UnifiedTreeCore { cache_actions: &mut Vec, mut insert_result: Option<&mut InsertResult>, pool_storage_result: Option<&PoolTransferResult>, - ) { - let node_id = self.arena.resolve(node_id); + ) -> Result<(), NodeAccessError> { + let node_id = self.arena.resolve(node_id)?; + self.validate_pool_transfer_handles(comp_xfers.values().flatten())?; + if let Some(insert_result) = insert_result.as_deref() { + if let Some(last_device_node_id) = insert_result.last_device_node_id { + self.arena.resolve(last_device_node_id)?; + } + if let Some(inserted_host_node) = insert_result.inserted_host_node { + self.arena.resolve(inserted_host_node)?; + } + } for (component_type, transfers) in comp_xfers { self.component_by_type_(component_type) .commit_hicache_transfer( @@ -3539,6 +3466,7 @@ impl UnifiedTreeCore { pool_storage_result, ); } + Ok(()) } /// Commit a successful backup to the node. @@ -3547,8 +3475,9 @@ impl UnifiedTreeCore { node_id: NodeId, host_indices: Tensor, comp_xfers: HashMap>, - ) { - let node_id = self.arena.resolve(node_id); + ) -> Result<(), NodeAccessError> { + let node_id = self.arena.resolve(node_id)?; + self.validate_pool_transfer_handles(comp_xfers.values().flatten())?; let mut cache_actions: Vec = Vec::new(); if host_indices.numel() > 0 { let kv_xfer = PoolTransfer { @@ -3581,6 +3510,7 @@ impl UnifiedTreeCore { } assert!(cache_actions.is_empty()); // BACKUP_HOST emits no actions self.update_full_coexisting_host_tracking_(node_id); + Ok(()) } /// Commit a successful H->D load-back onto the node; the SWA full->swa mapping @@ -3591,9 +3521,16 @@ impl UnifiedTreeCore { device_indices: Tensor, mut kv_xfer: PoolTransfer, comp_xfers: HashMap>, - ) -> Vec { + ) -> Result, NodeAccessError> { let anchor_id = node_id; - let node_id = self.arena.resolve(node_id); + let node_id = self.arena.resolve(node_id)?; + let loaded_node_indices = kv_xfer + .nodes_to_load + .iter() + .flatten() + .map(|&loaded_id| self.arena.resolve(loaded_id)) + .collect::, _>>()?; + self.validate_pool_transfer_handles(comp_xfers.values().flatten())?; let mut cache_actions: Vec = Vec::new(); kv_xfer.device_indices = Some(device_indices); let nodes_to_load = kv_xfer.nodes_to_load.clone(); @@ -3601,8 +3538,9 @@ impl UnifiedTreeCore { // Pin Full KV host slots against duplicate reclaim until the ack. // Auxiliary pools have independent host locks and may legitimately // load the same radix node under a different anchor. - for &pinned_id in nodes_to_load.iter().flatten() { - let pinned_idx = self.arena.resolve(pinned_id); + for (&pinned_id, &pinned_idx) in + nodes_to_load.iter().flatten().zip(&loaded_node_indices) + { let pinned = self.arena.node_mut(pinned_idx); assert!( pinned.load_back_pending_id.is_none_or(|id| id == anchor_id), @@ -3623,8 +3561,7 @@ impl UnifiedTreeCore { /* insert_result = */ None, /* pool_storage_result = */ None, ); - for loaded_id in nodes_to_load.unwrap_or_default() { - let loaded_idx = self.arena.resolve(loaded_id); + for loaded_idx in loaded_node_indices { self.record_store_event_(loaded_idx, StorageMedium::Gpu); } for (component_type, transfers) in comp_xfers { @@ -3640,7 +3577,7 @@ impl UnifiedTreeCore { ); } self.update_evictable_leaf_sets_(node_id); - cache_actions + Ok(cache_actions) } /// Finalize load-back state along the anchor's root path. @@ -3648,8 +3585,8 @@ impl UnifiedTreeCore { /// Write-back clears matching Full KV source pins. Write-through has no /// pins, but both policies refresh Full host/device duplicate tracking once /// the device copies are visible. - pub fn finish_load_back(&mut self, anchor_node_id: NodeId) { - let mut node_id = Some(self.arena.resolve(anchor_node_id)); + pub fn finish_load_back(&mut self, anchor_node_id: NodeId) -> Result<(), NodeAccessError> { + let mut node_id = Some(self.arena.resolve(anchor_node_id)?); while let Some(idx) = node_id { if self.arena.node(idx).is_root() { break; @@ -3666,6 +3603,7 @@ impl UnifiedTreeCore { self.update_full_coexisting_host_tracking_(idx); node_id = self.arena.node(idx).try_parent(); } + Ok(()) } /// Mark every node covered by one in-flight write-through backup, and return @@ -3674,10 +3612,17 @@ impl UnifiedTreeCore { &mut self, node_ids: Vec, ack_id: NodeId, - ) -> Vec { - let mut marked: Vec<(usize, NodeId)> = Vec::with_capacity(node_ids.len()); - for node_id in node_ids { - let node_idx = self.arena.resolve(node_id); + ) -> Result, NodeAccessError> { + let node_indices = node_ids + .into_iter() + .map(|node_id| { + self.arena + .resolve(node_id) + .map(|node_idx| (node_id, node_idx)) + }) + .collect::, _>>()?; + let mut marked: Vec<(usize, NodeId)> = Vec::with_capacity(node_indices.len()); + for (node_id, node_idx) in node_indices { let depth = self.depth_from_root_(node_idx); let node = self.arena.node_mut(node_idx); assert!( @@ -3690,7 +3635,7 @@ impl UnifiedTreeCore { marked.push((depth, node_id)); } marked.sort_unstable(); - marked.into_iter().map(|(_, node_id)| node_id).collect() + Ok(marked.into_iter().map(|(_, node_id)| node_id).collect()) } fn depth_from_root_(&self, node_idx: NodeIdx_) -> usize { @@ -3705,9 +3650,16 @@ impl UnifiedTreeCore { /// Clear the write-through-pending mark (when it matches ack_id) and record the /// host store event for each acked node. - pub fn finish_write_through(&mut self, node_ids: Vec, ack_id: usize) { - for node_id in node_ids { - let node_idx = self.arena.resolve(node_id); + pub fn finish_write_through( + &mut self, + node_ids: Vec, + ack_id: usize, + ) -> Result<(), NodeAccessError> { + let node_indices = node_ids + .into_iter() + .map(|node_id| self.arena.resolve(node_id)) + .collect::, _>>()?; + for node_idx in node_indices { let node = self.arena.node_mut(node_idx); if node.write_through_pending_id == Some(ack_id) { node.write_through_pending_id = None; @@ -3715,6 +3667,7 @@ impl UnifiedTreeCore { } self.record_store_event_(node_idx, StorageMedium::Cpu); } + Ok(()) } /// Store an auxiliary component's device value onto a node and restamp @@ -3724,10 +3677,11 @@ impl UnifiedTreeCore { node_id: NodeId, component_type: ComponentType, value: Tensor, - ) { + ) -> Result<(), NodeAccessError> { + let node_idx = self.arena.resolve(node_id)?; self.assert_component_enabled_(component_type); - let node_idx = self.arena.resolve(node_id); self.set_component_device_value_(node_idx, component_type, value); + Ok(()) } /// Slot-keyed aux store (internal): set the device value and restamp the LRU. @@ -3756,10 +3710,10 @@ impl UnifiedTreeCore { &self, node_id: NodeId, component_type: ComponentType, - ) -> Option<&Tensor> { + ) -> Result, NodeAccessError> { + let node_id = self.arena.resolve(node_id)?; self.assert_component_enabled_(component_type); - self.arena - .try_device_value(self.arena.resolve(node_id), component_type) + Ok(self.arena.try_device_value(node_id, component_type)) } /// Whether the component's data is device-evicted but host-backed. @@ -3767,11 +3721,11 @@ impl UnifiedTreeCore { &self, node_id: NodeId, component_type: ComponentType, - ) -> bool { + ) -> Result { + let node_idx = self.arena.resolve(node_id)?; self.assert_component_enabled_(component_type); - let node_idx = self.arena.resolve(node_id); - !self.arena.has_device_value(node_idx, component_type) - && self.arena.has_host_value(node_idx, component_type) + Ok(!self.arena.has_device_value(node_idx, component_type) + && self.arena.has_host_value(node_idx, component_type)) } /// Verify tree-structure, leaf-set, LRU, size, and ongoing-op invariants; raise @@ -3832,7 +3786,7 @@ impl UnifiedTreeCore { )); } for (&node_handle, hashes) in &self.salted_event_hashes { - let Some(node_id) = self.arena.try_resolve(node_handle) else { + let Ok(node_id) = self.arena.resolve(node_handle) else { errors.push(format!( "[Events] salted hashes reference freed node {node_handle}" )); @@ -4189,25 +4143,25 @@ impl UnifiedTreeCore { // ── PART 5: Ongoing Operations ── for &(op_id, node_id) in ongoing_write_through { - match self.arena.try_resolve(node_id) { - None => { + match self.arena.resolve(node_id) { + Err(_) => { errors.push(format!("[Ongoing] write_through node {op_id} not in tree")); } - Some(idx) if self.arena.device_lock_ref(idx, FULL) == 0 => { + Ok(idx) if self.arena.device_lock_ref(idx, FULL) == 0 => { errors.push(format!("[Ongoing] write_through node {op_id} lock_ref=0")); } - Some(_) => {} + Ok(_) => {} } } for &(op_id, node_id) in ongoing_load_back { - match self.arena.try_resolve(node_id) { - None => { + match self.arena.resolve(node_id) { + Err(_) => { errors.push(format!("[Ongoing] load_back node {op_id} not in tree")); } - Some(idx) if self.arena.device_lock_ref(idx, FULL) == 0 => { + Ok(idx) if self.arena.device_lock_ref(idx, FULL) == 0 => { errors.push(format!("[Ongoing] load_back node {op_id} lock_ref=0")); } - Some(_) => {} + Ok(_) => {} } } // Reject load-back pins that would survive their operation. @@ -4263,43 +4217,53 @@ impl UnifiedTreeCore { /// Whether the external node handle is currently live. pub fn inspect_contains_node(&self, node_id: NodeId) -> bool { - self.arena.try_resolve(node_id).is_some() + self.arena.resolve(node_id).is_ok() } /// The parent node's external handle, or None for the root. - pub fn inspect_get_parent_node_id(&self, node_id: NodeId) -> Option { - let node_id = self.arena.resolve(node_id); - self.arena + pub fn inspect_get_parent_node_id( + &self, + node_id: NodeId, + ) -> Result, NodeAccessError> { + let node_id = self.arena.resolve(node_id)?; + Ok(self + .arena .node(node_id) .try_parent() - .map(|parent_id| self.arena.node(parent_id).id) + .map(|parent_id| self.arena.node(parent_id).id)) } /// A materialized snapshot of the node's child handles. - pub fn inspect_get_child_node_ids(&self, node_id: NodeId) -> Vec { - let node_id = self.arena.resolve(node_id); - self.arena + pub fn inspect_get_child_node_ids( + &self, + node_id: NodeId, + ) -> Result, NodeAccessError> { + let node_id = self.arena.resolve(node_id)?; + Ok(self + .arena .node(node_id) .children .values() .map(|&child_id| self.arena.node(child_id).id) - .collect() + .collect()) } /// Logical radix-key length in key atoms. - pub fn inspect_get_node_key_length(&self, node_id: NodeId) -> usize { - self.arena.node(self.arena.resolve(node_id)).key.atom_len() + pub fn inspect_get_node_key_length(&self, node_id: NodeId) -> Result { + let node_id = self.arena.resolve(node_id)?; + Ok(self.arena.node(node_id).key.atom_len()) } /// Materialized raw token ids spanned by the node key. - pub fn inspect_get_node_token_ids(&self, node_id: NodeId) -> Vec { - K::raw_token_ids(self.arena.node(self.arena.resolve(node_id)).key.as_ref()).into_owned() + pub fn inspect_get_node_token_ids(&self, node_id: NodeId) -> Result, NodeAccessError> { + let node_id = self.arena.resolve(node_id)?; + Ok(K::raw_token_ids(self.arena.node(node_id).key.as_ref()).into_owned()) } /// Whether this core's key representation uses overlapping bigrams. - pub fn inspect_is_node_key_bigram(&self, node_id: NodeId) -> bool { - let node = self.arena.node(self.arena.resolve(node_id)); - !node.is_root() && K::IS_BIGRAM + pub fn inspect_is_node_key_bigram(&self, node_id: NodeId) -> Result { + let node_id = self.arena.resolve(node_id)?; + Ok(!self.arena.node(node_id).is_root() && K::IS_BIGRAM) } /// A shallow tensor snapshot of a component's host value. @@ -4307,12 +4271,14 @@ impl UnifiedTreeCore { &self, node_id: NodeId, component_type: ComponentType, - ) -> Option { + ) -> Result, NodeAccessError> { + let node_id = self.arena.resolve(node_id)?; self.assert_component_enabled_(component_type); - self.arena - .node(self.arena.resolve(node_id)) + Ok(self + .arena + .node(node_id) .try_host_value(component_type) - .map(Tensor::shallow_clone) + .map(Tensor::shallow_clone)) } /// A component's device lock count on a node. @@ -4320,23 +4286,25 @@ impl UnifiedTreeCore { &self, node_id: NodeId, component_type: ComponentType, - ) -> u32 { + ) -> Result { + let node_id = self.arena.resolve(node_id)?; self.assert_component_enabled_(component_type); - self.arena - .node(self.arena.resolve(node_id)) - .device_lock_ref(component_type) + Ok(self.arena.node(node_id).device_lock_ref(component_type)) } /// A node's accumulated match count. - pub fn inspect_get_node_hit_count(&self, node_id: NodeId) -> i64 { - self.arena.node(self.arena.resolve(node_id)).hit_count + pub fn inspect_get_node_hit_count(&self, node_id: NodeId) -> Result { + let node_id = self.arena.resolve(node_id)?; + Ok(self.arena.node(node_id).hit_count) } /// A node's in-flight write-through acknowledgement id. - pub fn inspect_get_write_through_pending_id(&self, node_id: NodeId) -> Option { - self.arena - .node(self.arena.resolve(node_id)) - .write_through_pending_id + pub fn inspect_get_write_through_pending_id( + &self, + node_id: NodeId, + ) -> Result, NodeAccessError> { + let node_id = self.arena.resolve(node_id)?; + Ok(self.arena.node(node_id).write_through_pending_id) } /// Whether a node is in a component's device LRU. @@ -4344,12 +4312,12 @@ impl UnifiedTreeCore { &self, node_id: NodeId, component_type: ComponentType, - ) -> bool { + ) -> Result { if self.try_component_by_type_(component_type).is_none() { - return false; + return Ok(false); } - self.device_lru_list(component_type) - .in_list(Some(self.arena.resolve(node_id))) + let node_id = self.arena.resolve(node_id)?; + Ok(self.device_lru_list(component_type).in_list(Some(node_id))) } /// Whether a node is in a component's host LRU. @@ -4357,12 +4325,12 @@ impl UnifiedTreeCore { &self, node_id: NodeId, component_type: ComponentType, - ) -> bool { + ) -> Result { if self.try_component_by_type_(component_type).is_none() { - return false; + return Ok(false); } - self.host_lru_list(component_type) - .in_list(Some(self.arena.resolve(node_id))) + let node_id = self.arena.resolve(node_id)?; + Ok(self.host_lru_list(component_type).in_list(Some(node_id))) } /// Materialize a component's device LRU from most to least recent. @@ -4383,28 +4351,28 @@ impl UnifiedTreeCore { /// Whether a live node belongs to the device-evictable leaf set. pub fn inspect_is_device_evictable_leaf(&self, node_id: NodeId) -> bool { self.arena - .try_resolve(node_id) - .is_some_and(|node_id| self.evictable_device_leaves.contains(node_id)) + .resolve(node_id) + .is_ok_and(|node_id| self.evictable_device_leaves.contains(node_id)) } /// Whether a live node belongs to the host-evictable leaf set. pub fn inspect_is_host_evictable_leaf(&self, node_id: NodeId) -> bool { self.arena - .try_resolve(node_id) - .is_some_and(|node_id| self.evictable_host_leaves.contains(node_id)) + .resolve(node_id) + .is_ok_and(|node_id| self.evictable_host_leaves.contains(node_id)) } /// Whether the node is currently eligible as a Full device leaf. - pub fn inspect_is_device_leaf(&self, node_id: NodeId) -> bool { - let node_id = self.arena.resolve(node_id); + pub fn inspect_is_device_leaf(&self, node_id: NodeId) -> Result { + let node_id = self.arena.resolve(node_id)?; let node = self.arena.node(node_id); if node.is_root() || node.evicted() || node.is_device_locked() { - return false; + return Ok(false); } - !node + Ok(!node .children .values() - .any(|&child_id| self.arena.has_device_value(child_id, FULL)) + .any(|&child_id| self.arena.has_device_value(child_id, FULL))) } /// Materialize every live tree node handle. @@ -4425,9 +4393,10 @@ impl UnifiedTreeCore { &mut self, node_id: NodeId, hash_values: Option>, - ) { - let node_id = self.arena.resolve(node_id); + ) -> Result<(), NodeAccessError> { + let node_id = self.arena.resolve(node_id)?; self.arena.node_mut(node_id).hash_value = hash_values; + Ok(()) } /// Replace a component's device value without updating bookkeeping. @@ -4436,13 +4405,14 @@ impl UnifiedTreeCore { node_id: NodeId, component_type: ComponentType, value: Option, - ) { + ) -> Result<(), NodeAccessError> { + let node_id = self.arena.resolve(node_id)?; self.assert_component_enabled_(component_type); - let node_id = self.arena.resolve(node_id); self.arena .node_mut(node_id) .state_mut_(ValueSlotIdx::device(component_type)) .value = value; + Ok(()) } /// Replace a component's host value without updating bookkeeping. @@ -4451,13 +4421,14 @@ impl UnifiedTreeCore { node_id: NodeId, component_type: ComponentType, value: Option, - ) { + ) -> Result<(), NodeAccessError> { + let node_id = self.arena.resolve(node_id)?; self.assert_component_enabled_(component_type); - let node_id = self.arena.resolve(node_id); self.arena .node_mut(node_id) .state_mut_(ValueSlotIdx::host(component_type)) .value = value; + Ok(()) } /// Replace a component's device lock count without updating bookkeeping. @@ -4466,12 +4437,13 @@ impl UnifiedTreeCore { node_id: NodeId, component_type: ComponentType, lock_ref: u32, - ) { + ) -> Result<(), NodeAccessError> { + let node_id = self.arena.resolve(node_id)?; self.assert_component_enabled_(component_type); - let node_id = self.arena.resolve(node_id); self.arena .node_mut(node_id) .set_lock_ref_(ValueSlotIdx::device(component_type), lock_ref); + Ok(()) } /// Remove a node from a component's device LRU. @@ -4479,11 +4451,12 @@ impl UnifiedTreeCore { &mut self, node_id: NodeId, component_type: ComponentType, - ) { + ) -> Result<(), NodeAccessError> { + let node_id = self.arena.resolve(node_id)?; self.assert_component_enabled_(component_type); - let node_id = self.arena.resolve(node_id); self.device_lru_list_mut(component_type) .remove_node(node_id); + Ok(()) } /// Insert a node as a component's most-recent host-LRU entry. @@ -4491,10 +4464,11 @@ impl UnifiedTreeCore { &mut self, node_id: NodeId, component_type: ComponentType, - ) { + ) -> Result<(), NodeAccessError> { + let node_id = self.arena.resolve(node_id)?; self.assert_component_enabled_(component_type); - let node_id = self.arena.resolve(node_id); self.host_lru_list_mut(component_type).insert_mru(node_id); + Ok(()) } /// Replace a component's evictable-device token count. @@ -4518,9 +4492,13 @@ impl UnifiedTreeCore { } /// Refresh Full device/host duplicate tracking for a node. - pub fn inspect_update_duplicate_tracking(&mut self, node_id: NodeId) { - let node_id = self.arena.resolve(node_id); + pub fn inspect_update_duplicate_tracking( + &mut self, + node_id: NodeId, + ) -> Result<(), NodeAccessError> { + let node_id = self.arena.resolve(node_id)?; self.update_full_coexisting_host_tracking_(node_id); + Ok(()) } /// Advance one suspended insert walk step without flushing its pending actions. @@ -4543,9 +4521,9 @@ impl UnifiedTreeCore { node_id: NodeId, component_type: ComponentType, target: EvictLayer, - ) -> EvictionStepResult { + ) -> Result { + let node_id = self.arena.resolve(node_id)?; self.assert_component_enabled_(component_type); - let node_id = self.arena.resolve(node_id); let mut result = EvictionStepResult::default(); self.evict_component_and_detach_lru_( node_id, @@ -4555,7 +4533,7 @@ impl UnifiedTreeCore { target, Some(&mut result.tracker), ); - result + Ok(result) } /// Validate component locks for a cascade without mutating the tree. @@ -4564,9 +4542,9 @@ impl UnifiedTreeCore { node_id: NodeId, trigger_component_type: ComponentType, target: EvictLayer, - ) -> Result<(), String> { + ) -> Result<(), TreeCoreRuntimeError> { + let node_id = self.arena.resolve(node_id)?; self.assert_component_enabled_(trigger_component_type); - let node_id = self.arena.resolve(node_id); let is_leaf = match target { EvictLayer::Device => self.evictable_device_leaves.contains(node_id), EvictLayer::Host => self.evictable_host_leaves.contains(node_id), @@ -4585,14 +4563,18 @@ impl UnifiedTreeCore { is_leaf, trigger_priority, trigger_internal_priority, - )?; + ) + .map_err(TreeCoreRuntimeError::InspectionAssertion)?; } Ok(()) } /// Delete childless tombstone ancestors starting at `node_id`. - pub fn inspect_cleanup_tombstone_ancestors(&mut self, node_id: NodeId) -> EvictionStepResult { - let node_id = self.arena.resolve(node_id); + pub fn inspect_cleanup_tombstone_ancestors( + &mut self, + node_id: NodeId, + ) -> Result { + let node_id = self.arena.resolve(node_id)?; let mut result = EvictionStepResult::default(); self.iteratively_delete_tombstone_leaf_( node_id, @@ -4600,7 +4582,7 @@ impl UnifiedTreeCore { &mut result.device_frees, &mut result.host_frees, ); - result + Ok(result) } /// Run one component's real match-result finalizer. @@ -4611,16 +4593,33 @@ impl UnifiedTreeCore { params: &MatchPrefixParams<'_, K>, value_chunks: &[Tensor], best_value_len: usize, - ) -> MatchResult { - self.component_by_type_(component_type) - .finalize_match_result_in_tree_core(self, result, params, value_chunks, best_value_len) + ) -> Result { + let last_device_node_idx = self.arena.resolve(result.last_device_node_id)?; + self.arena.resolve(result.last_host_node_id)?; + let best_match_node_idx = self.arena.resolve(result.best_match_node_id)?; + Ok(self + .component_by_type_(component_type) + .finalize_match_result_in_tree_core( + self, + result, + last_device_node_idx, + best_match_node_idx, + params, + value_chunks, + best_value_len, + )) } /// Build the ordered device-to-host backup node list. - pub fn inspect_build_backup_node_ids(&self, node_id: NodeId, write_back: bool) -> Vec { - let node_id = self.arena.resolve(node_id); - self.build_backup_kv_action_(self.arena.node(node_id), write_back) - .node_ids + pub fn inspect_build_backup_node_ids( + &self, + node_id: NodeId, + write_back: bool, + ) -> Result, NodeAccessError> { + let node_id = self.arena.resolve(node_id)?; + Ok(self + .build_backup_kv_action_(self.arena.node(node_id), write_back) + .node_ids) } } diff --git a/test/registered/unit/mem_cache/test_rust_tree_core_integration.py b/test/registered/unit/mem_cache/test_rust_tree_core_integration.py index 347c1502e364..fb835ba65221 100644 --- a/test/registered/unit/mem_cache/test_rust_tree_core_integration.py +++ b/test/registered/unit/mem_cache/test_rust_tree_core_integration.py @@ -29,6 +29,7 @@ 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 = { + "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 + ), + "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 + ), + "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.demote(stale_root), - lambda: core.build_hicache_transfers( - ComponentType.FULL, stale_root, CacheTransferPhase.BACKUP_STORAGE + 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, + {}, ), - lambda: core.build_load_back_spec(stale_root), - lambda: core.get_hash_values(stale_root), - lambda: core.dfs_weight_order([stale_root]), + ) + 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__])) From 7cb55002eb5131a071c0308a20f69146d9ed0666 Mon Sep 17 00:00:00 2001 From: Jialin Ouyang Date: Sun, 6 Sep 2026 23:45:45 -0700 Subject: [PATCH 5/6] [Rust TreeCore] Fix rebased write-through tests --- rust/sglang-radix-tree/src/tests/unified_tree_core.rs | 7 ++++--- .../unit/mem_cache/test_rust_tree_core_integration.py | 2 +- 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/rust/sglang-radix-tree/src/tests/unified_tree_core.rs b/rust/sglang-radix-tree/src/tests/unified_tree_core.rs index f2d0f62d03ad..f97611ee6229 100644 --- a/rust/sglang-radix-tree/src/tests/unified_tree_core.rs +++ b/rust/sglang-radix-tree/src/tests/unified_tree_core.rs @@ -2218,16 +2218,17 @@ fn mark_write_through_pending_stamps_one_ack_on_every_published_node() { for node_id in [parent, leaf] { assert_eq!( tc.arena - .node(tc.arena.resolve(node_id)) + .node(tc.arena.resolve(node_id).expect("live test node")) .write_through_pending_id, Some(leaf) ); } - tc.finish_write_through(vec![parent, leaf], /* ack_id = */ leaf); + tc.finish_write_through(vec![parent, leaf], /* ack_id = */ leaf) + .expect("live test nodes"); for node_id in [parent, leaf] { assert_eq!( tc.arena - .node(tc.arena.resolve(node_id)) + .node(tc.arena.resolve(node_id).expect("live test node")) .write_through_pending_id, None ); diff --git a/test/registered/unit/mem_cache/test_rust_tree_core_integration.py b/test/registered/unit/mem_cache/test_rust_tree_core_integration.py index fb835ba65221..38858e85dcb7 100644 --- a/test/registered/unit/mem_cache/test_rust_tree_core_integration.py +++ b/test/registered/unit/mem_cache/test_rust_tree_core_integration.py @@ -269,7 +269,7 @@ def test_stale_handle_operations_raise_key_error_without_poisoning_the_core(): "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], stale_root ), "finish_write_through": lambda: core.finish_write_through( [stale_root], stale_root From 316b4a4da488f0ff7eeecf2e95f69ba41dbdca48 Mon Sep 17 00:00:00 2001 From: Jialin Ouyang Date: Mon, 7 Sep 2026 00:46:02 -0700 Subject: [PATCH 6/6] [CI] Include proto in Rust extension cache save key --- .github/workflows/_pr-test-rust-ext-build.yml | 1 + 1 file changed, 1 insertion(+) diff --git a/.github/workflows/_pr-test-rust-ext-build.yml b/.github/workflows/_pr-test-rust-ext-build.yml index 7c4ddfc4eec5..9c2034f6a40d 100644 --- a/.github/workflows/_pr-test-rust-ext-build.yml +++ b/.github/workflows/_pr-test-rust-ext-build.yml @@ -326,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