feat(all): 重炼 crates/common 核心组件、上线 Web 运维看板与 Docker 容器化部署
This commit is contained in:
@@ -0,0 +1,29 @@
|
||||
[package]
|
||||
name = "server"
|
||||
version = "0.1.0"
|
||||
edition = "2021"
|
||||
|
||||
[dependencies]
|
||||
common = { path = "../common", default-features = false }
|
||||
mq = { path = "../mq" }
|
||||
axum.workspace = true
|
||||
tokio.workspace = true
|
||||
tokio-util.workspace = true
|
||||
tower-http.workspace = true
|
||||
tower.workspace = true
|
||||
serde.workspace = true
|
||||
serde_json.workspace = true
|
||||
serde_yaml.workspace = true
|
||||
rusqlite.workspace = true
|
||||
r2d2.workspace = true
|
||||
r2d2_sqlite.workspace = true
|
||||
tracing.workspace = true
|
||||
tracing-subscriber.workspace = true
|
||||
anyhow.workspace = true
|
||||
clap.workspace = true
|
||||
chrono.workspace = true
|
||||
uuid.workspace = true
|
||||
tempfile.workspace = true
|
||||
dotenvy.workspace = true
|
||||
|
||||
|
||||
@@ -0,0 +1,37 @@
|
||||
# server
|
||||
|
||||
> DCTS 中央调度与 REST API 服务端程序。
|
||||
|
||||
---
|
||||
|
||||
## 📦 模块概览
|
||||
|
||||
`server` 是基于 Axum 框架构建的高性能中央 Master 控制服务。
|
||||
|
||||
- **`main.rs`**:HTTP Router 初始化、命令行参数解析 (`clap`) 与后台掉线检测 / 任务超时重入循环。
|
||||
- **`db.rs`**:基于 SQLite + `r2d2` 的持久化数据层,管理网格点、节点及任务历史。
|
||||
- **`scheduler.rs`**:网格点自动展开生成器与状态更新管道。
|
||||
- **`api/`**:
|
||||
- `node.rs`:节点注册与心跳接口。
|
||||
- `task.rs`:任务 Claim 抢占与 Report 结果汇报。
|
||||
- `seed.rs`:`.7` 大气结构二进制种子下载。
|
||||
- `data.rs`:二进制运行依赖与谱线库分发。
|
||||
- `workflow.rs`:网格工作流 CRUD 与启动/停止控制。
|
||||
- `status.rs`:系统运行全貌状态上报。
|
||||
|
||||
---
|
||||
|
||||
## 🚀 Setup & Testing
|
||||
|
||||
### 编译与启动
|
||||
```bash
|
||||
cargo build -p server --release
|
||||
./target/release/server --workflow config.yaml --port 8080
|
||||
```
|
||||
|
||||
### 自动化测试
|
||||
```bash
|
||||
cargo test -p server
|
||||
```
|
||||
|
||||
详细 API 规范请参阅 [API Reference](../../docs/api_reference.md)。
|
||||
@@ -0,0 +1,92 @@
|
||||
use std::path::Path;
|
||||
use std::process::Command;
|
||||
|
||||
fn main() {
|
||||
// Re-run build script if dashboard files change
|
||||
println!("cargo:rerun-if-changed=../../dashboard/src");
|
||||
println!("cargo:rerun-if-changed=../../dashboard/package.json");
|
||||
println!("cargo:rerun-if-changed=../../dashboard/index.html");
|
||||
|
||||
// Skip building dashboard in Docker / CI environments if requested
|
||||
let skip_build = std::env::var("SKIP_DASHBOARD_BUILD")
|
||||
.map(|v| v == "1" || v == "true")
|
||||
.unwrap_or(false);
|
||||
|
||||
if skip_build {
|
||||
println!("cargo:warning=SKIP_DASHBOARD_BUILD 已启用,跳过前端 Dashboard 编译。");
|
||||
return;
|
||||
}
|
||||
|
||||
let dashboard_dir = Path::new("../../dashboard");
|
||||
if !dashboard_dir.exists() {
|
||||
println!("cargo:warning=未检测到 dashboard 目录,跳过前端构建。");
|
||||
return;
|
||||
}
|
||||
|
||||
let node_modules_exist = dashboard_dir.join("node_modules").exists();
|
||||
|
||||
if !node_modules_exist {
|
||||
println!("cargo:warning=未检测到 dashboard/node_modules,正在执行 npm install...");
|
||||
let status = Command::new("npm")
|
||||
.arg("install")
|
||||
.current_dir(dashboard_dir)
|
||||
.status();
|
||||
|
||||
match status {
|
||||
Ok(s) if s.success() => {
|
||||
println!("cargo:warning=前端依赖 (npm install) 执行成功。");
|
||||
}
|
||||
_ => {
|
||||
panic!("前端部署阶段:执行 'npm install' 失败!请确认环境中已安装且具备可用的 Node/NPM 工具。若意在进行服务端原生单独打包装订而完全不需要绑定前端静态页面资源,可以设置环境变量 SKIP_DASHBOARD_BUILD=1");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let dist_index = dashboard_dir.join("dist/index.html");
|
||||
let needs_build = if !dist_index.exists() {
|
||||
true
|
||||
} else {
|
||||
fn latest_mtime(dir: &Path) -> std::time::SystemTime {
|
||||
let mut max = std::fs::metadata(dir)
|
||||
.and_then(|m| m.modified())
|
||||
.unwrap_or(std::time::SystemTime::UNIX_EPOCH);
|
||||
if dir.is_dir() {
|
||||
if let Ok(entries) = std::fs::read_dir(dir) {
|
||||
for entry in entries.flatten() {
|
||||
let t = latest_mtime(&entry.path());
|
||||
if t > max {
|
||||
max = t;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
max
|
||||
}
|
||||
let dist_time = std::fs::metadata(&dist_index)
|
||||
.and_then(|m| m.modified())
|
||||
.unwrap_or(std::time::SystemTime::UNIX_EPOCH);
|
||||
latest_mtime(&dashboard_dir.join("src")) > dist_time
|
||||
|| latest_mtime(&dashboard_dir.join("package.json")) > dist_time
|
||||
|| latest_mtime(&dashboard_dir.join("index.html")) > dist_time
|
||||
};
|
||||
|
||||
if needs_build {
|
||||
println!("cargo:warning=正在打包前端 Dashboard 静态资源 (npm run build)...");
|
||||
let status = Command::new("npm")
|
||||
.arg("run")
|
||||
.arg("build")
|
||||
.current_dir(dashboard_dir)
|
||||
.status();
|
||||
|
||||
match status {
|
||||
Ok(s) if s.success() => {
|
||||
println!("cargo:warning=前端 Dashboard 静态资源打包成功。");
|
||||
}
|
||||
_ => {
|
||||
panic!("前端部署阶段:执行 'npm run build' 失败,无法编译生成面板资源包!未免造生功能存在组件缺憾和页面载入留白的严重程序构建散落碎片,本项预编译保护链早已锁严封口,并当即制停本次流程。");
|
||||
}
|
||||
}
|
||||
} else {
|
||||
println!("cargo:warning=检测到前端 Dashboard 静态资源已是最新,跳过构建。");
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,97 @@
|
||||
use axum::{
|
||||
body::Body,
|
||||
extract::Path as AxumPath,
|
||||
http::{header, StatusCode},
|
||||
response::IntoResponse,
|
||||
};
|
||||
use std::path::{Path, PathBuf};
|
||||
use tokio::fs::File;
|
||||
use tokio_util::io::ReaderStream;
|
||||
|
||||
pub async fn download_single_data_file(AxumPath(filename): AxumPath<String>) -> axum::response::Response {
|
||||
let safe_name = Path::new(&filename)
|
||||
.file_name()
|
||||
.map(|s| s.to_string_lossy().to_string())
|
||||
.unwrap_or_default();
|
||||
|
||||
if safe_name.is_empty() || safe_name.starts_with('.') {
|
||||
return (StatusCode::BAD_REQUEST, "无效的数据文件名").into_response();
|
||||
}
|
||||
|
||||
// 严苛白名单过滤:严防 `..`、特殊符号注入及路径穿透攻击
|
||||
if !safe_name.chars().all(|c| c.is_ascii_alphanumeric() || c == '.' || c == '_' || c == '-' || c == '+' || c == '@') {
|
||||
tracing::warn!("拦截到疑似非法字符构造的敏感及越界资源抓取行为: {}", safe_name);
|
||||
return (StatusCode::BAD_REQUEST, "参数非法,请求的文件包含系统不许可的危险专属占位或路径重定向字符").into_response();
|
||||
}
|
||||
|
||||
let rel_path = format!("assets/data/{}", safe_name);
|
||||
tracing::debug!("服务端处理数据文件下载请求: {}", safe_name);
|
||||
stream_asset_file(&rel_path, "application/octet-stream").await.into_response()
|
||||
}
|
||||
|
||||
pub async fn download_linelist() -> axum::response::Response {
|
||||
let linelist_path = std::env::var("DCTS_LINELIST_PATH").unwrap_or_else(|_| "assets/gfVIS99.dat".to_string());
|
||||
stream_asset_file(&linelist_path, "application/octet-stream").await.into_response()
|
||||
}
|
||||
|
||||
fn resolve_asset(rel_path: &str) -> Option<PathBuf> {
|
||||
if let Ok(base) = std::env::var("DCTS_ASSETS_DIR") {
|
||||
let p = Path::new(&base).join(rel_path);
|
||||
if p.exists() {
|
||||
return Some(p);
|
||||
}
|
||||
}
|
||||
|
||||
let p = Path::new(rel_path);
|
||||
if p.exists() {
|
||||
return Some(p.to_path_buf());
|
||||
}
|
||||
|
||||
if let Ok(exe_path) = std::env::current_exe() {
|
||||
if let Some(parent) = exe_path.parent() {
|
||||
let candidate1 = parent.join(rel_path);
|
||||
if candidate1.exists() {
|
||||
return Some(candidate1);
|
||||
}
|
||||
if let Some(grandparent) = parent.parent() {
|
||||
let candidate2 = grandparent.join(rel_path);
|
||||
if candidate2.exists() {
|
||||
return Some(candidate2);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
None
|
||||
}
|
||||
|
||||
async fn stream_asset_file(rel_path: &str, content_type: &'static str) -> impl IntoResponse {
|
||||
let resolved_path = match resolve_asset(rel_path) {
|
||||
Some(p) => p,
|
||||
None => return (StatusCode::NOT_FOUND, "资源数据文件不存在").into_response(),
|
||||
};
|
||||
|
||||
match File::open(&resolved_path).await {
|
||||
Ok(file) => {
|
||||
let stream = ReaderStream::new(file);
|
||||
let body = Body::from_stream(stream);
|
||||
|
||||
let filename = resolved_path
|
||||
.file_name()
|
||||
.unwrap_or_default()
|
||||
.to_string_lossy()
|
||||
.to_string();
|
||||
|
||||
let disposition = format!("attachment; filename=\"{}\"", filename);
|
||||
|
||||
let headers = [
|
||||
(header::CONTENT_TYPE, content_type.to_string()),
|
||||
(header::CONTENT_DISPOSITION, disposition),
|
||||
];
|
||||
|
||||
(headers, body).into_response()
|
||||
}
|
||||
Err(_) => (StatusCode::INTERNAL_SERVER_ERROR, "无法读取资源数据文件").into_response(),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,74 @@
|
||||
pub mod data;
|
||||
pub mod node;
|
||||
pub mod seed;
|
||||
pub mod status;
|
||||
pub mod task;
|
||||
pub mod workflow;
|
||||
|
||||
use axum::{
|
||||
extract::State,
|
||||
http::{header, Request, StatusCode},
|
||||
middleware::Next,
|
||||
response::IntoResponse,
|
||||
};
|
||||
use crate::db::Database;
|
||||
use crate::scheduler::GridScheduler;
|
||||
use mq::sqlite_queue::SqliteTaskQueue;
|
||||
use std::sync::Arc;
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct AppState {
|
||||
pub db: Database,
|
||||
pub queue: Arc<SqliteTaskQueue>,
|
||||
pub scheduler: Arc<GridScheduler>,
|
||||
pub results_dir: String,
|
||||
pub auth_token: Option<String>,
|
||||
}
|
||||
|
||||
/// 固定时间敏感字符串一致性核验函数,彻底消解时序测信道猜测危险
|
||||
fn constant_time_eq(a: &str, b: &str) -> bool {
|
||||
let a_bytes = a.as_bytes();
|
||||
let b_bytes = b.as_bytes();
|
||||
let mut diff = (a_bytes.len() ^ b_bytes.len()) as u64;
|
||||
// 遍历目标 secret (b_bytes) 的完整长度,使耗时仅受 server 预期 token 长度决定
|
||||
for (i, &y) in b_bytes.iter().enumerate() {
|
||||
let x = if i < a_bytes.len() { a_bytes[i] } else { 0 };
|
||||
diff |= (x ^ y) as u64;
|
||||
}
|
||||
diff == 0
|
||||
}
|
||||
|
||||
/// Axum 鉴权中间件:若 AppState 中配置了 auth_token 则强制校验 Bearer Token 或 X-API-Key
|
||||
pub async fn auth_middleware(
|
||||
State(state): State<AppState>,
|
||||
req: Request<axum::body::Body>,
|
||||
next: Next,
|
||||
) -> impl IntoResponse {
|
||||
if let Some(ref expected_token) = state.auth_token {
|
||||
let auth_header = req
|
||||
.headers()
|
||||
.get(header::AUTHORIZATION)
|
||||
.and_then(|v| v.to_str().ok());
|
||||
let api_key_header = req
|
||||
.headers()
|
||||
.get("x-api-key")
|
||||
.and_then(|v| v.to_str().ok());
|
||||
|
||||
let token_valid = match (auth_header, api_key_header) {
|
||||
(Some(auth), _) if auth.starts_with("Bearer ") => constant_time_eq(&auth[7..], expected_token),
|
||||
(Some(auth), _) => constant_time_eq(auth, expected_token),
|
||||
(_, Some(key)) => constant_time_eq(key, expected_token),
|
||||
_ => false,
|
||||
};
|
||||
|
||||
if !token_valid {
|
||||
return (
|
||||
StatusCode::UNAUTHORIZED,
|
||||
"Unauthorized: Invalid or missing authentication token",
|
||||
)
|
||||
.into_response();
|
||||
}
|
||||
}
|
||||
|
||||
next.run(req).await.into_response()
|
||||
}
|
||||
@@ -0,0 +1,25 @@
|
||||
use super::AppState;
|
||||
use axum::{extract::State, response::IntoResponse, Json};
|
||||
use common::models::{NodeHeartbeatRequest, NodeRegisterRequest};
|
||||
use serde_json::json;
|
||||
|
||||
pub async fn register_node(
|
||||
State(state): State<AppState>,
|
||||
Json(req): Json<NodeRegisterRequest>,
|
||||
) -> impl IntoResponse {
|
||||
match state.db.register_node(&req).await {
|
||||
Ok(_) => Json(json!({"status": "ok", "message": "节点注册成功"})),
|
||||
Err(e) => Json(json!({"status": "error", "message": e.to_string()})),
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
pub async fn heartbeat_node(
|
||||
State(state): State<AppState>,
|
||||
Json(req): Json<NodeHeartbeatRequest>,
|
||||
) -> impl IntoResponse {
|
||||
match state.db.heartbeat_node(&req).await {
|
||||
Ok(_) => Json(json!({"status": "ok"})),
|
||||
Err(e) => Json(json!({"status": "error", "message": e.to_string()})),
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,58 @@
|
||||
use super::AppState;
|
||||
use axum::{
|
||||
body::Body,
|
||||
extract::{Path as AxumPath, State},
|
||||
http::{header, StatusCode},
|
||||
response::Response,
|
||||
};
|
||||
use tokio::fs::File;
|
||||
use tokio_util::io::ReaderStream;
|
||||
use tracing::warn;
|
||||
|
||||
pub async fn download_seed(
|
||||
State(state): State<AppState>,
|
||||
AxumPath(name): AxumPath<String>,
|
||||
) -> Response {
|
||||
if name.is_empty() || name.starts_with('.') || !name.chars().all(|c| c.is_ascii_alphanumeric() || c == '.' || c == '_' || c == '-' || c == '+' || c == '@') {
|
||||
warn!("拒绝可能包含路径穿越或特别注入序列的非法种子下载请求: {}", name);
|
||||
return Response::builder()
|
||||
.status(StatusCode::BAD_REQUEST)
|
||||
.body(Body::from("非法的种子名称参数"))
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
let seed_file_path = std::path::Path::new(&state.results_dir)
|
||||
.join(&name)
|
||||
.join(format!("{}.7", name));
|
||||
|
||||
if !seed_file_path.is_file() {
|
||||
warn!("客户端请求的种子文件不存在: {}", seed_file_path.display());
|
||||
return Response::builder()
|
||||
.status(StatusCode::NOT_FOUND)
|
||||
.body(Body::from("请求的种子文件不存在"))
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
let file = match File::open(&seed_file_path).await {
|
||||
Ok(file) => file,
|
||||
Err(_) => {
|
||||
return Response::builder()
|
||||
.status(StatusCode::INTERNAL_SERVER_ERROR)
|
||||
.body(Body::from("无法打开种子文件"))
|
||||
.unwrap();
|
||||
}
|
||||
};
|
||||
|
||||
|
||||
let stream = ReaderStream::new(file);
|
||||
let body = Body::from_stream(stream);
|
||||
|
||||
Response::builder()
|
||||
.header(header::CONTENT_TYPE, "application/octet-stream")
|
||||
.header(
|
||||
header::CONTENT_DISPOSITION,
|
||||
format!("attachment; filename=\"{}.7\"", name),
|
||||
)
|
||||
.body(body)
|
||||
.unwrap()
|
||||
}
|
||||
@@ -0,0 +1,22 @@
|
||||
use super::AppState;
|
||||
use axum::{extract::State, response::IntoResponse, Json};
|
||||
use serde_json::json;
|
||||
|
||||
pub async fn get_status(State(state): State<AppState>) -> impl IntoResponse {
|
||||
let nodes = state.db.get_active_nodes().await.unwrap_or_default();
|
||||
let total_active_slots: i32 = nodes.iter().map(|n| n.active_slots).sum();
|
||||
let total_max_slots: i32 = nodes.iter().map(|n| n.max_slots).sum();
|
||||
|
||||
let grid_stats = state.db.get_grid_summary_stats().await.unwrap_or(serde_json::json!({
|
||||
"total": 0, "pending": 0, "running": 0, "converged": 0, "failed": 0
|
||||
}));
|
||||
|
||||
Json(json!({
|
||||
"status": "online",
|
||||
"nodes_online": nodes.len(),
|
||||
"total_active_slots": total_active_slots,
|
||||
"total_max_slots": total_max_slots,
|
||||
"nodes": nodes,
|
||||
"grid_stats": grid_stats,
|
||||
}))
|
||||
}
|
||||
@@ -0,0 +1,149 @@
|
||||
use super::AppState;
|
||||
use axum::{
|
||||
extract::{Multipart, State},
|
||||
response::IntoResponse,
|
||||
Json,
|
||||
};
|
||||
use common::models::{GridPointParams, ModelSummary, TaskReport, TaskStatus};
|
||||
use serde_json::json;
|
||||
use std::path::Path;
|
||||
use tokio::fs;
|
||||
use tracing::{info, warn};
|
||||
|
||||
use axum::http::StatusCode;
|
||||
|
||||
pub async fn claim_task(State(state): State<AppState>) -> impl IntoResponse {
|
||||
match state.queue.pop_task().await {
|
||||
Ok(Some(task)) => {
|
||||
if let Err(e) = state.db.mark_grid_point_running(&task.point_name).await {
|
||||
warn!("领用任务 {} 后同步变更为 running 状态遇到异常: {}. 后置 stale 定时自取检索引索将介入修复维护", task.task_id, e);
|
||||
}
|
||||
(StatusCode::OK, Json(json!({"status": "ok", "task": task}))).into_response()
|
||||
}
|
||||
Ok(None) => (StatusCode::OK, Json(json!({"status": "empty", "task": null}))).into_response(),
|
||||
Err(e) => (
|
||||
StatusCode::INTERNAL_SERVER_ERROR,
|
||||
Json(json!({"status": "error", "message": format!("领用任务失败: {}", e)})),
|
||||
)
|
||||
.into_response(),
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn report_task(
|
||||
State(state): State<AppState>,
|
||||
mut multipart: Multipart,
|
||||
) -> impl IntoResponse {
|
||||
let mut report_json: Option<TaskReport> = None;
|
||||
let mut seed_file_data: Option<Vec<u8>> = None;
|
||||
|
||||
while let Ok(Some(field)) = multipart.next_field().await {
|
||||
let field_name = field.name().unwrap_or("").to_string();
|
||||
if field_name == "report" {
|
||||
if let Ok(bytes) = field.bytes().await {
|
||||
if let Ok(report) = serde_json::from_slice::<TaskReport>(&bytes) {
|
||||
report_json = Some(report);
|
||||
}
|
||||
}
|
||||
} else if field_name == "seed_file" {
|
||||
if let Ok(bytes) = field.bytes().await {
|
||||
seed_file_data = Some(bytes.to_vec());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let report = match report_json {
|
||||
Some(r) => r,
|
||||
None => {
|
||||
return (
|
||||
StatusCode::BAD_REQUEST,
|
||||
Json(json!({"status": "error", "message": "请求中缺少 report 字段"})),
|
||||
)
|
||||
.into_response();
|
||||
}
|
||||
};
|
||||
|
||||
let name = report.point_name.clone();
|
||||
|
||||
if name.is_empty() || name.starts_with('.') || !name.chars().all(|c| c.is_ascii_alphanumeric() || c == '.' || c == '_' || c == '-' || c == '+' || c == '@') {
|
||||
warn!("拒绝可能包含路径穿越或特殊非常规编码号攻击的网格点名称请求: {}", name);
|
||||
return (
|
||||
StatusCode::BAD_REQUEST,
|
||||
Json(json!({"status": "error", "message": "非法的网格点名称参数"})),
|
||||
)
|
||||
.into_response();
|
||||
}
|
||||
|
||||
let params = match extract_params(&report) {
|
||||
Some(p) => p,
|
||||
None => {
|
||||
warn!("网格点 {} 汇报数据解析失败: 无法解析 params 或 summary_json", name);
|
||||
return (
|
||||
StatusCode::BAD_REQUEST,
|
||||
Json(json!({"status": "error", "message": "无法解析 params 或 summary_json"})),
|
||||
)
|
||||
.into_response();
|
||||
}
|
||||
};
|
||||
|
||||
// Record in DB
|
||||
if let Err(e) = state.db.record_task_report(&report).await {
|
||||
warn!("记录网格点 {} 任务结果到数据库失败: {}", name, e);
|
||||
return (
|
||||
StatusCode::INTERNAL_SERVER_ERROR,
|
||||
Json(json!({"status": "error", "message": format!("记录数据库失败: {}", e)})),
|
||||
)
|
||||
.into_response();
|
||||
}
|
||||
|
||||
// Clean up task from task_queue table to prevent queue DB bloat
|
||||
if let Err(e) = state.queue.remove_task(&report.task_id.to_string()).await {
|
||||
tracing::warn!("从任务队列中清理已上报任务记录 {} 失败: {}", report.task_id, e);
|
||||
}
|
||||
|
||||
// 采用原子写入模式保持 conv.json 与核心二进制数据完整落地后才揭晓真实文件名
|
||||
let model_dir = Path::new(&state.results_dir).join(&name);
|
||||
if fs::create_dir_all(&model_dir).await.is_ok() {
|
||||
let conv_tmp = model_dir.join(format!("conv.json.{}.tmp", uuid::Uuid::new_v4().simple()));
|
||||
let conv_path = model_dir.join("conv.json");
|
||||
if fs::write(&conv_tmp, &report.summary_json).await.is_ok() {
|
||||
let _ = fs::rename(&conv_tmp, &conv_path).await;
|
||||
}
|
||||
|
||||
// Save seed file .7 using atomic temporary writing strategy
|
||||
if report.converged && !report.atmosphere_has_nan {
|
||||
if let Some(bytes) = seed_file_data {
|
||||
let seed_tmp = model_dir.join(format!("{}.7.{}.tmp", name, uuid::Uuid::new_v4().simple()));
|
||||
let seed_path = model_dir.join(format!("{}.7", name));
|
||||
if fs::write(&seed_tmp, bytes).await.is_ok() {
|
||||
if fs::rename(&seed_tmp, &seed_path).await.is_ok() {
|
||||
info!("成功保持原子写入落地并保存网格点 {} 的收敛种子文件: {}", name, seed_path.display());
|
||||
let _ = state
|
||||
.db
|
||||
.insert_seed(¶ms, &seed_path.to_string_lossy())
|
||||
.await;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if report.status == TaskStatus::Failed || report.status == TaskStatus::Timeout || report.atmosphere_has_nan {
|
||||
// Task did not succeed -> check if seed_step fallback should be triggered
|
||||
info!("网格点 {} 计算未成功完成,检查种子回退机制...", name);
|
||||
if let Err(e) = state.scheduler.trigger_seed_step_fallback(¶ms).await {
|
||||
warn!("网格点 {} 触发种子回退机制失败: {}", name, e);
|
||||
}
|
||||
}
|
||||
|
||||
(StatusCode::OK, Json(json!({"status": "ok", "message": "上报成功"}))).into_response()
|
||||
}
|
||||
|
||||
fn extract_params(report: &TaskReport) -> Option<GridPointParams> {
|
||||
if let Some(ref p) = report.params {
|
||||
return Some(p.clone());
|
||||
}
|
||||
serde_json::from_str::<ModelSummary>(&report.summary_json)
|
||||
.ok()
|
||||
.map(|summary| summary.params)
|
||||
}
|
||||
|
||||
@@ -0,0 +1,255 @@
|
||||
use super::AppState;
|
||||
use axum::{
|
||||
extract::{Path as AxumPath, State},
|
||||
http::StatusCode,
|
||||
response::IntoResponse,
|
||||
Json,
|
||||
};
|
||||
use common::config::GridConfig;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use tracing::info;
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
pub struct CreateWorkflowRequest {
|
||||
pub name: String,
|
||||
pub description: Option<String>,
|
||||
pub config_yaml: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
pub struct ApiResponse<T> {
|
||||
pub success: bool,
|
||||
pub message: String,
|
||||
pub data: Option<T>,
|
||||
}
|
||||
|
||||
pub async fn list_workflows(State(state): State<AppState>) -> impl IntoResponse {
|
||||
match state.db.list_workflows().await {
|
||||
Ok(list) => (StatusCode::OK, Json(ApiResponse { success: true, message: "成功获取工作流列表".to_string(), data: Some(list) })),
|
||||
Err(e) => (StatusCode::INTERNAL_SERVER_ERROR, Json(ApiResponse { success: false, message: format!("获取工作流列表失败: {}", e), data: None })),
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn get_workflow(
|
||||
State(state): State<AppState>,
|
||||
AxumPath(name): AxumPath<String>,
|
||||
) -> impl IntoResponse {
|
||||
match state.db.get_workflow(&name).await {
|
||||
Ok(Some(item)) => (StatusCode::OK, Json(ApiResponse { success: true, message: "成功获取工作流详情".to_string(), data: Some(item) })),
|
||||
Ok(None) => (StatusCode::NOT_FOUND, Json(ApiResponse { success: false, message: format!("工作流 '{}' 未找到", name), data: None })),
|
||||
Err(e) => (StatusCode::INTERNAL_SERVER_ERROR, Json(ApiResponse { success: false, message: format!("获取工作流详情失败: {}", e), data: None })),
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn save_workflow(
|
||||
State(state): State<AppState>,
|
||||
Json(req): Json<CreateWorkflowRequest>,
|
||||
) -> impl IntoResponse {
|
||||
// Validate YAML config string
|
||||
if let Err(e) = serde_yaml::from_str::<GridConfig>(&req.config_yaml) {
|
||||
return (
|
||||
StatusCode::BAD_REQUEST,
|
||||
Json(ApiResponse::<()> {
|
||||
success: false,
|
||||
message: format!("无效的 YAML 配置: {}", e),
|
||||
data: None,
|
||||
}),
|
||||
);
|
||||
}
|
||||
|
||||
// 检查被编辑的工作流是否正处于激活运行中
|
||||
if let Ok(Some(existing)) = state.db.get_workflow(&req.name).await {
|
||||
if existing.status == "running" || existing.status == "initializing" {
|
||||
return (
|
||||
StatusCode::BAD_REQUEST,
|
||||
Json(ApiResponse::<()> {
|
||||
success: false,
|
||||
message: format!("工作流 '{}' 正处在运行或初始加载流程中,严禁原地覆写参数重设至 IDLE;如待变更参数请先调 API 显式触发停止后再保存", req.name),
|
||||
data: None,
|
||||
}),
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
match state.db.upsert_workflow(&req.name, req.description.as_deref(), &req.config_yaml, "idle").await {
|
||||
Ok(_) => {
|
||||
info!("成功注册/更新工作流配置: {}", req.name);
|
||||
(
|
||||
StatusCode::OK,
|
||||
Json(ApiResponse::<()> {
|
||||
success: true,
|
||||
message: format!("工作流 '{}' 保存成功", req.name),
|
||||
data: None,
|
||||
}),
|
||||
)
|
||||
}
|
||||
Err(e) => (
|
||||
StatusCode::INTERNAL_SERVER_ERROR,
|
||||
Json(ApiResponse::<()> {
|
||||
success: false,
|
||||
message: format!("保存工作流失败: {}", e),
|
||||
data: None,
|
||||
}),
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn delete_workflow(
|
||||
State(state): State<AppState>,
|
||||
AxumPath(name): AxumPath<String>,
|
||||
) -> impl IntoResponse {
|
||||
match state.db.delete_workflow(&name).await {
|
||||
Ok(_) => (
|
||||
StatusCode::OK,
|
||||
Json(ApiResponse::<()> {
|
||||
success: true,
|
||||
message: format!("工作流 '{}' 已删除", name),
|
||||
data: None,
|
||||
}),
|
||||
),
|
||||
Err(e) => (
|
||||
StatusCode::INTERNAL_SERVER_ERROR,
|
||||
Json(ApiResponse::<()> {
|
||||
success: false,
|
||||
message: format!("删除工作流失败: {}", e),
|
||||
data: None,
|
||||
}),
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn start_workflow(
|
||||
State(state): State<AppState>,
|
||||
AxumPath(name): AxumPath<String>,
|
||||
) -> impl IntoResponse {
|
||||
let item = match state.db.get_workflow(&name).await {
|
||||
Ok(Some(item)) => item,
|
||||
Ok(None) => {
|
||||
return (
|
||||
StatusCode::NOT_FOUND,
|
||||
Json(ApiResponse::<()> {
|
||||
success: false,
|
||||
message: format!("工作流 '{}' 未找到", name),
|
||||
data: None,
|
||||
}),
|
||||
)
|
||||
}
|
||||
Err(e) => {
|
||||
return (
|
||||
StatusCode::INTERNAL_SERVER_ERROR,
|
||||
Json(ApiResponse::<()> {
|
||||
success: false,
|
||||
message: format!("获取工作流失败: {}", e),
|
||||
data: None,
|
||||
}),
|
||||
)
|
||||
}
|
||||
};
|
||||
|
||||
if item.status == "running" || item.status == "initializing" {
|
||||
return (
|
||||
StatusCode::BAD_REQUEST,
|
||||
Json(ApiResponse::<()> {
|
||||
success: false,
|
||||
message: format!("工作流 '{}' 已处在初始建立状态中或者已处于运行状态,无需且不允许进行并行重置启动", name),
|
||||
data: None,
|
||||
}),
|
||||
);
|
||||
}
|
||||
|
||||
// 通过原子性抢占更新将状态切换为 initializing,拦截同名流上的多并发调用导致的双重加载破坏性竞态
|
||||
match state.db.transition_workflow_to_initializing(&name).await {
|
||||
Ok(false) => {
|
||||
return (
|
||||
StatusCode::CONFLICT,
|
||||
Json(ApiResponse::<()> {
|
||||
success: false,
|
||||
message: format!("工作流 '{}' 初始化抢占挂起异常,表明已在另一会话上下文中顺利推入启动通道", name),
|
||||
data: None,
|
||||
}),
|
||||
);
|
||||
}
|
||||
Err(e) => {
|
||||
return (
|
||||
StatusCode::INTERNAL_SERVER_ERROR,
|
||||
Json(ApiResponse::<()> {
|
||||
success: false,
|
||||
message: format!("原子化抢占和迁移工作流状态发生异常: {}", e),
|
||||
data: None,
|
||||
}),
|
||||
);
|
||||
}
|
||||
Ok(true) => {}
|
||||
}
|
||||
|
||||
let grid_cfg: GridConfig = match serde_yaml::from_str(&item.config_yaml) {
|
||||
Ok(cfg) => cfg,
|
||||
Err(e) => {
|
||||
let _ = state.db.update_workflow_status(&name, "idle").await;
|
||||
return (
|
||||
StatusCode::BAD_REQUEST,
|
||||
Json(ApiResponse::<()> {
|
||||
success: false,
|
||||
message: format!("解析工作流 YAML 发生语法或参数解析异常: {}", e),
|
||||
data: None,
|
||||
}),
|
||||
);
|
||||
}
|
||||
};
|
||||
|
||||
info!("成功占据独享启动权,开始启动工作流 '{}',系统进行 64/32 维深度平展开网格结构计算化推列并推送队列...", name);
|
||||
match state.scheduler.initialize_grid(&grid_cfg).await {
|
||||
Ok(_) => {
|
||||
let _ = state.db.update_workflow_status(&name, "running").await;
|
||||
let _ = state.scheduler.schedule_pending_tasks().await;
|
||||
(
|
||||
StatusCode::OK,
|
||||
Json(ApiResponse::<()> {
|
||||
success: true,
|
||||
message: format!("工作流 '{}' 建立与挂载成功并已接续排班", name),
|
||||
data: None,
|
||||
}),
|
||||
)
|
||||
}
|
||||
Err(e) => {
|
||||
let _ = state.db.update_workflow_status(&name, "idle").await;
|
||||
(
|
||||
StatusCode::INTERNAL_SERVER_ERROR,
|
||||
Json(ApiResponse::<()> {
|
||||
success: false,
|
||||
message: format!("展开与挂载初始化任务点到系统队列失败: {}", e),
|
||||
data: None,
|
||||
}),
|
||||
)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn stop_workflow(
|
||||
State(state): State<AppState>,
|
||||
AxumPath(name): AxumPath<String>,
|
||||
) -> impl IntoResponse {
|
||||
match state.db.update_workflow_status(&name, "paused").await {
|
||||
Ok(_) => {
|
||||
let _ = state.queue.clear_queue().await;
|
||||
let _ = state.db.reset_queued_grid_points_to_pending().await;
|
||||
(
|
||||
StatusCode::OK,
|
||||
Json(ApiResponse::<()> {
|
||||
success: true,
|
||||
message: format!("工作流 '{}' 已暂停,排队任务已暂停调度", name),
|
||||
data: None,
|
||||
}),
|
||||
)
|
||||
}
|
||||
Err(e) => (
|
||||
StatusCode::INTERNAL_SERVER_ERROR,
|
||||
Json(ApiResponse::<()> {
|
||||
success: false,
|
||||
message: format!("暂停工作流失败: {}", e),
|
||||
data: None,
|
||||
}),
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,873 @@
|
||||
use anyhow::{Context, Result};
|
||||
use common::models::{
|
||||
GridPointParams, GridPointStatus, NodeHeartbeatRequest, NodeInfo,
|
||||
NodeRegisterRequest, TaskReport, TaskStatus,
|
||||
};
|
||||
use r2d2::Pool;
|
||||
use r2d2_sqlite::SqliteConnectionManager;
|
||||
use rusqlite::params;
|
||||
use tracing::info;
|
||||
|
||||
#[derive(Debug)]
|
||||
struct SqliteCustomizer;
|
||||
|
||||
impl r2d2::CustomizeConnection<rusqlite::Connection, rusqlite::Error> for SqliteCustomizer {
|
||||
fn on_acquire(&self, conn: &mut rusqlite::Connection) -> Result<(), rusqlite::Error> {
|
||||
conn.pragma_update(None, "busy_timeout", 5000)?;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct SeedCacheItem {
|
||||
pub point_name: String,
|
||||
pub params: GridPointParams,
|
||||
pub file_path: String,
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct Database {
|
||||
pool: Pool<SqliteConnectionManager>,
|
||||
seed_cache: std::sync::Arc<tokio::sync::RwLock<Vec<SeedCacheItem>>>,
|
||||
}
|
||||
|
||||
impl Database {
|
||||
pub async fn new(db_path: &str) -> Result<Self> {
|
||||
let db_path_owned = db_path.to_string();
|
||||
let pool = tokio::task::spawn_blocking(move || -> Result<Pool<SqliteConnectionManager>> {
|
||||
if let Some(parent) = std::path::Path::new(&db_path_owned).parent() {
|
||||
let _ = std::fs::create_dir_all(parent);
|
||||
}
|
||||
let manager = SqliteConnectionManager::file(&db_path_owned);
|
||||
let pool = Pool::builder()
|
||||
.max_size(8)
|
||||
.connection_customizer(Box::new(SqliteCustomizer))
|
||||
.build(manager)
|
||||
.context("Failed to build SQLite main DB connection pool")?;
|
||||
|
||||
Ok(pool)
|
||||
})
|
||||
.await??;
|
||||
|
||||
let db = Self {
|
||||
pool,
|
||||
seed_cache: std::sync::Arc::new(tokio::sync::RwLock::new(Vec::new())),
|
||||
};
|
||||
db.init_tables().await?;
|
||||
db.reload_seed_cache().await?;
|
||||
Ok(db)
|
||||
}
|
||||
|
||||
async fn init_tables(&self) -> Result<()> {
|
||||
let pool = self.pool.clone();
|
||||
tokio::task::spawn_blocking(move || -> Result<()> {
|
||||
let conn = pool.get().map_err(|e| anyhow::anyhow!("DB Pool Error: {}", e))?;
|
||||
let _: String = conn.pragma_update_and_check(None, "journal_mode", "WAL", |r| r.get(0))?;
|
||||
conn.execute(
|
||||
"CREATE TABLE IF NOT EXISTS nodes (
|
||||
node_id TEXT PRIMARY KEY,
|
||||
host_name TEXT NOT NULL,
|
||||
max_slots INTEGER NOT NULL,
|
||||
active_slots INTEGER NOT NULL DEFAULT 0,
|
||||
status TEXT NOT NULL DEFAULT 'online',
|
||||
cpu_usage REAL NOT NULL DEFAULT 0.0,
|
||||
memory_usage REAL NOT NULL DEFAULT 0.0,
|
||||
last_heartbeat DATETIME NOT NULL
|
||||
);",
|
||||
[],
|
||||
)?;
|
||||
|
||||
conn.execute(
|
||||
"CREATE TABLE IF NOT EXISTS grid_points (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
name TEXT UNIQUE NOT NULL,
|
||||
teff REAL NOT NULL,
|
||||
logg REAL NOT NULL,
|
||||
loghe REAL NOT NULL,
|
||||
logc REAL NOT NULL,
|
||||
logn REAL NOT NULL,
|
||||
logo REAL NOT NULL,
|
||||
cno_sum REAL NOT NULL,
|
||||
wave INTEGER NOT NULL DEFAULT 0,
|
||||
status TEXT NOT NULL DEFAULT 'pending',
|
||||
attempt_count INTEGER NOT NULL DEFAULT 0,
|
||||
success_method TEXT
|
||||
);",
|
||||
[],
|
||||
)?;
|
||||
|
||||
conn.execute(
|
||||
"CREATE TABLE IF NOT EXISTS tasks (
|
||||
task_id TEXT PRIMARY KEY,
|
||||
point_name TEXT NOT NULL,
|
||||
node_id TEXT,
|
||||
task_type TEXT NOT NULL,
|
||||
seed_point_name TEXT,
|
||||
status TEXT NOT NULL DEFAULT 'pending',
|
||||
max_relc REAL,
|
||||
atmosphere_has_nan BOOLEAN NOT NULL DEFAULT 0,
|
||||
retry_count INTEGER NOT NULL DEFAULT 0,
|
||||
created_at DATETIME NOT NULL,
|
||||
started_at DATETIME,
|
||||
completed_at DATETIME,
|
||||
error_message TEXT
|
||||
);",
|
||||
[],
|
||||
)?;
|
||||
|
||||
conn.execute(
|
||||
"CREATE TABLE IF NOT EXISTS seeds (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
point_name TEXT UNIQUE NOT NULL,
|
||||
teff REAL NOT NULL,
|
||||
logg REAL NOT NULL,
|
||||
loghe REAL NOT NULL,
|
||||
logc REAL NOT NULL,
|
||||
logn REAL NOT NULL,
|
||||
logo REAL NOT NULL,
|
||||
file_path TEXT NOT NULL,
|
||||
is_clean BOOLEAN NOT NULL DEFAULT 1
|
||||
);",
|
||||
[],
|
||||
)?;
|
||||
|
||||
conn.execute(
|
||||
"CREATE TABLE IF NOT EXISTS workflows (
|
||||
name TEXT PRIMARY KEY,
|
||||
description TEXT,
|
||||
config_yaml TEXT NOT NULL,
|
||||
status TEXT NOT NULL DEFAULT 'idle',
|
||||
created_at DATETIME NOT NULL,
|
||||
updated_at DATETIME NOT NULL
|
||||
);",
|
||||
[],
|
||||
)?;
|
||||
|
||||
conn.execute(
|
||||
"CREATE INDEX IF NOT EXISTS idx_grid_points_status_wave ON grid_points(status, wave, cno_sum, teff);",
|
||||
[],
|
||||
)?;
|
||||
|
||||
conn.execute(
|
||||
"CREATE INDEX IF NOT EXISTS idx_seeds_is_clean ON seeds(is_clean);",
|
||||
[],
|
||||
)?;
|
||||
|
||||
info!("成功初始化 dcts.db 数据库结构表及索引");
|
||||
Ok(())
|
||||
})
|
||||
.await??;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
// --- Node operations ---
|
||||
pub async fn register_node(&self, req: &NodeRegisterRequest) -> Result<()> {
|
||||
let pool = self.pool.clone();
|
||||
let req_cloned = req.clone();
|
||||
|
||||
tokio::task::spawn_blocking(move || -> Result<()> {
|
||||
let conn = pool.get().map_err(|e| anyhow::anyhow!("DB Pool Error: {}", e))?;
|
||||
conn.execute(
|
||||
"INSERT INTO nodes (node_id, host_name, max_slots, status, last_heartbeat)
|
||||
VALUES (?1, ?2, ?3, 'online', datetime('now'))
|
||||
ON CONFLICT(node_id) DO UPDATE SET
|
||||
host_name = excluded.host_name,
|
||||
max_slots = excluded.max_slots,
|
||||
status = 'online',
|
||||
last_heartbeat = datetime('now')",
|
||||
params![req_cloned.node_id, req_cloned.host_name, req_cloned.max_slots],
|
||||
)?;
|
||||
Ok(())
|
||||
})
|
||||
.await??;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn heartbeat_node(&self, req: &NodeHeartbeatRequest) -> Result<()> {
|
||||
let pool = self.pool.clone();
|
||||
let req_cloned = req.clone();
|
||||
|
||||
tokio::task::spawn_blocking(move || -> Result<()> {
|
||||
let conn = pool.get().map_err(|e| anyhow::anyhow!("DB Pool Error: {}", e))?;
|
||||
conn.execute(
|
||||
"UPDATE nodes SET active_slots = ?1, cpu_usage = ?2, memory_usage = ?3, status = 'online', last_heartbeat = datetime('now')
|
||||
WHERE node_id = ?4",
|
||||
params![
|
||||
req_cloned.active_slots,
|
||||
req_cloned.cpu_usage,
|
||||
req_cloned.memory_usage,
|
||||
req_cloned.node_id
|
||||
],
|
||||
)?;
|
||||
Ok(())
|
||||
})
|
||||
.await??;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn get_active_nodes(&self) -> Result<Vec<NodeInfo>> {
|
||||
let pool = self.pool.clone();
|
||||
|
||||
tokio::task::spawn_blocking(move || -> Result<Vec<NodeInfo>> {
|
||||
let conn = pool.get().map_err(|e| anyhow::anyhow!("DB Pool Error: {}", e))?;
|
||||
let mut stmt = conn.prepare(
|
||||
"SELECT node_id, host_name, max_slots, active_slots, status, cpu_usage, memory_usage, strftime('%Y-%m-%dT%H:%M:%SZ', last_heartbeat) FROM nodes WHERE status = 'online'"
|
||||
)?;
|
||||
|
||||
let node_iter = stmt.query_map([], |r| {
|
||||
let hb_str: String = r.get(7)?;
|
||||
Ok(NodeInfo {
|
||||
node_id: r.get(0)?,
|
||||
host_name: r.get(1)?,
|
||||
max_slots: r.get(2)?,
|
||||
active_slots: r.get(3)?,
|
||||
status: r.get(4)?,
|
||||
cpu_usage: r.get(5)?,
|
||||
memory_usage: r.get(6)?,
|
||||
last_heartbeat: chrono::DateTime::parse_from_rfc3339(&hb_str)
|
||||
.map(|d| d.with_timezone(&chrono::Utc))
|
||||
.unwrap_or_else(|_| chrono::Utc::now()),
|
||||
})
|
||||
})?;
|
||||
|
||||
let mut nodes = Vec::new();
|
||||
for n in node_iter {
|
||||
nodes.push(n?);
|
||||
}
|
||||
Ok(nodes)
|
||||
})
|
||||
.await?
|
||||
}
|
||||
|
||||
// --- Grid Point & Task operations ---
|
||||
pub async fn upsert_grid_point(&self, params_in: &GridPointParams, wave: i32) -> Result<()> {
|
||||
let pool = self.pool.clone();
|
||||
let p = params_in.clone();
|
||||
let name = p.model_name();
|
||||
let cno_sum = p.cno_sum();
|
||||
|
||||
tokio::task::spawn_blocking(move || -> Result<()> {
|
||||
let conn = pool.get().map_err(|e| anyhow::anyhow!("DB Pool Error: {}", e))?;
|
||||
conn.execute(
|
||||
"INSERT INTO grid_points (name, teff, logg, loghe, logc, logn, logo, cno_sum, wave)
|
||||
VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9)
|
||||
ON CONFLICT(name) DO NOTHING",
|
||||
params![name, p.teff, p.logg, p.loghe, p.logc, p.logn, p.logo, cno_sum, wave],
|
||||
)?;
|
||||
Ok(())
|
||||
})
|
||||
.await??;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn get_pending_grid_points(&self) -> Result<Vec<(String, GridPointParams, i32)>> {
|
||||
self.get_pending_grid_points_limit(usize::MAX).await
|
||||
}
|
||||
|
||||
pub async fn get_pending_grid_points_limit(&self, limit: usize) -> Result<Vec<(String, GridPointParams, i32)>> {
|
||||
let pool = self.pool.clone();
|
||||
|
||||
tokio::task::spawn_blocking(move || -> Result<Vec<(String, GridPointParams, i32)>> {
|
||||
let conn = pool.get().map_err(|e| anyhow::anyhow!("DB Pool Error: {}", e))?;
|
||||
let mut stmt = conn.prepare(
|
||||
"SELECT name, teff, logg, loghe, logc, logn, logo, wave FROM grid_points WHERE status = 'pending' ORDER BY wave ASC, cno_sum ASC, teff ASC LIMIT ?1"
|
||||
)?;
|
||||
|
||||
let limit_param = if limit == usize::MAX { -1i64 } else { limit as i64 };
|
||||
let rows_iter = stmt.query_map([limit_param], |r| {
|
||||
Ok((
|
||||
r.get(0)?,
|
||||
GridPointParams {
|
||||
teff: r.get(1)?,
|
||||
logg: r.get(2)?,
|
||||
loghe: r.get(3)?,
|
||||
logc: r.get(4)?,
|
||||
logn: r.get(5)?,
|
||||
logo: r.get(6)?,
|
||||
},
|
||||
r.get(7)?,
|
||||
))
|
||||
})?;
|
||||
|
||||
let mut list = Vec::new();
|
||||
for r in rows_iter {
|
||||
list.push(r?);
|
||||
}
|
||||
Ok(list)
|
||||
})
|
||||
.await?
|
||||
}
|
||||
|
||||
pub async fn reset_queued_grid_points_to_pending(&self) -> Result<usize> {
|
||||
let pool = self.pool.clone();
|
||||
tokio::task::spawn_blocking(move || -> Result<usize> {
|
||||
let conn = pool.get().map_err(|e| anyhow::anyhow!("DB Pool Error: {}", e))?;
|
||||
// 只重置 queued 状态的任务为 pending。对于 running (正由 Worker 处理的项目),不可在系统重启或初始化时粗暴清零,让 Worker 正常完成汇报或触发心跳/超时自动逐回
|
||||
let rows = conn.execute(
|
||||
"UPDATE grid_points SET status = 'pending' WHERE status = 'queued'",
|
||||
[],
|
||||
)?;
|
||||
Ok(rows)
|
||||
})
|
||||
.await?
|
||||
}
|
||||
|
||||
pub async fn reset_specific_grid_points_to_pending(&self, names: &[String]) -> Result<usize> {
|
||||
if names.is_empty() {
|
||||
return Ok(0);
|
||||
}
|
||||
let pool = self.pool.clone();
|
||||
let names_owned = names.to_vec();
|
||||
tokio::task::spawn_blocking(move || -> Result<usize> {
|
||||
let mut conn = pool.get().map_err(|e| anyhow::anyhow!("DB Pool Error: {}", e))?;
|
||||
let tx = conn.transaction()?;
|
||||
let mut count = 0;
|
||||
for name in &names_owned {
|
||||
count += tx.execute(
|
||||
"UPDATE grid_points SET status = 'pending' WHERE name = ?1 AND status IN ('queued', 'running')",
|
||||
params![name],
|
||||
)?;
|
||||
}
|
||||
tx.commit()?;
|
||||
Ok(count)
|
||||
})
|
||||
.await?
|
||||
}
|
||||
|
||||
pub async fn update_grid_status(&self, name: &str, status: GridPointStatus) -> Result<()> {
|
||||
let pool = self.pool.clone();
|
||||
let name_owned = name.to_string();
|
||||
let status_str = status.to_string();
|
||||
|
||||
tokio::task::spawn_blocking(move || -> Result<()> {
|
||||
let conn = pool.get().map_err(|e| anyhow::anyhow!("DB Pool Error: {}", e))?;
|
||||
conn.execute(
|
||||
"UPDATE grid_points SET status = ?1 WHERE name = ?2",
|
||||
params![status_str, name_owned],
|
||||
)?;
|
||||
Ok(())
|
||||
})
|
||||
.await??;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn mark_grid_point_running(&self, name: &str) -> Result<()> {
|
||||
self.update_grid_status(name, GridPointStatus::Running).await
|
||||
}
|
||||
|
||||
pub async fn get_grid_point_status(&self, name: &str) -> Result<Option<(String, i32)>> {
|
||||
let pool = self.pool.clone();
|
||||
let name_owned = name.to_string();
|
||||
tokio::task::spawn_blocking(move || -> Result<Option<(String, i32)>> {
|
||||
let conn = pool.get().map_err(|e| anyhow::anyhow!("DB Pool Error: {}", e))?;
|
||||
let mut stmt = conn.prepare("SELECT status, attempt_count FROM grid_points WHERE name = ?1")?;
|
||||
let res = stmt.query_row(params![name_owned], |r| Ok((r.get(0)?, r.get(1)?)));
|
||||
match res {
|
||||
Ok(tuple) => Ok(Some(tuple)),
|
||||
Err(rusqlite::Error::QueryReturnedNoRows) => Ok(None),
|
||||
Err(e) => Err(e.into()),
|
||||
}
|
||||
})
|
||||
.await?
|
||||
}
|
||||
|
||||
pub async fn insert_task(&self, spec: &common::models::TaskSpec) -> Result<()> {
|
||||
let pool = self.pool.clone();
|
||||
let spec = spec.clone();
|
||||
tokio::task::spawn_blocking(move || -> Result<()> {
|
||||
let conn = pool.get().map_err(|e| anyhow::anyhow!("DB Pool Error: {}", e))?;
|
||||
let task_type_str = match spec.task_type {
|
||||
common::models::TaskType::ColdRun => "cold_run",
|
||||
common::models::TaskType::SeedStep => "seed_step",
|
||||
};
|
||||
conn.execute(
|
||||
"INSERT INTO tasks (task_id, point_name, task_type, seed_point_name, status, created_at)
|
||||
VALUES (?1, ?2, ?3, ?4, 'pending', datetime('now'))
|
||||
ON CONFLICT(task_id) DO UPDATE SET
|
||||
status = 'pending',
|
||||
seed_point_name = excluded.seed_point_name",
|
||||
params![
|
||||
spec.task_id.to_string(),
|
||||
spec.point_name,
|
||||
task_type_str,
|
||||
spec.seed_point_name
|
||||
],
|
||||
)?;
|
||||
Ok(())
|
||||
})
|
||||
.await??;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn record_task_report(&self, report: &TaskReport) -> Result<()> {
|
||||
let pool = self.pool.clone();
|
||||
let report_cloned = report.clone();
|
||||
let point_name = report.point_name.clone();
|
||||
let converged = report.converged;
|
||||
let atmo_has_nan = report.atmosphere_has_nan;
|
||||
|
||||
tokio::task::spawn_blocking(move || -> Result<()> {
|
||||
let mut conn = pool.get().map_err(|e| anyhow::anyhow!("DB Pool Error: {}", e))?;
|
||||
let tx = conn.transaction()?;
|
||||
|
||||
let status_str = match report_cloned.status {
|
||||
TaskStatus::Completed => "completed",
|
||||
TaskStatus::Failed => "failed",
|
||||
TaskStatus::Timeout => "timeout",
|
||||
_ => "pending",
|
||||
};
|
||||
|
||||
tx.execute(
|
||||
"UPDATE tasks SET status = ?1, node_id = ?2, max_relc = ?3, atmosphere_has_nan = ?4, completed_at = datetime('now'), error_message = ?5 WHERE task_id = ?6",
|
||||
params![
|
||||
status_str,
|
||||
report_cloned.node_id,
|
||||
report_cloned.max_relc,
|
||||
report_cloned.atmosphere_has_nan,
|
||||
report_cloned.error_message,
|
||||
report_cloned.task_id.to_string()
|
||||
],
|
||||
)?;
|
||||
|
||||
// 合并重试次数 +1 与查值操作至一条 atomic sql UPDATE RETURNING 语句,彻底杜绝多事务并发下的读写竞态;若失败则返回 i32::MAX 强制定向至 failed 回退保护
|
||||
let current_attempts: i32 = tx
|
||||
.query_row(
|
||||
"UPDATE grid_points SET attempt_count = attempt_count + 1 WHERE name = ?1 RETURNING attempt_count",
|
||||
params![point_name],
|
||||
|r| r.get(0),
|
||||
)
|
||||
.unwrap_or(i32::MAX);
|
||||
|
||||
if report_cloned.status == TaskStatus::Completed && converged && !atmo_has_nan {
|
||||
tx.execute(
|
||||
"UPDATE grid_points SET status = 'converged', success_method = (SELECT task_type FROM tasks WHERE task_id = ?1) WHERE name = ?2",
|
||||
params![report_cloned.task_id.to_string(), point_name],
|
||||
)?;
|
||||
} else {
|
||||
let max_attempts = 3;
|
||||
if current_attempts >= max_attempts {
|
||||
tx.execute(
|
||||
"UPDATE grid_points SET status = 'failed' WHERE name = ?1",
|
||||
params![point_name],
|
||||
)?;
|
||||
} else {
|
||||
tx.execute(
|
||||
"UPDATE grid_points SET status = 'pending' WHERE name = ?1",
|
||||
params![point_name],
|
||||
)?;
|
||||
}
|
||||
}
|
||||
|
||||
tx.commit()?;
|
||||
Ok(())
|
||||
})
|
||||
.await??;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn mark_stale_nodes_offline(&self, stale_sec: u64) -> Result<u64> {
|
||||
let pool = self.pool.clone();
|
||||
let count = tokio::task::spawn_blocking(move || -> Result<u64> {
|
||||
let conn = pool.get().map_err(|e| anyhow::anyhow!("DB Pool Error: {}", e))?;
|
||||
let rows = conn.execute(
|
||||
"UPDATE nodes SET status = 'offline' WHERE status = 'online' AND strftime('%s', 'now') - strftime('%s', last_heartbeat) > ?1",
|
||||
params![stale_sec as i64],
|
||||
)?;
|
||||
Ok(rows as u64)
|
||||
})
|
||||
.await??;
|
||||
|
||||
Ok(count)
|
||||
}
|
||||
|
||||
pub async fn reload_seed_cache(&self) -> Result<()> {
|
||||
let pool = self.pool.clone();
|
||||
let items = tokio::task::spawn_blocking(move || -> Result<Vec<SeedCacheItem>> {
|
||||
let conn = pool.get().map_err(|e| anyhow::anyhow!("DB Pool Error: {}", e))?;
|
||||
let mut stmt = conn.prepare(
|
||||
"SELECT point_name, teff, logg, loghe, logc, logn, logo, file_path FROM seeds WHERE is_clean = 1"
|
||||
)?;
|
||||
let rows = stmt.query_map([], |row| {
|
||||
Ok(SeedCacheItem {
|
||||
point_name: row.get(0)?,
|
||||
params: GridPointParams {
|
||||
teff: row.get(1)?,
|
||||
logg: row.get(2)?,
|
||||
loghe: row.get(3)?,
|
||||
logc: row.get(4)?,
|
||||
logn: row.get(5)?,
|
||||
logo: row.get(6)?,
|
||||
},
|
||||
file_path: row.get(7)?,
|
||||
})
|
||||
})?;
|
||||
let mut list = Vec::new();
|
||||
for r in rows {
|
||||
list.push(r?);
|
||||
}
|
||||
Ok(list)
|
||||
})
|
||||
.await??;
|
||||
|
||||
let mut lock = self.seed_cache.write().await;
|
||||
*lock = items;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn insert_seed(&self, params_in: &GridPointParams, file_path: &str) -> Result<()> {
|
||||
let pool = self.pool.clone();
|
||||
let p = params_in.clone();
|
||||
let name = p.model_name();
|
||||
let path_owned = file_path.to_string();
|
||||
|
||||
let name_db = name.clone();
|
||||
let path_db = path_owned.clone();
|
||||
let p_db = p.clone();
|
||||
tokio::task::spawn_blocking(move || -> Result<()> {
|
||||
let conn = pool.get().map_err(|e| anyhow::anyhow!("DB Pool Error: {}", e))?;
|
||||
conn.execute(
|
||||
"INSERT INTO seeds (point_name, teff, logg, loghe, logc, logn, logo, file_path)
|
||||
VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8)
|
||||
ON CONFLICT(point_name) DO UPDATE SET file_path = excluded.file_path",
|
||||
params![name_db, p_db.teff, p_db.logg, p_db.loghe, p_db.logc, p_db.logn, p_db.logo, path_db],
|
||||
)?;
|
||||
Ok(())
|
||||
})
|
||||
.await??;
|
||||
|
||||
let item = SeedCacheItem {
|
||||
point_name: name,
|
||||
params: p,
|
||||
file_path: path_owned,
|
||||
};
|
||||
let mut lock = self.seed_cache.write().await;
|
||||
if let Some(pos) = lock.iter().position(|x| x.point_name == item.point_name) {
|
||||
lock[pos] = item;
|
||||
} else {
|
||||
lock.push(item);
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn find_best_seed_from_db(
|
||||
&self,
|
||||
target: &GridPointParams,
|
||||
) -> Result<Option<common::seed_finder::SeedMatch>> {
|
||||
let lock = self.seed_cache.read().await;
|
||||
let mut exact_family: Option<(String, std::path::PathBuf, f64)> = None;
|
||||
let mut global_closest: Option<(String, std::path::PathBuf, f64)> = None;
|
||||
|
||||
for item in lock.iter() {
|
||||
let path = std::path::PathBuf::from(&item.file_path);
|
||||
let (is_exact, d) = common::seed_finder::calculate_seed_distance(&item.params, target);
|
||||
|
||||
if is_exact {
|
||||
if exact_family.is_none() || d < exact_family.as_ref().unwrap().2 {
|
||||
exact_family = Some((item.point_name.clone(), path, d));
|
||||
}
|
||||
} else if d <= common::seed_finder::MAX_GLOBAL_SEED_DISTANCE && (global_closest.is_none() || d < global_closest.as_ref().unwrap().2) {
|
||||
global_closest = Some((item.point_name.clone(), path, d));
|
||||
}
|
||||
}
|
||||
|
||||
if let Some((name, path, d)) = exact_family {
|
||||
Ok(Some(common::seed_finder::SeedMatch { name, path, distance: d }))
|
||||
} else if let Some((name, path, d)) = global_closest {
|
||||
Ok(Some(common::seed_finder::SeedMatch { name, path, distance: d }))
|
||||
} else {
|
||||
Ok(None)
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn upsert_workflow(
|
||||
&self,
|
||||
name: &str,
|
||||
description: Option<&str>,
|
||||
config_yaml: &str,
|
||||
status: &str,
|
||||
) -> Result<()> {
|
||||
let pool = self.pool.clone();
|
||||
let name_owned = name.to_string();
|
||||
let desc_owned = description.map(|s| s.to_string());
|
||||
let yaml_owned = config_yaml.to_string();
|
||||
let status_owned = status.to_string();
|
||||
|
||||
tokio::task::spawn_blocking(move || -> Result<()> {
|
||||
let conn = pool.get().map_err(|e| anyhow::anyhow!("DB Pool Error: {}", e))?;
|
||||
conn.execute(
|
||||
"INSERT INTO workflows (name, description, config_yaml, status, created_at, updated_at)
|
||||
VALUES (?1, ?2, ?3, ?4, datetime('now'), datetime('now'))
|
||||
ON CONFLICT(name) DO UPDATE SET
|
||||
description = excluded.description,
|
||||
config_yaml = excluded.config_yaml,
|
||||
status = CASE WHEN workflows.status = 'running' THEN workflows.status ELSE excluded.status END,
|
||||
updated_at = datetime('now')",
|
||||
params![name_owned, desc_owned, yaml_owned, status_owned],
|
||||
)?;
|
||||
Ok(())
|
||||
})
|
||||
.await??;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn list_workflows(&self) -> Result<Vec<WorkflowSummary>> {
|
||||
let pool = self.pool.clone();
|
||||
tokio::task::spawn_blocking(move || -> Result<Vec<WorkflowSummary>> {
|
||||
let conn = pool.get().map_err(|e| anyhow::anyhow!("DB Pool Error: {}", e))?;
|
||||
let mut stmt = conn.prepare(
|
||||
"SELECT name, description, status, created_at, updated_at FROM workflows ORDER BY updated_at DESC"
|
||||
)?;
|
||||
let rows = stmt.query_map([], |row| {
|
||||
Ok(WorkflowSummary {
|
||||
name: row.get(0)?,
|
||||
description: row.get(1)?,
|
||||
status: row.get(2)?,
|
||||
created_at: row.get(3)?,
|
||||
updated_at: row.get(4)?,
|
||||
})
|
||||
})?;
|
||||
|
||||
let mut list = Vec::new();
|
||||
for r in rows {
|
||||
list.push(r?);
|
||||
}
|
||||
Ok(list)
|
||||
})
|
||||
.await?
|
||||
}
|
||||
|
||||
pub async fn get_workflow(&self, name: &str) -> Result<Option<WorkflowItem>> {
|
||||
let pool = self.pool.clone();
|
||||
let name_owned = name.to_string();
|
||||
tokio::task::spawn_blocking(move || -> Result<Option<WorkflowItem>> {
|
||||
let conn = pool.get().map_err(|e| anyhow::anyhow!("DB Pool Error: {}", e))?;
|
||||
let mut stmt = conn.prepare(
|
||||
"SELECT name, description, config_yaml, status, created_at, updated_at FROM workflows WHERE name = ?1"
|
||||
)?;
|
||||
let row = stmt.query_row(params![name_owned], |row| {
|
||||
Ok(WorkflowItem {
|
||||
name: row.get(0)?,
|
||||
description: row.get(1)?,
|
||||
config_yaml: row.get(2)?,
|
||||
status: row.get(3)?,
|
||||
created_at: row.get(4)?,
|
||||
updated_at: row.get(5)?,
|
||||
})
|
||||
});
|
||||
|
||||
match row {
|
||||
Ok(item) => Ok(Some(item)),
|
||||
Err(rusqlite::Error::QueryReturnedNoRows) => Ok(None),
|
||||
Err(e) => Err(e.into()),
|
||||
}
|
||||
})
|
||||
.await?
|
||||
}
|
||||
|
||||
pub async fn update_workflow_status(&self, name: &str, status: &str) -> Result<()> {
|
||||
let pool = self.pool.clone();
|
||||
let name_owned = name.to_string();
|
||||
let status_owned = status.to_string();
|
||||
|
||||
tokio::task::spawn_blocking(move || -> Result<()> {
|
||||
let conn = pool.get().map_err(|e| anyhow::anyhow!("DB Pool Error: {}", e))?;
|
||||
conn.execute(
|
||||
"UPDATE workflows SET status = ?1, updated_at = datetime('now') WHERE name = ?2",
|
||||
params![status_owned, name_owned],
|
||||
)?;
|
||||
Ok(())
|
||||
})
|
||||
.await??;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn delete_workflow(&self, name: &str) -> Result<()> {
|
||||
let pool = self.pool.clone();
|
||||
let name_owned = name.to_string();
|
||||
|
||||
tokio::task::spawn_blocking(move || -> Result<()> {
|
||||
let conn = pool.get().map_err(|e| anyhow::anyhow!("DB Pool Error: {}", e))?;
|
||||
conn.execute("DELETE FROM workflows WHERE name = ?1", params![name_owned])?;
|
||||
Ok(())
|
||||
})
|
||||
.await??;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn has_running_workflow(&self) -> Result<bool> {
|
||||
let pool = self.pool.clone();
|
||||
tokio::task::spawn_blocking(move || -> Result<bool> {
|
||||
let conn = pool.get().map_err(|e| anyhow::anyhow!("DB Pool Error: {}", e))?;
|
||||
let count: i64 = conn.query_row(
|
||||
"SELECT COUNT(*) FROM workflows WHERE status = 'running'",
|
||||
[],
|
||||
|r| r.get(0),
|
||||
)?;
|
||||
Ok(count > 0)
|
||||
})
|
||||
.await?
|
||||
}
|
||||
|
||||
/// 原子切转工作流至 initializing 预占启动状态,杜绝高并发 POST /start 触发双重全量排队与重置网格竞态
|
||||
pub async fn transition_workflow_to_initializing(&self, name: &str) -> Result<bool> {
|
||||
let pool = self.pool.clone();
|
||||
let name_owned = name.to_string();
|
||||
let affected = tokio::task::spawn_blocking(move || -> Result<usize> {
|
||||
let conn = pool.get().map_err(|e| anyhow::anyhow!("DB Pool Error: {}", e))?;
|
||||
let count = conn.execute(
|
||||
"UPDATE workflows SET status = 'initializing', updated_at = datetime('now') WHERE name = ?1 AND status NOT IN ('running', 'initializing')",
|
||||
params![name_owned],
|
||||
)?;
|
||||
Ok(count)
|
||||
})
|
||||
.await??;
|
||||
Ok(affected > 0)
|
||||
}
|
||||
|
||||
/// 获取运行或启动态中的所有工作流 YAML 配置(替代原来低效 N 次循环与嵌套查询)
|
||||
pub async fn get_running_workflow_config_yamls(&self) -> Result<Vec<String>> {
|
||||
let pool = self.pool.clone();
|
||||
tokio::task::spawn_blocking(move || -> Result<Vec<String>> {
|
||||
let conn = pool.get().map_err(|e| anyhow::anyhow!("DB Pool Error: {}", e))?;
|
||||
let mut stmt = conn.prepare("SELECT config_yaml FROM workflows WHERE status IN ('running', 'initializing')")?;
|
||||
let rows = stmt.query_map([], |row| row.get(0))?;
|
||||
let mut list = Vec::new();
|
||||
for r in rows {
|
||||
list.push(r?);
|
||||
}
|
||||
Ok(list)
|
||||
})
|
||||
.await?
|
||||
}
|
||||
|
||||
pub async fn get_grid_summary_stats(&self) -> Result<serde_json::Value> {
|
||||
let pool = self.pool.clone();
|
||||
tokio::task::spawn_blocking(move || -> Result<serde_json::Value> {
|
||||
let conn = pool.get().map_err(|e| anyhow::anyhow!("DB Pool Error: {}", e))?;
|
||||
let total: i64 = conn.query_row("SELECT COUNT(*) FROM grid_points", [], |r| r.get(0)).unwrap_or(0);
|
||||
let pending: i64 = conn.query_row("SELECT COUNT(*) FROM grid_points WHERE status IN ('pending', 'queued')", [], |r| r.get(0)).unwrap_or(0);
|
||||
let running: i64 = conn.query_row("SELECT COUNT(*) FROM grid_points WHERE status = 'running'", [], |r| r.get(0)).unwrap_or(0);
|
||||
let converged: i64 = conn.query_row("SELECT COUNT(*) FROM grid_points WHERE status = 'converged'", [], |r| r.get(0)).unwrap_or(0);
|
||||
let failed: i64 = conn.query_row("SELECT COUNT(*) FROM grid_points WHERE status = 'failed'", [], |r| r.get(0)).unwrap_or(0);
|
||||
let cold_run_converged: i64 = conn.query_row("SELECT COUNT(*) FROM grid_points WHERE status = 'converged' AND success_method = 'cold_run'", [], |r| r.get(0)).unwrap_or(0);
|
||||
let seed_step_converged: i64 = conn.query_row("SELECT COUNT(*) FROM grid_points WHERE status = 'converged' AND success_method = 'seed_step'", [], |r| r.get(0)).unwrap_or(0);
|
||||
|
||||
Ok(serde_json::json!({
|
||||
"total": total,
|
||||
"pending": pending,
|
||||
"running": running,
|
||||
"converged": converged,
|
||||
"failed": failed,
|
||||
"cold_run_converged": cold_run_converged,
|
||||
"seed_step_converged": seed_step_converged,
|
||||
}))
|
||||
})
|
||||
.await?
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
|
||||
pub struct WorkflowSummary {
|
||||
pub name: String,
|
||||
pub description: Option<String>,
|
||||
pub status: String,
|
||||
pub created_at: String,
|
||||
pub updated_at: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
|
||||
pub struct WorkflowItem {
|
||||
pub name: String,
|
||||
pub description: Option<String>,
|
||||
pub config_yaml: String,
|
||||
pub status: String,
|
||||
pub created_at: String,
|
||||
pub updated_at: String,
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use uuid::Uuid;
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_db_node_and_grid_operations() {
|
||||
let temp_dir = tempfile::tempdir().unwrap();
|
||||
let db_path = temp_dir.path().join("test_db.db");
|
||||
let db = Database::new(&db_path.to_string_lossy()).await.unwrap();
|
||||
|
||||
// 1. Node registration & heartbeat
|
||||
let reg_req = NodeRegisterRequest {
|
||||
node_id: "node-test-1".to_string(),
|
||||
host_name: "localhost".to_string(),
|
||||
max_slots: 4,
|
||||
};
|
||||
db.register_node(®_req).await.unwrap();
|
||||
|
||||
let active_nodes = db.get_active_nodes().await.unwrap();
|
||||
assert_eq!(active_nodes.len(), 1);
|
||||
assert_eq!(active_nodes[0].node_id, "node-test-1");
|
||||
|
||||
let hb_req = NodeHeartbeatRequest {
|
||||
node_id: "node-test-1".to_string(),
|
||||
active_slots: 2,
|
||||
cpu_usage: 45.0,
|
||||
memory_usage: 60.0,
|
||||
};
|
||||
db.heartbeat_node(&hb_req).await.unwrap();
|
||||
|
||||
// 2. Grid points & task reports
|
||||
let params = GridPointParams {
|
||||
teff: 35000.0,
|
||||
logg: 5.5,
|
||||
loghe: -1.0,
|
||||
logc: -2.0,
|
||||
logn: -2.0,
|
||||
logo: -2.0,
|
||||
};
|
||||
db.upsert_grid_point(¶ms, 0).await.unwrap();
|
||||
|
||||
let pending = db.get_pending_grid_points().await.unwrap();
|
||||
assert_eq!(pending.len(), 1);
|
||||
assert_eq!(pending[0].0, params.model_name());
|
||||
|
||||
// Record successful report
|
||||
let report = TaskReport {
|
||||
task_id: Uuid::new_v4(),
|
||||
point_name: params.model_name(),
|
||||
params: Some(params.clone()),
|
||||
node_id: "node-test-1".to_string(),
|
||||
status: TaskStatus::Completed,
|
||||
converged: true,
|
||||
max_relc: Some(0.0001),
|
||||
atmosphere_has_nan: false,
|
||||
elapsed_sec: 15.0,
|
||||
error_message: None,
|
||||
summary_json: "{}".to_string(),
|
||||
};
|
||||
db.record_task_report(&report).await.unwrap();
|
||||
|
||||
// Check grid point is marked converged
|
||||
let pending_after = db.get_pending_grid_points().await.unwrap();
|
||||
assert_eq!(pending_after.len(), 0);
|
||||
|
||||
// 3. Workflow CRUD
|
||||
db.upsert_workflow("test_wf", Some("Test Workflow"), "grid:\n teff: [35000]", "idle").await.unwrap();
|
||||
let wf = db.get_workflow("test_wf").await.unwrap();
|
||||
assert!(wf.is_some());
|
||||
assert_eq!(wf.unwrap().name, "test_wf");
|
||||
|
||||
db.delete_workflow("test_wf").await.unwrap();
|
||||
assert!(db.get_workflow("test_wf").await.unwrap().is_none());
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,4 @@
|
||||
pub mod api;
|
||||
pub mod db;
|
||||
pub mod scheduler;
|
||||
|
||||
@@ -0,0 +1,162 @@
|
||||
use anyhow::Result;
|
||||
use server::api::{self, AppState};
|
||||
use server::db::Database;
|
||||
use server::scheduler::GridScheduler;
|
||||
|
||||
|
||||
use axum::{
|
||||
routing::{get, post},
|
||||
Router,
|
||||
};
|
||||
use clap::Parser;
|
||||
use common::config::ServerConfig;
|
||||
use common::logging::init_logging;
|
||||
use mq::sqlite_queue::SqliteTaskQueue;
|
||||
use std::net::SocketAddr;
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::sync::Arc;
|
||||
use tokio::time::{sleep, Duration};
|
||||
use tower_http::cors::CorsLayer;
|
||||
use tower_http::services::{ServeDir, ServeFile};
|
||||
use tracing::info;
|
||||
|
||||
#[derive(Parser, Debug)]
|
||||
#[command(name = "server", version = "0.1.0", about = "Distributed Computing TLUSTY/SYNSPEC (DCTS) Server")]
|
||||
struct CliArgs {
|
||||
/// Optional path to workflow configuration YAML file to auto-register on startup
|
||||
#[arg(short = 'w', long = "workflow")]
|
||||
workflow: Option<PathBuf>,
|
||||
|
||||
/// Listen port (overrides DCTS_PORT env var)
|
||||
#[arg(short = 'p', long = "port")]
|
||||
port: Option<u16>,
|
||||
}
|
||||
|
||||
#[tokio::main]
|
||||
async fn main() -> Result<()> {
|
||||
dotenvy::dotenv().ok();
|
||||
let _logging_guards = init_logging("server", "info,server=debug")?;
|
||||
|
||||
let cli = CliArgs::parse();
|
||||
|
||||
info!("启动 DCTS 服务端 (Distributed Computing TLUSTY/SYNSPEC Server)...");
|
||||
|
||||
let mut server_cfg = ServerConfig::default();
|
||||
if let Some(port) = cli.port {
|
||||
server_cfg.bind_addr = format!("0.0.0.0:{}", port);
|
||||
}
|
||||
if let Some(wf) = cli.workflow {
|
||||
server_cfg.grid_config = wf.to_string_lossy().to_string();
|
||||
}
|
||||
|
||||
let db = Database::new(&server_cfg.db_path).await?;
|
||||
let queue = Arc::new(SqliteTaskQueue::new(&server_cfg.queue_db_path).await?);
|
||||
let scheduler = Arc::new(GridScheduler::new(
|
||||
db.clone(),
|
||||
queue.clone(),
|
||||
server_cfg.results_dir.clone(),
|
||||
));
|
||||
|
||||
// Auto-register sdB_cno.yaml if exists and not yet in DB
|
||||
let default_wf_path = Path::new(&server_cfg.grid_config);
|
||||
if default_wf_path.is_file() {
|
||||
if let Ok(yaml_content) = std::fs::read_to_string(default_wf_path) {
|
||||
if let Err(e) = db.upsert_workflow(
|
||||
"sdB_cno",
|
||||
Some("sdB CNO 6D Stellar Atmosphere Grid"),
|
||||
&yaml_content,
|
||||
"idle",
|
||||
).await {
|
||||
tracing::warn!("预注册默认工作流失败: {}", e);
|
||||
} else {
|
||||
info!("已在数据库中成功预注册默认工作流 'sdB_cno'");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let state = AppState {
|
||||
db,
|
||||
queue: queue.clone(),
|
||||
scheduler: scheduler.clone(),
|
||||
results_dir: server_cfg.results_dir,
|
||||
auth_token: server_cfg.auth_token.clone(),
|
||||
};
|
||||
|
||||
// Background loop for stale task requeueing, offline node detection, and scheduler checking
|
||||
let bg_db = state.db.clone();
|
||||
let bg_queue = queue.clone();
|
||||
let bg_scheduler = scheduler.clone();
|
||||
let stale_sec = server_cfg.stale_sec;
|
||||
let node_stale_sec = server_cfg.node_stale_sec;
|
||||
|
||||
tokio::spawn(async move {
|
||||
loop {
|
||||
sleep(Duration::from_secs(30)).await;
|
||||
if let Ok(requeued_points) = bg_queue.requeue_stale_tasks(stale_sec).await {
|
||||
if !requeued_points.is_empty() {
|
||||
info!("重新将 {} 个超时/掉线任务放回待计算队列", requeued_points.len());
|
||||
let _ = bg_db.reset_specific_grid_points_to_pending(&requeued_points).await;
|
||||
}
|
||||
}
|
||||
if let Ok(offline) = bg_db.mark_stale_nodes_offline(node_stale_sec).await {
|
||||
if offline > 0 {
|
||||
info!("已标记 {} 个心跳超时的计算节点为离线状态", offline);
|
||||
}
|
||||
}
|
||||
if let Err(e) = bg_scheduler.schedule_pending_tasks().await {
|
||||
tracing::warn!("后台定时性任务调度检测失败: {}", e);
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
let api_router = Router::new()
|
||||
// Core Node & Task API
|
||||
.route("/node/register", post(api::node::register_node))
|
||||
.route("/node/heartbeat", post(api::node::heartbeat_node))
|
||||
.route("/task/claim", post(api::task::claim_task))
|
||||
.route("/task/report", post(api::task::report_task))
|
||||
.route("/seed/:name", get(api::seed::download_seed))
|
||||
.route("/status", get(api::status::get_status))
|
||||
// Static Data API
|
||||
.route("/data/file/*filename", get(api::data::download_single_data_file))
|
||||
.route("/data/linelist", get(api::data::download_linelist))
|
||||
// Workflow Management CRUD API
|
||||
.route("/workflows", get(api::workflow::list_workflows).post(api::workflow::save_workflow))
|
||||
.route("/workflows/:name", get(api::workflow::get_workflow).put(api::workflow::save_workflow).delete(api::workflow::delete_workflow))
|
||||
.route("/workflows/:name/start", post(api::workflow::start_workflow))
|
||||
.route("/workflows/:name/stop", post(api::workflow::stop_workflow));
|
||||
|
||||
let api_router = if state.auth_token.is_some() {
|
||||
info!("已为 DCTS 服务端 API 路由启用 Bearer Token / X-API-Key 访问控制鉴权");
|
||||
let auth_layer = axum::middleware::from_fn_with_state(state.clone(), api::auth_middleware);
|
||||
api_router.layer(auth_layer)
|
||||
} else {
|
||||
tracing::warn!("⚠️ 警告:未检测到 DCTS_AUTH_TOKEN 环境变量,服务端目前运行在【内网无鉴权模式】!所有 REST API 接口均为公开可访问状态。");
|
||||
api_router
|
||||
};
|
||||
|
||||
// Host Dashboard SPA static files from dashboard/dist if directory exists or fallback to index.html
|
||||
let serve_dir = ServeDir::new("dashboard/dist")
|
||||
.fallback(ServeFile::new("dashboard/dist/index.html"));
|
||||
|
||||
let app = Router::new()
|
||||
.nest("/api", api_router)
|
||||
.layer(CorsLayer::permissive())
|
||||
.fallback_service(serve_dir)
|
||||
.with_state(state);
|
||||
|
||||
let addr: SocketAddr = server_cfg.bind_addr.parse()?;
|
||||
info!("DCTS 服务端已在 http://{} 启动监听", addr);
|
||||
|
||||
let listener = tokio::net::TcpListener::bind(addr).await?;
|
||||
axum::serve(listener, app)
|
||||
.with_graceful_shutdown(async {
|
||||
let _ = tokio::signal::ctrl_c().await;
|
||||
info!("收到 Ctrl+C 终止信号,DCTS 服务端准备优雅关闭...");
|
||||
})
|
||||
.await?;
|
||||
|
||||
info!("DCTS 服务端已安全关闭。");
|
||||
Ok(())
|
||||
}
|
||||
|
||||
@@ -0,0 +1,254 @@
|
||||
use anyhow::Result;
|
||||
use common::config::GridConfig;
|
||||
use common::models::{GridPointParams, TaskSpec, TaskType};
|
||||
use mq::sqlite_queue::SqliteTaskQueue;
|
||||
use std::sync::Arc;
|
||||
use tracing::info;
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::db::Database;
|
||||
|
||||
pub struct GridScheduler {
|
||||
db: Database,
|
||||
queue: Arc<SqliteTaskQueue>,
|
||||
_results_dir: String,
|
||||
}
|
||||
|
||||
impl GridScheduler {
|
||||
pub fn new(db: Database, queue: Arc<SqliteTaskQueue>, results_dir: String) -> Self {
|
||||
Self {
|
||||
db,
|
||||
queue,
|
||||
_results_dir: results_dir,
|
||||
}
|
||||
}
|
||||
|
||||
/// Expands grid points from config and registers them into the database
|
||||
pub async fn initialize_grid(&self, cfg: &GridConfig) -> Result<()> {
|
||||
if let Err(e) = self.queue.clear_queue().await {
|
||||
tracing::warn!("初始化网格时清理闲置排队记录发生警告: {}", e);
|
||||
}
|
||||
if let Err(e) = self.db.reset_queued_grid_points_to_pending().await {
|
||||
tracing::warn!("重置网格状态到 pending 处理过程遇到异常: {}", e);
|
||||
}
|
||||
let mut points = Vec::new();
|
||||
|
||||
for &teff in &cfg.grid.teff {
|
||||
for &logg in &cfg.grid.logg {
|
||||
for &loghe in &cfg.grid.loghe {
|
||||
for &logc in &cfg.grid.logc {
|
||||
for &logn in &cfg.grid.logn {
|
||||
for &logo in &cfg.grid.logo {
|
||||
points.push(GridPointParams {
|
||||
teff,
|
||||
logg,
|
||||
loghe,
|
||||
logc,
|
||||
logn,
|
||||
logo,
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Sort by difficulty: cno_sum ASC -> teff ASC -> -logg -> loghe ASC
|
||||
points.sort_by(|a, b| {
|
||||
a.cno_sum()
|
||||
.partial_cmp(&b.cno_sum())
|
||||
.unwrap_or(std::cmp::Ordering::Equal)
|
||||
.then_with(|| a.teff.partial_cmp(&b.teff).unwrap_or(std::cmp::Ordering::Equal))
|
||||
.then_with(|| b.logg.partial_cmp(&a.logg).unwrap_or(std::cmp::Ordering::Equal))
|
||||
.then_with(|| a.loghe.partial_cmp(&b.loghe).unwrap_or(std::cmp::Ordering::Equal))
|
||||
});
|
||||
|
||||
// Group into Waves by cno_sum
|
||||
let mut current_cno: Option<f64> = None;
|
||||
let mut wave_idx = 0;
|
||||
|
||||
for pt in &points {
|
||||
let cno = pt.cno_sum();
|
||||
if let Some(cur) = current_cno {
|
||||
if (cur - cno).abs() > 1e-5 {
|
||||
wave_idx += 1;
|
||||
current_cno = Some(cno);
|
||||
}
|
||||
} else {
|
||||
current_cno = Some(cno);
|
||||
}
|
||||
|
||||
self.db.upsert_grid_point(pt, wave_idx).await?;
|
||||
}
|
||||
|
||||
info!("已在数据库中成功初始化并记录 {} 个恒星大气网格点", points.len());
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn get_active_timeout_sec(&self) -> u64 {
|
||||
if let Ok(yamls) = self.db.get_running_workflow_config_yamls().await {
|
||||
for yaml in yamls {
|
||||
if let Ok(cfg) = serde_yaml::from_str::<GridConfig>(&yaml) {
|
||||
return cfg.timeout_sec;
|
||||
}
|
||||
}
|
||||
}
|
||||
7200
|
||||
}
|
||||
|
||||
/// Enqueues pending grid points into MQ with active seed detection and batching
|
||||
pub async fn schedule_pending_tasks(&self) -> Result<usize> {
|
||||
if !self.db.has_running_workflow().await? {
|
||||
return Ok(0);
|
||||
}
|
||||
|
||||
let timeout_sec = self.get_active_timeout_sec().await;
|
||||
let batch_limit: usize = std::env::var("DCTS_BATCH_LIMIT")
|
||||
.ok()
|
||||
.and_then(|v| v.parse().ok())
|
||||
.unwrap_or(100);
|
||||
|
||||
// SQL 层直接附加 LIMIT = batch_limit 筛选,完全免除数万点位无谓内存反序列化和空耗对象释放开销
|
||||
let pending = self.db.get_pending_grid_points_limit(batch_limit).await?;
|
||||
let mut dispatched = 0;
|
||||
|
||||
for (name, params, _wave) in pending {
|
||||
|
||||
// Check if any seed is available in DB for active SeedStep scheduling
|
||||
let (task_type, seed_name) = match self.db.find_best_seed_from_db(¶ms).await {
|
||||
Ok(Some(seed_match)) => {
|
||||
info!("网格点 {} 匹配到数据库近邻种子 {} (距离: {:.2}),安排 SeedStep 热启动调度", name, seed_match.name, seed_match.distance);
|
||||
(TaskType::SeedStep, Some(seed_match.name))
|
||||
}
|
||||
_ => (TaskType::ColdRun, None),
|
||||
};
|
||||
|
||||
let task_spec = TaskSpec {
|
||||
task_id: Uuid::new_v4(),
|
||||
point_name: name.clone(),
|
||||
params,
|
||||
task_type,
|
||||
seed_point_name: seed_name,
|
||||
timeout_sec,
|
||||
};
|
||||
|
||||
self.db.insert_task(&task_spec).await?;
|
||||
// 采用先标记 DB 状态为 Queued 后发 MQ 的时序,防止推入 MQ 后数据库修改异常导向下一轮误重投
|
||||
self.db.update_grid_status(&name, common::models::GridPointStatus::Queued).await?;
|
||||
match self.queue.push_task(&task_spec).await {
|
||||
Ok(_) => {
|
||||
dispatched += 1;
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!("将任务 {} 推入 MQ 队列失败,回滚网格点状态: {}", name, e);
|
||||
let _ = self.db.update_grid_status(&name, common::models::GridPointStatus::Pending).await;
|
||||
let _ = self.queue.remove_task(&task_spec.task_id.to_string()).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if dispatched > 0 {
|
||||
info!("已成功将 {} 个待计算网格点推进任务队列", dispatched);
|
||||
}
|
||||
|
||||
Ok(dispatched)
|
||||
}
|
||||
|
||||
/// Triggers seed_step fallback for a failed point if a seed is available
|
||||
pub async fn trigger_seed_step_fallback(&self, params: &GridPointParams) -> Result<bool> {
|
||||
if !self.db.has_running_workflow().await? {
|
||||
return Ok(false);
|
||||
}
|
||||
|
||||
let name = params.model_name();
|
||||
if let Ok(Some((status, attempt_count))) = self.db.get_grid_point_status(&name).await {
|
||||
if status == "failed" || attempt_count >= 3 {
|
||||
info!("网格点 {} 已达到最大重试次数 ({}) 或处于 failed 状态,跳过种子热启动回退", name, attempt_count);
|
||||
return Ok(false);
|
||||
}
|
||||
}
|
||||
|
||||
let seed_match_opt = self.db.find_best_seed_from_db(params).await.ok().flatten();
|
||||
|
||||
if let Some(seed_match) = seed_match_opt {
|
||||
let timeout_sec = self.get_active_timeout_sec().await;
|
||||
let name = params.model_name();
|
||||
let task_spec = TaskSpec {
|
||||
task_id: Uuid::new_v4(),
|
||||
point_name: name.clone(),
|
||||
params: params.clone(),
|
||||
task_type: TaskType::SeedStep,
|
||||
seed_point_name: Some(seed_match.name.clone()),
|
||||
timeout_sec,
|
||||
};
|
||||
|
||||
self.db.insert_task(&task_spec).await?;
|
||||
self.db.update_grid_status(&name, common::models::GridPointStatus::Queued).await?;
|
||||
if let Err(e) = self.queue.push_task(&task_spec).await {
|
||||
let _ = self.db.update_grid_status(&name, common::models::GridPointStatus::Pending).await;
|
||||
let _ = self.queue.remove_task(&task_spec.task_id.to_string()).await;
|
||||
return Err(e);
|
||||
}
|
||||
info!("触发种子步进 (seed_step):网格点 {} 将使用 6 维近邻种子 {} 热启动重试", name, seed_match.name);
|
||||
Ok(true)
|
||||
} else {
|
||||
Ok(false)
|
||||
}
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use common::config::GridAxesConfig;
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_grid_scheduler_initialization_and_scheduling() {
|
||||
let temp_dir = tempfile::tempdir().unwrap();
|
||||
let db_path = temp_dir.path().join("sched_db.db");
|
||||
let queue_db_path = temp_dir.path().join("sched_queue.db");
|
||||
let results_dir = temp_dir.path().join("results");
|
||||
|
||||
let db = Database::new(&db_path.to_string_lossy()).await.unwrap();
|
||||
let queue = Arc::new(SqliteTaskQueue::new(&queue_db_path.to_string_lossy()).await.unwrap());
|
||||
let scheduler = GridScheduler::new(db.clone(), queue.clone(), results_dir.to_string_lossy().to_string());
|
||||
|
||||
let cfg = GridConfig {
|
||||
grid: GridAxesConfig {
|
||||
teff: vec![35000.0],
|
||||
logg: vec![5.5],
|
||||
loghe: vec![-1.0],
|
||||
logc: vec![-2.0],
|
||||
logn: vec![-2.0],
|
||||
logo: vec![-2.0],
|
||||
},
|
||||
chain: vec![],
|
||||
synspec: None,
|
||||
nworkers: 4,
|
||||
timeout_sec: 3600,
|
||||
resume: true,
|
||||
seed_step_fallback: true,
|
||||
results: None,
|
||||
itek_fallback: vec![],
|
||||
niter: Some(100),
|
||||
template: None,
|
||||
fort55: None,
|
||||
linelist: None,
|
||||
};
|
||||
|
||||
scheduler.initialize_grid(&cfg).await.unwrap();
|
||||
db.upsert_workflow("test_wf", None, "", "running").await.unwrap();
|
||||
|
||||
let pending = db.get_pending_grid_points().await.unwrap();
|
||||
assert_eq!(pending.len(), 1);
|
||||
|
||||
let dispatched = scheduler.schedule_pending_tasks().await.unwrap();
|
||||
assert_eq!(dispatched, 1);
|
||||
|
||||
let popped = queue.pop_task().await.unwrap();
|
||||
assert!(popped.is_some());
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,156 @@
|
||||
use axum::{
|
||||
body::Body,
|
||||
http::{Request, StatusCode},
|
||||
};
|
||||
use mq::sqlite_queue::SqliteTaskQueue;
|
||||
use server::{api::AppState, db::Database, scheduler::GridScheduler};
|
||||
use std::sync::Arc;
|
||||
use tower::ServiceExt; // for oneshot
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_server_api_flow() {
|
||||
let temp_dir = tempfile::tempdir().unwrap();
|
||||
let db_path = temp_dir.path().join("api_db.db");
|
||||
let queue_db_path = temp_dir.path().join("api_queue.db");
|
||||
let results_dir = temp_dir.path().join("results");
|
||||
std::fs::create_dir_all(&results_dir).unwrap();
|
||||
|
||||
let db = Database::new(&db_path.to_string_lossy()).await.unwrap();
|
||||
let queue = Arc::new(SqliteTaskQueue::new(&queue_db_path.to_string_lossy()).await.unwrap());
|
||||
let scheduler = Arc::new(GridScheduler::new(db.clone(), queue.clone(), results_dir.to_string_lossy().to_string()));
|
||||
|
||||
let state = AppState {
|
||||
db,
|
||||
queue,
|
||||
scheduler,
|
||||
results_dir: results_dir.to_string_lossy().to_string(),
|
||||
auth_token: None,
|
||||
};
|
||||
|
||||
let app = axum::Router::new()
|
||||
.route("/api/node/register", axum::routing::post(server::api::node::register_node))
|
||||
.route("/api/node/heartbeat", axum::routing::post(server::api::node::heartbeat_node))
|
||||
.route("/api/task/claim", axum::routing::post(server::api::task::claim_task))
|
||||
.route("/api/status", axum::routing::get(server::api::status::get_status))
|
||||
.route("/api/workflows", axum::routing::get(server::api::workflow::list_workflows).post(server::api::workflow::save_workflow))
|
||||
.with_state(state);
|
||||
|
||||
// 1. Check status API
|
||||
let response = app
|
||||
.clone()
|
||||
.oneshot(Request::builder().uri("/api/status").body(Body::empty()).unwrap())
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
|
||||
// 2. Register node API
|
||||
let reg_body = serde_json::json!({
|
||||
"node_id": "test-node-api",
|
||||
"host_name": "api-host",
|
||||
"max_slots": 8
|
||||
});
|
||||
let response = app
|
||||
.clone()
|
||||
.oneshot(
|
||||
Request::builder()
|
||||
.method("POST")
|
||||
.uri("/api/node/register")
|
||||
.header("content-type", "application/json")
|
||||
.body(Body::from(serde_json::to_vec(®_body).unwrap()))
|
||||
.unwrap(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
|
||||
// 3. Save Workflow API
|
||||
let wf_body = serde_json::json!({
|
||||
"name": "test_api_wf",
|
||||
"description": "Test Workflow Description",
|
||||
"config_yaml": "grid:\n teff: [35000]\n logg: [5.5]\n loghe: [-1]\n logc: [-2]\n logn: [-2]\n logo: [-2]"
|
||||
});
|
||||
let response = app
|
||||
.clone()
|
||||
.oneshot(
|
||||
Request::builder()
|
||||
.method("POST")
|
||||
.uri("/api/workflows")
|
||||
.header("content-type", "application/json")
|
||||
.body(Body::from(serde_json::to_vec(&wf_body).unwrap()))
|
||||
.unwrap(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_auth_middleware_scope_and_running_status() {
|
||||
let temp_dir = tempfile::tempdir().unwrap();
|
||||
let db_path = temp_dir.path().join("auth_db.db");
|
||||
let queue_db_path = temp_dir.path().join("auth_queue.db");
|
||||
let results_dir = temp_dir.path().join("results");
|
||||
std::fs::create_dir_all(&results_dir).unwrap();
|
||||
|
||||
let db = Database::new(&db_path.to_string_lossy()).await.unwrap();
|
||||
let queue = Arc::new(SqliteTaskQueue::new(&queue_db_path.to_string_lossy()).await.unwrap());
|
||||
let scheduler = Arc::new(GridScheduler::new(db.clone(), queue.clone(), results_dir.to_string_lossy().to_string()));
|
||||
|
||||
let state = AppState {
|
||||
db: db.clone(),
|
||||
queue: queue.clone(),
|
||||
scheduler,
|
||||
results_dir: results_dir.to_string_lossy().to_string(),
|
||||
auth_token: Some("secret_token_123".to_string()),
|
||||
};
|
||||
|
||||
let api_router = axum::Router::new()
|
||||
.route("/status", axum::routing::get(server::api::status::get_status));
|
||||
|
||||
let auth_layer = axum::middleware::from_fn_with_state(state.clone(), server::api::auth_middleware);
|
||||
let api_router = api_router.layer(auth_layer);
|
||||
|
||||
let app = axum::Router::new()
|
||||
.nest("/api", api_router)
|
||||
.with_state(state);
|
||||
|
||||
// Unauthenticated API request -> 401 Unauthorized
|
||||
let res = app
|
||||
.clone()
|
||||
.oneshot(Request::builder().uri("/api/status").body(Body::empty()).unwrap())
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(res.status(), StatusCode::UNAUTHORIZED);
|
||||
|
||||
// Authenticated API request -> 200 OK
|
||||
let res = app
|
||||
.clone()
|
||||
.oneshot(
|
||||
Request::builder()
|
||||
.uri("/api/status")
|
||||
.header("authorization", "Bearer secret_token_123")
|
||||
.body(Body::empty())
|
||||
.unwrap(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(res.status(), StatusCode::OK);
|
||||
|
||||
// Test mark_grid_point_running
|
||||
let params = common::models::GridPointParams {
|
||||
teff: 35000.0,
|
||||
logg: 5.5,
|
||||
loghe: -1.0,
|
||||
logc: -2.0,
|
||||
logn: -2.0,
|
||||
logo: -2.0,
|
||||
};
|
||||
db.upsert_grid_point(¶ms, 0).await.unwrap();
|
||||
db.mark_grid_point_running(¶ms.model_name()).await.unwrap();
|
||||
|
||||
let stats = db.get_grid_summary_stats().await.unwrap();
|
||||
assert_eq!(stats["running"], 1);
|
||||
}
|
||||
Reference in New Issue
Block a user