Refactor wandb parameters structure and update script to include params_wrapper for improved organization

Author: RichieHakim

batch_run/dispatcher.py

@@ -35,13 +35,15 @@
 
 
 params_template = {
-    "kwargs_wandb_init": {
-        "project": "face_rhythm",
-        "name": name_wandb,
-        "dir_save": dir_save,
-        "reinit": True,
+    "params_wrapper": {
+        "kwargs_wandb_init": {
+            "project": "face_rhythm",
+            "name": name_wandb,
+            "dir": dir_save,
+        },
+        "path_script": path_script,
+        "period_logger": 2,
     },
-    "path_script": path_script,
 
     "params_script": {
         "steps": [
@@ -269,9 +271,6 @@
 
 ## Prepare call str and set environment variables
 os.environ['WANDB_API_KEY'] = wandb_api_key
-# Initialize a wandb run with minimal settings.
-import wandb
-wandb.init(project='face_rhythm', name=name_wandb, config=params[0], reinit=True)
 
 ## define slurm SBATCH parameters
 # sbatch_config_list = \
@@ -320,12 +319,12 @@
 #SBATCH --account=kempner_bsabatini_lab  # The account name for the job.
 #SBATCH --job-name={name_slurm}          # Job name
 #SBATCH --output={path}                  # File to write: STDOUT (and STDERR if --error is not used)
-#SBATCH --partition=kempner_requeue      # Partition (job queue)
+#SBATCH --partition=kempner_requeue              # Partition (job queue)
 #SBATCH --gres=gpu:1                     # Number of GPUs
 #SBATCH -c 16                            # Number of cores (-c) on one node
 #SBATCH -n 1                             # Number of nodes (-n)
-#SBATCH --mem=64GB                       # Memory pool for all cores (see also --mem-per-cpu)
-#SBATCH --time=0-1:00:00                 # Runtime in D-HH:MM:SS
+#SBATCH --mem=128GB                      # Memory pool for all cores (see also --mem-per-cpu)
+#SBATCH --time=0-0:10:00                 # Runtime in D-HH:MM:SS
 #SBATCH --requeue                        # Requeue the job if it is preempted
 #SBATCH --export=WANDB_API_KEY           # Export the WANDB_API_KEY environment variable to the job
 
@@ -335,7 +334,7 @@
 echo "activating environment"
 source activate {name_env}
 
-echo "starting job with call: python $@"
+echo "starting job"
 python "$@"
 """ for path in paths_log]
 

scripts/wandb_script_wrapper.py

@@ -12,11 +12,12 @@
   1. Parses command-line arguments, expecting two positional arguments:
        - path_params: Path to a JSON file containing parameters for the target script. This
          JSON must include at minimum:
-           * "path_script": String path to the target script to be executed.
-           * "kwargs_wandb_init": (Optional) A dictionary of keyword arguments for 
-             wandb.init().
-           * "params_script" can be provided, which will be saved as a JSON to the output
-             directory and --path_params will be passed to the target script.
+            * "params_wrapper": Containing the following:
+                * "path_script": String path to the target script to be executed.
+                * "kwargs_wandb_init": (Optional) A dictionary of keyword arguments for 
+                    wandb.init().
+            * "params_script" can be provided, which will be saved as a JSON to the output
+                directory and --path_params will be passed to the target script.
        - directory_save: Directory path where output files (such as logs and saved parameters)
          will be stored. Will be passed to the target script as --directory_save.
 
@@ -37,11 +38,14 @@
   "/path/to/save_dir". The parameter JSON file should include entries like:
 
     {
-        "path_script": "/path/to/your_target_script.py",
-        "kwargs_wandb_init": {
-            "project": "face_rhythm",
-            "entity": "your_wandb_username",
-            "name": "example_run"
+        "params_wrapper": {
+            "path_script": "/path/to/your_target_script.py",
+            "kwargs_wandb_init": {
+                "project": "face_rhythm",
+                "entity": "your_wandb_username",
+                "name": "example_run"
+            },
+            "period_logger": 2,
         },
         "params_script": {
             "example_param": "value"
@@ -62,6 +66,8 @@
 import time
 import wandb
 import psutil
+import functools
+
 
 def stream_reader(pipe, log_label):
     """
@@ -113,11 +119,15 @@ def monitor_system_metrics(interval: int = 30):
     import json
     with open(path_params, 'r') as f:
         params = json.load(f)
+        
+    # Get sub parameters for wrapper
+    assert 'params_wrapper' in params, "Error: 'params_wrapper' is missing in the parameters file."
+    params_wrapper = params['params_wrapper']
 
     # Gather kwargs_wandb_init from the JSON file.
-    kwargs_wandb_init = params.get('kwargs_wandb_init', None)
+    kwargs_wandb_init = params_wrapper.get('kwargs_wandb_init', None)
     # Gather path_script from the JSON file. Error if missing
-    path_script = params.get('path_script', None)
+    path_script = params_wrapper.get('path_script', None)
     if path_script is None:
         print("Error: 'path_script' is missing in the parameters file.")
         sys.exit(1)
@@ -134,6 +144,9 @@ def monitor_system_metrics(interval: int = 30):
             json.dump(params_script, f)
     else:
         print("Warning: 'params_script' is not provided in the parameters file. Skipping saving parameters.")
+        
+    # Prepare call to monitor_system_metrics.
+    monitor_system_metrics = functools.partial(monitor_system_metrics, interval=params_wrapper.get('period_logger', 30))
 
     # Ensure WandB is installed.
     try: