wehub-resource-sync
313 lines
11 KiB
Python
313 lines
11 KiB
Python
import re
|
|
import threading
|
|
from subprocess import CalledProcessError
|
|
from typing import Any, Dict, List, Optional
|
|
|
|
from ray.autoscaler.node_provider import NodeProvider
|
|
|
|
|
|
class MockNode:
|
|
def __init__(
|
|
self,
|
|
node_id,
|
|
tags,
|
|
node_config,
|
|
node_type,
|
|
unique_ips=False,
|
|
resources=None,
|
|
labels=None,
|
|
):
|
|
self.node_id = str(node_id)
|
|
self.state = "pending"
|
|
self.tags = tags
|
|
self.external_ip = "1.2.3.4"
|
|
self.internal_ip = "172.0.0.{}".format(self.node_id)
|
|
if unique_ips:
|
|
self.external_ip = f"1.2.3.{self.node_id}"
|
|
|
|
self.node_config = node_config
|
|
self.node_type = node_type
|
|
self.created_in_main_thread = (
|
|
threading.current_thread() is threading.main_thread()
|
|
)
|
|
self.resources = resources or {}
|
|
self.labels = labels or {}
|
|
|
|
def matches(self, tags):
|
|
for k, v in tags.items():
|
|
if k not in self.tags or self.tags[k] != v:
|
|
return False
|
|
return True
|
|
|
|
|
|
class MockProcessRunner:
|
|
def __init__(self, fail_cmds=None, cmd_to_callback=None, print_out=False):
|
|
self.calls = []
|
|
self.cmd_to_callback = cmd_to_callback or {} # type: Dict[str, Callable]
|
|
self.print_out = print_out
|
|
self.fail_cmds = fail_cmds or []
|
|
self.call_response = {}
|
|
self.ready_to_run = threading.Event()
|
|
self.ready_to_run.set()
|
|
|
|
self.lock = threading.RLock()
|
|
|
|
def check_call(self, cmd, *args, **kwargs):
|
|
with self.lock:
|
|
self.ready_to_run.wait()
|
|
self.calls.append(cmd)
|
|
if self.print_out:
|
|
print(f">>>Process runner: Executing \n {str(cmd)}")
|
|
for token in self.cmd_to_callback:
|
|
if token in str(cmd):
|
|
# Trigger a callback if token is in cmd.
|
|
# Can be used to simulate background events during a node
|
|
# update (e.g. node disconnected).
|
|
callback = self.cmd_to_callback[token]
|
|
callback()
|
|
|
|
for token in self.fail_cmds:
|
|
if token in str(cmd):
|
|
raise CalledProcessError(1, token, "Failing command on purpose")
|
|
|
|
def check_output(self, cmd):
|
|
with self.lock:
|
|
self.check_call(cmd)
|
|
return_string = "command-output"
|
|
key_to_shrink = None
|
|
for pattern, response_list in self.call_response.items():
|
|
if pattern in str(cmd):
|
|
return_string = response_list[0]
|
|
key_to_shrink = pattern
|
|
break
|
|
if key_to_shrink:
|
|
self.call_response[key_to_shrink] = self.call_response[key_to_shrink][
|
|
1:
|
|
]
|
|
if len(self.call_response[key_to_shrink]) == 0:
|
|
del self.call_response[key_to_shrink]
|
|
|
|
return return_string.encode()
|
|
|
|
def assert_has_call(
|
|
self, ip: str, pattern: Optional[str] = None, exact: Optional[List[str]] = None
|
|
) -> bool:
|
|
"""Checks if the given value was called by this process runner.
|
|
|
|
NOTE: Either pattern or exact must be specified, not both!
|
|
|
|
Args:
|
|
ip: IP address of the node that the given call was executed on.
|
|
pattern: RegEx that matches one specific call.
|
|
exact: List of strings that when joined exactly match one call.
|
|
|
|
Returns:
|
|
``True`` when a matching call is found. Raises ``Exception`` when
|
|
no matching call is recorded.
|
|
"""
|
|
with self.lock:
|
|
assert bool(pattern) ^ bool(
|
|
exact
|
|
), "Must specify either a pattern or exact match."
|
|
debug_output = ""
|
|
if pattern is not None:
|
|
for cmd in self.command_history():
|
|
if ip in cmd:
|
|
debug_output += cmd
|
|
debug_output += "\n"
|
|
if re.search(pattern, cmd):
|
|
return True
|
|
else:
|
|
raise Exception(
|
|
f"Did not find [{pattern}] in [{debug_output}] for "
|
|
f"ip={ip}.\n\nFull output: {self.command_history()}"
|
|
)
|
|
elif exact is not None:
|
|
exact_cmd = " ".join(exact)
|
|
for cmd in self.command_history():
|
|
if ip in cmd:
|
|
debug_output += cmd
|
|
debug_output += "\n"
|
|
if cmd == exact_cmd:
|
|
return True
|
|
raise Exception(
|
|
f"Did not find [{exact_cmd}] in [{debug_output}] for "
|
|
f"ip={ip}.\n\nFull output: {self.command_history()}"
|
|
)
|
|
|
|
def assert_not_has_call(self, ip: str, pattern: str):
|
|
"""Ensure that the given regex pattern was never called."""
|
|
with self.lock:
|
|
out = ""
|
|
for cmd in self.command_history():
|
|
if ip in cmd:
|
|
out += cmd
|
|
out += "\n"
|
|
if re.search(pattern, out):
|
|
raise Exception("Found [{}] in [{}] for {}".format(pattern, out, ip))
|
|
else:
|
|
return True
|
|
|
|
def clear_history(self):
|
|
with self.lock:
|
|
self.calls = []
|
|
|
|
def command_history(self):
|
|
with self.lock:
|
|
return [" ".join(cmd) for cmd in self.calls]
|
|
|
|
def respond_to_call(self, pattern, response_list):
|
|
with self.lock:
|
|
self.call_response[pattern] = response_list
|
|
|
|
|
|
class MockProvider(NodeProvider):
|
|
def __init__(self, cache_stopped=False, unique_ips=False):
|
|
self.mock_nodes = {}
|
|
self.next_id = 0
|
|
self.throw = False
|
|
self.creation_error = None
|
|
self.termination_errors = None
|
|
self.fail_creates = False
|
|
self.ready_to_create = threading.Event()
|
|
self.ready_to_create.set()
|
|
self.cache_stopped = cache_stopped
|
|
self.unique_ips = unique_ips
|
|
self.fail_to_fetch_ip = False
|
|
self.safe_to_scale_flag = True
|
|
self.partical_success_count = None
|
|
# Many of these functions are called by node_launcher or updater in
|
|
# different threads. This can be treated as a global lock for
|
|
# everything.
|
|
self.lock = threading.Lock()
|
|
self.num_non_terminated_nodes_calls = 0
|
|
super().__init__(None, None)
|
|
|
|
def non_terminated_nodes(self, tag_filters):
|
|
self.num_non_terminated_nodes_calls += 1
|
|
with self.lock:
|
|
if self.throw:
|
|
raise Exception("oops")
|
|
return [
|
|
n.node_id
|
|
for n in self.mock_nodes.values()
|
|
if n.matches(tag_filters) and n.state not in ["stopped", "terminated"]
|
|
]
|
|
|
|
def non_terminated_node_ips(self, tag_filters):
|
|
with self.lock:
|
|
if self.throw:
|
|
raise Exception("oops")
|
|
return [
|
|
n.internal_ip
|
|
for n in self.mock_nodes.values()
|
|
if n.matches(tag_filters) and n.state not in ["stopped", "terminated"]
|
|
]
|
|
|
|
def is_running(self, node_id):
|
|
with self.lock:
|
|
return self.mock_nodes[node_id].state == "running"
|
|
|
|
def is_terminated(self, node_id):
|
|
if node_id is None:
|
|
# Circumvent test-cases where there's no head node.
|
|
return True
|
|
with self.lock:
|
|
return self.mock_nodes[node_id].state in ["stopped", "terminated"]
|
|
|
|
def node_tags(self, node_id):
|
|
if node_id is None:
|
|
# Circumvent test cases where there's no head node.
|
|
return {}
|
|
# Don't assume that node providers can retrieve tags from
|
|
# terminated nodes.
|
|
if self.is_terminated(node_id):
|
|
raise Exception(f"The node with id {node_id} has been terminated!")
|
|
with self.lock:
|
|
return self.mock_nodes[node_id].tags
|
|
|
|
def internal_ip(self, node_id):
|
|
if self.fail_to_fetch_ip:
|
|
raise Exception("Failed to fetch ip on purpose.")
|
|
if node_id is None:
|
|
# Circumvent test-cases where there's no head node.
|
|
return "mock"
|
|
with self.lock:
|
|
return self.mock_nodes[node_id].internal_ip
|
|
|
|
def external_ip(self, node_id):
|
|
with self.lock:
|
|
return self.mock_nodes[node_id].external_ip
|
|
|
|
def create_node(
|
|
self,
|
|
node_config: Dict[str, Any],
|
|
tags: Dict[str, str],
|
|
count: int,
|
|
_skip_wait=False,
|
|
) -> Dict[str, Any]:
|
|
return self.create_node_with_resources_and_labels(
|
|
node_config, tags, count, {}, {}, _skip_wait=_skip_wait
|
|
)
|
|
|
|
def create_node_with_resources_and_labels(
|
|
self, node_config, tags, count, resources, labels, _skip_wait=False
|
|
):
|
|
from ray.autoscaler.tags import TAG_RAY_USER_NODE_TYPE
|
|
|
|
if self.creation_error is not None:
|
|
raise self.creation_error
|
|
if not _skip_wait:
|
|
self.ready_to_create.wait()
|
|
if self.fail_creates:
|
|
return
|
|
|
|
created_nodes = {}
|
|
if self.partical_success_count is not None:
|
|
count = min(count, self.partical_success_count)
|
|
with self.lock:
|
|
if self.cache_stopped:
|
|
for node in self.mock_nodes.values():
|
|
if node.state == "stopped" and count > 0:
|
|
count -= 1
|
|
node.state = "pending"
|
|
node.tags.update(tags)
|
|
created_nodes[node.node_id] = node
|
|
for _ in range(count):
|
|
new_node = MockNode(
|
|
str(self.next_id),
|
|
tags.copy(),
|
|
node_config,
|
|
tags.get(TAG_RAY_USER_NODE_TYPE),
|
|
resources=resources,
|
|
labels=labels,
|
|
unique_ips=self.unique_ips,
|
|
)
|
|
self.mock_nodes[new_node.node_id] = new_node
|
|
created_nodes[new_node.node_id] = new_node
|
|
self.next_id += 1
|
|
return created_nodes
|
|
|
|
def set_node_tags(self, node_id, tags):
|
|
with self.lock:
|
|
self.mock_nodes[node_id].tags.update(tags)
|
|
|
|
def terminate_node(self, node_id):
|
|
with self.lock:
|
|
if self.termination_errors is not None:
|
|
raise self.termination_errors
|
|
|
|
if self.cache_stopped:
|
|
self.mock_nodes[node_id].state = "stopped"
|
|
else:
|
|
self.mock_nodes[node_id].state = "terminated"
|
|
|
|
def finish_starting_nodes(self):
|
|
with self.lock:
|
|
for node in self.mock_nodes.values():
|
|
if node.state == "pending":
|
|
node.state = "running"
|
|
|
|
def safe_to_scale(self):
|
|
return self.safe_to_scale_flag
|