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)
+55 -11
View File
@@ -17,18 +17,26 @@ logger = logging.getLogger(__name__)
class Resources(
namedtuple("Resources", [
"cpu", "gpu", "extra_cpu", "extra_gpu", "custom_resources",
"extra_custom_resources"
"cpu", "gpu", "memory", "object_store_memory", "extra_cpu",
"extra_gpu", "extra_memory", "extra_object_store_memory",
"custom_resources", "extra_custom_resources"
])):
"""Ray resources required to schedule a trial.
Attributes:
cpu (float): Number of CPUs to allocate to the trial.
gpu (float): Number of GPUs to allocate to the trial.
memory (float): Memory to reserve for the trial.
object_store_memory (float): Object store memory to reserve.
extra_cpu (float): Extra CPUs to reserve in case the trial needs to
launch additional Ray actors that use CPUs.
extra_gpu (float): Extra GPUs to reserve in case the trial needs to
launch additional Ray actors that use GPUs.
extra_memory (float): Memory to reserve for the trial launching
additional Ray actors that use memory.
extra_object_store_memory (float): Object store memory to reserve for
the trial launching additional Ray actors that use object store
memory.
custom_resources (dict): Mapping of resource to quantity to allocate
to the trial.
extra_custom_resources (dict): Extra custom resources to reserve in
@@ -42,8 +50,12 @@ class Resources(
def __new__(cls,
cpu,
gpu,
memory=0,
object_store_memory=0,
extra_cpu=0,
extra_gpu=0,
extra_memory=0,
extra_object_store_memory=0,
custom_resources=None,
extra_custom_resources=None):
custom_resources = custom_resources or {}
@@ -54,19 +66,32 @@ class Resources(
custom_resources.setdefault(value, 0)
extra_custom_resources.setdefault(value, 0)
all_values = [cpu, gpu, extra_cpu, extra_gpu]
all_values = [
cpu, gpu, memory, object_store_memory, extra_cpu, extra_gpu,
extra_memory, extra_object_store_memory
]
all_values += list(custom_resources.values())
all_values += list(extra_custom_resources.values())
assert len(custom_resources) == len(extra_custom_resources)
for entry in all_values:
assert isinstance(entry, Number), "Improper resource value."
return super(Resources,
cls).__new__(cls, cpu, gpu, extra_cpu, extra_gpu,
custom_resources, extra_custom_resources)
assert isinstance(entry, Number), ("Improper resource value.",
entry)
return super(Resources, cls).__new__(
cls, cpu, gpu, memory, object_store_memory, extra_cpu, extra_gpu,
extra_memory, extra_object_store_memory, custom_resources,
extra_custom_resources)
def summary_string(self):
summary = "{} CPUs, {} GPUs".format(self.cpu + self.extra_cpu,
self.gpu + self.extra_gpu)
if self.memory or self.extra_memory:
summary += ", {} GiB heap".format(
round((self.memory + self.extra_memory) / (1024**3), 2))
if self.object_store_memory or self.extra_object_store_memory:
summary += ", {} GiB objects".format(
round(
(self.object_store_memory + self.extra_object_store_memory)
/ (1024**3), 2))
custom_summary = ", ".join([
"{} {}".format(self.get_res_total(res), res)
for res in self.custom_resources
@@ -81,6 +106,12 @@ class Resources(
def gpu_total(self):
return self.gpu + self.extra_gpu
def memory_total(self):
return self.memory + self.extra_memory
def object_store_memory_total(self):
return self.object_store_memory + self.extra_object_store_memory
def get_res_total(self, key):
return self.custom_resources.get(
key, 0) + self.extra_custom_resources.get(key, 0)
@@ -98,8 +129,14 @@ class Resources(
def subtract(cls, original, to_remove):
cpu = original.cpu - to_remove.cpu
gpu = original.gpu - to_remove.gpu
memory = original.memory - to_remove.memory
object_store_memory = (
original.object_store_memory - to_remove.object_store_memory)
extra_cpu = original.extra_cpu - to_remove.extra_cpu
extra_gpu = original.extra_gpu - to_remove.extra_gpu
extra_memory = original.extra_memory - to_remove.extra_memory
extra_object_store_memory = (original.extra_object_store_memory -
to_remove.extra_object_store_memory)
all_resources = set(original.custom_resources).union(
set(to_remove.custom_resources))
new_custom_res = {
@@ -112,8 +149,9 @@ class Resources(
to_remove.extra_custom_resources.get(k, 0)
for k in all_resources
}
return Resources(cpu, gpu, extra_cpu, extra_gpu, new_custom_res,
extra_custom_res)
return Resources(cpu, gpu, memory, object_store_memory, extra_cpu,
extra_gpu, extra_memory, extra_object_store_memory,
new_custom_res, extra_custom_res)
def to_json(self):
return resources_to_json(self)
@@ -134,8 +172,10 @@ def json_to_resources(data):
"Unknown resource field {}, must be one of {}".format(
k, Resources._fields))
return Resources(
data.get("cpu", 1), data.get("gpu", 0), data.get("extra_cpu", 0),
data.get("extra_gpu", 0), data.get("custom_resources"),
data.get("cpu", 1), data.get("gpu", 0), data.get("memory", 0),
data.get("object_store_memory", 0), data.get("extra_cpu", 0),
data.get("extra_gpu", 0), data.get("extra_memory", 0),
data.get("extra_object_store_memory", 0), data.get("custom_resources"),
data.get("extra_custom_resources"))
@@ -145,8 +185,12 @@ def resources_to_json(resources):
return {
"cpu": resources.cpu,
"gpu": resources.gpu,
"memory": resources.memory,
"object_store_memory": resources.object_store_memory,
"extra_cpu": resources.extra_cpu,
"extra_gpu": resources.extra_gpu,
"extra_memory": resources.extra_memory,
"extra_object_store_memory": resources.extra_object_store_memory,
"custom_resources": resources.custom_resources.copy(),
"extra_custom_resources": resources.extra_custom_resources.copy()
}
+1 -1
View File
@@ -44,7 +44,7 @@ else:
class TrainableFunctionApiTest(unittest.TestCase):
def setUp(self):
ray.init(num_cpus=4, num_gpus=0, object_store_memory=int(1e8))
ray.init(num_cpus=4, num_gpus=0, object_store_memory=150 * 1024 * 1024)
def tearDown(self):
ray.shutdown()
+2
View File
@@ -178,6 +178,7 @@ class Trial(object):
self.result_logger = None
self.last_debug = 0
self.error_file = None
self.error_msg = None
self.num_failures = 0
self.custom_trial_name = None
@@ -270,6 +271,7 @@ class Trial(object):
with open(error_file, "w") as f:
f.write(error_msg)
self.error_file = error_file
self.error_msg = error_msg
def should_stop(self, result):
"""Whether the given result meets this trial's stopping criteria."""
+3 -3
View File
@@ -454,8 +454,8 @@ class TrialRunner(object):
def _memory_debug_string(self):
try:
import psutil
total_gb = psutil.virtual_memory().total / 1e9
used_gb = total_gb - psutil.virtual_memory().available / 1e9
total_gb = psutil.virtual_memory().total / (1024**3)
used_gb = total_gb - psutil.virtual_memory().available / (1024**3)
if used_gb > total_gb * 0.9:
warn = (": ***LOW MEMORY*** less than 10% of the memory on "
"this node is available for use. This can cause "
@@ -465,7 +465,7 @@ class TrialRunner(object):
"`object_store_memory` when calling `ray.init`.")
else:
warn = ""
return "Memory usage on this node: {}/{} GB{}".format(
return "Memory usage on this node: {}/{} GiB{}".format(
round(used_gb, 1), round(total_gb, 1), warn)
except ImportError:
return ("Unknown memory usage. Please run `pip install psutil` "