_config_module.py 31 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836
  1. import contextlib
  2. import copy
  3. import hashlib
  4. import importlib
  5. import inspect
  6. import io
  7. import os
  8. import pickle
  9. import tokenize
  10. import unittest
  11. from collections.abc import Callable
  12. from dataclasses import dataclass
  13. from types import FunctionType, ModuleType
  14. from typing import Any, Generic, NoReturn, Optional, TYPE_CHECKING, TypeVar
  15. from typing_extensions import deprecated
  16. from unittest import mock
  17. from torch._utils_internal import justknobs_check
  18. # Types saved/loaded in configs
  19. CONFIG_TYPES = (int, float, bool, type(None), str, list, set, tuple, dict)
  20. # Duplicated, because mypy needs these types statically
  21. T = TypeVar("T", bound=int | float | bool | str | list | set | tuple | dict | None)
  22. _UNSET_SENTINEL = object()
  23. @dataclass(kw_only=True)
  24. class _Config(Generic[T]):
  25. """Represents a config with richer behaviour than just a default value.
  26. ::
  27. i.e.
  28. foo = Config(justknob="//foo:bar", default=False)
  29. install_config_module(...)
  30. This configs must be installed with install_config_module to be used
  31. Precedence Order:
  32. alias: If set, the directly use the value of the alias.
  33. env_name_force: If set, this environment variable has precedence over
  34. everything after this.
  35. If multiple env variables are given, the precedence order is from
  36. left to right.
  37. user_override: If a user sets a value (i.e. foo.bar=True), that
  38. has precedence over everything after this.
  39. env_name_default: If set, this environment variable will override everything
  40. after this.
  41. If multiple env variables are given, the precedence order is from
  42. left to right.
  43. justknob: If this pytorch installation supports justknobs, that will
  44. override defaults, but will not override the user_override precedence.
  45. default: This value is the lowest precedence, and will be used if nothing is
  46. set.
  47. Environment Variables:
  48. These are interpreted to be either "0" or "1" to represent true and false.
  49. Arguments:
  50. justknob: the name of the feature / JK. In OSS this is unused.
  51. default: is the value to default this knob to in OSS.
  52. alias: The alias config to read instead.
  53. env_name_force: The environment variable, or list of, to read that is a FORCE
  54. environment variable. I.e. it overrides everything except for alias.
  55. env_name_default: The environment variable, or list of, to read that changes the
  56. default behaviour. I.e. user overrides take preference.
  57. """
  58. default: T | object
  59. justknob: str | None = None
  60. env_name_default: list[str] | None = None
  61. env_name_force: list[str] | None = None
  62. value_type: type | None = None
  63. alias: str | None = None
  64. def __post_init__(self) -> None:
  65. self.env_name_default = _Config.string_or_list_of_string_to_list(
  66. self.env_name_default
  67. )
  68. self.env_name_force = _Config.string_or_list_of_string_to_list(
  69. self.env_name_force
  70. )
  71. if self.alias is not None:
  72. if (
  73. self.default is not _UNSET_SENTINEL
  74. or self.justknob is not None
  75. or self.env_name_default is not None
  76. or self.env_name_force is not None
  77. ):
  78. raise AssertionError(
  79. "if alias is set, none of {default, justknob, \
  80. env_name_default and env_name_force} can be set"
  81. )
  82. @staticmethod
  83. def string_or_list_of_string_to_list(
  84. val: str | list[str] | None,
  85. ) -> list[str] | None:
  86. if val is None:
  87. return None
  88. if isinstance(val, str):
  89. return [val]
  90. if not isinstance(val, list):
  91. raise AssertionError(f"val is not a list, got {type(val)}")
  92. return val
  93. # In runtime, we unbox the Config[T] to a T, but typechecker cannot see this,
  94. # so in order to allow for this dynamic behavior to work correctly with
  95. # typechecking we are going to lie to the typechecker that Config[T] returns
  96. # a T.
  97. if TYPE_CHECKING:
  98. def Config(
  99. default: T | object = _UNSET_SENTINEL,
  100. justknob: str | None = None,
  101. env_name_default: str | list[str] | None = None,
  102. env_name_force: str | list[str] | None = None,
  103. value_type: type | None = None,
  104. alias: str | None = None,
  105. ) -> T: ...
  106. else:
  107. def Config(
  108. default: T | object = _UNSET_SENTINEL,
  109. justknob: str | None = None,
  110. env_name_default: str | list[str] | None = None,
  111. env_name_force: str | list[str] | None = None,
  112. value_type: type | None = None,
  113. alias: str | None = None,
  114. ) -> _Config[T]:
  115. return _Config(
  116. default=default,
  117. justknob=justknob,
  118. env_name_default=env_name_default,
  119. env_name_force=env_name_force,
  120. value_type=value_type,
  121. alias=alias,
  122. )
  123. def _read_env_variable(name: str) -> bool | str | None:
  124. value = os.environ.get(name)
  125. if value == "1":
  126. return True
  127. if value == "0":
  128. return False
  129. return value
  130. def install_config_module(module: ModuleType) -> None:
  131. """
  132. Converts a module-level config into a `ConfigModule()`.
  133. See _config_typing.pyi for instructions on how to get the converted module to typecheck.
  134. """
  135. class ConfigModuleInstance(ConfigModule):
  136. # __annotations__ is written to by Sphinx autodoc
  137. _bypass_keys = set({"_is_dirty", "_hash_digest", "__annotations__"})
  138. def visit(
  139. source: ModuleType | type,
  140. dest: ModuleType | SubConfigProxy,
  141. prefix: str,
  142. ) -> None:
  143. """Walk the module structure and move everything to module._config"""
  144. type_hints = inspect.get_annotations(source)
  145. for key, value in list(source.__dict__.items()):
  146. if (
  147. key.startswith("__")
  148. or isinstance(value, (ModuleType, FunctionType))
  149. or (
  150. hasattr(value, "__module__")
  151. and (
  152. value.__module__ == "typing"
  153. or value.__module__.startswith("collections.abc")
  154. )
  155. )
  156. # Handle from torch.utils._config_module import Config
  157. or (isinstance(value, type) and issubclass(value, _Config))
  158. ):
  159. continue
  160. name = f"{prefix}{key}"
  161. annotated_type = type_hints.get(key, None)
  162. if isinstance(value, CONFIG_TYPES):
  163. config[name] = _ConfigEntry(
  164. _Config(default=value, value_type=annotated_type)
  165. )
  166. if dest is module:
  167. delattr(module, key)
  168. elif isinstance(value, _Config):
  169. if annotated_type is not None and value.value_type is None:
  170. value.value_type = annotated_type
  171. config[name] = _ConfigEntry(value)
  172. if dest is module:
  173. delattr(module, key)
  174. elif isinstance(value, type):
  175. if value.__module__ != module.__name__:
  176. raise AssertionError(
  177. f"subconfig class {value} must be defined in module {module.__name__}"
  178. )
  179. # a subconfig with `class Blah:` syntax
  180. proxy = SubConfigProxy(module, f"{name}.")
  181. visit(value, proxy, f"{name}.")
  182. if dest is module:
  183. setattr(dest, key, proxy)
  184. else:
  185. dest.__dict__[key] = proxy
  186. else:
  187. raise AssertionError(f"Unhandled config {key}={value} ({type(value)})")
  188. config: dict[str, _ConfigEntry] = {}
  189. compile_ignored_keys = get_assignments_with_compile_ignored_comments(module)
  190. visit(module, module, "")
  191. module._config = config # type: ignore[attr-defined]
  192. module._compile_ignored_keys = compile_ignored_keys # type: ignore[attr-defined]
  193. module.__class__ = ConfigModuleInstance
  194. module._is_dirty = True # type: ignore[attr-defined]
  195. module._hash_digest = None # type: ignore[attr-defined]
  196. COMPILE_IGNORED_MARKER = "@compile_ignored"
  197. # Gets all the keys (i.e. assignments) with a @compile_ignored comment
  198. def get_assignments_with_compile_ignored_comments(module: ModuleType) -> set[str]:
  199. source_code = inspect.getsource(module)
  200. assignments = set()
  201. # Tokenize the source code to retrieve comments
  202. tokens = tokenize.tokenize(io.BytesIO(source_code.encode("utf-8")).readline)
  203. current_comment = "", -1
  204. prev_name = ""
  205. for token in tokens:
  206. if token.type == tokenize.COMMENT:
  207. prev_name = ""
  208. maybe_current = token.string.strip()
  209. if COMPILE_IGNORED_MARKER in maybe_current:
  210. if current_comment != ("", -1):
  211. raise AssertionError(f"unconsumed {COMPILE_IGNORED_MARKER}")
  212. current_comment = maybe_current, token.start[0]
  213. elif token.type == tokenize.NAME:
  214. # Only accept the first name token, to handle if you have
  215. # something like foo: Bar = ...
  216. if not prev_name:
  217. prev_name = token.string
  218. elif token.type == tokenize.OP and token.string == "=":
  219. # Check if the current assignment follows a comment
  220. # with COMPILE_IGNORED_MARKER
  221. if (
  222. COMPILE_IGNORED_MARKER in current_comment[0]
  223. and current_comment[1] == token.start[0] - 1
  224. ):
  225. assignments.add(prev_name)
  226. current_comment = "", -1 # reset
  227. prev_name = ""
  228. if current_comment != ("", -1):
  229. raise AssertionError(f"unconsumed {COMPILE_IGNORED_MARKER}")
  230. return assignments
  231. @dataclass
  232. class _ConfigEntry:
  233. # The default value specified in the configuration
  234. default: Any
  235. # The type of the configuration value
  236. value_type: type
  237. # The value specified by the user when they overrode the configuration
  238. # _UNSET_SENTINEL indicates the value is not set.
  239. user_override: Any = _UNSET_SENTINEL
  240. # The justknob to check for this config
  241. justknob: str | None = None
  242. # environment variables are read at install time
  243. env_value_force: Any = _UNSET_SENTINEL
  244. env_value_default: Any = _UNSET_SENTINEL
  245. # Used to work arounds bad assumptions in unittest.mock.patch
  246. # The code to blame is
  247. # https://github.com/python/cpython/blob/94a7a4e22fb8f567090514785c69e65298acca42/Lib/unittest/mock.py#L1637
  248. # Essentially, mock.patch requires, that if __dict__ isn't accessible
  249. # (which it isn't), that after delattr is called on the object, the
  250. # object must throw when hasattr is called. Otherwise, it doesn't call
  251. # setattr again.
  252. # Technically we'll have an intermediate state of hiding the config while
  253. # mock.patch is unpatching itself, but it calls setattr after the delete
  254. # call so the final state is correct. It's just very unintuitive.
  255. # upstream bug - python/cpython#126886
  256. hide: bool = False
  257. alias: str | None = None
  258. def __init__(self, config: _Config) -> None:
  259. self.default = config.default
  260. self.value_type = (
  261. config.value_type if config.value_type is not None else type(self.default)
  262. )
  263. self.justknob = config.justknob
  264. self.alias = config.alias
  265. if config.env_name_default is not None:
  266. for val in config.env_name_default:
  267. if (env_value := _read_env_variable(val)) is not None:
  268. self.env_value_default = env_value
  269. break
  270. if config.env_name_force is not None:
  271. for val in config.env_name_force:
  272. if (env_value := _read_env_variable(val)) is not None:
  273. self.env_value_force = env_value
  274. break
  275. # Ensure justknobs and envvars are allowlisted types
  276. if self.justknob is not None and self.default is not None:
  277. if not isinstance(self.default, bool):
  278. raise AssertionError(
  279. f"justknobs only support booleans, {self.default} is not a boolean"
  280. )
  281. if self.value_type is not None and (
  282. config.env_name_default is not None or config.env_name_force is not None
  283. ):
  284. if self.value_type not in (
  285. bool,
  286. str,
  287. Optional[bool],
  288. Optional[str],
  289. ):
  290. raise AssertionError(
  291. f"envvar configs only support (optional) booleans or strings, {self.value_type} is neither"
  292. )
  293. class ConfigModule(ModuleType):
  294. # NOTE: This should be kept in sync with _config_typing.pyi.
  295. # The actual configuration settings. E.g., torch._dynamo.config.debug
  296. # would live as "debug" in the key, and torch._inductor.config.triton.cudagraphs
  297. # maps as "triton.cudagraphs". See discussion on the class for meaning of various sub items
  298. _config: dict[str, _ConfigEntry]
  299. _bypass_keys: set[str]
  300. _compile_ignored_keys: set[str]
  301. _is_dirty: bool
  302. _hash_digest: bytes | None
  303. def __init__(self) -> None:
  304. raise NotImplementedError(
  305. f"use {__name__}.install_config_module(sys.modules[__name__])"
  306. )
  307. def __setattr__(self, name: str, value: object) -> None:
  308. if name in self._bypass_keys:
  309. super().__setattr__(name, value)
  310. elif name not in self._config:
  311. raise AttributeError(f"{self.__name__}.{name} does not exist")
  312. elif self._config[name].alias is not None:
  313. self._set_alias_val(self._config[name], value)
  314. else:
  315. self._config[name].user_override = value
  316. self._is_dirty = True
  317. self._config[name].hide = False
  318. def __getattr__(self, name: str) -> Any:
  319. try:
  320. config = self._config[name]
  321. if config.hide:
  322. raise AttributeError(f"{self.__name__}.{name} does not exist")
  323. alias_val = self._get_alias_val(config)
  324. if alias_val is not _UNSET_SENTINEL:
  325. return alias_val
  326. if config.env_value_force is not _UNSET_SENTINEL:
  327. return config.env_value_force
  328. if config.user_override is not _UNSET_SENTINEL:
  329. return config.user_override
  330. if config.env_value_default is not _UNSET_SENTINEL:
  331. return config.env_value_default
  332. if config.justknob is not None:
  333. # JK only supports bools and ints
  334. return justknobs_check(name=config.justknob, default=config.default)
  335. # Note that reference types can still be modified, so we
  336. # copy them to user_overrides in case the user overrides
  337. # them
  338. if isinstance(config.default, (list, set, dict)):
  339. config.user_override = copy.deepcopy(config.default)
  340. return config.user_override
  341. return config.default
  342. except KeyError as e:
  343. # make hasattr() work properly
  344. raise AttributeError(f"{self.__name__}.{name} does not exist") from e
  345. def __delattr__(self, name: str) -> None:
  346. self._is_dirty = True
  347. # must support delete because unittest.mock.patch deletes
  348. # then recreate things
  349. self._config[name].user_override = _UNSET_SENTINEL
  350. self._config[name].hide = True
  351. def _get_alias_module_and_name(
  352. self, entry: _ConfigEntry
  353. ) -> tuple[ModuleType, str] | None:
  354. alias = entry.alias
  355. if alias is None:
  356. return None
  357. module_name, constant_name = alias.rsplit(".", 1)
  358. try:
  359. module = importlib.import_module(module_name)
  360. except ImportError as e:
  361. raise AttributeError(f"config alias {alias} does not exist") from e
  362. return module, constant_name
  363. def _get_alias_val(self, entry: _ConfigEntry) -> Any:
  364. data = self._get_alias_module_and_name(entry)
  365. if data is None:
  366. return _UNSET_SENTINEL
  367. module, constant_name = data
  368. constant_value = getattr(module, constant_name)
  369. return constant_value
  370. def _set_alias_val(self, entry: _ConfigEntry, val: Any) -> None:
  371. data = self._get_alias_module_and_name(entry)
  372. if data is None:
  373. raise AssertionError(
  374. "alias data should not be None when setting alias value"
  375. )
  376. module, constant_name = data
  377. setattr(module, constant_name, val)
  378. def _is_default(self, name: str) -> bool:
  379. """
  380. Returns true if the config is at its default value.
  381. configs overridden by the env are not considered default.
  382. """
  383. config_val = self._config[name]
  384. # The config is not overridden by the user, and the env_value_default
  385. # is different from the default value (meaning user has set the env to
  386. # change the default value).
  387. not_set_env_default = (
  388. config_val.env_value_default is _UNSET_SENTINEL
  389. or config_val.env_value_default == config_val.default
  390. )
  391. not_set_env_force = (
  392. config_val.env_value_force is _UNSET_SENTINEL
  393. or config_val.env_value_force == config_val.default
  394. )
  395. unset = config_val.user_override is _UNSET_SENTINEL
  396. # Handle reference types specially to avoid spammy warnings
  397. if isinstance(config_val.default, (list, set, dict)):
  398. unset = unset or config_val.user_override == config_val.default
  399. return unset and not_set_env_default and not_set_env_force
  400. def _get_dict(
  401. self,
  402. ignored_keys: list[str] | None = None,
  403. ignored_prefixes: list[str] | None = None,
  404. skip_default: bool = False,
  405. ) -> dict[str, Any]:
  406. """Export a dictionary of current configuration keys and values.
  407. This function is design to provide a single point which handles
  408. accessing config options and exporting them into a dictionary.
  409. This is used by a number of different user facing export methods
  410. which all have slightly different semantics re: how and what to
  411. skip.
  412. If a config is aliased, it skips this config.
  413. Arguments:
  414. ignored_keys are keys that should not be exported.
  415. ignored_prefixes are prefixes that if a key matches should
  416. not be exported
  417. skip_default does two things. One if a key has not been modified
  418. it skips it.
  419. """
  420. config: dict[str, Any] = {}
  421. for key in self._config:
  422. if ignored_keys and key in ignored_keys:
  423. continue
  424. if ignored_prefixes:
  425. if any(key.startswith(prefix) for prefix in ignored_prefixes):
  426. continue
  427. if skip_default and self._is_default(key):
  428. continue
  429. if self._config[key].alias is not None:
  430. continue
  431. config[key] = copy.deepcopy(getattr(self, key))
  432. return config
  433. def get_type(self, config_name: str) -> type:
  434. return self._config[config_name].value_type
  435. def save_config(self) -> bytes:
  436. """Convert config to a pickled blob"""
  437. ignored_keys = getattr(self, "_save_config_ignore", [])
  438. return pickle.dumps(
  439. self._get_dict(ignored_keys=ignored_keys),
  440. protocol=2,
  441. )
  442. def save_config_portable(
  443. self, *, ignore_private_configs: bool = True
  444. ) -> dict[str, Any]:
  445. """Convert config to portable format"""
  446. prefixes = []
  447. if ignore_private_configs:
  448. prefixes.append("_")
  449. prefixes.extend(getattr(self, "_cache_config_ignore_prefix", []))
  450. return self._get_dict(ignored_prefixes=prefixes)
  451. def codegen_config(self) -> str:
  452. """Convert config to Python statements that replicate current config.
  453. This does NOT include config settings that are at default values.
  454. """
  455. # additional imports required
  456. imports = set()
  457. def get_module_name(func: Callable, add_dot: bool) -> str:
  458. module_name = func.__module__
  459. if module_name == "builtins":
  460. module_name = ""
  461. if add_dot and module_name != "":
  462. module_name += "."
  463. return module_name
  464. def add_import(func: Callable) -> None:
  465. module_name = get_module_name(func, False)
  466. if module_name:
  467. imports.add(module_name)
  468. def list_of_callables_to_string(v: list | set) -> list[str]:
  469. return [f"{get_module_name(item, True)}{item.__name__}" for item in v]
  470. def importable_callable(v: Any) -> bool:
  471. # functools.partial has no attributes below but is a callable
  472. return callable(v) and hasattr(v, "__module__") and hasattr(v, "__name__")
  473. def get_config_line(mod, k, v) -> str: # type: ignore[no-untyped-def]
  474. """
  475. Return a string version of the config line.
  476. Handle v when v is a callable, or a list/dict of callables. Add import statements for callables if necessary.
  477. We assume that the value of a single config won't be a mix of callables and non-callables.
  478. Example output:
  479. import logging
  480. import _warnings
  481. torch._dynamo.config.reorderable_logging_functions = { _warnings.warn, logging.warn, print }
  482. """
  483. if importable_callable(v):
  484. add_import(v)
  485. return f"{mod}.{k} = {get_module_name(v, True)}{v.__name__}"
  486. elif isinstance(v, (list, set)) and all(
  487. importable_callable(item) for item in v
  488. ):
  489. for item in v:
  490. add_import(item)
  491. v_list = list_of_callables_to_string(v)
  492. if isinstance(v, list):
  493. return f"{mod}.{k} = {v_list}"
  494. else:
  495. return f"{mod}.{k} = {{ {', '.join(v_list)} }}"
  496. else:
  497. return f"{mod}.{k} = {v!r}"
  498. lines = []
  499. mod = self.__name__
  500. for k, v in self._get_dict(
  501. ignored_keys=getattr(self, "_save_config_ignore", []), skip_default=True
  502. ).items():
  503. lines.append(get_config_line(mod, k, v))
  504. for import_name in imports:
  505. lines.insert(0, f"import {import_name}")
  506. return "\n".join(lines)
  507. def get_hash(self) -> bytes:
  508. """Hashes the configs that are not compile_ignored"""
  509. if self._is_dirty or self._hash_digest is None:
  510. dict_to_hash = self._get_dict(ignored_keys=list(self._compile_ignored_keys))
  511. string_to_hash = repr(sorted(dict_to_hash.items()))
  512. self._hash_digest = hashlib.md5(
  513. string_to_hash.encode("utf-8"), usedforsecurity=False
  514. ).digest()
  515. self._is_dirty = False
  516. return self._hash_digest
  517. @deprecated(
  518. "`config.to_dict()` has been deprecated. It no longer changes the underlying config."
  519. " use `config.get_config_copy()` instead if you just want a copy of the config, or "
  520. "config.load_config if you need mutable access",
  521. category=FutureWarning,
  522. )
  523. def to_dict(self) -> dict[str, Any]:
  524. return self.get_config_copy()
  525. @deprecated(
  526. "`config.shallow_copy_dict()` has been deprecated. It no longer changes the underlying config."
  527. " use `config.get_config_copy()` instead if you just want a copy of the config, or "
  528. "config.load_config if you need mutable access",
  529. category=FutureWarning,
  530. )
  531. def shallow_copy_dict(self) -> dict[str, Any]:
  532. return self.get_config_copy()
  533. def load_config(self, maybe_pickled_config: bytes | dict[str, Any]) -> None:
  534. """Restore from a prior call to save_config() or shallow_copy_dict()"""
  535. if not isinstance(maybe_pickled_config, dict):
  536. config = pickle.loads(maybe_pickled_config)
  537. else:
  538. config = maybe_pickled_config
  539. for k, v in config.items():
  540. if k in self._config:
  541. setattr(self, k, v)
  542. else:
  543. from torch._dynamo.utils import warn_once
  544. warn_once(f"key {k} with value {v} is not understood by this config")
  545. def get_config_copy(self) -> dict[str, Any]:
  546. return self._get_dict()
  547. def get_serializable_config_copy(self) -> dict[str, Any]:
  548. return self._get_dict(ignored_keys=getattr(self, "_save_config_ignore", []))
  549. def patch(
  550. self,
  551. arg1: str | dict[str, Any] | None = None,
  552. arg2: Any = None,
  553. **kwargs: dict[str, Any],
  554. ) -> "ContextDecorator":
  555. """
  556. Decorator and/or context manager to make temporary changes to a config.
  557. As a decorator:
  558. @config.patch("name", val)
  559. @config.patch(name1=val1, name2=val2)
  560. @config.patch({"name1": val1, "name2", val2})
  561. def foo(...):
  562. ...
  563. As a context manager:
  564. with config.patch("name", val):
  565. ...
  566. """
  567. changes: dict[str, Any]
  568. if arg1 is not None:
  569. if arg2 is not None:
  570. if not isinstance(arg1, str):
  571. raise AssertionError(
  572. "first argument must be a string when passing 2 positional args to patch"
  573. )
  574. # patch("key", True) syntax
  575. changes = {arg1: arg2}
  576. else:
  577. if not isinstance(arg1, dict):
  578. raise AssertionError(
  579. "first argument must be a dict when passing a single positional arg to patch"
  580. )
  581. # patch({"key": True}) syntax
  582. changes = arg1
  583. if kwargs:
  584. raise AssertionError(
  585. "cannot pass both positional and keyword arguments to patch"
  586. )
  587. else:
  588. # patch(key=True) syntax
  589. changes = kwargs
  590. if arg2 is not None:
  591. raise AssertionError(
  592. "second positional argument is only valid when first argument is a key string"
  593. )
  594. if not isinstance(changes, dict):
  595. raise AssertionError(f"expected `dict` got {type(changes)}")
  596. prior: dict[str, Any] = {}
  597. config = self
  598. class ConfigPatch(ContextDecorator):
  599. def __init__(self) -> None:
  600. self.changes = changes
  601. def __enter__(self) -> None:
  602. if prior:
  603. raise AssertionError(
  604. "prior should be empty when entering ConfigPatch"
  605. )
  606. for key in self.changes:
  607. # KeyError on invalid entry
  608. prior[key] = config.__getattr__(key)
  609. for k, v in self.changes.items():
  610. config.__setattr__(k, v)
  611. def __exit__(self, exc_type, exc_val, exc_tb): # type: ignore[no-untyped-def]
  612. for k, v in prior.items():
  613. config.__setattr__(k, v)
  614. prior.clear()
  615. return ConfigPatch()
  616. def _make_closure_patcher(self, **changes: dict[str, Any]) -> Any:
  617. """
  618. A lower-overhead version of patch() for things on the critical path.
  619. Usage:
  620. # do this off the critical path
  621. change_fn = config.make_closure_patcher(foo=True)
  622. ...
  623. revert = change_fn()
  624. try:
  625. ...
  626. finally:
  627. revert()
  628. """
  629. config = self._config
  630. def change() -> Callable[[], None]:
  631. prior = {k: config[k].user_override for k in changes}
  632. for k, v in changes.items():
  633. self._config[k].user_override = v
  634. def revert() -> None:
  635. for k, v in prior.items():
  636. self._config[k].user_override = v
  637. return revert
  638. return change
  639. class ContextDecorator(contextlib.ContextDecorator):
  640. """
  641. Same as contextlib.ContextDecorator, but with support for
  642. `unittest.TestCase`
  643. """
  644. def __enter__(self) -> None:
  645. raise NotImplementedError("NYI")
  646. def __exit__(self, exc_type, exc_val, exc_tb) -> NoReturn: # type: ignore[no-untyped-def]
  647. raise NotImplementedError("NYI")
  648. def __call__(self, func: Callable[[Any], Any]) -> Any:
  649. if isinstance(func, type) and issubclass(func, unittest.TestCase):
  650. class _TestCase(func): # type: ignore[valid-type, misc]
  651. @classmethod
  652. def setUpClass(cls) -> None:
  653. self.__enter__()
  654. try:
  655. super().setUpClass()
  656. except Exception:
  657. self.__exit__(None, None, None)
  658. raise
  659. @classmethod
  660. def tearDownClass(cls) -> None:
  661. try:
  662. super().tearDownClass()
  663. finally:
  664. self.__exit__(None, None, None)
  665. _TestCase.__name__ = func.__name__
  666. _TestCase.__qualname__ = func.__qualname__
  667. _TestCase.__module__ = func.__module__
  668. return _TestCase
  669. return super().__call__(func)
  670. class SubConfigProxy:
  671. """
  672. Shim to redirect to main config.
  673. `config.triton.cudagraphs` maps to _config["triton.cudagraphs"]
  674. """
  675. def __init__(self, config: object, prefix: str) -> None:
  676. # `super().__setattr__` to bypass custom `__setattr__`
  677. super().__setattr__("_config", config)
  678. super().__setattr__("_prefix", prefix)
  679. def __setattr__(self, name: str, value: object) -> None:
  680. return self._config.__setattr__(self._prefix + name, value)
  681. def __getattr__(self, name: str) -> Any:
  682. return self._config.__getattr__(self._prefix + name)
  683. def __delattr__(self, name: str) -> None:
  684. return self._config.__delattr__(self._prefix + name)
  685. def patch_object(obj: object, name: str, value: object) -> object:
  686. """
  687. Workaround `mock.patch.object` issue with ConfigModule
  688. """
  689. if isinstance(obj, ConfigModule):
  690. return obj.patch(name, value)
  691. return mock.patch.object(obj, name, value)
  692. def get_tristate_env(name: str, default: Any = None) -> bool | None:
  693. value = os.environ.get(name)
  694. if value == "1":
  695. return True
  696. if value == "0":
  697. return False
  698. return default
  699. def inherit_fields_from(parent_cls):
  700. def wrapper(child_cls):
  701. for k, v in parent_cls.__dict__.items():
  702. # copy fields that are not private and not overridden
  703. if not k.startswith("_") and k not in child_cls.__dict__:
  704. setattr(child_cls, k, v)
  705. return child_cls
  706. return wrapper