chore: import upstream snapshot with attribution

python/sglang/test/scripted_runtime/context/__init__.py

@@ -0,0 +1,3 @@
+from sglang.test.scripted_runtime.context.api import ScriptedContext
+
+__all__ = ["ScriptedContext"]

python/sglang/test/scripted_runtime/context/api.py

@@ -0,0 +1,164 @@
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING, Dict, List, Literal, Optional
+
+from sglang.test.scripted_runtime.context import (
+    engine,
+    lifecycle,
+    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:
+    from sglang.srt.managers.schedule_batch import Req
+    from sglang.test.scripted_runtime.background_http_poster import BackgroundHttpPoster
+    from sglang.test.scripted_runtime.req_handle import ScriptedReqHandle
+    from sglang.test.scripted_runtime.scheduler_hook import ScriptedSchedulerHook
+    from sglang.test.scripted_runtime.tokenizer_recv_proxy import (
+        ScriptedTokenizerRecvProxy,
+    )
+
+logger = logging.getLogger(__name__)
+
+
+class ScriptedContext:
+
+    def __init__(
+        self,
+        *,
+        scheduler_hook: ScriptedSchedulerHook,
+        tokenizer_recv_proxy: Optional[ScriptedTokenizerRecvProxy],
+        http_poster: BackgroundHttpPoster,
+    ) -> None:
+        assert (
+            scheduler_hook._is_driver
+        ), "ScriptedContext only exists on the driver rank"
+        self.scheduler = scheduler_hook.scheduler
+        self._scheduler_hook = scheduler_hook
+        self._tokenizer_recv_proxy = tokenizer_recv_proxy
+        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(
+        self,
+        *,
+        prompt_len: int,
+        max_new_tokens: int = 8,
+        rid: Optional[str] = None,
+        ignore_eos: bool = False,
+        priority: Optional[int] = None,
+        dp_rank: Optional[int] = None,
+        prompt_token: int = 1,
+        return_logprob: bool = False,
+        logprob_start_len: Optional[int] = None,
+        top_logprobs_num: Optional[int] = None,
+        stop_token_ids: Optional[List[int]] = None,
+        temperature: Optional[float] = None,
+        lora_path: Optional[str] = None,
+    ) -> ScriptedReqHandle:
+        return self._req_starter.start_req(
+            prompt_len=prompt_len,
+            max_new_tokens=max_new_tokens,
+            rid=rid,
+            ignore_eos=ignore_eos,
+            priority=priority,
+            dp_rank=dp_rank,
+            prompt_token=prompt_token,
+            return_logprob=return_logprob,
+            logprob_start_len=logprob_start_len,
+            top_logprobs_num=top_logprobs_num,
+            stop_token_ids=stop_token_ids,
+            temperature=temperature,
+            lora_path=lora_path,
+        )
+
+    def pause_generation(self, *, mode: Literal["retract", "in_place"]) -> None:
+        return lifecycle.pause_generation(self, mode=mode)
+
+    def continue_generation(self, *, torch_empty_cache: bool = False) -> None:
+        return lifecycle.continue_generation(self, torch_empty_cache=torch_empty_cache)
+
+    def abort_all(self) -> None:
+        return lifecycle.abort_all(self)
+
+    def abort(self, handle: ScriptedReqHandle, *, await_arrival: bool = True) -> None:
+        return lifecycle.abort(self, rid=handle.rid, await_arrival=await_arrival)
+
+    def flush_cache(self) -> None:
+        return lifecycle.flush_cache(self)
+
+    def evict_radix(self, *, prefix_tokens: Optional[List[int]]) -> None:
+        assert (
+            prefix_tokens is None
+        ), "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)
+
+    def get_all_node_lock_refs(self) -> Dict[int, int]:
+        return radix.get_all_node_lock_refs(self)
+
+    @property
+    def is_idle(self) -> bool:
+        return queries.is_idle(self)
+
+    @property
+    def is_fully_idle(self) -> bool:
+        return queries.is_fully_idle(self)
+
+    @property
+    def last_batch_forward_mode(self) -> Optional[str]:
+        return queries.last_batch_forward_mode(self)
+
+    def find_req_by_rid(self, rid: str) -> Optional[Req]:
+        return queries.find_req_by_rid(self, rid)
+
+    def is_finished(self, rid: str) -> bool:
+        return queries.is_finished(self, rid)
+
+    def is_chunking(self, rid: str) -> bool:
+        return queries.is_chunking(self, rid)
+
+    def status(self, rid: str) -> str:
+        return queries.status(self, rid)
+
+    def remaining_prompt_tokens(self, rid: str) -> int:
+        return queries.remaining_prompt_tokens(self, rid)
+
+    def list_active_reqs(self) -> List[Req]:
+        return queries.list_active_reqs(self)
+
+    def chunks_done(self, rid: str) -> int:
+        return queries.chunks_done(self, rid)
+
+    def chunked_parks(self, rid: str) -> int:
+        return queries.chunked_parks(self, rid)
+
+    def batch_composition(self) -> Dict[str, List[str]]:
+        return queries.batch_composition(self)
+
+    def engine_stats(self) -> Dict[str, int]:
+        return engine.engine_stats(self)

python/sglang/test/scripted_runtime/context/engine.py

@@ -0,0 +1,16 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Dict
+
+if TYPE_CHECKING:
+    from sglang.test.scripted_runtime.context.api import ScriptedContext
+
+
+def engine_stats(ctx: ScriptedContext) -> Dict[str, int]:
+    s = ctx.scheduler
+    return {
+        "kv_pool_free": s.token_to_kv_pool_allocator.available_size(),
+        "req_pool_free": s.req_to_token_pool.available_size(),
+        "req_pool_total": s.req_to_token_pool.size,
+        "page_size": s.page_size,
+    }

python/sglang/test/scripted_runtime/context/http_post.py

@@ -0,0 +1,52 @@
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING, Any, Callable, Dict, Optional
+
+if TYPE_CHECKING:
+    from sglang.test.scripted_runtime.context.api import ScriptedContext
+
+logger = logging.getLogger(__name__)
+
+RECV_MSG_ARRIVAL_TIMEOUT_S: float = 60.0
+
+
+def _http_post_and_await_recv_msg(
+    ctx: ScriptedContext,
+    *,
+    path: str,
+    json: Optional[Dict[str, Any]],
+    predicate: Callable[[Any], bool],
+    description: str,
+    timeout_s: float = RECV_MSG_ARRIVAL_TIMEOUT_S,
+) -> None:
+    _submit_post(ctx, path=path, json=json)
+    ctx._tokenizer_recv_proxy.wait_until_arrived(
+        predicate,
+        timeout_s=timeout_s,
+        description=description,
+    )
+
+
+def _http_post_fire_and_forget(
+    ctx: ScriptedContext,
+    *,
+    path: str,
+    json: Optional[Dict[str, Any]],
+) -> None:
+    _submit_post(ctx, path=path, json=json)
+
+
+def _submit_post(
+    ctx: ScriptedContext,
+    *,
+    path: str,
+    json: Optional[Dict[str, Any]],
+) -> None:
+    server_args = ctx.scheduler.server_args
+    url = f"http://{server_args.host}:{server_args.port}{path}"
+
+    async def _post() -> None:
+        await ctx._http_poster.post(url, json)
+
+    ctx._http_poster.submit_coro(_post())

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()

python/sglang/test/scripted_runtime/context/lifecycle.py

@@ -0,0 +1,85 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Literal
+
+from sglang.srt.managers.io_struct import (
+    AbortReq,
+    ContinueGenerationReqInput,
+    FlushCacheReqInput,
+    PauseGenerationReqInput,
+)
+from sglang.test.scripted_runtime.context.http_post import (
+    _http_post_and_await_recv_msg,
+    _http_post_fire_and_forget,
+)
+
+if TYPE_CHECKING:
+    from sglang.test.scripted_runtime.context.api import ScriptedContext
+
+
+def _await_control(
+    ctx: ScriptedContext,
+    *,
+    path: str,
+    json,
+    expect_type: type,
+    await_arrival: bool = True,
+) -> None:
+    if not await_arrival:
+        _http_post_fire_and_forget(ctx, path=path, json=json)
+        return
+    _http_post_and_await_recv_msg(
+        ctx,
+        path=path,
+        json=json,
+        predicate=lambda obj: isinstance(obj, expect_type),
+        description=expect_type.__name__,
+    )
+
+
+def pause_generation(
+    ctx: ScriptedContext, *, mode: Literal["retract", "in_place"]
+) -> None:
+    _await_control(
+        ctx,
+        path="/pause_generation",
+        json={"mode": mode},
+        expect_type=PauseGenerationReqInput,
+    )
+
+
+def continue_generation(ctx: ScriptedContext, *, torch_empty_cache: bool) -> None:
+    _await_control(
+        ctx,
+        path="/continue_generation",
+        json={"torch_empty_cache": torch_empty_cache},
+        expect_type=ContinueGenerationReqInput,
+    )
+
+
+def abort_all(ctx: ScriptedContext) -> None:
+    _await_control(
+        ctx,
+        path="/abort_request",
+        json={"rid": "", "abort_all": True},
+        expect_type=AbortReq,
+    )
+
+
+def abort(ctx: ScriptedContext, *, rid: str, await_arrival: bool = True) -> None:
+    _await_control(
+        ctx,
+        path="/abort_request",
+        json={"rid": rid, "abort_all": False},
+        expect_type=AbortReq,
+        await_arrival=await_arrival,
+    )
+
+
+def flush_cache(ctx: ScriptedContext) -> None:
+    _await_control(
+        ctx,
+        path="/flush_cache",
+        json=None,
+        expect_type=FlushCacheReqInput,
+    )

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

python/sglang/test/scripted_runtime/context/queries.py

@@ -0,0 +1,139 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Dict, Iterator, List, Optional
+
+if TYPE_CHECKING:
+    from sglang.srt.managers.schedule_batch import Req
+    from sglang.test.scripted_runtime.context.api import ScriptedContext
+
+
+def _get_all_reqs(ctx: ScriptedContext) -> Iterator[Req]:
+    s = ctx.scheduler
+    if s.chunked_req is not None:
+        yield s.chunked_req
+    yield from s.waiting_queue
+    if s.ps.pp_size > 1:
+        for mb in (*s.mbs, *s.last_mbs, *s.running_mbs):
+            if mb is not None:
+                yield from mb.reqs
+    else:
+        if s.running_batch is not None:
+            yield from s.running_batch.reqs
+        if s.last_batch is not None:
+            yield from s.last_batch.reqs
+
+
+def list_active_reqs(ctx: ScriptedContext) -> List[Req]:
+    return list(set(_get_all_reqs(ctx)))
+
+
+def batch_composition(ctx: ScriptedContext) -> Dict[str, List[str]]:
+    s = ctx.scheduler
+    chunked_rid = s.chunked_req.rid if s.chunked_req is not None else None
+    chunked = [chunked_rid] if chunked_rid is not None else []
+    running = (
+        [r.rid for r in s.running_batch.reqs] if s.running_batch is not None else []
+    )
+
+    prefill: List[str] = []
+    decode: List[str] = []
+    batch = s.last_batch
+    if batch is not None and not batch.is_empty() and batch.forward_mode is not None:
+        bucket = prefill if batch.forward_mode.is_extend() else decode
+        bucket.extend(r.rid for r in batch.reqs if r.rid != chunked_rid)
+
+    return {
+        "prefill": prefill,
+        "decode": decode,
+        "chunked": chunked,
+        "running": running,
+    }
+
+
+def is_idle(ctx: ScriptedContext) -> bool:
+    s = ctx.scheduler
+    return (
+        s.chunked_req is None
+        and len(s.waiting_queue) == 0
+        and (s.running_batch is None or s.running_batch.is_empty())
+    )
+
+
+def is_fully_idle(ctx: ScriptedContext) -> bool:
+    s = ctx.scheduler
+    return is_idle(ctx) and (s.last_batch is None or s.last_batch.is_empty())
+
+
+def last_batch_forward_mode(ctx: ScriptedContext) -> Optional[str]:
+    s = ctx.scheduler
+    if s.last_batch is not None and s.last_batch.forward_mode is not None:
+        return s.last_batch.forward_mode.name
+    return None
+
+
+def find_req_by_rid(ctx: ScriptedContext, rid: str) -> Optional[Req]:
+    req = next((r for r in _get_all_reqs(ctx) if r.rid == rid), None)
+    if req is not None:
+        ctx._seen_rids.add(rid)
+    return req
+
+
+def is_finished(ctx: ScriptedContext, rid: str) -> bool:
+    req = find_req_by_rid(ctx, rid)
+    if req is not None:
+        return req.finished()
+    if rid in ctx._seen_rids:
+        return True
+    # Fallback: if the req ran in a forward batch (recorded in _batch_log) but
+    # is now absent from all active scheduler sets, it must have finished.
+    # This catches requests that completed without ever being observed via
+    # find_req_by_rid (e.g. when Python short-circuit evaluation prevents the
+    # query while another request is still running).
+    log = ctx._scheduler_hook._batch_log
+    if any(rid in record.rids for record in log):
+        ctx._seen_rids.add(rid)
+        return True
+    return False
+
+
+def is_chunking(ctx: ScriptedContext, rid: str) -> bool:
+    s = ctx.scheduler
+    return s.chunked_req is not None and s.chunked_req.rid == rid
+
+
+def status(ctx: ScriptedContext, rid: str) -> str:
+    s = ctx.scheduler
+    if rid in {r.rid for r in s.waiting_queue}:
+        return "waiting"
+    req = find_req_by_rid(ctx, rid)
+    if req is not None:
+        return "finished" if req.finished() else "running"
+    if rid in ctx._seen_rids:
+        return "finished"
+    return "unknown"
+
+
+def remaining_prompt_tokens(ctx: ScriptedContext, rid: str) -> int:
+    req = find_req_by_rid(ctx, rid)
+    if req is None:
+        return 0
+    return max(0, len(req.origin_input_ids) - req.kv_committed_len)
+
+
+def chunks_done(ctx: ScriptedContext, rid: str) -> int:
+    log = ctx._scheduler_hook._batch_log
+    held = sum(1 for record in log if record.chunked_rid == rid and rid in record.rids)
+    if held == 0:
+        return 0
+    completed = any(
+        rid in record.extend_rids and record.chunked_rid != rid for record in log
+    )
+    return held + (1 if completed else 0)
+
+
+def chunked_parks(ctx: ScriptedContext, rid: str) -> int:
+    return sum(
+        1
+        for record in ctx._scheduler_hook._batch_log
+        if record.chunked_rid == rid and rid not in record.rids
+    )

python/sglang/test/scripted_runtime/context/radix.py

@@ -0,0 +1,34 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, Callable, Dict
+
+from sglang.srt.mem_cache.swa_radix_cache import TreeNode as SWATreeNode
+
+if TYPE_CHECKING:
+    from sglang.test.scripted_runtime.context.api import ScriptedContext
+
+
+def get_all_node_hit_counts(ctx: ScriptedContext) -> Dict[int, int]:
+    return _collect_node_attr(ctx, lambda node: node.hit_count)
+
+
+def get_all_node_lock_refs(ctx: ScriptedContext) -> Dict[int, int]:
+    return _collect_node_attr(ctx, _node_lock_ref)
+
+
+def _node_lock_ref(node: Any) -> int:
+    if isinstance(node, SWATreeNode):
+        return node.full_lock_ref + node.swa_lock_ref
+    return node.lock_ref
+
+
+def _collect_node_attr(
+    ctx: ScriptedContext, get_value: Callable[[Any], int]
+) -> Dict[int, int]:
+    values: Dict[int, int] = {}
+    stack = list(ctx.scheduler.tree_cache.root_node.children.values())
+    while stack:
+        node = stack.pop()
+        values[node.id] = get_value(node)
+        stack.extend(node.children.values())
+    return values

python/sglang/test/scripted_runtime/context/req_starter.py

@@ -0,0 +1,73 @@
+from __future__ import annotations
+
+import uuid
+from typing import TYPE_CHECKING, List, Optional
+
+from sglang.test.scripted_runtime.context.http_post import (
+    _http_post_and_await_recv_msg,
+)
+from sglang.test.scripted_runtime.req_handle import ScriptedReqHandle
+
+if TYPE_CHECKING:
+    from sglang.test.scripted_runtime.context.api import ScriptedContext
+
+
+class ScriptedContextReqStarter:
+    def __init__(self, ctx: ScriptedContext) -> None:
+        self._ctx = ctx
+        self._req_counter = 0
+
+    def start_req(
+        self,
+        *,
+        prompt_len: int,
+        max_new_tokens: int,
+        rid: Optional[str],
+        ignore_eos: bool,
+        priority: Optional[int],
+        dp_rank: Optional[int],
+        prompt_token: int = 1,
+        return_logprob: bool = False,
+        logprob_start_len: Optional[int] = None,
+        top_logprobs_num: Optional[int] = None,
+        stop_token_ids: Optional[List[int]] = None,
+        temperature: Optional[float] = None,
+        lora_path: Optional[str] = None,
+    ) -> ScriptedReqHandle:
+        ctx = self._ctx
+
+        if rid is None:
+            rid = f"scripted-{self._req_counter}-{uuid.uuid4().hex}"
+            self._req_counter += 1
+
+        sampling_params = {"max_new_tokens": max_new_tokens, "ignore_eos": ignore_eos}
+        if stop_token_ids is not None:
+            sampling_params["stop_token_ids"] = stop_token_ids
+        if temperature is not None:
+            sampling_params["temperature"] = temperature
+        payload = {
+            "input_ids": [prompt_token] * prompt_len,
+            "sampling_params": sampling_params,
+            "rid": rid,
+            "stream": True,
+        }
+        payload["return_logprob"] = return_logprob
+        if logprob_start_len is not None:
+            payload["logprob_start_len"] = logprob_start_len
+        if top_logprobs_num is not None:
+            payload["top_logprobs_num"] = top_logprobs_num
+        if priority is not None:
+            payload["priority"] = priority
+        if dp_rank is not None:
+            payload["routed_dp_rank"] = dp_rank
+        if lora_path is not None:
+            payload["lora_path"] = lora_path
+        _http_post_and_await_recv_msg(
+            ctx,
+            path="/generate",
+            json=payload,
+            predicate=lambda obj: getattr(obj, "rid", None) == rid,
+            description=f"request with rid {rid!r}",
+        )
+
+        return ScriptedReqHandle(rid=rid, context=ctx)