Adding SnapshotManager

submitit/__init__.py

@@ -16,5 +16,6 @@
 from .local.local import LocalJob as LocalJob
 from .slurm.slurm import SlurmExecutor as SlurmExecutor
 from .slurm.slurm import SlurmJob as SlurmJob
+from .snapshot import SnapshotManager
 
 __version__ = "1.1.5"

submitit/snapshot.py

@@ -0,0 +1,34 @@
+import fnmatch
+import os
+import shutil
+from subprocess import check_output
+from tempfile import NamedTemporaryFile
+
+def run_cmd(str_args, **kwargs):
+    return check_output(str_args, shell=True, **kwargs).decode("utf-8").strip()
+
+
+class SnapshotManager:
+    def __init__(
+        self, snapshot_dir: str, with_submodules: bool = False, exclude=None,
+    ):
+        self.snapshot_dir = snapshot_dir
+        self.original_dir = os.getcwd()
+        self.with_submodules = with_submodules
+        self.exclude = exclude or []
+
+    def __enter__(self):
+        self.original_dir = os.getcwd()
+        os.makedirs(self.snapshot_dir, exist_ok=True)
+        # Get the repository root
+        root_dir = run_cmd("git rev-parse --show-toplevel")
+        sub = "--recurse-submodules" if self.with_submodules else "-s"
+        # Get a list of all the checked in files that we can pass to rsync
+        with NamedTemporaryFile() as tfile:
+            # https://stackoverflow.com/a/51689219/4876946
+            files = run_cmd(f"git ls-files {sub} | grep -v ^16 | cut -f2- > {tfile.name}", cwd=root_dir).split("\n")
+            run_cmd(f"rsync -a --files-from={tfile.name} {root_dir} {self.snapshot_dir}")
+        os.chdir(self.snapshot_dir)
+
+    def __exit__(self, *args):
+        os.chdir(self.original_dir)

submitit/test_snapshot.py

@@ -0,0 +1,13 @@
+import os
+from tempfile import TemporaryDirectory
+
+from .snapshot import SnapshotManager
+
+
+def test_snapshot():
+    cwd = os.getcwd()
+    with TemporaryDirectory() as tdir:
+        with SnapshotManager(tdir):
+            assert os.getcwd() == tdir
+            assert os.path.exists(os.path.join(tdir, "submitit/test_snapshot.py"))
+    assert os.getcwd() == cwd