restart: resolve entries under --restartdir

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

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

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