# Copyright 2022 IBM, Red Hat
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
The config sub-module contains the definition of the ClusterConfiguration dataclass,
which is used to specify resource requirements and other details when creating a
Cluster object. It also defines WorkerGroup for multi-worker-group clusters.
"""
import pathlib
import re
import warnings
from dataclasses import dataclass, field, fields
from typing import Dict, Iterable, List, Optional, Union, get_args, get_origin
from kubernetes.client import V1Toleration, V1Volume, V1VolumeMount
dir = pathlib.Path(__file__).parent.parent.resolve()
# https://docs.ray.io/en/latest/ray-core/scheduling/accelerators.html
DEFAULT_RESOURCE_MAPPING = {
"nvidia.com/gpu": "GPU",
"intel.com/gpu": "GPU",
"amd.com/gpu": "GPU",
"aws.amazon.com/neuroncore": "neuron_cores",
"google.com/tpu": "TPU",
"habana.ai/gaudi": "HPU",
"huawei.com/Ascend910": "NPU",
"huawei.com/Ascend310": "NPU",
}
def _mapping_for_observed_resources(resource_names: Iterable[str]) -> Dict[str, str]:
"""Return custom mappings needed to validate observed resources."""
return {
name: name for name in resource_names if name not in DEFAULT_RESOURCE_MAPPING
}
def _gpu_resource_name(
resources: Dict[str, Union[str, int]],
) -> Optional[str]:
"""Return the first known GPU resource present in a resource mapping."""
return next(
(name for name in resources if DEFAULT_RESOURCE_MAPPING.get(name) == "GPU"),
None,
)
# Kubernetes metadata.name must be a lowercase RFC 1123 subdomain
_RFC1123_SUBDOMAIN = re.compile(
r"^[a-z0-9]([-a-z0-9]*[a-z0-9])?(\.[a-z0-9]([-a-z0-9]*[a-z0-9])?)*$"
)
# CPU quantity: whole cores, millicores (500m), or fractional (0.5)
_K8S_CPU_QUANTITY = re.compile(r"^(\d+m|\d+(\.\d+)?)$")
# Memory quantity: integer GB suffix added by SDK, or explicit K8s suffix (Gi, Mi, …)
_K8S_MEMORY_QUANTITY = re.compile(r"^\d+(\.\d+)?(Ki|Mi|Gi|Ti|Pi|Ei|K|M|G|T|P|E)?$")
def _validate_cluster_name(name: str):
"""Raise ValueError if name is not a valid Kubernetes metadata.name (RFC 1123 subdomain)."""
if not name or not _RFC1123_SUBDOMAIN.match(name):
raise ValueError(
"Cluster name must be a valid RFC 1123 subdomain "
"(lowercase, numbers, hyphens/dots; start and end with letter or number)."
)
def _validate_worker_group_name(group_name: str) -> None:
if not group_name or not str(group_name).strip():
raise ValueError(
"WorkerGroup group_name is required and cannot be empty or whitespace"
)
def _validate_cpu_quantity(
group_name: str, field_name: str, value: Union[int, str]
) -> None:
if isinstance(value, bool) or not isinstance(value, (int, str)):
raise ValueError(
f"WorkerGroup '{group_name}': {field_name}={value!r} must be an int or str"
)
if isinstance(value, int):
if value < 0:
raise ValueError(
f"WorkerGroup '{group_name}': {field_name}={value} must be >= 0"
)
return
if not _K8S_CPU_QUANTITY.match(value):
raise ValueError(
f"WorkerGroup '{group_name}': {field_name}={value!r} is not a valid "
"CPU quantity (examples: 1, 500m, 0.5)"
)
def _validate_memory_quantity(
group_name: str, field_name: str, value: Union[int, str]
) -> None:
if isinstance(value, bool) or not isinstance(value, (int, str)):
raise ValueError(
f"WorkerGroup '{group_name}': {field_name}={value!r} must be an int or str"
)
if isinstance(value, int):
if value < 0:
raise ValueError(
f"WorkerGroup '{group_name}': {field_name}={value} must be >= 0"
)
return
if not _K8S_MEMORY_QUANTITY.match(value):
raise ValueError(
f"WorkerGroup '{group_name}': {field_name}={value!r} is not a valid "
"memory quantity (examples: 8, 8G, 8Gi, 512Mi)"
)
[docs]
@dataclass
class WorkerGroup:
"""
Defines an additional worker group with its own resource profile.
Used in ClusterConfiguration.additional_worker_groups to create
heterogeneous clusters with multiple workerGroupSpecs in the
KubeRay RayCluster CR.
Args:
group_name:
Required. Unique name for this worker group. Maps to groupName in the CRD.
replicas:
Number of worker pods in this group.
min_replicas:
Minimum replicas for autoscaling. If None, equals replicas.
max_replicas:
Maximum replicas for autoscaling. If None, equals replicas.
cpu_requests:
CPU requests per worker pod.
cpu_limits:
CPU limits per worker pod.
memory_requests:
Memory requests per worker pod. Int values treated as GB.
memory_limits:
Memory limits per worker pod. Int values treated as GB.
gpu_type:
Extended resource key, e.g. "nvidia.com/gpu". Requires gpu_count.
gpu_count:
Number of GPUs per worker. Requires gpu_type.
extended_resource_requests:
Additional Kubernetes extended resources requested by each worker pod.
GPU resources may also be represented with ``gpu_type`` and
``gpu_count``.
image:
Container image. None inherits from ClusterConfiguration.image.
envs:
Per-group environment variables. Merged over cluster-level (group wins on conflict).
labels:
Pod labels for this worker group. Merged over cluster-level (group wins on conflict).
tolerations:
Pod tolerations. None inherits from ClusterConfiguration.worker_tolerations.
"""
group_name: str
replicas: int = 1
min_replicas: Optional[int] = None
max_replicas: Optional[int] = None
cpu_requests: Union[int, str] = 1
cpu_limits: Union[int, str] = 1
memory_requests: Union[int, str] = 3
memory_limits: Union[int, str] = 6
gpu_type: Optional[str] = None
gpu_count: Optional[int] = None
extended_resource_requests: Dict[str, Union[str, int]] = field(default_factory=dict)
image: Optional[str] = None
envs: Dict[str, str] = field(default_factory=dict)
labels: Dict[str, str] = field(default_factory=dict)
tolerations: Optional[List[V1Toleration]] = None
def __post_init__(self):
_validate_worker_group_name(self.group_name)
if self.replicas < 0:
raise ValueError(
f"WorkerGroup '{self.group_name}': replicas={self.replicas} must be >= 0"
)
if self.min_replicas is not None and self.min_replicas < 0:
raise ValueError(
f"WorkerGroup '{self.group_name}': min_replicas={self.min_replicas} must be >= 0"
)
if self.max_replicas is not None and self.max_replicas < 0:
raise ValueError(
f"WorkerGroup '{self.group_name}': max_replicas={self.max_replicas} must be >= 0"
)
if self.gpu_type and self.gpu_count is None:
raise ValueError(
f"WorkerGroup '{self.group_name}': gpu_count is required when gpu_type={self.gpu_type!r} is set"
)
if self.gpu_count is not None and not self.gpu_type:
raise ValueError(
f"WorkerGroup '{self.group_name}': gpu_type is required when gpu_count={self.gpu_count} is set"
)
if self.gpu_count is not None and self.gpu_count < 0:
raise ValueError(
f"WorkerGroup '{self.group_name}': gpu_count={self.gpu_count} must be >= 0"
)
if self.min_replicas is not None and self.max_replicas is not None:
if self.min_replicas > self.max_replicas:
raise ValueError(
f"WorkerGroup '{self.group_name}': min_replicas={self.min_replicas} "
f"cannot be greater than max_replicas={self.max_replicas}"
)
_validate_cpu_quantity(self.group_name, "cpu_requests", self.cpu_requests)
_validate_cpu_quantity(self.group_name, "cpu_limits", self.cpu_limits)
_validate_memory_quantity(
self.group_name, "memory_requests", self.memory_requests
)
_validate_memory_quantity(self.group_name, "memory_limits", self.memory_limits)
if not isinstance(self.extended_resource_requests, dict):
raise ValueError(
f"WorkerGroup '{self.group_name}': extended_resource_requests must be a dict"
)
for resource_name, resource_value in self.extended_resource_requests.items():
if not isinstance(resource_name, str) or not resource_name:
raise ValueError(
f"WorkerGroup '{self.group_name}': extended resource names must be non-empty strings"
)
if isinstance(resource_value, bool) or not isinstance(
resource_value, (int, str)
):
raise ValueError(
f"WorkerGroup '{self.group_name}': extended resource values must be ints or strings"
)
if isinstance(self.memory_requests, int):
self.memory_requests = f"{self.memory_requests}G"
if isinstance(self.memory_limits, int):
self.memory_limits = f"{self.memory_limits}G"
if isinstance(self.memory_requests, str) and self.memory_requests.isdecimal():
self.memory_requests = f"{self.memory_requests}G"
if isinstance(self.memory_limits, str) and self.memory_limits.isdecimal():
self.memory_limits = f"{self.memory_limits}G"
[docs]
@dataclass
class ClusterConfiguration:
"""
This dataclass is used to specify resource requirements and other details, and
is passed in as an argument when creating a Cluster object.
Args:
name:
The name of the cluster.
namespace:
The namespace in which the cluster should be created.
head_extended_resource_requests:
A dictionary of extended resource requests for the head node. ex: {"nvidia.com/gpu": 1}
head_tolerations:
List of tolerations for head nodes.
num_workers:
The number of workers to create.
worker_tolerations:
List of tolerations for worker nodes.
envs:
A dictionary of environment variables to set for the cluster.
image:
The image to use for the cluster.
image_pull_secrets:
A list of image pull secrets to use for the cluster.
write_to_file:
A boolean indicating whether to write the cluster configuration to a file.
verify_tls:
A boolean indicating whether to verify TLS when connecting to the cluster.
labels:
A dictionary of labels to apply to the cluster.
worker_extended_resource_requests:
A dictionary of extended resource requests for each worker. ex: {"nvidia.com/gpu": 1}
extended_resource_mapping:
A dictionary of custom resource mappings to map extended resource requests to RayCluster resource names
overwrite_default_resource_mapping:
A boolean indicating whether to overwrite the default resource mapping.
annotations:
A dictionary of annotations to apply to the cluster.
volumes:
A list of V1Volume objects to add to the Cluster
volume_mounts:
A list of V1VolumeMount objects to add to the Cluster
enable_autoscaling:
A boolean indicating whether to enable Ray in-tree autoscaling. When True,
min_workers and max_workers must also be set. With Kueue (local_queue or a
default LocalQueue), autoscaling is allowed only when Red Hat Build of Kueue
(Unmanaged) is >= 1.4; RHOAI-managed Kueue still blocks autoscaling.
min_workers:
Minimum number of workers when autoscaling is enabled. Maps to
workerGroupSpecs.replicas and workerGroupSpecs.minReplicas.
Required when enable_autoscaling is True.
max_workers:
Maximum number of workers when autoscaling is enabled. Maps to
workerGroupSpecs.maxReplicas. Required when enable_autoscaling is True.
enable_gcs_ft:
A boolean indicating whether to enable GCS fault tolerance.
enable_usage_stats:
A boolean indicating whether to capture and send Ray usage stats externally.
redis_address:
The address of the Redis server to use for GCS fault tolerance, required when enable_gcs_ft is True.
redis_password_secret:
Kubernetes secret reference containing Redis password. ex: {"name": "secret-name", "key": "password-key"}
external_storage_namespace:
The storage namespace to use for GCS fault tolerance. By default, KubeRay sets it to the UID of RayCluster.
"""
name: Optional[str] = None
namespace: Optional[str] = None
head_cpu_requests: Union[int, str] = 1
head_cpu_limits: Union[int, str] = 2
head_memory_requests: Union[int, str] = 5
head_memory_limits: Union[int, str] = 8
head_extended_resource_requests: Dict[str, Union[str, int]] = field(
default_factory=dict
)
head_tolerations: Optional[List[V1Toleration]] = None
worker_cpu_requests: Union[int, str] = 1
worker_cpu_limits: Union[int, str] = 1
num_workers: int = 1
worker_memory_requests: Union[int, str] = 3
worker_memory_limits: Union[int, str] = 6
worker_tolerations: Optional[List[V1Toleration]] = None
envs: Dict[str, str] = field(default_factory=dict)
image: str = ""
image_pull_secrets: List[str] = field(default_factory=list)
write_to_file: bool = False
verify_tls: bool = True
labels: Dict[str, str] = field(default_factory=dict)
worker_extended_resource_requests: Dict[str, Union[str, int]] = field(
default_factory=dict
)
extended_resource_mapping: Dict[str, str] = field(default_factory=dict)
overwrite_default_resource_mapping: bool = False
local_queue: Optional[str] = None
enable_autoscaling: bool = False
min_workers: Optional[int] = None
max_workers: Optional[int] = None
annotations: Dict[str, str] = field(default_factory=dict)
volumes: list[V1Volume] = field(default_factory=list)
volume_mounts: list[V1VolumeMount] = field(default_factory=list)
additional_worker_groups: List[WorkerGroup] = field(default_factory=list)
enable_gcs_ft: bool = False
enable_usage_stats: bool = False
redis_address: Optional[str] = None
redis_password_secret: Optional[Dict[str, str]] = None
external_storage_namespace: Optional[str] = None
def __post_init__(self):
if not self.verify_tls:
print(
"Warning: TLS verification has been disabled - Endpoint checks will be bypassed"
)
self._validate_autoscaling()
if self.enable_usage_stats:
self.envs["RAY_USAGE_STATS_ENABLED"] = "1"
else:
self.envs["RAY_USAGE_STATS_ENABLED"] = "0"
if self.enable_gcs_ft:
if not self.redis_address:
raise ValueError(
"redis_address must be provided when enable_gcs_ft is True"
)
if self.redis_password_secret and not isinstance(
self.redis_password_secret, dict
):
raise ValueError(
"redis_password_secret must be a dictionary with 'name' and 'key' fields"
)
if self.redis_password_secret and (
"name" not in self.redis_password_secret
or "key" not in self.redis_password_secret
):
raise ValueError(
"redis_password_secret must contain both 'name' and 'key' fields"
)
self._validate_types()
self._memory_to_string()
self._str_mem_no_unit_add_GB()
self._combine_extended_resource_mapping()
self._validate_extended_resource_requests(self.head_extended_resource_requests)
self._validate_extended_resource_requests(
self.worker_extended_resource_requests
)
if self.name is not None:
_validate_cluster_name(self.name)
self._validate_additional_worker_groups()
def _combine_extended_resource_mapping(self):
if overwritten := set(self.extended_resource_mapping.keys()).intersection(
DEFAULT_RESOURCE_MAPPING.keys()
):
if self.overwrite_default_resource_mapping:
warnings.warn(
f"Overwriting default resource mapping for {overwritten}",
UserWarning,
)
else:
raise ValueError(
f"Resource mapping already exists for {overwritten}, set overwrite_default_resource_mapping to True to overwrite"
)
self.extended_resource_mapping = {
**DEFAULT_RESOURCE_MAPPING,
**self.extended_resource_mapping,
}
def _validate_extended_resource_requests(self, extended_resources: Dict[str, int]):
for k in extended_resources.keys():
if k not in self.extended_resource_mapping.keys():
raise ValueError(
f"extended resource '{k}' not found in extended_resource_mapping, available resources are {list(self.extended_resource_mapping.keys())}, to add more supported resources use extended_resource_mapping. i.e. extended_resource_mapping = {{'{k}': 'FOO_BAR'}}"
)
def _validate_additional_worker_groups(self):
if not self.additional_worker_groups:
return
default_group_name = f"small-group-{self.name}" if self.name else None
names = set()
for wg in self.additional_worker_groups:
if not isinstance(wg, WorkerGroup):
raise TypeError(
f"additional_worker_groups entries must be WorkerGroup instances, got {type(wg)}"
)
if wg.group_name in names:
raise ValueError(
f"Duplicate worker group name: '{wg.group_name}' "
f"(each additional_worker_groups entry must have a unique group_name)"
)
if default_group_name and wg.group_name == default_group_name:
raise ValueError(
f"Worker group name '{wg.group_name}' conflicts with the default worker group name"
)
names.add(wg.group_name)
def _validate_autoscaling(self):
if self.enable_autoscaling:
if self.min_workers is None or self.max_workers is None:
raise ValueError(
"min_workers and max_workers must be provided when enable_autoscaling is True"
)
if self.min_workers < 0:
raise ValueError("min_workers must be >= 0")
if self.max_workers < self.min_workers:
raise ValueError("max_workers must be >= min_workers")
else:
if self.min_workers is not None or self.max_workers is not None:
warnings.warn(
"min_workers and max_workers are ignored when enable_autoscaling is False",
UserWarning,
)
def _str_mem_no_unit_add_GB(self):
if (
isinstance(self.worker_memory_requests, str)
and self.worker_memory_requests.isdecimal()
):
self.worker_memory_requests = f"{self.worker_memory_requests}G"
if (
isinstance(self.worker_memory_limits, str)
and self.worker_memory_limits.isdecimal()
):
self.worker_memory_limits = f"{self.worker_memory_limits}G"
def _memory_to_string(self):
if isinstance(self.head_memory_requests, int):
self.head_memory_requests = f"{self.head_memory_requests}G"
if isinstance(self.head_memory_limits, int):
self.head_memory_limits = f"{self.head_memory_limits}G"
if isinstance(self.worker_memory_requests, int):
self.worker_memory_requests = f"{self.worker_memory_requests}G"
if isinstance(self.worker_memory_limits, int):
self.worker_memory_limits = f"{self.worker_memory_limits}G"
def _validate_types(self):
"""Validate the types of all fields in the ClusterConfiguration dataclass."""
errors = []
for field_info in fields(self):
value = getattr(self, field_info.name)
expected_type = field_info.type
if not self._is_type(value, expected_type):
errors.append(f"'{field_info.name}' should be of type {expected_type}.")
if errors:
raise TypeError("Type validation failed:\n" + "\n".join(errors))
@staticmethod
def _is_type(value, expected_type):
"""Check if the value matches the expected type."""
def check_type(value, expected_type):
origin_type = get_origin(expected_type)
args = get_args(expected_type)
if origin_type is Union:
return any(check_type(value, union_type) for union_type in args)
if origin_type is list:
if value is not None:
return all(check_type(elem, args[0]) for elem in (value or []))
else:
return True
if origin_type is dict:
if value is not None:
return all(
check_type(k, args[0]) and check_type(v, args[1])
for k, v in value.items()
)
else:
return True
if origin_type is tuple:
return all(check_type(elem, etype) for elem, etype in zip(value, args))
if expected_type is int:
return isinstance(value, int) and not isinstance(value, bool)
if expected_type is bool:
return isinstance(value, bool)
return isinstance(value, expected_type)
return check_type(value, expected_type)