diff --git a/train_all_future.py b/train_all_future.py index e8d96ea..9d35ca5 100644 --- a/train_all_future.py +++ b/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}") diff --git a/train_next_step.py b/train_next_step.py index 71957c6..718bc9e 100644 --- a/train_next_step.py +++ b/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}") diff --git a/train_util.py b/train_util.py index b04fd45..c7e9830 100644 --- a/train_util.py +++ b/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")