File size: 12,642 Bytes
d1ceb73 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 381 382 383 384 |
# mypy: allow-untyped-defs
import contextlib
import copy
import hashlib
import inspect
import io
import pickle
import tokenize
import unittest
from types import FunctionType, ModuleType
from typing import Any, Dict, Optional, Set, Union
from typing_extensions import deprecated
from unittest import mock
# Types saved/loaded in configs
CONFIG_TYPES = (int, float, bool, type(None), str, list, set, tuple, dict)
def install_config_module(module):
"""
Converts a module-level config into a `ConfigModule()`.
See _config_typing.pyi for instructions on how to get the converted module to typecheck.
"""
class ConfigModuleInstance(ConfigModule):
_bypass_keys = set({"_is_dirty", "_hash_digest"})
def visit(source, dest, prefix):
"""Walk the module structure and move everything to module._config"""
for key, value in list(source.__dict__.items()):
if (
key.startswith("__")
or isinstance(value, (ModuleType, FunctionType))
or (hasattr(value, "__module__") and value.__module__ == "typing")
):
continue
name = f"{prefix}{key}"
if isinstance(value, CONFIG_TYPES):
config[name] = value
default[name] = value
if dest is module:
delattr(module, key)
elif isinstance(value, type):
assert value.__module__ == module.__name__
# a subconfig with `class Blah:` syntax
proxy = SubConfigProxy(module, f"{name}.")
visit(value, proxy, f"{name}.")
setattr(dest, key, proxy)
else:
raise AssertionError(f"Unhandled config {key}={value} ({type(value)})")
config: Dict[str, Any] = dict()
default: Dict[str, Any] = dict()
compile_ignored_keys = get_assignments_with_compile_ignored_comments(module)
visit(module, module, "")
module._config = config
module._default = default
module._allowed_keys = set(config.keys())
module._compile_ignored_keys = compile_ignored_keys
module.__class__ = ConfigModuleInstance
module._is_dirty = True
module._hash_digest = None
COMPILE_IGNORED_MARKER = "@compile_ignored"
# Gets all the keys (i.e. assignments) with a @compile_ignored comment
def get_assignments_with_compile_ignored_comments(module):
source_code = inspect.getsource(module)
assignments = set()
# Tokenize the source code to retrieve comments
tokens = tokenize.tokenize(io.BytesIO(source_code.encode("utf-8")).readline)
current_comment = "", -1
prev_name = ""
for token in tokens:
if token.type == tokenize.COMMENT:
prev_name = ""
maybe_current = token.string.strip()
if COMPILE_IGNORED_MARKER in maybe_current:
assert current_comment == (
"",
-1,
), f"unconsumed {COMPILE_IGNORED_MARKER}"
current_comment = maybe_current, token.start[0]
elif token.type == tokenize.NAME:
# Only accept the first name token, to handle if you have
# something like foo: Bar = ...
if not prev_name:
prev_name = token.string
elif token.type == tokenize.OP and token.string == "=":
# Check if the current assignment follows a comment
# with COMPILE_IGNORED_MARKER
if (
COMPILE_IGNORED_MARKER in current_comment[0]
and current_comment[1] == token.start[0] - 1
):
assignments.add(prev_name)
current_comment = "", -1 # reset
prev_name = ""
assert current_comment == ("", -1), f"unconsumed {COMPILE_IGNORED_MARKER}"
return assignments
class ConfigModule(ModuleType):
# NOTE: This should be kept in sync with _config_typing.pyi.
# The default values of the configuration settings. This can be used to
# determine if the config has been changed or not.
_default: Dict[str, Any]
# The actual configuration settings. E.g., torch._dynamo.config.debug
# would live as "debug" in the key, and torch._inductor.config.triton.cudagraphs
# maps as "triton.cudagraphs"
_config: Dict[str, Any]
_allowed_keys: Set[str]
_bypass_keys: Set[str]
_compile_ignored_keys: Set[str]
_is_dirty: bool
_hash_digest: Optional[bytes]
def __init__(self):
raise NotImplementedError(
f"use {__name__}.install_config_module(sys.modules[__name__])"
)
def __setattr__(self, name, value):
if name in self._bypass_keys:
super().__setattr__(name, value)
elif name not in self._allowed_keys:
raise AttributeError(f"{self.__name__}.{name} does not exist")
else:
self._config[name] = value
def __getattr__(self, name):
try:
return self._config[name]
except KeyError as e:
# make hasattr() work properly
raise AttributeError(f"{self.__name__}.{name} does not exist") from e
def __delattr__(self, name):
# must support delete because unittest.mock.patch deletes
# then recreate things
del self._config[name]
def save_config(self) -> bytes:
"""Convert config to a pickled blob"""
config = dict(self._config)
for key in config.get("_save_config_ignore", ()):
config.pop(key)
return pickle.dumps(config, protocol=2)
def save_config_portable(self) -> Dict[str, Any]:
"""Convert config to portable format"""
config: Dict[str, Any] = {}
for key in sorted(self._config):
if key.startswith("_"):
continue
if any(
key.startswith(e) for e in self._config["_cache_config_ignore_prefix"]
):
continue
config[key] = self._config[key]
return config
def codegen_config(self) -> str:
"""Convert config to Python statements that replicate current config.
This does NOT include config settings that are at default values.
"""
lines = []
mod = self.__name__
for k, v in self._config.items():
if k in self._config.get("_save_config_ignore", ()):
continue
if v == self._default[k]:
continue
lines.append(f"{mod}.{k} = {v!r}")
return "\n".join(lines)
def get_hash(self) -> bytes:
"""Hashes the configs that are not compile_ignored"""
if self._is_dirty or self._hash_digest is None:
dict_to_hash = {
k: v
for k, v in self._config.items()
if k not in self._compile_ignored_keys
}
string_to_hash = repr(sorted(dict_to_hash.items()))
self._hash_digest = hashlib.md5(string_to_hash.encode("utf-8")).digest()
self._is_dirty = False
return self._hash_digest
@deprecated(
"`config.to_dict()` has been deprecated. It may no longer change the underlying config."
" use `config.shallow_copy_dict()` or `config.get_config_copy()` instead",
category=FutureWarning,
)
def to_dict(self) -> Dict[str, Any]:
return self.shallow_copy_dict()
def shallow_copy_dict(self) -> Dict[str, Any]:
return {**self._config}
def load_config(self, maybe_pickled_config: Union[bytes, Dict[str, Any]]) -> None:
"""Restore from a prior call to save_config() or shallow_copy_dict()"""
if not isinstance(maybe_pickled_config, dict):
config = pickle.loads(maybe_pickled_config)
else:
config = maybe_pickled_config
self._config.update(config)
def get_config_copy(self) -> Dict[str, Any]:
return copy.deepcopy(self._config)
def patch(
self,
arg1: Optional[Union[str, Dict[str, Any]]] = None,
arg2: Any = None,
**kwargs,
):
"""
Decorator and/or context manager to make temporary changes to a config.
As a decorator:
@config.patch("name", val)
@config.patch(name1=val1, name2=val2)
@config.patch({"name1": val1, "name2", val2})
def foo(...):
...
As a context manager:
with config.patch("name", val):
...
"""
changes: Dict[str, Any]
if arg1 is not None:
if arg2 is not None:
assert isinstance(arg1, str)
# patch("key", True) syntax
changes = {arg1: arg2}
else:
assert isinstance(arg1, dict)
# patch({"key": True}) syntax
changes = arg1
assert not kwargs
else:
# patch(key=True) syntax
changes = kwargs
assert arg2 is None
assert isinstance(changes, dict), f"expected `dict` got {type(changes)}"
prior: Dict[str, Any] = {}
config = self
dirty = False
class ConfigPatch(ContextDecorator):
def __enter__(self):
assert not prior
nonlocal dirty
for key in changes.keys():
# KeyError on invalid entry
prior[key] = config._config[key]
dirty = key not in config._compile_ignored_keys
config._config.update(changes)
config._is_dirty = dirty
def __exit__(self, exc_type, exc_val, exc_tb):
nonlocal dirty
config._config.update(prior)
config._is_dirty = dirty
prior.clear()
return ConfigPatch()
def _make_closure_patcher(self, **changes):
"""
A lower-overhead version of patch() for things on the critical path.
Usage:
# do this off the critical path
change_fn = config.make_closure_patcher(foo=True)
...
revert = change_fn()
try:
...
finally:
revert()
"""
config = self._config
def change():
prior = {k: config[k] for k in changes}
config.update(changes)
def revert():
config.update(prior)
return revert
return change
class ContextDecorator(contextlib.ContextDecorator):
"""
Same as contextlib.ContextDecorator, but with support for
`unittest.TestCase`
"""
def __enter__(self):
raise NotImplementedError("NYI")
def __exit__(self, exc_type, exc_val, exc_tb):
raise NotImplementedError("NYI")
def __call__(self, func):
if isinstance(func, type) and issubclass(func, unittest.TestCase):
class _TestCase(func): # type: ignore[valid-type, misc]
@classmethod
def setUpClass(cls):
self.__enter__()
try:
super().setUpClass()
except Exception:
self.__exit__(None, None, None)
raise
@classmethod
def tearDownClass(cls):
try:
super().tearDownClass()
finally:
self.__exit__(None, None, None)
_TestCase.__name__ = func.__name__
_TestCase.__qualname__ = func.__qualname__
_TestCase.__module__ = func.__module__
return _TestCase
return super().__call__(func)
class SubConfigProxy:
"""
Shim to redirect to main config.
`config.triton.cudagraphs` maps to _config["triton.cudagraphs"]
"""
def __init__(self, config, prefix):
# `super().__setattr__` to bypass custom `__setattr__`
super().__setattr__("_config", config)
super().__setattr__("_prefix", prefix)
def __setattr__(self, name, value):
return self._config.__setattr__(self._prefix + name, value)
def __getattr__(self, name):
return self._config.__getattr__(self._prefix + name)
def __delattr__(self, name):
return self._config.__delattr__(self._prefix + name)
def patch_object(obj, name, value):
"""
Workaround `mock.patch.object` issue with ConfigModule
"""
if isinstance(obj, ConfigModule):
return obj.patch(name, value)
return mock.patch.object(obj, name, value)
|