config.py (17949B)
1 """ 2 Configuration using command line arguments and python files. 3 4 Configuration for t5 is first expressed by a schema, declaring options that are needed to describe the model. 5 These options need to be assigned a value using command line arguments. 6 Configuration values can be given as direct variable assignement or through a python module. 7 """ 8 9 from __future__ import annotations 10 11 import importlib.machinery 12 import importlib.util 13 import logging 14 import os 15 import pathlib 16 import types 17 from typing import Any, Iterator 18 19 20 logger: logging.Logger = logging.getLogger(__name__) 21 22 23 class Config(dict[str, Any]): 24 """ 25 A set of configuration values. 26 27 Config objects are dictionaries that can be access through `var.key` instead of `var["key"]`. 28 This class also defines functions for processing command-line arguments. 29 30 Attributes 31 ---------- 32 _root: list[str] 33 The location of the config root, empty for the root config. 34 _lazy_values: dict[str, str] 35 Values that have not been parsed yet because their type is not yet known (because the corresponding key was not declared). 36 _schema: dict[str, type] 37 Type constraints for configuration keys. 38 _help: dict[str, str] 39 Short message describing the purpose of each key. 40 _default: dict[str, Any] 41 Default value for configuration keys, a default does not always need to be declared. 42 _from_environment: set[str] 43 Set of keys that where imported from the environment instead of command line arguments. 44 45 Config 46 ------ 47 DATA_PATH: pathlib.Path 48 Path to the directory containing the datasets. 49 This is an environment variable that can be set by exporting `DATA_PATH`. 50 LOG_PATH: pathlib.Path 51 Path to where logs, metrics and metadata are saved. 52 This is an environment variable that can be set by exporting `LOG_PATH`. 53 MODEL_PATH: pathlib.Path 54 Path to where models are saved. 55 This is an environment variable that can be set by exporting `MODEL_PATH`. 56 """ 57 58 class _LazyMarkerType: 59 """Type used to mark lazy values in Config dictionaries.""" 60 61 def __init__(self: Config, argv: list[str] | None = None, root: list[str] | None = None, lazy_values: dict[str, str] | None = None) -> None: 62 """ 63 Initialize a configuration object with the given argument list. 64 65 Parameters 66 ---------- 67 argv: list[str], optional 68 Command line arguments to extract values from. 69 The first element is assumed to be the program name and is ignored. 70 root: list[str], optional 71 Location of this `Config` object inside the configuration tree. 72 For example the `Config` with root `["model", "generator"]` would contains values such as `model.generator.temperature`. 73 It is empty for the root object. 74 If the root starts with `_test_`, environment variables config keys will not be imported, this is to simplify unit tests. 75 lazy_values: dict[str, str], optional 76 Config values as dictionary of strings, useful when loading from json. 77 78 """ 79 super().__init__() 80 self._root: list[str] = root if root is not None else [] 81 self._lazy_values: dict[str, str] = {} 82 self._schema: dict[str, type] = {} 83 self._help: dict[str, str] = {} 84 self._default: dict[str, Any] = {} 85 self._from_environment: set[str] = set() 86 if argv is not None and lazy_values is not None: 87 msg: str = "both argv and lazy_values given" 88 raise ValueError(msg) 89 if argv is not None: 90 for argument in argv[1:]: 91 self.process_argument(argument) 92 if lazy_values is not None: 93 self._lazy_values.update(lazy_values) 94 95 # Do not require the environment variables to be set when running unit tests. 96 if len(self._root) == 0 or self._root[0] != "_test_": 97 self._import_environment("DATA_PATH", pathlib.Path, "Path to the directory containing the datasets.") 98 self._import_environment("LOG_PATH", pathlib.Path, "Path to where logs, metrics and metadata are saved.") 99 self._import_environment("MODEL_PATH", pathlib.Path, "Path to where models are saved.") 100 101 def __getattr__(self: Config, name: str) -> Any: 102 """Access config values as attributes.""" 103 if name in self.__dict__: 104 return self.__dict__[name] 105 if name not in self: 106 raise AttributeError(self.key_path(name)) 107 return self[name] 108 109 def pretty_string(self: Config) -> str: 110 """Return a colorful representation of the Config.""" 111 return "\n".join(self._pretty_color_lines()) 112 113 def _pretty_color_lines(self: Config, prefix: str = "") -> Iterator[str]: 114 """Describe the Config state as a text lines.""" 115 reset_color: str = "\033[0m" 116 root_color: str = "\033[91m" 117 key_color: str = "\033[92m" 118 value_color: str = "\033[94m" 119 help_color: str = "\033[2m" 120 121 root: str = ".".join(self._root) + " " if self._root else "" 122 yield f"{prefix}{root_color}{root}Config{reset_color}" 123 for key in sorted(self.keys()): 124 value: Any = self[key] 125 if isinstance(value, self._LazyMarkerType): 126 yield f"{prefix} {key_color}{key}{reset_color} {help_color}(LAZY){reset_color} = {value_color}{self._lazy_values[key]}{reset_color}" 127 elif isinstance(value, Config): 128 subconfig: Iterator[str] = value._pretty_color_lines(prefix=" ") 129 if self._help.get(key): 130 yield f"{prefix}{next(subconfig)} {help_color}({self._help[key]}){reset_color}" 131 yield from (f"{prefix}{line}" for line in subconfig) 132 else: 133 type_name: str = self._schema[key].__name__ if key in self._schema else "UNDECLARED" 134 default_msg: str = f" default: {self._default[key]}" if key in self.default else "" 135 help_msg: str = f" {help_color}({self._help[key]}){default_msg}{reset_color}" if self._help.get(key) else "" 136 yield f"{prefix} {key_color}{key}{reset_color} {help_color}({type_name}){reset_color} = {value_color}{value}{reset_color}{help_msg}" 137 138 def _import_environment(self: Config, name: str, cast: type, help_msg: str) -> None: 139 """Import an environment variable into the config namespace.""" 140 try: 141 self[name] = cast(os.environ[f"{name}"]) 142 self._schema[name] = cast 143 self._help[name] = help_msg 144 self._from_environment.add(name) 145 except AttributeError as exception: 146 exception.add_note(f"You need to set the environment variable {name}.") 147 exception.add_note(f"{name}: {help_msg}") 148 raise 149 150 def process_argument(self: Config, argument: str) -> None: 151 """ 152 Process a command line argument. 153 154 Three types of argument are accepted: 155 - Arguments starting with a `--` are untyped *lazy* arguments, that will be converted to the correct type once it is declared. 156 - Arguments starting with a `@` are eval arguments, they are executed as python code where other configuration keys can be directly accessed and modified as python variables. 157 - Other arguments are import arguments, they are treated as paths to python module files. 158 """ 159 if argument.startswith("--"): 160 self._process_lazy_argument(argument[2:]) 161 elif argument.startswith("@"): 162 self._process_eval_argument(argument[1:]) 163 else: 164 self._process_import_argument(argument) 165 166 def _resolve_config_path(self: Config, key: str) -> tuple[Config, str]: 167 """Resolve a path to the corresponding config object and key.""" 168 components: list[str] = key.split(".") 169 config: Config = self 170 for i, component in enumerate(components[:-1]): 171 if component not in config: 172 config[component] = Config(root=[*self._root, *components[:i+1]]) 173 config = config[component] 174 return config, components[-1] 175 176 @staticmethod 177 def _split_assignment(argument: str) -> tuple[str, str | bool]: 178 """ 179 Split an assignment into key and values. 180 181 The expected format is: 182 key(=value)? 183 If no value is provided, the key is assumed to be a boolean and True is considered to be its implicit value. 184 """ 185 operands: list[str] = argument.split("=", maxsplit=1) 186 if len(operands) == 1: 187 return argument, True 188 return operands[0], operands[1] 189 190 def _convert_dicts(self: Config) -> None: 191 """Convert dictionaries to Config objects.""" 192 for key, value in self.items(): 193 if isinstance(value, Config): 194 value._convert_dicts() 195 elif isinstance(value, dict): 196 subconfig = Config(root=[*self._root, key]) 197 for subkey, subvalue in value.items(): 198 subconfig[subkey] = subvalue 199 subconfig._convert_dicts() 200 self[key] = subconfig 201 202 def key_path(self: Config, key: str) -> str: 203 """Return the full path of a config key.""" 204 return ".".join([*self._root, key]) 205 206 def _process_lazy_argument(self: Config, argument: str) -> None: 207 """Process an assignment with a yet unknown type.""" 208 lhs: str 209 value: str | bool 210 lhs, value = self._split_assignment(argument) 211 212 config: Config 213 key: str 214 config, key = self._resolve_config_path(lhs) 215 if isinstance(value, bool): 216 config[key] = True 217 else: 218 config[key] = self._LazyMarkerType() 219 config._lazy_values[key] = value 220 221 def _process_eval_argument(self: Config, argument: str) -> None: 222 """ 223 Evaluate a python argument on the config dictionary. 224 225 When passing a string argument through the shell, it must be enclosed in quote (like all python string), which usually need to be escaped. 226 """ 227 exec(argument, self) 228 self.pop("__builtins__", None) 229 self._convert_dicts() 230 231 def _process_import_argument(self: Config, argument: str) -> None: 232 """ 233 Load python file argument. 234 235 The file is loaded in an independent context, all the variable defined in the file (even through import) are added to the configuration, with the exception of builtins and whole modules. 236 """ 237 spec: importlib.machinery.ModuleSpec | None = importlib.util.spec_from_file_location("argument", argument) 238 if spec is None or spec.loader is None: 239 raise ImportError(argument) 240 module: types.ModuleType = importlib.util.module_from_spec(spec) 241 spec.loader.exec_module(module) 242 key: str 243 value: Any 244 for key, value in vars(module).items(): 245 if ( 246 key not in module.__builtins__ 247 and not key.startswith("_") 248 and not isinstance(value, types.ModuleType) 249 ): 250 self[key] = value 251 self._convert_dicts() 252 253 def _resolve_lazy(self: Config, key: str) -> None: 254 """Cast lazily typed value to schema type.""" 255 if self._schema[key] in [str, int, float, bool, pathlib.Path]: 256 try: 257 self[key] = self._schema[key](self._lazy_values[key]) 258 self._lazy_values.pop(key) 259 except ValueError as exception: 260 raise InvalidValueError(self.key_path(key), self._schema[key], self._help[key], self[key]) from exception 261 else: 262 raise NotImplementedError(f"Lazy {self._schema[key]} config") 263 264 def declare_option(self: Config, full_key: str, type_constraint: type, help_msg: str, default: None | Any = None) -> None: 265 """ 266 Add an option to the configuration schema. 267 268 Parameters 269 ---------- 270 full_key: str 271 Configuration key to declare such as `model.generator.temperature`. 272 It is preferable to directly declare each option in its direct parent's config, so more often than not, this should be a dot-less key. 273 type_constraint: type 274 What is the type of the configuration key being declared. 275 help_msg: str 276 An explanation of what the configuration key is for. 277 default: type_constraint, optional 278 The value the configuration key will take if it is not assigned a value. 279 If a default is not given and the option was not defined by the command-line, this function will raise an exception. 280 281 Raises 282 ------ 283 MissingValueError 284 If the command-line arguments do not define the declared option when no default is provided. 285 MistypedValueError 286 The configuration key was defined by the command-line arguments with a different type from the one being declared. 287 InvalidValueError 288 The configuration key was lazily defined by the command-line argument but could not be converted to the declared type. 289 290 """ 291 config: Config 292 key: str 293 config, key = self._resolve_config_path(full_key) 294 295 config._schema[key] = type_constraint 296 config._help[key] = help_msg 297 if default is not None: 298 config._default[key] = default 299 300 if key not in config: 301 if default is None: 302 raise MissingValueError(config.key_path(key), type_constraint, help_msg) 303 config[key] = default 304 305 if isinstance(config[key], self._LazyMarkerType): 306 config._resolve_lazy(key) 307 if not isinstance(config[key], type_constraint): 308 raise MistypedValueError(config.key_path(key), type_constraint, help_msg, config[key]) 309 310 def check_schema(self: Config) -> None: 311 """Check whether the config adhere to the schema.""" 312 for key, value in self.items(): 313 if key not in self._schema: 314 raise UndeclaredValueError(self.key_path(key), value) 315 if isinstance(value, Config): 316 value.check_schema() 317 318 def log(self: Config) -> None: 319 """Write the config to log file.""" 320 for key in sorted(self): 321 value: Any = self[key] 322 if isinstance(value, Config): 323 value.log() 324 else: 325 logger.info("value %s %s", self.key_path(key), value) 326 327 def save(self: Config, output_path: pathlib.Path) -> None: 328 """Save the compiled configuration to a python file.""" 329 imports: set[str] 330 code: list[str] 331 imports, code = self._python_representation() 332 with output_path.open("w") as output: 333 print("\n".join(imports), file=output) 334 print("\n\n", end="", file=output) 335 print("\n".join(code), file=output) 336 337 def _python_representation(self: Config) -> tuple[set[str], list[str]]: 338 """Get a representation of the configuration as python code.""" 339 imports: set[str] = set() 340 code: list[str] = [] 341 for key in sorted(self): 342 if key in self._from_environment: 343 continue 344 value: Any = self[key] 345 full_key: str = key 346 if self._root: 347 indices: str = "".join(f"[\"{part}\"]" for part in [*self._root[1:], key]) 348 full_key = f"{self._root[0]}{indices}" 349 350 if isinstance(value, Config): 351 code.append(f"{full_key} = dict()") 352 sub_imports: set[str] 353 sub_code: list[str] 354 sub_imports, sub_code = value._python_representation() 355 imports.update(sub_imports) 356 code.extend(sub_code) 357 else: 358 if isinstance(value, pathlib.PosixPath): 359 imports.add("from pathlib import PosixPath") 360 code.append(f"{full_key} = {value!r}") 361 return imports, code 362 363 364 class ConfigError(Exception): 365 """Base Exception class for t5.config.""" 366 367 368 class MissingValueError(ConfigError): 369 """Exception raised when a needed config value is not provided as command line argument.""" 370 371 def __init__(self: MissingValueError, key: str, type_constraint: type, help_msg: str) -> None: 372 """Initiate a MissingValueError with given meta information.""" 373 super().__init__(f"Missing configuration value for: {key}\nType constraint: {type_constraint}\nHelp: {help_msg}") 374 375 376 class MistypedValueError(ConfigError): 377 """Exception raised when a config value is given the wrong type.""" 378 379 def __init__(self: MistypedValueError, key: str, type_constraint: type, help_msg: str, value: Any) -> None: 380 """Initiate a MistypedValueError with given meta information and wrong value.""" 381 super().__init__(f"Wrong type for configuration value: {key}\nType constraint: {type_constraint}\nActual type was {type(value)} from value {value}\nHelp: {help_msg}") 382 383 384 class InvalidValueError(ConfigError): 385 """Exception raised when a config value cannot be converted to the correct type.""" 386 387 def __init__(self: InvalidValueError, key: str, type_constraint: type, help_msg: str, value: Any) -> None: 388 """Initiate an InvalidValueError with given meta information and wrong value.""" 389 super().__init__(f"Cannot convert configuration value: {key}\nType constraint: {type_constraint}\nAssigned value: {value}\nHelp: {help_msg}") 390 391 392 class UndeclaredValueError(ConfigError): 393 """Exception raised when a config value is given but not declared by the program.""" 394 395 def __init__(self: UndeclaredValueError, key: str, value: Any) -> None: 396 """Initiate an UndeclaredValueError with given meta information.""" 397 super().__init__(f"Extra configuration value was provided but it is undeclared: {key}\nAssigned value: {value}")