ctm-dqn/utils.py

63 lines
1.8 KiB
Python

"""
Configuration utilities for loading and managing settings.
"""
import yaml
import os
import shutil
from datetime import datetime
from typing import Dict, Tuple
def load_config(config_path: str = "config.yaml") -> Dict:
"""Load configuration from YAML file."""
if not os.path.exists(config_path):
raise FileNotFoundError(f"Config file not found: {config_path}")
with open(config_path, "r") as f:
config = yaml.safe_load(f)
return config
def create_directories(config: Dict):
"""Create necessary directories for checkpoints and logs."""
checkpoint_dir = config["training"]["checkpoint_dir"]
log_dir = config["training"]["log_dir"]
os.makedirs(checkpoint_dir, exist_ok=True)
os.makedirs(log_dir, exist_ok=True)
def create_run_directory(config: Dict, config_path: str = "config.yaml") -> Tuple[str, str, str]:
"""
Create a timestamped run directory for this training session.
Args:
config: Configuration dictionary
config_path: Path to the config file
Returns:
Tuple of (run_dir, checkpoint_dir, log_dir)
"""
# Create timestamp
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
# Create run directory
run_dir = os.path.join("runs", f"run_{timestamp}")
os.makedirs(run_dir, exist_ok=True)
# Create subdirectories
checkpoint_dir = os.path.join(run_dir, "checkpoints")
log_dir = os.path.join(run_dir, "logs")
os.makedirs(checkpoint_dir, exist_ok=True)
os.makedirs(log_dir, exist_ok=True)
# Copy config file to run directory
config_copy_path = os.path.join(run_dir, "config.yaml")
shutil.copy(config_path, config_copy_path)
print(f"Created run directory: {run_dir}")
print(f"Config saved to: {config_copy_path}")
return run_dir, checkpoint_dir, log_dir