Add scripted-runtime harness core and wire scheduler/IPC hooks (#27411)

python/sglang/srt/environ.py

@@ -337,6 +337,11 @@ class Envs:
     SGLANG_TEST_PD_DISAGG_DEVICES = EnvStr(None)
     SGLANG_TEST_FORCE_OPTIMISTIC_PREFILL_RETRY_PROB = EnvFloat(0.0)
 
+    SGLANG_TEST_SCRIPTED_RUNTIME = EnvBool(False)
+    SGLANG_TEST_SCRIPTED_RUNTIME_IPC_ADDR = EnvStr(None)
+    SGLANG_TEST_SCRIPTED_RUNTIME_OUT_OF_BAND_ERROR_PATH = EnvStr(None)
+    SGLANG_TEST_SCRIPTED_RUNTIME_SYS_PATH_ENTRY = EnvStr(None)
+
     # Model Parallel
     SGLANG_USE_MESSAGE_QUEUE_BROADCASTER = EnvBool(True)
     SGLANG_ONE_VISIBLE_DEVICE_PER_PROCESS = EnvBool(False)

python/sglang/srt/managers/scheduler.py

@@ -545,6 +545,8 @@ class Scheduler(
         # Init the grammar backend for constrained generation
         self.init_grammar_manager()
 
+        self.maybe_init_scripted_scheduler_hook()
+
         self.init_request_receiver()
 
         self.init_dp_attn_adapter()
@@ -611,6 +613,7 @@ class Scheduler(
                 self.ps.attn_tp_rank == 0
                 or self.server_args.enable_metrics_for_all_schedulers
             ),
+            enable_scripted_runtime=envs.SGLANG_TEST_SCRIPTED_RUNTIME.get(),
         )
 
         self.load_snapshot_writer = None
@@ -1601,6 +1604,19 @@ class Scheduler(
     def init_grammar_manager(self) -> None:
         self.grammar_manager = GrammarManager(self)
 
+    def maybe_init_scripted_scheduler_hook(self) -> None:
+        if envs.SGLANG_TEST_SCRIPTED_RUNTIME.get():
+            from sglang.test.scripted_runtime.scheduler_hook import (
+                ScriptedSchedulerHook,
+            )
+
+            self.scripted_scheduler_hook = ScriptedSchedulerHook(
+                scheduler=self,
+                tokenizer_recv_proxy=self.ipc_channels.recv_from_tokenizer,
+            )
+        else:
+            self.scripted_scheduler_hook = None
+
     def init_request_receiver(self) -> None:
         self.request_receiver = SchedulerRequestReceiver(
             recv_from_tokenizer=self.ipc_channels.recv_from_tokenizer,
@@ -1623,6 +1639,7 @@ class Scheduler(
             get_last_forward_mode=lambda: (
                 self.last_batch.forward_mode if self.last_batch is not None else None
             ),
+            scripted_scheduler_hook=self.scripted_scheduler_hook,
         )
 
     def init_dp_attn_adapter(self) -> None:
@@ -2978,6 +2995,9 @@ class Scheduler(
         self.forward_ct += 1
         batch.forward_iter = self.forward_ct
 
+        if self.scripted_scheduler_hook is not None:
+            self.scripted_scheduler_hook.on_run_batch(batch)
+
         # Whether to run the profiler
         self.profiler_manager._profile_batch_predicate(batch)
         if self.forward_sleep_time is not None:

python/sglang/srt/managers/scheduler_components/ipc_channels.py

@@ -1,5 +1,5 @@
 from dataclasses import dataclass
-from typing import Optional
+from typing import TYPE_CHECKING, Optional, Union
 
 import zmq
 
@@ -7,10 +7,15 @@ from sglang.srt.managers.scheduler_components.output_sender import SenderWrapper
 from sglang.srt.server_args import PortArgs
 from sglang.srt.utils.network import get_zmq_socket
 
+if TYPE_CHECKING:
+    from sglang.test.scripted_runtime.tokenizer_recv_proxy import (
+        ScriptedTokenizerRecvProxy,
+    )
+
 
 @dataclass(frozen=True, slots=True, kw_only=True)
 class SchedulerIpcChannels:
-    recv_from_tokenizer: Optional[zmq.Socket]
+    recv_from_tokenizer: Union[zmq.Socket, "ScriptedTokenizerRecvProxy"]
     recv_from_rpc: Optional[zmq.Socket]
     send_to_tokenizer: SenderWrapper
     send_to_detokenizer: SenderWrapper
@@ -24,6 +29,7 @@ class SchedulerIpcChannels:
         is_rank_zero: bool,
         skip_tokenizer_init: bool,
         metrics_enabled: bool,
+        enable_scripted_runtime: bool,
     ) -> "SchedulerIpcChannels":
         context = zmq.Context(2)
 
@@ -31,6 +37,14 @@ class SchedulerIpcChannels:
             recv_from_tokenizer = get_zmq_socket(
                 context, zmq.PULL, port_args.scheduler_input_ipc_name, False
             )
+            if enable_scripted_runtime:
+                from sglang.test.scripted_runtime.tokenizer_recv_proxy import (
+                    ScriptedTokenizerRecvProxy,
+                )
+
+                recv_from_tokenizer = ScriptedTokenizerRecvProxy(
+                    underlying=recv_from_tokenizer
+                )
             recv_from_rpc = get_zmq_socket(
                 context, zmq.DEALER, port_args.rpc_ipc_name, False
             )

python/sglang/srt/managers/scheduler_components/request_receiver.py

@@ -34,11 +34,15 @@ if TYPE_CHECKING:
     from sglang.srt.configs.model_config import ModelConfig
     from sglang.srt.distributed.parallel_state_wrapper import ParallelState
     from sglang.srt.server_args import ServerArgs
+    from sglang.test.scripted_runtime.scheduler_hook import ScriptedSchedulerHook
+    from sglang.test.scripted_runtime.tokenizer_recv_proxy import (
+        ScriptedTokenizerRecvProxy,
+    )
 
 
 @dataclass(kw_only=True, slots=True, frozen=True)
 class SchedulerRequestReceiver:
-    recv_from_tokenizer: zmq.Socket
+    recv_from_tokenizer: Union[zmq.Socket, "ScriptedTokenizerRecvProxy"]
     recv_from_rpc: Optional[zmq.Socket]
     recv_skipper: Any
     input_blocker: Any
@@ -56,6 +60,7 @@ class SchedulerRequestReceiver:
     max_recv_per_poll: int
     stream_output: Callable[..., None]
     get_last_forward_mode: Callable[[], Any]
+    scripted_scheduler_hook: Optional["ScriptedSchedulerHook"] = None
 
     def recv_limit_reached(self, num_recv_reqs: int) -> bool:
         if self.max_recv_per_poll < 0:
@@ -67,6 +72,9 @@ class SchedulerRequestReceiver:
     ) -> List[Union[TokenizedGenerateReqInput, TokenizedEmbeddingReqInput, Any]]:
         """Receive results at tp_rank = 0 and broadcast it to all other TP ranks."""
 
+        if self.scripted_scheduler_hook is not None:
+            self.scripted_scheduler_hook.step()
+
         if self.recv_skipper is not None:
             if not self.recv_skipper.handle(self.get_last_forward_mode()):
                 return []

python/sglang/srt/utils/network.py

@@ -231,7 +231,7 @@ def config_socket(socket, socket_type: zmq.SocketType):
         set_send_opt()
     elif socket_type == zmq.PULL:
         set_recv_opt()
-    elif socket_type in [zmq.DEALER, zmq.REQ, zmq.REP]:
+    elif socket_type in [zmq.DEALER, zmq.REQ, zmq.REP, zmq.PAIR]:
         set_send_opt()
         set_recv_opt()
     else:

python/sglang/test/scripted_runtime/background_http_poster.py

@@ -0,0 +1,68 @@
+from __future__ import annotations
+
+import asyncio
+import logging
+import threading
+from concurrent.futures import Future
+from typing import Any, Coroutine, Optional
+
+import aiohttp
+
+logger = logging.getLogger(__name__)
+
+JOIN_TIMEOUT_S: float = 10.0
+
+
+class BackgroundHttpPoster:
+
+    def __init__(self) -> None:
+        self._session: Optional[aiohttp.ClientSession] = None
+        self._loop = asyncio.new_event_loop()
+        self._thread = threading.Thread(
+            target=self._run_loop, name="scripted-runtime-async", daemon=True
+        )
+        self._thread.start()
+
+    def _run_loop(self) -> None:
+        asyncio.set_event_loop(self._loop)
+        self._loop.run_forever()
+
+    def submit_coro(self, coro: Coroutine) -> None:
+        future = asyncio.run_coroutine_threadsafe(coro, self._loop)
+        future.add_done_callback(self._log_coro_exception)
+
+    @staticmethod
+    def _log_coro_exception(future: Future) -> None:
+        try:
+            future.result()
+        except asyncio.CancelledError:
+            pass
+        except Exception:
+            logger.exception("scripted_runtime: background async coroutine failed")
+
+    async def post(self, url: str, json: Any) -> None:
+        session = self._ensure_session()
+        async with session.post(url, json=json) as resp:
+            await resp.read()
+
+    def _ensure_session(self) -> aiohttp.ClientSession:
+        if self._session is None or self._session.closed:
+            self._session = aiohttp.ClientSession(
+                connector=aiohttp.TCPConnector(limit=0)
+            )
+        return self._session
+
+    def close(self) -> None:
+        try:
+            if self._session is not None:
+                future = asyncio.run_coroutine_threadsafe(
+                    self._session.close(), self._loop
+                )
+                future.result(timeout=JOIN_TIMEOUT_S)
+        except Exception:
+            logger.exception("scripted_runtime: failed to close aiohttp session")
+        try:
+            self._loop.call_soon_threadsafe(self._loop.stop)
+            self._thread.join(timeout=JOIN_TIMEOUT_S)
+        except Exception:
+            logger.exception("scripted_runtime: failed to stop background async loop")

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,142 @@
+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.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._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,
+        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,
+            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") -> None:
+        return lifecycle.abort(self, rid=handle.rid)
+
+    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 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,37 @@
+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:
+    server_args = ctx.scheduler.server_args
+    url = f"http://{server_args.host}:{server_args.port}{path}"
+
+    async def _post() -> None:
+        try:
+            await ctx._http_poster.post(url, json)
+        except Exception:  # noqa: BLE001 — fire-and-forget background POST
+            logger.exception("scripted_runtime: POST %s failed", path)
+
+    ctx._http_poster.submit_coro(_post())
+    ctx._tokenizer_recv_proxy.wait_until_arrived(
+        predicate,
+        timeout_s=timeout_s,
+        description=description,
+    )

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

@@ -0,0 +1,75 @@
+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,
+)
+
+if TYPE_CHECKING:
+    from sglang.test.scripted_runtime.context.api import ScriptedContext
+
+
+def _await_control(
+    ctx: "ScriptedContext", *, path: str, json, expect_type: type
+) -> None:
+    _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) -> None:
+    _await_control(
+        ctx,
+        path="/abort_request",
+        json={"rid": rid, "abort_all": False},
+        expect_type=AbortReq,
+    )
+
+
+def flush_cache(ctx: "ScriptedContext") -> None:
+    _await_control(
+        ctx,
+        path="/flush_cache",
+        json=None,
+        expect_type=FlushCacheReqInput,
+    )

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

@@ -0,0 +1,123 @@
+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.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 None:
+        return rid in ctx._seen_rids
+    return req.finished()
+
+
+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,67 @@
+from __future__ import annotations
+
+import uuid
+from typing import TYPE_CHECKING, 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,
+        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}
+        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)

python/sglang/test/scripted_runtime/http_server.py

@@ -0,0 +1,258 @@
+from __future__ import annotations
+
+import multiprocessing as mp
+import os
+import tempfile
+import time
+from pathlib import Path
+from typing import Any, Callable, Dict, Optional, Tuple
+
+import requests
+import zmq
+
+from sglang.srt.entrypoints.http_server import launch_server
+from sglang.srt.environ import envs
+from sglang.srt.server_args import ServerArgs
+from sglang.srt.utils.network import get_free_port, get_zmq_socket_on_host
+from sglang.test.scripted_runtime.io_struct import (
+    HookReady,
+    OutOfBandError,
+    RunScript,
+    ScriptFailed,
+    ScriptSucceeded,
+    Shutdown,
+)
+from sglang.test.scripted_runtime.utils import close_zmq_socket
+
+DEFAULT_RUN_TIMEOUT_S: float = 120.0
+SHUTDOWN_JOIN_TIMEOUT_S: float = 60.0
+LISTENER_ACCEPT_TIMEOUT_S: float = 300.0
+HTTP_READY_TIMEOUT_S: float = 300.0
+HTTP_READY_POLL_INTERVAL_S: float = 0.5
+SERVER_HOST: str = "127.0.0.1"
+
+
+class ScriptedHttpServer:
+
+    def __init__(
+        self,
+        *,
+        ctx: zmq.Context,
+        socket: zmq.Socket,
+        server_process: mp.process.BaseProcess,
+        out_of_band_error_path: Path,
+        http_port: int,
+    ) -> None:
+        self._ctx = ctx
+        self._socket = socket
+        self._server_process = server_process
+        self._out_of_band_error_path = out_of_band_error_path
+        self._base_url = f"http://{SERVER_HOST}:{http_port}"
+        self._shutdown_done = False
+        self._dirty: Optional[str] = None
+
+    @classmethod
+    def start(cls, **engine_kwargs: Any) -> "ScriptedHttpServer":
+        out_of_band_error_path = _create_oob_error_file()
+
+        ctx = zmq.Context()
+        dispatch_port, socket = get_zmq_socket_on_host(ctx, zmq.PAIR, host=SERVER_HOST)
+        server_process, http_port = _spawn_server_process(
+            endpoint=f"tcp://{SERVER_HOST}:{dispatch_port}",
+            out_of_band_error_path=out_of_band_error_path,
+            engine_kwargs=engine_kwargs,
+        )
+
+        self = cls(
+            ctx=ctx,
+            socket=socket,
+            server_process=server_process,
+            out_of_band_error_path=out_of_band_error_path,
+            http_port=http_port,
+        )
+        try:
+            self._await_handshake()
+            self._await_http_ready()
+        except BaseException:
+            self._teardown()
+            raise
+        return self
+
+    def execute_script(
+        self,
+        script_fn: Callable,
+        *,
+        args: Tuple[Any, ...] = (),
+        timeout_s: float = DEFAULT_RUN_TIMEOUT_S,
+    ) -> None:
+        if self._dirty:
+            raise RuntimeError(f"ScriptedHttpServer is dirty: {self._dirty}")
+
+        fn_path = f"{script_fn.__module__}:{script_fn.__qualname__}"
+        self._socket.send_pyobj(RunScript(fn_path=fn_path, args=args))
+
+        if not self._socket.poll(int(timeout_s * 1000)):
+            if not self._server_process.is_alive():
+                self._dirty = f"server process died before responding to {fn_path!r}"
+                raise RuntimeError(self._dirty)
+            self._dirty = f"script {fn_path!r} timed out after {timeout_s}s"
+            raise TimeoutError(self._dirty)
+
+        reply = self._socket.recv_pyobj()
+        match reply:
+            case ScriptFailed(traceback=tb):
+                raise AssertionError(f"scripted-runtime script failed:\n{tb}")
+            case ScriptSucceeded():
+                return
+            case _:
+                raise RuntimeError(
+                    f"scheduler hook replied with unexpected message {reply!r}"
+                )
+
+    def shutdown(self) -> None:
+        if self._shutdown_done:
+            return
+
+        fatal_error: Optional[OutOfBandError] = None
+        try:
+            try:
+                self._socket.send_pyobj(Shutdown())
+            except zmq.ZMQError:
+                pass
+
+            self._server_process.join(timeout=SHUTDOWN_JOIN_TIMEOUT_S)
+            self._terminate_process(self._server_process)
+            fatal_error = self._read_out_of_band_error()
+        finally:
+            self._teardown()
+            self._shutdown_done = True
+
+        if fatal_error:
+            raise AssertionError(
+                f"scripted-runtime server failed:\n{fatal_error.traceback}"
+            )
+
+    def _await_handshake(self) -> None:
+        if not self._socket.poll(int(LISTENER_ACCEPT_TIMEOUT_S * 1000)):
+            raise TimeoutError(
+                f"ScriptedHttpServer: HTTP server did not connect within "
+                f"{LISTENER_ACCEPT_TIMEOUT_S}s"
+            )
+
+        ready = self._socket.recv_pyobj()
+        if not isinstance(ready, HookReady):
+            raise RuntimeError(
+                f"ScriptedHttpServer: expected HookReady handshake, got {ready!r}"
+            )
+
+    def _await_http_ready(self) -> None:
+        # HookReady only means the scheduler dispatch loop started; the uvicorn
+        # entrypoint may not be bound yet. The first script runs
+        # _reset_engine_state, which POSTs to this server's own HTTP port, so
+        # block until the port is bound and routing before any script can run.
+        #
+        # Wait for *any* HTTP response, not status 200: in scripted mode the
+        # scheduler is driven by the script, so normal warmup never completes
+        # and /health stays 503 (server_status == Starting) for the whole run.
+        # A 503 still proves the socket is bound and routes are registered,
+        # which is all the control POSTs need.
+        url = f"{self._base_url}/health"
+        deadline = time.monotonic() + HTTP_READY_TIMEOUT_S
+        while time.monotonic() < deadline:
+            if not self._server_process.is_alive():
+                raise RuntimeError(
+                    "ScriptedHttpServer: server process died during HTTP startup"
+                )
+            try:
+                requests.get(url, timeout=2.0)
+                return
+            except requests.RequestException:
+                pass
+            time.sleep(HTTP_READY_POLL_INTERVAL_S)
+
+        raise TimeoutError(
+            f"ScriptedHttpServer: HTTP endpoint {url} not bound within "
+            f"{HTTP_READY_TIMEOUT_S}s"
+        )
+
+    def _teardown(self) -> None:
+        try:
+            close_zmq_socket(self._socket, self._ctx)
+        except Exception:  # noqa: BLE001 — best-effort cleanup
+            pass
+        self._terminate_process(self._server_process)
+        self._cleanup_files()
+
+    def _read_out_of_band_error(self) -> Optional[OutOfBandError]:
+        try:
+            text = self._out_of_band_error_path.read_text()
+        except OSError:
+            return None
+        text = text.strip()
+        return OutOfBandError.from_json(text) if text else None
+
+    def _cleanup_files(self) -> None:
+        try:
+            self._out_of_band_error_path.unlink()
+        except OSError:
+            pass
+
+    @staticmethod
+    def _terminate_process(process: mp.process.BaseProcess) -> None:
+        if not process.is_alive():
+            return
+        process.terminate()
+        process.join(timeout=10.0)
+        if process.is_alive():
+            process.kill()
+            process.join(timeout=10.0)
+
+
+def _create_oob_error_file() -> Path:
+    err_fd, err_path = tempfile.mkstemp(
+        prefix="scripted_runtime_oob_error_", suffix=".json"
+    )
+    os.close(err_fd)
+    return Path(err_path)
+
+
+def _spawn_server_process(
+    *,
+    endpoint: str,
+    out_of_band_error_path: Path,
+    engine_kwargs: Dict[str, Any],
+) -> Tuple[mp.process.BaseProcess, int]:
+    mp_ctx = mp.get_context("spawn")
+    launch_kwargs: Dict[str, Any] = dict(
+        host=SERVER_HOST,
+        port=get_free_port(),
+        kv_canary="raise",
+        kv_canary_real_data="partial",
+        kv_canary_sweep_interval=100,
+        disable_piecewise_cuda_graph=True,
+    )
+    launch_kwargs.update(engine_kwargs)
+    http_port = launch_kwargs["port"]
+    server_process = mp_ctx.Process(
+        target=_launch_scripted_http_server,
+        kwargs=launch_kwargs,
+        name="scripted-runtime-http-server",
+        daemon=False,
+    )
+
+    sys_path_entry = str(Path(__file__).resolve().parent)
+    with (
+        envs.SGLANG_TEST_SCRIPTED_RUNTIME.override(True),
+        envs.SGLANG_TEST_SCRIPTED_RUNTIME_IPC_ADDR.override(endpoint),
+        envs.SGLANG_TEST_SCRIPTED_RUNTIME_OUT_OF_BAND_ERROR_PATH.override(
+            str(out_of_band_error_path)
+        ),
+        envs.SGLANG_TEST_SCRIPTED_RUNTIME_SYS_PATH_ENTRY.override(sys_path_entry),
+    ):
+        server_process.start()
+
+    return server_process, http_port
+
+
+def _launch_scripted_http_server(**engine_kwargs: Any) -> None:
+    launch_server(ServerArgs(**engine_kwargs))

python/sglang/test/scripted_runtime/io_struct.py

@@ -0,0 +1,51 @@
+from __future__ import annotations
+
+import dataclasses
+import json
+from dataclasses import dataclass
+from typing import Any, Tuple, Union
+
+
+@dataclass(frozen=True, slots=True)
+class RunScript:
+
+    fn_path: str
+    args: Tuple[Any, ...] = ()
+
+
+@dataclass(frozen=True, slots=True)
+class Shutdown:
+    pass
+
+
+@dataclass(frozen=True, slots=True)
+class HookReady:
+    pass
+
+
+@dataclass(frozen=True, slots=True)
+class ScriptSucceeded:
+    pass
+
+
+@dataclass(frozen=True, slots=True)
+class ScriptFailed:
+
+    traceback: str
+
+
+@dataclass(frozen=True, slots=True)
+class OutOfBandError:
+
+    traceback: str
+
+    def to_json(self) -> str:
+        return json.dumps(dataclasses.asdict(self))
+
+    @classmethod
+    def from_json(cls, text: str) -> "OutOfBandError":
+        return cls(**json.loads(text))
+
+
+ScriptedCommand = Union[RunScript, Shutdown]
+ScriptedReply = Union[HookReady, ScriptSucceeded, ScriptFailed]

python/sglang/test/scripted_runtime/req_handle.py

@@ -0,0 +1,51 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Optional
+
+if TYPE_CHECKING:
+    from sglang.srt.managers.schedule_batch import Req
+    from sglang.test.scripted_runtime.context.api import ScriptedContext
+
+
+@dataclass(frozen=True, slots=True)
+class ScriptedReqHandle:
+    rid: str
+    context: "ScriptedContext"
+
+    @property
+    def req(self) -> Optional["Req"]:
+        return self.context.find_req_by_rid(self.rid)
+
+    @property
+    def finished(self) -> bool:
+        return self.context.is_finished(self.rid)
+
+    @property
+    def is_chunking(self) -> bool:
+        return self.context.is_chunking(self.rid)
+
+    @property
+    def chunks_done(self) -> int:
+        return self.context.chunks_done(self.rid)
+
+    @property
+    def status(self) -> str:
+        return self.context.status(self.rid)
+
+    @property
+    def remaining_prompt_tokens(self) -> int:
+        return self.context.remaining_prompt_tokens(self.rid)
+
+    @property
+    def kv_pages(self) -> int:
+        req = self.req
+        if req is None or req.req_pool_idx is None:
+            return 0
+        page_size = self.context.scheduler.page_size
+        return (req.kv_allocated_len + page_size - 1) // page_size
+
+    @property
+    def lock_refs(self) -> int:
+        node = self.req.last_node
+        return node.lock_ref if node is not None else 0

python/sglang/test/scripted_runtime/scheduler_hook.py

@@ -0,0 +1,190 @@
+from __future__ import annotations
+
+import logging
+import sys
+import traceback
+from dataclasses import dataclass
+from pathlib import Path
+from typing import TYPE_CHECKING, Generator, List, Optional, Tuple
+
+import zmq
+
+from sglang.srt.environ import envs
+from sglang.srt.utils.network import get_zmq_socket
+from sglang.test.scripted_runtime.background_http_poster import BackgroundHttpPoster
+from sglang.test.scripted_runtime.context import ScriptedContext
+from sglang.test.scripted_runtime.io_struct import (
+    HookReady,
+    OutOfBandError,
+    RunScript,
+    ScriptFailed,
+    ScriptSucceeded,
+    Shutdown,
+)
+from sglang.test.scripted_runtime.utils import (
+    close_zmq_socket,
+    ensure_script_importable,
+    resolve_fn,
+)
+
+if TYPE_CHECKING:
+    from sglang.srt.managers.scheduler import Scheduler
+    from sglang.test.scripted_runtime.tokenizer_recv_proxy import (
+        ScriptedTokenizerRecvProxy,
+    )
+
+logger = logging.getLogger(__name__)
+
+RESET_DRAIN_MAX_STEPS: int = 200
+
+
+@dataclass(frozen=True, slots=True)
+class ScriptedBatchRecord:
+    forward_iter: int
+    mode: Optional[str]
+    rids: Tuple[str, ...]
+    extend_rids: Tuple[str, ...]
+    chunked_rid: Optional[str]
+
+
+def _reset_engine_state(ctx: ScriptedContext) -> Generator:
+    scheduler = ctx.scheduler
+
+    ctx.abort_all()
+    for _ in range(RESET_DRAIN_MAX_STEPS):
+        yield
+        if (
+            scheduler.chunked_req is None
+            and len(scheduler.waiting_queue) == 0
+            and scheduler.running_batch.is_empty()
+        ):
+            break
+
+    server_args = scheduler.server_args
+    for _ in range(2 * (server_args.pp_size + server_args.pp_async_batch_depth)):
+        yield
+
+    ctx.flush_cache()
+    yield
+
+
+class ScriptedSchedulerHook:
+
+    def __init__(
+        self,
+        *,
+        scheduler: "Scheduler",
+        tokenizer_recv_proxy: Optional["ScriptedTokenizerRecvProxy"],
+    ) -> None:
+        self.scheduler = scheduler
+        self._is_driver = (
+            scheduler.ps.pp_rank == 0
+            and scheduler.ps.tp_rank == 0
+            and scheduler.ps.attn_cp_rank == 0
+        )
+        self._batch_log: List["ScriptedBatchRecord"] = []
+
+        if self._is_driver:
+            ensure_script_importable(
+                envs.SGLANG_TEST_SCRIPTED_RUNTIME_SYS_PATH_ENTRY.get()
+            )
+            self._http_poster: Optional[BackgroundHttpPoster] = BackgroundHttpPoster()
+            self._context: Optional[ScriptedContext] = ScriptedContext(
+                scheduler_hook=self,
+                tokenizer_recv_proxy=tokenizer_recv_proxy,
+                http_poster=self._http_poster,
+            )
+            self._script_fn_generator: Optional[Generator] = self._run_dispatch_loop()
+        else:
+            self._http_poster = None
+            self._context = None
+            self._script_fn_generator = None
+
+    def _run_dispatch_loop(self) -> Generator:
+        endpoint = envs.SGLANG_TEST_SCRIPTED_RUNTIME_IPC_ADDR.get()
+        ctx_zmq = zmq.Context()
+        socket = get_zmq_socket(ctx_zmq, zmq.PAIR, endpoint, bind=False)
+        try:
+            socket.send_pyobj(HookReady())
+            while True:
+                msg = socket.recv_pyobj()
+                match msg:
+                    case Shutdown():
+                        return
+                    case RunScript(fn_path=fn_path, args=args):
+                        fn = resolve_fn(fn_path)
+                        ctx = self._context
+                        yield from _reset_engine_state(ctx)
+                        self._batch_log.clear()
+                        sub_gen = fn(ctx, *args)
+                        try:
+                            yield from sub_gen
+                        except Exception:
+                            socket.send_pyobj(
+                                ScriptFailed(traceback=traceback.format_exc())
+                            )
+                        else:
+                            socket.send_pyobj(ScriptSucceeded())
+                    case _:
+                        raise ValueError(f"dispatch loop: unknown command {msg!r}")
+        finally:
+            close_zmq_socket(socket, ctx_zmq)
+            self._http_poster.close()
+
+    def on_run_batch(self, batch) -> None:
+        if not self._is_driver:
+            return
+        chunked = self.scheduler.chunked_req
+        self._batch_log.append(
+            ScriptedBatchRecord(
+                forward_iter=batch.forward_iter,
+                mode=(
+                    batch.forward_mode.name.lower()
+                    if batch.forward_mode is not None
+                    else None
+                ),
+                rids=tuple(r.rid for r in batch.reqs),
+                extend_rids=(
+                    tuple(r.rid for r in batch.reqs)
+                    if batch.forward_mode is not None and batch.forward_mode.is_extend()
+                    else ()
+                ),
+                chunked_rid=chunked.rid if chunked is not None else None,
+            )
+        )
+
+    def step(self) -> None:
+        if not self._is_driver:
+            return
+
+        done, exc_tb = _advance_generator(self._script_fn_generator)
+        if not done:
+            return
+
+        if exc_tb is not None:
+            _write_out_of_band_error(exc_tb)
+        sys.exit(0 if exc_tb is None else 1)
+
+
+def _write_out_of_band_error(exc_tb: str) -> None:
+    path = envs.SGLANG_TEST_SCRIPTED_RUNTIME_OUT_OF_BAND_ERROR_PATH.get()
+    if not path:
+        return
+    error = OutOfBandError(traceback=exc_tb or "<no traceback>")
+    try:
+        Path(path).write_text(error.to_json())
+    except OSError:
+        logger.exception(
+            "Failed to write scripted_runtime out-of-band error to %s", path
+        )
+
+
+def _advance_generator(generator: Generator) -> Tuple[bool, Optional[str]]:
+    try:
+        next(generator)
+        return False, None
+    except StopIteration:
+        return True, None
+    except Exception:
+        logger.exception("Failed to advance generator")
+        return True, traceback.format_exc()

python/sglang/test/scripted_runtime/test_case.py

@@ -0,0 +1,29 @@
+from __future__ import annotations
+
+from typing import Any, ClassVar, Dict
+
+from sglang.test.scripted_runtime.http_server import ScriptedHttpServer
+from sglang.test.test_utils import CustomTestCase
+
+
+class ScriptedTestCase(CustomTestCase):
+
+    ENGINE_KWARGS: ClassVar[Dict[str, Any]] = {}
+
+    server: ClassVar[ScriptedHttpServer]
+
+    @classmethod
+    def setUpClass(cls) -> None:
+        super().setUpClass()
+        if not cls.ENGINE_KWARGS:
+            raise AssertionError(
+                f"{cls.__name__} must set ENGINE_KWARGS to a non-empty dict"
+            )
+        cls.server = ScriptedHttpServer.start(**cls.ENGINE_KWARGS)
+
+    @classmethod
+    def tearDownClass(cls) -> None:
+        try:
+            cls.server.shutdown()
+        finally:
+            super().tearDownClass()

python/sglang/test/scripted_runtime/tokenizer_recv_proxy.py

@@ -0,0 +1,55 @@
+from __future__ import annotations
+
+import time
+from collections import deque
+from typing import Any, Callable
+
+import zmq
+
+
+class ScriptedTokenizerRecvProxy:
+
+    def __init__(self, *, underlying: zmq.Socket) -> None:
+        self._underlying = underlying
+        self._buffer: deque = deque()
+
+    def recv_pyobj(self, flags: int = 0) -> Any:
+        self._drain_underlying()
+
+        if self._buffer:
+            return self._buffer.popleft()
+
+        if flags & zmq.NOBLOCK:
+            raise zmq.ZMQError(zmq.EAGAIN, "Resource temporarily unavailable")
+        raise RuntimeError(
+            "ScriptedTokenizerRecvProxy.recv_pyobj: blocking recv is not supported"
+        )
+
+    def wait_until_arrived(
+        self,
+        predicate: Callable[[Any], bool],
+        *,
+        timeout_s: float,
+        description: str = "matching object",
+    ) -> None:
+        start_len = len(self._buffer)
+        deadline = time.monotonic() + timeout_s
+        while True:
+            self._drain_underlying()
+            for i, obj in enumerate(self._buffer):
+                if i >= start_len and predicate(obj):
+                    return
+            if time.monotonic() >= deadline:
+                raise TimeoutError(
+                    f"ScriptedTokenizerRecvProxy: no {description} arrived on the "
+                    f"recv_from_tokenizer socket within {timeout_s}s"
+                )
+            time.sleep(0.005)
+
+    def _drain_underlying(self) -> None:
+        while True:
+            try:
+                req = self._underlying.recv_pyobj(zmq.NOBLOCK)
+            except zmq.ZMQError:
+                break
+            self._buffer.append(req)

python/sglang/test/scripted_runtime/utils.py

@@ -0,0 +1,33 @@
+from __future__ import annotations
+
+import importlib
+import sys
+from typing import Callable, Optional
+
+import zmq
+
+
+def close_zmq_socket(socket: zmq.Socket, ctx: zmq.Context) -> None:
+    socket.setsockopt(zmq.LINGER, 0)
+    socket.close()
+    ctx.term()
+
+
+def ensure_script_importable(sys_path_entry: Optional[str]) -> None:
+    if sys_path_entry and sys_path_entry not in sys.path:
+        sys.path.insert(0, sys_path_entry)
+
+
+def resolve_fn(qualified: str) -> Callable:
+    module_name, sep, fn_name = qualified.partition(":")
+    if not sep or not module_name or not fn_name:
+        raise ValueError(
+            f"scripted-runtime fn path must be 'module.path:function_name', "
+            f"got {qualified!r}"
+        )
+    obj = importlib.import_module(module_name)
+    for part in fn_name.split("."):
+        obj = getattr(obj, part)
+    if not callable(obj):
+        raise TypeError(f"resolved object is not callable: {qualified!r} -> {obj!r}")
+    return obj