Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
40 changes: 33 additions & 7 deletions src/checksum.rs
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
use anyhow::{Context, Result, bail};
use anyhow::{bail, Context, Result};
use md5::{Digest, Md5};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
Expand Down Expand Up @@ -175,10 +175,7 @@ pub fn validate_checksums(
///
/// * `Result<()>` - Ok if the live version matches the stored one, Err otherwise
pub fn validate_haddock3_version_live(checksum_file: &Path) -> Result<()> {
let stored_content =
fs::read_to_string(checksum_file).context("Failed to read checksum file")?;
let stored_data: ChecksumData =
serde_json::from_str(&stored_content).context("Failed to parse checksum file")?;
let stored_data = read_checksum_data(checksum_file)?;

let live_version = crate::utils::get_haddock3_version()?;

Expand All @@ -193,6 +190,20 @@ pub fn validate_haddock3_version_live(checksum_file: &Path) -> Result<()> {
Ok(())
}

fn read_checksum_data(checksum_file: &Path) -> Result<ChecksumData> {
let stored_content =
fs::read_to_string(checksum_file).context("Failed to read checksum file")?;
serde_json::from_str(&stored_content).context("Failed to parse checksum file")
}

/// Returns the haddock3 version recorded in `checksum.json` when the benchmark started.
///
/// Used to embed the expected version into a generated SLURM job script so it can be
/// re-checked from inside the job itself, after `slurm_prologue` has run.
pub fn expected_haddock3_version(checksum_file: &Path) -> Result<String> {
Ok(read_checksum_data(checksum_file)?.haddock3_version)
}

fn find_modified(
current_checksums: HashMap<String, String>,
stored_checksums: HashMap<String, String>,
Expand Down Expand Up @@ -265,7 +276,7 @@ mod tests {
use super::*;
use std::collections::HashMap;
use std::io::Write;
use tempfile::{NamedTempFile, tempdir};
use tempfile::{tempdir, NamedTempFile};

#[test]
fn test_calculate_checksum() {
Expand Down Expand Up @@ -433,7 +444,7 @@ mod tests {
let mut stored = HashMap::new();
stored.insert("file1.txt".to_string(), "checksum1_changed".to_string()); // Changed
stored.insert("file4.txt".to_string(), "checksum4".to_string()); // Removed
// file2.txt and file3.txt are new
// file2.txt and file3.txt are new

// Find modified files
let error_msg = find_modified(current, stored);
Expand Down Expand Up @@ -570,4 +581,19 @@ mod tests {
let result = validate_haddock3_version_live(&checksum_file);
assert!(result.is_err());
}

#[test]
fn test_expected_haddock3_version() {
let temp_dir = tempdir().unwrap();
let checksum_file = temp_dir.path().join("checksum.json");

let data = ChecksumData {
files: HashMap::new(),
haddock3_version: "2026.3.0".to_string(),
};
fs::write(&checksum_file, serde_json::to_string_pretty(&data).unwrap()).unwrap();

let version = expected_haddock3_version(&checksum_file).unwrap();
assert_eq!(version, "2026.3.0");
}
}
38 changes: 37 additions & 1 deletion src/job.rs
Original file line number Diff line number Diff line change
Expand Up @@ -405,22 +405,34 @@ impl Job {
let mut file = fs::File::create(&job_script)?;

let absolute_wd = canonicalize(&self.wd).unwrap_or_else(|_| self.wd.to_path_buf());
let haddock3_exe = find_haddock3_executable()?;

// Write SLURM header
let header = self.generate_slurm_header();

let prologue = self.format_slurm_prologue();

// slurm_prologue can alter the environment (module load/conda activate/etc.)
// before haddock3 runs; re-check the version from inside the job itself so
// a swap isn't invisible to the login-node-only checks in checksum.rs.
let mut version_check = String::new();
if self.general.slurm_prologue.is_some() {
let checksum_file = self.general.work_dir.join("checksum.json");
let expected_version = crate::checksum::expected_haddock3_version(&checksum_file)?;
version_check = format_haddock3_version_check(&haddock3_exe, &expected_version);
}

// Write job body
let body = format!(
"cd {}\n{} {}\n",
absolute_wd.display(),
find_haddock3_executable()?,
haddock3_exe,
WORKFLOW_FILENAME
);

file.write_all(header.as_bytes())?;
file.write_all(prologue.as_bytes())?;
file.write_all(version_check.as_bytes())?;
file.write_all(body.as_bytes())?;

Ok(())
Expand Down Expand Up @@ -540,6 +552,21 @@ impl Job {
}
}

/// Renders a shell snippet that re-checks `haddock3_exe --version` against `expected_version`,
/// exiting the job script with an error on mismatch. Meant to run after `slurm_prologue`
/// commands, in the same shell state, since those can alter the environment haddock3 sees.
fn format_haddock3_version_check(haddock3_exe: &str, expected_version: &str) -> String {
format!(
r#"HADDOCK3_VERSION_OUTPUT="$({haddock3_exe} --version 2>&1)"
ACTUAL_HADDOCK3_VERSION="${{HADDOCK3_VERSION_OUTPUT##* - }}"
if [ "$ACTUAL_HADDOCK3_VERSION" != "{expected_version}" ]; then
echo "ERROR: haddock3 version mismatch after slurm_prologue: expected {expected_version}, found $ACTUAL_HADDOCK3_VERSION. Check slurm_prologue for environment changes (module load/conda activate/etc.)." >&2
exit 1
fi
"#
)
}

#[cfg(test)]
mod tests {
use super::*;
Expand Down Expand Up @@ -1073,6 +1100,15 @@ mod tests {
assert_eq!(job.format_slurm_prologue(), "module load foo\n");
}

#[test]
fn test_format_haddock3_version_check_contents() {
let check = format_haddock3_version_check("/opt/haddock3/bin/haddock3", "2026.3.0");
assert!(check.contains("/opt/haddock3/bin/haddock3 --version"));
assert!(check.contains("2026.3.0"));
assert!(check.contains("ACTUAL_HADDOCK3_VERSION"));
assert!(check.contains("exit 1"));
}

fn make_job_with_modules(modules: IndexMap<String, Value>) -> Job {
Job {
name: "test".to_string(),
Expand Down
Loading