Harden scheduled training state recovery

Author: root

src/scheduler.rs

@@ -11,6 +11,8 @@ use crate::profile::ModerationProfile;
 use crate::sample_rpc::SampleRpcConfig;
 use crate::training::{evaluate_training_need, fetch_training_sample_count_via_rpc};
 
+const RUNNING_STATUS_STALE_SECS: u64 = 2 * 60 * 60;
+
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub struct TrainingSubprocessCommand {
     pub program: String,
@@ -95,7 +97,7 @@ pub fn cooldown_allows_training(
     Ok(elapsed_secs >= cooldown_secs)
 }
 
-pub fn training_launch_allowed(profile: &ModerationProfile) -> Result<bool> {
+pub fn training_launch_allowed(profile: &ModerationProfile, now_unix_secs: u64) -> Result<bool> {
     let Some(status) = profile.training_status() else {
         return Ok(true);
     };
@@ -104,7 +106,46 @@ pub fn training_launch_allowed(profile: &ModerationProfile) -> Result<bool> {
         return Ok(true);
     };
 
-    Ok(state != "running")
+    if state != "running" {
+        return Ok(true);
+    }
+
+    let Some(timestamp) = status.get("timestamp").and_then(|value| value.as_u64()) else {
+        return Ok(true);
+    };
+    if now_unix_secs.saturating_sub(timestamp) >= RUNNING_STATUS_STALE_SECS {
+        return Ok(true);
+    }
+
+    let Some(pid) = status
+        .get("pid")
+        .and_then(|value| value.as_u64())
+        .and_then(|value| u32::try_from(value).ok())
+    else {
+        return Ok(false);
+    };
+
+    Ok(!training_process_matches_profile(pid, &profile.profile_name))
+}
+
+fn training_process_matches_profile(pid: u32, profile_name: &str) -> bool {
+    let cmdline_path = format!("/proc/{pid}/cmdline");
+    let Ok(raw) = fs::read(cmdline_path) else {
+        return false;
+    };
+    if raw.is_empty() {
+        return false;
+    }
+
+    let args = raw
+        .split(|byte| *byte == 0)
+        .filter(|part| !part.is_empty())
+        .filter_map(|part| std::str::from_utf8(part).ok())
+        .collect::<Vec<_>>();
+
+    args.windows(2).any(|window| {
+        window[0] == "train-profile" && window[1] == profile_name
+    })
 }
 
 pub fn build_training_subprocess_command(
@@ -139,7 +180,7 @@ pub async fn plan_training_round(
     let mut planned = Vec::new();
 
     for scanned in scan_profiles(root_dir)? {
-        if !training_launch_allowed(&scanned.profile)? {
+        if !training_launch_allowed(&scanned.profile, now_unix_secs)? {
             continue;
         }
 

src/training.rs

@@ -247,60 +247,110 @@ pub async fn run_training_subprocess_from_args(args: &[String]) -> Result<()> {
         .cloned()
         .ok_or_else(|| anyhow!("usage: {} train-profile <profile-name>", binary_name(args)))?;
     let root_dir = std::env::current_dir().context("failed to resolve current dir")?;
-    let settings = crate::config::Settings::load(&root_dir)?;
-    let rpc = SampleRpcConfig::from_settings(&settings)?;
-    let profile = ModerationProfile::load(&root_dir, &profile_name)?;
-
-    let initial_sample_count = fetch_training_sample_count_via_rpc(&rpc, &profile)
-        .await
-        .unwrap_or(0);
-    write_training_status(
-        &profile,
-        &serde_json::json!({
-            "status": "running",
-            "message": "training in progress",
-            "timestamp": current_unix_secs(),
-            "profile": profile.profile_name,
-            "model_type": profile.config.local_model_type,
-            "sample_count": initial_sample_count
-        }),
-    )?;
-
-    match run_profile_training(&rpc, &profile).await {
-        Ok(output) => {
-            write_training_status(
-                &profile,
-                &serde_json::json!({
-                    "status": "success",
-                    "message": "training completed",
-                    "timestamp": current_unix_secs(),
-                    "profile": profile.profile_name,
-                    "model_type": profile.config.local_model_type,
-                    "sample_count": output.sample_count,
-                    "pass_count": output.pass_count,
-                    "violation_count": output.violation_count,
-                    "runtime_json_path": output.runtime_json_path,
-                    "runtime_coef_path": output.runtime_coef_path,
-                    "model_marker_path": output.model_marker_path
-                }),
-            )?;
-            Ok(())
-        }
-        Err(err) => {
-            write_training_status(
-                &profile,
-                &serde_json::json!({
-                    "status": "failed",
-                    "message": format!("{err:#}"),
-                    "timestamp": current_unix_secs(),
-                    "profile": profile.profile_name,
-                    "model_type": profile.config.local_model_type,
-                    "sample_count": initial_sample_count
-                }),
-            )?;
-            Err(err)
+    let result = async {
+        let settings = crate::config::Settings::load(&root_dir)?;
+        let rpc = SampleRpcConfig::from_settings(&settings)?;
+        let profile = ModerationProfile::load(&root_dir, &profile_name)?;
+
+        let initial_sample_count = fetch_training_sample_count_via_rpc(&rpc, &profile)
+            .await
+            .unwrap_or(0);
+        write_training_status(
+            &profile,
+            &serde_json::json!({
+                "status": "running",
+                "message": "training in progress",
+                "timestamp": current_unix_secs(),
+                "started_at": current_unix_secs(),
+                "pid": std::process::id(),
+                "profile": profile.profile_name,
+                "model_type": profile.config.local_model_type,
+                "sample_count": initial_sample_count
+            }),
+        )?;
+
+        match run_profile_training(&rpc, &profile).await {
+            Ok(output) => {
+                write_training_status(
+                    &profile,
+                    &serde_json::json!({
+                        "status": "success",
+                        "message": "training completed",
+                        "timestamp": current_unix_secs(),
+                        "profile": profile.profile_name,
+                        "model_type": profile.config.local_model_type,
+                        "sample_count": output.sample_count,
+                        "pass_count": output.pass_count,
+                        "violation_count": output.violation_count,
+                        "runtime_json_path": output.runtime_json_path,
+                        "runtime_coef_path": output.runtime_coef_path,
+                        "model_marker_path": output.model_marker_path
+                    }),
+                )?;
+                Ok(())
+            }
+            Err(err) => {
+                write_training_status(
+                    &profile,
+                    &serde_json::json!({
+                        "status": "failed",
+                        "message": format!("{err:#}"),
+                        "timestamp": current_unix_secs(),
+                        "profile": profile.profile_name,
+                        "model_type": profile.config.local_model_type,
+                        "sample_count": initial_sample_count
+                    }),
+                )?;
+                Err(err)
+            }
         }
     }
+    .await;
+
+    if let Err(err) = &result {
+        let _ = write_training_failure_status(
+            &root_dir,
+            &profile_name,
+            format!("{err:#}"),
+        );
+    }
+
+    result
+}
+
+fn write_training_failure_status(
+    root_dir: &Path,
+    profile_name: &str,
+    message: String,
+) -> Result<()> {
+    let status_path = root_dir
+        .join("configs")
+        .join("mod_profiles")
+        .join(profile_name)
+        .join(".train_status.json");
+    if let Some(parent) = status_path.parent() {
+        fs::create_dir_all(parent)
+            .with_context(|| format!("failed to create profile dir {}", parent.display()))?;
+    }
+
+    let payload = serde_json::json!({
+        "status": "failed",
+        "message": message,
+        "timestamp": current_unix_secs(),
+        "profile": profile_name,
+    });
+    let tmp_path = status_path.with_extension("json.tmp");
+    let encoded = serde_json::to_vec_pretty(&payload).context("failed to encode training failure status")?;
+    fs::write(&tmp_path, encoded)
+        .with_context(|| format!("failed to write training failure status tmp {}", tmp_path.display()))?;
+    fs::rename(&tmp_path, &status_path).with_context(|| {
+        format!(
+            "failed to replace training failure status {} with {}",
+            status_path.display(),
+            tmp_path.display()
+        )
+    })?;
+    Ok(())
 }
 
 pub fn train_hashlinear_runtime(

tests/scheduler_tests.rs

@@ -204,7 +204,8 @@ fn scheduler_skips_profile_when_training_is_already_running() {
         }),
     );
 
-    let decision = scheduler::training_launch_allowed(&profile).expect("training launch decision");
+    let decision = scheduler::training_launch_allowed(&profile, current_unix_secs())
+        .expect("training launch decision");
 
     assert!(!decision);
 }