diff --git a/python/ray/autoscaler/_private/aws/config.py b/python/ray/autoscaler/_private/aws/config.py index 042504d24..8b865a3cb 100644 --- a/python/ray/autoscaler/_private/aws/config.py +++ b/python/ray/autoscaler/_private/aws/config.py @@ -591,6 +591,8 @@ def _create_security_group(config, vpc_id, group_name): GroupName=group_name, VpcId=vpc_id) security_group = _get_security_group(config, vpc_id, group_name) + cli_logger.doassert(security_group, + "Failed to create security group") # err msg cli_logger.verbose( "Created new security group {}", @@ -599,9 +601,6 @@ def _create_security_group(config, vpc_id, group_name): cli_logger.old_info( logger, "_create_security_group: Created new security group {} ({})", security_group.group_name, security_group.id) - - cli_logger.doassert(security_group, - "Failed to create security group") # err msg assert security_group, "Failed to create security group" return security_group diff --git a/python/ray/autoscaler/_private/commands.py b/python/ray/autoscaler/_private/commands.py index 948021db2..d01b48c3d 100644 --- a/python/ray/autoscaler/_private/commands.py +++ b/python/ray/autoscaler/_private/commands.py @@ -941,8 +941,11 @@ def rsync(config_file: str, target: Optional[str], override_cluster_name: Optional[str], down: bool, + ip_address: Optional[str] = None, + use_internal_ip: bool = False, no_config_cache: bool = False, - all_nodes: bool = False): + all_nodes: bool = False, + _runner=subprocess): """Rsyncs files. Arguments: @@ -951,6 +954,10 @@ def rsync(config_file: str, target: target dir override_cluster_name: set the name of the cluster down: whether we're syncing remote -> local + ip_address (str): Address of node. Raise Exception + if both ip_address and 'all_nodes' are provided. + use_internal_ip (bool): Whether the provided ip_address is + public or private. all_nodes: whether to sync worker nodes in addition to the head node """ if bool(source) != bool(target): @@ -960,6 +967,9 @@ def rsync(config_file: str, assert bool(source) == bool(target), ( "Must either provide both or neither source and target.") + if ip_address and all_nodes: + cli_logger.abort("Cannot provide both ip_address and 'all_nodes'.") + config = yaml.safe_load(open(config_file).read()) if override_cluster_name is not None: config["cluster_name"] = override_cluster_name @@ -973,47 +983,53 @@ def rsync(config_file: str, break provider = _get_node_provider(config["provider"], config["cluster_name"]) + + def rsync_to_node(node_id, is_head_node): + updater = NodeUpdaterThread( + node_id=node_id, + provider_config=config["provider"], + provider=provider, + auth_config=config["auth"], + cluster_name=config["cluster_name"], + file_mounts=config["file_mounts"], + initialization_commands=[], + setup_commands=[], + ray_start_commands=[], + runtime_hash="", + use_internal_ip=use_internal_ip, + process_runner=_runner, + file_mounts_contents_hash="", + is_head_node=is_head_node, + docker_config=config.get("docker")) + if down: + rsync = updater.rsync_down + else: + rsync = updater.rsync_up + + if source and target: + # print rsync progress for single file rsync + cmd_output_util.set_output_redirected(False) + set_rsync_silent(False) + rsync(source, target, is_file_mount) + else: + updater.sync_file_mounts(rsync) + try: nodes = [] - if all_nodes: - # technically we re-open the provider for no reason - # in get_worker_nodes but it's cleaner this way - # and _get_head_node does this too - nodes = _get_worker_nodes(config, override_cluster_name) - head_node = _get_head_node( config, config_file, override_cluster_name, create_if_needed=False) - - nodes += [head_node] + if ip_address: + nodes = [ + provider.get_node_id( + ip_address, use_internal_ip=use_internal_ip) + ] + else: + if all_nodes: + nodes = _get_worker_nodes(config, override_cluster_name) + nodes += [head_node] for node_id in nodes: - updater = NodeUpdaterThread( - node_id=node_id, - provider_config=config["provider"], - provider=provider, - auth_config=config["auth"], - cluster_name=config["cluster_name"], - file_mounts=config["file_mounts"], - initialization_commands=[], - setup_commands=[], - ray_start_commands=[], - runtime_hash="", - file_mounts_contents_hash="", - is_head_node=(node_id == head_node), - docker_config=config.get("docker")) - if down: - rsync = updater.rsync_down - else: - rsync = updater.rsync_up - - if source and target: - # print rsync progress for single file rsync - cmd_output_util.set_output_redirected(False) - set_rsync_silent(False) - - rsync(source, target, is_file_mount) - else: - updater.sync_file_mounts(rsync) + rsync_to_node(node_id, is_head_node=(node_id == head_node)) finally: provider.cleanup() diff --git a/python/ray/autoscaler/_private/kubernetes/node_provider.py b/python/ray/autoscaler/_private/kubernetes/node_provider.py index 22ff84c50..bed285a4b 100644 --- a/python/ray/autoscaler/_private/kubernetes/node_provider.py +++ b/python/ray/autoscaler/_private/kubernetes/node_provider.py @@ -66,6 +66,11 @@ class KubernetesNodeProvider(NodeProvider): pod = core_api().read_namespaced_pod_status(node_id, self.namespace) return pod.status.pod_ip + def get_node_id(self, ip_address, use_internal_ip=True) -> str: + if not use_internal_ip: + raise ValueError("Must use internal IPs with Kubernetes.") + return super().get_node_id(ip_address, use_internal_ip=use_internal_ip) + def set_node_tags(self, node_id, tags): pod = core_api().read_namespaced_pod_status(node_id, self.namespace) pod.metadata.labels.update(tags) diff --git a/python/ray/autoscaler/_private/providers.py b/python/ray/autoscaler/_private/providers.py index 34258ba51..329751916 100644 --- a/python/ray/autoscaler/_private/providers.py +++ b/python/ray/autoscaler/_private/providers.py @@ -1,5 +1,6 @@ import importlib import logging +import json import os from typing import Any, Dict @@ -7,6 +8,9 @@ import yaml logger = logging.getLogger(__name__) +# For caching provider instantiations across API calls of one python session +_provider_instances = {} + def _import_aws(provider_config): from ray.autoscaler._private.aws.node_provider import AWSNodeProvider @@ -131,18 +135,28 @@ def _load_class(path): def _get_node_provider(provider_config: Dict[str, Any], - cluster_name: str) -> Any: - """Retrieve a node provider. - - This returns access to an INTERNAL API. It is not allowed to call this - from any Ray package outside the autoscaler. - """ + cluster_name: str, + use_cache: bool = True) -> Any: importer = _NODE_PROVIDERS.get(provider_config["type"]) if importer is None: raise NotImplementedError("Unsupported node provider: {}".format( provider_config["type"])) provider_cls = importer(provider_config) - return provider_cls(provider_config, cluster_name) + provider_key = (json.dumps(provider_config, sort_keys=True), cluster_name) + if use_cache and provider_key in _provider_instances: + return _provider_instances[provider_key] + + new_provider = provider_cls(provider_config, cluster_name) + + if use_cache: + _provider_instances[provider_key] = new_provider + + return new_provider + + +def _clear_provider_cache(): + global _provider_instances + _provider_instances = {} def _get_default_config(provider_config): diff --git a/python/ray/autoscaler/_private/staroid/node_provider.py b/python/ray/autoscaler/_private/staroid/node_provider.py index 4748bc1f9..cde57fc3a 100644 --- a/python/ray/autoscaler/_private/staroid/node_provider.py +++ b/python/ray/autoscaler/_private/staroid/node_provider.py @@ -217,6 +217,11 @@ class StaroidNodeProvider(NodeProvider): pod = core_api.read_namespaced_pod_status(node_id, self.namespace) return pod.status.pod_ip + def get_node_id(self, ip_address, use_internal_ip=True) -> str: + if not use_internal_ip: + raise ValueError("Must use internal IPs with Staroid.") + return super().get_node_id(ip_address, use_internal_ip=use_internal_ip) + def set_node_tags(self, node_id, tags): kube_client = self.__cached[self.cluster_name]["kube_client"] core_api = client.CoreV1Api(kube_client) diff --git a/python/ray/autoscaler/node_provider.py b/python/ray/autoscaler/node_provider.py index b2e935915..cc37313c7 100644 --- a/python/ray/autoscaler/node_provider.py +++ b/python/ray/autoscaler/node_provider.py @@ -27,6 +27,8 @@ class NodeProvider: def __init__(self, provider_config, cluster_name): self.provider_config = provider_config self.cluster_name = cluster_name + self._internal_ip_cache = {} + self._external_ip_cache = {} def non_terminated_nodes(self, tag_filters): """Return a list of node ids filtered by the specified tags dict. @@ -62,6 +64,47 @@ class NodeProvider: """Returns the internal ip (Ray ip) of the given node.""" raise NotImplementedError + def get_node_id(self, ip_address, use_internal_ip=False) -> str: + """Returns the node_id given an IP address. + + Assumes ip-address is unique per node. + + Args: + ip_address (str): Address of node. + use_internal_ip (bool): Whether the ip address is + public or private. + + 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({}) + for node_id in all_nodes: + if use_internal_ip: + int_ip = self.internal_ip(node_id) + self._internal_ip_cache[int_ip] = node_id + else: + ext_ip = self.external_ip(node_id) + self._external_ip_cache[ext_ip] = 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, tags, count): """Creates a number of nodes within the namespace.""" raise NotImplementedError diff --git a/python/ray/autoscaler/sdk.py b/python/ray/autoscaler/sdk.py index 1083af76e..801bbb847 100644 --- a/python/ray/autoscaler/sdk.py +++ b/python/ray/autoscaler/sdk.py @@ -96,6 +96,8 @@ def rsync(cluster_config: Union[dict, str], source: str, target: str, down: bool, + ip_address: str = None, + use_internal_ip: bool = False, no_config_cache: bool = False): """Rsyncs files to or from the cluster. @@ -105,6 +107,9 @@ def rsync(cluster_config: Union[dict, str], source (str): rsync source argument. target (str): rsync target argument. down (bool): whether we're syncing remote -> local. + ip_address (str): Address of node. + use_internal_ip (bool): Whether the provided ip_address is + public or private. no_config_cache (bool): Whether to disable the config cache and fully resolve all environment settings from the Cloud provider again. @@ -117,6 +122,8 @@ def rsync(cluster_config: Union[dict, str], target=target, override_cluster_name=None, down=down, + ip_address=ip_address, + use_internal_ip=use_internal_ip, no_config_cache=no_config_cache, all_nodes=False) diff --git a/python/ray/tests/test_autoscaler.py b/python/ray/tests/test_autoscaler.py index b79ebd372..e6be2854f 100644 --- a/python/ray/tests/test_autoscaler.py +++ b/python/ray/tests/test_autoscaler.py @@ -5,6 +5,7 @@ import tempfile import threading import time import unittest +from unittest.mock import Mock import yaml import copy import sys @@ -13,11 +14,12 @@ from jsonschema.exceptions import ValidationError import ray import ray._private.services as services from ray.autoscaler._private.util import prepare_config, validate_config -from ray.autoscaler._private.commands import get_or_create_head_node +from ray.autoscaler._private import commands from ray.autoscaler._private.docker import DOCKER_MOUNT_PREFIX from ray.autoscaler._private.load_metrics import LoadMetrics from ray.autoscaler._private.autoscaler import StandardAutoscaler -from ray.autoscaler._private.providers import _NODE_PROVIDERS +from ray.autoscaler._private.providers import (_NODE_PROVIDERS, + _clear_provider_cache) from ray.autoscaler.tags import TAG_RAY_NODE_KIND, TAG_RAY_NODE_STATUS, \ STATUS_UP_TO_DATE, STATUS_UPDATE_FAILED, TAG_RAY_USER_NODE_TYPE from ray.autoscaler.node_provider import NodeProvider @@ -26,12 +28,16 @@ import pytest class MockNode: - def __init__(self, node_id, tags, node_config, node_type): + def __init__(self, node_id, tags, node_config, node_type, + unique_ips=False): self.node_id = 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 @@ -43,9 +49,9 @@ class MockNode: class MockProcessRunner: - def __init__(self, fail_cmds=[]): + def __init__(self, fail_cmds=None): self.calls = [] - self.fail_cmds = fail_cmds + self.fail_cmds = fail_cmds or [] self.call_response = {} def check_call(self, cmd, *args, **kwargs): @@ -77,33 +83,33 @@ class MockProcessRunner: "Must specify either a pattern or exact match." out = "" if pattern is not None: - for cmd in self.calls: - msg = " ".join(cmd) - if ip in msg: - out += msg + for cmd in self.command_history(): + if ip in cmd: + out += cmd out += "\n" if pattern in out: return True else: - raise Exception("Did not find [{}] in [{}] for {}".format( - pattern, out, ip)) - else: - for cmd in self.calls: - msg = " ".join(cmd) - if ip in msg: - out += msg + raise Exception( + f"Did not find [{pattern}] in [{out}] for ip={ip}." + f"\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: + out += cmd out += "\n" - if cmd == exact: + if cmd == exact_cmd: return True - raise Exception("Did not find {} in {} for {}".format( - exact, out, ip)) + raise Exception( + f"Did not find [{exact_cmd}] in [{out}] for ip={ip}." + f"\n\nFull output: {self.command_history()}") def assert_not_has_call(self, ip, pattern): out = "" - for cmd in self.calls: - msg = " ".join(cmd) - if ip in msg: - out += msg + for cmd in self.command_history(): + if ip in cmd: + out += cmd out += "\n" if pattern in out: raise Exception("Found [{}] in [{}] for {}".format( @@ -114,12 +120,15 @@ class MockProcessRunner: def clear_history(self): self.calls = [] + def command_history(self): + return [" ".join(cmd) for cmd in self.calls] + def respond_to_call(self, pattern, response_list): self.call_response[pattern] = response_list class MockProvider(NodeProvider): - def __init__(self, cache_stopped=False): + def __init__(self, cache_stopped=False, unique_ips=False): self.mock_nodes = {} self.next_id = 0 self.throw = False @@ -127,10 +136,12 @@ class MockProvider(NodeProvider): self.ready_to_create = threading.Event() self.ready_to_create.set() self.cache_stopped = cache_stopped + self.unique_ips = unique_ips # 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() + super().__init__(None, None) def non_terminated_nodes(self, tag_filters): with self.lock: @@ -180,8 +191,11 @@ class MockProvider(NodeProvider): node.tags.update(tags) for _ in range(count): self.mock_nodes[self.next_id] = MockNode( - self.next_id, tags.copy(), node_config, - tags.get(TAG_RAY_USER_NODE_TYPE)) + self.next_id, + tags.copy(), + node_config, + tags.get(TAG_RAY_USER_NODE_TYPE), + unique_ips=self.unique_ips) self.next_id += 1 def set_node_tags(self, node_id, tags): @@ -377,6 +391,7 @@ class AutoscalingTest(unittest.TestCase): def tearDown(self): self.provider = None del _NODE_PROVIDERS["mock"] + _clear_provider_cache() shutil.rmtree(self.tmpdir) ray.shutdown() @@ -457,7 +472,7 @@ class AutoscalingTest(unittest.TestCase): # Two initial calls to docker cp, one before run, two final calls to cp runner.respond_to_call(".State.Running", ["false", "false", "false", "true", "true"]) - get_or_create_head_node( + commands.get_or_create_head_node( SMALL_CLUSTER, config_path, no_restart=False, @@ -481,6 +496,100 @@ class AutoscalingTest(unittest.TestCase): "1.2.3.4", pattern="docker cp /tmp/ray_tmp_mount/~/ray_bootstrap_config.yaml") + @unittest.skipIf(sys.platform == "win32", "Failing on Windows.") + def testRsyncCommandWithDocker(self): + assert SMALL_CLUSTER["docker"]["container_name"] + config_path = self.write_config(SMALL_CLUSTER) + self.provider = MockProvider(unique_ips=True) + self.provider.create_node({}, {TAG_RAY_NODE_KIND: "head"}, 1) + self.provider.create_node({}, {TAG_RAY_NODE_KIND: "worker"}, 10) + self.provider.finish_starting_nodes() + ray.autoscaler.node_provider._get_node_provider = Mock( + return_value=self.provider) + ray.autoscaler._private.commands._bootstrap_config = Mock( + return_value=SMALL_CLUSTER) + runner = MockProcessRunner() + commands.rsync( + config_path, + source=config_path, + target="/tmp/test_path", + override_cluster_name=None, + down=True, + _runner=runner) + runner.assert_has_call("1.2.3.0", pattern="docker cp") + runner.assert_has_call("1.2.3.0", pattern="rsync") + runner.clear_history() + + commands.rsync( + config_path, + source=config_path, + target="/tmp/test_path", + override_cluster_name=None, + down=True, + ip_address="1.2.3.5", + _runner=runner) + runner.assert_has_call("1.2.3.5", pattern="docker cp") + runner.assert_has_call("1.2.3.5", pattern="rsync") + runner.clear_history() + + commands.rsync( + config_path, + source=config_path, + target="/tmp/test_path", + ip_address="172.0.0.4", + override_cluster_name=None, + down=True, + use_internal_ip=True, + _runner=runner) + runner.assert_has_call("172.0.0.4", pattern="docker cp") + runner.assert_has_call("172.0.0.4", pattern="rsync") + + @unittest.skipIf(sys.platform == "win32", "Failing on Windows.") + def testRsyncCommandWithoutDocker(self): + cluster_cfg = SMALL_CLUSTER.copy() + cluster_cfg["docker"] = {} + config_path = self.write_config(cluster_cfg) + self.provider = MockProvider(unique_ips=True) + self.provider.create_node({}, {TAG_RAY_NODE_KIND: "head"}, 1) + self.provider.create_node({}, {TAG_RAY_NODE_KIND: "worker"}, 10) + self.provider.finish_starting_nodes() + runner = MockProcessRunner() + ray.autoscaler.node_provider._get_node_provider = Mock( + return_value=self.provider) + ray.autoscaler._private.commands._bootstrap_config = Mock( + return_value=SMALL_CLUSTER) + commands.rsync( + config_path, + source=config_path, + target="/tmp/test_path", + override_cluster_name=None, + down=True, + _runner=runner) + runner.assert_has_call("1.2.3.0", pattern="rsync") + + commands.rsync( + config_path, + source=config_path, + target="/tmp/test_path", + override_cluster_name=None, + down=True, + ip_address="1.2.3.5", + _runner=runner) + runner.assert_has_call("1.2.3.5", pattern="rsync") + runner.clear_history() + + commands.rsync( + config_path, + source=config_path, + target="/tmp/test_path", + override_cluster_name=None, + down=True, + ip_address="172.0.0.4", + use_internal_ip=True, + _runner=runner) + runner.assert_has_call("172.0.0.4", pattern="rsync") + runner.clear_history() + def testScaleUp(self): config_path = self.write_config(SMALL_CLUSTER) self.provider = MockProvider() diff --git a/python/ray/tests/test_coordinator_server.py b/python/ray/tests/test_coordinator_server.py index 726ee9a80..6fb654e3e 100644 --- a/python/ray/tests/test_coordinator_server.py +++ b/python/ray/tests/test_coordinator_server.py @@ -60,8 +60,8 @@ class OnPremCoordinatorServerTest(unittest.TestCase): }, } provider_config = cluster_config["provider"] - node_provider = _get_node_provider(provider_config, - cluster_config["cluster_name"]) + node_provider = _get_node_provider( + provider_config, cluster_config["cluster_name"], use_cache=False) assert isinstance(node_provider, LocalNodeProvider) expected_workers = {} expected_workers[provider_config["head_ip"]] = { @@ -86,8 +86,8 @@ class OnPremCoordinatorServerTest(unittest.TestCase): # Test removing workers updates the cluster state. del expected_workers[provider_config["worker_ips"][0]] removed_ip = provider_config["worker_ips"].pop() - node_provider = _get_node_provider(provider_config, - cluster_config["cluster_name"]) + node_provider = _get_node_provider( + provider_config, cluster_config["cluster_name"], use_cache=False) workers = json.loads(open(state_save_path).read()) assert workers == expected_workers @@ -99,8 +99,8 @@ class OnPremCoordinatorServerTest(unittest.TestCase): "state": "terminated", } provider_config["worker_ips"].append(removed_ip) - node_provider = _get_node_provider(provider_config, - cluster_config["cluster_name"]) + node_provider = _get_node_provider( + provider_config, cluster_config["cluster_name"], use_cache=False) workers = json.loads(open(state_save_path).read()) assert workers == expected_workers @@ -163,8 +163,8 @@ class OnPremCoordinatorServerTest(unittest.TestCase): "worker_nodes": {}, } provider_config = cluster_config["provider"] - node_provider_1 = _get_node_provider(provider_config, - cluster_config["cluster_name"]) + node_provider_1 = _get_node_provider( + provider_config, cluster_config["cluster_name"], use_cache=False) assert isinstance(node_provider_1, CoordinatorSenderNodeProvider) assert not node_provider_1.non_terminated_nodes({}) @@ -190,8 +190,8 @@ class OnPremCoordinatorServerTest(unittest.TestCase): # Add another cluster. cluster_config["cluster_name"] = "random_name_2" provider_config = cluster_config["provider"] - node_provider_2 = _get_node_provider(provider_config, - cluster_config["cluster_name"]) + node_provider_2 = _get_node_provider( + provider_config, cluster_config["cluster_name"], use_cache=False) assert not node_provider_2.non_terminated_nodes({}) assert not node_provider_2.is_running(self.list_of_node_ips[1]) assert node_provider_2.is_terminated(self.list_of_node_ips[1]) @@ -212,8 +212,8 @@ class OnPremCoordinatorServerTest(unittest.TestCase): # Add another cluster (should fail because we only have two nodes). cluster_config["cluster_name"] = "random_name_3" provider_config = cluster_config["provider"] - node_provider_3 = _get_node_provider(provider_config, - cluster_config["cluster_name"]) + node_provider_3 = _get_node_provider( + provider_config, cluster_config["cluster_name"], use_cache=False) assert not node_provider_3.non_terminated_nodes(head_node_tags) head_node_tags[TAG_RAY_NODE_NAME] = "ray-{}-head".format( cluster_config["cluster_name"]) diff --git a/python/ray/tests/test_resource_demand_scheduler.py b/python/ray/tests/test_resource_demand_scheduler.py index 3e76c023e..680421e87 100644 --- a/python/ray/tests/test_resource_demand_scheduler.py +++ b/python/ray/tests/test_resource_demand_scheduler.py @@ -9,7 +9,8 @@ import copy import ray from ray.tests.test_autoscaler import SMALL_CLUSTER, MockProvider, \ MockProcessRunner -from ray.autoscaler._private.providers import _NODE_PROVIDERS +from ray.autoscaler._private.providers import (_NODE_PROVIDERS, + _clear_provider_cache) from ray.autoscaler._private.autoscaler import StandardAutoscaler from ray.autoscaler._private.load_metrics import LoadMetrics from ray.autoscaler._private.commands import get_or_create_head_node @@ -350,6 +351,7 @@ class AutoscalingTest(unittest.TestCase): def tearDown(self): self.provider = None del _NODE_PROVIDERS["mock"] + _clear_provider_cache() shutil.rmtree(self.tmpdir) ray.shutdown()