combine restart

examples/straka.yaml

@@ -59,6 +59,14 @@ scalar:
 outputs:
   - type: restart
     dt: 300.
+    # combine: true (default) writes one bundled <basename>.<NNNNN>.restart per
+    # dump via the root rank. Set `combine: false` for file-per-rank dumps
+    # (<basename>.block<rank>.<NNNNN>.restart) with no barrier or serial
+    # bundling -- recommended for large parallel runs on Lustre/GPFS scratch.
+    # Restart from either form with `--restart <file>` (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.

src/input/read_restart_file.cpp

@@ -1,4 +1,5 @@
 // C/C++
+#include <cctype>
 #include <filesystem>
 #include <fstream>
 #include <iomanip>
@@ -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<unsigned char>(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.

src/input/read_restart_file.hpp

@@ -18,4 +18,12 @@ using Variables = std::map<std::string, torch::Tensor>;
 
 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<N>." token in the filename;
+//! this rewrites <N> 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

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;

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) {

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: <basename>.<blockid>.<file_number>.part
-  std::string fname;
+  // shared stem: <output_dir>/<basename>.<blockid>.<file_number|final>
   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");
   }
 }
 

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)

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 ``<basename>.block<rank>.<num>.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 <basename>.block<rank>.<num>.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())