chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,369 @@
|
||||
import logging
|
||||
import os
|
||||
import time
|
||||
from collections import namedtuple
|
||||
from numbers import Number
|
||||
from typing import Any, Dict, Optional
|
||||
|
||||
import ray
|
||||
from ray._common.constants import NODE_ID_PREFIX
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
TUNE_STATE_REFRESH_PERIOD = 10 # Refresh resources every 10 s
|
||||
|
||||
|
||||
def _to_gb(n_bytes):
|
||||
return round(n_bytes / (1024**3), 2)
|
||||
|
||||
|
||||
class _Resources(
|
||||
namedtuple(
|
||||
"_Resources",
|
||||
[
|
||||
"cpu",
|
||||
"gpu",
|
||||
"memory",
|
||||
"object_store_memory",
|
||||
"extra_cpu",
|
||||
"extra_gpu",
|
||||
"extra_memory",
|
||||
"extra_object_store_memory",
|
||||
"custom_resources",
|
||||
"extra_custom_resources",
|
||||
"has_placement_group",
|
||||
],
|
||||
)
|
||||
):
|
||||
"""Ray resources required to schedule a trial.
|
||||
|
||||
Parameters:
|
||||
cpu: Number of CPUs to allocate to the trial.
|
||||
gpu: Number of GPUs to allocate to the trial.
|
||||
memory: Memory to reserve for the trial.
|
||||
object_store_memory: Object store memory to reserve.
|
||||
extra_cpu: Extra CPUs to reserve in case the trial needs to
|
||||
launch additional Ray actors that use CPUs.
|
||||
extra_gpu: Extra GPUs to reserve in case the trial needs to
|
||||
launch additional Ray actors that use GPUs.
|
||||
extra_memory: Memory to reserve for the trial launching
|
||||
additional Ray actors that use memory.
|
||||
extra_object_store_memory: Object store memory to reserve for
|
||||
the trial launching additional Ray actors that use object store
|
||||
memory.
|
||||
custom_resources: Mapping of resource to quantity to allocate
|
||||
to the trial.
|
||||
extra_custom_resources: Extra custom resources to reserve in
|
||||
case the trial needs to launch additional Ray actors that use
|
||||
any of these custom resources.
|
||||
has_placement_group: Bool indicating if the trial also
|
||||
has an associated placement group.
|
||||
|
||||
"""
|
||||
|
||||
__slots__ = ()
|
||||
|
||||
def __new__(
|
||||
cls,
|
||||
cpu: float,
|
||||
gpu: float,
|
||||
memory: float = 0,
|
||||
object_store_memory: float = 0.0,
|
||||
extra_cpu: float = 0.0,
|
||||
extra_gpu: float = 0.0,
|
||||
extra_memory: float = 0.0,
|
||||
extra_object_store_memory: float = 0.0,
|
||||
custom_resources: Optional[dict] = None,
|
||||
extra_custom_resources: Optional[dict] = None,
|
||||
has_placement_group: bool = False,
|
||||
):
|
||||
custom_resources = custom_resources or {}
|
||||
extra_custom_resources = extra_custom_resources or {}
|
||||
leftovers = set(custom_resources) ^ set(extra_custom_resources)
|
||||
|
||||
for value in leftovers:
|
||||
custom_resources.setdefault(value, 0)
|
||||
extra_custom_resources.setdefault(value, 0)
|
||||
|
||||
cpu = round(cpu, 2)
|
||||
gpu = round(gpu, 2)
|
||||
memory = round(memory, 2)
|
||||
object_store_memory = round(object_store_memory, 2)
|
||||
extra_cpu = round(extra_cpu, 2)
|
||||
extra_gpu = round(extra_gpu, 2)
|
||||
extra_memory = round(extra_memory, 2)
|
||||
extra_object_store_memory = round(extra_object_store_memory, 2)
|
||||
custom_resources = {
|
||||
resource: round(value, 2) for resource, value in custom_resources.items()
|
||||
}
|
||||
extra_custom_resources = {
|
||||
resource: round(value, 2)
|
||||
for resource, value in extra_custom_resources.items()
|
||||
}
|
||||
|
||||
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.", 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,
|
||||
has_placement_group,
|
||||
)
|
||||
|
||||
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
|
||||
if not res.startswith(NODE_ID_PREFIX)
|
||||
]
|
||||
)
|
||||
if custom_summary:
|
||||
summary += " ({})".format(custom_summary)
|
||||
return summary
|
||||
|
||||
def cpu_total(self):
|
||||
return self.cpu + self.extra_cpu
|
||||
|
||||
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
|
||||
)
|
||||
|
||||
def get(self, key):
|
||||
return self.custom_resources.get(key, 0)
|
||||
|
||||
def is_nonnegative(self):
|
||||
all_values = [self.cpu, self.gpu, self.extra_cpu, self.extra_gpu]
|
||||
all_values += list(self.custom_resources.values())
|
||||
all_values += list(self.extra_custom_resources.values())
|
||||
return all(v >= 0 for v in all_values)
|
||||
|
||||
@classmethod
|
||||
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 = {
|
||||
k: original.custom_resources.get(k, 0)
|
||||
- to_remove.custom_resources.get(k, 0)
|
||||
for k in all_resources
|
||||
}
|
||||
extra_custom_res = {
|
||||
k: original.extra_custom_resources.get(k, 0)
|
||||
- to_remove.extra_custom_resources.get(k, 0)
|
||||
for k in all_resources
|
||||
}
|
||||
return _Resources(
|
||||
cpu,
|
||||
gpu,
|
||||
memory,
|
||||
object_store_memory,
|
||||
extra_cpu,
|
||||
extra_gpu,
|
||||
extra_memory,
|
||||
extra_object_store_memory,
|
||||
new_custom_res,
|
||||
extra_custom_res,
|
||||
)
|
||||
|
||||
|
||||
class _ResourceUpdater:
|
||||
"""Periodic Resource updater for Tune.
|
||||
|
||||
Initially, all resources are set to 0. The updater will try to update resources
|
||||
when (1) init ResourceUpdater (2) call "update_avail_resources", "num_cpus"
|
||||
or "num_gpus".
|
||||
|
||||
The update takes effect when (1) Ray is initialized (2) the interval between
|
||||
this and last update is larger than "refresh_period"
|
||||
"""
|
||||
|
||||
def __init__(self, refresh_period: Optional[float] = None):
|
||||
self._avail_resources = _Resources(cpu=0, gpu=0)
|
||||
|
||||
if refresh_period is None:
|
||||
refresh_period = float(
|
||||
os.environ.get("TUNE_STATE_REFRESH_PERIOD", TUNE_STATE_REFRESH_PERIOD)
|
||||
)
|
||||
self._refresh_period = refresh_period
|
||||
self._last_resource_refresh = float("-inf")
|
||||
self.update_avail_resources()
|
||||
|
||||
def update_avail_resources(self, num_retries: int = 5, force: bool = False):
|
||||
if not ray.is_initialized():
|
||||
return
|
||||
if (
|
||||
time.time() - self._last_resource_refresh < self._refresh_period
|
||||
and not force
|
||||
):
|
||||
return
|
||||
logger.debug("Checking Ray cluster resources.")
|
||||
resources = None
|
||||
for i in range(num_retries):
|
||||
if i > 0:
|
||||
logger.warning(
|
||||
f"Cluster resources not detected or are 0. Attempt #{i + 1}...",
|
||||
)
|
||||
time.sleep(0.5)
|
||||
resources = ray.cluster_resources()
|
||||
if resources:
|
||||
break
|
||||
|
||||
if not resources:
|
||||
# NOTE: This hides the possibility that Ray may be waiting for
|
||||
# clients to connect.
|
||||
resources.setdefault("CPU", 0)
|
||||
resources.setdefault("GPU", 0)
|
||||
logger.warning(
|
||||
"Cluster resources cannot be detected or are 0. "
|
||||
"You can resume this experiment by passing in `resume=True` to `run`."
|
||||
)
|
||||
|
||||
resources = resources.copy()
|
||||
num_cpus = resources.pop("CPU", 0)
|
||||
num_gpus = resources.pop("GPU", 0)
|
||||
memory = resources.pop("memory", 0)
|
||||
object_store_memory = resources.pop("object_store_memory", 0)
|
||||
custom_resources = resources
|
||||
|
||||
self._avail_resources = _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()
|
||||
|
||||
def _get_used_avail_resources(self, total_allocated_resources: Dict[str, Any]):
|
||||
total_allocated_resources = total_allocated_resources.copy()
|
||||
|
||||
used_cpu = total_allocated_resources.pop("CPU", 0)
|
||||
total_cpu = self._avail_resources.cpu
|
||||
used_gpu = total_allocated_resources.pop("GPU", 0)
|
||||
total_gpu = self._avail_resources.gpu
|
||||
|
||||
custom_used_total = {
|
||||
name: (
|
||||
total_allocated_resources.get(name, 0.0),
|
||||
self._avail_resources.get_res_total(name),
|
||||
)
|
||||
for name in self._avail_resources.custom_resources
|
||||
if not name.startswith(NODE_ID_PREFIX)
|
||||
and (total_allocated_resources.get(name, 0.0) > 0 or "_group_" not in name)
|
||||
}
|
||||
return used_cpu, total_cpu, used_gpu, total_gpu, custom_used_total
|
||||
|
||||
def debug_string(self, total_allocated_resources: Dict[str, Any]) -> str:
|
||||
"""Returns a human readable message for printing to the console."""
|
||||
if self._last_resource_refresh > 0:
|
||||
(
|
||||
used_cpu,
|
||||
total_cpu,
|
||||
used_gpu,
|
||||
total_gpu,
|
||||
custom_used_total,
|
||||
) = self._get_used_avail_resources(total_allocated_resources)
|
||||
|
||||
if (
|
||||
used_cpu > total_cpu
|
||||
or used_gpu > total_gpu
|
||||
or any(used > total for (used, total) in custom_used_total.values())
|
||||
):
|
||||
# If any of the used resources are higher than what we currently think
|
||||
# is available, update our state and re-fetch
|
||||
self.update_avail_resources(force=True)
|
||||
(
|
||||
used_cpu,
|
||||
total_cpu,
|
||||
used_gpu,
|
||||
total_gpu,
|
||||
custom_used_total,
|
||||
) = self._get_used_avail_resources(total_allocated_resources)
|
||||
|
||||
status = (
|
||||
f"Logical resource usage: {used_cpu}/{total_cpu} CPUs, "
|
||||
f"{used_gpu}/{total_gpu} GPUs"
|
||||
)
|
||||
customs = ", ".join(
|
||||
f"{used}/{total} {name}"
|
||||
for name, (used, total) in custom_used_total.items()
|
||||
)
|
||||
|
||||
if customs:
|
||||
status += f" ({customs})"
|
||||
return status
|
||||
else:
|
||||
return "Logical resource usage: ?"
|
||||
|
||||
def get_num_cpus(self) -> int:
|
||||
self.update_avail_resources()
|
||||
return self._avail_resources.cpu
|
||||
|
||||
def get_num_gpus(self) -> int:
|
||||
self.update_avail_resources()
|
||||
return self._avail_resources.gpu
|
||||
|
||||
def __reduce__(self):
|
||||
# Do not need to serialize resources, because we can always
|
||||
# update it again. This also prevents keeping outdated resources
|
||||
# when deserialized.
|
||||
return _ResourceUpdater, (self._refresh_period,)
|
||||
Reference in New Issue
Block a user