feat(all): 重炼 crates/common 核心组件、上线 Web 运维看板与 Docker 容器化部署
This commit is contained in:
@@ -0,0 +1,254 @@
|
||||
use anyhow::Result;
|
||||
use common::config::GridConfig;
|
||||
use common::models::{GridPointParams, TaskSpec, TaskType};
|
||||
use mq::sqlite_queue::SqliteTaskQueue;
|
||||
use std::sync::Arc;
|
||||
use tracing::info;
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::db::Database;
|
||||
|
||||
pub struct GridScheduler {
|
||||
db: Database,
|
||||
queue: Arc<SqliteTaskQueue>,
|
||||
_results_dir: String,
|
||||
}
|
||||
|
||||
impl GridScheduler {
|
||||
pub fn new(db: Database, queue: Arc<SqliteTaskQueue>, results_dir: String) -> Self {
|
||||
Self {
|
||||
db,
|
||||
queue,
|
||||
_results_dir: results_dir,
|
||||
}
|
||||
}
|
||||
|
||||
/// Expands grid points from config and registers them into the database
|
||||
pub async fn initialize_grid(&self, cfg: &GridConfig) -> Result<()> {
|
||||
if let Err(e) = self.queue.clear_queue().await {
|
||||
tracing::warn!("初始化网格时清理闲置排队记录发生警告: {}", e);
|
||||
}
|
||||
if let Err(e) = self.db.reset_queued_grid_points_to_pending().await {
|
||||
tracing::warn!("重置网格状态到 pending 处理过程遇到异常: {}", e);
|
||||
}
|
||||
let mut points = Vec::new();
|
||||
|
||||
for &teff in &cfg.grid.teff {
|
||||
for &logg in &cfg.grid.logg {
|
||||
for &loghe in &cfg.grid.loghe {
|
||||
for &logc in &cfg.grid.logc {
|
||||
for &logn in &cfg.grid.logn {
|
||||
for &logo in &cfg.grid.logo {
|
||||
points.push(GridPointParams {
|
||||
teff,
|
||||
logg,
|
||||
loghe,
|
||||
logc,
|
||||
logn,
|
||||
logo,
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Sort by difficulty: cno_sum ASC -> teff ASC -> -logg -> loghe ASC
|
||||
points.sort_by(|a, b| {
|
||||
a.cno_sum()
|
||||
.partial_cmp(&b.cno_sum())
|
||||
.unwrap_or(std::cmp::Ordering::Equal)
|
||||
.then_with(|| a.teff.partial_cmp(&b.teff).unwrap_or(std::cmp::Ordering::Equal))
|
||||
.then_with(|| b.logg.partial_cmp(&a.logg).unwrap_or(std::cmp::Ordering::Equal))
|
||||
.then_with(|| a.loghe.partial_cmp(&b.loghe).unwrap_or(std::cmp::Ordering::Equal))
|
||||
});
|
||||
|
||||
// Group into Waves by cno_sum
|
||||
let mut current_cno: Option<f64> = None;
|
||||
let mut wave_idx = 0;
|
||||
|
||||
for pt in &points {
|
||||
let cno = pt.cno_sum();
|
||||
if let Some(cur) = current_cno {
|
||||
if (cur - cno).abs() > 1e-5 {
|
||||
wave_idx += 1;
|
||||
current_cno = Some(cno);
|
||||
}
|
||||
} else {
|
||||
current_cno = Some(cno);
|
||||
}
|
||||
|
||||
self.db.upsert_grid_point(pt, wave_idx).await?;
|
||||
}
|
||||
|
||||
info!("已在数据库中成功初始化并记录 {} 个恒星大气网格点", points.len());
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn get_active_timeout_sec(&self) -> u64 {
|
||||
if let Ok(yamls) = self.db.get_running_workflow_config_yamls().await {
|
||||
for yaml in yamls {
|
||||
if let Ok(cfg) = serde_yaml::from_str::<GridConfig>(&yaml) {
|
||||
return cfg.timeout_sec;
|
||||
}
|
||||
}
|
||||
}
|
||||
7200
|
||||
}
|
||||
|
||||
/// Enqueues pending grid points into MQ with active seed detection and batching
|
||||
pub async fn schedule_pending_tasks(&self) -> Result<usize> {
|
||||
if !self.db.has_running_workflow().await? {
|
||||
return Ok(0);
|
||||
}
|
||||
|
||||
let timeout_sec = self.get_active_timeout_sec().await;
|
||||
let batch_limit: usize = std::env::var("DCTS_BATCH_LIMIT")
|
||||
.ok()
|
||||
.and_then(|v| v.parse().ok())
|
||||
.unwrap_or(100);
|
||||
|
||||
// SQL 层直接附加 LIMIT = batch_limit 筛选，完全免除数万点位无谓内存反序列化和空耗对象释放开销
|
||||
let pending = self.db.get_pending_grid_points_limit(batch_limit).await?;
|
||||
let mut dispatched = 0;
|
||||
|
||||
for (name, params, _wave) in pending {
|
||||
|
||||
// Check if any seed is available in DB for active SeedStep scheduling
|
||||
let (task_type, seed_name) = match self.db.find_best_seed_from_db(¶ms).await {
|
||||
Ok(Some(seed_match)) => {
|
||||
info!("网格点 {} 匹配到数据库近邻种子 {} (距离: {:.2})，安排 SeedStep 热启动调度", name, seed_match.name, seed_match.distance);
|
||||
(TaskType::SeedStep, Some(seed_match.name))
|
||||
}
|
||||
_ => (TaskType::ColdRun, None),
|
||||
};
|
||||
|
||||
let task_spec = TaskSpec {
|
||||
task_id: Uuid::new_v4(),
|
||||
point_name: name.clone(),
|
||||
params,
|
||||
task_type,
|
||||
seed_point_name: seed_name,
|
||||
timeout_sec,
|
||||
};
|
||||
|
||||
self.db.insert_task(&task_spec).await?;
|
||||
// 采用先标记 DB 状态为 Queued 后发 MQ 的时序，防止推入 MQ 后数据库修改异常导向下一轮误重投
|
||||
self.db.update_grid_status(&name, common::models::GridPointStatus::Queued).await?;
|
||||
match self.queue.push_task(&task_spec).await {
|
||||
Ok(_) => {
|
||||
dispatched += 1;
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!("将任务 {} 推入 MQ 队列失败，回滚网格点状态: {}", name, e);
|
||||
let _ = self.db.update_grid_status(&name, common::models::GridPointStatus::Pending).await;
|
||||
let _ = self.queue.remove_task(&task_spec.task_id.to_string()).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if dispatched > 0 {
|
||||
info!("已成功将 {} 个待计算网格点推进任务队列", dispatched);
|
||||
}
|
||||
|
||||
Ok(dispatched)
|
||||
}
|
||||
|
||||
/// Triggers seed_step fallback for a failed point if a seed is available
|
||||
pub async fn trigger_seed_step_fallback(&self, params: &GridPointParams) -> Result<bool> {
|
||||
if !self.db.has_running_workflow().await? {
|
||||
return Ok(false);
|
||||
}
|
||||
|
||||
let name = params.model_name();
|
||||
if let Ok(Some((status, attempt_count))) = self.db.get_grid_point_status(&name).await {
|
||||
if status == "failed" || attempt_count >= 3 {
|
||||
info!("网格点 {} 已达到最大重试次数 ({}) 或处于 failed 状态，跳过种子热启动回退", name, attempt_count);
|
||||
return Ok(false);
|
||||
}
|
||||
}
|
||||
|
||||
let seed_match_opt = self.db.find_best_seed_from_db(params).await.ok().flatten();
|
||||
|
||||
if let Some(seed_match) = seed_match_opt {
|
||||
let timeout_sec = self.get_active_timeout_sec().await;
|
||||
let name = params.model_name();
|
||||
let task_spec = TaskSpec {
|
||||
task_id: Uuid::new_v4(),
|
||||
point_name: name.clone(),
|
||||
params: params.clone(),
|
||||
task_type: TaskType::SeedStep,
|
||||
seed_point_name: Some(seed_match.name.clone()),
|
||||
timeout_sec,
|
||||
};
|
||||
|
||||
self.db.insert_task(&task_spec).await?;
|
||||
self.db.update_grid_status(&name, common::models::GridPointStatus::Queued).await?;
|
||||
if let Err(e) = self.queue.push_task(&task_spec).await {
|
||||
let _ = self.db.update_grid_status(&name, common::models::GridPointStatus::Pending).await;
|
||||
let _ = self.queue.remove_task(&task_spec.task_id.to_string()).await;
|
||||
return Err(e);
|
||||
}
|
||||
info!("触发种子步进 (seed_step)：网格点 {} 将使用 6 维近邻种子 {} 热启动重试", name, seed_match.name);
|
||||
Ok(true)
|
||||
} else {
|
||||
Ok(false)
|
||||
}
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use common::config::GridAxesConfig;
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_grid_scheduler_initialization_and_scheduling() {
|
||||
let temp_dir = tempfile::tempdir().unwrap();
|
||||
let db_path = temp_dir.path().join("sched_db.db");
|
||||
let queue_db_path = temp_dir.path().join("sched_queue.db");
|
||||
let results_dir = temp_dir.path().join("results");
|
||||
|
||||
let db = Database::new(&db_path.to_string_lossy()).await.unwrap();
|
||||
let queue = Arc::new(SqliteTaskQueue::new(&queue_db_path.to_string_lossy()).await.unwrap());
|
||||
let scheduler = GridScheduler::new(db.clone(), queue.clone(), results_dir.to_string_lossy().to_string());
|
||||
|
||||
let cfg = GridConfig {
|
||||
grid: GridAxesConfig {
|
||||
teff: vec![35000.0],
|
||||
logg: vec![5.5],
|
||||
loghe: vec![-1.0],
|
||||
logc: vec![-2.0],
|
||||
logn: vec![-2.0],
|
||||
logo: vec![-2.0],
|
||||
},
|
||||
chain: vec![],
|
||||
synspec: None,
|
||||
nworkers: 4,
|
||||
timeout_sec: 3600,
|
||||
resume: true,
|
||||
seed_step_fallback: true,
|
||||
results: None,
|
||||
itek_fallback: vec![],
|
||||
niter: Some(100),
|
||||
template: None,
|
||||
fort55: None,
|
||||
linelist: None,
|
||||
};
|
||||
|
||||
scheduler.initialize_grid(&cfg).await.unwrap();
|
||||
db.upsert_workflow("test_wf", None, "", "running").await.unwrap();
|
||||
|
||||
let pending = db.get_pending_grid_points().await.unwrap();
|
||||
assert_eq!(pending.len(), 1);
|
||||
|
||||
let dispatched = scheduler.schedule_pending_tasks().await.unwrap();
|
||||
assert_eq!(dispatched, 1);
|
||||
|
||||
let popped = queue.pop_task().await.unwrap();
|
||||
assert!(popped.is_some());
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user