From 5a82db85f0567efa150891eb0e8c70f1218045ce Mon Sep 17 00:00:00 2001
From: fzyzcjy <5236035+fzyzcjy@users.noreply.github.com>
Date: Sat, 6 Jun 2026 09:07:40 +0800
Subject: [PATCH] Add scripted-runtime KV-pool and lock-ref exhauster
 primitives (#27412)

---
 .../test/scripted_runtime/context/api.py      | 18 +++++++
 .../context/kv_pool_exhauster.py              | 34 +++++++++++++
 .../context/lock_ref_exhauster.py             | 49 +++++++++++++++++++
 .../test/scripted_runtime/scheduler_hook.py   |  1 +
 4 files changed, 102 insertions(+)
 create mode 100644 python/sglang/test/scripted_runtime/context/kv_pool_exhauster.py
 create mode 100644 python/sglang/test/scripted_runtime/context/lock_ref_exhauster.py

diff --git a/python/sglang/test/scripted_runtime/context/api.py b/python/sglang/test/scripted_runtime/context/api.py
index 180db6a92..26090db1a 100644
--- a/python/sglang/test/scripted_runtime/context/api.py
+++ b/python/sglang/test/scripted_runtime/context/api.py
@@ -9,6 +9,12 @@ from sglang.test.scripted_runtime.context import (
     queries,
     radix,
 )
+from sglang.test.scripted_runtime.context.kv_pool_exhauster import (
+    ScriptedKvPoolExhauster,
+)
+from sglang.test.scripted_runtime.context.lock_ref_exhauster import (
+    ScriptedLockRefExhauster,
+)
 from sglang.test.scripted_runtime.context.req_starter import ScriptedContextReqStarter
 
 if TYPE_CHECKING:
@@ -41,6 +47,8 @@ class ScriptedContext:
         self._http_poster = http_poster
 
         self._seen_rids: set[str] = set()
+        self._kv_exhauster = ScriptedKvPoolExhauster(self.scheduler)
+        self._lock_ref_exhauster = ScriptedLockRefExhauster(self.scheduler)
         self._req_starter = ScriptedContextReqStarter(self)
 
     def start_req(
@@ -93,6 +101,16 @@ class ScriptedContext:
         ), "evict_radix currently supports only full eviction (prefix_tokens=None)"
         return lifecycle.flush_cache(self)
 
+    def exhaust_kv(self, *, leave_pages: int) -> None:
+        return self._kv_exhauster.exhaust(leave_pages=leave_pages)
+
+    def exhaust_lock_refs(self, *, leave_refs: int) -> None:
+        return self._lock_ref_exhauster.exhaust(leave_refs=leave_refs)
+
+    def _release_exhausted_pools(self) -> None:
+        self._kv_exhauster.release()
+        self._lock_ref_exhauster.release()
+
     def get_all_node_hit_counts(self) -> Dict[int, int]:
         return radix.get_all_node_hit_counts(self)
 
diff --git a/python/sglang/test/scripted_runtime/context/kv_pool_exhauster.py b/python/sglang/test/scripted_runtime/context/kv_pool_exhauster.py
new file mode 100644
index 000000000..695774bed
--- /dev/null
+++ b/python/sglang/test/scripted_runtime/context/kv_pool_exhauster.py
@@ -0,0 +1,34 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, List
+
+if TYPE_CHECKING:
+    import torch
+
+    from sglang.srt.managers.scheduler import Scheduler
+
+
+class ScriptedKvPoolExhauster:
+
+    def __init__(self, scheduler: "Scheduler") -> None:
+        self.scheduler = scheduler
+        self._held: List["torch.Tensor"] = []
+
+    def exhaust(self, *, leave_pages: int) -> None:
+        allocator = self.scheduler.token_to_kv_pool_allocator
+
+        leave_tokens = leave_pages * self.scheduler.page_size
+        need = allocator.available_size() - leave_tokens
+        if need <= 0:
+            return
+
+        held = allocator.alloc(need)
+        assert (
+            held is not None
+        ), f"exhaust_kv: allocator could not grab {need} tokens to create pressure"
+        self._held.append(held)
+
+    def release(self) -> None:
+        for held in self._held:
+            self.scheduler.token_to_kv_pool_allocator.free(held)
+        self._held.clear()
diff --git a/python/sglang/test/scripted_runtime/context/lock_ref_exhauster.py b/python/sglang/test/scripted_runtime/context/lock_ref_exhauster.py
new file mode 100644
index 000000000..f4257d7dc
--- /dev/null
+++ b/python/sglang/test/scripted_runtime/context/lock_ref_exhauster.py
@@ -0,0 +1,49 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, List
+
+from sglang.test.scripted_runtime.context.radix import _node_lock_ref
+
+if TYPE_CHECKING:
+    from sglang.srt.managers.scheduler import Scheduler
+
+
+class ScriptedLockRefExhauster:
+
+    def __init__(self, scheduler: "Scheduler") -> None:
+        self.scheduler = scheduler
+        self._locked: List[Any] = []
+
+    def exhaust(self, *, leave_refs: int) -> None:
+        tree_cache = self.scheduler.tree_cache
+        if tree_cache.disable:
+            return
+
+        while True:
+            evictable = self._evictable_nodes()
+            if len(evictable) <= leave_refs:
+                return
+
+            target = evictable[0]
+            tree_cache.inc_lock_ref(target)
+
+            newly_locked = [node for node in evictable if _node_lock_ref(node) > 0]
+            if not newly_locked:
+                return
+            self._locked.append(target)
+
+    def release(self) -> None:
+        tree_cache = self.scheduler.tree_cache
+        for node in self._locked:
+            tree_cache.dec_lock_ref(node)
+        self._locked.clear()
+
+    def _evictable_nodes(self) -> List[Any]:
+        evictable: List[Any] = []
+        stack = list(self.scheduler.tree_cache.root_node.children.values())
+        while stack:
+            node = stack.pop()
+            if _node_lock_ref(node) == 0:
+                evictable.append(node)
+            stack.extend(node.children.values())
+        return evictable
diff --git a/python/sglang/test/scripted_runtime/scheduler_hook.py b/python/sglang/test/scripted_runtime/scheduler_hook.py
index ef76322bc..ec46feb3c 100644
--- a/python/sglang/test/scripted_runtime/scheduler_hook.py
+++ b/python/sglang/test/scripted_runtime/scheduler_hook.py
@@ -50,6 +50,7 @@ class ScriptedBatchRecord:
 def _reset_engine_state(ctx: ScriptedContext) -> Generator:
     scheduler = ctx.scheduler
 
+    ctx._release_exhausted_pools()
     ctx.abort_all()
     for _ in range(RESET_DRAIN_MAX_STEPS):
         yield