wehub-resource-sync
251 lines
7.8 KiB
Python
251 lines
7.8 KiB
Python
import inspect
|
|
import logging
|
|
import types
|
|
from typing import TYPE_CHECKING, Any, Callable, Dict, Optional, Type, Union
|
|
|
|
import ray
|
|
from ray.tune.execution.placement_groups import (
|
|
PlacementGroupFactory,
|
|
resource_dict_to_pg_factory,
|
|
)
|
|
from ray.tune.registry import _ParameterRegistry
|
|
from ray.util.annotations import PublicAPI
|
|
|
|
if TYPE_CHECKING:
|
|
from ray.tune.trainable import Trainable
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
@PublicAPI(stability="beta")
|
|
def with_parameters(trainable: Union[Type["Trainable"], Callable], **kwargs):
|
|
"""Wrapper for trainables to pass arbitrary large data objects.
|
|
|
|
This wrapper function will store all passed parameters in the Ray
|
|
object store and retrieve them when calling the function. It can thus
|
|
be used to pass arbitrary data, even datasets, to Tune trainables.
|
|
|
|
This can also be used as an alternative to ``functools.partial`` to pass
|
|
default arguments to trainables.
|
|
|
|
When used with the function API, the trainable function is called with
|
|
the passed parameters as keyword arguments. When used with the class API,
|
|
the ``Trainable.setup()`` method is called with the respective kwargs.
|
|
|
|
If the data already exists in the object store (are instances of
|
|
ObjectRef), using ``tune.with_parameters()`` is not necessary. You can
|
|
instead pass the object refs to the training function via the ``config``
|
|
or use Python partials.
|
|
|
|
Args:
|
|
trainable: Trainable to wrap.
|
|
**kwargs: parameters to store in object store.
|
|
|
|
Function API example:
|
|
|
|
.. code-block:: python
|
|
|
|
from ray import tune
|
|
|
|
def train_fn(config, data=None):
|
|
for sample in data:
|
|
loss = update_model(sample)
|
|
tune.report(dict(loss=loss))
|
|
|
|
data = HugeDataset(download=True)
|
|
|
|
tuner = Tuner(
|
|
tune.with_parameters(train_fn, data=data),
|
|
# ...
|
|
)
|
|
tuner.fit()
|
|
|
|
Class API example:
|
|
|
|
.. code-block:: python
|
|
|
|
from ray import tune
|
|
|
|
class MyTrainable(tune.Trainable):
|
|
def setup(self, config, data=None):
|
|
self.data = data
|
|
self.iter = iter(self.data)
|
|
self.next_sample = next(self.iter)
|
|
|
|
def step(self):
|
|
loss = update_model(self.next_sample)
|
|
try:
|
|
self.next_sample = next(self.iter)
|
|
except StopIteration:
|
|
return {"loss": loss, done: True}
|
|
return {"loss": loss}
|
|
|
|
data = HugeDataset(download=True)
|
|
|
|
tuner = Tuner(
|
|
tune.with_parameters(MyTrainable, data=data),
|
|
# ...
|
|
)
|
|
|
|
Returns:
|
|
A wrapped trainable that has the provided ``kwargs`` injected via the
|
|
Ray object store at call time.
|
|
"""
|
|
from ray.tune.trainable import Trainable
|
|
|
|
if not callable(trainable) or (
|
|
inspect.isclass(trainable) and not issubclass(trainable, Trainable)
|
|
):
|
|
raise ValueError(
|
|
f"`tune.with_parameters() only works with function trainables "
|
|
f"or classes that inherit from `tune.Trainable()`. Got type: "
|
|
f"{type(trainable)}."
|
|
)
|
|
|
|
parameter_registry = _ParameterRegistry()
|
|
ray._private.worker._post_init_hooks.append(parameter_registry.flush)
|
|
|
|
# Objects are moved into the object store
|
|
prefix = f"{str(trainable)}_"
|
|
for k, v in kwargs.items():
|
|
parameter_registry.put(prefix + k, v)
|
|
|
|
trainable_name = getattr(trainable, "__name__", "tune_with_parameters")
|
|
keys = set(kwargs.keys())
|
|
|
|
if inspect.isclass(trainable):
|
|
# Class trainable
|
|
|
|
class _Inner(trainable):
|
|
def setup(self, config):
|
|
setup_kwargs = {}
|
|
for k in keys:
|
|
setup_kwargs[k] = parameter_registry.get(prefix + k)
|
|
super(_Inner, self).setup(config, **setup_kwargs)
|
|
|
|
trainable_with_params = _Inner
|
|
else:
|
|
# Function trainable
|
|
|
|
def inner(config):
|
|
fn_kwargs = {}
|
|
for k in keys:
|
|
fn_kwargs[k] = parameter_registry.get(prefix + k)
|
|
return trainable(config, **fn_kwargs)
|
|
|
|
trainable_with_params = inner
|
|
|
|
if hasattr(trainable, "__mixins__"):
|
|
trainable_with_params.__mixins__ = trainable.__mixins__
|
|
|
|
# If the trainable has been wrapped with `tune.with_resources`, we should
|
|
# keep the `_resources` attribute around
|
|
if hasattr(trainable, "_resources"):
|
|
trainable_with_params._resources = trainable._resources
|
|
|
|
trainable_with_params.__name__ = trainable_name
|
|
return trainable_with_params
|
|
|
|
|
|
@PublicAPI(stability="beta")
|
|
def with_resources(
|
|
trainable: Union[Type["Trainable"], Callable],
|
|
resources: Union[
|
|
Dict[str, float],
|
|
PlacementGroupFactory,
|
|
Callable[[dict], PlacementGroupFactory],
|
|
],
|
|
):
|
|
"""Wrapper for trainables to specify resource requests.
|
|
|
|
This wrapper allows specification of resource requirements for a specific
|
|
trainable. It will override potential existing resource requests (use
|
|
with caution!).
|
|
|
|
The main use case is to request resources for function trainables when used
|
|
with the Tuner() API.
|
|
|
|
Class trainables should usually just implement the ``default_resource_request()``
|
|
method.
|
|
|
|
Args:
|
|
trainable: Trainable to wrap.
|
|
resources: Resource dict, placement group factory, or callable that takes
|
|
in a config dict and returns a placement group factory.
|
|
|
|
Example:
|
|
|
|
.. code-block:: python
|
|
|
|
from ray import tune
|
|
from ray.tune.tuner import Tuner
|
|
|
|
def train_fn(config):
|
|
return len(ray.get_gpu_ids()) # Returns 2
|
|
|
|
tuner = Tuner(
|
|
tune.with_resources(train_fn, resources={"gpu": 2}),
|
|
# ...
|
|
)
|
|
results = tuner.fit()
|
|
|
|
Returns:
|
|
A trainable annotated with the requested resources so that Tune can
|
|
schedule trials accordingly.
|
|
"""
|
|
from ray.tune.trainable import Trainable
|
|
|
|
if not callable(trainable) or (
|
|
inspect.isclass(trainable) and not issubclass(trainable, Trainable)
|
|
):
|
|
raise ValueError(
|
|
f"`tune.with_resources() only works with function trainables "
|
|
f"or classes that inherit from `tune.Trainable()`. Got type: "
|
|
f"{type(trainable)}."
|
|
)
|
|
|
|
if isinstance(resources, PlacementGroupFactory):
|
|
pgf = resources
|
|
elif isinstance(resources, dict):
|
|
pgf = resource_dict_to_pg_factory(resources)
|
|
elif callable(resources):
|
|
pgf = resources
|
|
else:
|
|
raise ValueError(
|
|
f"Invalid resource type for `with_resources()`: {type(resources)}"
|
|
)
|
|
|
|
if not inspect.isclass(trainable):
|
|
if isinstance(trainable, types.MethodType):
|
|
# Methods cannot set arbitrary attributes, so we have to wrap them
|
|
def _trainable(config):
|
|
return trainable(config)
|
|
|
|
_trainable._resources = pgf
|
|
return _trainable
|
|
|
|
# Just set an attribute. This will be resolved later in `wrap_function()`.
|
|
try:
|
|
trainable._resources = pgf
|
|
except AttributeError as e:
|
|
raise RuntimeError(
|
|
"Could not use `tune.with_resources()` on the supplied trainable. "
|
|
"Wrap your trainable in a regular function before passing it "
|
|
"to Ray Tune."
|
|
) from e
|
|
else:
|
|
|
|
class ResourceTrainable(trainable):
|
|
@classmethod
|
|
def default_resource_request(
|
|
cls, config: Dict[str, Any]
|
|
) -> Optional[PlacementGroupFactory]:
|
|
if not isinstance(pgf, PlacementGroupFactory) and callable(pgf):
|
|
return pgf(config)
|
|
return pgf
|
|
|
|
ResourceTrainable.__name__ = trainable.__name__
|
|
trainable = ResourceTrainable
|
|
|
|
return trainable
|