config: simplify most of the config handling and fix some bugs (#138377)

This PR combines a number of cleanups in one PR. If any of the specific cleanups don't seem to make sense, let me know and I can remove them.

Cleanups

- This PR adds a set of test suites for the config module code, which handles basically all the APIs and ways it is used. Please let me know if you see anything critical that is not tested that I missed. This test suite is primarily used as the regression test suite for later changes in this diff. Note that there is some dynamo specific testing of the config module, but it isn't as verbose.
- I removed all internal usage of shallow_copy_dict. Those usages could all use the deep copy, and did not depend on the reference behavior of certain config values that shallow_copy_dict allows.
- I removed shallow copy semantics for configuration with a deprecation warning. I think this requires a release note, so hopefully I did that correctly. Let me know if we want to continue to expose shallow copy value semantics, but I just can't find a case where I expect anyone would want it. It also complicated later internal changes to the API (i.e. breaking apart various layers of the config changes).
- I fixed what I believe is a bug in how hashes are calculated on configs. In particular, if you got the hash, then made a config change, and then got the hash again, it would not update the hash. @oulgen, please let me know if I'm misunderstanding this behavior and it is desired.
- I switched our multiple implementations of iterating through the dictionary to a single one. This is primarily to make later changes easier, but it also makes it clear how inconsistent our various config ignoring options are. Let me know if people would be interested in me unifying the various options for ignoring config values.
- I updated the test patcher (not the performance critical one, just the normal one), to use __setattr__ and __getattr__ to remove direct API access to the underlying config fetcher.

For release notes, Not sure exactly how to communicate this, but something like
"ConfigModule.to_dict, and ConfigModule.shallow_copy_dict no longer retain their shallow copy semantics, which allowed reference values objects to be modified. If you wish to modify the config object, call load_config explicitly".

Pull Request resolved: https://github.com/pytorch/pytorch/pull/138377
Approved by: https://github.com/ezyang, https://github.com/jansel, https://github.com/jovianjaison
This commit is contained in:
Colin L. Rice
2024-10-22 13:40:26 +00:00
committed by PyTorch MergeBot
parent 1b61313acd
commit bb8bc7d6b3
8 changed files with 370 additions and 48 deletions

View File

@ -8,7 +8,7 @@ import tokenize
import unittest
import warnings
from types import FunctionType, ModuleType
from typing import Any, Callable, Dict, NoReturn, Optional, Set, Union
from typing import Any, Callable, Dict, List, NoReturn, Optional, Set, Union
from typing_extensions import deprecated
from unittest import mock
@ -144,6 +144,7 @@ class ConfigModule(ModuleType):
raise AttributeError(f"{self.__name__}.{name} does not exist")
else:
self._config[name] = value
self._is_dirty = True
def __getattr__(self, name: str) -> Any:
try:
@ -153,29 +154,61 @@ class ConfigModule(ModuleType):
raise AttributeError(f"{self.__name__}.{name} does not exist") from e
def __delattr__(self, name: str) -> None:
self._is_dirty = True
# must support delete because unittest.mock.patch deletes
# then recreate things
del self._config[name]
def _get_dict(
self,
ignored_keys: Optional[List[str]] = None,
ignored_prefixes: Optional[List[str]] = None,
skip_default: bool = False,
) -> Dict[str, Any]:
"""Export a dictionary of current configuration keys and values.
This function is design to provide a single point which handles
accessing config options and exporting them into a dictionary.
This is used by a number of different user facing export methods
which all have slightly different semantics re: how and what to
skip.
Arguments:
ignored_keys are keys that should not be exported.
ignored_prefixes are prefixes that if a key matches should
not be exported
skip_default does two things. One if a key has not been modified
it skips it. The other is it modified the logging behaviour
to match what codegen already did for modified skipped keys
"""
config: Dict[str, Any] = {}
for key in self._config:
if ignored_keys and key in ignored_keys:
if skip_default and self._config[key] != self._default[key]:
warnings.warn(
f"Skipping serialization of {key} value {self._config[key]}"
)
continue
if ignored_prefixes:
if any(key.startswith(prefix) for prefix in ignored_prefixes):
continue
if skip_default and self._config[key] == self._default[key]:
continue
config[key] = copy.deepcopy(self._config[key])
return config
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)
return pickle.dumps(
self._get_dict(ignored_keys=self._config.get("_save_config_ignore", ())),
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
prefixes = ["_"]
prefixes.extend(self._config["_cache_config_ignore_prefix"])
return self._get_dict(ignored_prefixes=prefixes)
def codegen_config(self) -> str:
"""Convert config to Python statements that replicate current config.
@ -183,39 +216,38 @@ class ConfigModule(ModuleType):
"""
lines = []
mod = self.__name__
for k, v in self._config.items():
if k in self._config.get("_save_config_ignore", ()):
if v != self._default[k]:
warnings.warn(f"Skipping serialization of {k} value {v}")
continue
if v == self._default[k]:
continue
for k, v in self._get_dict(
ignored_keys=self._config.get("_save_config_ignore"), skip_default=True
).items():
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
}
dict_to_hash = self._get_dict(ignored_keys=list(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",
"`config.to_dict()` has been deprecated. It no longer changes the underlying config."
" use `config.get_config_copy()` instead if you just want a copy of the config, or "
"config.load_config if you need mutable access",
category=FutureWarning,
)
def to_dict(self) -> Dict[str, Any]:
return self.shallow_copy_dict()
return self.get_config_copy()
@deprecated(
"`config.shallow_copy_dict()` has been deprecated. It no longer changes the underlying config."
" use `config.get_config_copy()` instead if you just want a copy of the config, or "
"config.load_config if you need mutable access",
category=FutureWarning,
)
def shallow_copy_dict(self) -> Dict[str, Any]:
return {**self._config}
return self.get_config_copy()
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()"""
@ -226,7 +258,7 @@ class ConfigModule(ModuleType):
self._config.update(config)
def get_config_copy(self) -> Dict[str, Any]:
return copy.deepcopy(self._config)
return self._get_dict()
def patch(
self,
@ -268,23 +300,19 @@ class ConfigModule(ModuleType):
assert isinstance(changes, dict), f"expected `dict` got {type(changes)}"
prior: Dict[str, Any] = {}
config = self
dirty = False
class ConfigPatch(ContextDecorator):
def __enter__(self) -> None:
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
prior[key] = config.__getattr__(key)
for k, v in changes.items():
config.__setattr__(k, v)
def __exit__(self, exc_type, exc_val, exc_tb): # type: ignore[no-untyped-def]
nonlocal dirty
config._config.update(prior)
config._is_dirty = dirty
for k, v in prior.items():
config.__setattr__(k, v)
prior.clear()
return ConfigPatch()