add the sglang-server rust implementation (api server, runtime, tokenizer manager, detokenizer)

rust/sglang-server/src/api_server.rs

@@ -0,0 +1,97 @@
+//! API server (axum / tokio). I/O-bound; own pinned multi-thread runtime. Only
+//! this module knows HTTP, so other protocols can mount the same `AppState`.
+//! `/generate` submits a `Request` then awaits one `Done` (unary) or relays SSE
+//! frames (`data: {json}` … `[DONE]`), byte-compatible with Python
+//! `http_server.generate_request`; `/server_info` reuses it for one control result.
+mod common;
+mod frame;
+mod guard;
+mod log;
+mod native_api;
+mod openai;
+mod submit;
+
+use std::sync::Arc;
+
+use axum::Router;
+
+use crate::runtime::ServerArgs;
+use crate::tokenizer_manager::ActivityCounter;
+use crate::tokenizer_manager::Senders;
+
+/// Shared handler state: the submit machinery (`senders`, `egress_buf`)
+/// + shared tokenizer.
+#[derive(Clone)]
+struct AppState {
+    senders: Senders,
+    egress_buf: usize,
+    server_args: Arc<ServerArgs>,
+    /// Egress heartbeat (bumped per drained ring frame).
+    egress_activity: ActivityCounter,
+}
+
+pub async fn serve(
+    listener: std::net::TcpListener,
+    senders: Senders,
+    egress_buf: usize,
+    server_args: Arc<ServerArgs>,
+    egress_activity: ActivityCounter,
+    // The SAME set ingress releases from — see `Ingress::on_abort`. Constructing a
+    // local one here would leave the api server admitting rids that nothing ever
+    // releases.
+    shutdown: flume::Receiver<()>,
+) {
+    let state = AppState {
+        senders,
+        egress_buf,
+        server_args: server_args.clone(),
+        egress_activity,
+    };
+    // Each endpoint module registers its own routes and merges here.
+    let app = Router::new()
+        .merge(common::routes())
+        .merge(native_api::routes())
+        .merge(openai::routes())
+        // TODO(auth): no API-key boundary yet. Python gates every route (except
+        // /health*, /metrics*, OPTIONS) via `add_api_key_middleware`; until ported,
+        // a configured `api_key` does NOT protect these routes.
+        //
+        // No body limit, matching the Python server.
+        .layer(axum::extract::DefaultBodyLimit::disable())
+        .with_state(state);
+    let app = log::apply(app, &server_args);
+
+    // The listener was already bound synchronously in `runtime::start` (so a port
+    // conflict fails startup); adopt it into the tokio reactor here.
+    let listener = match tokio::net::TcpListener::from_std(listener) {
+        Ok(l) => l,
+        Err(e) => {
+            tracing::error!(error = %e, "failed to adopt pre-bound listener");
+            return;
+        }
+    };
+    if let Ok(addr) = listener.local_addr() {
+        tracing::info!(%addr, "sglang-server api listening");
+    }
+    // Non-graceful shutdown: on the signal, stop accepting and RETURN without
+    // waiting for in-flight handlers (a `/generate` blocked on egress would wedge
+    // the join). Returning unwinds `block_on` in `runtime::start` → the api tokio
+    // runtime drops → detached handlers cancel → their `AbortGuard`s fire, release
+    // `Senders` clones → tok/detok channels close → workers exit. Full drain is
+    // deferred (see `request_shutdown`).
+    // `with_connect_info` exposes the peer address to the access-log middleware.
+    let serve = axum::serve(
+        listener,
+        app.into_make_service_with_connect_info::<std::net::SocketAddr>(),
+    );
+    tokio::select! {
+        r = serve => {
+            if let Err(e) = r {
+                tracing::error!(error = %e, "axum serve exited");
+            }
+        }
+        _ = shutdown.recv_async() => {
+            tracing::info!("shutdown: stopping accepts, aborting in-flight handlers");
+        }
+    }
+}

rust/sglang-server/src/api_server/common.rs

@@ -0,0 +1,226 @@
+//! Common control-plane endpoints — `/server_info`, `/get_model_info`
+//! (+ `/model_info` alias), plus the control-request submission path
+//! (`await_control_result`, on the shared `submit`). Data-plane endpoints (incl. `/health*`,
+//! which round-trips a generate probe) live in the sibling `native_api` and
+//! `openai` modules; the shared `AppState` lives in the parent
+//! `api_server` module.
+
+use axum::{
+    Router,
+    extract::State,
+    http::StatusCode,
+    response::{IntoResponse, Response},
+    routing::get,
+};
+
+use super::AppState;
+use super::guard::AbortGuard;
+use super::submit::submit;
+use crate::message::{ControlRequest, EgressItem, GetInternalStateReq, RequestKind};
+use crate::runtime::ServerArgs;
+
+/// The routes this module owns, mounted by `api_server::serve`.
+pub(super) fn routes() -> Router<AppState> {
+    Router::new()
+        // Control-plane: reuses the ingress FSM (no tokenization), returns one
+        // non-streamed JSON result. Adding one = a route line + its struct tag.
+        .route("/server_info", get(server_info))
+        // Static config, no scheduler round-trip. `/get_model_info` (+ `/model_info`
+        // alias).
+        .route("/get_model_info", get(model_info))
+        .route("/model_info", get(model_info))
+}
+
+/// Submit a control request through the ingress FSM (no tokenization) and await the
+/// scheduler's single msgpack result (a `structs.asdict` named map). Returns the
+/// raw bytes, or an error `Response` to return as-is.
+async fn await_control_result(
+    state: &AppState,
+    control: ControlRequest,
+) -> Result<bytes::Bytes, Response> {
+    let (rid, mut rx) = submit(state, RequestKind::Control(Box::new(control)), false).await?;
+    // Control requests register a detok entry like any other, and only
+    // `handle_result` removes it — so a request that never produces one (a stalled
+    // scheduler, a client that hangs up mid-await) leaves the entry behind. A
+    // monitor polling `/server_info` then leaks one `DetokState` per poll, forever.
+    // The guard deregisters on drop; it is disarmed below when the result lands.
+    let mut guard = AbortGuard::new(state.senders.clone(), rid.clone());
+    let received = rx.recv().await;
+    if received.is_some() {
+        guard.disarm(&rid); // completed normally — nothing to abort
+    }
+    match received {
+        Some(EgressItem::Control(bytes)) => Ok(bytes),
+        Some(EgressItem::Error(e)) => {
+            let code =
+                StatusCode::from_u16(e.http_status()).unwrap_or(StatusCode::INTERNAL_SERVER_ERROR);
+            Err((code, e.to_string()).into_response())
+        }
+        // A control request never receives generation frames.
+        Some(EgressItem::Frame(_)) | Some(EgressItem::Done(_)) => Err((
+            StatusCode::INTERNAL_SERVER_ERROR,
+            "unexpected generation output for control request",
+        )
+            .into_response()),
+        None => Err((StatusCode::from_u16(499).unwrap(), "request aborted").into_response()),
+    }
+}
+
+/// `GET /get_model_info` (+ `/model_info` alias) — static model metadata from
+/// `server_args` (no scheduler round-trip); `is_generation` always true.
+async fn model_info(State(state): State<AppState>) -> Response {
+    let sa = &state.server_args;
+    let body = serde_json::json!({
+        "model_path": sa.model_path,
+        "tokenizer_path": sa.tokenizer_path,
+        "is_generation": true,
+        // Python's `TokenizerManager` merges this into every request
+        // (`{**preferred, **client}`); this server has no equivalent yet, so
+        // `RustServer.launch` REFUSES to start when it is set. It can therefore
+        // only be null here — echoing it keeps the field's shape.
+        "preferred_sampling_params": sa.preferred_sampling_params,
+        "weight_version": serde_json::Value::Null,
+    });
+    (
+        StatusCode::OK,
+        [("content-type", "application/json")],
+        serde_json::to_vec(&body).unwrap_or_default(),
+    )
+        .into_response()
+}
+
+/// `GET /server_info` — surface only an allowlist ([`INTERNAL_STATE_ALLOWLIST`] +
+/// curated [`ServerArgs`] accessors), never the raw server-args dump (embeds
+/// `api_key`/`admin_api_key`; see [`shape_server_info`]).
+///
+/// TODO(server_info): Python also includes `kv_events`; add once plumbed.
+async fn server_info(State(state): State<AppState>) -> Response {
+    let bytes = match await_control_result(
+        &state,
+        ControlRequest::GetInternalStateReq(GetInternalStateReq::new(
+            crate::ids::Rid::new().to_string(),
+        )),
+    )
+    .await
+    {
+        Ok(b) => b,
+        Err(resp) => return resp,
+    };
+    match shape_server_info(&bytes, &state.server_args) {
+        Ok(json) => (StatusCode::OK, [("content-type", "application/json")], json).into_response(),
+        Err(e) => {
+            tracing::error!(error = %e, "server_info: shaping failed");
+            (
+                StatusCode::INTERNAL_SERVER_ERROR,
+                "bad server_info response",
+            )
+                .into_response()
+        }
+    }
+}
+
+/// Runtime-metric keys `get_internal_state` adds atop the server-args dump. We copy
+/// ONLY these out of `internal_state` (an allowlist), so the co-mingled
+/// `api_key`/`admin_api_key` can never reach the response.
+const INTERNAL_STATE_ALLOWLIST: &[&str] = &[
+    "last_gen_throughput",
+    "memory_usage",
+    "effective_max_running_requests_per_dp",
+    "avg_spec_accept_length",
+    "step_time_dict",
+];
+
+fn shape_server_info(msgpack: &[u8], server_args: &ServerArgs) -> Result<Vec<u8>, String> {
+    // GetInternalStateReqOutput asdict → `{ "internal_state": { server-args dump +
+    // metrics }, ... }`. Pull that inner map out (it is NOT safe to expose whole).
+    let mut obj: serde_json::Map<String, serde_json::Value> =
+        rmp_serde::from_slice(msgpack).map_err(|e| e.to_string())?;
+    let internal = match obj.remove("internal_state") {
+        Some(serde_json::Value::Object(m)) => m,
+        _ => serde_json::Map::new(),
+    };
+
+    // Copy only the allowlisted runtime metrics — never the raw server-args dump.
+    let mut state_out = serde_json::Map::new();
+    for &k in INTERNAL_STATE_ALLOWLIST {
+        match internal.get(k) {
+            Some(v) if !v.is_null() => {
+                state_out.insert(k.to_string(), v.clone());
+            }
+            _ => {}
+        }
+    }
+
+    // Top-level non-secret config from typed accessors (structurally can't surface
+    // a key field, unlike the raw dump).
+    let response = serde_json::json!({
+        "model_path": server_args.model_path,
+        "served_model_name": server_args.served_model_name,
+        "tokenizer_path": server_args.tokenizer_path,
+        "max_context_length": server_args.model_config.context_len,
+        "max_total_num_tokens": server_args.max_total_num_tokens,
+        "version": server_args.version,
+        "internal_states": [serde_json::Value::Object(state_out)],
+    });
+    serde_json::to_vec(&response).map_err(|e| e.to_string())
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    /// The scheduler's `internal_state` embeds the full server-args dump (incl.
+    /// `api_key`/`admin_api_key`). `/server_info` must surface only the allowlisted
+    /// runtime metrics + curated config — never the secrets — and must not re-nest
+    /// the dump under `internal_states[].internal_state`.
+    #[test]
+    fn shape_server_info_excludes_secrets_and_dump() {
+        // GetInternalStateReqOutput.asdict → { "internal_state": { …dump+metrics… } }.
+        let internal = rmpv::Value::Map(vec![
+            (
+                rmpv::Value::from("api_key"),
+                rmpv::Value::from("secret-token"),
+            ),
+            (
+                rmpv::Value::from("admin_api_key"),
+                rmpv::Value::from("admin-token"),
+            ),
+            (rmpv::Value::from("model_path"), rmpv::Value::from("/m")),
+            (
+                rmpv::Value::from("last_gen_throughput"),
+                rmpv::Value::from(1.5),
+            ),
+            (
+                rmpv::Value::from("effective_max_running_requests_per_dp"),
+                rmpv::Value::from(32),
+            ),
+        ]);
+        let outer = rmpv::Value::Map(vec![(rmpv::Value::from("internal_state"), internal)]);
+        let mut msgpack = Vec::new();
+        rmpv::encode::write_value(&mut msgpack, &outer).unwrap();
+
+        let sa =
+            ServerArgs::from_json(r#"{"model_path": "/m", "api_key": "secret-token"}"#).unwrap();
+        let out = shape_server_info(&msgpack, &sa).unwrap();
+        let text = String::from_utf8(out.clone()).unwrap();
+        // No secret leaks anywhere in the serialized response.
+        assert!(!text.contains("secret-token"), "api_key leaked: {text}");
+        assert!(
+            !text.contains("admin-token"),
+            "admin_api_key leaked: {text}"
+        );
+
+        let v: serde_json::Value = serde_json::from_slice(&out).unwrap();
+        // Allowlisted metric surfaced; the whole dump did not.
+        let state0 = &v["internal_states"][0];
+        assert_eq!(state0["last_gen_throughput"], 1.5);
+        assert_eq!(state0["effective_max_running_requests_per_dp"], 32);
+        assert!(
+            state0.get("internal_state").is_none(),
+            "must not re-nest the dump under internal_state"
+        );
+        assert!(state0.get("api_key").is_none());
+        // Curated top-level config comes from typed accessors, not the dump.
+        assert_eq!(v["model_path"], "/m");
+    }
+}