mirror of
https://github.com/wassname/ray.git
synced 2026-06-28 22:20:52 +08:00
@@ -6,10 +6,28 @@ import boto3
|
||||
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 @@ class AWSNodeProvider(NodeProvider):
|
||||
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 @@ class AWSNodeProvider(NodeProvider):
|
||||
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 @@ class AWSNodeProvider(NodeProvider):
|
||||
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 @@ class AWSNodeProvider(NodeProvider):
|
||||
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,
|
||||
|
||||
Reference in New Issue
Block a user