Merge pull request #254 from katzlabbrandeis/242-log-ram-usage

242 log ram usage during blech_run_process

Author: abuzarmahmood

blech_process.py

@@ -31,7 +31,6 @@
 
 from utils.blech_utils import imp_metadata, pipeline_graph_check
 import utils.blech_process_utils as bpu
-from utils import memory_monitor as mm
 import pylab as plt
 import json
 import sys
@@ -230,10 +229,6 @@
         cluster_handler.create_classifier_plots(classifier_handler)
 
 
-# Make file for dumping info about memory usage
-f= open(f'./memory_monitor_clustering/{electrode_num:02}.txt', 'w')
-print(mm.memory_usage_resource(), file=f)
-f.close()
 print(f'Electrode {electrode_num} complete.')
 
 # Write successful execution to log

blech_run_process.sh

@@ -40,8 +40,15 @@ if [ ! -d "$DIR" ]; then
     fi
 fi
 
+# Start RAM monitoring in background
+python3 utils/ram_monitor.py "$DIR" &
+RAM_MONITOR_PID=$!
+
 echo "Processing $DIR"
 for i in {1..10}; do
     echo Retry $i
     bash $DIR/temp/blech_process_parallel.sh
 done
+
+# Kill RAM monitor when done
+kill $RAM_MONITOR_PID

requirements/pip_requirements_base.txt

@@ -5,3 +5,4 @@ prefect==2.10.5
 pingouin==0.3
 tqdm
 gdown
+psutil>=6.1.0

utils/memory_monitor.py

@@ -0,0 +1,36 @@
+#!/usr/bin/env python3
+import psutil
+import time
+import sys
+import os
+from datetime import datetime
+
+def monitor_ram(output_dir):
+    """Monitor RAM usage and write to a log file"""
+    log_file = os.path.join(output_dir, "ram_usage.log")
+    
+    with open(log_file, 'a') as f:
+        f.write("Timestamp,RAM_Used_GB,RAM_Total_GB,RAM_Percent\n")
+        
+        while True:
+            try:
+                mem = psutil.virtual_memory()
+                used_gb = mem.used / (1024 * 1024 * 1024)  # Convert to GB
+                total_gb = mem.total / (1024 * 1024 * 1024)
+                timestamp = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
+                
+                log_line = f"{timestamp},{used_gb:.2f},{total_gb:.2f},{mem.percent}\n"
+                f.write(log_line)
+                f.flush()
+                
+                time.sleep(1)  # Log every 1 seconds
+                
+            except KeyboardInterrupt:
+                break
+
+if __name__ == "__main__":
+    if len(sys.argv) != 2:
+        print("Usage: python ram_monitor.py <output_directory>")
+        sys.exit(1)
+    
+    monitor_ram(sys.argv[1])