import json import socket import sys import os import argparse import logging import time from datetime import datetime import lmdb import dotenv import torch from torch.utils.data import DataLoader import torch.distributed as dist from experiment_setup import get_eval_functions from utils.evaluation import evaluate_model from utils.model_utils import get_model_config from utils.training_utils import get_training_config, get_data_ids from utils.data_utils import LMDBIterableDataset from utils.utils import get_variable_from_module, get_logger, convert_for_json from utils.training import train_model print(f"PID: {os.getpid()} on host: {socket.gethostname()}") dotenv.load_dotenv() logger = None def prepare_run(results_dir: str, lmdb_root_dir: str, base_model_configuration: dict, base_training_configuration: dict): # create model configuration from base logger.info("Creating model configuration") model_configuration = get_model_config(base_model_configuration, results_dir, lmdb_root_dir) feature_config = model_configuration["feature_config"] dataset_dir = f"{lmdb_root_dir}/{feature_config['feature_set_name']}" logger.info(f"Model configuration name: {model_configuration['id']}") # create training configuration logger.info("Creating training configuration") training_configuration = get_training_config(base_training_configuration, model_configuration) logger.info(f"Training configuration: {training_configuration['id']}") logger.info(f"Fetching data ids, limit: {item_limit}") train_ids, val_ids, test_ids = get_data_ids(model_configuration, training_configuration, dataset_dir, item_limit) logger.info(f"Train ids: {len(train_ids)}, Val ids: {len(val_ids)}, Test ids: {len(test_ids)}") return model_configuration, training_configuration, train_ids, val_ids, test_ids, dataset_dir def train( model_configuration: dict, training_configuration: dict, train_ids: list, val_ids: list, dataset_dir: str, log_dir: str, num_dataloader_workers: int = 1) -> None: # create training and validation loaders logger.info("Creating training loaders") train_loader = LMDBIterableDataset(dataset_dir, train_ids, model_configuration=model_configuration, batch_size=training_configuration["batch_size"]) logger.info("Creating validation loaders") val_loader = LMDBIterableDataset(dataset_dir, val_ids, model_configuration=model_configuration, batch_size=training_configuration["batch_size"]) # instantiate model logger.info("Creating model") lmdb_env = lmdb.open(dataset_dir, readonly=True) model_creation_fn = model_configuration["model_creation_fn"] model = model_creation_fn(model_configuration, train_ids[0], lmdb_env=lmdb_env) logger.info(f"Started training for model {model_configuration['id']}") train_model( model=model, model_configuration=model_configuration, training_configuration=training_configuration, train_dataset=train_loader, val_dataset=val_loader, log_dir=log_dir, logger=logger, num_dataloader_workers=num_dataloader_workers, ) def evaluate(model_configuration: dict, training_configuration: dict, test_ids: list) -> dict: logger.info(f"Evaluating model {model_configuration['id']}") eval_functions = get_eval_functions(model_configuration) results = evaluate_model(model_configuration, training_configuration, test_ids, eval_functions) # save results to file results_dir = training_configuration["training_dir"] if not os.path.exists(results_dir): os.makedirs(results_dir) results_file = os.path.join(results_dir, "evaluation_results.json") with open(results_file, "w") as f: json.dump(convert_for_json(results), f, indent=4) return results if __name__ == "__main__": # set up distributed training dist.init_process_group(backend="nccl", init_method="env://") # sleeping for a bit to allow all processes to initialize time.sleep(5) local_rank = torch.distributed.get_rank() world_size = torch.distributed.get_world_size() torch.cuda.set_device(local_rank) print(f"Rank {local_rank} initialized with world size {world_size}") parser = argparse.ArgumentParser(description="Training wrapper for model training") parser.add_argument("run_configuration_module", type=str, help="Path to the run configuration module") parser.add_argument("--run_configuration_variable", type=str, required=False, default="run_configuration", help="Name of the run configuration variable in the module") parser.add_argument("--results_dir", type=str, required=False, default=None, help="Directory to save the results") parser.add_argument("--lmdb_root_dir", type=str, required=False, default=None, help="Path to the lmdb directory") parser.add_argument("--log_dir", type=str, required=False, default=None, help="Directory to save the logs") parser.add_argument("--item_limit", type=int, required=False, default=None, help="Limit the number of items to process, default is None (no limit)") parser.add_argument("--num_dataloader_workers", type=int, required=False, default=1, help="Number of workers for the dataloader, default is 1") args = parser.parse_args() # load run configuration from module run_configuration = get_variable_from_module( # make sure to replace / with . and remove .py to get proper module tree module_path=args.run_configuration_module.replace(".py", "").replace("/", "."), variable_name=args.run_configuration_variable) # check for variables in run_configuration if not args.results_dir: if "base_results_dir" not in run_configuration: results_dir = os.getenv("RESULTS_ROOT_DIR") else: results_dir = run_configuration["base_results_dir"] else: results_dir = args.results_dir if results_dir is None: raise ValueError( "No results directory specified. Please set the RESULTS_ROOT_DIR environment variable or provide a results_dir argument.") if not args.lmdb_root_dir: if "base_lmdb_root_dir" not in run_configuration: lmdb_root_dir = os.getenv("LMDB_ROOT_DIR") else: lmdb_root_dir = run_configuration["base_lmdb_root_dir"] else: lmdb_root_dir = args.lmdb_root_dir if lmdb_root_dir is None: raise ValueError( "No LMDB root directory specified. Please set the LMDB_ROOT_DIR environment variable or provide a lmdb_root_dir argument.") if not args.log_dir: if "base_log_dir" not in run_configuration: log_dir = os.getenv("LOG_DIR") else: log_dir = run_configuration["log_dir"] else: log_dir = args.log_dir if log_dir is None: raise ValueError( "No log directory specified. Please set the LOG_DIR environment variable or provide a log_dir argument.") try: item_limit = args.item_limit if args.item_limit else run_configuration["item_limit"] if item_limit == -1: item_limit = None except: item_limit = None run_name = run_configuration["name"] runs = run_configuration["runs"] # set up run_id on rank 0 if local_rank == 0: run_id = f"{run_name}_{datetime.now().strftime('%Y%m%d_%H%M%S')}" if not os.path.exists(log_dir): os.makedirs(log_dir) if not os.path.exists(results_dir): os.makedirs(results_dir) else: run_id = None if torch.distributed.is_initialized(): torch.distributed.barrier() # broadcast run_id to all processes if dist.is_initialized(): print(f"Rank {local_rank}: Broadcasting run_id {run_id}") run_id_list = [run_id] torch.distributed.broadcast_object_list(run_id_list, src=0) run_id = run_id_list[0] # make sure run_id is a string run_id = str(run_id) print(f"Rank {local_rank}: fetched run_id {run_id}") # append run_id to results_dir and log_dir results_dir = os.path.join(results_dir, run_id) log_dir = os.path.join(log_dir, run_id) print(f"Rank {local_rank}: Results directory: {results_dir}") print(f"Rank {local_rank}: Log directory: {log_dir}") # create directories if they do not exist, only on the main process if torch.distributed.get_rank() == 0 or not dist.is_initialized(): print(f"Rank {local_rank}: Creating directories for run {run_name}") if not os.path.exists(results_dir): os.makedirs(results_dir) if not os.path.exists(log_dir): os.makedirs(log_dir) # sync all processes to make sure the directories are created if dist.is_initialized(): print(f"Rank {local_rank}: Waiting for all processes to create directories") dist.barrier() # set up logger for all processes logger = get_logger(module_name=run_name, filename=os.path.join(log_dir, f"main_{local_rank}.log")) logger.info(f"Rank {local_rank}: Starting run {run_name}") logger.info(f"Rank {local_rank}: Run ID: {run_id}") # print parameter values logger.info(f"Rank {local_rank}: Results directory: {results_dir}") logger.info(f"Rank {local_rank}: LMDB root directory: {lmdb_root_dir}") logger.info(f"Rank {local_rank}: Log directory: {log_dir}") logger.info(f"Rank {local_rank}: Item limit: {item_limit}") if not runs or len(runs) == 0: raise ValueError("No runs specified in the run configuration. Please provide a list of runs to train.") try: for run in runs: run_step_name = run["name"] run_description = run["description"] run_model_configuration = run["model_configuration"] run_training_configuration = run["training_configuration"] logger.info(f"Rank {local_rank}: Running: {run_step_name}") logger.info(f"Rank {local_rank}: Description: {run_description}") # prepare run, make sure rank 0 is the first to avoid race conditions if local_rank == 0: logger.info(f"Rank {local_rank}: Preparing run {run_step_name}") run_model_configuration, run_training_configuration, train_ids, val_ids, test_ids, dataset_dir = prepare_run( results_dir=results_dir, lmdb_root_dir=lmdb_root_dir, base_model_configuration=run_model_configuration, base_training_configuration=run_training_configuration ) logger.info(f"Rank {local_rank}: Finished preparing run {run_step_name}") # all ranks sync here if dist.is_initialized(): if local_rank != 0: logger.info(f"Rank {local_rank}: Waiting for rank 0 to finish preparing run {run_step_name}") dist.barrier() if local_rank != 0: logger.info( f"Rank {local_rank}: Finished waiting for rank 0 to finish preparing run {run_step_name}") if local_rank != 0: # after the barrier, rank 0 will have prepared the run run_model_configuration, run_training_configuration, train_ids, val_ids, test_ids, dataset_dir = prepare_run( results_dir=results_dir, lmdb_root_dir=lmdb_root_dir, base_model_configuration=run_model_configuration, base_training_configuration=run_training_configuration ) train( model_configuration=run_model_configuration, training_configuration=run_training_configuration, train_ids=train_ids, val_ids=val_ids, dataset_dir=dataset_dir, log_dir=log_dir, ) # sync after training if dist.is_initialized(): dist.barrier() logger.info(f"Rank {torch.distributed.get_rank()} finished training {run_step_name}") # run evaluation, only on rank 0 local_rank = torch.distributed.get_rank() if local_rank == 0: logger.info(f"Rank {local_rank} starting evaluation for {run_step_name}") results = evaluate( model_configuration=run_model_configuration, training_configuration=run_training_configuration, test_ids=test_ids, ) logger.info( f"Rank {local_rank} finished evaluation for {run_step_name} with model {run_model_configuration['id']}") logger.info(f"Results: {results}") # sync after evaluation if dist.is_initialized(): dist.barrier() finally: # clean up if dist.is_initialized(): dist.destroy_process_group()