wehub-resource-sync
876 lines
34 KiB
Python
876 lines
34 KiB
Python
import logging
|
|
import pickle
|
|
import warnings
|
|
from copy import deepcopy
|
|
from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Set, Tuple, Union
|
|
|
|
import numpy as np
|
|
|
|
from ray.air.execution.resources.request import _sum_bundles
|
|
from ray.tune.execution.placement_groups import PlacementGroupFactory
|
|
from ray.tune.experiment import Trial
|
|
from ray.tune.schedulers.trial_scheduler import FIFOScheduler, TrialScheduler
|
|
from ray.util.annotations import PublicAPI
|
|
|
|
if TYPE_CHECKING:
|
|
from ray.tune.execution.tune_controller import TuneController
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
@PublicAPI(stability="beta")
|
|
class DistributeResources:
|
|
"""This class creates a basic uniform resource allocation function.
|
|
|
|
The function naively balances free resources (CPUs and GPUs) between
|
|
trials, giving them all equal priority, ensuring that all resources
|
|
are always being used. The free resources will be placed in new bundles.
|
|
The function assumes that all bundles are equal (there is no "head"
|
|
bundle).
|
|
|
|
If for some reason a trial ends up with
|
|
more resources than there are free ones, it will adjust downwards.
|
|
It will also ensure that trial as at least as many resources as
|
|
it started with (``base_trial_resource``).
|
|
|
|
The function returns a new ``PlacementGroupFactory`` with updated
|
|
resource requirements, or None. If the returned
|
|
``PlacementGroupFactory`` is equal by value to the one the
|
|
trial has currently, the scheduler will skip the update process
|
|
internally (same with None).
|
|
|
|
If you wish to implement your own resource distribution logic,
|
|
you can do so by extending this class, as it provides several
|
|
generic methods. You can also implement a function instead.
|
|
|
|
Args:
|
|
add_bundles: If True, create new bundles from free resources.
|
|
Otherwise, spread them among base_trial_resource bundles.
|
|
increase_by: A dict with key-value
|
|
pairs representing an atomic unit of resources (name-amount)
|
|
the trial will be increased by. If not set, the trial will
|
|
increase by 1 CPU/GPU.
|
|
increase_by_times: If set to >=1 and ``increase_by`` is set,
|
|
the trial will increase by maximum of
|
|
``increase_by_times * increase_by`` resources. If set to <1,
|
|
no upper limit is set. Ignored if ``increase_by`` is not set.
|
|
reserve_resources: A dict of
|
|
resource_name-amount pairs representing the resources
|
|
that will not be allocated to resized trials.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
add_bundles: bool = False,
|
|
increase_by: Optional[Dict[str, float]] = None,
|
|
increase_by_times: int = -1,
|
|
reserve_resources: Optional[Dict[str, float]] = None,
|
|
):
|
|
self.add_bundles = add_bundles
|
|
self.increase_by = increase_by or {}
|
|
self.increase_by_times = increase_by_times
|
|
self.reserve_resources = reserve_resources or {}
|
|
|
|
def _validate(
|
|
self, base_trial_resource: PlacementGroupFactory, result: Dict[str, Any]
|
|
) -> bool:
|
|
"""Return False if we should keep the current resources outright."""
|
|
if not isinstance(base_trial_resource, PlacementGroupFactory):
|
|
raise ValueError(
|
|
f"{self.__class__.__name__} only supports PlacementGroupFactories."
|
|
)
|
|
|
|
if not self.add_bundles and len(base_trial_resource.bundles) > 1:
|
|
raise ValueError(
|
|
"If `add_bundles` is False, the number of bundles in "
|
|
"`resources_per_trial` must be 1 "
|
|
f"(got {len(base_trial_resource.bundles)})."
|
|
)
|
|
|
|
# Don't bother if this is just the first iteration
|
|
if result["training_iteration"] < 1:
|
|
return False
|
|
return True
|
|
|
|
def _get_total_available_resources(
|
|
self, tune_controller: "TuneController"
|
|
) -> Tuple[float, float]:
|
|
"""Get the number of CPUs and GPUs avaialble in total (not just free)"""
|
|
total_available_cpus = (
|
|
tune_controller._resource_updater.get_num_cpus()
|
|
- self.reserve_resources.get("CPU", 0)
|
|
)
|
|
total_available_gpus = (
|
|
tune_controller._resource_updater.get_num_gpus()
|
|
- self.reserve_resources.get("GPU", 0)
|
|
)
|
|
return total_available_cpus, total_available_gpus
|
|
|
|
def _get_used_cpus_and_gpus(self, t: Trial) -> Tuple[float, float]:
|
|
"""Check how many CPUs and GPUs a trial is using currently"""
|
|
return (
|
|
t.placement_group_factory.required_resources.get("CPU", 0),
|
|
t.placement_group_factory.required_resources.get("GPU", 0),
|
|
)
|
|
|
|
def _get_resources_from_bundles(
|
|
self, bundles: List[Dict[str, float]]
|
|
) -> Dict[str, float]:
|
|
"""Get total sums of resources in bundles"""
|
|
if not bundles:
|
|
return {"CPU": 0, "GPU": 0}
|
|
return _sum_bundles(bundles)
|
|
|
|
def _is_bundle_empty(self, bundle: Dict[str, float]) -> bool:
|
|
return not (bundle.get("CPU", 0) or bundle.get("GPU", 0))
|
|
|
|
def _add_two_bundles(
|
|
self,
|
|
bundles_a: List[Dict[str, float]],
|
|
bundles_b: List[Dict[str, float]],
|
|
increase_by: Dict[str, float],
|
|
limit_to_increase_by_times: bool,
|
|
max_increase_by_times: int = -1,
|
|
):
|
|
"""Add two bundles together.
|
|
|
|
If ``limit_to_increase_by_times`` is True, ``self.increase_by_times`` > 0
|
|
and ``max_increase_by_times`` > 0, ensure that the resulting number of
|
|
bundles is not above ``min(max_increase_by_times, self.increase_by_times)``.
|
|
|
|
If ``limit_to_increase_by_times`` is True and ``self.increase_by_times`` > 0,
|
|
ensure that the resulting number of bundles is not above
|
|
`self.increase_by_times``.
|
|
"""
|
|
if limit_to_increase_by_times:
|
|
if max_increase_by_times > 0 and self.increase_by_times > 0:
|
|
max_increase_by_times = min(
|
|
max_increase_by_times, self.increase_by_times
|
|
)
|
|
elif self.increase_by_times > 0:
|
|
max_increase_by_times = self.increase_by_times
|
|
|
|
if self.add_bundles:
|
|
bundles = [b for b in bundles_a if not self._is_bundle_empty(b)] + [
|
|
b for b in bundles_b if not self._is_bundle_empty(b)
|
|
]
|
|
if max_increase_by_times > 0:
|
|
bundles = bundles[:max_increase_by_times]
|
|
else:
|
|
bundles_a = bundles_a or [{}]
|
|
bundles_b = bundles_b or [{}]
|
|
bundles = [
|
|
{
|
|
"CPU": bundles_a[0].get("CPU", 0) + bundles_b[0].get("CPU", 0),
|
|
"GPU": bundles_a[0].get("GPU", 0) + bundles_b[0].get("GPU", 0),
|
|
}
|
|
]
|
|
if max_increase_by_times > 0:
|
|
bundles[0]["CPU"] = min(
|
|
bundles[0]["CPU"],
|
|
increase_by.get("CPU", 0) * max_increase_by_times,
|
|
)
|
|
bundles[0]["GPU"] = min(
|
|
bundles[0]["GPU"],
|
|
increase_by.get("GPU", 0) * max_increase_by_times,
|
|
)
|
|
|
|
return bundles
|
|
|
|
def _get_multiplier(
|
|
self,
|
|
increase_by: Dict[str, float],
|
|
cpus: float = 0,
|
|
gpus: float = 0,
|
|
max_multiplier: int = -1,
|
|
) -> int:
|
|
"""Get how many times ``increase_by`` bundles
|
|
occur in ``cpus`` and ``gpus``."""
|
|
if increase_by.get("CPU", 0) and increase_by.get("GPU", 0):
|
|
multiplier = min(
|
|
cpus // increase_by.get("CPU", 0),
|
|
gpus // increase_by.get("GPU", 0),
|
|
)
|
|
elif increase_by.get("GPU", 0):
|
|
multiplier = gpus // increase_by.get("GPU", 0)
|
|
else:
|
|
multiplier = cpus // increase_by.get("CPU", 0)
|
|
|
|
if max_multiplier > 0 and multiplier > 0:
|
|
multiplier = min(max_multiplier, multiplier)
|
|
return int(multiplier)
|
|
|
|
def _remove_bundles(
|
|
self,
|
|
bundles: List[Dict[str, float]],
|
|
increase_by: Dict[str, float],
|
|
multiplier: int,
|
|
) -> List[Dict[str, float]]:
|
|
"""Remove ``multiplier`` ``increase_by`` bundles from ``bundles``."""
|
|
multiplier = -abs(multiplier)
|
|
if self.add_bundles:
|
|
bundles = bundles[:multiplier]
|
|
else:
|
|
bundles = deepcopy(bundles)
|
|
bundles[0]["CPU"] += increase_by.get("CPU", 0) * multiplier
|
|
bundles[0]["GPU"] += increase_by.get("GPU", 0) * multiplier
|
|
bundles[0]["CPU"] = max(bundles[0]["CPU"], 0)
|
|
bundles[0]["GPU"] = max(bundles[0]["GPU"], 0)
|
|
return bundles
|
|
|
|
def _create_new_bundles(
|
|
self,
|
|
increase_by: Dict[str, float],
|
|
multiplier: int,
|
|
) -> List[Dict[str, float]]:
|
|
"""Create a list of new bundles containing ``increase_by`` * ``multiplier``."""
|
|
multiplier = abs(multiplier)
|
|
|
|
if self.add_bundles:
|
|
bundles = [increase_by] * int(multiplier)
|
|
else:
|
|
bundles = [{}]
|
|
bundles[0]["CPU"] = increase_by.get("CPU", 0) * multiplier
|
|
bundles[0]["GPU"] = increase_by.get("GPU", 0) * multiplier
|
|
|
|
return bundles
|
|
|
|
def _modify_bundles_with_free_resources(
|
|
self,
|
|
bundles: List[Dict[str, float]],
|
|
increase_by: Dict[str, float],
|
|
free_cpus: float,
|
|
free_gpus: float,
|
|
*,
|
|
max_multiplier: int = -1,
|
|
max_increase_by_times: int = -1,
|
|
):
|
|
"""Given free resources, increase/decrease the number of bundles in
|
|
``bundles``."""
|
|
multiplier = self._get_multiplier(
|
|
increase_by, free_cpus, free_gpus, max_multiplier
|
|
)
|
|
if multiplier < 0:
|
|
bundles = self._remove_bundles(bundles, increase_by, multiplier)
|
|
elif multiplier > 0:
|
|
bundles_to_add = self._create_new_bundles(increase_by, multiplier)
|
|
bundles = self._add_two_bundles(
|
|
bundles, bundles_to_add, increase_by, True, max_increase_by_times
|
|
)
|
|
return bundles
|
|
|
|
def _get_added_bundles(
|
|
self, bundles: List[Dict[str, float]], base_bundles: List[Dict[str, float]]
|
|
) -> List[Dict[str, float]]:
|
|
"""Return the difference between bundles and base_bundles"""
|
|
if self.add_bundles:
|
|
added_bundles = bundles[len(base_bundles) :]
|
|
else:
|
|
if not bundles:
|
|
bundles = [{"CPU": 0, "GPU": 0}]
|
|
if not base_bundles:
|
|
base_bundles = [{"CPU": 0, "GPU": 0}]
|
|
added_bundles = [
|
|
{
|
|
"CPU": bundles[0].get("CPU", 0) - base_bundles[0].get("CPU", 0),
|
|
"GPU": bundles[0].get("GPU", 0) - base_bundles[0].get("GPU", 0),
|
|
}
|
|
]
|
|
return added_bundles
|
|
|
|
def _are_bundles_below_limit(
|
|
self,
|
|
bundles: List[Dict[str, float]],
|
|
base_bundles: Optional[List[Dict[str, float]]] = None,
|
|
max_added_cpus: Optional[float] = None,
|
|
max_added_gpus: Optional[float] = None,
|
|
):
|
|
if not max_added_cpus:
|
|
if self.increase_by_times > 0:
|
|
max_added_cpus = self.increase_by.get("CPU", 0) * self.increase_by_times
|
|
else:
|
|
max_added_cpus = np.inf
|
|
if not max_added_gpus:
|
|
if self.increase_by_times > 0:
|
|
max_added_gpus = self.increase_by.get("GPU", 0) * self.increase_by_times
|
|
else:
|
|
max_added_gpus = np.inf
|
|
added_resources = self._get_resources_from_bundles(
|
|
self._get_added_bundles(bundles, base_bundles) if base_bundles else bundles
|
|
)
|
|
ret = (
|
|
added_resources.get("CPU", -np.inf) < max_added_cpus
|
|
or added_resources.get("GPU", -np.inf) < max_added_gpus
|
|
)
|
|
return ret
|
|
|
|
def _get_new_added_bundles(
|
|
self,
|
|
trial: Trial,
|
|
all_trials: List[Trial],
|
|
base_bundles: List[Dict[str, float]],
|
|
increase_by: Dict[str, float],
|
|
total_available_cpus: float,
|
|
total_available_gpus: float,
|
|
used_cpus: float,
|
|
used_gpus: float,
|
|
) -> List[Dict[str, float]]:
|
|
"""Returns updated added bundles."""
|
|
upper_limit_all_trials_bundles = [list() for _ in range(len(all_trials))]
|
|
|
|
free_cpus = total_available_cpus - used_cpus
|
|
free_gpus = total_available_gpus - used_gpus
|
|
|
|
base_resources = self._get_resources_from_bundles(base_bundles)
|
|
upper_limit_cpus_to_distribute = total_available_cpus - (
|
|
base_resources.get("CPU", 0) * len(all_trials)
|
|
)
|
|
upper_limit_gpus_to_distribute = total_available_gpus - (
|
|
base_resources.get("GPU", 0) * len(all_trials)
|
|
)
|
|
max_increase_by_times = 0
|
|
|
|
# First, calculate upper limits for uniform allocation
|
|
# This is done by simulating a clean slate scenario
|
|
# The loop runs until all resources are allocated or
|
|
# all trials are at their resource limits
|
|
i = 0
|
|
trials_at_limit = set()
|
|
while (
|
|
len(trials_at_limit) < len(all_trials)
|
|
# we have previously asserted that at least one resource has to be
|
|
# bigger than 0
|
|
and upper_limit_cpus_to_distribute >= increase_by.get("CPU", 0)
|
|
and upper_limit_gpus_to_distribute >= increase_by.get("GPU", 0)
|
|
):
|
|
idx = i % len(upper_limit_all_trials_bundles)
|
|
old_bundles = deepcopy(upper_limit_all_trials_bundles[idx])
|
|
upper_limit_all_trials_bundles[
|
|
idx
|
|
] = self._modify_bundles_with_free_resources(
|
|
upper_limit_all_trials_bundles[idx],
|
|
increase_by,
|
|
upper_limit_cpus_to_distribute,
|
|
upper_limit_gpus_to_distribute,
|
|
max_multiplier=1,
|
|
)
|
|
added_resources = self._get_resources_from_bundles(
|
|
self._get_added_bundles(
|
|
upper_limit_all_trials_bundles[idx], old_bundles
|
|
)
|
|
)
|
|
if not added_resources.get("CPU", 0) and not added_resources.get("GPU", 0):
|
|
trials_at_limit.add(idx)
|
|
elif idx == 0:
|
|
max_increase_by_times += 1
|
|
upper_limit_cpus_to_distribute -= added_resources.get("CPU", 0)
|
|
upper_limit_gpus_to_distribute -= added_resources.get("GPU", 0)
|
|
i += 1
|
|
|
|
# Add new resourcs, but only up to calculated upper limits
|
|
# (max_increase_by_times)
|
|
return self._modify_bundles_with_free_resources(
|
|
self._get_added_bundles(
|
|
trial.placement_group_factory.bundles, base_bundles
|
|
),
|
|
increase_by,
|
|
free_cpus,
|
|
free_gpus,
|
|
max_increase_by_times=max_increase_by_times,
|
|
)
|
|
|
|
def __call__(
|
|
self,
|
|
tune_controller: "TuneController",
|
|
trial: Trial,
|
|
result: Dict[str, Any],
|
|
scheduler: "ResourceChangingScheduler",
|
|
) -> Optional[PlacementGroupFactory]:
|
|
"""Run resource allocation logic.
|
|
|
|
Returns a new ``PlacementGroupFactory`` with updated
|
|
resource requirements, or None. If the returned
|
|
``PlacementGroupFactory`` is equal by value to the one the
|
|
trial has currently, the scheduler will skip the update process
|
|
internally (same with None).
|
|
|
|
Args:
|
|
tune_controller: Trial runner for this Tune run.
|
|
Can be used to obtain information about other trials.
|
|
trial: The trial to allocate new resources to.
|
|
result: The latest results of trial.
|
|
scheduler: The scheduler calling
|
|
the function.
|
|
|
|
Returns:
|
|
A new ``PlacementGroupFactory`` with updated resource requirements,
|
|
or None if the trial's resources should be left unchanged.
|
|
"""
|
|
# Get base trial resources as defined in
|
|
# ``tune.run(resources_per_trial)``
|
|
base_trial_resource = scheduler.base_trial_resources
|
|
|
|
if not self._validate(base_trial_resource=base_trial_resource, result=result):
|
|
return None
|
|
|
|
# default values if resources_per_trial is unspecified
|
|
if base_trial_resource is None:
|
|
base_trial_resource = PlacementGroupFactory([{"CPU": 1, "GPU": 0}])
|
|
|
|
if self.increase_by:
|
|
increase_by = self.increase_by
|
|
assert not self._is_bundle_empty(increase_by)
|
|
assert increase_by.get("CPU", 0) >= 0 and increase_by.get("GPU", 0) >= 0
|
|
elif self.add_bundles:
|
|
increase_by = base_trial_resource.bundles[-1]
|
|
elif base_trial_resource.bundles[0].get("GPU", 0):
|
|
increase_by = {"GPU": 1}
|
|
else:
|
|
increase_by = {"CPU": 1}
|
|
|
|
base_bundles = deepcopy(base_trial_resource.bundles)
|
|
|
|
(
|
|
total_available_cpus,
|
|
total_available_gpus,
|
|
) = self._get_total_available_resources(tune_controller=tune_controller)
|
|
|
|
all_trials = tune_controller.get_live_trials()
|
|
|
|
used_cpus_and_gpus = [self._get_used_cpus_and_gpus(t) for t in all_trials]
|
|
used_cpus, used_gpus = zip(*used_cpus_and_gpus)
|
|
used_cpus = sum(used_cpus)
|
|
used_gpus = sum(used_gpus)
|
|
|
|
added_bundles = self._get_new_added_bundles(
|
|
trial,
|
|
all_trials,
|
|
base_bundles,
|
|
increase_by,
|
|
total_available_cpus,
|
|
total_available_gpus,
|
|
used_cpus,
|
|
used_gpus,
|
|
)
|
|
|
|
new_bundles = self._add_two_bundles(
|
|
base_bundles, added_bundles, increase_by, False
|
|
)
|
|
|
|
pgf = PlacementGroupFactory(
|
|
new_bundles,
|
|
strategy=base_trial_resource.strategy,
|
|
*base_trial_resource._args,
|
|
**base_trial_resource._kwargs,
|
|
)
|
|
pgf._head_bundle_is_empty = base_trial_resource._head_bundle_is_empty
|
|
return pgf
|
|
|
|
|
|
@PublicAPI(stability="beta")
|
|
class DistributeResourcesToTopJob(DistributeResources):
|
|
"""This class creates a "TopJob" resource allocation function.
|
|
|
|
The function will assign all of the free resources to the best
|
|
performing trial (as defined by ``metric`` and ``mode``). The
|
|
previous best trials will not have their resources deallocated,
|
|
unless in the case outlined below.
|
|
|
|
If for some reason a trial ends up with
|
|
more resources than there are free ones, it will adjust downwards.
|
|
It will also ensure that trial as at least as many resources as
|
|
it started with (``base_trial_resource``).
|
|
|
|
The function returns a new ``PlacementGroupFactory`` with updated
|
|
resource requirements, or None. If the returned
|
|
``PlacementGroupFactory`` is equal by value to the one the
|
|
trial has currently, the scheduler will skip the update process
|
|
internally (same with None).
|
|
|
|
Args:
|
|
add_bundles: If True, create new bundles from free resources.
|
|
Otherwise, spread them among base_trial_resource bundles.
|
|
increase_by: A dict with key-value
|
|
pairs representing an atomic unit of resources (name-amount)
|
|
the trial will be increased by. If not set, the trial will
|
|
increase by 1 CPU/GPU.
|
|
increase_by_times: If set to >=1 and ``increase_by`` is set,
|
|
the trial will increase by maximum of
|
|
``increase_by_times * increase_by`` resources. If set to <1,
|
|
no upper limit is set. Ignored if ``increase_by`` is not set.
|
|
reserve_resources: A dict of
|
|
resource_name-amount pairs representing the resources
|
|
that will not be allocated to resized trials.
|
|
is that the attribute should increase monotonically.
|
|
metric: The training result objective value attribute. Stopping
|
|
procedures will use this attribute. If None, will use the metric
|
|
of the scheduler.
|
|
mode: One of {min, max}. Determines whether objective is
|
|
minimizing or maximizing the metric attribute. If None, will use the metric
|
|
of the scheduler.
|
|
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
add_bundles: bool = False,
|
|
increase_by: Optional[Dict[str, float]] = None,
|
|
increase_by_times: int = -1,
|
|
reserve_resources: Optional[Dict[str, float]] = None,
|
|
metric: Optional[str] = None,
|
|
mode: Optional[str] = None,
|
|
):
|
|
super().__init__(add_bundles, increase_by, increase_by_times, reserve_resources)
|
|
self.metric = metric
|
|
self.mode = mode
|
|
|
|
@property
|
|
def _metric_op(self) -> float:
|
|
if self.mode not in ("min", "max"):
|
|
raise ValueError("The mode parameter can only be either min or max.")
|
|
if self.mode == "max":
|
|
return 1.0
|
|
return -1.0
|
|
|
|
def _get_new_added_bundles(
|
|
self,
|
|
trial: Trial,
|
|
all_trials: List[Trial],
|
|
base_bundles: List[Dict[str, float]],
|
|
increase_by: Dict[str, float],
|
|
total_available_cpus: float,
|
|
total_available_gpus: float,
|
|
used_cpus: float,
|
|
used_gpus: float,
|
|
) -> List[Dict[str, float]]:
|
|
if self.metric is None:
|
|
raise ValueError(
|
|
"The metric parameter cannot be None. The parameter can be set in "
|
|
"either `DistributeResourcesToTopJob`, the base scheduler or in "
|
|
"`tune.TuneConfig()` (highest to lowest priority)."
|
|
)
|
|
|
|
free_cpus = total_available_cpus - used_cpus
|
|
free_gpus = total_available_gpus - used_gpus
|
|
|
|
sorted_trials = sorted(
|
|
all_trials,
|
|
key=lambda t: -self._metric_op * t.last_result.get(self.metric, np.inf),
|
|
)
|
|
|
|
added_bundles = self._get_added_bundles(
|
|
trial.placement_group_factory.bundles, base_bundles
|
|
)
|
|
|
|
best_trial = next(
|
|
(
|
|
t
|
|
for t in sorted_trials
|
|
if self._are_bundles_below_limit(
|
|
t.placement_group_factory.bundles, base_bundles
|
|
)
|
|
),
|
|
sorted_trials[0],
|
|
)
|
|
|
|
if (
|
|
trial.trial_id != best_trial.trial_id
|
|
# Only reduce resources here
|
|
and self._get_multiplier(increase_by, free_cpus, free_gpus) >= 0
|
|
):
|
|
return added_bundles
|
|
|
|
return self._modify_bundles_with_free_resources(
|
|
added_bundles,
|
|
increase_by,
|
|
free_cpus,
|
|
free_gpus,
|
|
)
|
|
|
|
|
|
_DistributeResourcesDefault = DistributeResources(add_bundles=False)
|
|
_DistributeResourcesDistributedDefault = DistributeResources(add_bundles=True)
|
|
|
|
|
|
@PublicAPI(stability="beta")
|
|
class ResourceChangingScheduler(TrialScheduler):
|
|
"""A utility scheduler to dynamically change resources of live trials.
|
|
|
|
.. versionadded:: 1.5.0
|
|
|
|
.. note::
|
|
Experimental. API may change in future releases.
|
|
|
|
The ResourceChangingScheduler works by wrapping around any other
|
|
scheduler and adjusting the resource requirements of live trials
|
|
in response to the decisions of the wrapped scheduler
|
|
through a user-specified ``resources_allocation_function``.
|
|
|
|
An example of such a function can be found in
|
|
:doc:`/tune/examples/includes/xgboost_dynamic_resources_example`.
|
|
|
|
If the functional API is used, the current trial resources can be obtained
|
|
by calling `tune.get_trial_resources()` inside the training function.
|
|
The function should be able to
|
|
:ref:`load and save checkpoints <tune-function-trainable-checkpointing>`
|
|
(the latter preferably every iteration).
|
|
|
|
If the Trainable (class) API is used, you can obtain the current trial
|
|
resources through the ``Trainable.trial_resources`` property.
|
|
|
|
Cannot be used if ``reuse_actors`` is True in ``tune.TuneConfig()``. A ValueError
|
|
will be raised in that case.
|
|
|
|
Args:
|
|
base_scheduler: The scheduler to provide decisions
|
|
about trials. If None, a default FIFOScheduler will be used.
|
|
resources_allocation_function: The callable used to change
|
|
live trial resource requiements during tuning. This callable
|
|
will be called on each trial as it finishes one step of training.
|
|
The callable must take four arguments: ``TrialRunner``, current
|
|
``Trial``, current result :class:`dict` and the
|
|
``ResourceChangingScheduler`` calling it. The callable must
|
|
return a ``PlacementGroupFactory``
|
|
or None (signifying no need for an update). If
|
|
``resources_allocation_function`` is None, no resource
|
|
requirements will be changed at any time.
|
|
By default, :class:`DistributeResources` will be used,
|
|
distributing available CPUs and GPUs over all running trials
|
|
in a robust way, without any prioritization.
|
|
|
|
Warning:
|
|
If the ``resources_allocation_function`` sets trial resource
|
|
requirements to values bigger than possible, the trial will
|
|
not run. Ensure that your callable accounts for that possibility
|
|
by setting upper limits. Consult :class:`DistributeResources`
|
|
to see how that may be done.
|
|
|
|
Example:
|
|
.. code-block:: python
|
|
|
|
base_scheduler = ASHAScheduler(max_t=16)
|
|
def my_resources_allocation_function(
|
|
tune_controller: "TuneController",
|
|
trial: Trial,
|
|
result: Dict[str, Any],
|
|
scheduler: "ResourceChangingScheduler"
|
|
) -> Optional[Union[PlacementGroupFactory, Resource]]:
|
|
# logic here
|
|
# usage of PlacementGroupFactory is strongly preferred
|
|
return PlacementGroupFactory(...)
|
|
scheduler = ResourceChangingScheduler(
|
|
base_scheduler,
|
|
my_resources_allocation_function
|
|
)
|
|
|
|
See :doc:`/tune/examples/includes/xgboost_dynamic_resources_example` for a
|
|
more detailed example.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
base_scheduler: Optional[TrialScheduler] = None,
|
|
resources_allocation_function: Optional[
|
|
Callable[
|
|
[
|
|
"TuneController",
|
|
Trial,
|
|
Dict[str, Any],
|
|
"ResourceChangingScheduler",
|
|
],
|
|
Optional[PlacementGroupFactory],
|
|
]
|
|
] = _DistributeResourcesDefault,
|
|
) -> None:
|
|
super().__init__()
|
|
if resources_allocation_function is None:
|
|
warnings.warn(
|
|
"`resources_allocation_function` is None. No resource "
|
|
"requirements will be changed at any time. Pass a "
|
|
"correctly defined function to enable functionality."
|
|
)
|
|
self._resources_allocation_function = resources_allocation_function
|
|
self._base_scheduler = base_scheduler or FIFOScheduler()
|
|
self._base_trial_resources: Optional[PlacementGroupFactory] = None
|
|
self._trials_to_reallocate: Dict[
|
|
Trial, Optional[Union[dict, PlacementGroupFactory]]
|
|
] = {}
|
|
self._reallocated_trial_ids: Set[str] = set()
|
|
self._metric = None
|
|
self._mode = None
|
|
|
|
@property
|
|
def metric(self):
|
|
return self._base_scheduler._metric
|
|
|
|
@property
|
|
def base_trial_resources(self) -> Optional[PlacementGroupFactory]:
|
|
return self._base_trial_resources
|
|
|
|
def set_search_properties(
|
|
self, metric: Optional[str], mode: Optional[str], **spec
|
|
) -> bool:
|
|
self._metric = metric
|
|
self._mode = mode
|
|
return self._base_scheduler.set_search_properties(metric, mode, **spec)
|
|
|
|
def on_trial_add(self, tune_controller: "TuneController", trial: Trial, **kwargs):
|
|
# use the first trial resources as the base
|
|
if self._base_trial_resources is None:
|
|
self._base_trial_resources = trial.placement_group_factory
|
|
# Raise error if the resources of a newly added trial don't match
|
|
# base resources, but allow trials that have already had their
|
|
# resources changed by ResourceChangingScheduler
|
|
# (those can be added again during loading from a checkpoint)
|
|
elif trial.trial_id not in self._reallocated_trial_ids:
|
|
trial_resources = trial.placement_group_factory
|
|
if trial_resources != self._base_trial_resources:
|
|
raise RuntimeError(
|
|
"ResourceChangingScheduler doesn't support trials with "
|
|
"varying base resources. First trial had "
|
|
f"{self._base_trial_resources}, trial {trial} has "
|
|
f"{trial_resources}."
|
|
)
|
|
|
|
return self._base_scheduler.on_trial_add(tune_controller, trial, **kwargs)
|
|
|
|
def on_trial_error(self, tune_controller: "TuneController", trial: Trial, **kwargs):
|
|
return self._base_scheduler.on_trial_error(tune_controller, trial, **kwargs)
|
|
|
|
def on_trial_result(
|
|
self, tune_controller: "TuneController", trial: Trial, result: Dict
|
|
) -> str:
|
|
base_scheduler_decision = self._base_scheduler.on_trial_result(
|
|
tune_controller, trial, result
|
|
)
|
|
if base_scheduler_decision == TrialScheduler.CONTINUE:
|
|
new_resources = self.reallocate_trial_resources_if_needed(
|
|
tune_controller, trial, result
|
|
)
|
|
if new_resources:
|
|
self._trials_to_reallocate[trial] = new_resources
|
|
return TrialScheduler.PAUSE
|
|
return base_scheduler_decision
|
|
|
|
def on_trial_complete(
|
|
self,
|
|
tune_controller: "TuneController",
|
|
trial: Trial,
|
|
result: Dict,
|
|
**kwargs,
|
|
):
|
|
return self._base_scheduler.on_trial_complete(
|
|
tune_controller, trial, result, **kwargs
|
|
)
|
|
|
|
def on_trial_remove(
|
|
self, tune_controller: "TuneController", trial: Trial, **kwargs
|
|
):
|
|
return self._base_scheduler.on_trial_remove(tune_controller, trial, **kwargs)
|
|
|
|
def choose_trial_to_run(
|
|
self, tune_controller: "TuneController", **kwargs
|
|
) -> Optional[Trial]:
|
|
if getattr(tune_controller, "_reuse_actors", False):
|
|
raise ValueError(
|
|
"ResourceChangingScheduler cannot be used with "
|
|
"`reuse_actors=True`. FIX THIS by setting "
|
|
"`reuse_actors=False` in `tune.TuneConfig()`."
|
|
)
|
|
|
|
any_resources_changed = False
|
|
|
|
new_trials_to_reallocate = {}
|
|
for trial, new_resources in self._trials_to_reallocate.items():
|
|
if trial.status == Trial.RUNNING:
|
|
new_trials_to_reallocate[trial] = new_resources
|
|
logger.debug(f"{trial} is still running, skipping for now")
|
|
continue
|
|
any_resources_changed = any_resources_changed or self.set_trial_resources(
|
|
trial, new_resources
|
|
)
|
|
self._trials_to_reallocate = new_trials_to_reallocate
|
|
|
|
trial = self._base_scheduler.choose_trial_to_run(tune_controller, **kwargs)
|
|
return trial
|
|
|
|
def debug_string(self) -> str:
|
|
return "(ResourceChangingScheduler) " f"{self._base_scheduler.debug_string()}"
|
|
|
|
def save(self, checkpoint_path: str):
|
|
save_object = self.__dict__
|
|
with open(checkpoint_path, "wb") as outputFile:
|
|
pickle.dump(save_object, outputFile)
|
|
|
|
def restore(self, checkpoint_path: str):
|
|
with open(checkpoint_path, "rb") as inputFile:
|
|
save_object = pickle.load(inputFile)
|
|
self.__dict__.update(save_object)
|
|
|
|
def set_trial_resources(
|
|
self, trial: Trial, new_resources: Union[Dict, PlacementGroupFactory]
|
|
) -> bool:
|
|
"""Returns True if new_resources were set."""
|
|
if new_resources:
|
|
logger.info(
|
|
f"Setting trial {trial} resource to {new_resources} "
|
|
f"with {new_resources._bundles}"
|
|
)
|
|
trial.placement_group_factory = None
|
|
trial.update_resources(new_resources)
|
|
# keep track of all trials which had their resources changed
|
|
self._reallocated_trial_ids.add(trial.trial_id)
|
|
return True
|
|
return False
|
|
|
|
def _are_resources_the_same(
|
|
self,
|
|
trial: Trial,
|
|
new_resources,
|
|
) -> bool:
|
|
"""Returns True if trial's resources are value equal to new_resources.
|
|
|
|
Only checks for PlacementGroupFactories at this moment.
|
|
"""
|
|
if (
|
|
isinstance(new_resources, PlacementGroupFactory)
|
|
and trial.placement_group_factory == new_resources
|
|
):
|
|
logger.debug(
|
|
f"{trial} PGF "
|
|
f"{trial.placement_group_factory.required_resources}"
|
|
f" and {new_resources.required_resources}"
|
|
f" are the same, skipping"
|
|
)
|
|
return True
|
|
else:
|
|
return False
|
|
|
|
def reallocate_trial_resources_if_needed(
|
|
self, tune_controller: "TuneController", trial: Trial, result: Dict
|
|
) -> Optional[Union[dict, PlacementGroupFactory]]:
|
|
"""Calls user defined resources_allocation_function. If the returned
|
|
resources are not none and not the same as currently present, returns
|
|
them. Otherwise, returns None."""
|
|
if self._resources_allocation_function is None:
|
|
return None
|
|
|
|
if not getattr(self._resources_allocation_function, "metric", None):
|
|
self._resources_allocation_function.metric = getattr(
|
|
self._base_scheduler, "_metric", self._metric
|
|
)
|
|
if not getattr(self._resources_allocation_function, "mode", None):
|
|
self._resources_allocation_function.mode = getattr(
|
|
self._base_scheduler, "_mode", self._mode
|
|
)
|
|
|
|
new_resources = self._resources_allocation_function(
|
|
tune_controller, trial, result, self
|
|
)
|
|
|
|
# if we can check if the new resources are the same,
|
|
# we do that here and skip resource allocation
|
|
if new_resources and not self._are_resources_the_same(trial, new_resources):
|
|
return new_resources
|
|
return None
|