experiment.py (6198B)
1 """Toolbox for handling experiment metadata.""" 2 3 from __future__ import annotations 4 5 import abc 6 import hashlib 7 import logging 8 import multiprocessing 9 import pathlib 10 import pickle 11 import re 12 import signal 13 import subprocess 14 import sys 15 import time 16 import uuid 17 from typing import Any 18 19 import codecarbon 20 import pynvml 21 22 from t5.config import Config 23 24 25 MAXIMUM_FILENAME_LENGTH: int = 256 26 27 logger: logging.Logger = logging.getLogger(__name__) 28 29 30 class BaseExperiment(abc.ABC): 31 """ 32 Base class for running an experiment. 33 34 Calling run on an instance of this class will call the `init()` then `main()` functions. 35 The sole purpose of this class is to make an experiment "prettier": it displays config values, store a diff of the repo in the experiment logdir, etc. 36 37 Config 38 ------ 39 log_level: str 40 Set log level for stderr (one of DEBUG, INFO, WARNING, ERROR, CRITICAL) 41 42 """ 43 44 def __init__(self: BaseExperiment) -> None: 45 """Create a new experiment, reading the config from argv.""" 46 self.config = Config(sys.argv) 47 self.version: str = self.get_repository_version() 48 self.set_expdir_path() 49 self.setup_logging() 50 self.save_patch() 51 self.save_raw_config() 52 self.hook_signals() 53 self.init() 54 self.check_save_config() 55 56 def __call__(self: BaseExperiment) -> None: 57 """Run the experiment.""" 58 self.run() 59 60 @abc.abstractmethod 61 def init(self: BaseExperiment) -> None: 62 """Initialize experiment (e.g. create model).""" 63 64 @abc.abstractmethod 65 def run(self: BaseExperiment) -> None: 66 """Run experiment (e.g. train model).""" 67 68 def set_expdir_path(self: BaseExperiment) -> None: 69 """Set paths to the experiment log and model directories.""" 70 args: str = " ".join(sys.argv[1:]) 71 stime: str = time.strftime("%FT%H:%M:%S") 72 directory_name: str = f"AEC {self.version} {args} {stime}".replace("/", "_") 73 if len(directory_name) > MAXIMUM_FILENAME_LENGTH: 74 unique_suffix: str = hashlib.sha256(directory_name.encode("utf-8")).hexdigest()[:8] 75 directory_name = f"{directory_name[:MAXIMUM_FILENAME_LENGTH-len(unique_suffix)-1]} {unique_suffix}" 76 self.logdir: pathlib.Path = self.config.LOG_PATH / directory_name 77 self.modeldir: pathlib.Path = self.config.MODEL_PATH / directory_name 78 79 @staticmethod 80 def get_repository_version() -> str: 81 """Get the git repository version.""" 82 source_directory = pathlib.Path(__file__).parents[0] 83 result: subprocess.CompletedProcess[str] = subprocess.run( 84 ["git", "rev-parse", "HEAD"], 85 stdout=subprocess.PIPE, stderr=subprocess.DEVNULL, encoding="utf-8", 86 cwd=source_directory, shell=False, check=False) 87 88 if result.returncode != 0: 89 return "release" 90 commit_hash: str = result.stdout.strip()[:8] 91 92 result = subprocess.run( 93 ["git", "status", "--porcelain"], 94 stdout=subprocess.PIPE, stderr=subprocess.DEVNULL, encoding="utf-8", 95 cwd=source_directory, shell=False, check=True) 96 if re.search(r"^ M ", result.stdout, re.MULTILINE): 97 return f"{commit_hash}+" 98 return commit_hash 99 100 def setup_logging(self: BaseExperiment) -> None: 101 """Initialise the logging system.""" 102 print(f"Log directory is {self.logdir}", file=sys.stderr) 103 self.logdir.mkdir(parents=True, exist_ok=False) 104 root_logger: logging.Logger = logging.getLogger() 105 root_logger.setLevel(logging.NOTSET) 106 107 stderr_handler: logging.Handler = logging.StreamHandler() 108 self.config.declare_option("log_level", str, "Set log level for stderr (one of DEBUG, INFO, WARNING, ERROR, CRITICAL)", "WARNING") 109 stderr_handler.setLevel(getattr(logging, self.config.log_level)) 110 stderr_handler.setFormatter(logging.Formatter("%(levelname)s (%(name)s) %(message)s")) 111 stderr_handler.addFilter(logging.Filter("t5")) 112 root_logger.addHandler(stderr_handler) 113 114 logfile_handler: logging.Handler = logging.FileHandler(self.logdir / "log") 115 logfile_handler.setLevel(logging.INFO) 116 logfile_handler.setFormatter(logging.Formatter("%(asctime)s\t%(levelname)s:%(name)s:%(message)s")) 117 root_logger.addHandler(logfile_handler) 118 119 def save_patch(self: BaseExperiment) -> None: 120 """Save the diff between the current code and the last commit.""" 121 if not self.version.endswith("+"): 122 # The code is a release or wasn't modified 123 return 124 125 logging.warning("Uncommited changes detected, saving patch to logdir.") 126 source_directory = pathlib.Path(__file__).parents[0] 127 output_path: pathlib.Path = self.logdir / "uncommited.patch" 128 with output_path.open("w") as outfile: 129 subprocess.run( 130 ["git", "diff", "HEAD"], 131 stdout=outfile, stderr=subprocess.DEVNULL, encoding="utf-8", 132 cwd=source_directory, shell=False, check=True) 133 134 def save_raw_config(self: BaseExperiment) -> None: 135 """Save the unchecked configuration in a machine format.""" 136 with (self.logdir / "config.pkl").open("wb") as file: 137 pickle.dump(self.config, file) 138 139 def check_save_config(self: BaseExperiment) -> None: 140 """Check the config is valid and save it in a human-readable format.""" 141 self.config.check_schema() 142 self.config.log() 143 self.config.save(self.logdir / "config.py") 144 145 def hook_signals(self: BaseExperiment) -> None: 146 """Change the behavior of SIGINT (^C) to change a variable `self.interrupted' before killing the process.""" 147 self.interrupted: bool = False 148 149 def handler(*_: Any) -> None: 150 if multiprocessing.current_process().name != "MainProcess": 151 return 152 153 logging.critical("Interrupted, attempting to stop execution gracefully.") 154 logging.critical("NEXT ^C WILL KILL THE PROCESS!") 155 self.interrupted = True 156 signal.signal(signal.SIGINT, signal.SIG_DFL) 157 158 signal.signal(signal.SIGINT, handler)