import os from typing import Any, Dict, List, Optional, Tuple from dora import get_xp, hydra_main import hydra import lightning as L import rootutils import torch from lightning import Callback, LightningDataModule, LightningModule, Trainer from lightning.pytorch.loggers import Logger from omegaconf import DictConfig import logging from pathlib import Path rootutils.setup_root(__file__, indicator=".project-root", pythonpath=True) # ------------------------------------------------------------------------------------ # # the setup_root above is equivalent to: # - adding project root dir to PYTHONPATH # (so you don't need to force user to install project as a package) # (necessary before importing any local modules e.g. `from gdr import utils`) # - setting up PROJECT_ROOT environment variable # (which is used as a base for paths in "configs/paths/default.yaml") # (this way all filepaths are the same no matter where you run the code) # - loading environment variables from ".env" in root dir # # you can remove it if you: # 1. either install project as a package or move entry files to project root dir # 2. set `root_dir` to "." in "configs/paths/default.yaml" # # more info: https://github.com/ashleve/rootutils # ------------------------------------------------------------------------------------ # from steerable_retrieval.utils import ( RankedLogger, extras, get_metric_value, instantiate_callbacks, instantiate_loggers, log_hyperparameters, register_resolvers, task_wrapper, ) log = RankedLogger(__name__, rank_zero_only=True) register_resolvers() @task_wrapper def train(cfg: DictConfig) -> Tuple[Dict[str, Any], Dict[str, Any]]: """Trains the model. Can additionally evaluate on a testset, using best weights obtained during training. This method is wrapped in optional @task_wrapper decorator, that controls the behavior during failure. Useful for multiruns, saving info about the crash, etc. :param cfg: A DictConfig configuration composed by Hydra. :return: A tuple with metrics and dict with all instantiated objects. """ # set seed for random number generators in pytorch, numpy and python.random if cfg.get("seed"): L.seed_everything(cfg.seed, workers=True) log.info(f"Instantiating datamodule <{cfg.data._target_}>") datamodule = hydra.utils.instantiate(cfg.data) log.info(f"Instantiating model <{cfg.model._target_}>") model: LightningModule = hydra.utils.instantiate(cfg.model) # model.xp = get_xp() log.info("Instantiating callbacks...") callbacks: List[Callback] = instantiate_callbacks(cfg.get("callbacks")) log.info(f"Callbacks: {callbacks}") log.info("Instantiating loggers...") logger: List[Logger] = instantiate_loggers(cfg.get("logger")) log.info(f"Instantiating trainer <{cfg.trainer._target_}>") trainer: Trainer = hydra.utils.instantiate(cfg.trainer, logger=logger, callbacks=callbacks) object_dict = { "cfg": cfg, "datamodule": datamodule, "model": model, "callbacks": callbacks, "logger": logger, "trainer": trainer, } if logger: log.info("Logging hyperparameters!") log_hyperparameters(object_dict) # automatically resume from latest checkpoint if exists and ckpt_path not manually specified # TODO: discuss cfg.resume, this is anti-dora but maybe it's useful ckpt_path = cfg.get("ckpt_path") cfg.resume = cfg.resume or os.environ.get("USE_MPI") if '/opt/ml/' in cfg.paths.ckpt_dir: was_s3 = True else: was_s3 = False logging.info("="*100) # logging.info(os.listdir('/opt/ml/input/data')) if os.path.exists('/opt/ml/input/data') else logging.info("No data found in /opt/ml/input/data") # log tree of /opt/ml/input/data def tree_str( path=".", max_depth=None, max_files=2, ignore={".git", "__pycache__"} ): lines = [] path = Path(path) def _walk(p, prefix="", level=0): if max_depth is not None and level > max_depth: return entries = [e for e in p.iterdir() if e.name not in ignore] dirs = sorted((e for e in entries if e.is_dir()), key=lambda x: x.name.lower()) files = sorted((e for e in entries if e.is_file()), key=lambda x: x.name.lower()) shown_files = files[:max_files] omitted_files = len(files) - len(shown_files) combined = dirs + shown_files for i, entry in enumerate(combined): is_last = i == len(combined) - 1 connector = "└── " if is_last else "├── " lines.append(prefix + connector + entry.name) if entry.is_dir(): extension = " " if is_last else "│ " _walk(entry, prefix + extension, level + 1) if omitted_files > 0: lines.append(prefix + f"└── … ({omitted_files} more files)") _walk(path) return "\n".join(lines) logging.info("="*100) data_root = cfg.paths.get('data_dir') # config-provided data path (was hardcoded /opt/ml on SageMaker) if data_root and os.path.exists(data_root): logging.info(tree_str(data_root)) logging.info("="*100) if os.path.exists(cfg.paths.ckpt_dir) and cfg.resume: candidates = [os.path.join(cfg.paths.ckpt_dir, ckpt_file) for ckpt_file in os.listdir(cfg.paths.ckpt_dir) if ckpt_file.endswith(".ckpt")] if candidates: # get the last modified ckpt else get last.ckpt, reason is that s3 downloads are not in order of creation # ckpt_path = max(candidates, key=os.path.getmtime) if not was_s3 else if was_s3: ckpt_path = os.path.join(cfg.paths.ckpt_dir, "last.ckpt") if "last.ckpt" not in os.listdir(cfg.paths.ckpt_dir): log.warning("last.ckpt not found in s3 ckpt_dir. Training from scratch!") ckpt_path = None else: ckpt_path = max(candidates, key=os.path.getmtime) log.info(f"Resuming from checkpoint {ckpt_path}...") # ckpt_path = os.path.join(cfg.paths.ckpt_dir, "last.ckpt") if "last.ckpt" in os.listdir(cfg.paths.ckpt_dir) else None log.info(f"Resuming from checkpoint {ckpt_path}...") else: log.info(ckpt_path, "is empty. Training from scratch!") trainer.true_accumulate_grad_batches, trainer.accumulate_grad_batches = trainer.accumulate_grad_batches, 1 model.gradient_clip_val, trainer.gradient_clip_val = trainer.gradient_clip_val, None if cfg.get("train"): log.info("Starting training!") with torch.backends.cuda.sdp_kernel(enable_flash=True, enable_math=True, enable_mem_efficient=True): trainer.fit(model=model, datamodule=datamodule, ckpt_path=ckpt_path) train_metrics = trainer.callback_metrics if cfg.get("test"): log.info("Starting testing!") # Get best checkpoint path if checkpoint callback exists if hasattr(trainer, 'checkpoint_callback') and trainer.checkpoint_callback is not None: ckpt_path = trainer.checkpoint_callback.best_model_path if ckpt_path == "": log.warning("Best ckpt not found! Using current weights for testing...") ckpt_path = None else: log.warning("No checkpoint callback found! Using current weights for testing...") ckpt_path = None trainer.test(model=model, datamodule=datamodule, ckpt_path=ckpt_path) log.info(f"Best ckpt path: {ckpt_path}") test_metrics = trainer.callback_metrics # merge train and test metrics metric_dict = {**train_metrics, **test_metrics} return metric_dict, object_dict return {}, object_dict @hydra_main(version_base="1.3", config_path="../configs", config_name="train.yaml") def main(cfg: DictConfig) -> Optional[float]: """Main entry point for training. :param cfg: DictConfig configuration composed by Hydra. :return: Optional[float] with optimized metric value. """ # handle A100 GPUs if torch.cuda.is_available() and ("A100" in torch.cuda.get_device_name() or "A5000" in torch.cuda.get_device_name()): torch.set_float32_matmul_precision("high") # avoid annoying multiprocessing errors torch.multiprocessing.set_sharing_strategy('file_system') # prevent annoying warning os.environ["TOKENIZERS_PARALLELISM"] = "false" # apply extra utilities # (e.g. ask for tags if none are provided in cfg, print cfg tree, etc.) extras(cfg) # train the model metric_dict, _ = train(cfg) # safely retrieve metric value for hydra-based hyperparameter optimization metric_value = get_metric_value( metric_dict=metric_dict, metric_name=cfg.get("optimized_metric") ) # return optimized metric return metric_value if __name__ == "__main__": main()