wehub-resource-sync
474 lines
19 KiB
Python
474 lines
19 KiB
Python
import json
|
|
import logging
|
|
import os
|
|
from typing import Dict, Optional, Tuple
|
|
|
|
import ray
|
|
import ray._private.ray_constants as ray_constants
|
|
from ray._common.constants import HEAD_NODE_RESOURCE_NAME, NODE_ID_PREFIX
|
|
from ray._common.utils import RESOURCE_CONSTRAINT_PREFIX
|
|
from ray._private import accelerators
|
|
from ray._private.accelerators import AcceleratorManager
|
|
from ray._private.resource_isolation_config import ResourceIsolationConfig
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class ResourceAndLabelSpec:
|
|
"""Represents the resource and label configuration passed to a raylet.
|
|
|
|
All fields can be None. Before starting services, resolve() should be
|
|
called to return a ResourceAndLabelSpec with unknown values filled in with
|
|
merged values based on the local machine and user specifications.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
num_cpus: Optional[int] = None,
|
|
num_gpus: Optional[int] = None,
|
|
memory: Optional[float] = None,
|
|
object_store_memory: Optional[float] = None,
|
|
resources: Optional[Dict[str, float]] = None,
|
|
labels: Optional[Dict[str, str]] = None,
|
|
):
|
|
"""
|
|
Initialize a ResourceAndLabelSpec
|
|
|
|
Args:
|
|
num_cpus: The CPUs allocated for this raylet.
|
|
num_gpus: The GPUs allocated for this raylet.
|
|
memory: The memory allocated for this raylet.
|
|
object_store_memory: The object store memory allocated for this raylet.
|
|
resources: The custom resources allocated for this raylet.
|
|
labels: The labels associated with this node. Labels can be used along
|
|
with resources for scheduling.
|
|
"""
|
|
self.num_cpus = num_cpus
|
|
self.num_gpus = num_gpus
|
|
self.memory = memory
|
|
self.object_store_memory = object_store_memory
|
|
self.resources = resources
|
|
self.labels = labels
|
|
self._is_resolved = False
|
|
|
|
def resolved(self) -> bool:
|
|
"""Returns if resolve() has been called for this ResourceAndLabelSpec
|
|
and default values are filled out."""
|
|
return self._is_resolved
|
|
|
|
def _all_fields_set(self) -> bool:
|
|
"""Returns whether all fields in this ResourceAndLabelSpec are not None."""
|
|
return all(
|
|
v is not None
|
|
for v in (
|
|
self.num_cpus,
|
|
self.num_gpus,
|
|
self.memory,
|
|
self.object_store_memory,
|
|
self.resources,
|
|
self.labels,
|
|
)
|
|
)
|
|
|
|
def to_resource_dict(self):
|
|
"""Returns a dict suitable to pass to raylet initialization.
|
|
|
|
This renames num_cpus / num_gpus to "CPU" / "GPU",
|
|
and check types and values.
|
|
"""
|
|
assert self.resolved()
|
|
|
|
resources = dict(
|
|
self.resources,
|
|
CPU=self.num_cpus,
|
|
GPU=self.num_gpus,
|
|
memory=int(self.memory),
|
|
object_store_memory=int(self.object_store_memory),
|
|
)
|
|
|
|
resources = {
|
|
resource_label: resource_quantity
|
|
for resource_label, resource_quantity in resources.items()
|
|
if resource_quantity != 0
|
|
}
|
|
|
|
# Check types.
|
|
for resource_label, resource_quantity in resources.items():
|
|
assert isinstance(resource_quantity, int) or isinstance(
|
|
resource_quantity, float
|
|
), (
|
|
f"{resource_label} ({type(resource_quantity)}): " f"{resource_quantity}"
|
|
)
|
|
if (
|
|
isinstance(resource_quantity, float)
|
|
and not resource_quantity.is_integer()
|
|
):
|
|
raise ValueError(
|
|
"Resource quantities must all be whole numbers. "
|
|
"Violated by resource '{}' in {}.".format(resource_label, resources)
|
|
)
|
|
if resource_quantity < 0:
|
|
raise ValueError(
|
|
"Resource quantities must be nonnegative. "
|
|
"Violated by resource '{}' in {}.".format(resource_label, resources)
|
|
)
|
|
if resource_quantity > ray_constants.MAX_RESOURCE_QUANTITY:
|
|
raise ValueError(
|
|
"Resource quantities must be at most {}. "
|
|
"Violated by resource '{}' in {}.".format(
|
|
ray_constants.MAX_RESOURCE_QUANTITY, resource_label, resources
|
|
)
|
|
)
|
|
|
|
return resources
|
|
|
|
def resolve(
|
|
self,
|
|
is_head: bool,
|
|
node_ip_address: Optional[str] = None,
|
|
resource_isolation_config: Optional[ResourceIsolationConfig] = None,
|
|
) -> "ResourceAndLabelSpec":
|
|
"""Fills out this ResourceAndLabelSpec instance with merged values from system defaults and user specification.
|
|
|
|
Args:
|
|
is_head: Whether this is the head node.
|
|
node_ip_address: The IP address of the node that we are on.
|
|
This is used to automatically create a node id resource.
|
|
resource_isolation_config: Optional resource isolation config. When
|
|
enabled and memory is not explicitly set, the system reserved
|
|
memory for resource isolation is subtracted from available user memory.
|
|
|
|
Returns:
|
|
ResourceAndLabelSpec: This instance with all fields resolved.
|
|
"""
|
|
|
|
self._resolve_resources(is_head=is_head, node_ip_address=node_ip_address)
|
|
|
|
# Resolve accelerator-specific resources
|
|
(
|
|
accelerator_manager,
|
|
num_accelerators,
|
|
) = ResourceAndLabelSpec._get_current_node_accelerator(
|
|
self.num_gpus, self.resources
|
|
)
|
|
self._resolve_accelerator_resources(accelerator_manager, num_accelerators)
|
|
|
|
# Default num_gpus value if unset by user and unable to auto-detect.
|
|
if self.num_gpus is None:
|
|
self.num_gpus = 0
|
|
|
|
# Resolve and merge node labels from all sources (params, env, and default).
|
|
self._resolve_labels(accelerator_manager)
|
|
|
|
# Resolve memory resources
|
|
self._resolve_memory_resources(resource_isolation_config)
|
|
|
|
self._is_resolved = True
|
|
assert self._all_fields_set()
|
|
return self
|
|
|
|
@staticmethod
|
|
def _load_env_resources() -> Dict[str, float]:
|
|
"""Load resource overrides from the environment, if present."""
|
|
env_resources = {}
|
|
env_string = os.getenv(ray_constants.RESOURCES_ENVIRONMENT_VARIABLE)
|
|
if env_string:
|
|
try:
|
|
env_resources = json.loads(env_string)
|
|
except Exception:
|
|
logger.exception(f"Failed to load {env_string}")
|
|
raise
|
|
logger.debug(f"Autoscaler overriding resources: {env_resources}.")
|
|
return env_resources
|
|
|
|
@staticmethod
|
|
def _merge_resources(env_dict: Dict[str, float], params_dict: Dict[str, float]):
|
|
"""Merge environment and Ray param-provided resources, with env values taking precedence.
|
|
Returns separated special case params (CPU/GPU/memory) and the merged resource dict.
|
|
"""
|
|
num_cpus = env_dict.pop("CPU", None)
|
|
num_gpus = env_dict.pop("GPU", None)
|
|
memory = env_dict.pop("memory", None)
|
|
object_store_memory = env_dict.pop("object_store_memory", None)
|
|
|
|
result = params_dict.copy()
|
|
result.update(env_dict)
|
|
|
|
for key in set(env_dict.keys()).intersection(params_dict or {}):
|
|
if params_dict[key] != env_dict[key]:
|
|
logger.warning(
|
|
f"Autoscaler is overriding your resource: {key}: "
|
|
f"{params_dict[key]} with {env_dict[key]}."
|
|
)
|
|
|
|
return num_cpus, num_gpus, memory, object_store_memory, result
|
|
|
|
def _resolve_resources(
|
|
self, is_head: bool, node_ip_address: Optional[str] = None
|
|
) -> None:
|
|
"""Resolve CPU, GPU, and custom resources. Merges resources from environment,
|
|
Ray params, and defaults in that order of precedence."""
|
|
|
|
# Load environment override resources and merge with resources passed
|
|
# in from Ray Params. Separates special case params if found in env.
|
|
env_resources = ResourceAndLabelSpec._load_env_resources()
|
|
(
|
|
num_cpus,
|
|
num_gpus,
|
|
memory,
|
|
object_store_memory,
|
|
merged_resources,
|
|
) = ResourceAndLabelSpec._merge_resources(env_resources, self.resources or {})
|
|
|
|
self.num_cpus = self.num_cpus if num_cpus is None else num_cpus
|
|
self.num_gpus = self.num_gpus if num_gpus is None else num_gpus
|
|
self.memory = self.memory if memory is None else memory
|
|
self.object_store_memory = (
|
|
self.object_store_memory
|
|
if object_store_memory is None
|
|
else object_store_memory
|
|
)
|
|
self.resources = merged_resources
|
|
|
|
if node_ip_address is None:
|
|
node_ip_address = ray.util.get_node_ip_address()
|
|
|
|
# Automatically create a node id resource on each node. This is
|
|
# queryable with ray._private.state.node_ids() and
|
|
# ray._private.state.current_node_id().
|
|
self.resources[NODE_ID_PREFIX + node_ip_address] = 1.0
|
|
|
|
# Automatically create a head node resource.
|
|
if HEAD_NODE_RESOURCE_NAME in self.resources:
|
|
raise ValueError(
|
|
f"{HEAD_NODE_RESOURCE_NAME}"
|
|
" is a reserved resource name, use another name instead."
|
|
)
|
|
if is_head:
|
|
self.resources[HEAD_NODE_RESOURCE_NAME] = 1.0
|
|
|
|
# Auto-detect CPU count if not explicitly set
|
|
if self.num_cpus is None:
|
|
self.num_cpus = ray._private.utils.get_num_cpus()
|
|
|
|
@staticmethod
|
|
def _load_env_labels() -> Dict[str, str]:
|
|
env_override_labels = {}
|
|
env_override_labels_string = os.getenv(
|
|
ray_constants.LABELS_ENVIRONMENT_VARIABLE
|
|
)
|
|
if env_override_labels_string:
|
|
try:
|
|
env_override_labels = json.loads(env_override_labels_string)
|
|
except Exception:
|
|
logger.exception(f"Failed to load {env_override_labels_string}")
|
|
raise
|
|
logger.info(f"Autoscaler overriding labels: {env_override_labels}.")
|
|
|
|
return env_override_labels
|
|
|
|
@staticmethod
|
|
def _get_default_labels(
|
|
accelerator_manager: Optional[AcceleratorManager],
|
|
) -> Dict[str, str]:
|
|
default_labels = {}
|
|
|
|
# Get environment variables populated from K8s Pod Spec
|
|
node_group = os.environ.get(ray._raylet.NODE_TYPE_NAME_ENV, "")
|
|
market_type = os.environ.get(ray._raylet.NODE_MARKET_TYPE_ENV, "")
|
|
availability_region = os.environ.get(ray._raylet.NODE_REGION_ENV, "")
|
|
availability_zone = os.environ.get(ray._raylet.NODE_ZONE_ENV, "")
|
|
|
|
# Map environment variables to default ray node labels
|
|
if market_type:
|
|
default_labels[ray._raylet.RAY_NODE_MARKET_TYPE_KEY] = market_type
|
|
if node_group:
|
|
default_labels[ray._raylet.RAY_NODE_GROUP_KEY] = node_group
|
|
if availability_zone:
|
|
default_labels[ray._raylet.RAY_NODE_ZONE_KEY] = availability_zone
|
|
if availability_region:
|
|
default_labels[ray._raylet.RAY_NODE_REGION_KEY] = availability_region
|
|
|
|
# Get accelerator type from AcceleratorManager
|
|
if accelerator_manager:
|
|
accelerator_type = accelerator_manager.get_current_node_accelerator_type()
|
|
if accelerator_type:
|
|
default_labels[
|
|
ray._raylet.RAY_NODE_ACCELERATOR_TYPE_KEY
|
|
] = accelerator_type
|
|
|
|
# Set TPU specific default labels to enable multi-host scheduling.
|
|
if accelerator_manager.get_resource_name() == "TPU":
|
|
tpu_labels = accelerator_manager.get_current_node_accelerator_labels()
|
|
if tpu_labels:
|
|
default_labels.update(tpu_labels)
|
|
|
|
return default_labels
|
|
|
|
def _resolve_labels(
|
|
self, accelerator_manager: Optional[AcceleratorManager]
|
|
) -> None:
|
|
"""Resolve and merge environment override, user-input from params, and Ray default
|
|
labels in that order of precedence."""
|
|
|
|
# Start with a dictionary filled out with Ray default labels
|
|
merged = ResourceAndLabelSpec._get_default_labels(accelerator_manager)
|
|
|
|
# Merge user-specified labels from Ray params
|
|
for key, val in (self.labels or {}).items():
|
|
if key in merged and merged[key] != val:
|
|
logger.warning(
|
|
f"User label is overriding Ray default label: {key}: "
|
|
f"{key}: {merged[key]} to "
|
|
f"{key}: {self.labels[key]}."
|
|
)
|
|
merged[key] = val
|
|
|
|
# Merge autoscaler override labels from environment
|
|
env_labels = ResourceAndLabelSpec._load_env_labels()
|
|
for key, val in (env_labels or {}).items():
|
|
if key in merged and merged[key] != val:
|
|
logger.warning(
|
|
"Autoscaler is overriding your label:"
|
|
f"{key}: {merged[key]} to "
|
|
f"{key}: {env_labels[key]}."
|
|
)
|
|
merged[key] = val
|
|
|
|
self.labels = merged
|
|
|
|
def _resolve_accelerator_resources(self, accelerator_manager, num_accelerators):
|
|
"""Detect and update accelerator resources on a node."""
|
|
if not accelerator_manager:
|
|
return
|
|
|
|
accelerator_resource_name = accelerator_manager.get_resource_name()
|
|
visible_accelerator_ids = (
|
|
accelerator_manager.get_current_process_visible_accelerator_ids()
|
|
)
|
|
|
|
# Check that the number of accelerators that the raylet wants doesn't
|
|
# exceed the amount allowed by visible accelerator ids.
|
|
if (
|
|
num_accelerators is not None
|
|
and visible_accelerator_ids is not None
|
|
and num_accelerators > len(visible_accelerator_ids)
|
|
):
|
|
raise ValueError(
|
|
f"Attempting to start raylet with {num_accelerators} "
|
|
f"{accelerator_resource_name}, "
|
|
f"but {accelerator_manager.get_visible_accelerator_ids_env_var()} "
|
|
f"contains {visible_accelerator_ids}."
|
|
)
|
|
|
|
if accelerator_resource_name == "GPU":
|
|
self.num_gpus = num_accelerators
|
|
else:
|
|
self.resources[accelerator_resource_name] = num_accelerators
|
|
|
|
accelerator_type = accelerator_manager.get_current_node_accelerator_type()
|
|
if accelerator_type:
|
|
self.resources[f"{RESOURCE_CONSTRAINT_PREFIX}{accelerator_type}"] = 1
|
|
additional_resources = (
|
|
accelerator_manager.get_current_node_additional_resources()
|
|
)
|
|
if additional_resources:
|
|
self.resources.update(additional_resources)
|
|
|
|
def _resolve_memory_resources(
|
|
self,
|
|
resource_isolation_config: Optional[ResourceIsolationConfig] = None,
|
|
):
|
|
"""
|
|
Resolves logical and object store memory resources if not
|
|
explicitly set.
|
|
|
|
Args:
|
|
resource_isolation_config: Optional resource isolation config. When
|
|
enabled and memory is not explicitly set, the system reserved
|
|
memory for resource isolation is subtracted from available user memory.
|
|
"""
|
|
# Choose a default object store size.
|
|
system_memory = ray._common.utils.get_system_memory()
|
|
if (
|
|
resource_isolation_config is not None
|
|
and resource_isolation_config.is_enabled()
|
|
):
|
|
available_memory_bytes = (
|
|
system_memory
|
|
- resource_isolation_config.system_reserved_memory
|
|
- ray_constants.DEFAULT_USER_PHYSICAL_LOGICAL_MEMORY_LIMIT_BUFFER_BYTES
|
|
)
|
|
else:
|
|
available_memory_bytes = ray._private.utils.estimate_available_memory()
|
|
if self.object_store_memory is None:
|
|
self.object_store_memory = ray._private.utils.resolve_object_store_memory(
|
|
available_memory_bytes
|
|
)
|
|
|
|
memory = self.memory
|
|
if memory is None:
|
|
memory = available_memory_bytes - self.object_store_memory
|
|
if memory < 100e6 and memory < 0.05 * system_memory:
|
|
raise ValueError(
|
|
"After taking into account object store and redis memory "
|
|
"usage, the amount of memory on this node available for "
|
|
"tasks and actors ({} GB) is less than {}% of total. "
|
|
"You can adjust these settings with "
|
|
"ray.init(memory=<bytes>, "
|
|
"object_store_memory=<bytes>).".format(
|
|
round(memory / 1e9, 2), int(100 * (memory / system_memory))
|
|
)
|
|
)
|
|
|
|
self.memory = memory
|
|
|
|
@staticmethod
|
|
def _get_current_node_accelerator(
|
|
num_gpus: Optional[int], resources: Dict[str, float]
|
|
) -> Tuple[AcceleratorManager, int]:
|
|
"""
|
|
Returns the AcceleratorManager and accelerator count for the accelerator
|
|
associated with this node. This assumes each node has at most one accelerator type.
|
|
If no accelerators are present, returns None.
|
|
|
|
The resolved accelerator count uses num_gpus (for GPUs) or resources if set, and
|
|
otherwise falls back to the count auto-detected by the AcceleratorManager. The
|
|
resolved accelerator count is capped by the number of visible accelerators.
|
|
|
|
Args:
|
|
num_gpus: GPU count (if provided by user).
|
|
resources: Resource dictionary containing custom resource keys.
|
|
|
|
Returns:
|
|
Tuple[Optional[AcceleratorManager], int]: A tuple containing the accelerator
|
|
manager (or None) the final resolved accelerator count.
|
|
"""
|
|
for resource_name in accelerators.get_all_accelerator_resource_names():
|
|
accelerator_manager = accelerators.get_accelerator_manager_for_resource(
|
|
resource_name
|
|
)
|
|
if accelerator_manager is None:
|
|
continue
|
|
# Respect configured value for GPUs if set
|
|
if resource_name == "GPU":
|
|
num_accelerators = num_gpus
|
|
else:
|
|
num_accelerators = resources.get(resource_name)
|
|
if num_accelerators is None:
|
|
num_accelerators = (
|
|
accelerator_manager.get_current_node_num_accelerators()
|
|
)
|
|
visible_accelerator_ids = (
|
|
accelerator_manager.get_current_process_visible_accelerator_ids()
|
|
)
|
|
if visible_accelerator_ids is not None:
|
|
num_accelerators = min(
|
|
num_accelerators, len(visible_accelerator_ids)
|
|
)
|
|
|
|
if num_accelerators > 0:
|
|
return accelerator_manager, num_accelerators
|
|
|
|
return None, 0
|