chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,244 @@
|
||||
import hashlib
|
||||
from collections import defaultdict
|
||||
from typing import Any, Dict, Tuple
|
||||
|
||||
from ray.tune.search.sample import Categorical, Domain, Function
|
||||
from ray.tune.search.variant_generator import assign_value
|
||||
from ray.util.annotations import DeveloperAPI
|
||||
|
||||
ID_HASH_LENGTH = 8
|
||||
|
||||
|
||||
def create_resolvers_map():
|
||||
return defaultdict(list)
|
||||
|
||||
|
||||
def _id_hash(path_tuple):
|
||||
"""Compute a hash for the specific placeholder based on its path."""
|
||||
return hashlib.sha256(str(path_tuple).encode("utf-8")).hexdigest()[:ID_HASH_LENGTH]
|
||||
|
||||
|
||||
class _FunctionResolver:
|
||||
"""Replaced value for function typed objects."""
|
||||
|
||||
TOKEN = "__fn_ph"
|
||||
|
||||
def __init__(self, hash, fn):
|
||||
self.hash = hash
|
||||
self._fn = fn
|
||||
|
||||
def resolve(self, config: Dict):
|
||||
"""Some functions take a resolved spec dict as input.
|
||||
|
||||
Note: Function placeholders are independently sampled during
|
||||
resolution. Therefore their random states are not restored.
|
||||
"""
|
||||
return self._fn.sample(config=config)
|
||||
|
||||
def get_placeholder(self) -> str:
|
||||
return (self.TOKEN, self.hash)
|
||||
|
||||
|
||||
class _RefResolver:
|
||||
"""Replaced value for all other non-primitive objects."""
|
||||
|
||||
TOKEN = "__ref_ph"
|
||||
|
||||
def __init__(self, hash, value):
|
||||
self.hash = hash
|
||||
self._value = value
|
||||
|
||||
def resolve(self):
|
||||
return self._value
|
||||
|
||||
def get_placeholder(self) -> str:
|
||||
return (self.TOKEN, self.hash)
|
||||
|
||||
|
||||
def _is_primitive(x):
|
||||
"""Returns True if x is a primitive type.
|
||||
|
||||
Primitive types are int, float, str, bool, and None.
|
||||
"""
|
||||
return isinstance(x, (int, float, str, bool)) or x is None
|
||||
|
||||
|
||||
@DeveloperAPI
|
||||
def inject_placeholders(
|
||||
config: Any,
|
||||
resolvers: defaultdict,
|
||||
id_prefix: Tuple = (),
|
||||
path_prefix: Tuple = (),
|
||||
) -> Dict:
|
||||
"""Replaces reference objects contained by a config dict with placeholders.
|
||||
|
||||
Given a config dict, this function replaces all reference objects contained
|
||||
by this dict with placeholder strings. It recursively expands nested dicts
|
||||
and lists, and properly handles Tune native search objects such as Categorical
|
||||
and Function.
|
||||
This makes sure the config dict only contains primitive typed values, which
|
||||
can then be handled by different search algorithms.
|
||||
|
||||
A few details about id_prefix and path_prefix. Consider the following config,
|
||||
where "param1" is a simple grid search of 3 tuples.
|
||||
|
||||
config = {
|
||||
"param1": tune.grid_search([
|
||||
(Cat, None, None),
|
||||
(None, Dog, None),
|
||||
(None, None, Fish),
|
||||
]),
|
||||
}
|
||||
|
||||
We will replace the 3 objects contained with placeholders. And after trial
|
||||
expansion, the config may look like this:
|
||||
|
||||
config = {
|
||||
"param1": (None, (placeholder, hash), None)
|
||||
}
|
||||
|
||||
Now you need 2 pieces of information to resolve the placeholder. One is the
|
||||
path of ("param1", 1), which tells you that the first element of the tuple
|
||||
under "param1" key is a placeholder that needs to be resolved.
|
||||
The other is the mapping from the placeholder to the actual object. In this
|
||||
case hash -> Dog.
|
||||
|
||||
id and path prefixes serve exactly this purpose here. The difference between
|
||||
these two is that id_prefix is the location of the value in the pre-injected
|
||||
config tree. So if a value is the second option in a grid_search, it gets an
|
||||
id part of 1. Injected placeholders all get unique id prefixes. path prefix
|
||||
identifies a placeholder in the expanded config tree. So for example, all
|
||||
options of a single grid_search will get the same path prefix. This is how
|
||||
we know which location has a placeholder to be resolved in the post-expansion
|
||||
tree.
|
||||
|
||||
Args:
|
||||
config: The config dict to replace references in.
|
||||
resolvers: A dict from path to replaced objects.
|
||||
id_prefix: The prefix to prepend to id every single placeholders.
|
||||
path_prefix: The prefix to prepend to every path identifying
|
||||
potential locations of placeholders in an expanded tree.
|
||||
|
||||
Returns:
|
||||
The config with all references replaced.
|
||||
"""
|
||||
if isinstance(config, dict) and "grid_search" in config and len(config) == 1:
|
||||
config["grid_search"] = [
|
||||
# Different options gets different id prefixes.
|
||||
# But we should omit appending to path_prefix because after expansion,
|
||||
# this level will not be there.
|
||||
inject_placeholders(choice, resolvers, id_prefix + (i,), path_prefix)
|
||||
for i, choice in enumerate(config["grid_search"])
|
||||
]
|
||||
return config
|
||||
elif isinstance(config, dict):
|
||||
return {
|
||||
k: inject_placeholders(v, resolvers, id_prefix + (k,), path_prefix + (k,))
|
||||
for k, v in config.items()
|
||||
}
|
||||
elif isinstance(config, list):
|
||||
return [
|
||||
inject_placeholders(elem, resolvers, id_prefix + (i,), path_prefix + (i,))
|
||||
for i, elem in enumerate(config)
|
||||
]
|
||||
elif isinstance(config, tuple):
|
||||
return tuple(
|
||||
inject_placeholders(elem, resolvers, id_prefix + (i,), path_prefix + (i,))
|
||||
for i, elem in enumerate(config)
|
||||
)
|
||||
elif _is_primitive(config):
|
||||
# Primitive types.
|
||||
return config
|
||||
elif isinstance(config, Categorical):
|
||||
config.categories = [
|
||||
# Different options gets different id prefixes.
|
||||
# But we should omit appending to path_prefix because after expansion,
|
||||
# this level will not be there.
|
||||
inject_placeholders(choice, resolvers, id_prefix + (i,), path_prefix)
|
||||
for i, choice in enumerate(config.categories)
|
||||
]
|
||||
return config
|
||||
elif isinstance(config, Function):
|
||||
# Function type.
|
||||
id_hash = _id_hash(id_prefix)
|
||||
v = _FunctionResolver(id_hash, config)
|
||||
resolvers[path_prefix].append(v)
|
||||
return v.get_placeholder()
|
||||
elif not isinstance(config, Domain):
|
||||
# Other non-search space reference objects, dataset, actor handle, etc.
|
||||
id_hash = _id_hash(id_prefix)
|
||||
v = _RefResolver(id_hash, config)
|
||||
resolvers[path_prefix].append(v)
|
||||
return v.get_placeholder()
|
||||
else:
|
||||
# All the other cases, do nothing.
|
||||
return config
|
||||
|
||||
|
||||
def _get_placeholder(config: Any, prefix: Tuple, path: Tuple):
|
||||
if not path:
|
||||
return prefix, config
|
||||
|
||||
key = path[0]
|
||||
if isinstance(config, tuple):
|
||||
if config[0] in (_FunctionResolver.TOKEN, _RefResolver.TOKEN):
|
||||
# Found a matching placeholder.
|
||||
# Note that we do not require that the full path are consumed before
|
||||
# declaring a match. Because this placeholder may be part of a nested
|
||||
# search space. For example, the following config:
|
||||
# config = {
|
||||
# "param1": tune.grid_search([
|
||||
# tune.grid_search([Object1, 2, 3]),
|
||||
# tune.grid_search([Object2, 5, 6]),
|
||||
# ]),
|
||||
# }
|
||||
# will result in placeholders under path ("param1", 0, 0).
|
||||
# After expansion though, the choosen placeholder will live under path
|
||||
# ("param1", 0) like this: config = {"param1": (Placeholder1, 2, 3)}
|
||||
return prefix, config
|
||||
elif key < len(config):
|
||||
return _get_placeholder(
|
||||
config[key], prefix=prefix + (path[0],), path=path[1:]
|
||||
)
|
||||
elif (isinstance(config, dict) and key in config) or (
|
||||
isinstance(config, list) and key < len(config)
|
||||
):
|
||||
# Expand config tree recursively.
|
||||
return _get_placeholder(config[key], prefix=prefix + (path[0],), path=path[1:])
|
||||
|
||||
# Can not find a matching placeholder.
|
||||
return None, None
|
||||
|
||||
|
||||
@DeveloperAPI
|
||||
def resolve_placeholders(config: Any, replaced: defaultdict):
|
||||
"""Replaces placeholders contained by a config dict with the original values.
|
||||
|
||||
Args:
|
||||
config: The config to replace placeholders in.
|
||||
replaced: A dict from path to replaced objects.
|
||||
"""
|
||||
|
||||
def __resolve(resolver_type, args):
|
||||
for path, resolvers in replaced.items():
|
||||
assert resolvers
|
||||
|
||||
if not isinstance(resolvers[0], resolver_type):
|
||||
continue
|
||||
|
||||
prefix, ph = _get_placeholder(config, (), path)
|
||||
if not ph:
|
||||
# Represents an unchosen value. Just skip.
|
||||
continue
|
||||
|
||||
for resolver in resolvers:
|
||||
if resolver.hash != ph[1]:
|
||||
continue
|
||||
# Found the matching resolver.
|
||||
assign_value(config, prefix, resolver.resolve(*args))
|
||||
|
||||
# RefResolvers first.
|
||||
__resolve(_RefResolver, args=())
|
||||
# Functions need to be resolved after RefResolvers, in case they are
|
||||
# referencing values from the RefResolvers.
|
||||
__resolve(_FunctionResolver, args=(config,))
|
||||
Reference in New Issue
Block a user