From dd176387c651d026b0090587a52936c73709dd2c Mon Sep 17 00:00:00 2001 From: fzyzcjy <5236035+fzyzcjy@users.noreply.github.com> Date: Sat, 6 Jun 2026 09:07:01 +0800 Subject: [PATCH] Add scripted-runtime harness core and wire scheduler/IPC hooks (#27411) --- .claude/rules/modify-component-must-read.md | 1 + .../skills/scripted-runtime-notes/SKILL.md | 29 ++ python/sglang/srt/environ.py | 5 + python/sglang/srt/managers/scheduler.py | 20 ++ .../scheduler_components/ipc_channels.py | 18 +- .../scheduler_components/request_receiver.py | 10 +- python/sglang/srt/utils/network.py | 2 +- .../sglang/test/scripted_runtime/__init__.py | 0 .../background_http_poster.py | 68 +++++ .../test/scripted_runtime/context/__init__.py | 3 + .../test/scripted_runtime/context/api.py | 142 ++++++++++ .../test/scripted_runtime/context/engine.py | 16 ++ .../scripted_runtime/context/http_post.py | 37 +++ .../scripted_runtime/context/lifecycle.py | 75 +++++ .../test/scripted_runtime/context/queries.py | 123 +++++++++ .../test/scripted_runtime/context/radix.py | 34 +++ .../scripted_runtime/context/req_starter.py | 67 +++++ .../test/scripted_runtime/http_server.py | 258 ++++++++++++++++++ .../sglang/test/scripted_runtime/io_struct.py | 51 ++++ .../test/scripted_runtime/req_handle.py | 51 ++++ .../test/scripted_runtime/scheduler_hook.py | 190 +++++++++++++ .../sglang/test/scripted_runtime/test_case.py | 29 ++ .../scripted_runtime/tokenizer_recv_proxy.py | 55 ++++ python/sglang/test/scripted_runtime/utils.py | 33 +++ 24 files changed, 1313 insertions(+), 4 deletions(-) create mode 100644 .claude/skills/scripted-runtime-notes/SKILL.md create mode 100644 python/sglang/test/scripted_runtime/__init__.py create mode 100644 python/sglang/test/scripted_runtime/background_http_poster.py create mode 100644 python/sglang/test/scripted_runtime/context/__init__.py create mode 100644 python/sglang/test/scripted_runtime/context/api.py create mode 100644 python/sglang/test/scripted_runtime/context/engine.py create mode 100644 python/sglang/test/scripted_runtime/context/http_post.py create mode 100644 python/sglang/test/scripted_runtime/context/lifecycle.py create mode 100644 python/sglang/test/scripted_runtime/context/queries.py create mode 100644 python/sglang/test/scripted_runtime/context/radix.py create mode 100644 python/sglang/test/scripted_runtime/context/req_starter.py create mode 100644 python/sglang/test/scripted_runtime/http_server.py create mode 100644 python/sglang/test/scripted_runtime/io_struct.py create mode 100644 python/sglang/test/scripted_runtime/req_handle.py create mode 100644 python/sglang/test/scripted_runtime/scheduler_hook.py create mode 100644 python/sglang/test/scripted_runtime/test_case.py create mode 100644 python/sglang/test/scripted_runtime/tokenizer_recv_proxy.py create mode 100644 python/sglang/test/scripted_runtime/utils.py diff --git a/.claude/rules/modify-component-must-read.md b/.claude/rules/modify-component-must-read.md index abacb07bc..64003e7c3 100644 --- a/.claude/rules/modify-component-must-read.md +++ b/.claude/rules/modify-component-must-read.md @@ -5,3 +5,4 @@ Before modifying the following components, read the listed skill first. - **Speculative decoding code** (anything under `python/sglang/srt/speculative/`, related attention backends, scheduler accumulators, IPC fields, observability metrics, or CLI flags) → [`speculative-naming`](../skills/speculative-naming/SKILL.md) - **`Scheduler` / `TokenizerManager` / `ModelRunner` `__init__`** (`python/sglang/srt/managers/scheduler.py`, `python/sglang/srt/managers/tokenizer_manager.py`, `python/sglang/srt/model_executor/model_runner.py`) → [`large-class-init-style`](../skills/large-class-init-style/SKILL.md) - **Environment variables** (adding, renaming, or reviewing any `SGLANG_*` env var, migrating a legacy `SGL_*` alias, or touching `python/sglang/srt/environ.py`) → [`env-var-conventions`](../skills/env-var-conventions/SKILL.md) +- **Scripted runtime** (anything related to the scripted runtime) → [`scripted-runtime-notes`](../skills/scripted-runtime-notes/SKILL.md) diff --git a/.claude/skills/scripted-runtime-notes/SKILL.md b/.claude/skills/scripted-runtime-notes/SKILL.md new file mode 100644 index 000000000..8ad89d951 --- /dev/null +++ b/.claude/skills/scripted-runtime-notes/SKILL.md @@ -0,0 +1,29 @@ +--- +name: scripted-runtime-notes +description: Requirements for the SGLang scripted runtime, chiefly when to add (vs not add) a harness API. Use for anything related to the scripted runtime. +--- + +# Scripted Runtime — Notes + +Notes for anything related to the SGLang scripted runtime. + +## When to Add an API + +Tests read `r.req.*` and `t._scheduler.*` directly — there is no encapsulation boundary. A thin wrapper buys zero isolation; it only grows the surface. + +Add an API only if it does real work: + +1. **Control primitive** — drives the engine through a real path (`start_req`, `pause_generation`, `abort`, `evict_radix`, `exhaust_kv`). Reuse the real path; never hand-mutate state. +2. **Hook-backed** — value cannot be read from a snapshot; accumulate via `scheduler_hook.on_run_batch` or the recv proxy (`chunks_done`). Read-only; never monkey-patch; never add `*_count` to `srt/`. +3. **Multi-structure derivation, widely reused** — scans `chunked_req` + `waiting_queue` + `running_batch` + `last_batch` (`is_idle`, `status`, `batch_composition`). + +Else: don't. Read `r.req.X` / `t._scheduler.X` in the test; inline single-use accessors. + +Never: + +- Weaken an assertion to fit a missing probe. +- Probe implementation details ("field non-None", "which branch ran") — assert the consequence. + +## Other Tips + +- **Engine-self-driven behavior: drive the real loop, don't call the private.** Never synchronously call a scheduler private (e.g. `scheduler._abort_on_waiting_timeout()`) from the harness/test — it runs at the wrong loop phase, bypasses the ordered `recv_requests` → `process_input_requests` injection, and can fire in states the real loop never reaches (e.g. while paused). For sweeps the engine runs itself (timeout/idle), enable the config/env and advance the loop with `yield`. diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 1fc5b4b23..dc1e063aa 100644 --- a/python/sglang/srt/environ.py +++ b/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) diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 7c1e4ec32..1ea37269f 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/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: diff --git a/python/sglang/srt/managers/scheduler_components/ipc_channels.py b/python/sglang/srt/managers/scheduler_components/ipc_channels.py index 246655bf0..43da8ccf6 100644 --- a/python/sglang/srt/managers/scheduler_components/ipc_channels.py +++ b/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 ) diff --git a/python/sglang/srt/managers/scheduler_components/request_receiver.py b/python/sglang/srt/managers/scheduler_components/request_receiver.py index 0c5b83b67..608aece5f 100644 --- a/python/sglang/srt/managers/scheduler_components/request_receiver.py +++ b/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 [] diff --git a/python/sglang/srt/utils/network.py b/python/sglang/srt/utils/network.py index 500cb486e..a78034ed5 100644 --- a/python/sglang/srt/utils/network.py +++ b/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: diff --git a/python/sglang/test/scripted_runtime/__init__.py b/python/sglang/test/scripted_runtime/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/python/sglang/test/scripted_runtime/background_http_poster.py b/python/sglang/test/scripted_runtime/background_http_poster.py new file mode 100644 index 000000000..671d22e1a --- /dev/null +++ b/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") diff --git a/python/sglang/test/scripted_runtime/context/__init__.py b/python/sglang/test/scripted_runtime/context/__init__.py new file mode 100644 index 000000000..f8e5e4efc --- /dev/null +++ b/python/sglang/test/scripted_runtime/context/__init__.py @@ -0,0 +1,3 @@ +from sglang.test.scripted_runtime.context.api import ScriptedContext + +__all__ = ["ScriptedContext"] diff --git a/python/sglang/test/scripted_runtime/context/api.py b/python/sglang/test/scripted_runtime/context/api.py new file mode 100644 index 000000000..180db6a92 --- /dev/null +++ b/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) diff --git a/python/sglang/test/scripted_runtime/context/engine.py b/python/sglang/test/scripted_runtime/context/engine.py new file mode 100644 index 000000000..88a37ea53 --- /dev/null +++ b/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, + } diff --git a/python/sglang/test/scripted_runtime/context/http_post.py b/python/sglang/test/scripted_runtime/context/http_post.py new file mode 100644 index 000000000..447759589 --- /dev/null +++ b/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, + ) diff --git a/python/sglang/test/scripted_runtime/context/lifecycle.py b/python/sglang/test/scripted_runtime/context/lifecycle.py new file mode 100644 index 000000000..d66c838a9 --- /dev/null +++ b/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, + ) diff --git a/python/sglang/test/scripted_runtime/context/queries.py b/python/sglang/test/scripted_runtime/context/queries.py new file mode 100644 index 000000000..3b114f1cd --- /dev/null +++ b/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 + ) diff --git a/python/sglang/test/scripted_runtime/context/radix.py b/python/sglang/test/scripted_runtime/context/radix.py new file mode 100644 index 000000000..6e3294981 --- /dev/null +++ b/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 diff --git a/python/sglang/test/scripted_runtime/context/req_starter.py b/python/sglang/test/scripted_runtime/context/req_starter.py new file mode 100644 index 000000000..ff9ee8cab --- /dev/null +++ b/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) diff --git a/python/sglang/test/scripted_runtime/http_server.py b/python/sglang/test/scripted_runtime/http_server.py new file mode 100644 index 000000000..19360e5c7 --- /dev/null +++ b/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)) diff --git a/python/sglang/test/scripted_runtime/io_struct.py b/python/sglang/test/scripted_runtime/io_struct.py new file mode 100644 index 000000000..d5c6caded --- /dev/null +++ b/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] diff --git a/python/sglang/test/scripted_runtime/req_handle.py b/python/sglang/test/scripted_runtime/req_handle.py new file mode 100644 index 000000000..945d763cf --- /dev/null +++ b/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 diff --git a/python/sglang/test/scripted_runtime/scheduler_hook.py b/python/sglang/test/scripted_runtime/scheduler_hook.py new file mode 100644 index 000000000..ef76322bc --- /dev/null +++ b/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 "") + 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() diff --git a/python/sglang/test/scripted_runtime/test_case.py b/python/sglang/test/scripted_runtime/test_case.py new file mode 100644 index 000000000..491ab5e09 --- /dev/null +++ b/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() diff --git a/python/sglang/test/scripted_runtime/tokenizer_recv_proxy.py b/python/sglang/test/scripted_runtime/tokenizer_recv_proxy.py new file mode 100644 index 000000000..e3fe19c20 --- /dev/null +++ b/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) diff --git a/python/sglang/test/scripted_runtime/utils.py b/python/sglang/test/scripted_runtime/utils.py new file mode 100644 index 000000000..38de51c95 --- /dev/null +++ b/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