chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,304 @@
|
||||
import logging
|
||||
from types import ModuleType
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
from ray.autoscaler._private.command_runner import DockerCommandRunner, SSHCommandRunner
|
||||
from ray.autoscaler.command_runner import CommandRunnerInterface
|
||||
from ray.util.annotations import DeveloperAPI
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@DeveloperAPI
|
||||
class NodeProvider:
|
||||
"""Interface for getting and returning nodes from a Cloud.
|
||||
|
||||
**Important**: This is an INTERNAL API that is only exposed for the purpose
|
||||
of implementing custom node providers. It is not allowed to call into
|
||||
NodeProvider methods from any Ray package outside the autoscaler, only to
|
||||
define new implementations of NodeProvider for use with the "external" node
|
||||
provider option.
|
||||
|
||||
NodeProviders are namespaced by the `cluster_name` parameter; they only
|
||||
operate on nodes within that namespace.
|
||||
|
||||
Nodes may be in one of three states: {pending, running, terminated}. Nodes
|
||||
appear immediately once started by `create_node`, and transition
|
||||
immediately to terminated when `terminate_node` is called.
|
||||
|
||||
Threading and concurrency:
|
||||
- The autoscaler calls the following methods from multiple threads
|
||||
(NodeLauncher, NodeUpdaterThread, autoscaler main loop, and
|
||||
NodeProviderAdapter executors).
|
||||
- These methods MUST be thread-safe:
|
||||
non_terminated_nodes, is_running, is_terminated, node_tags, internal_ip,
|
||||
external_ip, get_node_id, create_node/create_node_with_resources_and_labels,
|
||||
set_node_tags, terminate_node/terminate_nodes.
|
||||
|
||||
TODO (rueian): make sure all the existing implementations are thread-safe.
|
||||
"""
|
||||
|
||||
def __init__(self, provider_config: Dict[str, Any], cluster_name: str) -> None:
|
||||
self.provider_config = provider_config
|
||||
self.cluster_name = cluster_name
|
||||
self._internal_ip_cache: Dict[str, str] = {}
|
||||
self._external_ip_cache: Dict[str, str] = {}
|
||||
|
||||
def is_readonly(self) -> bool:
|
||||
"""Returns whether this provider is readonly.
|
||||
|
||||
Readonly node providers do not allow nodes to be created or terminated.
|
||||
"""
|
||||
return False
|
||||
|
||||
def non_terminated_nodes(self, tag_filters: Dict[str, str]) -> List[str]:
|
||||
"""Return a list of node ids filtered by the specified tags dict.
|
||||
|
||||
This list must not include terminated nodes. For performance reasons,
|
||||
providers are allowed to cache the result of a call to
|
||||
non_terminated_nodes() to serve single-node queries
|
||||
(e.g. is_running(node_id)). This means that non_terminate_nodes() must
|
||||
be called again to refresh results.
|
||||
|
||||
Args:
|
||||
tag_filters: Tag key/value pairs that nodes must match to be
|
||||
included in the result.
|
||||
|
||||
Returns:
|
||||
A list of node ids matching the given tag filters.
|
||||
|
||||
Examples:
|
||||
>>> from ray.autoscaler.node_provider import NodeProvider
|
||||
>>> from ray.autoscaler.tags import TAG_RAY_NODE_KIND
|
||||
>>> provider = NodeProvider(...) # doctest: +SKIP
|
||||
>>> provider.non_terminated_nodes( # doctest: +SKIP
|
||||
... {TAG_RAY_NODE_KIND: "worker"})
|
||||
["node-1", "node-2"]
|
||||
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
def nodes_for_teardown(self, tag_filters: Dict[str, str]) -> List[str]:
|
||||
"""Return all node ids matching tag_filters, including terminated nodes.
|
||||
|
||||
Used during teardown to ensure cleanup of external resources (e.g.
|
||||
Docker containers) on nodes whose state may not be accurately tracked
|
||||
by this provider instance. For example, LocalNodeProvider on the
|
||||
machine invoking ``ray down`` may show workers as terminated even
|
||||
though the head node's autoscaler started them and their Docker
|
||||
containers are still running.
|
||||
|
||||
The default delegates to non_terminated_nodes(), which is correct for
|
||||
cloud providers that always query live infrastructure state. Providers
|
||||
that maintain state locally should override this to include all known
|
||||
nodes regardless of recorded state.
|
||||
"""
|
||||
return self.non_terminated_nodes(tag_filters)
|
||||
|
||||
def is_running(self, node_id: str) -> bool:
|
||||
"""Return whether the specified node is running."""
|
||||
raise NotImplementedError
|
||||
|
||||
def is_terminated(self, node_id: str) -> bool:
|
||||
"""Return whether the specified node is terminated."""
|
||||
raise NotImplementedError
|
||||
|
||||
def node_tags(self, node_id: str) -> Dict[str, str]:
|
||||
"""Returns the tags of the given node (string dict)."""
|
||||
raise NotImplementedError
|
||||
|
||||
def external_ip(self, node_id: str) -> str:
|
||||
"""Returns the external ip of the given node."""
|
||||
raise NotImplementedError
|
||||
|
||||
def internal_ip(self, node_id: str) -> str:
|
||||
"""Returns the internal ip (Ray ip) of the given node."""
|
||||
raise NotImplementedError
|
||||
|
||||
def get_node_id(self, ip_address: str, use_internal_ip: bool = False) -> str:
|
||||
"""Returns the node_id given an IP address.
|
||||
|
||||
Assumes ip-address is unique per node.
|
||||
|
||||
Args:
|
||||
ip_address: Address of node.
|
||||
use_internal_ip: Whether the ip address is
|
||||
public or private.
|
||||
|
||||
Returns:
|
||||
The node id corresponding to the given IP address.
|
||||
|
||||
Raises:
|
||||
ValueError: If not found.
|
||||
"""
|
||||
|
||||
def find_node_id():
|
||||
if use_internal_ip:
|
||||
return self._internal_ip_cache.get(ip_address)
|
||||
else:
|
||||
return self._external_ip_cache.get(ip_address)
|
||||
|
||||
if not find_node_id():
|
||||
all_nodes = self.non_terminated_nodes({})
|
||||
ip_func = self.internal_ip if use_internal_ip else self.external_ip
|
||||
ip_cache = (
|
||||
self._internal_ip_cache if use_internal_ip else self._external_ip_cache
|
||||
)
|
||||
for node_id in all_nodes:
|
||||
ip_cache[ip_func(node_id)] = node_id
|
||||
|
||||
if not find_node_id():
|
||||
if use_internal_ip:
|
||||
known_msg = f"Worker internal IPs: {list(self._internal_ip_cache)}"
|
||||
else:
|
||||
known_msg = f"Worker external IP: {list(self._external_ip_cache)}"
|
||||
raise ValueError(f"ip {ip_address} not found. " + known_msg)
|
||||
|
||||
return find_node_id()
|
||||
|
||||
def create_node(
|
||||
self, node_config: Dict[str, Any], tags: Dict[str, str], count: int
|
||||
) -> Optional[Dict[str, Any]]:
|
||||
"""Creates a number of nodes within the namespace.
|
||||
|
||||
Optionally returns a mapping from created node ids to node metadata.
|
||||
|
||||
Optionally may throw a
|
||||
ray.autoscaler.node_launch_exception.NodeLaunchException which the
|
||||
autoscaler may use to provide additional functionality such as
|
||||
observability.
|
||||
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
def create_node_with_resources_and_labels(
|
||||
self,
|
||||
node_config: Dict[str, Any],
|
||||
tags: Dict[str, str],
|
||||
count: int,
|
||||
resources: Dict[str, float],
|
||||
labels: Dict[str, str],
|
||||
) -> Optional[Dict[str, Any]]:
|
||||
"""Create nodes with a given resource and label config.
|
||||
|
||||
This is the method actually called by the autoscaler. Prefer to
|
||||
implement this when possible directly, otherwise it delegates to the
|
||||
create_node() implementation.
|
||||
|
||||
Optionally may throw a ray.autoscaler.node_launch_exception.NodeLaunchException.
|
||||
"""
|
||||
return self.create_node(node_config, tags, count)
|
||||
|
||||
def set_node_tags(self, node_id: str, tags: Dict[str, str]) -> None:
|
||||
"""Sets the tag values (string dict) for the specified node."""
|
||||
raise NotImplementedError
|
||||
|
||||
def terminate_node(self, node_id: str) -> Optional[Dict[str, Any]]:
|
||||
"""Terminates the specified node.
|
||||
|
||||
Optionally return a mapping from deleted node ids to node
|
||||
metadata.
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
def terminate_nodes(self, node_ids: List[str]) -> Optional[Dict[str, Any]]:
|
||||
"""Terminates a set of nodes.
|
||||
|
||||
May be overridden with a batch method, which optionally may return a
|
||||
mapping from deleted node ids to node metadata.
|
||||
"""
|
||||
for node_id in node_ids:
|
||||
logger.info("NodeProvider: {}: Terminating node".format(node_id))
|
||||
self.terminate_node(node_id)
|
||||
return None
|
||||
|
||||
@property
|
||||
def max_terminate_nodes(self) -> Optional[int]:
|
||||
"""The maximum number of nodes which can be terminated in one single
|
||||
API request. By default, this is "None", which means that the node
|
||||
provider's underlying API allows infinite requests to be terminated
|
||||
with one request.
|
||||
|
||||
For example, AWS only allows 1000 nodes to be terminated
|
||||
at once; to terminate more, we must issue multiple separate API
|
||||
requests. If the limit is infinity, then simply set this to None.
|
||||
|
||||
This may be overridden. The value may be useful when overriding the
|
||||
"terminate_nodes" method.
|
||||
"""
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def bootstrap_config(cluster_config: Dict[str, Any]) -> Dict[str, Any]:
|
||||
"""Bootstraps the cluster config by adding env defaults if needed."""
|
||||
return cluster_config
|
||||
|
||||
def get_command_runner(
|
||||
self,
|
||||
log_prefix: str,
|
||||
node_id: str,
|
||||
auth_config: Dict[str, Any],
|
||||
cluster_name: str,
|
||||
process_runner: ModuleType,
|
||||
use_internal_ip: bool,
|
||||
docker_config: Optional[Dict[str, Any]] = None,
|
||||
) -> CommandRunnerInterface:
|
||||
"""Returns the CommandRunner class used to perform SSH commands.
|
||||
|
||||
Args:
|
||||
log_prefix: stores "NodeUpdater: {}: ".format(<node_id>). Used
|
||||
to print progress in the CommandRunner.
|
||||
node_id: the node ID.
|
||||
auth_config: the authentication configs from the autoscaler
|
||||
yaml file.
|
||||
cluster_name: the name of the cluster.
|
||||
process_runner: the module to use to run the commands
|
||||
in the CommandRunner. E.g., subprocess.
|
||||
use_internal_ip: whether the node_id belongs to an internal ip
|
||||
or external ip.
|
||||
docker_config: If set, the docker information of the docker
|
||||
container that commands should be run on.
|
||||
|
||||
Returns:
|
||||
A CommandRunner instance for the node.
|
||||
"""
|
||||
common_args = {
|
||||
"log_prefix": log_prefix,
|
||||
"node_id": node_id,
|
||||
"provider": self,
|
||||
"auth_config": auth_config,
|
||||
"cluster_name": cluster_name,
|
||||
"process_runner": process_runner,
|
||||
"use_internal_ip": use_internal_ip,
|
||||
}
|
||||
if docker_config and docker_config["container_name"] != "":
|
||||
return DockerCommandRunner(docker_config, **common_args)
|
||||
else:
|
||||
return SSHCommandRunner(**common_args)
|
||||
|
||||
def prepare_for_head_node(self, cluster_config: Dict[str, Any]) -> Dict[str, Any]:
|
||||
"""Returns a new cluster config with custom configs for head node."""
|
||||
return cluster_config
|
||||
|
||||
@staticmethod
|
||||
def fillout_available_node_types_resources(
|
||||
cluster_config: Dict[str, Any]
|
||||
) -> Dict[str, Any]:
|
||||
"""Fills out missing "resources" field for available_node_types."""
|
||||
return cluster_config
|
||||
|
||||
def safe_to_scale(self) -> bool:
|
||||
"""Optional condition to determine if it's safe to proceed with an autoscaling
|
||||
update. Can be used to wait for convergence of state managed by an external
|
||||
cluster manager.
|
||||
|
||||
Called by the autoscaler immediately after non_terminated_nodes().
|
||||
If False is returned, the autoscaler will abort the update.
|
||||
"""
|
||||
return True
|
||||
|
||||
def post_process(self) -> None:
|
||||
"""This optional method is executed at the end of
|
||||
StandardAutoscaler._update().
|
||||
"""
|
||||
pass
|
||||
Reference in New Issue
Block a user