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

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 = (
-        f"{args.time_mode}_exponential_next_token_{args.target_mode}_"
-        f"gap_{args.no_event_interval_years:g}y_{timestamp}"
+    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}")