1use 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
27fn 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
37fn 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
119pub 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 .nest("/api", api_router())
131 .fallback(frontend::serve_frontend)
133 .with_state(state)
134 .layer(cors)
135}
136
137pub 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 info!("Cleaning up interactive notebooks...");
161 state_shutdown_all.jupyter.shutdown_all().await;
162
163 Ok(())
164}