wehub-resource-sync
562 lines
21 KiB
Python
562 lines
21 KiB
Python
import copy
|
|
import logging
|
|
from abc import ABC, abstractmethod
|
|
from dataclasses import dataclass, field
|
|
from enum import Enum
|
|
from pathlib import Path
|
|
from typing import Any, Dict, List, Optional
|
|
|
|
import yaml
|
|
|
|
from ray._common.utils import binary_to_hex
|
|
from ray._private.ray_constants import env_integer
|
|
from ray._raylet import GcsClient
|
|
from ray.autoscaler._private.constants import (
|
|
AUTOSCALER_MAX_CONCURRENT_LAUNCHES,
|
|
DEFAULT_UPSCALING_SPEED,
|
|
DISABLE_LAUNCH_CONFIG_CHECK_KEY,
|
|
DISABLE_NODE_UPDATERS_KEY,
|
|
)
|
|
from ray.autoscaler._private.kuberay.autoscaling_config import AutoscalingConfigProducer
|
|
from ray.autoscaler._private.monitor import BASE_READONLY_CONFIG
|
|
from ray.autoscaler._private.util import (
|
|
format_readonly_node_type,
|
|
hash_launch_conf,
|
|
hash_runtime_conf,
|
|
prepare_config,
|
|
validate_config,
|
|
)
|
|
from ray.autoscaler.v2.schema import NodeType
|
|
from ray.autoscaler.v2.sdk import get_cluster_resource_state
|
|
from ray.autoscaler.v2.utils import is_head_node
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class Provider(Enum):
|
|
UNKNOWN = 0
|
|
ALIYUN = 1
|
|
AWS = 2
|
|
AZURE = 3
|
|
GCP = 4
|
|
KUBERAY = 5
|
|
LOCAL = 6
|
|
READ_ONLY = 7
|
|
|
|
|
|
class IConfigReader(ABC):
|
|
"""An interface for reading Autoscaling config.
|
|
|
|
A utility class that reads autoscaling configs from various sources:
|
|
- File
|
|
- In-memory dict
|
|
- Remote config service (e.g. KubeRay's config)
|
|
|
|
Example:
|
|
reader = FileConfigReader("path/to/config.yaml")
|
|
# Get the recently cached config.
|
|
config = reader.get_cached_autoscaling_config()
|
|
|
|
...
|
|
# Refresh the cached config.
|
|
reader.refresh_cached_autoscaling_config()
|
|
config = reader.get_cached_autoscaling_config()
|
|
|
|
"""
|
|
|
|
@abstractmethod
|
|
def get_cached_autoscaling_config(self) -> "AutoscalingConfig":
|
|
"""Returns the recently read autoscaling config.
|
|
|
|
Returns:
|
|
AutoscalingConfig: The recently read autoscaling config.
|
|
"""
|
|
pass
|
|
|
|
@abstractmethod
|
|
def refresh_cached_autoscaling_config(self):
|
|
"""Read the config from the source."""
|
|
pass
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class InstanceReconcileConfig:
|
|
# The timeout for waiting for a REQUESTED instance to be ALLOCATED.
|
|
request_status_timeout_s: int = env_integer(
|
|
"RAY_AUTOSCALER_RECONCILE_REQUEST_STATUS_TIMEOUT_S", 10 * 60
|
|
)
|
|
# The timeout for waiting for a ALLOCATED instance to be RAY_RUNNING.
|
|
allocate_status_timeout_s: int = env_integer(
|
|
"RAY_AUTOSCALER_RECONCILE_ALLOCATE_STATUS_TIMEOUT_S", 60 * 60
|
|
)
|
|
# The timeout for waiting for a RAY_INSTALLING instance to be RAY_RUNNING.
|
|
ray_install_status_timeout_s: int = env_integer(
|
|
"RAY_AUTOSCALER_RECONCILE_RAY_INSTALL_STATUS_TIMEOUT_S", 30 * 60
|
|
)
|
|
# The timeout for waiting for a TERMINATING instance to be TERMINATED.
|
|
terminating_status_timeout_s: int = env_integer(
|
|
"RAY_AUTOSCALER_RECONCILE_TERMINATING_STATUS_TIMEOUT_S", 300
|
|
)
|
|
# The timeout for waiting for a RAY_STOP_REQUESTED instance
|
|
# to be RAY_STOPPING or RAY_STOPPED.
|
|
ray_stop_requested_status_timeout_s: int = env_integer(
|
|
"RAY_AUTOSCALER_RECONCILE_RAY_STOP_REQUESTED_STATUS_TIMEOUT_S", 300
|
|
)
|
|
# The interval for raise a warning when an instance in transient status
|
|
# is not updated for a long time.
|
|
transient_status_warn_interval_s: int = env_integer(
|
|
"RAY_AUTOSCALER_RECONCILE_TRANSIENT_STATUS_WARN_INTERVAL_S", 90
|
|
)
|
|
# The number of times to retry requesting to allocate an instance.
|
|
max_num_retry_request_to_allocate: int = env_integer(
|
|
"RAY_AUTOSCALER_RECONCILE_MAX_NUM_RETRY_REQUEST_TO_ALLOCATE", 3
|
|
)
|
|
|
|
|
|
@dataclass
|
|
class NodeTypeConfig:
|
|
"""
|
|
NodeTypeConfig is the helper class to provide node type specific configs.
|
|
This maps to subset of the `available_node_types` field in the
|
|
autoscaling config.
|
|
"""
|
|
|
|
# Node type name
|
|
name: NodeType
|
|
# The minimal number of worker nodes to be launched for this node type.
|
|
min_worker_nodes: int
|
|
# The maximal number of worker nodes can be launched for this node type.
|
|
max_worker_nodes: int
|
|
# Idle timeout seconds for worker nodes of this node type.
|
|
idle_timeout_s: Optional[float] = None
|
|
# The priority of the worker group. Higher value means the group will be scaled up first if everything else is equal.
|
|
priority: int = 0
|
|
# The total resources on the node.
|
|
resources: Dict[str, float] = field(default_factory=dict)
|
|
# The labels on the node.
|
|
labels: Dict[str, str] = field(default_factory=dict)
|
|
# The node config's launch config hash. It's calculated from the auth
|
|
# config, and the node's config in the `AutoscalingConfig` for the node
|
|
# type when launching the node. It's used to detect config changes.
|
|
launch_config_hash: str = ""
|
|
|
|
def __post_init__(self):
|
|
assert self.min_worker_nodes <= self.max_worker_nodes
|
|
assert self.min_worker_nodes >= 0
|
|
|
|
|
|
class AutoscalingConfig:
|
|
"""
|
|
AutoscalingConfig is the helper class to provide autoscaling
|
|
related configs.
|
|
|
|
# TODO(rickyx):
|
|
1. Move the config validation logic here.
|
|
2. Deprecate the ray-schema.json for validation because it's
|
|
static thus not possible to validate the config with interdependency
|
|
of each other.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
configs: Dict[str, Any],
|
|
skip_content_hash: bool = False,
|
|
) -> None:
|
|
"""Initialize the autoscaling config.
|
|
|
|
Args:
|
|
configs: The raw configs dict.
|
|
skip_content_hash: Whether to skip file mounts/ray command hash
|
|
calculation.
|
|
"""
|
|
self._sync_continuously = False
|
|
self.update_configs(configs, skip_content_hash)
|
|
|
|
def update_configs(self, configs: Dict[str, Any], skip_content_hash: bool) -> None:
|
|
self._configs = prepare_config(configs)
|
|
validate_config(self._configs)
|
|
if skip_content_hash:
|
|
return
|
|
self._calculate_hashes()
|
|
self._sync_continuously = self._configs.get(
|
|
"generate_file_mounts_contents_hash", True
|
|
)
|
|
|
|
def _calculate_hashes(self) -> None:
|
|
logger.info("Calculating hashes for file mounts and ray commands.")
|
|
self._runtime_hash, self._file_mounts_contents_hash = hash_runtime_conf(
|
|
self._configs.get("file_mounts", {}),
|
|
self._configs.get("cluster_synced_files", []),
|
|
[
|
|
self._configs.get("worker_setup_commands", []),
|
|
self._configs.get("worker_start_ray_commands", []),
|
|
],
|
|
generate_file_mounts_contents_hash=self._configs.get(
|
|
"generate_file_mounts_contents_hash", True
|
|
),
|
|
)
|
|
|
|
def get_cloud_node_config(self, ray_node_type: NodeType) -> Dict[str, Any]:
|
|
return copy.deepcopy(
|
|
self.get_node_type_specific_config(ray_node_type, "node_config") or {}
|
|
)
|
|
|
|
def get_docker_config(self, ray_node_type: NodeType) -> Dict[str, Any]:
|
|
"""
|
|
Return the docker config for the specified node type.
|
|
If it's a head node, the image will be chosen in the following order:
|
|
1. Node specific docker image.
|
|
2. The 'docker' config's 'head_image' field.
|
|
3. The 'docker' config's 'image' field.
|
|
If it's a worker node, the image will be chosen in the following order:
|
|
1. Node specific docker image.
|
|
2. The 'docker' config's 'worker_image' field.
|
|
3. The 'docker' config's 'image' field.
|
|
"""
|
|
# TODO(rickyx): It's unfortunate we have multiple fields in ray-schema.json
|
|
# that can specify docker images. We should consolidate them.
|
|
docker_config = copy.deepcopy(self._configs.get("docker", {}))
|
|
node_specific_docker_config = self._configs["available_node_types"][
|
|
ray_node_type
|
|
].get("docker", {})
|
|
# Override the global docker config with node specific docker config.
|
|
docker_config.update(node_specific_docker_config)
|
|
|
|
if self._configs.get("head_node_type") == ray_node_type:
|
|
if "head_image" in docker_config:
|
|
logger.info(
|
|
"Overwriting image={} by head_image({}) for head node docker.".format( # noqa: E501
|
|
docker_config["image"], docker_config["head_image"]
|
|
)
|
|
)
|
|
docker_config["image"] = docker_config["head_image"]
|
|
else:
|
|
if "worker_image" in docker_config:
|
|
logger.info(
|
|
"Overwriting image={} by worker_image({}) for worker node docker.".format( # noqa: E501
|
|
docker_config["image"], docker_config["worker_image"]
|
|
)
|
|
)
|
|
docker_config["image"] = docker_config["worker_image"]
|
|
|
|
# These fields should be merged.
|
|
docker_config.pop("head_image", None)
|
|
docker_config.pop("worker_image", None)
|
|
return docker_config
|
|
|
|
def get_worker_start_ray_commands(self) -> List[str]:
|
|
return self._configs.get("worker_start_ray_commands", [])
|
|
|
|
def get_head_setup_commands(self) -> List[str]:
|
|
return self._configs.get("head_setup_commands", [])
|
|
|
|
def get_head_start_ray_commands(self) -> List[str]:
|
|
return self._configs.get("head_start_ray_commands", [])
|
|
|
|
def get_worker_setup_commands(self, ray_node_type: NodeType) -> List[str]:
|
|
"""
|
|
Return the worker setup commands for the specified node type.
|
|
|
|
If the node type specific worker setup commands are not specified,
|
|
return the global worker setup commands.
|
|
"""
|
|
worker_setup_command = self.get_node_type_specific_config(
|
|
ray_node_type, "worker_setup_commands"
|
|
)
|
|
if worker_setup_command is None:
|
|
# Return global worker setup commands if node type specific
|
|
# worker setup commands are not specified.
|
|
logger.info(
|
|
"Using global worker setup commands for {}".format(ray_node_type)
|
|
)
|
|
return self._configs.get("worker_setup_commands", [])
|
|
return worker_setup_command
|
|
|
|
def get_initialization_commands(self, ray_node_type: NodeType) -> List[str]:
|
|
"""
|
|
Return the initialization commands for the specified node type.
|
|
|
|
If the node type specific initialization commands are not specified,
|
|
return the global initialization commands.
|
|
"""
|
|
initialization_command = self.get_node_type_specific_config(
|
|
ray_node_type, "initialization_commands"
|
|
)
|
|
if initialization_command is None:
|
|
logger.info(
|
|
"Using global initialization commands for {}".format(ray_node_type)
|
|
)
|
|
return self._configs.get("initialization_commands", [])
|
|
return initialization_command
|
|
|
|
def get_node_type_specific_config(
|
|
self, ray_node_type: NodeType, config_name: str
|
|
) -> Optional[Any]:
|
|
node_specific_config = self._configs["available_node_types"].get(
|
|
ray_node_type, {}
|
|
)
|
|
return node_specific_config.get(config_name, None)
|
|
|
|
def get_node_resources(self, ray_node_type: NodeType) -> Dict[str, float]:
|
|
return copy.deepcopy(
|
|
self.get_node_type_specific_config(ray_node_type, "resources") or {}
|
|
)
|
|
|
|
def get_node_labels(self, ray_node_type: NodeType) -> Dict[str, str]:
|
|
return copy.deepcopy(
|
|
self.get_node_type_specific_config(ray_node_type, "labels") or {}
|
|
)
|
|
|
|
def get_config(self, config_name, default=None) -> Any:
|
|
return self._configs.get(config_name, default)
|
|
|
|
def get_provider_instance_type(self, ray_node_type: NodeType) -> str:
|
|
provider = self.provider
|
|
node_config = (
|
|
self.get_node_type_specific_config(ray_node_type, "node_config") or {}
|
|
)
|
|
if provider in [Provider.AWS, Provider.ALIYUN]:
|
|
return node_config.get("InstanceType", "")
|
|
elif provider == Provider.AZURE:
|
|
return node_config.get("azure_arm_parameters", {}).get("vmSize", "")
|
|
elif provider == Provider.GCP:
|
|
return node_config.get("machineType", "")
|
|
elif provider in [Provider.KUBERAY, Provider.LOCAL, Provider.UNKNOWN]:
|
|
return ""
|
|
else:
|
|
raise ValueError(f"Unknown provider {provider}")
|
|
|
|
def get_node_type_configs(self) -> Dict[NodeType, NodeTypeConfig]:
|
|
"""
|
|
Returns the node type configs from the `available_node_types` field.
|
|
|
|
Returns:
|
|
Dict[NodeType, NodeTypeConfig]: The node type configs.
|
|
"""
|
|
available_node_types = self._configs.get("available_node_types", {})
|
|
if not available_node_types:
|
|
return None
|
|
node_type_configs = {}
|
|
auth_config = self._configs.get("auth", {})
|
|
head_node_type = self.get_head_node_type()
|
|
assert head_node_type
|
|
for node_type, node_config in available_node_types.items():
|
|
launch_config_hash = hash_launch_conf(
|
|
node_config.get("node_config", {}), auth_config
|
|
)
|
|
max_workers_nodes = node_config.get("max_workers", 0)
|
|
if head_node_type == node_type:
|
|
max_workers_nodes += 1
|
|
|
|
node_type_configs[node_type] = NodeTypeConfig(
|
|
name=node_type,
|
|
min_worker_nodes=node_config.get("min_workers", 0),
|
|
max_worker_nodes=max_workers_nodes,
|
|
idle_timeout_s=node_config.get("idle_timeout_s", None),
|
|
priority=node_config.get("priority", 0),
|
|
resources=node_config.get("resources", {}),
|
|
labels=node_config.get("labels", {}),
|
|
launch_config_hash=launch_config_hash,
|
|
)
|
|
return node_type_configs
|
|
|
|
def get_head_node_type(self) -> NodeType:
|
|
"""
|
|
Returns the head node type.
|
|
|
|
If there is only one node type, return the only node type as the head
|
|
node type.
|
|
If there are multiple node types, return the head node type specified
|
|
in the config.
|
|
"""
|
|
available_node_types = self._configs.get("available_node_types", {})
|
|
if len(available_node_types) == 1:
|
|
return list(available_node_types.keys())[0]
|
|
return self._configs.get("head_node_type")
|
|
|
|
def get_max_num_worker_nodes(self) -> Optional[int]:
|
|
return self.get_config("max_workers", None)
|
|
|
|
def get_max_num_nodes(self) -> Optional[int]:
|
|
max_num_workers = self.get_max_num_worker_nodes()
|
|
if max_num_workers is not None:
|
|
return max_num_workers + 1 # For head node
|
|
return None
|
|
|
|
def get_raw_config_mutable(self) -> Dict[str, Any]:
|
|
return self._configs
|
|
|
|
def get_upscaling_speed(self) -> float:
|
|
return self.get_config("upscaling_speed", DEFAULT_UPSCALING_SPEED)
|
|
|
|
def get_max_concurrent_launches(self) -> int:
|
|
return AUTOSCALER_MAX_CONCURRENT_LAUNCHES
|
|
|
|
def disable_node_updaters(self) -> bool:
|
|
provider_config = self._configs.get("provider", {})
|
|
return provider_config.get(DISABLE_NODE_UPDATERS_KEY, False)
|
|
|
|
def get_idle_timeout_s(self) -> Optional[float]:
|
|
"""
|
|
Returns the idle timeout in seconds if present in config, otherwise None.
|
|
"""
|
|
idle_timeout_s = self.get_config("idle_timeout_minutes", None)
|
|
return idle_timeout_s * 60 if idle_timeout_s is not None else None
|
|
|
|
def disable_launch_config_check(self) -> bool:
|
|
provider_config = self.get_provider_config()
|
|
return provider_config.get(DISABLE_LAUNCH_CONFIG_CHECK_KEY, True)
|
|
|
|
def get_instance_reconcile_config(self) -> InstanceReconcileConfig:
|
|
# TODO(rickyx): we need a way to customize these configs,
|
|
# either extending the current ray-schema.json, or just use another
|
|
# schema validation paths.
|
|
return InstanceReconcileConfig()
|
|
|
|
def get_provider_config(self) -> Dict[str, Any]:
|
|
return self._configs.get("provider", {})
|
|
|
|
def dump(self) -> str:
|
|
return yaml.safe_dump(self._configs)
|
|
|
|
@property
|
|
def provider(self) -> Provider:
|
|
provider_str = self._configs.get("provider", {}).get("type", "")
|
|
if provider_str == "local":
|
|
return Provider.LOCAL
|
|
elif provider_str == "aws":
|
|
return Provider.AWS
|
|
elif provider_str == "azure":
|
|
return Provider.AZURE
|
|
elif provider_str == "gcp":
|
|
return Provider.GCP
|
|
elif provider_str == "aliyun":
|
|
return Provider.ALIYUN
|
|
elif provider_str == "kuberay":
|
|
return Provider.KUBERAY
|
|
elif provider_str == "readonly":
|
|
return Provider.READ_ONLY
|
|
else:
|
|
return Provider.UNKNOWN
|
|
|
|
@property
|
|
def runtime_hash(self) -> str:
|
|
if not hasattr(self, "_runtime_hash"):
|
|
self._calculate_hashes()
|
|
return self._runtime_hash
|
|
|
|
@property
|
|
def file_mounts_contents_hash(self) -> str:
|
|
if not hasattr(self, "_file_mounts_contents_hash"):
|
|
self._calculate_hashes()
|
|
return self._file_mounts_contents_hash
|
|
|
|
|
|
class FileConfigReader(IConfigReader):
|
|
"""A class that reads cluster config from a yaml file."""
|
|
|
|
def __init__(self, config_file: str, skip_content_hash: bool = True) -> None:
|
|
"""Initialize the file config reader.
|
|
|
|
Args:
|
|
config_file: The path to the config file.
|
|
skip_content_hash: Whether to skip file mounts/ray command
|
|
hash calculation. Default to True.
|
|
"""
|
|
self._config_file_path = Path(config_file).resolve()
|
|
self._skip_content_hash = skip_content_hash
|
|
self._cached_config = self._read()
|
|
|
|
def _read(self) -> AutoscalingConfig:
|
|
with open(self._config_file_path) as f:
|
|
config = yaml.safe_load(f.read())
|
|
return AutoscalingConfig(config, skip_content_hash=self._skip_content_hash)
|
|
|
|
def get_cached_autoscaling_config(self) -> AutoscalingConfig:
|
|
"""
|
|
Returns:
|
|
AutoscalingConfig: The autoscaling config.
|
|
"""
|
|
|
|
return self._cached_config
|
|
|
|
def refresh_cached_autoscaling_config(self):
|
|
self._cached_config = self._read()
|
|
|
|
|
|
class KubeRayConfigReader(IConfigReader):
|
|
"""A class that reads cluster config from a K8s RayCluster CR."""
|
|
|
|
def __init__(self, config_producer: AutoscalingConfigProducer):
|
|
self._config_producer = config_producer
|
|
self._cached_config = self._generate_configs_from_k8s()
|
|
|
|
def _generate_configs_from_k8s(self) -> AutoscalingConfig:
|
|
return AutoscalingConfig(self._config_producer())
|
|
|
|
def get_cached_autoscaling_config(self) -> AutoscalingConfig:
|
|
"""
|
|
Returns:
|
|
AutoscalingConfig: The autoscaling config.
|
|
"""
|
|
return self._cached_config
|
|
|
|
def refresh_cached_autoscaling_config(self):
|
|
"""
|
|
Reads the configs from the K8s RayCluster CR.
|
|
|
|
This reads from the K8s API server every time to pick up changes.
|
|
"""
|
|
self._cached_config = self._generate_configs_from_k8s()
|
|
|
|
|
|
class ReadOnlyProviderConfigReader(IConfigReader):
|
|
"""A class that reads cluster config for a read-only provider.
|
|
|
|
This is used for laptop mode / manual cluster setup modes, in order to
|
|
provide status reporting in the same way for users."""
|
|
|
|
def __init__(self, gcs_address: str):
|
|
self._configs = BASE_READONLY_CONFIG
|
|
self._gcs_client = GcsClient(address=gcs_address)
|
|
|
|
def refresh_cached_autoscaling_config(self) -> AutoscalingConfig:
|
|
# Update the config with node types from GCS.
|
|
ray_cluster_resource_state = get_cluster_resource_state(self._gcs_client)
|
|
|
|
# Format each node type's config from the running nodes.
|
|
available_node_types = {}
|
|
|
|
head_node_type = None
|
|
for node_state in ray_cluster_resource_state.node_states:
|
|
node_type = node_state.ray_node_type_name
|
|
if not node_type:
|
|
node_type = format_readonly_node_type(binary_to_hex(node_state.node_id))
|
|
|
|
if is_head_node(node_state):
|
|
head_node_type = node_type
|
|
|
|
if node_type not in available_node_types:
|
|
available_node_types[node_type] = {
|
|
"resources": dict(node_state.total_resources),
|
|
"min_workers": 0,
|
|
"max_workers": 0 if is_head_node(node_state) else 1,
|
|
"node_config": {},
|
|
}
|
|
elif not is_head_node(node_state):
|
|
available_node_types[node_type]["max_workers"] += 1
|
|
|
|
if available_node_types:
|
|
self._configs["available_node_types"].update(available_node_types)
|
|
self._configs["max_workers"] = sum(
|
|
v["max_workers"] for v in available_node_types.values()
|
|
)
|
|
assert head_node_type, "Head node type should be found."
|
|
self._configs["head_node_type"] = head_node_type
|
|
|
|
# Don't idle terminated nodes in read-only mode.
|
|
self._configs.pop("idle_timeout_minutes", None)
|
|
|
|
def get_cached_autoscaling_config(self) -> AutoscalingConfig:
|
|
return AutoscalingConfig(self._configs, skip_content_hash=True)
|