sgl-project--sglang/python/sglang/test/scripted_runtime/http_server.py

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.managers.io_struct import sock_recv, sock_send, wrap_as_pickle
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__}"
        sock_send(self._socket, wrap_as_pickle(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 = sock_recv(self._socket)
        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:
                sock_send(self._socket, wrap_as_pickle(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 = sock_recv(self._socket)
        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.
        #
        # Poll /model_info rather than /health: once the server reports Up,
        # /health runs the generation-based health check (a real probe request
        # through the scheduler), which cannot complete while the scheduler
        # waits for scripts between RunScript commands. Any /model_info
        # response proves the socket is bound and routes are registered, which
        # is all the control POSTs need.
        url = f"{self._base_url}/model_info"
        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_prefill_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))