Refactor run directory creation to use a unique naming function in training scripts

train_all_future.py

@@ -15,7 +15,6 @@ import json
 import logging
 import math
 import time
-from datetime import datetime
 from pathlib import Path
 from typing import Any, Dict
 
@@ -32,6 +31,7 @@ from models import DeepHealth
 from targets import CHECKUP_IDX, PAD_IDX
 from train_util import (
     configure_torch_for_training,
+    create_unique_run_dir,
     load_extra_info_types_file,
     resolve_device,
     save_checkpoint,
@@ -295,10 +295,9 @@ def main() -> None:
     device = resolve_device(args.device)
     configure_torch_for_training(device)
 
-    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
-    run_name = f"{args.time_mode}_{args.dist_mode}_all_future_pure_disease_{timestamp}"
-    run_dir = Path("runs") / run_name
-    run_dir.mkdir(parents=True, exist_ok=False)
+    run_dir, run_name = create_unique_run_dir(
+        lambda timestamp: f"{args.time_mode}_{args.dist_mode}_all_future_pure_disease_{timestamp}"
+    )
     logger = setup_logging(run_dir)
 
     logger.info(f"Starting all-future training run: {run_name}")

train_next_step.py

@@ -12,7 +12,6 @@ import json
 import logging
 import math
 import time
-from datetime import datetime
 from pathlib import Path
 from typing import Any, Dict
 
@@ -30,6 +29,7 @@ from readouts import build_readout
 from targets import CHECKUP_IDX, NO_EVENT_IDX, PAD_IDX
 from train_util import (
     configure_torch_for_training,
+    create_unique_run_dir,
     load_extra_info_types_file,
     resolve_device,
     save_checkpoint,
@@ -319,13 +319,12 @@ def main() -> None:
     device = resolve_device(args.device)
     configure_torch_for_training(device)
 
-    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
-    run_name = (
+    run_dir, run_name = create_unique_run_dir(
+        lambda timestamp: (
             f"{args.time_mode}_exponential_next_token_{args.target_mode}_"
             f"gap_{args.no_event_interval_years:g}y_{timestamp}"
         )
-    run_dir = Path("runs") / run_name
-    run_dir.mkdir(parents=True, exist_ok=False)
+    )
     logger = setup_logging(run_dir)
 
     logger.info(f"Starting next-step training run: {run_name}")

train_util.py

@@ -3,6 +3,8 @@ from __future__ import annotations
 import json
 import logging
 import sys
+import time
+from datetime import datetime
 from pathlib import Path
 from typing import Any, Dict, Iterable, Tuple
 
@@ -15,6 +17,18 @@ from dataset import AllFutureHealthDataset, HealthDataset
 from models import DeepHealth
 
 
+def create_unique_run_dir(name_fn, runs_root: Path = Path("runs")) -> tuple[Path, str]:
+    while True:
+        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+        run_name = name_fn(timestamp)
+        run_dir = runs_root / run_name
+        try:
+            run_dir.mkdir(parents=True, exist_ok=False)
+            return run_dir, run_name
+        except FileExistsError:
+            time.sleep(1.0)
+
+
 def setup_logging(run_dir: Path) -> logging.Logger:
     run_dir.mkdir(parents=True, exist_ok=True)
     logger = logging.getLogger("DeepHealth")