wehub-resource-sync
245 lines
8.6 KiB
Python
245 lines
8.6 KiB
Python
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,))
|