Ignore .env and wandb and other files during code archive creation.

Resolves #1.

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,38 @@ 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] = []):
+    """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):
     """Upload code to GCS that entrypoint should retrieve.
 
@@ -50,8 +84,11 @@ def upload_code(project_id: str, experiment_id: str):
     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(