wehub-resource-sync
330 lines
12 KiB
Python
330 lines
12 KiB
Python
import json
|
|
import logging
|
|
import os
|
|
import socket
|
|
from threading import RLock
|
|
|
|
from filelock import FileLock
|
|
|
|
from ray.autoscaler._private.local.config import (
|
|
LOCAL_CLUSTER_NODE_TYPE,
|
|
bootstrap_local,
|
|
get_lock_path,
|
|
get_state_path,
|
|
)
|
|
from ray.autoscaler.node_provider import NodeProvider
|
|
from ray.autoscaler.tags import (
|
|
NODE_KIND_HEAD,
|
|
NODE_KIND_WORKER,
|
|
STATUS_UP_TO_DATE,
|
|
TAG_RAY_NODE_KIND,
|
|
TAG_RAY_NODE_NAME,
|
|
TAG_RAY_NODE_STATUS,
|
|
TAG_RAY_USER_NODE_TYPE,
|
|
)
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
filelock_logger = logging.getLogger("filelock")
|
|
filelock_logger.setLevel(logging.WARNING)
|
|
|
|
|
|
class ClusterState:
|
|
def __init__(self, lock_path, save_path, provider_config):
|
|
self.lock = RLock()
|
|
os.makedirs(os.path.dirname(lock_path), exist_ok=True)
|
|
self.file_lock = FileLock(lock_path)
|
|
self.save_path = save_path
|
|
|
|
with self.lock:
|
|
with self.file_lock:
|
|
if os.path.exists(self.save_path):
|
|
workers = json.loads(open(self.save_path).read())
|
|
head_config = workers.get(provider_config["head_ip"])
|
|
if (
|
|
not head_config
|
|
or head_config.get("tags", {}).get(TAG_RAY_NODE_KIND)
|
|
!= NODE_KIND_HEAD
|
|
):
|
|
workers = {}
|
|
logger.info("Head IP changed - recreating cluster.")
|
|
else:
|
|
workers = {}
|
|
logger.info(
|
|
"ClusterState: Loaded cluster state: {}".format(list(workers))
|
|
)
|
|
for worker_ip in provider_config["worker_ips"]:
|
|
if worker_ip not in workers:
|
|
workers[worker_ip] = {
|
|
"tags": {TAG_RAY_NODE_KIND: NODE_KIND_WORKER},
|
|
"state": "terminated",
|
|
}
|
|
else:
|
|
assert (
|
|
workers[worker_ip]["tags"][TAG_RAY_NODE_KIND]
|
|
== NODE_KIND_WORKER
|
|
)
|
|
if provider_config["head_ip"] not in workers:
|
|
workers[provider_config["head_ip"]] = {
|
|
"tags": {TAG_RAY_NODE_KIND: NODE_KIND_HEAD},
|
|
"state": "terminated",
|
|
}
|
|
else:
|
|
assert (
|
|
workers[provider_config["head_ip"]]["tags"][TAG_RAY_NODE_KIND]
|
|
== NODE_KIND_HEAD
|
|
)
|
|
# Relevant when a user reduces the number of workers
|
|
# without changing the headnode.
|
|
list_of_node_ips = list(provider_config["worker_ips"])
|
|
list_of_node_ips.append(provider_config["head_ip"])
|
|
for worker_ip in list(workers):
|
|
if worker_ip not in list_of_node_ips:
|
|
del workers[worker_ip]
|
|
|
|
# Set external head ip, if provided by user.
|
|
# Necessary if calling `ray up` from outside the network.
|
|
# Refer to LocalNodeProvider.external_ip function.
|
|
external_head_ip = provider_config.get("external_head_ip")
|
|
if external_head_ip:
|
|
head = workers[provider_config["head_ip"]]
|
|
head["external_ip"] = external_head_ip
|
|
|
|
assert len(workers) == len(provider_config["worker_ips"]) + 1
|
|
with open(self.save_path, "w") as f:
|
|
logger.debug(
|
|
"ClusterState: Writing cluster state: {}".format(workers)
|
|
)
|
|
f.write(json.dumps(workers))
|
|
|
|
def get(self):
|
|
with self.lock:
|
|
with self.file_lock:
|
|
workers = json.loads(open(self.save_path).read())
|
|
return workers
|
|
|
|
def put(self, worker_id, info):
|
|
assert "tags" in info
|
|
assert "state" in info
|
|
with self.lock:
|
|
with self.file_lock:
|
|
workers = self.get()
|
|
workers[worker_id] = info
|
|
with open(self.save_path, "w") as f:
|
|
logger.info(
|
|
"ClusterState: "
|
|
"Writing cluster state: {}".format(list(workers))
|
|
)
|
|
f.write(json.dumps(workers))
|
|
|
|
|
|
class OnPremCoordinatorState(ClusterState):
|
|
"""Generates & updates the state file of CoordinatorSenderNodeProvider.
|
|
|
|
Unlike ClusterState, which generates a cluster specific file with
|
|
predefined head and worker ips, OnPremCoordinatorState overwrites
|
|
ClusterState's __init__ function to generate and manage a unified
|
|
file of the status of all the nodes for multiple clusters.
|
|
"""
|
|
|
|
def __init__(self, lock_path, save_path, list_of_node_ips):
|
|
self.lock = RLock()
|
|
self.file_lock = FileLock(lock_path)
|
|
self.save_path = save_path
|
|
|
|
with self.lock:
|
|
with self.file_lock:
|
|
if os.path.exists(self.save_path):
|
|
nodes = json.loads(open(self.save_path).read())
|
|
else:
|
|
nodes = {}
|
|
logger.info(
|
|
"OnPremCoordinatorState: "
|
|
"Loaded on prem coordinator state: {}".format(nodes)
|
|
)
|
|
|
|
# Filter removed node ips.
|
|
for node_ip in list(nodes):
|
|
if node_ip not in list_of_node_ips:
|
|
del nodes[node_ip]
|
|
|
|
for node_ip in list_of_node_ips:
|
|
if node_ip not in nodes:
|
|
nodes[node_ip] = {
|
|
"tags": {},
|
|
"state": "terminated",
|
|
}
|
|
assert len(nodes) == len(list_of_node_ips)
|
|
with open(self.save_path, "w") as f:
|
|
logger.info(
|
|
"OnPremCoordinatorState: "
|
|
"Writing on prem coordinator state: {}".format(nodes)
|
|
)
|
|
f.write(json.dumps(nodes))
|
|
|
|
|
|
class LocalNodeProvider(NodeProvider):
|
|
"""NodeProvider for private/local clusters.
|
|
|
|
`node_id` is overloaded to also be `node_ip` in this class.
|
|
|
|
When `cluster_name` is provided, it manages a single cluster in a cluster
|
|
specific state file. But when `cluster_name` is None, it manages multiple
|
|
clusters in a unified state file that requires each node to be tagged with
|
|
TAG_RAY_CLUSTER_NAME in create and non_terminated_nodes function calls to
|
|
associate each node with the right cluster.
|
|
|
|
The current use case of managing multiple clusters is by
|
|
OnPremCoordinatorServer which receives node provider HTTP requests
|
|
from CoordinatorSenderNodeProvider and uses LocalNodeProvider to get
|
|
the responses.
|
|
"""
|
|
|
|
def __init__(self, provider_config, cluster_name):
|
|
NodeProvider.__init__(self, provider_config, cluster_name)
|
|
|
|
if cluster_name:
|
|
lock_path = get_lock_path(cluster_name)
|
|
state_path = get_state_path(cluster_name)
|
|
self.state = ClusterState(
|
|
lock_path,
|
|
state_path,
|
|
provider_config,
|
|
)
|
|
self.use_coordinator = False
|
|
else:
|
|
# LocalNodeProvider with a coordinator server.
|
|
self.state = OnPremCoordinatorState(
|
|
"/tmp/coordinator.lock",
|
|
"/tmp/coordinator.state",
|
|
provider_config["list_of_node_ips"],
|
|
)
|
|
self.use_coordinator = True
|
|
|
|
def non_terminated_nodes(self, tag_filters):
|
|
workers = self.state.get()
|
|
matching_ips = []
|
|
for worker_ip, info in workers.items():
|
|
if info["state"] == "terminated":
|
|
continue
|
|
ok = True
|
|
for k, v in tag_filters.items():
|
|
if info["tags"].get(k) != v:
|
|
ok = False
|
|
break
|
|
if ok:
|
|
matching_ips.append(worker_ip)
|
|
return matching_ips
|
|
|
|
def nodes_for_teardown(self, tag_filters):
|
|
"""Return all known node ids matching tag_filters regardless of state.
|
|
|
|
The local state file on the machine invoking ``ray down`` may show
|
|
workers as terminated because only the head node's autoscaler updates
|
|
them to running. During teardown we still need to reach these nodes
|
|
to stop their Docker containers.
|
|
|
|
Nodes that have already been fully torn down (i.e. terminate_node was
|
|
called, setting teardown_complete) are excluded so they are not
|
|
targeted again on subsequent teardown passes.
|
|
"""
|
|
workers = self.state.get()
|
|
return [
|
|
worker_ip
|
|
for worker_ip, info in workers.items()
|
|
if all(info["tags"].get(k) == v for k, v in tag_filters.items())
|
|
and not info.get("teardown_complete", False)
|
|
]
|
|
|
|
def is_running(self, node_id):
|
|
return self.state.get()[node_id]["state"] == "running"
|
|
|
|
def is_terminated(self, node_id):
|
|
return not self.is_running(node_id)
|
|
|
|
def node_tags(self, node_id):
|
|
return self.state.get()[node_id]["tags"]
|
|
|
|
def external_ip(self, node_id):
|
|
"""Returns an external ip if the user has supplied one.
|
|
Otherwise, use the same logic as internal_ip below.
|
|
|
|
This can be used to call ray up from outside the network, for example
|
|
if the Ray cluster exists in an AWS VPC and we're interacting with
|
|
the cluster from a laptop (where using an internal_ip will not work).
|
|
|
|
Useful for debugging the local node provider with cloud VMs."""
|
|
|
|
node_state = self.state.get()[node_id]
|
|
ext_ip = node_state.get("external_ip")
|
|
if ext_ip:
|
|
return ext_ip
|
|
else:
|
|
return socket.gethostbyname(node_id)
|
|
|
|
def internal_ip(self, node_id):
|
|
return socket.gethostbyname(node_id)
|
|
|
|
def set_node_tags(self, node_id, tags):
|
|
with self.state.lock:
|
|
with self.state.file_lock:
|
|
info = self.state.get()[node_id]
|
|
info["tags"].update(tags)
|
|
self.state.put(node_id, info)
|
|
|
|
def create_node(self, node_config, tags, count):
|
|
"""Creates min(count, currently available) nodes."""
|
|
node_type = tags[TAG_RAY_NODE_KIND]
|
|
with self.state.lock:
|
|
with self.state.file_lock:
|
|
workers = self.state.get()
|
|
for node_id, info in workers.items():
|
|
if info["state"] == "terminated" and (
|
|
self.use_coordinator
|
|
or info["tags"][TAG_RAY_NODE_KIND] == node_type
|
|
):
|
|
info["tags"] = tags
|
|
info["state"] = "running"
|
|
info.pop("teardown_complete", None)
|
|
self.state.put(node_id, info)
|
|
count = count - 1
|
|
if count == 0:
|
|
return
|
|
|
|
def terminate_node(self, node_id):
|
|
workers = self.state.get()
|
|
info = workers[node_id]
|
|
info["state"] = "terminated"
|
|
info["teardown_complete"] = True
|
|
self.state.put(node_id, info)
|
|
|
|
@staticmethod
|
|
def bootstrap_config(cluster_config):
|
|
return bootstrap_local(cluster_config)
|
|
|
|
|
|
def record_local_head_state_if_needed(local_provider: LocalNodeProvider) -> None:
|
|
"""This function is called on the Ray head from StandardAutoscaler.reset
|
|
to record the head node's own existence in the cluster state file.
|
|
|
|
This is necessary because `provider.create_node` in
|
|
`commands.get_or_create_head_node` records the head state on the
|
|
cluster-launching machine but not on the head.
|
|
"""
|
|
head_ip = local_provider.provider_config["head_ip"]
|
|
cluster_name = local_provider.cluster_name
|
|
# If the head node is not marked as created in the cluster state file,
|
|
if head_ip not in local_provider.non_terminated_nodes({}):
|
|
# These tags are based on the ones in commands.get_or_create_head_node;
|
|
# keep in sync.
|
|
head_tags = {
|
|
TAG_RAY_NODE_KIND: NODE_KIND_HEAD,
|
|
TAG_RAY_USER_NODE_TYPE: LOCAL_CLUSTER_NODE_TYPE,
|
|
TAG_RAY_NODE_NAME: "ray-{}-head".format(cluster_name),
|
|
TAG_RAY_NODE_STATUS: STATUS_UP_TO_DATE,
|
|
}
|
|
# Mark the head node as created in the cluster state file.
|
|
local_provider.create_node(node_config={}, tags=head_tags, count=1)
|
|
|
|
assert head_ip in local_provider.non_terminated_nodes({})
|