import argparse import re import os import datetime import torch import random import yaml import wandb import shutil import copy import time import torch.distributed as dist import numpy as np from pathlib import Path from omegaconf import OmegaConf from easydict import EasyDict as edict def process_overrides(overrides): """Normalize CLI override strings so that spaces around '=' are removed. Args: overrides (list[str]): Raw override tokens from argparse, which may contain spaces around the '=' separator. Returns: list[str]: Overrides reformatted as 'key=value' strings. """ combined = ' '.join(overrides) # Use regex to identify and fix patterns like 'param = value' to 'param=value' fixed_string = re.sub(r'(\S+)\s*=\s*(\S+)', r'\1=\2', combined) # Split the fixed string back into a list, preserving properly formatted args processed = re.findall(r'[^\s=]+=\S+|\S+', fixed_string) return processed def init_config(): """Parse command-line arguments, load the YAML config, apply CLI overrides, and return it. Returns: edict: Merged and resolved configuration as an EasyDict. """ parser = argparse.ArgumentParser() parser.add_argument("--config", "-c", required=True) parser.add_argument("overrides", nargs="*") # Capture all "key=value" args args = parser.parse_args() config = OmegaConf.load(args.config) processed_overrides = process_overrides(args.overrides) cli_overrides = OmegaConf.from_cli(processed_overrides) # Merge configs (with type-safe automatic conversion) config = OmegaConf.merge(config, cli_overrides) config = OmegaConf.to_container(config, resolve=True) config = edict(config) return config def init_distributed(seed=42): """Initialize distributed training environment and set random seeds for reproducibility. Args: seed (int): Base random seed. Each process derives its own seed as seed + global_rank to ensure different random states per worker. Returns: edict: Dictionary with attribute access containing: - local_rank: GPU rank within the current node - global_rank: Global rank of the process - world_size: Total number of processes - device: The CUDA device assigned to this process - is_main_process: Flag to identify the main process - seed: The random seed used for this process """ global_rank = int(os.environ["RANK"]) world_size = int(os.environ["WORLD_SIZE"]) local_rank = int(os.environ["LOCAL_RANK"]) dist.init_process_group( backend="nccl", timeout=datetime.timedelta(seconds=3600) ) device = torch.device(f"cuda:{local_rank}") torch.cuda.set_device(device) # Each process gets a different seed derived from the base seed process_seed = seed + global_rank torch.manual_seed(process_seed) torch.cuda.manual_seed(process_seed) torch.cuda.manual_seed_all(process_seed) np.random.seed(process_seed) random.seed(process_seed) # Use deterministic algorithms and disable benchmarking for stability torch.backends.cudnn.benchmark = False torch.backends.cudnn.deterministic = True return edict({ 'local_rank': local_rank, 'global_rank': global_rank, 'world_size': world_size, 'device': device, 'is_main_process': global_rank == 0, 'seed': process_seed }) def local_backup_src_code( src_dir, dst_dir, max_size_MB=4.0, extension_to_backup=(".py", ".yaml", ".sh", ".bash", ".json"), exclude_dirs=("wandb", ".git", "checkpoints", "experiments"), verbose=True, ): """Back up source code files from src_dir to dst_dir, enforcing a total size limit. Args: src_dir: Source directory to backup. dst_dir: Destination directory for backups. max_size_MB (float): Maximum total size allowed for backup in MB. extension_to_backup (tuple[str]): File extensions to include in backup. exclude_dirs (tuple[str]): Directories to exclude from backup. verbose (bool): Whether to print progress information. Returns: tuple[int, int]: (num_files_backed_up, total_size_in_bytes). Raises: ValueError: If total size exceeds max_size_MB. """ start_time = time.time() src_path = Path(src_dir).resolve() dst_path = Path(dst_dir).resolve() extension_set = set(extension_to_backup) ignore_paths = {(src_path / d).resolve() for d in exclude_dirs} max_bytes = int(max_size_MB * 1024 * 1024) if not src_path.exists(): raise FileNotFoundError(f"Source directory does not exist: {src_path}") files = [] total_size = 0 for dirpath, dirnames, filenames in os.walk(src_path): current_path = Path(dirpath).resolve() if any(parent in ignore_paths for parent in current_path.parents) or current_path in ignore_paths: dirnames.clear() continue for filename in filenames: file_ext = os.path.splitext(filename)[1] if file_ext not in extension_set: continue src_file = current_path / filename rel_path = current_path.relative_to(src_path) dst_file = dst_path / rel_path / filename try: file_size = src_file.stat().st_size total_size += file_size files.append((src_file, dst_file, file_size)) except (FileNotFoundError, PermissionError) as e: if verbose: print(f"Warning: Could not access {src_file}: {e}") if total_size > max_bytes: if verbose: print(f"Size limit exceeded: {total_size / (1024*1024):.2f} MB > {max_size_MB} MB") print("Largest files:") for src_file, _, size in sorted(files, key=lambda x: x[2], reverse=True)[:5]: print(f"{src_file}: {size / 1024:.1f} KB") raise ValueError(f"Size limit exceeded: {total_size / (1024*1024):.2f} MB > {max_size_MB} MB") if verbose: print(f"Backing up {len(files)} files ({total_size / (1024*1024):.2f} MB)") dst_path.mkdir(parents=True, exist_ok=True) successful_copies = 0 for src_file, dst_file, _ in files: try: dst_file.parent.mkdir(parents=True, exist_ok=True) shutil.copy2(src_file, dst_file) successful_copies += 1 except Exception as e: if verbose: print(f"Error copying {src_file} to {dst_file}: {e}") elapsed_time = time.time() - start_time if verbose: print(f"Backup completed: {successful_copies}/{len(files)} files copied in {elapsed_time:.2f} seconds") return successful_copies, total_size def init_wandb_and_backup(config): """Initialize W&B, back up source code, and save the resolved config to disk. Args: config (edict): Resolved training configuration. Must contain config.training.api_key_path, config.training.wandb_project, config.training.wandb_exp_name, and config.training.checkpoint_dir. """ assert os.path.exists( config.training.api_key_path ), f"API key file does not exist: {config.training.api_key_path}" api_keys = edict(yaml.safe_load(open(config.training.api_key_path, "r"))) assert api_keys.wandb is not None, "Wandb API key not found in api key file" os.environ["WANDB_API_KEY"] = api_keys.wandb config_copy = copy.deepcopy(config) wandb.init( project=config.training.wandb_project, name=config.training.wandb_exp_name, config=config_copy, ) cur_dir = os.path.dirname(os.path.realpath(__file__)) trgt_dir = os.path.join(config.training.checkpoint_dir, "src", os.path.basename(cur_dir)) os.makedirs(trgt_dir, exist_ok=True) extension_to_backup = (".py", ".yaml", ".sh", ".bash", ".json") exclude_dirs = ("wandb", ".git", "checkpoints", "experiments") local_backup_src_code(cur_dir, trgt_dir, extension_to_backup=extension_to_backup, exclude_dirs=exclude_dirs) config_save_path = os.path.join(config.training.checkpoint_dir, "config.yaml") with open(config_save_path, 'w') as f: yaml.dump(dict(config), f) wandb.run.log_code( trgt_dir, include_fn=lambda path: path.endswith(extension_to_backup), )