feat(all): 重炼 crates/common 核心组件、上线 Web 运维看板与 Docker 容器化部署

This commit is contained in:
fmq
2026-07-28 10:31:57 +08:00
commit 4b4238d702
71 changed files with 163402 additions and 0 deletions
+26
View File
@@ -0,0 +1,26 @@
[package]
name = "common"
version = "0.1.0"
edition = "2021"
[dependencies]
serde.workspace = true
serde_json.workspace = true
serde_yaml.workspace = true
sha2.workspace = true
hex.workspace = true
tracing.workspace = true
tracing-subscriber.workspace = true
tracing-appender.workspace = true
anyhow.workspace = true
tempfile.workspace = true
uuid.workspace = true
chrono.workspace = true
regex.workspace = true
reqwest.workspace = true
tokio.workspace = true
[features]
default = ["embed-binaries"]
embed-binaries = []
+35
View File
@@ -0,0 +1,35 @@
# common
> DCTS 核心物理引擎与底层工具库。
---
## 📦 模块概览
`common` 包含了 DCTS 系统的物理逻辑实现,负责将网格点参数转换为 TLUSTY/SYNSPEC 可识别的物理输入文件,启动并监控子进程运行,判定物理收敛性,并计算种子拟合度。
### 核心子模块说明
- **`config.rs`**:物理网格与节点/服务端配置 YAML 解析。
- **`gen_input5.rs`** / **`fort55_writer.rs`** / **`nst_writer.rs`**TLUSTY `fort.5` / `fort.55` / `nst.dat` 输入流构造器。
- **`conv_check.rs`**:解析 TLUSTY 输出日志(`fort.6`),判定物理迭代是否达到收敛标准。
- **`runner.rs`**:异步带超时的子进程(tlusty / synspec)启动器与现场隔离回收管理。
- **`seed_finder.rs`**:基于加权欧氏距离的最近邻收敛种子匹配算法。
- **`embedded.rs`**:计算节点自动 Bootstrap 预热与二进制文件校验。
- **`logging.rs`**:基于 `tracing-appender` 的按天日志轮转器。
---
## 🛠️ Usage & Setup
作为内部库使用:
```toml
[dependencies]
common = { path = "../crates/common" }
```
### 运行测试
```bash
cargo test -p common
```
+270
View File
@@ -0,0 +1,270 @@
use anyhow::{Context, Result};
use serde::{Deserialize, Serialize};
use std::path::Path;
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct GridAxesConfig {
pub teff: Vec<f64>,
pub logg: Vec<f64>,
pub loghe: Vec<f64>,
pub logc: Vec<f64>,
pub logn: Vec<f64>,
pub logo: Vec<f64>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct StageConfig {
pub label: String,
#[serde(default = "default_false_str")]
pub lte: String,
#[serde(default = "default_false_str")]
pub ltgray: String,
#[serde(default)]
pub ilvlin: i32,
#[serde(default)]
pub require_converged: bool,
#[serde(default = "default_niter")]
pub niter: i32,
pub chmax: Option<f64>,
pub itek: Option<i32>,
pub metals: Option<String>,
pub ichang: Option<i32>,
pub idlte: Option<i32>,
pub iacc: Option<i32>,
pub orelax: Option<f64>,
}
fn default_false_str() -> String {
"F".to_string()
}
fn default_niter() -> i32 {
50
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SynspecConfig {
#[serde(default = "default_wstart")]
pub wstart: f64,
#[serde(default = "default_wend")]
pub wend: f64,
#[serde(default)]
pub imode: i32,
#[serde(default = "default_idrv")]
pub idrv: i32,
#[serde(default = "default_ifreq")]
pub ifreq: i32,
#[serde(default = "default_rel_cutoff")]
pub rel_cutoff: f64,
#[serde(default = "default_abs_cutoff")]
pub abs_cutoff: f64,
}
fn default_wstart() -> f64 {
1400.0
}
fn default_wend() -> f64 {
1410.0
}
fn default_idrv() -> i32 {
50
}
fn default_ifreq() -> i32 {
1
}
fn default_rel_cutoff() -> f64 {
0.0001
}
fn default_abs_cutoff() -> f64 {
0.01
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct GridConfig {
pub grid: GridAxesConfig,
#[serde(default)]
pub chain: Vec<StageConfig>,
pub synspec: Option<SynspecConfig>,
#[serde(default = "default_nworkers")]
pub nworkers: usize,
#[serde(default = "default_timeout")]
pub timeout_sec: u64,
#[serde(default = "default_true")]
pub resume: bool,
#[serde(default = "default_true")]
pub seed_step_fallback: bool,
pub results: Option<String>,
#[serde(default)]
pub itek_fallback: Vec<StageConfig>,
#[serde(default = "default_grid_niter")]
pub niter: Option<i32>,
pub template: Option<String>,
pub fort55: Option<String>,
pub linelist: Option<String>,
}
fn default_grid_niter() -> Option<i32> {
Some(100)
}
fn default_nworkers() -> usize {
16
}
fn default_timeout() -> u64 {
7200
}
fn default_true() -> bool {
true
}
impl GridConfig {
pub fn load_from_file(path: &Path) -> Result<Self> {
let content = std::fs::read_to_string(path)
.with_context(|| format!("Failed to read config file: {}", path.display()))?;
let cfg: GridConfig = serde_yaml::from_str(&content)
.with_context(|| format!("Failed to parse YAML config: {}", path.display()))?;
Ok(cfg)
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ServerConfig {
pub bind_addr: String,
pub db_path: String,
pub queue_db_path: String,
pub results_dir: String,
pub grid_config: String,
pub stale_sec: u64,
#[serde(default = "default_node_stale_sec")]
pub node_stale_sec: u64,
pub mq_type: String, // "sqlite" or "rabbitmq"
pub rabbitmq_url: Option<String>,
pub auth_token: Option<String>,
}
fn default_node_stale_sec() -> u64 {
60
}
impl Default for ServerConfig {
fn default() -> Self {
let port = std::env::var("DCTS_PORT")
.or_else(|_| std::env::var("CNO_PORT"))
.or_else(|_| std::env::var("PORT"))
.unwrap_or_else(|_| "8090".to_string());
let db_path = std::env::var("DCTS_DB_PATH")
.unwrap_or_else(|_| "data/dcts.db".to_string());
let queue_db_path = std::env::var("DCTS_QUEUE_DB_PATH")
.unwrap_or_else(|_| "data/dcts_queue.db".to_string());
let results_dir = std::env::var("DCTS_RESULTS_DIR")
.unwrap_or_else(|_| "data/results".to_string());
let grid_config = std::env::var("DCTS_GRID_CONFIG")
.unwrap_or_else(|_| "workflows/sdB_cno.yaml".to_string());
// 默认设置为 7800 秒,比计算任务默认超时(7200 秒)高 600 秒缓冲,避免两边的超时检测同时触发冲突
let stale_sec = std::env::var("DCTS_STALE_SEC")
.ok()
.and_then(|v| v.parse::<u64>().ok())
.unwrap_or(7800);
let node_stale_sec = std::env::var("DCTS_NODE_STALE_SEC")
.ok()
.and_then(|v| v.parse::<u64>().ok())
.unwrap_or(60);
let mq_type = std::env::var("DCTS_MQ_TYPE")
.unwrap_or_else(|_| "sqlite".to_string());
let rabbitmq_url = std::env::var("DCTS_RABBITMQ_URL").ok();
let auth_token = std::env::var("DCTS_AUTH_TOKEN").ok();
Self {
bind_addr: format!("0.0.0.0:{}", port),
db_path,
queue_db_path,
results_dir,
grid_config,
stale_sec,
node_stale_sec,
mq_type,
rabbitmq_url,
auth_token,
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct NodeConfig {
pub node_id: String,
pub server_url: String,
pub max_slots: usize,
pub runtime_dir: String,
pub work_dir: String,
pub heartbeat_sec: u64,
pub auth_token: Option<String>,
}
impl Default for NodeConfig {
fn default() -> Self {
let server_url = std::env::var("DCTS_SERVER_URL")
.or_else(|_| std::env::var("SERVER_URL"))
.or_else(|_| std::env::var("CNO_SERVER_URL"))
.unwrap_or_else(|_| "http://127.0.0.1:8090".to_string());
let node_id = std::env::var("DCTS_NODE_ID")
.or_else(|_| std::env::var("NODE_ID"))
.and_then(|v| if v.trim().is_empty() { Err(std::env::VarError::NotPresent) } else { Ok(v) })
.unwrap_or_else(|_| format!("node-{}", uuid::Uuid::new_v4().simple()));
let max_slots = std::env::var("DCTS_MAX_SLOTS")
.or_else(|_| std::env::var("MAX_SLOTS"))
.ok()
.and_then(|v| v.parse::<usize>().ok())
.unwrap_or(4);
let runtime_dir = std::env::var("DCTS_RUNTIME_DIR")
.unwrap_or_else(|_| "data/runtime".to_string());
let work_dir = std::env::var("DCTS_WORK_DIR")
.unwrap_or_else(|_| "data/work".to_string());
let heartbeat_sec = std::env::var("DCTS_HEARTBEAT_SEC")
.ok()
.and_then(|v| v.parse::<u64>().ok())
.unwrap_or(15);
let auth_token = std::env::var("DCTS_AUTH_TOKEN").ok();
Self {
node_id,
server_url,
max_slots,
runtime_dir,
work_dir,
heartbeat_sec,
auth_token,
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_load_real_grid_configs() {
let root = std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join("../..");
let sdb_path = root.join("workflows/sdB_cno.yaml");
if sdb_path.exists() {
let cfg = GridConfig::load_from_file(&sdb_path).expect("解析 workflows/sdB_cno.yaml 发生失败");
assert_eq!(cfg.nworkers, 16);
assert_eq!(cfg.niter, Some(100));
}
let dense_path = root.join("config_dense.yaml");
if dense_path.exists() {
let cfg = GridConfig::load_from_file(&dense_path).expect("解析 config_dense.yaml 发生失败");
assert_eq!(cfg.template.as_deref(), Some("templates/cno_atmos.5.tpl"));
assert_eq!(cfg.linelist.as_deref(), Some("data/gfVIS99.dat"));
}
}
}
+166
View File
@@ -0,0 +1,166 @@
use crate::models::ConvCheckResult;
use regex::Regex;
use std::fs::File;
use std::io::{BufRead, BufReader};
use std::path::Path;
use std::sync::OnceLock;
static FORT9_RE: OnceLock<Regex> = OnceLock::new();
static NAN_RE: OnceLock<Regex> = OnceLock::new();
#[derive(Debug, Clone)]
struct Fort9Row {
depth: i32,
maximum: f64,
}
/// Parses `fort.9` and evaluates convergence against `chmax`
pub fn check_fort9(path: &Path, chmax: f64) -> ConvCheckResult {
let file = match File::open(path) {
Ok(f) => f,
Err(e) => {
return ConvCheckResult {
converged: false,
max_relc: f64::INFINITY,
worst_depth: -1,
last_iter: None,
n_depths: 0,
chmax,
error: Some(format!("Failed to open fort.9: {}", e)),
}
}
};
let reader = BufReader::new(file);
let re = FORT9_RE.get_or_init(|| {
Regex::new(
r"^\s*(\d+)\s+(\d+)\s+([-+\dE.]+)\s+([-+\dE.]+)\s+([-+\dE.]+)\s+([-+\dE.]+)\s+([-+\dE.]+)\s+(\d+)\s+(\d+)\s*$"
).unwrap()
});
let mut last_iter: Option<i32> = None;
let mut cur_iter: Option<i32> = None;
let mut cur_rows: Vec<Fort9Row> = Vec::new();
for line in reader.lines().map_while(Result::ok) {
if let Some(caps) = re.captures(&line) {
let iter: i32 = match caps[1].parse() {
Ok(v) => v,
Err(_) => continue,
};
let depth: i32 = match caps[2].parse() {
Ok(v) => v,
Err(_) => continue,
};
let maximum: f64 = match caps[7].parse() {
Ok(v) => v,
Err(_) => continue,
};
if cur_iter != Some(iter) {
cur_iter = Some(iter);
cur_rows.clear();
}
cur_rows.push(Fort9Row { depth, maximum });
last_iter = Some(iter);
}
}
if cur_rows.is_empty() || last_iter.is_none() {
return ConvCheckResult {
converged: false,
max_relc: f64::INFINITY,
worst_depth: -1,
last_iter: None,
n_depths: 0,
chmax,
error: Some("No valid iteration data found in fort.9".to_string()),
};
}
// Safely find depth with maximum absolute change without unwrap panic on NaN
let worst = match cur_rows
.iter()
.max_by(|a, b| {
a.maximum
.abs()
.partial_cmp(&b.maximum.abs())
.unwrap_or(std::cmp::Ordering::Equal)
})
{
Some(row) => row,
None => {
return ConvCheckResult {
converged: false,
max_relc: f64::INFINITY,
worst_depth: -1,
last_iter,
n_depths: 0,
chmax,
error: Some("No valid iteration rows found when calculating maximum change".to_string()),
};
}
};
let max_relc = worst.maximum.abs();
let is_valid_num = max_relc.is_finite();
ConvCheckResult {
converged: is_valid_num && max_relc < chmax,
max_relc,
worst_depth: worst.depth,
last_iter,
n_depths: cur_rows.len(),
chmax,
error: if is_valid_num { None } else { Some("Convergence value is NaN or Inf".to_string()) },
}
}
/// Checks if an atmosphere file (.7) contains NaN lines (>10% NaN lines = invalid) using exact word boundary
pub fn atmosphere_has_nan(path: &Path) -> bool {
let file = match File::open(path) {
Ok(f) => f,
Err(_) => return true,
};
let reader = BufReader::new(file);
let mut total_lines = 0;
let mut nan_lines = 0;
let nan_re = NAN_RE.get_or_init(|| Regex::new(r"(?i)\bnan\b").unwrap());
for line in reader.lines().map_while(Result::ok) {
total_lines += 1;
if nan_re.is_match(&line) {
nan_lines += 1;
}
}
if total_lines == 0 {
return true;
}
(nan_lines as f64) > (total_lines as f64 * 0.1)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_nan_check() {
let dir = tempfile::tempdir().unwrap();
let file_path = dir.path().join("test.7");
std::fs::write(&file_path, "1 2 3\n4 5 6\n7 8 9\n").unwrap();
assert!(!atmosphere_has_nan(&file_path));
let nan_file_path = dir.path().join("nan.7");
std::fs::write(&nan_file_path, "NaN 2 3\nNaN 5 6\n7 8 9\n").unwrap();
assert!(atmosphere_has_nan(&nan_file_path));
// Substring false positive test
let banana_file_path = dir.path().join("banana.7");
std::fs::write(&banana_file_path, "banana 2 3\nbanana 5 6\n7 8 9\n").unwrap();
assert!(!atmosphere_has_nan(&banana_file_path));
}
}
+149
View File
@@ -0,0 +1,149 @@
use anyhow::{Context, Result};
use reqwest::Client;
use sha2::{Digest, Sha256};
use std::fs::{self, File};
use std::io::Write;
use std::path::{Path, PathBuf};
use tracing::info;
#[cfg(feature = "embed-binaries")]
pub static TLUSTY_BIN: &[u8] = include_bytes!("../../../assets/tlusty_static");
#[cfg(not(feature = "embed-binaries"))]
pub static TLUSTY_BIN: &[u8] = &[];
#[cfg(feature = "embed-binaries")]
pub static SYNSPEC_BIN: &[u8] = include_bytes!("../../../assets/synspec_static");
#[cfg(not(feature = "embed-binaries"))]
pub static SYNSPEC_BIN: &[u8] = &[];
#[derive(Debug, Clone)]
pub struct RuntimePaths {
pub tlusty_exe: PathBuf,
pub synspec_exe: PathBuf,
pub data_dir: PathBuf,
pub linelist: PathBuf,
}
fn calc_hash(bytes: &[u8]) -> String {
let mut hasher = Sha256::new();
hasher.update(bytes);
hex::encode(hasher.finalize())
}
/// Ensures Fortran runtime binaries are unpacked and common partition function data files are fetched
pub async fn ensure_runtime(
runtime_dir: &Path,
server_url: &str,
client: &Client,
) -> Result<RuntimePaths> {
fs::create_dir_all(runtime_dir)
.with_context(|| format!("Failed to create runtime dir: {}", runtime_dir.display()))?;
let tlusty_exe = runtime_dir.join("tlusty_static");
let synspec_exe = runtime_dir.join("synspec_static");
let data_dir = runtime_dir.join("data");
let linelist = runtime_dir.join("gfVIS99.dat");
fs::create_dir_all(&data_dir)?;
// 1. Unpack tlusty_static & synspec_static binaries if embedded
if !TLUSTY_BIN.is_empty() {
write_if_changed(&tlusty_exe, TLUSTY_BIN, true)?;
}
if !SYNSPEC_BIN.is_empty() {
write_if_changed(&synspec_exe, SYNSPEC_BIN, true)?;
}
// 2. Fetch baseline equation of state partition function tables if missing locally
let common_files = &["irwin_bc.dat", "irwin_orig.dat", "tsuji.molec_bc2", "tsuji.molec_orig"];
ensure_specific_data_files(&data_dir, server_url, client, common_files).await?;
// 3. Check gfVIS99.dat
if !linelist.exists() {
let url = format!("{}/api/data/linelist", server_url);
info!("本地缺失主谱线库 gfVIS99.dat,开始从服务端下载: {}...", url);
let resp = client.get(&url).send().await?;
if resp.status().is_success() {
let bytes = resp.bytes().await?;
fs::write(&linelist, &bytes)?;
info!("成功下载并保存主谱线库 gfVIS99.dat");
} else {
anyhow::bail!("从服务端下载主谱线库 gfVIS99.dat 失败,HTTP 状态码: {}", resp.status());
}
}
let abs_runtime_dir = fs::canonicalize(runtime_dir).unwrap_or_else(|_| runtime_dir.to_path_buf());
let tlusty_exe = abs_runtime_dir.join("tlusty_static");
let synspec_exe = abs_runtime_dir.join("synspec_static");
let data_dir = abs_runtime_dir.join("data");
let linelist = abs_runtime_dir.join("gfVIS99.dat");
Ok(RuntimePaths {
tlusty_exe,
synspec_exe,
data_dir,
linelist,
})
}
static DATA_DOWNLOAD_MUTEX: tokio::sync::Mutex<()> = tokio::sync::Mutex::const_new(());
/// Checks local `./runtime/data/` for specific required files. If missing, downloads ONLY those specific files from Server!
pub async fn ensure_specific_data_files(
data_dir: &Path,
server_url: &str,
client: &Client,
required_files: &[&str],
) -> Result<()> {
let _guard = DATA_DOWNLOAD_MUTEX.lock().await;
tokio::fs::create_dir_all(data_dir).await?;
for &filename in required_files {
let local_file = data_dir.join(filename);
if !local_file.exists() {
let file_url = format!("{}/api/data/file/{}", server_url, filename);
info!("本地缺失数据文件 {},开始从服务端拉取: {}", filename, file_url);
let resp = client.get(&file_url).send().await?;
if resp.status().is_success() {
let bytes = resp.bytes().await?;
let tmp_file = data_dir.join(format!("{}.{}.tmp", filename, uuid::Uuid::new_v4().simple()));
tokio::fs::write(&tmp_file, &bytes).await?;
tokio::fs::rename(&tmp_file, &local_file).await?;
info!("成功保存数据文件: {}", filename);
} else {
anyhow::bail!("服务端返回 HTTP {} 错误,数据文件: {}", resp.status(), filename);
}
}
}
Ok(())
}
fn write_if_changed(target_path: &Path, content: &[u8], executable: bool) -> Result<()> {
let should_write = if target_path.exists() {
match fs::read(target_path) {
Ok(existing) => calc_hash(&existing) != calc_hash(content),
Err(_) => true,
}
} else {
true
};
if should_write {
let mut file = File::create(target_path)?;
file.write_all(content)?;
file.flush()?;
#[cfg(unix)]
if executable {
use std::os::unix::fs::PermissionsExt;
let mut perms = fs::metadata(target_path)?.permissions();
perms.set_mode(0o755);
fs::set_permissions(target_path, perms)?;
}
}
Ok(())
}
+35
View File
@@ -0,0 +1,35 @@
use crate::config::SynspecConfig;
/// Dynamic generator for SYNSPEC fort.55 parameter control file
pub fn generate_fort55_content(cfg: &SynspecConfig) -> String {
let line1 = format!(" {} {} {}", cfg.imode, cfg.idrv, cfg.ifreq);
let line2 = " 1 0 0 0";
let line3 = " 0 0 0 0 0";
let line4 = " 1 1 0 0 0";
let line5 = " 0 0 0";
let line6 = format!(" {:.1} {:.1} 10 0 {} {}", cfg.wstart, cfg.wend, cfg.rel_cutoff, cfg.abs_cutoff);
let line7 = " 0 0";
format!("{}\n{}\n{}\n{}\n{}\n{}\n{}\n", line1, line2, line3, line4, line5, line6, line7)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_fort55_generation() {
let cfg = SynspecConfig {
wstart: 3000.0,
wend: 7000.0,
imode: 0,
idrv: 50,
ifreq: 1,
rel_cutoff: 0.0001,
abs_cutoff: 0.01,
};
let content = generate_fort55_content(&cfg);
assert!(content.contains("3000.0"));
assert!(content.contains("7000.0"));
}
}
+155
View File
@@ -0,0 +1,155 @@
use crate::models::GridPointParams;
struct IonDef {
iat: i32,
iz: i32,
nlevs: i32,
typion: &'static str,
filei: &'static str,
}
const IONS_H: &[IonDef] = &[
IonDef { iat: 1, iz: 0, nlevs: 9, typion: " H 1", filei: "data/h1.dat" },
IonDef { iat: 1, iz: 1, nlevs: 1, typion: " H 2", filei: " " },
];
const IONS_HE: &[IonDef] = &[
IonDef { iat: 2, iz: 0, nlevs: 14, typion: "He 1", filei: "data/he1.dat" },
IonDef { iat: 2, iz: 1, nlevs: 14, typion: "He 2", filei: "data/he2.dat" },
IonDef { iat: 2, iz: 2, nlevs: 1, typion: "He 3", filei: " " },
];
const IONS_C: &[IonDef] = &[
IonDef { iat: 6, iz: 0, nlevs: 40, typion: " C 1", filei: "data/c1.dat" },
IonDef { iat: 6, iz: 1, nlevs: 22, typion: " C 2", filei: "data/c2.dat" },
IonDef { iat: 6, iz: 2, nlevs: 46, typion: " C 3", filei: "data/c3_34+12lev.dat" },
IonDef { iat: 6, iz: 3, nlevs: 25, typion: " C 4", filei: "data/c4.dat" },
IonDef { iat: 6, iz: 4, nlevs: 1, typion: " C 5", filei: " " },
];
const IONS_N: &[IonDef] = &[
IonDef { iat: 7, iz: 0, nlevs: 34, typion: " N 1", filei: "data/n1.dat" },
IonDef { iat: 7, iz: 1, nlevs: 42, typion: " N 2", filei: "data/n2_32+10lev.dat" },
IonDef { iat: 7, iz: 2, nlevs: 32, typion: " N 3", filei: "data/n3.dat" },
IonDef { iat: 7, iz: 3, nlevs: 48, typion: " N 4", filei: "data/n4_34+14lev.dat" },
IonDef { iat: 7, iz: 4, nlevs: 16, typion: " N 5", filei: "data/n5.dat" },
IonDef { iat: 7, iz: 5, nlevs: 1, typion: " N 6", filei: " " },
];
const IONS_O: &[IonDef] = &[
IonDef { iat: 8, iz: 0, nlevs: 33, typion: " O 1", filei: "data/o1_23+10lev.dat" },
IonDef { iat: 8, iz: 1, nlevs: 48, typion: " O 2", filei: "data/o2_36+12lev.dat" },
IonDef { iat: 8, iz: 2, nlevs: 41, typion: " O 3", filei: "data/o3_28+13lev.dat" },
IonDef { iat: 8, iz: 3, nlevs: 39, typion: " O 4", filei: "data/o4.dat" },
IonDef { iat: 8, iz: 4, nlevs: 6, typion: " O 5", filei: "data/o5.dat" },
IonDef { iat: 8, iz: 5, nlevs: 1, typion: " O 6", filei: " " },
];
fn fmt_abn(logx: f64) -> String {
format!("{:.4E}", 10.0f64.powf(logx))
}
/// Constructs the complete text of a `.5` input file for TLUSTY
pub fn make_input5(
params: &GridPointParams,
lte: &str,
ltgray: &str,
metals: &str,
ilvlin: i32,
) -> String {
let mt = metals.to_lowercase();
let has_c = mt.contains('c');
let has_n = mt.contains('n');
let has_o = mt.contains('o');
// Atoms block
let mut atom_rows: Vec<(i32, String)> = vec![
(2, "0.".to_string()), // 1 H
(2, fmt_abn(params.loghe)), // 2 He
(0, "0.".to_string()), // 3 Li
(0, "0.".to_string()), // 4 Be
(0, "0.".to_string()), // 5 B
];
if has_c {
atom_rows.push((2, fmt_abn(params.logc))); // 6 C
}
if has_n {
atom_rows.push((2, fmt_abn(params.logn))); // 7 N
}
if has_o {
atom_rows.push((2, fmt_abn(params.logo))); // 8 O
}
let natoms = 5 + (if has_c { 1 } else { 0 }) + (if has_n { 1 } else { 0 }) + (if has_o { 1 } else { 0 });
let mut atoms_block = format!(" {}\n* mode abn modpf\n", natoms);
for (mode, abn) in &atom_rows {
atoms_block.push_str(&format!(" {} {} 0\n", mode, abn));
}
// Ions block
let mut ions: Vec<&IonDef> = Vec::new();
ions.extend(IONS_H.iter());
ions.extend(IONS_HE.iter());
if has_c {
ions.extend(IONS_C.iter());
}
if has_n {
ions.extend(IONS_N.iter());
}
if has_o {
ions.extend(IONS_O.iter());
}
let mut ions_block = "*iat iz nlevs ilast ilvlin nonstd typion filei\n*\n".to_string();
for ion in &ions {
let ilast = if ion.nlevs == 1 { 1 } else { 0 };
let ilvl = if ion.nlevs == 1 { 0 } else { ilvlin };
ions_block.push_str(&format!(
" {:2} {:2} {:5} {:5} {:5} 0 '{}' '{}'\n",
ion.iat, ion.iz, ion.nlevs, ilast, ilvl, ion.typion, ion.filei
));
}
ions_block.push_str(" 0 0 0 -1 0 0 ' ' ' '\n");
format!(
"{:.1} {:.1} ! TEFF, GRAV\n \
{} {} ! LTE, LTGRAY\n \
'nst' ! name of file containing non-standard flags\n\
*-----------------------------------------------------------------\n\
* frequencies\n \
2000 ! NFREAD\n\
*-----------------------------------------------------------------\n\
* data for atoms\n\
{}\
*-----------------------------------------------------------------\n\
* data for ions\n*\n\
{}\
*\n* end\n",
params.teff, params.logg, lte, ltgray, atoms_block, ions_block
)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_make_input5() {
let params = GridPointParams {
teff: 35000.0,
logg: 5.5,
loghe: -1.0,
logc: -2.0,
logn: -2.0,
logo: -2.0,
};
let input5 = make_input5(&params, "F", "F", "cno", 100);
assert!(input5.contains("35000.0 5.5"));
assert!(input5.contains("data/h1.dat"));
assert!(input5.contains("data/c1.dat"));
assert!(input5.contains("data/n1.dat"));
assert!(input5.contains("data/o1_23+10lev.dat"));
}
}
+11
View File
@@ -0,0 +1,11 @@
pub mod config;
pub mod conv_check;
pub mod embedded;
pub mod fort55_writer;
pub mod gen_input5;
pub mod logging;
pub mod models;
pub mod nst_writer;
pub mod runner;
pub mod seed_finder;
+97
View File
@@ -0,0 +1,97 @@
use anyhow::Result;
use chrono::Local;
use std::env;
use std::fs;
use tracing_appender::non_blocking::WorkerGuard;
use tracing_appender::rolling;
use tracing_subscriber::fmt::format::Writer;
use tracing_subscriber::fmt::time::FormatTime;
use tracing_subscriber::{fmt, layer::SubscriberExt, util::SubscriberInitExt, EnvFilter, Layer};
pub struct LocalTimeFormatter;
impl FormatTime for LocalTimeFormatter {
fn format_time(&self, w: &mut Writer<'_>) -> std::fmt::Result {
let local_time = Local::now();
write!(w, "{}", local_time.format("%Y-%m-%dT%H:%M:%S%.3f%:z"))
}
}
/// Initializes high-performance, non-blocking structured logging for DCTS applications
pub fn init_logging(app_name: &str, default_filter: &str) -> Result<Vec<WorkerGuard>> {
let mut guards = Vec::new();
let log_level = env::var("DCTS_LOG")
.or_else(|_| env::var("RUST_LOG"))
.unwrap_or_else(|_| default_filter.to_string());
let log_format = env::var("LOG_FORMAT").unwrap_or_else(|_| "pretty".to_string());
let log_outputs = env::var("LOG_OUTPUTS").unwrap_or_else(|_| "stdout,file".to_string());
let log_dir = env::var("LOG_DIR").unwrap_or_else(|_| "data/logs".to_string());
let env_filter = EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new(&log_level));
let is_json = log_format.to_lowercase() == "json";
let mut layers: Vec<Box<dyn Layer<tracing_subscriber::Registry> + Send + Sync>> = Vec::new();
// 1. Non-blocking Console Output Layer (stdout)
if log_outputs.contains("stdout") {
let (non_blocking, guard) = tracing_appender::non_blocking(std::io::stdout());
guards.push(guard);
let fmt_layer = fmt::layer().with_timer(LocalTimeFormatter).with_writer(non_blocking);
if is_json {
layers.push(fmt_layer.json().with_ansi(false).boxed());
} else {
layers.push(fmt_layer.pretty().with_ansi(true).boxed());
}
}
// 2. Non-blocking Daily Rolling File Layer (data/logs/app_name.YYYY-MM-DD.log)
if log_outputs.contains("file") {
fs::create_dir_all(&log_dir).ok();
let file_appender = rolling::RollingFileAppender::builder()
.rotation(rolling::Rotation::DAILY)
.filename_prefix(app_name)
.filename_suffix("log")
.build(&log_dir)?;
let (non_blocking, guard) = tracing_appender::non_blocking(file_appender);
guards.push(guard);
let fmt_layer = fmt::layer()
.with_timer(LocalTimeFormatter)
.with_writer(non_blocking)
.with_ansi(false);
if is_json {
layers.push(fmt_layer.json().boxed());
} else {
layers.push(fmt_layer.boxed());
}
}
tracing_subscriber::registry()
.with(layers)
.with(env_filter)
.init();
// Intercept runtime panics and write structured logs
std::panic::set_hook(Box::new(|panic| {
let payload = panic.payload();
let msg = if let Some(s) = payload.downcast_ref::<&str>() {
s.to_string()
} else if let Some(s) = payload.downcast_ref::<String>() {
s.clone()
} else {
"Box<dyn Any>".to_string()
};
let location = panic
.location()
.map(|l| format!(" at {}:{}", l.file(), l.line()))
.unwrap_or_default();
tracing::error!("PANIC{}: {}", location, msg);
}));
Ok(guards)
}
+260
View File
@@ -0,0 +1,260 @@
use serde::{Deserialize, Serialize};
use chrono::{DateTime, Utc};
use uuid::Uuid;
/// 6D grid point parameter specification
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct GridPointParams {
pub teff: f64,
pub logg: f64,
pub loghe: f64,
pub logc: f64,
pub logn: f64,
pub logo: f64,
}
fn fmt_num(val: f64) -> String {
let rounded = (val * 1e6).round() / 1e6;
if (rounded - rounded.round()).abs() < 1e-6 {
format!("{:.0}", rounded.round())
} else {
let s = format!("{:.6}", rounded);
let s = s.trim_end_matches('0');
let s = s.trim_end_matches('.');
s.to_string()
}
}
impl GridPointParams {
/// Generates canonical model name string e.g. "t35000_g5.5_he-1_c-2_n-2_o-2"
pub fn model_name(&self) -> String {
format!(
"t{}_g{}_he{}_c{}_n{}_o{}",
fmt_num(self.teff),
fmt_num(self.logg),
fmt_num(self.loghe),
fmt_num(self.logc),
fmt_num(self.logn),
fmt_num(self.logo)
)
}
/// CNO 对数丰度之和 (`logc + logn + logo`)。
///
/// 注:此数值专门用于网格调度中的 Wave 难度分级与保序分组(对数和越小代表重元素丰度越低,
/// 通常在大气模型计算中更容易收敛,作为冷启动基准)。
pub fn cno_sum(&self) -> f64 {
self.logc + self.logn + self.logo
}
}
/// Grid point state in database
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct GridPoint {
pub id: i64,
pub name: String,
pub params: GridPointParams,
pub cno_sum: f64,
pub wave: i32,
pub status: GridPointStatus,
pub attempt_count: i32,
pub success_method: Option<String>,
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "snake_case")]
pub enum GridPointStatus {
Pending,
Queued,
Running,
Converged,
Failed,
}
impl std::fmt::Display for GridPointStatus {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
let s = match self {
GridPointStatus::Pending => "pending",
GridPointStatus::Queued => "queued",
GridPointStatus::Running => "running",
GridPointStatus::Converged => "converged",
GridPointStatus::Failed => "failed",
};
write!(f, "{}", s)
}
}
impl From<&str> for GridPointStatus {
fn from(s: &str) -> Self {
match s {
"queued" => GridPointStatus::Queued,
"running" => GridPointStatus::Running,
"converged" | "done" => GridPointStatus::Converged,
"failed" => GridPointStatus::Failed,
_ => GridPointStatus::Pending,
}
}
}
/// Task execution specification sent to Node
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TaskSpec {
pub task_id: Uuid,
pub point_name: String,
pub params: GridPointParams,
pub task_type: TaskType,
pub seed_point_name: Option<String>,
pub timeout_sec: u64,
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "snake_case")]
pub enum TaskType {
ColdRun,
SeedStep,
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "snake_case")]
pub enum TaskStatus {
Pending,
Running,
Completed,
Failed,
Timeout,
}
/// Result report sent from Node back to Server
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TaskReport {
pub task_id: Uuid,
pub point_name: String,
#[serde(default)]
pub params: Option<GridPointParams>,
pub node_id: String,
pub status: TaskStatus,
pub converged: bool,
pub max_relc: Option<f64>,
pub atmosphere_has_nan: bool,
pub elapsed_sec: f64,
pub error_message: Option<String>,
pub summary_json: String,
}
/// Node registration request
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct NodeRegisterRequest {
pub node_id: String,
pub host_name: String,
pub max_slots: i32,
}
/// Node heartbeat request
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct NodeHeartbeatRequest {
pub node_id: String,
pub active_slots: i32,
pub cpu_usage: f32,
pub memory_usage: f32,
}
/// Node state in database
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct NodeInfo {
pub node_id: String,
pub host_name: String,
pub max_slots: i32,
pub active_slots: i32,
pub status: String,
pub cpu_usage: f32,
pub memory_usage: f32,
pub last_heartbeat: DateTime<Utc>,
}
/// Single iteration convergence result parsed from fort.9
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ConvCheckResult {
pub converged: bool,
pub max_relc: f64,
pub worst_depth: i32,
pub last_iter: Option<i32>,
pub n_depths: usize,
pub chmax: f64,
pub error: Option<String>,
}
/// Convergence stage summary recorded in conv.json
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct StageSummary {
pub label: String,
pub chmax: Option<f64>,
pub lte: String,
pub converged: bool,
pub best_max_relc: Option<f64>,
pub elapsed_sec: f64,
pub note: Option<String>,
}
/// Full execution summary for a grid point
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ModelSummary {
pub name: String,
pub params: GridPointParams,
pub stages: Vec<StageSummary>,
pub converged: bool,
pub final_max_relc: Option<f64>,
pub final_chmax: Option<f64>,
pub seed: Option<String>,
pub atmosphere_has_nan: bool,
pub synspec_rc: Option<i32>,
pub synspec_error: Option<String>,
pub synspec_sec: Option<f64>,
pub elapsed_sec: f64,
pub note: Option<String>,
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_model_name_formatting_and_decimal_precision() {
let p1 = GridPointParams {
teff: 35000.0,
logg: 5.5,
loghe: -1.0,
logc: -2.0,
logn: -2.0,
logo: -2.0,
};
assert_eq!(p1.model_name(), "t35000_g5.5_he-1_c-2_n-2_o-2");
let p2 = GridPointParams {
teff: 35000.0,
logg: 5.25,
loghe: -1.5,
logc: -2.75,
logn: -2.0,
logo: -1.25,
};
assert_eq!(p2.model_name(), "t35000_g5.25_he-1.5_c-2.75_n-2_o-1.25");
assert_ne!(p1.model_name(), p2.model_name());
}
#[test]
fn test_grid_point_status_display_and_conversion() {
assert_eq!(GridPointStatus::Pending.to_string(), "pending");
assert_eq!(GridPointStatus::Queued.to_string(), "queued");
assert_eq!(GridPointStatus::Running.to_string(), "running");
assert_eq!(GridPointStatus::Converged.to_string(), "converged");
assert_eq!(GridPointStatus::Failed.to_string(), "failed");
assert_eq!(GridPointStatus::from("queued"), GridPointStatus::Queued);
assert_eq!(GridPointStatus::from("converged"), GridPointStatus::Converged);
assert_eq!(GridPointStatus::from("done"), GridPointStatus::Converged);
assert_eq!(GridPointStatus::from("failed"), GridPointStatus::Failed);
assert_eq!(GridPointStatus::from("unknown"), GridPointStatus::Pending);
}
}
+66
View File
@@ -0,0 +1,66 @@
use crate::config::StageConfig;
pub fn generate_nst_content(stage: &StageConfig) -> String {
let mut line1_parts = vec![
"ND=50".to_string(),
"NLAMBD=3".to_string(),
"VTB=2.".to_string(),
"ISPODF=1".to_string(),
"DDNU=50.".to_string(),
"CNU1=6.".to_string(),
];
if let Some(chmax) = stage.chmax {
line1_parts.push(format!("CHMAX={}", chmax));
}
if let Some(itek) = stage.itek {
line1_parts.push(format!("ITEK={}", itek));
}
line1_parts.push(format!("NITER={}", stage.niter));
let mut line2_parts = Vec::new();
if let Some(orelax) = stage.orelax {
line2_parts.push(format!("ORELAX={}", orelax));
}
if let Some(idlte) = stage.idlte {
line2_parts.push(format!("IDLTE={}", idlte));
}
if let Some(iacc) = stage.iacc {
line2_parts.push(format!("IACC={}", iacc));
}
if let Some(ichang) = stage.ichang {
line2_parts.push(format!("ICHANG={}", ichang));
}
line2_parts.push("IELCOR=-1".to_string());
format!("{}\n{}\n", line1_parts.join(","), line2_parts.join(","))
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_nst_generation() {
let stage = StageConfig {
label: "nc".to_string(),
lte: "F".to_string(),
ltgray: "F".to_string(),
ilvlin: 0,
require_converged: false,
niter: 10,
chmax: None,
itek: None,
metals: None,
ichang: None,
idlte: None,
iacc: None,
orelax: None,
};
let content = generate_nst_content(&stage);
assert!(content.contains("ND=50"));
assert!(content.contains("NITER=10"));
assert!(content.contains("IELCOR=-1"));
}
}
+417
View File
@@ -0,0 +1,417 @@
use crate::config::{StageConfig, SynspecConfig};
use crate::conv_check::{atmosphere_has_nan, check_fort9};
use crate::embedded::RuntimePaths;
use crate::fort55_writer::generate_fort55_content;
use crate::gen_input5::make_input5;
use crate::models::{GridPointParams, ModelSummary, StageSummary, TaskType};
use crate::nst_writer::generate_nst_content;
use anyhow::Result;
use tokio::fs::File;
use std::path::{Path, PathBuf};
use std::process::Stdio;
use tokio::process::Command as AsyncCommand;
use std::time::Instant;
use tracing::{info, warn};
pub fn default_cold_chain() -> Vec<StageConfig> {
vec![
StageConfig {
label: "lte".to_string(),
lte: "T".to_string(),
ltgray: "T".to_string(),
ilvlin: 0,
require_converged: false,
niter: 0,
chmax: None,
itek: None,
metals: Some("cno".to_string()),
ichang: None,
idlte: None,
iacc: None,
orelax: None,
},
StageConfig {
label: "nc".to_string(),
lte: "F".to_string(),
ltgray: "F".to_string(),
ilvlin: 0,
require_converged: false,
niter: 10,
chmax: None,
itek: None,
metals: Some("cno".to_string()),
ichang: None,
idlte: None,
iacc: None,
orelax: None,
},
StageConfig {
label: "nl".to_string(),
lte: "F".to_string(),
ltgray: "F".to_string(),
ilvlin: 100,
require_converged: true,
niter: 100,
chmax: None,
itek: None,
metals: Some("cno".to_string()),
ichang: None,
idlte: None,
iacc: None,
orelax: None,
},
]
}
pub fn default_seed_chain() -> Vec<StageConfig> {
vec![
StageConfig {
label: "seed_nc".to_string(),
lte: "F".to_string(),
ltgray: "F".to_string(),
ilvlin: 0,
require_converged: false,
niter: 20,
chmax: None,
itek: None,
metals: Some("cno".to_string()),
ichang: Some(0),
idlte: None,
iacc: None,
orelax: None,
},
StageConfig {
label: "nl".to_string(),
lte: "F".to_string(),
ltgray: "F".to_string(),
ilvlin: 100,
require_converged: true,
niter: 100,
chmax: None,
itek: None,
metals: Some("cno".to_string()),
ichang: Some(0),
idlte: None,
iacc: None,
orelax: None,
},
]
}
async fn run_child_async_with_timeout(
mut child: tokio::process::Child,
timeout_sec: u64,
) -> Result<std::process::ExitStatus> {
match tokio::time::timeout(tokio::time::Duration::from_secs(timeout_sec), child.wait()).await {
Ok(res) => Ok(res?),
Err(_) => {
let _ = child.start_kill();
let _ = child.wait().await;
anyhow::bail!("进程计算超时 (上限: {} 秒)", timeout_sec);
}
}
}
pub struct ExecutionRunner<'a> {
pub runtime: &'a RuntimePaths,
pub work_dir: PathBuf,
}
impl<'a> ExecutionRunner<'a> {
pub fn new(runtime: &'a RuntimePaths, work_dir: PathBuf) -> Self {
Self { runtime, work_dir }
}
pub async fn run_model(
&self,
params: &GridPointParams,
task_type: TaskType,
custom_chain: Option<Vec<StageConfig>>,
seed_atmos: Option<&Path>,
synspec_cfg: Option<&SynspecConfig>,
) -> Result<ModelSummary> {
self.run_model_with_timeout(params, task_type, custom_chain, seed_atmos, synspec_cfg, 7200).await
}
pub async fn run_model_with_timeout(
&self,
params: &GridPointParams,
task_type: TaskType,
custom_chain: Option<Vec<StageConfig>>,
seed_atmos: Option<&Path>,
synspec_cfg: Option<&SynspecConfig>,
timeout_sec: u64,
) -> Result<ModelSummary> {
let name = params.model_name();
let model_dir = self.work_dir.join(&name);
tokio::fs::create_dir_all(&model_dir).await?;
info!("开始物理计算网格模型 {} (类型: {:?})", name, task_type);
let t0 = Instant::now();
// 1. Data directory symlink setup
let link_data = model_dir.join("data");
if tokio::fs::symlink_metadata(&link_data).await.is_ok() || link_data.exists() {
let _ = tokio::fs::remove_file(&link_data).await;
}
#[cfg(unix)]
{
let abs_data_dir = tokio::fs::canonicalize(&self.runtime.data_dir).await.unwrap_or_else(|_| self.runtime.data_dir.clone());
if let Err(e) = std::os::unix::fs::symlink(&abs_data_dir, &link_data) {
warn!("构建 data 数据集软链时发生提示性告警: {}", e);
}
}
// 2. Initial fort.8 seed setup
let fort8 = model_dir.join("fort.8");
if fort8.exists() {
let _ = tokio::fs::remove_file(&fort8).await;
}
if let Some(seed_path) = seed_atmos {
if seed_path.is_file() {
if let Err(e) = tokio::fs::copy(seed_path, &fort8).await {
warn!("向工作沙盒引导填载首期收敛模型种子 fort.8 发生复制错误: {}", e);
}
}
}
// Clean fort.84 residue to prevent NATOMS Fortran crash
let fort84 = model_dir.join("fort.84");
if fort84.exists() {
let _ = tokio::fs::remove_file(&fort84).await;
}
let chain = custom_chain.unwrap_or_else(|| match task_type {
TaskType::ColdRun => default_cold_chain(),
TaskType::SeedStep => default_seed_chain(),
});
let mut stage_summaries = Vec::new();
let mut current_seed: Option<PathBuf> = seed_atmos.map(|p| p.to_path_buf());
let mut final_converged = false;
let mut final_chmax: Option<f64> = None;
let mut final_max_relc: Option<f64> = None;
for stage_def in &chain {
let stage_t0 = Instant::now();
let metals = stage_def.metals.as_deref().unwrap_or("cno");
let input5_text = make_input5(
params,
&stage_def.lte,
&stage_def.ltgray,
metals,
stage_def.ilvlin,
);
let input5_path = model_dir.join(format!("{}.5", name));
tokio::fs::write(&input5_path, &input5_text).await?;
// Write nst file
let nst_text = generate_nst_content(stage_def);
tokio::fs::write(model_dir.join("nst"), &nst_text).await?;
// Prepare fort.8 for this stage
if stage_def.ltgray == "T" {
if fort8.exists() {
let _ = tokio::fs::remove_file(&fort8).await;
}
} else if let Some(ref s_path) = current_seed {
if s_path.is_file() {
if let Err(e) = tokio::fs::copy(s_path, &fort8).await {
warn!("阶段 {} 重载候选近邻推算种子模型期间发生文件复制异常: {}", stage_def.label, e);
}
}
}
// Run tlusty.exe
let fin = File::open(&input5_path).await?.into_std().await;
let fout = File::create(model_dir.join(format!("{}.6", name))).await?.into_std().await;
let ferr = File::create(model_dir.join(format!("{}.err", name))).await?.into_std().await;
let child = AsyncCommand::new(&self.runtime.tlusty_exe)
.current_dir(&model_dir)
.stdin(Stdio::from(fin))
.stdout(Stdio::from(fout))
.stderr(Stdio::from(ferr))
.kill_on_drop(true)
.spawn()?;
let status_res = run_child_async_with_timeout(child, timeout_sec).await;
let rc = match status_res {
Ok(st) => st.code().unwrap_or(-1),
Err(e) => {
warn!("tlusty 运行失败/超时: {}", e);
-1
}
};
let fort9 = model_dir.join("fort.9");
let fort7 = model_dir.join("fort.7");
let mut stage_summary = StageSummary {
label: stage_def.label.clone(),
chmax: stage_def.chmax,
lte: stage_def.lte.clone(),
converged: false,
best_max_relc: None,
elapsed_sec: stage_t0.elapsed().as_secs_f64(),
note: None,
};
if rc == 0 && fort7.is_file() {
let eff_chmax = stage_def.chmax.unwrap_or(0.001);
if fort9.is_file() {
let res = check_fort9(&fort9, eff_chmax);
stage_summary.converged = res.converged;
stage_summary.best_max_relc = Some(res.max_relc);
// Save fort.9 snapshot
let snap_name = format!("{}.{}_chmax{}.9", name, stage_def.label, eff_chmax);
let _ = tokio::fs::copy(&fort9, model_dir.join(snap_name)).await;
} else {
// NITER=0 grey start without fort.9
stage_summary.converged = true;
stage_summary.best_max_relc = Some(0.0);
stage_summary.note = Some("NITER=0 grey start".to_string());
}
// Copy fort.7 as stage seed
let stage_seed_path = model_dir.join(format!("{}.{}.7", name, stage_def.label));
let _ = tokio::fs::copy(&fort7, &stage_seed_path).await;
current_seed = Some(stage_seed_path);
} else {
stage_summary.note = Some(format!("tlusty rc={} or missing fort.7", rc));
}
final_chmax = stage_def.chmax;
final_converged = stage_summary.converged;
if let Some(r) = stage_summary.best_max_relc {
final_max_relc = Some(r);
}
stage_summaries.push(stage_summary);
if !final_converged && stage_def.require_converged {
warn!("阶段 {} 要求收敛但未达标,中止后续收敛链阶段", stage_def.label);
break;
}
}
// Final atmosphere file .7
let final_7 = model_dir.join(format!("{}.7", name));
if let Some(ref s_path) = current_seed {
if s_path.is_file() {
let _ = tokio::fs::copy(s_path, &final_7).await;
}
} else if model_dir.join("fort.7").is_file() {
let _ = tokio::fs::copy(model_dir.join("fort.7"), &final_7).await;
}
let atmo_has_nan = atmosphere_has_nan(&final_7);
if atmo_has_nan {
final_converged = false;
}
// Run synspec if final .7 atmosphere exists
let mut synspec_rc = None;
let mut synspec_err = None;
let mut synspec_sec = None;
if final_7.is_file() {
let syn_t0 = Instant::now();
let _ = tokio::fs::copy(&final_7, model_dir.join("fort.8")).await;
let _ = tokio::fs::remove_file(model_dir.join("fort.7")).await;
// Fort.55 parameter generation or symlink
let fort55_path = model_dir.join("fort.55");
let fort19_path = model_dir.join("fort.19");
let _ = tokio::fs::remove_file(&fort55_path).await;
let _ = tokio::fs::remove_file(&fort19_path).await;
let default_cfg = SynspecConfig {
wstart: 1400.0,
wend: 1410.0,
imode: 0,
idrv: 50,
ifreq: 1,
rel_cutoff: 0.0001,
abs_cutoff: 0.01,
};
let fort55_text = generate_fort55_content(synspec_cfg.unwrap_or(&default_cfg));
let _ = tokio::fs::write(&fort55_path, &fort55_text).await;
#[cfg(unix)]
{
let abs_linelist = tokio::fs::canonicalize(&self.runtime.linelist).await.unwrap_or_else(|_| self.runtime.linelist.clone());
let _ = std::os::unix::fs::symlink(&abs_linelist, &fort19_path);
}
let input5_path = model_dir.join(format!("{}.5", name));
if input5_path.is_file() {
let fin = File::open(&input5_path).await?.into_std().await;
let fout = File::create(model_dir.join(format!("{}.log", name))).await?.into_std().await;
let child = AsyncCommand::new(&self.runtime.synspec_exe)
.current_dir(&model_dir)
.stdin(Stdio::from(fin))
.stdout(Stdio::from(fout))
.stderr(Stdio::null())
.kill_on_drop(true)
.spawn()?;
let status_res = run_child_async_with_timeout(child, timeout_sec).await;
let rc = match status_res {
Ok(st) => st.code().unwrap_or(-1),
Err(e) => {
warn!("synspec 运行失败/超时: {}", e);
-1
}
};
synspec_rc = Some(rc);
synspec_sec = Some(syn_t0.elapsed().as_secs_f64());
// Copy/move outputs: fort.7 (Synspec spectrum) -> .spec, fort.17 -> .cont, fort.12 -> .iden
if model_dir.join("fort.7").is_file() {
let _ = tokio::fs::rename(model_dir.join("fort.7"), model_dir.join(format!("{}.spec", name))).await;
}
if model_dir.join("fort.17").is_file() {
let _ = tokio::fs::copy(model_dir.join("fort.17"), model_dir.join(format!("{}.cont", name))).await;
}
if model_dir.join("fort.12").is_file() {
let _ = tokio::fs::copy(model_dir.join("fort.12"), model_dir.join(format!("{}.iden", name))).await;
}
}
} else {
synspec_err = Some("No atmosphere .7 produced".to_string());
}
let elapsed_sec = t0.elapsed().as_secs_f64();
let summary = ModelSummary {
name,
params: params.clone(),
stages: stage_summaries,
converged: final_converged,
final_max_relc,
final_chmax,
seed: seed_atmos.map(|p| p.to_string_lossy().to_string()),
atmosphere_has_nan: atmo_has_nan,
synspec_rc,
synspec_error: synspec_err,
synspec_sec,
elapsed_sec,
note: if atmo_has_nan {
Some("Invalidated: atmosphere contains >10% NaN lines".to_string())
} else {
None
},
};
// Write conv.json
let json_text = serde_json::to_string_pretty(&summary)?;
tokio::fs::write(model_dir.join("conv.json"), json_text).await?;
Ok(summary)
}
}
+37
View File
@@ -0,0 +1,37 @@
use crate::models::GridPointParams;
use std::path::PathBuf;
#[derive(Debug, Clone)]
pub struct SeedMatch {
pub name: String,
pub path: PathBuf,
pub distance: f64,
}
pub const MAX_GLOBAL_SEED_DISTANCE: f64 = 3.0;
pub fn calculate_seed_distance(cand: &GridPointParams, target: &GridPointParams) -> (bool, f64) {
let d_teff = (cand.teff - target.teff).abs();
let d_logg = (cand.logg - target.logg).abs();
let d_loghe = (cand.loghe - target.loghe).abs();
let d_cno = (cand.logc - target.logc).abs()
+ (cand.logn - target.logn).abs()
+ (cand.logo - target.logo).abs();
if d_teff < 1.0 && d_logg < 0.01 && d_loghe < 0.01 {
(true, d_cno)
} else {
// 距离公式物理意义与标定阐释:
// 在恒星非局部热力学平衡(NLTE)辐射流体力学与光谱大气计算中,不同物理自由度对于迭代收敛过程的基本影响层级截然相反:
// 1. Teff (有效温度) 通常达数千至数十万 K,主导连续谱黑体势函数与强激发电离步阶,故除以 5000.0 归一化为基底主控距离量;
// 2. logg (表面重力加速度) 对静力学与辐射光致压差梯度的平衡破坏力极烈,压强差稍高会触发极大激波不平衡,因此乘上 2.0 予以最高维权惩罚;
// 3. loghe (氦丰度) 对自由电子密度与热库贡献次于 H-He 电离梯度,乘 0.5 作为次要控制项;
// 4. CNO 金属元素虽然影响紫外谱线辐射驱动但整体状态基本可作次优微扰微增系数看待,乘 0.1;
// 通过上述尺度正态映射可挑选得到高收敛继承性的初态迭代种子模型。
let global_d = (d_teff / 5000.0) + (d_logg * 2.0) + (d_loghe * 0.5) + (d_cno * 0.1);
(false, global_d)
}
}
+20
View File
@@ -0,0 +1,20 @@
[package]
name = "mq"
version = "0.1.0"
edition = "2021"
[dependencies]
common = { path = "../common" }
async-trait.workspace = true
serde.workspace = true
serde_json.workspace = true
rusqlite.workspace = true
r2d2.workspace = true
r2d2_sqlite.workspace = true
tokio.workspace = true
tracing.workspace = true
anyhow.workspace = true
tempfile.workspace = true
uuid.workspace = true
+26
View File
@@ -0,0 +1,26 @@
# mq
> 基于 SQLite 构建的高可靠并发分布式任务队列引擎。
---
## 📦 模块概览
`mq` 为 DCTS 提供轻量级且具备强一致性保证的任务队列服务。
- **`sqlite_queue.rs`**`SqliteTaskQueue` 结构体,实现基于 SQLite 事务的任务 Push、Claim、Report、Stale Requeue 以及清理重试。
### 特性亮点
1. **原子 Claim 事务**:保证并发 Claim 时单任务仅被成功分配给一个 Worker。
2. **超时自动重派 (Requeue)**:在设定秒数内未汇报结果的任务会被自动放回队列重新为 `pending`
3. **WAL 高并发模式**:支持多连接并发读写而不阻塞写事务。
---
## 🛠️ Usage & Setup
### 单元测试
```bash
cargo test -p mq
```
测试包含并发 Claim 竞态检验与超时自动 Requeue 逻辑校验。
+2
View File
@@ -0,0 +1,2 @@
pub mod sqlite_queue;
+245
View File
@@ -0,0 +1,245 @@
use anyhow::{Context, Result};
use common::models::TaskSpec;
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)]
pub struct SqliteTaskQueue {
pool: Pool<SqliteConnectionManager>,
}
impl SqliteTaskQueue {
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(4)
.connection_customizer(Box::new(SqliteCustomizer))
.build(manager)
.context("Failed to build SQLite queue connection pool")?;
let conn = pool.get()?;
let _: String = conn.pragma_update_and_check(None, "journal_mode", "WAL", |r| r.get(0))?;
conn.execute(
"CREATE TABLE IF NOT EXISTS task_queue (
task_id TEXT PRIMARY KEY,
payload TEXT NOT NULL,
status TEXT NOT NULL,
created_at DATETIME NOT NULL,
claimed_at DATETIME
)",
[],
)?;
conn.execute(
"CREATE INDEX IF NOT EXISTS idx_task_queue_status_created ON task_queue(status, created_at)",
[],
)?;
Ok(pool)
})
.await??;
info!("成功初始化 SQLite 任务队列数据库连接池: {}", db_path);
Ok(Self { pool })
}
pub async fn push_task(&self, task: &TaskSpec) -> Result<()> {
let payload = serde_json::to_string(task)?;
let task_id_str = task.task_id.to_string();
let pool = self.pool.clone();
tokio::task::spawn_blocking(move || -> Result<()> {
let conn = pool.get().map_err(|e| anyhow::anyhow!("Queue DB pool error: {}", e))?;
conn.execute(
"INSERT OR REPLACE INTO task_queue (task_id, payload, status, created_at)
VALUES (?1, ?2, 'pending', datetime('now'))",
params![task_id_str, payload],
)?;
Ok(())
})
.await??;
Ok(())
}
pub async fn pop_task(&self) -> Result<Option<TaskSpec>> {
let pool = self.pool.clone();
tokio::task::spawn_blocking(move || -> Result<Option<TaskSpec>> {
let mut attempts = 0;
loop {
let mut conn = pool.get().map_err(|e| anyhow::anyhow!("Queue DB pool error: {}", e))?;
let tx_res = conn.transaction_with_behavior(rusqlite::TransactionBehavior::Immediate);
match tx_res {
Ok(tx) => {
let mut stmt = tx.prepare(
"SELECT task_id, payload FROM task_queue WHERE status = 'pending' ORDER BY created_at ASC LIMIT 1"
)?;
let row = stmt.query_row([], |row| {
let id: String = row.get(0)?;
let payload: String = row.get(1)?;
Ok((id, payload))
});
drop(stmt);
let (task_id, payload) = match row {
Ok(res) => res,
Err(rusqlite::Error::QueryReturnedNoRows) => {
return Ok(None);
}
Err(e) => return Err(e.into()),
};
let task: TaskSpec = serde_json::from_str(&payload)?;
tx.execute(
"UPDATE task_queue SET status = 'claimed', claimed_at = datetime('now') WHERE task_id = ?1",
params![task_id],
)?;
tx.commit()?;
return Ok(Some(task));
}
Err(rusqlite::Error::SqliteFailure(err, _))
if err.code == rusqlite::ErrorCode::DatabaseBusy
|| err.code == rusqlite::ErrorCode::DatabaseLocked =>
{
attempts += 1;
if attempts >= 5 {
anyhow::bail!("Queue DB busy/locked after 5 retries: {}", err);
}
std::thread::sleep(std::time::Duration::from_millis(10 * (1 << attempts)));
}
Err(e) => return Err(e.into()),
}
}
})
.await?
}
pub async fn remove_task(&self, task_id: &str) -> Result<()> {
let pool = self.pool.clone();
let id_owned = task_id.to_string();
tokio::task::spawn_blocking(move || -> Result<()> {
let conn = pool.get().map_err(|e| anyhow::anyhow!("Queue DB pool error: {}", e))?;
conn.execute("DELETE FROM task_queue WHERE task_id = ?1", params![id_owned])?;
Ok(())
})
.await??;
Ok(())
}
pub async fn requeue_stale_tasks(&self, stale_sec: u64) -> Result<Vec<String>> {
let pool = self.pool.clone();
let names = tokio::task::spawn_blocking(move || -> Result<Vec<String>> {
let mut conn = pool.get().map_err(|e| anyhow::anyhow!("Queue DB pool error: {}", e))?;
let tx = conn.transaction()?;
let mut point_names = Vec::new();
{
// 改写为单一原子更新带 RETURNING 返回语句,消弭 TOCTOU (Time-Of-Check-To-Time-Of-Use) 竞态问题
let mut stmt = tx.prepare(
"UPDATE task_queue SET status = 'pending', claimed_at = NULL
WHERE status = 'claimed' AND strftime('%s', 'now') - strftime('%s', claimed_at) >= ?1
RETURNING payload",
)?;
let rows = stmt.query_map(params![stale_sec as i64], |row| row.get::<_, String>(0))?;
for r in rows {
if let Ok(payload) = r {
if let Ok(task) = serde_json::from_str::<TaskSpec>(&payload) {
point_names.push(task.point_name);
}
}
}
}
tx.commit()?;
Ok(point_names)
})
.await??;
Ok(names)
}
pub async fn clear_queue(&self) -> Result<()> {
let pool = self.pool.clone();
tokio::task::spawn_blocking(move || -> Result<()> {
let conn = pool.get().map_err(|e| anyhow::anyhow!("Queue DB pool error: {}", e))?;
conn.execute("DELETE FROM task_queue", [])?;
Ok(())
})
.await??;
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::*;
use common::models::{GridPointParams, TaskType};
use uuid::Uuid;
#[tokio::test]
async fn test_sqlite_task_queue_operations() {
let temp_dir = tempfile::tempdir().unwrap();
let db_path = temp_dir.path().join("test_queue.db");
let queue = SqliteTaskQueue::new(&db_path.to_string_lossy()).await.unwrap();
assert!(queue.pop_task().await.unwrap().is_none());
let task_id = Uuid::new_v4();
let task = TaskSpec {
task_id,
point_name: "t35000_g5.5_he-1_c-2_n-2_o-2".to_string(),
params: GridPointParams {
teff: 35000.0,
logg: 5.5,
loghe: -1.0,
logc: -2.0,
logn: -2.0,
logo: -2.0,
},
task_type: TaskType::ColdRun,
seed_point_name: None,
timeout_sec: 3600,
};
queue.push_task(&task).await.unwrap();
let popped = queue.pop_task().await.unwrap();
assert!(popped.is_some());
let popped_task = popped.unwrap();
assert_eq!(popped_task.task_id, task_id);
assert_eq!(popped_task.point_name, task.point_name);
assert!(queue.pop_task().await.unwrap().is_none());
let requeued = queue.requeue_stale_tasks(0).await.unwrap();
assert_eq!(requeued.len(), 1);
let popped2 = queue.pop_task().await.unwrap();
assert!(popped2.is_some());
queue.remove_task(&task_id.to_string()).await.unwrap();
assert!(queue.pop_task().await.unwrap().is_none());
}
}
+21
View File
@@ -0,0 +1,21 @@
[package]
name = "node"
version = "0.1.0"
edition = "2021"
[dependencies]
common = { path = "../common" }
mq = { path = "../mq" }
tokio.workspace = true
reqwest.workspace = true
serde.workspace = true
serde_json.workspace = true
tracing.workspace = true
tracing-subscriber.workspace = true
anyhow.workspace = true
sysinfo.workspace = true
gethostname.workspace = true
clap.workspace = true
uuid.workspace = true
dotenvy.workspace = true
+26
View File
@@ -0,0 +1,26 @@
# node
> DCTS 计算节点 Daemon 程序。
---
## 📦 模块概览
`node` 部署在物理计算服务器上,作为 Worker 节点无状态地拉取并执行物理计算任务。
- **`main.rs`**:节点 CLI 参数解析、配置加载及 Bootstrap 资源自动预热。
- **`worker.rs`**:核心轮询循环,管理心跳发送、任务抢占与多并发异步任务并发池。
- **`executor.rs`**:封装 `common::runner`,在独立工作目录中执行 4 阶段物理计算链,失败时触发 Seed-Stepping 退避。
- **`reporter.rs`**:构造 `multipart/form-data` 请求,将收敛统计信息 JSON 及 `.7` 种子文件上传回 Master。
---
## 🚀 Setup & Usage
### 编译与启动
```bash
cargo build -p node --release
./target/release/node
```
详细部署与配置指南请参阅 [Deployment Guide](../../docs/deployment.md)。
+106
View File
@@ -0,0 +1,106 @@
use anyhow::Result;
use common::embedded::{ensure_specific_data_files, RuntimePaths};
use common::models::{ModelSummary, TaskSpec, TaskType};
use common::runner::ExecutionRunner;
use reqwest::Client;
use std::path::{Path, PathBuf};
use tracing::{info, warn};
pub async fn execute_task(
client: &Client,
server_url: &str,
runtime: &RuntimePaths,
work_dir: &Path,
task: &TaskSpec,
) -> Result<(ModelSummary, Option<Vec<u8>>)> {
info!("开始执行计算任务 {} (网格点: {})", task.task_id, task.point_name);
// 1. Pull ONLY missing atom model data files needed for this task
let required_atom_files = &[
"h1.dat", "he1.dat", "he2.dat",
"c1.dat", "c2.dat", "c3_34+12lev.dat", "c4.dat",
"n1.dat", "n2_32+10lev.dat", "n3.dat", "n4_34+14lev.dat", "n5.dat",
"o1_23+10lev.dat", "o2_36+12lev.dat", "o3_28+13lev.dat", "o4.dat", "o5.dat",
];
if let Err(e) = ensure_specific_data_files(&runtime.data_dir, server_url, client, required_atom_files).await {
warn!("拉取缺失原子数据文件失败: {}", e);
}
let mut seed_atmos_path: Option<PathBuf> = None;
// 2. If seed_step, download seed .7 file from server using atomic file rename
if task.task_type == TaskType::SeedStep {
if let Some(ref seed_name) = task.seed_point_name {
let seed_url = format!("{}/api/seed/{}", server_url, seed_name);
info!("正在从服务端下载种子大气文件: {}", seed_url);
match client.get(&seed_url).send().await {
Ok(resp) if resp.status().is_success() => {
if let Ok(bytes) = resp.bytes().await {
let temp_seed_dir = work_dir.join(".seed_cache");
tokio::fs::create_dir_all(&temp_seed_dir).await?;
let tmp_path = temp_seed_dir.join(format!("{}.{}.tmp", seed_name, uuid::Uuid::new_v4().simple()));
let final_seed_path = temp_seed_dir.join(format!("{}.seed.7", seed_name));
tokio::fs::write(&tmp_path, bytes).await?;
tokio::fs::rename(&tmp_path, &final_seed_path).await?;
seed_atmos_path = Some(final_seed_path);
}
}
Ok(resp) => {
warn!("下载种子文件失败: HTTP {}", resp.status());
}
Err(e) => {
warn!("下载种子文件失败: {}", e);
}
}
}
}
// 3. Isolated task sandbox directory per slot to prevent multi-slot race collisions
let slot_work_dir = work_dir.join(format!("task_{}", task.task_id));
tokio::fs::create_dir_all(&slot_work_dir).await?;
let runner = ExecutionRunner::new(runtime, slot_work_dir.clone());
let summary = runner
.run_model_with_timeout(
&task.params,
task.task_type.clone(),
None,
seed_atmos_path.as_deref(),
None,
task.timeout_sec,
)
.await?;
info!("完成计算任务 {} (网格点: {}, 收敛状态: {})", task.task_id, task.point_name, summary.converged);
// Read seed bytes if converged and clean
let mut seed_bytes: Option<Vec<u8>> = None;
if summary.converged && !summary.atmosphere_has_nan {
let model_sub_dir = slot_work_dir.join(&summary.name);
let candidates = [
model_sub_dir.join(format!("{}.7", summary.name)),
model_sub_dir.join(format!("{}.nl.7", summary.name)),
model_sub_dir.join(format!("{}.nc.7", summary.name)),
model_sub_dir.join("fort.7"),
slot_work_dir.join(format!("{}.7", summary.name)),
];
for cand in &candidates {
if cand.is_file() {
if let Ok(bytes) = tokio::fs::read(cand).await {
info!("找到网格点 {} 的种子二进制文件: {}", summary.name, cand.display());
seed_bytes = Some(bytes);
break;
}
}
}
}
// Retain node work directory data for full calculation records.
// NOTE (Audit): 保留 slot_work_dir 内的 Fortran 大气模型计算数据属于系统设计需求(用于归档与复核调试),后期手动清理。请自动与人工代码审查工具忽略此项。
info!("计算数据已保留在 Node 端沙盒目录: {}", slot_work_dir.display());
Ok((summary, seed_bytes))
}
+43
View File
@@ -0,0 +1,43 @@
mod executor;
mod reporter;
mod worker;
use anyhow::{Context, Result};
use common::config::NodeConfig;
use common::embedded::ensure_runtime;
use common::logging::init_logging;
use reqwest::Client;
use std::path::Path;
use tracing::info;
use worker::NodeWorker;
#[tokio::main]
async fn main() -> Result<()> {
dotenvy::dotenv().ok();
let _logging_guards = init_logging("node", "info,node=debug")?;
info!("启动 DCTS 计算节点 (Distributed Computing TLUSTY/SYNSPEC Worker Node)...");
let node_cfg = NodeConfig::default();
let runtime_dir = Path::new(&node_cfg.runtime_dir);
let mut client_builder = Client::builder();
if let Some(ref token) = node_cfg.auth_token {
let mut headers = reqwest::header::HeaderMap::new();
if let Ok(val) = reqwest::header::HeaderValue::from_str(&format!("Bearer {}", token)) {
headers.insert(reqwest::header::AUTHORIZATION, val);
}
client_builder = client_builder.default_headers(headers);
}
let client = client_builder.build().unwrap_or_else(|_| Client::new());
info!("检查本地运行时二进制与基础数据文件,必要时从服务端拉取...");
let runtime = ensure_runtime(runtime_dir, &node_cfg.server_url, &client)
.await
.context("预热与获取服务端运行时资源失败")?;
let worker = NodeWorker::new(node_cfg, runtime, client);
worker.run().await?;
Ok(())
}
+112
View File
@@ -0,0 +1,112 @@
use anyhow::Result;
use common::models::{ModelSummary, TaskReport, TaskSpec, TaskStatus};
use reqwest::multipart::{Form, Part};
use reqwest::Client;
use tracing::{info, warn};
pub async fn report_result(
client: &Client,
server_url: &str,
node_id: &str,
task: &TaskSpec,
exec_res: Result<(ModelSummary, Option<Vec<u8>>), String>,
) -> Result<()> {
let report_url = format!("{}/api/task/report", server_url);
let (status, converged, max_relc, atmo_has_nan, elapsed_sec, err_msg, summary_json, seed_bytes) = match exec_res {
Ok((s, s_bytes)) => (
if s.converged {
TaskStatus::Completed
} else {
TaskStatus::Failed
},
s.converged,
s.final_max_relc,
s.atmosphere_has_nan,
s.elapsed_sec,
s.note.clone(),
serde_json::to_string(&s).unwrap_or_default(),
s_bytes,
),
Err(e) => (
TaskStatus::Failed,
false,
None,
false,
0.0,
Some(e.clone()),
serde_json::json!({"error": e}).to_string(),
None,
),
};
let report = TaskReport {
task_id: task.task_id,
point_name: task.point_name.clone(),
params: Some(task.params.clone()),
node_id: node_id.to_string(),
status,
converged,
max_relc,
atmosphere_has_nan: atmo_has_nan,
elapsed_sec,
error_message: err_msg,
summary_json,
};
let report_bytes = serde_json::to_vec(&report)?;
let seed_file_name = format!("{}.7", task.point_name);
let max_attempts = 8;
for attempt in 1..=max_attempts {
let mut form = Form::new().part(
"report",
Part::bytes(report_bytes.clone()).mime_str("application/json")?,
);
if converged && !atmo_has_nan {
if let Some(ref bytes) = seed_bytes {
let part = Part::bytes(bytes.clone())
.file_name(seed_file_name.clone())
.mime_str("application/octet-stream")?;
form = form.part("seed_file", part);
}
}
match client.post(&report_url).multipart(form).send().await {
Ok(resp) if resp.status().is_success() => {
info!(
"成功向服务端上报任务 {} (网格点: {}) 的计算结果",
task.task_id, task.point_name
);
return Ok(());
}
Ok(resp) => {
warn!(
"向服务端上报任务 {} 结果失败 (尝试 {}/{}): HTTP {}",
task.task_id, attempt, max_attempts, resp.status()
);
}
Err(e) => {
warn!(
"向服务端上报任务 {} 结果网络异常 (尝试 {}/{}): {}",
task.task_id, attempt, max_attempts, e
);
}
}
if attempt < max_attempts {
let backoff_secs = (1 << (attempt - 1)).min(60);
let backoff = std::time::Duration::from_secs(backoff_secs);
tokio::time::sleep(backoff).await;
}
}
anyhow::bail!(
"连续 {} 次向服务端上报任务 {} 结果均失败",
max_attempts,
task.task_id
)
}
+235
View File
@@ -0,0 +1,235 @@
use crate::executor::execute_task;
use crate::reporter::report_result;
use anyhow::Result;
use common::config::NodeConfig;
use common::embedded::RuntimePaths;
use common::models::{NodeHeartbeatRequest, NodeRegisterRequest, TaskSpec};
use reqwest::Client;
use serde_json::Value;
use std::path::PathBuf;
use std::sync::atomic::{AtomicI32, Ordering};
use std::sync::Arc;
use tokio::time::{sleep, Duration};
use tracing::{info, warn};
pub struct NodeWorker {
config: NodeConfig,
client: Client,
runtime: RuntimePaths,
active_slots: Arc<AtomicI32>,
}
impl NodeWorker {
pub fn new(config: NodeConfig, runtime: RuntimePaths, client: Client) -> Self {
Self {
config,
client,
runtime,
active_slots: Arc::new(AtomicI32::new(0)),
}
}
pub async fn register(&self) -> Result<()> {
info!("正在向服务端 {} 注册计算节点 {}...", self.config.server_url, self.config.node_id);
let req = NodeRegisterRequest {
node_id: self.config.node_id.clone(),
host_name: gethostname::gethostname().to_string_lossy().to_string(),
max_slots: self.config.max_slots as i32,
};
let resp = self
.client
.post(format!("{}/api/node/register", self.config.server_url))
.json(&req)
.send()
.await?;
if !resp.status().is_success() {
anyhow::bail!("向服务端注册节点失败,HTTP 状态码: {}", resp.status());
}
Ok(())
}
pub async fn run(&self) -> Result<()> {
self.register().await?;
info!("计算节点已激活,最大并行 Slot 槽位数: {}", self.config.max_slots);
// Start background heartbeat loop
let hb_client = self.client.clone();
let hb_url = format!("{}/api/node/heartbeat", self.config.server_url);
let hb_node_id = self.config.node_id.clone();
let hb_slots = self.active_slots.clone();
let hb_interval = self.config.heartbeat_sec;
tokio::spawn(async move {
let sys_arc = std::sync::Arc::new(std::sync::Mutex::new(sysinfo::System::new_all()));
{
let s = sys_arc.clone();
let _ = tokio::task::spawn_blocking(move || {
if let Ok(mut sys) = s.lock() {
sys.refresh_cpu();
}
}).await;
}
sleep(Duration::from_millis(200)).await;
{
let s = sys_arc.clone();
let _ = tokio::task::spawn_blocking(move || {
if let Ok(mut sys) = s.lock() {
sys.refresh_cpu();
}
}).await;
}
loop {
sleep(Duration::from_secs(hb_interval)).await;
let s = sys_arc.clone();
let (cpu_usage, memory_usage) = tokio::task::spawn_blocking(move || {
let mut sys = match s.lock() {
Ok(guard) => guard,
Err(_) => return (0.0, 0.0),
};
sys.refresh_cpu();
sys.refresh_memory();
let cpu_usage = sys.global_cpu_info().cpu_usage();
let mem_total = sys.total_memory() as f32;
let mem_used = sys.used_memory() as f32;
let memory_usage = if mem_total > 0.0 { (mem_used / mem_total) * 100.0 } else { 0.0 };
(cpu_usage, memory_usage)
})
.await
.unwrap_or((0.0, 0.0));
let active = hb_slots.load(Ordering::Relaxed);
let req = NodeHeartbeatRequest {
node_id: hb_node_id.clone(),
active_slots: active,
cpu_usage,
memory_usage,
};
let _ = hb_client.post(&hb_url).json(&req).send().await;
}
});
let work_dir = PathBuf::from(&self.config.work_dir);
tokio::fs::create_dir_all(&work_dir).await?;
let shutting_down = Arc::new(std::sync::atomic::AtomicBool::new(false));
let shutdown_signal = shutting_down.clone();
tokio::spawn(async move {
if tokio::signal::ctrl_c().await.is_ok() {
info!("收到 Ctrl+C 终止信号,停止领用新任务,准备优雅退出 (再次按 Ctrl+C 可强制立即退出)...");
shutdown_signal.store(true, Ordering::SeqCst);
// 二次 Ctrl+C 强行立即退出
if tokio::signal::ctrl_c().await.is_ok() {
warn!("再次收到 Ctrl+C 终止信号,强行立即中断退出!");
std::process::exit(130);
}
}
});
let mut was_disconnected = false;
// 带有优雅退出信号响应的任务领用主循环
loop {
if shutting_down.load(Ordering::Relaxed) {
break;
}
let active = self.active_slots.load(Ordering::Relaxed);
if (active as usize) < self.config.max_slots {
match self.claim_task().await {
Ok(Some(task)) => {
if was_disconnected {
info!("与服务端恢复网络连接,已自动重新上线并开始领用计算任务!");
was_disconnected = false;
}
self.active_slots.fetch_add(1, Ordering::SeqCst);
let client = self.client.clone();
let server_url = self.config.server_url.clone();
let node_id = self.config.node_id.clone();
let runtime = self.runtime.clone();
let work_dir = work_dir.clone();
let slots_counter = self.active_slots.clone();
tokio::spawn(async move {
let res = execute_task(&client, &server_url, &runtime, &work_dir, &task)
.await
.map_err(|e| e.to_string());
if let Err(e) = report_result(&client, &server_url, &node_id, &task, res).await {
warn!("向服务端上报任务 {} 计算结果失败: {}", task.task_id, e);
}
slots_counter.fetch_sub(1, Ordering::SeqCst);
});
}
Ok(None) => {
if was_disconnected {
info!("与服务端恢复网络连接,已自动重新上线 (当前暂无排队任务)。");
was_disconnected = false;
}
sleep(Duration::from_secs(5)).await;
}
Err(e) => {
was_disconnected = true;
warn!("向服务端请求领用计算任务时出错: {}", e);
sleep(Duration::from_secs(10)).await;
}
}
} else {
sleep(Duration::from_secs(2)).await;
}
}
// 等待在途任务完结(最多等待 30 秒)
if self.active_slots.load(Ordering::SeqCst) > 0 {
info!(
"正在等待 {} 个在途计算任务优雅完结 (上限 30 秒,按二次 Ctrl+C 可强行中断)...",
self.active_slots.load(Ordering::SeqCst)
);
}
let start_wait = std::time::Instant::now();
let mut last_log_time = std::time::Instant::now();
while self.active_slots.load(Ordering::SeqCst) > 0 {
if start_wait.elapsed().as_secs() >= 30 {
warn!("在途任务等待超时 (30s),强制退出节点");
break;
}
if last_log_time.elapsed().as_secs() >= 5 {
info!(
"仍在等待 {} 个在途计算任务完结...",
self.active_slots.load(Ordering::SeqCst)
);
last_log_time = std::time::Instant::now();
}
sleep(Duration::from_millis(500)).await;
}
info!("DCTS 计算节点安全退出。");
Ok(())
}
async fn claim_task(&self) -> Result<Option<TaskSpec>> {
let claim_url = format!("{}/api/task/claim", self.config.server_url);
let resp = self.client.post(&claim_url).send().await?;
if !resp.status().is_success() {
return Ok(None);
}
let json: Value = resp.json().await?;
if json["status"] == "ok" && !json["task"].is_null() {
let task: TaskSpec = serde_json::from_value(json["task"].clone())?;
Ok(Some(task))
} else {
Ok(None)
}
}
}
+29
View File
@@ -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
+37
View File
@@ -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)。
+92
View File
@@ -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 静态资源已是最新,跳过构建。");
}
}
+97
View File
@@ -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(),
}
}
+74
View File
@@ -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()
}
+25
View File
@@ -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()})),
}
}
+58
View File
@@ -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()
}
+22
View File
@@ -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,
}))
}
+149
View File
@@ -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(&params, &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(&params).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)
}
+255
View File
@@ -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,
}),
),
}
}
+873
View File
@@ -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(&reg_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(&params, 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());
}
}
+4
View File
@@ -0,0 +1,4 @@
pub mod api;
pub mod db;
pub mod scheduler;
+162
View File
@@ -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(())
}
+254
View File
@@ -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(&params).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());
}
}
+156
View File
@@ -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(&reg_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(&params, 0).await.unwrap();
db.mark_grid_point_running(&params.model_name()).await.unwrap();
let stats = db.get_grid_summary_stats().await.unwrap();
assert_eq!(stats["running"], 1);
}