frozenleaves/pytorch
1
0
Fork 0
mirror of https://github.com/pytorch/pytorch.git synced 2025-10-20 12:54:11 +08:00
Code Issues Packages Projects Releases Wiki Activity

Add config alias (#142088)

Browse Source 
Pull Request resolved: https://github.com/pytorch/pytorch/pull/142088
Approved by: https://github.com/c00w
This commit is contained in:
Oguz Ulgen
2024-12-14 09:24:12 -08:00
committed by PyTorch MergeBot
parent 1b6b86fad7
commit 17b71e5d6a
4 changed files with 119 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files

78
torch/utils/_config_module.py
View File

@ -1,6 +1,7 @@
import contextlib
import copy
import hashlib
import importlib
import inspect
import io
import os
@ -20,6 +21,7 @@ from typing import (
NoReturn,
Optional,
Set,
Tuple,
TYPE_CHECKING,
TypeVar,
Union,
@ -38,6 +40,9 @@ CONFIG_TYPES = (int, float, bool, type(None), str, list, set, tuple, dict)
T = TypeVar("T", bound=Union[int, float, bool, None, str, list, set, tuple, dict])
_UNSET_SENTINEL = object()
@dataclass
class _Config(Generic[T]):
"""Represents a config with richer behaviour than just a default value.
@ -49,7 +54,9 @@ class _Config(Generic[T]):
This configs must be installed with install_config_module to be used
Precedence Order:
env_name_force: If set, this environment variable overrides everything
alias: If set, the directly use the value of the alias.
env_name_force: If set, this environment variable has precedence over
everything after this.
user_override: If a user sets a value (i.e. foo.bar=True), that
has precedence over everything after this.
env_name_default: If set, this environment variable will override everything
@ -65,25 +72,28 @@ class _Config(Generic[T]):
Arguments:
justknob: the name of the feature / JK. In OSS this is unused.
default: is the value to default this knob to in OSS.
alias: The alias config to read instead.
env_name_force: The environment variable to read that is a FORCE
environment variable. I.e. it overrides everything
environment variable. I.e. it overrides everything except for alias.
env_name_default: The environment variable to read that changes the
default behaviour. I.e. user overrides take preference.
"""
default: T
default: Union[T, object]
justknob: Optional[str] = None
env_name_default: Optional[str] = None
env_name_force: Optional[str] = None
value_type: Optional[type] = None
alias: Optional[str] = None
def __init__(
self,
default: T,
default: Union[T, object] = _UNSET_SENTINEL,
justknob: Optional[str] = None,
env_name_default: Optional[str] = None,
env_name_force: Optional[str] = None,
value_type: Optional[type] = None,
alias: Optional[str] = None,
):
# python 3.9 does not support kw_only on the dataclass :(.
self.default = default
@ -91,10 +101,18 @@ class _Config(Generic[T]):
self.env_name_default = env_name_default
self.env_name_force = env_name_force
self.value_type = value_type
self.alias = alias
if self.justknob is not None:
assert isinstance(
self.default, bool
), f"justknobs only support booleans, {self.default} is not a boolean"
if self.alias is not None:
assert (
default is _UNSET_SENTINEL
and justknob is None
and env_name_default is None
and env_name_force is None
), "if alias is set, default, justknob or env var cannot be set"
# In runtime, we unbox the Config[T] to a T, but typechecker cannot see this,
@ -104,24 +122,28 @@ class _Config(Generic[T]):
if TYPE_CHECKING:
def Config(
default: T,
default: Union[T, object] = _UNSET_SENTINEL,
justknob: Optional[str] = None,
env_name_default: Optional[str] = None,
env_name_force: Optional[str] = None,
value_type: Optional[type] = None,
alias: Optional[str] = None,
) -> T:
...
else:
def Config(
default: T,
default: Union[T, object] = _UNSET_SENTINEL,
justknob: Optional[str] = None,
env_name_default: Optional[str] = None,
env_name_force: Optional[str] = None,
value_type: Optional[type] = None,
alias: Optional[str] = None,
) -> _Config[T]:
return _Config(default, justknob, env_name_default, env_name_force, value_type)
return _Config(
default, justknob, env_name_default, env_name_force, value_type, alias
)
def _read_env_variable(name: str) -> Optional[bool]:
@ -243,9 +265,6 @@ def get_assignments_with_compile_ignored_comments(module: ModuleType) -> Set[str
return assignments
_UNSET_SENTINEL = object()
@dataclass
class _ConfigEntry:
# The default value specified in the configuration
@ -272,6 +291,7 @@ class _ConfigEntry:
# call so the final state is correct. It's just very unintuitive.
# upstream bug - python/cpython#126886
hide: bool = False
alias: Optional[str] = None
def __init__(self, config: _Config):
self.default = config.default
@ -279,6 +299,7 @@ class _ConfigEntry:
config.value_type if config.value_type is not None else type(self.default)
)
self.justknob = config.justknob
self.alias = config.alias
if config.env_name_default is not None:
if (env_value := _read_env_variable(config.env_name_default)) is not None:
self.env_value_default = env_value
@ -309,6 +330,8 @@ class ConfigModule(ModuleType):
super().__setattr__(name, value)
elif name not in self._config:
raise AttributeError(f"{self.__name__}.{name} does not exist")
elif self._config[name].alias is not None:
self._set_alias_val(self._config[name], value)
else:
self._config[name].user_override = value
self._is_dirty = True
@ -321,6 +344,10 @@ class ConfigModule(ModuleType):
if config.hide:
raise AttributeError(f"{self.__name__}.{name} does not exist")
alias_val = self._get_alias_val(config)
if alias_val is not _UNSET_SENTINEL:
return alias_val
if config.env_value_force is not _UNSET_SENTINEL:
return config.env_value_force
@ -353,6 +380,33 @@ class ConfigModule(ModuleType):
self._config[name].user_override = _UNSET_SENTINEL
self._config[name].hide = True
def _get_alias_module_and_name(
self, entry: _ConfigEntry
) -> Optional[Tuple[ModuleType, str]]:
alias = entry.alias
if alias is None:
return None
module_name, constant_name = alias.rsplit(".", 1)
try:
module = importlib.import_module(module_name)
except ImportError as e:
raise AttributeError("config alias {alias} does not exist") from e
return module, constant_name
def _get_alias_val(self, entry: _ConfigEntry) -> Any:
data = self._get_alias_module_and_name(entry)
if data is None:
return _UNSET_SENTINEL
module, constant_name = data
constant_value = getattr(module, constant_name)
return constant_value
def _set_alias_val(self, entry: _ConfigEntry, val: Any) -> None:
data = self._get_alias_module_and_name(entry)
assert data is not None
module, constant_name = data
setattr(module, constant_name, val)
def _is_default(self, name: str) -> bool:
return self._config[name].user_override is _UNSET_SENTINEL
@ -369,6 +423,7 @@ class ConfigModule(ModuleType):
This is used by a number of different user facing export methods
which all have slightly different semantics re: how and what to
skip.
If a config is aliased, it skips this config.
Arguments:
ignored_keys are keys that should not be exported.
@ -391,7 +446,10 @@ class ConfigModule(ModuleType):
continue
if skip_default and self._is_default(key):
continue
if self._config[key].alias is not None:
continue
config[key] = copy.deepcopy(getattr(self, key))
return config
def get_type(self, config_name: str) -> type: