101 lines
3.6 KiB
Rust
101 lines
3.6 KiB
Rust
//! 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 pd_bootstrap;
|
|
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: submission handles, immutable server configuration,
|
|
/// and the API-owned chat formatter.
|
|
#[derive(Clone)]
|
|
struct AppState {
|
|
senders: Senders,
|
|
egress_buf: usize,
|
|
server_args: Arc<ServerArgs>,
|
|
chat_formatter: Option<openai::ChatFormatter>,
|
|
/// 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 chat_formatter = openai::load_chat_support(&server_args);
|
|
let state = AppState {
|
|
senders,
|
|
egress_buf,
|
|
server_args: server_args.clone(),
|
|
chat_formatter,
|
|
egress_activity,
|
|
};
|
|
// Each endpoint module registers its own routes and merges here.
|
|
let mut 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);
|
|
if server_args.enable_pd_bootstrap() {
|
|
// Merged after `with_state` (the registry carries its own state) and
|
|
// before `log::apply`, so bootstrap traffic shows in the access log.
|
|
let (bootstrap_routes, sweeper) = pd_bootstrap::router_and_sweeper();
|
|
tokio::spawn(sweeper); // cancelled with the runtime on shutdown
|
|
app = app.merge(bootstrap_routes);
|
|
tracing::info!("PD KV bootstrap registry mounted on the api listener");
|
|
}
|
|
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;
|
|
}
|
|
};
|
|
// `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");
|
|
}
|
|
}
|
|
}
|