diff --git a/python/ray/autoscaler/aws/node_provider.py b/python/ray/autoscaler/aws/node_provider.py index 9a887df7a55f..ced3601cd15b 100644 --- a/python/ray/autoscaler/aws/node_provider.py +++ b/python/ray/autoscaler/aws/node_provider.py @@ -6,10 +6,28 @@ from botocore.config import Config from ray.autoscaler.node_provider import NodeProvider -from ray.autoscaler.tags import TAG_RAY_CLUSTER_NAME +from ray.autoscaler.tags import TAG_RAY_CLUSTER_NAME, TAG_RAY_NODE_NAME from ray.ray_constants import BOTO_MAX_RETRIES +def to_aws_format(tags): + """Convert the Ray node name tag to the AWS-specific 'Name' tag.""" + + if TAG_RAY_NODE_NAME in tags: + tags["Name"] = tags[TAG_RAY_NODE_NAME] + del tags[TAG_RAY_NODE_NAME] + return tags + + +def from_aws_format(tags): + """Convert the AWS-specific 'Name' tag to the Ray node name tag.""" + + if "Name" in tags: + tags[TAG_RAY_NODE_NAME] = tags["Name"] + del tags["Name"] + return tags + + class AWSNodeProvider(NodeProvider): def __init__(self, provider_config, cluster_name): NodeProvider.__init__(self, provider_config, cluster_name) @@ -26,6 +44,7 @@ def __init__(self, provider_config, cluster_name): self.external_ip_cache = {} def nodes(self, tag_filters): + tag_filters = to_aws_format(tag_filters) filters = [ { "Name": "instance-state-name", @@ -59,7 +78,7 @@ def node_tags(self, node_id): tags = {} for tag in node.tags: tags[tag["Key"]] = tag["Value"] - return tags + return from_aws_format(tags) def external_ip(self, node_id): if node_id in self.external_ip_cache: @@ -80,6 +99,7 @@ def internal_ip(self, node_id): return ip def set_node_tags(self, node_id, tags): + tags = to_aws_format(tags) node = self._node(node_id) tag_pairs = [] for k, v in tags.items(): @@ -90,6 +110,7 @@ def set_node_tags(self, node_id, tags): node.create_tags(Tags=tag_pairs) def create_node(self, node_config, tags, count): + tags = to_aws_format(tags) conf = node_config.copy() tag_pairs = [{ "Key": TAG_RAY_CLUSTER_NAME,