Skip to main content

expman_server/api/
mod.rs

1//! API module: Axum router, helpers, and handler submodules.
2
3use std::net::SocketAddr;
4use std::path::PathBuf;
5
6use axum::{
7    routing::{get, post},
8    Router,
9};
10use tower_http::cors::{Any, CorsLayer};
11use tracing::info;
12
13use self::state::AppState;
14
15pub use self::state::ServerConfig;
16
17mod artifacts;
18mod experiments;
19mod frontend;
20mod jupyter_handlers;
21pub(crate) mod jupyter_service;
22mod metrics;
23mod runs;
24pub mod state;
25mod stats;
26
27// ─── Helpers ─────────────────────────────────────────────────────────────
28
29fn run_dir(base: &std::path::Path, exp: &str, run: &str) -> PathBuf {
30    base.join(exp).join(run)
31}
32
33fn exp_dir(base: &std::path::Path, exp: &str) -> PathBuf {
34    base.join(exp)
35}
36
37// ─── Router ──────────────────────────────────────────────────────────────
38
39fn api_router() -> Router<AppState> {
40    Router::new()
41        .route("/experiments", get(experiments::list_experiments))
42        .route("/experiments/{exp}/runs", get(runs::list_runs))
43        .route(
44            "/experiments/{exp}/metadata",
45            get(experiments::get_experiment_metadata)
46                .patch(experiments::update_experiment_metadata),
47        )
48        .route(
49            "/experiments/{exp}/runs/{run}/metrics",
50            get(metrics::get_metrics),
51        )
52        .route(
53            "/run/{exp}/{run}/stream/vectors",
54            get(metrics::stream_vectors),
55        )
56        .route(
57            "/experiments/{exp}/runs/{run}/log/stream",
58            get(metrics::stream_log),
59        )
60        .route(
61            "/experiments/{exp}/runs/{run}/config",
62            get(metrics::get_config),
63        )
64        .route(
65            "/experiments/{exp}/runs/{run}/metadata",
66            get(runs::get_run_metadata).patch(runs::update_run_metadata),
67        )
68        .route(
69            "/experiments/{exp}/runs/{run}/artifacts",
70            get(artifacts::list_artifacts),
71        )
72        .route(
73            "/experiments/{exp}/runs/{run}/artifacts/content",
74            get(artifacts::get_artifact_content),
75        )
76        .route("/experiments/{exp}/stats", get(stats::get_experiment_stats))
77        .route("/config", get(stats::get_server_config))
78        .route("/stats", get(stats::get_global_stats))
79        .route(
80            "/jupyter/available",
81            get(jupyter_handlers::available_jupyter),
82        )
83        .route(
84            "/experiments/{exp}/runs/{run}/jupyter/start",
85            post(jupyter_handlers::start_jupyter),
86        )
87        .route(
88            "/experiments/{exp}/runs/{run}/jupyter/stop",
89            post(jupyter_handlers::stop_jupyter),
90        )
91        .route(
92            "/experiments/{exp}/runs/{run}/jupyter/status",
93            get(jupyter_handlers::status_jupyter),
94        )
95        .route(
96            "/experiments/{exp}/runs/{run}/jupyter/notebook",
97            get(jupyter_handlers::get_jupyter_notebook)
98                .post(jupyter_handlers::create_jupyter_notebook),
99        )
100        .route(
101            "/experiments/{exp}/jupyter/start",
102            post(jupyter_handlers::start_multi_jupyter),
103        )
104        .route(
105            "/experiments/{exp}/jupyter/stop",
106            post(jupyter_handlers::stop_multi_jupyter),
107        )
108        .route(
109            "/experiments/{exp}/jupyter/status",
110            get(jupyter_handlers::status_multi_jupyter),
111        )
112        .route(
113            "/experiments/{exp}/jupyter/notebook",
114            get(jupyter_handlers::get_multi_jupyter_notebook)
115                .post(jupyter_handlers::create_multi_jupyter_notebook),
116        )
117}
118
119// ─── Public API ──────────────────────────────────────────────────────────
120
121/// Build the Axum router with all routes.
122pub fn build_router(state: AppState) -> Router {
123    let cors = CorsLayer::new()
124        .allow_origin(Any)
125        .allow_methods(Any)
126        .allow_headers(Any);
127
128    Router::new()
129        // API routes
130        .nest("/api", api_router())
131        // Frontend: serve embedded static files
132        .fallback(frontend::serve_frontend)
133        .with_state(state)
134        .layer(cors)
135}
136
137/// Start the server on the given address.
138pub async fn serve(config: ServerConfig) -> anyhow::Result<()> {
139    let state = AppState::new(config.base_dir.clone());
140    let state_shutdown_all = state.clone();
141    let state_shutdown_token = state.clone();
142    let app = build_router(state);
143
144    let addr: SocketAddr = format!("{}:{}", config.host, config.port).parse()?;
145    info!("ExpMan dashboard at http://{}", addr);
146
147    let listener = tokio::net::TcpListener::bind(addr).await?;
148
149    axum::serve(listener, app)
150        .with_graceful_shutdown(async move {
151            tokio::signal::ctrl_c()
152                .await
153                .expect("failed to install CTRL+C handler");
154            info!("Shutting down ExpMan server...");
155            state_shutdown_token.shutdown_token.cancel();
156        })
157        .await?;
158
159    // Cleanup all Jupyter instances
160    info!("Cleaning up interactive notebooks...");
161    state_shutdown_all.jupyter.shutdown_all().await;
162
163    Ok(())
164}