Ray, Tune, and RLlib support for memory, object_store_memory options (#5226)

This commit is contained in:
Eric Liang
2019-08-22 14:01:10 +08:00
committed by Robert Nishihara
parent c852213b83
commit e2e30ca507
40 changed files with 1006 additions and 296 deletions
+51 -15
View File
@@ -11,6 +11,8 @@ import time
import traceback
import ray
from ray import ray_constants
from ray.resource_spec import ResourceSpec
from ray.tune.error import AbortTrialExecution
from ray.tune.logger import NoopLogger
from ray.tune.trial import Trial, Checkpoint
@@ -61,7 +63,7 @@ class RayTrialExecutor(TrialExecutor):
logger.info("Initializing Ray automatically."
"For cluster usage or custom Ray initialization, "
"call `ray.init(...)` before `tune.run`.")
ray.init(object_store_memory=int(1e8))
ray.init()
if ray.is_initialized():
self._update_avail_resources()
@@ -85,6 +87,8 @@ class RayTrialExecutor(TrialExecutor):
cls = ray.remote(
num_cpus=trial.resources.cpu,
num_gpus=trial.resources.gpu,
memory=trial.resources.memory,
object_store_memory=trial.resources.object_store_memory,
resources=trial.resources.custom_resources)(
trial._get_trainable_cls())
@@ -360,6 +364,9 @@ class RayTrialExecutor(TrialExecutor):
self._committed_resources = Resources(
committed.cpu + resources.cpu_total(),
committed.gpu + resources.gpu_total(),
committed.memory + resources.memory_total(),
committed.object_store_memory +
resources.object_store_memory_total(),
custom_resources=custom_resources)
def _return_resources(self, resources):
@@ -388,8 +395,7 @@ class RayTrialExecutor(TrialExecutor):
# TODO(rliaw): Remove this when local mode is fixed.
# https://github.com/ray-project/ray/issues/4147
logger.debug("Using resources for local machine.")
resources = ray.services.check_and_update_resources(
None, None, None)
resources = ResourceSpec().resolve(True).to_resource_dict()
if not resources:
logger.warning(
"Cluster resources not detected or are 0. Retrying...")
@@ -407,10 +413,17 @@ class RayTrialExecutor(TrialExecutor):
resources = resources.copy()
num_cpus = resources.pop("CPU", 0)
num_gpus = resources.pop("GPU", 0)
memory = ray_constants.from_memory_units(resources.pop("memory", 0))
object_store_memory = ray_constants.from_memory_units(
resources.pop("object_store_memory", 0))
custom_resources = resources
self._avail_resources = Resources(
int(num_cpus), int(num_gpus), custom_resources=custom_resources)
int(num_cpus),
int(num_gpus),
memory=int(memory),
object_store_memory=int(object_store_memory),
custom_resources=custom_resources)
self._last_resource_refresh = time.time()
self._resources_initialized = True
@@ -429,7 +442,10 @@ class RayTrialExecutor(TrialExecutor):
have_space = (
resources.cpu_total() <= currently_available.cpu
and resources.gpu_total() <= currently_available.gpu and all(
and resources.gpu_total() <= currently_available.gpu
and resources.memory_total() <= currently_available.memory
and resources.object_store_memory_total() <=
currently_available.object_store_memory and all(
resources.get_res_total(res) <= currently_available.get(res)
for res in resources.custom_resources))
@@ -438,11 +454,15 @@ class RayTrialExecutor(TrialExecutor):
can_overcommit = self._queue_trials
if (resources.cpu_total() > 0 and currently_available.cpu <= 0) or \
(resources.gpu_total() > 0 and currently_available.gpu <= 0) or \
any((resources.get_res_total(res_name) > 0
and currently_available.get(res_name) <= 0)
for res_name in resources.custom_resources):
if ((resources.cpu_total() > 0 and currently_available.cpu <= 0)
or (resources.gpu_total() > 0 and currently_available.gpu <= 0)
or
(resources.memory_total() > 0 and currently_available.memory <= 0)
or (resources.object_store_memory_total() > 0
and currently_available.object_store_memory <= 0) or any(
(resources.get_res_total(res_name) > 0
and currently_available.get(res_name) <= 0)
for res_name in resources.custom_resources)):
can_overcommit = False # requested resource is already saturated
if can_overcommit:
@@ -461,9 +481,17 @@ class RayTrialExecutor(TrialExecutor):
"""Returns a human readable message for printing to the console."""
if self._resources_initialized:
status = "Resources requested: {}/{} CPUs, {}/{} GPUs".format(
self._committed_resources.cpu, self._avail_resources.cpu,
self._committed_resources.gpu, self._avail_resources.gpu)
status = ("Resources requested: {}/{} CPUs, {}/{} GPUs, "
"{}/{} GiB heap, {}/{} GiB objects".format(
self._committed_resources.cpu,
self._avail_resources.cpu,
self._committed_resources.gpu,
self._avail_resources.gpu,
_to_gb(self._committed_resources.memory),
_to_gb(self._avail_resources.memory),
_to_gb(
self._committed_resources.object_store_memory),
_to_gb(self._avail_resources.object_store_memory)))
customs = ", ".join([
"{}/{} {}".format(
self._committed_resources.get_res_total(name),
@@ -480,8 +508,12 @@ class RayTrialExecutor(TrialExecutor):
"""Returns a string describing the total resources available."""
if self._resources_initialized:
res_str = "{} CPUs, {} GPUs".format(self._avail_resources.cpu,
self._avail_resources.gpu)
res_str = ("{} CPUs, {} GPUs, "
"{} GiB heap, {} GiB objects".format(
self._avail_resources.cpu,
self._avail_resources.gpu,
_to_gb(self._avail_resources.memory),
_to_gb(self._avail_resources.object_store_memory)))
if self._avail_resources.custom_resources:
custom = ", ".join(
"{} {}".format(
@@ -589,3 +621,7 @@ class RayTrialExecutor(TrialExecutor):
return ray.get(
trial.runner.export_model.remote(trial.export_formats))
return {}
def _to_gb(n_bytes):
return round(n_bytes / (1024**3), 2)