From 8a3cde013c259d13dadcb7a5613d80b60c0864f7 Mon Sep 17 00:00:00 2001 From: zoeyzyhu Date: Wed, 27 May 2026 13:13:53 -0400 Subject: [PATCH] combine restart --- examples/straka.yaml | 8 ++ src/input/read_restart_file.cpp | 28 +++++ src/input/read_restart_file.hpp | 8 ++ src/mesh/meshblock.cpp | 4 + src/output/combine_restart.cpp | 16 ++- src/output/restart.cpp | 55 ++++++---- tests/CMakeLists.txt | 15 +++ tests/run_restart_uncombined.py | 177 ++++++++++++++++++++++++++++++++ 8 files changed, 291 insertions(+), 20 deletions(-) create mode 100644 tests/run_restart_uncombined.py diff --git a/examples/straka.yaml b/examples/straka.yaml index f95ddae7..b96caf8d 100644 --- a/examples/straka.yaml +++ b/examples/straka.yaml @@ -59,6 +59,14 @@ scalar: outputs: - type: restart dt: 300. + # combine: true (default) writes one bundled ..restart per + # dump via the root rank. Set `combine: false` for file-per-rank dumps + # (.block..restart) with no barrier or serial + # bundling -- recommended for large parallel runs on Lustre/GPFS scratch. + # Restart from either form with `--restart ` (pass any one per-rank + # file; each rank loads its own). On Lustre, `lfs setstripe` the run dir so + # bundles (combined) or per-rank files (uncombined) spread across OSTs. + combine: true - type: netcdf variables: [prim, uov, scalar_prim] dt: 300. diff --git a/src/input/read_restart_file.cpp b/src/input/read_restart_file.cpp index 1b66b544..75966a55 100644 --- a/src/input/read_restart_file.cpp +++ b/src/input/read_restart_file.cpp @@ -1,4 +1,5 @@ // C/C++ +#include #include #include #include @@ -220,6 +221,33 @@ static Variables load_pt_from_bundle(std::string const& path, int block_rank) { return {}; } +std::string restart_path_for_rank(std::string const& path, int block_rank) { + fs::path p(path); + std::string fname = p.filename().string(); + + static const std::string tok = ".block"; + size_t pos = fname.find(tok); + if (pos == std::string::npos) return path; + + size_t dstart = pos + tok.size(); + size_t dend = dstart; + while (dend < fname.size() && + std::isdigit(static_cast(fname[dend]))) { + ++dend; + } + + // Require at least one digit followed by a '.', e.g. ".block3." — this avoids + // rewriting an unrelated basename that merely contains the text "block". + if (dend == dstart || dend >= fname.size() || fname[dend] != '.') { + return path; + } + + std::string rewritten = + fname.substr(0, dstart) + std::to_string(block_rank) + fname.substr(dend); + p.replace_filename(rewritten); + return p.string(); +} + Variables load_restart(std::string const& path, int block_rank) { // Dispatch based on whether `path` is a restart bundle or a single tensor // dump. diff --git a/src/input/read_restart_file.hpp b/src/input/read_restart_file.hpp index ee95774a..42d3647e 100644 --- a/src/input/read_restart_file.hpp +++ b/src/input/read_restart_file.hpp @@ -18,4 +18,12 @@ using Variables = std::map; Variables load_restart(std::string const& path, int block_rank = get_rank()); +//! \brief Resolve a restart path for a specific block rank. +//! +//! File-per-rank (uncombined) dumps embed a ".block." token in the filename; +//! this rewrites to `block_rank` so every rank loads its own file. Paths +//! without that token (combined bundles or single dumps) are returned +//! unchanged. +std::string restart_path_for_rank(std::string const& path, int block_rank); + } // namespace snap diff --git a/src/mesh/meshblock.cpp b/src/mesh/meshblock.cpp index 3631c161..00a035c2 100644 --- a/src/mesh/meshblock.cpp +++ b/src/mesh/meshblock.cpp @@ -973,6 +973,10 @@ int MeshBlockImpl::check_redo(Variables& vars) { } double MeshBlockImpl::_init_from_restart(Variables& vars, std::string fname) { + // For file-per-rank (uncombined) dumps, resolve the path to this block's own + // file. Combined bundles and single dumps are returned unchanged. + fname = restart_path_for_rank(fname, options->layout()->rank()); + std::filesystem::path restart_path(fname); if (!restart_path.is_absolute() && !std::filesystem::exists(restart_path)) { restart_path = std::filesystem::path(options->output_dir()) / fname; diff --git a/src/output/combine_restart.cpp b/src/output/combine_restart.cpp index 60507cbb..3e06ecb9 100644 --- a/src/output/combine_restart.cpp +++ b/src/output/combine_restart.cpp @@ -138,11 +138,23 @@ void RestartOutput::combine_blocks(MeshBlockImpl* pmb, bool final_write) { file_list.push_back(std::string(glob_result.gl_pathv[i])); } - remove(outfile.c_str()); if (file_list.size() == 1) { + // std::rename atomically replaces an existing destination on POSIX. err = std::rename(file_list.front().c_str(), outfile.c_str()); } else { - err = make_restart_bundle(outfile, file_list); + // Bundle into a temp file, then atomically rename over the destination so + // a crash mid-bundle leaves the previous ".restart" intact. + std::string tmp = outfile + ".tmp"; + err = make_restart_bundle(tmp, file_list); + if (err == 0) { + std::error_code rn_ec; + std::filesystem::rename(tmp, outfile, rn_ec); + if (rn_ec) { + std::error_code rm_ec; + std::filesystem::remove(tmp, rm_ec); + err = -1; + } + } } if (err) { diff --git a/src/output/restart.cpp b/src/output/restart.cpp index ece35959..64f21578 100644 --- a/src/output/restart.cpp +++ b/src/output/restart.cpp @@ -15,10 +15,28 @@ namespace snap { +namespace { +//! \brief Write tensors to a temporary file and atomically rename it into +//! place, so a crash mid-write can never corrupt an existing restart file. +void save_tensors_atomic(Variables const& vars, std::string const& final_path) { + std::string tmp_path = final_path + ".tmp"; + kintera::save_tensors(vars, tmp_path); + std::error_code ec; + std::filesystem::rename(tmp_path, final_path, ec); + if (ec) { + std::error_code rm_ec; + std::filesystem::remove(tmp_path, rm_ec); + throw std::runtime_error("Failed to finalize restart file '" + final_path + + "': " + ec.message()); + } +} +} // namespace + RestartOutput::RestartOutput(OutputOptions const& options_) : OutputType(options_) { - // restart files are always combined - options->combine(true); + // Combined output (a single bundled ".restart" per dump) is the default. + // Set `combine: false` in the output block to write one ".restart" per rank + // with no cross-rank barrier or serial bundling (parallel-friendly at scale). } void RestartOutput::write_output_file(MeshBlockImpl* pmb, Variables const& vars, @@ -51,37 +69,38 @@ void RestartOutput::write_output_file(MeshBlockImpl* pmb, Variables const& vars, out_vars["file_number"] = torch::tensor(output_file_numbers, torch::kInt64); out_vars["next_time"] = torch::tensor(output_next_times, torch::kFloat64); - // create filename: ...part - std::string fname; + // shared stem: /.. char number[6]; snprintf(number, sizeof(number), "%05d", file_number); char blockid[12]; snprintf(blockid, sizeof(blockid), "block%d", pmb->options->layout()->rank()); - fname.assign(pmb->options->output_dir()); - fname.append("/"); - fname.append(pmb->options->basename()); - fname.append("."); - fname.append(blockid); - fname.append("."); - if (final_write) { - fname.append("final"); - } else { - fname.append(number); - } - fname.append(".part"); + std::string stem; + stem.assign(pmb->options->output_dir()); + stem.append("/"); + stem.append(pmb->options->basename()); + stem.append("."); + stem.append(blockid); + stem.append("."); + stem.append(final_write ? "final" : number); - // save to disk + // ensure the output directory exists std::error_code ec; std::filesystem::create_directories(pmb->options->output_dir(), ec); if (ec) { throw std::runtime_error("Failed to create output directory '" + pmb->options->output_dir() + "': " + ec.message()); } - kintera::save_tensors(out_vars, fname); if (options->combine()) { + // Each rank writes a per-block ".part"; the root rank bundles all parts + // into one combined ".restart" file. + save_tensors_atomic(out_vars, stem + ".part"); combine_blocks(pmb, final_write); + } else { + // File-per-rank: each rank writes its own ".restart" directly. On restart + // every rank loads its own file (the block id is rewritten to its rank). + save_tensors_atomic(out_vars, stem + ".restart"); } } diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 2a558bb0..9b4ad177 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -185,6 +185,21 @@ if (NOT APPLE) TIMEOUT 600 LABELS "restart;examples" ) + + add_test( + NAME test_restart_uncombined + COMMAND ${Python3_EXECUTABLE} + ${CMAKE_CURRENT_SOURCE_DIR}/run_restart_uncombined.py + --build-dir ${CMAKE_BINARY_DIR} + --build-type release + ) + + set_tests_properties(test_restart_uncombined PROPERTIES + WORKING_DIRECTORY ${CMAKE_BINARY_DIR}/tests + SKIP_RETURN_CODE 125 + TIMEOUT 600 + LABELS "restart;examples" + ) endif() if (DEFINED FULL_TESTS) diff --git a/tests/run_restart_uncombined.py b/tests/run_restart_uncombined.py new file mode 100644 index 00000000..7586dbc4 --- /dev/null +++ b/tests/run_restart_uncombined.py @@ -0,0 +1,177 @@ +#!/usr/bin/env python3 +"""Verify file-per-rank (uncombined) restart reproduces an uninterrupted run. + +This exercises the ``combine: false`` restart path, where every rank writes its +own ``.block..restart`` file with no cross-rank barrier or +serial bundling. On restart, each rank loads its own file (the block id in the +path is rewritten to the local rank), so passing any one of the per-rank files +to ``--restart`` resumes all blocks. +""" +import argparse +import os +import shutil +import subprocess +import sys +from pathlib import Path + +import netCDF4 +import numpy as np + +SKIP_CODE = 125 + +try: + import yaml +except Exception as exc: # pragma: no cover - dependency guard + print(f"Skipping test_restart_uncombined: yaml import failed: {exc}") + sys.exit(SKIP_CODE) + + +def run(cmd, cwd: Path, env=None) -> None: + print(f"+ (cd {cwd} && {' '.join(cmd)})") + subprocess.run(cmd, cwd=cwd, env=env, check=True) + + +def prepare_case(base_yaml: Path, case_dir: Path, blocks_per_process: int) -> Path: + if case_dir.exists(): + shutil.rmtree(case_dir) + case_dir.mkdir(parents=True) + + with base_yaml.open("r") as f: + config = yaml.safe_load(f) + + dist = config.setdefault("distribute", {}) + dist["backend"] = "gloo" + dist["blocks_per_process"] = blocks_per_process + + integration = config.setdefault("integration", {}) + integration["tlim"] = 0.0 + integration["nlim"] = 0 + integration["ncycle_out"] = 0 + + config["output_dir"] = "." + config["outputs"] = [ + # File-per-rank restart: no root-rank bundling. + {"type": "restart", "dt": 0.0, "combine": False}, + {"type": "netcdf", "variables": ["prim"], "dt": 0.0}, + ] + + target_yaml = case_dir / base_yaml.name + with target_yaml.open("w") as f: + yaml.safe_dump(config, f) + return target_yaml + + +def combine_output(case_dir: Path) -> Path: + nc_files = sorted(case_dir.glob("*.nc")) + if not nc_files: + raise FileNotFoundError(f"No NetCDF output found in {case_dir}") + return nc_files[-1] + + +def find_per_rank_restart(case_dir: Path) -> Path: + # Uncombined dumps are named .block..restart. Confirm a + # per-rank file exists for each block, then return one for --restart (each + # rank rewrites the block id to its own rank on load). + restart_files = sorted(case_dir.glob("*.block*.restart")) + if not restart_files: + raise FileNotFoundError(f"No per-rank restart file found in {case_dir}") + return restart_files[0] + + +def compare_netcdf(path_a: Path, path_b: Path) -> None: + with netCDF4.Dataset(path_a, "r") as data_a, netCDF4.Dataset(path_b, "r") as data_b: + vars_a = {k: np.asarray(v[:]) for k, v in data_a.variables.items() if v.ndim > 0} + vars_b = {k: np.asarray(v[:]) for k, v in data_b.variables.items() if v.ndim > 0} + + if vars_a.keys() != vars_b.keys(): + raise ValueError(f"Variable mismatch: {vars_a.keys()} vs {vars_b.keys()}") + + for name in vars_a: + diff = np.abs(vars_a[name] - vars_b[name]) + max_abs = float(diff.max(initial=0.0)) + if max_abs != 0.0: + raise ValueError(f"{name} differs after restart (max abs diff {max_abs})") + + +def main() -> int: + parser = argparse.ArgumentParser() + parser.add_argument("--build-dir", required=True) + parser.add_argument("--build-type", required=True) + args = parser.parse_args() + + build_dir = Path(args.build_dir).resolve() + bin_dir = build_dir / "bin" + tests_dir = build_dir / "tests" + repo_root = Path(__file__).resolve().parent.parent + + torchrun = shutil.which("torchrun") + if torchrun is None: + print("Skipping test_restart_uncombined: torchrun not found") + return SKIP_CODE + + name = "straka" + blocks_per_process = 2 + exe = bin_dir / f"{name}.{args.build_type}" + if not exe.exists(): + raise FileNotFoundError(f"missing executable {exe}") + + base_yaml = repo_root / "examples" / "straka.yaml" + if not base_yaml.exists(): + raise FileNotFoundError(f"missing input file {base_yaml}") + + env = os.environ.copy() + env["BACKEND"] = "gloo" + py_paths = [str(repo_root / "python"), str(repo_root)] + existing = env.get("PYTHONPATH") + if existing: + py_paths.append(existing) + env["PYTHONPATH"] = ":".join(py_paths) + + case_dir = tests_dir / f"restart_uncombined_{name}" + yaml_path = prepare_case(base_yaml, case_dir, blocks_per_process) + + run( + [ + torchrun, + "--no-python", + "--nproc-per-node=1", + str(exe), + str(yaml_path), + ], + cwd=case_dir, + env=env, + ) + base_nc = combine_output(case_dir) + restart_file = find_per_rank_restart(case_dir) + + restart_dir = tests_dir / f"restart_uncombined_{name}_from_restart" + if restart_dir.exists(): + shutil.rmtree(restart_dir) + restart_dir.mkdir(parents=True) + restart_yaml = restart_dir / yaml_path.name + shutil.copy2(yaml_path, restart_yaml) + # Copy every per-rank restart file so each block can load its own. + for part in case_dir.glob("*.block*.restart"): + shutil.copy2(part, restart_dir / part.name) + + run( + [ + torchrun, + "--no-python", + "--nproc-per-node=1", + str(exe), + str(restart_yaml), + "--restart", + str((restart_dir / restart_file.name).resolve()), + ], + cwd=restart_dir, + env=env, + ) + restarted_nc = combine_output(restart_dir) + compare_netcdf(base_nc, restarted_nc) + + return 0 + + +if __name__ == "__main__": + raise SystemExit(main())