import argparse import datetime import glob import inspect import os import sys from inspect import Parameter from typing import Union import numpy as np import pytorch_lightning as pl import torch import torchvision import wandb from matplotlib import pyplot as plt from natsort import natsorted from omegaconf import OmegaConf from packaging import version from PIL import Image from pytorch_lightning import seed_everything from pytorch_lightning.callbacks import Callback from pytorch_lightning.loggers import WandbLogger from pytorch_lightning.trainer import Trainer from pytorch_lightning.utilities import rank_zero_only from sgm.util import exists, instantiate_from_config, isheatmap MULTINODE_HACKS = True def default_trainer_args(): argspec = dict(inspect.signature(Trainer.__init__).parameters) argspec.pop("self") default_args = { param: argspec[param].default for param in argspec if argspec[param] != Parameter.empty } return default_args def get_parser(**parser_kwargs): def str2bool(v): if isinstance(v, bool): return v if v.lower() in ("yes", "true", "t", "y", "1"): return True elif v.lower() in ("no", "false", "f", "n", "0"): return False else: raise argparse.ArgumentTypeError("Boolean value expected.") parser = argparse.ArgumentParser(**parser_kwargs) parser.add_argument( "-n", "--name", type=str, const=True, default="", nargs="?", help="postfix for logdir", ) parser.add_argument( "--no_date", type=str2bool, nargs="?", const=True, default=False, help="if True, skip date generation for logdir and only use naming via opt.base or opt.name (+ opt.postfix, optionally)", ) parser.add_argument( "-r", "--resume", type=str, const=True, default="", nargs="?", help="resume from logdir or checkpoint in logdir", ) parser.add_argument( "-b", "--base", nargs="*", metavar="base_config.yaml", help="paths to base configs. Loaded from left-to-right. " "Parameters can be overwritten or added with command-line options of the form `--key value`.", default=list(), ) parser.add_argument( "-t", "--train", type=str2bool, const=True, default=True, nargs="?", help="train", ) parser.add_argument( "--no-test", type=str2bool, const=True, default=False, nargs="?", help="disable test", ) parser.add_argument( "-p", "--project", help="name of new or path to existing project" ) parser.add_argument( "-d", "--debug", type=str2bool, nargs="?", const=True, default=False, help="enable post-mortem debugging", ) parser.add_argument( "-s", "--seed", type=int, default=23, help="seed for seed_everything", ) parser.add_argument( "-f", "--postfix", type=str, default="", help="post-postfix for default name", ) parser.add_argument( "--projectname", type=str, default="stablediffusion", ) parser.add_argument( "-l", "--logdir", type=str, default="logs", help="directory for logging dat shit", ) parser.add_argument( "--scale_lr", type=str2bool, nargs="?", const=True, default=False, help="scale base-lr by ngpu * batch_size * n_accumulate", ) parser.add_argument( "--legacy_naming", type=str2bool, nargs="?", const=True, default=False, help="name run based on config file name if true, else by whole path", ) parser.add_argument( "--enable_tf32", type=str2bool, nargs="?", const=True, default=False, help="enables the TensorFloat32 format both for matmuls and cuDNN for pytorch 1.12", ) parser.add_argument( "--startup", type=str, default=None, help="Startuptime from distributed script", ) parser.add_argument( "--wandb", type=str2bool, nargs="?", const=True, default=False, # TODO: later default to True help="log to wandb", ) parser.add_argument( "--no_base_name", type=str2bool, nargs="?", const=True, default=False, # TODO: later default to True help="log to wandb", ) if version.parse(torch.__version__) >= version.parse("2.0.0"): parser.add_argument( "--resume_from_checkpoint", type=str, default=None, help="single checkpoint file to resume from", ) default_args = default_trainer_args() for key in default_args: parser.add_argument("--" + key, default=default_args[key]) return parser def get_checkpoint_name(logdir): ckpt = os.path.join(logdir, "checkpoints", "last**.ckpt") ckpt = natsorted(glob.glob(ckpt)) print('available "last" checkpoints:') print(ckpt) if len(ckpt) > 1: print("got most recent checkpoint") ckpt = sorted(ckpt, key=lambda x: os.path.getmtime(x))[-1] print(f"Most recent ckpt is {ckpt}") with open(os.path.join(logdir, "most_recent_ckpt.txt"), "w") as f: f.write(ckpt + "\n") try: version = int(ckpt.split("/")[-1].split("-v")[-1].split(".")[0]) except Exception as e: print("version confusion but not bad") print(e) version = 1 # version = last_version + 1 else: # in this case, we only have one "last.ckpt" ckpt = ckpt[0] version = 1 melk_ckpt_name = f"last-v{version}.ckpt" print(f"Current melk ckpt name: {melk_ckpt_name}") return ckpt, melk_ckpt_name class SetupCallback(Callback): def __init__( self, resume, now, logdir, ckptdir, cfgdir, config, lightning_config, debug, ckpt_name=None, ): super().__init__() self.resume = resume self.now = now self.logdir = logdir self.ckptdir = ckptdir self.cfgdir = cfgdir self.config = config self.lightning_config = lightning_config self.debug = debug self.ckpt_name = ckpt_name def on_exception(self, trainer: pl.Trainer, pl_module, exception): if not self.debug and trainer.global_rank == 0: print("Summoning checkpoint.") if self.ckpt_name is None: ckpt_path = os.path.join(self.ckptdir, "last.ckpt") else: ckpt_path = os.path.join(self.ckptdir, self.ckpt_name) trainer.save_checkpoint(ckpt_path) def on_fit_start(self, trainer, pl_module): if trainer.global_rank == 0: # Create logdirs and save configs os.makedirs(self.logdir, exist_ok=True) os.makedirs(self.ckptdir, exist_ok=True) os.makedirs(self.cfgdir, exist_ok=True) if "callbacks" in self.lightning_config: if ( "metrics_over_trainsteps_checkpoint" in self.lightning_config["callbacks"] ): os.makedirs( os.path.join(self.ckptdir, "trainstep_checkpoints"), exist_ok=True, ) print("Project config") print(OmegaConf.to_yaml(self.config)) if MULTINODE_HACKS: import time time.sleep(5) OmegaConf.save( self.config, os.path.join(self.cfgdir, "{}-project.yaml".format(self.now)), ) print("Lightning config") print(OmegaConf.to_yaml(self.lightning_config)) OmegaConf.save( OmegaConf.create({"lightning": self.lightning_config}), os.path.join(self.cfgdir, "{}-lightning.yaml".format(self.now)), ) else: # ModelCheckpoint callback created log directory --- remove it if not MULTINODE_HACKS and not self.resume and os.path.exists(self.logdir): dst, name = os.path.split(self.logdir) dst = os.path.join(dst, "child_runs", name) os.makedirs(os.path.split(dst)[0], exist_ok=True) try: os.rename(self.logdir, dst) except FileNotFoundError: pass class ImageLogger(Callback): def __init__( self, batch_frequency, max_images, clamp=True, increase_log_steps=True, rescale=True, disabled=False, log_on_batch_idx=False, log_first_step=False, log_images_kwargs=None, log_before_first_step=False, enable_autocast=True, ): super().__init__() self.enable_autocast = enable_autocast self.rescale = rescale self.batch_freq = batch_frequency self.max_images = max_images self.log_steps = [2**n for n in range(int(np.log2(self.batch_freq)) + 1)] if not increase_log_steps: self.log_steps = [self.batch_freq] self.clamp = clamp self.disabled = disabled self.log_on_batch_idx = log_on_batch_idx self.log_images_kwargs = log_images_kwargs if log_images_kwargs else {} self.log_first_step = log_first_step self.log_before_first_step = log_before_first_step @rank_zero_only def log_local( self, save_dir, split, images, global_step, current_epoch, batch_idx, pl_module: Union[None, pl.LightningModule] = None, ): root = os.path.join(save_dir, "images", split) for k in images: if isheatmap(images[k]): fig, ax = plt.subplots() ax = ax.matshow( images[k].cpu().numpy(), cmap="hot", interpolation="lanczos" ) plt.colorbar(ax) plt.axis("off") filename = "{}_gs-{:06}_e-{:06}_b-{:06}.png".format( k, global_step, current_epoch, batch_idx ) os.makedirs(root, exist_ok=True) path = os.path.join(root, filename) plt.savefig(path) plt.close() # TODO: support wandb else: grid = torchvision.utils.make_grid(images[k], nrow=4) if self.rescale: grid = (grid + 1.0) / 2.0 # -1,1 -> 0,1; c,h,w grid = grid.transpose(0, 1).transpose(1, 2).squeeze(-1) grid = grid.numpy() grid = (grid * 255).astype(np.uint8) filename = "{}_gs-{:06}_e-{:06}_b-{:06}.png".format( k, global_step, current_epoch, batch_idx ) path = os.path.join(root, filename) os.makedirs(os.path.split(path)[0], exist_ok=True) img = Image.fromarray(grid) img.save(path) if exists(pl_module): assert isinstance( pl_module.logger, WandbLogger ), "logger_log_image only supports WandbLogger currently" pl_module.logger.log_image( key=f"{split}/{k}", images=[ img, ], step=pl_module.global_step, ) @rank_zero_only def log_img(self, pl_module, batch, batch_idx, split="train"): check_idx = batch_idx if self.log_on_batch_idx else pl_module.global_step if ( self.check_frequency(check_idx) and hasattr(pl_module, "log_images") # batch_idx % self.batch_freq == 0 and callable(pl_module.log_images) and # batch_idx > 5 and self.max_images > 0 ): logger = type(pl_module.logger) is_train = pl_module.training if is_train: pl_module.eval() gpu_autocast_kwargs = { "enabled": self.enable_autocast, # torch.is_autocast_enabled(), "dtype": torch.get_autocast_gpu_dtype(), "cache_enabled": torch.is_autocast_cache_enabled(), } with torch.no_grad(), torch.cuda.amp.autocast(**gpu_autocast_kwargs): images = pl_module.log_images( batch, split=split, **self.log_images_kwargs ) for k in images: N = min(images[k].shape[0], self.max_images) if not isheatmap(images[k]): images[k] = images[k][:N] if isinstance(images[k], torch.Tensor): images[k] = images[k].detach().float().cpu() if self.clamp and not isheatmap(images[k]): images[k] = torch.clamp(images[k], -1.0, 1.0) self.log_local( pl_module.logger.save_dir, split, images, pl_module.global_step, pl_module.current_epoch, batch_idx, pl_module=pl_module if isinstance(pl_module.logger, WandbLogger) else None, ) if is_train: pl_module.train() def check_frequency(self, check_idx): if ((check_idx % self.batch_freq) == 0 or (check_idx in self.log_steps)) and ( check_idx > 0 or self.log_first_step ): try: self.log_steps.pop(0) except IndexError as e: print(e) pass return True return False @rank_zero_only def on_train_batch_end(self, trainer, pl_module, outputs, batch, batch_idx): if not self.disabled and (pl_module.global_step > 0 or self.log_first_step): self.log_img(pl_module, batch, batch_idx, split="train") @rank_zero_only def on_train_batch_start(self, trainer, pl_module, batch, batch_idx): if self.log_before_first_step and pl_module.global_step == 0: print(f"{self.__class__.__name__}: logging before training") self.log_img(pl_module, batch, batch_idx, split="train") @rank_zero_only def on_validation_batch_end( self, trainer, pl_module, outputs, batch, batch_idx, *args, **kwargs ): if not self.disabled and pl_module.global_step > 0: self.log_img(pl_module, batch, batch_idx, split="val") if hasattr(pl_module, "calibrate_grad_norm"): if ( pl_module.calibrate_grad_norm and batch_idx % 25 == 0 ) and batch_idx > 0: self.log_gradients(trainer, pl_module, batch_idx=batch_idx) @rank_zero_only def init_wandb(save_dir, opt, config, group_name, name_str): print(f"setting WANDB_DIR to {save_dir}") os.makedirs(save_dir, exist_ok=True) os.environ["WANDB_DIR"] = save_dir if opt.debug: wandb.init(project=opt.projectname, mode="offline", group=group_name) else: wandb.init( project=opt.projectname, config=config, settings=wandb.Settings(code_dir="./sgm"), group=group_name, name=name_str, ) if __name__ == "__main__": # custom parser to specify config files, train, test and debug mode, # postfix, resume. # `--key value` arguments are interpreted as arguments to the trainer. # `nested.key=value` arguments are interpreted as config parameters. # configs are merged from left-to-right followed by command line parameters. # model: # base_learning_rate: float # target: path to lightning module # params: # key: value # data: # target: main.DataModuleFromConfig # params: # batch_size: int # wrap: bool # train: # target: path to train dataset # params: # key: value # validation: # target: path to validation dataset # params: # key: value # test: # target: path to test dataset # params: # key: value # lightning: (optional, has sane defaults and can be specified on cmdline) # trainer: # additional arguments to trainer # logger: # logger to instantiate # modelcheckpoint: # modelcheckpoint to instantiate # callbacks: # callback1: # target: importpath # params: # key: value now = datetime.datetime.now().strftime("%Y-%m-%dT%H-%M-%S") # add cwd for convenience and to make classes in this file available when # running as `python main.py` # (in particular `main.DataModuleFromConfig`) sys.path.append(os.getcwd()) parser = get_parser() opt, unknown = parser.parse_known_args() if opt.name and opt.resume: raise ValueError( "-n/--name and -r/--resume cannot be specified both." "If you want to resume training in a new log folder, " "use -n/--name in combination with --resume_from_checkpoint" ) melk_ckpt_name = None name = None if opt.resume: if not os.path.exists(opt.resume): raise ValueError("Cannot find {}".format(opt.resume)) if os.path.isfile(opt.resume): paths = opt.resume.split("/") # idx = len(paths)-paths[::-1].index("logs")+1 # logdir = "/".join(paths[:idx]) logdir = "/".join(paths[:-2]) ckpt = opt.resume _, melk_ckpt_name = get_checkpoint_name(logdir) else: assert os.path.isdir(opt.resume), opt.resume logdir = opt.resume.rstrip("/") ckpt, melk_ckpt_name = get_checkpoint_name(logdir) print("#" * 100) print(f'Resuming from checkpoint "{ckpt}"') print("#" * 100) opt.resume_from_checkpoint = ckpt base_configs = sorted(glob.glob(os.path.join(logdir, "configs/*.yaml"))) opt.base = base_configs + opt.base _tmp = logdir.split("/") nowname = _tmp[-1] else: if opt.name: name = "_" + opt.name elif opt.base: if opt.no_base_name: name = "" else: if opt.legacy_naming: cfg_fname = os.path.split(opt.base[0])[-1] cfg_name = os.path.splitext(cfg_fname)[0] else: assert "configs" in os.path.split(opt.base[0])[0], os.path.split( opt.base[0] )[0] cfg_path = os.path.split(opt.base[0])[0].split(os.sep)[ os.path.split(opt.base[0])[0].split(os.sep).index("configs") + 1 : ] # cut away the first one (we assert all configs are in "configs") cfg_name = os.path.splitext(os.path.split(opt.base[0])[-1])[0] cfg_name = "-".join(cfg_path) + f"-{cfg_name}" name = "_" + cfg_name else: name = "" if not opt.no_date: nowname = now + name + opt.postfix else: nowname = name + opt.postfix if nowname.startswith("_"): nowname = nowname[1:] logdir = os.path.join(opt.logdir, nowname) print(f"LOGDIR: {logdir}") ckptdir = os.path.join(logdir, "checkpoints") cfgdir = os.path.join(logdir, "configs") seed_everything(opt.seed, workers=True) # move before model init, in case a torch.compile(...) is called somewhere if opt.enable_tf32: # pt_version = version.parse(torch.__version__) torch.backends.cuda.matmul.allow_tf32 = True torch.backends.cudnn.allow_tf32 = True print(f"Enabling TF32 for PyTorch {torch.__version__}") else: print(f"Using default TF32 settings for PyTorch {torch.__version__}:") print( f"torch.backends.cuda.matmul.allow_tf32={torch.backends.cuda.matmul.allow_tf32}" ) print(f"torch.backends.cudnn.allow_tf32={torch.backends.cudnn.allow_tf32}") try: # init and save configs configs = [OmegaConf.load(cfg) for cfg in opt.base] cli = OmegaConf.from_dotlist(unknown) config = OmegaConf.merge(*configs, cli) lightning_config = config.pop("lightning", OmegaConf.create()) # merge trainer cli with config trainer_config = lightning_config.get("trainer", OmegaConf.create()) # default to gpu trainer_config["accelerator"] = "gpu" # standard_args = default_trainer_args() for k in standard_args: if getattr(opt, k) != standard_args[k]: trainer_config[k] = getattr(opt, k) ckpt_resume_path = opt.resume_from_checkpoint if not "devices" in trainer_config and trainer_config["accelerator"] != "gpu": del trainer_config["accelerator"] cpu = True else: gpuinfo = trainer_config["devices"] print(f"Running on GPUs {gpuinfo}") cpu = False trainer_opt = argparse.Namespace(**trainer_config) lightning_config.trainer = trainer_config # model model = instantiate_from_config(config.model) # trainer and callbacks trainer_kwargs = dict() # default logger configs default_logger_cfgs = { "wandb": { "target": "pytorch_lightning.loggers.WandbLogger", "params": { "name": nowname, # "save_dir": logdir, "offline": opt.debug, "id": nowname, "project": opt.projectname, "log_model": False, # "dir": logdir, }, }, "csv": { "target": "pytorch_lightning.loggers.CSVLogger", "params": { "name": "testtube", # hack for sbord fanatics "save_dir": logdir, }, }, } default_logger_cfg = default_logger_cfgs["wandb" if opt.wandb else "csv"] if opt.wandb: # TODO change once leaving "swiffer" config directory try: group_name = nowname.split(now)[-1].split("-")[1] except: group_name = nowname default_logger_cfg["params"]["group"] = group_name init_wandb( os.path.join(os.getcwd(), logdir), opt=opt, group_name=group_name, config=config, name_str=nowname, ) if "logger" in lightning_config: logger_cfg = lightning_config.logger else: logger_cfg = OmegaConf.create() logger_cfg = OmegaConf.merge(default_logger_cfg, logger_cfg) trainer_kwargs["logger"] = instantiate_from_config(logger_cfg) # modelcheckpoint - use TrainResult/EvalResult(checkpoint_on=metric) to # specify which metric is used to determine best models default_modelckpt_cfg = { "target": "pytorch_lightning.callbacks.ModelCheckpoint", "params": { "dirpath": ckptdir, "filename": "{epoch:06}", "verbose": True, "save_last": True, }, } if hasattr(model, "monitor"): print(f"Monitoring {model.monitor} as checkpoint metric.") default_modelckpt_cfg["params"]["monitor"] = model.monitor default_modelckpt_cfg["params"]["save_top_k"] = 3 if "modelcheckpoint" in lightning_config: modelckpt_cfg = lightning_config.modelcheckpoint else: modelckpt_cfg = OmegaConf.create() modelckpt_cfg = OmegaConf.merge(default_modelckpt_cfg, modelckpt_cfg) print(f"Merged modelckpt-cfg: \n{modelckpt_cfg}") # https://pytorch-lightning.readthedocs.io/en/stable/extensions/strategy.html # default to ddp if not further specified default_strategy_config = {"target": "pytorch_lightning.strategies.DDPStrategy"} if "strategy" in lightning_config: strategy_cfg = lightning_config.strategy else: strategy_cfg = OmegaConf.create() default_strategy_config["params"] = { "find_unused_parameters": False, # "static_graph": True, # "ddp_comm_hook": default.fp16_compress_hook # TODO: experiment with this, also for DDPSharded } strategy_cfg = OmegaConf.merge(default_strategy_config, strategy_cfg) print( f"strategy config: \n ++++++++++++++ \n {strategy_cfg} \n ++++++++++++++ " ) trainer_kwargs["strategy"] = instantiate_from_config(strategy_cfg) # add callback which sets up log directory default_callbacks_cfg = { "setup_callback": { "target": "main.SetupCallback", "params": { "resume": opt.resume, "now": now, "logdir": logdir, "ckptdir": ckptdir, "cfgdir": cfgdir, "config": config, "lightning_config": lightning_config, "debug": opt.debug, "ckpt_name": melk_ckpt_name, }, }, "image_logger": { "target": "main.ImageLogger", "params": {"batch_frequency": 1000, "max_images": 4, "clamp": True}, }, "learning_rate_logger": { "target": "pytorch_lightning.callbacks.LearningRateMonitor", "params": { "logging_interval": "step", # "log_momentum": True }, }, } if version.parse(pl.__version__) >= version.parse("1.4.0"): default_callbacks_cfg.update({"checkpoint_callback": modelckpt_cfg}) if "callbacks" in lightning_config: callbacks_cfg = lightning_config.callbacks else: callbacks_cfg = OmegaConf.create() if "metrics_over_trainsteps_checkpoint" in callbacks_cfg: print( "Caution: Saving checkpoints every n train steps without deleting. This might require some free space." ) default_metrics_over_trainsteps_ckpt_dict = { "metrics_over_trainsteps_checkpoint": { "target": "pytorch_lightning.callbacks.ModelCheckpoint", "params": { "dirpath": os.path.join(ckptdir, "trainstep_checkpoints"), "filename": "{epoch:06}-{step:09}", "verbose": True, "save_top_k": -1, "every_n_train_steps": 10000, "save_weights_only": True, }, } } default_callbacks_cfg.update(default_metrics_over_trainsteps_ckpt_dict) callbacks_cfg = OmegaConf.merge(default_callbacks_cfg, callbacks_cfg) if "ignore_keys_callback" in callbacks_cfg and ckpt_resume_path is not None: callbacks_cfg.ignore_keys_callback.params["ckpt_path"] = ckpt_resume_path elif "ignore_keys_callback" in callbacks_cfg: del callbacks_cfg["ignore_keys_callback"] trainer_kwargs["callbacks"] = [ instantiate_from_config(callbacks_cfg[k]) for k in callbacks_cfg ] if not "plugins" in trainer_kwargs: trainer_kwargs["plugins"] = list() # cmd line trainer args (which are in trainer_opt) have always priority over config-trainer-args (which are in trainer_kwargs) trainer_opt = vars(trainer_opt) trainer_kwargs = { key: val for key, val in trainer_kwargs.items() if key not in trainer_opt } trainer = Trainer(**trainer_opt, **trainer_kwargs) trainer.logdir = logdir ### # data data = instantiate_from_config(config.data) # NOTE according to https://pytorch-lightning.readthedocs.io/en/latest/datamodules.html # calling these ourselves should not be necessary but it is. # lightning still takes care of proper multiprocessing though data.prepare_data() # data.setup() print("#### Data #####") try: for k in data.datasets: print( f"{k}, {data.datasets[k].__class__.__name__}, {len(data.datasets[k])}" ) except: print("datasets not yet initialized.") # configure learning rate if "batch_size" in config.data.params: bs, base_lr = config.data.params.batch_size, config.model.base_learning_rate else: bs, base_lr = ( config.data.params.train.loader.batch_size, config.model.base_learning_rate, ) if not cpu: ngpu = len(lightning_config.trainer.devices.strip(",").split(",")) else: ngpu = 1 if "accumulate_grad_batches" in lightning_config.trainer: accumulate_grad_batches = lightning_config.trainer.accumulate_grad_batches else: accumulate_grad_batches = 1 print(f"accumulate_grad_batches = {accumulate_grad_batches}") lightning_config.trainer.accumulate_grad_batches = accumulate_grad_batches if opt.scale_lr: model.learning_rate = accumulate_grad_batches * ngpu * bs * base_lr print( "Setting learning rate to {:.2e} = {} (accumulate_grad_batches) * {} (num_gpus) * {} (batchsize) * {:.2e} (base_lr)".format( model.learning_rate, accumulate_grad_batches, ngpu, bs, base_lr ) ) else: model.learning_rate = base_lr print("++++ NOT USING LR SCALING ++++") print(f"Setting learning rate to {model.learning_rate:.2e}") # allow checkpointing via USR1 def melk(*args, **kwargs): # run all checkpoint hooks if trainer.global_rank == 0: print("Summoning checkpoint.") if melk_ckpt_name is None: ckpt_path = os.path.join(ckptdir, "last.ckpt") else: ckpt_path = os.path.join(ckptdir, melk_ckpt_name) trainer.save_checkpoint(ckpt_path) def divein(*args, **kwargs): if trainer.global_rank == 0: import pudb pudb.set_trace() import signal signal.signal(signal.SIGUSR1, melk) signal.signal(signal.SIGUSR2, divein) # run if opt.train: try: trainer.fit(model, data, ckpt_path=ckpt_resume_path) except Exception: if not opt.debug: melk() raise if not opt.no_test and not trainer.interrupted: trainer.test(model, data) except RuntimeError as err: if MULTINODE_HACKS: import datetime import os import socket import requests device = os.environ.get("CUDA_VISIBLE_DEVICES", "?") hostname = socket.gethostname() ts = datetime.datetime.utcnow().strftime("%Y-%m-%d %H:%M:%S") resp = requests.get("http://169.254.169.254/latest/meta-data/instance-id") print( f"ERROR at {ts} on {hostname}/{resp.text} (CUDA_VISIBLE_DEVICES={device}): {type(err).__name__}: {err}", flush=True, ) raise err except Exception: if opt.debug and trainer.global_rank == 0: try: import pudb as debugger except ImportError: import pdb as debugger debugger.post_mortem() raise finally: # move newly created debug project to debug_runs if opt.debug and not opt.resume and trainer.global_rank == 0: dst, name = os.path.split(logdir) dst = os.path.join(dst, "debug_runs", name) os.makedirs(os.path.split(dst)[0], exist_ok=True) os.rename(logdir, dst) if opt.wandb: wandb.finish() # if trainer.global_rank == 0: # print(trainer.profiler.summary())