Merge pull request #51 from allenai/favyen/issue1

Ignore files when creating code archive to launch Beaker job

Author: favyen2

rslp/launcher_lib.py

@@ -4,12 +4,14 @@
 import os
 import shutil
 import tempfile
+import zipfile
 
 import yaml
 from google.cloud import storage
 
 CODE_BLOB_PATH = "projects/{project_id}/{experiment_id}/code.zip"
 WANDB_ID_BLOB_PATH = "projects/{project_id}/{experiment_id}/wandb_id"
+CODE_EXCLUDES = [".env", "wandb", "rslp/__pycache__"]
 
 bucket = None
 
@@ -38,6 +40,40 @@ def get_project_and_experiment(config_path: str) -> tuple[str, str]:
     return project_id, experiment_id
 
 
+def make_archive(
+    zip_filename: str, root_dir: str, exclude_prefixes: list[str] = []
+) -> None:
+    """Create a zip archive of the contents of root_dir.
+
+    The paths in the zip archive will be relative to root_dir.
+
+    This is similar to shutil.make_archive but it allows specifying a list of prefixes
+    that should not be added to the zip archive.
+
+    Args:
+        zip_filename: the filename to save archive under.
+        root_dir: the directory to create archive of.
+        exclude_prefixes: a list of prefixes to exclude from the archive. If the
+            relative path of a file from root_dir starts with one of the prefixes, then
+            it will not be added to the resulting archive.
+    """
+
+    def should_exclude(rel_path: str) -> bool:
+        for prefix in exclude_prefixes:
+            if rel_path.startswith(prefix):
+                return True
+        return False
+
+    with zipfile.ZipFile(zip_filename, "w", zipfile.ZIP_DEFLATED) as zipf:
+        for root, _, files in os.walk(root_dir):
+            for fname in files:
+                full_path = os.path.join(root, fname)
+                rel_path = os.path.relpath(full_path, start=root_dir)
+                if should_exclude(rel_path):
+                    continue
+                zipf.write(full_path, arcname=rel_path)
+
+
 def upload_code(project_id: str, experiment_id: str) -> None:
     """Upload code to GCS that entrypoint should retrieve.
 
@@ -50,8 +86,11 @@ def upload_code(project_id: str, experiment_id: str) -> None:
     bucket = _get_bucket()
     with tempfile.TemporaryDirectory() as tmpdirname:
         print("creating archive of current code state")
-        zip_fname = shutil.make_archive(
-            os.path.join(tmpdirname, "archive"), "zip", root_dir="."
+        zip_fname = os.path.join(tmpdirname, "archive.zip")
+        make_archive(
+            zip_fname,
+            root_dir=".",
+            exclude_prefixes=CODE_EXCLUDES,
         )
         print("uploading archive")
         blob_path = CODE_BLOB_PATH.format(
@@ -73,7 +112,7 @@ def download_code(project_id: str, experiment_id: str) -> None:
     """
     bucket = _get_bucket()
     with tempfile.TemporaryDirectory() as tmpdirname:
-        print("downloading code acrhive")
+        print("downloading code archive")
         blob_path = CODE_BLOB_PATH.format(
             project_id=project_id, experiment_id=experiment_id
         )

tests/unit/test_launcher_lib.py

@@ -0,0 +1,32 @@
+import pathlib
+import zipfile
+
+from rslp.launcher_lib import make_archive
+
+
+def test_make_archive(tmp_path: pathlib.Path) -> None:
+    # Make sure make_archive correctly ignores the passed prefixes.
+    # We make sure it works with exactly matching file as well as a subdirectory.
+    exclude_prefixes = [
+        "ignored_file",
+        "dir/ignored_subdir",
+    ]
+    root_dir = tmp_path / "root"
+    (root_dir / "dir" / "ignored_subdir").mkdir(parents=True, exist_ok=True)
+    (root_dir / "dir" / "okay_subdir").mkdir(parents=True, exist_ok=True)
+    (root_dir / "okay_file1").touch()
+    (root_dir / "ignored_file").touch()
+    (root_dir / "dir" / "okay_file2").touch()
+    (root_dir / "dir" / "ignored_subdir" / "also_ignored").touch()
+    (root_dir / "dir" / "okay_subdir" / "okay_file3").touch()
+
+    zip_fname = str(tmp_path / "archive.zip")
+
+    make_archive(zip_fname, str(root_dir), exclude_prefixes=exclude_prefixes)
+    with zipfile.ZipFile(zip_fname) as zipf:
+        members = zipf.namelist()
+
+    assert "okay_file1" in members
+    assert "dir/okay_file2" in members
+    assert "dir/okay_subdir/okay_file3" in members
+    assert len(members) == 3