chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,53 @@
|
||||
from dataclasses import dataclass
|
||||
|
||||
from ray.air.config import (
|
||||
CheckpointConfig as _CheckpointConfig,
|
||||
FailureConfig as _FailureConfig,
|
||||
RunConfig as _RunConfig,
|
||||
)
|
||||
from ray.train.constants import (
|
||||
V2_MIGRATION_GUIDE_MESSAGE,
|
||||
_v2_migration_warnings_enabled,
|
||||
)
|
||||
from ray.train.utils import _copy_doc, _log_deprecation_warning
|
||||
|
||||
# NOTE: This is just a pass-through wrapper around `ray.tune.RunConfig`
|
||||
# in order to detect whether the import module was correct (e.g. `ray.tune.RunConfig`).
|
||||
|
||||
|
||||
@dataclass
|
||||
@_copy_doc(_CheckpointConfig)
|
||||
class CheckpointConfig(_CheckpointConfig):
|
||||
pass
|
||||
|
||||
|
||||
@dataclass
|
||||
@_copy_doc(_FailureConfig)
|
||||
class FailureConfig(_FailureConfig):
|
||||
pass
|
||||
|
||||
|
||||
@dataclass
|
||||
@_copy_doc(_RunConfig)
|
||||
class RunConfig(_RunConfig):
|
||||
def __post_init__(self):
|
||||
self.checkpoint_config = self.checkpoint_config or CheckpointConfig()
|
||||
self.failure_config = self.failure_config or FailureConfig()
|
||||
|
||||
super().__post_init__()
|
||||
|
||||
if not isinstance(self.checkpoint_config, CheckpointConfig):
|
||||
if _v2_migration_warnings_enabled():
|
||||
_log_deprecation_warning(
|
||||
"The `CheckpointConfig` class should be imported from `ray.tune` "
|
||||
"when passing it to the Tuner. Please update your imports."
|
||||
f"{V2_MIGRATION_GUIDE_MESSAGE}"
|
||||
)
|
||||
|
||||
if not isinstance(self.failure_config, FailureConfig):
|
||||
if _v2_migration_warnings_enabled():
|
||||
_log_deprecation_warning(
|
||||
"The `FailureConfig` class should be imported from `ray.tune` "
|
||||
"when passing it to the Tuner. Please update your imports."
|
||||
f"{V2_MIGRATION_GUIDE_MESSAGE}"
|
||||
)
|
||||
@@ -0,0 +1,33 @@
|
||||
import contextlib
|
||||
import traceback
|
||||
|
||||
import ray
|
||||
|
||||
|
||||
def _deserialize_and_fully_execute_if_needed(serialized_ds: bytes):
|
||||
ds = ray.data.Dataset.deserialize_lineage(serialized_ds)
|
||||
return ds
|
||||
|
||||
|
||||
def _reduce(ds: ray.data.Dataset):
|
||||
tb_list = traceback.format_list(traceback.extract_stack())
|
||||
_already_in_out_of_band_serialization = False
|
||||
for tb in tb_list:
|
||||
# TODO(xwjiang): Let's make this less hacky.
|
||||
if "serialize_lineage" in tb:
|
||||
_already_in_out_of_band_serialization = True
|
||||
break
|
||||
if not _already_in_out_of_band_serialization and ds.has_serializable_lineage():
|
||||
return _deserialize_and_fully_execute_if_needed, (ds.serialize_lineage(),)
|
||||
else:
|
||||
return ds.__reduce__()
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def out_of_band_serialize_dataset():
|
||||
context = ray._private.worker.global_worker.get_serialization_context()
|
||||
try:
|
||||
context._register_cloudpickle_reducer(ray.data.Dataset, _reduce)
|
||||
yield
|
||||
finally:
|
||||
context._unregister_cloudpickle_reducer(ray.data.Dataset)
|
||||
@@ -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,))
|
||||
@@ -0,0 +1,65 @@
|
||||
from sklearn.datasets import load_breast_cancer
|
||||
|
||||
from ray import tune
|
||||
from ray.data import Dataset, Datasource, ReadTask, read_datasource
|
||||
from ray.data.block import BlockMetadata
|
||||
from ray.tune.impl.utils import execute_dataset
|
||||
|
||||
|
||||
# TODO(xwjiang): Enable this when Clark's out-of-band-serialization is landed.
|
||||
class TestDatasource(Datasource):
|
||||
def prepare_read(self, parallelism: int, **read_args):
|
||||
import pyarrow as pa
|
||||
|
||||
def load_data():
|
||||
data_raw = load_breast_cancer(as_frame=True)
|
||||
dataset_df = data_raw["data"]
|
||||
dataset_df["target"] = data_raw["target"]
|
||||
return [pa.Table.from_pandas(dataset_df)]
|
||||
|
||||
meta = BlockMetadata(
|
||||
num_rows=None,
|
||||
size_bytes=None,
|
||||
input_files=None,
|
||||
exec_stats=None,
|
||||
)
|
||||
return [ReadTask(load_data, meta)]
|
||||
|
||||
|
||||
def gen_dataset_func() -> Dataset:
|
||||
test_datasource = TestDatasource()
|
||||
return read_datasource(test_datasource)
|
||||
|
||||
|
||||
def test_grid_search():
|
||||
ds1 = gen_dataset_func().lazy().map(lambda x: x)
|
||||
ds2 = gen_dataset_func().lazy().map(lambda x: x)
|
||||
assert not ds1._has_computed_output()
|
||||
assert not ds2._has_computed_output()
|
||||
param_space = {"train_dataset": tune.grid_search([ds1, ds2])}
|
||||
execute_dataset(param_space)
|
||||
executed_ds = param_space["train_dataset"]["grid_search"]
|
||||
assert len(executed_ds) == 2
|
||||
assert executed_ds[0]._has_computed_output()
|
||||
assert executed_ds[1]._has_computed_output()
|
||||
|
||||
|
||||
def test_choice():
|
||||
ds1 = gen_dataset_func().lazy().map(lambda x: x)
|
||||
ds2 = gen_dataset_func().lazy().map(lambda x: x)
|
||||
assert not ds1._has_computed_output()
|
||||
assert not ds2._has_computed_output()
|
||||
param_space = {"train_dataset": tune.choice([ds1, ds2])}
|
||||
execute_dataset(param_space)
|
||||
executed_ds = param_space["train_dataset"].categories
|
||||
assert len(executed_ds) == 2
|
||||
assert executed_ds[0]._has_computed_output()
|
||||
assert executed_ds[1]._has_computed_output()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
|
||||
sys.exit(pytest.main(["-v", "-x", __file__]))
|
||||
@@ -0,0 +1,698 @@
|
||||
import copy
|
||||
import io
|
||||
import logging
|
||||
import math
|
||||
from pathlib import Path
|
||||
from typing import (
|
||||
TYPE_CHECKING,
|
||||
Any,
|
||||
Callable,
|
||||
Dict,
|
||||
List,
|
||||
Optional,
|
||||
Tuple,
|
||||
Type,
|
||||
Union,
|
||||
)
|
||||
|
||||
import pyarrow.fs
|
||||
|
||||
import ray.cloudpickle as pickle
|
||||
import ray.train
|
||||
import ray.tune
|
||||
from ray.air._internal.uri_utils import URI
|
||||
from ray.air._internal.usage import AirEntrypoint
|
||||
from ray.train._internal.storage import StorageContext, get_fs_and_path
|
||||
from ray.train.constants import (
|
||||
V2_MIGRATION_GUIDE_MESSAGE,
|
||||
_v2_migration_warnings_enabled,
|
||||
)
|
||||
from ray.train.utils import _log_deprecation_warning
|
||||
from ray.tune import (
|
||||
Experiment,
|
||||
ExperimentAnalysis,
|
||||
ResumeConfig,
|
||||
RunConfig,
|
||||
TuneConfig,
|
||||
TuneError,
|
||||
)
|
||||
from ray.tune.registry import is_function_trainable
|
||||
from ray.tune.result_grid import ResultGrid
|
||||
from ray.tune.trainable import Trainable
|
||||
from ray.tune.tune import _Config, run
|
||||
from ray.tune.utils import flatten_dict
|
||||
from ray.util import inspect_serializability
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ray.train.trainer import BaseTrainer
|
||||
from ray.util.queue import Queue
|
||||
|
||||
|
||||
_TUNER_PKL = "tuner.pkl"
|
||||
_TRAINABLE_KEY = "_trainable"
|
||||
_CONVERTED_TRAINABLE_KEY = "_converted_trainable"
|
||||
_PARAM_SPACE_KEY = "_param_space"
|
||||
_EXPERIMENT_ANALYSIS_KEY = "_experiment_analysis"
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
TrainableType = Union[str, Callable, Type[Trainable]]
|
||||
TrainableTypeOrTrainer = Union[TrainableType, "BaseTrainer"]
|
||||
|
||||
|
||||
class TunerInternal:
|
||||
"""The real implementation behind external facing ``Tuner``.
|
||||
|
||||
The external facing ``Tuner`` multiplexes between local Tuner and remote Tuner
|
||||
depending on whether in Ray client mode.
|
||||
|
||||
In Ray client mode, external ``Tuner`` wraps ``TunerInternal`` into a remote actor,
|
||||
which is guaranteed to be placed on head node.
|
||||
|
||||
``TunerInternal`` can be constructed from fresh, in which case, ``trainable`` needs
|
||||
to be provided, together with optional ``param_space``, ``tune_config`` and
|
||||
``run_config``.
|
||||
|
||||
It can also be restored from a previous failed run (given ``restore_path``).
|
||||
|
||||
Args:
|
||||
restore_path: The path from where the Tuner can be restored. If provided, None
|
||||
of the rest args are needed.
|
||||
storage_filesystem: A custom ``pyarrow.fs.FileSystem`` corresponding
|
||||
to ``restore_path``. This may be necessary if the original
|
||||
experiment used a custom filesystem.
|
||||
resume_config: Resume config to configure which trials to continue.
|
||||
trainable: The trainable to be tuned.
|
||||
param_space: Search space of the tuning job.
|
||||
One thing to note is that both preprocessor and dataset can be tuned here.
|
||||
tune_config: Tuning algorithm specific configs.
|
||||
Refer to ray.tune.tune_config.TuneConfig for more info.
|
||||
run_config: Runtime configuration that is specific to individual trials.
|
||||
If passed, this will overwrite the run config passed to the Trainer,
|
||||
if applicable. Refer to ray.tune.RunConfig for more info.
|
||||
_tuner_kwargs: Internal. Extra kwargs forwarded to ``tune.run`` when
|
||||
this Tuner is fit.
|
||||
_entrypoint: Internal. Tracks which user-facing entrypoint constructed
|
||||
this Tuner so that warnings and errors can be specialized.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
restore_path: str = None,
|
||||
storage_filesystem: Optional[pyarrow.fs.FileSystem] = None,
|
||||
resume_config: Optional[ResumeConfig] = None,
|
||||
trainable: Optional[TrainableTypeOrTrainer] = None,
|
||||
param_space: Optional[Dict[str, Any]] = None,
|
||||
tune_config: Optional[TuneConfig] = None,
|
||||
run_config: Optional[RunConfig] = None,
|
||||
_tuner_kwargs: Optional[Dict] = None,
|
||||
_entrypoint: AirEntrypoint = AirEntrypoint.TUNER,
|
||||
):
|
||||
from ray.train.trainer import BaseTrainer
|
||||
|
||||
if isinstance(trainable, BaseTrainer):
|
||||
if _v2_migration_warnings_enabled():
|
||||
_log_deprecation_warning(
|
||||
"The Ray Train + Ray Tune integration has been reworked. "
|
||||
"Passing a Trainer to the Tuner is deprecated and will be removed "
|
||||
"in a future release. "
|
||||
f"{V2_MIGRATION_GUIDE_MESSAGE}"
|
||||
)
|
||||
|
||||
run_config = self._choose_run_config(
|
||||
tuner_run_config=run_config,
|
||||
trainer=trainable,
|
||||
param_space=param_space,
|
||||
)
|
||||
|
||||
self._tune_config = tune_config or TuneConfig()
|
||||
self._run_config = copy.copy(run_config) or RunConfig()
|
||||
self._entrypoint = _entrypoint
|
||||
|
||||
# Restore from Tuner checkpoint.
|
||||
if restore_path:
|
||||
self._restore_from_path_or_uri(
|
||||
path_or_uri=restore_path,
|
||||
trainable=trainable,
|
||||
overwrite_param_space=param_space,
|
||||
resume_config=resume_config,
|
||||
storage_filesystem=storage_filesystem,
|
||||
)
|
||||
return
|
||||
|
||||
# Start from fresh
|
||||
if not trainable:
|
||||
raise TuneError("You need to provide a trainable to tune.")
|
||||
|
||||
if self._entrypoint == AirEntrypoint.TUNER and not isinstance(
|
||||
self._run_config, ray.tune.RunConfig
|
||||
):
|
||||
if _v2_migration_warnings_enabled():
|
||||
_log_deprecation_warning(
|
||||
"The `RunConfig` class should be imported from `ray.tune` "
|
||||
"when passing it to the Tuner. Please update your imports. "
|
||||
f"{V2_MIGRATION_GUIDE_MESSAGE}"
|
||||
)
|
||||
|
||||
self.trainable = trainable
|
||||
assert self.converted_trainable
|
||||
self._validate_trainable(self.converted_trainable)
|
||||
|
||||
self.param_space = param_space
|
||||
|
||||
self._resume_config = None
|
||||
self._is_restored = False
|
||||
self._tuner_kwargs = copy.deepcopy(_tuner_kwargs) or {}
|
||||
self._experiment_analysis = None
|
||||
|
||||
self._run_config.name = (
|
||||
self._run_config.name
|
||||
or StorageContext.get_experiment_dir_name(self.converted_trainable)
|
||||
)
|
||||
# The storage context here is only used to access the resolved
|
||||
# storage fs and experiment path, in order to avoid duplicating that logic.
|
||||
# This is NOT the storage context object that gets passed to remote workers.
|
||||
storage = StorageContext(
|
||||
storage_path=self._run_config.storage_path,
|
||||
experiment_dir_name=self._run_config.name,
|
||||
storage_filesystem=self._run_config.storage_filesystem,
|
||||
)
|
||||
|
||||
fs = storage.storage_filesystem
|
||||
fs.create_dir(storage.experiment_fs_path)
|
||||
with fs.open_output_stream(
|
||||
Path(storage.experiment_fs_path, _TUNER_PKL).as_posix()
|
||||
) as f:
|
||||
f.write(pickle.dumps(self.__getstate__()))
|
||||
|
||||
def get_run_config(self) -> RunConfig:
|
||||
return self._run_config
|
||||
|
||||
# For Jupyter output with Ray Client
|
||||
def set_run_config_and_remote_string_queue(
|
||||
self, run_config: RunConfig, string_queue: "Queue"
|
||||
):
|
||||
self._run_config = run_config
|
||||
self._tuner_kwargs["_remote_string_queue"] = string_queue
|
||||
|
||||
def clear_remote_string_queue(self):
|
||||
self._tuner_kwargs.pop("_remote_string_queue", None)
|
||||
|
||||
def _expected_utilization(self, cpus_per_trial, cpus_total):
|
||||
num_samples = self._tune_config.num_samples
|
||||
if num_samples < 0: # TODO: simplify this in Tune
|
||||
num_samples = math.inf
|
||||
concurrent_trials = self._tune_config.max_concurrent_trials or 0
|
||||
if concurrent_trials < 1: # TODO: simplify this in Tune
|
||||
concurrent_trials = math.inf
|
||||
|
||||
actual_concurrency = min(
|
||||
(
|
||||
(cpus_total // cpus_per_trial) if cpus_per_trial else 0,
|
||||
num_samples,
|
||||
concurrent_trials,
|
||||
)
|
||||
)
|
||||
return (actual_concurrency * cpus_per_trial) / (cpus_total + 0.001)
|
||||
|
||||
def _validate_trainable(
|
||||
self,
|
||||
trainable: TrainableType,
|
||||
required_trainable_name: Optional[str] = None,
|
||||
) -> None:
|
||||
"""Determines whether or not the trainable is valid.
|
||||
|
||||
This includes checks on the serializability of the trainable, as well
|
||||
asserting that the trainable name is as expected on restoration.
|
||||
|
||||
This trainable name validation is needed due to an implementation detail
|
||||
where the trainable name (which is differently generated depending on
|
||||
the trainable type) is saved in the Trial metadata and needs to match
|
||||
upon restoration. This does not affect the typical path, since `Tuner.restore`
|
||||
expects the exact same trainable (which will have the same name).
|
||||
|
||||
Args:
|
||||
trainable: The trainable to validate.
|
||||
required_trainable_name: If provided, the trainable's generated
|
||||
name must match this value; used on restoration to detect a
|
||||
trainable swap.
|
||||
|
||||
Raises:
|
||||
ValueError: if the trainable name does not match or if the trainable
|
||||
is not serializable.
|
||||
"""
|
||||
try:
|
||||
pickle.dumps(trainable)
|
||||
except TypeError as e:
|
||||
sio = io.StringIO()
|
||||
inspect_serializability(trainable, print_file=sio)
|
||||
msg = (
|
||||
"The provided trainable is not serializable, which is a requirement "
|
||||
"since the trainable is serialized and deserialized when transferred "
|
||||
"to remote workers. See below for a trace of the non-serializable "
|
||||
"objects that were found in your trainable:\n"
|
||||
f"{sio.getvalue()}"
|
||||
)
|
||||
raise TypeError(msg) from e
|
||||
|
||||
if not required_trainable_name:
|
||||
return
|
||||
|
||||
trainable_name = Experiment.get_trainable_name(trainable)
|
||||
|
||||
if trainable_name != required_trainable_name:
|
||||
raise ValueError(
|
||||
"Invalid `trainable` input to `Tuner.restore()`. To fix this error, "
|
||||
"pass in the same trainable that was used to initialize the Tuner. "
|
||||
"Got a trainable with identifier "
|
||||
f"'{trainable_name}' but expected '{required_trainable_name}'."
|
||||
)
|
||||
|
||||
def _set_trainable_on_restore(
|
||||
self, trainable: TrainableType, old_trainable_name: Optional[str]
|
||||
):
|
||||
from ray.train.base_trainer import BaseTrainer
|
||||
|
||||
self.trainable = trainable
|
||||
assert self.converted_trainable
|
||||
self._validate_trainable(
|
||||
trainable=self.converted_trainable,
|
||||
required_trainable_name=old_trainable_name,
|
||||
)
|
||||
|
||||
if isinstance(self.trainable, BaseTrainer):
|
||||
# Log a warning in case the user tries to modify the
|
||||
# `RunConfig` from the Trainer
|
||||
trainer: BaseTrainer = self.trainable
|
||||
|
||||
# Only log if the Trainer has a non-default RunConfig
|
||||
if trainer.run_config != RunConfig():
|
||||
logger.warning(
|
||||
"The Tune experiment will restore using the original run's "
|
||||
"`RunConfig`. If you made any changes to the `RunConfig` "
|
||||
"within the Trainer you passed into `Tuner.restore`, "
|
||||
"they will be ignored in the resumed run."
|
||||
)
|
||||
|
||||
trainer.run_config = self._run_config
|
||||
|
||||
def _validate_param_space_on_restore(
|
||||
self,
|
||||
new_param_space: Dict[str, Any],
|
||||
flattened_param_space_keys: Optional[List[str]],
|
||||
) -> None:
|
||||
"""Determines whether the (optionally) re-specified `param_space` is valid.
|
||||
|
||||
This method performs very loose validation on the new param_space to
|
||||
prevent users from trying to specify new hyperparameters to tune over.
|
||||
|
||||
Args:
|
||||
new_param_space: The newly provided search space to validate.
|
||||
flattened_param_space_keys: Sorted flat keys of the original
|
||||
``param_space``. ``None`` skips validation for backwards
|
||||
compatibility.
|
||||
|
||||
Raises:
|
||||
ValueError: if not all keys match the original param_space.
|
||||
"""
|
||||
if flattened_param_space_keys is None:
|
||||
# Backwards compatibility: skip validation
|
||||
return
|
||||
|
||||
keys = sorted(flatten_dict(new_param_space).keys())
|
||||
if keys != flattened_param_space_keys:
|
||||
raise ValueError(
|
||||
"Invalid `param_space` input to `Tuner.restore()`. To fix this error, "
|
||||
"pass in the same `param_space` that was used to initialize the Tuner. "
|
||||
"Only re-specify the `param_space` to refresh Ray object references "
|
||||
"that no longer exist due to restoring from a new Ray cluster session. "
|
||||
"It should not be used to introduce new hyperparameters to tune."
|
||||
f"\n\nGot: {keys}\nExpected: {flattened_param_space_keys}"
|
||||
)
|
||||
|
||||
def _set_param_space_on_restore(
|
||||
self,
|
||||
param_space: Optional[Dict[str, Any]],
|
||||
flattened_param_space_keys: Optional[List[str]],
|
||||
):
|
||||
self.param_space = param_space
|
||||
|
||||
if self.param_space is not None:
|
||||
# param_space = None -> use the original param_space
|
||||
self._validate_param_space_on_restore(
|
||||
new_param_space=self.param_space,
|
||||
flattened_param_space_keys=flattened_param_space_keys,
|
||||
)
|
||||
|
||||
def _load_tuner_state(
|
||||
self, tuner_state: Dict[str, Any]
|
||||
) -> Tuple[Optional[str], Optional[List[str]]]:
|
||||
"""Loads Tuner state from the previously saved `tuner.pkl`.
|
||||
|
||||
Args:
|
||||
tuner_state: Deserialized contents of the `tuner.pkl` saved during
|
||||
the original Tuner initialization.
|
||||
|
||||
Returns:
|
||||
tuple: of `(old_trainable_name, flattened_param_space_keys)` used for
|
||||
validating the re-specified `trainable` and `param_space`.
|
||||
"""
|
||||
# NOTE: These are magic keys used for validating restore args.
|
||||
old_trainable_name = tuner_state.pop("__trainable_name", None)
|
||||
flattened_param_space_keys = tuner_state.pop(
|
||||
"__flattened_param_space_keys", None
|
||||
)
|
||||
|
||||
self.__setstate__(tuner_state)
|
||||
|
||||
return old_trainable_name, flattened_param_space_keys
|
||||
|
||||
def _restore_from_path_or_uri(
|
||||
self,
|
||||
path_or_uri: str,
|
||||
trainable: TrainableTypeOrTrainer,
|
||||
overwrite_param_space: Optional[Dict[str, Any]],
|
||||
resume_config: ResumeConfig,
|
||||
storage_filesystem: Optional[pyarrow.fs.FileSystem],
|
||||
):
|
||||
fs, fs_path = get_fs_and_path(path_or_uri, storage_filesystem)
|
||||
with fs.open_input_file(Path(fs_path, _TUNER_PKL).as_posix()) as f:
|
||||
tuner_state = pickle.loads(f.readall())
|
||||
|
||||
old_trainable_name, flattened_param_space_keys = self._load_tuner_state(
|
||||
tuner_state
|
||||
)
|
||||
|
||||
# Perform validation and set the re-specified `trainable` and `param_space`
|
||||
self._set_trainable_on_restore(
|
||||
trainable=trainable, old_trainable_name=old_trainable_name
|
||||
)
|
||||
self._set_param_space_on_restore(
|
||||
param_space=overwrite_param_space,
|
||||
flattened_param_space_keys=flattened_param_space_keys,
|
||||
)
|
||||
|
||||
# Update RunConfig to reflect changes in the experiment directory
|
||||
path_or_uri_obj = URI(path_or_uri)
|
||||
|
||||
# Infer the `storage_path` and run `name` of the restored run using the
|
||||
# experiment directory.
|
||||
# Ex: ~/ray_results/exp_name -> ~/ray_results, exp_name
|
||||
# Ex: s3://bucket/exp_name -> s3://bucket, exp_name
|
||||
self._run_config.name = path_or_uri_obj.name
|
||||
self._run_config.storage_path = str(path_or_uri_obj.parent)
|
||||
# Update the storage_filesystem with the one passed in on restoration, if any.
|
||||
self._run_config.storage_filesystem = storage_filesystem
|
||||
|
||||
# Load the experiment results at the point where it left off.
|
||||
try:
|
||||
self._experiment_analysis = ExperimentAnalysis(
|
||||
experiment_checkpoint_path=path_or_uri,
|
||||
default_metric=self._tune_config.metric,
|
||||
default_mode=self._tune_config.mode,
|
||||
storage_filesystem=storage_filesystem,
|
||||
)
|
||||
except Exception:
|
||||
self._experiment_analysis = None
|
||||
|
||||
self._resume_config = resume_config
|
||||
self._is_restored = True
|
||||
|
||||
def _choose_run_config(
|
||||
self,
|
||||
tuner_run_config: Optional[RunConfig],
|
||||
trainer: "BaseTrainer",
|
||||
param_space: Optional[Dict[str, Any]],
|
||||
) -> RunConfig:
|
||||
"""Chooses which `RunConfig` to use when multiple can be passed in
|
||||
through a Trainer or the Tuner itself.
|
||||
|
||||
Args:
|
||||
tuner_run_config: The run config passed into the Tuner constructor.
|
||||
trainer: The Trainer instance to use with Tune, which may have
|
||||
a RunConfig specified by the user.
|
||||
param_space: The param space passed to the Tuner.
|
||||
|
||||
Returns:
|
||||
The resolved ``RunConfig`` to use for the Tune experiment.
|
||||
|
||||
Raises:
|
||||
ValueError: if the `run_config` is specified as a hyperparameter.
|
||||
"""
|
||||
if param_space and "run_config" in param_space:
|
||||
raise ValueError(
|
||||
"`RunConfig` cannot be tuned as part of the `param_space`! "
|
||||
"Move the run config to be a parameter of the `Tuner`: "
|
||||
"Tuner(..., run_config=RunConfig(...))"
|
||||
)
|
||||
|
||||
# Both Tuner RunConfig + Trainer RunConfig --> prefer Tuner RunConfig
|
||||
if tuner_run_config and trainer.run_config != ray.train.RunConfig():
|
||||
logger.info(
|
||||
"A `RunConfig` was passed to both the `Tuner` and the "
|
||||
f"`{trainer.__class__.__name__}`. The run config passed to "
|
||||
"the `Tuner` is the one that will be used."
|
||||
)
|
||||
return tuner_run_config
|
||||
|
||||
# No Tuner RunConfig -> pass the Trainer config through
|
||||
# This returns either a user-specified config, or the default RunConfig
|
||||
# if nothing was provided to both the Trainer or Tuner.
|
||||
if not tuner_run_config:
|
||||
return trainer.run_config
|
||||
|
||||
# Tuner RunConfig + No Trainer RunConfig --> Use the Tuner config
|
||||
return tuner_run_config
|
||||
|
||||
def _process_scaling_config(self) -> None:
|
||||
"""Converts ``self._param_space["scaling_config"]`` to a dict.
|
||||
|
||||
The dict is converted back to a dataclass by the Trainer, after the
|
||||
Tune search specification is resolved.
|
||||
"""
|
||||
# TODO: introduce `ray.tune.sample.TuneableDataclass` and allow Tune to
|
||||
# natively resolve specs with dataclasses.
|
||||
scaling_config = self._param_space.get("scaling_config")
|
||||
if not isinstance(scaling_config, ray.train.ScalingConfig):
|
||||
return
|
||||
self._param_space["scaling_config"] = scaling_config.__dict__.copy()
|
||||
|
||||
@property
|
||||
def trainable(self) -> TrainableTypeOrTrainer:
|
||||
return self._trainable
|
||||
|
||||
@property
|
||||
def converted_trainable(self) -> TrainableType:
|
||||
return self._converted_trainable
|
||||
|
||||
@trainable.setter
|
||||
def trainable(self, trainable: TrainableTypeOrTrainer):
|
||||
self._trainable = trainable
|
||||
self._converted_trainable = self._convert_trainable(trainable)
|
||||
|
||||
@property
|
||||
def param_space(self) -> Optional[Dict[str, Any]]:
|
||||
return self._param_space
|
||||
|
||||
@param_space.setter
|
||||
def param_space(self, param_space: Optional[Dict[str, Any]]):
|
||||
# Handle any configs that adhere to the `to_dict` interface.
|
||||
# Ex: AlgorithmConfig from RLlib
|
||||
if isinstance(param_space, _Config):
|
||||
param_space = param_space.to_dict()
|
||||
|
||||
if not isinstance(param_space, dict) and param_space is not None:
|
||||
raise ValueError(
|
||||
"The `param_space` passed to the `Tuner` must be a dict. "
|
||||
f"Got '{type(param_space)}' instead."
|
||||
)
|
||||
|
||||
self._param_space = param_space
|
||||
|
||||
if param_space:
|
||||
self._process_scaling_config()
|
||||
|
||||
def _convert_trainable(self, trainable: TrainableTypeOrTrainer) -> TrainableType:
|
||||
"""Converts a Trainer to a Tune trainable and saves the converted
|
||||
trainable. If not using a Trainer, this leaves the trainable as is."""
|
||||
from ray.train.trainer import BaseTrainer
|
||||
|
||||
return (
|
||||
trainable.as_trainable()
|
||||
if isinstance(trainable, BaseTrainer)
|
||||
else trainable
|
||||
)
|
||||
|
||||
def fit(self) -> ResultGrid:
|
||||
trainable = self.converted_trainable
|
||||
param_space = copy.deepcopy(self.param_space)
|
||||
if not self._is_restored:
|
||||
analysis = self._fit_internal(trainable, param_space)
|
||||
else:
|
||||
analysis = self._fit_resume(trainable, param_space)
|
||||
|
||||
self._experiment_analysis = analysis
|
||||
|
||||
return ResultGrid(self._experiment_analysis)
|
||||
|
||||
def get_results(self) -> ResultGrid:
|
||||
if not self._experiment_analysis:
|
||||
raise RuntimeError(
|
||||
"Can't return results as experiment has not been run, yet. "
|
||||
"Call `Tuner.fit()` to run the experiment first."
|
||||
)
|
||||
return ResultGrid(self._experiment_analysis)
|
||||
|
||||
def _get_tune_run_arguments(self, trainable: TrainableType) -> Dict[str, Any]:
|
||||
"""Get tune.run arguments common for both new and resumed runs."""
|
||||
# Avoid overwriting the originally configured checkpoint config.
|
||||
checkpoint_config = copy.deepcopy(self._run_config.checkpoint_config)
|
||||
|
||||
if checkpoint_config.checkpoint_frequency:
|
||||
# Function trainables (and thus most of our trainers) usually don't handle
|
||||
# this argument.
|
||||
handle_checkpoint_freq = getattr(
|
||||
trainable, "_handles_checkpoint_freq", None
|
||||
)
|
||||
if handle_checkpoint_freq is False:
|
||||
# If we specifically know this trainable doesn't support the
|
||||
# argument, raise an error
|
||||
raise ValueError(
|
||||
"You passed `checkpoint_frequency="
|
||||
f"{checkpoint_config.checkpoint_frequency}` to your "
|
||||
"CheckpointConfig, but this trainer does not support "
|
||||
"this argument. If you passed in a Trainer that takes in a "
|
||||
"custom training loop, you will need to "
|
||||
"report a checkpoint every `checkpoint_frequency` iterations "
|
||||
"within your training loop using "
|
||||
"`ray.tune.report(metrics=..., checkpoint=...)` "
|
||||
"to get this behavior."
|
||||
)
|
||||
elif handle_checkpoint_freq is True:
|
||||
# If we specifically support it, it's handled in the training loop,
|
||||
# so we disable tune's bookkeeping.
|
||||
checkpoint_config.checkpoint_frequency = 0
|
||||
# Otherwise, the trainable is not a Trainer and we just keep the
|
||||
# user-supplied value.
|
||||
# Function trainables will raise a runtime error later if set > 0
|
||||
if checkpoint_config.checkpoint_at_end is not None:
|
||||
# Again, function trainables usually don't handle this argument.
|
||||
handle_cp_at_end = getattr(trainable, "_handles_checkpoint_at_end", None)
|
||||
if handle_cp_at_end is False:
|
||||
# If we specifically know we don't support it, raise an error.
|
||||
raise ValueError(
|
||||
"You passed `checkpoint_at_end="
|
||||
f"{checkpoint_config.checkpoint_at_end}` "
|
||||
"to your CheckpointConfig, but this trainer does not support "
|
||||
"this argument. If you passed in a Trainer that takes in a "
|
||||
"custom training loop, you should include one last call to "
|
||||
"`ray.tune.report(metrics=..., checkpoint=...)` "
|
||||
"at the end of your training loop to get this behavior."
|
||||
)
|
||||
elif handle_cp_at_end is True:
|
||||
# If we specifically support it, it's handled in the training loop,
|
||||
# so we disable tune's internal bookkeeping.
|
||||
checkpoint_config.checkpoint_at_end = False
|
||||
# If this is a user-defined trainable, just keep the value
|
||||
# Function trainables will raise a runtime error later if set to True
|
||||
else:
|
||||
# Set default to False for function trainables and True for everything else
|
||||
if is_function_trainable(trainable):
|
||||
checkpoint_config.checkpoint_at_end = False
|
||||
else:
|
||||
checkpoint_config.checkpoint_at_end = True
|
||||
|
||||
return dict(
|
||||
storage_path=self._run_config.storage_path,
|
||||
storage_filesystem=self._run_config.storage_filesystem,
|
||||
name=self._run_config.name,
|
||||
mode=self._tune_config.mode,
|
||||
metric=self._tune_config.metric,
|
||||
callbacks=self._run_config.callbacks,
|
||||
sync_config=self._run_config.sync_config,
|
||||
stop=self._run_config.stop,
|
||||
max_failures=self._run_config.failure_config.max_failures,
|
||||
checkpoint_config=checkpoint_config,
|
||||
raise_on_failed_trial=False,
|
||||
fail_fast=(self._run_config.failure_config.fail_fast),
|
||||
progress_reporter=self._run_config.progress_reporter,
|
||||
verbose=self._run_config.verbose,
|
||||
reuse_actors=self._tune_config.reuse_actors,
|
||||
max_concurrent_trials=self._tune_config.max_concurrent_trials,
|
||||
time_budget_s=self._tune_config.time_budget_s,
|
||||
trial_name_creator=self._tune_config.trial_name_creator,
|
||||
trial_dirname_creator=self._tune_config.trial_dirname_creator,
|
||||
_entrypoint=self._entrypoint,
|
||||
# Deprecated
|
||||
chdir_to_trial_dir=self._tune_config.chdir_to_trial_dir,
|
||||
)
|
||||
|
||||
def _fit_internal(
|
||||
self, trainable: TrainableType, param_space: Optional[Dict[str, Any]]
|
||||
) -> ExperimentAnalysis:
|
||||
"""Fitting for a fresh Tuner."""
|
||||
args = {
|
||||
**self._get_tune_run_arguments(trainable),
|
||||
**dict(
|
||||
run_or_experiment=trainable,
|
||||
config=param_space,
|
||||
num_samples=self._tune_config.num_samples,
|
||||
search_alg=self._tune_config.search_alg,
|
||||
scheduler=self._tune_config.scheduler,
|
||||
log_to_file=self._run_config.log_to_file,
|
||||
),
|
||||
**self._tuner_kwargs,
|
||||
}
|
||||
analysis = run(
|
||||
**args,
|
||||
)
|
||||
self.clear_remote_string_queue()
|
||||
return analysis
|
||||
|
||||
def _fit_resume(
|
||||
self, trainable: TrainableType, param_space: Optional[Dict[str, Any]]
|
||||
) -> ExperimentAnalysis:
|
||||
"""Fitting for a restored Tuner."""
|
||||
assert self._resume_config
|
||||
|
||||
args = {
|
||||
**self._get_tune_run_arguments(trainable),
|
||||
**dict(
|
||||
run_or_experiment=trainable,
|
||||
config=param_space,
|
||||
resume_config=self._resume_config,
|
||||
search_alg=self._tune_config.search_alg,
|
||||
scheduler=self._tune_config.scheduler,
|
||||
),
|
||||
**self._tuner_kwargs,
|
||||
}
|
||||
analysis = run(**args)
|
||||
self.clear_remote_string_queue()
|
||||
return analysis
|
||||
|
||||
def __getstate__(self):
|
||||
state = self.__dict__.copy()
|
||||
state["_tuner_kwargs"] = state["_tuner_kwargs"].copy()
|
||||
state["_tuner_kwargs"].pop("_remote_string_queue", None)
|
||||
state.pop(_TRAINABLE_KEY, None)
|
||||
trainable = state.pop(_CONVERTED_TRAINABLE_KEY, None)
|
||||
param_space = state.pop(_PARAM_SPACE_KEY, None)
|
||||
state.pop(_EXPERIMENT_ANALYSIS_KEY, None)
|
||||
|
||||
state["__trainable_name"] = (
|
||||
Experiment.get_trainable_name(trainable) if trainable else None
|
||||
)
|
||||
state["__flattened_param_space_keys"] = (
|
||||
sorted(flatten_dict(param_space).keys())
|
||||
if param_space is not None
|
||||
else None
|
||||
)
|
||||
|
||||
return state
|
||||
|
||||
def __setstate__(self, state):
|
||||
# Make sure the magic metadata gets removed first.
|
||||
state.pop("__flattened_param_space_keys", None)
|
||||
state.pop("__trainable_name", None)
|
||||
|
||||
self.__dict__.update(state)
|
||||
Reference in New Issue
Block a user