separate archiving from terminating ec2 for more frequent output dir sync with s3; save test metrics too

Author: codinglabsong

src/aging_gan/train.py

@@ -4,7 +4,8 @@
 import torch.nn as nn
 import torch.optim as optim
 import wandb
-import time
+import uuid
+import json
 from accelerate import Accelerator
 from torch.optim.lr_scheduler import LambdaLR
 from tqdm import tqdm
@@ -21,7 +22,7 @@
 )
 from aging_gan.data import prepare_dataset
 from aging_gan.model import initialize_models
-from aging_gan.utils import archive_and_terminate
+from aging_gan.utils import terminate_ec2, archive_ec2
 
 logger = logging.getLogger(__name__)
 
@@ -106,7 +107,7 @@ def parse_args() -> argparse.Namespace:
     p.add_argument(
         "--num_workers",
         type=int,
-        default=2,
+        default=3,
         help="Number of workers for dataloaders.",
     )
     p.add_argument(
@@ -120,6 +121,12 @@ def parse_args() -> argparse.Namespace:
         action="store_true",
         help="Upload outputs/ into s3 bucket and terminate current ec2.",
     )
+    p.add_argument(
+        "--s3_bucket_name",
+        type=str,
+        default="aging-gan",
+        help="Name of your s3 bucket to sync outputs. Effective only when --archive_and_terminate_ec2.",
+    )
 
     p.add_argument("--wandb_project", type=str, default="aging-gan")
 
@@ -160,7 +167,9 @@ def initialize_optimizers(cfg, G, F, DX, DY):
 
 
 def initialize_loss_functions(
-    lambda_adv_value: float = 2.0, lambda_cyc_value: float = 10.0, lambda_id_value: float = 7.0
+    lambda_adv_value: float = 2.0,
+    lambda_cyc_value: float = 10.0,
+    lambda_id_value: float = 7.0,
 ):
     mse = nn.MSELoss()
     l1 = nn.L1Loss()
@@ -500,8 +509,10 @@ def main() -> None:
 
     # ---------- Run Initialization ----------
     # wandb
+    run_name = uuid.uuid4().hex[:8]
     wandb.init(
         project=cfg.wandb_project,
+        name=run_name,
         config={
             k: v for k, v in vars(cfg).items() if not k.startswith("_")
         },  # drop python's or "private" framework-internal attributes
@@ -644,6 +655,9 @@ def main() -> None:
                 sched_DY,  # schedulers
                 "best",
             )
+            # upload outputs to s3 bucket
+            if cfg.archive_and_terminate_ec2:
+                archive_ec2(bucket=cfg.s3_bucket_name, prefix=f"outputs/run-{run_name}")
         # save the latest checkpoint
         if epoch % 5 == 0:
             save_checkpoint(
@@ -662,6 +676,9 @@ def main() -> None:
                 sched_DY,  # schedulers
                 "current",
             )
+            # upload outputs to s3 bucket
+            if cfg.archive_and_terminate_ec2:
+                archive_ec2(bucket=cfg.s3_bucket_name, prefix=f"outputs/run-{run_name}")
 
     # ---------- Test ----------
     if cfg.do_test:
@@ -702,17 +719,25 @@ def main() -> None:
         )
         logger.info(f"Test metrics (best.pth):\n{test_metrics}")
         wandb.log(test_metrics)
+        # write metrics out
+        out_dir = Path(__file__).resolve().parents[2] / "outputs" / "metrics"
+        out_dir.mkdir(parents=True, exist_ok=True)
+        metrics_path = out_dir / "test_metrics.json"
+        with metrics_path.open("w") as f:
+            json.dump(test_metrics, f, indent=4)
+        print(f"Saved test metrics to {metrics_path}")
+        # upload outputs to s3 bucket
+        if cfg.archive_and_terminate_ec2:
+            archive_ec2(bucket=cfg.s3_bucket_name, prefix=f"outputs/run-{run_name}")
     else:
         logger.info("Skipping test evaluation...")
 
     # Finished
     logger.info("Finished run.")
 
-    # upload outputs to s3 bucket and terminate ec2 instance
+    # terminate ec2 instance
     if cfg.archive_and_terminate_ec2:
-        archive_and_terminate(
-            bucket="aging-gan", prefix=f"outputs/run-{int(time.time())}"
-        )
+        terminate_ec2()
 
 
 if __name__ == "__main__":

src/aging_gan/utils.py

@@ -5,6 +5,7 @@
 import numpy as np
 import torch
 import matplotlib.pyplot as plt
+import subprocess
 import boto3
 import time
 from dotenv import load_dotenv
@@ -103,7 +104,7 @@ def generate_and_save_samples(
     device: torch.device,
     num_samples: int = 8,
 ):
-    # grab batches until num_samples   
+    # grab batches until num_samples
     collected = []
     for imgs, _ in val_loader:
         collected.append(imgs)
@@ -112,9 +113,9 @@ def generate_and_save_samples(
 
     if not collected:
         raise ValueError("Validation loader is empty.")
-    
+
     inputs = torch.cat(collected, dim=0)[:num_samples].to(device)
-    
+
     with torch.no_grad():
         outputs = generator(inputs)
 
@@ -149,29 +150,35 @@ def generate_and_save_samples(
     plt.close(fig)
 
 
-def archive_and_terminate(
+def archive_ec2(
     bucket: str,
-    prefix: str = "outputs/boto3/",
+    prefix: str = "outputs",
 ) -> None:
+    """Syncs everything under ./outputs to `s3://{bucket}/{prefix}/`."""
+    # Upload
+    out_root = Path(__file__).resolve().parents[2] / "outputs"
+    print(f"Uploading {out_root} -> s3://{bucket}/{prefix}/ ...")
+    cmd = [
+        "aws",
+        "s3",
+        "sync",
+        str(out_root),
+        f"s3://{bucket}/{prefix}",
+        "--only-show-errors",  # quieter logging
+    ]
+    subprocess.run(cmd, check=True)
+    print("S3 sync complete")
+
+
+def terminate_ec2() -> None:
     """
-    1. Recursively uploads everything under ./outputs to `s3://{bucket}/{prefix}/`.
-    2. Calls the EC2 API to terminate *this* instance.
+    Calls the EC2 API to terminate this instance.
 
     The instance must run with an IAM role that can:
         s3:PutObject   on   arn:aws:s3:::{bucket}/*
         ec2:TerminateInstances on itself (resource‑level ARN)
     """
-    # Upload
-    s3 = boto3.client("s3")
-    out_root = Path(__file__).resolve().parents[2] / "outputs/"
-    print(f"Uploading {out_root} -> s3://{bucket}/{prefix}/ ...")
-    for fp in out_root.rglob("*"):
-        if fp.is_file():
-            key = f"{prefix}/{fp.relative_to(out_root)}"
-            s3.upload_file(str(fp), bucket, key)
-    print("S3 sync complete")
-
-    # ---------- 2. Gather instance metadata (IMDSv2) ---------------------------------------
+    # Gather instance metadata (IMDSv2)
     token = requests.put(
         "http://169.254.169.254/latest/api/token",
         headers={"X-aws-ec2-metadata-token-ttl-seconds": "300"},
@@ -190,7 +197,7 @@ def archive_and_terminate(
     ).text
     print(f"Terminating {instance_id} in {region}")
 
-    # ---------- 3. Terminate self ----------------------------------------------------------
+    # Terminate self
     ec2 = boto3.client("ec2", region_name=region)
     ec2.terminate_instances(InstanceIds=[instance_id])
     print("Termination request sent - instance will shut down shortly")