From 7aafe1b3d3b6895ce5f6488b8b902b3b116c8263 Mon Sep 17 00:00:00 2001 From: root Date: Fri, 17 Jul 2026 15:31:02 +0200 Subject: [PATCH] restart: resolve entries under restart directory --- bin/mana_restart | 56 +++++++++++++++++++-------- ci/test-mana-restart.py | 85 +++++++++++++++++++++++++++++++++++++++++ ci/unit-test.sh | 3 ++ 3 files changed, 129 insertions(+), 15 deletions(-) create mode 100755 ci/test-mana-restart.py diff --git a/bin/mana_restart b/bin/mana_restart index 806d66c13..deeb3f5d9 100755 --- a/bin/mana_restart +++ b/bin/mana_restart @@ -30,22 +30,48 @@ if "--timing" in dmtcp_flags: os.environ["MANA_TIMING"] = "1" dmtcp_flags.remove("--timing") if "--restartdir" not in dmtcp_flags: - dmtcp_flags.append("--restartdir ./") -if "--restartdir" in dmtcp_flags: - ckptdir_path = dmtcp_flags[dmtcp_flags.index("--restartdir") + 1] - if not os.path.exists(ckptdir_path): - print("mana_restart: --restartdir " + ckptdir_path + - ": Restart directory doesn't exist") + dmtcp_flags.extend(["--restartdir", "./"]) + +restartdir_index = dmtcp_flags.index("--restartdir") + +if restartdir_index + 1 >= len(dmtcp_flags): + print("mana_restart: --restartdir requires a directory") + sys.exit(1) + +restartdir_path = os.path.abspath( + os.path.expanduser(dmtcp_flags[restartdir_index + 1]) +) + +if not os.path.isdir(restartdir_path): + print( + "mana_restart: --restartdir " + + restartdir_path + + ": Restart directory doesn't exist" + ) sys.exit(1) - for rank_dir in os.listdir(ckptdir_path): - if (os.path.isfile(rank_dir)): - continue - for fname in os.listdir(rank_dir): - if fname.endswith('.tmp'): - print("mana_restart: --restartdir " + ckptdir_path + - ": Restart directory has .tmp files. ") - print("Previous checkpoint may be incomplete.") - sys.exit(1) + +for rank_entry in os.scandir(restartdir_path): + if not rank_entry.is_dir(): + continue + for checkpoint_entry in os.scandir(rank_entry.path): + if checkpoint_entry.name.endswith(".tmp"): + print( + "mana_restart: --restartdir " + + restartdir_path + + ": Restart directory has .tmp files." + ) + print("Previous checkpoint may be incomplete.") + sys.exit(1) + +dmtcp_flags[restartdir_index + 1] = restartdir_path + +# lower-half replaces this pair with a rank-specific checkpoint image. +# Keep it after all other DMTCP options so the image is the final +# positional argument. +restartdir_pair = dmtcp_flags[restartdir_index:restartdir_index + 2] +del dmtcp_flags[restartdir_index:restartdir_index + 2] +dmtcp_flags.extend(restartdir_pair) + if "--ckptdir" in dmtcp_flags: ckptdir_path = dmtcp_flags[dmtcp_flags.index("--ckptdir") + 1] if not os.path.exists(ckptdir_path): diff --git a/ci/test-mana-restart.py b/ci/test-mana-restart.py new file mode 100755 index 000000000..fe977a056 --- /dev/null +++ b/ci/test-mana-restart.py @@ -0,0 +1,85 @@ +#!/usr/bin/env python3 +import os +import shutil +import subprocess +import tempfile +import unittest +from pathlib import Path + +ROOT = Path(__file__).resolve().parents[1] +WRAPPER = ROOT / "bin" / "mana_restart" + +class ManaRestartTest(unittest.TestCase): + def setUp(self): + self.temp = tempfile.TemporaryDirectory(prefix="mana-restart-test-") + self.root = Path(self.temp.name) + self.fake_bin = self.root / "fake-mana" / "bin" + self.fake_bin.mkdir(parents=True) + self.home = self.root / "home" + self.home.mkdir() + self.cwd = self.root / "unrelated" + self.cwd.mkdir() + self.restart = self.root / "restart" + (self.restart / "ckpt_rank_0").mkdir(parents=True) + (self.restart / "ckpt_rank_1").mkdir(parents=True) + (self.restart / "ckpt_rank_0" / "ckpt_0.dmtcp").touch() + (self.restart / "ckpt_rank_1" / "ckpt_1.dmtcp").touch() + self.capture = self.root / "args.txt" + shutil.copy2(WRAPPER, self.fake_bin / "mana_restart") + os.chmod(self.fake_bin / "mana_restart", 0o755) + self._exe(self.fake_bin / "dmtcp_command", "#!/bin/sh\nexit 0\n") + self._exe(self.fake_bin / "lower-half", '#!/bin/sh\nprintf "%s\\n" "$@" > "$MANA_TEST_CAPTURE"\nexit 0\n') + (self.home / ".mana.rc").write_text("Host: localhost\nPort: 7780\n", encoding="utf-8") + + def tearDown(self): + self.temp.cleanup() + + @staticmethod + def _exe(path, content): + path.write_text(content, encoding="utf-8") + os.chmod(path, 0o755) + + def run_wrapper(self, *args): + env = os.environ.copy() + env["HOME"] = str(self.home) + env["MANA_TEST_CAPTURE"] = str(self.capture) + env.pop("SLURM_JOB_ID", None) + return subprocess.run([str(self.fake_bin / "mana_restart"), *args], cwd=self.cwd, env=env, text=True, capture_output=True, check=False, timeout=30) + + def captured(self): + return self.capture.read_text(encoding="utf-8").splitlines() + + def test_outside_current_directory(self): + result = self.run_wrapper("--verbose", "--restartdir", str(self.restart)) + self.assertEqual(result.returncode, 0, result.stdout + result.stderr) + args = self.captured() + i = args.index("--restartdir") + self.assertEqual(args[i + 1], str(self.restart.resolve())) + + def test_tmp_detection(self): + (self.restart / "ckpt_rank_1" / "incomplete.tmp").touch() + result = self.run_wrapper("--restartdir", str(self.restart)) + self.assertNotEqual(result.returncode, 0) + self.assertIn("Restart directory has .tmp files", result.stdout + result.stderr) + + def test_missing_value(self): + result = self.run_wrapper("--restartdir") + self.assertNotEqual(result.returncode, 0) + self.assertIn("requires a directory", result.stdout + result.stderr) + + def test_restartdir_after_other_options(self): + result = self.run_wrapper("--restartdir", str(self.restart), "--ckptdir", str(self.restart)) + self.assertEqual(result.returncode, 0, result.stdout + result.stderr) + args = self.captured() + self.assertGreater(args.index("--restartdir"), args.index("--ckptdir")) + + def test_default_current_directory(self): + (self.cwd / "ckpt_rank_0").mkdir() + result = self.run_wrapper() + self.assertEqual(result.returncode, 0, result.stdout + result.stderr) + args = self.captured() + i = args.index("--restartdir") + self.assertEqual(args[i + 1], str(self.cwd.resolve())) + +if __name__ == "__main__": + unittest.main() diff --git a/ci/unit-test.sh b/ci/unit-test.sh index 31d424764..b269bcbbc 100644 --- a/ci/unit-test.sh +++ b/ci/unit-test.sh @@ -21,3 +21,6 @@ cd $SCRIPT_DIR/../mpi-proxy-split/unit-test make || exit 1 make clean make check || exit 1 + +# Additional wrapper regression test. +python3 ci/test-mana-restart.py -v || exit 1