add ram logging

README.md

@@ -76,9 +76,17 @@ do not collide.
 Each run automatically:
 - evaluates on **both** WILDS splits: `id_test` (ID) and `test` (OOD), so the
   generalization gap can be measured;
-- records the WILDS metrics, the wall-clock **training time**, and the number of
+- records the WILDS metrics, the wall-clock **training time**, the number of
-  **epochs run** (accounting for early stopping) into the per-split metrics JSON
+  **epochs run** (accounting for early stopping), and the **effective memory
-  and into `params.yaml` in the eval folder.
+  usage** into the per-split metrics JSON and into `params.yaml` in the eval
+  folder.
+
+Effective memory is captured as a high-water mark during training:
+- `peak_host_ram_gb`: peak process RSS (`resource.getrusage`), to compare
+  against the SLURM `mem_per_gpu` request (e.g. 180G) and right-size future jobs.
+  With `tasks_per_node: 1` this reflects the whole training worker.
+- `peak_gpu_alloc_gb` / `peak_gpu_reserved_gb`: peak GPU VRAM
+  (`torch.cuda.max_memory_allocated` / `max_memory_reserved`).
 
 The four leaderboard columns are: Test ID Macro F1, Test ID Avg Acc,
 Test OOD Macro F1, Test OOD Avg Acc (headline metric: `F1-macro_all`).
@@ -107,8 +115,9 @@ python3 tools/aggregate_seeds.py --root experiment_logs/eval-wilds
 
 Outputs:
 - `experiment_logs/seed-runs/<model>/summary.json` (per-seed rows + mean/std for
-  all metrics, training time, epochs, and ID-OOD generalization gap)
+  all metrics, training time, epochs, peak memory, and ID-OOD generalization gap)
-- `experiment_logs/seed-runs/summary_all.csv` (one row per model, paper-ready)
+- `experiment_logs/seed-runs/summary_all.csv` (one row per model, paper-ready;
+  includes `peak_host_ram_gb_mean/std` and `peak_gpu_alloc_gb_mean/std`)
 
 ## License
 

src/train_supervised.py

@@ -1,5 +1,6 @@
 import os
 import random
+import resource
 import shutil
 import sys
 import json
@@ -53,6 +54,33 @@ def _format_hms(seconds):
     return f"{h:02d}:{m:02d}:{s:02d}"
 
 
+def _peak_host_ram_gb():
+    """Peak resident set size (RSS) of this process, in GB.
+
+    Uses resource.getrusage(RUSAGE_SELF).ru_maxrss, which on Linux is reported
+    in kilobytes. This is the process high-water mark; with tasks_per_node=1 it
+    reflects the whole training worker. Compare against the SLURM mem request
+    (e.g. 180G) to right-size future jobs.
+    """
+    try:
+        maxrss_kb = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
+        return float(maxrss_kb) / (1024.0 * 1024.0)
+    except (ValueError, OSError):
+        return None
+
+
+def _peak_gpu_mem_gb(device):
+    """Peak allocated and reserved GPU memory (GB) since the last reset."""
+    if not torch.cuda.is_available():
+        return None, None
+    try:
+        alloc = torch.cuda.max_memory_allocated(device) / 1e9
+        reserved = torch.cuda.max_memory_reserved(device) / 1e9
+        return float(alloc), float(reserved)
+    except (RuntimeError, ValueError):
+        return None, None
+
+
 def strip_module_prefix(state_dict):
     if not any(k.startswith("module.") for k in state_dict.keys()):
         return state_dict
@@ -446,6 +474,9 @@ def main(args, resume_preempt=False):
             if is_best:
                 torch.save(save_dict, best_path)
 
+    if torch.cuda.is_available():
+        torch.cuda.reset_peak_memory_stats(device)
+
     train_start_time = time.perf_counter()
     epochs_run = 0
     early_stopped = False
@@ -535,12 +566,17 @@ def main(args, resume_preempt=False):
             break
 
     train_time_seconds = time.perf_counter() - train_start_time
+    peak_host_ram_gb = _peak_host_ram_gb()
+    peak_gpu_alloc_gb, peak_gpu_reserved_gb = _peak_gpu_mem_gb(device)
     if rank == 0:
         logger.info(
             f"Training finished: epochs_run={epochs_run} "
             f"early_stopped={early_stopped} "
             f"train_time={_format_hms(train_time_seconds)} "
-            f"({train_time_seconds:.1f}s)"
+            f"({train_time_seconds:.1f}s) "
+            f"peak_host_ram_gb={peak_host_ram_gb} "
+            f"peak_gpu_alloc_gb={peak_gpu_alloc_gb} "
+            f"peak_gpu_reserved_gb={peak_gpu_reserved_gb}"
         )
 
     if early_stopper.enabled and early_stopper.restore_best_weights:
@@ -615,6 +651,9 @@ def main(args, resume_preempt=False):
             "best_epoch": int(early_stopper.best_epoch),
             "early_stopped": bool(early_stopped),
             "best_val_loss": float(early_stopper.best_metric),
+            "peak_host_ram_gb": peak_host_ram_gb,
+            "peak_gpu_alloc_gb": peak_gpu_alloc_gb,
+            "peak_gpu_reserved_gb": peak_gpu_reserved_gb,
         }
 
         # Fold run_info into each split's metrics JSON so an aggregator can read
@@ -650,6 +689,9 @@ def main(args, resume_preempt=False):
                     "epochs_run": int(epochs_run),
                     "configured_epochs": int(o_args["epochs"]),
                     "early_stopped": bool(early_stopped),
+                    "peak_host_ram_gb": peak_host_ram_gb,
+                    "peak_gpu_alloc_gb": peak_gpu_alloc_gb,
+                    "peak_gpu_reserved_gb": peak_gpu_reserved_gb,
                     "eval_metrics_id_test": (
                         eval_results.get("id_test", {}).get("metrics")
                         if eval_results.get("id_test")

tools/aggregate_seeds.py

@@ -192,6 +192,12 @@ def main():
         "train_time_seconds_std",
         "epochs_run_mean",
         "epochs_run_std",
+        "peak_host_ram_gb_mean",
+        "peak_host_ram_gb_std",
+        "peak_gpu_alloc_gb_mean",
+        "peak_gpu_alloc_gb_std",
+        "peak_gpu_reserved_gb_mean",
+        "peak_gpu_reserved_gb_std",
     ]
 
     for group_name in sorted(groups.keys()):
@@ -228,6 +234,9 @@ def main():
                     "configured_epochs": r["run_info"].get("configured_epochs"),
                     "best_epoch": r["run_info"].get("best_epoch"),
                     "early_stopped": r["run_info"].get("early_stopped"),
+                    "peak_host_ram_gb": r["run_info"].get("peak_host_ram_gb"),
+                    "peak_gpu_alloc_gb": r["run_info"].get("peak_gpu_alloc_gb"),
+                    "peak_gpu_reserved_gb": r["run_info"].get("peak_gpu_reserved_gb"),
                 }
             )
 
@@ -245,8 +254,14 @@ def main():
 
         time_vals = [r["run_info"].get("train_time_seconds") for r in records]
         epoch_vals = [r["run_info"].get("epochs_run") for r in records]
+        host_ram_vals = [r["run_info"].get("peak_host_ram_gb") for r in records]
+        gpu_alloc_vals = [r["run_info"].get("peak_gpu_alloc_gb") for r in records]
+        gpu_reserved_vals = [r["run_info"].get("peak_gpu_reserved_gb") for r in records]
         time_mean, time_std, _ = _mean_std(time_vals)
         epoch_mean, epoch_std, _ = _mean_std(epoch_vals)
+        host_ram_mean, host_ram_std, _ = _mean_std(host_ram_vals)
+        gpu_alloc_mean, gpu_alloc_std, _ = _mean_std(gpu_alloc_vals)
+        gpu_reserved_mean, gpu_reserved_std, _ = _mean_std(gpu_reserved_vals)
 
         # Headline (leaderboard) numbers.
         id_f1_mean, id_f1_std, _ = _mean_std(metric_values("id_metrics", args.primary))
@@ -275,6 +290,14 @@ def main():
                 "train_time_seconds": {"mean": time_mean, "std": time_std},
                 "epochs_run": {"mean": epoch_mean, "std": epoch_std},
             },
+            "resources": {
+                "peak_host_ram_gb": {"mean": host_ram_mean, "std": host_ram_std},
+                "peak_gpu_alloc_gb": {"mean": gpu_alloc_mean, "std": gpu_alloc_std},
+                "peak_gpu_reserved_gb": {
+                    "mean": gpu_reserved_mean,
+                    "std": gpu_reserved_std,
+                },
+            },
             "id_metrics_aggregated": id_agg,
             "ood_metrics_aggregated": ood_agg,
             "per_seed": per_seed,
@@ -303,6 +326,12 @@ def main():
                 "train_time_seconds_std": time_std,
                 "epochs_run_mean": epoch_mean,
                 "epochs_run_std": epoch_std,
+                "peak_host_ram_gb_mean": host_ram_mean,
+                "peak_host_ram_gb_std": host_ram_std,
+                "peak_gpu_alloc_gb_mean": gpu_alloc_mean,
+                "peak_gpu_alloc_gb_std": gpu_alloc_std,
+                "peak_gpu_reserved_gb_mean": gpu_reserved_mean,
+                "peak_gpu_reserved_gb_std": gpu_reserved_std,
             }
         )
 
@@ -318,7 +347,8 @@ def main():
     header = (
         f"{'model':<40} {'seeds':>5}  "
         f"{'ID MacroF1':>18}  {'ID AvgAcc':>18}  "
-        f"{'OOD MacroF1':>18}  {'OOD AvgAcc':>18}  {'gap':>8}"
+        f"{'OOD MacroF1':>18}  {'OOD AvgAcc':>18}  {'gap':>8}  "
+        f"{'peakRAM_GB':>14}  {'peakVRAM_GB':>14}"
     )
     print(header)
     print("-" * len(header))
@@ -332,7 +362,9 @@ def main():
             f"{_fmt(row['id_avg_acc_mean'], row['id_avg_acc_std']):>18}  "
             f"{_fmt(row['ood_macro_f1_mean'], row['ood_macro_f1_std']):>18}  "
             f"{_fmt(row['ood_avg_acc_mean'], row['ood_avg_acc_std']):>18}  "
-            f"{gap_str:>8}"
+            f"{gap_str:>8}  "
+            f"{_fmt(row['peak_host_ram_gb_mean'], row['peak_host_ram_gb_std']):>14}  "
+            f"{_fmt(row['peak_gpu_alloc_gb_mean'], row['peak_gpu_alloc_gb_std']):>14}"
         )
 
     print(f"\nWrote per-model summaries to: {args.out}/<model>/summary.json")