chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,137 @@
|
||||
{
|
||||
"$schema": "https://schema.management.azure.com/schemas/2019-04-01/deploymentTemplate.json#",
|
||||
"contentVersion": "1.0.0.0",
|
||||
"parameters": {
|
||||
"clusterId": {
|
||||
"type": "string",
|
||||
"metadata": {
|
||||
"description": "Unique string appended to resource names to isolate resources from different ray clusters."
|
||||
}
|
||||
},
|
||||
"subnet": {
|
||||
"type": "string",
|
||||
"metadata": {
|
||||
"description": "Subnet parameters."
|
||||
}
|
||||
},
|
||||
"msiName": {
|
||||
"type": "string",
|
||||
"metadata": {
|
||||
"description": "Managed service identity."
|
||||
}
|
||||
},
|
||||
"msiResourceGroup": {
|
||||
"type": "string",
|
||||
"metadata": {
|
||||
"description": "Managed service identity resource group."
|
||||
}
|
||||
},
|
||||
"createMsi": {
|
||||
"type": "bool",
|
||||
"defaultValue": "true"
|
||||
},
|
||||
"roleAssignmentGuid": {
|
||||
"type": "string",
|
||||
"metadata": {
|
||||
"description": "Deterministic resource name for the MSI role assignment."
|
||||
}
|
||||
}
|
||||
},
|
||||
"variables": {
|
||||
"contributor": "[subscriptionResourceId('Microsoft.Authorization/roleDefinitions', 'b24988ac-6180-42a0-ab88-20f7382dd24c')]",
|
||||
"location": "[resourceGroup().location]",
|
||||
"roleAssignmentName": "[concat('ray-', parameters('clusterId'), '-ra')]",
|
||||
"nsgName": "[concat('ray-', parameters('clusterId'), '-nsg')]",
|
||||
"nsg": "[resourceId('Microsoft.Network/networkSecurityGroups', variables('nsgName'))]",
|
||||
"vnetName": "[concat('ray-', parameters('clusterId'), '-vnet')]",
|
||||
"subnetName": "[concat('ray-', parameters('clusterId'), '-subnet')]"
|
||||
},
|
||||
"resources": [
|
||||
{
|
||||
"condition": "[parameters('createMsi')]",
|
||||
"type": "Microsoft.ManagedIdentity/userAssignedIdentities",
|
||||
"apiVersion": "2024-11-30",
|
||||
"location": "[variables('location')]",
|
||||
"name": "[parameters('msiName')]"
|
||||
},
|
||||
{
|
||||
"type": "Microsoft.Authorization/roleAssignments",
|
||||
"apiVersion": "2022-04-01",
|
||||
"name": "[parameters('roleAssignmentGuid')]",
|
||||
"properties": {
|
||||
"principalId": "[reference(resourceId(parameters('msiResourceGroup'), 'Microsoft.ManagedIdentity/userAssignedIdentities', parameters('msiName')), '2024-11-30').principalId]",
|
||||
"roleDefinitionId": "[variables('contributor')]",
|
||||
"scope": "[resourceGroup().id]",
|
||||
"principalType": "ServicePrincipal",
|
||||
"description": "[concat('Ray autoscaler cluster ', parameters('clusterId'))]"
|
||||
},
|
||||
"dependsOn": [
|
||||
"[parameters('msiName')]"
|
||||
]
|
||||
},
|
||||
{
|
||||
"type": "Microsoft.Network/networkSecurityGroups",
|
||||
"apiVersion": "2024-10-01",
|
||||
"name": "[variables('nsgName')]",
|
||||
"location": "[variables('location')]",
|
||||
"properties": {
|
||||
"securityRules": [
|
||||
{
|
||||
"name": "SSH",
|
||||
"properties": {
|
||||
"priority": 1000,
|
||||
"protocol": "Tcp",
|
||||
"access": "Allow",
|
||||
"direction": "Inbound",
|
||||
"sourceAddressPrefix": "*",
|
||||
"sourcePortRange": "*",
|
||||
"destinationAddressPrefix": "*",
|
||||
"destinationPortRange": "22"
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
},
|
||||
{
|
||||
"type": "Microsoft.Network/virtualNetworks",
|
||||
"apiVersion": "2024-10-01",
|
||||
"name": "[variables('vnetName')]",
|
||||
"location": "[variables('location')]",
|
||||
"properties": {
|
||||
"addressSpace": {
|
||||
"addressPrefixes": [
|
||||
"[parameters('subnet')]"
|
||||
]
|
||||
},
|
||||
"subnets": [
|
||||
{
|
||||
"name": "[variables('subnetName')]",
|
||||
"properties": {
|
||||
"addressPrefix": "[parameters('subnet')]",
|
||||
"networkSecurityGroup": {
|
||||
"id": "[variables('nsg')]"
|
||||
}
|
||||
}
|
||||
}
|
||||
]
|
||||
},
|
||||
"dependsOn": [
|
||||
"[variables('nsg')]"
|
||||
]
|
||||
}
|
||||
],
|
||||
"outputs": {
|
||||
"subnet": {
|
||||
"type": "string",
|
||||
"value": "[resourceId('Microsoft.Network/virtualNetworks/subnets', variables('vnetName'), variables('subnetName'))]"
|
||||
},
|
||||
"nsg": {
|
||||
"type": "string",
|
||||
"value": "[variables('nsg')]"
|
||||
},
|
||||
"msi": {
|
||||
"type": "string",
|
||||
"value": "[resourceId(parameters('msiResourceGroup'), 'Microsoft.ManagedIdentity/userAssignedIdentities', parameters('msiName'))]"
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,318 @@
|
||||
{
|
||||
"$schema": "https://schema.management.azure.com/schemas/2019-04-01/deploymentTemplate.json#",
|
||||
"contentVersion": "1.0.0.0",
|
||||
"parameters": {
|
||||
"vmName": {
|
||||
"type": "string",
|
||||
"metadata": {
|
||||
"description": "The name of you Virtual Machine."
|
||||
}
|
||||
},
|
||||
"adminUsername": {
|
||||
"type": "string",
|
||||
"metadata": {
|
||||
"description": "Username for the Virtual Machine."
|
||||
}
|
||||
},
|
||||
"publicKey": {
|
||||
"type": "securestring",
|
||||
"metadata": {
|
||||
"description": "SSH Key for the Virtual Machine"
|
||||
}
|
||||
},
|
||||
"imagePublisher": {
|
||||
"type": "string",
|
||||
"metadata": {
|
||||
"description": "The publisher of the VM image"
|
||||
},
|
||||
"defaultValue": ""
|
||||
},
|
||||
"imageOffer": {
|
||||
"type": "string",
|
||||
"metadata": {
|
||||
"description": "The offer of the VM image"
|
||||
},
|
||||
"defaultValue": ""
|
||||
},
|
||||
"imageSku": {
|
||||
"type": "string",
|
||||
"metadata": {
|
||||
"description": "The sku of the VM image"
|
||||
},
|
||||
"defaultValue": ""
|
||||
},
|
||||
"imageVersion": {
|
||||
"type": "string",
|
||||
"metadata": {
|
||||
"description": "The version of the VM image"
|
||||
},
|
||||
"defaultValue": ""
|
||||
},
|
||||
"imageId": {
|
||||
"type": "string",
|
||||
"metadata": {
|
||||
"description": "The resource id of the VM image. If provided, it will override the imagePublisher, imageOffer, imageSku and imageVersion parameters."
|
||||
},
|
||||
"defaultValue": ""
|
||||
},
|
||||
"vmSize": {
|
||||
"type": "string",
|
||||
"metadata": {
|
||||
"description": "The size of the VM"
|
||||
}
|
||||
},
|
||||
"vmTags": {
|
||||
"type": "object",
|
||||
"metadata": {
|
||||
"description": "Tags for the VM"
|
||||
}
|
||||
},
|
||||
"vmCount": {
|
||||
"type": "int",
|
||||
"metadata": {
|
||||
"description": "Number of VMs to deploy"
|
||||
}
|
||||
},
|
||||
"osDiskSize": {
|
||||
"type": "int",
|
||||
"defaultValue": 0,
|
||||
"metadata": {
|
||||
"description": "Size of the OS disk in GB"
|
||||
}
|
||||
},
|
||||
"provisionPublicIp": {
|
||||
"type": "bool",
|
||||
"defaultValue": true,
|
||||
"metadata": {
|
||||
"description": "If true creates a public ip"
|
||||
}
|
||||
},
|
||||
"priority": {
|
||||
"type": "string",
|
||||
"defaultValue": "Regular",
|
||||
"metadata": {
|
||||
"description": "Specifies the priority for the virtual machine."
|
||||
}
|
||||
},
|
||||
"evictionPolicy": {
|
||||
"type": "string",
|
||||
"defaultValue": "Delete",
|
||||
"metadata": {
|
||||
"description": "Specifies the eviction policy for the virtual machine."
|
||||
}
|
||||
},
|
||||
"billingProfile": {
|
||||
"type": "object",
|
||||
"defaultValue": {},
|
||||
"metadata": {
|
||||
"description": "Specifies the maximum price to pay for Azure Spot VM."
|
||||
}
|
||||
},
|
||||
"msi": {
|
||||
"type": "string",
|
||||
"metadata": {
|
||||
"description": "Managed service identity resource id."
|
||||
}
|
||||
},
|
||||
"nsg": {
|
||||
"type": "string",
|
||||
"metadata": {
|
||||
"description": "Network security group resource id."
|
||||
}
|
||||
},
|
||||
"subnet": {
|
||||
"type": "string",
|
||||
"metadata": {
|
||||
"descriptions": "Subnet resource id."
|
||||
}
|
||||
},
|
||||
"enableAcceleratedNetworking": {
|
||||
"type": "bool",
|
||||
"defaultValue": false,
|
||||
"metadata": {
|
||||
"descriptions": "Whether to enable accelerated networking."
|
||||
}
|
||||
},
|
||||
"zones": {
|
||||
"type": "array",
|
||||
"defaultValue": [],
|
||||
"metadata": {
|
||||
"description": "Availability zones for the virtual machine. If empty, no zones will be specified."
|
||||
}
|
||||
}
|
||||
},
|
||||
"variables": {
|
||||
"location": "[resourceGroup().location]",
|
||||
"useZones": "[greater(length(parameters('zones')), 0)]",
|
||||
"networkInterfaceNamePrivate": "[concat(parameters('vmName'), '-nic')]",
|
||||
"networkInterfaceNamePublic": "[concat(parameters('vmName'), '-nic-public')]",
|
||||
"networkInterfaceName": "[if(parameters('provisionPublicIp'), variables('networkInterfaceNamePublic'), variables('networkInterfaceNamePrivate'))]",
|
||||
"networkIpConfig": "[guid(resourceGroup().id, parameters('vmName'))]",
|
||||
"osDiskType": "Standard_LRS",
|
||||
"publicIpAddressName": "[concat(parameters('vmName'), '-ip')]"
|
||||
},
|
||||
"resources": [
|
||||
{
|
||||
"type": "Microsoft.Network/networkInterfaces",
|
||||
"apiVersion": "2024-07-01",
|
||||
"name": "[concat(variables('networkInterfaceNamePublic'), copyIndex())]",
|
||||
"location": "[variables('location')]",
|
||||
"dependsOn": [
|
||||
"[resourceId('Microsoft.Network/publicIpAddresses/', concat(variables('publicIpAddressName'), copyIndex()))]"
|
||||
],
|
||||
"copy": {
|
||||
"name": "NICPublicCopy",
|
||||
"count": "[parameters('vmCount')]"
|
||||
},
|
||||
"properties": {
|
||||
"ipConfigurations": [
|
||||
{
|
||||
"name": "[variables('networkIpConfig')]",
|
||||
"properties": {
|
||||
"subnet": {
|
||||
"id": "[parameters('subnet')]"
|
||||
},
|
||||
"privateIPAllocationMethod": "Dynamic",
|
||||
"publicIpAddress": {
|
||||
"id": "[resourceId('Microsoft.Network/publicIPAddresses', concat(variables('publicIPAddressName'), copyIndex()))]"
|
||||
}
|
||||
}
|
||||
}
|
||||
],
|
||||
"networkSecurityGroup": {
|
||||
"id": "[parameters('nsg')]"
|
||||
},
|
||||
"enableAcceleratedNetworking": "[parameters('enableAcceleratedNetworking')]"
|
||||
},
|
||||
"condition": "[parameters('provisionPublicIp')]"
|
||||
},
|
||||
{
|
||||
"type": "Microsoft.Network/networkInterfaces",
|
||||
"apiVersion": "2024-07-01",
|
||||
"name": "[concat(variables('networkInterfaceNamePrivate'), copyIndex())]",
|
||||
"location": "[variables('location')]",
|
||||
"copy": {
|
||||
"name": "NICPrivateCopy",
|
||||
"count": "[parameters('vmCount')]"
|
||||
},
|
||||
"properties": {
|
||||
"ipConfigurations": [
|
||||
{
|
||||
"name": "[variables('networkIpConfig')]",
|
||||
"properties": {
|
||||
"subnet": {
|
||||
"id": "[parameters('subnet')]"
|
||||
},
|
||||
"privateIPAllocationMethod": "Dynamic"
|
||||
}
|
||||
}
|
||||
],
|
||||
"networkSecurityGroup": {
|
||||
"id": "[parameters('nsg')]"
|
||||
},
|
||||
"enableAcceleratedNetworking": "[parameters('enableAcceleratedNetworking')]"
|
||||
},
|
||||
"condition": "[not(parameters('provisionPublicIp'))]"
|
||||
},
|
||||
{
|
||||
"type": "Microsoft.Network/publicIpAddresses",
|
||||
"apiVersion": "2024-07-01",
|
||||
"name": "[concat(variables('publicIpAddressName'), copyIndex())]",
|
||||
"location": "[variables('location')]",
|
||||
"properties": {
|
||||
"publicIpAllocationMethod": "Static",
|
||||
"publicIPAddressVersion": "IPv4"
|
||||
},
|
||||
"copy": {
|
||||
"name": "PublicIpCopy",
|
||||
"count": "[parameters('vmCount')]"
|
||||
},
|
||||
"zones": "[if(variables('useZones'), createArray(string(parameters('zones')[mod(copyIndex(), length(parameters('zones')))])), json('null'))]",
|
||||
"sku": {
|
||||
"name": "Standard",
|
||||
"tier": "Regional"
|
||||
},
|
||||
"condition": "[parameters('provisionPublicIp')]"
|
||||
},
|
||||
{
|
||||
"type": "Microsoft.Compute/virtualMachines",
|
||||
"apiVersion": "2024-07-01",
|
||||
"name": "[concat(parameters('vmName'), copyIndex())]",
|
||||
"location": "[variables('location')]",
|
||||
"dependsOn": [
|
||||
"[resourceId('Microsoft.Network/networkInterfaces/', concat(variables('networkInterfaceName'), copyIndex()))]"
|
||||
],
|
||||
"copy": {
|
||||
"name": "VmCopy",
|
||||
"count": "[parameters('vmCount')]"
|
||||
},
|
||||
"tags": "[parameters('vmTags')]",
|
||||
"zones": "[if(variables('useZones'), createArray(string(parameters('zones')[mod(copyIndex(), length(parameters('zones')))])), json('null'))]",
|
||||
"properties": {
|
||||
"hardwareProfile": {
|
||||
"vmSize": "[parameters('vmSize')]"
|
||||
},
|
||||
"storageProfile": {
|
||||
"osDisk": {
|
||||
"createOption": "fromImage",
|
||||
"diskSizeGB": "[if(equals(parameters('osDiskSize'), 0), json('null'), parameters('osDiskSize'))]",
|
||||
"managedDisk": {
|
||||
"storageAccountType": "[variables('osDiskType')]"
|
||||
}
|
||||
},
|
||||
"imageReference": "[if(equals(parameters('imageId'), ''), json(concat('{\"publisher\":\"', parameters('imagePublisher'), '\",\"offer\":\"', parameters('imageOffer'), '\",\"sku\":\"', parameters('imageSku'), '\",\"version\":\"', parameters('imageVersion'), '\"}')), json(concat('{\"id\":\"', parameters('imageId'), '\"}')))]"
|
||||
},
|
||||
"networkProfile": {
|
||||
"networkInterfaces": [
|
||||
{
|
||||
"id": "[resourceId('Microsoft.Network/networkInterfaces', concat(variables('networkInterfaceName'), copyIndex()))]"
|
||||
}
|
||||
]
|
||||
},
|
||||
"osProfile": {
|
||||
"computerName": "[concat(parameters('vmName'), copyIndex())]",
|
||||
"adminUsername": "[parameters('adminUsername')]",
|
||||
"adminPassword": "[parameters('publicKey')]",
|
||||
"linuxConfiguration": {
|
||||
"disablePasswordAuthentication": true,
|
||||
"ssh": {
|
||||
"publicKeys": [
|
||||
{
|
||||
"path": "[concat('/home/', parameters('adminUsername'), '/.ssh/authorized_keys')]",
|
||||
"keyData": "[parameters('publicKey')]"
|
||||
}
|
||||
]
|
||||
}
|
||||
}
|
||||
},
|
||||
"priority": "[parameters('priority')]",
|
||||
"evictionPolicy": "[if(equals(parameters('priority'), 'Spot'), parameters('evictionPolicy'), '')]",
|
||||
"billingProfile": "[parameters('billingProfile')]"
|
||||
},
|
||||
"identity": {
|
||||
"type": "UserAssigned",
|
||||
"userAssignedIdentities": {
|
||||
"[parameters('msi')]": {
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
],
|
||||
"outputs": {
|
||||
"publicIp": {
|
||||
"type": "array",
|
||||
"copy": {
|
||||
"count": "[parameters('vmCount')]",
|
||||
"input": "[reference(concat(variables('publicIpAddressName'), copyIndex())).ipAddress]"
|
||||
},
|
||||
"condition": "[parameters('provisionPublicIp')]"
|
||||
},
|
||||
"privateIp": {
|
||||
"type": "array",
|
||||
"copy": {
|
||||
"count": "[parameters('vmCount')]",
|
||||
"input": "[reference(concat(variables('networkInterfaceName'), copyIndex())).ipConfigurations[0].properties.privateIPAddress]"
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,586 @@
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import random
|
||||
import time
|
||||
from hashlib import md5, sha256
|
||||
from pathlib import Path
|
||||
from typing import Any, Callable
|
||||
from uuid import UUID
|
||||
|
||||
from azure.common.credentials import get_cli_profile
|
||||
from azure.core.exceptions import HttpResponseError, ResourceNotFoundError
|
||||
from azure.identity import AzureCliCredential
|
||||
from azure.mgmt.resource import ResourceManagementClient
|
||||
from azure.mgmt.resource.resources.models import DeploymentMode
|
||||
|
||||
from ray.autoscaler._private.util import (
|
||||
generate_rsa_key_pair,
|
||||
generate_ssh_key_name,
|
||||
generate_ssh_key_paths,
|
||||
)
|
||||
|
||||
# Built-in Azure Contributor role definition ID used for role assignments.
|
||||
CONTRIBUTOR_ROLE_DEFINITION_ID = "b24988ac-6180-42a0-ab88-20f7382dd24c"
|
||||
|
||||
UNIQUE_ID_LEN = 4
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def get_azure_sdk_function(client: Any, function_name: str) -> Callable:
|
||||
"""Retrieve a callable function from Azure SDK client object.
|
||||
|
||||
Newer versions of the various client SDKs renamed function names to
|
||||
have a begin_ prefix. This function supports both the old and new
|
||||
versions of the SDK by first trying the old name and falling back to
|
||||
the prefixed new name.
|
||||
"""
|
||||
func = getattr(
|
||||
client, function_name, getattr(client, f"begin_{function_name}", None)
|
||||
)
|
||||
if func is None:
|
||||
raise AttributeError(
|
||||
"'{obj}' object has no {func} or begin_{func} attribute".format(
|
||||
obj={client.__name__}, func=function_name
|
||||
)
|
||||
)
|
||||
return func
|
||||
|
||||
|
||||
def bootstrap_azure(config):
|
||||
config = _configure_key_pair(config)
|
||||
config = _configure_resource_group(config)
|
||||
return config
|
||||
|
||||
|
||||
def _configure_resource_group(config):
|
||||
# TODO: look at availability sets
|
||||
# https://docs.microsoft.com/en-us/azure/virtual-machines/windows/tutorial-availability-sets
|
||||
subscription_id = config["provider"].get("subscription_id")
|
||||
if subscription_id is None:
|
||||
subscription_id = get_cli_profile().get_subscription_id()
|
||||
resource_client = ResourceManagementClient(AzureCliCredential(), subscription_id)
|
||||
config["provider"]["subscription_id"] = subscription_id
|
||||
logger.info("Using subscription id: %s", subscription_id)
|
||||
|
||||
assert (
|
||||
"resource_group" in config["provider"]
|
||||
), "Provider config must include resource_group field"
|
||||
resource_group = config["provider"]["resource_group"]
|
||||
|
||||
assert (
|
||||
"location" in config["provider"]
|
||||
), "Provider config must include location field"
|
||||
params = {"location": config["provider"]["location"]}
|
||||
|
||||
if "tags" in config["provider"]:
|
||||
params["tags"] = config["provider"]["tags"]
|
||||
|
||||
logger.info("Creating/Updating resource group: %s", resource_group)
|
||||
rg_create_or_update = get_azure_sdk_function(
|
||||
client=resource_client.resource_groups, function_name="create_or_update"
|
||||
)
|
||||
rg_create_or_update(resource_group_name=resource_group, parameters=params)
|
||||
|
||||
# load the template file
|
||||
current_path = Path(__file__).parent
|
||||
template_path = current_path.joinpath("azure-config-template.json")
|
||||
with open(template_path, "r") as template_fp:
|
||||
template = json.load(template_fp)
|
||||
|
||||
logger.info("Using cluster name: %s", config["cluster_name"])
|
||||
|
||||
# set unique id for resources in this cluster
|
||||
unique_id = config["provider"].get("unique_id")
|
||||
if unique_id is None:
|
||||
hasher = sha256()
|
||||
hasher.update(config["provider"]["resource_group"].encode("utf-8"))
|
||||
unique_id = hasher.hexdigest()[:UNIQUE_ID_LEN]
|
||||
else:
|
||||
unique_id = str(unique_id)
|
||||
config["provider"]["unique_id"] = unique_id
|
||||
logger.info("Using unique id: %s", unique_id)
|
||||
cluster_id = "{}-{}".format(config["cluster_name"], unique_id)
|
||||
role_assignment_name = f"ray-{cluster_id}-ra"
|
||||
role_assignment_guid = _generate_arm_guid(role_assignment_name)
|
||||
role_assignment_resource_id = (
|
||||
f"/subscriptions/{subscription_id}/resourceGroups/{resource_group}/providers"
|
||||
f"/Microsoft.Authorization/roleAssignments/{role_assignment_guid}"
|
||||
)
|
||||
|
||||
subnet_mask = config["provider"].get("subnet_mask")
|
||||
if subnet_mask is None:
|
||||
# choose a random subnet, skipping most common value of 0
|
||||
random.seed(unique_id)
|
||||
subnet_mask = "10.{}.0.0/16".format(random.randint(1, 254))
|
||||
logger.info("Using subnet mask: %s", subnet_mask)
|
||||
|
||||
# Copy over properties from existing subnet.
|
||||
# Addresses issue (https://github.com/Azure/azure-quickstart-templates/issues/2786)
|
||||
# where existing subnet properties will get overwritten unless explicitly specified
|
||||
# during multiple deployments even if vnet/subnet do not change.
|
||||
# May eventually be fixed by passing empty subnet list if they already exist:
|
||||
# https://techcommunity.microsoft.com/t5/azure-networking-blog/azure-virtual-network-now-supports-updates-without-subnet/ba-p/4067952
|
||||
list_by_rg = get_azure_sdk_function(
|
||||
client=resource_client.resources, function_name="list_by_resource_group"
|
||||
)
|
||||
existing_vnets = list(
|
||||
list_by_rg(
|
||||
resource_group,
|
||||
f"substringof('{unique_id}', name) and "
|
||||
"resourceType eq 'Microsoft.Network/virtualNetworks'",
|
||||
)
|
||||
)
|
||||
if len(existing_vnets) > 0:
|
||||
vnid = existing_vnets[0].id
|
||||
get_by_id = get_azure_sdk_function(
|
||||
client=resource_client.resources, function_name="get_by_id"
|
||||
)
|
||||
|
||||
# Query for supported API versions for Microsoft.Network/virtualNetworks
|
||||
# because resource_client.DEFAULT_API_VERSION is not always supported.
|
||||
# (Example: "2024-11-01" is the default at the time of this writing)
|
||||
# Use "2024-10-01" as a fallback if we can't determine the latest stable version.
|
||||
vnet_api_version = "2024-10-01"
|
||||
try:
|
||||
# Get supported API versions for Microsoft.Network provider
|
||||
providers = resource_client.providers.get("Microsoft.Network")
|
||||
vnet_resource_type = next(
|
||||
(
|
||||
rt
|
||||
for rt in providers.resource_types
|
||||
if rt.resource_type == "virtualNetworks"
|
||||
),
|
||||
None,
|
||||
)
|
||||
if vnet_resource_type and vnet_resource_type.api_versions:
|
||||
stable_versions = [
|
||||
v for v in vnet_resource_type.api_versions if "preview" not in v
|
||||
]
|
||||
versions_to_consider = (
|
||||
stable_versions or vnet_resource_type.api_versions
|
||||
)
|
||||
vnet_api_version = sorted(versions_to_consider)[-1]
|
||||
logger.info(
|
||||
"Using API version: %s for virtualNetworks", vnet_api_version
|
||||
)
|
||||
else:
|
||||
logger.warning(
|
||||
"Could not determine supported API versions for virtualNetworks, using fallback version %s",
|
||||
vnet_api_version,
|
||||
)
|
||||
except Exception as e:
|
||||
logger.warning(
|
||||
"Failed to query Microsoft.Network provider: %s. Using fallback API version 2024-10-01",
|
||||
str(e),
|
||||
)
|
||||
|
||||
subnet = get_by_id(vnid, vnet_api_version).properties["subnets"][0]
|
||||
template_vnet = next(
|
||||
(
|
||||
rs
|
||||
for rs in template["resources"]
|
||||
if rs["type"] == "Microsoft.Network/virtualNetworks"
|
||||
),
|
||||
None,
|
||||
)
|
||||
if template_vnet is not None:
|
||||
template_subnets = template_vnet["properties"].get("subnets")
|
||||
if template_subnets is not None:
|
||||
template_subnets[0]["properties"].update(subnet["properties"])
|
||||
|
||||
# Get or create an MSI name and resource group.
|
||||
# Defaults to current resource group if not provided.
|
||||
use_existing_msi = _is_shared_msi(config["provider"])
|
||||
msi_resource_group = config["provider"].get("msi_resource_group", resource_group)
|
||||
msi_name = config["provider"].get("msi_name", f"ray-{cluster_id}-msi")
|
||||
logger.info(
|
||||
"Using msi_name: %s from msi_resource_group: %s", msi_name, msi_resource_group
|
||||
)
|
||||
|
||||
existing_principal_id = None
|
||||
if not use_existing_msi:
|
||||
# When creating a MSI for managing Azure resources, we first need to clean up
|
||||
# any role assignments from the previous MSI's principal ID to avoid
|
||||
# orphaned permissions when the MSI gets recreated with a new principal ID.
|
||||
# Role assignments cannot be updated, only created/removed, so cleanup is required.
|
||||
msi_resource_id = (
|
||||
f"/subscriptions/{subscription_id}/resourceGroups/{msi_resource_group}"
|
||||
f"/providers/Microsoft.ManagedIdentity/userAssignedIdentities/{msi_name}"
|
||||
)
|
||||
try:
|
||||
get_identity = get_azure_sdk_function(
|
||||
client=resource_client.resources, function_name="get_by_id"
|
||||
)
|
||||
existing_msi = get_identity(msi_resource_id, "2023-01-31")
|
||||
existing_principal_id = getattr(existing_msi, "properties", {}).get(
|
||||
"principalId"
|
||||
)
|
||||
except ResourceNotFoundError:
|
||||
existing_principal_id = None
|
||||
except Exception as exc:
|
||||
logger.warning(
|
||||
"Failed to query MSI %s for existing principal: %s",
|
||||
msi_name,
|
||||
exc,
|
||||
)
|
||||
|
||||
if existing_principal_id:
|
||||
logger.info(
|
||||
"Removing existing role assignments for MSI principal %s before recreation",
|
||||
existing_principal_id,
|
||||
)
|
||||
_delete_role_assignments_for_principal(
|
||||
resource_client,
|
||||
resource_group,
|
||||
existing_principal_id,
|
||||
)
|
||||
|
||||
delete_role_assignment = get_azure_sdk_function(
|
||||
client=resource_client.resources, function_name="delete_by_id"
|
||||
)
|
||||
get_role_assignment = get_azure_sdk_function(
|
||||
client=resource_client.resources, function_name="get_by_id"
|
||||
)
|
||||
|
||||
role_assignment_known_missing = False
|
||||
initial_query_failed = False
|
||||
try:
|
||||
get_role_assignment(
|
||||
role_assignment_resource_id,
|
||||
"2022-04-01",
|
||||
)
|
||||
except ResourceNotFoundError:
|
||||
role_assignment_known_missing = True
|
||||
logger.debug(
|
||||
"Role assignment %s not found before MSI creation; skipping deletion",
|
||||
role_assignment_guid,
|
||||
)
|
||||
except Exception as exc:
|
||||
logger.warning(
|
||||
"Failed to query role assignment %s before deletion: %s",
|
||||
role_assignment_guid,
|
||||
exc,
|
||||
)
|
||||
initial_query_failed = True
|
||||
|
||||
if not role_assignment_known_missing:
|
||||
try:
|
||||
delete_lro = delete_role_assignment(
|
||||
resource_id=role_assignment_resource_id,
|
||||
api_version="2022-04-01",
|
||||
)
|
||||
if hasattr(delete_lro, "wait"):
|
||||
delete_lro.wait()
|
||||
logger.info(
|
||||
"Deleted existing role assignment %s before recreating MSI",
|
||||
role_assignment_guid,
|
||||
)
|
||||
|
||||
if initial_query_failed:
|
||||
logger.debug(
|
||||
"Retrying verification for role assignment %s",
|
||||
role_assignment_guid,
|
||||
)
|
||||
|
||||
if not _wait_for_role_assignment_deletion(
|
||||
get_role_assignment,
|
||||
role_assignment_resource_id,
|
||||
role_assignment_guid,
|
||||
):
|
||||
logger.warning(
|
||||
"Role assignment %s not confirmed deleted",
|
||||
role_assignment_guid,
|
||||
)
|
||||
except ResourceNotFoundError:
|
||||
logger.debug(
|
||||
"Role assignment %s disappeared before deletion attempt; continuing",
|
||||
role_assignment_guid,
|
||||
)
|
||||
except Exception as e:
|
||||
logger.warning(
|
||||
"Failed to delete role assignment %s before MSI creation: %s",
|
||||
role_assignment_guid,
|
||||
e,
|
||||
)
|
||||
|
||||
parameters = {
|
||||
"properties": {
|
||||
"mode": DeploymentMode.incremental,
|
||||
"template": template,
|
||||
"parameters": {
|
||||
"subnet": {"value": subnet_mask},
|
||||
"clusterId": {"value": cluster_id},
|
||||
"msiName": {"value": msi_name},
|
||||
"msiResourceGroup": {"value": msi_resource_group},
|
||||
"createMsi": {"value": not use_existing_msi},
|
||||
"roleAssignmentGuid": {"value": role_assignment_guid},
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
create_or_update = get_azure_sdk_function(
|
||||
client=resource_client.deployments, function_name="create_or_update"
|
||||
)
|
||||
outputs = (
|
||||
create_or_update(
|
||||
resource_group_name=resource_group,
|
||||
deployment_name="ray-config",
|
||||
parameters=parameters,
|
||||
)
|
||||
.result()
|
||||
.properties.outputs
|
||||
)
|
||||
|
||||
# append output resource ids to be used with vm creation
|
||||
config["provider"]["msi"] = outputs["msi"]["value"]
|
||||
config["provider"]["nsg"] = outputs["nsg"]["value"]
|
||||
config["provider"]["subnet"] = outputs["subnet"]["value"]
|
||||
|
||||
return config
|
||||
|
||||
|
||||
def _configure_key_pair(config):
|
||||
"""
|
||||
Configure SSH keypair. Use user specified custom paths, otherwise,
|
||||
generate Ray-specific keypair in this format: "ray-autoscaler_azure_{region}_{resource_group}_{ssh_user}_{index}"
|
||||
"""
|
||||
ssh_user = config["auth"]["ssh_user"]
|
||||
public_key = None
|
||||
|
||||
# Check if user specified custom SSH key paths
|
||||
user_specified_private_key = "ssh_private_key" in config["auth"]
|
||||
user_specified_public_key = "ssh_public_key" in config["auth"]
|
||||
|
||||
# Validate that the user either specfied both keys or none, but not just one
|
||||
if user_specified_private_key != user_specified_public_key:
|
||||
if user_specified_private_key:
|
||||
missing_key, specified_key = "ssh_public_key", "ssh_private_key"
|
||||
else:
|
||||
missing_key, specified_key = "ssh_private_key", "ssh_public_key"
|
||||
raise ValueError(
|
||||
f"{specified_key} is specified but {missing_key} is missing. "
|
||||
"Both SSH key paths must be specified together, or omit both from "
|
||||
"your config to use auto-generated keys."
|
||||
)
|
||||
|
||||
if user_specified_private_key and user_specified_public_key:
|
||||
# User specified custom paths
|
||||
private_key_path = Path(config["auth"]["ssh_private_key"]).expanduser()
|
||||
public_key_path = Path(config["auth"]["ssh_public_key"]).expanduser()
|
||||
|
||||
# Validate that user-specified keys exist
|
||||
missing_keys = []
|
||||
if not private_key_path.is_file():
|
||||
missing_keys.append(f"ssh_private_key: {private_key_path}")
|
||||
if not public_key_path.is_file():
|
||||
missing_keys.append(f"ssh_public_key: {public_key_path}")
|
||||
|
||||
if missing_keys:
|
||||
raise ValueError(
|
||||
"SSH key files from config do not exist: {}. "
|
||||
"Please create the keys or remove the custom paths from your config "
|
||||
"to use auto-generated keys.".format(", ".join(missing_keys))
|
||||
)
|
||||
logger.info(
|
||||
"Using specified SSH keys from config: {} and {}".format(
|
||||
private_key_path, public_key_path
|
||||
)
|
||||
)
|
||||
|
||||
with open(public_key_path, "r") as f:
|
||||
public_key = f.read()
|
||||
else:
|
||||
# Generate Ray-specific keys
|
||||
region = config["provider"]["location"]
|
||||
resource_group = config["provider"]["resource_group"]
|
||||
|
||||
# Generate single deterministic key name for this configuration
|
||||
key_name = generate_ssh_key_name(
|
||||
"azure", None, region, resource_group, ssh_user
|
||||
)
|
||||
public_key_path, private_key_path = generate_ssh_key_paths(key_name)
|
||||
|
||||
# Check if this key pair already exists
|
||||
if os.path.exists(private_key_path) and os.path.exists(public_key_path):
|
||||
logger.info(
|
||||
"Using existing Ray-specific SSH keys: {} and {}".format(
|
||||
private_key_path, public_key_path
|
||||
)
|
||||
)
|
||||
with open(public_key_path, "r") as f:
|
||||
public_key = f.read()
|
||||
else:
|
||||
# Create a key pair since it doesn't exist locally
|
||||
logger.info(
|
||||
"Generating new Ray-specific SSH key pair at {} and {}".format(
|
||||
private_key_path, public_key_path
|
||||
)
|
||||
)
|
||||
os.makedirs(os.path.dirname(private_key_path), exist_ok=True)
|
||||
public_key, private_key = generate_rsa_key_pair()
|
||||
with open(
|
||||
private_key_path,
|
||||
"w",
|
||||
opener=lambda path, flags: os.open(path, flags, 0o600),
|
||||
) as f:
|
||||
f.write(private_key)
|
||||
with open(public_key_path, "w") as f:
|
||||
f.write(public_key)
|
||||
|
||||
assert os.path.exists(
|
||||
private_key_path
|
||||
), "Private key file {} not found for user {}".format(
|
||||
private_key_path, ssh_user
|
||||
)
|
||||
|
||||
config["auth"]["ssh_private_key"] = str(private_key_path)
|
||||
# Remove public key path because bootstrap config must only contain paths that exist on head node
|
||||
config["auth"].pop("ssh_public_key", None)
|
||||
|
||||
for node_type in config["available_node_types"].values():
|
||||
azure_arm_parameters = node_type["node_config"].setdefault(
|
||||
"azure_arm_parameters", {}
|
||||
)
|
||||
azure_arm_parameters["adminUsername"] = ssh_user
|
||||
azure_arm_parameters["publicKey"] = public_key
|
||||
|
||||
return config
|
||||
|
||||
|
||||
def _delete_role_assignments_for_principal(
|
||||
resource_client: ResourceManagementClient,
|
||||
resource_group: str,
|
||||
principal_id: str,
|
||||
) -> None:
|
||||
"""Delete all role assignments in the resource group for the given principal.
|
||||
|
||||
Uses the generic ResourceManagementClient so we avoid depending on
|
||||
azure-mgmt-authorization. All role assignments associated with the
|
||||
provided principal ID are removed.
|
||||
"""
|
||||
|
||||
if not principal_id:
|
||||
return
|
||||
|
||||
list_by_rg = get_azure_sdk_function(
|
||||
client=resource_client.resources, function_name="list_by_resource_group"
|
||||
)
|
||||
delete_role_assignment = get_azure_sdk_function(
|
||||
client=resource_client.resources, function_name="delete_by_id"
|
||||
)
|
||||
|
||||
try:
|
||||
assignments = list(
|
||||
list_by_rg(
|
||||
resource_group,
|
||||
"resourceType eq 'Microsoft.Authorization/roleAssignments'",
|
||||
)
|
||||
)
|
||||
logger.debug(
|
||||
"Found %d role assignments in resource group %s while cleaning up principal %s",
|
||||
len(assignments),
|
||||
resource_group,
|
||||
principal_id,
|
||||
)
|
||||
except HttpResponseError as exc:
|
||||
logger.warning(
|
||||
"Failed to enumerate role assignments for resource group %s: %s",
|
||||
resource_group,
|
||||
exc,
|
||||
)
|
||||
return
|
||||
|
||||
for assignment in assignments:
|
||||
properties = getattr(assignment, "properties", {}) or {}
|
||||
logger.debug(
|
||||
"Inspecting role assignment %s with principalId=%s roleDefinitionId=%s",
|
||||
getattr(assignment, "name", "<unknown>"),
|
||||
properties.get("principalId"),
|
||||
properties.get("roleDefinitionId"),
|
||||
)
|
||||
if properties.get("principalId") != principal_id:
|
||||
continue
|
||||
|
||||
try:
|
||||
delete_lro = delete_role_assignment(
|
||||
resource_id=assignment.id,
|
||||
api_version="2022-04-01",
|
||||
)
|
||||
if hasattr(delete_lro, "wait"):
|
||||
delete_lro.wait()
|
||||
logger.info(
|
||||
"Deleted existing role assignment %s for principal %s",
|
||||
assignment.name,
|
||||
principal_id,
|
||||
)
|
||||
except ResourceNotFoundError:
|
||||
logger.debug(
|
||||
"Role assignment %s not found while processing principal %s",
|
||||
assignment.name,
|
||||
principal_id,
|
||||
)
|
||||
except Exception as exc:
|
||||
logger.warning(
|
||||
"Failed to delete role assignment %s for principal %s: %s",
|
||||
assignment.name,
|
||||
principal_id,
|
||||
exc,
|
||||
)
|
||||
|
||||
|
||||
def _wait_for_role_assignment_deletion(
|
||||
get_role_assignment: Callable[..., Any],
|
||||
resource_id: str,
|
||||
role_assignment_guid: str,
|
||||
*,
|
||||
max_attempts: int = 10,
|
||||
delay_seconds: int = 2,
|
||||
) -> bool:
|
||||
"""Poll until a role assignment disappears after deletion.
|
||||
|
||||
Returns True if the assignment is confirmed deleted, False otherwise.
|
||||
Logs detailed progress to aid troubleshooting when transient errors occur.
|
||||
"""
|
||||
|
||||
for attempt in range(1, max_attempts + 1):
|
||||
try:
|
||||
get_role_assignment(resource_id, "2022-04-01")
|
||||
except ResourceNotFoundError:
|
||||
return True
|
||||
except Exception as exc: # noqa: BLE001
|
||||
logger.debug(
|
||||
"Attempt %d/%d to verify removal of role assignment %s failed: %s",
|
||||
attempt,
|
||||
max_attempts,
|
||||
role_assignment_guid,
|
||||
exc,
|
||||
)
|
||||
else:
|
||||
logger.debug(
|
||||
"Role assignment %s still present after deletion (attempt %d/%d)",
|
||||
role_assignment_guid,
|
||||
attempt,
|
||||
max_attempts,
|
||||
)
|
||||
|
||||
if attempt < max_attempts:
|
||||
time.sleep(delay_seconds)
|
||||
|
||||
return False
|
||||
|
||||
|
||||
def _generate_arm_guid(*values: Any) -> str:
|
||||
"""Replicates ARM template guid() function for creating deterministic IDs."""
|
||||
|
||||
concatenated = "".join(str(v) for v in values)
|
||||
return str(UUID(md5(concatenated.encode("utf-8")).hexdigest()))
|
||||
|
||||
|
||||
def _is_shared_msi(provider_config: dict) -> bool:
|
||||
"""Checks if the provider config specifies a shared/pre-existing MSI."""
|
||||
return (
|
||||
provider_config.get("msi_name") is not None
|
||||
and provider_config.get("msi_resource_group") is not None
|
||||
)
|
||||
@@ -0,0 +1,952 @@
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import time
|
||||
from concurrent.futures import Future, ThreadPoolExecutor
|
||||
from pathlib import Path
|
||||
from threading import RLock
|
||||
from typing import Dict, List, Optional
|
||||
from uuid import uuid4
|
||||
|
||||
from azure.common.credentials import get_cli_profile
|
||||
from azure.core.exceptions import ResourceNotFoundError
|
||||
from azure.identity import DefaultAzureCredential
|
||||
from azure.mgmt.compute import ComputeManagementClient
|
||||
from azure.mgmt.network import NetworkManagementClient
|
||||
from azure.mgmt.resource import ResourceManagementClient
|
||||
from azure.mgmt.resource.resources.models import DeploymentMode
|
||||
|
||||
from ray._common.usage.usage_lib import get_cloud_from_metadata_requests
|
||||
from ray.autoscaler._private._azure.config import (
|
||||
_delete_role_assignments_for_principal,
|
||||
_generate_arm_guid,
|
||||
_is_shared_msi,
|
||||
bootstrap_azure,
|
||||
get_azure_sdk_function,
|
||||
)
|
||||
from ray.autoscaler._private.constants import (
|
||||
AUTOSCALER_NODE_START_WAIT_S,
|
||||
AUTOSCALER_NODE_TERMINATE_WAIT_S,
|
||||
MAX_PARALLEL_SHUTDOWN_WORKERS,
|
||||
)
|
||||
from ray.autoscaler.node_provider import NodeProvider
|
||||
from ray.autoscaler.tags import (
|
||||
NODE_KIND_HEAD,
|
||||
TAG_RAY_CLUSTER_NAME,
|
||||
TAG_RAY_LAUNCH_CONFIG,
|
||||
TAG_RAY_NODE_KIND,
|
||||
TAG_RAY_NODE_NAME,
|
||||
TAG_RAY_USER_NODE_TYPE,
|
||||
)
|
||||
|
||||
VM_NAME_MAX_LEN = 64
|
||||
UNIQUE_ID_LEN = 4
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
azure_logger = logging.getLogger("azure.core.pipeline.policies.http_logging_policy")
|
||||
azure_logger.setLevel(logging.WARNING)
|
||||
|
||||
|
||||
def synchronized(f):
|
||||
def wrapper(self, *args, **kwargs):
|
||||
self.lock.acquire()
|
||||
try:
|
||||
return f(self, *args, **kwargs)
|
||||
finally:
|
||||
self.lock.release()
|
||||
|
||||
return wrapper
|
||||
|
||||
|
||||
class AzureNodeProvider(NodeProvider):
|
||||
"""Node Provider for Azure
|
||||
|
||||
This provider assumes Azure credentials are set by running ``az login``
|
||||
and the default subscription is configured through ``az account``
|
||||
or set in the ``provider`` field of the autoscaler configuration.
|
||||
|
||||
Nodes may be in one of three states: {pending, running, terminated}. Nodes
|
||||
appear immediately once started by ``create_node``, and transition
|
||||
immediately to terminated when ``terminate_node`` is called.
|
||||
"""
|
||||
|
||||
def __init__(self, provider_config, cluster_name):
|
||||
NodeProvider.__init__(self, provider_config, cluster_name)
|
||||
subscription_id = provider_config.get("subscription_id")
|
||||
if subscription_id is None:
|
||||
# Get subscription from logged in azure profile
|
||||
# if it isn't provided in the provider_config
|
||||
# so operations like `get-head-ip` will work
|
||||
subscription_id = get_cli_profile().get_subscription_id()
|
||||
logger.info(
|
||||
"subscription_id not found in provider config, falling back "
|
||||
f"to subscription_id from the logged in azure profile: {subscription_id}"
|
||||
)
|
||||
self.cache_stopped_nodes = provider_config.get("cache_stopped_nodes", True)
|
||||
|
||||
# Detect cloud environment to optimize Azure credential chain.
|
||||
# On non-Azure clouds (AWS, GCP), skip Azure-specific auth methods
|
||||
# (managed identity, workload identity) to avoid IMDS timeout delays / failures.
|
||||
detected_cloud = get_cloud_from_metadata_requests()
|
||||
on_azure = detected_cloud == "azure"
|
||||
|
||||
if on_azure:
|
||||
logger.info(
|
||||
"Initializing Azure node provider for Azure infrastructure "
|
||||
"running on Azure cloud environment"
|
||||
)
|
||||
else:
|
||||
logger.info(
|
||||
f"Initializing Azure node provider for Azure infrastructure "
|
||||
f"but detected this is running on a '{detected_cloud}' environment. "
|
||||
f"Skipping Azure-specific authentication methods to avoid timeouts."
|
||||
)
|
||||
|
||||
credential = DefaultAzureCredential(
|
||||
exclude_shared_token_cache_credential=True,
|
||||
exclude_managed_identity_credential=not on_azure,
|
||||
exclude_workload_identity_credential=not on_azure,
|
||||
)
|
||||
self.compute_client = ComputeManagementClient(credential, subscription_id)
|
||||
self.network_client = NetworkManagementClient(credential, subscription_id)
|
||||
self.resource_client = ResourceManagementClient(credential, subscription_id)
|
||||
|
||||
self.lock = RLock()
|
||||
|
||||
# cache node objects
|
||||
self.cached_nodes = {}
|
||||
|
||||
# Cache terminating node operations
|
||||
self.terminating_nodes: dict[str, Future] = {}
|
||||
self.termination_executor = ThreadPoolExecutor(
|
||||
max_workers=MAX_PARALLEL_SHUTDOWN_WORKERS
|
||||
)
|
||||
|
||||
@synchronized
|
||||
def _get_filtered_nodes(self, tag_filters):
|
||||
# add cluster name filter to only get nodes from this cluster
|
||||
cluster_tag_filters = {**tag_filters, TAG_RAY_CLUSTER_NAME: self.cluster_name}
|
||||
|
||||
def match_tags(tags):
|
||||
for k, v in cluster_tag_filters.items():
|
||||
if tags.get(k) != v:
|
||||
return False
|
||||
return True
|
||||
|
||||
vms = self.compute_client.virtual_machines.list(
|
||||
resource_group_name=self.provider_config["resource_group"]
|
||||
)
|
||||
|
||||
nodes = [self._extract_metadata(vm) for vm in vms]
|
||||
self.cached_nodes = {node["name"]: node for node in nodes}
|
||||
|
||||
# Update terminating nodes list by removing nodes that
|
||||
# have finished termination.
|
||||
self.terminating_nodes = {
|
||||
k: v for k, v in self.terminating_nodes.items() if not v.done()
|
||||
}
|
||||
|
||||
return {
|
||||
k: v
|
||||
for k, v in self.cached_nodes.items()
|
||||
if v["tags"] is not None and match_tags(v["tags"])
|
||||
}
|
||||
|
||||
def _extract_metadata(self, vm):
|
||||
# get tags
|
||||
metadata = {"name": vm.name, "tags": vm.tags, "status": ""}
|
||||
|
||||
# get status
|
||||
resource_group = self.provider_config["resource_group"]
|
||||
try:
|
||||
instance = self.compute_client.virtual_machines.instance_view(
|
||||
resource_group_name=resource_group, vm_name=vm.name
|
||||
).as_dict()
|
||||
except ResourceNotFoundError:
|
||||
return metadata
|
||||
|
||||
for status in instance["statuses"]:
|
||||
# If ProvisioningState is "failed" (e.g.,
|
||||
# ProvisioningState/failed/RetryableError), we can get a third
|
||||
# string here, so we need to limit to the first two outputs.
|
||||
code, state = status["code"].split("/")[:2]
|
||||
# skip provisioning status
|
||||
if code == "PowerState":
|
||||
metadata["status"] = state
|
||||
break
|
||||
|
||||
# get ip data
|
||||
nic_id = vm.network_profile.network_interfaces[0].id
|
||||
metadata["nic_name"] = nic_id.split("/")[-1]
|
||||
nic = self.network_client.network_interfaces.get(
|
||||
resource_group_name=resource_group,
|
||||
network_interface_name=metadata["nic_name"],
|
||||
)
|
||||
ip_config = nic.ip_configurations[0]
|
||||
|
||||
# Get public IP if not using internal IPs or if this is the
|
||||
# head node and use_external_head_ip is True
|
||||
if not self.provider_config.get("use_internal_ips", False) or (
|
||||
self.provider_config.get("use_external_head_ip", False)
|
||||
and metadata["tags"][TAG_RAY_NODE_KIND] == NODE_KIND_HEAD
|
||||
):
|
||||
public_ip_id = ip_config.public_ip_address.id
|
||||
metadata["public_ip_name"] = public_ip_id.split("/")[-1]
|
||||
public_ip = self.network_client.public_ip_addresses.get(
|
||||
resource_group_name=resource_group,
|
||||
public_ip_address_name=metadata["public_ip_name"],
|
||||
)
|
||||
metadata["external_ip"] = public_ip.ip_address
|
||||
|
||||
metadata["internal_ip"] = ip_config.private_ip_address
|
||||
|
||||
return metadata
|
||||
|
||||
def _get_zones_for_vm_size(self, vm_size, location):
|
||||
"""Get usable availability zones for a given VM size in a specific location."""
|
||||
try:
|
||||
# Note: Azure ResourceSKUs API filters don't work reliably(?), so we query all SKUs
|
||||
# and filter in code. Each SKU object represents one location for the VM size.
|
||||
skus = self.compute_client.resource_skus.list()
|
||||
|
||||
for sku in skus:
|
||||
if sku.name == vm_size and sku.location_info:
|
||||
# Each SKU object represents one location, check if it matches our target
|
||||
for location_info in sku.location_info:
|
||||
if location_info.location.lower() == location.lower():
|
||||
zones = location_info.zones if location_info.zones else []
|
||||
logger.debug(
|
||||
f"Found {vm_size} in {location} with zones: {zones}"
|
||||
)
|
||||
return sorted(zones)
|
||||
|
||||
logger.warning(f"No zones found for {vm_size} in {location}")
|
||||
return [] # No zones available for this VM size
|
||||
except Exception as e:
|
||||
logger.warning(
|
||||
f"Failed to get zones for VM size {vm_size} in location {location}: {str(e)}"
|
||||
)
|
||||
return []
|
||||
|
||||
def _parse_availability_zones(
|
||||
self, availability_zone_config: Optional[str]
|
||||
) -> Optional[List[str]]:
|
||||
"""Parse availability_zone configuration from comma-separated string format.
|
||||
|
||||
Args:
|
||||
availability_zone_config: Can be:
|
||||
- String: comma-separated zones like "1,2,3"
|
||||
- "none": explicitly disable zones
|
||||
- "auto": let Azure automatically pick zones
|
||||
- None: no zones specified (defaults to letting Azure pick)
|
||||
|
||||
Returns:
|
||||
List of zone strings, or None if zones explicitly disabled, or [] if auto/unspecified
|
||||
"""
|
||||
if availability_zone_config is None:
|
||||
return [] # Auto - let Azure pick
|
||||
|
||||
# Handle string format (AWS-style comma-separated)
|
||||
if isinstance(availability_zone_config, str):
|
||||
# Strip whitespace and split by comma
|
||||
zones = [zone.strip() for zone in availability_zone_config.split(",")]
|
||||
|
||||
# Handle special cases (case-insensitive)
|
||||
if len(zones) == 1:
|
||||
zone_lower = zones[0].lower()
|
||||
if zone_lower in ["none", "null"]:
|
||||
return None # Explicitly disabled
|
||||
elif zone_lower == "auto":
|
||||
return [] # Auto - let Azure pick
|
||||
|
||||
# Handle empty string or whitespace-only
|
||||
if not zones or all(not zone for zone in zones):
|
||||
return [] # Auto - let Azure pick
|
||||
return zones
|
||||
|
||||
# Unsupported format
|
||||
raise ValueError(
|
||||
f"availability_zone must be a string, got {type(availability_zone_config).__name__}: {availability_zone_config!r}"
|
||||
)
|
||||
|
||||
def _validate_zones_for_node_pool(self, zones, location, vm_size):
|
||||
"""Validate that the specified zones are available for the given VM size in the location."""
|
||||
# Special case: zones=None means explicitly disabled availability zones
|
||||
if zones is None:
|
||||
logger.info(
|
||||
"Zones explicitly disabled with 'none' - will create VM without an availability zone"
|
||||
)
|
||||
return None # Special return value to indicate "no zones by choice"
|
||||
|
||||
vm_zones = self._get_zones_for_vm_size(vm_size, location)
|
||||
|
||||
available_zones = set(vm_zones)
|
||||
if not available_zones:
|
||||
logger.warning("No zones available for this VM size and location")
|
||||
return []
|
||||
|
||||
if zones:
|
||||
requested_zones = {str(z) for z in zones}
|
||||
intersection = sorted(available_zones.intersection(requested_zones))
|
||||
return intersection
|
||||
|
||||
return sorted(available_zones)
|
||||
|
||||
def stopped_nodes(self, tag_filters):
|
||||
"""Return a list of stopped node ids filtered by the specified tags dict."""
|
||||
nodes = self._get_filtered_nodes(tag_filters=tag_filters)
|
||||
return [k for k, v in nodes.items() if v["status"].startswith("deallocat")]
|
||||
|
||||
def non_terminated_nodes(self, tag_filters: Dict[str, str]) -> List[str]:
|
||||
"""Return a list of node ids filtered by the specified tags dict.
|
||||
|
||||
This list must not include terminated nodes. For performance reasons,
|
||||
providers are allowed to cache the result of a call to nodes() to
|
||||
serve single-node queries (e.g. is_running(node_id)). This means that
|
||||
nodes() must be called again to refresh results.
|
||||
|
||||
Args:
|
||||
tag_filters: Tags to filter nodes by.
|
||||
|
||||
Returns:
|
||||
List of non-terminated node ids matching the tag filters.
|
||||
|
||||
Examples:
|
||||
>>> from ray.autoscaler.tags import TAG_RAY_NODE_KIND
|
||||
>>> provider = ... # doctest: +SKIP
|
||||
>>> provider.non_terminated_nodes( # doctest: +SKIP
|
||||
... {TAG_RAY_NODE_KIND: "worker"})
|
||||
["node-1", "node-2"]
|
||||
"""
|
||||
nodes = self._get_filtered_nodes(tag_filters=tag_filters)
|
||||
return [
|
||||
k
|
||||
for k, v in nodes.items()
|
||||
if not v["status"].startswith("deallocat") or k in self.terminating_nodes
|
||||
]
|
||||
|
||||
def is_running(self, node_id):
|
||||
"""Return whether the specified node is running."""
|
||||
# always get current status
|
||||
node = self._get_node(node_id=node_id)
|
||||
return node["status"] == "running"
|
||||
|
||||
def is_terminated(self, node_id):
|
||||
"""Return whether the specified node is terminated."""
|
||||
# always get current status
|
||||
node = self._get_node(node_id=node_id)
|
||||
return node["status"].startswith("deallocat")
|
||||
|
||||
def node_tags(self, node_id):
|
||||
"""Returns the tags of the given node (string dict)."""
|
||||
return self._get_cached_node(node_id=node_id)["tags"]
|
||||
|
||||
def external_ip(self, node_id):
|
||||
"""Returns the external ip of the given node."""
|
||||
ip = (
|
||||
self._get_cached_node(node_id=node_id)["external_ip"]
|
||||
or self._get_node(node_id=node_id)["external_ip"]
|
||||
)
|
||||
return ip
|
||||
|
||||
def internal_ip(self, node_id):
|
||||
"""Returns the internal ip (Ray ip) of the given node."""
|
||||
ip = (
|
||||
self._get_cached_node(node_id=node_id)["internal_ip"]
|
||||
or self._get_node(node_id=node_id)["internal_ip"]
|
||||
)
|
||||
return ip
|
||||
|
||||
def create_node(self, node_config, tags, count):
|
||||
resource_group = self.provider_config["resource_group"]
|
||||
|
||||
if self.cache_stopped_nodes:
|
||||
VALIDITY_TAGS = [
|
||||
TAG_RAY_CLUSTER_NAME,
|
||||
TAG_RAY_NODE_KIND,
|
||||
TAG_RAY_LAUNCH_CONFIG,
|
||||
TAG_RAY_USER_NODE_TYPE,
|
||||
]
|
||||
filters = {tag: tags[tag] for tag in VALIDITY_TAGS if tag in tags}
|
||||
reuse_nodes = self.stopped_nodes(filters)[:count]
|
||||
logger.info(
|
||||
f"Reusing nodes {list(reuse_nodes)}. "
|
||||
"To disable reuse, set `cache_stopped_nodes: False` "
|
||||
"under `provider` in the cluster configuration.",
|
||||
)
|
||||
start = get_azure_sdk_function(
|
||||
client=self.compute_client.virtual_machines, function_name="start"
|
||||
)
|
||||
for node_id in reuse_nodes:
|
||||
start(resource_group_name=resource_group, vm_name=node_id).wait()
|
||||
self.set_node_tags(node_id, tags)
|
||||
count -= len(reuse_nodes)
|
||||
|
||||
if count:
|
||||
self._create_node(node_config, tags, count)
|
||||
|
||||
def _create_node(self, node_config, tags, count):
|
||||
"""Creates a number of nodes within the namespace."""
|
||||
resource_group = self.provider_config["resource_group"]
|
||||
location = self.provider_config["location"]
|
||||
vm_size = node_config["azure_arm_parameters"]["vmSize"]
|
||||
|
||||
# Determine availability zones with precedence: node-level > provider-level
|
||||
# Check for "availability_zone" field in node config first
|
||||
node_availability_zone = node_config.get("azure_arm_parameters", {}).get(
|
||||
"availability_zone"
|
||||
)
|
||||
# Then check provider-level "availability_zone"
|
||||
provider_availability_zone = self.provider_config.get("availability_zone")
|
||||
|
||||
requested_zones = []
|
||||
zone_source = "default"
|
||||
|
||||
# Precedence: node availability_zone > provider availability_zone
|
||||
if node_availability_zone is not None:
|
||||
requested_zones = self._parse_availability_zones(node_availability_zone)
|
||||
zone_source = "node config availability_zone"
|
||||
elif provider_availability_zone is not None:
|
||||
requested_zones = self._parse_availability_zones(provider_availability_zone)
|
||||
zone_source = "provider availability_zone"
|
||||
|
||||
logger.info(f"Requested zones from {zone_source}: {requested_zones}")
|
||||
|
||||
# Get actually available zones for this VM size
|
||||
available_zones = self._validate_zones_for_node_pool(
|
||||
requested_zones, location, vm_size
|
||||
)
|
||||
|
||||
# Handle explicit zone disabling
|
||||
zones_explicitly_disabled = available_zones is None
|
||||
|
||||
if requested_zones and not zones_explicitly_disabled and not available_zones:
|
||||
raise ValueError(
|
||||
f"No available zones for VM size {vm_size} in {location}. "
|
||||
f"Requested: {requested_zones}, but none are available for this VM size."
|
||||
)
|
||||
|
||||
# load the template file
|
||||
current_path = Path(__file__).parent
|
||||
template_path = current_path.joinpath("azure-vm-template.json")
|
||||
with open(template_path, "r") as template_fp:
|
||||
template = json.load(template_fp)
|
||||
|
||||
# get the tags
|
||||
config_tags = node_config.get("tags", {}).copy()
|
||||
config_tags.update(tags)
|
||||
config_tags[TAG_RAY_CLUSTER_NAME] = self.cluster_name
|
||||
|
||||
deployment_name = "{node}-{unique_id}-{vm_id}".format(
|
||||
node=config_tags.get(TAG_RAY_NODE_NAME, "node"),
|
||||
unique_id=self.provider_config["unique_id"],
|
||||
vm_id=uuid4().hex[:UNIQUE_ID_LEN],
|
||||
)[:VM_NAME_MAX_LEN]
|
||||
|
||||
template_params = node_config["azure_arm_parameters"].copy()
|
||||
# Remove availability_zone from template params since ARM template expects "zones"
|
||||
template_params.pop("availability_zone", None)
|
||||
# Use deployment_name for the vmName template parameter since
|
||||
# the template will append copyIndex() for each VM that gets created
|
||||
# to guarantee uniqueness.
|
||||
template_params["vmName"] = deployment_name
|
||||
# Provision public IP if not using internal IPs or if this is the
|
||||
# head node and use_external_head_ip is True
|
||||
template_params["provisionPublicIp"] = not self.provider_config.get(
|
||||
"use_internal_ips", False
|
||||
) or (
|
||||
self.provider_config.get("use_external_head_ip", False)
|
||||
and config_tags[TAG_RAY_NODE_KIND] == NODE_KIND_HEAD
|
||||
)
|
||||
template_params["vmTags"] = config_tags
|
||||
template_params["vmCount"] = count
|
||||
template_params["msi"] = self.provider_config["msi"]
|
||||
template_params["nsg"] = self.provider_config["nsg"]
|
||||
template_params["subnet"] = self.provider_config["subnet"]
|
||||
|
||||
# Add zone information based on availability and requested zones
|
||||
if zones_explicitly_disabled:
|
||||
# User explicitly disabled zones with ["None"]
|
||||
template_params["zones"] = []
|
||||
logger.info(
|
||||
f"Creating {count} VMs with zones explicitly disabled (no availability zone)"
|
||||
)
|
||||
elif available_zones:
|
||||
# Pass the list of available zones to the template
|
||||
template_params["zones"] = available_zones
|
||||
logger.info(
|
||||
f"Creating {count} VMs, distributed across availability zones: {available_zones}"
|
||||
)
|
||||
else:
|
||||
# For non-zonal deployments (no zones available), use empty array
|
||||
template_params["zones"] = []
|
||||
logger.info(f"Creating {count} VMs without specific availability zone")
|
||||
|
||||
parameters = {
|
||||
"properties": {
|
||||
"mode": DeploymentMode.incremental,
|
||||
"template": template,
|
||||
"parameters": {
|
||||
key: {"value": value} for key, value in template_params.items()
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
# TODO: we could get the private/public ips back directly
|
||||
create_or_update = get_azure_sdk_function(
|
||||
client=self.resource_client.deployments,
|
||||
function_name="create_or_update",
|
||||
)
|
||||
create_or_update(
|
||||
resource_group_name=resource_group,
|
||||
deployment_name=deployment_name,
|
||||
parameters=parameters,
|
||||
).wait(timeout=AUTOSCALER_NODE_START_WAIT_S)
|
||||
|
||||
@synchronized
|
||||
def set_node_tags(self, node_id, tags):
|
||||
"""Sets the tag values (string dict) for the specified node."""
|
||||
node_tags = self._get_cached_node(node_id)["tags"]
|
||||
node_tags.update(tags)
|
||||
update = get_azure_sdk_function(
|
||||
client=self.compute_client.virtual_machines, function_name="update"
|
||||
)
|
||||
update(
|
||||
resource_group_name=self.provider_config["resource_group"],
|
||||
vm_name=node_id,
|
||||
parameters={"tags": node_tags},
|
||||
)
|
||||
self.cached_nodes[node_id]["tags"] = node_tags
|
||||
|
||||
def terminate_node(self, node_id):
|
||||
"""Terminates the specified node. This will delete the VM and
|
||||
associated resources (NIC, IP, Storage) for the specified node."""
|
||||
|
||||
resource_group = self.provider_config["resource_group"]
|
||||
|
||||
if self.cache_stopped_nodes:
|
||||
try:
|
||||
# stop machine and leave all resources
|
||||
logger.info(
|
||||
f"Stopping instance {node_id}"
|
||||
"(to fully terminate instead, "
|
||||
"set `cache_stopped_nodes: False` "
|
||||
"under `provider` in the cluster configuration)"
|
||||
)
|
||||
stop = get_azure_sdk_function(
|
||||
client=self.compute_client.virtual_machines,
|
||||
function_name="deallocate",
|
||||
)
|
||||
stop(resource_group_name=resource_group, vm_name=node_id)
|
||||
except Exception as e:
|
||||
logger.warning("Failed to stop VM: {}".format(e))
|
||||
|
||||
# If node_id is in terminating nodes dict, it's already terminating
|
||||
# Otherwise, kick off termination and add it to the dict
|
||||
elif node_id not in self.terminating_nodes:
|
||||
self.terminating_nodes[node_id] = self.termination_executor.submit(
|
||||
self._delete_node_and_resources, resource_group, node_id
|
||||
)
|
||||
|
||||
def _delete_node_and_resources(self, resource_group, node_id):
|
||||
try:
|
||||
vm = self.compute_client.virtual_machines.get(
|
||||
resource_group_name=resource_group, vm_name=node_id
|
||||
)
|
||||
except ResourceNotFoundError as e:
|
||||
# Node no longer exists
|
||||
logger.warning("Failed to delete VM: {}".format(e))
|
||||
return
|
||||
|
||||
# Gather dependent disks
|
||||
disks = set()
|
||||
if vm.storage_profile is not None and vm.storage_profile.data_disks is not None:
|
||||
for d in vm.storage_profile.data_disks:
|
||||
if d.name is not None:
|
||||
disks.add(d.name)
|
||||
if (
|
||||
vm.storage_profile is not None
|
||||
and vm.storage_profile.os_disk is not None
|
||||
and vm.storage_profile.os_disk.name is not None
|
||||
):
|
||||
disks.add(vm.storage_profile.os_disk.name)
|
||||
|
||||
# Gather dependent NICs and public IPs
|
||||
nics = set()
|
||||
ips = set()
|
||||
if (
|
||||
vm.network_profile is not None
|
||||
and vm.network_profile.network_interfaces is not None
|
||||
):
|
||||
for nint in vm.network_profile.network_interfaces:
|
||||
if nint.id is not None:
|
||||
nic_name = nint.id.split("/")[-1]
|
||||
nics.add(nic_name)
|
||||
# Get public IP if not using internal IPs or if this is the
|
||||
# head node and use_external_head_ip is True
|
||||
if not self.provider_config.get("use_internal_ips", False) or (
|
||||
self.provider_config.get("use_external_head_ip", False)
|
||||
and vm.tags[TAG_RAY_NODE_KIND] == NODE_KIND_HEAD
|
||||
):
|
||||
nic = self.network_client.network_interfaces.get(
|
||||
resource_group_name=resource_group,
|
||||
network_interface_name=nic_name,
|
||||
)
|
||||
if nic.ip_configurations is not None:
|
||||
for ipc in nic.ip_configurations:
|
||||
if ipc.public_ip_address.id is not None:
|
||||
ips.add(ipc.public_ip_address.id.split("/")[-1])
|
||||
|
||||
# Delete VM
|
||||
st = time.monotonic()
|
||||
delete = get_azure_sdk_function(
|
||||
client=self.compute_client.virtual_machines,
|
||||
function_name="delete",
|
||||
)
|
||||
try:
|
||||
delete(resource_group_name=resource_group, vm_name=node_id).wait(
|
||||
timeout=AUTOSCALER_NODE_TERMINATE_WAIT_S
|
||||
)
|
||||
except Exception as e:
|
||||
logger.warning("Failed to delete VM: {}".format(e))
|
||||
|
||||
# Delete disks (no need to wait for these, but gather the LROs for end)
|
||||
disk_lros = []
|
||||
delete = get_azure_sdk_function(
|
||||
client=self.compute_client.disks, function_name="delete"
|
||||
)
|
||||
for d in disks:
|
||||
try:
|
||||
disk_lros.append(
|
||||
delete(
|
||||
resource_group_name=resource_group,
|
||||
disk_name=d,
|
||||
)
|
||||
)
|
||||
except Exception as e:
|
||||
logger.warning("Failed to delete disk: {}".format(e))
|
||||
|
||||
# Delete NICs
|
||||
nic_lros = []
|
||||
delete = get_azure_sdk_function(
|
||||
client=self.network_client.network_interfaces, function_name="delete"
|
||||
)
|
||||
for n in nics:
|
||||
try:
|
||||
nic_lros.append(
|
||||
delete(
|
||||
resource_group_name=resource_group,
|
||||
network_interface_name=n,
|
||||
)
|
||||
)
|
||||
except Exception as e:
|
||||
logger.warning("Failed to delete NIC: {}".format(e))
|
||||
|
||||
while (
|
||||
not all(nlro.done() for nlro in nic_lros)
|
||||
and (time.monotonic() - st) < AUTOSCALER_NODE_TERMINATE_WAIT_S
|
||||
):
|
||||
time.sleep(0.1)
|
||||
|
||||
# Delete Public IPs
|
||||
delete = get_azure_sdk_function(
|
||||
client=self.network_client.public_ip_addresses,
|
||||
function_name="delete",
|
||||
)
|
||||
ip_lros = []
|
||||
for ip in ips:
|
||||
try:
|
||||
ip_lros.append(
|
||||
delete(
|
||||
resource_group_name=resource_group,
|
||||
public_ip_address_name=ip,
|
||||
)
|
||||
)
|
||||
except Exception as e:
|
||||
logger.warning("Failed to delete public IP: {}".format(e))
|
||||
|
||||
while (
|
||||
not all(dlro.done() for dlro in disk_lros)
|
||||
and (time.monotonic() - st) < AUTOSCALER_NODE_TERMINATE_WAIT_S
|
||||
):
|
||||
time.sleep(0.1)
|
||||
while (
|
||||
not all(iplro.done() for iplro in ip_lros)
|
||||
and (time.monotonic() - st) < AUTOSCALER_NODE_TERMINATE_WAIT_S
|
||||
):
|
||||
time.sleep(0.1)
|
||||
|
||||
def cleanup_cluster_resources(self):
|
||||
"""Delete shared cluster infrastructure (MSI, NSG, Subnet, VNet)."""
|
||||
|
||||
resource_group = self.provider_config["resource_group"]
|
||||
|
||||
msi_principal_id = self._cleanup_managed_identity(
|
||||
resource_group, self.provider_config.get("msi")
|
||||
)
|
||||
|
||||
subnet_id = self.provider_config.get("subnet")
|
||||
vnet_name = self._cleanup_subnet(resource_group, subnet_id)
|
||||
|
||||
nsg_id = self.provider_config.get("nsg")
|
||||
self._cleanup_nsg(resource_group, nsg_id)
|
||||
|
||||
self._cleanup_vnet(resource_group, subnet_id, vnet_name)
|
||||
|
||||
self._cleanup_role_assignments(resource_group, msi_principal_id)
|
||||
|
||||
self._prune_provider_config_entries()
|
||||
self._cleanup_config_cache()
|
||||
|
||||
@staticmethod
|
||||
def _get_resource_name_from_id(resource_id: Optional[str]) -> Optional[str]:
|
||||
if resource_id:
|
||||
return resource_id.split("/")[-1]
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def _retry_delete(delete_fn, max_retries: int = 5, initial_delay: int = 2):
|
||||
"""Retry a delete operation with exponential backoff."""
|
||||
|
||||
delay = initial_delay
|
||||
for attempt in range(max_retries):
|
||||
try:
|
||||
return delete_fn()
|
||||
except Exception as exc: # noqa: BLE001
|
||||
error_msg = str(exc)
|
||||
if "InUse" in error_msg and attempt < max_retries - 1:
|
||||
logger.info(
|
||||
"Resource still in use, retrying in %ss (attempt %s/%s)...",
|
||||
delay,
|
||||
attempt + 1,
|
||||
max_retries,
|
||||
)
|
||||
time.sleep(delay)
|
||||
delay *= 2
|
||||
else:
|
||||
raise
|
||||
|
||||
def _cleanup_managed_identity(
|
||||
self, resource_group: str, msi_id: Optional[str]
|
||||
) -> Optional[str]:
|
||||
if not msi_id:
|
||||
return None
|
||||
|
||||
msi_name = self._get_resource_name_from_id(msi_id)
|
||||
if not msi_name:
|
||||
return None
|
||||
|
||||
msi_principal_id: Optional[str] = None
|
||||
try:
|
||||
get_identity = get_azure_sdk_function(
|
||||
client=self.resource_client.resources,
|
||||
function_name="get_by_id",
|
||||
)
|
||||
existing_msi = get_identity(msi_id, "2023-01-31")
|
||||
msi_principal_id = getattr(existing_msi, "properties", {}).get(
|
||||
"principalId"
|
||||
)
|
||||
except ResourceNotFoundError:
|
||||
msi_principal_id = None
|
||||
except Exception as exc: # noqa: BLE001
|
||||
logger.warning(
|
||||
"Failed to query MSI %s for principal ID prior to deletion: %s",
|
||||
msi_name,
|
||||
exc,
|
||||
)
|
||||
|
||||
if _is_shared_msi(self.provider_config):
|
||||
return msi_principal_id
|
||||
|
||||
try:
|
||||
logger.info("Deleting Managed Service Identity: %s", msi_name)
|
||||
delete = get_azure_sdk_function(
|
||||
client=self.resource_client.resources,
|
||||
function_name="delete_by_id",
|
||||
)
|
||||
delete(resource_id=msi_id, api_version="2023-01-31").wait()
|
||||
logger.info("Successfully deleted MSI: %s", msi_name)
|
||||
except ResourceNotFoundError:
|
||||
logger.info("MSI %s not found, may have been already deleted", msi_name)
|
||||
except Exception as exc: # noqa: BLE001
|
||||
logger.warning("Failed to delete MSI %s: %s", msi_name, exc)
|
||||
|
||||
return msi_principal_id
|
||||
|
||||
def _cleanup_subnet(
|
||||
self, resource_group: str, subnet_id: Optional[str]
|
||||
) -> Optional[str]:
|
||||
if not subnet_id:
|
||||
return None
|
||||
|
||||
subnet_name = self._get_resource_name_from_id(subnet_id)
|
||||
vnet_name: Optional[str] = None
|
||||
|
||||
if subnet_id and "/virtualNetworks/" in subnet_id:
|
||||
parts = subnet_id.split("/")
|
||||
vnet_idx = parts.index("virtualNetworks")
|
||||
if vnet_idx + 1 < len(parts):
|
||||
vnet_name = parts[vnet_idx + 1]
|
||||
|
||||
if not subnet_name or not vnet_name:
|
||||
return None
|
||||
|
||||
try:
|
||||
logger.info("Deleting Subnet: %s in VNet: %s", subnet_name, vnet_name)
|
||||
|
||||
def delete_subnet():
|
||||
delete = get_azure_sdk_function(
|
||||
client=self.network_client.subnets, function_name="delete"
|
||||
)
|
||||
delete(
|
||||
resource_group_name=resource_group,
|
||||
virtual_network_name=vnet_name,
|
||||
subnet_name=subnet_name,
|
||||
).wait()
|
||||
|
||||
self._retry_delete(delete_subnet)
|
||||
logger.info("Successfully deleted Subnet: %s", subnet_name)
|
||||
except ResourceNotFoundError:
|
||||
logger.info(
|
||||
"Subnet %s not found, may have been already deleted", subnet_name
|
||||
)
|
||||
except Exception as exc: # noqa: BLE001
|
||||
logger.warning("Failed to delete Subnet %s: %s", subnet_name, exc)
|
||||
|
||||
return vnet_name
|
||||
|
||||
def _cleanup_nsg(self, resource_group: str, nsg_id: Optional[str]) -> None:
|
||||
if not nsg_id:
|
||||
return
|
||||
|
||||
nsg_name = self._get_resource_name_from_id(nsg_id)
|
||||
if not nsg_name:
|
||||
return
|
||||
|
||||
try:
|
||||
logger.info("Deleting Network Security Group: %s", nsg_name)
|
||||
|
||||
def delete_nsg():
|
||||
delete = get_azure_sdk_function(
|
||||
client=self.network_client.network_security_groups,
|
||||
function_name="delete",
|
||||
)
|
||||
delete(
|
||||
resource_group_name=resource_group,
|
||||
network_security_group_name=nsg_name,
|
||||
).wait()
|
||||
|
||||
self._retry_delete(delete_nsg)
|
||||
logger.info("Successfully deleted NSG: %s", nsg_name)
|
||||
except ResourceNotFoundError:
|
||||
logger.info("NSG %s not found, may have been already deleted", nsg_name)
|
||||
except Exception as exc: # noqa: BLE001
|
||||
logger.warning("Failed to delete NSG %s: %s", nsg_name, exc)
|
||||
|
||||
def _cleanup_vnet(
|
||||
self,
|
||||
resource_group: str,
|
||||
subnet_id: Optional[str],
|
||||
vnet_name: Optional[str],
|
||||
) -> None:
|
||||
if not subnet_id or not vnet_name:
|
||||
return
|
||||
|
||||
try:
|
||||
logger.info("Deleting Virtual Network: %s", vnet_name)
|
||||
|
||||
def delete_vnet():
|
||||
delete = get_azure_sdk_function(
|
||||
client=self.network_client.virtual_networks,
|
||||
function_name="delete",
|
||||
)
|
||||
delete(
|
||||
resource_group_name=resource_group,
|
||||
virtual_network_name=vnet_name,
|
||||
).wait()
|
||||
|
||||
self._retry_delete(delete_vnet)
|
||||
logger.info("Successfully deleted VNet: %s", vnet_name)
|
||||
except ResourceNotFoundError:
|
||||
logger.info("VNet %s not found, may have been already deleted", vnet_name)
|
||||
except Exception as exc: # noqa: BLE001
|
||||
logger.warning("Failed to delete VNet %s: %s", vnet_name, exc)
|
||||
|
||||
def _cleanup_role_assignments(
|
||||
self, resource_group: str, msi_principal_id: Optional[str]
|
||||
) -> None:
|
||||
subscription_id = self.provider_config.get("subscription_id")
|
||||
unique_id = self.provider_config.get("unique_id")
|
||||
if not subscription_id or not unique_id:
|
||||
return
|
||||
|
||||
cluster_id = f"{self.cluster_name}-{unique_id}"
|
||||
role_assignment_name = f"ray-{cluster_id}-ra"
|
||||
role_assignment_guid = _generate_arm_guid(role_assignment_name)
|
||||
role_assignment_id = (
|
||||
f"/subscriptions/{subscription_id}/resourceGroups/{resource_group}/providers"
|
||||
f"/Microsoft.Authorization/roleAssignments/{role_assignment_guid}"
|
||||
)
|
||||
|
||||
if msi_principal_id:
|
||||
_delete_role_assignments_for_principal(
|
||||
self.resource_client, resource_group, msi_principal_id
|
||||
)
|
||||
|
||||
delete_role_assignment = get_azure_sdk_function(
|
||||
client=self.resource_client.resources, function_name="delete_by_id"
|
||||
)
|
||||
try:
|
||||
delete_lro = delete_role_assignment(
|
||||
resource_id=role_assignment_id,
|
||||
api_version="2022-04-01",
|
||||
)
|
||||
if hasattr(delete_lro, "wait"):
|
||||
delete_lro.wait()
|
||||
logger.info(
|
||||
"Deleted role assignment %s for cluster %s",
|
||||
role_assignment_guid,
|
||||
self.cluster_name,
|
||||
)
|
||||
except ResourceNotFoundError:
|
||||
logger.debug(
|
||||
"Role assignment %s not found during cleanup", role_assignment_guid
|
||||
)
|
||||
except Exception as exc: # noqa: BLE001
|
||||
logger.warning(
|
||||
"Failed to delete role assignment %s: %s",
|
||||
role_assignment_guid,
|
||||
exc,
|
||||
)
|
||||
|
||||
def _prune_provider_config_entries(self) -> None:
|
||||
for key in ("msi", "nsg", "subnet"):
|
||||
self.provider_config.pop(key, None)
|
||||
|
||||
def _cleanup_config_cache(self) -> None:
|
||||
cache_path = self.provider_config.get("_config_cache_path")
|
||||
if not cache_path:
|
||||
return
|
||||
|
||||
try:
|
||||
if os.path.exists(cache_path):
|
||||
os.remove(cache_path)
|
||||
logger.info(
|
||||
"Deleted cached Ray config at %s after resource cleanup",
|
||||
cache_path,
|
||||
)
|
||||
except Exception as exc: # noqa: BLE001
|
||||
logger.warning("Failed to delete cached Ray config %s: %s", cache_path, exc)
|
||||
finally:
|
||||
self.provider_config.pop("_config_cache_path", None)
|
||||
|
||||
def _get_node(self, node_id):
|
||||
self._get_filtered_nodes({}) # Side effect: updates cache
|
||||
return self.cached_nodes[node_id]
|
||||
|
||||
def _get_cached_node(self, node_id):
|
||||
return self.cached_nodes.get(node_id) or self._get_node(node_id=node_id)
|
||||
|
||||
@staticmethod
|
||||
def bootstrap_config(cluster_config):
|
||||
return bootstrap_azure(cluster_config)
|
||||
@@ -0,0 +1,116 @@
|
||||
import logging
|
||||
import os
|
||||
import stat
|
||||
|
||||
from ray.autoscaler._private.aliyun.utils import AcsClient
|
||||
|
||||
# instance status
|
||||
PENDING = "Pending"
|
||||
RUNNING = "Running"
|
||||
STARTING = "Starting"
|
||||
STOPPING = "Stopping"
|
||||
STOPPED = "Stopped"
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def bootstrap_aliyun(config):
|
||||
# print(config["provider"])
|
||||
# create vpc
|
||||
_get_or_create_vpc(config)
|
||||
# create security group id
|
||||
_get_or_create_security_group(config)
|
||||
# create vswitch
|
||||
_get_or_create_vswitch(config)
|
||||
# create key pair
|
||||
_get_or_import_key_pair(config)
|
||||
# print(config["provider"])
|
||||
return config
|
||||
|
||||
|
||||
def _client(config):
|
||||
return AcsClient(
|
||||
access_key=config["provider"].get("access_key"),
|
||||
access_key_secret=config["provider"].get("access_key_secret"),
|
||||
region_id=config["provider"]["region"],
|
||||
max_retries=1,
|
||||
)
|
||||
|
||||
|
||||
def _get_or_create_security_group(config):
|
||||
cli = _client(config)
|
||||
security_groups = cli.describe_security_groups(vpc_id=config["provider"]["vpc_id"])
|
||||
if security_groups is not None and len(security_groups) > 0:
|
||||
config["provider"]["security_group_id"] = security_groups[0]["SecurityGroupId"]
|
||||
return config
|
||||
|
||||
security_group_id = cli.create_security_group(vpc_id=config["provider"]["vpc_id"])
|
||||
|
||||
for rule in config["provider"].get("security_group_rule", {}):
|
||||
cli.authorize_security_group(
|
||||
security_group_id=security_group_id,
|
||||
port_range=rule["port_range"],
|
||||
source_cidr_ip=rule["source_cidr_ip"],
|
||||
ip_protocol=rule["ip_protocol"],
|
||||
)
|
||||
config["provider"]["security_group_id"] = security_group_id
|
||||
return
|
||||
|
||||
|
||||
def _get_or_create_vpc(config):
|
||||
cli = _client(config)
|
||||
vpcs = cli.describe_vpcs()
|
||||
if vpcs is not None and len(vpcs) > 0:
|
||||
config["provider"]["vpc_id"] = vpcs[0].get("VpcId")
|
||||
return
|
||||
|
||||
vpc_id = cli.create_vpc()
|
||||
if vpc_id is not None:
|
||||
config["provider"]["vpc_id"] = vpc_id
|
||||
|
||||
|
||||
def _get_or_create_vswitch(config):
|
||||
cli = _client(config)
|
||||
vswitches = cli.describe_v_switches(vpc_id=config["provider"]["vpc_id"])
|
||||
if vswitches is not None and len(vswitches) > 0:
|
||||
config["provider"]["v_switch_id"] = vswitches[0].get("VSwitchId")
|
||||
return
|
||||
|
||||
v_switch_id = cli.create_v_switch(
|
||||
vpc_id=config["provider"]["vpc_id"],
|
||||
zone_id=config["provider"]["zone_id"],
|
||||
cidr_block=config["provider"]["cidr_block"],
|
||||
)
|
||||
|
||||
if v_switch_id is not None:
|
||||
config["provider"]["v_switch_id"] = v_switch_id
|
||||
|
||||
|
||||
def _get_or_import_key_pair(config):
|
||||
cli = _client(config)
|
||||
key_name = config["provider"].get("key_name", "ray")
|
||||
key_path = os.path.expanduser("~/.ssh/{}".format(key_name))
|
||||
keypairs = cli.describe_key_pairs(key_pair_name=key_name)
|
||||
|
||||
if keypairs is not None and len(keypairs) > 0:
|
||||
if "ssh_private_key" not in config["auth"]:
|
||||
logger.info(
|
||||
"{} keypair exists, use {} as local ssh key".format(key_name, key_path)
|
||||
)
|
||||
config["auth"]["ssh_private_key"] = key_path
|
||||
else:
|
||||
if "ssh_private_key" not in config["auth"]:
|
||||
# create new keypair
|
||||
resp = cli.create_key_pair(key_pair_name=key_name)
|
||||
if resp is not None:
|
||||
with open(key_path, "w+") as f:
|
||||
f.write(resp.get("PrivateKeyBody"))
|
||||
os.chmod(key_path, stat.S_IRUSR)
|
||||
config["auth"]["ssh_private_key"] = key_path
|
||||
else:
|
||||
public_key_file = config["auth"]["ssh_private_key"] + ".pub"
|
||||
# create new keypair, from local file
|
||||
with open(public_key_file) as f:
|
||||
public_key = f.readline().strip("\n")
|
||||
cli.import_key_pair(key_pair_name=key_name, public_key_body=public_key)
|
||||
return
|
||||
@@ -0,0 +1,324 @@
|
||||
import logging
|
||||
import random
|
||||
import threading
|
||||
import time
|
||||
from collections import defaultdict
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
from ray.autoscaler._private.aliyun.config import (
|
||||
PENDING,
|
||||
RUNNING,
|
||||
STOPPED,
|
||||
STOPPING,
|
||||
bootstrap_aliyun,
|
||||
)
|
||||
from ray.autoscaler._private.aliyun.utils import AcsClient
|
||||
from ray.autoscaler._private.cli_logger import cli_logger
|
||||
from ray.autoscaler._private.constants import BOTO_MAX_RETRIES
|
||||
from ray.autoscaler._private.log_timer import LogTimer
|
||||
from ray.autoscaler.node_provider import NodeProvider
|
||||
from ray.autoscaler.tags import (
|
||||
TAG_RAY_CLUSTER_NAME,
|
||||
TAG_RAY_LAUNCH_CONFIG,
|
||||
TAG_RAY_NODE_KIND,
|
||||
TAG_RAY_NODE_NAME,
|
||||
TAG_RAY_NODE_STATUS,
|
||||
TAG_RAY_USER_NODE_TYPE,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
TAG_BATCH_DELAY = 1
|
||||
STOPPING_NODE_DELAY = 1
|
||||
|
||||
|
||||
class AliyunNodeProvider(NodeProvider):
|
||||
def __init__(self, provider_config, cluster_name):
|
||||
NodeProvider.__init__(self, provider_config, cluster_name)
|
||||
self.cache_stopped_nodes = provider_config.get("cache_stopped_nodes", True)
|
||||
self.acs = AcsClient(
|
||||
access_key=provider_config["access_key"],
|
||||
access_key_secret=provider_config["access_key_secret"],
|
||||
region_id=provider_config["region"],
|
||||
max_retries=BOTO_MAX_RETRIES,
|
||||
)
|
||||
|
||||
# Try availability zones round-robin, starting from random offset
|
||||
self.subnet_idx = random.randint(0, 100)
|
||||
|
||||
# Tags that we believe to actually be on the node.
|
||||
self.tag_cache = {}
|
||||
# Tags that we will soon upload.
|
||||
self.tag_cache_pending = defaultdict(dict)
|
||||
# Number of threads waiting for a batched tag update.
|
||||
self.batch_thread_count = 0
|
||||
self.batch_update_done = threading.Event()
|
||||
self.batch_update_done.set()
|
||||
self.ready_for_new_batch = threading.Event()
|
||||
self.ready_for_new_batch.set()
|
||||
self.tag_cache_lock = threading.Lock()
|
||||
self.count_lock = threading.Lock()
|
||||
|
||||
# Cache of node objects from the last nodes() call. This avoids
|
||||
# excessive DescribeInstances requests.
|
||||
self.cached_nodes = {}
|
||||
|
||||
def non_terminated_nodes(self, tag_filters: Dict[str, str]) -> List[str]:
|
||||
tags = [
|
||||
{
|
||||
"Key": TAG_RAY_CLUSTER_NAME,
|
||||
"Value": self.cluster_name,
|
||||
},
|
||||
]
|
||||
for k, v in tag_filters.items():
|
||||
tags.append(
|
||||
{
|
||||
"Key": k,
|
||||
"Value": v,
|
||||
}
|
||||
)
|
||||
|
||||
instances = self.acs.describe_instances(tags=tags)
|
||||
non_terminated_instance = []
|
||||
for instance in instances:
|
||||
if instance.get("Status") == RUNNING or instance.get("Status") == PENDING:
|
||||
non_terminated_instance.append(instance.get("InstanceId"))
|
||||
self.cached_nodes[instance.get("InstanceId")] = instance
|
||||
return non_terminated_instance
|
||||
|
||||
def is_running(self, node_id: str) -> bool:
|
||||
instances = self.acs.describe_instances(instance_ids=[node_id])
|
||||
if instances is not None:
|
||||
instance = instances[0]
|
||||
return instance.get("Status") == "Running"
|
||||
cli_logger.error("Invalid node id: %s", node_id)
|
||||
return False
|
||||
|
||||
def is_terminated(self, node_id: str) -> bool:
|
||||
instances = self.acs.describe_instances(instance_ids=[node_id])
|
||||
if instances is not None:
|
||||
assert len(instances) == 1
|
||||
instance = instances[0]
|
||||
return instance.get("Status") == "Stopped"
|
||||
cli_logger.error("Invalid node id: %s", node_id)
|
||||
return False
|
||||
|
||||
def node_tags(self, node_id: str) -> Dict[str, str]:
|
||||
instances = self.acs.describe_instances(instance_ids=[node_id])
|
||||
if instances is not None:
|
||||
assert len(instances) == 1
|
||||
instance = instances[0]
|
||||
if instance.get("Tags") is not None:
|
||||
node_tags = dict()
|
||||
for tag in instance.get("Tags").get("Tag"):
|
||||
node_tags[tag.get("TagKey")] = tag.get("TagValue")
|
||||
return node_tags
|
||||
return dict()
|
||||
|
||||
def external_ip(self, node_id: str) -> str:
|
||||
while True:
|
||||
instances = self.acs.describe_instances(instance_ids=[node_id])
|
||||
if instances is not None:
|
||||
assert len(instances)
|
||||
instance = instances[0]
|
||||
if (
|
||||
instance.get("PublicIpAddress") is not None
|
||||
and instance.get("PublicIpAddress").get("IpAddress") is not None
|
||||
):
|
||||
if len(instance.get("PublicIpAddress").get("IpAddress")) > 0:
|
||||
return instance.get("PublicIpAddress").get("IpAddress")[0]
|
||||
cli_logger.error("PublicIpAddress attribute is not exist. %s" % instance)
|
||||
time.sleep(STOPPING_NODE_DELAY)
|
||||
|
||||
def internal_ip(self, node_id: str) -> str:
|
||||
while True:
|
||||
instances = self.acs.describe_instances(instance_ids=[node_id])
|
||||
if instances is not None:
|
||||
assert len(instances) == 1
|
||||
instance = instances[0]
|
||||
if (
|
||||
instance.get("VpcAttributes") is not None
|
||||
and instance.get("VpcAttributes").get("PrivateIpAddress")
|
||||
is not None
|
||||
and len(
|
||||
instance.get("VpcAttributes")
|
||||
.get("PrivateIpAddress")
|
||||
.get("IpAddress")
|
||||
)
|
||||
> 0
|
||||
):
|
||||
return (
|
||||
instance.get("VpcAttributes")
|
||||
.get("PrivateIpAddress")
|
||||
.get("IpAddress")[0]
|
||||
)
|
||||
cli_logger.error("InnerIpAddress attribute is not exist. %s" % instance)
|
||||
time.sleep(STOPPING_NODE_DELAY)
|
||||
|
||||
def set_node_tags(self, node_id: str, tags: Dict[str, str]) -> None:
|
||||
is_batching_thread = False
|
||||
with self.tag_cache_lock:
|
||||
if not self.tag_cache_pending:
|
||||
is_batching_thread = True
|
||||
# Wait for threads in the last batch to exit
|
||||
self.ready_for_new_batch.wait()
|
||||
self.ready_for_new_batch.clear()
|
||||
self.batch_update_done.clear()
|
||||
self.tag_cache_pending[node_id].update(tags)
|
||||
|
||||
if is_batching_thread:
|
||||
time.sleep(TAG_BATCH_DELAY)
|
||||
with self.tag_cache_lock:
|
||||
self._update_node_tags()
|
||||
self.batch_update_done.set()
|
||||
|
||||
with self.count_lock:
|
||||
self.batch_thread_count += 1
|
||||
self.batch_update_done.wait()
|
||||
|
||||
with self.count_lock:
|
||||
self.batch_thread_count -= 1
|
||||
if self.batch_thread_count == 0:
|
||||
self.ready_for_new_batch.set()
|
||||
|
||||
def _update_node_tags(self):
|
||||
batch_updates = defaultdict(list)
|
||||
|
||||
for node_id, tags in self.tag_cache_pending.items():
|
||||
for x in tags.items():
|
||||
batch_updates[x].append(node_id)
|
||||
self.tag_cache[node_id] = tags
|
||||
|
||||
self.tag_cache_pending = defaultdict(dict)
|
||||
|
||||
self._create_tags(batch_updates)
|
||||
|
||||
def _create_tags(self, batch_updates):
|
||||
|
||||
for (k, v), node_ids in batch_updates.items():
|
||||
m = "Set tag {}={} on {}".format(k, v, node_ids)
|
||||
with LogTimer("AliyunNodeProvider: {}".format(m)):
|
||||
if k == TAG_RAY_NODE_NAME:
|
||||
k = "Name"
|
||||
|
||||
self.acs.tag_resource(node_ids, [{"Key": k, "Value": v}])
|
||||
|
||||
def create_node(
|
||||
self, node_config: Dict[str, Any], tags: Dict[str, str], count: int
|
||||
) -> Optional[Dict[str, Any]]:
|
||||
filter_tags = [
|
||||
{
|
||||
"Key": TAG_RAY_CLUSTER_NAME,
|
||||
"Value": self.cluster_name,
|
||||
},
|
||||
{"Key": TAG_RAY_NODE_KIND, "Value": tags[TAG_RAY_NODE_KIND]},
|
||||
{"Key": TAG_RAY_USER_NODE_TYPE, "Value": tags[TAG_RAY_USER_NODE_TYPE]},
|
||||
{"Key": TAG_RAY_LAUNCH_CONFIG, "Value": tags[TAG_RAY_LAUNCH_CONFIG]},
|
||||
{"Key": TAG_RAY_NODE_NAME, "Value": tags[TAG_RAY_NODE_NAME]},
|
||||
]
|
||||
|
||||
reused_nodes_dict = {}
|
||||
if self.cache_stopped_nodes:
|
||||
reuse_nodes_candidate = self.acs.describe_instances(tags=filter_tags)
|
||||
if reuse_nodes_candidate:
|
||||
with cli_logger.group("Stopping instances to reuse"):
|
||||
reuse_node_ids = []
|
||||
for node in reuse_nodes_candidate:
|
||||
node_id = node.get("InstanceId")
|
||||
status = node.get("Status")
|
||||
if status != STOPPING and status != STOPPED:
|
||||
continue
|
||||
if status == STOPPING:
|
||||
# wait for node stopped
|
||||
while (
|
||||
self.acs.describe_instances(instance_ids=[node_id])[
|
||||
0
|
||||
].get("Status")
|
||||
== STOPPING
|
||||
):
|
||||
logging.info("wait for %s stop" % node_id)
|
||||
time.sleep(STOPPING_NODE_DELAY)
|
||||
# logger.info("reuse %s" % node_id)
|
||||
reuse_node_ids.append(node_id)
|
||||
reused_nodes_dict[node.get("InstanceId")] = node
|
||||
self.acs.start_instance(node_id)
|
||||
self.tag_cache[node_id] = node.get("Tags")
|
||||
self.set_node_tags(node_id, tags)
|
||||
if len(reuse_node_ids) == count:
|
||||
break
|
||||
count -= len(reuse_node_ids)
|
||||
|
||||
created_nodes_dict = {}
|
||||
if count > 0:
|
||||
filter_tags.append(
|
||||
{"Key": TAG_RAY_NODE_STATUS, "Value": tags[TAG_RAY_NODE_STATUS]}
|
||||
)
|
||||
instance_id_sets = self.acs.run_instances(
|
||||
instance_type=node_config["InstanceType"],
|
||||
image_id=node_config["ImageId"],
|
||||
tags=filter_tags,
|
||||
amount=count,
|
||||
vswitch_id=self.provider_config["v_switch_id"],
|
||||
security_group_id=self.provider_config["security_group_id"],
|
||||
key_pair_name=self.provider_config["key_name"],
|
||||
)
|
||||
instances = self.acs.describe_instances(instance_ids=instance_id_sets)
|
||||
|
||||
if instances is not None:
|
||||
for instance in instances:
|
||||
created_nodes_dict[instance.get("InstanceId")] = instance
|
||||
|
||||
all_created_nodes = reused_nodes_dict
|
||||
all_created_nodes.update(created_nodes_dict)
|
||||
return all_created_nodes
|
||||
|
||||
def terminate_node(self, node_id: str) -> None:
|
||||
logger.info("terminate node: %s" % node_id)
|
||||
if self.cache_stopped_nodes:
|
||||
logger.info(
|
||||
"Stopping instance {} (to terminate instead, "
|
||||
"set `cache_stopped_nodes: False` "
|
||||
"under `provider` in the cluster configuration)"
|
||||
).format(node_id)
|
||||
self.acs.stop_instance(node_id)
|
||||
else:
|
||||
self.acs.delete_instance(node_id)
|
||||
|
||||
def terminate_nodes(self, node_ids: List[str]) -> None:
|
||||
if not node_ids:
|
||||
return
|
||||
if self.cache_stopped_nodes:
|
||||
logger.info(
|
||||
"Stopping instances {} (to terminate instead, "
|
||||
"set `cache_stopped_nodes: False` "
|
||||
"under `provider` in the cluster configuration)".format(node_ids)
|
||||
)
|
||||
|
||||
self.acs.stop_instances(node_ids)
|
||||
else:
|
||||
self.acs.delete_instances(node_ids)
|
||||
|
||||
def _get_node(self, node_id):
|
||||
"""Refresh and get info for this node, updating the cache."""
|
||||
self.non_terminated_nodes({}) # Side effect: updates cache
|
||||
|
||||
if node_id in self.cached_nodes:
|
||||
return self.cached_nodes[node_id]
|
||||
|
||||
# Node not in {pending, running} -- retry with a point query. This
|
||||
# usually means the node was recently preempted or terminated.
|
||||
matches = self.acs.describe_instances(instance_ids=[node_id])
|
||||
|
||||
assert len(matches) == 1, "Invalid instance id {}".format(node_id)
|
||||
return matches[0]
|
||||
|
||||
def _get_cached_node(self, node_id):
|
||||
"""Return node info from cache if possible, otherwise fetches it."""
|
||||
if node_id in self.cached_nodes:
|
||||
return self.cached_nodes[node_id]
|
||||
|
||||
return self._get_node(node_id)
|
||||
|
||||
@staticmethod
|
||||
def bootstrap_config(cluster_config):
|
||||
return bootstrap_aliyun(cluster_config)
|
||||
@@ -0,0 +1,521 @@
|
||||
import json
|
||||
import logging
|
||||
from typing import List, Optional
|
||||
|
||||
from aliyunsdkcore import client
|
||||
from aliyunsdkcore.acs_exception.exceptions import ClientException, ServerException
|
||||
from aliyunsdkecs.request.v20140526.AllocatePublicIpAddressRequest import (
|
||||
AllocatePublicIpAddressRequest,
|
||||
)
|
||||
from aliyunsdkecs.request.v20140526.AuthorizeSecurityGroupRequest import (
|
||||
AuthorizeSecurityGroupRequest,
|
||||
)
|
||||
from aliyunsdkecs.request.v20140526.CreateInstanceRequest import CreateInstanceRequest
|
||||
from aliyunsdkecs.request.v20140526.CreateKeyPairRequest import CreateKeyPairRequest
|
||||
from aliyunsdkecs.request.v20140526.CreateSecurityGroupRequest import (
|
||||
CreateSecurityGroupRequest,
|
||||
)
|
||||
from aliyunsdkecs.request.v20140526.CreateVpcRequest import CreateVpcRequest
|
||||
from aliyunsdkecs.request.v20140526.CreateVSwitchRequest import CreateVSwitchRequest
|
||||
from aliyunsdkecs.request.v20140526.DeleteInstanceRequest import DeleteInstanceRequest
|
||||
from aliyunsdkecs.request.v20140526.DeleteInstancesRequest import DeleteInstancesRequest
|
||||
from aliyunsdkecs.request.v20140526.DeleteKeyPairsRequest import DeleteKeyPairsRequest
|
||||
from aliyunsdkecs.request.v20140526.DescribeInstancesRequest import (
|
||||
DescribeInstancesRequest,
|
||||
)
|
||||
from aliyunsdkecs.request.v20140526.DescribeKeyPairsRequest import (
|
||||
DescribeKeyPairsRequest,
|
||||
)
|
||||
from aliyunsdkecs.request.v20140526.DescribeSecurityGroupsRequest import (
|
||||
DescribeSecurityGroupsRequest,
|
||||
)
|
||||
from aliyunsdkecs.request.v20140526.DescribeVpcsRequest import DescribeVpcsRequest
|
||||
from aliyunsdkecs.request.v20140526.DescribeVSwitchesRequest import (
|
||||
DescribeVSwitchesRequest,
|
||||
)
|
||||
from aliyunsdkecs.request.v20140526.ImportKeyPairRequest import ImportKeyPairRequest
|
||||
from aliyunsdkecs.request.v20140526.RunInstancesRequest import RunInstancesRequest
|
||||
from aliyunsdkecs.request.v20140526.StartInstanceRequest import StartInstanceRequest
|
||||
from aliyunsdkecs.request.v20140526.StopInstanceRequest import StopInstanceRequest
|
||||
from aliyunsdkecs.request.v20140526.StopInstancesRequest import StopInstancesRequest
|
||||
from aliyunsdkecs.request.v20140526.TagResourcesRequest import TagResourcesRequest
|
||||
|
||||
|
||||
class AcsClient:
|
||||
"""
|
||||
A wrapper around Aliyun SDK. We use this wrapper in aliyun node provider.
|
||||
|
||||
Parameters:
|
||||
access_key: The AccessKey ID of your aliyun account.
|
||||
access_key_secret: The AccessKey secret of your aliyun account.
|
||||
region_id: A region is a geographic area where a data center resides.
|
||||
Region_id is the ID of region (e.g., cn-hangzhou,
|
||||
us-west-1, etc.)
|
||||
max_retries: The maximum number of retries each connection.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
access_key: str,
|
||||
access_key_secret: str,
|
||||
region_id: str,
|
||||
max_retries: int,
|
||||
):
|
||||
self.cli = client.AcsClient(
|
||||
ak=access_key,
|
||||
secret=access_key_secret,
|
||||
max_retry_time=max_retries,
|
||||
region_id=region_id,
|
||||
)
|
||||
|
||||
def describe_instances(
|
||||
self,
|
||||
tags: Optional[List[dict]] = None,
|
||||
instance_ids: Optional[List[str]] = None,
|
||||
) -> Optional[list]:
|
||||
"""Query the details of one or more Elastic Compute Service (ECS) instances.
|
||||
|
||||
Args:
|
||||
tags: The tags of the instance.
|
||||
instance_ids: The IDs of ECS instances.
|
||||
|
||||
Returns:
|
||||
ECS instance list.
|
||||
"""
|
||||
request = DescribeInstancesRequest()
|
||||
if tags is not None:
|
||||
request.set_Tags(tags)
|
||||
if instance_ids is not None:
|
||||
request.set_InstanceIds(instance_ids)
|
||||
response = self._send_request(request)
|
||||
if response is not None:
|
||||
instance_list = response.get("Instances").get("Instance")
|
||||
return instance_list
|
||||
return None
|
||||
|
||||
def create_instance(
|
||||
self,
|
||||
instance_type: str,
|
||||
image_id: str,
|
||||
tags: List[dict],
|
||||
key_pair_name: str,
|
||||
optimized: str = "optimized",
|
||||
instance_charge_type: str = "PostPaid",
|
||||
spot_strategy: str = "SpotWithPriceLimit",
|
||||
internet_charge_type: str = "PayByTraffic",
|
||||
internet_max_bandwidth_out: int = 5,
|
||||
) -> Optional[str]:
|
||||
"""Create a subscription or pay-as-you-go ECS instance.
|
||||
|
||||
Args:
|
||||
instance_type: The instance type of the ECS.
|
||||
image_id: The ID of the image used to create the instance.
|
||||
tags: The tags of the instance.
|
||||
key_pair_name: The name of the key pair to be bound to
|
||||
the instance.
|
||||
optimized: Specifies whether the instance is I/O optimized.
|
||||
instance_charge_type: The billing method of the instance.
|
||||
Default value: PostPaid.
|
||||
spot_strategy: The preemption policy for the pay-as-you-go
|
||||
instance.
|
||||
internet_charge_type: The billing method for network usage.
|
||||
Default value: PayByTraffic.
|
||||
internet_max_bandwidth_out: The maximum inbound public
|
||||
bandwidth. Unit: Mbit/s.
|
||||
|
||||
Returns:
|
||||
The created instance ID.
|
||||
"""
|
||||
request = CreateInstanceRequest()
|
||||
request.set_InstanceType(instance_type)
|
||||
request.set_ImageId(image_id)
|
||||
request.set_IoOptimized(optimized)
|
||||
request.set_InstanceChargeType(instance_charge_type)
|
||||
request.set_SpotStrategy(spot_strategy)
|
||||
request.set_InternetChargeType(internet_charge_type)
|
||||
request.set_InternetMaxBandwidthOut(internet_max_bandwidth_out)
|
||||
request.set_KeyPairName(key_pair_name)
|
||||
request.set_Tags(tags)
|
||||
|
||||
response = self._send_request(request)
|
||||
if response is not None:
|
||||
instance_id = response.get("InstanceId")
|
||||
logging.info("instance %s created task submit successfully.", instance_id)
|
||||
return instance_id
|
||||
logging.error("instance created failed.")
|
||||
return None
|
||||
|
||||
def run_instances(
|
||||
self,
|
||||
instance_type: str,
|
||||
image_id: str,
|
||||
tags: List[dict],
|
||||
security_group_id: str,
|
||||
vswitch_id: str,
|
||||
key_pair_name: str,
|
||||
amount: int = 1,
|
||||
optimized: str = "optimized",
|
||||
instance_charge_type: str = "PostPaid",
|
||||
spot_strategy: str = "SpotWithPriceLimit",
|
||||
internet_charge_type: str = "PayByTraffic",
|
||||
internet_max_bandwidth_out: int = 1,
|
||||
) -> Optional[List[str]]:
|
||||
"""Create one or more pay-as-you-go or subscription ECS instances.
|
||||
|
||||
Args:
|
||||
instance_type: The instance type of the ECS.
|
||||
image_id: The ID of the image used to create the instance.
|
||||
tags: The tags of the instance.
|
||||
security_group_id: The ID of the security group to which to
|
||||
assign the instance. Instances in the same security group
|
||||
can communicate with each other.
|
||||
vswitch_id: The ID of the vSwitch to which to connect
|
||||
the instance.
|
||||
key_pair_name: The name of the key pair to be bound to
|
||||
the instance.
|
||||
amount: The number of instances that you want to create.
|
||||
optimized: Specifies whether the instance is I/O optimized.
|
||||
instance_charge_type: The billing method of the instance.
|
||||
Default value: PostPaid.
|
||||
spot_strategy: The preemption policy for the pay-as-you-go
|
||||
instance.
|
||||
internet_charge_type: The billing method for network usage.
|
||||
Default value: PayByTraffic.
|
||||
internet_max_bandwidth_out: The maximum inbound public
|
||||
bandwidth. Unit: Mbit/s.
|
||||
|
||||
Returns:
|
||||
The created instance IDs.
|
||||
"""
|
||||
request = RunInstancesRequest()
|
||||
request.set_InstanceType(instance_type)
|
||||
request.set_ImageId(image_id)
|
||||
request.set_IoOptimized(optimized)
|
||||
request.set_InstanceChargeType(instance_charge_type)
|
||||
request.set_SpotStrategy(spot_strategy)
|
||||
request.set_InternetChargeType(internet_charge_type)
|
||||
request.set_InternetMaxBandwidthOut(internet_max_bandwidth_out)
|
||||
request.set_Tags(tags)
|
||||
request.set_Amount(amount)
|
||||
request.set_SecurityGroupId(security_group_id)
|
||||
request.set_VSwitchId(vswitch_id)
|
||||
request.set_KeyPairName(key_pair_name)
|
||||
|
||||
response = self._send_request(request)
|
||||
if response is not None:
|
||||
instance_ids = response.get("InstanceIdSets").get("InstanceIdSet")
|
||||
return instance_ids
|
||||
logging.error("instance created failed.")
|
||||
return None
|
||||
|
||||
def create_security_group(self, vpc_id: str) -> Optional[str]:
|
||||
"""Create a security group.
|
||||
|
||||
Args:
|
||||
vpc_id: The ID of the VPC in which to create the security group.
|
||||
|
||||
Returns:
|
||||
The created security group ID.
|
||||
"""
|
||||
request = CreateSecurityGroupRequest()
|
||||
request.set_VpcId(vpc_id)
|
||||
response = self._send_request(request)
|
||||
if response is not None:
|
||||
security_group_id = response.get("SecurityGroupId")
|
||||
return security_group_id
|
||||
return None
|
||||
|
||||
def describe_security_groups(
|
||||
self,
|
||||
vpc_id: Optional[str] = None,
|
||||
tags: Optional[List[dict]] = None,
|
||||
) -> Optional[list]:
|
||||
"""Query basic information of security groups.
|
||||
|
||||
Args:
|
||||
vpc_id: The ID of the VPC to which the security group belongs.
|
||||
tags: The tags of the security group.
|
||||
|
||||
Returns:
|
||||
Security group list.
|
||||
"""
|
||||
request = DescribeSecurityGroupsRequest()
|
||||
if vpc_id is not None:
|
||||
request.set_VpcId(vpc_id)
|
||||
if tags is not None:
|
||||
request.set_Tags(tags)
|
||||
response = self._send_request(request)
|
||||
if response is not None:
|
||||
security_groups = response.get("SecurityGroups").get("SecurityGroup")
|
||||
return security_groups
|
||||
logging.error("describe security group failed.")
|
||||
return None
|
||||
|
||||
def authorize_security_group(
|
||||
self,
|
||||
ip_protocol: str,
|
||||
port_range: str,
|
||||
security_group_id: str,
|
||||
source_cidr_ip: str,
|
||||
) -> None:
|
||||
"""Create an inbound security group rule.
|
||||
|
||||
Args:
|
||||
ip_protocol: The transport layer protocol.
|
||||
port_range: The range of destination ports relevant to
|
||||
the transport layer protocol.
|
||||
security_group_id: The ID of the destination security group.
|
||||
source_cidr_ip: The range of source IPv4 addresses.
|
||||
CIDR blocks and IPv4 addresses are supported.
|
||||
"""
|
||||
request = AuthorizeSecurityGroupRequest()
|
||||
request.set_IpProtocol(ip_protocol)
|
||||
request.set_PortRange(port_range)
|
||||
request.set_SecurityGroupId(security_group_id)
|
||||
request.set_SourceCidrIp(source_cidr_ip)
|
||||
self._send_request(request)
|
||||
|
||||
def create_v_switch(
|
||||
self, vpc_id: str, zone_id: str, cidr_block: str
|
||||
) -> Optional[str]:
|
||||
"""Create vSwitches to divide the VPC into one or more subnets.
|
||||
|
||||
Args:
|
||||
vpc_id: The ID of the VPC to which the VSwitch belongs.
|
||||
zone_id: The ID of the zone to which the target VSwitch belongs.
|
||||
cidr_block: The CIDR block of the VSwitch.
|
||||
|
||||
Returns:
|
||||
The created VSwitch ID, or None if the request failed.
|
||||
"""
|
||||
request = CreateVSwitchRequest()
|
||||
request.set_ZoneId(zone_id)
|
||||
request.set_VpcId(vpc_id)
|
||||
request.set_CidrBlock(cidr_block)
|
||||
response = self._send_request(request)
|
||||
if response is not None:
|
||||
return response.get("VSwitchId")
|
||||
else:
|
||||
logging.error("create_v_switch vpc_id %s failed.", vpc_id)
|
||||
return None
|
||||
|
||||
def create_vpc(self) -> Optional[str]:
|
||||
"""Create a virtual private cloud (VPC).
|
||||
|
||||
Returns:
|
||||
The created VPC ID.
|
||||
"""
|
||||
request = CreateVpcRequest()
|
||||
response = self._send_request(request)
|
||||
if response is not None:
|
||||
return response.get("VpcId")
|
||||
return None
|
||||
|
||||
def describe_vpcs(self) -> Optional[list]:
|
||||
"""Query one or more VPCs in a region.
|
||||
|
||||
Returns:
|
||||
VPC list.
|
||||
"""
|
||||
request = DescribeVpcsRequest()
|
||||
response = self._send_request(request)
|
||||
if response is not None:
|
||||
return response.get("Vpcs").get("Vpc")
|
||||
return None
|
||||
|
||||
def tag_resource(
|
||||
self,
|
||||
resource_ids: List[str],
|
||||
tags: List[dict],
|
||||
resource_type: str = "instance",
|
||||
) -> None:
|
||||
"""Create and bind tags to specified ECS resources.
|
||||
|
||||
Args:
|
||||
resource_ids: The IDs of N resources.
|
||||
tags: The tags of the resource.
|
||||
resource_type: The type of the resource.
|
||||
"""
|
||||
request = TagResourcesRequest()
|
||||
request.set_Tags(tags)
|
||||
request.set_ResourceType(resource_type)
|
||||
request.set_ResourceIds(resource_ids)
|
||||
response = self._send_request(request)
|
||||
if response is not None:
|
||||
logging.info("instance %s create tag successfully.", resource_ids)
|
||||
else:
|
||||
logging.error("instance %s create tag failed.", resource_ids)
|
||||
|
||||
def start_instance(self, instance_id: str) -> None:
|
||||
"""Start an ECS instance.
|
||||
|
||||
Args:
|
||||
instance_id: The ECS instance ID.
|
||||
"""
|
||||
request = StartInstanceRequest()
|
||||
request.set_InstanceId(instance_id)
|
||||
response = self._send_request(request)
|
||||
|
||||
if response is not None:
|
||||
logging.info("instance %s start successfully.", instance_id)
|
||||
else:
|
||||
logging.error("instance %s start failed.", instance_id)
|
||||
|
||||
def stop_instance(self, instance_id: str, force_stop: bool = False) -> None:
|
||||
"""Stop an ECS instance that is in the Running state.
|
||||
|
||||
Args:
|
||||
instance_id: The ECS instance ID.
|
||||
force_stop: Specifies whether to forcibly stop the instance.
|
||||
"""
|
||||
request = StopInstanceRequest()
|
||||
request.set_InstanceId(instance_id)
|
||||
request.set_ForceStop(force_stop)
|
||||
logging.info("Stop %s command submit successfully.", instance_id)
|
||||
self._send_request(request)
|
||||
|
||||
def stop_instances(
|
||||
self, instance_ids: List[str], stopped_mode: str = "StopCharging"
|
||||
) -> None:
|
||||
"""Stop one or more ECS instances that are in the Running state.
|
||||
|
||||
Args:
|
||||
instance_ids: The IDs of instances.
|
||||
stopped_mode: Specifies whether billing for the instance
|
||||
continues after the instance is stopped.
|
||||
"""
|
||||
request = StopInstancesRequest()
|
||||
request.set_InstanceIds(instance_ids)
|
||||
request.set_StoppedMode(stopped_mode)
|
||||
response = self._send_request(request)
|
||||
if response is None:
|
||||
logging.error("stop_instances failed")
|
||||
|
||||
def delete_instance(self, instance_id: str) -> None:
|
||||
"""Release a pay-as-you-go instance or an expired subscription instance.
|
||||
|
||||
Args:
|
||||
instance_id: The ID of the instance that you want to release.
|
||||
"""
|
||||
request = DeleteInstanceRequest()
|
||||
request.set_InstanceId(instance_id)
|
||||
request.set_Force(True)
|
||||
logging.info("Delete %s command submit successfully", instance_id)
|
||||
self._send_request(request)
|
||||
|
||||
def delete_instances(self, instance_ids: List[str]) -> None:
|
||||
"""Release one or more pay-as-you-go or expired subscription instances.
|
||||
|
||||
Args:
|
||||
instance_ids: The IDs of instances that you want to release.
|
||||
"""
|
||||
request = DeleteInstancesRequest()
|
||||
request.set_Force(True)
|
||||
request.set_InstanceIds(instance_ids)
|
||||
self._send_request(request)
|
||||
|
||||
def allocate_public_address(self, instance_id: str) -> Optional[str]:
|
||||
"""Assign a public IP address to an ECS instance.
|
||||
|
||||
Args:
|
||||
instance_id: The ID of the instance to which you want to
|
||||
assign a public IP address.
|
||||
|
||||
Returns:
|
||||
The assigned IP address.
|
||||
"""
|
||||
request = AllocatePublicIpAddressRequest()
|
||||
request.set_InstanceId(instance_id)
|
||||
response = self._send_request(request)
|
||||
if response is not None:
|
||||
return response.get("IpAddress")
|
||||
|
||||
def create_key_pair(self, key_pair_name: str) -> Optional[dict]:
|
||||
"""Create an SSH key pair.
|
||||
|
||||
Args:
|
||||
key_pair_name: The name of the key pair.
|
||||
|
||||
Returns:
|
||||
The created keypair data.
|
||||
"""
|
||||
request = CreateKeyPairRequest()
|
||||
request.set_KeyPairName(key_pair_name)
|
||||
response = self._send_request(request)
|
||||
if response is not None:
|
||||
logging.info("Create Key Pair %s Successfully", response.get("KeyPairId"))
|
||||
return response
|
||||
else:
|
||||
logging.error("Create Key Pair Failed")
|
||||
return None
|
||||
|
||||
def import_key_pair(self, key_pair_name: str, public_key_body: str) -> None:
|
||||
"""Import the public key of an RSA-encrypted key pair.
|
||||
|
||||
Args:
|
||||
key_pair_name: The name of the key pair.
|
||||
public_key_body: The public key of the key pair.
|
||||
"""
|
||||
request = ImportKeyPairRequest()
|
||||
request.set_KeyPairName(key_pair_name)
|
||||
request.set_PublicKeyBody(public_key_body)
|
||||
self._send_request(request)
|
||||
|
||||
def delete_key_pairs(self, key_pair_names: List[str]) -> None:
|
||||
"""Delete one or more SSH key pairs.
|
||||
|
||||
Args:
|
||||
key_pair_names: The names of the key pairs.
|
||||
"""
|
||||
request = DeleteKeyPairsRequest()
|
||||
request.set_KeyPairNames(key_pair_names)
|
||||
self._send_request(request)
|
||||
|
||||
def describe_key_pairs(self, key_pair_name: Optional[str] = None) -> Optional[list]:
|
||||
"""Query one or more key pairs.
|
||||
|
||||
Args:
|
||||
key_pair_name: The name of the key pair.
|
||||
|
||||
Returns:
|
||||
Key pair list, or None if the request failed.
|
||||
"""
|
||||
request = DescribeKeyPairsRequest()
|
||||
if key_pair_name is not None:
|
||||
request.set_KeyPairName(key_pair_name)
|
||||
response = self._send_request(request)
|
||||
if response is not None:
|
||||
return response.get("KeyPairs").get("KeyPair")
|
||||
else:
|
||||
return None
|
||||
|
||||
def describe_v_switches(self, vpc_id: Optional[str] = None) -> Optional[list]:
|
||||
"""Query one or more VSwitches.
|
||||
|
||||
Args:
|
||||
vpc_id: The ID of the VPC to which the VSwitch belongs.
|
||||
|
||||
Returns:
|
||||
VSwitch list.
|
||||
"""
|
||||
request = DescribeVSwitchesRequest()
|
||||
if vpc_id is not None:
|
||||
request.set_VpcId(vpc_id)
|
||||
response = self._send_request(request)
|
||||
if response is not None:
|
||||
return response.get("VSwitches").get("VSwitch")
|
||||
else:
|
||||
logging.error("Describe VSwitches Failed.")
|
||||
return None
|
||||
|
||||
def _send_request(self, request):
|
||||
"""send open api request"""
|
||||
request.set_accept_format("json")
|
||||
try:
|
||||
response_str = self.cli.do_action_with_exception(request)
|
||||
response_detail = json.loads(response_str)
|
||||
return response_detail
|
||||
except (ClientException, ServerException) as e:
|
||||
logging.error(request.get_action_name())
|
||||
logging.error(e)
|
||||
return None
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,821 @@
|
||||
import copy
|
||||
import hashlib
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import time
|
||||
from enum import Enum
|
||||
from typing import Any, Callable, Dict, List, Union
|
||||
|
||||
import botocore
|
||||
|
||||
from ray.autoscaler._private.aws.utils import client_cache, resource_cache
|
||||
from ray.autoscaler.tags import NODE_KIND_HEAD, TAG_RAY_CLUSTER_NAME, TAG_RAY_NODE_KIND
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
RAY = "ray-autoscaler"
|
||||
CLOUDWATCH_RAY_INSTANCE_PROFILE = RAY + "-cloudwatch-v1"
|
||||
CLOUDWATCH_RAY_IAM_ROLE = RAY + "-cloudwatch-v1"
|
||||
CLOUDWATCH_AGENT_INSTALLED_AMI_TAG = "T6Iq2faj"
|
||||
CLOUDWATCH_AGENT_INSTALLED_TAG = "cloudwatch-agent-installed"
|
||||
CLOUDWATCH_CONFIG_HASH_TAG_BASE = "cloudwatch-config-hash"
|
||||
|
||||
|
||||
class CloudwatchConfigType(str, Enum):
|
||||
AGENT = "agent"
|
||||
DASHBOARD = "dashboard"
|
||||
ALARM = "alarm"
|
||||
|
||||
|
||||
class CloudwatchHelper:
|
||||
def __init__(
|
||||
self, provider_config: Dict[str, Any], node_id: str, cluster_name: str
|
||||
) -> None:
|
||||
self.node_id = node_id
|
||||
self.cluster_name = cluster_name
|
||||
self.provider_config = provider_config
|
||||
region = provider_config["region"]
|
||||
self.ec2_resource = resource_cache("ec2", region)
|
||||
self.ec2_client = self.ec2_resource.meta.client
|
||||
self.ssm_client = client_cache("ssm", region)
|
||||
cloudwatch_resource = resource_cache("cloudwatch", region)
|
||||
self.cloudwatch_client = cloudwatch_resource.meta.client
|
||||
self.CLOUDWATCH_CONFIG_TYPE_TO_CONFIG_VARIABLE_REPLACE_FUNC: Dict[
|
||||
str, Callable
|
||||
] = {
|
||||
CloudwatchConfigType.AGENT.value: self._replace_cwa_config_vars,
|
||||
CloudwatchConfigType.DASHBOARD.value: self._replace_dashboard_config_vars,
|
||||
CloudwatchConfigType.ALARM.value: self._load_config_file,
|
||||
}
|
||||
self.CLOUDWATCH_CONFIG_TYPE_TO_UPDATE_FUNC_HEAD_NODE: Dict[str, Callable] = {
|
||||
CloudwatchConfigType.AGENT.value: self._restart_cloudwatch_agent,
|
||||
CloudwatchConfigType.DASHBOARD.value: self._put_cloudwatch_dashboard,
|
||||
CloudwatchConfigType.ALARM.value: self._put_cloudwatch_alarm,
|
||||
}
|
||||
self.CLOUDWATCH_CONFIG_TYPE_TO_UPDATE_FUNC_WORKER_NODE: Dict[str, Callable] = {
|
||||
CloudwatchConfigType.AGENT.value: self._restart_cloudwatch_agent,
|
||||
CloudwatchConfigType.ALARM.value: self._put_cloudwatch_alarm,
|
||||
}
|
||||
|
||||
def update_from_config(self, is_head_node: bool) -> None:
|
||||
"""Discovers and applies CloudWatch config updates as required.
|
||||
|
||||
Args:
|
||||
is_head_node: whether this node is the head node.
|
||||
"""
|
||||
for config_type in CloudwatchConfigType:
|
||||
if CloudwatchHelper.cloudwatch_config_exists(
|
||||
self.provider_config, config_type.value
|
||||
):
|
||||
self._update_cloudwatch_config(config_type.value, is_head_node)
|
||||
|
||||
def _ec2_health_check_waiter(self, node_id: str) -> None:
|
||||
# wait for all EC2 instance checks to complete
|
||||
try:
|
||||
logger.info(
|
||||
"Waiting for EC2 instance health checks to complete before "
|
||||
"configuring Unified Cloudwatch Agent. This may take a few "
|
||||
"minutes..."
|
||||
)
|
||||
waiter = self.ec2_client.get_waiter("instance_status_ok")
|
||||
waiter.wait(InstanceIds=[node_id])
|
||||
except botocore.exceptions.WaiterError as e:
|
||||
logger.error(
|
||||
"Failed while waiting for EC2 instance checks to complete: {}".format(
|
||||
e.message
|
||||
)
|
||||
)
|
||||
raise e
|
||||
|
||||
def _update_cloudwatch_config(self, config_type: str, is_head_node: bool) -> None:
|
||||
"""
|
||||
check whether update operations are needed in
|
||||
cloudwatch related configs
|
||||
"""
|
||||
cwa_installed = self._setup_cwa()
|
||||
param_name = self._get_ssm_param_name(config_type)
|
||||
if cwa_installed:
|
||||
if is_head_node:
|
||||
cw_config_ssm = self._set_cloudwatch_ssm_config_param(
|
||||
param_name, config_type
|
||||
)
|
||||
cur_cw_config_hash = self._sha256_hash_file(config_type)
|
||||
ssm_cw_config_hash = self._sha256_hash_json(cw_config_ssm)
|
||||
# check if user updated cloudwatch related config files.
|
||||
# if so, perform corresponding actions.
|
||||
if cur_cw_config_hash != ssm_cw_config_hash:
|
||||
logger.info(
|
||||
"Cloudwatch {} config file has changed.".format(config_type)
|
||||
)
|
||||
self._upload_config_to_ssm_and_set_hash_tag(config_type)
|
||||
self.CLOUDWATCH_CONFIG_TYPE_TO_UPDATE_FUNC_HEAD_NODE.get(
|
||||
config_type
|
||||
)()
|
||||
else:
|
||||
head_node_hash = self._get_head_node_config_hash(config_type)
|
||||
cur_node_hash = self._get_cur_node_config_hash(config_type)
|
||||
if head_node_hash != cur_node_hash:
|
||||
logger.info(
|
||||
"Cloudwatch {} config file has changed.".format(config_type)
|
||||
)
|
||||
update_func = (
|
||||
self.CLOUDWATCH_CONFIG_TYPE_TO_UPDATE_FUNC_WORKER_NODE.get(
|
||||
config_type
|
||||
)
|
||||
)
|
||||
if update_func:
|
||||
update_func()
|
||||
self._update_cloudwatch_hash_tag_value(
|
||||
self.node_id, head_node_hash, config_type
|
||||
)
|
||||
|
||||
def _put_cloudwatch_dashboard(self) -> Dict[str, Any]:
|
||||
"""put dashboard to cloudwatch console"""
|
||||
|
||||
cloudwatch_config = self.provider_config["cloudwatch"]
|
||||
dashboard_config = cloudwatch_config.get("dashboard", {})
|
||||
dashboard_name_cluster = dashboard_config.get("name", self.cluster_name)
|
||||
dashboard_name = self.cluster_name + "-" + dashboard_name_cluster
|
||||
|
||||
widgets = self._replace_dashboard_config_vars(
|
||||
CloudwatchConfigType.DASHBOARD.value
|
||||
)
|
||||
|
||||
response = self.cloudwatch_client.put_dashboard(
|
||||
DashboardName=dashboard_name, DashboardBody=json.dumps({"widgets": widgets})
|
||||
)
|
||||
issue_count = len(response.get("DashboardValidationMessages", []))
|
||||
if issue_count > 0:
|
||||
for issue in response.get("DashboardValidationMessages"):
|
||||
logging.error(
|
||||
"Error in dashboard config: {} - {}".format(
|
||||
issue["Message"], issue["DataPath"]
|
||||
)
|
||||
)
|
||||
raise Exception(
|
||||
"Errors in dashboard configuration: {} issues raised".format(
|
||||
issue_count
|
||||
)
|
||||
)
|
||||
else:
|
||||
logger.info("Successfully put dashboard to CloudWatch console")
|
||||
return response
|
||||
|
||||
def _put_cloudwatch_alarm(self) -> None:
|
||||
"""put CloudWatch metric alarms read from config"""
|
||||
param_name = self._get_ssm_param_name(CloudwatchConfigType.ALARM.value)
|
||||
data = json.loads(self._get_ssm_param(param_name))
|
||||
for item in data:
|
||||
item_out = copy.deepcopy(item)
|
||||
self._replace_all_config_variables(
|
||||
item_out,
|
||||
self.node_id,
|
||||
self.cluster_name,
|
||||
self.provider_config["region"],
|
||||
)
|
||||
self.cloudwatch_client.put_metric_alarm(**item_out)
|
||||
logger.info("Successfully put alarms to CloudWatch console")
|
||||
|
||||
def _send_command_to_node(
|
||||
self, document_name: str, parameters: Dict[str, List[str]], node_id: str
|
||||
) -> Dict[str, Any]:
|
||||
"""send SSM command to the given nodes"""
|
||||
logger.debug(
|
||||
"Sending SSM command to {} node(s). Document name: {}. "
|
||||
"Parameters: {}.".format(node_id, document_name, parameters)
|
||||
)
|
||||
response = self.ssm_client.send_command(
|
||||
InstanceIds=[node_id],
|
||||
DocumentName=document_name,
|
||||
Parameters=parameters,
|
||||
MaxConcurrency="1",
|
||||
MaxErrors="0",
|
||||
)
|
||||
return response
|
||||
|
||||
def _ssm_command_waiter(
|
||||
self,
|
||||
document_name: str,
|
||||
parameters: Dict[str, List[str]],
|
||||
node_id: str,
|
||||
retry_failed: bool = True,
|
||||
) -> Dict[str, Any]:
|
||||
"""wait for SSM command to complete on all cluster nodes"""
|
||||
|
||||
# This waiter differs from the built-in SSM.Waiter by
|
||||
# optimistically waiting for the command invocation to
|
||||
# exist instead of failing immediately, and by resubmitting
|
||||
# any failed command until all retry attempts are exhausted
|
||||
# by default.
|
||||
response = self._send_command_to_node(document_name, parameters, node_id)
|
||||
command_id = response["Command"]["CommandId"]
|
||||
|
||||
cloudwatch_config = self.provider_config["cloudwatch"]
|
||||
agent_retryer_config = cloudwatch_config.get(
|
||||
CloudwatchConfigType.AGENT.value
|
||||
).get("retryer", {})
|
||||
max_attempts = agent_retryer_config.get("max_attempts", 120)
|
||||
delay_seconds = agent_retryer_config.get("delay_seconds", 30)
|
||||
num_attempts = 0
|
||||
cmd_invocation_res = {}
|
||||
while True:
|
||||
num_attempts += 1
|
||||
logger.debug(
|
||||
"Listing SSM command ID {} invocations on node {}".format(
|
||||
command_id, node_id
|
||||
)
|
||||
)
|
||||
response = self.ssm_client.list_command_invocations(
|
||||
CommandId=command_id,
|
||||
InstanceId=node_id,
|
||||
)
|
||||
cmd_invocations = response["CommandInvocations"]
|
||||
if not cmd_invocations:
|
||||
logger.debug(
|
||||
"SSM Command ID {} invocation does not exist. If "
|
||||
"the command was just started, it may take a "
|
||||
"few seconds to register.".format(command_id)
|
||||
)
|
||||
else:
|
||||
if len(cmd_invocations) > 1:
|
||||
logger.warning(
|
||||
"Expected to find 1 SSM command invocation with "
|
||||
"ID {} on node {} but found {}: {}".format(
|
||||
command_id,
|
||||
node_id,
|
||||
len(cmd_invocations),
|
||||
cmd_invocations,
|
||||
)
|
||||
)
|
||||
cmd_invocation = cmd_invocations[0]
|
||||
if cmd_invocation["Status"] == "Success":
|
||||
logger.debug(
|
||||
"SSM Command ID {} completed successfully.".format(command_id)
|
||||
)
|
||||
cmd_invocation_res[node_id] = True
|
||||
break
|
||||
if num_attempts >= max_attempts:
|
||||
logger.error(
|
||||
"Max attempts for command {} exceeded on node {}".format(
|
||||
command_id, node_id
|
||||
)
|
||||
)
|
||||
raise botocore.exceptions.WaiterError(
|
||||
name="ssm_waiter",
|
||||
reason="Max attempts exceeded",
|
||||
last_response=cmd_invocation,
|
||||
)
|
||||
if cmd_invocation["Status"] == "Failed":
|
||||
logger.debug(f"SSM Command ID {command_id} failed.")
|
||||
if retry_failed:
|
||||
logger.debug(f"Retrying in {delay_seconds} seconds.")
|
||||
response = self._send_command_to_node(
|
||||
document_name, parameters, node_id
|
||||
)
|
||||
command_id = response["Command"]["CommandId"]
|
||||
logger.debug(
|
||||
"Sent SSM command ID {} to node {}".format(
|
||||
command_id, node_id
|
||||
)
|
||||
)
|
||||
else:
|
||||
logger.debug(f"Ignoring Command ID {command_id} failure.")
|
||||
cmd_invocation_res[node_id] = False
|
||||
break
|
||||
time.sleep(delay_seconds)
|
||||
|
||||
return cmd_invocation_res
|
||||
|
||||
def _replace_config_variables(
|
||||
self, string: str, node_id: str, cluster_name: str, region: str
|
||||
) -> str:
|
||||
"""
|
||||
replace known config variable occurrences in the input string
|
||||
does not replace variables with undefined or empty strings
|
||||
"""
|
||||
|
||||
if node_id:
|
||||
string = string.replace("{instance_id}", node_id)
|
||||
if cluster_name:
|
||||
string = string.replace("{cluster_name}", cluster_name)
|
||||
if region:
|
||||
string = string.replace("{region}", region)
|
||||
return string
|
||||
|
||||
def _replace_all_config_variables(
|
||||
self,
|
||||
collection: Union[Dict[str, Any], str],
|
||||
node_id: str,
|
||||
cluster_name: str,
|
||||
region: str,
|
||||
) -> Union[str, Dict[str, Any]]:
|
||||
"""
|
||||
Replace known config variable occurrences in the input collection.
|
||||
The input collection must be either a dict or list.
|
||||
Returns a tuple consisting of the output collection and the number of
|
||||
modified strings in the collection (which is not necessarily equal to
|
||||
the number of variables replaced).
|
||||
"""
|
||||
|
||||
for key in collection:
|
||||
if type(collection) is dict:
|
||||
value = collection.get(key)
|
||||
index_key = key
|
||||
elif type(collection) is list:
|
||||
value = key
|
||||
index_key = collection.index(key)
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Can't replace CloudWatch config variables "
|
||||
f"in unsupported collection type: {type(collection)}."
|
||||
f"Please check your CloudWatch JSON config files."
|
||||
)
|
||||
if type(value) is str:
|
||||
collection[index_key] = self._replace_config_variables(
|
||||
value, node_id, cluster_name, region
|
||||
)
|
||||
elif type(value) is dict or type(value) is list:
|
||||
collection[index_key] = self._replace_all_config_variables(
|
||||
value, node_id, cluster_name, region
|
||||
)
|
||||
return collection
|
||||
|
||||
def _load_config_file(self, config_type: str) -> Dict[str, Any]:
|
||||
"""load JSON config file"""
|
||||
cloudwatch_config = self.provider_config["cloudwatch"]
|
||||
json_config_file_section = cloudwatch_config.get(config_type, {})
|
||||
json_config_file_path = json_config_file_section.get("config", {})
|
||||
json_config_path = os.path.abspath(json_config_file_path)
|
||||
with open(json_config_path) as f:
|
||||
data = json.load(f)
|
||||
return data
|
||||
|
||||
def _set_cloudwatch_ssm_config_param(
|
||||
self, parameter_name: str, config_type: str
|
||||
) -> str:
|
||||
"""
|
||||
get cloudwatch config for the given param and config type from SSM
|
||||
if it exists, put it in the SSM param store if not
|
||||
"""
|
||||
try:
|
||||
parameter_value = self._get_ssm_param(parameter_name)
|
||||
except botocore.exceptions.ClientError as e:
|
||||
if e.response["Error"]["Code"] == "ParameterNotFound":
|
||||
logger.info(
|
||||
"Cloudwatch {} config file is not found "
|
||||
"at SSM parameter store. "
|
||||
"Checking for Unified CloudWatch Agent installation".format(
|
||||
config_type
|
||||
)
|
||||
)
|
||||
return self._get_default_empty_config_file_hash()
|
||||
else:
|
||||
logger.info(
|
||||
"Failed to fetch Unified CloudWatch Agent config from SSM "
|
||||
"parameter store."
|
||||
)
|
||||
logger.error(e)
|
||||
raise e
|
||||
return parameter_value
|
||||
|
||||
def _get_default_empty_config_file_hash(self):
|
||||
default_cw_config = "{}"
|
||||
parameter_value = self._sha256_hash_json(default_cw_config)
|
||||
return parameter_value
|
||||
|
||||
def _get_ssm_param(self, parameter_name: str) -> str:
|
||||
"""
|
||||
get the SSM parameter value associated with the given parameter name
|
||||
"""
|
||||
response = self.ssm_client.get_parameter(Name=parameter_name)
|
||||
logger.info("Successfully fetch ssm parameter: {}".format(parameter_name))
|
||||
res = response.get("Parameter", {})
|
||||
cwa_parameter = res.get("Value", {})
|
||||
return cwa_parameter
|
||||
|
||||
def _sha256_hash_json(self, value: str) -> str:
|
||||
"""calculate the json string sha256 hash"""
|
||||
sha256_hash = hashlib.new("sha256")
|
||||
binary_value = value.encode("utf-8")
|
||||
sha256_hash.update(binary_value)
|
||||
sha256_res = sha256_hash.hexdigest()
|
||||
return sha256_res
|
||||
|
||||
def _sha256_hash_file(self, config_type: str) -> str:
|
||||
"""calculate the config file sha256 hash"""
|
||||
config = self.CLOUDWATCH_CONFIG_TYPE_TO_CONFIG_VARIABLE_REPLACE_FUNC.get(
|
||||
config_type
|
||||
)(config_type)
|
||||
value = json.dumps(config)
|
||||
sha256_res = self._sha256_hash_json(value)
|
||||
return sha256_res
|
||||
|
||||
def _upload_config_to_ssm_and_set_hash_tag(self, config_type: str):
|
||||
data = self.CLOUDWATCH_CONFIG_TYPE_TO_CONFIG_VARIABLE_REPLACE_FUNC.get(
|
||||
config_type
|
||||
)(config_type)
|
||||
sha256_hash_value = self._sha256_hash_file(config_type)
|
||||
self._upload_config_to_ssm(data, config_type)
|
||||
self._update_cloudwatch_hash_tag_value(
|
||||
self.node_id, sha256_hash_value, config_type
|
||||
)
|
||||
|
||||
def _add_cwa_installed_tag(self, node_id: str) -> None:
|
||||
self.ec2_client.create_tags(
|
||||
Resources=[node_id],
|
||||
Tags=[{"Key": CLOUDWATCH_AGENT_INSTALLED_TAG, "Value": "True"}],
|
||||
)
|
||||
logger.info(
|
||||
"Successfully add Unified CloudWatch Agent installed "
|
||||
"tag on {}".format(node_id)
|
||||
)
|
||||
|
||||
def _update_cloudwatch_hash_tag_value(
|
||||
self, node_id: str, sha256_hash_value: str, config_type: str
|
||||
):
|
||||
hash_key_value = "-".join([CLOUDWATCH_CONFIG_HASH_TAG_BASE, config_type])
|
||||
self.ec2_client.create_tags(
|
||||
Resources=[node_id],
|
||||
Tags=[{"Key": hash_key_value, "Value": sha256_hash_value}],
|
||||
)
|
||||
logger.info(
|
||||
"Successfully update cloudwatch {} hash tag on {}".format(
|
||||
config_type, node_id
|
||||
)
|
||||
)
|
||||
|
||||
def _get_ssm_param_name(self, config_type: str) -> str:
|
||||
"""return the parameter name for cloudwatch configs"""
|
||||
ssm_config_param_name = "AmazonCloudWatch-" + "ray_{}_config_{}".format(
|
||||
config_type, self.cluster_name
|
||||
)
|
||||
return ssm_config_param_name
|
||||
|
||||
def _put_ssm_param(self, parameter: Dict[str, Any], parameter_name: str) -> None:
|
||||
"""upload cloudwatch config to the SSM parameter store"""
|
||||
self.ssm_client.put_parameter(
|
||||
Name=parameter_name,
|
||||
Type="String",
|
||||
Value=json.dumps(parameter),
|
||||
Overwrite=True,
|
||||
Tier="Intelligent-Tiering",
|
||||
)
|
||||
|
||||
def _upload_config_to_ssm(self, param: Dict[str, Any], config_type: str):
|
||||
param_name = self._get_ssm_param_name(config_type)
|
||||
self._put_ssm_param(param, param_name)
|
||||
|
||||
def _replace_cwa_config_vars(self, config_type: str) -> Dict[str, Any]:
|
||||
"""
|
||||
replace {instance_id}, {region}, {cluster_name}
|
||||
variable occurrences in Unified Cloudwatch Agent config file
|
||||
"""
|
||||
cwa_config = self._load_config_file(config_type)
|
||||
self._replace_all_config_variables(
|
||||
cwa_config,
|
||||
self.node_id,
|
||||
self.cluster_name,
|
||||
self.provider_config["region"],
|
||||
)
|
||||
return cwa_config
|
||||
|
||||
def _replace_dashboard_config_vars(self, config_type: str) -> List[str]:
|
||||
"""
|
||||
replace known variable occurrences in CloudWatch Dashboard config file
|
||||
"""
|
||||
data = self._load_config_file(config_type)
|
||||
widgets = []
|
||||
for item in data:
|
||||
item_out = self._replace_all_config_variables(
|
||||
item,
|
||||
self.node_id,
|
||||
self.cluster_name,
|
||||
self.provider_config["region"],
|
||||
)
|
||||
widgets.append(item_out)
|
||||
return widgets
|
||||
|
||||
def _replace_alarm_config_vars(self, config_type: str) -> List[str]:
|
||||
"""
|
||||
replace {instance_id}, {region}, {cluster_name}
|
||||
variable occurrences in cloudwatch alarm config file
|
||||
"""
|
||||
data = self._load_config_file(config_type)
|
||||
param_data = []
|
||||
for item in data:
|
||||
item_out = copy.deepcopy(item)
|
||||
self._replace_all_config_variables(
|
||||
item_out,
|
||||
self.node_id,
|
||||
self.cluster_name,
|
||||
self.provider_config["region"],
|
||||
)
|
||||
param_data.append(item_out)
|
||||
return param_data
|
||||
|
||||
def _restart_cloudwatch_agent(self) -> None:
|
||||
"""restart Unified CloudWatch Agent"""
|
||||
cwa_param_name = self._get_ssm_param_name(CloudwatchConfigType.AGENT.value)
|
||||
logger.info(
|
||||
"Restarting Unified CloudWatch Agent package on node {}.".format(
|
||||
self.node_id
|
||||
)
|
||||
)
|
||||
self._stop_cloudwatch_agent()
|
||||
self._start_cloudwatch_agent(cwa_param_name)
|
||||
|
||||
def _stop_cloudwatch_agent(self) -> None:
|
||||
"""stop Unified CloudWatch Agent"""
|
||||
logger.info(
|
||||
"Stopping Unified CloudWatch Agent package on node {}.".format(self.node_id)
|
||||
)
|
||||
parameters_stop_cwa = {
|
||||
"action": ["stop"],
|
||||
"mode": ["ec2"],
|
||||
}
|
||||
# don't retry failed stop commands
|
||||
# (there's not always an agent to stop)
|
||||
self._ssm_command_waiter(
|
||||
"AmazonCloudWatch-ManageAgent",
|
||||
parameters_stop_cwa,
|
||||
self.node_id,
|
||||
False,
|
||||
)
|
||||
logger.info("Unified CloudWatch Agent stopped on node {}.".format(self.node_id))
|
||||
|
||||
def _start_cloudwatch_agent(self, cwa_param_name: str) -> None:
|
||||
"""start Unified CloudWatch Agent"""
|
||||
logger.info(
|
||||
"Starting Unified CloudWatch Agent package on node {}.".format(self.node_id)
|
||||
)
|
||||
parameters_start_cwa = {
|
||||
"action": ["configure"],
|
||||
"mode": ["ec2"],
|
||||
"optionalConfigurationSource": ["ssm"],
|
||||
"optionalConfigurationLocation": [cwa_param_name],
|
||||
"optionalRestart": ["yes"],
|
||||
}
|
||||
self._ssm_command_waiter(
|
||||
"AmazonCloudWatch-ManageAgent", parameters_start_cwa, self.node_id
|
||||
)
|
||||
logger.info(
|
||||
"Unified CloudWatch Agent started successfully on node {}.".format(
|
||||
self.node_id
|
||||
)
|
||||
)
|
||||
|
||||
def _setup_cwa(self) -> bool:
|
||||
cwa_installed = self._check_cwa_installed_ec2_tag()
|
||||
if cwa_installed == "False":
|
||||
res_cwa_installed = self._ensure_cwa_installed_ssm(self.node_id)
|
||||
return res_cwa_installed
|
||||
else:
|
||||
return True
|
||||
|
||||
def _get_head_node_config_hash(self, config_type: str) -> str:
|
||||
hash_key_value = "-".join([CLOUDWATCH_CONFIG_HASH_TAG_BASE, config_type])
|
||||
filters = copy.deepcopy(
|
||||
self._get_current_cluster_session_nodes(self.cluster_name)
|
||||
)
|
||||
filters.append(
|
||||
{
|
||||
"Name": "tag:{}".format(TAG_RAY_NODE_KIND),
|
||||
"Values": [NODE_KIND_HEAD],
|
||||
}
|
||||
)
|
||||
try:
|
||||
instance = list(self.ec2_resource.instances.filter(Filters=filters))
|
||||
assert len(instance) == 1, "More than 1 head node found!"
|
||||
for tag in instance[0].tags:
|
||||
if tag["Key"] == hash_key_value:
|
||||
return tag["Value"]
|
||||
except botocore.exceptions.ClientError as e:
|
||||
logger.warning(
|
||||
"{} Error caught when getting value of {} tag on head node".format(
|
||||
e.response["Error"], hash_key_value
|
||||
)
|
||||
)
|
||||
|
||||
def _get_cur_node_config_hash(self, config_type: str) -> str:
|
||||
hash_key_value = "-".join([CLOUDWATCH_CONFIG_HASH_TAG_BASE, config_type])
|
||||
try:
|
||||
response = self.ec2_client.describe_instances(InstanceIds=[self.node_id])
|
||||
reservations = response["Reservations"]
|
||||
message = "More than 1 response received from describing current node"
|
||||
assert len(reservations) == 1, message
|
||||
instances = reservations[0]["Instances"]
|
||||
assert len(reservations) == 1, message
|
||||
tags = instances[0]["Tags"]
|
||||
hash_value = self._get_default_empty_config_file_hash()
|
||||
for tag in tags:
|
||||
if tag["Key"] == hash_key_value:
|
||||
logger.info(
|
||||
"Successfully get cloudwatch {} hash tag value from "
|
||||
"node {}".format(config_type, self.node_id)
|
||||
)
|
||||
hash_value = tag["Value"]
|
||||
return hash_value
|
||||
except botocore.exceptions.ClientError as e:
|
||||
logger.warning(
|
||||
"{} Error caught when getting hash tag {} tag".format(
|
||||
e.response["Error"], hash_key_value
|
||||
)
|
||||
)
|
||||
|
||||
def _ensure_cwa_installed_ssm(self, node_id: str) -> bool:
|
||||
"""
|
||||
Check if Unified Cloudwatch Agent is installed via ssm run command.
|
||||
If not, notify user to use an AMI with
|
||||
the Unified CloudWatch Agent installed.
|
||||
"""
|
||||
logger.info(
|
||||
"Checking Unified Cloudwatch Agent status on node {}".format(node_id)
|
||||
)
|
||||
parameters_status_cwa = {
|
||||
"action": ["status"],
|
||||
"mode": ["ec2"],
|
||||
}
|
||||
self._ec2_health_check_waiter(node_id)
|
||||
cmd_invocation_res = self._ssm_command_waiter(
|
||||
"AmazonCloudWatch-ManageAgent", parameters_status_cwa, node_id, False
|
||||
)
|
||||
cwa_installed = cmd_invocation_res.get(node_id, False)
|
||||
if not cwa_installed:
|
||||
logger.warning(
|
||||
"Unified CloudWatch Agent not installed on {}. "
|
||||
"Ray logs, metrics not picked up. "
|
||||
"Please use an AMI with Unified CloudWatch Agent installed.".format(
|
||||
node_id
|
||||
)
|
||||
)
|
||||
return False
|
||||
else:
|
||||
return True
|
||||
|
||||
def _get_current_cluster_session_nodes(self, cluster_name: str) -> List[dict]:
|
||||
filters = [
|
||||
{
|
||||
"Name": "instance-state-name",
|
||||
"Values": ["pending", "running"],
|
||||
},
|
||||
{
|
||||
"Name": "tag:{}".format(TAG_RAY_CLUSTER_NAME),
|
||||
"Values": [cluster_name],
|
||||
},
|
||||
]
|
||||
return filters
|
||||
|
||||
def _check_cwa_installed_ec2_tag(self) -> List[str]:
|
||||
"""
|
||||
Filtering all nodes to get nodes
|
||||
without Unified CloudWatch Agent installed
|
||||
"""
|
||||
try:
|
||||
response = self.ec2_client.describe_instances(InstanceIds=[self.node_id])
|
||||
reservations = response["Reservations"]
|
||||
message = "More than 1 response received from describing current node"
|
||||
assert len(reservations) == 1, message
|
||||
instances = reservations[0]["Instances"]
|
||||
assert len(instances) == 1, message
|
||||
tags = instances[0]["Tags"]
|
||||
cwa_installed = str(False)
|
||||
for tag in tags:
|
||||
if tag["Key"] == CLOUDWATCH_AGENT_INSTALLED_TAG:
|
||||
logger.info(
|
||||
"Unified CloudWatch Agent is installed on "
|
||||
"node {}".format(self.node_id)
|
||||
)
|
||||
cwa_installed = tag["Value"]
|
||||
return cwa_installed
|
||||
except botocore.exceptions.ClientError as e:
|
||||
logger.warning(
|
||||
"{} Error caught when getting Unified CloudWatch Agent "
|
||||
"status based on {} tag".format(
|
||||
e.response["Error"], CLOUDWATCH_AGENT_INSTALLED_TAG
|
||||
)
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def resolve_instance_profile_name(
|
||||
config: Dict[str, Any], default_instance_profile_name: str
|
||||
) -> str:
|
||||
"""Get default cloudwatch instance profile name.
|
||||
|
||||
Args:
|
||||
config: provider section of cluster config file.
|
||||
default_instance_profile_name: default ray instance profile name.
|
||||
|
||||
Returns:
|
||||
default cloudwatch instance profile name if cloudwatch config file
|
||||
exists.
|
||||
default ray instance profile name if cloudwatch config file
|
||||
doesn't exist.
|
||||
"""
|
||||
cwa_cfg_exists = CloudwatchHelper.cloudwatch_config_exists(
|
||||
config, CloudwatchConfigType.AGENT.value
|
||||
)
|
||||
return (
|
||||
CLOUDWATCH_RAY_INSTANCE_PROFILE
|
||||
if cwa_cfg_exists
|
||||
else default_instance_profile_name
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def resolve_iam_role_name(
|
||||
config: Dict[str, Any], default_iam_role_name: str
|
||||
) -> str:
|
||||
"""Get default cloudwatch iam role name.
|
||||
|
||||
Args:
|
||||
config: provider section of cluster config file.
|
||||
default_iam_role_name: default ray iam role name.
|
||||
|
||||
Returns:
|
||||
default cloudwatch iam role name if cloudwatch config file exists.
|
||||
default ray iam role name if cloudwatch config file doesn't exist.
|
||||
"""
|
||||
cwa_cfg_exists = CloudwatchHelper.cloudwatch_config_exists(
|
||||
config, CloudwatchConfigType.AGENT.value
|
||||
)
|
||||
return CLOUDWATCH_RAY_IAM_ROLE if cwa_cfg_exists else default_iam_role_name
|
||||
|
||||
@staticmethod
|
||||
def resolve_policy_arns(
|
||||
config: Dict[str, Any], iam: Any, default_policy_arns: List[str]
|
||||
) -> List[str]:
|
||||
"""Attach necessary AWS policies for CloudWatch related operations.
|
||||
|
||||
Args:
|
||||
config: provider section of cluster config file.
|
||||
iam: AWS iam resource.
|
||||
default_policy_arns: List of default ray AWS policies.
|
||||
|
||||
Returns:
|
||||
list of policy arns including additional policies for CloudWatch
|
||||
related operations if cloudwatch agent config is specifed in
|
||||
cluster config file.
|
||||
"""
|
||||
cwa_cfg_exists = CloudwatchHelper.cloudwatch_config_exists(
|
||||
config, CloudwatchConfigType.AGENT.value
|
||||
)
|
||||
if cwa_cfg_exists:
|
||||
cloudwatch_managed_policy = {
|
||||
"Version": "2012-10-17",
|
||||
"Statement": [
|
||||
{
|
||||
"Effect": "Allow",
|
||||
"Action": [
|
||||
"ssm:SendCommand",
|
||||
"ssm:ListCommandInvocations",
|
||||
"iam:PassRole",
|
||||
],
|
||||
"Resource": "*",
|
||||
}
|
||||
],
|
||||
}
|
||||
iam_client = iam.meta.client
|
||||
iam_client.create_policy(
|
||||
PolicyName="CloudwatchManagedPolicies",
|
||||
PolicyDocument=json.dumps(cloudwatch_managed_policy),
|
||||
)
|
||||
sts_client = client_cache("sts", config["region"])
|
||||
account_id = sts_client.get_caller_identity().get("Account")
|
||||
managed_policy_arn = (
|
||||
"arn:aws:iam::{}:policy/CloudwatchManagedPolicies".format(account_id)
|
||||
)
|
||||
policy_waiter = iam_client.get_waiter("policy_exists")
|
||||
policy_waiter.wait(
|
||||
PolicyArn=managed_policy_arn,
|
||||
WaiterConfig={"Delay": 2, "MaxAttempts": 200},
|
||||
)
|
||||
new_policy_arns = copy.copy(default_policy_arns)
|
||||
new_policy_arns.extend(
|
||||
[
|
||||
"arn:aws:iam::aws:policy/CloudWatchAgentAdminPolicy",
|
||||
"arn:aws:iam::aws:policy/AmazonSSMManagedInstanceCore",
|
||||
managed_policy_arn,
|
||||
]
|
||||
)
|
||||
return new_policy_arns
|
||||
else:
|
||||
return default_policy_arns
|
||||
|
||||
@staticmethod
|
||||
def cloudwatch_config_exists(config: Dict[str, Any], config_type: str) -> bool:
|
||||
"""Check if CloudWatch configuration was specified by the user
|
||||
in their cluster config file.
|
||||
|
||||
Specifically, this function checks if a CloudWatch config file is
|
||||
specified by the user in their cluster config file.
|
||||
|
||||
Args:
|
||||
config: provider section of cluster config file.
|
||||
config_type: type of CloudWatch config file.
|
||||
|
||||
Returns:
|
||||
True if config file is specified by user.
|
||||
False if config file is not specified.
|
||||
"""
|
||||
cfg = config.get("cloudwatch", {}).get(config_type, {}).get("config")
|
||||
return bool(cfg)
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,715 @@
|
||||
import copy
|
||||
import logging
|
||||
import sys
|
||||
import threading
|
||||
import time
|
||||
from collections import OrderedDict, defaultdict
|
||||
from typing import Any, Dict, List
|
||||
|
||||
import botocore
|
||||
from boto3.resources.base import ServiceResource
|
||||
|
||||
import ray
|
||||
import ray._private.ray_constants as ray_constants
|
||||
from ray.autoscaler._private.aws.cloudwatch.cloudwatch_helper import (
|
||||
CLOUDWATCH_AGENT_INSTALLED_AMI_TAG,
|
||||
CLOUDWATCH_AGENT_INSTALLED_TAG,
|
||||
CloudwatchHelper,
|
||||
)
|
||||
from ray.autoscaler._private.aws.config import bootstrap_aws
|
||||
from ray.autoscaler._private.aws.utils import (
|
||||
boto_exception_handler,
|
||||
client_cache,
|
||||
resource_cache,
|
||||
)
|
||||
from ray.autoscaler._private.cli_logger import cf, cli_logger
|
||||
from ray.autoscaler._private.constants import BOTO_CREATE_MAX_RETRIES, BOTO_MAX_RETRIES
|
||||
from ray.autoscaler._private.log_timer import LogTimer
|
||||
from ray.autoscaler.node_launch_exception import NodeLaunchException
|
||||
from ray.autoscaler.node_provider import NodeProvider
|
||||
from ray.autoscaler.tags import (
|
||||
TAG_RAY_CLUSTER_NAME,
|
||||
TAG_RAY_LAUNCH_CONFIG,
|
||||
TAG_RAY_NODE_KIND,
|
||||
TAG_RAY_NODE_NAME,
|
||||
TAG_RAY_USER_NODE_TYPE,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
TAG_BATCH_DELAY = 1
|
||||
LIST_RETRY_DELAY_SEC = 1
|
||||
|
||||
|
||||
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
|
||||
|
||||
|
||||
def make_ec2_resource(region, max_retries, aws_credentials=None) -> ServiceResource:
|
||||
"""Make client, retrying requests up to `max_retries`."""
|
||||
aws_credentials = aws_credentials or {}
|
||||
return resource_cache("ec2", region, max_retries, **aws_credentials)
|
||||
|
||||
|
||||
def list_ec2_instances(
|
||||
region: str, aws_credentials: Dict[str, Any] = None
|
||||
) -> List[Dict[str, Any]]:
|
||||
"""Get all instance-types/resources available in the user's AWS region.
|
||||
|
||||
Args:
|
||||
region: the region of the AWS provider. e.g., "us-west-2".
|
||||
aws_credentials: AWS credentials to use for the boto3 client.
|
||||
|
||||
Returns:
|
||||
A list of instances. An example of one element in the list:
|
||||
{'InstanceType': 'm5a.xlarge', 'ProcessorInfo':
|
||||
{'SupportedArchitectures': ['x86_64'], 'SustainedClockSpeedInGhz':
|
||||
2.5},'VCpuInfo': {'DefaultVCpus': 4, 'DefaultCores': 2,
|
||||
'DefaultThreadsPerCore': 2, 'ValidCores': [2],
|
||||
'ValidThreadsPerCore': [1, 2]}, 'MemoryInfo': {'SizeInMiB': 16384},
|
||||
...}
|
||||
|
||||
"""
|
||||
final_instance_types = []
|
||||
aws_credentials = aws_credentials or {}
|
||||
ec2 = client_cache("ec2", region, BOTO_MAX_RETRIES, **aws_credentials)
|
||||
instance_types = ec2.describe_instance_types()
|
||||
final_instance_types.extend(copy.deepcopy(instance_types["InstanceTypes"]))
|
||||
while "NextToken" in instance_types:
|
||||
instance_types = ec2.describe_instance_types(
|
||||
NextToken=instance_types["NextToken"]
|
||||
)
|
||||
final_instance_types.extend(copy.deepcopy(instance_types["InstanceTypes"]))
|
||||
|
||||
return final_instance_types
|
||||
|
||||
|
||||
class AWSNodeProvider(NodeProvider):
|
||||
max_terminate_nodes = 1000
|
||||
|
||||
def __init__(self, provider_config, cluster_name):
|
||||
NodeProvider.__init__(self, provider_config, cluster_name)
|
||||
self.cache_stopped_nodes = provider_config.get("cache_stopped_nodes", True)
|
||||
aws_credentials = provider_config.get("aws_credentials")
|
||||
|
||||
self.ec2 = make_ec2_resource(
|
||||
region=provider_config["region"],
|
||||
max_retries=BOTO_MAX_RETRIES,
|
||||
aws_credentials=aws_credentials,
|
||||
)
|
||||
self.ec2_fail_fast = make_ec2_resource(
|
||||
region=provider_config["region"],
|
||||
max_retries=0,
|
||||
aws_credentials=aws_credentials,
|
||||
)
|
||||
|
||||
# Tags that we believe to actually be on EC2.
|
||||
self.tag_cache = {}
|
||||
# Tags that we will soon upload.
|
||||
self.tag_cache_pending = defaultdict(dict)
|
||||
# Number of threads waiting for a batched tag update.
|
||||
self.batch_thread_count = 0
|
||||
self.batch_update_done = threading.Event()
|
||||
self.batch_update_done.set()
|
||||
self.ready_for_new_batch = threading.Event()
|
||||
self.ready_for_new_batch.set()
|
||||
self.tag_cache_lock = threading.Lock()
|
||||
self.count_lock = threading.Lock()
|
||||
# Prevent concurrent create_node calls to get the same stopped/stopping node to reuse.
|
||||
self._reuse_node_lock = threading.Lock()
|
||||
|
||||
# Cache of node objects from the last nodes() call. This avoids
|
||||
# excessive DescribeInstances requests.
|
||||
self.cached_nodes = {}
|
||||
|
||||
def non_terminated_nodes(self, tag_filters):
|
||||
# Note that these filters are acceptable because they are set on
|
||||
# node initialization, and so can never be sitting in the cache.
|
||||
tag_filters = to_aws_format(tag_filters)
|
||||
filters = [
|
||||
{
|
||||
"Name": "instance-state-name",
|
||||
"Values": ["pending", "running"],
|
||||
},
|
||||
{
|
||||
"Name": "tag:{}".format(TAG_RAY_CLUSTER_NAME),
|
||||
"Values": [self.cluster_name],
|
||||
},
|
||||
]
|
||||
for k, v in tag_filters.items():
|
||||
filters.append(
|
||||
{
|
||||
"Name": "tag:{}".format(k),
|
||||
"Values": [v],
|
||||
}
|
||||
)
|
||||
|
||||
with boto_exception_handler("Failed to fetch running instances from AWS."):
|
||||
nodes = list(self.ec2.instances.filter(Filters=filters))
|
||||
|
||||
# Populate the tag cache with initial information if necessary
|
||||
for node in nodes:
|
||||
if node.id in self.tag_cache:
|
||||
continue
|
||||
|
||||
self.tag_cache[node.id] = from_aws_format(
|
||||
{x["Key"]: x["Value"] for x in node.tags}
|
||||
)
|
||||
|
||||
self.cached_nodes = {node.id: node for node in nodes}
|
||||
return [node.id for node in nodes]
|
||||
|
||||
def is_running(self, node_id):
|
||||
node = self._get_cached_node(node_id)
|
||||
return node.state["Name"] == "running"
|
||||
|
||||
def is_terminated(self, node_id):
|
||||
node = self._get_cached_node(node_id)
|
||||
state = node.state["Name"]
|
||||
return state not in ["running", "pending"]
|
||||
|
||||
def node_tags(self, node_id):
|
||||
with self.tag_cache_lock:
|
||||
d1 = self.tag_cache[node_id]
|
||||
d2 = self.tag_cache_pending.get(node_id, {})
|
||||
return dict(d1, **d2)
|
||||
|
||||
def external_ip(self, node_id):
|
||||
node = self._get_cached_node(node_id)
|
||||
|
||||
if node.public_ip_address is None:
|
||||
node = self._get_node(node_id)
|
||||
|
||||
return node.public_ip_address
|
||||
|
||||
def internal_ip(self, node_id):
|
||||
node = self._get_cached_node(node_id)
|
||||
|
||||
if node.private_ip_address is None:
|
||||
node = self._get_node(node_id)
|
||||
|
||||
return node.private_ip_address
|
||||
|
||||
def set_node_tags(self, node_id, tags):
|
||||
is_batching_thread = False
|
||||
with self.tag_cache_lock:
|
||||
if not self.tag_cache_pending:
|
||||
is_batching_thread = True
|
||||
# Wait for threads in the last batch to exit
|
||||
self.ready_for_new_batch.wait()
|
||||
self.ready_for_new_batch.clear()
|
||||
self.batch_update_done.clear()
|
||||
self.tag_cache_pending[node_id].update(tags)
|
||||
|
||||
if is_batching_thread:
|
||||
time.sleep(TAG_BATCH_DELAY)
|
||||
with self.tag_cache_lock:
|
||||
self._update_node_tags()
|
||||
self.batch_update_done.set()
|
||||
|
||||
with self.count_lock:
|
||||
self.batch_thread_count += 1
|
||||
self.batch_update_done.wait()
|
||||
|
||||
with self.count_lock:
|
||||
self.batch_thread_count -= 1
|
||||
if self.batch_thread_count == 0:
|
||||
self.ready_for_new_batch.set()
|
||||
|
||||
def _update_node_tags(self):
|
||||
batch_updates = defaultdict(list)
|
||||
|
||||
for node_id, tags in self.tag_cache_pending.items():
|
||||
for x in tags.items():
|
||||
batch_updates[x].append(node_id)
|
||||
self.tag_cache[node_id].update(tags)
|
||||
|
||||
self.tag_cache_pending = defaultdict(dict)
|
||||
|
||||
self._create_tags(batch_updates)
|
||||
|
||||
def _create_tags(self, batch_updates):
|
||||
for (k, v), node_ids in batch_updates.items():
|
||||
m = "Set tag {}={} on {}".format(k, v, node_ids)
|
||||
with LogTimer("AWSNodeProvider: {}".format(m)):
|
||||
if k == TAG_RAY_NODE_NAME:
|
||||
k = "Name"
|
||||
self.ec2.meta.client.create_tags(
|
||||
Resources=node_ids,
|
||||
Tags=[{"Key": k, "Value": v}],
|
||||
)
|
||||
|
||||
def create_node(self, node_config, tags, count) -> Dict[str, Any]:
|
||||
"""Creates instances.
|
||||
|
||||
Returns dict mapping instance id to ec2.Instance object for the created
|
||||
instances.
|
||||
"""
|
||||
# sort tags by key to support deterministic unit test stubbing
|
||||
tags = OrderedDict(sorted(copy.deepcopy(tags).items()))
|
||||
|
||||
reused_nodes_dict = {}
|
||||
# Try to reuse previously stopped nodes with compatible configs
|
||||
if self.cache_stopped_nodes:
|
||||
# TODO(ekl) this is breaking the abstraction boundary a little by
|
||||
# peeking into the tag set.
|
||||
filters = [
|
||||
{
|
||||
"Name": "instance-state-name",
|
||||
"Values": ["stopped", "stopping"],
|
||||
},
|
||||
{
|
||||
"Name": "tag:{}".format(TAG_RAY_CLUSTER_NAME),
|
||||
"Values": [self.cluster_name],
|
||||
},
|
||||
{
|
||||
"Name": "tag:{}".format(TAG_RAY_NODE_KIND),
|
||||
"Values": [tags[TAG_RAY_NODE_KIND]],
|
||||
},
|
||||
{
|
||||
"Name": "tag:{}".format(TAG_RAY_LAUNCH_CONFIG),
|
||||
"Values": [tags[TAG_RAY_LAUNCH_CONFIG]],
|
||||
},
|
||||
]
|
||||
# This tag may not always be present.
|
||||
if TAG_RAY_USER_NODE_TYPE in tags:
|
||||
filters.append(
|
||||
{
|
||||
"Name": "tag:{}".format(TAG_RAY_USER_NODE_TYPE),
|
||||
"Values": [tags[TAG_RAY_USER_NODE_TYPE]],
|
||||
}
|
||||
)
|
||||
|
||||
with self._reuse_node_lock:
|
||||
reuse_nodes = list(self.ec2.instances.filter(Filters=filters))[:count]
|
||||
reuse_node_ids = [n.id for n in reuse_nodes]
|
||||
reused_nodes_dict = {n.id: n for n in reuse_nodes}
|
||||
if reuse_nodes:
|
||||
cli_logger.print(
|
||||
# todo: handle plural vs singular?
|
||||
"Reusing nodes {}. "
|
||||
"To disable reuse, set `cache_stopped_nodes: False` "
|
||||
"under `provider` in the cluster configuration.",
|
||||
cli_logger.render_list(reuse_node_ids),
|
||||
)
|
||||
|
||||
# todo: timed?
|
||||
with cli_logger.group("Stopping instances to reuse"):
|
||||
for node in reuse_nodes:
|
||||
self.tag_cache[node.id] = from_aws_format(
|
||||
{x["Key"]: x["Value"] for x in node.tags}
|
||||
)
|
||||
if node.state["Name"] == "stopping":
|
||||
cli_logger.print(
|
||||
"Waiting for instance {} to stop", node.id
|
||||
)
|
||||
node.wait_until_stopped()
|
||||
|
||||
self.ec2.meta.client.start_instances(InstanceIds=reuse_node_ids)
|
||||
for node_id in reuse_node_ids:
|
||||
self.set_node_tags(node_id, tags)
|
||||
count -= len(reuse_node_ids)
|
||||
|
||||
created_nodes_dict = {}
|
||||
if count:
|
||||
created_nodes_dict = self._create_node(node_config, tags, count)
|
||||
|
||||
all_created_nodes = reused_nodes_dict
|
||||
all_created_nodes.update(created_nodes_dict)
|
||||
return all_created_nodes
|
||||
|
||||
@staticmethod
|
||||
def _merge_tag_specs(
|
||||
tag_specs: List[Dict[str, Any]], user_tag_specs: List[Dict[str, Any]]
|
||||
) -> None:
|
||||
"""
|
||||
Merges user-provided node config tag specifications into a base
|
||||
list of node provider tag specifications. The base list of
|
||||
node provider tag specs is modified in-place.
|
||||
|
||||
This allows users to add tags and override values of existing
|
||||
tags with their own, and only applies to the resource type
|
||||
"instance". All other resource types are appended to the list of
|
||||
tag specs.
|
||||
|
||||
Args:
|
||||
tag_specs: base node provider tag specs
|
||||
user_tag_specs: user's node config tag specs
|
||||
"""
|
||||
|
||||
for user_tag_spec in user_tag_specs:
|
||||
if user_tag_spec["ResourceType"] == "instance":
|
||||
for user_tag in user_tag_spec["Tags"]:
|
||||
exists = False
|
||||
for tag in tag_specs[0]["Tags"]:
|
||||
if user_tag["Key"] == tag["Key"]:
|
||||
exists = True
|
||||
tag["Value"] = user_tag["Value"]
|
||||
break
|
||||
if not exists:
|
||||
tag_specs[0]["Tags"] += [user_tag]
|
||||
else:
|
||||
tag_specs += [user_tag_spec]
|
||||
|
||||
def _create_node(self, node_config, tags, count):
|
||||
created_nodes_dict = {}
|
||||
|
||||
tags = to_aws_format(tags)
|
||||
conf = node_config.copy()
|
||||
|
||||
tag_pairs = [
|
||||
{
|
||||
"Key": TAG_RAY_CLUSTER_NAME,
|
||||
"Value": self.cluster_name,
|
||||
}
|
||||
]
|
||||
for k, v in tags.items():
|
||||
tag_pairs.append(
|
||||
{
|
||||
"Key": k,
|
||||
"Value": v,
|
||||
}
|
||||
)
|
||||
if CloudwatchHelper.cloudwatch_config_exists(self.provider_config, "agent"):
|
||||
cwa_installed = self._check_ami_cwa_installation(node_config)
|
||||
if cwa_installed:
|
||||
tag_pairs.extend(
|
||||
[
|
||||
{
|
||||
"Key": CLOUDWATCH_AGENT_INSTALLED_TAG,
|
||||
"Value": "True",
|
||||
}
|
||||
]
|
||||
)
|
||||
tag_specs = [
|
||||
{
|
||||
"ResourceType": "instance",
|
||||
"Tags": tag_pairs,
|
||||
}
|
||||
]
|
||||
user_tag_specs = conf.get("TagSpecifications", [])
|
||||
AWSNodeProvider._merge_tag_specs(tag_specs, user_tag_specs)
|
||||
|
||||
# SubnetIds is not a real config key: we must resolve to a
|
||||
# single SubnetId before invoking the AWS API.
|
||||
subnet_ids = conf.pop("SubnetIds")
|
||||
|
||||
# update config with min/max node counts and tag specs
|
||||
conf.update({"MinCount": 1, "MaxCount": count, "TagSpecifications": tag_specs})
|
||||
|
||||
# Try to always launch in the first listed subnet.
|
||||
subnet_idx = 0
|
||||
cli_logger_tags = {}
|
||||
# NOTE: This ensures that we try ALL availability zones before
|
||||
# throwing an error.
|
||||
max_tries = max(BOTO_CREATE_MAX_RETRIES, len(subnet_ids))
|
||||
for attempt in range(1, max_tries + 1):
|
||||
try:
|
||||
if "NetworkInterfaces" in conf:
|
||||
net_ifs = conf["NetworkInterfaces"]
|
||||
# remove security group IDs previously copied from network
|
||||
# interfaces (create_instances call fails otherwise)
|
||||
conf.pop("SecurityGroupIds", None)
|
||||
cli_logger_tags["network_interfaces"] = str(net_ifs)
|
||||
else:
|
||||
subnet_id = subnet_ids[subnet_idx % len(subnet_ids)]
|
||||
conf["SubnetId"] = subnet_id
|
||||
cli_logger_tags["subnet_id"] = subnet_id
|
||||
|
||||
created = self.ec2_fail_fast.create_instances(**conf)
|
||||
created_nodes_dict = {n.id: n for n in created}
|
||||
|
||||
# todo: timed?
|
||||
# todo: handle plurality?
|
||||
with cli_logger.group(
|
||||
"Launched {} nodes", count, _tags=cli_logger_tags
|
||||
):
|
||||
for instance in created:
|
||||
# NOTE(maximsmol): This is needed for mocking
|
||||
# boto3 for tests. This is likely a bug in moto
|
||||
# but AWS docs don't seem to say.
|
||||
# You can patch moto/ec2/responses/instances.py
|
||||
# to fix this (add <stateReason> to EC2_RUN_INSTANCES)
|
||||
|
||||
# The correct value is technically
|
||||
# {"code": "0", "Message": "pending"}
|
||||
state_reason = "pending"
|
||||
if instance.state_reason:
|
||||
state_reason = (
|
||||
instance.state_reason["Message"] or state_reason
|
||||
)
|
||||
|
||||
cli_logger.print(
|
||||
"Launched instance {}",
|
||||
instance.instance_id,
|
||||
_tags=dict(
|
||||
state=instance.state["Name"],
|
||||
info=state_reason,
|
||||
),
|
||||
)
|
||||
break
|
||||
except botocore.exceptions.ClientError as exc:
|
||||
# Launch failure may be due to instance type availability in
|
||||
# the given AZ
|
||||
subnet_idx += 1
|
||||
if attempt == max_tries:
|
||||
try:
|
||||
exc = NodeLaunchException(
|
||||
category=exc.response["Error"]["Code"],
|
||||
description=exc.response["Error"]["Message"],
|
||||
src_exc_info=sys.exc_info(),
|
||||
)
|
||||
except Exception:
|
||||
# In theory, all ClientError's we expect to get should
|
||||
# have these fields, but just in case we can't parse
|
||||
# it, it's fine, just throw the original error.
|
||||
logger.warning("Couldn't parse exception.", exc)
|
||||
pass
|
||||
cli_logger.abort(
|
||||
"Failed to launch instances. Max attempts exceeded.",
|
||||
exc=exc,
|
||||
)
|
||||
else:
|
||||
cli_logger.warning(
|
||||
"create_instances: Attempt failed with {}, retrying.", exc
|
||||
)
|
||||
|
||||
return created_nodes_dict
|
||||
|
||||
def terminate_node(self, node_id):
|
||||
node = self._get_cached_node(node_id)
|
||||
if self.cache_stopped_nodes:
|
||||
if node.spot_instance_request_id:
|
||||
cli_logger.print(
|
||||
"Terminating instance {} "
|
||||
+ cf.dimmed("(cannot stop spot instances, only terminate)"),
|
||||
node_id,
|
||||
) # todo: show node name?
|
||||
node.terminate()
|
||||
else:
|
||||
cli_logger.print(
|
||||
"Stopping instance {} "
|
||||
+ cf.dimmed(
|
||||
"(to terminate instead, "
|
||||
"set `cache_stopped_nodes: False` "
|
||||
"under `provider` in the cluster configuration)"
|
||||
),
|
||||
node_id,
|
||||
) # todo: show node name?
|
||||
node.stop()
|
||||
else:
|
||||
node.terminate()
|
||||
|
||||
# TODO (Alex): We are leaking the tag cache here. Naively, we would
|
||||
# want to just remove the cache entry here, but terminating can be
|
||||
# asyncrhonous or error, which would result in a use after free error.
|
||||
# If this leak becomes bad, we can garbage collect the tag cache when
|
||||
# the node cache is updated.
|
||||
|
||||
def _check_ami_cwa_installation(self, config):
|
||||
response = self.ec2.meta.client.describe_images(ImageIds=[config["ImageId"]])
|
||||
cwa_installed = False
|
||||
images = response.get("Images")
|
||||
if images:
|
||||
assert len(images) == 1, (
|
||||
f"Expected to find only 1 AMI with the given ID, "
|
||||
f"but found {len(images)}."
|
||||
)
|
||||
image_name = images[0].get("Name", "")
|
||||
if CLOUDWATCH_AGENT_INSTALLED_AMI_TAG in image_name:
|
||||
cwa_installed = True
|
||||
return cwa_installed
|
||||
|
||||
def terminate_nodes(self, node_ids):
|
||||
if not node_ids:
|
||||
return
|
||||
|
||||
terminate_instances_func = self.ec2.meta.client.terminate_instances
|
||||
stop_instances_func = self.ec2.meta.client.stop_instances
|
||||
|
||||
# In some cases, this function stops some nodes, but terminates others.
|
||||
# Each of these requires a different EC2 API call. So, we use the
|
||||
# "nodes_to_terminate" dict below to keep track of exactly which API
|
||||
# call will be used to stop/terminate which set of nodes. The key is
|
||||
# the function to use, and the value is the list of nodes to terminate
|
||||
# with that function.
|
||||
nodes_to_terminate = {terminate_instances_func: [], stop_instances_func: []}
|
||||
|
||||
if self.cache_stopped_nodes:
|
||||
spot_ids = []
|
||||
on_demand_ids = []
|
||||
|
||||
for node_id in node_ids:
|
||||
if self._get_cached_node(node_id).spot_instance_request_id:
|
||||
spot_ids += [node_id]
|
||||
else:
|
||||
on_demand_ids += [node_id]
|
||||
|
||||
if on_demand_ids:
|
||||
# todo: show node names?
|
||||
cli_logger.print(
|
||||
"Stopping instances {} "
|
||||
+ cf.dimmed(
|
||||
"(to terminate instead, "
|
||||
"set `cache_stopped_nodes: False` "
|
||||
"under `provider` in the cluster configuration)"
|
||||
),
|
||||
cli_logger.render_list(on_demand_ids),
|
||||
)
|
||||
|
||||
if spot_ids:
|
||||
cli_logger.print(
|
||||
"Terminating instances {} "
|
||||
+ cf.dimmed("(cannot stop spot instances, only terminate)"),
|
||||
cli_logger.render_list(spot_ids),
|
||||
)
|
||||
|
||||
nodes_to_terminate[stop_instances_func] = on_demand_ids
|
||||
nodes_to_terminate[terminate_instances_func] = spot_ids
|
||||
else:
|
||||
nodes_to_terminate[terminate_instances_func] = node_ids
|
||||
|
||||
max_terminate_nodes = (
|
||||
self.max_terminate_nodes
|
||||
if self.max_terminate_nodes is not None
|
||||
else len(node_ids)
|
||||
)
|
||||
|
||||
for terminate_func, nodes in nodes_to_terminate.items():
|
||||
for start in range(0, len(nodes), max_terminate_nodes):
|
||||
terminate_func(InstanceIds=nodes[start : start + max_terminate_nodes])
|
||||
|
||||
def _get_node(self, node_id):
|
||||
"""Refresh and get info for this node, updating the cache."""
|
||||
self.non_terminated_nodes({}) # Side effect: updates cache
|
||||
|
||||
if node_id in self.cached_nodes:
|
||||
return self.cached_nodes[node_id]
|
||||
|
||||
# Node not in {pending, running} -- retry with a point query. This
|
||||
# usually means the node was recently preempted or terminated.
|
||||
# The EC2 API is eventually consistent. This means that an instance
|
||||
# might not be immediately visible. So we need to retry the query a few times.
|
||||
# See: https://docs.aws.amazon.com/ec2/latest/devguide/eventual-consistency.html
|
||||
# and https://github.com/ray-project/ray/issues/51861
|
||||
for attempts in range(max(BOTO_MAX_RETRIES, 1)): # at least try once.
|
||||
matches = list(self.ec2.instances.filter(InstanceIds=[node_id]))
|
||||
if len(matches) == 1:
|
||||
return matches[0]
|
||||
cli_logger.warning(
|
||||
"Attempt to fetch EC2 instances that have instance ID {}. Got {} matching EC2 instances. Will retry after {} second. This is retry number {}, and the maximum number of retries is {}.",
|
||||
node_id,
|
||||
len(matches),
|
||||
LIST_RETRY_DELAY_SEC,
|
||||
attempts + 1,
|
||||
BOTO_MAX_RETRIES,
|
||||
)
|
||||
time.sleep(LIST_RETRY_DELAY_SEC)
|
||||
raise AssertionError("Invalid instance id {}".format(node_id))
|
||||
|
||||
def _get_cached_node(self, node_id):
|
||||
"""Return node info from cache if possible, otherwise fetches it."""
|
||||
if node_id in self.cached_nodes:
|
||||
return self.cached_nodes[node_id]
|
||||
|
||||
return self._get_node(node_id)
|
||||
|
||||
@staticmethod
|
||||
def bootstrap_config(cluster_config):
|
||||
return bootstrap_aws(cluster_config)
|
||||
|
||||
@staticmethod
|
||||
def fillout_available_node_types_resources(
|
||||
cluster_config: Dict[str, Any]
|
||||
) -> Dict[str, Any]:
|
||||
"""Fills out missing "resources" field for available_node_types."""
|
||||
if "available_node_types" not in cluster_config:
|
||||
return cluster_config
|
||||
cluster_config = copy.deepcopy(cluster_config)
|
||||
|
||||
instances_list = list_ec2_instances(
|
||||
cluster_config["provider"]["region"],
|
||||
cluster_config["provider"].get("aws_credentials"),
|
||||
)
|
||||
instances_dict = {
|
||||
instance["InstanceType"]: instance for instance in instances_list
|
||||
}
|
||||
available_node_types = cluster_config["available_node_types"]
|
||||
head_node_type = cluster_config["head_node_type"]
|
||||
for node_type in available_node_types:
|
||||
instance_type = available_node_types[node_type]["node_config"][
|
||||
"InstanceType"
|
||||
]
|
||||
if instance_type in instances_dict:
|
||||
cpus = instances_dict[instance_type]["VCpuInfo"]["DefaultVCpus"]
|
||||
|
||||
autodetected_resources = {"CPU": cpus}
|
||||
if node_type != head_node_type:
|
||||
# we only autodetect worker node type memory resource
|
||||
memory_total = instances_dict[instance_type]["MemoryInfo"][
|
||||
"SizeInMiB"
|
||||
]
|
||||
memory_total = int(memory_total) * 1024 * 1024
|
||||
prop = 1 - ray_constants.DEFAULT_OBJECT_STORE_MEMORY_PROPORTION
|
||||
memory_resources = int(memory_total * prop)
|
||||
autodetected_resources["memory"] = memory_resources
|
||||
|
||||
for (
|
||||
accelerator_manager
|
||||
) in ray._private.accelerators.get_all_accelerator_managers():
|
||||
num_accelerators = (
|
||||
accelerator_manager.get_ec2_instance_num_accelerators(
|
||||
instance_type, instances_dict
|
||||
)
|
||||
)
|
||||
accelerator_type = (
|
||||
accelerator_manager.get_ec2_instance_accelerator_type(
|
||||
instance_type, instances_dict
|
||||
)
|
||||
)
|
||||
if num_accelerators:
|
||||
autodetected_resources[
|
||||
accelerator_manager.get_resource_name()
|
||||
] = num_accelerators
|
||||
if accelerator_type:
|
||||
autodetected_resources[
|
||||
f"accelerator_type:{accelerator_type}"
|
||||
] = 1
|
||||
|
||||
autodetected_resources.update(
|
||||
available_node_types[node_type].get("resources", {})
|
||||
)
|
||||
if autodetected_resources != available_node_types[node_type].get(
|
||||
"resources", {}
|
||||
):
|
||||
available_node_types[node_type][
|
||||
"resources"
|
||||
] = autodetected_resources
|
||||
logger.debug(
|
||||
"Updating the resources of {} to {}.".format(
|
||||
node_type, autodetected_resources
|
||||
)
|
||||
)
|
||||
else:
|
||||
raise ValueError(
|
||||
"Instance type "
|
||||
+ instance_type
|
||||
+ " is not available in AWS region: "
|
||||
+ cluster_config["provider"]["region"]
|
||||
+ "."
|
||||
)
|
||||
return cluster_config
|
||||
@@ -0,0 +1,181 @@
|
||||
from collections import defaultdict
|
||||
from functools import lru_cache
|
||||
|
||||
import boto3
|
||||
from boto3.exceptions import ResourceNotExistsError
|
||||
from boto3.resources.base import ServiceResource
|
||||
from botocore.client import BaseClient
|
||||
from botocore.config import Config
|
||||
|
||||
from ray.autoscaler._private.cli_logger import cf, cli_logger
|
||||
from ray.autoscaler._private.constants import BOTO_MAX_RETRIES
|
||||
|
||||
|
||||
class LazyDefaultDict(defaultdict):
|
||||
"""
|
||||
LazyDefaultDict(default_factory[, ...]) --> dict with default factory
|
||||
|
||||
The default factory is call with the key argument to produce
|
||||
a new value when a key is not present, in __getitem__ only.
|
||||
A LazyDefaultDict compares equal to a dict with the same items.
|
||||
All remaining arguments are treated the same as if they were
|
||||
passed to the dict constructor, including keyword arguments.
|
||||
"""
|
||||
|
||||
def __missing__(self, key):
|
||||
"""
|
||||
__missing__(key) # Called by __getitem__ for missing key; pseudo-code:
|
||||
if self.default_factory is None: raise KeyError((key,))
|
||||
self[key] = value = self.default_factory(key)
|
||||
return value
|
||||
"""
|
||||
self[key] = self.default_factory(key)
|
||||
return self[key]
|
||||
|
||||
|
||||
def handle_boto_error(exc, msg, *args, **kwargs):
|
||||
error_code = None
|
||||
error_info = None
|
||||
# todo: not sure if these exceptions always have response
|
||||
if hasattr(exc, "response"):
|
||||
error_info = exc.response.get("Error", None)
|
||||
if error_info is not None:
|
||||
error_code = error_info.get("Code", None)
|
||||
|
||||
generic_message_args = [
|
||||
"{}\nError code: {}",
|
||||
msg.format(*args, **kwargs),
|
||||
cf.bold(error_code),
|
||||
]
|
||||
|
||||
# apparently
|
||||
# ExpiredTokenException
|
||||
# ExpiredToken
|
||||
# RequestExpired
|
||||
# are all the same pretty much
|
||||
credentials_expiration_codes = [
|
||||
"ExpiredTokenException",
|
||||
"ExpiredToken",
|
||||
"RequestExpired",
|
||||
]
|
||||
|
||||
if error_code in credentials_expiration_codes:
|
||||
# "An error occurred (ExpiredToken) when calling the
|
||||
# GetInstanceProfile operation: The security token
|
||||
# included in the request is expired"
|
||||
|
||||
# "An error occurred (RequestExpired) when calling the
|
||||
# DescribeKeyPairs operation: Request has expired."
|
||||
|
||||
token_command = (
|
||||
"aws sts get-session-token "
|
||||
"--serial-number arn:aws:iam::"
|
||||
+ cf.underlined("ROOT_ACCOUNT_ID")
|
||||
+ ":mfa/"
|
||||
+ cf.underlined("AWS_USERNAME")
|
||||
+ " --token-code "
|
||||
+ cf.underlined("TWO_FACTOR_AUTH_CODE")
|
||||
)
|
||||
|
||||
secret_key_var = (
|
||||
"export AWS_SECRET_ACCESS_KEY = "
|
||||
+ cf.underlined("REPLACE_ME")
|
||||
+ " # found at Credentials.SecretAccessKey"
|
||||
)
|
||||
session_token_var = (
|
||||
"export AWS_SESSION_TOKEN = "
|
||||
+ cf.underlined("REPLACE_ME")
|
||||
+ " # found at Credentials.SessionToken"
|
||||
)
|
||||
access_key_id_var = (
|
||||
"export AWS_ACCESS_KEY_ID = "
|
||||
+ cf.underlined("REPLACE_ME")
|
||||
+ " # found at Credentials.AccessKeyId"
|
||||
)
|
||||
|
||||
# fixme: replace with a Github URL that points
|
||||
# to our repo
|
||||
aws_session_script_url = (
|
||||
"https://gist.github.com/maximsmol/a0284e1d97b25d417bd9ae02e5f450cf"
|
||||
)
|
||||
|
||||
cli_logger.verbose_error(*generic_message_args)
|
||||
cli_logger.verbose(vars(exc))
|
||||
|
||||
cli_logger.panic("Your AWS session has expired.")
|
||||
cli_logger.newline()
|
||||
cli_logger.panic("You can request a new one using")
|
||||
cli_logger.panic(cf.bold(token_command))
|
||||
cli_logger.panic("then expose it to Ray by setting")
|
||||
cli_logger.panic(cf.bold(secret_key_var))
|
||||
cli_logger.panic(cf.bold(session_token_var))
|
||||
cli_logger.panic(cf.bold(access_key_id_var))
|
||||
cli_logger.newline()
|
||||
cli_logger.panic("You can find a script that automates this at:")
|
||||
cli_logger.panic(cf.underlined(aws_session_script_url))
|
||||
# Do not re-raise the exception here because it looks awful
|
||||
# and we already print all the info in verbose
|
||||
cli_logger.abort()
|
||||
|
||||
# todo: any other errors that we should catch separately?
|
||||
|
||||
cli_logger.panic(*generic_message_args)
|
||||
cli_logger.newline()
|
||||
with cli_logger.verbatim_error_ctx("Boto3 error:"):
|
||||
cli_logger.verbose("{}", str(vars(exc)))
|
||||
cli_logger.panic("{}", str(exc))
|
||||
cli_logger.abort()
|
||||
|
||||
|
||||
def boto_exception_handler(msg, *args, **kwargs):
|
||||
# todo: implement timer
|
||||
class ExceptionHandlerContextManager:
|
||||
def __enter__(self):
|
||||
pass
|
||||
|
||||
def __exit__(self, type, value, tb):
|
||||
import botocore
|
||||
|
||||
if type is botocore.exceptions.ClientError:
|
||||
handle_boto_error(value, msg, *args, **kwargs)
|
||||
|
||||
return ExceptionHandlerContextManager()
|
||||
|
||||
|
||||
@lru_cache()
|
||||
def resource_cache(
|
||||
name, region, max_retries=BOTO_MAX_RETRIES, **kwargs
|
||||
) -> ServiceResource:
|
||||
cli_logger.verbose(
|
||||
"Creating AWS resource `{}` in `{}`", cf.bold(name), cf.bold(region)
|
||||
)
|
||||
kwargs.setdefault(
|
||||
"config",
|
||||
Config(retries={"max_attempts": max_retries}),
|
||||
)
|
||||
return boto3.resource(
|
||||
name,
|
||||
region,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
|
||||
@lru_cache()
|
||||
def client_cache(name, region, max_retries=BOTO_MAX_RETRIES, **kwargs) -> BaseClient:
|
||||
try:
|
||||
# try to re-use a client from the resource cache first
|
||||
return resource_cache(name, region, max_retries, **kwargs).meta.client
|
||||
except ResourceNotExistsError:
|
||||
# fall back for clients without an associated resource
|
||||
cli_logger.verbose(
|
||||
"Creating AWS client `{}` in `{}`", cf.bold(name), cf.bold(region)
|
||||
)
|
||||
kwargs.setdefault(
|
||||
"config",
|
||||
Config(retries={"max_attempts": max_retries}),
|
||||
)
|
||||
return boto3.client(
|
||||
name,
|
||||
region,
|
||||
**kwargs,
|
||||
)
|
||||
@@ -0,0 +1,814 @@
|
||||
"""Logger implementing the Command Line Interface.
|
||||
|
||||
A replacement for the standard Python `logging` API
|
||||
designed for implementing a better CLI UX for the cluster launcher.
|
||||
|
||||
Supports color, bold text, italics, underlines, etc.
|
||||
(depending on TTY features)
|
||||
as well as indentation and other structured output.
|
||||
"""
|
||||
import inspect
|
||||
import logging
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
from contextlib import contextmanager
|
||||
from functools import wraps
|
||||
from typing import Any, Callable, Dict, List, Optional, Tuple
|
||||
|
||||
import click
|
||||
import colorama
|
||||
|
||||
# Import ray first to use the bundled colorama
|
||||
import ray # noqa: F401
|
||||
|
||||
if sys.platform == "win32":
|
||||
import msvcrt
|
||||
else:
|
||||
import select
|
||||
|
||||
|
||||
class _ColorfulMock:
|
||||
def __init__(self):
|
||||
# do not do any color work
|
||||
self.identity = lambda x: x
|
||||
|
||||
self.colorful = self
|
||||
self.colormode = None
|
||||
|
||||
self.NO_COLORS = None
|
||||
self.ANSI_8_COLORS = None
|
||||
|
||||
def disable(self):
|
||||
pass
|
||||
|
||||
@contextmanager
|
||||
def with_style(self, x):
|
||||
class IdentityClass:
|
||||
def __getattr__(self, name):
|
||||
return lambda y: y
|
||||
|
||||
yield IdentityClass()
|
||||
|
||||
def __getattr__(self, name):
|
||||
if name == "with_style":
|
||||
return self.with_style
|
||||
|
||||
return self.identity
|
||||
|
||||
|
||||
try:
|
||||
import colorful as _cf
|
||||
from colorful.core import ColorfulString
|
||||
|
||||
_cf.use_8_ansi_colors()
|
||||
except ModuleNotFoundError:
|
||||
# We mock Colorful to restrict the colors used for consistency
|
||||
# anyway, so we also allow for not having colorful at all.
|
||||
# If the Ray Core dependency on colorful is ever removed,
|
||||
# the CliLogger code will still work.
|
||||
class ColorfulString:
|
||||
pass
|
||||
|
||||
_cf = _ColorfulMock()
|
||||
|
||||
|
||||
# We want to only allow specific formatting
|
||||
# to prevent people from accidentally making bad looking color schemes.
|
||||
#
|
||||
# This is especially important since most will look bad on either light
|
||||
# or dark themes.
|
||||
class _ColorfulProxy:
|
||||
_proxy_allowlist = [
|
||||
"disable",
|
||||
"reset",
|
||||
"bold",
|
||||
"italic",
|
||||
"underlined",
|
||||
# used instead of `gray` as `dimmed` adapts to
|
||||
# both light and dark themes
|
||||
"dimmed",
|
||||
"dodgerBlue", # group
|
||||
"limeGreen", # success
|
||||
"red", # error
|
||||
"orange", # warning
|
||||
"skyBlue", # label
|
||||
"magenta", # syntax highlighting key words and symbols
|
||||
"yellow", # syntax highlighting strings
|
||||
]
|
||||
|
||||
def __getattr__(self, name):
|
||||
res = getattr(_cf, name)
|
||||
if callable(res) and name not in _ColorfulProxy._proxy_allowlist:
|
||||
raise ValueError(
|
||||
"Usage of the colorful method '" + name + "' is forbidden "
|
||||
"by the proxy to keep a consistent color scheme. "
|
||||
"Check `cli_logger.py` for allowed methods"
|
||||
)
|
||||
return res
|
||||
|
||||
|
||||
cf = _ColorfulProxy()
|
||||
|
||||
colorama.init(strip=False)
|
||||
|
||||
|
||||
def _external_caller_info():
|
||||
"""Get the info from the caller frame.
|
||||
|
||||
Used to override the logging function and line number with the correct
|
||||
ones. See the comment on _patched_makeRecord for more info.
|
||||
"""
|
||||
|
||||
frame = inspect.currentframe()
|
||||
caller = frame
|
||||
levels = 0
|
||||
while caller.f_code.co_filename == __file__:
|
||||
caller = caller.f_back
|
||||
levels += 1
|
||||
return {
|
||||
"lineno": caller.f_lineno,
|
||||
"filename": os.path.basename(caller.f_code.co_filename),
|
||||
}
|
||||
|
||||
|
||||
def _format_msg(
|
||||
msg: str,
|
||||
*args: Any,
|
||||
no_format: bool = None,
|
||||
_tags: Dict[str, Any] = None,
|
||||
_numbered: Tuple[str, int, int] = None,
|
||||
**kwargs: Any,
|
||||
):
|
||||
"""Formats a message for printing.
|
||||
|
||||
Renders `msg` using the built-in `str.format` and the passed-in
|
||||
`*args` and `**kwargs`.
|
||||
|
||||
Args:
|
||||
msg: The message to format.
|
||||
*args: `.format` arguments for `msg`.
|
||||
no_format: If `no_format` is `True`, `.format` will not be called on
|
||||
the message. Useful if the output is user-provided or may otherwise
|
||||
contain an unexpected formatting string (e.g. "{}").
|
||||
_tags: key-value pairs to display at the end of the message in
|
||||
square brackets. If a tag is set to `True`, it is printed without
|
||||
the value, the presence of the tag treated as a "flag".
|
||||
|
||||
E.g. `_format_msg("hello", _tags=dict(from=mom, signed=True))`
|
||||
`hello [from=Mom, signed]`
|
||||
_numbered: `(brackets, i, n)` The `brackets` string is composed of
|
||||
two "bracket" characters, `i` is the index, `n` is the total.
|
||||
The string `{i}/{n}` surrounded by the "brackets" is
|
||||
prepended to the message. This is used to number steps in a
|
||||
procedure, with different brackets specifying different major tasks.
|
||||
|
||||
E.g. `_format_msg("hello", _numbered=("[]", 0, 5))`
|
||||
`[0/5] hello`
|
||||
**kwargs: `.format` keyword arguments for `msg`.
|
||||
|
||||
Returns:
|
||||
The formatted message.
|
||||
"""
|
||||
|
||||
if isinstance(msg, str) or isinstance(msg, ColorfulString):
|
||||
tags_str = ""
|
||||
if _tags is not None:
|
||||
tags_list = []
|
||||
for k, v in _tags.items():
|
||||
if v is True:
|
||||
tags_list += [k]
|
||||
continue
|
||||
if v is False:
|
||||
continue
|
||||
|
||||
tags_list += [k + "=" + v]
|
||||
if tags_list:
|
||||
tags_str = cf.reset(cf.dimmed(" [{}]".format(", ".join(tags_list))))
|
||||
|
||||
numbering_str = ""
|
||||
if _numbered is not None:
|
||||
chars, i, n = _numbered
|
||||
numbering_str = cf.dimmed(chars[0] + str(i) + "/" + str(n) + chars[1]) + " "
|
||||
|
||||
if no_format:
|
||||
# todo: throw if given args/kwargs?
|
||||
return numbering_str + msg + tags_str
|
||||
return numbering_str + msg.format(*args, **kwargs) + tags_str
|
||||
|
||||
if kwargs:
|
||||
raise ValueError("We do not support printing kwargs yet.")
|
||||
|
||||
res = [msg, *args]
|
||||
res = [str(x) for x in res]
|
||||
return ", ".join(res)
|
||||
|
||||
|
||||
# TODO: come up with a plan to unify logging.
|
||||
# formatter = logging.Formatter(
|
||||
# # TODO(maximsmol): figure out the required log level padding
|
||||
# # width automatically
|
||||
# fmt="[{asctime}] {levelname:6} {message}",
|
||||
# datefmt="%x %X",
|
||||
# # We want alignment on our level names
|
||||
# style="{")
|
||||
|
||||
|
||||
def _isatty():
|
||||
"""More robust check for interactive terminal/tty."""
|
||||
try:
|
||||
# https://stackoverflow.com/questions/6108330/
|
||||
# checking-for-interactive-shell-in-a-python-script
|
||||
return sys.__stdin__.isatty()
|
||||
except Exception:
|
||||
# sometimes this can fail due to closed output
|
||||
# either way, no-tty is generally safe fallback.
|
||||
return False
|
||||
|
||||
|
||||
class _CliLogger:
|
||||
"""Singleton class for CLI logging.
|
||||
|
||||
Without calling 'cli_logger.configure', the CLILogger will default
|
||||
to 'record' style logging.
|
||||
|
||||
Attributes:
|
||||
color_mode: Can be "true", "false", or "auto". Enables or disables
|
||||
`colorful`. If `color_mode` is "auto", is set to `not stdout.isatty()`
|
||||
indent_level: The current indentation level. All messages will
|
||||
be indented by prepending `" " * indent_level`
|
||||
_verbosity: Output verbosity. Low verbosity will disable
|
||||
`verbose` and `very_verbose` messages.
|
||||
"""
|
||||
|
||||
color_mode: str
|
||||
# color_mode: Union[Literal["auto"], Literal["false"], Literal["true"]]
|
||||
indent_level: int
|
||||
interactive: bool
|
||||
VALID_LOG_STYLES = ("auto", "record", "pretty")
|
||||
|
||||
_autodetected_cf_colormode: int
|
||||
|
||||
def __init__(self):
|
||||
self.indent_level = 0
|
||||
|
||||
self._verbosity = 0
|
||||
self._verbosity_overriden = False
|
||||
self._color_mode = "auto"
|
||||
self._log_style = "record"
|
||||
self.pretty = False
|
||||
self.interactive = False
|
||||
|
||||
# store whatever colorful has detected for future use if
|
||||
# the color ouput is toggled (colorful detects # of supported colors,
|
||||
# so it has some non-trivial logic to determine this)
|
||||
self._autodetected_cf_colormode = cf.colorful.colormode
|
||||
self.set_format()
|
||||
|
||||
def set_format(self, format_tmpl=None):
|
||||
if not format_tmpl:
|
||||
from ray.autoscaler._private.constants import LOGGER_FORMAT
|
||||
|
||||
format_tmpl = LOGGER_FORMAT
|
||||
self._formatter = logging.Formatter(format_tmpl)
|
||||
|
||||
def configure(self, log_style=None, color_mode=None, verbosity=None):
|
||||
"""Configures the logger according to values."""
|
||||
if log_style is not None:
|
||||
self._set_log_style(log_style)
|
||||
|
||||
if color_mode is not None:
|
||||
self._set_color_mode(color_mode)
|
||||
|
||||
if verbosity is not None:
|
||||
self._set_verbosity(verbosity)
|
||||
|
||||
self.detect_colors()
|
||||
|
||||
@property
|
||||
def log_style(self):
|
||||
return self._log_style
|
||||
|
||||
def _set_log_style(self, x):
|
||||
"""Configures interactivity and formatting."""
|
||||
self._log_style = x.lower()
|
||||
self.interactive = _isatty()
|
||||
|
||||
if self._log_style == "auto":
|
||||
self.pretty = _isatty()
|
||||
elif self._log_style == "record":
|
||||
self.pretty = False
|
||||
self._set_color_mode("false")
|
||||
elif self._log_style == "pretty":
|
||||
self.pretty = True
|
||||
|
||||
@property
|
||||
def color_mode(self):
|
||||
return self._color_mode
|
||||
|
||||
def _set_color_mode(self, x):
|
||||
self._color_mode = x.lower()
|
||||
self.detect_colors()
|
||||
|
||||
@property
|
||||
def verbosity(self):
|
||||
if self._verbosity_overriden:
|
||||
return self._verbosity
|
||||
elif not self.pretty:
|
||||
return 999
|
||||
return self._verbosity
|
||||
|
||||
def _set_verbosity(self, x):
|
||||
self._verbosity = x
|
||||
self._verbosity_overriden = True
|
||||
|
||||
def detect_colors(self):
|
||||
"""Update color output settings.
|
||||
|
||||
Parse the `color_mode` string and optionally disable or force-enable
|
||||
color output
|
||||
(8-color ANSI if no terminal detected to be safe) in colorful.
|
||||
"""
|
||||
if self.color_mode == "true":
|
||||
if self._autodetected_cf_colormode != cf.NO_COLORS:
|
||||
cf.colormode = self._autodetected_cf_colormode
|
||||
else:
|
||||
cf.colormode = cf.ANSI_8_COLORS
|
||||
return
|
||||
if self.color_mode == "false":
|
||||
cf.disable()
|
||||
return
|
||||
if self.color_mode == "auto":
|
||||
# colorful autodetects tty settings
|
||||
return
|
||||
|
||||
raise ValueError("Invalid log color setting: " + self.color_mode)
|
||||
|
||||
def newline(self):
|
||||
"""Print a line feed."""
|
||||
self.print("")
|
||||
|
||||
def _print(
|
||||
self,
|
||||
msg: str,
|
||||
_level_str: str = "INFO",
|
||||
_linefeed: bool = True,
|
||||
end: str = None,
|
||||
) -> None:
|
||||
"""Proxy for printing messages.
|
||||
|
||||
Args:
|
||||
msg: Message to print.
|
||||
_level_str: Log level label used by the non-pretty formatter.
|
||||
_linefeed: If `_linefeed` is `False` no linefeed is printed at
|
||||
the end of the message.
|
||||
end: String appended after the last value, passed to print().
|
||||
"""
|
||||
if self.pretty:
|
||||
rendered_message = " " * self.indent_level + msg
|
||||
else:
|
||||
if msg.strip() == "":
|
||||
return
|
||||
caller_info = _external_caller_info()
|
||||
record = logging.LogRecord(
|
||||
name="cli",
|
||||
# We override the level name later
|
||||
# TODO(maximsmol): give approximate level #s to our log levels
|
||||
level=0,
|
||||
# The user-facing logs do not need this information anyway
|
||||
# and it would be very tedious to extract since _print
|
||||
# can be at varying depths in the call stack
|
||||
# TODO(maximsmol): do it anyway to be extra
|
||||
pathname=caller_info["filename"],
|
||||
lineno=caller_info["lineno"],
|
||||
msg=msg,
|
||||
args={},
|
||||
# No exception
|
||||
exc_info=None,
|
||||
)
|
||||
record.levelname = _level_str
|
||||
rendered_message = self._formatter.format(record)
|
||||
|
||||
# We aren't using standard python logging convention, so we hardcode
|
||||
# the log levels for now.
|
||||
if _level_str in ["WARNING", "ERROR", "PANIC"]:
|
||||
stream = sys.stderr
|
||||
else:
|
||||
stream = sys.stdout
|
||||
|
||||
if not _linefeed:
|
||||
stream.write(rendered_message)
|
||||
stream.flush()
|
||||
return
|
||||
|
||||
kwargs = {"end": end}
|
||||
print(rendered_message, file=stream, **kwargs)
|
||||
|
||||
def indented(self):
|
||||
"""Context manager that starts an indented block of output."""
|
||||
cli_logger = self
|
||||
|
||||
class IndentedContextManager:
|
||||
def __enter__(self):
|
||||
cli_logger.indent_level += 1
|
||||
|
||||
def __exit__(self, type, value, tb):
|
||||
cli_logger.indent_level -= 1
|
||||
|
||||
return IndentedContextManager()
|
||||
|
||||
def group(self, msg: str, *args: Any, **kwargs: Any):
|
||||
"""Print a group title in a special color and start an indented block.
|
||||
|
||||
For arguments, see `_format_msg`.
|
||||
"""
|
||||
self.print(cf.dodgerBlue(msg), *args, **kwargs)
|
||||
|
||||
return self.indented()
|
||||
|
||||
def verbatim_error_ctx(self, msg: str, *args: Any, **kwargs: Any):
|
||||
"""Context manager for printing multi-line error messages.
|
||||
|
||||
Displays a start sequence "!!! {optional message}"
|
||||
and a matching end sequence "!!!".
|
||||
|
||||
The string "!!!" can be used as a "tombstone" for searching.
|
||||
|
||||
For arguments, see `_format_msg`.
|
||||
"""
|
||||
cli_logger = self
|
||||
|
||||
class VerbatimErorContextManager:
|
||||
def __enter__(self):
|
||||
cli_logger.error(cf.bold("!!! ") + "{}", msg, *args, **kwargs)
|
||||
|
||||
def __exit__(self, type, value, tb):
|
||||
cli_logger.error(cf.bold("!!!"))
|
||||
|
||||
return VerbatimErorContextManager()
|
||||
|
||||
def labeled_value(self, key: str, msg: str, *args: Any, **kwargs: Any):
|
||||
"""Displays a key-value pair with special formatting.
|
||||
|
||||
Args:
|
||||
key: Label that is prepended to the message.
|
||||
msg: Message to print after the label. See `_format_msg`.
|
||||
*args: Additional positional args forwarded to `_format_msg`.
|
||||
**kwargs: Additional keyword args forwarded to `_format_msg`.
|
||||
"""
|
||||
self._print(cf.skyBlue(key) + ": " + _format_msg(cf.bold(msg), *args, **kwargs))
|
||||
|
||||
def verbose(self, msg: str, *args: Any, **kwargs: Any):
|
||||
"""Prints a message if verbosity is not 0.
|
||||
|
||||
For arguments, see `_format_msg`.
|
||||
"""
|
||||
if self.verbosity > 0:
|
||||
self.print(msg, *args, _level_str="VINFO", **kwargs)
|
||||
|
||||
def verbose_warning(self, msg, *args, **kwargs):
|
||||
"""Prints a formatted warning if verbosity is not 0.
|
||||
|
||||
For arguments, see `_format_msg`.
|
||||
"""
|
||||
if self.verbosity > 0:
|
||||
self._warning(msg, *args, _level_str="VWARN", **kwargs)
|
||||
|
||||
def verbose_error(self, msg: str, *args: Any, **kwargs: Any):
|
||||
"""Logs an error if verbosity is not 0.
|
||||
|
||||
For arguments, see `_format_msg`.
|
||||
"""
|
||||
if self.verbosity > 0:
|
||||
self._error(msg, *args, _level_str="VERR", **kwargs)
|
||||
|
||||
def very_verbose(self, msg: str, *args: Any, **kwargs: Any):
|
||||
"""Prints if verbosity is > 1.
|
||||
|
||||
For arguments, see `_format_msg`.
|
||||
"""
|
||||
if self.verbosity > 1:
|
||||
self.print(msg, *args, _level_str="VVINFO", **kwargs)
|
||||
|
||||
def success(self, msg: str, *args: Any, **kwargs: Any):
|
||||
"""Prints a formatted success message.
|
||||
|
||||
For arguments, see `_format_msg`.
|
||||
"""
|
||||
self.print(cf.limeGreen(msg), *args, _level_str="SUCC", **kwargs)
|
||||
|
||||
def _warning(self, msg: str, *args: Any, _level_str: str = None, **kwargs: Any):
|
||||
"""Prints a formatted warning message.
|
||||
|
||||
For arguments, see `_format_msg`.
|
||||
"""
|
||||
if _level_str is None:
|
||||
raise ValueError("Log level not set.")
|
||||
self.print(cf.orange(msg), *args, _level_str=_level_str, **kwargs)
|
||||
|
||||
def warning(self, *args, **kwargs):
|
||||
self._warning(*args, _level_str="WARN", **kwargs)
|
||||
|
||||
def _error(self, msg: str, *args: Any, _level_str: str = None, **kwargs: Any):
|
||||
"""Prints a formatted error message.
|
||||
|
||||
For arguments, see `_format_msg`.
|
||||
"""
|
||||
if _level_str is None:
|
||||
raise ValueError("Log level not set.")
|
||||
self.print(cf.red(msg), *args, _level_str=_level_str, **kwargs)
|
||||
|
||||
def error(self, *args, **kwargs):
|
||||
self._error(*args, _level_str="ERR", **kwargs)
|
||||
|
||||
def panic(self, *args, **kwargs):
|
||||
self._error(*args, _level_str="PANIC", **kwargs)
|
||||
|
||||
# Fine to expose _level_str here, since this is a general log function.
|
||||
def print(
|
||||
self,
|
||||
msg: str,
|
||||
*args: Any,
|
||||
_level_str: str = "INFO",
|
||||
end: str = None,
|
||||
**kwargs: Any,
|
||||
):
|
||||
"""Prints a message.
|
||||
|
||||
For arguments, see `_format_msg`.
|
||||
"""
|
||||
self._print(_format_msg(msg, *args, **kwargs), _level_str=_level_str, end=end)
|
||||
|
||||
def info(self, msg: str, no_format=True, *args, **kwargs):
|
||||
self.print(msg, no_format=no_format, *args, **kwargs)
|
||||
|
||||
def abort(
|
||||
self, msg: Optional[str] = None, *args: Any, exc: Any = None, **kwargs: Any
|
||||
):
|
||||
"""Prints an error and aborts execution.
|
||||
|
||||
Print an error and throw an exception to terminate the program
|
||||
(the exception will not print a message).
|
||||
"""
|
||||
if msg is not None:
|
||||
self._error(msg, *args, _level_str="PANIC", **kwargs)
|
||||
|
||||
if exc is not None:
|
||||
raise exc
|
||||
|
||||
exc_cls = click.ClickException
|
||||
if self.pretty:
|
||||
exc_cls = SilentClickException
|
||||
|
||||
if msg is None:
|
||||
msg = "Exiting due to cli_logger.abort()"
|
||||
raise exc_cls(msg)
|
||||
|
||||
def doassert(self, val: bool, msg: str, *args: Any, **kwargs: Any):
|
||||
"""Handle assertion without throwing a scary exception.
|
||||
|
||||
Args:
|
||||
val: Value to check.
|
||||
msg: Message to print on assertion failure. See `_format_msg`.
|
||||
*args: Additional positional args forwarded to `_format_msg`.
|
||||
**kwargs: Additional keyword args forwarded to `_format_msg`.
|
||||
"""
|
||||
if not val:
|
||||
exc = None
|
||||
if not self.pretty:
|
||||
exc = AssertionError()
|
||||
|
||||
# TODO(maximsmol): rework asserts so that we get the expression
|
||||
# that triggered the assert
|
||||
# to do this, install a global try-catch
|
||||
# for AssertionError and raise them normally
|
||||
self.abort(msg, *args, exc=exc, **kwargs)
|
||||
|
||||
def render_list(self, xs: List[str], separator: str = cf.reset(", ")):
|
||||
"""Render a list of bolded values using a non-bolded separator."""
|
||||
return separator.join([str(cf.bold(x)) for x in xs])
|
||||
|
||||
def confirm(
|
||||
self,
|
||||
yes: bool,
|
||||
msg: str,
|
||||
*args: Any,
|
||||
_abort: bool = False,
|
||||
_default: bool = False,
|
||||
_timeout_s: Optional[float] = None,
|
||||
**kwargs: Any,
|
||||
):
|
||||
"""Display a confirmation dialog.
|
||||
|
||||
Valid answers are "y/yes/true/1" and "n/no/false/0".
|
||||
|
||||
Args:
|
||||
yes: If `yes` is `True` the dialog will default to "yes"
|
||||
and continue without waiting for user input.
|
||||
msg: The prompt message to display. See `_format_msg`.
|
||||
*args: Additional positional args forwarded to `_format_msg`.
|
||||
_abort: If `_abort` is `True`, "no" means aborting the program.
|
||||
_default: The default action to take if the user just presses
|
||||
enter with no input.
|
||||
_timeout_s: If user has no input within _timeout_s seconds,
|
||||
the default action is taken. None means no timeout.
|
||||
**kwargs: Additional keyword args forwarded to `_format_msg`.
|
||||
|
||||
Returns:
|
||||
True if the user confirmed (or `yes` was set); False otherwise.
|
||||
"""
|
||||
should_abort = _abort
|
||||
default = _default
|
||||
|
||||
if not self.interactive and not yes:
|
||||
# no formatting around --yes here since this is non-interactive
|
||||
self.error(
|
||||
"This command requires user confirmation. "
|
||||
"When running non-interactively, supply --yes to skip."
|
||||
)
|
||||
raise ValueError("Non-interactive confirm without --yes.")
|
||||
|
||||
if default:
|
||||
yn_str = "Y/n"
|
||||
else:
|
||||
yn_str = "y/N"
|
||||
|
||||
confirm_str = cf.underlined("Confirm [" + yn_str + "]:") + " "
|
||||
|
||||
rendered_message = _format_msg(msg, *args, **kwargs)
|
||||
# the rendered message ends with ascii coding
|
||||
if rendered_message and not msg.endswith("\n"):
|
||||
rendered_message += " "
|
||||
|
||||
msg_len = len(rendered_message.split("\n")[-1])
|
||||
complete_str = rendered_message + confirm_str
|
||||
|
||||
if yes:
|
||||
self._print(complete_str + "y " + cf.dimmed("[automatic, due to --yes]"))
|
||||
return True
|
||||
|
||||
self._print(complete_str, _linefeed=False)
|
||||
|
||||
res = None
|
||||
yes_answers = ["y", "yes", "true", "1"]
|
||||
no_answers = ["n", "no", "false", "0"]
|
||||
try:
|
||||
while True:
|
||||
if _timeout_s is None:
|
||||
ans = sys.stdin.readline()
|
||||
elif sys.platform == "win32":
|
||||
# Windows doesn't support select
|
||||
start_time = time.time()
|
||||
ans = ""
|
||||
while True:
|
||||
if (time.time() - start_time) >= _timeout_s:
|
||||
self.newline()
|
||||
ans = "\n"
|
||||
break
|
||||
elif msvcrt.kbhit():
|
||||
ch = msvcrt.getwch()
|
||||
if ch in ("\n", "\r"):
|
||||
self.newline()
|
||||
ans = ans + "\n"
|
||||
break
|
||||
elif ch == "\b":
|
||||
if ans:
|
||||
ans = ans[:-1]
|
||||
# Emulate backspace erasing
|
||||
print("\b \b", end="", flush=True)
|
||||
else:
|
||||
ans = ans + ch
|
||||
print(ch, end="", flush=True)
|
||||
else:
|
||||
time.sleep(0.1)
|
||||
else:
|
||||
ready, _, _ = select.select([sys.stdin], [], [], _timeout_s)
|
||||
if not ready:
|
||||
self.newline()
|
||||
ans = "\n"
|
||||
else:
|
||||
ans = sys.stdin.readline()
|
||||
|
||||
ans = ans.lower()
|
||||
|
||||
if ans == "\n":
|
||||
res = default
|
||||
break
|
||||
|
||||
ans = ans.strip()
|
||||
if ans in yes_answers:
|
||||
res = True
|
||||
break
|
||||
if ans in no_answers:
|
||||
res = False
|
||||
break
|
||||
|
||||
indent = " " * msg_len
|
||||
self.error(
|
||||
"{}Invalid answer: {}. Expected {} or {}",
|
||||
indent,
|
||||
cf.bold(ans.strip()),
|
||||
self.render_list(yes_answers, "/"),
|
||||
self.render_list(no_answers, "/"),
|
||||
)
|
||||
self._print(indent + confirm_str, _linefeed=False)
|
||||
except KeyboardInterrupt:
|
||||
self.newline()
|
||||
res = default
|
||||
|
||||
if not res and should_abort:
|
||||
# todo: make sure we tell the user if they
|
||||
# need to do cleanup
|
||||
self._print("Exiting...")
|
||||
raise SilentClickException(
|
||||
"Exiting due to the response to confirm(should_abort=True)."
|
||||
)
|
||||
|
||||
return res
|
||||
|
||||
def prompt(self, msg: str, *args, **kwargs):
|
||||
"""Prompt the user for some text input.
|
||||
|
||||
Args:
|
||||
msg: The message to display to the user before the prompt.
|
||||
*args: Additional positional args forwarded to `_format_msg`.
|
||||
**kwargs: Additional keyword args forwarded to `_format_msg`.
|
||||
|
||||
Returns:
|
||||
The string entered by the user.
|
||||
"""
|
||||
complete_str = cf.underlined(msg)
|
||||
rendered_message = _format_msg(complete_str, *args, **kwargs)
|
||||
# the rendered message ends with ascii coding
|
||||
if rendered_message and not msg.endswith("\n"):
|
||||
rendered_message += " "
|
||||
self._print(rendered_message, linefeed=False)
|
||||
|
||||
res = ""
|
||||
try:
|
||||
ans = sys.stdin.readline()
|
||||
ans = ans.lower()
|
||||
res = ans.strip()
|
||||
except KeyboardInterrupt:
|
||||
self.newline()
|
||||
|
||||
return res
|
||||
|
||||
def flush(self):
|
||||
sys.stdout.flush()
|
||||
sys.stderr.flush()
|
||||
|
||||
|
||||
class SilentClickException(click.ClickException):
|
||||
"""`ClickException` that does not print a message.
|
||||
|
||||
Some of our tooling relies on catching ClickException in particular.
|
||||
|
||||
However the default prints a message, which is undesirable since we expect
|
||||
our code to log errors manually using `cli_logger.error()` to allow for
|
||||
colors and other formatting.
|
||||
"""
|
||||
|
||||
def __init__(self, message: str):
|
||||
super(SilentClickException, self).__init__(message)
|
||||
|
||||
def show(self, file=None):
|
||||
pass
|
||||
|
||||
|
||||
cli_logger = _CliLogger()
|
||||
|
||||
CLICK_LOGGING_OPTIONS = [
|
||||
click.option(
|
||||
"--log-style",
|
||||
required=False,
|
||||
type=click.Choice(cli_logger.VALID_LOG_STYLES, case_sensitive=False),
|
||||
default="auto",
|
||||
help=(
|
||||
"If 'pretty', outputs with formatting and color. If 'record', "
|
||||
"outputs record-style without formatting. "
|
||||
"'auto' defaults to 'pretty', and disables pretty logging "
|
||||
"if stdin is *not* a TTY."
|
||||
),
|
||||
),
|
||||
click.option(
|
||||
"--log-color",
|
||||
required=False,
|
||||
type=click.Choice(["auto", "false", "true"], case_sensitive=False),
|
||||
default="auto",
|
||||
help=("Use color logging. Auto enables color logging if stdout is a TTY."),
|
||||
),
|
||||
click.option("-v", "--verbose", default=None, count=True),
|
||||
]
|
||||
|
||||
|
||||
def add_click_logging_options(f: Callable) -> Callable:
|
||||
for option in reversed(CLICK_LOGGING_OPTIONS):
|
||||
f = option(f)
|
||||
|
||||
@wraps(f)
|
||||
def wrapper(*args, log_style=None, log_color=None, verbose=None, **kwargs):
|
||||
cli_logger.configure(log_style, log_color, verbose)
|
||||
return f(*args, **kwargs)
|
||||
|
||||
return wrapper
|
||||
+40
@@ -0,0 +1,40 @@
|
||||
#!/usr/bin/env python
|
||||
|
||||
# This is an executable script that runs an example of every single CliLogger
|
||||
# function for demonstration purposes. Primarily useful for tuning color and
|
||||
# other formatting.
|
||||
|
||||
from ray.autoscaler._private.cli_logger import cf, cli_logger
|
||||
|
||||
cli_logger.configure(log_style="auto", verbosity=999)
|
||||
|
||||
cli_logger.print(cf.bold("Bold ") + cf.italic("Italic ") + cf.underlined("Underlined"))
|
||||
cli_logger.labeled_value("Label", "value")
|
||||
cli_logger.print("List: {}", cli_logger.render_list([1, 2, 3]))
|
||||
cli_logger.newline()
|
||||
cli_logger.very_verbose("Very verbose")
|
||||
cli_logger.verbose("Verbose")
|
||||
cli_logger.verbose_warning("Verbose warning")
|
||||
cli_logger.verbose_error("Verbose error")
|
||||
cli_logger.print("Info")
|
||||
cli_logger.success("Success")
|
||||
cli_logger.warning("Warning")
|
||||
cli_logger.error("Error")
|
||||
cli_logger.newline()
|
||||
try:
|
||||
cli_logger.abort("Abort")
|
||||
except Exception:
|
||||
pass
|
||||
try:
|
||||
cli_logger.doassert(False, "Assert")
|
||||
except Exception:
|
||||
pass
|
||||
cli_logger.newline()
|
||||
cli_logger.confirm(True, "example")
|
||||
cli_logger.newline()
|
||||
with cli_logger.indented():
|
||||
cli_logger.print("Indented")
|
||||
with cli_logger.group("Group"):
|
||||
cli_logger.print("Group contents")
|
||||
with cli_logger.verbatim_error_ctx("Verbtaim error"):
|
||||
cli_logger.print("Error contents")
|
||||
@@ -0,0 +1,45 @@
|
||||
"""Lightweight CLI output helpers for cluster launcher context notes.
|
||||
|
||||
These functions produce informational output that disambiguates head-node
|
||||
commands from local-machine commands. They have NO dependency on the ``ray``
|
||||
C++ runtime and accept ``cli_logger`` / ``cf`` as explicit parameters so they
|
||||
can be tested in isolation.
|
||||
"""
|
||||
|
||||
|
||||
def print_next_steps_context_note(cli_logger, cf):
|
||||
"""Print a dimmed note at the top of the 'Next steps' block.
|
||||
|
||||
Informs the user that the commands below are intended for the head node
|
||||
or for machines within the cluster network.
|
||||
"""
|
||||
cli_logger.print(cf.dimmed("Note: The following commands are intended for use on"))
|
||||
cli_logger.print(cf.dimmed("the head node or within the cluster network."))
|
||||
cli_logger.newline()
|
||||
|
||||
|
||||
def print_head_node_context_separator(cli_logger, cf):
|
||||
"""Print a visual separator and context note after head-node output.
|
||||
|
||||
Used by ``ray up`` to separate the streamed ``ray start`` output from
|
||||
the local-machine commands that follow.
|
||||
"""
|
||||
cli_logger.print(cf.dimmed("-" * 60))
|
||||
cli_logger.print(
|
||||
cf.dimmed("Note: The output above is from the head node (via `ray start`).")
|
||||
)
|
||||
cli_logger.print(
|
||||
cf.dimmed(" Commands shown in 'Next steps' may only work from the head node")
|
||||
)
|
||||
cli_logger.print(cf.dimmed(" or from within the cluster network."))
|
||||
cli_logger.newline()
|
||||
|
||||
|
||||
def print_head_node_context_separator_if_needed(ray_start_commands, cli_logger, cf):
|
||||
"""Print the head-node separator when ``ray start`` output was shown."""
|
||||
if ray_start_commands:
|
||||
print_head_node_context_separator(cli_logger, cf)
|
||||
|
||||
|
||||
# Group heading used by ``ray up`` for the local commands section.
|
||||
USEFUL_COMMANDS_HEADING = "Useful commands for your local machine:"
|
||||
@@ -0,0 +1,652 @@
|
||||
import os
|
||||
import re
|
||||
import subprocess
|
||||
import sys
|
||||
import tarfile
|
||||
import tempfile
|
||||
import threading
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from contextlib import contextmanager
|
||||
from typing import Any, Iterator, List, Optional, Sequence, Tuple
|
||||
|
||||
import yaml
|
||||
|
||||
import ray # noqa: F401
|
||||
from ray.autoscaler._private.cli_logger import cli_logger
|
||||
from ray.autoscaler._private.providers import _get_node_provider
|
||||
from ray.autoscaler.tags import NODE_KIND_HEAD, NODE_KIND_WORKER, TAG_RAY_NODE_KIND
|
||||
|
||||
# Import psutil after ray so the packaged version is used.
|
||||
import psutil
|
||||
|
||||
MAX_PARALLEL_SSH_WORKERS = 8
|
||||
DEFAULT_SSH_USER = "ubuntu"
|
||||
DEFAULT_SSH_KEYS = ["~/ray_bootstrap_key.pem", "~/.ssh/ray-autoscaler_2_us-west-2.pem"]
|
||||
|
||||
|
||||
class CommandFailed(RuntimeError):
|
||||
pass
|
||||
|
||||
|
||||
class LocalCommandFailed(CommandFailed):
|
||||
pass
|
||||
|
||||
|
||||
class RemoteCommandFailed(CommandFailed):
|
||||
pass
|
||||
|
||||
|
||||
class GetParameters:
|
||||
def __init__(
|
||||
self,
|
||||
logs: bool = True,
|
||||
debug_state: bool = True,
|
||||
pip: bool = True,
|
||||
processes: bool = True,
|
||||
processes_verbose: bool = True,
|
||||
processes_list: Optional[List[Tuple[str, bool]]] = None,
|
||||
):
|
||||
self.logs = logs
|
||||
self.debug_state = debug_state
|
||||
self.pip = pip
|
||||
self.processes = processes
|
||||
self.processes_verbose = processes_verbose
|
||||
self.processes_list = processes_list
|
||||
|
||||
|
||||
class Node:
|
||||
"""Node (as in "machine")"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
host: str,
|
||||
ssh_user: str = "ubuntu",
|
||||
ssh_key: str = "~/ray_bootstrap_key.pem",
|
||||
docker_container: Optional[str] = None,
|
||||
is_head: bool = False,
|
||||
):
|
||||
self.host = host
|
||||
self.ssh_user = ssh_user
|
||||
self.ssh_key = ssh_key
|
||||
self.docker_container = docker_container
|
||||
self.is_head = is_head
|
||||
|
||||
|
||||
class Archive:
|
||||
"""Archive object to collect and compress files into a single file.
|
||||
|
||||
Objects of this class can be passed around to different data collection
|
||||
functions. These functions can use the :meth:`subdir` method to add
|
||||
files to a sub directory of the archive.
|
||||
|
||||
"""
|
||||
|
||||
def __init__(self, file: Optional[str] = None):
|
||||
self.file = file or tempfile.mkstemp(prefix="ray_logs_", suffix=".tar.gz")[1]
|
||||
self.tar = None
|
||||
self._lock = threading.Lock()
|
||||
|
||||
@property
|
||||
def is_open(self):
|
||||
return bool(self.tar)
|
||||
|
||||
def open(self):
|
||||
self.tar = tarfile.open(self.file, "w:gz")
|
||||
|
||||
def close(self):
|
||||
self.tar.close()
|
||||
self.tar = None
|
||||
|
||||
def __enter__(self):
|
||||
self.open()
|
||||
return self
|
||||
|
||||
def __exit__(self, exc_type, exc_val, exc_tb):
|
||||
self.close()
|
||||
|
||||
@contextmanager
|
||||
def subdir(self, subdir: str, root: Optional[str] = "/") -> Iterator[Any]:
|
||||
"""Open a context to add files to the archive.
|
||||
|
||||
Example:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
with Archive("file.tar.gz") as archive:
|
||||
with archive.subdir("logfiles", root="/tmp/logs") as sd:
|
||||
# Will be added as `logfiles/nested/file.txt`
|
||||
sd.add("/tmp/logs/nested/file.txt")
|
||||
|
||||
Args:
|
||||
subdir: Subdir to which to add files to. Calling the
|
||||
``add(path)`` command will place files into the ``subdir``
|
||||
directory of the archive.
|
||||
root: Root path. Files without an explicit ``arcname``
|
||||
will be named relatively to this path.
|
||||
|
||||
Yields:
|
||||
Any: A context object that can be used to add files to the archive.
|
||||
"""
|
||||
root = os.path.abspath(root)
|
||||
|
||||
class _Context:
|
||||
@staticmethod
|
||||
def add(path: str, arcname: Optional[str] = None):
|
||||
path = os.path.abspath(path)
|
||||
arcname = arcname or os.path.join(subdir, os.path.relpath(path, root))
|
||||
|
||||
self._lock.acquire()
|
||||
self.tar.add(path, arcname=arcname)
|
||||
self._lock.release()
|
||||
|
||||
yield _Context()
|
||||
|
||||
|
||||
###
|
||||
# Functions to gather logs and information on the local node
|
||||
###
|
||||
|
||||
|
||||
def get_local_ray_logs(
|
||||
archive: Archive,
|
||||
exclude: Optional[Sequence[str]] = None,
|
||||
session_log_dir: str = "/tmp/ray/session_latest",
|
||||
) -> Archive:
|
||||
"""Copy local log files into an archive.
|
||||
|
||||
Args:
|
||||
archive: Archive object to add log files to.
|
||||
exclude: Sequence of regex patterns. Files that match
|
||||
any of these patterns will not be included in the archive.
|
||||
session_log_dir: Path to the Ray session files. Defaults to
|
||||
``/tmp/ray/session_latest``
|
||||
|
||||
Returns:
|
||||
Open archive object.
|
||||
|
||||
"""
|
||||
if not archive.is_open:
|
||||
archive.open()
|
||||
|
||||
exclude = exclude or []
|
||||
|
||||
session_log_dir = os.path.join(os.path.expanduser(session_log_dir), "logs")
|
||||
|
||||
with archive.subdir("logs", root=session_log_dir) as sd:
|
||||
for root, dirs, files in os.walk(session_log_dir):
|
||||
for file in files:
|
||||
file_path = os.path.join(root, file)
|
||||
rel_path = os.path.relpath(file_path, start=session_log_dir)
|
||||
# Skip file if it matches any pattern in `exclude`
|
||||
if any(re.match(pattern, rel_path) for pattern in exclude):
|
||||
continue
|
||||
sd.add(file_path)
|
||||
|
||||
return archive
|
||||
|
||||
|
||||
def get_local_debug_state(
|
||||
archive: Archive, session_dir: str = "/tmp/ray/session_latest"
|
||||
) -> Archive:
|
||||
"""Copy local log files into an archive.
|
||||
|
||||
Args:
|
||||
archive: Archive object to add log files to.
|
||||
session_dir: Path to the Ray session files. Defaults to
|
||||
``/tmp/ray/session_latest``
|
||||
|
||||
Returns:
|
||||
Open archive object.
|
||||
|
||||
"""
|
||||
if not archive.is_open:
|
||||
archive.open()
|
||||
|
||||
session_dir = os.path.expanduser(session_dir)
|
||||
debug_state_file = os.path.join(session_dir, "logs/debug_state.txt")
|
||||
|
||||
if not os.path.exists(debug_state_file):
|
||||
raise LocalCommandFailed("No `debug_state.txt` file found.")
|
||||
|
||||
with archive.subdir("", root=session_dir) as sd:
|
||||
sd.add(debug_state_file)
|
||||
|
||||
return archive
|
||||
|
||||
|
||||
def get_local_pip_packages(archive: Archive):
|
||||
"""Get currently installed pip packages and write into an archive.
|
||||
|
||||
Args:
|
||||
archive: Archive object to add meta files to.
|
||||
|
||||
Returns:
|
||||
Open archive object.
|
||||
"""
|
||||
if not archive.is_open:
|
||||
archive.open()
|
||||
|
||||
try:
|
||||
from pip._internal.operations import freeze
|
||||
except ImportError: # pip < 10.0
|
||||
from pip.operations import freeze
|
||||
|
||||
with tempfile.NamedTemporaryFile("wt") as fp:
|
||||
for line in freeze.freeze():
|
||||
fp.writelines([line, "\n"])
|
||||
|
||||
fp.flush()
|
||||
with archive.subdir("") as sd:
|
||||
sd.add(fp.name, "pip_packages.txt")
|
||||
|
||||
return archive
|
||||
|
||||
|
||||
def get_local_ray_processes(
|
||||
archive: Archive,
|
||||
processes: Optional[List[Tuple[str, bool]]] = None,
|
||||
verbose: bool = False,
|
||||
):
|
||||
"""Get the status of all the relevant ray processes.
|
||||
Args:
|
||||
archive: Archive object to add process info files to.
|
||||
processes: List of processes to get information on. The first
|
||||
element of the tuple is a string to filter by, and the second
|
||||
element is a boolean indicating if we should filter by command
|
||||
name (True) or command line including parameters (False)
|
||||
verbose: If True, show entire executable command line.
|
||||
If False, show just the first term.
|
||||
Returns:
|
||||
Open archive object.
|
||||
"""
|
||||
if not processes:
|
||||
# local import to avoid circular dependencies
|
||||
from ray.autoscaler._private.constants import RAY_PROCESSES
|
||||
|
||||
processes = RAY_PROCESSES
|
||||
|
||||
process_infos = []
|
||||
for process in psutil.process_iter(["pid", "name", "cmdline", "status"]):
|
||||
try:
|
||||
with process.oneshot():
|
||||
cmdline = " ".join(process.cmdline())
|
||||
process_infos.append(
|
||||
(
|
||||
{
|
||||
"executable": cmdline
|
||||
if verbose
|
||||
else cmdline.split("--", 1)[0][:-1],
|
||||
"name": process.name(),
|
||||
"pid": process.pid,
|
||||
"status": process.status(),
|
||||
},
|
||||
process.cmdline(),
|
||||
)
|
||||
)
|
||||
except Exception as exc:
|
||||
raise LocalCommandFailed(exc) from exc
|
||||
|
||||
relevant_processes = {}
|
||||
for process_dict, cmdline in process_infos:
|
||||
for keyword, filter_by_cmd in processes:
|
||||
if filter_by_cmd:
|
||||
corpus = process_dict["name"]
|
||||
else:
|
||||
corpus = subprocess.list2cmdline(cmdline)
|
||||
if keyword in corpus and process_dict["pid"] not in relevant_processes:
|
||||
relevant_processes[process_dict["pid"]] = process_dict
|
||||
|
||||
with tempfile.NamedTemporaryFile("wt") as fp:
|
||||
for line in relevant_processes.values():
|
||||
fp.writelines([yaml.dump(line), "\n"])
|
||||
|
||||
fp.flush()
|
||||
with archive.subdir("meta") as sd:
|
||||
sd.add(fp.name, "process_info.txt")
|
||||
|
||||
return archive
|
||||
|
||||
|
||||
def get_all_local_data(archive: Archive, parameters: GetParameters):
|
||||
"""Get all local data.
|
||||
|
||||
Gets:
|
||||
- The Ray logs of the latest session
|
||||
- The currently installed pip packages
|
||||
|
||||
Args:
|
||||
archive: Archive object to add meta files to.
|
||||
parameters: Parameters (settings) for getting data.
|
||||
|
||||
Returns:
|
||||
Open archive object.
|
||||
"""
|
||||
if not archive.is_open:
|
||||
archive.open()
|
||||
|
||||
if parameters.logs:
|
||||
try:
|
||||
get_local_ray_logs(archive=archive)
|
||||
except LocalCommandFailed as exc:
|
||||
cli_logger.error(exc)
|
||||
if parameters.debug_state:
|
||||
try:
|
||||
get_local_debug_state(archive=archive)
|
||||
except LocalCommandFailed as exc:
|
||||
cli_logger.error(exc)
|
||||
if parameters.pip:
|
||||
try:
|
||||
get_local_pip_packages(archive=archive)
|
||||
except LocalCommandFailed as exc:
|
||||
cli_logger.error(exc)
|
||||
if parameters.processes:
|
||||
try:
|
||||
get_local_ray_processes(
|
||||
archive=archive,
|
||||
processes=parameters.processes_list,
|
||||
verbose=parameters.processes_verbose,
|
||||
)
|
||||
except LocalCommandFailed as exc:
|
||||
cli_logger.error(exc)
|
||||
|
||||
return archive
|
||||
|
||||
|
||||
###
|
||||
# Functions to invoke remote scripts and gather data from remote nodes
|
||||
###
|
||||
|
||||
|
||||
def _wrap(items: List[str], quotes="'"):
|
||||
return f"{quotes}{' '.join(items)}{quotes}"
|
||||
|
||||
|
||||
def create_and_get_archive_from_remote_node(
|
||||
remote_node: Node, parameters: GetParameters, script_path: str = "ray"
|
||||
) -> Optional[str]:
|
||||
"""Create an archive containing logs on a remote node and transfer.
|
||||
|
||||
This will call ``ray local-dump --stream`` on the remote
|
||||
node. The resulting file will be saved locally in a temporary file and
|
||||
returned.
|
||||
|
||||
Args:
|
||||
remote_node: Remote node to gather archive from.
|
||||
parameters: Parameters (settings) for getting data.
|
||||
script_path: Path to this script on the remote node.
|
||||
|
||||
Returns:
|
||||
Path to a temporary file containing the node's collected data.
|
||||
|
||||
"""
|
||||
cmd = [
|
||||
"ssh",
|
||||
"-o StrictHostKeyChecking=no",
|
||||
"-o UserKnownHostsFile=/dev/null",
|
||||
"-o LogLevel=ERROR",
|
||||
"-i",
|
||||
remote_node.ssh_key,
|
||||
f"{remote_node.ssh_user}@{remote_node.host}",
|
||||
]
|
||||
|
||||
if remote_node.docker_container:
|
||||
cmd += [
|
||||
"docker",
|
||||
"exec",
|
||||
remote_node.docker_container,
|
||||
]
|
||||
|
||||
collect_cmd = [script_path, "local-dump", "--stream"]
|
||||
collect_cmd += ["--logs"] if parameters.logs else ["--no-logs"]
|
||||
collect_cmd += ["--debug-state"] if parameters.debug_state else ["--no-debug-state"]
|
||||
collect_cmd += ["--pip"] if parameters.pip else ["--no-pip"]
|
||||
collect_cmd += ["--processes"] if parameters.processes else ["--no-processes"]
|
||||
if parameters.processes:
|
||||
collect_cmd += (
|
||||
["--processes-verbose"]
|
||||
if parameters.processes_verbose
|
||||
else ["--no-proccesses-verbose"]
|
||||
)
|
||||
|
||||
cmd += ["/bin/bash", "-c", _wrap(collect_cmd, quotes='"')]
|
||||
|
||||
cat = "node" if not remote_node.is_head else "head"
|
||||
|
||||
cli_logger.print(f"Collecting data from remote node: {remote_node.host}")
|
||||
tmp = tempfile.mkstemp(prefix=f"ray_{cat}_{remote_node.host}_", suffix=".tar.gz")[1]
|
||||
with open(tmp, "wb") as fp:
|
||||
try:
|
||||
subprocess.check_call(cmd, stdout=fp, stderr=sys.stderr)
|
||||
except subprocess.CalledProcessError as exc:
|
||||
raise RemoteCommandFailed(
|
||||
f"Gathering logs from remote node failed: {' '.join(cmd)}"
|
||||
) from exc
|
||||
|
||||
return tmp
|
||||
|
||||
|
||||
def create_and_add_remote_data_to_local_archive(
|
||||
archive: Archive, remote_node: Node, parameters: GetParameters
|
||||
):
|
||||
"""Create and get data from remote node and add to local archive.
|
||||
|
||||
Args:
|
||||
archive: Archive object to add remote data to.
|
||||
remote_node: Remote node to gather archive from.
|
||||
parameters: Parameters (settings) for getting data.
|
||||
|
||||
Returns:
|
||||
Open archive object.
|
||||
"""
|
||||
tmp = create_and_get_archive_from_remote_node(remote_node, parameters)
|
||||
|
||||
if not archive.is_open:
|
||||
archive.open()
|
||||
|
||||
cat = "node" if not remote_node.is_head else "head"
|
||||
|
||||
with archive.subdir("", root=os.path.dirname(tmp)) as sd:
|
||||
sd.add(tmp, arcname=f"ray_{cat}_{remote_node.host}.tar.gz")
|
||||
|
||||
return archive
|
||||
|
||||
|
||||
def create_and_add_local_data_to_local_archive(
|
||||
archive: Archive, parameters: GetParameters
|
||||
):
|
||||
"""Create and get data from this node and add to archive.
|
||||
|
||||
Args:
|
||||
archive: Archive object to add remote data to.
|
||||
parameters: Parameters (settings) for getting data.
|
||||
|
||||
Returns:
|
||||
Open archive object.
|
||||
"""
|
||||
with Archive() as local_data_archive:
|
||||
get_all_local_data(local_data_archive, parameters)
|
||||
|
||||
if not archive.is_open:
|
||||
archive.open()
|
||||
|
||||
with archive.subdir("", root=os.path.dirname(local_data_archive.file)) as sd:
|
||||
sd.add(local_data_archive.file, arcname="local_node.tar.gz")
|
||||
|
||||
os.remove(local_data_archive.file)
|
||||
|
||||
return archive
|
||||
|
||||
|
||||
def create_archive_for_remote_nodes(
|
||||
archive: Archive, remote_nodes: Sequence[Node], parameters: GetParameters
|
||||
):
|
||||
"""Create an archive combining data from the remote nodes.
|
||||
|
||||
This will parallelize calls to get data from remote nodes.
|
||||
|
||||
Args:
|
||||
archive: Archive object to add remote data to.
|
||||
remote_nodes: Sequence of remote nodes.
|
||||
parameters: Parameters (settings) for getting data.
|
||||
|
||||
Returns:
|
||||
Open archive object.
|
||||
|
||||
"""
|
||||
if not archive.is_open:
|
||||
archive.open()
|
||||
|
||||
with ThreadPoolExecutor(max_workers=MAX_PARALLEL_SSH_WORKERS) as executor:
|
||||
for remote_node in remote_nodes:
|
||||
executor.submit(
|
||||
create_and_add_remote_data_to_local_archive,
|
||||
archive=archive,
|
||||
remote_node=remote_node,
|
||||
parameters=parameters,
|
||||
)
|
||||
|
||||
return archive
|
||||
|
||||
|
||||
def create_archive_for_local_and_remote_nodes(
|
||||
archive: Archive, remote_nodes: Sequence[Node], parameters: GetParameters
|
||||
):
|
||||
"""Create an archive combining data from the local and remote nodes.
|
||||
|
||||
This will parallelize calls to get data from remote nodes.
|
||||
|
||||
Args:
|
||||
archive: Archive object to add data to.
|
||||
remote_nodes: Sequence of remote nodes.
|
||||
parameters: Parameters (settings) for getting data.
|
||||
|
||||
Returns:
|
||||
Open archive object.
|
||||
|
||||
"""
|
||||
if not archive.is_open:
|
||||
archive.open()
|
||||
|
||||
try:
|
||||
create_and_add_local_data_to_local_archive(archive, parameters)
|
||||
except CommandFailed as exc:
|
||||
cli_logger.error(exc)
|
||||
|
||||
create_archive_for_remote_nodes(archive, remote_nodes, parameters)
|
||||
|
||||
cli_logger.print(
|
||||
f"Collected data from local node and {len(remote_nodes)} " f"remote nodes."
|
||||
)
|
||||
return archive
|
||||
|
||||
|
||||
###
|
||||
# Ray cluster info
|
||||
###
|
||||
def get_info_from_ray_cluster_config(
|
||||
cluster_config: str,
|
||||
) -> Tuple[List[str], str, str, Optional[str], Optional[str]]:
|
||||
"""Get information from Ray cluster config.
|
||||
|
||||
Return list of host IPs, ssh user, ssh key file, and optional docker
|
||||
container.
|
||||
|
||||
Args:
|
||||
cluster_config: Path to ray cluster config.
|
||||
|
||||
Returns:
|
||||
Tuple of list of host IPs, ssh user name, ssh key file path,
|
||||
optional docker container name, optional cluster name.
|
||||
"""
|
||||
from ray.autoscaler._private.commands import _bootstrap_config
|
||||
|
||||
cli_logger.print(
|
||||
f"Retrieving cluster information from ray cluster file: " f"{cluster_config}"
|
||||
)
|
||||
|
||||
cluster_config = os.path.expanduser(cluster_config)
|
||||
|
||||
config = yaml.safe_load(open(cluster_config).read())
|
||||
config = _bootstrap_config(config, no_config_cache=True)
|
||||
|
||||
provider = _get_node_provider(config["provider"], config["cluster_name"])
|
||||
head_nodes = provider.non_terminated_nodes({TAG_RAY_NODE_KIND: NODE_KIND_HEAD})
|
||||
worker_nodes = provider.non_terminated_nodes({TAG_RAY_NODE_KIND: NODE_KIND_WORKER})
|
||||
|
||||
hosts = [provider.external_ip(node) for node in head_nodes + worker_nodes]
|
||||
ssh_user = config["auth"]["ssh_user"]
|
||||
ssh_key = config["auth"]["ssh_private_key"]
|
||||
|
||||
docker = None
|
||||
docker_config = config.get("docker", None)
|
||||
if docker_config:
|
||||
docker = docker_config.get("container_name", None)
|
||||
|
||||
cluster_name = config.get("cluster_name", None)
|
||||
|
||||
return hosts, ssh_user, ssh_key, docker, cluster_name
|
||||
|
||||
|
||||
def _info_from_params(
|
||||
cluster: Optional[str] = None,
|
||||
host: Optional[str] = None,
|
||||
ssh_user: Optional[str] = None,
|
||||
ssh_key: Optional[str] = None,
|
||||
docker: Optional[str] = None,
|
||||
):
|
||||
"""Parse command line arguments.
|
||||
|
||||
Note: This returns a list of hosts, not a comma separated string!
|
||||
"""
|
||||
if not host and not cluster:
|
||||
bootstrap_config = os.path.expanduser("~/ray_bootstrap_config.yaml")
|
||||
if os.path.exists(bootstrap_config):
|
||||
cluster = bootstrap_config
|
||||
cli_logger.warning(
|
||||
f"Detected cluster config file at {cluster}. "
|
||||
f"If this is incorrect, specify with "
|
||||
f"`ray cluster-dump <config>`"
|
||||
)
|
||||
elif cluster:
|
||||
cluster = os.path.expanduser(cluster)
|
||||
|
||||
cluster_name = None
|
||||
|
||||
if cluster:
|
||||
h, u, k, d, cluster_name = get_info_from_ray_cluster_config(cluster)
|
||||
|
||||
ssh_user = ssh_user or u
|
||||
ssh_key = ssh_key or k
|
||||
docker = docker or d
|
||||
hosts = host.split(",") if host else h
|
||||
|
||||
if not hosts:
|
||||
raise LocalCommandFailed(
|
||||
f"Invalid cluster file or cluster has no running nodes: " f"{cluster}"
|
||||
)
|
||||
elif host:
|
||||
hosts = host.split(",")
|
||||
else:
|
||||
raise LocalCommandFailed(
|
||||
"You need to either specify a `<cluster_config>` or `--host`."
|
||||
)
|
||||
|
||||
if not ssh_user:
|
||||
ssh_user = DEFAULT_SSH_USER
|
||||
cli_logger.warning(
|
||||
f"Using default SSH user `{ssh_user}`. "
|
||||
f"If this is incorrect, specify with `--ssh-user <user>`"
|
||||
)
|
||||
|
||||
if not ssh_key:
|
||||
for cand_key in DEFAULT_SSH_KEYS:
|
||||
cand_key_file = os.path.expanduser(cand_key)
|
||||
if os.path.exists(cand_key_file):
|
||||
ssh_key = cand_key_file
|
||||
cli_logger.warning(
|
||||
f"Auto detected SSH key file: {ssh_key}. "
|
||||
f"If this is incorrect, specify with `--ssh-key <key>`"
|
||||
)
|
||||
break
|
||||
|
||||
return cluster, hosts, ssh_user, ssh_key, docker, cluster_name
|
||||
@@ -0,0 +1,966 @@
|
||||
import hashlib
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import subprocess
|
||||
import sys
|
||||
import time
|
||||
from getpass import getuser
|
||||
from shlex import quote
|
||||
from typing import Dict, List, Optional, Tuple
|
||||
|
||||
import click
|
||||
|
||||
from ray._private.ray_constants import DEFAULT_OBJECT_STORE_MAX_MEMORY_BYTES
|
||||
from ray.autoscaler._private.cli_logger import cf, cli_logger
|
||||
from ray.autoscaler._private.constants import (
|
||||
AUTOSCALER_NODE_SSH_INTERVAL_S,
|
||||
AUTOSCALER_NODE_START_WAIT_S,
|
||||
DEFAULT_OBJECT_STORE_MEMORY_PROPORTION,
|
||||
)
|
||||
from ray.autoscaler._private.docker import (
|
||||
check_bind_mounts_cmd,
|
||||
check_docker_image,
|
||||
check_docker_running_cmd,
|
||||
docker_start_cmds,
|
||||
with_docker_exec,
|
||||
)
|
||||
from ray.autoscaler._private.log_timer import LogTimer
|
||||
from ray.autoscaler._private.subprocess_output_util import (
|
||||
ProcessRunnerError,
|
||||
is_output_redirected,
|
||||
run_cmd_redirected,
|
||||
)
|
||||
from ray.autoscaler.command_runner import CommandRunnerInterface
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# How long to wait for a node to start, in seconds
|
||||
HASH_MAX_LENGTH = 10
|
||||
KUBECTL_RSYNC = os.path.join(
|
||||
os.path.dirname(os.path.abspath(__file__)), "_kubernetes/kubectl-rsync.sh"
|
||||
)
|
||||
MAX_HOME_RETRIES = 3
|
||||
HOME_RETRY_DELAY_S = 5
|
||||
|
||||
_config = {"use_login_shells": True, "silent_rsync": True}
|
||||
|
||||
|
||||
def is_rsync_silent():
|
||||
return _config["silent_rsync"]
|
||||
|
||||
|
||||
def set_rsync_silent(val):
|
||||
"""Choose whether to silence rsync output.
|
||||
|
||||
Most commands will want to list rsync'd files themselves rather than
|
||||
print the default rsync spew.
|
||||
"""
|
||||
_config["silent_rsync"] = val
|
||||
|
||||
|
||||
def is_using_login_shells():
|
||||
return _config["use_login_shells"]
|
||||
|
||||
|
||||
def set_using_login_shells(val: bool):
|
||||
"""Choose between login and non-interactive shells.
|
||||
|
||||
Non-interactive shells have the benefit of receiving less output from
|
||||
subcommands (since progress bars and TTY control codes are not printed).
|
||||
Sometimes this can be significant since e.g. `pip install` prints
|
||||
hundreds of progress bar lines when downloading.
|
||||
|
||||
Login shells have the benefit of working very close to how a proper bash
|
||||
session does, regarding how scripts execute and how the environment is
|
||||
setup. This is also how all commands were ran in the past. The only reason
|
||||
to use login shells over non-interactive shells is if you need some weird
|
||||
and non-robust tool to work.
|
||||
|
||||
Args:
|
||||
val: If true, login shells will be used to run all commands.
|
||||
"""
|
||||
_config["use_login_shells"] = val
|
||||
|
||||
|
||||
def _with_environment_variables(cmd: str, environment_variables: Dict[str, object]):
|
||||
"""Prepend environment variables to a shell command.
|
||||
|
||||
Args:
|
||||
cmd: The base command.
|
||||
environment_variables: The set of environment variables. If an
|
||||
environment variable value is a dict, it will automatically be
|
||||
converted to a one line yaml string.
|
||||
|
||||
Returns:
|
||||
The base command prefixed with `export` statements for each variable.
|
||||
"""
|
||||
|
||||
as_strings = []
|
||||
for key, val in environment_variables.items():
|
||||
val = json.dumps(val, separators=(",", ":"))
|
||||
s = "export {}={};".format(key, quote(val))
|
||||
as_strings.append(s)
|
||||
all_vars = "".join(as_strings)
|
||||
return all_vars + cmd
|
||||
|
||||
|
||||
def _with_interactive(cmd):
|
||||
force_interactive = (
|
||||
f"source ~/.bashrc; "
|
||||
f"export OMP_NUM_THREADS=1 PYTHONWARNINGS=ignore && ({cmd})"
|
||||
)
|
||||
return ["bash", "--login", "-c", "-i", quote(force_interactive)]
|
||||
|
||||
|
||||
class SSHOptions:
|
||||
def __init__(self, ssh_key, control_path=None, **kwargs):
|
||||
self.ssh_key = ssh_key
|
||||
self.arg_dict = {
|
||||
# Supresses initial fingerprint verification.
|
||||
"StrictHostKeyChecking": "no",
|
||||
# SSH IP and fingerprint pairs no longer added to known_hosts.
|
||||
# This is to remove a "REMOTE HOST IDENTIFICATION HAS CHANGED"
|
||||
# warning if a new node has the same IP as a previously
|
||||
# deleted node, because the fingerprints will not match in
|
||||
# that case.
|
||||
"UserKnownHostsFile": os.devnull,
|
||||
# Try fewer extraneous key pairs.
|
||||
"IdentitiesOnly": "yes",
|
||||
# Abort if port forwarding fails (instead of just printing to
|
||||
# stderr).
|
||||
"ExitOnForwardFailure": "yes",
|
||||
# Quickly kill the connection if network connection breaks (as
|
||||
# opposed to hanging/blocking).
|
||||
"ServerAliveInterval": 5,
|
||||
"ServerAliveCountMax": 3,
|
||||
}
|
||||
if control_path:
|
||||
if sys.platform == "win32":
|
||||
# Don't set any control path options on Windows
|
||||
pass
|
||||
else:
|
||||
self.arg_dict.update(
|
||||
{
|
||||
"ControlMaster": "auto",
|
||||
"ControlPath": "{}/%C".format(control_path),
|
||||
"ControlPersist": "10s",
|
||||
}
|
||||
)
|
||||
|
||||
self.arg_dict.update(kwargs)
|
||||
|
||||
def to_ssh_options_list(self, *, timeout=60):
|
||||
self.arg_dict["ConnectTimeout"] = "{}s".format(timeout)
|
||||
ssh_key_option = ["-i", self.ssh_key] if self.ssh_key else []
|
||||
return ssh_key_option + [
|
||||
x
|
||||
for y in (
|
||||
["-o", "{}={}".format(k, v)]
|
||||
for k, v in self.arg_dict.items()
|
||||
if v is not None
|
||||
)
|
||||
for x in y
|
||||
]
|
||||
|
||||
|
||||
class SSHCommandRunner(CommandRunnerInterface):
|
||||
def __init__(
|
||||
self,
|
||||
log_prefix,
|
||||
node_id,
|
||||
provider,
|
||||
auth_config,
|
||||
cluster_name,
|
||||
process_runner,
|
||||
use_internal_ip,
|
||||
):
|
||||
|
||||
ssh_control_hash = hashlib.sha256(cluster_name.encode()).hexdigest()
|
||||
ssh_user_hash = hashlib.sha256(getuser().encode()).hexdigest()
|
||||
if sys.platform == "win32":
|
||||
# Disable SSH control paths on Windows - currently using it cause socket errors
|
||||
ssh_control_path = None
|
||||
else:
|
||||
ssh_control_path = "/tmp/ray_ssh_{}/{}".format(
|
||||
ssh_user_hash[:HASH_MAX_LENGTH], ssh_control_hash[:HASH_MAX_LENGTH]
|
||||
)
|
||||
|
||||
self.cluster_name = cluster_name
|
||||
self.log_prefix = log_prefix
|
||||
self.process_runner = process_runner
|
||||
self.node_id = node_id
|
||||
self.use_internal_ip = use_internal_ip
|
||||
self.provider = provider
|
||||
self.ssh_private_key = auth_config.get("ssh_private_key")
|
||||
self.ssh_user = auth_config["ssh_user"]
|
||||
self.ssh_control_path = ssh_control_path
|
||||
self.ssh_ip = None
|
||||
self.ssh_proxy_command = auth_config.get("ssh_proxy_command", None)
|
||||
self.ssh_options = SSHOptions(
|
||||
self.ssh_private_key,
|
||||
self.ssh_control_path,
|
||||
ProxyCommand=self.ssh_proxy_command,
|
||||
)
|
||||
|
||||
def _get_node_ip(self):
|
||||
if self.use_internal_ip:
|
||||
return self.provider.internal_ip(self.node_id)
|
||||
else:
|
||||
return self.provider.external_ip(self.node_id)
|
||||
|
||||
def _wait_for_ip(self, deadline):
|
||||
# if we have IP do not print waiting info
|
||||
ip = self._get_node_ip()
|
||||
if ip is not None:
|
||||
cli_logger.labeled_value("Fetched IP", ip)
|
||||
return ip
|
||||
|
||||
interval = AUTOSCALER_NODE_SSH_INTERVAL_S
|
||||
with cli_logger.group("Waiting for IP"):
|
||||
while time.time() < deadline and not self.provider.is_terminated(
|
||||
self.node_id
|
||||
):
|
||||
ip = self._get_node_ip()
|
||||
if ip is not None:
|
||||
cli_logger.labeled_value("Received", ip)
|
||||
return ip
|
||||
cli_logger.print(
|
||||
"Not yet available, retrying in {} seconds", cf.bold(str(interval))
|
||||
)
|
||||
time.sleep(interval)
|
||||
|
||||
return None
|
||||
|
||||
def _set_ssh_ip_if_required(self):
|
||||
if self.ssh_ip is not None:
|
||||
return
|
||||
|
||||
# We assume that this never changes.
|
||||
# I think that's reasonable.
|
||||
deadline = time.time() + AUTOSCALER_NODE_START_WAIT_S
|
||||
with LogTimer(self.log_prefix + "Got IP"):
|
||||
ip = self._wait_for_ip(deadline)
|
||||
|
||||
cli_logger.doassert(ip is not None, "Could not get node IP.") # todo: msg
|
||||
assert ip is not None, "Unable to find IP of node"
|
||||
|
||||
self.ssh_ip = ip
|
||||
|
||||
# This should run before any SSH commands and therefore ensure that
|
||||
# the ControlPath directory exists, allowing SSH to maintain
|
||||
# persistent sessions later on.
|
||||
if self.ssh_control_path is not None:
|
||||
try:
|
||||
os.makedirs(self.ssh_control_path, mode=0o700, exist_ok=True)
|
||||
except OSError as e:
|
||||
cli_logger.warning("{}", str(e)) # todo: msg
|
||||
|
||||
def _run_helper(
|
||||
self,
|
||||
final_cmd: List[str],
|
||||
with_output: bool = False,
|
||||
exit_on_fail: bool = False,
|
||||
silent: bool = False,
|
||||
):
|
||||
"""Run a command that was already setup with SSH and `bash` settings.
|
||||
|
||||
Args:
|
||||
final_cmd: Full command to run. Should include SSH options and
|
||||
other processing that we do.
|
||||
with_output: If `with_output` is `True`, command stdout will be
|
||||
captured and returned.
|
||||
exit_on_fail: If `exit_on_fail` is `True`, the process will exit
|
||||
if the command fails (exits with a code other than 0).
|
||||
silent: If true, the command output will be silenced.
|
||||
|
||||
Returns:
|
||||
Captured stdout bytes when `with_output` is True, otherwise the
|
||||
return value of ``run_cmd_redirected`` (typically the return code).
|
||||
|
||||
Raises:
|
||||
ProcessRunnerError: If using new log style and disabled
|
||||
login shells.
|
||||
click.ClickException: If using login shells.
|
||||
"""
|
||||
try:
|
||||
# For now, if the output is needed we just skip the new logic.
|
||||
# In the future we could update the new logic to support
|
||||
# capturing output, but it is probably not needed.
|
||||
if not with_output:
|
||||
return run_cmd_redirected(
|
||||
final_cmd,
|
||||
process_runner=self.process_runner,
|
||||
silent=silent,
|
||||
use_login_shells=is_using_login_shells(),
|
||||
)
|
||||
else:
|
||||
return self.process_runner.check_output(final_cmd)
|
||||
except subprocess.CalledProcessError as e:
|
||||
joined_cmd = " ".join(final_cmd)
|
||||
if not is_using_login_shells():
|
||||
raise ProcessRunnerError(
|
||||
"Command failed",
|
||||
"ssh_command_failed",
|
||||
code=e.returncode,
|
||||
command=joined_cmd,
|
||||
)
|
||||
|
||||
if exit_on_fail:
|
||||
raise click.ClickException(
|
||||
"Command failed:\n\n {}\n".format(joined_cmd)
|
||||
) from None
|
||||
else:
|
||||
fail_msg = "SSH command failed."
|
||||
if is_output_redirected():
|
||||
fail_msg += " See above for the output from the failure."
|
||||
raise click.ClickException(fail_msg) from None
|
||||
finally:
|
||||
# Do our best to flush output to terminal.
|
||||
# See https://github.com/ray-project/ray/pull/19473.
|
||||
sys.stdout.flush()
|
||||
sys.stderr.flush()
|
||||
|
||||
def run(
|
||||
self,
|
||||
cmd: Optional[str] = None,
|
||||
timeout: int = 120,
|
||||
exit_on_fail: bool = False,
|
||||
port_forward: Optional[List[Tuple[int, int]]] = None,
|
||||
with_output: bool = False,
|
||||
environment_variables: Optional[Dict[str, object]] = None,
|
||||
run_env: str = "auto", # Unused argument.
|
||||
ssh_options_override_ssh_key: str = "",
|
||||
shutdown_after_run: bool = False,
|
||||
silent: bool = False,
|
||||
) -> str:
|
||||
if shutdown_after_run:
|
||||
cmd += "; sudo shutdown -h now"
|
||||
|
||||
if ssh_options_override_ssh_key:
|
||||
if self.ssh_proxy_command:
|
||||
ssh_options = SSHOptions(
|
||||
ssh_options_override_ssh_key, ProxyCommand=self.ssh_proxy_command
|
||||
)
|
||||
else:
|
||||
ssh_options = SSHOptions(ssh_options_override_ssh_key)
|
||||
else:
|
||||
ssh_options = self.ssh_options
|
||||
|
||||
assert isinstance(
|
||||
ssh_options, SSHOptions
|
||||
), "ssh_options must be of type SSHOptions, got {}".format(type(ssh_options))
|
||||
|
||||
self._set_ssh_ip_if_required()
|
||||
|
||||
if is_using_login_shells():
|
||||
ssh = ["ssh", "-tt"]
|
||||
else:
|
||||
ssh = ["ssh"]
|
||||
|
||||
if port_forward:
|
||||
with cli_logger.group("Forwarding ports"):
|
||||
if not isinstance(port_forward, list):
|
||||
port_forward = [port_forward]
|
||||
for local, remote in port_forward:
|
||||
cli_logger.verbose(
|
||||
"Forwarding port {} to port {} on localhost.",
|
||||
cf.bold(local),
|
||||
cf.bold(remote),
|
||||
) # todo: msg
|
||||
ssh += ["-L", "{}:localhost:{}".format(local, remote)]
|
||||
|
||||
final_cmd = (
|
||||
ssh
|
||||
+ ssh_options.to_ssh_options_list(timeout=timeout)
|
||||
+ ["{}@{}".format(self.ssh_user, self.ssh_ip)]
|
||||
)
|
||||
if cmd:
|
||||
if environment_variables:
|
||||
cmd = _with_environment_variables(cmd, environment_variables)
|
||||
if is_using_login_shells():
|
||||
final_cmd += _with_interactive(cmd)
|
||||
else:
|
||||
final_cmd += [cmd]
|
||||
else:
|
||||
# We do this because `-o ControlMaster` causes the `-N` flag to
|
||||
# still create an interactive shell in some ssh versions.
|
||||
final_cmd.append("while true; do sleep 86400; done")
|
||||
|
||||
cli_logger.verbose("Running `{}`", cf.bold(cmd))
|
||||
with cli_logger.indented():
|
||||
cli_logger.very_verbose(
|
||||
"Full command is `{}`", cf.bold(" ".join(final_cmd))
|
||||
)
|
||||
|
||||
if cli_logger.verbosity > 0:
|
||||
with cli_logger.indented():
|
||||
return self._run_helper(
|
||||
final_cmd, with_output, exit_on_fail, silent=silent
|
||||
)
|
||||
else:
|
||||
return self._run_helper(final_cmd, with_output, exit_on_fail, silent=silent)
|
||||
|
||||
def _create_rsync_filter_args(self, options):
|
||||
rsync_excludes = options.get("rsync_exclude") or []
|
||||
rsync_filters = options.get("rsync_filter") or []
|
||||
|
||||
exclude_args = [
|
||||
["--exclude", rsync_exclude] for rsync_exclude in rsync_excludes
|
||||
]
|
||||
filter_args = [
|
||||
["--filter", "dir-merge,- {}".format(rsync_filter)]
|
||||
for rsync_filter in rsync_filters
|
||||
]
|
||||
|
||||
# Combine and flatten the two lists
|
||||
return [arg for args_list in exclude_args + filter_args for arg in args_list]
|
||||
|
||||
def run_rsync_up(self, source, target, options=None):
|
||||
self._set_ssh_ip_if_required()
|
||||
options = options or {}
|
||||
|
||||
# on windows use scp -r instead of rsync
|
||||
if sys.platform == "win32":
|
||||
# Use scp as fallback for Windows
|
||||
command = ["scp", "-r"]
|
||||
command += self.ssh_options.to_ssh_options_list(timeout=120)
|
||||
command += [source, "{}@{}:{}".format(self.ssh_user, self.ssh_ip, target)]
|
||||
else:
|
||||
command = ["rsync"]
|
||||
command += [
|
||||
"--rsh",
|
||||
subprocess.list2cmdline(
|
||||
["ssh"] + self.ssh_options.to_ssh_options_list(timeout=120)
|
||||
),
|
||||
]
|
||||
command += ["-avz"]
|
||||
command += self._create_rsync_filter_args(options=options)
|
||||
command += [source, "{}@{}:{}".format(self.ssh_user, self.ssh_ip, target)]
|
||||
|
||||
cli_logger.verbose("Running `{}`", cf.bold(" ".join(command)))
|
||||
self._run_helper(command, silent=is_rsync_silent())
|
||||
|
||||
def run_rsync_down(self, source, target, options=None):
|
||||
self._set_ssh_ip_if_required()
|
||||
|
||||
# on Windows use scp -r instead of rsync
|
||||
if sys.platform == "win32":
|
||||
# Use scp as fallback for Windows
|
||||
command = ["scp", "-r"]
|
||||
command += self.ssh_options.to_ssh_options_list(timeout=120)
|
||||
command += ["{}@{}:{}".format(self.ssh_user, self.ssh_ip, source), target]
|
||||
else:
|
||||
command = ["rsync"]
|
||||
command += [
|
||||
"--rsh",
|
||||
subprocess.list2cmdline(
|
||||
["ssh"] + self.ssh_options.to_ssh_options_list(timeout=120)
|
||||
),
|
||||
]
|
||||
command += ["-avz"]
|
||||
command += self._create_rsync_filter_args(options=options)
|
||||
command += ["{}@{}:{}".format(self.ssh_user, self.ssh_ip, source), target]
|
||||
|
||||
cli_logger.verbose("Running `{}`", cf.bold(" ".join(command)))
|
||||
self._run_helper(command, silent=is_rsync_silent())
|
||||
|
||||
def remote_shell_command_str(self):
|
||||
if self.ssh_private_key:
|
||||
return "ssh -o IdentitiesOnly=yes -i {} {}@{}\n".format(
|
||||
self.ssh_private_key, self.ssh_user, self.ssh_ip
|
||||
)
|
||||
else:
|
||||
return "ssh -o IdentitiesOnly=yes {}@{}\n".format(
|
||||
self.ssh_user, self.ssh_ip
|
||||
)
|
||||
|
||||
|
||||
class DockerCommandRunner(CommandRunnerInterface):
|
||||
def __init__(self, docker_config, **common_args):
|
||||
self.ssh_command_runner = SSHCommandRunner(**common_args)
|
||||
self.container_name = docker_config["container_name"]
|
||||
self.docker_config = docker_config
|
||||
self.home_dir = None
|
||||
self.initialized = False
|
||||
# Optionally use 'podman' instead of 'docker'
|
||||
use_podman = docker_config.get("use_podman", False)
|
||||
self.docker_cmd = "podman" if use_podman else "docker"
|
||||
|
||||
def run(
|
||||
self,
|
||||
cmd: Optional[str] = None,
|
||||
timeout: int = 120,
|
||||
exit_on_fail: bool = False,
|
||||
port_forward: Optional[List[Tuple[int, int]]] = None,
|
||||
with_output: bool = False,
|
||||
environment_variables: Optional[Dict[str, object]] = None,
|
||||
run_env: str = "auto",
|
||||
ssh_options_override_ssh_key: str = "",
|
||||
shutdown_after_run: bool = False,
|
||||
) -> str:
|
||||
if run_env == "auto":
|
||||
run_env = (
|
||||
"host"
|
||||
if (not bool(cmd) or cmd.find(self.docker_cmd) == 0)
|
||||
else self.docker_cmd
|
||||
)
|
||||
|
||||
if environment_variables:
|
||||
cmd = _with_environment_variables(cmd, environment_variables)
|
||||
|
||||
if run_env == self.docker_cmd:
|
||||
cmd = self._docker_expand_user(cmd, any_char=True)
|
||||
if is_using_login_shells():
|
||||
cmd = " ".join(_with_interactive(cmd))
|
||||
cmd = with_docker_exec(
|
||||
[cmd],
|
||||
container_name=self.container_name,
|
||||
with_interactive=is_using_login_shells(),
|
||||
docker_cmd=self.docker_cmd,
|
||||
)[0]
|
||||
|
||||
if shutdown_after_run:
|
||||
# sudo shutdown should run after `with_docker_exec` command above
|
||||
cmd += "; sudo shutdown -h now"
|
||||
# Do not pass shutdown_after_run argument to ssh_command_runner.run()
|
||||
# since it is handled above.
|
||||
return self.ssh_command_runner.run(
|
||||
cmd,
|
||||
timeout=timeout,
|
||||
exit_on_fail=exit_on_fail,
|
||||
port_forward=port_forward,
|
||||
with_output=with_output,
|
||||
ssh_options_override_ssh_key=ssh_options_override_ssh_key,
|
||||
)
|
||||
|
||||
def run_rsync_up(self, source, target, options=None):
|
||||
options = options or {}
|
||||
host_destination = os.path.join(
|
||||
self._get_docker_host_mount_location(self.ssh_command_runner.cluster_name),
|
||||
target.lstrip("/"),
|
||||
)
|
||||
host_mount_location = os.path.dirname(host_destination.rstrip("/"))
|
||||
if sys.platform == "win32":
|
||||
# fix paths if running on Windows
|
||||
source = source.replace("\\", "/")
|
||||
host_mount_location = host_mount_location.replace("\\", "/")
|
||||
host_destination = host_destination.replace("\\", "/")
|
||||
|
||||
self.ssh_command_runner.run(
|
||||
f"mkdir -p {host_mount_location} && chown -R "
|
||||
f"{self.ssh_command_runner.ssh_user} {host_mount_location}",
|
||||
silent=is_rsync_silent(),
|
||||
)
|
||||
|
||||
self.ssh_command_runner.run_rsync_up(source, host_destination, options=options)
|
||||
if self._check_container_status() and not options.get(
|
||||
"docker_mount_if_possible", False
|
||||
):
|
||||
if os.path.isdir(source):
|
||||
# Adding a "." means that docker copies the *contents*
|
||||
# Without it, docker copies the source *into* the target
|
||||
host_destination += "/."
|
||||
|
||||
# This path may not exist inside the container. This ensures
|
||||
# that the path is created!
|
||||
prefix = with_docker_exec(
|
||||
[
|
||||
"mkdir -p {}".format(
|
||||
os.path.dirname(self._docker_expand_user(target))
|
||||
)
|
||||
],
|
||||
container_name=self.container_name,
|
||||
with_interactive=is_using_login_shells(),
|
||||
docker_cmd=self.docker_cmd,
|
||||
)[0]
|
||||
|
||||
self.ssh_command_runner.run(
|
||||
"{} && rsync -e '{} exec -i' -avz {} {}:{}".format(
|
||||
prefix,
|
||||
self.docker_cmd,
|
||||
host_destination,
|
||||
self.container_name,
|
||||
self._docker_expand_user(target),
|
||||
),
|
||||
silent=is_rsync_silent(),
|
||||
)
|
||||
|
||||
def run_rsync_down(self, source, target, options=None):
|
||||
options = options or {}
|
||||
host_source = os.path.join(
|
||||
self._get_docker_host_mount_location(self.ssh_command_runner.cluster_name),
|
||||
source.lstrip("/"),
|
||||
)
|
||||
host_mount_location = os.path.dirname(host_source.rstrip("/"))
|
||||
# Convert Windows paths to Unix-style for remote commands
|
||||
host_mount_location_unix = host_mount_location.replace("\\", "/")
|
||||
self.ssh_command_runner.run(
|
||||
f"mkdir -p {host_mount_location_unix} && chown -R "
|
||||
f"{self.ssh_command_runner.ssh_user} {host_mount_location_unix}",
|
||||
silent=is_rsync_silent(),
|
||||
)
|
||||
if source[-1] == "/":
|
||||
source += "."
|
||||
# Adding a "." means that docker copies the *contents*
|
||||
# Without it, docker copies the source *into* the target
|
||||
if not options.get("docker_mount_if_possible", False):
|
||||
# NOTE: `--delete` is okay here because the container is the source
|
||||
# of truth.
|
||||
self.ssh_command_runner.run(
|
||||
"rsync -e '{} exec -i' -avz --delete {}:{} {}".format(
|
||||
self.docker_cmd,
|
||||
self.container_name,
|
||||
self._docker_expand_user(source),
|
||||
host_source.replace(
|
||||
"\\", "/"
|
||||
), # Convert Windows paths to Unix-style for rsync
|
||||
),
|
||||
silent=is_rsync_silent(),
|
||||
)
|
||||
self.ssh_command_runner.run_rsync_down(host_source, target, options=options)
|
||||
|
||||
def remote_shell_command_str(self):
|
||||
inner_str = (
|
||||
self.ssh_command_runner.remote_shell_command_str()
|
||||
.replace("ssh", "ssh -tt", 1)
|
||||
.strip("\n")
|
||||
)
|
||||
return inner_str + " {} exec -it {} /bin/bash\n".format(
|
||||
self.docker_cmd, self.container_name
|
||||
)
|
||||
|
||||
def _check_docker_installed(self):
|
||||
no_exist = "NoExist"
|
||||
output = self.ssh_command_runner.run(
|
||||
f"command -v {self.docker_cmd} || echo '{no_exist}'", with_output=True
|
||||
)
|
||||
cleaned_output = output.decode().strip()
|
||||
if no_exist in cleaned_output or "docker" not in cleaned_output:
|
||||
if self.docker_cmd == "docker":
|
||||
install_commands = [
|
||||
"curl -fsSL https://get.docker.com -o get-docker.sh",
|
||||
"sudo sh get-docker.sh",
|
||||
"sudo usermod -aG docker $USER",
|
||||
"sudo systemctl restart docker -f",
|
||||
]
|
||||
else:
|
||||
install_commands = [
|
||||
"sudo apt-get update",
|
||||
"sudo apt-get -y install podman",
|
||||
]
|
||||
|
||||
logger.error(
|
||||
f"{self.docker_cmd.capitalize()} not installed. You can "
|
||||
f"install {self.docker_cmd.capitalize()} by adding the "
|
||||
"following commands to 'initialization_commands':\n"
|
||||
+ "\n".join(install_commands)
|
||||
)
|
||||
|
||||
def _check_container_status(self):
|
||||
if self.initialized:
|
||||
return True
|
||||
output = (
|
||||
self.ssh_command_runner.run(
|
||||
check_docker_running_cmd(self.container_name, self.docker_cmd),
|
||||
with_output=True,
|
||||
)
|
||||
.decode("utf-8")
|
||||
.strip()
|
||||
)
|
||||
# Checks for the false positive where "true" is in the container name
|
||||
return "true" in output.lower() and "no such object" not in output.lower()
|
||||
|
||||
def _docker_expand_user(self, string, any_char=False):
|
||||
user_pos = string.find("~")
|
||||
if user_pos > -1:
|
||||
if self.home_dir is None:
|
||||
self.home_dir = (
|
||||
self.ssh_command_runner.run(
|
||||
f"{self.docker_cmd} exec {self.container_name} "
|
||||
"printenv HOME",
|
||||
with_output=True,
|
||||
)
|
||||
.decode("utf-8")
|
||||
.strip()
|
||||
)
|
||||
|
||||
if any_char:
|
||||
return string.replace("~/", self.home_dir + "/")
|
||||
|
||||
elif not any_char and user_pos == 0:
|
||||
return string.replace("~", self.home_dir, 1)
|
||||
|
||||
return string
|
||||
|
||||
def _check_if_container_restart_is_needed(
|
||||
self, image: str, cleaned_bind_mounts: Dict[str, str]
|
||||
) -> bool:
|
||||
re_init_required = False
|
||||
running_image = (
|
||||
self.run(
|
||||
check_docker_image(self.container_name, self.docker_cmd),
|
||||
with_output=True,
|
||||
run_env="host",
|
||||
)
|
||||
.decode("utf-8")
|
||||
.strip()
|
||||
)
|
||||
if running_image != image:
|
||||
cli_logger.error(
|
||||
"A container with name {} is running image {} instead "
|
||||
+ "of {} (which was provided in the YAML)",
|
||||
self.container_name,
|
||||
running_image,
|
||||
image,
|
||||
)
|
||||
mounts = (
|
||||
self.run(
|
||||
check_bind_mounts_cmd(self.container_name, self.docker_cmd),
|
||||
with_output=True,
|
||||
run_env="host",
|
||||
)
|
||||
.decode("utf-8")
|
||||
.strip()
|
||||
)
|
||||
try:
|
||||
active_mounts = json.loads(mounts)
|
||||
active_remote_mounts = {
|
||||
mnt["Destination"].strip("/") for mnt in active_mounts
|
||||
}
|
||||
# Ignore ray bootstrap files.
|
||||
requested_remote_mounts = {
|
||||
self._docker_expand_user(remote).strip("/")
|
||||
for remote in cleaned_bind_mounts.keys()
|
||||
}
|
||||
unfulfilled_mounts = requested_remote_mounts - active_remote_mounts
|
||||
if unfulfilled_mounts:
|
||||
re_init_required = True
|
||||
cli_logger.warning(
|
||||
"This Docker Container is already running. "
|
||||
"Restarting the Docker container on "
|
||||
"this node to pick up the following file_mounts {}",
|
||||
unfulfilled_mounts,
|
||||
)
|
||||
except json.JSONDecodeError:
|
||||
cli_logger.verbose(
|
||||
"Unable to check if file_mounts specified in the YAML "
|
||||
"differ from those on the running container."
|
||||
)
|
||||
return re_init_required
|
||||
|
||||
def run_init(
|
||||
self, *, as_head: bool, file_mounts: Dict[str, str], sync_run_yet: bool
|
||||
):
|
||||
BOOTSTRAP_MOUNTS = ["~/ray_bootstrap_config.yaml", "~/ray_bootstrap_key.pem"]
|
||||
|
||||
specific_image = self.docker_config.get(
|
||||
f"{'head' if as_head else 'worker'}_image", self.docker_config.get("image")
|
||||
)
|
||||
|
||||
self._check_docker_installed()
|
||||
if self.docker_config.get("pull_before_run", True):
|
||||
assert specific_image, (
|
||||
"Image must be included in config if " + "pull_before_run is specified"
|
||||
)
|
||||
self.run(
|
||||
"{} pull {}".format(self.docker_cmd, specific_image), run_env="host"
|
||||
)
|
||||
else:
|
||||
self.run(
|
||||
f"{self.docker_cmd} image inspect {specific_image} "
|
||||
"1> /dev/null 2>&1 || "
|
||||
f"{self.docker_cmd} pull {specific_image}"
|
||||
)
|
||||
|
||||
# Bootstrap files cannot be bind mounted because docker opens the
|
||||
# underlying inode. When the file is switched, docker becomes outdated.
|
||||
cleaned_bind_mounts = file_mounts.copy()
|
||||
for mnt in BOOTSTRAP_MOUNTS:
|
||||
cleaned_bind_mounts.pop(mnt, None)
|
||||
|
||||
docker_run_executed = False
|
||||
|
||||
container_running = self._check_container_status()
|
||||
requires_re_init = False
|
||||
if container_running:
|
||||
requires_re_init = self._check_if_container_restart_is_needed(
|
||||
specific_image, cleaned_bind_mounts
|
||||
)
|
||||
if requires_re_init:
|
||||
docker_stop_cmd = f"{self.docker_cmd} stop {self.container_name}"
|
||||
logger.info("Executing Docker command: %s", docker_stop_cmd)
|
||||
self.run(docker_stop_cmd, run_env="host")
|
||||
|
||||
if (not container_running) or requires_re_init:
|
||||
if not sync_run_yet:
|
||||
# Do not start the actual image as we need to run file_sync
|
||||
# first to ensure that all folders are created with the
|
||||
# correct ownership. Docker will create the folders with
|
||||
# `root` as the owner.
|
||||
return True
|
||||
# Get home directory
|
||||
image_env = (
|
||||
self.ssh_command_runner.run(
|
||||
f"{self.docker_cmd} "
|
||||
+ "inspect -f '{{json .Config.Env}}' "
|
||||
+ specific_image,
|
||||
with_output=True,
|
||||
)
|
||||
.decode()
|
||||
.strip()
|
||||
)
|
||||
home_directory = "/root"
|
||||
try:
|
||||
for env_var in json.loads(image_env):
|
||||
if env_var.startswith("HOME="):
|
||||
home_directory = env_var.split("HOME=")[1]
|
||||
break
|
||||
except json.JSONDecodeError as e:
|
||||
cli_logger.error(
|
||||
"Unable to deserialize `image_env` to Python object. "
|
||||
f"The `image_env` is:\n{image_env}"
|
||||
)
|
||||
raise e
|
||||
|
||||
user_docker_run_options = self.docker_config.get(
|
||||
"run_options", []
|
||||
) + self.docker_config.get(
|
||||
f"{'head' if as_head else 'worker'}_run_options", []
|
||||
)
|
||||
start_command = docker_start_cmds(
|
||||
self.ssh_command_runner.ssh_user,
|
||||
specific_image,
|
||||
cleaned_bind_mounts,
|
||||
self.container_name,
|
||||
self._configure_runtime(
|
||||
self._auto_configure_shm(user_docker_run_options)
|
||||
),
|
||||
self.ssh_command_runner.cluster_name,
|
||||
home_directory,
|
||||
self.docker_cmd,
|
||||
)
|
||||
self.run(start_command, run_env="host")
|
||||
docker_run_executed = True
|
||||
|
||||
# Explicitly copy in ray bootstrap files.
|
||||
for mount in BOOTSTRAP_MOUNTS:
|
||||
if mount in file_mounts:
|
||||
if not sync_run_yet:
|
||||
# NOTE(ilr) This rsync is needed because when starting from
|
||||
# a stopped instance, /tmp may be deleted and `run_init`
|
||||
# is called before the first `file_sync` happens
|
||||
self.run_rsync_up(file_mounts[mount], mount)
|
||||
self.ssh_command_runner.run(
|
||||
"rsync -e '{cmd} exec -i' -avz {src} {container}:{dst}".format(
|
||||
cmd=self.docker_cmd,
|
||||
src=os.path.join(
|
||||
self._get_docker_host_mount_location(
|
||||
self.ssh_command_runner.cluster_name
|
||||
),
|
||||
mount,
|
||||
).replace(
|
||||
"\\", "/"
|
||||
), # Convert Windows paths to Unix-style for rsync
|
||||
container=self.container_name,
|
||||
dst=self._docker_expand_user(mount),
|
||||
)
|
||||
)
|
||||
try:
|
||||
# Check if the current user has read permission.
|
||||
# If they do not, try to change ownership!
|
||||
self.run(
|
||||
f"cat {mount} >/dev/null 2>&1 || "
|
||||
f"sudo chown $(id -u):$(id -g) {mount}"
|
||||
)
|
||||
except Exception:
|
||||
lsl_string = (
|
||||
self.run(f"ls -l {mount}", with_output=True)
|
||||
.decode("utf-8")
|
||||
.strip()
|
||||
)
|
||||
# The string is of format <Permission> <Links>
|
||||
# <Owner> <Group> <Size> <Date> <Name>
|
||||
permissions = lsl_string.split(" ")[0]
|
||||
owner = lsl_string.split(" ")[2]
|
||||
group = lsl_string.split(" ")[3]
|
||||
current_user = (
|
||||
self.run("whoami", with_output=True).decode("utf-8").strip()
|
||||
)
|
||||
cli_logger.warning(
|
||||
f"File ({mount}) is owned by user:{owner} and group:"
|
||||
f"{group} with permissions ({permissions}). The "
|
||||
f"current user ({current_user}) does not have "
|
||||
"permission to read these files, and Ray may not be "
|
||||
"able to autoscale. This can be resolved by "
|
||||
"installing `sudo` in your container, or adding a "
|
||||
f"command like 'chown {current_user} {mount}' to "
|
||||
"your `setup_commands`."
|
||||
)
|
||||
self.initialized = True
|
||||
return docker_run_executed
|
||||
|
||||
def _configure_runtime(self, run_options: List[str]) -> List[str]:
|
||||
if self.docker_config.get("disable_automatic_runtime_detection"):
|
||||
return run_options
|
||||
|
||||
runtime_output = (
|
||||
self.ssh_command_runner.run(
|
||||
f"{self.docker_cmd} " + "info -f '{{.Runtimes}}' ", with_output=True
|
||||
)
|
||||
.decode()
|
||||
.strip()
|
||||
)
|
||||
if "nvidia-container-runtime" in runtime_output:
|
||||
try:
|
||||
self.ssh_command_runner.run("nvidia-smi", with_output=False)
|
||||
return run_options + ["--runtime=nvidia"]
|
||||
except Exception as e:
|
||||
logger.warning(
|
||||
"NVIDIA Container Runtime is present, but no GPUs found."
|
||||
)
|
||||
logger.debug(f"nvidia-smi error: {e}")
|
||||
return run_options
|
||||
|
||||
return run_options
|
||||
|
||||
def _auto_configure_shm(self, run_options: List[str]) -> List[str]:
|
||||
if self.docker_config.get("disable_shm_size_detection"):
|
||||
return run_options
|
||||
for run_opt in run_options:
|
||||
if "--shm-size" in run_opt:
|
||||
logger.info(
|
||||
"Bypassing automatic SHM-Detection because of "
|
||||
f"`run_option`: {run_opt}"
|
||||
)
|
||||
return run_options
|
||||
try:
|
||||
shm_output = (
|
||||
self.ssh_command_runner.run(
|
||||
"cat /proc/meminfo || true", with_output=True
|
||||
)
|
||||
.decode()
|
||||
.strip()
|
||||
)
|
||||
available_memory = int(
|
||||
[ln for ln in shm_output.split("\n") if "MemAvailable" in ln][
|
||||
0
|
||||
].split()[1]
|
||||
)
|
||||
available_memory_bytes = available_memory * 1024
|
||||
# Overestimate SHM size by 10%
|
||||
shm_size = min(
|
||||
(available_memory_bytes * DEFAULT_OBJECT_STORE_MEMORY_PROPORTION * 1.1),
|
||||
DEFAULT_OBJECT_STORE_MAX_MEMORY_BYTES,
|
||||
)
|
||||
return run_options + [f"--shm-size='{shm_size}b'"]
|
||||
except Exception as e:
|
||||
logger.warning(f"Received error while trying to auto-compute SHM size {e}")
|
||||
return run_options
|
||||
|
||||
def _get_docker_host_mount_location(self, cluster_name: str) -> str:
|
||||
"""Return the docker host mount directory location."""
|
||||
# Imported here due to circular dependency in imports.
|
||||
from ray.autoscaler.sdk import get_docker_host_mount_location
|
||||
|
||||
return get_docker_host_mount_location(cluster_name)
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,164 @@
|
||||
import os
|
||||
import sys
|
||||
|
||||
from ray._private.ray_constants import ( # noqa F401
|
||||
AGENT_PROCESS_TYPE_DASHBOARD_AGENT,
|
||||
AGENT_PROCESS_TYPE_RUNTIME_ENV_AGENT,
|
||||
AUTOSCALER_RESOURCE_REQUEST_CHANNEL,
|
||||
DEFAULT_OBJECT_STORE_MEMORY_PROPORTION,
|
||||
LABELS_ENVIRONMENT_VARIABLE,
|
||||
LOGGER_FORMAT,
|
||||
RESOURCES_ENVIRONMENT_VARIABLE,
|
||||
)
|
||||
|
||||
|
||||
def env_integer(key, default):
|
||||
if key in os.environ:
|
||||
val = os.environ[key]
|
||||
if val == "inf":
|
||||
return sys.maxsize
|
||||
else:
|
||||
return int(val)
|
||||
return default
|
||||
|
||||
|
||||
# Whether autoscaler cluster status logging is enabled. Set to 0 disable.
|
||||
AUTOSCALER_STATUS_LOG = env_integer("RAY_ENABLE_CLUSTER_STATUS_LOG", 1)
|
||||
|
||||
# The name of the environment variable for plugging in a utilization scorer.
|
||||
AUTOSCALER_UTILIZATION_SCORER_KEY = "RAY_AUTOSCALER_UTILIZATION_SCORER"
|
||||
|
||||
# Whether to avoid launching GPU nodes for CPU only tasks.
|
||||
AUTOSCALER_CONSERVE_GPU_NODES = env_integer("AUTOSCALER_CONSERVE_GPU_NODES", 1)
|
||||
|
||||
# How long to wait for a node to start and terminate, in seconds.
|
||||
AUTOSCALER_NODE_START_WAIT_S = env_integer("AUTOSCALER_NODE_START_WAIT_S", 900)
|
||||
AUTOSCALER_NODE_TERMINATE_WAIT_S = env_integer("AUTOSCALER_NODE_TERMINATE_WAIT_S", 900)
|
||||
|
||||
# Interval at which to check if node SSH became available.
|
||||
AUTOSCALER_NODE_SSH_INTERVAL_S = env_integer("AUTOSCALER_NODE_SSH_INTERVAL_S", 5)
|
||||
|
||||
# Abort autoscaling if more than this number of errors are encountered. This
|
||||
# is a safety feature to prevent e.g. runaway node launches.
|
||||
AUTOSCALER_MAX_NUM_FAILURES = env_integer("AUTOSCALER_MAX_NUM_FAILURES", 5)
|
||||
|
||||
# The maximum number of nodes to launch in a single request.
|
||||
# Multiple requests may be made for this batch size, up to
|
||||
# the limit of AUTOSCALER_MAX_CONCURRENT_LAUNCHES.
|
||||
AUTOSCALER_MAX_LAUNCH_BATCH = env_integer("AUTOSCALER_MAX_LAUNCH_BATCH", 5)
|
||||
|
||||
# Max number of nodes to launch at a time.
|
||||
AUTOSCALER_MAX_CONCURRENT_LAUNCHES = env_integer(
|
||||
"AUTOSCALER_MAX_CONCURRENT_LAUNCHES", 10
|
||||
)
|
||||
|
||||
# Default upscaling speed for the autoscaler. This specifies how many nodes
|
||||
# to request at a time, where the desired number to upscale is
|
||||
# min(1, upscaling_speed * current_num_nodes)
|
||||
# e.g. 1.0 means to request enough nodes to double
|
||||
# the cluster size in each round of requests.
|
||||
# When the upscaling speed is 0.0, the autoscaler will request 1 node.
|
||||
DEFAULT_UPSCALING_SPEED = 0.0
|
||||
|
||||
# Interval at which to perform autoscaling updates.
|
||||
AUTOSCALER_UPDATE_INTERVAL_S = env_integer("AUTOSCALER_UPDATE_INTERVAL_S", 5)
|
||||
|
||||
# The autoscaler will attempt to restart Ray on nodes it hasn't heard from
|
||||
# in more than this interval.
|
||||
AUTOSCALER_HEARTBEAT_TIMEOUT_S = env_integer("AUTOSCALER_HEARTBEAT_TIMEOUT_S", 30)
|
||||
# The maximum number of nodes (including failed nodes) that the autoscaler will
|
||||
# track for logging purposes.
|
||||
AUTOSCALER_MAX_NODES_TRACKED = 1500
|
||||
|
||||
AUTOSCALER_MAX_FAILURES_DISPLAYED = 20
|
||||
|
||||
AUTOSCALER_NODE_AVAILABILITY_MAX_STALENESS_S = env_integer(
|
||||
"AUTOSCALER_NODE_AVAILABILITY_MAX_STALENESS_S", 30 * 60
|
||||
)
|
||||
|
||||
# The window during which the recoverable cloud resource availability score linearly recovers from 0.0 to 1.0.
|
||||
RAY_AUTOSCALER_AVAILABILITY_RECOVERY_S = env_integer(
|
||||
"RAY_AUTOSCALER_AVAILABILITY_RECOVERY_S", 600
|
||||
)
|
||||
|
||||
AUTOSCALER_REPORT_PER_NODE_STATUS = (
|
||||
env_integer("AUTOSCALER_REPORT_PER_NODE_STATUS", 1) == 1
|
||||
)
|
||||
|
||||
# The maximum allowed resource demand vector size to guarantee the resource
|
||||
# demand scheduler bin packing algorithm takes a reasonable amount of time
|
||||
# to run.
|
||||
AUTOSCALER_MAX_RESOURCE_DEMAND_VECTOR_SIZE = env_integer(
|
||||
"AUTOSCALER_MAX_RESOURCE_DEMAND_VECTOR_SIZE", 1000
|
||||
)
|
||||
|
||||
# Port that autoscaler prometheus metrics will be exported to
|
||||
AUTOSCALER_METRIC_PORT = env_integer("AUTOSCALER_METRIC_PORT", 44217)
|
||||
|
||||
# The minimum number of nodes to launch concurrently.
|
||||
AUTOSCALER_UPSCALING_INITIAL_NUM_NODES = 5
|
||||
|
||||
# Max number of retries to AWS (default is 5, time increases exponentially)
|
||||
BOTO_MAX_RETRIES = env_integer("BOTO_MAX_RETRIES", 12)
|
||||
# Max number of retries to create an EC2 node (retry different subnet)
|
||||
BOTO_CREATE_MAX_RETRIES = env_integer("BOTO_CREATE_MAX_RETRIES", 5)
|
||||
|
||||
# ray home path in the container image
|
||||
RAY_HOME = "/home/ray"
|
||||
|
||||
# The order of this list matters! `scripts.py` kills the ray processes in order of this
|
||||
# list. Think twice when you add to this list.
|
||||
# Invariants:
|
||||
# RAYLET must be the first in the list.
|
||||
# GCS SERVER must be the last in the list.
|
||||
RAY_PROCESSES = [
|
||||
# The first element is the substring to filter.
|
||||
# The second element, if True, is to filter ps results by command name
|
||||
# (only the first 15 charactors of the executable name on Linux);
|
||||
# if False, is to filter ps results by command with all its arguments.
|
||||
# See STANDARD FORMAT SPECIFIERS section of
|
||||
# http://man7.org/linux/man-pages/man1/ps.1.html
|
||||
# about comm and args. This can help avoid killing non-ray processes.
|
||||
# Format:
|
||||
# Keyword to filter, filter by command (True)/filter by args (False)
|
||||
["raylet", True],
|
||||
["plasma_store", True],
|
||||
["monitor.py", False],
|
||||
["ray.util.client.server", False],
|
||||
["default_worker.py", False], # Python worker.
|
||||
["setup_worker.py", False], # Python environment setup worker.
|
||||
# For mac osx, setproctitle doesn't change the process name returned
|
||||
# by psutil but only cmdline.
|
||||
[
|
||||
"ray::",
|
||||
sys.platform != "darwin",
|
||||
], # Python worker. TODO(mehrdadn): Fix for Windows
|
||||
["io.ray.runtime.runner.worker.DefaultWorker", False], # Java worker.
|
||||
["log_monitor.py", False],
|
||||
[AGENT_PROCESS_TYPE_DASHBOARD_AGENT, False],
|
||||
[os.path.join("dashboard", "dashboard.py"), False],
|
||||
[AGENT_PROCESS_TYPE_RUNTIME_ENV_AGENT, False],
|
||||
# On Windows, setproctitle does not change the process name or command line
|
||||
# visible to psutil, so the "ray::DashboardAgent" / "ray::RuntimeEnvAgent"
|
||||
# keywords above will never match. Add fallback entries that match the
|
||||
# actual script paths in the command line. See
|
||||
# https://github.com/ray-project/ray/issues/61452
|
||||
*(
|
||||
[
|
||||
[os.path.join("dashboard", "agent.py"), False],
|
||||
[os.path.join("runtime_env", "agent", "main.py"), False],
|
||||
]
|
||||
if sys.platform == "win32"
|
||||
else []
|
||||
),
|
||||
["ray_process_reaper.py", False],
|
||||
["gcs_server", True],
|
||||
]
|
||||
|
||||
# Max Concurrent SSH Calls to stop Docker
|
||||
MAX_PARALLEL_SHUTDOWN_WORKERS = env_integer("MAX_PARALLEL_SHUTDOWN_WORKERS", 50)
|
||||
|
||||
DISABLE_NODE_UPDATERS_KEY = "disable_node_updaters"
|
||||
DISABLE_LAUNCH_CONFIG_CHECK_KEY = "disable_launch_config_check"
|
||||
FOREGROUND_NODE_LAUNCH_KEY = "foreground_node_launch"
|
||||
WORKER_LIVENESS_CHECK_KEY = "worker_liveness_check"
|
||||
@@ -0,0 +1,130 @@
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict
|
||||
|
||||
from ray.autoscaler._private.cli_logger import cli_logger
|
||||
|
||||
try: # py3
|
||||
from shlex import quote
|
||||
except ImportError: # py2
|
||||
from pipes import quote
|
||||
|
||||
|
||||
def _check_docker_file_mounts(file_mounts: Dict[str, str]) -> None:
|
||||
"""Checks if files are passed as file_mounts. This is a problem for Docker
|
||||
based clusters because when a file is bind-mounted in Docker, updates to
|
||||
the file on the host do not always propagate to the container. Using
|
||||
directories is recommended.
|
||||
"""
|
||||
for remote, local in file_mounts.items():
|
||||
if Path(local).is_file():
|
||||
cli_logger.warning(
|
||||
f"File Mount: ({remote}:{local}) refers to a file.\n To ensure"
|
||||
" this mount updates properly, please use a directory."
|
||||
)
|
||||
|
||||
|
||||
def validate_docker_config(config: Dict[str, Any]) -> None:
|
||||
"""Checks whether the Docker configuration is valid."""
|
||||
if "docker" not in config:
|
||||
return
|
||||
|
||||
_check_docker_file_mounts(config.get("file_mounts", {}))
|
||||
|
||||
docker_image = config["docker"].get("image")
|
||||
cname = config["docker"].get("container_name")
|
||||
|
||||
head_docker_image = config["docker"].get("head_image", docker_image)
|
||||
|
||||
worker_docker_image = config["docker"].get("worker_image", docker_image)
|
||||
|
||||
image_present = docker_image or (head_docker_image and worker_docker_image)
|
||||
if (not cname) and (not image_present):
|
||||
return
|
||||
else:
|
||||
assert cname and image_present, "Must provide a container & image name"
|
||||
|
||||
return None
|
||||
|
||||
|
||||
def with_docker_exec(
|
||||
cmds, container_name, docker_cmd, env_vars=None, with_interactive=False
|
||||
):
|
||||
assert docker_cmd, "Must provide docker command"
|
||||
env_str = ""
|
||||
if env_vars:
|
||||
env_str = " ".join(["-e {env}=${env}".format(env=env) for env in env_vars])
|
||||
return [
|
||||
"{docker_cmd} exec {interactive} {env} {container} /bin/bash -c {cmd} ".format(
|
||||
docker_cmd=docker_cmd,
|
||||
interactive="-it" if with_interactive else "",
|
||||
env=env_str,
|
||||
container=container_name,
|
||||
cmd=quote(cmd),
|
||||
)
|
||||
for cmd in cmds
|
||||
]
|
||||
|
||||
|
||||
def _check_helper(cname, template, docker_cmd):
|
||||
return " ".join(
|
||||
[docker_cmd, "inspect", "-f", "'{{" + template + "}}'", cname, "||", "true"]
|
||||
)
|
||||
|
||||
|
||||
def check_docker_running_cmd(cname, docker_cmd):
|
||||
return _check_helper(cname, ".State.Running", docker_cmd)
|
||||
|
||||
|
||||
def check_bind_mounts_cmd(cname, docker_cmd):
|
||||
return _check_helper(cname, "json .Mounts", docker_cmd)
|
||||
|
||||
|
||||
def check_docker_image(cname, docker_cmd):
|
||||
return _check_helper(cname, ".Config.Image", docker_cmd)
|
||||
|
||||
|
||||
def docker_start_cmds(
|
||||
user,
|
||||
image,
|
||||
mount_dict,
|
||||
container_name,
|
||||
user_options,
|
||||
cluster_name,
|
||||
home_directory,
|
||||
docker_cmd,
|
||||
):
|
||||
# Imported here due to circular dependency.
|
||||
from ray.autoscaler.sdk import get_docker_host_mount_location
|
||||
|
||||
docker_mount_prefix = get_docker_host_mount_location(cluster_name)
|
||||
mount = {f"{docker_mount_prefix}/{dst}": dst for dst in mount_dict}
|
||||
|
||||
mount_flags = " ".join(
|
||||
[
|
||||
"-v {src}:{dest}".format(src=k, dest=v.replace("~/", home_directory + "/"))
|
||||
for k, v in mount.items()
|
||||
]
|
||||
)
|
||||
|
||||
# for click, used in ray cli
|
||||
env_vars = {"LC_ALL": "C.UTF-8", "LANG": "C.UTF-8"}
|
||||
env_flags = " ".join(
|
||||
["-e {name}={val}".format(name=k, val=v) for k, v in env_vars.items()]
|
||||
)
|
||||
|
||||
user_options_str = " ".join(user_options)
|
||||
docker_run = [
|
||||
docker_cmd,
|
||||
"run",
|
||||
"--rm",
|
||||
"--name {}".format(container_name),
|
||||
"-d",
|
||||
"-it",
|
||||
mount_flags,
|
||||
env_flags,
|
||||
user_options_str,
|
||||
"--net=host",
|
||||
image,
|
||||
"bash",
|
||||
]
|
||||
return " ".join(docker_run)
|
||||
@@ -0,0 +1,75 @@
|
||||
import time
|
||||
from threading import RLock
|
||||
from typing import Any, Callable, Dict, List
|
||||
|
||||
|
||||
class EventSummarizer:
|
||||
"""Utility that aggregates related log messages to reduce log spam."""
|
||||
|
||||
def __init__(self):
|
||||
self.events_by_key: Dict[str, int] = {}
|
||||
# Messages to send in next summary batch.
|
||||
self.messages_to_send: List[str] = []
|
||||
# Tracks TTL of messages. A message will not be re-sent once it is
|
||||
# added here, until its TTL expires.
|
||||
self.throttled_messages: Dict[str, float] = {}
|
||||
|
||||
# Event summarizer is used by the main thread and
|
||||
# by node launcher child threads.
|
||||
self.lock = RLock()
|
||||
|
||||
def add(
|
||||
self, template: str, *, quantity: Any, aggregate: Callable[[Any, Any], Any]
|
||||
) -> None:
|
||||
"""Add a log message, which will be combined by template.
|
||||
|
||||
Args:
|
||||
template: Format string with one placeholder for quantity.
|
||||
quantity: Quantity to aggregate.
|
||||
aggregate: Aggregation function used to combine the
|
||||
quantities. The result is inserted into the template to
|
||||
produce the final log message.
|
||||
"""
|
||||
with self.lock:
|
||||
# Enforce proper sentence structure.
|
||||
if not template.endswith("."):
|
||||
template += "."
|
||||
if template in self.events_by_key:
|
||||
self.events_by_key[template] = aggregate(
|
||||
self.events_by_key[template], quantity
|
||||
)
|
||||
else:
|
||||
self.events_by_key[template] = quantity
|
||||
|
||||
def add_once_per_interval(self, message: str, key: str, interval_s: int):
|
||||
"""Add a log message, which is throttled once per interval by a key.
|
||||
|
||||
Args:
|
||||
message: The message to log.
|
||||
key: The key to use to deduplicate the message.
|
||||
interval_s: Throttling interval in seconds.
|
||||
"""
|
||||
with self.lock:
|
||||
if key not in self.throttled_messages:
|
||||
self.throttled_messages[key] = time.time() + interval_s
|
||||
self.messages_to_send.append(message)
|
||||
|
||||
def summary(self) -> List[str]:
|
||||
"""Generate the aggregated log summary of all added events."""
|
||||
with self.lock:
|
||||
out = []
|
||||
for template, quantity in self.events_by_key.items():
|
||||
out.append(template.format(quantity))
|
||||
out.extend(self.messages_to_send)
|
||||
return out
|
||||
|
||||
def clear(self) -> None:
|
||||
"""Clear the events added."""
|
||||
with self.lock:
|
||||
self.events_by_key.clear()
|
||||
self.messages_to_send.clear()
|
||||
# Expire any messages that have reached their TTL. This allows them
|
||||
# to be sent again.
|
||||
for k, t in list(self.throttled_messages.items()):
|
||||
if time.time() > t:
|
||||
del self.throttled_messages[k]
|
||||
@@ -0,0 +1,106 @@
|
||||
from enum import Enum, auto
|
||||
from typing import Any, Callable, Dict, List, Optional, Union
|
||||
|
||||
from ray.autoscaler._private.cli_logger import cli_logger
|
||||
|
||||
|
||||
class CreateClusterEvent(Enum):
|
||||
"""Events to track in ray.autoscaler.sdk.create_or_update_cluster.
|
||||
|
||||
Attributes:
|
||||
up_started : Invoked at the beginning of create_or_update_cluster.
|
||||
ssh_keypair_downloaded : Invoked when the ssh keypair is downloaded.
|
||||
cluster_booting_started : Invoked when when the cluster booting starts.
|
||||
acquiring_new_head_node : Invoked before the head node is acquired.
|
||||
head_node_acquired : Invoked after the head node is acquired.
|
||||
ssh_control_acquired : Invoked when the node is being updated.
|
||||
run_initialization_cmd : Invoked before all initialization
|
||||
commands are called and again before each initialization command.
|
||||
run_setup_cmd : Invoked before all setup commands are
|
||||
called and again before each setup command.
|
||||
start_ray_runtime : Invoked before ray start commands are run.
|
||||
start_ray_runtime_completed : Invoked after ray start commands
|
||||
are run.
|
||||
cluster_booting_completed : Invoked after cluster booting
|
||||
is completed.
|
||||
"""
|
||||
|
||||
up_started = auto()
|
||||
ssh_keypair_downloaded = auto()
|
||||
cluster_booting_started = auto()
|
||||
acquiring_new_head_node = auto()
|
||||
head_node_acquired = auto()
|
||||
ssh_control_acquired = auto()
|
||||
run_initialization_cmd = auto()
|
||||
run_setup_cmd = auto()
|
||||
start_ray_runtime = auto()
|
||||
start_ray_runtime_completed = auto()
|
||||
cluster_booting_completed = auto()
|
||||
|
||||
|
||||
class _EventSystem:
|
||||
"""Event system that handles storing and calling callbacks for events.
|
||||
|
||||
Attributes:
|
||||
callback_map (Dict[str, List[Callable]]) : Stores list of callbacks
|
||||
for events when registered.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
self.callback_map = {}
|
||||
|
||||
def add_callback_handler(
|
||||
self,
|
||||
event: str,
|
||||
callback: Union[Callable[[Dict], None], List[Callable[[Dict], None]]],
|
||||
):
|
||||
"""Stores callback handler for event.
|
||||
|
||||
Args:
|
||||
event: Event that callback should be called on. See
|
||||
CreateClusterEvent for details on the events available to be
|
||||
registered against.
|
||||
callback: Callable object that is invoked
|
||||
when specified event occurs.
|
||||
"""
|
||||
if event not in CreateClusterEvent.__members__.values():
|
||||
cli_logger.warning(
|
||||
f"{event} is not currently tracked, and this"
|
||||
" callback will not be invoked."
|
||||
)
|
||||
|
||||
self.callback_map.setdefault(event, []).extend(
|
||||
[callback] if type(callback) is not list else callback
|
||||
)
|
||||
|
||||
def execute_callback(
|
||||
self, event: CreateClusterEvent, event_data: Optional[Dict[str, Any]] = None
|
||||
):
|
||||
"""Executes all callbacks for event.
|
||||
|
||||
Args:
|
||||
event: Event that is invoked. See CreateClusterEvent
|
||||
for details on the available events.
|
||||
event_data: Argument that is passed to each
|
||||
callable object stored for this particular event.
|
||||
"""
|
||||
if event_data is None:
|
||||
event_data = {}
|
||||
|
||||
event_data["event_name"] = event
|
||||
if event in self.callback_map:
|
||||
for callback in self.callback_map[event]:
|
||||
callback(event_data)
|
||||
|
||||
def clear_callbacks_for_event(self, event: str):
|
||||
"""Clears stored callable objects for event.
|
||||
|
||||
Args:
|
||||
event: Event that has callable objects stored in map.
|
||||
See CreateClusterEvent for details on the available events.
|
||||
"""
|
||||
if event in self.callback_map:
|
||||
del self.callback_map[event]
|
||||
|
||||
|
||||
global_event_system = _EventSystem()
|
||||
@@ -0,0 +1,91 @@
|
||||
import os
|
||||
import subprocess
|
||||
from typing import Dict, List, Tuple
|
||||
|
||||
from ray.autoscaler._private.docker import with_docker_exec
|
||||
from ray.autoscaler.command_runner import CommandRunnerInterface
|
||||
|
||||
|
||||
class FakeDockerCommandRunner(CommandRunnerInterface):
|
||||
"""Command runner for the fke docker multinode cluster.
|
||||
|
||||
This command runner uses ``docker exec`` and ``docker cp`` to
|
||||
run commands and copy files, respectively.
|
||||
|
||||
The regular ``DockerCommandRunner`` is made for use in SSH settings
|
||||
where Docker runs on a remote hose. In contrast, this command runner
|
||||
does not wrap the docker commands in ssh calls.
|
||||
"""
|
||||
|
||||
def __init__(self, docker_config, **common_args):
|
||||
self.container_name = docker_config["container_name"]
|
||||
self.docker_config = docker_config
|
||||
self.home_dir = None
|
||||
self.initialized = False
|
||||
# Optionally use 'podman' instead of 'docker'
|
||||
use_podman = docker_config.get("use_podman", False)
|
||||
self.docker_cmd = "podman" if use_podman else "docker"
|
||||
|
||||
def _run_shell(self, cmd: str, timeout: int = 120) -> str:
|
||||
return subprocess.check_output(
|
||||
cmd, shell=True, timeout=timeout, encoding="utf-8"
|
||||
)
|
||||
|
||||
def run(
|
||||
self,
|
||||
cmd: str = None,
|
||||
timeout: int = 120,
|
||||
exit_on_fail: bool = False,
|
||||
port_forward: List[Tuple[int, int]] = None,
|
||||
with_output: bool = False,
|
||||
environment_variables: Dict[str, object] = None,
|
||||
run_env: str = "auto",
|
||||
ssh_options_override_ssh_key: str = "",
|
||||
shutdown_after_run: bool = False,
|
||||
) -> str:
|
||||
prefix = with_docker_exec(
|
||||
[cmd],
|
||||
container_name=self.container_name,
|
||||
with_interactive=False,
|
||||
docker_cmd=self.docker_cmd,
|
||||
)[0]
|
||||
return self._run_shell(prefix)
|
||||
|
||||
def run_init(
|
||||
self, *, as_head: bool, file_mounts: Dict[str, str], sync_run_yet: bool
|
||||
):
|
||||
pass
|
||||
|
||||
def remote_shell_command_str(self):
|
||||
return "{} exec -it {} bash".format(self.docker_cmd, self.container_name)
|
||||
|
||||
def run_rsync_down(self, source, target, options=None):
|
||||
docker_dir = os.path.dirname(self._docker_expand_user(source))
|
||||
|
||||
self._run_shell(f"docker cp {self.container_name}:{docker_dir} {target}")
|
||||
|
||||
def run_rsync_up(self, source, target, options=None):
|
||||
docker_dir = os.path.dirname(self._docker_expand_user(target))
|
||||
self.run(cmd=f"mkdir -p {docker_dir}")
|
||||
|
||||
self._run_shell(f"docker cp {source} {self.container_name}:{docker_dir}")
|
||||
|
||||
def _docker_expand_user(self, string, any_char=False):
|
||||
user_pos = string.find("~")
|
||||
if user_pos > -1:
|
||||
if self.home_dir is None:
|
||||
self.home_dir = self._run_shell(
|
||||
with_docker_exec(
|
||||
["printenv HOME"],
|
||||
container_name=self.container_name,
|
||||
docker_cmd=self.docker_cmd,
|
||||
)
|
||||
).strip()
|
||||
|
||||
if any_char:
|
||||
return string.replace("~/", self.home_dir + "/")
|
||||
|
||||
elif not any_char and user_pos == 0:
|
||||
return string.replace("~", self.home_dir, 1)
|
||||
|
||||
return string
|
||||
@@ -0,0 +1,246 @@
|
||||
"""Fake multinode docker monitoring script.
|
||||
|
||||
This script is the "docker compose server" for the fake_multinode
|
||||
provider using Docker compose. It should be started before running
|
||||
`RAY_FAKE_CLUSTER=1 ray up <cluster_config>`.
|
||||
|
||||
This script reads the volume directory from a supplied fake multinode
|
||||
docker cluster config file.
|
||||
It then waits until a docker-compose.yaml file is created in the same
|
||||
directory, which is done by the `ray up` command.
|
||||
|
||||
It then watches for changes in the docker-compose.yaml file and runs
|
||||
`docker compose up` whenever changes are detected. This will start docker
|
||||
containers as requested by the autoscaler.
|
||||
|
||||
Generally, the docker-compose.yaml will be mounted in the head node of the
|
||||
cluster, which will then continue to change it according to the autoscaler
|
||||
requirements.
|
||||
|
||||
Additionally, this script monitors the docker container status using
|
||||
`docker status` and writes it into a `status.json`. This information is
|
||||
again used by the autoscaler to determine if any nodes have died.
|
||||
"""
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
import shutil
|
||||
import subprocess
|
||||
import time
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
import yaml
|
||||
|
||||
|
||||
def _read_yaml(path: str):
|
||||
with open(path, "rt") as f:
|
||||
return yaml.safe_load(f)
|
||||
|
||||
|
||||
def _update_docker_compose(
|
||||
docker_compose_path: str, project_name: str, status: Optional[Dict[str, Any]]
|
||||
) -> bool:
|
||||
docker_compose_config = _read_yaml(docker_compose_path)
|
||||
|
||||
if not docker_compose_config:
|
||||
print("Docker compose currently empty")
|
||||
return False
|
||||
|
||||
cmd = ["up", "-d"]
|
||||
if status and len(status) > 0:
|
||||
cmd += ["--no-recreate"]
|
||||
|
||||
shutdown = False
|
||||
if not docker_compose_config["services"]:
|
||||
# If no more nodes, run `down` instead of `up`
|
||||
print("Shutting down nodes")
|
||||
cmd = ["down"]
|
||||
shutdown = True
|
||||
try:
|
||||
subprocess.check_call(
|
||||
["docker", "compose", "-f", docker_compose_path, "-p", project_name]
|
||||
+ cmd
|
||||
+ [
|
||||
"--remove-orphans",
|
||||
]
|
||||
)
|
||||
except Exception as e:
|
||||
print(f"Ran into error when updating docker compose: {e}")
|
||||
# Ignore error
|
||||
|
||||
return shutdown
|
||||
|
||||
|
||||
def _get_ip(
|
||||
project_name: str,
|
||||
container_name: str,
|
||||
override_network: Optional[str] = None,
|
||||
retry_times: int = 3,
|
||||
) -> Optional[str]:
|
||||
network = override_network or f"{project_name}_ray_local"
|
||||
|
||||
cmd = [
|
||||
"docker",
|
||||
"inspect",
|
||||
"-f",
|
||||
'"{{ .NetworkSettings.Networks' f".{network}.IPAddress" ' }}"',
|
||||
f"{container_name}",
|
||||
]
|
||||
for i in range(retry_times):
|
||||
try:
|
||||
ip_address = subprocess.check_output(cmd, encoding="utf-8")
|
||||
except Exception:
|
||||
time.sleep(1)
|
||||
else:
|
||||
return ip_address.strip().strip('"').strip('\\"')
|
||||
return None
|
||||
|
||||
|
||||
def _update_docker_status(
|
||||
docker_compose_path: str, project_name: str, docker_status_path: str
|
||||
):
|
||||
data_str = ""
|
||||
try:
|
||||
data_str = (
|
||||
subprocess.check_output(
|
||||
[
|
||||
"docker",
|
||||
"compose",
|
||||
"-f",
|
||||
docker_compose_path,
|
||||
"-p",
|
||||
project_name,
|
||||
"ps",
|
||||
"--format",
|
||||
"json",
|
||||
]
|
||||
)
|
||||
.decode("utf-8")
|
||||
.strip()
|
||||
.split("\n")
|
||||
)
|
||||
data: List[Dict[str, str]] = []
|
||||
for line in data_str:
|
||||
line = line.strip()
|
||||
if line:
|
||||
data.append(json.loads(line))
|
||||
except Exception as e:
|
||||
print(f"Ran into error when fetching status: {e}")
|
||||
print(f"docker compose ps output: {data_str}")
|
||||
return None
|
||||
|
||||
status = {}
|
||||
for container in data:
|
||||
node_id = container["Service"]
|
||||
container_name = container["Name"]
|
||||
if container["State"] == "running":
|
||||
ip = _get_ip(project_name, container_name)
|
||||
else:
|
||||
ip = ""
|
||||
container["IP"] = ip
|
||||
status[node_id] = container
|
||||
|
||||
with open(docker_status_path, "wt") as f:
|
||||
json.dump(status, f)
|
||||
|
||||
return status
|
||||
|
||||
|
||||
def monitor_docker(
|
||||
docker_compose_path: str,
|
||||
status_path: str,
|
||||
project_name: str,
|
||||
update_interval: float = 1.0,
|
||||
):
|
||||
while not os.path.exists(docker_compose_path):
|
||||
# Wait until cluster is created
|
||||
time.sleep(0.5)
|
||||
|
||||
print("Docker compose config detected, starting status monitoring")
|
||||
|
||||
# Make sure this is always writeable from inside the containers
|
||||
os.chmod(docker_compose_path, 0o777)
|
||||
|
||||
docker_config = {"force_update": True}
|
||||
|
||||
# Force update
|
||||
next_update = time.monotonic() - 1.0
|
||||
shutdown = False
|
||||
status = None
|
||||
|
||||
# Loop:
|
||||
# If the config changed, update cluster.
|
||||
# Every `update_interval` seconds, update docker status.
|
||||
while not shutdown:
|
||||
new_docker_config = _read_yaml(docker_compose_path)
|
||||
if new_docker_config != docker_config:
|
||||
# Update cluster
|
||||
shutdown = _update_docker_compose(docker_compose_path, project_name, status)
|
||||
|
||||
# Force status update
|
||||
next_update = time.monotonic() - 1.0
|
||||
|
||||
if time.monotonic() > next_update:
|
||||
# Update docker status
|
||||
status = _update_docker_status(
|
||||
docker_compose_path, project_name, status_path
|
||||
)
|
||||
next_update = time.monotonic() + update_interval
|
||||
|
||||
docker_config = new_docker_config
|
||||
time.sleep(0.1)
|
||||
|
||||
print("Cluster shut down, terminating monitoring script.")
|
||||
|
||||
|
||||
def start_monitor(config_file: str):
|
||||
cluster_config = _read_yaml(config_file)
|
||||
|
||||
provider_config = cluster_config["provider"]
|
||||
assert provider_config["type"] == "fake_multinode_docker", (
|
||||
f"The docker monitor only works with providers of type "
|
||||
f"`fake_multinode_docker`, got `{provider_config['type']}`"
|
||||
)
|
||||
|
||||
project_name = provider_config["project_name"]
|
||||
|
||||
volume_dir = provider_config["shared_volume_dir"]
|
||||
os.makedirs(volume_dir, mode=0o755, exist_ok=True)
|
||||
|
||||
# Create bootstrap config
|
||||
bootstrap_config_path = os.path.join(volume_dir, "bootstrap_config.yaml")
|
||||
shutil.copy(config_file, bootstrap_config_path)
|
||||
|
||||
# These two files usually don't exist, yet
|
||||
docker_compose_config_path = os.path.join(volume_dir, "docker-compose.yaml")
|
||||
|
||||
docker_status_path = os.path.join(volume_dir, "status.json")
|
||||
|
||||
if os.path.exists(docker_compose_config_path):
|
||||
# We wait until this file exists, so remove it if it exists
|
||||
# from a previous run.
|
||||
os.remove(docker_compose_config_path)
|
||||
|
||||
if os.path.exists(docker_status_path):
|
||||
os.remove(docker_status_path)
|
||||
# Create empty file so it can be mounted
|
||||
with open(docker_status_path, "wt") as f:
|
||||
f.write("{}")
|
||||
|
||||
print(
|
||||
f"Starting monitor process. Please start Ray cluster with:\n"
|
||||
f" RAY_FAKE_CLUSTER=1 ray up {config_file}"
|
||||
)
|
||||
monitor_docker(docker_compose_config_path, docker_status_path, project_name)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument(
|
||||
"config_file",
|
||||
help="Path to cluster config file containing a fake docker "
|
||||
"cluster configuration.",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
start_monitor(args.config_file)
|
||||
@@ -0,0 +1,59 @@
|
||||
# Example command to start a cluster with this config:
|
||||
#
|
||||
# RAY_FAKE_CLUSTER=1 ray start --autoscaling-config=example.yaml --head --block
|
||||
#
|
||||
# Alternatively, you can programmatically create a fake autoscaling cluster
|
||||
# using ray.cluster_utils.AutoscalingCluster.
|
||||
cluster_name: fake_multinode
|
||||
max_workers: 8
|
||||
provider:
|
||||
type: fake_multinode
|
||||
# This must be true since the nodes share the same ip!
|
||||
use_node_id_as_ip: True
|
||||
disable_node_updaters: True
|
||||
disable_launch_config_check: True
|
||||
available_node_types:
|
||||
ray.head.default:
|
||||
# You must set this manually to your "head" node resources!! The head
|
||||
# node is launched via `ray start` and hence the autoscaler cannot
|
||||
# configure its resources. The resources specified for its node type
|
||||
# must line up with what Ray detects/is configured with on start.
|
||||
resources:
|
||||
CPU: 8 # <-- set this to num CPUs used/detected in `ray start`
|
||||
GPU: 0 # <-- set this to num GPUs used/detected in `ray start`
|
||||
node_config: {}
|
||||
max_workers: 0
|
||||
ray.worker.cpu:
|
||||
resources:
|
||||
CPU: 1
|
||||
object_store_memory: 1000000000
|
||||
node_config: {}
|
||||
min_workers: 0
|
||||
max_workers: 4
|
||||
ray.worker.gpu:
|
||||
resources:
|
||||
CPU: 4
|
||||
GPU: 1
|
||||
object_store_memory: 1000000000
|
||||
node_config: {}
|
||||
min_workers: 0
|
||||
max_workers: 2
|
||||
head_node_type: ray.head.default
|
||||
upscaling_speed: 1.0
|
||||
idle_timeout_minutes: 0.1
|
||||
#
|
||||
# !!! Configurations below are not supported in fake cluster mode !!!
|
||||
#
|
||||
auth: {}
|
||||
docker: {}
|
||||
initialization_commands: []
|
||||
setup_commands: []
|
||||
head_setup_commands: []
|
||||
worker_setup_commands: []
|
||||
head_start_ray_commands: []
|
||||
worker_start_ray_commands: []
|
||||
file_mounts: {}
|
||||
cluster_synced_files: []
|
||||
file_mounts_sync_continuously: false
|
||||
rsync_exclude: []
|
||||
rsync_filter: []
|
||||
@@ -0,0 +1,85 @@
|
||||
# This is an example config file to start a local
|
||||
# multi-node cluster using Docker compose.
|
||||
# It requires the ``docker compose`` plugin to be installed:
|
||||
# https://docs.docker.com/compose/cli-command/#installing-compose-v2
|
||||
# The resulting cluster will consist of docker containers
|
||||
# scheduled via docker compose. These containers behave just like
|
||||
# regular Ray nodes, have their own IPs, and can SSH into each other.
|
||||
# They are mostly used to test multi-node setups and autoscaling on
|
||||
# a single node.
|
||||
|
||||
# Example command to start a cluster with this config:
|
||||
#
|
||||
# python docker_monitor.py example_docker.yaml &
|
||||
# RAY_FAKE_DOCKER=1 ray up -y example_docker.yaml
|
||||
cluster_name: fake_docker
|
||||
max_workers: 8
|
||||
provider:
|
||||
type: fake_multinode_docker
|
||||
disable_launch_config_check: True
|
||||
disable_node_updaters: True
|
||||
# Docker-compose config
|
||||
project_name: fake_docker
|
||||
image: rayproject/ray:nightly
|
||||
shared_volume_dir: /tmp/fake_docker
|
||||
|
||||
# For now, this has to be set here separately again:
|
||||
head_resources:
|
||||
CPU: 4
|
||||
GPU: 0
|
||||
|
||||
auth:
|
||||
ssh_user: ubuntu
|
||||
|
||||
available_node_types:
|
||||
ray.head.default:
|
||||
# You must set this manually to your "head" node resources!! The head
|
||||
# node is launched via `ray start` and hence the autoscaler cannot
|
||||
# configure its resources. The resources specified for its node type
|
||||
# must line up with what Ray detects/is configured with on start.
|
||||
resources:
|
||||
CPU: 4
|
||||
GPU: 0
|
||||
node_config: {}
|
||||
max_workers: 0
|
||||
ray.worker.cpu:
|
||||
resources:
|
||||
CPU: 2
|
||||
object_store_memory: 1000000000
|
||||
node_config: {}
|
||||
min_workers: 1
|
||||
max_workers: 4
|
||||
ray.worker.gpu:
|
||||
resources:
|
||||
CPU: 4
|
||||
GPU: 1
|
||||
object_store_memory: 1000000000
|
||||
node_config: {}
|
||||
min_workers: 1
|
||||
max_workers: 2
|
||||
head_node_type: ray.head.default
|
||||
upscaling_speed: 1.0
|
||||
idle_timeout_minutes: 0.1
|
||||
|
||||
# The start commands currently don't work - docker doesn't seem to like docker exec
|
||||
# and Ray only works when including it in the docker-compose command
|
||||
head_start_ray_commands: []
|
||||
worker_start_ray_commands: []
|
||||
|
||||
# The docker config is currently not propagated to the node provider config.
|
||||
# Thus, docker-specific configuration is expected to go into the provider part
|
||||
# as demonstrated above.
|
||||
docker: {}
|
||||
|
||||
#
|
||||
# !!! Configurations below are not supported in fake cluster mode !!!
|
||||
#
|
||||
initialization_commands: []
|
||||
setup_commands: []
|
||||
head_setup_commands: []
|
||||
worker_setup_commands: []
|
||||
file_mounts: {}
|
||||
cluster_synced_files: []
|
||||
file_mounts_sync_continuously: false
|
||||
rsync_exclude: []
|
||||
rsync_filter: []
|
||||
@@ -0,0 +1,730 @@
|
||||
import copy
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import subprocess
|
||||
import sys
|
||||
import time
|
||||
from threading import RLock
|
||||
from types import ModuleType
|
||||
from typing import Any, Dict, Optional
|
||||
|
||||
import yaml
|
||||
|
||||
import ray
|
||||
import ray._private.ray_constants as ray_constants
|
||||
from ray._common.network_utils import build_address
|
||||
from ray.autoscaler._private.fake_multi_node.command_runner import (
|
||||
FakeDockerCommandRunner,
|
||||
)
|
||||
from ray.autoscaler.command_runner import CommandRunnerInterface
|
||||
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__)
|
||||
|
||||
# We generate the node ids deterministically in the fake node provider, so that
|
||||
# we can associate launched nodes with their resource reports. IDs increment
|
||||
# starting with fffff*00000 for the head node, fffff*00001, etc. for workers.
|
||||
FAKE_HEAD_NODE_ID = "fffffffffffffffffffffffffffffffffffffffffffffffffff00000"
|
||||
FAKE_HEAD_NODE_TYPE = "ray.head.default"
|
||||
|
||||
FAKE_DOCKER_DEFAULT_GCS_PORT = 16379
|
||||
FAKE_DOCKER_DEFAULT_OBJECT_MANAGER_PORT = 18076
|
||||
FAKE_DOCKER_DEFAULT_CLIENT_PORT = 10002
|
||||
|
||||
DOCKER_COMPOSE_SKELETON = {
|
||||
"services": {},
|
||||
"networks": {"ray_local": {}},
|
||||
}
|
||||
|
||||
DOCKER_NODE_SKELETON = {
|
||||
"networks": ["ray_local"],
|
||||
"mem_limit": "3000m",
|
||||
"mem_reservation": "3000m",
|
||||
"shm_size": "1200m",
|
||||
"volumes": [],
|
||||
}
|
||||
|
||||
DOCKER_HEAD_CMD = (
|
||||
'bash -c "'
|
||||
"sudo mkdir -p {volume_dir} && "
|
||||
"sudo chmod 777 {volume_dir} && "
|
||||
"touch {volume_dir}/.in_docker && "
|
||||
"sudo chown -R ray:users /cluster/node && "
|
||||
"sudo chmod -R 777 /cluster/node && "
|
||||
"sudo chown -R ray:users /cluster/shared && "
|
||||
"sudo chmod -R 777 /cluster/shared && "
|
||||
"sudo chmod 700 ~/.ssh && "
|
||||
"sudo chmod 600 ~/.ssh/authorized_keys && "
|
||||
"sudo chmod 600 ~/ray_bootstrap_key.pem && "
|
||||
"sudo chown ray:users "
|
||||
"~/.ssh ~/.ssh/authorized_keys ~/ray_bootstrap_key.pem && "
|
||||
"{ensure_ssh} && "
|
||||
"sleep 1 && "
|
||||
"RAY_FAKE_CLUSTER=1 ray start --head "
|
||||
"--autoscaling-config=~/ray_bootstrap_config.yaml "
|
||||
"--object-manager-port=8076 "
|
||||
"--num-cpus {num_cpus} "
|
||||
"--num-gpus {num_gpus} "
|
||||
# "--resources='{resources}' "
|
||||
'--block"'
|
||||
)
|
||||
|
||||
DOCKER_WORKER_CMD = (
|
||||
'bash -c "'
|
||||
"sudo mkdir -p {volume_dir} && "
|
||||
"sudo chmod 777 {volume_dir} && "
|
||||
"touch {volume_dir}/.in_docker && "
|
||||
"sudo chown -R ray:users /cluster/node && "
|
||||
"sudo chmod -R 777 /cluster/node && "
|
||||
"sudo chmod 700 ~/.ssh && "
|
||||
"sudo chmod 600 ~/.ssh/authorized_keys && "
|
||||
"sudo chown ray:users ~/.ssh ~/.ssh/authorized_keys && "
|
||||
"{ensure_ssh} && "
|
||||
"sleep 1 && "
|
||||
f"ray start --address={FAKE_HEAD_NODE_ID}:6379 "
|
||||
"--object-manager-port=8076 "
|
||||
"--num-cpus {num_cpus} "
|
||||
"--num-gpus {num_gpus} "
|
||||
# "--resources='{resources}' "
|
||||
'--block"'
|
||||
)
|
||||
|
||||
|
||||
def host_dir(container_dir: str):
|
||||
"""Replace local dir with potentially different host dir.
|
||||
|
||||
E.g. in docker-in-docker environments, the host dir might be
|
||||
different to the mounted directory in the container.
|
||||
|
||||
This method will do a simple global replace to adjust the paths.
|
||||
"""
|
||||
ray_tempdir = os.environ.get("RAY_TEMPDIR", None)
|
||||
ray_hostdir = os.environ.get("RAY_HOSTDIR", None)
|
||||
|
||||
if not ray_tempdir or not ray_hostdir:
|
||||
return container_dir
|
||||
|
||||
return container_dir.replace(ray_tempdir, ray_hostdir)
|
||||
|
||||
|
||||
def create_node_spec(
|
||||
head: bool,
|
||||
docker_image: str,
|
||||
mounted_cluster_dir: str,
|
||||
mounted_node_dir: str,
|
||||
num_cpus: int = 2,
|
||||
num_gpus: int = 0,
|
||||
resources: Optional[Dict] = None,
|
||||
env_vars: Optional[Dict] = None,
|
||||
host_gcs_port: int = 16379,
|
||||
host_object_manager_port: int = 18076,
|
||||
host_client_port: int = 10002,
|
||||
volume_dir: Optional[str] = None,
|
||||
node_state_path: Optional[str] = None,
|
||||
docker_status_path: Optional[str] = None,
|
||||
docker_compose_path: Optional[str] = None,
|
||||
bootstrap_config_path: Optional[str] = None,
|
||||
private_key_path: Optional[str] = None,
|
||||
public_key_path: Optional[str] = None,
|
||||
):
|
||||
node_spec = copy.deepcopy(DOCKER_NODE_SKELETON)
|
||||
node_spec["image"] = docker_image
|
||||
|
||||
bootstrap_cfg_path_on_container = "/home/ray/ray_bootstrap_config.yaml"
|
||||
bootstrap_key_path_on_container = "/home/ray/ray_bootstrap_key.pem"
|
||||
|
||||
resources = resources or {}
|
||||
|
||||
ensure_ssh = (
|
||||
(
|
||||
"((sudo apt update && sudo apt install -y openssh-server && "
|
||||
"sudo service ssh start) || true)"
|
||||
)
|
||||
if not bool(int(os.environ.get("RAY_HAS_SSH", "0") or "0"))
|
||||
else "sudo service ssh start"
|
||||
)
|
||||
|
||||
cmd_kwargs = dict(
|
||||
ensure_ssh=ensure_ssh,
|
||||
num_cpus=num_cpus,
|
||||
num_gpus=num_gpus,
|
||||
resources=json.dumps(resources, indent=None),
|
||||
volume_dir=volume_dir,
|
||||
autoscaling_config=bootstrap_cfg_path_on_container,
|
||||
)
|
||||
|
||||
env_vars = env_vars or {}
|
||||
|
||||
# Set to "auto" to mount current autoscaler directory to nodes for dev
|
||||
fake_cluster_dev_dir = os.environ.get("FAKE_CLUSTER_DEV", "")
|
||||
if fake_cluster_dev_dir:
|
||||
if fake_cluster_dev_dir == "auto":
|
||||
local_ray_dir = os.path.dirname(ray.__file__)
|
||||
else:
|
||||
local_ray_dir = fake_cluster_dev_dir
|
||||
os.environ["FAKE_CLUSTER_DEV"] = local_ray_dir
|
||||
|
||||
mj = sys.version_info.major
|
||||
mi = sys.version_info.minor
|
||||
|
||||
fake_modules_str = os.environ.get("FAKE_CLUSTER_DEV_MODULES", "autoscaler")
|
||||
fake_modules = fake_modules_str.split(",")
|
||||
|
||||
docker_ray_dir = f"/home/ray/anaconda3/lib/python{mj}.{mi}/site-packages/ray"
|
||||
|
||||
node_spec["volumes"] += [
|
||||
f"{local_ray_dir}/{module}:{docker_ray_dir}/{module}:ro"
|
||||
for module in fake_modules
|
||||
]
|
||||
env_vars["FAKE_CLUSTER_DEV"] = local_ray_dir
|
||||
env_vars["FAKE_CLUSTER_DEV_MODULES"] = fake_modules_str
|
||||
os.environ["FAKE_CLUSTER_DEV_MODULES"] = fake_modules_str
|
||||
|
||||
if head:
|
||||
node_spec["command"] = DOCKER_HEAD_CMD.format(**cmd_kwargs)
|
||||
# Expose ports so we can connect to the cluster from outside
|
||||
node_spec["ports"] = [
|
||||
f"{host_gcs_port}:{ray_constants.DEFAULT_PORT}",
|
||||
f"{host_object_manager_port}:8076",
|
||||
f"{host_client_port}:10001",
|
||||
]
|
||||
# Mount status and config files for the head node
|
||||
node_spec["volumes"] += [
|
||||
f"{host_dir(node_state_path)}:{node_state_path}",
|
||||
f"{host_dir(docker_status_path)}:{docker_status_path}",
|
||||
f"{host_dir(docker_compose_path)}:{docker_compose_path}",
|
||||
f"{host_dir(bootstrap_config_path)}:" f"{bootstrap_cfg_path_on_container}",
|
||||
f"{host_dir(private_key_path)}:{bootstrap_key_path_on_container}",
|
||||
]
|
||||
|
||||
# Create file if it does not exist on local filesystem
|
||||
for filename in [node_state_path, docker_status_path, bootstrap_config_path]:
|
||||
if not os.path.exists(filename):
|
||||
with open(filename, "wt") as f:
|
||||
f.write("{}")
|
||||
else:
|
||||
node_spec["command"] = DOCKER_WORKER_CMD.format(**cmd_kwargs)
|
||||
node_spec["depends_on"] = [FAKE_HEAD_NODE_ID]
|
||||
|
||||
# Mount shared directories and ssh access keys
|
||||
node_spec["volumes"] += [
|
||||
f"{host_dir(mounted_cluster_dir)}:/cluster/shared",
|
||||
f"{host_dir(mounted_node_dir)}:/cluster/node",
|
||||
f"{host_dir(public_key_path)}:/home/ray/.ssh/authorized_keys",
|
||||
]
|
||||
|
||||
# Pass these environment variables (to the head node)
|
||||
# These variables are propagated by the `docker compose` command.
|
||||
env_vars.setdefault("RAY_HAS_SSH", os.environ.get("RAY_HAS_SSH", ""))
|
||||
env_vars.setdefault("RAY_TEMPDIR", os.environ.get("RAY_TEMPDIR", ""))
|
||||
env_vars.setdefault("RAY_HOSTDIR", os.environ.get("RAY_HOSTDIR", ""))
|
||||
|
||||
node_spec["environment"] = [f"{k}={v}" for k, v in env_vars.items()]
|
||||
|
||||
return node_spec
|
||||
|
||||
|
||||
class FakeMultiNodeProvider(NodeProvider):
|
||||
"""A node provider that implements multi-node on a single machine.
|
||||
|
||||
This is used for laptop mode testing of autoscaling functionality."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
provider_config: Dict[str, Any],
|
||||
cluster_name: str,
|
||||
):
|
||||
"""Initialize the fake multi-node provider.
|
||||
|
||||
Args:
|
||||
provider_config: Configuration for the provider.
|
||||
cluster_name: Name of the cluster.
|
||||
"""
|
||||
|
||||
NodeProvider.__init__(self, provider_config, cluster_name)
|
||||
self.lock = RLock()
|
||||
if "RAY_FAKE_CLUSTER" not in os.environ:
|
||||
raise RuntimeError(
|
||||
"FakeMultiNodeProvider requires ray to be started with "
|
||||
"RAY_FAKE_CLUSTER=1 ray start ..."
|
||||
)
|
||||
# GCS address to use for the cluster
|
||||
self._gcs_address = provider_config.get("gcs_address", None)
|
||||
# Head node id
|
||||
self._head_node_id = provider_config.get("head_node_id", FAKE_HEAD_NODE_ID)
|
||||
# Whether to launch multiple nodes at once, or one by one regardless of
|
||||
# the count (default)
|
||||
self._launch_multiple = provider_config.get("launch_multiple", False)
|
||||
|
||||
# These are injected errors for testing purposes. If not None,
|
||||
# these will be raised on `create_node_with_resources_and_labels`` and
|
||||
# `terminate_node``, respectively.
|
||||
self._creation_error = None
|
||||
self._termination_errors = None
|
||||
|
||||
self._nodes = {
|
||||
self._head_node_id: {
|
||||
"tags": {
|
||||
TAG_RAY_NODE_KIND: NODE_KIND_HEAD,
|
||||
TAG_RAY_USER_NODE_TYPE: FAKE_HEAD_NODE_TYPE,
|
||||
TAG_RAY_NODE_NAME: self._head_node_id,
|
||||
TAG_RAY_NODE_STATUS: STATUS_UP_TO_DATE,
|
||||
}
|
||||
},
|
||||
}
|
||||
self._next_node_id = 0
|
||||
|
||||
def _next_hex_node_id(self):
|
||||
self._next_node_id += 1
|
||||
base = "fffffffffffffffffffffffffffffffffffffffffffffffffff"
|
||||
return base + str(self._next_node_id).zfill(5)
|
||||
|
||||
def non_terminated_nodes(self, tag_filters):
|
||||
with self.lock:
|
||||
nodes = []
|
||||
for node_id in self._nodes:
|
||||
tags = self.node_tags(node_id)
|
||||
ok = True
|
||||
for k, v in tag_filters.items():
|
||||
if tags.get(k) != v:
|
||||
ok = False
|
||||
if ok:
|
||||
nodes.append(node_id)
|
||||
|
||||
return nodes
|
||||
|
||||
def is_running(self, node_id):
|
||||
with self.lock:
|
||||
return node_id in self._nodes
|
||||
|
||||
def is_terminated(self, node_id):
|
||||
with self.lock:
|
||||
return node_id not in self._nodes
|
||||
|
||||
def node_tags(self, node_id):
|
||||
with self.lock:
|
||||
return self._nodes[node_id]["tags"]
|
||||
|
||||
def _get_ip(self, node_id: str) -> Optional[str]:
|
||||
return node_id
|
||||
|
||||
def external_ip(self, node_id):
|
||||
return self._get_ip(node_id)
|
||||
|
||||
def internal_ip(self, node_id):
|
||||
return self._get_ip(node_id)
|
||||
|
||||
def set_node_tags(self, node_id, tags):
|
||||
raise AssertionError("Readonly node provider cannot be updated")
|
||||
|
||||
def create_node_with_resources_and_labels(
|
||||
self, node_config, tags, count, resources, labels
|
||||
):
|
||||
if self._creation_error:
|
||||
raise self._creation_error
|
||||
|
||||
if self._launch_multiple:
|
||||
for _ in range(count):
|
||||
self._create_node_with_resources_and_labels(
|
||||
node_config, tags, count, resources, labels
|
||||
)
|
||||
else:
|
||||
self._create_node_with_resources_and_labels(
|
||||
node_config, tags, count, resources, labels
|
||||
)
|
||||
|
||||
def _create_node_with_resources_and_labels(
|
||||
self, node_config, tags, count, resources, labels
|
||||
):
|
||||
# This function calls `pop`. To avoid side effects, we make a
|
||||
# copy of `resources`.
|
||||
resources = copy.deepcopy(resources)
|
||||
with self.lock:
|
||||
node_type = tags[TAG_RAY_USER_NODE_TYPE]
|
||||
next_id = self._next_hex_node_id()
|
||||
ray_params = ray._private.parameter.RayParams(
|
||||
min_worker_port=0,
|
||||
max_worker_port=0,
|
||||
dashboard_port=None,
|
||||
dashboard_agent_listen_port=0,
|
||||
num_cpus=resources.pop("CPU", 0),
|
||||
num_gpus=resources.pop("GPU", 0),
|
||||
object_store_memory=resources.pop("object_store_memory", None),
|
||||
resources=resources,
|
||||
labels=labels,
|
||||
redis_address=build_address(
|
||||
ray._private.services.get_node_ip_address(), 6379
|
||||
)
|
||||
if not self._gcs_address
|
||||
else self._gcs_address,
|
||||
gcs_address=build_address(
|
||||
ray._private.services.get_node_ip_address(), 6379
|
||||
)
|
||||
if not self._gcs_address
|
||||
else self._gcs_address,
|
||||
env_vars={
|
||||
"RAY_OVERRIDE_NODE_ID_FOR_TESTING": next_id,
|
||||
"RAY_CLOUD_INSTANCE_ID": next_id,
|
||||
"RAY_NODE_TYPE_NAME": node_type,
|
||||
ray_constants.RESOURCES_ENVIRONMENT_VARIABLE: json.dumps(resources),
|
||||
ray_constants.LABELS_ENVIRONMENT_VARIABLE: json.dumps(labels),
|
||||
},
|
||||
)
|
||||
node = ray._private.node.Node(
|
||||
ray_params, head=False, shutdown_at_exit=False, spawn_reaper=False
|
||||
)
|
||||
all_tags = {
|
||||
TAG_RAY_NODE_KIND: NODE_KIND_WORKER,
|
||||
TAG_RAY_USER_NODE_TYPE: node_type,
|
||||
TAG_RAY_NODE_NAME: next_id,
|
||||
TAG_RAY_NODE_STATUS: STATUS_UP_TO_DATE,
|
||||
}
|
||||
all_tags.update(tags)
|
||||
self._nodes[next_id] = {
|
||||
"tags": all_tags,
|
||||
"node": node,
|
||||
}
|
||||
|
||||
def terminate_node(self, node_id):
|
||||
with self.lock:
|
||||
if self._termination_errors:
|
||||
raise self._termination_errors
|
||||
|
||||
try:
|
||||
node = self._nodes.pop(node_id)
|
||||
except Exception as e:
|
||||
raise e
|
||||
|
||||
self._terminate_node(node)
|
||||
|
||||
def _terminate_node(self, node):
|
||||
node["node"].kill_all_processes(check_alive=False, allow_graceful=True)
|
||||
|
||||
@staticmethod
|
||||
def bootstrap_config(cluster_config):
|
||||
return cluster_config
|
||||
|
||||
############################
|
||||
# Test only methods
|
||||
############################
|
||||
def _test_set_creation_error(self, e: Exception):
|
||||
"""Set an error that will be raised on
|
||||
create_node_with_resources_and_labels."""
|
||||
self._creation_error = e
|
||||
|
||||
def _test_add_termination_errors(self, e: Exception):
|
||||
"""Set an error that will be raised on terminate_node."""
|
||||
self._termination_errors = e
|
||||
|
||||
|
||||
class FakeMultiNodeDockerProvider(FakeMultiNodeProvider):
|
||||
"""A node provider that implements multi-node on a single machine.
|
||||
|
||||
This is used for laptop mode testing of multi node functionality
|
||||
where each node has their own FS and IP."""
|
||||
|
||||
def __init__(self, provider_config, cluster_name):
|
||||
super(FakeMultiNodeDockerProvider, self).__init__(provider_config, cluster_name)
|
||||
|
||||
fake_head = copy.deepcopy(self._nodes)
|
||||
|
||||
self._project_name = self.provider_config["project_name"]
|
||||
self._docker_image = self.provider_config["image"]
|
||||
|
||||
self._host_gcs_port = self.provider_config.get(
|
||||
"host_gcs_port", FAKE_DOCKER_DEFAULT_GCS_PORT
|
||||
)
|
||||
self._host_object_manager_port = self.provider_config.get(
|
||||
"host_object_manager_port", FAKE_DOCKER_DEFAULT_OBJECT_MANAGER_PORT
|
||||
)
|
||||
self._host_client_port = self.provider_config.get(
|
||||
"host_client_port", FAKE_DOCKER_DEFAULT_CLIENT_PORT
|
||||
)
|
||||
|
||||
self._head_resources = self.provider_config["head_resources"]
|
||||
|
||||
# subdirs:
|
||||
# - ./shared (shared filesystem)
|
||||
# - ./nodes/<node_id> (node-specific mounted filesystem)
|
||||
self._volume_dir = self.provider_config["shared_volume_dir"]
|
||||
self._mounted_cluster_dir = os.path.join(self._volume_dir, "shared")
|
||||
|
||||
if not self.in_docker_container:
|
||||
# Only needed on host
|
||||
os.makedirs(self._mounted_cluster_dir, mode=0o755, exist_ok=True)
|
||||
|
||||
self._boostrap_config_path = os.path.join(
|
||||
self._volume_dir, "bootstrap_config.yaml"
|
||||
)
|
||||
|
||||
self._private_key_path = os.path.join(self._volume_dir, "bootstrap_key.pem")
|
||||
self._public_key_path = os.path.join(self._volume_dir, "bootstrap_key.pem.pub")
|
||||
|
||||
if not self.in_docker_container:
|
||||
# Create private key
|
||||
if not os.path.exists(self._private_key_path):
|
||||
subprocess.check_call(
|
||||
f'ssh-keygen -b 2048 -t rsa -q -N "" '
|
||||
f"-f {self._private_key_path}",
|
||||
shell=True,
|
||||
)
|
||||
|
||||
# Create public key
|
||||
if not os.path.exists(self._public_key_path):
|
||||
subprocess.check_call(
|
||||
f"ssh-keygen -y "
|
||||
f"-f {self._private_key_path} "
|
||||
f"> {self._public_key_path}",
|
||||
shell=True,
|
||||
)
|
||||
|
||||
self._docker_compose_config_path = os.path.join(
|
||||
self._volume_dir, "docker-compose.yaml"
|
||||
)
|
||||
self._docker_compose_config = None
|
||||
|
||||
self._node_state_path = os.path.join(self._volume_dir, "nodes.json")
|
||||
self._docker_status_path = os.path.join(self._volume_dir, "status.json")
|
||||
|
||||
self._load_node_state()
|
||||
if FAKE_HEAD_NODE_ID not in self._nodes:
|
||||
# Reset
|
||||
self._nodes = copy.deepcopy(fake_head)
|
||||
|
||||
self._nodes[FAKE_HEAD_NODE_ID][
|
||||
"node_spec"
|
||||
] = self._create_node_spec_with_resources(
|
||||
head=True, node_id=FAKE_HEAD_NODE_ID, resources=self._head_resources
|
||||
)
|
||||
self._possibly_terminated_nodes = dict()
|
||||
|
||||
self._cleanup_interval = provider_config.get("cleanup_interval", 9.5)
|
||||
|
||||
self._docker_status = {}
|
||||
|
||||
self._update_docker_compose_config()
|
||||
self._update_docker_status()
|
||||
self._save_node_state()
|
||||
|
||||
@property
|
||||
def in_docker_container(self):
|
||||
return os.path.exists(os.path.join(self._volume_dir, ".in_docker"))
|
||||
|
||||
def _create_node_spec_with_resources(
|
||||
self, head: bool, node_id: str, resources: Dict[str, Any]
|
||||
):
|
||||
resources = resources.copy()
|
||||
|
||||
# Create shared directory
|
||||
node_dir = os.path.join(self._volume_dir, "nodes", node_id)
|
||||
os.makedirs(node_dir, mode=0o777, exist_ok=True)
|
||||
|
||||
resource_str = json.dumps(resources, indent=None)
|
||||
|
||||
return create_node_spec(
|
||||
head=head,
|
||||
docker_image=self._docker_image,
|
||||
mounted_cluster_dir=self._mounted_cluster_dir,
|
||||
mounted_node_dir=node_dir,
|
||||
num_cpus=resources.pop("CPU", 0),
|
||||
num_gpus=resources.pop("GPU", 0),
|
||||
host_gcs_port=self._host_gcs_port,
|
||||
host_object_manager_port=self._host_object_manager_port,
|
||||
host_client_port=self._host_client_port,
|
||||
resources=resources,
|
||||
env_vars={
|
||||
"RAY_OVERRIDE_NODE_ID_FOR_TESTING": node_id,
|
||||
ray_constants.RESOURCES_ENVIRONMENT_VARIABLE: resource_str,
|
||||
**self.provider_config.get("env_vars", {}),
|
||||
},
|
||||
volume_dir=self._volume_dir,
|
||||
node_state_path=self._node_state_path,
|
||||
docker_status_path=self._docker_status_path,
|
||||
docker_compose_path=self._docker_compose_config_path,
|
||||
bootstrap_config_path=self._boostrap_config_path,
|
||||
public_key_path=self._public_key_path,
|
||||
private_key_path=self._private_key_path,
|
||||
)
|
||||
|
||||
def _load_node_state(self) -> bool:
|
||||
if not os.path.exists(self._node_state_path):
|
||||
return False
|
||||
try:
|
||||
with open(self._node_state_path, "rt") as f:
|
||||
nodes = json.load(f)
|
||||
except Exception:
|
||||
return False
|
||||
if not nodes:
|
||||
return False
|
||||
self._nodes = nodes
|
||||
return True
|
||||
|
||||
def _save_node_state(self):
|
||||
with open(self._node_state_path, "wt") as f:
|
||||
json.dump(self._nodes, f)
|
||||
|
||||
# Make sure this is always writeable from inside the containers
|
||||
if not self.in_docker_container:
|
||||
# Only chmod from the outer container
|
||||
os.chmod(self._node_state_path, 0o777)
|
||||
|
||||
def _update_docker_compose_config(self):
|
||||
config = copy.deepcopy(DOCKER_COMPOSE_SKELETON)
|
||||
config["services"] = {}
|
||||
for node_id, node in self._nodes.items():
|
||||
config["services"][node_id] = node["node_spec"]
|
||||
|
||||
with open(self._docker_compose_config_path, "wt") as f:
|
||||
yaml.safe_dump(config, f)
|
||||
|
||||
def _update_docker_status(self):
|
||||
if not os.path.exists(self._docker_status_path):
|
||||
return
|
||||
with open(self._docker_status_path, "rt") as f:
|
||||
self._docker_status = json.load(f)
|
||||
|
||||
def _update_nodes(self):
|
||||
for node_id in list(self._nodes):
|
||||
if not self._is_docker_running(node_id):
|
||||
self._possibly_terminated_nodes.setdefault(node_id, time.monotonic())
|
||||
else:
|
||||
self._possibly_terminated_nodes.pop(node_id, None)
|
||||
self._cleanup_nodes()
|
||||
|
||||
def _cleanup_nodes(self):
|
||||
for node_id, timestamp in list(self._possibly_terminated_nodes.items()):
|
||||
if time.monotonic() > timestamp + self._cleanup_interval:
|
||||
if not self._is_docker_running(node_id):
|
||||
self._nodes.pop(node_id, None)
|
||||
self._possibly_terminated_nodes.pop(node_id, None)
|
||||
self._save_node_state()
|
||||
|
||||
def _container_name(self, node_id):
|
||||
node_status = self._docker_status.get(node_id, {})
|
||||
timeout = time.monotonic() + 60
|
||||
while not node_status:
|
||||
if time.monotonic() > timeout:
|
||||
raise RuntimeError(f"Container for {node_id} never became available.")
|
||||
time.sleep(1)
|
||||
self._update_docker_status()
|
||||
node_status = self._docker_status.get(node_id, {})
|
||||
|
||||
return node_status["Name"]
|
||||
|
||||
def _is_docker_running(self, node_id):
|
||||
self._update_docker_status()
|
||||
|
||||
return self._docker_status.get(node_id, {}).get("State", None) == "running"
|
||||
|
||||
def non_terminated_nodes(self, tag_filters):
|
||||
self._update_nodes()
|
||||
return super(FakeMultiNodeDockerProvider, self).non_terminated_nodes(
|
||||
tag_filters
|
||||
)
|
||||
|
||||
def is_running(self, node_id):
|
||||
with self.lock:
|
||||
self._update_nodes()
|
||||
|
||||
return node_id in self._nodes and self._is_docker_running(node_id)
|
||||
|
||||
def is_terminated(self, node_id):
|
||||
with self.lock:
|
||||
self._update_nodes()
|
||||
|
||||
return node_id not in self._nodes and not self._is_docker_running(node_id)
|
||||
|
||||
def get_command_runner(
|
||||
self,
|
||||
log_prefix: str,
|
||||
node_id: str,
|
||||
auth_config: Dict[str, Any],
|
||||
cluster_name: str,
|
||||
process_runner: ModuleType,
|
||||
use_internal_ip: bool,
|
||||
docker_config: Optional[Dict[str, Any]] = None,
|
||||
) -> CommandRunnerInterface:
|
||||
if self.in_docker_container:
|
||||
return super(FakeMultiNodeProvider, self).get_command_runner(
|
||||
log_prefix,
|
||||
node_id,
|
||||
auth_config,
|
||||
cluster_name,
|
||||
process_runner,
|
||||
use_internal_ip,
|
||||
)
|
||||
|
||||
# Else, host command runner:
|
||||
common_args = {
|
||||
"log_prefix": log_prefix,
|
||||
"node_id": node_id,
|
||||
"provider": self,
|
||||
"auth_config": auth_config,
|
||||
"cluster_name": cluster_name,
|
||||
"process_runner": process_runner,
|
||||
"use_internal_ip": use_internal_ip,
|
||||
}
|
||||
|
||||
docker_config["container_name"] = self._container_name(node_id)
|
||||
docker_config["image"] = self._docker_image
|
||||
|
||||
return FakeDockerCommandRunner(docker_config, **common_args)
|
||||
|
||||
def _get_ip(self, node_id: str) -> Optional[str]:
|
||||
for i in range(3):
|
||||
self._update_docker_status()
|
||||
ip = self._docker_status.get(node_id, {}).get("IP", None)
|
||||
if ip:
|
||||
return ip
|
||||
time.sleep(3)
|
||||
return None
|
||||
|
||||
def set_node_tags(self, node_id, tags):
|
||||
assert node_id in self._nodes
|
||||
self._nodes[node_id]["tags"].update(tags)
|
||||
|
||||
def create_node_with_resources_and_labels(
|
||||
self, node_config, tags, count, resources, labels
|
||||
):
|
||||
with self.lock:
|
||||
is_head = tags[TAG_RAY_NODE_KIND] == NODE_KIND_HEAD
|
||||
|
||||
if is_head:
|
||||
next_id = FAKE_HEAD_NODE_ID
|
||||
else:
|
||||
next_id = self._next_hex_node_id()
|
||||
|
||||
self._nodes[next_id] = {
|
||||
"tags": tags,
|
||||
"node_spec": self._create_node_spec_with_resources(
|
||||
head=is_head, node_id=next_id, resources=resources
|
||||
),
|
||||
}
|
||||
self._update_docker_compose_config()
|
||||
self._save_node_state()
|
||||
|
||||
def create_node(
|
||||
self, node_config: Dict[str, Any], tags: Dict[str, str], count: int
|
||||
) -> Optional[Dict[str, Any]]:
|
||||
resources = self._head_resources
|
||||
return self.create_node_with_resources_and_labels(
|
||||
node_config, tags, count, resources, {}
|
||||
)
|
||||
|
||||
def _terminate_node(self, node):
|
||||
self._update_docker_compose_config()
|
||||
self._save_node_state()
|
||||
|
||||
@staticmethod
|
||||
def bootstrap_config(cluster_config):
|
||||
return cluster_config
|
||||
@@ -0,0 +1,399 @@
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import random
|
||||
import shutil
|
||||
import subprocess
|
||||
import sys
|
||||
import tempfile
|
||||
import threading
|
||||
import time
|
||||
from typing import Any, Dict, Optional
|
||||
|
||||
import yaml
|
||||
|
||||
import ray
|
||||
from ray._common.network_utils import build_address
|
||||
from ray._private.dict import deep_update
|
||||
from ray.autoscaler._private.fake_multi_node.node_provider import (
|
||||
FAKE_DOCKER_DEFAULT_CLIENT_PORT,
|
||||
FAKE_DOCKER_DEFAULT_GCS_PORT,
|
||||
)
|
||||
from ray.util.queue import Empty, Queue
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
DEFAULT_DOCKER_IMAGE = "rayproject/ray:nightly-py{major}{minor}-cpu"
|
||||
|
||||
|
||||
class ResourcesNotReadyError(RuntimeError):
|
||||
pass
|
||||
|
||||
|
||||
class DockerCluster:
|
||||
"""Docker cluster wrapper.
|
||||
|
||||
Creates a directory for starting a fake multinode docker cluster.
|
||||
|
||||
Includes APIs to update the cluster config as needed in tests,
|
||||
and to start and connect to the cluster.
|
||||
"""
|
||||
|
||||
def __init__(self, config: Optional[Dict[str, Any]] = None):
|
||||
self._base_config_file = os.path.join(
|
||||
os.path.dirname(__file__), "example_docker.yaml"
|
||||
)
|
||||
self._tempdir = None
|
||||
self._config_file = None
|
||||
self._nodes_file = None
|
||||
self._nodes = {}
|
||||
self._status_file = None
|
||||
self._status = {}
|
||||
self._partial_config = config
|
||||
self._cluster_config = None
|
||||
self._docker_image = None
|
||||
|
||||
self._monitor_script = os.path.join(
|
||||
os.path.dirname(__file__), "docker_monitor.py"
|
||||
)
|
||||
self._monitor_process = None
|
||||
|
||||
self._execution_thread = None
|
||||
self._execution_event = threading.Event()
|
||||
self._execution_queue = None
|
||||
|
||||
@property
|
||||
def config_file(self):
|
||||
return self._config_file
|
||||
|
||||
@property
|
||||
def cluster_config(self):
|
||||
return self._cluster_config
|
||||
|
||||
@property
|
||||
def cluster_dir(self):
|
||||
return self._tempdir
|
||||
|
||||
@property
|
||||
def gcs_port(self):
|
||||
return self._cluster_config.get("provider", {}).get(
|
||||
"host_gcs_port", FAKE_DOCKER_DEFAULT_GCS_PORT
|
||||
)
|
||||
|
||||
@property
|
||||
def client_port(self):
|
||||
return self._cluster_config.get("provider", {}).get(
|
||||
"host_client_port", FAKE_DOCKER_DEFAULT_CLIENT_PORT
|
||||
)
|
||||
|
||||
def connect(self, client: bool = True, timeout: int = 120, **init_kwargs):
|
||||
"""Connect to the docker-compose Ray cluster.
|
||||
|
||||
Assumes the cluster is at RAY_TESTHOST (defaults to
|
||||
``127.0.0.1``).
|
||||
|
||||
Args:
|
||||
client: If True, uses Ray client to connect to the
|
||||
cluster. If False, uses GCS to connect to the cluster.
|
||||
timeout: Connection timeout in seconds.
|
||||
**init_kwargs: kwargs to pass to ``ray.init()``.
|
||||
|
||||
"""
|
||||
host = os.environ.get("RAY_TESTHOST", "127.0.0.1")
|
||||
|
||||
if client:
|
||||
port = self.client_port
|
||||
address = f"ray://{build_address(host, port)}"
|
||||
else:
|
||||
port = self.gcs_port
|
||||
address = build_address(host, port)
|
||||
|
||||
timeout_at = time.monotonic() + timeout
|
||||
while time.monotonic() < timeout_at:
|
||||
try:
|
||||
ray.init(address, **init_kwargs)
|
||||
self.wait_for_resources({"CPU": 1})
|
||||
except ResourcesNotReadyError:
|
||||
time.sleep(1)
|
||||
continue
|
||||
else:
|
||||
break
|
||||
|
||||
try:
|
||||
ray.cluster_resources()
|
||||
except Exception as e:
|
||||
raise RuntimeError(f"Timed out connecting to Ray: {e}")
|
||||
|
||||
def remote_execution_api(self) -> "RemoteAPI":
|
||||
"""Create an object to control cluster state from within the cluster."""
|
||||
self._execution_queue = Queue(actor_options={"num_cpus": 0})
|
||||
stop_event = self._execution_event
|
||||
|
||||
def entrypoint():
|
||||
while not stop_event.is_set():
|
||||
try:
|
||||
cmd, kwargs = self._execution_queue.get(timeout=1)
|
||||
except Empty:
|
||||
continue
|
||||
|
||||
if cmd == "kill_node":
|
||||
self.kill_node(**kwargs)
|
||||
|
||||
self._execution_thread = threading.Thread(target=entrypoint)
|
||||
self._execution_thread.start()
|
||||
|
||||
return RemoteAPI(self._execution_queue)
|
||||
|
||||
@staticmethod
|
||||
def wait_for_resources(resources: Dict[str, float], timeout: int = 60):
|
||||
"""Wait until Ray cluster resources are available
|
||||
|
||||
Args:
|
||||
resources: Minimum resources needed before
|
||||
this function returns.
|
||||
timeout: Timeout in seconds.
|
||||
|
||||
"""
|
||||
timeout = time.monotonic() + timeout
|
||||
|
||||
available = ray.cluster_resources()
|
||||
while any(available.get(k, 0.0) < v for k, v in resources.items()):
|
||||
if time.monotonic() > timeout:
|
||||
raise ResourcesNotReadyError(
|
||||
f"Timed out waiting for resources: {resources}"
|
||||
)
|
||||
time.sleep(1)
|
||||
available = ray.cluster_resources()
|
||||
|
||||
def update_config(self, config: Optional[Dict[str, Any]] = None):
|
||||
"""Update autoscaling config.
|
||||
|
||||
Does a deep update of the base config with a new configuration.
|
||||
This can change autoscaling behavior.
|
||||
|
||||
Args:
|
||||
config: Partial config to update current
|
||||
config with.
|
||||
|
||||
"""
|
||||
assert self._tempdir, "Call setup() first"
|
||||
|
||||
config = config or {}
|
||||
|
||||
if config:
|
||||
self._partial_config = config
|
||||
|
||||
if not config.get("provider", {}).get("image"):
|
||||
# No image specified, trying to parse from buildkite
|
||||
docker_image = os.environ.get("RAY_DOCKER_IMAGE", None)
|
||||
|
||||
if not docker_image:
|
||||
# If still no docker image, use one according to Python version
|
||||
mj = sys.version_info.major
|
||||
mi = sys.version_info.minor
|
||||
|
||||
docker_image = DEFAULT_DOCKER_IMAGE.format(major=mj, minor=mi)
|
||||
|
||||
self._docker_image = docker_image
|
||||
|
||||
with open(self._base_config_file, "rt") as f:
|
||||
cluster_config = yaml.safe_load(f)
|
||||
|
||||
if self._partial_config:
|
||||
deep_update(cluster_config, self._partial_config, new_keys_allowed=True)
|
||||
|
||||
if self._docker_image:
|
||||
cluster_config["provider"]["image"] = self._docker_image
|
||||
|
||||
cluster_config["provider"]["shared_volume_dir"] = self._tempdir
|
||||
|
||||
self._cluster_config = cluster_config
|
||||
|
||||
with open(self._config_file, "wt") as f:
|
||||
yaml.safe_dump(self._cluster_config, f)
|
||||
|
||||
logging.info(f"Updated cluster config to: {self._cluster_config}")
|
||||
|
||||
def maybe_pull_image(self):
|
||||
if self._docker_image:
|
||||
try:
|
||||
images_str = subprocess.check_output(
|
||||
f"docker image inspect {self._docker_image}", shell=True
|
||||
)
|
||||
images = json.loads(images_str)
|
||||
except Exception as e:
|
||||
logger.error(f"Error inspecting image {self._docker_image}: {e}")
|
||||
return
|
||||
|
||||
if not images:
|
||||
try:
|
||||
subprocess.check_call(
|
||||
f"docker pull {self._docker_image}", shell=True
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(f"Error pulling image {self._docker_image}: {e}")
|
||||
|
||||
def setup(self):
|
||||
"""Setup docker compose cluster environment.
|
||||
|
||||
Creates the temporary directory, writes the initial config file,
|
||||
and pulls the docker image, if required.
|
||||
"""
|
||||
self._tempdir = tempfile.mkdtemp(dir=os.environ.get("RAY_TEMPDIR", None))
|
||||
os.chmod(self._tempdir, 0o777)
|
||||
self._config_file = os.path.join(self._tempdir, "cluster.yaml")
|
||||
self._nodes_file = os.path.join(self._tempdir, "nodes.json")
|
||||
self._status_file = os.path.join(self._tempdir, "status.json")
|
||||
self.update_config()
|
||||
self.maybe_pull_image()
|
||||
|
||||
def teardown(self, keep_dir: bool = False):
|
||||
"""Tear down docker compose cluster environment.
|
||||
|
||||
Args:
|
||||
keep_dir: If True, cluster directory
|
||||
will not be removed after termination.
|
||||
"""
|
||||
if not keep_dir:
|
||||
shutil.rmtree(self._tempdir)
|
||||
self._tempdir = None
|
||||
self._config_file = None
|
||||
|
||||
def _start_monitor(self):
|
||||
self._monitor_process = subprocess.Popen(
|
||||
[sys.executable, self._monitor_script, self.config_file]
|
||||
)
|
||||
time.sleep(2)
|
||||
|
||||
def _stop_monitor(self):
|
||||
if self._monitor_process:
|
||||
self._monitor_process.wait(timeout=30)
|
||||
if self._monitor_process.poll() is None:
|
||||
self._monitor_process.terminate()
|
||||
self._monitor_process = None
|
||||
|
||||
def start(self):
|
||||
"""Start docker compose cluster.
|
||||
|
||||
Starts the monitor process and runs ``ray up``.
|
||||
"""
|
||||
self._start_monitor()
|
||||
|
||||
subprocess.check_call(
|
||||
f"RAY_FAKE_CLUSTER=1 ray up -y {self.config_file}", shell=True
|
||||
)
|
||||
|
||||
def stop(self):
|
||||
"""Stop docker compose cluster.
|
||||
|
||||
Runs ``ray down`` and stops the monitor process.
|
||||
"""
|
||||
if ray.is_initialized:
|
||||
ray.shutdown()
|
||||
|
||||
subprocess.check_call(
|
||||
f"RAY_FAKE_CLUSTER=1 ray down -y {self.config_file}", shell=True
|
||||
)
|
||||
|
||||
self._stop_monitor()
|
||||
self._execution_event.set()
|
||||
|
||||
def _update_nodes(self):
|
||||
with open(self._nodes_file, "rt") as f:
|
||||
self._nodes = json.load(f)
|
||||
|
||||
def _update_status(self):
|
||||
with open(self._status_file, "rt") as f:
|
||||
self._status = json.load(f)
|
||||
|
||||
def _get_node(
|
||||
self,
|
||||
node_id: Optional[str] = None,
|
||||
num: Optional[int] = None,
|
||||
rand: Optional[str] = None,
|
||||
) -> str:
|
||||
self._update_nodes()
|
||||
if node_id:
|
||||
assert (
|
||||
not num and not rand
|
||||
), "Only provide either `node_id`, `num`, or `random`."
|
||||
elif num:
|
||||
assert (
|
||||
not node_id and not rand
|
||||
), "Only provide either `node_id`, `num`, or `random`."
|
||||
base = "fffffffffffffffffffffffffffffffffffffffffffffffffff"
|
||||
node_id = base + str(num).zfill(5)
|
||||
elif rand:
|
||||
assert (
|
||||
not node_id and not num
|
||||
), "Only provide either `node_id`, `num`, or `random`."
|
||||
assert rand in [
|
||||
"worker",
|
||||
"any",
|
||||
], "`random` must be one of ['worker', 'any']"
|
||||
choices = list(self._nodes.keys())
|
||||
if rand == "worker":
|
||||
choices.remove(
|
||||
"fffffffffffffffffffffffffffffffffffffffffffffffffff00000"
|
||||
)
|
||||
# Else: any
|
||||
node_id = random.choice(choices)
|
||||
|
||||
assert node_id in self._nodes, f"Node with ID {node_id} is not in active nodes."
|
||||
return node_id
|
||||
|
||||
def _get_docker_container(self, node_id: str) -> Optional[str]:
|
||||
self._update_status()
|
||||
node_status = self._status.get(node_id)
|
||||
if not node_status:
|
||||
return None
|
||||
|
||||
return node_status["Name"]
|
||||
|
||||
def kill_node(
|
||||
self,
|
||||
node_id: Optional[str] = None,
|
||||
num: Optional[int] = None,
|
||||
rand: Optional[str] = None,
|
||||
):
|
||||
"""Kill node.
|
||||
|
||||
If ``node_id`` is given, kill that node.
|
||||
|
||||
If ``num`` is given, construct node_id from this number, and kill
|
||||
that node.
|
||||
|
||||
If ``rand`` is given (as either ``worker`` or ``any``), kill a random
|
||||
node.
|
||||
"""
|
||||
node_id = self._get_node(node_id=node_id, num=num, rand=rand)
|
||||
container = self._get_docker_container(node_id=node_id)
|
||||
subprocess.check_call(f"docker kill {container}", shell=True)
|
||||
|
||||
|
||||
class RemoteAPI:
|
||||
"""Remote API to control cluster state from within cluster tasks.
|
||||
|
||||
This API uses a Ray queue to interact with an execution thread on the
|
||||
host machine that will execute commands passed to the queue.
|
||||
|
||||
Instances of this class can be serialized and passed to Ray remote actors
|
||||
to interact with cluster state (but they can also be used outside actors).
|
||||
|
||||
The API subset is limited to specific commands.
|
||||
|
||||
Args:
|
||||
queue: Ray queue to push command instructions to.
|
||||
|
||||
"""
|
||||
|
||||
def __init__(self, queue: Queue):
|
||||
self._queue = queue
|
||||
|
||||
def kill_node(
|
||||
self,
|
||||
node_id: Optional[str] = None,
|
||||
num: Optional[int] = None,
|
||||
rand: Optional[str] = None,
|
||||
):
|
||||
self._queue.put(("kill_node", dict(node_id=node_id, num=num, rand=rand)))
|
||||
@@ -0,0 +1,874 @@
|
||||
import copy
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import re
|
||||
import time
|
||||
from functools import partial, reduce
|
||||
|
||||
import google_auth_httplib2
|
||||
import googleapiclient
|
||||
import httplib2
|
||||
from google.oauth2 import service_account
|
||||
from google.oauth2.credentials import Credentials as OAuthCredentials
|
||||
from googleapiclient import discovery, errors
|
||||
|
||||
from ray._private.accelerators import TPUAcceleratorManager, tpu
|
||||
from ray.autoscaler._private.gcp.node import MAX_POLLS, POLL_INTERVAL, GCPNodeType
|
||||
from ray.autoscaler._private.util import (
|
||||
check_legacy_fields,
|
||||
generate_rsa_key_pair,
|
||||
generate_ssh_key_name,
|
||||
generate_ssh_key_paths,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
VERSION = "v1"
|
||||
TPU_VERSION = "v2alpha" # change once v2 is stable
|
||||
|
||||
RAY = "ray-autoscaler"
|
||||
DEFAULT_SERVICE_ACCOUNT_ID = RAY + "-sa-" + VERSION
|
||||
SERVICE_ACCOUNT_EMAIL_TEMPLATE = "{account_id}@{project_id}.iam.gserviceaccount.com"
|
||||
DEFAULT_SERVICE_ACCOUNT_CONFIG = {
|
||||
"displayName": "Ray Autoscaler Service Account ({})".format(VERSION),
|
||||
}
|
||||
|
||||
# Those roles will be always added.
|
||||
# NOTE: `serviceAccountUser` allows the head node to create workers with
|
||||
# a serviceAccount. `roleViewer` allows the head node to run bootstrap_gcp.
|
||||
DEFAULT_SERVICE_ACCOUNT_ROLES = [
|
||||
"roles/storage.objectAdmin",
|
||||
"roles/compute.admin",
|
||||
"roles/iam.serviceAccountUser",
|
||||
"roles/iam.roleViewer",
|
||||
]
|
||||
# Those roles will only be added if there are TPU nodes defined in config.
|
||||
TPU_SERVICE_ACCOUNT_ROLES = ["roles/tpu.admin"]
|
||||
|
||||
# If there are TPU nodes in config, this field will be set
|
||||
# to True in config["provider"].
|
||||
HAS_TPU_PROVIDER_FIELD = "_has_tpus"
|
||||
|
||||
# NOTE: iam.serviceAccountUser allows the Head Node to create worker nodes
|
||||
# with ServiceAccounts.
|
||||
|
||||
|
||||
def tpu_accelerator_config_to_type(accelerator_config: dict) -> str:
|
||||
"""Convert a provided accelerator_config to accelerator_type.
|
||||
|
||||
Args:
|
||||
accelerator_config: A dictionary defining the spec of a
|
||||
TPU accelerator. The dictionary should consist of
|
||||
the keys 'type', indicating the TPU chip type, and
|
||||
'topology', indicating the topology of the TPU.
|
||||
|
||||
Returns:
|
||||
A string, accelerator_type, e.g. "v4-8".
|
||||
|
||||
"""
|
||||
generation = accelerator_config["type"].lower()
|
||||
topology = accelerator_config["topology"]
|
||||
# Reduce e.g. "2x2x2" to 8
|
||||
chip_dimensions = [int(chip_count) for chip_count in topology.split("x")]
|
||||
num_chips = reduce(lambda x, y: x * y, chip_dimensions)
|
||||
|
||||
# V5LitePod is rendered as "V5LITE_POD" in accelerator configuration but
|
||||
# accelerator type uses a format like "v5litepod-{cores}", so we need
|
||||
# to manually convert the string here.
|
||||
if generation == "v5lite_pod":
|
||||
generation = "v5litepod"
|
||||
|
||||
num_cores = tpu.get_tpu_cores_per_chip(generation) * num_chips
|
||||
|
||||
return f"{generation}-{num_cores}"
|
||||
|
||||
|
||||
def _validate_tpu_config(node: dict):
|
||||
"""Validate the provided node with TPU support.
|
||||
|
||||
If the config is malformed, users will run into an error but this function
|
||||
will raise the error at config parsing time. This only tests very simple assertions.
|
||||
|
||||
Raises: `ValueError` in case the input is malformed.
|
||||
|
||||
"""
|
||||
if "acceleratorType" in node and "acceleratorConfig" in node:
|
||||
raise ValueError(
|
||||
"For TPU usage, acceleratorType and acceleratorConfig "
|
||||
"cannot both be set."
|
||||
)
|
||||
if "acceleratorType" in node:
|
||||
accelerator_type = node["acceleratorType"]
|
||||
if not TPUAcceleratorManager.is_valid_tpu_accelerator_type(accelerator_type):
|
||||
raise ValueError(
|
||||
"`acceleratorType` should match v(generation)-(cores/chips). "
|
||||
f"Got {accelerator_type}."
|
||||
)
|
||||
else: # "acceleratorConfig" in node
|
||||
accelerator_config = node["acceleratorConfig"]
|
||||
if "type" not in accelerator_config or "topology" not in accelerator_config:
|
||||
raise ValueError(
|
||||
"acceleratorConfig expects 'type' and 'topology'. "
|
||||
f"Got {accelerator_config}"
|
||||
)
|
||||
generation = node["acceleratorConfig"]["type"]
|
||||
topology = node["acceleratorConfig"]["topology"]
|
||||
|
||||
generation_pattern = re.compile(r"^V\d+[a-zA-Z]*$")
|
||||
topology_pattern = re.compile(r"^\d+x\d+(x\d+)?$")
|
||||
|
||||
if generation != "V5LITE_POD" and not generation_pattern.match(generation):
|
||||
raise ValueError(f"type should match V(generation). Got {generation}.")
|
||||
if generation == "V2" or generation == "V3":
|
||||
raise ValueError(
|
||||
f"acceleratorConfig is not supported on V2/V3 TPUs. Got {generation}."
|
||||
)
|
||||
if not topology_pattern.match(topology):
|
||||
raise ValueError(
|
||||
f"topology should be of form axbxc or axb. Got {topology}."
|
||||
)
|
||||
|
||||
|
||||
def _get_num_tpu_chips(node: dict) -> int:
|
||||
chips = 0
|
||||
if "acceleratorType" in node:
|
||||
accelerator_type = node["acceleratorType"]
|
||||
# `acceleratorType` is typically v{generation}-{cores}
|
||||
cores = int(accelerator_type.split("-")[1])
|
||||
chips = cores / tpu.get_tpu_cores_per_chip(accelerator_type)
|
||||
if "acceleratorConfig" in node:
|
||||
topology = node["acceleratorConfig"]["topology"]
|
||||
# `topology` is typically {chips}x{chips}x{chips}
|
||||
# Multiply all dimensions together to get total number of chips
|
||||
chips = 1
|
||||
for dim in topology.split("x"):
|
||||
chips *= int(dim)
|
||||
return chips
|
||||
|
||||
|
||||
def _is_single_host_tpu(node: dict) -> bool:
|
||||
accelerator_type = ""
|
||||
if "acceleratorType" in node:
|
||||
accelerator_type = node["acceleratorType"]
|
||||
else:
|
||||
accelerator_type = tpu_accelerator_config_to_type(node["acceleratorConfig"])
|
||||
return _get_num_tpu_chips(node) <= tpu.get_num_tpu_visible_chips_per_host(
|
||||
accelerator_type
|
||||
)
|
||||
|
||||
|
||||
def get_node_type(node: dict) -> GCPNodeType:
|
||||
"""Returns node type based on the keys in ``node``.
|
||||
|
||||
This is a very simple check. If we have a ``machineType`` key,
|
||||
this is a Compute instance. If we don't have a ``machineType`` key,
|
||||
but we have ``acceleratorType``, this is a TPU. Otherwise, it's
|
||||
invalid and an exception is raised.
|
||||
|
||||
This works for both node configs and API returned nodes.
|
||||
"""
|
||||
|
||||
if (
|
||||
"machineType" not in node
|
||||
and "acceleratorType" not in node
|
||||
and "acceleratorConfig" not in node
|
||||
):
|
||||
raise ValueError(
|
||||
"Invalid node. For a Compute instance, 'machineType' is required."
|
||||
"For a TPU instance, 'acceleratorType' OR 'acceleratorConfig' and "
|
||||
f"no 'machineType' is required. Got {list(node)}."
|
||||
)
|
||||
|
||||
if "machineType" not in node and (
|
||||
"acceleratorType" in node or "acceleratorConfig" in node
|
||||
):
|
||||
_validate_tpu_config(node)
|
||||
if not _is_single_host_tpu(node):
|
||||
# Remove once proper autoscaling support is added.
|
||||
logger.warning(
|
||||
"TPU pod detected. Note that while the cluster launcher can create "
|
||||
"multiple TPU pods, proper autoscaling will not work as expected, "
|
||||
"as all hosts in a TPU pod need to execute the same program. "
|
||||
"Proceed with caution."
|
||||
)
|
||||
return GCPNodeType.TPU
|
||||
return GCPNodeType.COMPUTE
|
||||
|
||||
|
||||
def wait_for_crm_operation(operation, crm):
|
||||
"""Poll for cloud resource manager operation until finished."""
|
||||
logger.info(
|
||||
"wait_for_crm_operation: "
|
||||
"Waiting for operation {} to finish...".format(operation)
|
||||
)
|
||||
|
||||
for _ in range(MAX_POLLS):
|
||||
result = crm.operations().get(name=operation["name"]).execute()
|
||||
if "error" in result:
|
||||
raise Exception(result["error"])
|
||||
|
||||
if "done" in result and result["done"]:
|
||||
logger.info("wait_for_crm_operation: Operation done.")
|
||||
break
|
||||
|
||||
time.sleep(POLL_INTERVAL)
|
||||
|
||||
return result
|
||||
|
||||
|
||||
def wait_for_compute_global_operation(project_name, operation, compute):
|
||||
"""Poll for global compute operation until finished."""
|
||||
logger.info(
|
||||
"wait_for_compute_global_operation: "
|
||||
"Waiting for operation {} to finish...".format(operation["name"])
|
||||
)
|
||||
|
||||
for _ in range(MAX_POLLS):
|
||||
result = (
|
||||
compute.globalOperations()
|
||||
.get(
|
||||
project=project_name,
|
||||
operation=operation["name"],
|
||||
)
|
||||
.execute()
|
||||
)
|
||||
if "error" in result:
|
||||
raise Exception(result["error"])
|
||||
|
||||
if result["status"] == "DONE":
|
||||
logger.info("wait_for_compute_global_operation: Operation done.")
|
||||
break
|
||||
|
||||
time.sleep(POLL_INTERVAL)
|
||||
|
||||
return result
|
||||
|
||||
|
||||
def _has_tpus_in_node_configs(config: dict) -> bool:
|
||||
"""Check if any nodes in config are TPUs."""
|
||||
node_configs = [
|
||||
node_type["node_config"]
|
||||
for node_type in config["available_node_types"].values()
|
||||
]
|
||||
return any(get_node_type(node) == GCPNodeType.TPU for node in node_configs)
|
||||
|
||||
|
||||
def _is_head_node_a_tpu(config: dict) -> bool:
|
||||
"""Check if the head node is a TPU."""
|
||||
node_configs = {
|
||||
node_id: node_type["node_config"]
|
||||
for node_id, node_type in config["available_node_types"].items()
|
||||
}
|
||||
return get_node_type(node_configs[config["head_node_type"]]) == GCPNodeType.TPU
|
||||
|
||||
|
||||
def build_request(http, *args, **kwargs):
|
||||
new_http = google_auth_httplib2.AuthorizedHttp(
|
||||
http.credentials, http=httplib2.Http()
|
||||
)
|
||||
return googleapiclient.http.HttpRequest(new_http, *args, **kwargs)
|
||||
|
||||
|
||||
def _create_crm(gcp_credentials=None):
|
||||
return discovery.build(
|
||||
"cloudresourcemanager",
|
||||
"v1",
|
||||
credentials=gcp_credentials,
|
||||
requestBuilder=build_request,
|
||||
cache_discovery=False,
|
||||
)
|
||||
|
||||
|
||||
def _create_iam(gcp_credentials=None):
|
||||
return discovery.build(
|
||||
"iam",
|
||||
"v1",
|
||||
credentials=gcp_credentials,
|
||||
requestBuilder=build_request,
|
||||
cache_discovery=False,
|
||||
)
|
||||
|
||||
|
||||
def _create_compute(gcp_credentials=None):
|
||||
return discovery.build(
|
||||
"compute",
|
||||
"v1",
|
||||
credentials=gcp_credentials,
|
||||
requestBuilder=build_request,
|
||||
cache_discovery=False,
|
||||
)
|
||||
|
||||
|
||||
def _create_tpu(gcp_credentials=None):
|
||||
return discovery.build(
|
||||
"tpu",
|
||||
TPU_VERSION,
|
||||
credentials=gcp_credentials,
|
||||
requestBuilder=build_request,
|
||||
cache_discovery=False,
|
||||
discoveryServiceUrl="https://tpu.googleapis.com/$discovery/rest",
|
||||
)
|
||||
|
||||
|
||||
def construct_clients_from_provider_config(provider_config):
|
||||
"""
|
||||
Attempt to fetch and parse the JSON GCP credentials from the provider
|
||||
config yaml file.
|
||||
|
||||
tpu resource (the last element of the tuple) will be None if
|
||||
`_has_tpus` in provider config is not set or False.
|
||||
"""
|
||||
gcp_credentials = provider_config.get("gcp_credentials")
|
||||
if gcp_credentials is None:
|
||||
logger.debug(
|
||||
"gcp_credentials not found in cluster yaml file. "
|
||||
"Falling back to GOOGLE_APPLICATION_CREDENTIALS "
|
||||
"environment variable."
|
||||
)
|
||||
tpu_resource = (
|
||||
_create_tpu()
|
||||
if provider_config.get(HAS_TPU_PROVIDER_FIELD, False)
|
||||
else None
|
||||
)
|
||||
# If gcp_credentials is None, then discovery.build will search for
|
||||
# credentials in the local environment.
|
||||
return _create_crm(), _create_iam(), _create_compute(), tpu_resource
|
||||
|
||||
assert (
|
||||
"type" in gcp_credentials
|
||||
), "gcp_credentials cluster yaml field missing 'type' field."
|
||||
assert (
|
||||
"credentials" in gcp_credentials
|
||||
), "gcp_credentials cluster yaml field missing 'credentials' field."
|
||||
|
||||
cred_type = gcp_credentials["type"]
|
||||
credentials_field = gcp_credentials["credentials"]
|
||||
|
||||
if cred_type == "service_account":
|
||||
# If parsing the gcp_credentials failed, then the user likely made a
|
||||
# mistake in copying the credentials into the config yaml.
|
||||
try:
|
||||
service_account_info = json.loads(credentials_field)
|
||||
except json.decoder.JSONDecodeError:
|
||||
raise RuntimeError(
|
||||
"gcp_credentials found in cluster yaml file but "
|
||||
"formatted improperly."
|
||||
)
|
||||
credentials = service_account.Credentials.from_service_account_info(
|
||||
service_account_info
|
||||
)
|
||||
elif cred_type == "credentials_token":
|
||||
# Otherwise the credentials type must be credentials_token.
|
||||
credentials = OAuthCredentials(credentials_field)
|
||||
|
||||
tpu_resource = (
|
||||
_create_tpu(credentials)
|
||||
if provider_config.get(HAS_TPU_PROVIDER_FIELD, False)
|
||||
else None
|
||||
)
|
||||
|
||||
return (
|
||||
_create_crm(credentials),
|
||||
_create_iam(credentials),
|
||||
_create_compute(credentials),
|
||||
tpu_resource,
|
||||
)
|
||||
|
||||
|
||||
def bootstrap_gcp(config):
|
||||
config = copy.deepcopy(config)
|
||||
check_legacy_fields(config)
|
||||
# Used internally to store head IAM role.
|
||||
config["head_node"] = {}
|
||||
|
||||
# Check if we have any TPUs defined, and if so,
|
||||
# insert that information into the provider config
|
||||
if _has_tpus_in_node_configs(config):
|
||||
config["provider"][HAS_TPU_PROVIDER_FIELD] = True
|
||||
|
||||
crm, iam, compute, tpu = construct_clients_from_provider_config(config["provider"])
|
||||
|
||||
config = _configure_project(config, crm)
|
||||
config = _configure_iam_role(config, crm, iam)
|
||||
config = _configure_key_pair(config, compute)
|
||||
config = _configure_subnet(config, compute)
|
||||
|
||||
return config
|
||||
|
||||
|
||||
def _configure_project(config, crm):
|
||||
"""Setup a Google Cloud Platform Project.
|
||||
|
||||
Google Compute Platform organizes all the resources, such as storage
|
||||
buckets, users, and instances under projects. This is different from
|
||||
aws ec2 where everything is global.
|
||||
"""
|
||||
config = copy.deepcopy(config)
|
||||
|
||||
project_id = config["provider"].get("project_id")
|
||||
assert config["provider"]["project_id"] is not None, (
|
||||
"'project_id' must be set in the 'provider' section of the autoscaler"
|
||||
" config. Notice that the project id must be globally unique."
|
||||
)
|
||||
project = _get_project(project_id, crm)
|
||||
|
||||
if project is None:
|
||||
# Project not found, try creating it
|
||||
_create_project(project_id, crm)
|
||||
project = _get_project(project_id, crm)
|
||||
|
||||
assert project is not None, "Failed to create project"
|
||||
assert (
|
||||
project["lifecycleState"] == "ACTIVE"
|
||||
), "Project status needs to be ACTIVE, got {}".format(project["lifecycleState"])
|
||||
|
||||
config["provider"]["project_id"] = project["projectId"]
|
||||
|
||||
return config
|
||||
|
||||
|
||||
def _configure_iam_role(config, crm, iam):
|
||||
"""Setup a gcp service account with IAM roles.
|
||||
|
||||
Creates a gcp service acconut and binds IAM roles which allow it to control
|
||||
control storage/compute services. Specifically, the head node needs to have
|
||||
an IAM role that allows it to create further gce instances and store items
|
||||
in google cloud storage.
|
||||
|
||||
TODO: Allow the name/id of the service account to be configured
|
||||
"""
|
||||
config = copy.deepcopy(config)
|
||||
|
||||
email = SERVICE_ACCOUNT_EMAIL_TEMPLATE.format(
|
||||
account_id=DEFAULT_SERVICE_ACCOUNT_ID,
|
||||
project_id=config["provider"]["project_id"],
|
||||
)
|
||||
service_account = _get_service_account(email, config, iam)
|
||||
|
||||
if service_account is None:
|
||||
logger.info(
|
||||
"_configure_iam_role: "
|
||||
"Creating new service account {}".format(DEFAULT_SERVICE_ACCOUNT_ID)
|
||||
)
|
||||
|
||||
service_account = _create_service_account(
|
||||
DEFAULT_SERVICE_ACCOUNT_ID, DEFAULT_SERVICE_ACCOUNT_CONFIG, config, iam
|
||||
)
|
||||
|
||||
assert service_account is not None, "Failed to create service account"
|
||||
|
||||
if config["provider"].get(HAS_TPU_PROVIDER_FIELD, False):
|
||||
roles = DEFAULT_SERVICE_ACCOUNT_ROLES + TPU_SERVICE_ACCOUNT_ROLES
|
||||
else:
|
||||
roles = DEFAULT_SERVICE_ACCOUNT_ROLES
|
||||
|
||||
_add_iam_policy_binding(service_account, roles, crm)
|
||||
|
||||
config["head_node"]["serviceAccounts"] = [
|
||||
{
|
||||
"email": service_account["email"],
|
||||
# NOTE: The amount of access is determined by the scope + IAM
|
||||
# role of the service account. Even if the cloud-platform scope
|
||||
# gives (scope) access to the whole cloud-platform, the service
|
||||
# account is limited by the IAM rights specified below.
|
||||
"scopes": ["https://www.googleapis.com/auth/cloud-platform"],
|
||||
}
|
||||
]
|
||||
|
||||
return config
|
||||
|
||||
|
||||
def _configure_key_pair(config, compute):
|
||||
"""Configure SSH access, using an existing key pair if possible.
|
||||
|
||||
Creates a project-wide ssh key that can be used to access all the instances
|
||||
unless explicitly prohibited by instance config.
|
||||
|
||||
The ssh-keys created by ray are of format:
|
||||
|
||||
[USERNAME]:ssh-rsa [KEY_VALUE] [USERNAME]
|
||||
|
||||
where:
|
||||
|
||||
[USERNAME] is the user for the SSH key, specified in the config.
|
||||
[KEY_VALUE] is the public SSH key value.
|
||||
"""
|
||||
config = copy.deepcopy(config)
|
||||
|
||||
if "ssh_private_key" in config["auth"]:
|
||||
return config
|
||||
|
||||
ssh_user = config["auth"]["ssh_user"]
|
||||
|
||||
project = compute.projects().get(project=config["provider"]["project_id"]).execute()
|
||||
|
||||
# Key pairs associated with project meta data. The key pairs are general,
|
||||
# and not just ssh keys.
|
||||
ssh_keys_str = next(
|
||||
(
|
||||
item
|
||||
for item in project["commonInstanceMetadata"].get("items", [])
|
||||
if item["key"] == "ssh-keys"
|
||||
),
|
||||
{},
|
||||
).get("value", "")
|
||||
|
||||
ssh_keys = ssh_keys_str.split("\n") if ssh_keys_str else []
|
||||
|
||||
# Try a few times to get or create a good key pair.
|
||||
key_found = False
|
||||
for i in range(10):
|
||||
key_name = generate_ssh_key_name(
|
||||
"gcp",
|
||||
i,
|
||||
config["provider"]["region"],
|
||||
config["provider"]["project_id"],
|
||||
ssh_user,
|
||||
)
|
||||
public_key_path, private_key_path = generate_ssh_key_paths(key_name)
|
||||
|
||||
for ssh_key in ssh_keys:
|
||||
key_parts = ssh_key.split(" ")
|
||||
if len(key_parts) != 3:
|
||||
continue
|
||||
|
||||
if key_parts[2] == ssh_user and os.path.exists(private_key_path):
|
||||
# Found a key
|
||||
key_found = True
|
||||
break
|
||||
|
||||
# Writing the new ssh key to the filesystem fails if the ~/.ssh
|
||||
# directory doesn't already exist.
|
||||
os.makedirs(os.path.expanduser("~/.ssh"), exist_ok=True)
|
||||
|
||||
# Create a key since it doesn't exist locally or in GCP
|
||||
if not key_found and not os.path.exists(private_key_path):
|
||||
logger.info(
|
||||
"_configure_key_pair: Creating new key pair {}".format(key_name)
|
||||
)
|
||||
public_key, private_key = generate_rsa_key_pair()
|
||||
|
||||
for attempt in range(MAX_POLLS):
|
||||
try:
|
||||
_create_project_ssh_key_pair(project, public_key, ssh_user, compute)
|
||||
break
|
||||
except errors.HttpError as e:
|
||||
if e.resp.status != 412 or attempt == MAX_POLLS - 1:
|
||||
raise
|
||||
logger.warning(
|
||||
"GCP project metadata update conflict for %s (%s); retrying",
|
||||
config["provider"]["project_id"],
|
||||
e,
|
||||
)
|
||||
time.sleep(POLL_INTERVAL)
|
||||
project = (
|
||||
compute.projects()
|
||||
.get(project=config["provider"]["project_id"])
|
||||
.execute()
|
||||
)
|
||||
|
||||
# Create the directory if it doesn't exists
|
||||
private_key_dir = os.path.dirname(private_key_path)
|
||||
os.makedirs(private_key_dir, exist_ok=True)
|
||||
|
||||
# We need to make sure to _create_ the file with the right
|
||||
# permissions. In order to do that we need to change the default
|
||||
# os.open behavior to include the mode we want.
|
||||
with open(
|
||||
private_key_path,
|
||||
"w",
|
||||
opener=partial(os.open, mode=0o600),
|
||||
) as f:
|
||||
f.write(private_key)
|
||||
|
||||
with open(public_key_path, "w") as f:
|
||||
f.write(public_key)
|
||||
|
||||
key_found = True
|
||||
|
||||
break
|
||||
|
||||
if key_found:
|
||||
break
|
||||
|
||||
assert key_found, "SSH keypair for user {} not found for {}".format(
|
||||
ssh_user, private_key_path
|
||||
)
|
||||
assert os.path.exists(
|
||||
private_key_path
|
||||
), "Private key file {} not found for user {}".format(private_key_path, ssh_user)
|
||||
|
||||
logger.info(
|
||||
"_configure_key_pair: "
|
||||
"Private key not specified in config, using"
|
||||
"{}".format(private_key_path)
|
||||
)
|
||||
|
||||
config["auth"]["ssh_private_key"] = private_key_path
|
||||
|
||||
return config
|
||||
|
||||
|
||||
def _configure_subnet(config, compute):
|
||||
"""Pick a reasonable subnet if not specified by the config."""
|
||||
config = copy.deepcopy(config)
|
||||
|
||||
node_configs = [
|
||||
node_type["node_config"]
|
||||
for node_type in config["available_node_types"].values()
|
||||
]
|
||||
# Rationale: avoid subnet lookup if the network is already
|
||||
# completely manually configured
|
||||
|
||||
# networkInterfaces is compute, networkConfig is TPU
|
||||
if all(
|
||||
"networkInterfaces" in node_config or "networkConfig" in node_config
|
||||
for node_config in node_configs
|
||||
):
|
||||
return config
|
||||
|
||||
subnets = _list_subnets(config, compute)
|
||||
|
||||
if not subnets:
|
||||
raise NotImplementedError("Should be able to create subnet.")
|
||||
|
||||
# TODO: make sure that we have usable subnet. Maybe call
|
||||
# compute.subnetworks().listUsable? For some reason it didn't
|
||||
# work out-of-the-box
|
||||
default_subnet = subnets[0]
|
||||
|
||||
default_interfaces = [
|
||||
{
|
||||
"subnetwork": default_subnet["selfLink"],
|
||||
"accessConfigs": [
|
||||
{
|
||||
"name": "External NAT",
|
||||
"type": "ONE_TO_ONE_NAT",
|
||||
}
|
||||
],
|
||||
}
|
||||
]
|
||||
|
||||
for node_config in node_configs:
|
||||
# The not applicable key will be removed during node creation
|
||||
|
||||
# compute
|
||||
if "networkInterfaces" not in node_config:
|
||||
node_config["networkInterfaces"] = copy.deepcopy(default_interfaces)
|
||||
# TPU
|
||||
if "networkConfig" not in node_config:
|
||||
node_config["networkConfig"] = copy.deepcopy(default_interfaces)[0]
|
||||
node_config["networkConfig"].pop("accessConfigs")
|
||||
|
||||
return config
|
||||
|
||||
|
||||
def _list_subnets(config, compute):
|
||||
response = (
|
||||
compute.subnetworks()
|
||||
.list(
|
||||
project=config["provider"]["project_id"],
|
||||
region=config["provider"]["region"],
|
||||
)
|
||||
.execute()
|
||||
)
|
||||
|
||||
return response["items"]
|
||||
|
||||
|
||||
def _get_subnet(config, subnet_id, compute):
|
||||
subnet = (
|
||||
compute.subnetworks()
|
||||
.get(
|
||||
project=config["provider"]["project_id"],
|
||||
region=config["provider"]["region"],
|
||||
subnetwork=subnet_id,
|
||||
)
|
||||
.execute()
|
||||
)
|
||||
|
||||
return subnet
|
||||
|
||||
|
||||
def _get_project(project_id, crm):
|
||||
try:
|
||||
project = crm.projects().get(projectId=project_id).execute()
|
||||
except errors.HttpError as e:
|
||||
if e.resp.status != 403:
|
||||
raise
|
||||
project = None
|
||||
|
||||
return project
|
||||
|
||||
|
||||
def _create_project(project_id, crm):
|
||||
operation = (
|
||||
crm.projects()
|
||||
.create(body={"projectId": project_id, "name": project_id})
|
||||
.execute()
|
||||
)
|
||||
|
||||
result = wait_for_crm_operation(operation, crm)
|
||||
|
||||
return result
|
||||
|
||||
|
||||
def _get_service_account(account, config, iam):
|
||||
project_id = config["provider"]["project_id"]
|
||||
full_name = "projects/{project_id}/serviceAccounts/{account}".format(
|
||||
project_id=project_id, account=account
|
||||
)
|
||||
try:
|
||||
service_account = iam.projects().serviceAccounts().get(name=full_name).execute()
|
||||
except errors.HttpError as e:
|
||||
if e.resp.status != 404:
|
||||
raise
|
||||
service_account = None
|
||||
|
||||
return service_account
|
||||
|
||||
|
||||
def _create_service_account(account_id, account_config, config, iam):
|
||||
project_id = config["provider"]["project_id"]
|
||||
|
||||
service_account = (
|
||||
iam.projects()
|
||||
.serviceAccounts()
|
||||
.create(
|
||||
name="projects/{project_id}".format(project_id=project_id),
|
||||
body={
|
||||
"accountId": account_id,
|
||||
"serviceAccount": account_config,
|
||||
},
|
||||
)
|
||||
.execute()
|
||||
)
|
||||
|
||||
return service_account
|
||||
|
||||
|
||||
def _add_iam_policy_binding(service_account, roles, crm):
|
||||
"""Add new IAM roles for the service account."""
|
||||
project_id = service_account["projectId"]
|
||||
email = service_account["email"]
|
||||
member_id = "serviceAccount:" + email
|
||||
|
||||
policy = (
|
||||
crm.projects()
|
||||
.getIamPolicy(
|
||||
resource=project_id, body={"options": {"requestedPolicyVersion": 3}}
|
||||
)
|
||||
.execute()
|
||||
)
|
||||
|
||||
already_configured = True
|
||||
for role in roles:
|
||||
role_exists = False
|
||||
for binding in policy["bindings"]:
|
||||
if binding["role"] == role:
|
||||
if member_id not in binding["members"]:
|
||||
binding["members"].append(member_id)
|
||||
already_configured = False
|
||||
role_exists = True
|
||||
|
||||
if not role_exists:
|
||||
already_configured = False
|
||||
policy["bindings"].append(
|
||||
{
|
||||
"members": [member_id],
|
||||
"role": role,
|
||||
}
|
||||
)
|
||||
|
||||
if already_configured:
|
||||
# In some managed environments, an admin needs to grant the
|
||||
# roles, so only call setIamPolicy if needed.
|
||||
return
|
||||
|
||||
result = (
|
||||
crm.projects()
|
||||
.setIamPolicy(
|
||||
resource=project_id,
|
||||
body={
|
||||
"policy": policy,
|
||||
},
|
||||
)
|
||||
.execute()
|
||||
)
|
||||
|
||||
return result
|
||||
|
||||
|
||||
def _create_project_ssh_key_pair(project, public_key, ssh_user, compute):
|
||||
"""Inserts an ssh-key into project commonInstanceMetadata"""
|
||||
|
||||
key_parts = public_key.split(" ")
|
||||
|
||||
# Sanity checks to make sure that the generated key matches expectation
|
||||
assert len(key_parts) == 2, key_parts
|
||||
assert key_parts[0] == "ssh-rsa", key_parts
|
||||
|
||||
new_ssh_meta = "{ssh_user}:ssh-rsa {key_value} {ssh_user}".format(
|
||||
ssh_user=ssh_user, key_value=key_parts[1]
|
||||
)
|
||||
|
||||
common_instance_metadata = project["commonInstanceMetadata"]
|
||||
items = common_instance_metadata.get("items", [])
|
||||
|
||||
ssh_keys_i = next(
|
||||
(i for i, item in enumerate(items) if item["key"] == "ssh-keys"), None
|
||||
)
|
||||
|
||||
if ssh_keys_i is None:
|
||||
items.append({"key": "ssh-keys", "value": new_ssh_meta})
|
||||
else:
|
||||
ssh_keys = items[ssh_keys_i]
|
||||
ssh_keys["value"] += "\n" + new_ssh_meta
|
||||
items[ssh_keys_i] = ssh_keys
|
||||
|
||||
common_instance_metadata["items"] = items
|
||||
|
||||
try:
|
||||
operation = (
|
||||
compute.projects()
|
||||
.setCommonInstanceMetadata(
|
||||
project=project["name"], body=common_instance_metadata
|
||||
)
|
||||
.execute()
|
||||
)
|
||||
response = wait_for_compute_global_operation(
|
||||
project["name"], operation, compute
|
||||
)
|
||||
except errors.HttpError as e: # Only trim when explicitly opted in.
|
||||
if (
|
||||
e.resp.status != 413
|
||||
or ssh_keys_i is None
|
||||
or os.environ.get("RAY_TRIM_AUTOSCALER_SSH_KEYS") != "1"
|
||||
):
|
||||
raise
|
||||
logger.warning(
|
||||
"GCP project metadata size limit reached for %s (%s); "
|
||||
"removing half of ssh keys and retrying",
|
||||
project["name"],
|
||||
e,
|
||||
)
|
||||
ssh_keys_value = items[ssh_keys_i].get("value", "")
|
||||
ssh_keys = [key for key in ssh_keys_value.splitlines() if key.strip()]
|
||||
trim_start = ( # If our new key is the only one, this preserves it.
|
||||
len(ssh_keys) // 2
|
||||
)
|
||||
items[ssh_keys_i]["value"] = "\n".join(ssh_keys[trim_start:])
|
||||
common_instance_metadata["items"] = items
|
||||
operation = (
|
||||
compute.projects()
|
||||
.setCommonInstanceMetadata(
|
||||
project=project["name"], body=common_instance_metadata
|
||||
)
|
||||
.execute()
|
||||
)
|
||||
response = wait_for_compute_global_operation(
|
||||
project["name"], operation, compute
|
||||
)
|
||||
|
||||
return response
|
||||
@@ -0,0 +1,856 @@
|
||||
"""Abstractions around GCP resources and nodes.
|
||||
|
||||
The logic has been abstracted away here to allow for different GCP resources
|
||||
(API endpoints), which can differ widely, making it impossible to use
|
||||
the same logic for everything.
|
||||
|
||||
Classes inheriting from ``GCPResource`` represent different GCP resources -
|
||||
API endpoints that allow for nodes to be created, removed, listed and
|
||||
otherwise managed. Those classes contain methods abstracting GCP REST API
|
||||
calls.
|
||||
Each resource has a corresponding node type, represented by a
|
||||
class inheriting from ``GCPNode``. Those classes are essentially dicts
|
||||
with some extra methods. The instances of those classes will be created
|
||||
from API responses.
|
||||
|
||||
The ``GCPNodeType`` enum is a lightweight way to classify nodes.
|
||||
|
||||
Currently, Compute and TPU resources & nodes are supported.
|
||||
|
||||
In order to add support for new resources, create classes inheriting from
|
||||
``GCPResource`` and ``GCPNode``, update the ``GCPNodeType`` enum,
|
||||
update the ``_generate_node_name`` method and finally update the
|
||||
node provider.
|
||||
"""
|
||||
|
||||
import abc
|
||||
import logging
|
||||
import re
|
||||
import time
|
||||
from collections import UserDict
|
||||
from copy import deepcopy
|
||||
from enum import Enum
|
||||
from functools import wraps
|
||||
from typing import Any, Dict, List, Optional, Tuple, Union
|
||||
from uuid import uuid4
|
||||
|
||||
import httplib2
|
||||
from google_auth_httplib2 import AuthorizedHttp
|
||||
from googleapiclient.discovery import Resource
|
||||
from googleapiclient.errors import HttpError
|
||||
|
||||
from ray.autoscaler.tags import TAG_RAY_CLUSTER_NAME, TAG_RAY_NODE_NAME
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
INSTANCE_NAME_MAX_LEN = 64
|
||||
INSTANCE_NAME_UUID_LEN = 8
|
||||
MAX_POLLS = 12
|
||||
# TPUs take a long while to respond, so we increase the MAX_POLLS
|
||||
# considerably - this probably could be smaller
|
||||
# TPU deletion uses MAX_POLLS
|
||||
MAX_POLLS_TPU = MAX_POLLS * 8
|
||||
POLL_INTERVAL = 5
|
||||
|
||||
|
||||
def _retry_on_exception(
|
||||
exception: Union[Exception, Tuple[Exception]],
|
||||
regex: Optional[str] = None,
|
||||
max_retries: int = MAX_POLLS,
|
||||
retry_interval_s: int = POLL_INTERVAL,
|
||||
):
|
||||
"""Retry a function call n-times for as long as it throws an exception."""
|
||||
|
||||
def dec(func):
|
||||
@wraps(func)
|
||||
def wrapper(*args, **kwargs):
|
||||
def try_catch_exc():
|
||||
try:
|
||||
value = func(*args, **kwargs)
|
||||
return value
|
||||
except Exception as e:
|
||||
if not isinstance(e, exception) or (
|
||||
regex and not re.search(regex, str(e))
|
||||
):
|
||||
raise e
|
||||
return e
|
||||
|
||||
for _ in range(max_retries):
|
||||
ret = try_catch_exc()
|
||||
if not isinstance(ret, Exception):
|
||||
break
|
||||
time.sleep(retry_interval_s)
|
||||
if isinstance(ret, Exception):
|
||||
raise ret
|
||||
return ret
|
||||
|
||||
return wrapper
|
||||
|
||||
return dec
|
||||
|
||||
|
||||
def _generate_node_name(labels: dict, node_suffix: str) -> str:
|
||||
"""Generate node name from labels and suffix.
|
||||
|
||||
This is required so that the correct resource can be selected
|
||||
when the only information autoscaler has is the name of the node.
|
||||
|
||||
The suffix is expected to be one of 'compute' or 'tpu'
|
||||
(as in ``GCPNodeType``).
|
||||
"""
|
||||
name_label = labels[TAG_RAY_NODE_NAME]
|
||||
assert len(name_label) <= (INSTANCE_NAME_MAX_LEN - INSTANCE_NAME_UUID_LEN - 1), (
|
||||
name_label,
|
||||
len(name_label),
|
||||
)
|
||||
return f"{name_label}-{uuid4().hex[:INSTANCE_NAME_UUID_LEN]}-{node_suffix}"
|
||||
|
||||
|
||||
class GCPNodeType(Enum):
|
||||
"""Enum for GCP node types (compute & tpu)"""
|
||||
|
||||
COMPUTE = "compute"
|
||||
TPU = "tpu"
|
||||
|
||||
@staticmethod
|
||||
def from_gcp_node(node: "GCPNode"):
|
||||
"""Return GCPNodeType based on ``node``'s class"""
|
||||
if isinstance(node, GCPTPUNode):
|
||||
return GCPNodeType.TPU
|
||||
if isinstance(node, GCPComputeNode):
|
||||
return GCPNodeType.COMPUTE
|
||||
raise TypeError(f"Wrong GCPNode type {type(node)}.")
|
||||
|
||||
@staticmethod
|
||||
def name_to_type(name: str):
|
||||
"""Provided a node name, determine the type.
|
||||
|
||||
This expects the name to be in format '[NAME]-[UUID]-[TYPE]',
|
||||
where [TYPE] is either 'compute' or 'tpu'.
|
||||
"""
|
||||
return GCPNodeType(name.split("-")[-1])
|
||||
|
||||
|
||||
class GCPNode(UserDict, metaclass=abc.ABCMeta):
|
||||
"""Abstraction around compute and tpu nodes"""
|
||||
|
||||
NON_TERMINATED_STATUSES = None
|
||||
RUNNING_STATUSES = None
|
||||
STATUS_FIELD = None
|
||||
|
||||
def __init__(self, base_dict: dict, resource: "GCPResource", **kwargs) -> None:
|
||||
super().__init__(base_dict, **kwargs)
|
||||
self.resource = resource
|
||||
assert isinstance(self.resource, GCPResource)
|
||||
|
||||
def is_running(self) -> bool:
|
||||
return self.get(self.STATUS_FIELD) in self.RUNNING_STATUSES
|
||||
|
||||
def is_terminated(self) -> bool:
|
||||
return self.get(self.STATUS_FIELD) not in self.NON_TERMINATED_STATUSES
|
||||
|
||||
@abc.abstractmethod
|
||||
def get_labels(self) -> dict:
|
||||
return
|
||||
|
||||
@abc.abstractmethod
|
||||
def get_external_ip(self) -> str:
|
||||
return
|
||||
|
||||
@abc.abstractmethod
|
||||
def get_internal_ip(self) -> str:
|
||||
return
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return f"<{self.__class__.__name__}: {self.get('name')}>"
|
||||
|
||||
|
||||
class GCPComputeNode(GCPNode):
|
||||
"""Abstraction around compute nodes"""
|
||||
|
||||
# https://cloud.google.com/compute/docs/instances/instance-life-cycle
|
||||
NON_TERMINATED_STATUSES = {"PROVISIONING", "STAGING", "RUNNING"}
|
||||
TERMINATED_STATUSES = {"TERMINATED", "SUSPENDED"}
|
||||
RUNNING_STATUSES = {"RUNNING"}
|
||||
STATUS_FIELD = "status"
|
||||
|
||||
def get_labels(self) -> dict:
|
||||
return self.get("labels", {})
|
||||
|
||||
def get_external_ip(self) -> str:
|
||||
return (
|
||||
self.get("networkInterfaces", [{}])[0]
|
||||
.get("accessConfigs", [{}])[0]
|
||||
.get("natIP", None)
|
||||
)
|
||||
|
||||
def get_internal_ip(self) -> str:
|
||||
return self.get("networkInterfaces", [{}])[0].get("networkIP")
|
||||
|
||||
|
||||
class GCPTPUNode(GCPNode):
|
||||
"""Abstraction around tpu nodes"""
|
||||
|
||||
# https://cloud.google.com/tpu/docs/reference/rest/v2alpha1/projects.locations.nodes#State
|
||||
|
||||
NON_TERMINATED_STATUSES = {"CREATING", "STARTING", "RESTARTING", "READY"}
|
||||
RUNNING_STATUSES = {"READY"}
|
||||
STATUS_FIELD = "state"
|
||||
|
||||
def get_labels(self) -> dict:
|
||||
return self.get("labels", {})
|
||||
|
||||
@property
|
||||
def num_workers(self) -> int:
|
||||
return len(self.get("networkEndpoints", [{}]))
|
||||
|
||||
def get_external_ips(self) -> List[str]:
|
||||
return self.get("networkEndpoints", [{}])
|
||||
|
||||
def get_external_ip(self, worker_index: int = 0) -> str:
|
||||
return (
|
||||
self.get_external_ips()[worker_index]
|
||||
.get("accessConfig", {})
|
||||
.get("externalIp", None)
|
||||
)
|
||||
|
||||
def get_internal_ips(self) -> List[str]:
|
||||
return self.get("networkEndpoints", [{}])
|
||||
|
||||
def get_internal_ip(self, worker_index: int = 0) -> str:
|
||||
return self.get_internal_ips()[worker_index].get("ipAddress", None)
|
||||
|
||||
|
||||
class GCPResource(metaclass=abc.ABCMeta):
|
||||
"""Abstraction around compute and TPU resources"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
resource: Resource,
|
||||
project_id: str,
|
||||
availability_zone: str,
|
||||
cluster_name: str,
|
||||
) -> None:
|
||||
self.resource = resource
|
||||
self.project_id = project_id
|
||||
self.availability_zone = availability_zone
|
||||
self.cluster_name = cluster_name
|
||||
|
||||
@abc.abstractmethod
|
||||
def get_new_authorized_http(self, http: AuthorizedHttp) -> AuthorizedHttp:
|
||||
"""Generate a new AuthorizedHttp object with the given credentials."""
|
||||
return
|
||||
|
||||
@abc.abstractmethod
|
||||
def wait_for_operation(
|
||||
self,
|
||||
operation: dict,
|
||||
max_polls: int = MAX_POLLS,
|
||||
poll_interval: int = POLL_INTERVAL,
|
||||
) -> dict:
|
||||
"""Waits a preset amount of time for operation to complete."""
|
||||
return None
|
||||
|
||||
@abc.abstractmethod
|
||||
def list_instances(
|
||||
self,
|
||||
label_filters: Optional[dict] = None,
|
||||
is_terminated: bool = False,
|
||||
) -> List["GCPNode"]:
|
||||
"""Returns a filtered list of all instances.
|
||||
|
||||
The filter removes all terminated instances and, if ``label_filters``
|
||||
are provided, all instances which labels are not matching the
|
||||
ones provided.
|
||||
"""
|
||||
return
|
||||
|
||||
@abc.abstractmethod
|
||||
def get_instance(self, node_id: str) -> "GCPNode":
|
||||
"""Returns a single instance."""
|
||||
return
|
||||
|
||||
@abc.abstractmethod
|
||||
def set_labels(
|
||||
self, node: GCPNode, labels: dict, wait_for_operation: bool = True
|
||||
) -> dict:
|
||||
"""Sets labels on an instance and returns result.
|
||||
|
||||
Completely replaces the labels dictionary."""
|
||||
return
|
||||
|
||||
@abc.abstractmethod
|
||||
def create_instance(
|
||||
self, base_config: dict, labels: dict, wait_for_operation: bool = True
|
||||
) -> Tuple[dict, str]:
|
||||
"""Creates a single instance and returns result.
|
||||
|
||||
Returns a tuple of (result, node_name).
|
||||
"""
|
||||
return
|
||||
|
||||
def create_instances(
|
||||
self,
|
||||
base_config: dict,
|
||||
labels: dict,
|
||||
count: int,
|
||||
wait_for_operation: bool = True,
|
||||
) -> List[Tuple[dict, str]]:
|
||||
"""Creates multiple instances and returns result.
|
||||
|
||||
Returns a list of tuples of (result, node_name).
|
||||
"""
|
||||
operations = [
|
||||
self.create_instance(base_config, labels, wait_for_operation=False)
|
||||
for i in range(count)
|
||||
]
|
||||
|
||||
if wait_for_operation:
|
||||
results = [
|
||||
(self.wait_for_operation(operation), node_name)
|
||||
for operation, node_name in operations
|
||||
]
|
||||
else:
|
||||
results = operations
|
||||
|
||||
return results
|
||||
|
||||
@abc.abstractmethod
|
||||
def delete_instance(self, node_id: str, wait_for_operation: bool = True) -> dict:
|
||||
"""Deletes an instance and returns result."""
|
||||
return
|
||||
|
||||
@abc.abstractmethod
|
||||
def stop_instance(self, node_id: str, wait_for_operation: bool = True) -> dict:
|
||||
"""Deletes an instance and returns result."""
|
||||
return
|
||||
|
||||
@abc.abstractmethod
|
||||
def start_instance(self, node_id: str, wait_for_operation: bool = True) -> dict:
|
||||
"""Starts a single instance and returns result."""
|
||||
return
|
||||
|
||||
|
||||
class GCPCompute(GCPResource):
|
||||
"""Abstraction around GCP compute resource"""
|
||||
|
||||
def get_new_authorized_http(self, http: AuthorizedHttp) -> AuthorizedHttp:
|
||||
"""Generate a new AuthorizedHttp object with the given credentials."""
|
||||
new_http = AuthorizedHttp(http.credentials, http=httplib2.Http())
|
||||
return new_http
|
||||
|
||||
def wait_for_operation(
|
||||
self,
|
||||
operation: dict,
|
||||
max_polls: int = MAX_POLLS,
|
||||
poll_interval: int = POLL_INTERVAL,
|
||||
) -> dict:
|
||||
"""Poll for compute zone operation until finished."""
|
||||
logger.info(
|
||||
"wait_for_compute_zone_operation: "
|
||||
f"Waiting for operation {operation['name']} to finish..."
|
||||
)
|
||||
|
||||
for _ in range(max_polls):
|
||||
result = (
|
||||
self.resource.zoneOperations()
|
||||
.get(
|
||||
project=self.project_id,
|
||||
operation=operation["name"],
|
||||
zone=self.availability_zone,
|
||||
)
|
||||
.execute(http=self.get_new_authorized_http(self.resource._http))
|
||||
)
|
||||
if "error" in result:
|
||||
raise Exception(result["error"])
|
||||
|
||||
if result["status"] == "DONE":
|
||||
logger.info(
|
||||
"wait_for_compute_zone_operation: "
|
||||
f"Operation {operation['name']} finished."
|
||||
)
|
||||
break
|
||||
|
||||
time.sleep(poll_interval)
|
||||
|
||||
return result
|
||||
|
||||
def list_instances(
|
||||
self,
|
||||
label_filters: Optional[dict] = None,
|
||||
is_terminated: bool = False,
|
||||
) -> List[GCPComputeNode]:
|
||||
label_filters = label_filters or {}
|
||||
|
||||
if label_filters:
|
||||
label_filter_expr = (
|
||||
"("
|
||||
+ " AND ".join(
|
||||
[
|
||||
"(labels.{key} = {value})".format(key=key, value=value)
|
||||
for key, value in label_filters.items()
|
||||
]
|
||||
)
|
||||
+ ")"
|
||||
)
|
||||
else:
|
||||
label_filter_expr = ""
|
||||
|
||||
statuses = (
|
||||
GCPComputeNode.TERMINATED_STATUSES
|
||||
if is_terminated
|
||||
else GCPComputeNode.NON_TERMINATED_STATUSES
|
||||
)
|
||||
|
||||
instance_state_filter_expr = (
|
||||
"("
|
||||
+ " OR ".join(
|
||||
["(status = {status})".format(status=status) for status in statuses]
|
||||
)
|
||||
+ ")"
|
||||
)
|
||||
|
||||
cluster_name_filter_expr = "(labels.{key} = {value})".format(
|
||||
key=TAG_RAY_CLUSTER_NAME, value=self.cluster_name
|
||||
)
|
||||
|
||||
# TPU VMs spawn accompanying Compute Instances that must be filtered out,
|
||||
# else this results in duplicated nodes.
|
||||
tpu_negation_filter_expr = "(NOT labels.{label}:*)".format(label="tpu_cores")
|
||||
|
||||
not_empty_filters = [
|
||||
f
|
||||
for f in [
|
||||
label_filter_expr,
|
||||
instance_state_filter_expr,
|
||||
cluster_name_filter_expr,
|
||||
tpu_negation_filter_expr,
|
||||
]
|
||||
if f
|
||||
]
|
||||
|
||||
filter_expr = " AND ".join(not_empty_filters)
|
||||
|
||||
response = (
|
||||
self.resource.instances()
|
||||
.list(
|
||||
project=self.project_id,
|
||||
zone=self.availability_zone,
|
||||
filter=filter_expr,
|
||||
)
|
||||
.execute(http=self.get_new_authorized_http(self.resource._http))
|
||||
)
|
||||
|
||||
instances = response.get("items", [])
|
||||
return [GCPComputeNode(i, self) for i in instances]
|
||||
|
||||
def get_instance(self, node_id: str) -> GCPComputeNode:
|
||||
instance = (
|
||||
self.resource.instances()
|
||||
.get(
|
||||
project=self.project_id,
|
||||
zone=self.availability_zone,
|
||||
instance=node_id,
|
||||
)
|
||||
.execute()
|
||||
)
|
||||
|
||||
return GCPComputeNode(instance, self)
|
||||
|
||||
def set_labels(
|
||||
self, node: GCPComputeNode, labels: dict, wait_for_operation: bool = True
|
||||
) -> dict:
|
||||
body = {
|
||||
"labels": dict(node["labels"], **labels),
|
||||
"labelFingerprint": node["labelFingerprint"],
|
||||
}
|
||||
node_id = node["name"]
|
||||
operation = (
|
||||
self.resource.instances()
|
||||
.setLabels(
|
||||
project=self.project_id,
|
||||
zone=self.availability_zone,
|
||||
instance=node_id,
|
||||
body=body,
|
||||
)
|
||||
.execute(http=self.get_new_authorized_http(self.resource._http))
|
||||
)
|
||||
|
||||
if wait_for_operation:
|
||||
result = self.wait_for_operation(operation)
|
||||
else:
|
||||
result = operation
|
||||
|
||||
return result
|
||||
|
||||
def _convert_resources_to_urls(
|
||||
self, configuration_dict: Dict[str, Any]
|
||||
) -> Dict[str, Any]:
|
||||
"""Ensures that resources are in their full URL form.
|
||||
|
||||
GCP expects machineType and acceleratorType to be a full URL (e.g.
|
||||
`zones/us-west1/machineTypes/n1-standard-2`) instead of just the
|
||||
type (`n1-standard-2`)
|
||||
|
||||
Args:
|
||||
configuration_dict: Dict of options that will be passed to GCP
|
||||
Returns:
|
||||
Input dictionary, but with possibly expanding `machineType` and
|
||||
`acceleratorType`.
|
||||
"""
|
||||
configuration_dict = deepcopy(configuration_dict)
|
||||
existing_machine_type = configuration_dict["machineType"]
|
||||
if not re.search(".*/machineTypes/.*", existing_machine_type):
|
||||
configuration_dict[
|
||||
"machineType"
|
||||
] = "zones/{zone}/machineTypes/{machine_type}".format(
|
||||
zone=self.availability_zone,
|
||||
machine_type=configuration_dict["machineType"],
|
||||
)
|
||||
|
||||
for accelerator in configuration_dict.get("guestAccelerators", []):
|
||||
gpu_type = accelerator["acceleratorType"]
|
||||
if not re.search(".*/acceleratorTypes/.*", gpu_type):
|
||||
accelerator[
|
||||
"acceleratorType"
|
||||
] = "projects/{project}/zones/{zone}/acceleratorTypes/{accelerator}".format( # noqa: E501
|
||||
project=self.project_id,
|
||||
zone=self.availability_zone,
|
||||
accelerator=gpu_type,
|
||||
)
|
||||
|
||||
return configuration_dict
|
||||
|
||||
def create_instance(
|
||||
self, base_config: dict, labels: dict, wait_for_operation: bool = True
|
||||
) -> Tuple[dict, str]:
|
||||
|
||||
config = self._convert_resources_to_urls(base_config)
|
||||
# removing TPU-specific default key set in config.py
|
||||
config.pop("networkConfig", None)
|
||||
name = _generate_node_name(labels, GCPNodeType.COMPUTE.value)
|
||||
|
||||
labels = dict(config.get("labels", {}), **labels)
|
||||
|
||||
config.update(
|
||||
{
|
||||
"labels": dict(labels, **{TAG_RAY_CLUSTER_NAME: self.cluster_name}),
|
||||
"name": name,
|
||||
}
|
||||
)
|
||||
|
||||
# Allow Google Compute Engine instance templates.
|
||||
#
|
||||
# Config example:
|
||||
#
|
||||
# ...
|
||||
# node_config:
|
||||
# sourceInstanceTemplate: global/instanceTemplates/worker-16
|
||||
# machineType: e2-standard-16
|
||||
# ...
|
||||
#
|
||||
# node_config parameters override matching template parameters, if any.
|
||||
#
|
||||
# https://cloud.google.com/compute/docs/instance-templates
|
||||
# https://cloud.google.com/compute/docs/reference/rest/v1/instances/insert
|
||||
source_instance_template = config.pop("sourceInstanceTemplate", None)
|
||||
|
||||
operation = (
|
||||
self.resource.instances()
|
||||
.insert(
|
||||
project=self.project_id,
|
||||
zone=self.availability_zone,
|
||||
sourceInstanceTemplate=source_instance_template,
|
||||
body=config,
|
||||
)
|
||||
.execute(http=self.get_new_authorized_http(self.resource._http))
|
||||
)
|
||||
|
||||
if wait_for_operation:
|
||||
result = self.wait_for_operation(operation)
|
||||
else:
|
||||
result = operation
|
||||
|
||||
return result, name
|
||||
|
||||
def delete_instance(self, node_id: str, wait_for_operation: bool = True) -> dict:
|
||||
operation = (
|
||||
self.resource.instances()
|
||||
.delete(
|
||||
project=self.project_id,
|
||||
zone=self.availability_zone,
|
||||
instance=node_id,
|
||||
)
|
||||
.execute()
|
||||
)
|
||||
|
||||
if wait_for_operation:
|
||||
result = self.wait_for_operation(operation)
|
||||
else:
|
||||
result = operation
|
||||
|
||||
return result
|
||||
|
||||
def stop_instance(self, node_id: str, wait_for_operation: bool = True) -> dict:
|
||||
operation = (
|
||||
self.resource.instances()
|
||||
.stop(
|
||||
project=self.project_id,
|
||||
zone=self.availability_zone,
|
||||
instance=node_id,
|
||||
)
|
||||
.execute()
|
||||
)
|
||||
|
||||
if wait_for_operation:
|
||||
result = self.wait_for_operation(operation)
|
||||
else:
|
||||
result = operation
|
||||
return result
|
||||
|
||||
def start_instance(self, node_id: str, wait_for_operation: bool = True) -> dict:
|
||||
|
||||
operation = (
|
||||
self.resource.instances()
|
||||
.start(
|
||||
project=self.project_id,
|
||||
zone=self.availability_zone,
|
||||
instance=node_id,
|
||||
)
|
||||
.execute(http=self.get_new_authorized_http(self.resource._http))
|
||||
)
|
||||
|
||||
if wait_for_operation:
|
||||
result = self.wait_for_operation(operation)
|
||||
else:
|
||||
result = operation
|
||||
return result
|
||||
|
||||
|
||||
class GCPTPU(GCPResource):
|
||||
"""Abstraction around GCP TPU resource"""
|
||||
|
||||
# node names already contain the path, but this is required for `parent`
|
||||
# arguments
|
||||
@property
|
||||
def path(self):
|
||||
return f"projects/{self.project_id}/locations/{self.availability_zone}"
|
||||
|
||||
def get_new_authorized_http(self, http: AuthorizedHttp) -> AuthorizedHttp:
|
||||
"""Generate a new AuthorizedHttp object with the given credentials."""
|
||||
new_http = AuthorizedHttp(http.credentials, http=httplib2.Http())
|
||||
return new_http
|
||||
|
||||
def wait_for_operation(
|
||||
self,
|
||||
operation: dict,
|
||||
max_polls: int = MAX_POLLS_TPU,
|
||||
poll_interval: int = POLL_INTERVAL,
|
||||
) -> dict:
|
||||
"""Poll for TPU operation until finished."""
|
||||
logger.info(
|
||||
"wait_for_tpu_operation: "
|
||||
f"Waiting for operation {operation['name']} to finish..."
|
||||
)
|
||||
|
||||
for _ in range(max_polls):
|
||||
result = (
|
||||
self.resource.projects()
|
||||
.locations()
|
||||
.operations()
|
||||
.get(name=f"{operation['name']}")
|
||||
.execute(http=self.get_new_authorized_http(self.resource._http))
|
||||
)
|
||||
if "error" in result:
|
||||
raise Exception(result["error"])
|
||||
|
||||
if "response" in result:
|
||||
logger.info(
|
||||
"wait_for_tpu_operation: "
|
||||
f"Operation {operation['name']} finished."
|
||||
)
|
||||
break
|
||||
|
||||
time.sleep(poll_interval)
|
||||
|
||||
return result
|
||||
|
||||
def list_instances(
|
||||
self,
|
||||
label_filters: Optional[dict] = None,
|
||||
is_terminated: bool = False,
|
||||
) -> List[GCPTPUNode]:
|
||||
response = (
|
||||
self.resource.projects()
|
||||
.locations()
|
||||
.nodes()
|
||||
.list(parent=self.path)
|
||||
.execute(http=self.get_new_authorized_http(self.resource._http))
|
||||
)
|
||||
|
||||
instances = response.get("nodes", [])
|
||||
instances = [GCPTPUNode(i, self) for i in instances]
|
||||
|
||||
# filter_expr cannot be passed directly to API
|
||||
# so we need to filter the results ourselves
|
||||
|
||||
# same logic as in GCPCompute.list_instances
|
||||
label_filters = label_filters or {}
|
||||
label_filters[TAG_RAY_CLUSTER_NAME] = self.cluster_name
|
||||
|
||||
def filter_instance(instance: GCPTPUNode) -> bool:
|
||||
if instance.is_terminated():
|
||||
return False
|
||||
|
||||
labels = instance.get_labels()
|
||||
if label_filters:
|
||||
for key, value in label_filters.items():
|
||||
if key not in labels:
|
||||
return False
|
||||
if value != labels[key]:
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
instances = list(filter(filter_instance, instances))
|
||||
|
||||
return instances
|
||||
|
||||
def get_instance(self, node_id: str) -> GCPTPUNode:
|
||||
instance = (
|
||||
self.resource.projects()
|
||||
.locations()
|
||||
.nodes()
|
||||
.get(name=node_id)
|
||||
.execute(http=self.get_new_authorized_http(self.resource._http))
|
||||
)
|
||||
|
||||
return GCPTPUNode(instance, self)
|
||||
|
||||
# this sometimes fails without a clear reason, so we retry it
|
||||
# MAX_POLLS times
|
||||
@_retry_on_exception(HttpError, "unable to queue the operation")
|
||||
def set_labels(
|
||||
self, node: GCPTPUNode, labels: dict, wait_for_operation: bool = True
|
||||
) -> dict:
|
||||
body = {
|
||||
"labels": dict(node["labels"], **labels),
|
||||
}
|
||||
update_mask = "labels"
|
||||
|
||||
operation = (
|
||||
self.resource.projects()
|
||||
.locations()
|
||||
.nodes()
|
||||
.patch(
|
||||
name=node["name"],
|
||||
updateMask=update_mask,
|
||||
body=body,
|
||||
)
|
||||
.execute(http=self.get_new_authorized_http(self.resource._http))
|
||||
)
|
||||
|
||||
if wait_for_operation:
|
||||
result = self.wait_for_operation(operation)
|
||||
else:
|
||||
result = operation
|
||||
|
||||
return result
|
||||
|
||||
def create_instance(
|
||||
self, base_config: dict, labels: dict, wait_for_operation: bool = True
|
||||
) -> Tuple[dict, str]:
|
||||
config = base_config.copy()
|
||||
# removing Compute-specific default key set in config.py
|
||||
config.pop("networkInterfaces", None)
|
||||
name = _generate_node_name(labels, GCPNodeType.TPU.value)
|
||||
|
||||
labels = dict(config.get("labels", {}), **labels)
|
||||
|
||||
config.update(
|
||||
{
|
||||
"labels": dict(labels, **{TAG_RAY_CLUSTER_NAME: self.cluster_name}),
|
||||
}
|
||||
)
|
||||
|
||||
if "networkConfig" not in config:
|
||||
config["networkConfig"] = {}
|
||||
if "enableExternalIps" not in config["networkConfig"]:
|
||||
# this is required for SSH to work, per google documentation
|
||||
# https://cloud.google.com/tpu/docs/users-guide-tpu-vm#create-curl
|
||||
config["networkConfig"]["enableExternalIps"] = True
|
||||
|
||||
# replace serviceAccounts with serviceAccount, and scopes with scope
|
||||
# this is necessary for the head node to work
|
||||
# see here: https://tpu.googleapis.com/$discovery/rest?version=v2alpha1
|
||||
if "serviceAccounts" in config:
|
||||
config["serviceAccount"] = config.pop("serviceAccounts")[0]
|
||||
config["serviceAccount"]["scope"] = config["serviceAccount"].pop("scopes")
|
||||
|
||||
operation = (
|
||||
self.resource.projects()
|
||||
.locations()
|
||||
.nodes()
|
||||
.create(
|
||||
parent=self.path,
|
||||
body=config,
|
||||
nodeId=name,
|
||||
)
|
||||
.execute(http=self.get_new_authorized_http(self.resource._http))
|
||||
)
|
||||
|
||||
if wait_for_operation:
|
||||
result = self.wait_for_operation(operation)
|
||||
else:
|
||||
result = operation
|
||||
|
||||
return result, name
|
||||
|
||||
def delete_instance(self, node_id: str, wait_for_operation: bool = True) -> dict:
|
||||
operation = (
|
||||
self.resource.projects()
|
||||
.locations()
|
||||
.nodes()
|
||||
.delete(name=node_id)
|
||||
.execute(http=self.get_new_authorized_http(self.resource._http))
|
||||
)
|
||||
|
||||
# No need to increase MAX_POLLS for deletion
|
||||
if wait_for_operation:
|
||||
result = self.wait_for_operation(operation, max_polls=MAX_POLLS)
|
||||
else:
|
||||
result = operation
|
||||
|
||||
return result
|
||||
|
||||
def stop_instance(self, node_id: str, wait_for_operation: bool = True) -> dict:
|
||||
operation = (
|
||||
self.resource.projects()
|
||||
.locations()
|
||||
.nodes()
|
||||
.stop(name=node_id)
|
||||
.execute(http=self.get_new_authorized_http(self.resource._http))
|
||||
)
|
||||
|
||||
if wait_for_operation:
|
||||
result = self.wait_for_operation(operation, max_polls=MAX_POLLS)
|
||||
else:
|
||||
result = operation
|
||||
|
||||
return result
|
||||
|
||||
def start_instance(self, node_id: str, wait_for_operation: bool = True) -> dict:
|
||||
operation = (
|
||||
self.resource.projects()
|
||||
.locations()
|
||||
.nodes()
|
||||
.start(name=node_id)
|
||||
.execute(http=self.get_new_authorized_http(self.resource._http))
|
||||
)
|
||||
|
||||
if wait_for_operation:
|
||||
result = self.wait_for_operation(operation, max_polls=MAX_POLLS)
|
||||
else:
|
||||
result = operation
|
||||
|
||||
return result
|
||||
@@ -0,0 +1,350 @@
|
||||
import copy
|
||||
import logging
|
||||
import time
|
||||
from functools import wraps
|
||||
from threading import RLock
|
||||
from types import ModuleType
|
||||
from typing import Any, Dict, List, Optional, Tuple
|
||||
|
||||
import googleapiclient
|
||||
|
||||
from ray.autoscaler._private.gcp.config import (
|
||||
bootstrap_gcp,
|
||||
construct_clients_from_provider_config,
|
||||
get_node_type,
|
||||
tpu_accelerator_config_to_type,
|
||||
)
|
||||
|
||||
# The logic has been abstracted away here to allow for different GCP resources
|
||||
# (API endpoints), which can differ widely, making it impossible to use
|
||||
# the same logic for everything.
|
||||
from ray.autoscaler._private.gcp.node import (
|
||||
GCPTPU, # noqa
|
||||
GCPCompute,
|
||||
GCPNode,
|
||||
GCPNodeType,
|
||||
GCPResource,
|
||||
)
|
||||
from ray.autoscaler._private.gcp.tpu_command_runner import TPUCommandRunner
|
||||
from ray.autoscaler.command_runner import CommandRunnerInterface
|
||||
from ray.autoscaler.node_provider import NodeProvider
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _retry(method, max_tries=5, backoff_s=1):
|
||||
"""Retry decorator for methods of GCPNodeProvider.
|
||||
|
||||
Upon catching BrokenPipeError, API clients are rebuilt and
|
||||
decorated methods are retried.
|
||||
|
||||
Work-around for https://github.com/ray-project/ray/issues/16072.
|
||||
Based on https://github.com/kubeflow/pipelines/pull/5250/files.
|
||||
"""
|
||||
|
||||
@wraps(method)
|
||||
def method_with_retries(self, *args, **kwargs):
|
||||
try_count = 0
|
||||
while try_count < max_tries:
|
||||
try:
|
||||
return method(self, *args, **kwargs)
|
||||
except BrokenPipeError:
|
||||
logger.warning("Caught a BrokenPipeError. Retrying.")
|
||||
try_count += 1
|
||||
if try_count < max_tries:
|
||||
self._construct_clients()
|
||||
time.sleep(backoff_s)
|
||||
else:
|
||||
raise
|
||||
|
||||
return method_with_retries
|
||||
|
||||
|
||||
class GCPNodeProvider(NodeProvider):
|
||||
def __init__(self, provider_config: dict, cluster_name: str):
|
||||
NodeProvider.__init__(self, provider_config, cluster_name)
|
||||
self.lock = RLock()
|
||||
self._construct_clients()
|
||||
self.cache_stopped_nodes = provider_config.get("cache_stopped_nodes", False)
|
||||
|
||||
# Cache of node objects from the last nodes() call. This avoids
|
||||
# excessive DescribeInstances requests.
|
||||
self.cached_nodes: Dict[str, GCPNode] = {}
|
||||
|
||||
def _construct_clients(self):
|
||||
_, _, compute, tpu = construct_clients_from_provider_config(
|
||||
self.provider_config
|
||||
)
|
||||
|
||||
# Dict of different resources provided by GCP.
|
||||
# At this moment - Compute and TPUs
|
||||
self.resources: Dict[GCPNodeType, GCPResource] = {}
|
||||
|
||||
# Compute is always required
|
||||
self.resources[GCPNodeType.COMPUTE] = GCPCompute(
|
||||
compute,
|
||||
self.provider_config["project_id"],
|
||||
self.provider_config["availability_zone"],
|
||||
self.cluster_name,
|
||||
)
|
||||
|
||||
# if there are no TPU nodes defined in config, tpu will be None.
|
||||
if tpu is not None:
|
||||
self.resources[GCPNodeType.TPU] = GCPTPU(
|
||||
tpu,
|
||||
self.provider_config["project_id"],
|
||||
self.provider_config["availability_zone"],
|
||||
self.cluster_name,
|
||||
)
|
||||
|
||||
def _get_resource_depending_on_node_name(self, node_name: str) -> GCPResource:
|
||||
"""Return the resource responsible for the node, based on node_name.
|
||||
|
||||
This expects the name to be in format '[NAME]-[UUID]-[TYPE]',
|
||||
where [TYPE] is either 'compute' or 'tpu' (see ``GCPNodeType``).
|
||||
"""
|
||||
return self.resources[GCPNodeType.name_to_type(node_name)]
|
||||
|
||||
@_retry
|
||||
def non_terminated_nodes(self, tag_filters: dict):
|
||||
with self.lock:
|
||||
instances = []
|
||||
|
||||
for resource in self.resources.values():
|
||||
node_instances = resource.list_instances(tag_filters)
|
||||
instances += node_instances
|
||||
|
||||
# Note: All the operations use "name" as the unique instance id
|
||||
self.cached_nodes = {i["name"]: i for i in instances}
|
||||
return [i["name"] for i in instances]
|
||||
|
||||
def is_running(self, node_id: str):
|
||||
with self.lock:
|
||||
node = self._get_cached_node(node_id)
|
||||
return node.is_running()
|
||||
|
||||
def is_terminated(self, node_id: str):
|
||||
with self.lock:
|
||||
node = self._get_cached_node(node_id)
|
||||
return node.is_terminated()
|
||||
|
||||
def node_tags(self, node_id: str):
|
||||
with self.lock:
|
||||
node = self._get_cached_node(node_id)
|
||||
return node.get_labels()
|
||||
|
||||
@_retry
|
||||
def set_node_tags(self, node_id: str, tags: dict):
|
||||
with self.lock:
|
||||
labels = tags
|
||||
node = self._get_node(node_id)
|
||||
|
||||
resource = self._get_resource_depending_on_node_name(node_id)
|
||||
|
||||
result = resource.set_labels(node=node, labels=labels)
|
||||
|
||||
return result
|
||||
|
||||
def external_ip(self, node_id: str):
|
||||
with self.lock:
|
||||
node = self._get_cached_node(node_id)
|
||||
|
||||
ip = node.get_external_ip()
|
||||
if ip is None:
|
||||
node = self._get_node(node_id)
|
||||
ip = node.get_external_ip()
|
||||
|
||||
return ip
|
||||
|
||||
def internal_ip(self, node_id: str):
|
||||
with self.lock:
|
||||
node = self._get_cached_node(node_id)
|
||||
|
||||
ip = node.get_internal_ip()
|
||||
if ip is None:
|
||||
node = self._get_node(node_id)
|
||||
ip = node.get_internal_ip()
|
||||
|
||||
return ip
|
||||
|
||||
@_retry
|
||||
def create_node(self, base_config: dict, tags: dict, count: int) -> Dict[str, dict]:
|
||||
"""Creates instances.
|
||||
|
||||
Returns dict mapping instance id to each create operation result for the created
|
||||
instances.
|
||||
"""
|
||||
with self.lock:
|
||||
labels = tags # gcp uses "labels" instead of aws "tags"
|
||||
|
||||
node_type = get_node_type(base_config)
|
||||
resource = self.resources[node_type]
|
||||
|
||||
all_nodes = {}
|
||||
if self.cache_stopped_nodes:
|
||||
filters = {
|
||||
"ray-node-name": labels["ray-node-name"],
|
||||
"ray-node-type": labels["ray-node-type"],
|
||||
"ray-user-node-type": labels["ray-user-node-type"],
|
||||
}
|
||||
reuse_nodes = resource.list_instances(filters, True)[:count]
|
||||
if reuse_nodes:
|
||||
reused_nodes_dict = {
|
||||
n["name"]: resource.start_instance(n["name"])
|
||||
for n in reuse_nodes
|
||||
}
|
||||
all_nodes.update(reused_nodes_dict)
|
||||
count -= len(reuse_nodes)
|
||||
|
||||
if count > 0:
|
||||
results: List[Tuple[dict, str]] = resource.create_instances(
|
||||
base_config, labels, count
|
||||
)
|
||||
|
||||
created_nodes_dict = {
|
||||
instance_id: result for result, instance_id in results
|
||||
}
|
||||
all_nodes.update(created_nodes_dict)
|
||||
|
||||
return all_nodes
|
||||
|
||||
def _thread_unsafe_terminate_node(self, node_id: str):
|
||||
# Assumes the global lock is held for the duration of this operation.
|
||||
# The lock may be held by a different thread if in `terminate_nodes()` case.
|
||||
logger.info("NodeProvider: {}: Terminating node".format(node_id))
|
||||
resource = self._get_resource_depending_on_node_name(node_id)
|
||||
try:
|
||||
result = resource.delete_instance(
|
||||
node_id=node_id,
|
||||
)
|
||||
except googleapiclient.errors.HttpError as http_error:
|
||||
if http_error.resp.status == 404:
|
||||
logger.warning(
|
||||
f"Tried to delete the node with id {node_id} "
|
||||
"but it was already gone."
|
||||
)
|
||||
result = None
|
||||
else:
|
||||
raise http_error from None
|
||||
return result
|
||||
|
||||
@_retry
|
||||
def terminate_node(self, node_id: str):
|
||||
with self.lock:
|
||||
resource = self._get_resource_depending_on_node_name(node_id)
|
||||
try:
|
||||
if self.cache_stopped_nodes:
|
||||
node = self._get_cached_node(node_id)
|
||||
if node.is_running():
|
||||
result = resource.stop_instance(node_id=node_id)
|
||||
else:
|
||||
result = None
|
||||
else:
|
||||
result = resource.delete_instance(
|
||||
node_id=node_id,
|
||||
)
|
||||
except googleapiclient.errors.HttpError as http_error:
|
||||
if http_error.resp.status == 404:
|
||||
logger.warning(
|
||||
f"Tried to delete the node with id {node_id} "
|
||||
"but it was already gone."
|
||||
)
|
||||
else:
|
||||
raise http_error from None
|
||||
return result
|
||||
|
||||
@_retry
|
||||
def _get_node(self, node_id: str) -> GCPNode:
|
||||
self.non_terminated_nodes({}) # Side effect: updates cache
|
||||
|
||||
with self.lock:
|
||||
if node_id in self.cached_nodes:
|
||||
return self.cached_nodes[node_id]
|
||||
|
||||
resource = self._get_resource_depending_on_node_name(node_id)
|
||||
instance = resource.get_instance(node_id=node_id)
|
||||
|
||||
return instance
|
||||
|
||||
def _get_cached_node(self, node_id: str) -> GCPNode:
|
||||
if node_id in self.cached_nodes:
|
||||
return self.cached_nodes[node_id]
|
||||
|
||||
return self._get_node(node_id)
|
||||
|
||||
@staticmethod
|
||||
def bootstrap_config(cluster_config):
|
||||
return bootstrap_gcp(cluster_config)
|
||||
|
||||
@staticmethod
|
||||
def fillout_available_node_types_resources(
|
||||
cluster_config: Dict[str, Any]
|
||||
) -> Dict[str, Any]:
|
||||
"""Fill out TPU resources to the cluster config.
|
||||
|
||||
To enable TPU pod autoscaling, we provide the TPU accelerator
|
||||
type as a resource that only exists on worker 0 of the pod slice.
|
||||
For instance, a v4-16 should have the resource labels:
|
||||
worker 0: resources = {"TPU": 4, "TPU-v4-16-head": 1}
|
||||
worker 1: resources = {"TPU": 4}
|
||||
|
||||
For the autoscaler to correctly process the demands of
|
||||
creating a new TPU pod, then the autoscaler must know what
|
||||
a TPU pod is in the form of the TPU accelerator resource.
|
||||
|
||||
Therefore we fill out TPU pods appropriately by providing the
|
||||
expected resource which we can deduce from the cluster config.
|
||||
|
||||
"""
|
||||
if "available_node_types" not in cluster_config:
|
||||
return cluster_config
|
||||
cluster_config = copy.deepcopy(cluster_config)
|
||||
available_node_types = cluster_config["available_node_types"]
|
||||
for node_type in available_node_types:
|
||||
node_config = available_node_types[node_type]["node_config"]
|
||||
if get_node_type(node_config) == GCPNodeType.TPU:
|
||||
autodetected_resources = {}
|
||||
accelerator_type = ""
|
||||
if "acceleratorType" in node_config:
|
||||
accelerator_type = node_config["acceleratorType"]
|
||||
elif "acceleratorConfig" in node_config:
|
||||
accelerator_type = tpu_accelerator_config_to_type(
|
||||
node_config["acceleratorConfig"]
|
||||
)
|
||||
if not accelerator_type:
|
||||
continue
|
||||
autodetected_resources[f"TPU-{accelerator_type}-head"] = 1
|
||||
available_node_types[node_type]["resources"].update(
|
||||
autodetected_resources
|
||||
)
|
||||
return cluster_config
|
||||
|
||||
def get_command_runner(
|
||||
self,
|
||||
log_prefix: str,
|
||||
node_id: str,
|
||||
auth_config: Dict[str, Any],
|
||||
cluster_name: str,
|
||||
process_runner: ModuleType,
|
||||
use_internal_ip: bool,
|
||||
docker_config: Optional[Dict[str, Any]] = None,
|
||||
) -> CommandRunnerInterface:
|
||||
"""Returns a TPU command runner as applicable."""
|
||||
resource = self._get_resource_depending_on_node_name(node_id)
|
||||
instance = resource.get_instance(node_id)
|
||||
common_args = {
|
||||
"docker_config": docker_config,
|
||||
"log_prefix": log_prefix,
|
||||
"node_id": node_id,
|
||||
"auth_config": auth_config,
|
||||
"cluster_name": cluster_name,
|
||||
"process_runner": process_runner,
|
||||
"use_internal_ip": use_internal_ip,
|
||||
}
|
||||
if (
|
||||
GCPNodeType.TPU in self.resources
|
||||
and resource == self.resources[GCPNodeType.TPU]
|
||||
):
|
||||
return TPUCommandRunner(instance=instance, provider=self, **common_args)
|
||||
else:
|
||||
return super().get_command_runner(**common_args)
|
||||
@@ -0,0 +1,329 @@
|
||||
"""Command runners specific to TPU VM pods.
|
||||
|
||||
TPU VM pods may contain multiple hosts, each including attached TPU chips and
|
||||
associated internal/external IP addresses.
|
||||
|
||||
To support TPU VM pods, we represent entire TPU pods as "Ray Nodes", meaning
|
||||
that TPU pods will need to run the operations specified in `CommandRunnerInterface`
|
||||
N times, where N denotes the number of hosts that comprise a TPU pod.
|
||||
|
||||
To maintain feature completeness, we simply wrap the existing `SSHCommandRunner` and
|
||||
`DockerCommandRunner` and run them as batched calls.
|
||||
|
||||
"""
|
||||
import copy
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from types import ModuleType
|
||||
from typing import Any, Dict, Optional
|
||||
|
||||
from ray._private import ray_constants
|
||||
from ray.autoscaler._private.command_runner import DockerCommandRunner, SSHCommandRunner
|
||||
from ray.autoscaler._private.gcp.node import GCPTPUNode
|
||||
from ray.autoscaler.command_runner import CommandRunnerInterface
|
||||
from ray.autoscaler.node_provider import NodeProvider
|
||||
|
||||
|
||||
class TPUVMSSHCommandRunner(SSHCommandRunner):
|
||||
"""An SSH command runner with overwritten IP address calls."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
internal_ip: str,
|
||||
external_ip: str,
|
||||
worker_id: int,
|
||||
accelerator_type: str,
|
||||
*args,
|
||||
**kwargs,
|
||||
):
|
||||
self._internal_ip = internal_ip
|
||||
self._external_ip = external_ip
|
||||
self._worker_id = worker_id
|
||||
self._accelerator_type = accelerator_type
|
||||
super().__init__(*args, **kwargs)
|
||||
|
||||
def _get_node_ip(self) -> str:
|
||||
if self.use_internal_ip:
|
||||
return self._internal_ip
|
||||
else:
|
||||
return self._external_ip
|
||||
|
||||
def run(
|
||||
self,
|
||||
cmd,
|
||||
timeout=120,
|
||||
exit_on_fail=False,
|
||||
port_forward=None,
|
||||
with_output=False,
|
||||
environment_variables: Dict[str, object] = None,
|
||||
run_env="auto", # Unused argument.
|
||||
ssh_options_override_ssh_key="",
|
||||
shutdown_after_run=False,
|
||||
) -> str:
|
||||
"""Override the SSH run for TPU VM pods.
|
||||
|
||||
Main functionality here we need to inject is to intercept the resources
|
||||
provided by the node_provider TPU node type fillout.
|
||||
|
||||
node_provider will provide a resource "TPU-{TPU_POD_TYPE}-head" which:
|
||||
1) allows application developers to target worker 0 of an arbitary TPU pod, and
|
||||
2) signals to the autoscaler how to address the demand for more TPU pods.
|
||||
|
||||
Without this intercept, then all workers of a TPU pod will have the
|
||||
"TPU-{TPU_POD_TYPE}-head" resource which will violate functionality (1)
|
||||
above.
|
||||
|
||||
"""
|
||||
|
||||
if environment_variables:
|
||||
environment_variables = _maybe_remove_head_resource(
|
||||
environment_variables, self._worker_id, self._accelerator_type
|
||||
)
|
||||
|
||||
return super().run(
|
||||
cmd=cmd,
|
||||
timeout=timeout,
|
||||
exit_on_fail=exit_on_fail,
|
||||
port_forward=port_forward,
|
||||
with_output=with_output,
|
||||
environment_variables=environment_variables,
|
||||
run_env=run_env,
|
||||
ssh_options_override_ssh_key=ssh_options_override_ssh_key,
|
||||
shutdown_after_run=shutdown_after_run,
|
||||
)
|
||||
|
||||
|
||||
class TPUVMDockerCommandRunner(DockerCommandRunner):
|
||||
"""A Docker command runner with overwritten IP addresses."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
docker_config: Dict[str, Any],
|
||||
internal_ip: str,
|
||||
external_ip: str,
|
||||
worker_id: int,
|
||||
accelerator_type: str,
|
||||
**common_args,
|
||||
):
|
||||
super().__init__(docker_config=docker_config, **common_args)
|
||||
self._worker_id = worker_id
|
||||
self._accelerator_type = accelerator_type
|
||||
|
||||
self.ssh_command_runner = TPUVMSSHCommandRunner(
|
||||
internal_ip=internal_ip,
|
||||
external_ip=external_ip,
|
||||
worker_id=worker_id,
|
||||
accelerator_type=accelerator_type,
|
||||
**common_args,
|
||||
)
|
||||
|
||||
def run(
|
||||
self,
|
||||
cmd,
|
||||
timeout=120,
|
||||
exit_on_fail=False,
|
||||
port_forward=None,
|
||||
with_output=False,
|
||||
environment_variables: Optional[Dict[str, object]] = None,
|
||||
run_env="auto",
|
||||
ssh_options_override_ssh_key="",
|
||||
shutdown_after_run=False,
|
||||
):
|
||||
if environment_variables:
|
||||
environment_variables = _maybe_remove_head_resource(
|
||||
environment_variables, self._worker_id, self._accelerator_type
|
||||
)
|
||||
return super().run(
|
||||
cmd,
|
||||
timeout,
|
||||
exit_on_fail,
|
||||
port_forward,
|
||||
with_output,
|
||||
environment_variables,
|
||||
run_env,
|
||||
ssh_options_override_ssh_key,
|
||||
shutdown_after_run,
|
||||
)
|
||||
|
||||
|
||||
class TPUCommandRunner(CommandRunnerInterface):
|
||||
"""A TPU pod command runner."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
instance: GCPTPUNode,
|
||||
log_prefix: str,
|
||||
node_id: str,
|
||||
auth_config: Dict[str, Any],
|
||||
provider: NodeProvider,
|
||||
cluster_name: str,
|
||||
process_runner: ModuleType,
|
||||
use_internal_ip: bool,
|
||||
docker_config: Optional[Dict[str, Any]] = None,
|
||||
):
|
||||
def create_command_runner(
|
||||
worker_id: int, accelerator_type: str, internal_ip: str, external_ip: str
|
||||
) -> CommandRunnerInterface:
|
||||
"""Returns the correct base command runner."""
|
||||
|
||||
common_args = {
|
||||
"internal_ip": internal_ip,
|
||||
"external_ip": external_ip,
|
||||
"worker_id": worker_id,
|
||||
"accelerator_type": accelerator_type,
|
||||
"log_prefix": "[tpu_worker_{}] ".format(worker_id) + log_prefix,
|
||||
"node_id": node_id,
|
||||
"provider": provider,
|
||||
"auth_config": auth_config,
|
||||
"cluster_name": cluster_name,
|
||||
"process_runner": process_runner,
|
||||
"use_internal_ip": use_internal_ip,
|
||||
}
|
||||
if docker_config and docker_config["container_name"] != "":
|
||||
return TPUVMDockerCommandRunner(
|
||||
docker_config=docker_config, **common_args
|
||||
)
|
||||
else:
|
||||
return TPUVMSSHCommandRunner(**common_args)
|
||||
|
||||
self._command_runners = []
|
||||
self._num_workers = instance.num_workers
|
||||
for i in range(self._num_workers):
|
||||
self._command_runners.append(
|
||||
create_command_runner(
|
||||
worker_id=i,
|
||||
accelerator_type=instance.get("acceleratorType"),
|
||||
internal_ip=instance.get_internal_ip(i),
|
||||
external_ip=instance.get_external_ip(i),
|
||||
)
|
||||
)
|
||||
|
||||
@property
|
||||
def num_connections(self) -> int:
|
||||
"""Return the number of active connections allowed at a time.
|
||||
|
||||
We occasionally see issues where too many concurrent connections may lead to
|
||||
failed SSH connections when there are too many TPU hosts.
|
||||
|
||||
We utilize this property to cap the maximum number of active connections
|
||||
at a time until a proper fix is found.
|
||||
|
||||
"""
|
||||
num_max_concurrent_active_connections = ray_constants.env_integer(
|
||||
ray_constants.RAY_TPU_MAX_CONCURRENT_CONNECTIONS_ENV_VAR, default=16
|
||||
)
|
||||
return min(self._num_workers, num_max_concurrent_active_connections)
|
||||
|
||||
def run(
|
||||
self,
|
||||
cmd,
|
||||
timeout=120,
|
||||
exit_on_fail=False,
|
||||
port_forward=None,
|
||||
with_output=False,
|
||||
environment_variables: Dict[str, object] = None,
|
||||
run_env="auto", # Unused argument.
|
||||
ssh_options_override_ssh_key="",
|
||||
shutdown_after_run=False,
|
||||
) -> str:
|
||||
with ThreadPoolExecutor(self.num_connections) as executor:
|
||||
results = executor.map(
|
||||
lambda i: self._command_runners[i].run(
|
||||
cmd=cmd,
|
||||
timeout=timeout,
|
||||
exit_on_fail=exit_on_fail,
|
||||
port_forward=port_forward,
|
||||
with_output=with_output,
|
||||
environment_variables=copy.deepcopy(environment_variables),
|
||||
run_env=run_env,
|
||||
ssh_options_override_ssh_key=ssh_options_override_ssh_key,
|
||||
shutdown_after_run=shutdown_after_run,
|
||||
),
|
||||
range(self._num_workers),
|
||||
)
|
||||
# Note: the `run` abstract function may return a string representing
|
||||
# representing command output, but this result is rarely used - especially
|
||||
# if the node is a worker (which a TPU pod is).
|
||||
# We return only the results from worker 0 which may not always be expected.
|
||||
return list(results)[0]
|
||||
|
||||
def run_rsync_up(self, *args, **kwargs) -> None:
|
||||
with ThreadPoolExecutor(self.num_connections) as executor:
|
||||
executor.map(
|
||||
lambda i: self._command_runners[i].run_rsync_up(*args, **kwargs),
|
||||
range(self._num_workers),
|
||||
)
|
||||
|
||||
def run_rsync_down(self, *args: Any, **kwargs: Any) -> None:
|
||||
"""Rsync files down from the cluster node.
|
||||
|
||||
Args:
|
||||
*args: Forwarded to each per-worker ``run_rsync_down`` call.
|
||||
Includes the (remote) source path and (local) target path.
|
||||
**kwargs: Forwarded to each per-worker ``run_rsync_down`` call.
|
||||
"""
|
||||
with ThreadPoolExecutor(self.num_connections) as executor:
|
||||
executor.map(
|
||||
lambda i: self._command_runners[i].run_rsync_down(*args, **kwargs),
|
||||
range(self._num_workers),
|
||||
)
|
||||
|
||||
def remote_shell_command_str(self) -> str:
|
||||
"""Return the command the user can use to open a shell."""
|
||||
# Note: this function is rarely used if the node is a worker.
|
||||
# We return only the results from worker 0 which may not always be expected.
|
||||
return self._command_runners[0].remote_shell_command_str()
|
||||
|
||||
def run_init(self, *args: Any, **kwargs: Any) -> Optional[bool]:
|
||||
"""Used to run extra initialization commands.
|
||||
|
||||
Args:
|
||||
*args: Forwarded to each per-worker ``run_init`` call. Includes
|
||||
``as_head``, ``file_mounts``, and ``sync_run_yet``.
|
||||
**kwargs: Forwarded to each per-worker ``run_init`` call.
|
||||
|
||||
Returns:
|
||||
Whether initialization is necessary on any worker.
|
||||
"""
|
||||
with ThreadPoolExecutor(self.num_connections) as executor:
|
||||
results = executor.map(
|
||||
lambda i: self._command_runners[i].run_init(*args, **kwargs),
|
||||
range(self._num_workers),
|
||||
)
|
||||
# Note: the `run_init` abstract function may return a bool representing
|
||||
# whether initialization is necessary, but this result is rarely used -
|
||||
# especially if the node is a worker (which a TPU pod is).
|
||||
# Here we return whether any workers require initialization, which may not be
|
||||
# the expected result.
|
||||
return any(results)
|
||||
|
||||
|
||||
def _maybe_remove_head_resource(
|
||||
environment_variables: Dict[str, Any], worker_id: int, accelerator_type: str
|
||||
):
|
||||
"""
|
||||
node_provider will provide a resource "TPU-{TPU_POD_TYPE}-head" which:
|
||||
1) allows application developers to target worker 0 of an arbitary TPU pod, and
|
||||
2) signals to the autoscaler how to address the demand for more TPU pods.
|
||||
|
||||
Without this intercept, then all workers of a TPU pod will have the
|
||||
"TPU-{TPU_POD_TYPE}-head" resource which will violate functionality (1)
|
||||
above.
|
||||
"""
|
||||
resources = environment_variables.get(
|
||||
ray_constants.RESOURCES_ENVIRONMENT_VARIABLE, None
|
||||
)
|
||||
|
||||
if resources:
|
||||
# For TPU pod support, we need to ensure that the
|
||||
# tpu pod resource type only propagates to worker 0.
|
||||
if worker_id != 0:
|
||||
tpu_pod_resource_type = f"TPU-{accelerator_type}-head"
|
||||
if tpu_pod_resource_type in resources:
|
||||
resources = copy.copy(resources)
|
||||
resources.pop(tpu_pod_resource_type, None)
|
||||
environment_variables = {
|
||||
**environment_variables,
|
||||
ray_constants.RESOURCES_ENVIRONMENT_VARIABLE: resources,
|
||||
}
|
||||
|
||||
return environment_variables
|
||||
@@ -0,0 +1,555 @@
|
||||
import decimal
|
||||
import json
|
||||
import logging
|
||||
import time
|
||||
from itertools import chain
|
||||
from typing import Any, Dict, Optional
|
||||
|
||||
import requests
|
||||
|
||||
from ray._private.label_utils import (
|
||||
validate_node_label_syntax,
|
||||
)
|
||||
from ray.autoscaler._private.constants import (
|
||||
DISABLE_LAUNCH_CONFIG_CHECK_KEY,
|
||||
DISABLE_NODE_UPDATERS_KEY,
|
||||
FOREGROUND_NODE_LAUNCH_KEY,
|
||||
WORKER_LIVENESS_CHECK_KEY,
|
||||
)
|
||||
from ray.autoscaler._private.kuberay import node_provider, utils
|
||||
from ray.autoscaler._private.util import validate_config
|
||||
from ray.util.debug import log_once
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
AUTOSCALER_OPTIONS_KEY = "autoscalerOptions"
|
||||
IDLE_SECONDS_KEY = "idleTimeoutSeconds"
|
||||
UPSCALING_KEY = "upscalingMode"
|
||||
UPSCALING_VALUE_AGGRESSIVE = "Aggressive"
|
||||
UPSCALING_VALUE_DEFAULT = "Default"
|
||||
UPSCALING_VALUE_CONSERVATIVE = "Conservative"
|
||||
|
||||
MAX_RAYCLUSTER_FETCH_TRIES = 5
|
||||
RAYCLUSTER_FETCH_RETRY_S = 5
|
||||
|
||||
GKE_TPU_TOPOLOGY_LABEL = "cloud.google.com/gke-tpu-topology"
|
||||
GKE_TPU_ACCELERATOR_LABEL = "cloud.google.com/gke-tpu-accelerator"
|
||||
|
||||
# Logical group name for the KubeRay head group.
|
||||
# Used as the name of the "head node type" by the autoscaler.
|
||||
_HEAD_GROUP_NAME = "headgroup"
|
||||
|
||||
|
||||
class AutoscalingConfigProducer:
|
||||
"""Produces an autoscaling config by reading data from the RayCluster CR.
|
||||
|
||||
Used to fetch the autoscaling config at the beginning of each autoscaler iteration.
|
||||
|
||||
In the context of Ray deployment on Kubernetes, the autoscaling config is an
|
||||
internal interface.
|
||||
|
||||
The autoscaling config carries the strict subset of RayCluster CR data required by
|
||||
the autoscaler to make scaling decisions; in particular, the autoscaling config does
|
||||
not carry pod configuration data.
|
||||
|
||||
This class is the only public object in this file.
|
||||
"""
|
||||
|
||||
def __init__(self, ray_cluster_name, ray_cluster_namespace):
|
||||
self.kubernetes_api_client = node_provider.KubernetesHttpApiClient(
|
||||
namespace=ray_cluster_namespace
|
||||
)
|
||||
self._ray_cr_path = f"rayclusters/{ray_cluster_name}"
|
||||
|
||||
def __call__(self):
|
||||
ray_cr = self._fetch_ray_cr_from_k8s_with_retries()
|
||||
autoscaling_config = _derive_autoscaling_config_from_ray_cr(ray_cr)
|
||||
return autoscaling_config
|
||||
|
||||
def _fetch_ray_cr_from_k8s_with_retries(self) -> Dict[str, Any]:
|
||||
"""Fetch the RayCluster CR by querying the K8s API server.
|
||||
|
||||
Retry on HTTPError for robustness, in particular to protect autoscaler
|
||||
initialization.
|
||||
"""
|
||||
for i in range(1, MAX_RAYCLUSTER_FETCH_TRIES + 1):
|
||||
try:
|
||||
return self.kubernetes_api_client.get(self._ray_cr_path)
|
||||
except requests.HTTPError as e:
|
||||
if i < MAX_RAYCLUSTER_FETCH_TRIES:
|
||||
logger.exception(
|
||||
"Failed to fetch RayCluster CR from K8s. Retrying."
|
||||
)
|
||||
time.sleep(RAYCLUSTER_FETCH_RETRY_S)
|
||||
else:
|
||||
raise e from None
|
||||
|
||||
# This branch is inaccessible. Raise to satisfy mypy.
|
||||
raise AssertionError
|
||||
|
||||
|
||||
def _derive_autoscaling_config_from_ray_cr(ray_cr: Dict[str, Any]) -> Dict[str, Any]:
|
||||
provider_config = _generate_provider_config(ray_cr["metadata"]["namespace"])
|
||||
|
||||
available_node_types = _generate_available_node_types_from_ray_cr_spec(
|
||||
ray_cr["spec"]
|
||||
)
|
||||
|
||||
# The autoscaler expects a global max workers field. We set it to the sum of
|
||||
# node type max workers.
|
||||
global_max_workers = sum(
|
||||
node_type["max_workers"] for node_type in available_node_types.values()
|
||||
)
|
||||
|
||||
# Legacy autoscaling fields carry no information but are required for compatibility.
|
||||
legacy_autoscaling_fields = _generate_legacy_autoscaling_config_fields()
|
||||
|
||||
# Process autoscaler options.
|
||||
autoscaler_options = ray_cr["spec"].get(AUTOSCALER_OPTIONS_KEY, {})
|
||||
if IDLE_SECONDS_KEY in autoscaler_options:
|
||||
idle_timeout_minutes = autoscaler_options[IDLE_SECONDS_KEY] / 60.0
|
||||
else:
|
||||
idle_timeout_minutes = 1.0
|
||||
|
||||
if autoscaler_options.get(UPSCALING_KEY) == UPSCALING_VALUE_CONSERVATIVE:
|
||||
upscaling_speed = 1 # Rate-limit upscaling if "Conservative" is set by user.
|
||||
# This elif is redudant but included for clarity.
|
||||
elif autoscaler_options.get(UPSCALING_KEY) == UPSCALING_VALUE_DEFAULT:
|
||||
upscaling_speed = 1000 # i.e. big, no rate-limiting by default
|
||||
# This elif is redudant but included for clarity.
|
||||
elif autoscaler_options.get(UPSCALING_KEY) == UPSCALING_VALUE_AGGRESSIVE:
|
||||
upscaling_speed = 1000
|
||||
else:
|
||||
upscaling_speed = 1000
|
||||
|
||||
autoscaling_config = {
|
||||
"provider": provider_config,
|
||||
"cluster_name": ray_cr["metadata"]["name"],
|
||||
"head_node_type": _HEAD_GROUP_NAME,
|
||||
"available_node_types": available_node_types,
|
||||
"max_workers": global_max_workers,
|
||||
# Should consider exposing `idleTimeoutMinutes` in the RayCluster CRD,
|
||||
# under an `autoscaling` field.
|
||||
"idle_timeout_minutes": idle_timeout_minutes,
|
||||
# Should consider exposing `upscalingSpeed` in the RayCluster CRD,
|
||||
# under an `autoscaling` field.
|
||||
"upscaling_speed": upscaling_speed,
|
||||
**legacy_autoscaling_fields,
|
||||
}
|
||||
|
||||
# Make sure the config is readable by the autoscaler.
|
||||
validate_config(autoscaling_config)
|
||||
|
||||
return autoscaling_config
|
||||
|
||||
|
||||
def _generate_provider_config(ray_cluster_namespace: str) -> Dict[str, Any]:
|
||||
"""Generates the `provider` field of the autoscaling config, which carries data
|
||||
required to instantiate the KubeRay node provider.
|
||||
"""
|
||||
return {
|
||||
"type": "kuberay",
|
||||
"namespace": ray_cluster_namespace,
|
||||
DISABLE_NODE_UPDATERS_KEY: True,
|
||||
DISABLE_LAUNCH_CONFIG_CHECK_KEY: True,
|
||||
FOREGROUND_NODE_LAUNCH_KEY: True,
|
||||
WORKER_LIVENESS_CHECK_KEY: False,
|
||||
}
|
||||
|
||||
|
||||
def _generate_legacy_autoscaling_config_fields() -> Dict[str, Any]:
|
||||
"""Generates legacy autoscaling config fields required for compatibiliy."""
|
||||
return {
|
||||
"file_mounts": {},
|
||||
"cluster_synced_files": [],
|
||||
"file_mounts_sync_continuously": False,
|
||||
"initialization_commands": [],
|
||||
"setup_commands": [],
|
||||
"head_setup_commands": [],
|
||||
"worker_setup_commands": [],
|
||||
"head_start_ray_commands": [],
|
||||
"worker_start_ray_commands": [],
|
||||
"auth": {},
|
||||
}
|
||||
|
||||
|
||||
def _generate_available_node_types_from_ray_cr_spec(
|
||||
ray_cr_spec: Dict[str, Any],
|
||||
) -> Dict[str, Any]:
|
||||
"""Formats autoscaler "available_node_types" field based on the Ray CR's group
|
||||
specs.
|
||||
"""
|
||||
headGroupSpec = ray_cr_spec["headGroupSpec"]
|
||||
return {
|
||||
_HEAD_GROUP_NAME: _node_type_from_group_spec(headGroupSpec, is_head=True),
|
||||
**{
|
||||
worker_group_spec["groupName"]: _node_type_from_group_spec(
|
||||
worker_group_spec, is_head=False
|
||||
)
|
||||
for worker_group_spec in ray_cr_spec["workerGroupSpecs"]
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def _node_type_from_group_spec(
|
||||
group_spec: Dict[str, Any], is_head: bool
|
||||
) -> Dict[str, Any]:
|
||||
"""Converts CR group spec to autoscaler node type."""
|
||||
group_name = _HEAD_GROUP_NAME if is_head else group_spec["groupName"]
|
||||
|
||||
if is_head:
|
||||
# The head node type has no workers because the head is not a worker.
|
||||
min_workers = max_workers = 0
|
||||
else:
|
||||
# `minReplicas` and `maxReplicas` are required fields for each workerGroupSpec.
|
||||
# numOfHosts specifies the number of workers per replica in KubeRay v1.1+.
|
||||
min_workers = group_spec["minReplicas"] * group_spec.get("numOfHosts", 1)
|
||||
max_workers = group_spec["maxReplicas"] * group_spec.get("numOfHosts", 1)
|
||||
|
||||
resources = _get_ray_resources_from_group_spec(group_spec, is_head)
|
||||
labels = _get_labels_from_group_spec(group_spec, group_name)
|
||||
|
||||
node_type = {
|
||||
"min_workers": min_workers,
|
||||
"max_workers": max_workers,
|
||||
# `node_config` is a legacy field required for compatibility.
|
||||
# Pod config data is required by the operator but not by the autoscaler.
|
||||
"node_config": {},
|
||||
"resources": resources,
|
||||
"labels": labels,
|
||||
}
|
||||
|
||||
idle_timeout_s = group_spec.get(IDLE_SECONDS_KEY)
|
||||
if idle_timeout_s is not None:
|
||||
node_type["idle_timeout_s"] = float(idle_timeout_s)
|
||||
|
||||
return node_type
|
||||
|
||||
|
||||
def _get_ray_resources_from_group_spec(
|
||||
group_spec: Dict[str, Any], is_head: bool
|
||||
) -> Dict[str, int]:
|
||||
"""
|
||||
Infers Ray resources from group `Resources` field, rayStartCommands, or K8s limits.
|
||||
The resources extracted are used in autoscaling calculations.
|
||||
"""
|
||||
# Set resources from top-level group 'Resources' field if it exists.
|
||||
group_resources = group_spec.get("resources", {})
|
||||
|
||||
ray_start_params = group_spec.get("rayStartParams", {})
|
||||
# In KubeRay, Ray container is always the first application container of a Ray Pod.
|
||||
k8s_resources = group_spec["template"]["spec"]["containers"][0].get("resources", {})
|
||||
group_name = _HEAD_GROUP_NAME if is_head else group_spec["groupName"]
|
||||
|
||||
num_cpus = _get_num_cpus(
|
||||
group_resources, ray_start_params, k8s_resources, group_name
|
||||
)
|
||||
num_gpus = _get_num_gpus(
|
||||
group_resources, ray_start_params, k8s_resources, group_name
|
||||
)
|
||||
custom_resource_dict = _get_custom_resources(
|
||||
group_resources, ray_start_params, group_name
|
||||
)
|
||||
num_tpus = _get_num_tpus(group_resources, custom_resource_dict, k8s_resources)
|
||||
memory = _get_memory(group_resources, ray_start_params, k8s_resources)
|
||||
|
||||
# It's not allowed to use object store memory as a resource request, so we don't
|
||||
# add that to the autoscaler's resources annotations.
|
||||
|
||||
resources = {}
|
||||
|
||||
assert isinstance(num_cpus, int)
|
||||
resources["CPU"] = num_cpus
|
||||
|
||||
if num_gpus is not None:
|
||||
resources["GPU"] = num_gpus
|
||||
|
||||
if num_tpus is not None:
|
||||
# Add TPU Ray resource if not already added by ray_start_params,
|
||||
# but specified in k8s_resource_limits.
|
||||
if "TPU" not in custom_resource_dict:
|
||||
resources["TPU"] = num_tpus
|
||||
|
||||
"""Add TPU head resource, similar to the GCP node_provider.
|
||||
Sets the Ray resource TPU-{...}-head to ensure the Ray autoscaler
|
||||
has sufficient resources to make scaling decisions.
|
||||
TPU worker groups treat each TPU podslice as a replica, with `NumOfHosts`
|
||||
specifying the number of workers per slice. Each replica of a TPU worker
|
||||
group has one TPU head.
|
||||
|
||||
For example, a v4-16 worker group with 2 replicas should have the following
|
||||
resource labels on worker 0 of each replica:
|
||||
worker 0: resources = {"TPU": 4, "TPU-v4-16-head": 1}
|
||||
"""
|
||||
if (
|
||||
"nodeSelector" in group_spec["template"]["spec"]
|
||||
and GKE_TPU_TOPOLOGY_LABEL in group_spec["template"]["spec"]["nodeSelector"]
|
||||
and GKE_TPU_ACCELERATOR_LABEL
|
||||
in group_spec["template"]["spec"]["nodeSelector"]
|
||||
):
|
||||
topology = group_spec["template"]["spec"]["nodeSelector"][
|
||||
GKE_TPU_TOPOLOGY_LABEL
|
||||
]
|
||||
accelerator = group_spec["template"]["spec"]["nodeSelector"][
|
||||
GKE_TPU_ACCELERATOR_LABEL
|
||||
]
|
||||
accelerator_type = utils.tpu_node_selectors_to_type(topology, accelerator)
|
||||
if accelerator_type:
|
||||
resources[f"TPU-{accelerator_type}-head"] = 1
|
||||
else:
|
||||
logger.error(
|
||||
f"Pods using TPUs require both `{GKE_TPU_TOPOLOGY_LABEL}` and `{GKE_TPU_ACCELERATOR_LABEL}` node selectors. "
|
||||
"See https://docs.ray.io/en/latest/cluster/kubernetes/user-guides/tpu.html#configuring-ray-pods-for-tpu-usage "
|
||||
"and https://cloud.google.com/kubernetes-engine/docs/how-to/tpus."
|
||||
)
|
||||
|
||||
if memory is not None:
|
||||
resources["memory"] = memory
|
||||
|
||||
resources.update(custom_resource_dict)
|
||||
|
||||
return resources
|
||||
|
||||
|
||||
def _get_labels_from_group_spec(
|
||||
group_spec: Dict[str, Any], group_name: str = ""
|
||||
) -> Dict[str, str]:
|
||||
"""
|
||||
Parses Ray node labels for the autoscaling config based on the following
|
||||
priority:
|
||||
1. Top-level `labels` field in the group spec.
|
||||
2. `labels` field in `rayStartParams`.
|
||||
|
||||
Args:
|
||||
group_spec: The group specification dictionary.
|
||||
group_name: The name of the group (used in warning messages).
|
||||
|
||||
Returns:
|
||||
A dictionary of labels for the node type.
|
||||
"""
|
||||
labels_dict = {}
|
||||
|
||||
ray_start_params = group_spec.get("rayStartParams", {})
|
||||
labels_str = ray_start_params.get("labels")
|
||||
# Use a unique log_once key per group to ensure each group's warning is shown.
|
||||
log_once_key = (
|
||||
f"raystartparams_labels_warning_{group_name}"
|
||||
if group_name
|
||||
else "raystartparams_labels_warning"
|
||||
)
|
||||
if labels_str and log_once(log_once_key):
|
||||
if group_name:
|
||||
logger.warning(
|
||||
f"Ignoring labels: {labels_str} set in rayStartParams for group "
|
||||
f"'{group_name}'. Group labels are supported in the top-level "
|
||||
"Labels field starting in KubeRay v1.5"
|
||||
)
|
||||
else:
|
||||
logger.warning(
|
||||
f"Ignoring labels: {labels_str} set in rayStartParams. "
|
||||
"Group labels are supported in the top-level Labels field "
|
||||
"starting in KubeRay v1.5"
|
||||
)
|
||||
|
||||
# Check for top-level structured Labels field.
|
||||
if "labels" in group_spec and isinstance(group_spec.get("labels"), dict):
|
||||
labels_dict = group_spec.get("labels")
|
||||
# Validate node labels follow expected Kubernetes label syntax.
|
||||
validate_node_label_syntax(labels_dict)
|
||||
|
||||
return labels_dict
|
||||
|
||||
|
||||
def _get_num_cpus(
|
||||
group_resources: Dict[str, str],
|
||||
ray_start_params: Dict[str, str],
|
||||
k8s_resources: Dict[str, Dict[str, str]],
|
||||
group_name: str,
|
||||
) -> int:
|
||||
"""Get CPU annotation from `resources` field, ray_start_params or k8s_resources,
|
||||
with priority for `resources` field.
|
||||
"""
|
||||
if "CPU" in group_resources:
|
||||
if "num-cpus" in ray_start_params:
|
||||
logger.warning(
|
||||
f"'CPU' specified in both the top-level 'resources' field and in 'rayStartParams'. "
|
||||
f"Using the value from 'resources': {group_resources['CPU']}."
|
||||
)
|
||||
return _round_up_k8s_quantity(group_resources["CPU"])
|
||||
if "num-cpus" in ray_start_params:
|
||||
return int(ray_start_params["num-cpus"])
|
||||
elif "cpu" in k8s_resources.get("limits", {}):
|
||||
cpu_quantity: str = k8s_resources["limits"]["cpu"]
|
||||
return _round_up_k8s_quantity(cpu_quantity)
|
||||
elif "cpu" in k8s_resources.get("requests", {}):
|
||||
cpu_quantity: str = k8s_resources["requests"]["cpu"]
|
||||
return _round_up_k8s_quantity(cpu_quantity)
|
||||
else:
|
||||
# Getting the number of CPUs is important, so raise an error if we can't do it.
|
||||
raise ValueError(
|
||||
f"Autoscaler failed to detect `CPU` resources for group {group_name}."
|
||||
"\nSet the `--num-cpus` rayStartParam and/or "
|
||||
"the CPU resource limit for the Ray container."
|
||||
)
|
||||
|
||||
|
||||
def _get_memory(
|
||||
group_resources: Dict[str, str],
|
||||
ray_start_params: Dict[str, str],
|
||||
k8s_resources: Dict[str, Dict[str, str]],
|
||||
) -> Optional[int]:
|
||||
"""Get memory resource annotation from `resources` field, ray_start_params or k8s_resources,
|
||||
with priority for `resources` field.
|
||||
"""
|
||||
if "memory" in group_resources:
|
||||
if "memory" in ray_start_params:
|
||||
logger.warning(
|
||||
f"'memory' specified in both the top-level 'resources' field and in 'rayStartParams'. "
|
||||
f"Using the value from 'resources': {group_resources['memory']}."
|
||||
)
|
||||
return _round_up_k8s_quantity(group_resources["memory"])
|
||||
if "memory" in ray_start_params:
|
||||
return int(ray_start_params["memory"])
|
||||
elif "memory" in k8s_resources.get("limits", {}):
|
||||
memory_quantity: str = k8s_resources["limits"]["memory"]
|
||||
return _round_up_k8s_quantity(memory_quantity)
|
||||
elif "memory" in k8s_resources.get("requests", {}):
|
||||
memory_quantity: str = k8s_resources["requests"]["memory"]
|
||||
return _round_up_k8s_quantity(memory_quantity)
|
||||
return None
|
||||
|
||||
|
||||
def _get_num_gpus(
|
||||
group_resources: Dict[str, str],
|
||||
ray_start_params: Dict[str, str],
|
||||
k8s_resources: Dict[str, Dict[str, str]],
|
||||
group_name: str,
|
||||
) -> Optional[int]:
|
||||
"""Get GPU resource annotation from `resources` field, ray_start_params or k8s_resources,
|
||||
with priority for `resources` field.
|
||||
"""
|
||||
if "GPU" in group_resources:
|
||||
if "num-gpus" in ray_start_params:
|
||||
logger.warning(
|
||||
f"'GPU' specified in both the top-level 'resources' field and in 'rayStartParams'. "
|
||||
f"Using the value from 'resources': {group_resources['GPU']}."
|
||||
)
|
||||
return _round_up_k8s_quantity(group_resources["GPU"])
|
||||
elif "num-gpus" in ray_start_params:
|
||||
return int(ray_start_params["num-gpus"])
|
||||
else:
|
||||
for key, resource_quantity in chain(
|
||||
k8s_resources.get("limits", {}).items(),
|
||||
k8s_resources.get("requests", {}).items(),
|
||||
):
|
||||
# e.g. nvidia.com/gpu
|
||||
if key.endswith("gpu"):
|
||||
# Typically, this is a string representing an interger, e.g. "1".
|
||||
# Convert to int, making no assumptions on the resource_quantity,
|
||||
# besides that it's valid as a K8s resource quantity.
|
||||
num_gpus = _round_up_k8s_quantity(resource_quantity)
|
||||
if num_gpus > 0:
|
||||
# Only one GPU type supported for now, break out on first
|
||||
# "/gpu" match.
|
||||
return num_gpus
|
||||
return None
|
||||
|
||||
|
||||
def _get_num_tpus(
|
||||
group_resources: Dict[str, str],
|
||||
custom_resource_dict: Dict[str, int],
|
||||
k8s_resources: Dict[str, Dict[str, str]],
|
||||
) -> Optional[int]:
|
||||
"""Get TPU custom resource annotation from `resources` field, custom_resource_dict in ray_start_params,
|
||||
or k8s_resources, with priority for `resources` field.
|
||||
"""
|
||||
if "TPU" in group_resources:
|
||||
return _round_up_k8s_quantity(group_resources["TPU"])
|
||||
elif "TPU" in custom_resource_dict:
|
||||
return custom_resource_dict["TPU"]
|
||||
else:
|
||||
for typ in ["limits", "requests"]:
|
||||
tpu_resource_quantity = k8s_resources.get(typ, {}).get("google.com/tpu")
|
||||
if tpu_resource_quantity is not None:
|
||||
# Typically, this is a string representing an integer, e.g. "1".
|
||||
# Convert to int, making no assumptions on the tpu_resource_quantity,
|
||||
# besides that it's valid as a K8s resource quantity.
|
||||
num_tpus = _round_up_k8s_quantity(tpu_resource_quantity)
|
||||
if num_tpus > 0:
|
||||
return num_tpus
|
||||
return None
|
||||
|
||||
|
||||
def _round_up_k8s_quantity(quantity: str) -> int:
|
||||
"""Rounds a Kubernetes resource quantity up to the nearest integer.
|
||||
|
||||
Args:
|
||||
quantity: Resource quantity as a string in the canonical K8s form.
|
||||
|
||||
Returns:
|
||||
The quantity, rounded up, as an integer.
|
||||
"""
|
||||
resource_decimal: decimal.Decimal = utils.parse_quantity(quantity)
|
||||
rounded = resource_decimal.to_integral_value(rounding=decimal.ROUND_UP)
|
||||
return int(rounded)
|
||||
|
||||
|
||||
def _get_custom_resources(
|
||||
group_resources: Dict[str, str], ray_start_params: Dict[str, Any], group_name: str
|
||||
) -> Dict[str, int]:
|
||||
"""Format custom resources based on the group `resources` field or `resources` Ray start param.
|
||||
|
||||
Currently, the value of the rayStartParam `resources` field must
|
||||
be formatted as follows:
|
||||
'"{\"Custom1\": 1, \"Custom2\": 5}"'.
|
||||
|
||||
This method first converts the input to a correctly formatted
|
||||
json string and then loads that json string to a dict.
|
||||
"""
|
||||
# If the top-level `resources` field is defined, use it as the exclusive source.
|
||||
if group_resources:
|
||||
if "resources" in ray_start_params:
|
||||
logger.warning(
|
||||
f"custom resources specified in both the top-level 'resources' field and in 'rayStartParams'. "
|
||||
f"Using the values from 'resources': {group_resources}."
|
||||
)
|
||||
standard_keys = {"CPU", "GPU", "TPU", "memory"}
|
||||
try:
|
||||
custom_resources = {
|
||||
k: _round_up_k8s_quantity(v)
|
||||
for k, v in group_resources.items()
|
||||
if k not in standard_keys
|
||||
}
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
f"Error reading `resource` for group {group_name}."
|
||||
" For the correct format, refer to example configuration at "
|
||||
"https://github.com/ray-project/ray/blob/master/python/"
|
||||
"ray/autoscaler/kuberay/ray-cluster.complete.yaml."
|
||||
)
|
||||
raise e
|
||||
return custom_resources
|
||||
|
||||
# Otherwise, check rayStartParams.
|
||||
if "resources" not in ray_start_params:
|
||||
return {}
|
||||
resources_string = ray_start_params["resources"]
|
||||
try:
|
||||
# Drop the extra pair of quotes and remove the backslash escapes.
|
||||
# resources_json should be a json string.
|
||||
resources_json = resources_string[1:-1].replace("\\", "")
|
||||
# Load a dict from the json string.
|
||||
resources = json.loads(resources_json)
|
||||
assert isinstance(resources, dict)
|
||||
for key, value in resources.items():
|
||||
assert isinstance(key, str)
|
||||
assert isinstance(value, int)
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
f"Error reading `resource` rayStartParam for group {group_name}."
|
||||
" For the correct format, refer to example configuration at "
|
||||
"https://github.com/ray-project/ray/blob/master/python/"
|
||||
"ray/autoscaler/kuberay/ray-cluster.complete.yaml."
|
||||
)
|
||||
raise e
|
||||
return resources
|
||||
@@ -0,0 +1,558 @@
|
||||
import datetime
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
from abc import ABC, abstractmethod
|
||||
from collections import defaultdict
|
||||
from typing import Any, Dict, List, Optional, Tuple, Union
|
||||
|
||||
import requests
|
||||
|
||||
from ray._common.network_utils import build_address
|
||||
from ray.autoscaler._private.constants import WORKER_LIVENESS_CHECK_KEY
|
||||
from ray.autoscaler._private.util import NodeID, NodeIP, NodeKind, NodeStatus, NodeType
|
||||
from ray.autoscaler.batching_node_provider import (
|
||||
BatchingNodeProvider,
|
||||
NodeData,
|
||||
ScaleRequest,
|
||||
)
|
||||
from ray.autoscaler.tags import (
|
||||
NODE_KIND_HEAD,
|
||||
NODE_KIND_WORKER,
|
||||
STATUS_UP_TO_DATE,
|
||||
STATUS_UPDATE_FAILED,
|
||||
TAG_RAY_USER_NODE_TYPE,
|
||||
)
|
||||
|
||||
# Key for KubeRay label that identifies a Ray pod as head or worker.
|
||||
KUBERAY_LABEL_KEY_KIND = "ray.io/node-type"
|
||||
# Key for KubeRay label that identifies the worker group (autoscaler node type) of a
|
||||
# Ray pod.
|
||||
KUBERAY_LABEL_KEY_TYPE = "ray.io/group"
|
||||
|
||||
# These should be synced with:
|
||||
# https://github.com/ray-project/kuberay/blob/f2d94ffe213dd8f69481b09c474047cb899fa73b/ray-operator/apis/ray/v1/raycluster_types.go#L165-L171 # noqa
|
||||
# Kind label value indicating the pod is the head.
|
||||
KUBERAY_KIND_HEAD = "head"
|
||||
# Kind label value indicating the pod is the worker.
|
||||
KUBERAY_KIND_WORKER = "worker"
|
||||
|
||||
# KubeRay CRD version
|
||||
KUBERAY_CRD_VER = os.getenv("KUBERAY_CRD_VER", "v1alpha1")
|
||||
|
||||
KUBERAY_REQUEST_TIMEOUT_S = int(os.getenv("KUBERAY_REQUEST_TIMEOUT_S", 60))
|
||||
|
||||
RAY_HEAD_POD_NAME = os.getenv("RAY_HEAD_POD_NAME")
|
||||
|
||||
# https://kubernetes.io/docs/tasks/run-application/access-api-from-pod
|
||||
# While running in a Pod, your container can create an HTTPS URL for the
|
||||
# Kubernetes API server by fetching the KUBERNETES_SERVICE_HOST and
|
||||
# KUBERNETES_SERVICE_PORT_HTTPS environment variables.
|
||||
KUBERNETES_SERVICE_HOST = os.getenv(
|
||||
"KUBERNETES_SERVICE_HOST", "https://kubernetes.default"
|
||||
)
|
||||
KUBERNETES_SERVICE_PORT = os.getenv("KUBERNETES_SERVICE_PORT_HTTPS", "443")
|
||||
KUBERNETES_HOST = build_address(KUBERNETES_SERVICE_HOST, KUBERNETES_SERVICE_PORT)
|
||||
# Key for GKE label that identifies which multi-host replica a pod belongs to
|
||||
REPLICA_INDEX_KEY = "replicaIndex"
|
||||
|
||||
TOKEN_REFRESH_PERIOD = datetime.timedelta(minutes=1)
|
||||
|
||||
# Design:
|
||||
|
||||
# Each modification the autoscaler wants to make is posted to the API server goal state
|
||||
# (e.g. if the autoscaler wants to scale up, it increases the number of
|
||||
# replicas of the worker group it wants to scale, if it wants to scale down
|
||||
# it decreases the number of replicas and adds the exact pods that should be
|
||||
# terminated to the scaleStrategy).
|
||||
|
||||
# KubeRayNodeProvider inherits from BatchingNodeProvider.
|
||||
# Thus, the autoscaler's create and terminate requests are batched into a single
|
||||
# Scale Request object which is submitted at the end of autoscaler update.
|
||||
# KubeRay node provider converts the ScaleRequest into a RayCluster CR patch
|
||||
# and applies the patch in the submit_scale_request method.
|
||||
|
||||
# To reduce potential for race conditions, KubeRayNodeProvider
|
||||
# aborts the autoscaler update if the operator has not yet processed workersToDelete -
|
||||
# see KubeRayNodeProvider.safe_to_scale().
|
||||
# Once it is confirmed that workersToDelete have been cleaned up, KubeRayNodeProvider
|
||||
# clears the workersToDelete list.
|
||||
|
||||
|
||||
# Note: Log handlers set up in autoscaling monitor entrypoint.
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def node_data_from_pod(pod: Dict[str, Any]) -> NodeData:
|
||||
"""Converts a Ray pod extracted from K8s into Ray NodeData.
|
||||
NodeData is processed by BatchingNodeProvider.
|
||||
"""
|
||||
kind, type = kind_and_type(pod)
|
||||
status = status_tag(pod)
|
||||
ip = pod_ip(pod)
|
||||
replica_index = _replica_index_label(pod)
|
||||
return NodeData(
|
||||
kind=kind, type=type, replica_index=replica_index, status=status, ip=ip
|
||||
)
|
||||
|
||||
|
||||
def kind_and_type(pod: Dict[str, Any]) -> Tuple[NodeKind, NodeType]:
|
||||
"""Determine Ray node kind (head or workers) and node type (worker group name)
|
||||
from a Ray pod's labels.
|
||||
"""
|
||||
labels = pod["metadata"]["labels"]
|
||||
kind = (
|
||||
NODE_KIND_HEAD
|
||||
if labels[KUBERAY_LABEL_KEY_KIND] == KUBERAY_KIND_HEAD
|
||||
else NODE_KIND_WORKER
|
||||
)
|
||||
type = labels[KUBERAY_LABEL_KEY_TYPE]
|
||||
return kind, type
|
||||
|
||||
|
||||
def _replica_index_label(pod: Dict[str, Any]) -> Optional[str]:
|
||||
"""Returns the replicaIndex label for a Pod in a multi-host TPU worker group.
|
||||
The replicaIndex label is set by the GKE TPU Ray webhook and is of
|
||||
the form {$WORKER_GROUP_NAME-$REPLICA_INDEX} where $REPLICA_INDEX
|
||||
is an integer from 0 to Replicas-1.
|
||||
"""
|
||||
labels = pod["metadata"]["labels"]
|
||||
return labels.get(REPLICA_INDEX_KEY, None)
|
||||
|
||||
|
||||
def pod_ip(pod: Dict[str, Any]) -> NodeIP:
|
||||
return pod["status"].get("podIP", "IP not yet assigned")
|
||||
|
||||
|
||||
def status_tag(pod: Dict[str, Any]) -> NodeStatus:
|
||||
"""Convert pod state to Ray autoscaler node status.
|
||||
|
||||
See the doc string of the class
|
||||
batching_node_provider.NodeData for the semantics of node status.
|
||||
"""
|
||||
if (
|
||||
"containerStatuses" not in pod["status"]
|
||||
or not pod["status"]["containerStatuses"]
|
||||
):
|
||||
return "pending"
|
||||
|
||||
state = pod["status"]["containerStatuses"][0]["state"]
|
||||
|
||||
if "pending" in state:
|
||||
return "pending"
|
||||
if "running" in state:
|
||||
return STATUS_UP_TO_DATE
|
||||
if "waiting" in state:
|
||||
return "waiting"
|
||||
if "terminated" in state:
|
||||
return STATUS_UPDATE_FAILED
|
||||
raise ValueError("Unexpected container state.")
|
||||
|
||||
|
||||
def worker_delete_patch(group_index: str, workers_to_delete: List[NodeID]):
|
||||
path = f"/spec/workerGroupSpecs/{group_index}/scaleStrategy"
|
||||
value = {"workersToDelete": workers_to_delete}
|
||||
return replace_patch(path, value)
|
||||
|
||||
|
||||
def worker_replica_patch(group_index: str, target_replicas: int):
|
||||
path = f"/spec/workerGroupSpecs/{group_index}/replicas"
|
||||
value = target_replicas
|
||||
return replace_patch(path, value)
|
||||
|
||||
|
||||
def replace_patch(path: str, value: Any) -> Dict[str, Any]:
|
||||
return {"op": "replace", "path": path, "value": value}
|
||||
|
||||
|
||||
def load_k8s_secrets() -> Tuple[Dict[str, str], str]:
|
||||
"""
|
||||
Loads secrets needed to access K8s resources.
|
||||
|
||||
Returns:
|
||||
headers: Headers with K8s access token
|
||||
verify: Path to certificate
|
||||
"""
|
||||
with open("/var/run/secrets/kubernetes.io/serviceaccount/token") as secret:
|
||||
token = secret.read()
|
||||
|
||||
headers = {
|
||||
"Authorization": "Bearer " + token,
|
||||
}
|
||||
verify = "/var/run/secrets/kubernetes.io/serviceaccount/ca.crt"
|
||||
|
||||
return headers, verify
|
||||
|
||||
|
||||
def url_from_resource(
|
||||
namespace: str,
|
||||
path: str,
|
||||
kuberay_crd_version: str = KUBERAY_CRD_VER,
|
||||
kubernetes_host: str = KUBERNETES_HOST,
|
||||
) -> str:
|
||||
"""Convert resource path to REST URL for Kubernetes API server.
|
||||
|
||||
Args:
|
||||
namespace: The K8s namespace of the resource
|
||||
path: The part of the resource path that starts with the resource type.
|
||||
Supported resource types are "pods" and "rayclusters".
|
||||
kuberay_crd_version: The API version of the KubeRay CRD.
|
||||
Looks like "v1alpha1", "v1".
|
||||
kubernetes_host: The host of the Kubernetes API server.
|
||||
Uses $KUBERNETES_SERVICE_HOST and
|
||||
$KUBERNETES_SERVICE_PORT to construct the kubernetes_host if not provided.
|
||||
|
||||
When set by Kubernetes,
|
||||
$KUBERNETES_SERVICE_HOST could be an IP address. That's why the https
|
||||
scheme is added here.
|
||||
|
||||
Defaults to "https://kubernetes.default:443".
|
||||
|
||||
Returns:
|
||||
The REST URL for the resource.
|
||||
"""
|
||||
if kubernetes_host.startswith("http://"):
|
||||
raise ValueError("Kubernetes host must be accessed over HTTPS.")
|
||||
if not kubernetes_host.startswith("https://"):
|
||||
kubernetes_host = "https://" + kubernetes_host
|
||||
if path.startswith("pods"):
|
||||
api_group = "/api/v1"
|
||||
elif path.startswith("rayclusters"):
|
||||
api_group = "/apis/ray.io/" + kuberay_crd_version
|
||||
else:
|
||||
raise NotImplementedError("Tried to access unknown entity at {}".format(path))
|
||||
return kubernetes_host + api_group + "/namespaces/" + namespace + "/" + path
|
||||
|
||||
|
||||
def _worker_group_index(raycluster: Dict[str, Any], group_name: str) -> int:
|
||||
"""Extract worker group index from RayCluster."""
|
||||
group_names = [
|
||||
spec["groupName"] for spec in raycluster["spec"].get("workerGroupSpecs", [])
|
||||
]
|
||||
return group_names.index(group_name)
|
||||
|
||||
|
||||
def _worker_group_max_replicas(
|
||||
raycluster: Dict[str, Any], group_index: int
|
||||
) -> Optional[int]:
|
||||
"""Extract the maxReplicas of a worker group.
|
||||
|
||||
If maxReplicas is unset, return None, to be interpreted as "no constraint".
|
||||
At time of writing, it should be impossible for maxReplicas to be unset, but it's
|
||||
better to handle this anyway.
|
||||
"""
|
||||
return raycluster["spec"]["workerGroupSpecs"][group_index].get("maxReplicas")
|
||||
|
||||
|
||||
def _worker_group_replicas(raycluster: Dict[str, Any], group_index: int):
|
||||
# 1 is the default replicas value used by the KubeRay operator
|
||||
return raycluster["spec"]["workerGroupSpecs"][group_index].get("replicas", 1)
|
||||
|
||||
|
||||
def _worker_group_num_of_hosts(raycluster: Dict[str, Any], group_index: int):
|
||||
# 1 is the default numOfHosts value used by the KubeRay operator
|
||||
return raycluster["spec"]["workerGroupSpecs"][group_index].get("numOfHosts", 1)
|
||||
|
||||
|
||||
class IKubernetesHttpApiClient(ABC):
|
||||
"""
|
||||
An interface for a Kubernetes HTTP API client.
|
||||
|
||||
This interface could be used to mock the Kubernetes API client in tests.
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
def get(self, path: str) -> Dict[str, Any]:
|
||||
"""Wrapper for REST GET of resource with proper headers."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def patch(
|
||||
self,
|
||||
path: str,
|
||||
payload: Union[List[Dict[str, Any]], Dict[str, Any]],
|
||||
content_type: str = "application/json-patch+json",
|
||||
) -> Dict[str, Any]:
|
||||
"""Wrapper for REST PATCH of resource with proper headers."""
|
||||
pass
|
||||
|
||||
|
||||
class KubernetesHttpApiClient(IKubernetesHttpApiClient):
|
||||
def __init__(self, namespace: str, kuberay_crd_version: str = KUBERAY_CRD_VER):
|
||||
self._kuberay_crd_version = kuberay_crd_version
|
||||
self._namespace = namespace
|
||||
self._token_expires_at = datetime.datetime.now() + TOKEN_REFRESH_PERIOD
|
||||
self._headers, self._verify = None, None
|
||||
|
||||
def _get_refreshed_headers_and_verify(self):
|
||||
if (datetime.datetime.now() >= self._token_expires_at) or (
|
||||
self._headers is None or self._verify is None
|
||||
):
|
||||
logger.info("Refreshing K8s API client token and certs.")
|
||||
self._headers, self._verify = load_k8s_secrets()
|
||||
self._token_expires_at = datetime.datetime.now() + TOKEN_REFRESH_PERIOD
|
||||
return self._headers, self._verify
|
||||
else:
|
||||
return self._headers, self._verify
|
||||
|
||||
def get(self, path: str) -> Dict[str, Any]:
|
||||
"""Wrapper for REST GET of resource with proper headers.
|
||||
|
||||
Args:
|
||||
path: The part of the resource path that starts with the resource type.
|
||||
|
||||
Returns:
|
||||
The JSON response of the GET request.
|
||||
|
||||
Raises:
|
||||
HTTPError: If the GET request fails.
|
||||
"""
|
||||
url = url_from_resource(
|
||||
namespace=self._namespace,
|
||||
path=path,
|
||||
kuberay_crd_version=self._kuberay_crd_version,
|
||||
)
|
||||
|
||||
headers, verify = self._get_refreshed_headers_and_verify()
|
||||
result = requests.get(
|
||||
url,
|
||||
headers=headers,
|
||||
timeout=KUBERAY_REQUEST_TIMEOUT_S,
|
||||
verify=verify,
|
||||
)
|
||||
if not result.status_code == 200:
|
||||
result.raise_for_status()
|
||||
return result.json()
|
||||
|
||||
def patch(
|
||||
self,
|
||||
path: str,
|
||||
payload: Union[List[Dict[str, Any]], Dict[str, Any]],
|
||||
content_type: str = "application/json-patch+json",
|
||||
) -> Dict[str, Any]:
|
||||
"""Wrapper for REST PATCH of resource with proper headers
|
||||
|
||||
Args:
|
||||
path: The part of the resource path that starts with the resource type.
|
||||
payload: The patch payload, either a JSON Patch list or a
|
||||
strategic-merge patch object.
|
||||
content_type: The content type of the merge strategy.
|
||||
|
||||
Returns:
|
||||
The JSON response of the PATCH request.
|
||||
|
||||
Raises:
|
||||
HTTPError: If the PATCH request fails.
|
||||
"""
|
||||
url = url_from_resource(
|
||||
namespace=self._namespace,
|
||||
path=path,
|
||||
kuberay_crd_version=self._kuberay_crd_version,
|
||||
)
|
||||
headers, verify = self._get_refreshed_headers_and_verify()
|
||||
result = requests.patch(
|
||||
url,
|
||||
json.dumps(payload),
|
||||
headers={**headers, "Content-type": content_type},
|
||||
timeout=KUBERAY_REQUEST_TIMEOUT_S,
|
||||
verify=verify,
|
||||
)
|
||||
if not result.status_code == 200:
|
||||
result.raise_for_status()
|
||||
return result.json()
|
||||
|
||||
|
||||
class KubeRayNodeProvider(BatchingNodeProvider): # type: ignore
|
||||
def __init__(
|
||||
self,
|
||||
provider_config: Dict[str, Any],
|
||||
cluster_name: str,
|
||||
):
|
||||
logger.info("Creating KubeRayNodeProvider.")
|
||||
self.namespace = provider_config["namespace"]
|
||||
self.cluster_name = cluster_name
|
||||
|
||||
self.k8s_api_client = KubernetesHttpApiClient(self.namespace)
|
||||
|
||||
assert (
|
||||
provider_config.get(WORKER_LIVENESS_CHECK_KEY, True) is False
|
||||
), f"To use KubeRayNodeProvider, must set `{WORKER_LIVENESS_CHECK_KEY}:False`."
|
||||
BatchingNodeProvider.__init__(self, provider_config, cluster_name)
|
||||
|
||||
def get_node_data(self) -> Dict[NodeID, NodeData]:
|
||||
"""Queries K8s for pods in the RayCluster. Converts that pod data into a
|
||||
map of pod name to Ray NodeData, as required by BatchingNodeProvider.
|
||||
"""
|
||||
# Store the raycluster CR
|
||||
self._raycluster = self._get(f"rayclusters/{self.cluster_name}")
|
||||
|
||||
# Get the pods resource version.
|
||||
# Specifying a resource version in list requests is important for scalability:
|
||||
# https://kubernetes.io/docs/reference/using-api/api-concepts/#semantics-for-get-and-list
|
||||
resource_version = self._get_pods_resource_version()
|
||||
if resource_version:
|
||||
logger.info(
|
||||
f"Listing pods for RayCluster {self.cluster_name}"
|
||||
f" in namespace {self.namespace}"
|
||||
f" at pods resource version >= {resource_version}."
|
||||
)
|
||||
|
||||
# Filter pods by cluster_name.
|
||||
label_selector = requests.utils.quote(f"ray.io/cluster={self.cluster_name}")
|
||||
|
||||
resource_path = f"pods?labelSelector={label_selector}"
|
||||
if resource_version:
|
||||
resource_path += (
|
||||
f"&resourceVersion={resource_version}"
|
||||
+ "&resourceVersionMatch=NotOlderThan"
|
||||
)
|
||||
|
||||
pod_list = self._get(resource_path)
|
||||
fetched_resource_version = pod_list["metadata"]["resourceVersion"]
|
||||
logger.info(
|
||||
f"Fetched pod data at resource version" f" {fetched_resource_version}."
|
||||
)
|
||||
|
||||
# Extract node data from the pod list.
|
||||
node_data_dict = {}
|
||||
for pod in pod_list["items"]:
|
||||
# Kubernetes sets metadata.deletionTimestamp immediately after admitting a
|
||||
# request to delete an object. Full removal of the object may take some time
|
||||
# after the deletion timestamp is set. See link for details:
|
||||
# https://kubernetes.io/docs/reference/using-api/api-concepts/#resource-deletion
|
||||
if "deletionTimestamp" in pod["metadata"]:
|
||||
# Ignore pods marked for termination.
|
||||
continue
|
||||
pod_name = pod["metadata"]["name"]
|
||||
node_data_dict[pod_name] = node_data_from_pod(pod)
|
||||
return node_data_dict
|
||||
|
||||
def submit_scale_request(self, scale_request: ScaleRequest):
|
||||
"""Converts the scale request generated by BatchingNodeProvider into
|
||||
a patch that modifies the RayCluster CR's replicas and/or workersToDelete
|
||||
fields. Then submits the patch to the K8s API server.
|
||||
"""
|
||||
# Transform the scale request into a patch payload.
|
||||
patch_payload = self._scale_request_to_patch_payload(
|
||||
scale_request, self._raycluster
|
||||
)
|
||||
|
||||
# Submit the patch to K8s.
|
||||
logger.info(
|
||||
"Autoscaler is submitting the following patch to RayCluster "
|
||||
f"{self.cluster_name} in namespace {self.namespace}."
|
||||
)
|
||||
logger.info(patch_payload)
|
||||
self._submit_raycluster_patch(patch_payload)
|
||||
|
||||
def safe_to_scale(self) -> bool:
|
||||
"""Returns False iff non_terminated_nodes contains any pods in the RayCluster's
|
||||
workersToDelete lists.
|
||||
|
||||
Explanation:
|
||||
If there are any workersToDelete which are non-terminated,
|
||||
we should wait for the operator to do its job and delete those
|
||||
pods. Therefore, we back off the autoscaler update.
|
||||
|
||||
If, on the other hand, all of the workersToDelete have already been cleaned up,
|
||||
then we patch away the workersToDelete lists and return True.
|
||||
In the future, we may consider having the operator clean up workersToDelete
|
||||
on it own:
|
||||
https://github.com/ray-project/kuberay/issues/733
|
||||
|
||||
Note (Dmitri):
|
||||
It is stylistically bad that this function has a side effect.
|
||||
"""
|
||||
# Get the list of nodes.
|
||||
node_set = set(self.node_data_dict.keys())
|
||||
worker_groups = self._raycluster["spec"].get("workerGroupSpecs", [])
|
||||
|
||||
# Accumulates the indices of worker groups with non-empty workersToDelete
|
||||
non_empty_worker_group_indices = []
|
||||
|
||||
for group_index, worker_group in enumerate(worker_groups):
|
||||
workersToDelete = worker_group.get("scaleStrategy", {}).get(
|
||||
"workersToDelete", []
|
||||
)
|
||||
if workersToDelete:
|
||||
non_empty_worker_group_indices.append(group_index)
|
||||
for worker in workersToDelete:
|
||||
if worker in node_set:
|
||||
# The operator hasn't removed this worker yet. Abort
|
||||
# the autoscaler update.
|
||||
logger.warning(f"Waiting for operator to remove worker {worker}.")
|
||||
return False
|
||||
|
||||
# All required workersToDelete have been removed.
|
||||
# Clean up the workersToDelete field.
|
||||
patch_payload = []
|
||||
for group_index in non_empty_worker_group_indices:
|
||||
patch = worker_delete_patch(group_index, workers_to_delete=[])
|
||||
patch_payload.append(patch)
|
||||
if patch_payload:
|
||||
logger.info("Cleaning up workers to delete.")
|
||||
logger.info(f"Submitting patch {patch_payload}.")
|
||||
self._submit_raycluster_patch(patch_payload)
|
||||
|
||||
# It's safe to proceed with the autoscaler update.
|
||||
return True
|
||||
|
||||
def _get_pods_resource_version(self) -> str:
|
||||
"""
|
||||
Extract a recent pods resource version by reading the head pod's
|
||||
metadata.resourceVersion of the response.
|
||||
"""
|
||||
if not RAY_HEAD_POD_NAME:
|
||||
return None
|
||||
pod_resp = self._get(f"pods/{RAY_HEAD_POD_NAME}")
|
||||
return pod_resp["metadata"]["resourceVersion"]
|
||||
|
||||
def _scale_request_to_patch_payload(
|
||||
self, scale_request: ScaleRequest, raycluster: Dict[str, Any]
|
||||
) -> List[Dict[str, Any]]:
|
||||
"""Converts autoscaler scale request into a RayCluster CR patch payload."""
|
||||
patch_payload = []
|
||||
# Collect patches for replica counts.
|
||||
for node_type, target_replicas in scale_request.desired_num_workers.items():
|
||||
group_index = _worker_group_index(raycluster, node_type)
|
||||
group_max_replicas = _worker_group_max_replicas(raycluster, group_index)
|
||||
# Cap the replica count to maxReplicas.
|
||||
if group_max_replicas is not None and group_max_replicas < target_replicas:
|
||||
logger.warning(
|
||||
"Autoscaler attempted to create "
|
||||
+ "more than maxReplicas pods of type {}.".format(node_type)
|
||||
)
|
||||
target_replicas = group_max_replicas
|
||||
# Check if we need to change the target count.
|
||||
if target_replicas == _worker_group_replicas(raycluster, group_index):
|
||||
# No patch required.
|
||||
continue
|
||||
# Need to patch replica count. Format the patch and add it to the payload.
|
||||
patch = worker_replica_patch(group_index, target_replicas)
|
||||
patch_payload.append(patch)
|
||||
|
||||
# Maps node_type to nodes to delete for that group.
|
||||
deletion_groups = defaultdict(list)
|
||||
for worker in scale_request.workers_to_delete:
|
||||
node_type = self.node_tags(worker)[TAG_RAY_USER_NODE_TYPE]
|
||||
deletion_groups[node_type].append(worker)
|
||||
|
||||
for node_type, workers_to_delete in deletion_groups.items():
|
||||
group_index = _worker_group_index(raycluster, node_type)
|
||||
patch = worker_delete_patch(group_index, workers_to_delete)
|
||||
patch_payload.append(patch)
|
||||
|
||||
return patch_payload
|
||||
|
||||
def _submit_raycluster_patch(self, patch_payload: List[Dict[str, Any]]):
|
||||
"""Submits a patch to modify a RayCluster CR."""
|
||||
path = "rayclusters/{}".format(self.cluster_name)
|
||||
self._patch(path, patch_payload)
|
||||
|
||||
def _get(self, path: str) -> Dict[str, Any]:
|
||||
"""Wrapper for REST GET of resource with proper headers."""
|
||||
return self.k8s_api_client.get(path)
|
||||
|
||||
def _patch(self, path: str, payload: List[Dict[str, Any]]) -> Dict[str, Any]:
|
||||
"""Wrapper for REST PATCH of resource with proper headers."""
|
||||
return self.k8s_api_client.patch(path, payload)
|
||||
@@ -0,0 +1,143 @@
|
||||
import logging
|
||||
import os
|
||||
import subprocess
|
||||
import time
|
||||
|
||||
import ray
|
||||
from ray._common.network_utils import build_address
|
||||
from ray._common.ray_constants import (
|
||||
LOGGING_ROTATE_BACKUP_COUNT,
|
||||
LOGGING_ROTATE_BYTES,
|
||||
)
|
||||
from ray._common.utils import env_integer, try_to_create_directory
|
||||
from ray._private import ray_constants
|
||||
from ray._private.ray_logging import setup_component_logger
|
||||
from ray._private.services import get_node_ip_address
|
||||
from ray._private.utils import get_all_node_info_until_retrieved
|
||||
from ray._raylet import GcsClient
|
||||
from ray.autoscaler._private.kuberay.autoscaling_config import AutoscalingConfigProducer
|
||||
from ray.autoscaler._private.monitor import Monitor
|
||||
from ray.autoscaler.v2.instance_manager.config import KubeRayConfigReader
|
||||
from ray.autoscaler.v2.utils import is_autoscaler_v2
|
||||
from ray.core.generated.gcs_service_pb2 import GetAllNodeInfoRequest
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
BACKOFF_S = 5
|
||||
|
||||
|
||||
def _get_log_dir(gcs_client: GcsClient) -> str:
|
||||
head_node_selector = GetAllNodeInfoRequest.NodeSelector()
|
||||
head_node_selector.is_head_node = True
|
||||
|
||||
# We need to wait until head node's raylet is registered in GCS.
|
||||
node_infos = get_all_node_info_until_retrieved(
|
||||
gcs_client,
|
||||
node_selectors=[head_node_selector],
|
||||
)
|
||||
|
||||
node_info = next(iter(node_infos))
|
||||
temp_dir = getattr(node_info, "temp_dir", None)
|
||||
if temp_dir is None:
|
||||
raise Exception(
|
||||
"Node temp_dir was not found in NodeInfo. did the head node's raylet start successfully?"
|
||||
)
|
||||
return os.path.join(temp_dir, ray._private.ray_constants.SESSION_LATEST, "logs")
|
||||
|
||||
|
||||
def run_kuberay_autoscaler(cluster_name: str, cluster_namespace: str):
|
||||
"""Wait until the Ray head container is ready. Then start the autoscaler."""
|
||||
head_ip = get_node_ip_address()
|
||||
ray_address = build_address(head_ip, 6379)
|
||||
while True:
|
||||
try:
|
||||
# Autoscaler Ray version might not exactly match GCS version, so skip the
|
||||
# version check when checking GCS status.
|
||||
subprocess.check_call(
|
||||
[
|
||||
"ray",
|
||||
"health-check",
|
||||
"--address",
|
||||
ray_address,
|
||||
"--skip-version-check",
|
||||
]
|
||||
)
|
||||
logger.info("The Ray head is ready. Starting the autoscaler.")
|
||||
break
|
||||
except subprocess.CalledProcessError:
|
||||
logger.warning(
|
||||
f"The Ray head is not ready. Will check again in {BACKOFF_S} seconds."
|
||||
)
|
||||
time.sleep(BACKOFF_S)
|
||||
|
||||
gcs_client = GcsClient(ray_address)
|
||||
log_dir = _get_log_dir(gcs_client)
|
||||
|
||||
# The Ray head container sets up the log directory. Thus, we set up logging
|
||||
# only after the Ray head is ready.
|
||||
_setup_logging(log_dir)
|
||||
|
||||
# autoscaling_config_producer reads the RayCluster CR from K8s and uses the CR
|
||||
# to output an autoscaling config.
|
||||
autoscaling_config_producer = AutoscalingConfigProducer(
|
||||
cluster_name, cluster_namespace
|
||||
)
|
||||
|
||||
if is_autoscaler_v2(fetch_from_server=True, gcs_client=gcs_client):
|
||||
from ray.autoscaler.v2.monitor import AutoscalerMonitor as MonitorV2
|
||||
|
||||
MonitorV2(
|
||||
address=gcs_client.address,
|
||||
config_reader=KubeRayConfigReader(autoscaling_config_producer),
|
||||
log_dir=log_dir,
|
||||
monitor_ip=head_ip,
|
||||
).run()
|
||||
else:
|
||||
Monitor(
|
||||
address=gcs_client.address,
|
||||
# The `autoscaling_config` arg can be a dict or a `Callable: () -> dict`.
|
||||
# In this case, it's a callable.
|
||||
autoscaling_config=autoscaling_config_producer,
|
||||
monitor_ip=head_ip,
|
||||
# Let the autoscaler process exit after it hits 5 exceptions.
|
||||
# (See ray.autoscaler._private.constants.AUTOSCALER_MAX_NUM_FAILURES.)
|
||||
# Kubernetes will then restart the autoscaler container.
|
||||
retry_on_failure=False,
|
||||
).run()
|
||||
|
||||
|
||||
def _setup_logging(log_dir: str) -> None:
|
||||
"""Log to autoscaler log file
|
||||
(typically, /tmp/ray/session_latest/logs/monitor.*)
|
||||
|
||||
Also log to pod stdout (logs viewable with `kubectl logs <head-pod> -c autoscaler`).
|
||||
|
||||
Args:
|
||||
log_dir: The path to the log directory.
|
||||
"""
|
||||
# The director should already exist, but try (safely) to create it just in case.
|
||||
try_to_create_directory(log_dir)
|
||||
|
||||
# Write logs at info level to monitor.log.
|
||||
max_bytes = env_integer("RAY_ROTATION_MAX_BYTES", LOGGING_ROTATE_BYTES)
|
||||
backup_count = env_integer("RAY_ROTATION_BACKUP_COUNT", LOGGING_ROTATE_BACKUP_COUNT)
|
||||
setup_component_logger(
|
||||
logging_level=ray_constants.LOGGER_LEVEL,
|
||||
logging_format=ray_constants.LOGGER_FORMAT,
|
||||
log_dir=log_dir,
|
||||
filename=ray_constants.MONITOR_LOG_FILE_NAME, # monitor.log
|
||||
max_bytes=max_bytes,
|
||||
backup_count=backup_count,
|
||||
)
|
||||
|
||||
# For the autoscaler, the root logger _also_ needs to write to stderr, not just
|
||||
# ray_constants.MONITOR_LOG_FILE_NAME.
|
||||
level = logging.getLevelName(ray_constants.LOGGER_LEVEL.upper())
|
||||
stderr_handler = logging._StderrHandler()
|
||||
stderr_handler.setFormatter(logging.Formatter(ray_constants.LOGGER_FORMAT))
|
||||
stderr_handler.setLevel(level)
|
||||
logging.root.setLevel(level)
|
||||
logging.root.addHandler(stderr_handler)
|
||||
|
||||
# The stdout handler was set up in the Ray CLI entry point.
|
||||
# See ray.scripts.scripts::cli().
|
||||
@@ -0,0 +1,112 @@
|
||||
# Source:
|
||||
# https://github.com/kubernetes-client/python/blob/master/kubernetes/utils/quantity.py
|
||||
from decimal import Decimal, InvalidOperation
|
||||
from typing import Optional, Union
|
||||
|
||||
from ray._private.accelerators.tpu import (
|
||||
get_num_chips_from_topology,
|
||||
get_tpu_cores_per_chip,
|
||||
)
|
||||
|
||||
# Mapping used to get generation for TPU-{accelerator}-head resource
|
||||
# https://cloud.google.com/kubernetes-engine/docs/how-to/tpus#run
|
||||
gke_tpu_accelerator_to_generation = {
|
||||
"tpu-v4-podslice": "v4",
|
||||
"tpu-v5-lite-device": "v5litepod",
|
||||
"tpu-v5-lite-podslice": "v5litepod",
|
||||
"tpu-v5p-slice": "v5p",
|
||||
"tpu-v6e-slice": "v6e",
|
||||
"tpu7x": "v7x",
|
||||
}
|
||||
|
||||
|
||||
def parse_quantity(quantity: Union[str, int, float, Decimal]) -> Decimal:
|
||||
"""Parse kubernetes canonical form quantity like 200Mi to a decimal number.
|
||||
Supported SI suffixes:
|
||||
base1024: Ki | Mi | Gi | Ti | Pi | Ei
|
||||
base1000: n | u | m | "" | k | M | G | T | P | E
|
||||
|
||||
See
|
||||
https://github.com/kubernetes/apimachinery/blob/master/pkg/api/resource/quantity.go
|
||||
|
||||
Args:
|
||||
quantity: kubernetes canonical form quantity.
|
||||
|
||||
Returns:
|
||||
The parsed quantity as a decimal number.
|
||||
|
||||
Raises:
|
||||
ValueError: On invalid or unknown input
|
||||
"""
|
||||
if isinstance(quantity, (int, float, Decimal)):
|
||||
return Decimal(quantity)
|
||||
|
||||
exponents = {
|
||||
"n": -3,
|
||||
"u": -2,
|
||||
"m": -1,
|
||||
"K": 1,
|
||||
"k": 1,
|
||||
"M": 2,
|
||||
"G": 3,
|
||||
"T": 4,
|
||||
"P": 5,
|
||||
"E": 6,
|
||||
}
|
||||
|
||||
quantity = str(quantity)
|
||||
number = quantity
|
||||
suffix = None
|
||||
if len(quantity) >= 2 and quantity[-1] == "i":
|
||||
if quantity[-2] in exponents:
|
||||
number = quantity[:-2]
|
||||
suffix = quantity[-2:]
|
||||
elif len(quantity) >= 1 and quantity[-1] in exponents:
|
||||
number = quantity[:-1]
|
||||
suffix = quantity[-1:]
|
||||
|
||||
try:
|
||||
number = Decimal(number)
|
||||
except InvalidOperation:
|
||||
raise ValueError("Invalid number format: {}".format(number))
|
||||
|
||||
if suffix is None:
|
||||
return number
|
||||
|
||||
if suffix.endswith("i"):
|
||||
base = 1024
|
||||
elif len(suffix) == 1:
|
||||
base = 1000
|
||||
else:
|
||||
raise ValueError("{} has unknown suffix".format(quantity))
|
||||
|
||||
# handle SI inconsistency
|
||||
if suffix == "ki":
|
||||
raise ValueError("{} has unknown suffix".format(quantity))
|
||||
|
||||
if suffix[0] not in exponents:
|
||||
raise ValueError("{} has unknown suffix".format(quantity))
|
||||
|
||||
exponent = Decimal(exponents[suffix[0]])
|
||||
return number * (base**exponent)
|
||||
|
||||
|
||||
def tpu_node_selectors_to_type(topology: str, accelerator: str) -> Optional[str]:
|
||||
"""Convert Kubernetes gke-tpu nodeSelectors to TPU accelerator_type
|
||||
for a kuberay TPU worker group.
|
||||
Args:
|
||||
topology: value of the cloud.google.com/gke-tpu-topology Kubernetes
|
||||
nodeSelector, describes the physical topology of the TPU podslice.
|
||||
accelerator: value of the cloud.google.com/gke-tpu-accelerator nodeSelector,
|
||||
the name of the TPU accelerator, e.g. tpu-v4-podslice
|
||||
Returns:
|
||||
A string, accelerator_type, e.g. "v4-8".
|
||||
"""
|
||||
if topology and accelerator:
|
||||
generation = gke_tpu_accelerator_to_generation.get(accelerator)
|
||||
if not generation:
|
||||
return None
|
||||
num_chips = get_num_chips_from_topology(topology)
|
||||
num_cores = num_chips * get_tpu_cores_per_chip(generation)
|
||||
return f"{generation}-{num_cores}"
|
||||
return None
|
||||
@@ -0,0 +1,37 @@
|
||||
import logging
|
||||
|
||||
from ray._private.ray_constants import DEBUG_AUTOSCALING_STATUS_LEGACY
|
||||
from ray.experimental.internal_kv import _internal_kv_initialized, _internal_kv_put
|
||||
|
||||
"""This file provides legacy support for the old info string in order to
|
||||
ensure the dashboard's `api/cluster_status` does not break backwards
|
||||
compatibilty.
|
||||
"""
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def legacy_log_info_string(autoscaler, nodes):
|
||||
tmp = "Cluster status: "
|
||||
tmp += info_string(autoscaler, nodes)
|
||||
tmp += "\n"
|
||||
tmp += autoscaler.load_metrics.info_string()
|
||||
tmp += "\n"
|
||||
tmp += autoscaler.resource_demand_scheduler.debug_string(
|
||||
nodes,
|
||||
autoscaler.pending_launches.breakdown(),
|
||||
autoscaler.load_metrics.get_resource_utilization(),
|
||||
)
|
||||
if _internal_kv_initialized():
|
||||
_internal_kv_put(DEBUG_AUTOSCALING_STATUS_LEGACY, tmp, overwrite=True)
|
||||
logger.debug(tmp)
|
||||
|
||||
|
||||
def info_string(autoscaler, nodes):
|
||||
suffix = ""
|
||||
if autoscaler.updaters:
|
||||
suffix += " ({} updating)".format(len(autoscaler.updaters))
|
||||
if autoscaler.num_failed_updates:
|
||||
suffix += " ({} failed to update)".format(len(autoscaler.num_failed_updates))
|
||||
|
||||
return "{} nodes{}".format(len(nodes), suffix)
|
||||
@@ -0,0 +1,383 @@
|
||||
import logging
|
||||
import time
|
||||
from collections import Counter
|
||||
from functools import reduce
|
||||
from typing import Any, Callable, Dict, List, Optional
|
||||
|
||||
from ray._private.gcs_utils import PlacementGroupTableData
|
||||
from ray.autoscaler._private.constants import (
|
||||
AUTOSCALER_MAX_RESOURCE_DEMAND_VECTOR_SIZE,
|
||||
AUTOSCALER_REPORT_PER_NODE_STATUS,
|
||||
)
|
||||
from ray.autoscaler._private.util import (
|
||||
DictCount,
|
||||
LoadMetricsSummary,
|
||||
NodeIP,
|
||||
ResourceDict,
|
||||
)
|
||||
from ray.core.generated.common_pb2 import PlacementStrategy
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def add_resources(dict1: Dict[str, float], dict2: Dict[str, float]) -> Dict[str, float]:
|
||||
"""Add the values in two dictionaries.
|
||||
|
||||
Args:
|
||||
dict1: The first dictionary.
|
||||
dict2: The second dictionary.
|
||||
|
||||
Returns:
|
||||
A new dictionary (inputs remain unmodified).
|
||||
"""
|
||||
new_dict = dict1.copy()
|
||||
for k, v in dict2.items():
|
||||
new_dict[k] = v + new_dict.get(k, 0)
|
||||
return new_dict
|
||||
|
||||
|
||||
def freq_of_dicts(
|
||||
dicts: List[Dict],
|
||||
serializer: Optional[Callable[[Dict], Any]] = None,
|
||||
deserializer: Callable[[Any], Any] = dict,
|
||||
) -> DictCount:
|
||||
"""Count a list of dictionaries (or unhashable types).
|
||||
|
||||
This is somewhat annoying because mutable data structures aren't hashable,
|
||||
and set/dict keys must be hashable.
|
||||
|
||||
Args:
|
||||
dicts: A list of dictionaries to be counted.
|
||||
serializer: A custom serialization function. The output type
|
||||
must be hashable. The default serializer converts a dictionary into
|
||||
a frozenset of KV pairs.
|
||||
deserializer: A custom deserialization function. See the
|
||||
serializer for information about the intermediate type. For
|
||||
dictionaries the output type is the same as the input type.
|
||||
|
||||
Returns:
|
||||
A list of tuples. Each entry in the list is a tuple containing a unique
|
||||
entry from `dicts` and its corresponding frequency count.
|
||||
"""
|
||||
if serializer is None:
|
||||
serializer = lambda d: frozenset(d.items()) # noqa: E731
|
||||
|
||||
freqs = Counter(serializer(d) for d in dicts)
|
||||
as_list = []
|
||||
for as_set, count in freqs.items():
|
||||
as_list.append((deserializer(as_set), count))
|
||||
return as_list
|
||||
|
||||
|
||||
class LoadMetrics:
|
||||
"""Container for cluster load metrics.
|
||||
|
||||
Metrics here are updated from raylet heartbeats. The autoscaler
|
||||
queries these metrics to determine when to scale up, and which nodes
|
||||
can be removed.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
self.last_heartbeat_time_by_ip = {}
|
||||
self.static_resources_by_ip = {}
|
||||
self.dynamic_resources_by_ip = {}
|
||||
self.node_id_by_ip = {}
|
||||
self.waiting_bundles = []
|
||||
self.infeasible_bundles = []
|
||||
self.pending_placement_groups = []
|
||||
self.resource_requests = []
|
||||
self.ray_nodes_last_used_time_by_ip = {}
|
||||
|
||||
def __bool__(self):
|
||||
"""A load metrics instance is Falsey iff the autoscaler process
|
||||
has not received a resource message from the GCS.
|
||||
"""
|
||||
return bool(self.node_id_by_ip)
|
||||
|
||||
def update(
|
||||
self,
|
||||
ip: str,
|
||||
node_id: bytes,
|
||||
static_resources: Dict[str, Dict],
|
||||
dynamic_resources: Dict[str, Dict],
|
||||
node_idle_duration_s: float,
|
||||
waiting_bundles: List[Dict[str, float]] = None,
|
||||
infeasible_bundles: List[Dict[str, float]] = None,
|
||||
pending_placement_groups: List[PlacementGroupTableData] = None,
|
||||
):
|
||||
self.static_resources_by_ip[ip] = static_resources
|
||||
self.node_id_by_ip[ip] = node_id
|
||||
|
||||
if not waiting_bundles:
|
||||
waiting_bundles = []
|
||||
if not infeasible_bundles:
|
||||
infeasible_bundles = []
|
||||
if not pending_placement_groups:
|
||||
pending_placement_groups = []
|
||||
|
||||
# We are not guaranteed to have a corresponding dynamic resource
|
||||
# for every static resource because dynamic resources are based on
|
||||
# the available resources in the heartbeat, which does not exist
|
||||
# if it is zero. Thus, we have to update dynamic resources here.
|
||||
dynamic_resources_update = dynamic_resources.copy()
|
||||
for resource_name, capacity in self.static_resources_by_ip[ip].items():
|
||||
if resource_name not in dynamic_resources_update:
|
||||
dynamic_resources_update[resource_name] = 0.0
|
||||
self.dynamic_resources_by_ip[ip] = dynamic_resources_update
|
||||
|
||||
now = time.time()
|
||||
self.ray_nodes_last_used_time_by_ip[ip] = now - node_idle_duration_s
|
||||
self.last_heartbeat_time_by_ip[ip] = now
|
||||
self.waiting_bundles = waiting_bundles
|
||||
self.infeasible_bundles = infeasible_bundles
|
||||
self.pending_placement_groups = pending_placement_groups
|
||||
|
||||
def mark_active(self, ip):
|
||||
assert ip is not None, "IP should be known at this time"
|
||||
logger.debug("Node {} is newly setup, treating as active".format(ip))
|
||||
self.last_heartbeat_time_by_ip[ip] = time.time()
|
||||
|
||||
def prune_active_ips(self, active_ips: List[str]):
|
||||
"""The Raylet ips stored by LoadMetrics are obtained by polling
|
||||
the GCS in Monitor.update_load_metrics().
|
||||
|
||||
On the other hand, the autoscaler gets a list of node ips from
|
||||
its NodeProvider.
|
||||
|
||||
This method removes from LoadMetrics the ips unknown to the autoscaler.
|
||||
|
||||
Args:
|
||||
active_ips: The node ips known to the autoscaler.
|
||||
"""
|
||||
active_ips = set(active_ips)
|
||||
|
||||
def prune(mapping, should_log):
|
||||
unwanted_ips = set(mapping) - active_ips
|
||||
for unwanted_ip in unwanted_ips:
|
||||
if should_log:
|
||||
logger.info("LoadMetrics: " f"Removed ip: {unwanted_ip}.")
|
||||
del mapping[unwanted_ip]
|
||||
if unwanted_ips and should_log:
|
||||
logger.info(
|
||||
"LoadMetrics: "
|
||||
"Removed {} stale ip mappings: {} not in {}".format(
|
||||
len(unwanted_ips), unwanted_ips, active_ips
|
||||
)
|
||||
)
|
||||
assert not (unwanted_ips & set(mapping))
|
||||
|
||||
prune(self.ray_nodes_last_used_time_by_ip, should_log=True)
|
||||
prune(self.static_resources_by_ip, should_log=False)
|
||||
prune(self.node_id_by_ip, should_log=False)
|
||||
prune(self.dynamic_resources_by_ip, should_log=False)
|
||||
prune(self.last_heartbeat_time_by_ip, should_log=False)
|
||||
|
||||
def get_node_resources(self):
|
||||
"""Return a list of node resources (static resource sizes).
|
||||
|
||||
Returns:
|
||||
An iterable of node resource dicts.
|
||||
|
||||
Example:
|
||||
>>> from ray.autoscaler._private.load_metrics import LoadMetrics
|
||||
>>> metrics = LoadMetrics(...) # doctest: +SKIP
|
||||
>>> metrics.get_node_resources() # doctest: +SKIP
|
||||
[{"CPU": 1}, {"CPU": 4, "GPU": 8}] # for two different nodes
|
||||
"""
|
||||
return self.static_resources_by_ip.values()
|
||||
|
||||
def get_static_node_resources_by_ip(self) -> Dict[NodeIP, ResourceDict]:
|
||||
"""Return a dict of node resources for every node ip.
|
||||
|
||||
Returns:
|
||||
A mapping from node IP to its static resource dict.
|
||||
|
||||
Example:
|
||||
>>> from ray.autoscaler._private.load_metrics import LoadMetrics
|
||||
>>> metrics = LoadMetrics(...) # doctest: +SKIP
|
||||
>>> metrics.get_static_node_resources_by_ip() # doctest: +SKIP
|
||||
{127.0.0.1: {"CPU": 1}, 127.0.0.2: {"CPU": 4, "GPU": 8}}
|
||||
"""
|
||||
return self.static_resources_by_ip
|
||||
|
||||
def get_resource_utilization(self):
|
||||
return self.dynamic_resources_by_ip
|
||||
|
||||
def _get_resource_usage(self):
|
||||
resources_used = {}
|
||||
resources_total = {}
|
||||
for ip, max_resources in self.static_resources_by_ip.items():
|
||||
avail_resources = self.dynamic_resources_by_ip[ip]
|
||||
for resource_id, amount in max_resources.items():
|
||||
used = amount - avail_resources[resource_id]
|
||||
if resource_id not in resources_used:
|
||||
resources_used[resource_id] = 0.0
|
||||
resources_total[resource_id] = 0.0
|
||||
resources_used[resource_id] += used
|
||||
resources_total[resource_id] += amount
|
||||
used = max(0, used)
|
||||
|
||||
return resources_used, resources_total
|
||||
|
||||
def get_resource_demand_vector(self, clip=True):
|
||||
if clip:
|
||||
# Bound the total number of bundles to
|
||||
# 2xMAX_RESOURCE_DEMAND_VECTOR_SIZE. This guarantees the resource
|
||||
# demand scheduler bin packing algorithm takes a reasonable amount
|
||||
# of time to run.
|
||||
return (
|
||||
self.waiting_bundles[:AUTOSCALER_MAX_RESOURCE_DEMAND_VECTOR_SIZE]
|
||||
+ self.infeasible_bundles[:AUTOSCALER_MAX_RESOURCE_DEMAND_VECTOR_SIZE]
|
||||
)
|
||||
else:
|
||||
return self.waiting_bundles + self.infeasible_bundles
|
||||
|
||||
def get_resource_requests(self):
|
||||
return self.resource_requests
|
||||
|
||||
def get_pending_placement_groups(self):
|
||||
return self.pending_placement_groups
|
||||
|
||||
def resources_avail_summary(self) -> str:
|
||||
"""Return a concise string of cluster size to report to event logs.
|
||||
|
||||
For example, "3 CPUs, 4 GPUs".
|
||||
"""
|
||||
total_resources = (
|
||||
reduce(add_resources, self.static_resources_by_ip.values())
|
||||
if self.static_resources_by_ip
|
||||
else {}
|
||||
)
|
||||
out = "{} CPUs".format(int(total_resources.get("CPU", 0)))
|
||||
if "GPU" in total_resources:
|
||||
out += ", {} GPUs".format(int(total_resources["GPU"]))
|
||||
if "TPU" in total_resources:
|
||||
out += ", {} TPUs".format(int(total_resources["TPU"]))
|
||||
return out
|
||||
|
||||
def summary(self):
|
||||
available_resources = (
|
||||
reduce(add_resources, self.dynamic_resources_by_ip.values())
|
||||
if self.dynamic_resources_by_ip
|
||||
else {}
|
||||
)
|
||||
total_resources = (
|
||||
reduce(add_resources, self.static_resources_by_ip.values())
|
||||
if self.static_resources_by_ip
|
||||
else {}
|
||||
)
|
||||
usage_dict = {}
|
||||
for key in total_resources:
|
||||
if key in ["memory", "object_store_memory"]:
|
||||
total = total_resources[key]
|
||||
available = available_resources[key]
|
||||
usage_dict[key] = (total - available, total)
|
||||
else:
|
||||
total = total_resources[key]
|
||||
usage_dict[key] = (total - available_resources[key], total)
|
||||
|
||||
summarized_demand_vector = freq_of_dicts(
|
||||
self.get_resource_demand_vector(clip=False)
|
||||
)
|
||||
summarized_resource_requests = freq_of_dicts(self.get_resource_requests())
|
||||
|
||||
def placement_group_serializer(pg):
|
||||
bundles = tuple(
|
||||
frozenset(bundle.unit_resources.items()) for bundle in pg.bundles
|
||||
)
|
||||
return (bundles, pg.strategy)
|
||||
|
||||
def placement_group_deserializer(pg_tuple):
|
||||
# We marshal this as a dictionary so that we can easily json.dumps
|
||||
# it later.
|
||||
# TODO (Alex): Would there be a benefit to properly
|
||||
# marshalling this (into a protobuf)?
|
||||
bundles = list(map(dict, pg_tuple[0]))
|
||||
return {
|
||||
"bundles": freq_of_dicts(bundles),
|
||||
"strategy": PlacementStrategy.Name(pg_tuple[1]),
|
||||
}
|
||||
|
||||
summarized_placement_groups = freq_of_dicts(
|
||||
self.get_pending_placement_groups(),
|
||||
serializer=placement_group_serializer,
|
||||
deserializer=placement_group_deserializer,
|
||||
)
|
||||
nodes_summary = freq_of_dicts(self.static_resources_by_ip.values())
|
||||
|
||||
usage_by_node = None
|
||||
if AUTOSCALER_REPORT_PER_NODE_STATUS:
|
||||
usage_by_node = {}
|
||||
for ip, totals in self.static_resources_by_ip.items():
|
||||
available = self.dynamic_resources_by_ip.get(ip, {})
|
||||
usage_by_node[ip] = {}
|
||||
for resource, total in totals.items():
|
||||
usage_by_node[ip][resource] = (
|
||||
total - available.get(resource, 0),
|
||||
total,
|
||||
)
|
||||
|
||||
return LoadMetricsSummary(
|
||||
usage=usage_dict,
|
||||
resource_demand=summarized_demand_vector,
|
||||
pg_demand=summarized_placement_groups,
|
||||
request_demand=summarized_resource_requests,
|
||||
node_types=nodes_summary,
|
||||
usage_by_node=usage_by_node,
|
||||
)
|
||||
|
||||
def set_resource_requests(self, requested_resources):
|
||||
if requested_resources is not None:
|
||||
assert isinstance(requested_resources, list), requested_resources
|
||||
self.resource_requests = [
|
||||
request for request in requested_resources if len(request) > 0
|
||||
]
|
||||
|
||||
def info_string(self):
|
||||
return " - " + "\n - ".join(
|
||||
["{}: {}".format(k, v) for k, v in sorted(self._info().items())]
|
||||
)
|
||||
|
||||
def _info(self):
|
||||
resources_used, resources_total = self._get_resource_usage()
|
||||
|
||||
now = time.time()
|
||||
idle_times = [now - t for t in self.ray_nodes_last_used_time_by_ip.values()]
|
||||
heartbeat_times = [now - t for t in self.last_heartbeat_time_by_ip.values()]
|
||||
most_delayed_heartbeats = sorted(
|
||||
self.last_heartbeat_time_by_ip.items(), key=lambda pair: pair[1]
|
||||
)[:5]
|
||||
most_delayed_heartbeats = {ip: (now - t) for ip, t in most_delayed_heartbeats}
|
||||
|
||||
def format_resource(key, value):
|
||||
if key in ["object_store_memory", "memory"]:
|
||||
return "{} GiB".format(round(value / (1024 * 1024 * 1024), 2))
|
||||
else:
|
||||
return round(value, 2)
|
||||
|
||||
return {
|
||||
"ResourceUsage": ", ".join(
|
||||
[
|
||||
"{}/{} {}".format(
|
||||
format_resource(rid, resources_used[rid]),
|
||||
format_resource(rid, resources_total[rid]),
|
||||
rid,
|
||||
)
|
||||
for rid in sorted(resources_used)
|
||||
if not rid.startswith("node:")
|
||||
]
|
||||
),
|
||||
"NodeIdleSeconds": "Min={} Mean={} Max={}".format(
|
||||
int(min(idle_times)) if idle_times else -1,
|
||||
int(float(sum(idle_times)) / len(idle_times)) if idle_times else -1,
|
||||
int(max(idle_times)) if idle_times else -1,
|
||||
),
|
||||
"TimeSinceLastHeartbeat": "Min={} Mean={} Max={}".format(
|
||||
int(min(heartbeat_times)) if heartbeat_times else -1,
|
||||
int(float(sum(heartbeat_times)) / len(heartbeat_times))
|
||||
if heartbeat_times
|
||||
else -1,
|
||||
int(max(heartbeat_times)) if heartbeat_times else -1,
|
||||
),
|
||||
"MostDelayedHeartbeats": most_delayed_heartbeats,
|
||||
}
|
||||
@@ -0,0 +1,15 @@
|
||||
import importlib
|
||||
|
||||
|
||||
def load_function_or_class(path):
|
||||
"""Load a function or class at runtime given a full path.
|
||||
|
||||
Example of the path: mypkg.mysubpkg.myclass
|
||||
"""
|
||||
class_data = path.split(".")
|
||||
if len(class_data) < 2:
|
||||
raise ValueError("You need to pass a valid path like mymodule.provider_class")
|
||||
module_path = ".".join(class_data[:-1])
|
||||
fn_or_class_str = class_data[-1]
|
||||
module = importlib.import_module(module_path)
|
||||
return getattr(module, fn_or_class_str)
|
||||
@@ -0,0 +1,131 @@
|
||||
import copy
|
||||
import os
|
||||
from typing import Any, Dict, Tuple
|
||||
|
||||
from ray._common.utils import get_default_ray_temp_dir
|
||||
from ray.autoscaler._private.cli_logger import cli_logger
|
||||
|
||||
unsupported_field_message = "The field {} is not supported for on-premise clusters."
|
||||
|
||||
LOCAL_CLUSTER_NODE_TYPE = "local.cluster.node"
|
||||
|
||||
|
||||
def prepare_local(config: Dict[str, Any]) -> Tuple[Dict[str, Any], bool]:
|
||||
"""
|
||||
Prepare local cluster config for ingestion by cluster launcher and
|
||||
autoscaler.
|
||||
"""
|
||||
config = copy.deepcopy(config)
|
||||
for field in "head_node", "worker_nodes", "available_node_types":
|
||||
if config.get(field):
|
||||
# If the config already contains the internal node type, it's been prepared via ray up already hence return as-is.
|
||||
if (
|
||||
field == "available_node_types"
|
||||
and LOCAL_CLUSTER_NODE_TYPE in config.get(field, {})
|
||||
):
|
||||
return config, False
|
||||
err_msg = unsupported_field_message.format(field)
|
||||
cli_logger.abort(err_msg)
|
||||
# We use a config with a single node type for on-prem clusters.
|
||||
# Resources internally detected by Ray are not overridden by the autoscaler
|
||||
# (see NodeProvider.do_update)
|
||||
config["available_node_types"] = {
|
||||
LOCAL_CLUSTER_NODE_TYPE: {"node_config": {}, "resources": {}}
|
||||
}
|
||||
config["head_node_type"] = LOCAL_CLUSTER_NODE_TYPE
|
||||
if "coordinator_address" in config["provider"]:
|
||||
config = prepare_coordinator(config)
|
||||
else:
|
||||
config = prepare_manual(config)
|
||||
return config, True
|
||||
|
||||
|
||||
def prepare_coordinator(config: Dict[str, Any]) -> Dict[str, Any]:
|
||||
config = copy.deepcopy(config)
|
||||
# User should explicitly set the max number of workers for the coordinator
|
||||
# to allocate.
|
||||
if "max_workers" not in config:
|
||||
cli_logger.abort(
|
||||
"The field `max_workers` is required when using an "
|
||||
"automatically managed on-premise cluster."
|
||||
)
|
||||
node_type = config["available_node_types"][LOCAL_CLUSTER_NODE_TYPE]
|
||||
# The autoscaler no longer uses global `min_workers`.
|
||||
# Move `min_workers` to the node_type config.
|
||||
node_type["min_workers"] = config.pop("min_workers", 0)
|
||||
node_type["max_workers"] = config["max_workers"]
|
||||
return config
|
||||
|
||||
|
||||
def prepare_manual(config: Dict[str, Any]) -> Dict[str, Any]:
|
||||
"""Validates and sets defaults for configs of manually managed on-prem
|
||||
clusters.
|
||||
|
||||
- Checks for presence of required `worker_ips` and `head_ips` fields.
|
||||
- Defaults min and max workers to the number of `worker_ips`.
|
||||
- Caps min and max workers at the number of `worker_ips`.
|
||||
- Writes min and max worker info into the single worker node type.
|
||||
"""
|
||||
config = copy.deepcopy(config)
|
||||
if ("worker_ips" not in config["provider"]) or (
|
||||
"head_ip" not in config["provider"]
|
||||
):
|
||||
cli_logger.abort(
|
||||
"Please supply a `head_ip` and list of `worker_ips`. "
|
||||
"Alternatively, supply a `coordinator_address`."
|
||||
)
|
||||
num_ips = len(config["provider"]["worker_ips"])
|
||||
node_type = config["available_node_types"][LOCAL_CLUSTER_NODE_TYPE]
|
||||
# Default to keeping all provided ips in the cluster.
|
||||
config.setdefault("max_workers", num_ips)
|
||||
|
||||
# The autoscaler no longer uses global `min_workers`.
|
||||
# We will move `min_workers` to the node_type config.
|
||||
min_workers = config.pop("min_workers", num_ips)
|
||||
max_workers = config["max_workers"]
|
||||
|
||||
if min_workers > num_ips:
|
||||
cli_logger.warning(
|
||||
f"The value of `min_workers` supplied ({min_workers}) is greater"
|
||||
f" than the number of available worker ips ({num_ips})."
|
||||
f" Setting `min_workers={num_ips}`."
|
||||
)
|
||||
node_type["min_workers"] = num_ips
|
||||
else:
|
||||
node_type["min_workers"] = min_workers
|
||||
|
||||
if max_workers > num_ips:
|
||||
cli_logger.warning(
|
||||
f"The value of `max_workers` supplied ({max_workers}) is greater"
|
||||
f" than the number of available worker ips ({num_ips})."
|
||||
f" Setting `max_workers={num_ips}`."
|
||||
)
|
||||
node_type["max_workers"] = num_ips
|
||||
config["max_workers"] = num_ips
|
||||
else:
|
||||
node_type["max_workers"] = max_workers
|
||||
|
||||
if max_workers < num_ips:
|
||||
cli_logger.warning(
|
||||
f"The value of `max_workers` supplied ({max_workers}) is less"
|
||||
f" than the number of available worker ips ({num_ips})."
|
||||
f" At most {max_workers} Ray worker nodes will connect to the cluster."
|
||||
)
|
||||
|
||||
return config
|
||||
|
||||
|
||||
def get_lock_path(cluster_name: str) -> str:
|
||||
return os.path.join(
|
||||
get_default_ray_temp_dir(), "cluster-{}.lock".format(cluster_name)
|
||||
)
|
||||
|
||||
|
||||
def get_state_path(cluster_name: str) -> str:
|
||||
return os.path.join(
|
||||
get_default_ray_temp_dir(), "cluster-{}.state".format(cluster_name)
|
||||
)
|
||||
|
||||
|
||||
def bootstrap_local(config: Dict[str, Any]) -> Dict[str, Any]:
|
||||
return config
|
||||
@@ -0,0 +1,110 @@
|
||||
import json
|
||||
import logging
|
||||
from http.client import RemoteDisconnected
|
||||
|
||||
from ray.autoscaler.node_provider import NodeProvider
|
||||
from ray.autoscaler.tags import TAG_RAY_CLUSTER_NAME
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class CoordinatorSenderNodeProvider(NodeProvider):
|
||||
"""NodeProvider for automatically managed private/local clusters.
|
||||
|
||||
The cluster management is handled by a remote coordinating server.
|
||||
The server listens on <coordinator_address>, therefore, the address
|
||||
should be provided in the provider section in the cluster config.
|
||||
The server receieves HTTP requests from this class and uses
|
||||
LocalNodeProvider to get their responses.
|
||||
"""
|
||||
|
||||
def __init__(self, provider_config, cluster_name):
|
||||
NodeProvider.__init__(self, provider_config, cluster_name)
|
||||
self.coordinator_address = provider_config["coordinator_address"]
|
||||
|
||||
def _get_http_response(self, request):
|
||||
headers = {
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
request_message = json.dumps(request).encode()
|
||||
http_coordinator_address = "http://" + self.coordinator_address
|
||||
|
||||
try:
|
||||
import requests # `requests` is not part of stdlib.
|
||||
from requests.exceptions import ConnectionError
|
||||
|
||||
r = requests.get(
|
||||
http_coordinator_address,
|
||||
data=request_message,
|
||||
headers=headers,
|
||||
timeout=None,
|
||||
)
|
||||
except (RemoteDisconnected, ConnectionError):
|
||||
logger.exception(
|
||||
"Could not connect to: "
|
||||
+ http_coordinator_address
|
||||
+ ". Did you run python coordinator_server.py"
|
||||
+ " --ips <list_of_node_ips> --host <HOST> --port <PORT>?"
|
||||
)
|
||||
raise
|
||||
except ImportError:
|
||||
logger.exception(
|
||||
"Not all Ray Autoscaler dependencies were found. "
|
||||
"In Ray 1.4+, the Ray CLI, autoscaler, and dashboard will "
|
||||
'only be usable via `pip install "ray[default]"`. Please '
|
||||
"update your install command."
|
||||
)
|
||||
raise
|
||||
|
||||
response = r.json()
|
||||
return response
|
||||
|
||||
def non_terminated_nodes(self, tag_filters):
|
||||
# Only get the non terminated nodes associated with this cluster name.
|
||||
tag_filters[TAG_RAY_CLUSTER_NAME] = self.cluster_name
|
||||
request = {"type": "non_terminated_nodes", "args": (tag_filters,)}
|
||||
return self._get_http_response(request)
|
||||
|
||||
def is_running(self, node_id):
|
||||
request = {"type": "is_running", "args": (node_id,)}
|
||||
return self._get_http_response(request)
|
||||
|
||||
def is_terminated(self, node_id):
|
||||
request = {"type": "is_terminated", "args": (node_id,)}
|
||||
return self._get_http_response(request)
|
||||
|
||||
def node_tags(self, node_id):
|
||||
request = {"type": "node_tags", "args": (node_id,)}
|
||||
return self._get_http_response(request)
|
||||
|
||||
def external_ip(self, node_id):
|
||||
request = {"type": "external_ip", "args": (node_id,)}
|
||||
response = self._get_http_response(request)
|
||||
return response
|
||||
|
||||
def internal_ip(self, node_id):
|
||||
request = {"type": "internal_ip", "args": (node_id,)}
|
||||
response = self._get_http_response(request)
|
||||
return response
|
||||
|
||||
def create_node(self, node_config, tags, count):
|
||||
# Tag the newly created node with this cluster name. Helps to get
|
||||
# the right nodes when calling non_terminated_nodes.
|
||||
tags[TAG_RAY_CLUSTER_NAME] = self.cluster_name
|
||||
request = {
|
||||
"type": "create_node",
|
||||
"args": (node_config, tags, count),
|
||||
}
|
||||
self._get_http_response(request)
|
||||
|
||||
def set_node_tags(self, node_id, tags):
|
||||
request = {"type": "set_node_tags", "args": (node_id, tags)}
|
||||
self._get_http_response(request)
|
||||
|
||||
def terminate_node(self, node_id):
|
||||
request = {"type": "terminate_node", "args": (node_id,)}
|
||||
self._get_http_response(request)
|
||||
|
||||
def terminate_nodes(self, node_ids):
|
||||
request = {"type": "terminate_nodes", "args": (node_ids,)}
|
||||
self._get_http_response(request)
|
||||
@@ -0,0 +1,329 @@
|
||||
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({})
|
||||
@@ -0,0 +1,33 @@
|
||||
import datetime
|
||||
import logging
|
||||
|
||||
from ray.autoscaler._private.cli_logger import cli_logger
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class LogTimer:
|
||||
def __init__(self, message, show_status=False):
|
||||
self._message = message
|
||||
self._show_status = show_status
|
||||
|
||||
def __enter__(self):
|
||||
self._start_time = datetime.datetime.utcnow()
|
||||
|
||||
def __exit__(self, *error_vals):
|
||||
if cli_logger.log_style != "record":
|
||||
return
|
||||
|
||||
td = datetime.datetime.utcnow() - self._start_time
|
||||
status = ""
|
||||
if self._show_status:
|
||||
status = "failed" if any(error_vals) else "succeeded"
|
||||
cli_logger.print(
|
||||
" ".join(
|
||||
[
|
||||
self._message,
|
||||
status,
|
||||
"[LogTimer={:.0f}ms]".format(td.total_seconds() * 1000),
|
||||
]
|
||||
)
|
||||
)
|
||||
@@ -0,0 +1,743 @@
|
||||
"""Autoscaler monitoring loop daemon."""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import signal
|
||||
import sys
|
||||
import time
|
||||
import traceback
|
||||
from collections import Counter
|
||||
from dataclasses import asdict
|
||||
from typing import Any, Callable, Dict, List, Optional, Tuple, Union
|
||||
|
||||
import ray
|
||||
import ray._private.ray_constants as ray_constants
|
||||
from ray._common.network_utils import (
|
||||
build_address,
|
||||
get_localhost_ip,
|
||||
is_localhost,
|
||||
parse_address,
|
||||
)
|
||||
from ray._common.ray_constants import (
|
||||
LOGGING_ROTATE_BACKUP_COUNT,
|
||||
LOGGING_ROTATE_BYTES,
|
||||
)
|
||||
from ray._private import logging_utils
|
||||
from ray._private.event.event_logger import get_event_logger
|
||||
from ray._private.ray_logging import setup_component_logger
|
||||
from ray._raylet import GcsClient
|
||||
from ray.autoscaler._private.autoscaler import StandardAutoscaler
|
||||
from ray.autoscaler._private.commands import teardown_cluster
|
||||
from ray.autoscaler._private.constants import (
|
||||
AUTOSCALER_MAX_RESOURCE_DEMAND_VECTOR_SIZE,
|
||||
AUTOSCALER_METRIC_PORT,
|
||||
AUTOSCALER_UPDATE_INTERVAL_S,
|
||||
DISABLE_LAUNCH_CONFIG_CHECK_KEY,
|
||||
)
|
||||
from ray.autoscaler._private.event_summarizer import EventSummarizer
|
||||
from ray.autoscaler._private.load_metrics import LoadMetrics
|
||||
from ray.autoscaler._private.prom_metrics import AutoscalerPrometheusMetrics
|
||||
from ray.autoscaler._private.util import format_readonly_node_type
|
||||
from ray.autoscaler.v2.sdk import get_cluster_resource_state
|
||||
from ray.core.generated import gcs_pb2
|
||||
from ray.core.generated.event_pb2 import Event as RayEvent
|
||||
from ray.experimental.internal_kv import (
|
||||
_initialize_internal_kv,
|
||||
_internal_kv_del,
|
||||
_internal_kv_get,
|
||||
_internal_kv_initialized,
|
||||
_internal_kv_put,
|
||||
)
|
||||
|
||||
try:
|
||||
import prometheus_client
|
||||
except ImportError:
|
||||
prometheus_client = None
|
||||
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def parse_resource_demands(
|
||||
resource_load_by_shape: "gcs_pb2.ResourceLoad",
|
||||
) -> Tuple[List[Dict], List[Dict]]:
|
||||
"""Handle the message.resource_load_by_shape protobuf for the demand
|
||||
based autoscaling. Catch and log all exceptions so this doesn't
|
||||
interfere with the utilization based autoscaler until we're confident
|
||||
this is stable. Worker queue backlogs are added to the appropriate
|
||||
resource demand vector.
|
||||
|
||||
Args:
|
||||
resource_load_by_shape: The resource demands in protobuf form or None.
|
||||
|
||||
Returns:
|
||||
Tuple of (Waiting bundles both ready and feasible, and Infeasible bundles).
|
||||
"""
|
||||
waiting_bundles, infeasible_bundles = [], []
|
||||
try:
|
||||
for resource_demand_pb in list(resource_load_by_shape.resource_demands):
|
||||
request_shape = dict(resource_demand_pb.shape)
|
||||
for _ in range(resource_demand_pb.num_ready_requests_queued):
|
||||
waiting_bundles.append(request_shape)
|
||||
for _ in range(resource_demand_pb.num_infeasible_requests_queued):
|
||||
infeasible_bundles.append(request_shape)
|
||||
|
||||
# Infeasible and ready states for tasks are (logically)
|
||||
# mutually exclusive.
|
||||
if resource_demand_pb.num_infeasible_requests_queued > 0:
|
||||
backlog_queue = infeasible_bundles
|
||||
else:
|
||||
backlog_queue = waiting_bundles
|
||||
for _ in range(resource_demand_pb.backlog_size):
|
||||
backlog_queue.append(request_shape)
|
||||
if (
|
||||
len(waiting_bundles + infeasible_bundles)
|
||||
> AUTOSCALER_MAX_RESOURCE_DEMAND_VECTOR_SIZE
|
||||
):
|
||||
break
|
||||
except Exception:
|
||||
logger.exception("Failed to parse resource demands.")
|
||||
|
||||
return waiting_bundles, infeasible_bundles
|
||||
|
||||
|
||||
# Readonly provider config (e.g., for laptop mode, manually setup clusters).
|
||||
BASE_READONLY_CONFIG = {
|
||||
"cluster_name": "default",
|
||||
"max_workers": 0,
|
||||
"upscaling_speed": 1.0,
|
||||
"docker": {},
|
||||
"idle_timeout_minutes": 0,
|
||||
"provider": {
|
||||
"type": "readonly",
|
||||
"use_node_id_as_ip": True, # For emulated multi-node on laptop.
|
||||
DISABLE_LAUNCH_CONFIG_CHECK_KEY: True, # No launch check.
|
||||
},
|
||||
"auth": {},
|
||||
"available_node_types": {
|
||||
"ray.head.default": {"resources": {}, "node_config": {}, "max_workers": 0}
|
||||
},
|
||||
"head_node_type": "ray.head.default",
|
||||
"file_mounts": {},
|
||||
"cluster_synced_files": [],
|
||||
"file_mounts_sync_continuously": False,
|
||||
"rsync_exclude": [],
|
||||
"rsync_filter": [],
|
||||
"initialization_commands": [],
|
||||
"setup_commands": [],
|
||||
"head_setup_commands": [],
|
||||
"worker_setup_commands": [],
|
||||
"head_start_ray_commands": [],
|
||||
"worker_start_ray_commands": [],
|
||||
}
|
||||
|
||||
|
||||
class Monitor:
|
||||
"""Autoscaling monitor.
|
||||
|
||||
This process periodically collects stats from the GCS and triggers
|
||||
autoscaler updates.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
address: str,
|
||||
autoscaling_config: Union[str, Callable[[], Dict[str, Any]]],
|
||||
log_dir: str = None,
|
||||
prefix_cluster_info: bool = False,
|
||||
monitor_ip: Optional[str] = None,
|
||||
retry_on_failure: bool = True,
|
||||
):
|
||||
self.gcs_address = address
|
||||
worker = ray._private.worker.global_worker
|
||||
# TODO: eventually plumb ClusterID through to here
|
||||
self.gcs_client = GcsClient(address=self.gcs_address)
|
||||
|
||||
_initialize_internal_kv(self.gcs_client)
|
||||
|
||||
if monitor_ip:
|
||||
monitor_addr = build_address(monitor_ip, AUTOSCALER_METRIC_PORT)
|
||||
self.gcs_client.internal_kv_put(
|
||||
b"AutoscalerMetricsAddress", monitor_addr.encode(), True, None
|
||||
)
|
||||
self._session_name = self.get_session_name(self.gcs_client)
|
||||
logger.info(f"session_name: {self._session_name}")
|
||||
worker.mode = 0
|
||||
head_node_ip = parse_address(self.gcs_address)[0]
|
||||
|
||||
self.load_metrics = LoadMetrics()
|
||||
self.last_avail_resources = None
|
||||
self.event_summarizer = EventSummarizer()
|
||||
self.prefix_cluster_info = prefix_cluster_info
|
||||
self.retry_on_failure = retry_on_failure
|
||||
self.autoscaling_config = autoscaling_config
|
||||
self.autoscaler = None
|
||||
# If set, we are in a manually created cluster (non-autoscaling) and
|
||||
# simply mirroring what the GCS tells us the cluster node types are.
|
||||
self.readonly_config = None
|
||||
|
||||
if log_dir:
|
||||
try:
|
||||
self.event_logger = get_event_logger(
|
||||
RayEvent.SourceType.AUTOSCALER, log_dir
|
||||
)
|
||||
except Exception:
|
||||
self.event_logger = None
|
||||
else:
|
||||
self.event_logger = None
|
||||
|
||||
self.prom_metrics = AutoscalerPrometheusMetrics(session_name=self._session_name)
|
||||
|
||||
if monitor_ip and prometheus_client:
|
||||
# If monitor_ip wasn't passed in, then don't attempt to start the
|
||||
# metric server to keep behavior identical to before metrics were
|
||||
# introduced
|
||||
try:
|
||||
logger.info(
|
||||
"Starting autoscaler metrics server on port {}".format(
|
||||
AUTOSCALER_METRIC_PORT
|
||||
)
|
||||
)
|
||||
kwargs = (
|
||||
{"addr": get_localhost_ip()} if is_localhost(head_node_ip) else {}
|
||||
)
|
||||
prometheus_client.start_http_server(
|
||||
port=AUTOSCALER_METRIC_PORT,
|
||||
registry=self.prom_metrics.registry,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
# Reset some gauges, since we don't know which labels have
|
||||
# leaked if the autoscaler was restarted.
|
||||
self.prom_metrics.pending_nodes.clear()
|
||||
self.prom_metrics.active_nodes.clear()
|
||||
except Exception:
|
||||
logger.exception(
|
||||
"An exception occurred while starting the metrics server."
|
||||
)
|
||||
elif not prometheus_client:
|
||||
logger.warning(
|
||||
"`prometheus_client` not found, so metrics will not be exported."
|
||||
)
|
||||
|
||||
logger.info("Monitor: Started")
|
||||
|
||||
def _initialize_autoscaler(self):
|
||||
if self.autoscaling_config:
|
||||
autoscaling_config = self.autoscaling_config
|
||||
else:
|
||||
# This config mirrors the current setup of the manually created
|
||||
# cluster. Each node gets its own unique node type.
|
||||
self.readonly_config = BASE_READONLY_CONFIG
|
||||
|
||||
# Note that the "available_node_types" of the config can change.
|
||||
def get_latest_readonly_config():
|
||||
return self.readonly_config
|
||||
|
||||
autoscaling_config = get_latest_readonly_config
|
||||
self.autoscaler = StandardAutoscaler(
|
||||
autoscaling_config,
|
||||
self.load_metrics,
|
||||
self.gcs_client,
|
||||
self._session_name,
|
||||
prefix_cluster_info=self.prefix_cluster_info,
|
||||
event_summarizer=self.event_summarizer,
|
||||
prom_metrics=self.prom_metrics,
|
||||
)
|
||||
|
||||
def update_load_metrics(self):
|
||||
"""Fetches resource usage data from GCS and updates load metrics."""
|
||||
|
||||
# TODO(jinbum-kim): Still needed since some fields aren't in cluster_resource_state.
|
||||
# Remove after v1 autoscaler fully migrates to get_cluster_resource_state().
|
||||
# ref: https://github.com/ray-project/ray/pull/57130
|
||||
response = self.gcs_client.get_all_resource_usage(timeout=60)
|
||||
resources_batch_data = response.resource_usage_data
|
||||
log_resource_batch_data_if_desired(resources_batch_data)
|
||||
|
||||
# This is a workaround to get correct idle_duration_ms
|
||||
# from "get_cluster_resource_state"
|
||||
# ref: https://github.com/ray-project/ray/pull/48519#issuecomment-2481659346
|
||||
cluster_resource_state = get_cluster_resource_state(self.gcs_client)
|
||||
ray_node_states = cluster_resource_state.node_states
|
||||
ray_nodes_idle_duration_ms_by_id = {
|
||||
node.node_id: node.idle_duration_ms for node in ray_node_states
|
||||
}
|
||||
|
||||
# Tell the readonly node provider what nodes to report.
|
||||
if self.readonly_config:
|
||||
new_nodes = []
|
||||
for msg in list(cluster_resource_state.node_states):
|
||||
node_id = msg.node_id.hex()
|
||||
new_nodes.append((node_id, msg.node_ip_address))
|
||||
self.autoscaler.provider._set_nodes(new_nodes)
|
||||
|
||||
waiting_bundles, infeasible_bundles = parse_resource_demands(
|
||||
resources_batch_data.resource_load_by_shape
|
||||
)
|
||||
|
||||
pending_placement_groups = list(
|
||||
resources_batch_data.placement_group_load.placement_group_data
|
||||
)
|
||||
|
||||
mirror_node_types = {}
|
||||
for resource_message in cluster_resource_state.node_states:
|
||||
node_id = resource_message.node_id
|
||||
# Generate node type config based on GCS reported node list.
|
||||
if self.readonly_config:
|
||||
# Keep prefix in sync with ReadonlyNodeProvider.
|
||||
node_type = format_readonly_node_type(node_id.hex())
|
||||
resources = {}
|
||||
for k, v in resource_message.total_resources.items():
|
||||
resources[k] = v
|
||||
mirror_node_types[node_type] = {
|
||||
"resources": resources,
|
||||
"node_config": {},
|
||||
"max_workers": 1,
|
||||
}
|
||||
total_resources = dict(resource_message.total_resources)
|
||||
available_resources = dict(resource_message.available_resources)
|
||||
|
||||
use_node_id_as_ip = self.autoscaler is not None and self.autoscaler.config[
|
||||
"provider"
|
||||
].get("use_node_id_as_ip", False)
|
||||
|
||||
# "use_node_id_as_ip" is a hack meant to address situations in
|
||||
# which there's more than one Ray node residing at a given ip.
|
||||
# TODO (Dmitri): Stop using ips as node identifiers.
|
||||
# https://github.com/ray-project/ray/issues/19086
|
||||
if use_node_id_as_ip:
|
||||
peloton_id = total_resources.get("NODE_ID_AS_RESOURCE")
|
||||
# Legacy support https://github.com/ray-project/ray/pull/17312
|
||||
if peloton_id is not None:
|
||||
ip = str(int(peloton_id))
|
||||
else:
|
||||
ip = node_id.hex()
|
||||
else:
|
||||
ip = resource_message.node_ip_address
|
||||
|
||||
idle_duration_s = 0.0
|
||||
if node_id in ray_nodes_idle_duration_ms_by_id:
|
||||
idle_duration_s = ray_nodes_idle_duration_ms_by_id[node_id] / 1000
|
||||
else:
|
||||
logger.warning(
|
||||
f"node_id {node_id} not found in ray_nodes_idle_duration_ms_by_id"
|
||||
)
|
||||
|
||||
self.load_metrics.update(
|
||||
ip,
|
||||
node_id,
|
||||
total_resources,
|
||||
available_resources,
|
||||
idle_duration_s,
|
||||
waiting_bundles,
|
||||
infeasible_bundles,
|
||||
pending_placement_groups,
|
||||
)
|
||||
if self.readonly_config:
|
||||
self.readonly_config["available_node_types"].update(mirror_node_types)
|
||||
|
||||
def get_session_name(self, gcs_client: GcsClient) -> Optional[str]:
|
||||
"""Obtain the session name from the GCS.
|
||||
|
||||
If the GCS doesn't respond, session name is considered None.
|
||||
In this case, the metrics reported from the monitor won't have
|
||||
the correct session name.
|
||||
"""
|
||||
if not _internal_kv_initialized():
|
||||
return None
|
||||
|
||||
session_name = gcs_client.internal_kv_get(
|
||||
b"session_name",
|
||||
ray_constants.KV_NAMESPACE_SESSION,
|
||||
timeout=10,
|
||||
)
|
||||
|
||||
if session_name:
|
||||
session_name = session_name.decode()
|
||||
|
||||
return session_name
|
||||
|
||||
def update_resource_requests(self):
|
||||
"""Fetches resource requests from the internal KV and updates load."""
|
||||
if not _internal_kv_initialized():
|
||||
return
|
||||
data = _internal_kv_get(
|
||||
ray._private.ray_constants.AUTOSCALER_RESOURCE_REQUEST_CHANNEL
|
||||
)
|
||||
if data:
|
||||
try:
|
||||
resource_request = json.loads(data)
|
||||
self.load_metrics.set_resource_requests(resource_request)
|
||||
except Exception:
|
||||
logger.exception("Error parsing resource requests")
|
||||
|
||||
def _run(self):
|
||||
"""Run the monitor loop."""
|
||||
|
||||
while True:
|
||||
try:
|
||||
gcs_request_start_time = time.time()
|
||||
self.update_load_metrics()
|
||||
gcs_request_time = time.time() - gcs_request_start_time
|
||||
self.update_resource_requests()
|
||||
self.update_event_summary()
|
||||
load_metrics_summary = self.load_metrics.summary()
|
||||
status = {
|
||||
"gcs_request_time": gcs_request_time,
|
||||
"time": time.time(),
|
||||
"monitor_pid": os.getpid(),
|
||||
}
|
||||
|
||||
if self.autoscaler and not self.load_metrics:
|
||||
# load_metrics is Falsey iff we haven't collected any
|
||||
# resource messages from the GCS, which can happen at startup if
|
||||
# the GCS hasn't yet received data from the Raylets.
|
||||
# In this case, do not do an autoscaler update.
|
||||
# Wait to get load metrics.
|
||||
logger.info(
|
||||
"Autoscaler has not yet received load metrics. Waiting."
|
||||
)
|
||||
elif self.autoscaler:
|
||||
# Process autoscaling actions
|
||||
update_start_time = time.time()
|
||||
self.autoscaler.update()
|
||||
status["autoscaler_update_time"] = time.time() - update_start_time
|
||||
autoscaler_summary = self.autoscaler.summary()
|
||||
try:
|
||||
self.emit_metrics(
|
||||
load_metrics_summary,
|
||||
autoscaler_summary,
|
||||
self.autoscaler.all_node_types,
|
||||
)
|
||||
except Exception:
|
||||
logger.exception("Error emitting metrics")
|
||||
|
||||
if autoscaler_summary:
|
||||
status["autoscaler_report"] = asdict(autoscaler_summary)
|
||||
status[
|
||||
"non_terminated_nodes_time"
|
||||
] = (
|
||||
self.autoscaler.non_terminated_nodes.non_terminated_nodes_time # noqa: E501
|
||||
)
|
||||
|
||||
for msg in self.event_summarizer.summary():
|
||||
# Need to prefix each line of the message for the lines to
|
||||
# get pushed to the driver logs.
|
||||
for line in msg.split("\n"):
|
||||
logger.info(
|
||||
"{}{}".format(
|
||||
ray_constants.LOG_PREFIX_EVENT_SUMMARY, line
|
||||
)
|
||||
)
|
||||
if self.event_logger:
|
||||
self.event_logger.info(line)
|
||||
|
||||
self.event_summarizer.clear()
|
||||
|
||||
status["load_metrics_report"] = asdict(load_metrics_summary)
|
||||
as_json = json.dumps(status)
|
||||
if _internal_kv_initialized():
|
||||
_internal_kv_put(
|
||||
ray_constants.DEBUG_AUTOSCALING_STATUS, as_json, overwrite=True
|
||||
)
|
||||
except Exception:
|
||||
# By default, do not exit the monitor on failure.
|
||||
if self.retry_on_failure:
|
||||
logger.exception("Monitor: Execution exception. Trying again...")
|
||||
else:
|
||||
raise
|
||||
|
||||
# Wait for a autoscaler update interval before processing the next
|
||||
# round of messages.
|
||||
time.sleep(AUTOSCALER_UPDATE_INTERVAL_S)
|
||||
|
||||
def emit_metrics(self, load_metrics_summary, autoscaler_summary, node_types):
|
||||
if autoscaler_summary is None:
|
||||
return None
|
||||
|
||||
for resource_name in ["CPU", "GPU", "TPU"]:
|
||||
_, total = load_metrics_summary.usage.get(resource_name, (0, 0))
|
||||
pending = autoscaler_summary.pending_resources.get(resource_name, 0)
|
||||
self.prom_metrics.cluster_resources.labels(
|
||||
resource=resource_name,
|
||||
SessionName=self.prom_metrics.session_name,
|
||||
).set(total)
|
||||
self.prom_metrics.pending_resources.labels(
|
||||
resource=resource_name,
|
||||
SessionName=self.prom_metrics.session_name,
|
||||
).set(pending)
|
||||
|
||||
pending_node_count = Counter()
|
||||
for _, node_type, _ in autoscaler_summary.pending_nodes:
|
||||
pending_node_count[node_type] += 1
|
||||
|
||||
for node_type, count in autoscaler_summary.pending_launches.items():
|
||||
pending_node_count[node_type] += count
|
||||
|
||||
for node_type in node_types:
|
||||
count = pending_node_count[node_type]
|
||||
self.prom_metrics.pending_nodes.labels(
|
||||
SessionName=self.prom_metrics.session_name,
|
||||
NodeType=node_type,
|
||||
).set(count)
|
||||
|
||||
for node_type in node_types:
|
||||
count = autoscaler_summary.active_nodes.get(node_type, 0)
|
||||
self.prom_metrics.active_nodes.labels(
|
||||
SessionName=self.prom_metrics.session_name,
|
||||
NodeType=node_type,
|
||||
).set(count)
|
||||
|
||||
failed_node_counts = Counter()
|
||||
for _, node_type in autoscaler_summary.failed_nodes:
|
||||
failed_node_counts[node_type] += 1
|
||||
|
||||
# NOTE: This metric isn't reset with monitor resets. This means it will
|
||||
# only be updated when the autoscaler' node tracker remembers failed
|
||||
# nodes. If the node type failure is evicted from the autoscaler, the
|
||||
# metric may not update for a while.
|
||||
for node_type, count in failed_node_counts.items():
|
||||
self.prom_metrics.recently_failed_nodes.labels(
|
||||
SessionName=self.prom_metrics.session_name,
|
||||
NodeType=node_type,
|
||||
).set(count)
|
||||
|
||||
def update_event_summary(self):
|
||||
"""Report the current size of the cluster.
|
||||
|
||||
To avoid log spam, only cluster size changes (CPU, GPU or TPU count change)
|
||||
are reported to the event summarizer. The event summarizer will report
|
||||
only the latest cluster size per batch.
|
||||
"""
|
||||
avail_resources = self.load_metrics.resources_avail_summary()
|
||||
if not self.readonly_config and avail_resources != self.last_avail_resources:
|
||||
self.event_summarizer.add(
|
||||
"Resized to {}.", # e.g., Resized to 100 CPUs, 4 GPUs, 4 TPUs.
|
||||
quantity=avail_resources,
|
||||
aggregate=lambda old, new: new,
|
||||
)
|
||||
self.last_avail_resources = avail_resources
|
||||
|
||||
def destroy_autoscaler_workers(self):
|
||||
"""Cleanup the autoscaler, in case of an exception in the run() method.
|
||||
|
||||
We kill the worker nodes, but retain the head node in order to keep
|
||||
logs around, keeping costs minimal. This monitor process runs on the
|
||||
head node anyway, so this is more reliable."""
|
||||
|
||||
if self.autoscaler is None:
|
||||
return # Nothing to clean up.
|
||||
|
||||
if self.autoscaling_config is None:
|
||||
# This is a logic error in the program. Can't do anything.
|
||||
logger.error("Monitor: Cleanup failed due to lack of autoscaler config.")
|
||||
return
|
||||
|
||||
logger.info("Monitor: Exception caught. Taking down workers...")
|
||||
clean = False
|
||||
while not clean:
|
||||
try:
|
||||
teardown_cluster(
|
||||
config_file=self.autoscaling_config,
|
||||
yes=True, # Non-interactive.
|
||||
workers_only=True, # Retain head node for logs.
|
||||
override_cluster_name=None,
|
||||
keep_min_workers=True, # Retain minimal amount of workers.
|
||||
)
|
||||
clean = True
|
||||
logger.info("Monitor: Workers taken down.")
|
||||
except Exception:
|
||||
logger.error("Monitor: Cleanup exception. Trying again...")
|
||||
time.sleep(2)
|
||||
|
||||
def _handle_failure(self, error):
|
||||
if (
|
||||
self.autoscaler is not None
|
||||
and os.environ.get("RAY_AUTOSCALER_FATESHARE_WORKERS", "") == "1"
|
||||
):
|
||||
self.autoscaler.kill_workers()
|
||||
# Take down autoscaler workers if necessary.
|
||||
self.destroy_autoscaler_workers()
|
||||
|
||||
# Something went wrong, so push an error to all current and future
|
||||
# drivers.
|
||||
message = f"The autoscaler failed with the following error:\n{error}"
|
||||
if _internal_kv_initialized():
|
||||
_internal_kv_put(
|
||||
ray_constants.DEBUG_AUTOSCALING_ERROR, message, overwrite=True
|
||||
)
|
||||
from ray._private.utils import publish_error_to_driver
|
||||
|
||||
publish_error_to_driver(
|
||||
ray_constants.MONITOR_DIED_ERROR,
|
||||
message,
|
||||
gcs_client=self.gcs_client,
|
||||
)
|
||||
|
||||
def _signal_handler(self, sig, frame):
|
||||
try:
|
||||
self._handle_failure(
|
||||
f"Terminated with signal {sig}\n"
|
||||
+ "".join(traceback.format_stack(frame))
|
||||
)
|
||||
except Exception:
|
||||
logger.exception("Monitor: Failure in signal handler.")
|
||||
sys.exit(sig + 128)
|
||||
|
||||
def run(self):
|
||||
# Register signal handlers for autoscaler termination.
|
||||
# Signals will not be received on windows
|
||||
signal.signal(signal.SIGINT, self._signal_handler)
|
||||
signal.signal(signal.SIGTERM, self._signal_handler)
|
||||
try:
|
||||
if _internal_kv_initialized():
|
||||
# Delete any previous autoscaling errors.
|
||||
_internal_kv_del(ray_constants.DEBUG_AUTOSCALING_ERROR)
|
||||
self._initialize_autoscaler()
|
||||
self._run()
|
||||
except Exception:
|
||||
logger.exception("Error in monitor loop")
|
||||
self._handle_failure(traceback.format_exc())
|
||||
raise
|
||||
|
||||
|
||||
def log_resource_batch_data_if_desired(
|
||||
resources_batch_data: gcs_pb2.ResourceUsageBatchData,
|
||||
) -> None:
|
||||
if os.getenv("AUTOSCALER_LOG_RESOURCE_BATCH_DATA") == "1":
|
||||
logger.info("Logging raw resource message pulled from GCS.")
|
||||
logger.info(resources_batch_data)
|
||||
logger.info("Done logging raw resource message.")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser(
|
||||
description=("Parse GCS server for the monitor to connect to.")
|
||||
)
|
||||
parser.add_argument(
|
||||
"--gcs-address", required=False, type=str, help="The address (ip:port) of GCS."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--autoscaling-config",
|
||||
required=False,
|
||||
type=str,
|
||||
help="the path to the autoscaling config file",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--logging-level",
|
||||
required=False,
|
||||
type=str,
|
||||
default=ray_constants.LOGGER_LEVEL,
|
||||
choices=ray_constants.LOGGER_LEVEL_CHOICES,
|
||||
help=ray_constants.LOGGER_LEVEL_HELP,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--logging-format",
|
||||
required=False,
|
||||
type=str,
|
||||
default=ray_constants.LOGGER_FORMAT,
|
||||
help=ray_constants.LOGGER_FORMAT_HELP,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--logging-filename",
|
||||
required=False,
|
||||
type=str,
|
||||
default=ray_constants.MONITOR_LOG_FILE_NAME,
|
||||
help="Specify the name of log file, "
|
||||
"log to stdout if set empty, default is "
|
||||
f'"{ray_constants.MONITOR_LOG_FILE_NAME}"',
|
||||
)
|
||||
parser.add_argument(
|
||||
"--logs-dir",
|
||||
required=True,
|
||||
type=str,
|
||||
help="Specify the path of the temporary directory used by Ray processes.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--logging-rotate-bytes",
|
||||
required=False,
|
||||
type=int,
|
||||
default=LOGGING_ROTATE_BYTES,
|
||||
help="Specify the max bytes for rotating "
|
||||
"log file, default is "
|
||||
f"{LOGGING_ROTATE_BYTES} bytes.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--logging-rotate-backup-count",
|
||||
required=False,
|
||||
type=int,
|
||||
default=LOGGING_ROTATE_BACKUP_COUNT,
|
||||
help="Specify the backup count of rotated log file, default is "
|
||||
f"{LOGGING_ROTATE_BACKUP_COUNT}.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--monitor-ip",
|
||||
required=False,
|
||||
type=str,
|
||||
default=None,
|
||||
help="The IP address of the machine hosting the monitor process.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--stdout-filepath",
|
||||
required=False,
|
||||
type=str,
|
||||
default="",
|
||||
help="The filepath to dump monitor stdout.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--stderr-filepath",
|
||||
required=False,
|
||||
type=str,
|
||||
default="",
|
||||
help="The filepath to dump monitor stderr.",
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
# Disable log rotation for windows, because NTFS doesn't allow file deletion when there're multiple owners or borrowers, which happens to be how ray accesses log files.
|
||||
logging_rotation_bytes = args.logging_rotate_bytes if sys.platform != "win32" else 0
|
||||
logging_rotation_backup_count = (
|
||||
args.logging_rotate_backup_count if sys.platform != "win32" else 1
|
||||
)
|
||||
setup_component_logger(
|
||||
logging_level=args.logging_level,
|
||||
logging_format=args.logging_format,
|
||||
log_dir=args.logs_dir,
|
||||
filename=args.logging_filename,
|
||||
max_bytes=logging_rotation_bytes,
|
||||
backup_count=logging_rotation_backup_count,
|
||||
)
|
||||
|
||||
# Setup stdout/stderr redirect files if redirection enabled.
|
||||
logging_utils.redirect_stdout_stderr_if_needed(
|
||||
args.stdout_filepath,
|
||||
args.stderr_filepath,
|
||||
logging_rotation_bytes,
|
||||
logging_rotation_backup_count,
|
||||
)
|
||||
|
||||
logger.info(f"Starting monitor using ray installation: {ray.__file__}")
|
||||
logger.info(f"Ray version: {ray.__version__}")
|
||||
logger.info(f"Ray commit: {ray.__commit__}")
|
||||
logger.info(f"Monitor started with command: {sys.argv}")
|
||||
|
||||
if args.autoscaling_config:
|
||||
autoscaling_config = os.path.expanduser(args.autoscaling_config)
|
||||
else:
|
||||
autoscaling_config = None
|
||||
|
||||
bootstrap_address = args.gcs_address
|
||||
if bootstrap_address is None:
|
||||
raise ValueError("--gcs-address must be set!")
|
||||
|
||||
monitor = Monitor(
|
||||
bootstrap_address,
|
||||
autoscaling_config,
|
||||
log_dir=args.logs_dir,
|
||||
monitor_ip=args.monitor_ip,
|
||||
)
|
||||
|
||||
monitor.run()
|
||||
@@ -0,0 +1,221 @@
|
||||
import copy
|
||||
import logging
|
||||
import operator
|
||||
import threading
|
||||
import time
|
||||
import traceback
|
||||
from typing import Any, Dict, Optional
|
||||
|
||||
from ray.autoscaler._private.node_provider_availability_tracker import (
|
||||
NodeProviderAvailabilityTracker,
|
||||
)
|
||||
from ray.autoscaler._private.prom_metrics import AutoscalerPrometheusMetrics
|
||||
from ray.autoscaler._private.util import hash_launch_conf
|
||||
from ray.autoscaler.node_launch_exception import NodeLaunchException
|
||||
from ray.autoscaler.tags import (
|
||||
NODE_KIND_WORKER,
|
||||
STATUS_UNINITIALIZED,
|
||||
TAG_RAY_LAUNCH_CONFIG,
|
||||
TAG_RAY_NODE_KIND,
|
||||
TAG_RAY_NODE_NAME,
|
||||
TAG_RAY_NODE_STATUS,
|
||||
TAG_RAY_USER_NODE_TYPE,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class BaseNodeLauncher:
|
||||
"""Launches Ray nodes in the main thread using
|
||||
`BaseNodeLauncher.launch_node()`.
|
||||
|
||||
This is a superclass of NodeLauncher, which launches nodes asynchronously
|
||||
in the background.
|
||||
|
||||
By default, the subclass NodeLauncher is used to launch nodes in subthreads.
|
||||
That behavior can be flagged off in the provider config by setting
|
||||
`foreground_node_launch: True`; the autoscaler will then makes blocking calls to
|
||||
BaseNodeLauncher.launch_node() in the main thread.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
provider,
|
||||
pending,
|
||||
event_summarizer,
|
||||
node_provider_availability_tracker: NodeProviderAvailabilityTracker,
|
||||
session_name: Optional[str] = None,
|
||||
prom_metrics=None,
|
||||
node_types=None,
|
||||
index=None,
|
||||
*args,
|
||||
**kwargs,
|
||||
):
|
||||
self.pending = pending
|
||||
self.event_summarizer = event_summarizer
|
||||
self.node_provider_availability_tracker = node_provider_availability_tracker
|
||||
self.prom_metrics = prom_metrics or AutoscalerPrometheusMetrics(
|
||||
session_name=session_name
|
||||
)
|
||||
self.provider = provider
|
||||
self.node_types = node_types
|
||||
self.index = str(index) if index is not None else ""
|
||||
|
||||
def launch_node(
|
||||
self, config: Dict[str, Any], count: int, node_type: str
|
||||
) -> Optional[Dict]:
|
||||
self.log("Got {} nodes to launch.".format(count))
|
||||
created_nodes = self._launch_node(config, count, node_type)
|
||||
self.pending.dec(node_type, count)
|
||||
return created_nodes
|
||||
|
||||
def _launch_node(
|
||||
self, config: Dict[str, Any], count: int, node_type: str
|
||||
) -> Optional[Dict]:
|
||||
if self.node_types:
|
||||
assert node_type, node_type
|
||||
|
||||
# The `worker_nodes` field is deprecated in favor of per-node-type
|
||||
# node_configs. We allow it for backwards-compatibility.
|
||||
launch_config = copy.deepcopy(config.get("worker_nodes", {}))
|
||||
if node_type:
|
||||
launch_config.update(
|
||||
config["available_node_types"][node_type]["node_config"]
|
||||
)
|
||||
resources = copy.deepcopy(
|
||||
config["available_node_types"][node_type]["resources"]
|
||||
)
|
||||
labels = copy.deepcopy(
|
||||
config["available_node_types"][node_type].get("labels", {})
|
||||
)
|
||||
launch_hash = hash_launch_conf(launch_config, config["auth"])
|
||||
node_config = copy.deepcopy(config.get("worker_nodes", {}))
|
||||
node_tags = {
|
||||
TAG_RAY_NODE_NAME: "ray-{}-worker".format(config["cluster_name"]),
|
||||
TAG_RAY_NODE_KIND: NODE_KIND_WORKER,
|
||||
TAG_RAY_NODE_STATUS: STATUS_UNINITIALIZED,
|
||||
TAG_RAY_LAUNCH_CONFIG: launch_hash,
|
||||
}
|
||||
# A custom node type is specified; set the tag in this case, and also
|
||||
# merge the configs. We merge the configs instead of overriding, so
|
||||
# that the bootstrapped per-cloud properties are preserved.
|
||||
# TODO(ekl) this logic is duplicated in commands.py (keep in sync)
|
||||
if node_type:
|
||||
node_tags[TAG_RAY_USER_NODE_TYPE] = node_type
|
||||
node_config.update(launch_config)
|
||||
|
||||
node_launch_start_time = time.time()
|
||||
|
||||
error_msg = None
|
||||
full_exception = None
|
||||
created_nodes = {}
|
||||
try:
|
||||
created_nodes = self.provider.create_node_with_resources_and_labels(
|
||||
node_config, node_tags, count, resources, labels
|
||||
)
|
||||
except NodeLaunchException as node_launch_exception:
|
||||
self.node_provider_availability_tracker.update_node_availability(
|
||||
node_type, int(node_launch_start_time), node_launch_exception
|
||||
)
|
||||
|
||||
if node_launch_exception.src_exc_info is not None:
|
||||
full_exception = "\n".join(
|
||||
traceback.format_exception(*node_launch_exception.src_exc_info)
|
||||
)
|
||||
|
||||
error_msg = (
|
||||
f"Failed to launch {{}} node(s) of type {node_type}. "
|
||||
f"({node_launch_exception.category}): "
|
||||
f"{node_launch_exception.description}"
|
||||
)
|
||||
except Exception:
|
||||
error_msg = f"Failed to launch {{}} node(s) of type {node_type}."
|
||||
full_exception = traceback.format_exc()
|
||||
else:
|
||||
# Record some metrics/observability information when a node is launched.
|
||||
launch_time = time.time() - node_launch_start_time
|
||||
for _ in range(count):
|
||||
# Note: when launching multiple nodes we observe the time it
|
||||
# took all nodes to launch for each node. For example, if 4
|
||||
# nodes were created in 25 seconds, we would observe the 25
|
||||
# second create time 4 times.
|
||||
self.prom_metrics.worker_create_node_time.observe(launch_time)
|
||||
self.prom_metrics.started_nodes.inc(count)
|
||||
self.node_provider_availability_tracker.update_node_availability(
|
||||
node_type=node_type,
|
||||
timestamp=int(node_launch_start_time),
|
||||
node_launch_exception=None,
|
||||
)
|
||||
|
||||
if error_msg is not None:
|
||||
self.event_summarizer.add(
|
||||
error_msg,
|
||||
quantity=count,
|
||||
aggregate=operator.add,
|
||||
)
|
||||
self.log(error_msg)
|
||||
self.prom_metrics.node_launch_exceptions.inc()
|
||||
self.prom_metrics.failed_create_nodes.inc(count)
|
||||
else:
|
||||
self.log("Launching {} nodes, type {}.".format(count, node_type))
|
||||
self.event_summarizer.add(
|
||||
"Adding {} node(s) of type " + str(node_type) + ".",
|
||||
quantity=count,
|
||||
aggregate=operator.add,
|
||||
)
|
||||
|
||||
if full_exception is not None:
|
||||
self.log(full_exception)
|
||||
|
||||
return created_nodes
|
||||
|
||||
def log(self, statement):
|
||||
# launcher_class is "BaseNodeLauncher", or "NodeLauncher" if called
|
||||
# from that subclass.
|
||||
launcher_class: str = type(self).__name__
|
||||
prefix = "{}{}:".format(launcher_class, self.index)
|
||||
logger.info(prefix + " {}".format(statement))
|
||||
|
||||
|
||||
class NodeLauncher(BaseNodeLauncher, threading.Thread):
|
||||
"""Launches nodes asynchronously in the background."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
provider,
|
||||
queue,
|
||||
pending,
|
||||
event_summarizer,
|
||||
node_provider_availability_tracker,
|
||||
session_name: Optional[str] = None,
|
||||
prom_metrics=None,
|
||||
node_types=None,
|
||||
index=None,
|
||||
*thread_args,
|
||||
**thread_kwargs,
|
||||
):
|
||||
self.queue = queue
|
||||
BaseNodeLauncher.__init__(
|
||||
self,
|
||||
provider=provider,
|
||||
pending=pending,
|
||||
event_summarizer=event_summarizer,
|
||||
session_name=session_name,
|
||||
node_provider_availability_tracker=node_provider_availability_tracker,
|
||||
prom_metrics=prom_metrics,
|
||||
node_types=node_types,
|
||||
index=index,
|
||||
)
|
||||
threading.Thread.__init__(self, *thread_args, **thread_kwargs)
|
||||
|
||||
def run(self):
|
||||
"""Collects launch data from queue populated by StandardAutoscaler.
|
||||
Launches nodes in a background thread.
|
||||
|
||||
Overrides threading.Thread.run().
|
||||
NodeLauncher.start() executes this loop in a background thread.
|
||||
"""
|
||||
while True:
|
||||
config, count, node_type = self.queue.get()
|
||||
# launch_node is implemented in BaseNodeLauncher
|
||||
self.launch_node(config, count, node_type)
|
||||
@@ -0,0 +1,165 @@
|
||||
import threading
|
||||
import time
|
||||
from dataclasses import dataclass
|
||||
from typing import Callable, Dict, Optional, Tuple
|
||||
|
||||
from ray.autoscaler._private.constants import (
|
||||
AUTOSCALER_NODE_AVAILABILITY_MAX_STALENESS_S,
|
||||
)
|
||||
from ray.autoscaler.node_launch_exception import NodeLaunchException
|
||||
|
||||
|
||||
@dataclass
|
||||
class UnavailableNodeInformation:
|
||||
category: str
|
||||
description: str
|
||||
|
||||
|
||||
@dataclass
|
||||
class NodeAvailabilityRecord:
|
||||
node_type: str
|
||||
is_available: bool
|
||||
last_checked_timestamp: float
|
||||
unavailable_node_information: Optional[UnavailableNodeInformation]
|
||||
|
||||
|
||||
@dataclass
|
||||
class NodeAvailabilitySummary:
|
||||
node_availabilities: Dict[
|
||||
str, NodeAvailabilityRecord
|
||||
] # Mapping from node type to node availability record.
|
||||
|
||||
@classmethod
|
||||
def from_fields(cls, **fields) -> Optional["NodeAvailabilitySummary"]:
|
||||
"""Implement marshalling from nested fields. pydantic isn't a core dependency
|
||||
so we're implementing this by hand instead."""
|
||||
parsed = {}
|
||||
|
||||
node_availabilites_dict = fields.get("node_availabilities", {})
|
||||
|
||||
for node_type, node_availability_record_dict in node_availabilites_dict.items():
|
||||
unavailable_information_dict = node_availability_record_dict.pop(
|
||||
"unavailable_node_information", None
|
||||
)
|
||||
unavaiable_information = None
|
||||
if unavailable_information_dict is not None:
|
||||
unavaiable_information = UnavailableNodeInformation(
|
||||
**unavailable_information_dict
|
||||
)
|
||||
|
||||
parsed[node_type] = NodeAvailabilityRecord(
|
||||
unavailable_node_information=unavaiable_information,
|
||||
**node_availability_record_dict,
|
||||
)
|
||||
|
||||
return NodeAvailabilitySummary(node_availabilities=parsed)
|
||||
|
||||
def __eq__(self, other: "NodeAvailabilitySummary"):
|
||||
return self.node_availabilities == other.node_availabilities
|
||||
|
||||
def __bool__(self) -> bool:
|
||||
return bool(self.node_availabilities)
|
||||
|
||||
|
||||
class NodeProviderAvailabilityTracker:
|
||||
"""A thread safe, TTL cache of node provider availability. We don't use
|
||||
cachetools.TTLCache because it always sets the expiration time relative to
|
||||
insertion time, but in our case, we want entries to expire relative to when
|
||||
the node creation was attempted (and entries aren't necessarily added in
|
||||
order). We want the entries to expire because the information grows stale
|
||||
over time.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
timer: Callable[[], float] = time.time,
|
||||
ttl: float = AUTOSCALER_NODE_AVAILABILITY_MAX_STALENESS_S,
|
||||
):
|
||||
"""A cache that tracks the availability of nodes and throw away
|
||||
entries which have grown too stale.
|
||||
|
||||
Args:
|
||||
timer: A function that returns the current time in seconds.
|
||||
ttl: The ttl from the insertion timestamp of an entry.
|
||||
"""
|
||||
self.timer = timer
|
||||
self.ttl = ttl
|
||||
# Mapping from node type to (eviction_time, record)
|
||||
self.store: Dict[str, Tuple[float, NodeAvailabilityRecord]] = {}
|
||||
# A global lock to simplify thread safety handling.
|
||||
self.lock = threading.RLock()
|
||||
|
||||
def _update_node_availability_requires_lock(
|
||||
self,
|
||||
node_type: str,
|
||||
timestamp: int,
|
||||
node_launch_exception: Optional[NodeLaunchException],
|
||||
) -> None:
|
||||
if node_launch_exception is None:
|
||||
record = NodeAvailabilityRecord(
|
||||
node_type=node_type,
|
||||
is_available=True,
|
||||
last_checked_timestamp=timestamp,
|
||||
unavailable_node_information=None,
|
||||
)
|
||||
else:
|
||||
info = UnavailableNodeInformation(
|
||||
category=node_launch_exception.category,
|
||||
description=node_launch_exception.description,
|
||||
)
|
||||
record = NodeAvailabilityRecord(
|
||||
node_type=node_type,
|
||||
is_available=False,
|
||||
last_checked_timestamp=timestamp,
|
||||
unavailable_node_information=info,
|
||||
)
|
||||
|
||||
expiration_time = timestamp + self.ttl
|
||||
|
||||
# TODO (Alex): In theory it would be nice to make this dictionary
|
||||
# ordered by expiration time, unfortunately that's a bit difficult
|
||||
# since `update_node_availability` can be called with out of order
|
||||
# timestamps.
|
||||
self.store[node_type] = (expiration_time, record)
|
||||
|
||||
self._remove_old_entries()
|
||||
|
||||
def update_node_availability(
|
||||
self,
|
||||
node_type: str,
|
||||
timestamp: int,
|
||||
node_launch_exception: Optional[NodeLaunchException],
|
||||
) -> None:
|
||||
"""
|
||||
Update the availability and details of a single ndoe type.
|
||||
|
||||
Args:
|
||||
node_type: The node type.
|
||||
timestamp: The timestamp that this information is accurate as of.
|
||||
node_launch_exception: Details about why the node launch failed. If
|
||||
empty, the node type will be considered available."""
|
||||
with self.lock:
|
||||
self._update_node_availability_requires_lock(
|
||||
node_type, timestamp, node_launch_exception
|
||||
)
|
||||
|
||||
def summary(self) -> NodeAvailabilitySummary:
|
||||
"""
|
||||
Returns a summary of node availabilities and their staleness.
|
||||
|
||||
Returns
|
||||
A summary of node availabilities and their staleness.
|
||||
"""
|
||||
with self.lock:
|
||||
self._remove_old_entries()
|
||||
return NodeAvailabilitySummary(
|
||||
{node_type: record for node_type, (_, record) in self.store.items()}
|
||||
)
|
||||
|
||||
def _remove_old_entries(self):
|
||||
"""Remove any expired entries from the cache."""
|
||||
cur_time = self.timer()
|
||||
with self.lock:
|
||||
for key, (expiration_time, _) in list(self.store.items()):
|
||||
if expiration_time < cur_time:
|
||||
del self.store[key]
|
||||
@@ -0,0 +1,77 @@
|
||||
from typing import List, Set, Tuple
|
||||
|
||||
from ray.autoscaler._private import constants
|
||||
|
||||
|
||||
class NodeTracker:
|
||||
"""Map nodes to their corresponding logs.
|
||||
|
||||
We need to be a little careful here. At an given point in time, node_id <->
|
||||
ip can be interchangeably used, but the node_id -> ip relation is not
|
||||
bijective _across time_ since IP addresses can be reused. Therefore, we
|
||||
should treat node_id as the only unique identifier.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
# Mapping from node_id -> (ip, node type, stdout_path, process runner)
|
||||
self.node_mapping = {}
|
||||
|
||||
# A quick, inefficient FIFO cache implementation.
|
||||
self.lru_order = []
|
||||
|
||||
def _add_node_mapping(self, node_id: str, value: str):
|
||||
if node_id in self.node_mapping:
|
||||
return
|
||||
|
||||
assert len(self.lru_order) == len(self.node_mapping)
|
||||
if len(self.lru_order) >= constants.AUTOSCALER_MAX_NODES_TRACKED:
|
||||
# The LRU eviction case
|
||||
node_id = self.lru_order.pop(0)
|
||||
del self.node_mapping[node_id]
|
||||
|
||||
self.node_mapping[node_id] = value
|
||||
self.lru_order.append(node_id)
|
||||
|
||||
def track(self, node_id: str, ip: str, node_type: str):
|
||||
"""
|
||||
Begin to track a new node.
|
||||
|
||||
Args:
|
||||
node_id: The node id.
|
||||
ip: The node ip address.
|
||||
node_type: The node type.
|
||||
"""
|
||||
if node_id not in self.node_mapping:
|
||||
self._add_node_mapping(node_id, (ip, node_type))
|
||||
|
||||
def untrack(self, node_id: str):
|
||||
"""Gracefully stop tracking a node. If a node is intentionally removed from
|
||||
the cluster, we should stop tracking it so we don't mistakenly mark it
|
||||
as failed.
|
||||
|
||||
Args:
|
||||
node_id: The node id which failed.
|
||||
"""
|
||||
if node_id in self.node_mapping:
|
||||
self.lru_order.remove(node_id)
|
||||
del self.node_mapping[node_id]
|
||||
|
||||
def get_all_failed_node_info(
|
||||
self, non_failed_ids: Set[str]
|
||||
) -> List[Tuple[str, str]]:
|
||||
"""Get the information about all failed nodes. A failed node is any node which
|
||||
we began to track that is not pending or alive (i.e. not failed).
|
||||
|
||||
Args:
|
||||
non_failed_ids: Nodes are failed unless they are in this set.
|
||||
|
||||
Returns:
|
||||
List[Tuple[str, str]]: A list of tuples. Each tuple is the ip
|
||||
address and type of a failed node.
|
||||
"""
|
||||
failed_nodes = self.node_mapping.keys() - non_failed_ids
|
||||
failed_info = []
|
||||
# Returning the list in order is important for display purposes.
|
||||
for node_id in filter(lambda node_id: node_id in failed_nodes, self.lru_order):
|
||||
failed_info.append(self.node_mapping[node_id])
|
||||
return failed_info
|
||||
@@ -0,0 +1,292 @@
|
||||
from typing import Optional
|
||||
|
||||
|
||||
class NullMetric:
|
||||
"""Mock metric class to be used in case of prometheus_client import error."""
|
||||
|
||||
def set(self, *args, **kwargs):
|
||||
pass
|
||||
|
||||
def observe(self, *args, **kwargs):
|
||||
pass
|
||||
|
||||
def inc(self, *args, **kwargs):
|
||||
pass
|
||||
|
||||
def labels(self, *args, **kwargs):
|
||||
return self
|
||||
|
||||
def clear(self):
|
||||
pass
|
||||
|
||||
|
||||
try:
|
||||
|
||||
from prometheus_client import CollectorRegistry, Counter, Gauge, Histogram
|
||||
|
||||
# The metrics in this class should be kept in sync with
|
||||
# python/ray/tests/test_metrics_agent.py
|
||||
class AutoscalerPrometheusMetrics:
|
||||
def __init__(
|
||||
self, session_name: str = None, registry: Optional[CollectorRegistry] = None
|
||||
):
|
||||
self.registry: CollectorRegistry = registry or CollectorRegistry(
|
||||
auto_describe=True
|
||||
)
|
||||
self._session_name = session_name
|
||||
# Buckets: 5 seconds, 10 seconds, 20 seconds, 30 seconds,
|
||||
# 45 seconds, 1 minute, 1.5 minutes, 2 minutes,
|
||||
# 3 minutes, 4 minutes, 5 minutes, 6 minutes,
|
||||
# 8 minutes, 10 minutes, 12 minutes, 15 minutes
|
||||
# 20 minutes, 25 minutes, 30 minutes
|
||||
# used for both worker launch time and worker update time
|
||||
histogram_buckets = [
|
||||
5,
|
||||
10,
|
||||
20,
|
||||
30,
|
||||
45,
|
||||
60,
|
||||
90,
|
||||
120,
|
||||
180,
|
||||
240,
|
||||
300,
|
||||
360,
|
||||
480,
|
||||
600,
|
||||
720,
|
||||
900,
|
||||
1200,
|
||||
1500,
|
||||
1800,
|
||||
]
|
||||
# Buckets: .01 seconds to 1000 seconds.
|
||||
# Used for autoscaler update time.
|
||||
update_time_buckets = [0.01, 0.1, 1, 10, 100, 1000]
|
||||
self.worker_create_node_time: Histogram = Histogram(
|
||||
"worker_create_node_time_seconds",
|
||||
"Worker launch time. This is the time it takes for a call to "
|
||||
"a node provider's create_node method to return. Note that "
|
||||
"when nodes are launched in batches, the launch time for that "
|
||||
"batch will be observed once for *each* node in that batch. "
|
||||
"For example, if 8 nodes are launched in 3 minutes, a launch "
|
||||
"time of 3 minutes will be observed 8 times.",
|
||||
labelnames=("SessionName",),
|
||||
unit="seconds",
|
||||
namespace="autoscaler",
|
||||
registry=self.registry,
|
||||
buckets=histogram_buckets,
|
||||
).labels(SessionName=session_name)
|
||||
self.worker_update_time: Histogram = Histogram(
|
||||
"worker_update_time_seconds",
|
||||
"Worker update time. This is the time between when an updater "
|
||||
"thread begins executing and when it exits successfully. This "
|
||||
"metric only observes times for successful updates.",
|
||||
labelnames=("SessionName",),
|
||||
unit="seconds",
|
||||
namespace="autoscaler",
|
||||
registry=self.registry,
|
||||
buckets=histogram_buckets,
|
||||
).labels(SessionName=session_name)
|
||||
self.update_time: Histogram = Histogram(
|
||||
"update_time",
|
||||
"Autoscaler update time. This is the time for an autoscaler "
|
||||
"update iteration to complete.",
|
||||
labelnames=("SessionName",),
|
||||
unit="seconds",
|
||||
namespace="autoscaler",
|
||||
registry=self.registry,
|
||||
buckets=update_time_buckets,
|
||||
).labels(SessionName=session_name)
|
||||
self.pending_nodes: Gauge = Gauge(
|
||||
"pending_nodes",
|
||||
"Number of nodes pending to be started.",
|
||||
labelnames=(
|
||||
"NodeType",
|
||||
"SessionName",
|
||||
),
|
||||
unit="nodes",
|
||||
namespace="autoscaler",
|
||||
registry=self.registry,
|
||||
)
|
||||
self.active_nodes: Gauge = Gauge(
|
||||
"active_nodes",
|
||||
"Number of nodes in the cluster.",
|
||||
labelnames=(
|
||||
"NodeType",
|
||||
"SessionName",
|
||||
),
|
||||
unit="nodes",
|
||||
namespace="autoscaler",
|
||||
registry=self.registry,
|
||||
)
|
||||
self.recently_failed_nodes = Gauge(
|
||||
"recently_failed_nodes",
|
||||
"The number of recently failed nodes. This count could reset "
|
||||
"at undefined times.",
|
||||
labelnames=(
|
||||
"NodeType",
|
||||
"SessionName",
|
||||
),
|
||||
unit="nodes",
|
||||
namespace="autoscaler",
|
||||
registry=self.registry,
|
||||
)
|
||||
self.started_nodes: Counter = Counter(
|
||||
"started_nodes",
|
||||
"Number of nodes started.",
|
||||
labelnames=("SessionName",),
|
||||
unit="nodes",
|
||||
namespace="autoscaler",
|
||||
registry=self.registry,
|
||||
).labels(SessionName=session_name)
|
||||
self.stopped_nodes: Counter = Counter(
|
||||
"stopped_nodes",
|
||||
"Number of nodes stopped.",
|
||||
labelnames=("SessionName",),
|
||||
unit="nodes",
|
||||
namespace="autoscaler",
|
||||
registry=self.registry,
|
||||
).labels(SessionName=session_name)
|
||||
self.updating_nodes: Gauge = Gauge(
|
||||
"updating_nodes",
|
||||
"Number of nodes in the process of updating.",
|
||||
labelnames=("SessionName",),
|
||||
unit="nodes",
|
||||
namespace="autoscaler",
|
||||
registry=self.registry,
|
||||
).labels(SessionName=session_name)
|
||||
self.recovering_nodes: Gauge = Gauge(
|
||||
"recovering_nodes",
|
||||
"Number of nodes in the process of recovering.",
|
||||
labelnames=("SessionName",),
|
||||
unit="nodes",
|
||||
namespace="autoscaler",
|
||||
registry=self.registry,
|
||||
).labels(SessionName=session_name)
|
||||
self.running_workers: Gauge = Gauge(
|
||||
"running_workers",
|
||||
"Number of worker nodes running.",
|
||||
labelnames=("SessionName",),
|
||||
unit="nodes",
|
||||
namespace="autoscaler",
|
||||
registry=self.registry,
|
||||
).labels(SessionName=session_name)
|
||||
self.failed_create_nodes: Counter = Counter(
|
||||
"failed_create_nodes",
|
||||
"Number of nodes that failed to be created due to an "
|
||||
"exception in the node provider's create_node method.",
|
||||
labelnames=("SessionName",),
|
||||
unit="nodes",
|
||||
namespace="autoscaler",
|
||||
registry=self.registry,
|
||||
).labels(SessionName=session_name)
|
||||
self.failed_updates: Counter = Counter(
|
||||
"failed_updates",
|
||||
"Number of failed worker node updates.",
|
||||
labelnames=("SessionName",),
|
||||
unit="updates",
|
||||
namespace="autoscaler",
|
||||
registry=self.registry,
|
||||
).labels(SessionName=session_name)
|
||||
self.successful_updates: Counter = Counter(
|
||||
"successful_updates",
|
||||
"Number of succesfful worker node updates.",
|
||||
labelnames=("SessionName",),
|
||||
unit="updates",
|
||||
namespace="autoscaler",
|
||||
registry=self.registry,
|
||||
).labels(SessionName=session_name)
|
||||
self.failed_recoveries: Counter = Counter(
|
||||
"failed_recoveries",
|
||||
"Number of failed node recoveries.",
|
||||
labelnames=("SessionName",),
|
||||
unit="recoveries",
|
||||
namespace="autoscaler",
|
||||
registry=self.registry,
|
||||
).labels(SessionName=session_name)
|
||||
self.successful_recoveries: Counter = Counter(
|
||||
"successful_recoveries",
|
||||
"Number of successful node recoveries.",
|
||||
labelnames=("SessionName",),
|
||||
unit="recoveries",
|
||||
namespace="autoscaler",
|
||||
registry=self.registry,
|
||||
).labels(SessionName=session_name)
|
||||
self.update_loop_exceptions: Counter = Counter(
|
||||
"update_loop_exceptions",
|
||||
"Number of exceptions raised in the update loop of the autoscaler.",
|
||||
labelnames=("SessionName",),
|
||||
unit="exceptions",
|
||||
namespace="autoscaler",
|
||||
registry=self.registry,
|
||||
).labels(SessionName=session_name)
|
||||
self.node_launch_exceptions: Counter = Counter(
|
||||
"node_launch_exceptions",
|
||||
"Number of exceptions raised while launching nodes.",
|
||||
labelnames=("SessionName",),
|
||||
unit="exceptions",
|
||||
namespace="autoscaler",
|
||||
registry=self.registry,
|
||||
).labels(SessionName=session_name)
|
||||
self.reset_exceptions: Counter = Counter(
|
||||
"reset_exceptions",
|
||||
"Number of exceptions raised while resetting the autoscaler.",
|
||||
labelnames=("SessionName",),
|
||||
unit="exceptions",
|
||||
namespace="autoscaler",
|
||||
registry=self.registry,
|
||||
).labels(SessionName=session_name)
|
||||
self.config_validation_exceptions: Counter = Counter(
|
||||
"config_validation_exceptions",
|
||||
"Number of exceptions raised while validating the config "
|
||||
"during a reset.",
|
||||
labelnames=("SessionName",),
|
||||
unit="exceptions",
|
||||
namespace="autoscaler",
|
||||
registry=self.registry,
|
||||
).labels(SessionName=session_name)
|
||||
self.drain_node_exceptions: Counter = Counter(
|
||||
"drain_node_exceptions",
|
||||
"Number of exceptions raised when making a DrainNode rpc"
|
||||
"prior to node termination.",
|
||||
labelnames=("SessionName",),
|
||||
unit="exceptions",
|
||||
namespace="autoscaler",
|
||||
registry=self.registry,
|
||||
).labels(SessionName=session_name)
|
||||
# This represents the autoscaler's view of essentially
|
||||
# `ray.cluster_resources()`, it may be slightly different from the
|
||||
# core metric from an eventual consistency perspective.
|
||||
self.cluster_resources: Gauge = Gauge(
|
||||
"cluster_resources",
|
||||
"Total logical resources in the cluster.",
|
||||
labelnames=("resource", "SessionName"),
|
||||
unit="resources",
|
||||
namespace="autoscaler",
|
||||
registry=self.registry,
|
||||
)
|
||||
# This represents the pending launches + nodes being set up for the
|
||||
# autoscaler.
|
||||
self.pending_resources: Gauge = Gauge(
|
||||
"pending_resources",
|
||||
"Pending logical resources in the cluster.",
|
||||
labelnames=("resource", "SessionName"),
|
||||
unit="resources",
|
||||
namespace="autoscaler",
|
||||
registry=self.registry,
|
||||
)
|
||||
|
||||
@property
|
||||
def session_name(self):
|
||||
return self._session_name
|
||||
|
||||
except ImportError:
|
||||
|
||||
class AutoscalerPrometheusMetrics(object):
|
||||
def __init__(self, session_name: str = None):
|
||||
pass
|
||||
|
||||
def __getattr__(self, attr):
|
||||
return NullMetric()
|
||||
@@ -0,0 +1,299 @@
|
||||
import copy
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
from typing import Any, Dict
|
||||
|
||||
import yaml
|
||||
|
||||
from ray.autoscaler._private.loader import load_function_or_class
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# For caching provider instantiations across API calls of one python session
|
||||
_provider_instances = {}
|
||||
|
||||
# Minimal config for compatibility with legacy-style external configs.
|
||||
MINIMAL_EXTERNAL_CONFIG = {
|
||||
"available_node_types": {
|
||||
"ray.head.default": {},
|
||||
"ray.worker.default": {},
|
||||
},
|
||||
"head_node_type": "ray.head.default",
|
||||
"head_node": {},
|
||||
"worker_nodes": {},
|
||||
}
|
||||
|
||||
|
||||
def _import_aws(provider_config):
|
||||
try:
|
||||
# boto3 and botocore are imported in multiple places in the codebase,
|
||||
# so we just import them here to ensure that they are installed.
|
||||
import boto3 # noqa: F401
|
||||
except ImportError as e:
|
||||
raise ImportError(
|
||||
"The Ray AWS VM launcher requires the AWS SDK for Python (Boto3) "
|
||||
"to be installed. You can install it with `pip install boto3`."
|
||||
) from e
|
||||
|
||||
from ray.autoscaler._private.aws.node_provider import AWSNodeProvider
|
||||
|
||||
return AWSNodeProvider
|
||||
|
||||
|
||||
def _import_gcp(provider_config):
|
||||
try:
|
||||
import googleapiclient # noqa: F401
|
||||
except ImportError as e:
|
||||
raise ImportError(
|
||||
"The Ray GCP VM launcher requires the Google API Client to be installed. "
|
||||
"You can install it with `pip install google-api-python-client`."
|
||||
) from e
|
||||
|
||||
from ray.autoscaler._private.gcp.node_provider import GCPNodeProvider
|
||||
|
||||
return GCPNodeProvider
|
||||
|
||||
|
||||
def _import_azure(provider_config):
|
||||
from ray.autoscaler._private._azure.node_provider import AzureNodeProvider
|
||||
|
||||
return AzureNodeProvider
|
||||
|
||||
|
||||
def _import_vsphere(provider_config):
|
||||
from ray.autoscaler._private.vsphere.node_provider import VsphereWcpNodeProvider
|
||||
|
||||
return VsphereWcpNodeProvider
|
||||
|
||||
|
||||
def _import_local(provider_config):
|
||||
if "coordinator_address" in provider_config:
|
||||
from ray.autoscaler._private.local.coordinator_node_provider import (
|
||||
CoordinatorSenderNodeProvider,
|
||||
)
|
||||
|
||||
return CoordinatorSenderNodeProvider
|
||||
else:
|
||||
from ray.autoscaler._private.local.node_provider import LocalNodeProvider
|
||||
|
||||
return LocalNodeProvider
|
||||
|
||||
|
||||
def _import_readonly(provider_config):
|
||||
from ray.autoscaler._private.readonly.node_provider import ReadOnlyNodeProvider
|
||||
|
||||
return ReadOnlyNodeProvider
|
||||
|
||||
|
||||
def _import_fake_multinode(provider_config):
|
||||
from ray.autoscaler._private.fake_multi_node.node_provider import (
|
||||
FakeMultiNodeProvider,
|
||||
)
|
||||
|
||||
return FakeMultiNodeProvider
|
||||
|
||||
|
||||
def _import_fake_multinode_docker(provider_config):
|
||||
from ray.autoscaler._private.fake_multi_node.node_provider import (
|
||||
FakeMultiNodeDockerProvider,
|
||||
)
|
||||
|
||||
return FakeMultiNodeDockerProvider
|
||||
|
||||
|
||||
def _import_kuberay(provider_config):
|
||||
from ray.autoscaler._private.kuberay.node_provider import KubeRayNodeProvider
|
||||
|
||||
return KubeRayNodeProvider
|
||||
|
||||
|
||||
def _import_aliyun(provider_config):
|
||||
from ray.autoscaler._private.aliyun.node_provider import AliyunNodeProvider
|
||||
|
||||
return AliyunNodeProvider
|
||||
|
||||
|
||||
def _import_spark(provider_config):
|
||||
from ray.autoscaler._private.spark.node_provider import SparkNodeProvider
|
||||
|
||||
return SparkNodeProvider
|
||||
|
||||
|
||||
def _load_fake_multinode_defaults_config():
|
||||
import ray.autoscaler._private.fake_multi_node as ray_fake_multinode
|
||||
|
||||
return os.path.join(os.path.dirname(ray_fake_multinode.__file__), "example.yaml")
|
||||
|
||||
|
||||
def _load_read_only_defaults_config():
|
||||
import ray.autoscaler._private.readonly as ray_readonly
|
||||
|
||||
return os.path.join(os.path.dirname(ray_readonly.__file__), "defaults.yaml")
|
||||
|
||||
|
||||
def _load_fake_multinode_docker_defaults_config():
|
||||
import ray.autoscaler._private.fake_multi_node as ray_fake_multinode
|
||||
|
||||
return os.path.join(
|
||||
os.path.dirname(ray_fake_multinode.__file__), "example_docker.yaml"
|
||||
)
|
||||
|
||||
|
||||
def _load_local_defaults_config():
|
||||
import ray.autoscaler.local as ray_local
|
||||
|
||||
return os.path.join(os.path.dirname(ray_local.__file__), "defaults.yaml")
|
||||
|
||||
|
||||
def _load_aws_defaults_config():
|
||||
import ray.autoscaler.aws as ray_aws
|
||||
|
||||
return os.path.join(os.path.dirname(ray_aws.__file__), "defaults.yaml")
|
||||
|
||||
|
||||
def _load_vsphere_defaults_config():
|
||||
import ray.autoscaler.vsphere as ray_vsphere
|
||||
|
||||
return os.path.join(os.path.dirname(ray_vsphere.__file__), "defaults.yaml")
|
||||
|
||||
|
||||
def _load_gcp_defaults_config():
|
||||
import ray.autoscaler.gcp as ray_gcp
|
||||
|
||||
return os.path.join(os.path.dirname(ray_gcp.__file__), "defaults.yaml")
|
||||
|
||||
|
||||
def _load_azure_defaults_config():
|
||||
import ray.autoscaler.azure as ray_azure
|
||||
|
||||
return os.path.join(os.path.dirname(ray_azure.__file__), "defaults.yaml")
|
||||
|
||||
|
||||
def _load_aliyun_defaults_config():
|
||||
import ray.autoscaler.aliyun as ray_aliyun
|
||||
|
||||
return os.path.join(os.path.dirname(ray_aliyun.__file__), "defaults.yaml")
|
||||
|
||||
|
||||
def _import_external(provider_config):
|
||||
provider_cls = load_function_or_class(path=provider_config["module"])
|
||||
return provider_cls
|
||||
|
||||
|
||||
_NODE_PROVIDERS = {
|
||||
"local": _import_local,
|
||||
"fake_multinode": _import_fake_multinode,
|
||||
"fake_multinode_docker": _import_fake_multinode_docker,
|
||||
"readonly": _import_readonly,
|
||||
"aws": _import_aws,
|
||||
"gcp": _import_gcp,
|
||||
"vsphere": _import_vsphere,
|
||||
"azure": _import_azure,
|
||||
"kuberay": _import_kuberay,
|
||||
"aliyun": _import_aliyun,
|
||||
"external": _import_external, # Import an external module
|
||||
"spark": _import_spark,
|
||||
}
|
||||
|
||||
_PROVIDER_PRETTY_NAMES = {
|
||||
"readonly": "Readonly (Manual Cluster Setup)",
|
||||
"fake_multinode": "Fake Multinode",
|
||||
"fake_multinode_docker": "Fake Multinode Docker",
|
||||
"local": "Local",
|
||||
"aws": "AWS",
|
||||
"gcp": "GCP",
|
||||
"azure": "Azure",
|
||||
"kuberay": "KubeRay",
|
||||
"aliyun": "Aliyun",
|
||||
"external": "External",
|
||||
"vsphere": "vSphere",
|
||||
"spark": "Spark",
|
||||
}
|
||||
|
||||
_DEFAULT_CONFIGS = {
|
||||
"fake_multinode": _load_fake_multinode_defaults_config,
|
||||
"fake_multinode_docker": _load_fake_multinode_docker_defaults_config,
|
||||
"local": _load_local_defaults_config,
|
||||
"aws": _load_aws_defaults_config,
|
||||
"gcp": _load_gcp_defaults_config,
|
||||
"azure": _load_azure_defaults_config,
|
||||
"aliyun": _load_aliyun_defaults_config,
|
||||
"vsphere": _load_vsphere_defaults_config,
|
||||
"readonly": _load_read_only_defaults_config,
|
||||
}
|
||||
|
||||
|
||||
def _get_node_provider_cls(provider_config: Dict[str, Any]):
|
||||
"""Get the node provider class for a given provider config.
|
||||
|
||||
Note that this may be used by private node providers that proxy methods to
|
||||
built-in node providers, so we should maintain backwards compatibility.
|
||||
|
||||
Args:
|
||||
provider_config: provider section of the autoscaler config.
|
||||
|
||||
Returns:
|
||||
NodeProvider class
|
||||
"""
|
||||
importer = _NODE_PROVIDERS.get(provider_config["type"])
|
||||
if importer is None:
|
||||
raise NotImplementedError(
|
||||
"Unsupported node provider: {}".format(provider_config["type"])
|
||||
)
|
||||
return importer(provider_config)
|
||||
|
||||
|
||||
def _get_node_provider(
|
||||
provider_config: Dict[str, Any], cluster_name: str, use_cache: bool = True
|
||||
) -> Any:
|
||||
"""Get the instantiated node provider for a given provider config.
|
||||
|
||||
Note that this may be used by private node providers that proxy methods to
|
||||
built-in node providers, so we should maintain backwards compatibility.
|
||||
|
||||
Args:
|
||||
provider_config: provider section of the autoscaler config.
|
||||
cluster_name: cluster name from the autoscaler config.
|
||||
use_cache: whether or not to use a cached definition if available. If
|
||||
False, the returned object will also not be stored in the cache.
|
||||
|
||||
Returns:
|
||||
NodeProvider
|
||||
"""
|
||||
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]
|
||||
|
||||
provider_cls = _get_node_provider_cls(provider_config)
|
||||
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):
|
||||
"""Retrieve a node provider.
|
||||
|
||||
This is an INTERNAL API. It is not allowed to call this from any Ray
|
||||
package outside the autoscaler.
|
||||
"""
|
||||
if provider_config["type"] == "external":
|
||||
return copy.deepcopy(MINIMAL_EXTERNAL_CONFIG)
|
||||
load_config = _DEFAULT_CONFIGS.get(provider_config["type"])
|
||||
if load_config is None:
|
||||
raise NotImplementedError(
|
||||
"Unsupported node provider: {}".format(provider_config["type"])
|
||||
)
|
||||
path_to_default = load_config()
|
||||
with open(path_to_default) as f:
|
||||
defaults = yaml.safe_load(f)
|
||||
|
||||
return defaults
|
||||
@@ -0,0 +1,31 @@
|
||||
cluster_name: default
|
||||
max_workers: 0
|
||||
provider:
|
||||
type: readonly
|
||||
# This must be true since the nodes share the same ip!
|
||||
use_node_id_as_ip: True
|
||||
disable_node_updaters: True
|
||||
disable_launch_config_check: True
|
||||
available_node_types:
|
||||
ray.head.default:
|
||||
resources: {}
|
||||
node_config: {}
|
||||
max_workers: 0
|
||||
head_node_type: ray.head.default
|
||||
upscaling_speed: 1.0
|
||||
#
|
||||
# !!! Configurations below are not supported in fake cluster mode !!!
|
||||
#
|
||||
auth: {}
|
||||
docker: {}
|
||||
initialization_commands: []
|
||||
setup_commands: []
|
||||
head_setup_commands: []
|
||||
worker_setup_commands: []
|
||||
head_start_ray_commands: []
|
||||
worker_start_ray_commands: []
|
||||
file_mounts: {}
|
||||
cluster_synced_files: []
|
||||
file_mounts_sync_continuously: false
|
||||
rsync_exclude: []
|
||||
rsync_filter: []
|
||||
@@ -0,0 +1,80 @@
|
||||
from typing import List, Tuple
|
||||
|
||||
from ray.autoscaler._private.util import format_readonly_node_type
|
||||
from ray.autoscaler.node_provider import NodeProvider
|
||||
from ray.autoscaler.tags import (
|
||||
NODE_KIND_HEAD,
|
||||
STATUS_UP_TO_DATE,
|
||||
TAG_RAY_NODE_KIND,
|
||||
TAG_RAY_NODE_NAME,
|
||||
TAG_RAY_NODE_STATUS,
|
||||
TAG_RAY_USER_NODE_TYPE,
|
||||
)
|
||||
|
||||
|
||||
class ReadOnlyNodeProvider(NodeProvider):
|
||||
"""A node provider that merely reports the current cluster state.
|
||||
|
||||
This is used for laptop mode / manual cluster setup modes, in order to
|
||||
provide status reporting in the same way for users."""
|
||||
|
||||
def __init__(self, provider_config, cluster_name):
|
||||
NodeProvider.__init__(self, provider_config, cluster_name)
|
||||
self.nodes = {}
|
||||
|
||||
def is_readonly(self):
|
||||
return True
|
||||
|
||||
def _set_nodes(self, nodes: List[Tuple[str, str]]):
|
||||
"""Update the set of nodes in the cluster.
|
||||
|
||||
Args:
|
||||
nodes: List of (node_id, node_manager_address) tuples.
|
||||
"""
|
||||
new_nodes = {}
|
||||
for node_id, node_manager_address in nodes:
|
||||
# We make up a fake node type for each node (since each node
|
||||
# could have its own unique configuration).
|
||||
new_nodes[node_id] = {
|
||||
# Keep prefix in sync with node config gen in monitor.py
|
||||
"node_type": format_readonly_node_type(node_id),
|
||||
"ip": node_manager_address,
|
||||
}
|
||||
self.nodes = new_nodes
|
||||
|
||||
def non_terminated_nodes(self, tag_filters):
|
||||
return list(self.nodes.keys())
|
||||
|
||||
def is_running(self, node_id):
|
||||
return node_id in self.nodes
|
||||
|
||||
def is_terminated(self, node_id):
|
||||
return node_id not in self.nodes
|
||||
|
||||
def node_tags(self, node_id):
|
||||
tags = {
|
||||
TAG_RAY_NODE_KIND: NODE_KIND_HEAD,
|
||||
TAG_RAY_USER_NODE_TYPE: self.nodes[node_id]["node_type"],
|
||||
TAG_RAY_NODE_NAME: node_id,
|
||||
TAG_RAY_NODE_STATUS: STATUS_UP_TO_DATE,
|
||||
}
|
||||
return tags
|
||||
|
||||
def external_ip(self, node_id):
|
||||
return node_id
|
||||
|
||||
def internal_ip(self, node_id):
|
||||
return node_id
|
||||
|
||||
def set_node_tags(self, node_id, tags):
|
||||
raise AssertionError("Readonly node provider cannot be updated")
|
||||
|
||||
def create_node(self, node_config, tags, count):
|
||||
raise AssertionError("Readonly node provider cannot be updated")
|
||||
|
||||
def terminate_node(self, node_id):
|
||||
raise AssertionError("Readonly node provider cannot be updated")
|
||||
|
||||
@staticmethod
|
||||
def bootstrap_config(cluster_config):
|
||||
return cluster_config
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,250 @@
|
||||
import json
|
||||
import logging
|
||||
import sys
|
||||
from threading import RLock
|
||||
from typing import Any, Dict, Optional
|
||||
|
||||
import requests
|
||||
|
||||
from ray._common.network_utils import build_address
|
||||
from ray.autoscaler.node_launch_exception import NodeLaunchException
|
||||
from ray.autoscaler.node_provider import NodeProvider
|
||||
from ray.autoscaler.tags import (
|
||||
NODE_KIND_HEAD,
|
||||
NODE_KIND_WORKER,
|
||||
STATUS_SETTING_UP,
|
||||
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__)
|
||||
|
||||
HEAD_NODE_ID = 0
|
||||
HEAD_NODE_TYPE = "ray.head.default"
|
||||
|
||||
|
||||
class SparkNodeProvider(NodeProvider):
|
||||
"""A node provider that implements provider for nodes of Ray on spark."""
|
||||
|
||||
def __init__(self, provider_config, cluster_name):
|
||||
NodeProvider.__init__(self, provider_config, cluster_name)
|
||||
self.lock = RLock()
|
||||
|
||||
self._nodes = {
|
||||
str(HEAD_NODE_ID): {
|
||||
"tags": {
|
||||
TAG_RAY_NODE_KIND: NODE_KIND_HEAD,
|
||||
TAG_RAY_USER_NODE_TYPE: HEAD_NODE_TYPE,
|
||||
TAG_RAY_NODE_NAME: HEAD_NODE_ID,
|
||||
TAG_RAY_NODE_STATUS: STATUS_UP_TO_DATE,
|
||||
}
|
||||
},
|
||||
}
|
||||
self._next_node_id = 0
|
||||
|
||||
self.ray_head_ip = self.provider_config["ray_head_ip"]
|
||||
# The port of spark job server. We send http request to spark job server
|
||||
# to launch spark jobs, ray worker nodes are launched by spark task in
|
||||
# spark jobs.
|
||||
spark_job_server_port = self.provider_config["spark_job_server_port"]
|
||||
self.spark_job_server_url = (
|
||||
f"http://{build_address(self.ray_head_ip, spark_job_server_port)}"
|
||||
)
|
||||
self.ray_head_port = self.provider_config["ray_head_port"]
|
||||
# The unique id for the Ray on spark cluster.
|
||||
self.cluster_id = self.provider_config["cluster_unique_id"]
|
||||
|
||||
def get_next_node_id(self):
|
||||
with self.lock:
|
||||
self._next_node_id += 1
|
||||
return self._next_node_id
|
||||
|
||||
def non_terminated_nodes(self, tag_filters):
|
||||
with self.lock:
|
||||
nodes = []
|
||||
|
||||
died_nodes = []
|
||||
for node_id in self._nodes:
|
||||
if node_id == str(HEAD_NODE_ID):
|
||||
status = "running"
|
||||
else:
|
||||
status = self._query_node_status(node_id)
|
||||
|
||||
if status == "running":
|
||||
if (
|
||||
self._nodes[node_id]["tags"][TAG_RAY_NODE_STATUS]
|
||||
== STATUS_SETTING_UP
|
||||
):
|
||||
self._nodes[node_id]["tags"][
|
||||
TAG_RAY_NODE_STATUS
|
||||
] = STATUS_UP_TO_DATE
|
||||
logger.info(
|
||||
f"Spark node provider node {node_id} starts running."
|
||||
)
|
||||
|
||||
if status == "terminated":
|
||||
died_nodes.append(node_id)
|
||||
else:
|
||||
tags = self.node_tags(node_id)
|
||||
ok = True
|
||||
for k, v in tag_filters.items():
|
||||
if tags.get(k) != v:
|
||||
ok = False
|
||||
if ok:
|
||||
nodes.append(node_id)
|
||||
|
||||
for died_node_id in died_nodes:
|
||||
self._nodes.pop(died_node_id)
|
||||
|
||||
return nodes
|
||||
|
||||
def _query_node_status(self, node_id):
|
||||
spark_job_group_id = self._gen_spark_job_group_id(node_id)
|
||||
|
||||
response = requests.post(
|
||||
url=self.spark_job_server_url + "/query_task_status",
|
||||
json={"spark_job_group_id": spark_job_group_id},
|
||||
)
|
||||
response.raise_for_status()
|
||||
|
||||
decoded_resp = response.content.decode("utf-8")
|
||||
json_res = json.loads(decoded_resp)
|
||||
return json_res["status"]
|
||||
|
||||
def is_running(self, node_id):
|
||||
with self.lock:
|
||||
return (
|
||||
node_id in self._nodes
|
||||
and self._nodes[node_id]["tags"][TAG_RAY_NODE_STATUS]
|
||||
== STATUS_UP_TO_DATE
|
||||
)
|
||||
|
||||
def is_terminated(self, node_id):
|
||||
with self.lock:
|
||||
return node_id not in self._nodes
|
||||
|
||||
def node_tags(self, node_id):
|
||||
with self.lock:
|
||||
return self._nodes[node_id]["tags"]
|
||||
|
||||
def _get_ip(self, node_id: str) -> Optional[str]:
|
||||
return node_id
|
||||
|
||||
def external_ip(self, node_id):
|
||||
return self._get_ip(node_id)
|
||||
|
||||
def internal_ip(self, node_id):
|
||||
return self._get_ip(node_id)
|
||||
|
||||
def set_node_tags(self, node_id, tags):
|
||||
assert node_id in self._nodes
|
||||
self._nodes[node_id]["tags"].update(tags)
|
||||
|
||||
def create_node(
|
||||
self, node_config: Dict[str, Any], tags: Dict[str, str], count: int
|
||||
) -> Optional[Dict[str, Any]]:
|
||||
raise AssertionError("This method should not be called.")
|
||||
|
||||
def _gen_spark_job_group_id(self, node_id):
|
||||
return (
|
||||
f"ray-cluster-{self.ray_head_port}-{self.cluster_id}"
|
||||
f"-worker-node-{node_id}"
|
||||
)
|
||||
|
||||
def create_node_with_resources_and_labels(
|
||||
self, node_config, tags, count, resources, labels
|
||||
):
|
||||
for _ in range(count):
|
||||
self._create_node_with_resources_and_labels(
|
||||
node_config, tags, resources, labels
|
||||
)
|
||||
|
||||
def _create_node_with_resources_and_labels(
|
||||
self, node_config, tags, resources, labels
|
||||
):
|
||||
from ray.util.spark.cluster_init import _append_resources_config
|
||||
|
||||
with self.lock:
|
||||
resources = resources.copy()
|
||||
node_type = tags[TAG_RAY_USER_NODE_TYPE]
|
||||
# NOTE:
|
||||
# "NODE_ID_AS_RESOURCE" value must be an integer,
|
||||
# but `node_id` used by autoscaler must be a string.
|
||||
node_id = str(self.get_next_node_id())
|
||||
resources["NODE_ID_AS_RESOURCE"] = int(node_id)
|
||||
|
||||
conf = self.provider_config.copy()
|
||||
|
||||
num_cpus_per_node = resources.pop("CPU")
|
||||
num_gpus_per_node = resources.pop("GPU")
|
||||
heap_memory_per_node = resources.pop("memory")
|
||||
object_store_memory_per_node = resources.pop("object_store_memory")
|
||||
|
||||
conf["worker_node_options"] = _append_resources_config(
|
||||
conf["worker_node_options"], resources
|
||||
)
|
||||
response = requests.post(
|
||||
url=self.spark_job_server_url + "/create_node",
|
||||
json={
|
||||
"spark_job_group_id": self._gen_spark_job_group_id(node_id),
|
||||
"spark_job_group_desc": (
|
||||
"This job group is for spark job which runs the Ray "
|
||||
f"cluster worker node {node_id} connecting to ray "
|
||||
f"head node {build_address(self.ray_head_ip, self.ray_head_port)}"
|
||||
),
|
||||
"using_stage_scheduling": conf["using_stage_scheduling"],
|
||||
"ray_head_ip": self.ray_head_ip,
|
||||
"ray_head_port": self.ray_head_port,
|
||||
"ray_temp_dir": conf["ray_temp_dir"],
|
||||
"num_cpus_per_node": num_cpus_per_node,
|
||||
"num_gpus_per_node": num_gpus_per_node,
|
||||
"heap_memory_per_node": heap_memory_per_node,
|
||||
"object_store_memory_per_node": object_store_memory_per_node,
|
||||
"worker_node_options": conf["worker_node_options"],
|
||||
"collect_log_to_path": conf["collect_log_to_path"],
|
||||
"node_id": resources["NODE_ID_AS_RESOURCE"],
|
||||
},
|
||||
)
|
||||
|
||||
try:
|
||||
# Spark job server is locally launched, if spark job server request
|
||||
# failed, it is unlikely network error but probably unrecoverable
|
||||
# error, so we make it fast-fail.
|
||||
response.raise_for_status()
|
||||
except Exception:
|
||||
raise NodeLaunchException(
|
||||
"Node creation failure",
|
||||
f"Starting ray worker node {node_id} failed",
|
||||
sys.exc_info(),
|
||||
)
|
||||
|
||||
self._nodes[node_id] = {
|
||||
"tags": {
|
||||
TAG_RAY_NODE_KIND: NODE_KIND_WORKER,
|
||||
TAG_RAY_USER_NODE_TYPE: node_type,
|
||||
TAG_RAY_NODE_NAME: node_id,
|
||||
TAG_RAY_NODE_STATUS: STATUS_SETTING_UP,
|
||||
},
|
||||
}
|
||||
logger.info(f"Spark node provider creates node {node_id}.")
|
||||
|
||||
def terminate_node(self, node_id):
|
||||
if node_id in self._nodes:
|
||||
response = requests.post(
|
||||
url=self.spark_job_server_url + "/terminate_node",
|
||||
json={"spark_job_group_id": self._gen_spark_job_group_id(node_id)},
|
||||
)
|
||||
response.raise_for_status()
|
||||
|
||||
with self.lock:
|
||||
if node_id in self._nodes:
|
||||
self._nodes.pop(node_id)
|
||||
|
||||
logger.info(f"Spark node provider terminates node {node_id}")
|
||||
|
||||
@staticmethod
|
||||
def bootstrap_config(cluster_config):
|
||||
return cluster_config
|
||||
@@ -0,0 +1,244 @@
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import threading
|
||||
import time
|
||||
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
||||
from pathlib import Path
|
||||
|
||||
from pyspark.util import inheritable_thread_target
|
||||
|
||||
from ray.util.spark.cluster_init import _start_ray_worker_nodes
|
||||
|
||||
|
||||
class SparkJobServerRequestHandler(BaseHTTPRequestHandler):
|
||||
def setup(self) -> None:
|
||||
super().setup()
|
||||
self._handler_lock = threading.RLock()
|
||||
self._created_node_id_set = set()
|
||||
self._logger = logging.getLogger(__name__)
|
||||
if "RAY_ON_SPARK_JOB_SERVER_VERBOSE" in os.environ:
|
||||
self._logger.setLevel(logging.DEBUG)
|
||||
else:
|
||||
self._logger.setLevel(logging.WARN)
|
||||
|
||||
def _set_headers(self):
|
||||
self.send_response(200)
|
||||
self.send_header("Content-type", "application/json")
|
||||
self.end_headers()
|
||||
|
||||
def handle_POST(self, path, data):
|
||||
path_parts = Path(path).parts[1:]
|
||||
|
||||
spark_job_group_id = data["spark_job_group_id"]
|
||||
|
||||
if path_parts[0] == "create_node":
|
||||
assert len(path_parts) == 1, f"Illegal request path: {path}"
|
||||
spark_job_group_desc = data["spark_job_group_desc"]
|
||||
using_stage_scheduling = data["using_stage_scheduling"]
|
||||
ray_head_ip = data["ray_head_ip"]
|
||||
ray_head_port = data["ray_head_port"]
|
||||
ray_temp_dir = data["ray_temp_dir"]
|
||||
num_cpus_per_node = data["num_cpus_per_node"]
|
||||
num_gpus_per_node = data["num_gpus_per_node"]
|
||||
heap_memory_per_node = data["heap_memory_per_node"]
|
||||
object_store_memory_per_node = data["object_store_memory_per_node"]
|
||||
worker_node_options = data["worker_node_options"]
|
||||
collect_log_to_path = data["collect_log_to_path"]
|
||||
node_id = data["node_id"]
|
||||
self._created_node_id_set.add(node_id)
|
||||
|
||||
def start_ray_worker_thread_fn():
|
||||
try:
|
||||
err_msg = _start_ray_worker_nodes(
|
||||
spark_job_server=self.server,
|
||||
spark_job_group_id=spark_job_group_id,
|
||||
spark_job_group_desc=spark_job_group_desc,
|
||||
num_worker_nodes=1,
|
||||
using_stage_scheduling=using_stage_scheduling,
|
||||
ray_head_ip=ray_head_ip,
|
||||
ray_head_port=ray_head_port,
|
||||
ray_temp_dir=ray_temp_dir,
|
||||
num_cpus_per_node=num_cpus_per_node,
|
||||
num_gpus_per_node=num_gpus_per_node,
|
||||
heap_memory_per_node=heap_memory_per_node,
|
||||
object_store_memory_per_node=object_store_memory_per_node,
|
||||
worker_node_options=worker_node_options,
|
||||
collect_log_to_path=collect_log_to_path,
|
||||
node_id=node_id,
|
||||
)
|
||||
if err_msg:
|
||||
self._logger.warning(
|
||||
f"Spark job {spark_job_group_id} hosting Ray worker node "
|
||||
f"launching failed, error:\n{err_msg}"
|
||||
)
|
||||
except Exception:
|
||||
if spark_job_group_id in self.server.task_status_dict:
|
||||
self.server.task_status_dict.pop(spark_job_group_id)
|
||||
|
||||
msg = (
|
||||
f"Spark job {spark_job_group_id} hosting Ray worker node exit."
|
||||
)
|
||||
if self._logger.level > logging.DEBUG:
|
||||
self._logger.warning(
|
||||
f"{msg} To see details, you can set "
|
||||
"'RAY_ON_SPARK_JOB_SERVER_VERBOSE' environmental variable "
|
||||
"to '1' before calling 'ray.util.spark.setup_ray_cluster'."
|
||||
)
|
||||
else:
|
||||
# This branch is only for debugging Ray-on-Spark purpose.
|
||||
# User can configure 'RAY_ON_SPARK_JOB_SERVER_VERBOSE'
|
||||
# environment variable to make the spark job server logging
|
||||
# showing full exception stack here.
|
||||
self._logger.debug(msg, exc_info=True)
|
||||
|
||||
threading.Thread(
|
||||
target=inheritable_thread_target(start_ray_worker_thread_fn),
|
||||
args=(),
|
||||
daemon=True,
|
||||
).start()
|
||||
|
||||
self.server.task_status_dict[spark_job_group_id] = "pending"
|
||||
return {}
|
||||
|
||||
elif path_parts[0] == "check_node_id_availability":
|
||||
node_id = data["node_id"]
|
||||
with self._handler_lock:
|
||||
if node_id in self._created_node_id_set:
|
||||
# If the node with the node id has been created,
|
||||
# it shouldn't be created twice so fail fast here.
|
||||
# The case happens when a Ray node is down unexpected
|
||||
# caused by spark worker node down and spark tries to
|
||||
# reschedule the spark task, so it triggers node
|
||||
# creation with duplicated node id.
|
||||
return {"available": False}
|
||||
else:
|
||||
self._created_node_id_set.add(node_id)
|
||||
return {"available": True}
|
||||
|
||||
elif path_parts[0] == "terminate_node":
|
||||
assert len(path_parts) == 1, f"Illegal request path: {path}"
|
||||
self.server.spark.sparkContext.cancelJobGroup(spark_job_group_id)
|
||||
if spark_job_group_id in self.server.task_status_dict:
|
||||
self.server.task_status_dict.pop(spark_job_group_id)
|
||||
return {}
|
||||
|
||||
elif path_parts[0] == "notify_task_launched":
|
||||
if spark_job_group_id in self.server.task_status_dict:
|
||||
# Note that if `spark_job_group_id` not in task_status_dict,
|
||||
# the task has been terminated
|
||||
self.server.task_status_dict[spark_job_group_id] = "running"
|
||||
self._logger.info(f"Spark task in {spark_job_group_id} has started.")
|
||||
return {}
|
||||
|
||||
elif path_parts[0] == "query_task_status":
|
||||
if spark_job_group_id in self.server.task_status_dict:
|
||||
return {"status": self.server.task_status_dict[spark_job_group_id]}
|
||||
else:
|
||||
return {"status": "terminated"}
|
||||
|
||||
elif path_parts[0] == "query_last_worker_err":
|
||||
return {"last_worker_err": self.server.last_worker_error}
|
||||
|
||||
else:
|
||||
raise ValueError(f"Illegal request path: {path}")
|
||||
|
||||
def do_POST(self):
|
||||
"""Reads post request body"""
|
||||
self._set_headers()
|
||||
content_len = int(self.headers["content-length"])
|
||||
content_type = self.headers["content-type"]
|
||||
assert content_type == "application/json"
|
||||
path = self.path
|
||||
post_body = self.rfile.read(content_len).decode("utf-8")
|
||||
post_body_json = json.loads(post_body)
|
||||
with self._handler_lock:
|
||||
response_body_json = self.handle_POST(path, post_body_json)
|
||||
response_body = json.dumps(response_body_json)
|
||||
self.wfile.write(response_body.encode("utf-8"))
|
||||
|
||||
def log_request(self, code="-", size="-"):
|
||||
# Make logs less verbose.
|
||||
pass
|
||||
|
||||
|
||||
class SparkJobServer(ThreadingHTTPServer):
|
||||
"""
|
||||
High level design:
|
||||
|
||||
1. In Ray on spark autoscaling mode, How to start and terminate Ray worker node ?
|
||||
|
||||
It uses spark job to launch Ray worker node,
|
||||
and each spark job contains only one spark task, the corresponding spark task
|
||||
creates Ray worker node as subprocess.
|
||||
When autoscaler request terminating specific Ray worker node, it cancels
|
||||
corresponding spark job to trigger Ray worker node termination.
|
||||
Because we can only cancel spark job not spark task when we need to scale
|
||||
down a Ray worker node. So we have to have one spark job for each Ray worker node.
|
||||
|
||||
2. How to create / cancel spark job from spark node provider?
|
||||
|
||||
Spark node provider runs in autoscaler process that is different process
|
||||
than the one that executes "setup_ray_cluster" API. User calls "setup_ray_cluster"
|
||||
API in spark application driver node, and the semantic is "setup_ray_cluster"
|
||||
requests spark resources from this spark application.
|
||||
Internally, "setup_ray_cluster" should use "spark session" instance to request
|
||||
spark application resources. But spark node provider runs in another python
|
||||
process, in order to share spark session to the separate NodeProvider process,
|
||||
it sets up a spark job server that runs inside spark application driver process
|
||||
(the process that calls "setup_ray_cluster" API), and in NodeProvider process,
|
||||
it sends RPC request to the spark job server for creating spark jobs in the
|
||||
spark application.
|
||||
Note that we cannot create another spark session in NodeProvider process,
|
||||
because if doing so, it means we create another spark application, and then
|
||||
it causes NodeProvider requests resources belonging to the new spark application,
|
||||
but we need to ensure all requested spark resources belong to
|
||||
the original spark application that calls "setup_ray_cluster" API.
|
||||
|
||||
Note:
|
||||
The server must inherit ThreadingHTTPServer because request handler uses
|
||||
the active spark session in current process to create spark jobs, so all request
|
||||
handler must be running in current process.
|
||||
"""
|
||||
|
||||
def __init__(self, server_address, spark, ray_node_custom_env):
|
||||
super().__init__(server_address, SparkJobServerRequestHandler)
|
||||
self.spark = spark
|
||||
|
||||
# For ray on spark autoscaling mode,
|
||||
# for each ray worker node, we create an individual spark job
|
||||
# to launch it, the corresponding spark job has only one
|
||||
# spark task that starts ray worker node, and the spark job
|
||||
# is assigned with a unique spark job group ID that is used
|
||||
# to cancel this spark job (i.e., kill corresponding ray worker node).
|
||||
# Each spark task has status of pending, running, or terminated.
|
||||
# the task_status_dict key is spark job group id,
|
||||
# and value is the corresponding spark task status.
|
||||
# each spark task holds a ray worker node.
|
||||
self.task_status_dict = {}
|
||||
self.last_worker_error = None
|
||||
self.ray_node_custom_env = ray_node_custom_env
|
||||
|
||||
def shutdown(self) -> None:
|
||||
super().shutdown()
|
||||
for spark_job_group_id in list(self.task_status_dict.keys()):
|
||||
self.spark.sparkContext.cancelJobGroup(spark_job_group_id)
|
||||
# Sleep 1 second to wait for all spark job cancellation
|
||||
# The spark job cancellation will do things asyncly in a background thread,
|
||||
# On Databricks platform, when detaching a notebook, it triggers SIGTERM
|
||||
# and then sigterm handler triggers Ray cluster shutdown, without sleep,
|
||||
# after the SIGTERM handler execution the process is killed and then
|
||||
# these cancelling spark job background threads are killed.
|
||||
time.sleep(1)
|
||||
|
||||
|
||||
def _start_spark_job_server(host, port, spark, ray_node_custom_env):
|
||||
server = SparkJobServer((host, port), spark, ray_node_custom_env)
|
||||
|
||||
def run_server():
|
||||
server.serve_forever()
|
||||
|
||||
server_thread = threading.Thread(target=run_server, daemon=True)
|
||||
server_thread.start()
|
||||
|
||||
return server
|
||||
@@ -0,0 +1,446 @@
|
||||
import os
|
||||
import re
|
||||
import subprocess
|
||||
import sys
|
||||
import tempfile
|
||||
import time
|
||||
from typing import IO, Any, List, Optional
|
||||
|
||||
from ray.autoscaler._private.cli_logger import cf, cli_logger
|
||||
|
||||
CONN_REFUSED_PATIENCE = 30 # how long to wait for sshd to run
|
||||
|
||||
_redirect_output = False # Whether to log command output to a temporary file
|
||||
_allow_interactive = True # whether to pass on stdin to running commands.
|
||||
|
||||
|
||||
def is_output_redirected():
|
||||
return _redirect_output
|
||||
|
||||
|
||||
def set_output_redirected(val: bool):
|
||||
"""Choose between logging to a temporary file and to `sys.stdout`.
|
||||
|
||||
The default is to log to a file.
|
||||
|
||||
Args:
|
||||
val: If true, subprocess output will be redirected to
|
||||
a temporary file.
|
||||
"""
|
||||
global _redirect_output
|
||||
_redirect_output = val
|
||||
|
||||
|
||||
def does_allow_interactive():
|
||||
return _allow_interactive
|
||||
|
||||
|
||||
def set_allow_interactive(val: bool):
|
||||
"""Choose whether to pass on stdin to running commands.
|
||||
|
||||
The default is to pipe stdin and close it immediately.
|
||||
|
||||
Args:
|
||||
val: If true, stdin will be passed to commands.
|
||||
"""
|
||||
global _allow_interactive
|
||||
_allow_interactive = val
|
||||
|
||||
|
||||
class ProcessRunnerError(Exception):
|
||||
def __init__(self, msg, msg_type, code=None, command=None, special_case=None):
|
||||
super(ProcessRunnerError, self).__init__(
|
||||
"{} (discovered={}): type={}, code={}, command={}".format(
|
||||
msg, special_case, msg_type, code, command
|
||||
)
|
||||
)
|
||||
|
||||
self.msg_type = msg_type
|
||||
self.code = code
|
||||
self.command = command
|
||||
|
||||
self.special_case = special_case
|
||||
|
||||
|
||||
_ssh_output_regexes = {
|
||||
"known_host_update": re.compile(
|
||||
r"\s*Warning: Permanently added '.+' \(.+\) " r"to the list of known hosts.\s*"
|
||||
),
|
||||
"connection_closed": re.compile(r"\s*Shared connection to .+ closed.\s*"),
|
||||
"timeout": re.compile(
|
||||
r"\s*ssh: connect to host .+ port .+: " r"Operation timed out\s*"
|
||||
),
|
||||
"conn_refused": re.compile(
|
||||
r"\s*ssh: connect to host .+ port .+: Connection refused\s*"
|
||||
)
|
||||
# todo: check for other connection failures for better error messages?
|
||||
}
|
||||
|
||||
|
||||
def _read_subprocess_stream(
|
||||
f: IO[str], output_file: Optional[IO[str]], is_stdout: bool = False
|
||||
) -> Optional[str]:
|
||||
"""Read and process a subprocess output stream.
|
||||
|
||||
The goal is to find error messages and respond to them in a clever way.
|
||||
Currently just used for SSH messages (CONN_REFUSED, TIMEOUT, etc.), so
|
||||
the user does not get confused by these.
|
||||
|
||||
Ran in a thread each for both `stdout` and `stderr` to
|
||||
allow for cross-platform asynchronous IO.
|
||||
|
||||
Note: `select`-based IO is another option, but Windows has
|
||||
no support for `select`ing pipes, and Linux support varies somewhat.
|
||||
Spefically, Older *nix systems might also have quirks in how they
|
||||
handle `select` on pipes.
|
||||
|
||||
Args:
|
||||
f: File object for the stream.
|
||||
output_file: File object to which filtered output is written.
|
||||
is_stdout:
|
||||
When `is_stdout` is `False`, the stream is assumed to
|
||||
be `stderr`. Different error message detectors are used,
|
||||
and the output is displayed to the user unless it matches
|
||||
a special case (e.g. SSH timeout), in which case this is
|
||||
left up to the caller.
|
||||
|
||||
Returns:
|
||||
The detected special case name (e.g. "ssh_timeout") or None.
|
||||
"""
|
||||
|
||||
detected_special_case = None
|
||||
while True:
|
||||
# ! Readline here is crucial.
|
||||
# ! Normal `read()` will block until EOF instead of until
|
||||
# something is available.
|
||||
line = f.readline()
|
||||
|
||||
if line is None or line == "":
|
||||
# EOF
|
||||
break
|
||||
|
||||
if line[-1] == "\n":
|
||||
line = line[:-1]
|
||||
|
||||
if not is_stdout:
|
||||
if _ssh_output_regexes["connection_closed"].fullmatch(line) is not None:
|
||||
# Do not log "connection closed" messages which SSH
|
||||
# puts in stderr for no reason.
|
||||
#
|
||||
# They are never errors since the connection will
|
||||
# close no matter whether the command succeeds or not.
|
||||
continue
|
||||
|
||||
if _ssh_output_regexes["timeout"].fullmatch(line) is not None:
|
||||
# Timeout is not really an error but rather a special
|
||||
# condition. It should be handled by the caller, since
|
||||
# network conditions/nodes in the early stages of boot
|
||||
# are expected to sometimes cause connection timeouts.
|
||||
if detected_special_case is not None:
|
||||
raise ValueError(
|
||||
"Bug: ssh_timeout conflicts with another "
|
||||
"special codition: " + detected_special_case
|
||||
)
|
||||
|
||||
detected_special_case = "ssh_timeout"
|
||||
continue
|
||||
|
||||
if _ssh_output_regexes["conn_refused"].fullmatch(line) is not None:
|
||||
# Connection refused is not really an error but
|
||||
# rather a special condition. It should be handled by
|
||||
# the caller, since network conditions/nodes in the
|
||||
# early stages of boot are expected to sometimes cause
|
||||
# CONN_REFUSED.
|
||||
if detected_special_case is not None:
|
||||
raise ValueError(
|
||||
"Bug: ssh_conn_refused conflicts with another "
|
||||
"special codition: " + detected_special_case
|
||||
)
|
||||
|
||||
detected_special_case = "ssh_conn_refused"
|
||||
continue
|
||||
|
||||
if _ssh_output_regexes["known_host_update"].fullmatch(line) is not None:
|
||||
# Since we ignore SSH host control anyway
|
||||
# (-o UserKnownHostsFile=/dev/null),
|
||||
# we should silence the host control warnings.
|
||||
continue
|
||||
|
||||
cli_logger.error(line)
|
||||
|
||||
if output_file is not None and output_file != subprocess.DEVNULL:
|
||||
output_file.write(line + "\n")
|
||||
|
||||
return detected_special_case
|
||||
|
||||
|
||||
def _run_and_process_output(
|
||||
cmd: List[str],
|
||||
stdout_file: Optional[IO[str]],
|
||||
process_runner: Any = subprocess,
|
||||
stderr_file: Optional[IO[str]] = None,
|
||||
use_login_shells: bool = False,
|
||||
):
|
||||
"""Run a command and process its output for special cases.
|
||||
|
||||
Calls a standard 'check_call' if process_runner is not subprocess.
|
||||
|
||||
Specifically, run all command output through regex to detect
|
||||
error conditions and filter out non-error messages that went to stderr
|
||||
anyway (SSH writes ALL of its "system" messages to stderr even if they
|
||||
are not actually errors).
|
||||
|
||||
Args:
|
||||
cmd: Command to run.
|
||||
stdout_file: File to redirect stdout to.
|
||||
process_runner: Used for command execution. Assumed to have
|
||||
'check_call' and 'check_output' inplemented.
|
||||
stderr_file: File to redirect stderr to.
|
||||
use_login_shells: Whether to disable special output processing because
|
||||
an interactive login shell is in use.
|
||||
|
||||
Returns:
|
||||
The return code of the executed process.
|
||||
|
||||
Implementation notes:
|
||||
1. `use_login_shells` disables special processing
|
||||
If we run interactive apps, output processing will likely get
|
||||
overwhelmed with the interactive output elements.
|
||||
Thus, we disable output processing for login shells. This makes
|
||||
the logging experience considerably worse, but it only degrades
|
||||
to old-style logging.
|
||||
|
||||
For example, `pip install` outputs HUNDREDS of progress-bar lines
|
||||
when downloading a package, and we have to
|
||||
read + regex + write all of them.
|
||||
|
||||
After all, even just printing output to console can often slow
|
||||
down a fast-printing app, and we do more than just print, and
|
||||
all that from Python, which is much slower than C regarding
|
||||
stream processing.
|
||||
|
||||
2. `stdin=PIPE` for subprocesses
|
||||
Do not inherit stdin as it messes with bash signals
|
||||
(ctrl-C for SIGINT) and these commands aren't supposed to
|
||||
take input anyway.
|
||||
|
||||
3. `ThreadPoolExecutor` without the `Pool`
|
||||
We use `ThreadPoolExecutor` to create futures from threads.
|
||||
Threads are never reused.
|
||||
|
||||
This approach allows us to have no custom synchronization by
|
||||
off-loading the return value and exception passing to the
|
||||
standard library (`ThreadPoolExecutor` internals).
|
||||
|
||||
This instance will be `shutdown()` ASAP so it's fine to
|
||||
create one in such a weird place.
|
||||
|
||||
The code is thus 100% thread-safe as long as the stream readers
|
||||
are read-only except for return values and possible exceptions.
|
||||
"""
|
||||
stdin_overwrite = subprocess.PIPE
|
||||
# This already should be validated in a higher place of the stack.
|
||||
assert not (
|
||||
does_allow_interactive() and is_output_redirected()
|
||||
), "Cannot redirect output while in interactive mode."
|
||||
if process_runner != subprocess or (
|
||||
does_allow_interactive() and not is_output_redirected()
|
||||
):
|
||||
stdin_overwrite = None
|
||||
|
||||
# See implementation note #1
|
||||
|
||||
if use_login_shells or process_runner != subprocess:
|
||||
return process_runner.check_call(
|
||||
cmd,
|
||||
# See implementation note #2
|
||||
stdin=stdin_overwrite,
|
||||
stdout=stdout_file,
|
||||
stderr=stderr_file,
|
||||
)
|
||||
|
||||
with subprocess.Popen(
|
||||
cmd,
|
||||
# See implementation note #2
|
||||
stdin=stdin_overwrite,
|
||||
stdout=subprocess.PIPE,
|
||||
stderr=subprocess.PIPE,
|
||||
bufsize=1, # line buffering
|
||||
universal_newlines=True, # text mode outputs
|
||||
) as p:
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
|
||||
# Closing stdin might be necessary to signal EOF to some
|
||||
# apps (they might get stuck waiting for input forever otherwise).
|
||||
p.stdin.close()
|
||||
|
||||
# See implementation note #3
|
||||
with ThreadPoolExecutor(max_workers=2) as executor:
|
||||
stdout_future = executor.submit(
|
||||
_read_subprocess_stream, p.stdout, stdout_file, is_stdout=True
|
||||
)
|
||||
stderr_future = executor.submit(
|
||||
_read_subprocess_stream, p.stderr, stderr_file, is_stdout=False
|
||||
)
|
||||
# Wait for completion.
|
||||
executor.shutdown()
|
||||
|
||||
# Update `p.returncode`
|
||||
p.poll()
|
||||
|
||||
detected_special_case = stdout_future.result()
|
||||
if stderr_future.result() is not None:
|
||||
if detected_special_case is not None:
|
||||
# This might some day need to be changed.
|
||||
# We should probably make sure the two special cases
|
||||
# are compatible then and that we can handle both by
|
||||
# e.g. reporting both to the caller.
|
||||
raise ValueError(
|
||||
"Bug: found a special case in both stdout and "
|
||||
"stderr. This is not valid behavior at the time "
|
||||
"of writing this code."
|
||||
)
|
||||
detected_special_case = stderr_future.result()
|
||||
|
||||
if p.returncode > 0:
|
||||
# Process failed, but not due to a signal, since signals
|
||||
# set the exit code to a negative value.
|
||||
raise ProcessRunnerError(
|
||||
"Command failed",
|
||||
"ssh_command_failed",
|
||||
code=p.returncode,
|
||||
command=cmd,
|
||||
special_case=detected_special_case,
|
||||
)
|
||||
elif p.returncode < 0:
|
||||
# Process failed due to a signal, since signals
|
||||
# set the exit code to a negative value.
|
||||
raise ProcessRunnerError(
|
||||
"Command failed",
|
||||
"ssh_command_failed",
|
||||
code=p.returncode,
|
||||
command=cmd,
|
||||
special_case="died_to_signal",
|
||||
)
|
||||
|
||||
return p.returncode
|
||||
|
||||
|
||||
def run_cmd_redirected(
|
||||
cmd: List[str],
|
||||
process_runner: Any = subprocess,
|
||||
silent: bool = False,
|
||||
use_login_shells: bool = False,
|
||||
):
|
||||
"""Run a command and optionally redirect output to a file.
|
||||
|
||||
Args:
|
||||
cmd: Command to run.
|
||||
process_runner: Process runner used for executing commands.
|
||||
silent: If true, the command output will be silenced completely
|
||||
(redirected to /dev/null), unless verbose logging
|
||||
is enabled. Use this for running utility commands like
|
||||
rsync.
|
||||
use_login_shells: Whether to disable special output processing because
|
||||
an interactive login shell is in use.
|
||||
|
||||
Returns:
|
||||
The return code of the executed process.
|
||||
"""
|
||||
if silent and cli_logger.verbosity < 1:
|
||||
return _run_and_process_output(
|
||||
cmd,
|
||||
process_runner=process_runner,
|
||||
stdout_file=process_runner.DEVNULL,
|
||||
stderr_file=process_runner.DEVNULL,
|
||||
use_login_shells=use_login_shells,
|
||||
)
|
||||
|
||||
if not is_output_redirected():
|
||||
return _run_and_process_output(
|
||||
cmd,
|
||||
process_runner=process_runner,
|
||||
stdout_file=sys.stdout,
|
||||
stderr_file=sys.stderr,
|
||||
use_login_shells=use_login_shells,
|
||||
)
|
||||
else:
|
||||
tmpfile_path = os.path.join(
|
||||
tempfile.gettempdir(), "ray-up-{}-{}.txt".format(cmd[0], time.time())
|
||||
)
|
||||
with open(
|
||||
tmpfile_path,
|
||||
mode="w",
|
||||
# line buffering
|
||||
buffering=1,
|
||||
) as tmp:
|
||||
cli_logger.verbose("Command stdout is redirected to {}", cf.bold(tmp.name))
|
||||
|
||||
return _run_and_process_output(
|
||||
cmd,
|
||||
process_runner=process_runner,
|
||||
stdout_file=tmp,
|
||||
stderr_file=tmp,
|
||||
use_login_shells=use_login_shells,
|
||||
)
|
||||
|
||||
|
||||
def handle_ssh_fails(
|
||||
e: ProcessRunnerError,
|
||||
first_conn_refused_time: Optional[float],
|
||||
retry_interval: float,
|
||||
) -> Optional[float]:
|
||||
"""Handle SSH system failures coming from a subprocess.
|
||||
|
||||
Args:
|
||||
e: The `ProcessRunnerException` to handle.
|
||||
first_conn_refused_time:
|
||||
The time (as reported by this function) or None,
|
||||
indicating the last time a CONN_REFUSED error was caught.
|
||||
|
||||
After exceeding a patience value, the program will be aborted
|
||||
since SSH will likely never recover.
|
||||
retry_interval: The interval after which the command will be retried,
|
||||
used here just to inform the user.
|
||||
|
||||
Returns:
|
||||
The updated `first_conn_refused_time`, or None if not applicable.
|
||||
"""
|
||||
if e.msg_type != "ssh_command_failed":
|
||||
return
|
||||
|
||||
if e.special_case == "ssh_conn_refused":
|
||||
if (
|
||||
first_conn_refused_time is not None
|
||||
and time.time() - first_conn_refused_time > CONN_REFUSED_PATIENCE
|
||||
):
|
||||
cli_logger.error(
|
||||
"SSH connection was being refused "
|
||||
"for {} seconds. Head node assumed "
|
||||
"unreachable.",
|
||||
cf.bold(str(CONN_REFUSED_PATIENCE)),
|
||||
)
|
||||
cli_logger.abort(
|
||||
"Check the node's firewall settings "
|
||||
"and the cloud network configuration."
|
||||
)
|
||||
|
||||
cli_logger.warning("SSH connection was refused.")
|
||||
cli_logger.warning(
|
||||
"This might mean that the SSH daemon is "
|
||||
"still setting up, or that "
|
||||
"the host is inaccessible (e.g. due to "
|
||||
"a firewall)."
|
||||
)
|
||||
|
||||
return time.time()
|
||||
|
||||
if e.special_case in ["ssh_timeout", "ssh_conn_refused"]:
|
||||
cli_logger.print(
|
||||
"SSH still not available, retrying in {} seconds.",
|
||||
cf.bold(str(retry_interval)),
|
||||
)
|
||||
else:
|
||||
raise e
|
||||
|
||||
return first_conn_refused_time
|
||||
@@ -0,0 +1,576 @@
|
||||
import logging
|
||||
import os
|
||||
import subprocess
|
||||
import time
|
||||
import traceback
|
||||
from threading import Thread
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
import click
|
||||
|
||||
from ray._common.usage import usage_constants, usage_lib
|
||||
from ray.autoscaler._private import subprocess_output_util as cmd_output_util
|
||||
from ray.autoscaler._private.cli_logger import cf, cli_logger
|
||||
from ray.autoscaler._private.command_runner import (
|
||||
AUTOSCALER_NODE_START_WAIT_S,
|
||||
ProcessRunnerError,
|
||||
)
|
||||
from ray.autoscaler._private.constants import (
|
||||
LABELS_ENVIRONMENT_VARIABLE,
|
||||
RESOURCES_ENVIRONMENT_VARIABLE,
|
||||
)
|
||||
from ray.autoscaler._private.event_system import CreateClusterEvent, global_event_system
|
||||
from ray.autoscaler._private.log_timer import LogTimer
|
||||
from ray.autoscaler.tags import (
|
||||
STATUS_SETTING_UP,
|
||||
STATUS_SYNCING_FILES,
|
||||
STATUS_UP_TO_DATE,
|
||||
STATUS_UPDATE_FAILED,
|
||||
STATUS_WAITING_FOR_SSH,
|
||||
TAG_RAY_FILE_MOUNTS_CONTENTS,
|
||||
TAG_RAY_NODE_STATUS,
|
||||
TAG_RAY_RUNTIME_CONFIG,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
NUM_SETUP_STEPS = 7
|
||||
READY_CHECK_INTERVAL = 5
|
||||
|
||||
|
||||
class NodeUpdater:
|
||||
"""A process for syncing files and running init commands on a node.
|
||||
|
||||
Arguments:
|
||||
node_id: the Node ID
|
||||
provider_config: Provider section of autoscaler yaml
|
||||
provider: NodeProvider Class
|
||||
auth_config: Auth section of autoscaler yaml
|
||||
cluster_name: the name of the cluster.
|
||||
file_mounts: Map of remote to local paths
|
||||
initialization_commands: Commands run before container launch
|
||||
setup_commands: Commands run before ray starts
|
||||
ray_start_commands: Commands to start ray
|
||||
runtime_hash: Used to check for config changes
|
||||
file_mounts_contents_hash: Used to check for changes to file mounts
|
||||
is_head_node: Whether to use head start/setup commands
|
||||
node_resources: Optional resources string that the raylet should
|
||||
advertise on startup.
|
||||
node_labels: Optional labels for the node.
|
||||
cluster_synced_files: Files synced from the head to worker nodes after
|
||||
the head node finishes its setup.
|
||||
rsync_options: Extra options related to the rsync command.
|
||||
process_runner: the module to use to run the commands
|
||||
in the CommandRunner. E.g., subprocess.
|
||||
use_internal_ip: Wwhether the node_id belongs to an internal ip
|
||||
or external ip.
|
||||
docker_config: Docker section of autoscaler yaml
|
||||
restart_only: Whether to skip setup commands & just restart ray
|
||||
for_recovery: True if updater is for a recovering node. Only used for
|
||||
metric tracking.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
node_id: str,
|
||||
provider_config: dict,
|
||||
provider: Any,
|
||||
auth_config: dict,
|
||||
cluster_name: str,
|
||||
file_mounts: dict,
|
||||
initialization_commands: list,
|
||||
setup_commands: list,
|
||||
ray_start_commands: list,
|
||||
runtime_hash: str,
|
||||
file_mounts_contents_hash: str,
|
||||
is_head_node: bool,
|
||||
node_resources: Optional[str] = None,
|
||||
node_labels: Optional[Dict[str, str]] = None,
|
||||
cluster_synced_files: Optional[List[str]] = None,
|
||||
rsync_options: Optional[Dict[str, Any]] = None,
|
||||
process_runner: Any = subprocess,
|
||||
use_internal_ip: bool = False,
|
||||
docker_config: Optional[Dict[str, Any]] = None,
|
||||
restart_only: bool = False,
|
||||
for_recovery: bool = False,
|
||||
):
|
||||
self.log_prefix = "NodeUpdater: {}: ".format(node_id)
|
||||
# Three cases:
|
||||
# 1) use_internal_ip arg is True -> use_internal_ip is True
|
||||
# 2) worker node -> use value of provider_config["use_internal_ips"]
|
||||
# 3) head node -> use value of provider_config["use_internal_ips"] unless
|
||||
# overridden by provider_config["use_external_head_ip"]
|
||||
use_internal_ip = use_internal_ip or (
|
||||
provider_config.get("use_internal_ips", False)
|
||||
and not (
|
||||
is_head_node and provider_config.get("use_external_head_ip", False)
|
||||
)
|
||||
)
|
||||
self.cmd_runner = provider.get_command_runner(
|
||||
self.log_prefix,
|
||||
node_id,
|
||||
auth_config,
|
||||
cluster_name,
|
||||
process_runner,
|
||||
use_internal_ip,
|
||||
docker_config,
|
||||
)
|
||||
|
||||
self.daemon = True
|
||||
self.node_id = node_id
|
||||
self.provider_type = provider_config.get("type")
|
||||
self.provider = provider
|
||||
# Some node providers don't specify empty structures as
|
||||
# defaults. Better to be defensive.
|
||||
file_mounts = file_mounts or {}
|
||||
self.file_mounts = {
|
||||
remote: os.path.expanduser(local) for remote, local in file_mounts.items()
|
||||
}
|
||||
|
||||
self.initialization_commands = initialization_commands
|
||||
self.setup_commands = setup_commands
|
||||
self.ray_start_commands = ray_start_commands
|
||||
self.node_resources = node_resources
|
||||
self.node_labels = node_labels
|
||||
self.runtime_hash = runtime_hash
|
||||
self.file_mounts_contents_hash = file_mounts_contents_hash
|
||||
# TODO (Alex): This makes the assumption that $HOME on the head and
|
||||
# worker nodes is the same. Also note that `cluster_synced_files` is
|
||||
# set on the head -> worker updaters only (so `expanduser` is only run
|
||||
# on the head node).
|
||||
cluster_synced_files = cluster_synced_files or []
|
||||
self.cluster_synced_files = [
|
||||
os.path.expanduser(path) for path in cluster_synced_files
|
||||
]
|
||||
self.rsync_options = rsync_options or {}
|
||||
self.auth_config = auth_config
|
||||
self.is_head_node = is_head_node
|
||||
self.docker_config = docker_config
|
||||
self.restart_only = restart_only
|
||||
self.update_time = None
|
||||
self.for_recovery = for_recovery
|
||||
|
||||
def run(self):
|
||||
update_start_time = time.time()
|
||||
if (
|
||||
cmd_output_util.does_allow_interactive()
|
||||
and cmd_output_util.is_output_redirected()
|
||||
):
|
||||
# this is most probably a bug since the user has no control
|
||||
# over these settings
|
||||
msg = (
|
||||
"Output was redirected for an interactive command. "
|
||||
"Either do not pass `--redirect-command-output` "
|
||||
"or also pass in `--use-normal-shells`."
|
||||
)
|
||||
cli_logger.abort(msg)
|
||||
|
||||
try:
|
||||
with LogTimer(
|
||||
self.log_prefix + "Applied config {}".format(self.runtime_hash)
|
||||
):
|
||||
self.do_update()
|
||||
except Exception as e:
|
||||
self.provider.set_node_tags(
|
||||
self.node_id, {TAG_RAY_NODE_STATUS: STATUS_UPDATE_FAILED}
|
||||
)
|
||||
cli_logger.error("New status: {}", cf.bold(STATUS_UPDATE_FAILED))
|
||||
|
||||
cli_logger.error("!!!")
|
||||
if hasattr(e, "cmd"):
|
||||
stderr_output = getattr(e, "stderr", "No stderr available")
|
||||
cli_logger.error(
|
||||
"Setup command `{}` failed with exit code {}. stderr: {}",
|
||||
cf.bold(e.cmd),
|
||||
e.returncode,
|
||||
stderr_output,
|
||||
)
|
||||
else:
|
||||
cli_logger.verbose_error("Exception details: {}", str(vars(e)))
|
||||
full_traceback = traceback.format_exc()
|
||||
cli_logger.error("Full traceback: {}", full_traceback)
|
||||
# todo: handle this better somehow?
|
||||
cli_logger.error("Error message: {}", str(e))
|
||||
cli_logger.error("!!!")
|
||||
cli_logger.newline()
|
||||
|
||||
if isinstance(e, click.ClickException):
|
||||
# todo: why do we ignore this here
|
||||
return
|
||||
raise
|
||||
|
||||
tags_to_set = {
|
||||
TAG_RAY_NODE_STATUS: STATUS_UP_TO_DATE,
|
||||
TAG_RAY_RUNTIME_CONFIG: self.runtime_hash,
|
||||
}
|
||||
if self.file_mounts_contents_hash is not None:
|
||||
tags_to_set[TAG_RAY_FILE_MOUNTS_CONTENTS] = self.file_mounts_contents_hash
|
||||
|
||||
self.provider.set_node_tags(self.node_id, tags_to_set)
|
||||
cli_logger.labeled_value("New status", STATUS_UP_TO_DATE)
|
||||
|
||||
self.update_time = time.time() - update_start_time
|
||||
self.exitcode = 0
|
||||
|
||||
def sync_file_mounts(self, sync_cmd, step_numbers=(0, 2)):
|
||||
# step_numbers is (# of previous steps, total steps)
|
||||
previous_steps, total_steps = step_numbers
|
||||
|
||||
nolog_paths = []
|
||||
if cli_logger.verbosity == 0:
|
||||
nolog_paths = ["~/ray_bootstrap_key.pem", "~/ray_bootstrap_config.yaml"]
|
||||
|
||||
def do_sync(remote_path, local_path, allow_non_existing_paths=False):
|
||||
if allow_non_existing_paths and not os.path.exists(local_path):
|
||||
cli_logger.print("sync: {} does not exist. Skipping.", local_path)
|
||||
# Ignore missing source files. In the future we should support
|
||||
# the --delete-missing-args command to delete files that have
|
||||
# been removed
|
||||
return
|
||||
|
||||
assert os.path.exists(local_path), local_path
|
||||
|
||||
if os.path.isdir(local_path):
|
||||
if not local_path.endswith("/"):
|
||||
local_path += "/"
|
||||
if not remote_path.endswith("/"):
|
||||
remote_path += "/"
|
||||
|
||||
with LogTimer(
|
||||
self.log_prefix + "Synced {} to {}".format(local_path, remote_path)
|
||||
):
|
||||
is_docker = (
|
||||
self.docker_config and self.docker_config["container_name"] != ""
|
||||
)
|
||||
if not is_docker:
|
||||
# The DockerCommandRunner handles this internally.
|
||||
self.cmd_runner.run(
|
||||
"mkdir -p {}".format(os.path.dirname(remote_path)),
|
||||
run_env="host",
|
||||
)
|
||||
sync_cmd(local_path, remote_path, docker_mount_if_possible=True)
|
||||
|
||||
if remote_path not in nolog_paths:
|
||||
# todo: timed here?
|
||||
cli_logger.print(
|
||||
"{} from {}", cf.bold(remote_path), cf.bold(local_path)
|
||||
)
|
||||
|
||||
# Rsync file mounts
|
||||
with cli_logger.group(
|
||||
"Processing file mounts", _numbered=("[]", previous_steps + 1, total_steps)
|
||||
):
|
||||
for remote_path, local_path in self.file_mounts.items():
|
||||
do_sync(remote_path, local_path)
|
||||
previous_steps += 1
|
||||
|
||||
if self.cluster_synced_files:
|
||||
with cli_logger.group(
|
||||
"Processing worker file mounts",
|
||||
_numbered=("[]", previous_steps + 1, total_steps),
|
||||
):
|
||||
cli_logger.print("synced files: {}", str(self.cluster_synced_files))
|
||||
for path in self.cluster_synced_files:
|
||||
do_sync(path, path, allow_non_existing_paths=True)
|
||||
previous_steps += 1
|
||||
else:
|
||||
cli_logger.print(
|
||||
"No worker file mounts to sync",
|
||||
_numbered=("[]", previous_steps + 1, total_steps),
|
||||
)
|
||||
|
||||
def wait_ready(self, deadline):
|
||||
with cli_logger.group(
|
||||
"Waiting for SSH to become available", _numbered=("[]", 1, NUM_SETUP_STEPS)
|
||||
):
|
||||
with LogTimer(self.log_prefix + "Got remote shell"):
|
||||
cli_logger.print("Running `{}` as a test.", cf.bold("uptime"))
|
||||
first_conn_refused_time = None
|
||||
while True:
|
||||
if time.time() > deadline:
|
||||
raise Exception("wait_ready timeout exceeded.")
|
||||
if self.provider.is_terminated(self.node_id):
|
||||
raise Exception(
|
||||
"wait_ready aborting because node "
|
||||
"detected as terminated."
|
||||
)
|
||||
|
||||
try:
|
||||
# Run outside of the container
|
||||
self.cmd_runner.run("uptime", timeout=10, run_env="host")
|
||||
cli_logger.success("Success.")
|
||||
return True
|
||||
except ProcessRunnerError as e:
|
||||
first_conn_refused_time = cmd_output_util.handle_ssh_fails(
|
||||
e,
|
||||
first_conn_refused_time,
|
||||
retry_interval=READY_CHECK_INTERVAL,
|
||||
)
|
||||
time.sleep(READY_CHECK_INTERVAL)
|
||||
except Exception as e:
|
||||
# TODO(maximsmol): we should not be ignoring
|
||||
# exceptions if they get filtered properly
|
||||
# (new style log + non-interactive shells)
|
||||
#
|
||||
# however threading this configuration state
|
||||
# is a pain and I'm leaving it for later
|
||||
|
||||
retry_str = "(" + str(e) + ")"
|
||||
if hasattr(e, "cmd"):
|
||||
if isinstance(e.cmd, str):
|
||||
cmd_ = e.cmd
|
||||
elif isinstance(e.cmd, list):
|
||||
cmd_ = " ".join(e.cmd)
|
||||
else:
|
||||
logger.debug(
|
||||
f"e.cmd type ({type(e.cmd)}) not list or str."
|
||||
)
|
||||
cmd_ = str(e.cmd)
|
||||
retry_str = "(Exit Status {}): {}".format(
|
||||
e.returncode, cmd_
|
||||
)
|
||||
|
||||
cli_logger.print(
|
||||
"SSH still not available {}, retrying in {} seconds.",
|
||||
cf.dimmed(retry_str),
|
||||
cf.bold(str(READY_CHECK_INTERVAL)),
|
||||
)
|
||||
|
||||
time.sleep(READY_CHECK_INTERVAL)
|
||||
|
||||
def do_update(self):
|
||||
self.provider.set_node_tags(
|
||||
self.node_id, {TAG_RAY_NODE_STATUS: STATUS_WAITING_FOR_SSH}
|
||||
)
|
||||
cli_logger.labeled_value("New status", STATUS_WAITING_FOR_SSH)
|
||||
|
||||
deadline = time.time() + AUTOSCALER_NODE_START_WAIT_S
|
||||
self.wait_ready(deadline)
|
||||
global_event_system.execute_callback(CreateClusterEvent.ssh_control_acquired)
|
||||
|
||||
node_tags = self.provider.node_tags(self.node_id)
|
||||
logger.debug("Node tags: {}".format(str(node_tags)))
|
||||
|
||||
if self.provider_type == "aws" and self.provider.provider_config:
|
||||
from ray.autoscaler._private.aws.cloudwatch.cloudwatch_helper import (
|
||||
CloudwatchHelper,
|
||||
)
|
||||
|
||||
CloudwatchHelper(
|
||||
self.provider.provider_config, self.node_id, self.provider.cluster_name
|
||||
).update_from_config(self.is_head_node)
|
||||
|
||||
if node_tags.get(TAG_RAY_RUNTIME_CONFIG) == self.runtime_hash:
|
||||
# When resuming from a stopped instance the runtime_hash may be the
|
||||
# same, but the container will not be started.
|
||||
init_required = self.cmd_runner.run_init(
|
||||
as_head=self.is_head_node,
|
||||
file_mounts=self.file_mounts,
|
||||
sync_run_yet=False,
|
||||
)
|
||||
if init_required:
|
||||
node_tags[TAG_RAY_RUNTIME_CONFIG] += "-invalidate"
|
||||
# This ensures that `setup_commands` are not removed
|
||||
self.restart_only = False
|
||||
|
||||
if self.restart_only:
|
||||
self.setup_commands = []
|
||||
|
||||
# runtime_hash will only change whenever the user restarts
|
||||
# or updates their cluster with `get_or_create_head_node`
|
||||
if node_tags.get(TAG_RAY_RUNTIME_CONFIG) == self.runtime_hash and (
|
||||
not self.file_mounts_contents_hash
|
||||
or node_tags.get(TAG_RAY_FILE_MOUNTS_CONTENTS)
|
||||
== self.file_mounts_contents_hash
|
||||
):
|
||||
# todo: we lie in the confirmation message since
|
||||
# full setup might be cancelled here
|
||||
cli_logger.print(
|
||||
"Configuration already up to date, "
|
||||
"skipping file mounts, initalization and setup commands.",
|
||||
_numbered=("[]", "2-6", NUM_SETUP_STEPS),
|
||||
)
|
||||
|
||||
else:
|
||||
cli_logger.print(
|
||||
"Updating cluster configuration.", _tags=dict(hash=self.runtime_hash)
|
||||
)
|
||||
|
||||
self.provider.set_node_tags(
|
||||
self.node_id, {TAG_RAY_NODE_STATUS: STATUS_SYNCING_FILES}
|
||||
)
|
||||
cli_logger.labeled_value("New status", STATUS_SYNCING_FILES)
|
||||
self.sync_file_mounts(self.rsync_up, step_numbers=(1, NUM_SETUP_STEPS))
|
||||
|
||||
# Only run setup commands if runtime_hash has changed because
|
||||
# we don't want to run setup_commands every time the head node
|
||||
# file_mounts folders have changed.
|
||||
if node_tags.get(TAG_RAY_RUNTIME_CONFIG) != self.runtime_hash:
|
||||
# Run init commands
|
||||
self.provider.set_node_tags(
|
||||
self.node_id, {TAG_RAY_NODE_STATUS: STATUS_SETTING_UP}
|
||||
)
|
||||
cli_logger.labeled_value("New status", STATUS_SETTING_UP)
|
||||
|
||||
if self.initialization_commands:
|
||||
with cli_logger.group(
|
||||
"Running initialization commands",
|
||||
_numbered=("[]", 4, NUM_SETUP_STEPS),
|
||||
):
|
||||
global_event_system.execute_callback(
|
||||
CreateClusterEvent.run_initialization_cmd
|
||||
)
|
||||
with LogTimer(
|
||||
self.log_prefix + "Initialization commands",
|
||||
show_status=True,
|
||||
):
|
||||
for cmd in self.initialization_commands:
|
||||
global_event_system.execute_callback(
|
||||
CreateClusterEvent.run_initialization_cmd,
|
||||
{"command": cmd},
|
||||
)
|
||||
try:
|
||||
# Overriding the existing SSHOptions class
|
||||
# with a new SSHOptions class that uses
|
||||
# this ssh_private_key as its only __init__
|
||||
# argument.
|
||||
# Run outside docker.
|
||||
self.cmd_runner.run(
|
||||
cmd,
|
||||
ssh_options_override_ssh_key=self.auth_config.get( # noqa: E501
|
||||
"ssh_private_key"
|
||||
),
|
||||
run_env="host",
|
||||
)
|
||||
except ProcessRunnerError as e:
|
||||
if e.msg_type == "ssh_command_failed":
|
||||
cli_logger.error("Failed.")
|
||||
cli_logger.error("See above for stderr.")
|
||||
|
||||
raise click.ClickException(
|
||||
"Initialization command failed."
|
||||
) from None
|
||||
else:
|
||||
cli_logger.print(
|
||||
"No initialization commands to run.",
|
||||
_numbered=("[]", 4, NUM_SETUP_STEPS),
|
||||
)
|
||||
with cli_logger.group(
|
||||
"Initializing command runner",
|
||||
# todo: fix command numbering
|
||||
_numbered=("[]", 5, NUM_SETUP_STEPS),
|
||||
):
|
||||
self.cmd_runner.run_init(
|
||||
as_head=self.is_head_node,
|
||||
file_mounts=self.file_mounts,
|
||||
sync_run_yet=True,
|
||||
)
|
||||
if self.setup_commands:
|
||||
with cli_logger.group(
|
||||
"Running setup commands",
|
||||
# todo: fix command numbering
|
||||
_numbered=("[]", 6, NUM_SETUP_STEPS),
|
||||
):
|
||||
global_event_system.execute_callback(
|
||||
CreateClusterEvent.run_setup_cmd
|
||||
)
|
||||
with LogTimer(
|
||||
self.log_prefix + "Setup commands", show_status=True
|
||||
):
|
||||
total = len(self.setup_commands)
|
||||
for i, cmd in enumerate(self.setup_commands):
|
||||
global_event_system.execute_callback(
|
||||
CreateClusterEvent.run_setup_cmd, {"command": cmd}
|
||||
)
|
||||
if cli_logger.verbosity == 0 and len(cmd) > 30:
|
||||
cmd_to_print = cf.bold(cmd[:30]) + "..."
|
||||
else:
|
||||
cmd_to_print = cf.bold(cmd)
|
||||
|
||||
cli_logger.print(
|
||||
"{}", cmd_to_print, _numbered=("()", i, total)
|
||||
)
|
||||
|
||||
try:
|
||||
# Runs in the container if docker is in use
|
||||
self.cmd_runner.run(cmd, run_env="auto")
|
||||
except ProcessRunnerError as e:
|
||||
if e.msg_type == "ssh_command_failed":
|
||||
cli_logger.error("Failed.")
|
||||
cli_logger.error("See above for stderr.")
|
||||
|
||||
raise click.ClickException("Setup command failed.")
|
||||
else:
|
||||
cli_logger.print(
|
||||
"No setup commands to run.",
|
||||
_numbered=("[]", 6, NUM_SETUP_STEPS),
|
||||
)
|
||||
|
||||
with cli_logger.group(
|
||||
"Starting the Ray runtime", _numbered=("[]", 7, NUM_SETUP_STEPS)
|
||||
):
|
||||
global_event_system.execute_callback(CreateClusterEvent.start_ray_runtime)
|
||||
with LogTimer(self.log_prefix + "Ray start commands", show_status=True):
|
||||
for cmd in self.ray_start_commands:
|
||||
env_vars = {}
|
||||
if self.is_head_node:
|
||||
if usage_lib.usage_stats_enabled():
|
||||
env_vars[usage_constants.USAGE_STATS_ENABLED_ENV_VAR] = 1
|
||||
else:
|
||||
# Disable usage stats collection in the cluster.
|
||||
env_vars[usage_constants.USAGE_STATS_ENABLED_ENV_VAR] = 0
|
||||
|
||||
# Add a resource override env variable if needed.
|
||||
# Local NodeProvider doesn't need resource and label override.
|
||||
if self.provider_type != "local":
|
||||
if self.node_resources:
|
||||
env_vars[
|
||||
RESOURCES_ENVIRONMENT_VARIABLE
|
||||
] = self.node_resources
|
||||
if self.node_labels:
|
||||
env_vars[LABELS_ENVIRONMENT_VARIABLE] = self.node_labels
|
||||
|
||||
try:
|
||||
old_redirected = cmd_output_util.is_output_redirected()
|
||||
cmd_output_util.set_output_redirected(False)
|
||||
# Runs in the container if docker is in use
|
||||
self.cmd_runner.run(
|
||||
cmd, environment_variables=env_vars, run_env="auto"
|
||||
)
|
||||
cmd_output_util.set_output_redirected(old_redirected)
|
||||
except ProcessRunnerError as e:
|
||||
if e.msg_type == "ssh_command_failed":
|
||||
cli_logger.error("Failed.")
|
||||
cli_logger.error("See above for stderr.")
|
||||
|
||||
raise click.ClickException("Start command failed.")
|
||||
global_event_system.execute_callback(
|
||||
CreateClusterEvent.start_ray_runtime_completed
|
||||
)
|
||||
|
||||
def rsync_up(self, source, target, docker_mount_if_possible=False):
|
||||
options = {}
|
||||
options["docker_mount_if_possible"] = docker_mount_if_possible
|
||||
options["rsync_exclude"] = self.rsync_options.get("rsync_exclude")
|
||||
options["rsync_filter"] = self.rsync_options.get("rsync_filter")
|
||||
self.cmd_runner.run_rsync_up(source, target, options=options)
|
||||
cli_logger.verbose(
|
||||
"`rsync`ed {} (local) to {} (remote)", cf.bold(source), cf.bold(target)
|
||||
)
|
||||
|
||||
def rsync_down(self, source, target, docker_mount_if_possible=False):
|
||||
options = {}
|
||||
options["docker_mount_if_possible"] = docker_mount_if_possible
|
||||
options["rsync_exclude"] = self.rsync_options.get("rsync_exclude")
|
||||
options["rsync_filter"] = self.rsync_options.get("rsync_filter")
|
||||
self.cmd_runner.run_rsync_down(source, target, options=options)
|
||||
cli_logger.verbose(
|
||||
"`rsync`ed {} (remote) to {} (local)", cf.bold(source), cf.bold(target)
|
||||
)
|
||||
|
||||
|
||||
class NodeUpdaterThread(NodeUpdater, Thread):
|
||||
def __init__(self, *args, **kwargs):
|
||||
Thread.__init__(self)
|
||||
NodeUpdater.__init__(self, *args, **kwargs)
|
||||
self.exitcode = -1
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,45 @@
|
||||
# Ray on vSphere Architecture Guide
|
||||
|
||||
To support ray on vSphere, the implementation has been added into [python/ray/autoscaler/_private/vsphere](../vsphere) directory. The following sections will explain the vSphere terminologies used in the code and also explain the whole code flow.
|
||||
|
||||
|
||||
# vSphere Terminologies
|
||||
## [OVF file](https://techdocs.broadcom.com/us/en/vmware-cis/vsphere/vsphere/8-0/vsphere-virtual-machine-administration-guide-8-0.html)
|
||||
OVF format is a packaging and distribution format for virtual machines. It is a standard which can be used to describe the VM metadata. We use the OVF files to create the virtual machines which will act as Ray head and worker node.
|
||||
|
||||
## VI Admin
|
||||
|
||||
The term VI stands for [Virtual Infrastructure](https://techdocs.broadcom.com/us/en/vmware-cis/vsphere/vsphere/8-0/vsphere-virtual-machine-administration-guide-8-0/introduction-to-vmware-vsphere-virtual-machinesvsphere-vm-admin/virtual-machines-and-the-virtual-infrastructurevsphere-vm-admin.html).
|
||||
|
||||
A VI Admin is used to describe a persona that manages the lifecycle of VMware infrastructure. VI Admins engage in a range of activities. A subset of them are listed below:
|
||||
1. Provisioning [ESXi](https://www.vmware.com/in/products/esxi-and-esx.html) (Hypervisor developed by VMware) hosts.
|
||||
2. Provisioning a vSphere infrastructure.
|
||||
3. Managing lifecycle of VMs.
|
||||
4. Provisioning [vSAN](https://docs.vmware.com/en/VMware-vSAN/index.html) storage.
|
||||
|
||||
# Code Flow
|
||||
|
||||
## Node Creation on `ray up`
|
||||
The following sections explain the code flow in a sequential manner. The execution is triggered from the moment user executed `ray up` command
|
||||
|
||||
### Inject private Key ([config.py](./config.py))
|
||||
During running `ray up`, the private key is injected into `config["auth"]["ssh_private_key"]`. The bootstrap machine (where the `ray up` command is executed) and the head node subsequently use this key to SSH onto the ray worker nodes.
|
||||
|
||||
### Update vSphere Configs ([config.py](./config.py))
|
||||
Used to make sure that the user has created the YAML file with valid configs.
|
||||
|
||||
### Create Nodes ([node_provider.py](./cluster_operator_client.py))
|
||||
|
||||
#### Call `create_node`
|
||||
Starts the creation of nodes with `create_node` function, which internally calls `_create_node`.
|
||||
|
||||
## Autoscaling
|
||||
|
||||
### Get and create nodes ([node_provider.py](./cluster_operator_client.py))
|
||||
The autoscaler can find the currently running nodes with `non_terminated_nodes` function and can request for new nodes by calling `create_node` function.
|
||||
|
||||
### Fetch node IPs ([node_provider.py](./cluster_operator_client.py))
|
||||
The autoscaler can use `external_ip` or `internal_ip` function to fetch a node's IP.
|
||||
|
||||
## Cluster tear down ([node_provider.py](./cluster_operator_client.py))
|
||||
`terminate_nodes` function gets called on ray down command's execution. It deletes all the nodes.
|
||||
@@ -0,0 +1,685 @@
|
||||
import base64
|
||||
import ipaddress
|
||||
import logging
|
||||
import os
|
||||
import random
|
||||
import string
|
||||
from enum import Enum
|
||||
from threading import RLock
|
||||
from typing import Any, Dict, Optional, Tuple
|
||||
|
||||
from kubernetes import client, config
|
||||
|
||||
from ray.autoscaler.tags import (
|
||||
NODE_KIND_HEAD,
|
||||
NODE_KIND_WORKER,
|
||||
STATUS_SETTING_UP,
|
||||
STATUS_UNINITIALIZED,
|
||||
STATUS_UP_TO_DATE,
|
||||
TAG_RAY_CLUSTER_NAME,
|
||||
TAG_RAY_NODE_KIND,
|
||||
TAG_RAY_NODE_NAME,
|
||||
TAG_RAY_NODE_STATUS,
|
||||
TAG_RAY_USER_NODE_TYPE,
|
||||
)
|
||||
|
||||
# Design:
|
||||
|
||||
# Each modification the autoscaler wants to make is posted to the API server's desired
|
||||
# state (e.g. if the autoscaler wants to scale up, it adds VM name to the desired
|
||||
# worker list it wants to scale, if it wants to scale down it removes the name from
|
||||
# the list).
|
||||
|
||||
# VMRay CRD
|
||||
VMRAY_CRD_VER = os.getenv("VMRAY_CRD_VER", "v1alpha1")
|
||||
VMRAY_GROUP = "vmray.broadcom.com"
|
||||
VMRAYCLUSTER_PLURAL = "vmrayclusters"
|
||||
|
||||
# VirtualMachineService CRD
|
||||
VMSERVICE_CRD_VER = os.getenv("VMSERVICE_CRD_VER", "v1alpha2")
|
||||
VMSERVICE_GROUP = "vmoperator.vmware.com"
|
||||
VMSERVICE_PLURAL = "virtualmachineservices"
|
||||
|
||||
SERVICE_ACCOUNT_TOKEN = os.getenv("SVC_ACCOUNT_TOKEN", None)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
cur_path = os.path.dirname(__file__)
|
||||
|
||||
|
||||
class VMNodeStatus(Enum):
|
||||
INITIALIZED = "initialized"
|
||||
RUNNING = "running"
|
||||
FAIL = "failure"
|
||||
|
||||
|
||||
class KubernetesHttpApiClient(object):
|
||||
def __init__(self, ca_cert: str, api_server: str):
|
||||
token = SERVICE_ACCOUNT_TOKEN
|
||||
# If SERVICE_ACCOUNT_TOKEN not present, use local
|
||||
# ~/.kube/config file. Active context will be used.
|
||||
# This is useful when Ray CLI are used and local autoscaler needs
|
||||
# communicate with the k8s API server
|
||||
# If the token is present then use that for communication.
|
||||
if not token:
|
||||
self.client = client.ApiClient(config.load_kube_config())
|
||||
else:
|
||||
configuration = client.Configuration()
|
||||
configuration.api_key["authorization"] = token
|
||||
configuration.api_key_prefix["authorization"] = "Bearer"
|
||||
configuration.host = f"https://{api_server}"
|
||||
if ca_cert:
|
||||
configuration.ssl_ca_cert = ca_cert
|
||||
else:
|
||||
configuration.verify_ssl = False
|
||||
self.client = client.ApiClient(configuration)
|
||||
|
||||
# Use customObjectsApi to access custom resources
|
||||
self.custom_object_api = client.CustomObjectsApi(self.client)
|
||||
|
||||
|
||||
class ClusterOperatorClient(KubernetesHttpApiClient):
|
||||
def __init__(
|
||||
self,
|
||||
cluster_name: str,
|
||||
provider_config: Dict[str, Any],
|
||||
cluster_config: Dict[str, Any],
|
||||
):
|
||||
self.cluster_name = cluster_name
|
||||
self.vmraycluster_nounce = None
|
||||
self.max_worker_nodes = None
|
||||
|
||||
self.vsphere_config = provider_config["vsphere_config"]
|
||||
|
||||
self.namespace = self.vsphere_config["namespace"]
|
||||
self.k8s_api_client = KubernetesHttpApiClient(
|
||||
self.vsphere_config.get("ca_cert"),
|
||||
self.vsphere_config.get("api_server"),
|
||||
)
|
||||
self.lock = RLock()
|
||||
|
||||
if cluster_config:
|
||||
self.max_worker_nodes = cluster_config["max_workers"]
|
||||
self.head_setup_commands = cluster_config["head_setup_commands"]
|
||||
self.available_node_types = cluster_config["available_node_types"]
|
||||
self.head_node_type = cluster_config["head_node_type"]
|
||||
|
||||
# docker configurations.
|
||||
self.provider_auth = cluster_config["auth"]
|
||||
self.docker = cluster_config["docker"]
|
||||
|
||||
# create docker login info secret, if it exists.
|
||||
docker_auth_secret_name = self._create_docker_auth_secrets()
|
||||
if docker_auth_secret_name:
|
||||
self.docker_config = {
|
||||
"auth_secret_name": docker_auth_secret_name,
|
||||
}
|
||||
else:
|
||||
self.docker_config = None
|
||||
else:
|
||||
self._set_max_worker_nodes()
|
||||
self._create_tls_secrets()
|
||||
|
||||
def _create_docker_auth_secrets(self):
|
||||
docker_auth_secret_name = self.cluster_name + "-docker-auth"
|
||||
docker_auth = self.vsphere_config.get("docker_auth", {})
|
||||
username = docker_auth.get("username", None)
|
||||
password = docker_auth.get("password", None)
|
||||
kp = {}
|
||||
if username and password:
|
||||
kp["username"] = username
|
||||
kp["password"] = password
|
||||
registry = docker_auth.get("registry", None)
|
||||
if registry:
|
||||
kp["registry"] = registry
|
||||
self._create_secret(self.namespace, docker_auth_secret_name, kp)
|
||||
return docker_auth_secret_name
|
||||
return None
|
||||
|
||||
def _create_tls_secrets(self):
|
||||
# If token is passed that means its instance of autoscaler
|
||||
# running inside the head node, so validate if tls server cert
|
||||
# and key are available then create a secret with their
|
||||
# value.
|
||||
tls_enabled = os.environ.get("RAY_USE_TLS", None) == "1"
|
||||
if not SERVICE_ACCOUNT_TOKEN or not tls_enabled:
|
||||
return
|
||||
|
||||
tls_cert = None
|
||||
tls_key = None
|
||||
|
||||
cert_path = os.environ.get("RAY_TLS_SERVER_CERT", None)
|
||||
key_path = os.environ.get("RAY_TLS_SERVER_KEY", None)
|
||||
if cert_path:
|
||||
with open(cert_path) as f:
|
||||
tls_cert = f.read()
|
||||
if key_path:
|
||||
with open(key_path) as f:
|
||||
tls_key = f.read()
|
||||
|
||||
if tls_cert and tls_key:
|
||||
kp = {"tls.crt": tls_cert, "tls.key": tls_key}
|
||||
self._create_secret(self.namespace, self.cluster_name + "-tls", kp)
|
||||
|
||||
def list_vms(self, tag_filters: Dict[str, str]) -> Tuple[list, dict]:
|
||||
"""Queries K8s for VMs in the RayCluster and filter them as per
|
||||
tags provided in the tag_filters.
|
||||
"""
|
||||
logger.info(f"Getting nodes using tags \n{tag_filters}")
|
||||
tag_cache = {}
|
||||
|
||||
filters = tag_filters.copy()
|
||||
# Use Ray cluster name to get resources
|
||||
if TAG_RAY_CLUSTER_NAME not in tag_filters:
|
||||
filters[TAG_RAY_CLUSTER_NAME] = self.cluster_name
|
||||
nodes = []
|
||||
vmray_cluster_response = self._get_cluster_response()
|
||||
if not vmray_cluster_response:
|
||||
return nodes, tag_cache
|
||||
|
||||
vmray_cluster_status = vmray_cluster_response.get("status", {})
|
||||
if not vmray_cluster_status:
|
||||
return nodes, tag_cache
|
||||
|
||||
vmray_cluster_spec = vmray_cluster_response.get("spec", {})
|
||||
|
||||
# Check for a head node
|
||||
if NODE_KIND_HEAD in tag_filters.values() or not tag_filters:
|
||||
head_node_status = vmray_cluster_status.get("head_node_status", {})
|
||||
# head node found
|
||||
if head_node_status:
|
||||
node_id = self._get_head_name()
|
||||
nodes.append(node_id)
|
||||
|
||||
# Setting head node status
|
||||
status = head_node_status.get("vm_status", None)
|
||||
head_nt = vmray_cluster_spec["head_node"]["node_type"]
|
||||
tag_cache[node_id] = self._set_tags(
|
||||
node_id, NODE_KIND_HEAD, head_nt, status, filters
|
||||
)
|
||||
# Check current worker nodes
|
||||
if NODE_KIND_WORKER in tag_filters.values() or not tag_filters:
|
||||
current_workers = vmray_cluster_status.get("current_workers", {})
|
||||
desired_workers = vmray_cluster_spec.get("autoscaler_desired_workers", {})
|
||||
|
||||
# worker nodes found
|
||||
for worker in current_workers.keys():
|
||||
nodes.append(worker)
|
||||
# setting worker node status
|
||||
status = current_workers[worker].get("vm_status", None)
|
||||
node_type = desired_workers.get(worker, "")
|
||||
|
||||
tag_cache[worker] = self._set_tags(
|
||||
worker, NODE_KIND_WORKER, node_type, status, filters
|
||||
)
|
||||
|
||||
# List VMs from the desired workers' list
|
||||
for worker in desired_workers.keys():
|
||||
if worker in current_workers.keys():
|
||||
continue
|
||||
nodes.append(worker)
|
||||
node_type = desired_workers.get(worker, "")
|
||||
tag_cache[worker] = self._set_tags(
|
||||
worker, NODE_KIND_WORKER, node_type, STATUS_SETTING_UP, filters
|
||||
)
|
||||
|
||||
logger.info(
|
||||
f"Non terminated nodes {nodes}, Tags for these are: {tag_cache}"
|
||||
)
|
||||
return nodes, tag_cache
|
||||
|
||||
def is_vm_power_on(self, node_id: str) -> bool:
|
||||
"""Check current vm list. If its state is Running then return
|
||||
true else false."""
|
||||
node = self._get_node(node_id)
|
||||
if node:
|
||||
return node.get("vm_status", None) == VMNodeStatus.RUNNING.value
|
||||
logger.info(f"VM {node_id} not found")
|
||||
return False
|
||||
|
||||
def is_vm_creating(self, node_id: str) -> bool:
|
||||
"""Check current vm list. If its state is INITIALIZED then return
|
||||
true else false."""
|
||||
node = self._get_node(node_id)
|
||||
if node:
|
||||
return node.get("vm_status", None) == VMNodeStatus.INITIALIZED.value
|
||||
logger.info(f"VM {node_id} is not yet initialized")
|
||||
return False
|
||||
|
||||
def set_node_tags(self, tags: Dict[str, str]) -> None:
|
||||
"""
|
||||
Not required
|
||||
"""
|
||||
pass
|
||||
|
||||
def get_vm_external_ip(self, node_id: str) -> Optional[str]:
|
||||
"""Check current worker list and get the external ip."""
|
||||
node = {}
|
||||
# For a Ray head node, return external IP of the VMService
|
||||
# Ray head node is not accessible directly.
|
||||
if node_id == self._get_head_name():
|
||||
ingress = self._get_vm_service_ingress()
|
||||
for item in ingress:
|
||||
if "ip" in item.keys():
|
||||
node = item
|
||||
break
|
||||
else:
|
||||
worker_node = self._get_node(node_id)
|
||||
if (
|
||||
worker_node
|
||||
and worker_node.get("vm_status", None) == VMNodeStatus.RUNNING.value
|
||||
):
|
||||
node = worker_node
|
||||
ip = node.get("ip", None)
|
||||
# Validate returned IP
|
||||
if ip and _is_ipv4(ip):
|
||||
return ip
|
||||
logger.info(
|
||||
f"External IPv4 address: {ip} of VM: {node_id}"
|
||||
f"is either invalid or not available"
|
||||
)
|
||||
return None
|
||||
|
||||
def delete_node(self, node_id: str) -> None:
|
||||
"""Remove name of the vm from the desired worker list and patch
|
||||
the VmRayCluster CR"""
|
||||
with self.lock:
|
||||
vmray_cluster_response = self._get_cluster_response()
|
||||
vmray_cluster_spec = vmray_cluster_response.get("spec", {})
|
||||
|
||||
# Get desired workers
|
||||
desired_workers = vmray_cluster_spec.get("autoscaler_desired_workers", {})
|
||||
logger.info(f"Current desired workers: {desired_workers}")
|
||||
|
||||
# remove the node from the desired workers list
|
||||
if node_id in desired_workers:
|
||||
|
||||
# By default it follow patch application of `merge-patch+json`
|
||||
# so we need to remove the node ids by making them null.
|
||||
# refs:
|
||||
# 1. https://kubernetes.io/docs/tasks/manage-kubernetes-objects/
|
||||
# update-api-object-kubectl-patch/#use-a-json-merge-patch-to-update-a-deployment
|
||||
# 2. https://github.com/kubernetes-client/python/blob/master/kubernetes/
|
||||
# client/api/custom_objects_api.py#L3106
|
||||
payload = {"spec": {"autoscaler_desired_workers": {node_id: None}}}
|
||||
|
||||
logger.info(f"Deleting VM {node_id} | payload: {payload}")
|
||||
self.k8s_api_client.custom_object_api.patch_namespaced_custom_object(
|
||||
VMRAY_GROUP,
|
||||
VMRAY_CRD_VER,
|
||||
self.namespace,
|
||||
VMRAYCLUSTER_PLURAL,
|
||||
self.cluster_name,
|
||||
payload,
|
||||
async_req=False,
|
||||
)
|
||||
elif node_id == self._get_head_name():
|
||||
# Handle case to delete a head node
|
||||
# Delete VMRayCluster which will delete head node
|
||||
# as well as associated secrets and other resources.
|
||||
self.k8s_api_client.custom_object_api.delete_namespaced_custom_object(
|
||||
VMRAY_GROUP,
|
||||
VMRAY_CRD_VER,
|
||||
self.namespace,
|
||||
VMRAYCLUSTER_PLURAL,
|
||||
self.cluster_name,
|
||||
)
|
||||
|
||||
def create_nodes(
|
||||
self,
|
||||
tags: Dict[str, str],
|
||||
to_be_launched_node_count: int,
|
||||
node_config: Dict[str, Any],
|
||||
) -> Optional[Dict[str, Any]]:
|
||||
"""Ask cluster operator to create worker VMs"""
|
||||
logger.info(
|
||||
f"Creating {to_be_launched_node_count} nodes with tags: {tags}"
|
||||
f"and with config: {node_config}"
|
||||
)
|
||||
created_nodes_dict = {}
|
||||
with self.lock:
|
||||
if to_be_launched_node_count > 0:
|
||||
new_desired_workers = {}
|
||||
|
||||
new_vm_names = {}
|
||||
for _ in range(to_be_launched_node_count):
|
||||
name = self._create_node_name(tags[TAG_RAY_NODE_NAME])
|
||||
new_vm_names[name] = tags[TAG_RAY_USER_NODE_TYPE]
|
||||
|
||||
# Create a head node
|
||||
# Autoscaler sends a tag
|
||||
# head_node_tags[TAG_RAY_NODE_NAME] = "ray-{}-head".format(
|
||||
# config["cluster_name"])
|
||||
if "head" in tags[TAG_RAY_NODE_NAME]:
|
||||
# head node will be created as a part of VMRayCluster CR
|
||||
self._create_ssh_secret()
|
||||
self._create_vmraycluster()
|
||||
else:
|
||||
# Once VMRayCluster CR is created, update it to create worker
|
||||
# nodes.
|
||||
vmray_cluster_response = self._get_cluster_response()
|
||||
vmray_cluster_spec = vmray_cluster_response.get("spec", {})
|
||||
|
||||
# get desired workers
|
||||
desired_workers = vmray_cluster_spec.get(
|
||||
"autoscaler_desired_workers", {}
|
||||
)
|
||||
|
||||
# If workers are present in both the list then it shows stable
|
||||
# state for the cluster.
|
||||
# Append new VM names with existing one
|
||||
if desired_workers:
|
||||
new_desired_workers.update(desired_workers)
|
||||
|
||||
new_desired_workers.update(new_vm_names)
|
||||
logger.info(f"New desired state will be {new_desired_workers}")
|
||||
if len(new_desired_workers) > self.max_worker_nodes:
|
||||
logger.warning(
|
||||
"Autoscaler attempted to create more than max_workers VMs."
|
||||
)
|
||||
return created_nodes_dict
|
||||
|
||||
payload = {
|
||||
"spec": {"autoscaler_desired_workers": new_desired_workers}
|
||||
}
|
||||
|
||||
custom_api = self.k8s_api_client.custom_object_api
|
||||
custom_api.patch_namespaced_custom_object(
|
||||
VMRAY_GROUP,
|
||||
VMRAY_CRD_VER,
|
||||
self.namespace,
|
||||
VMRAYCLUSTER_PLURAL,
|
||||
self.cluster_name,
|
||||
payload,
|
||||
async_req=False,
|
||||
)
|
||||
for vm in new_vm_names:
|
||||
created_nodes_dict[vm] = vm
|
||||
return created_nodes_dict
|
||||
|
||||
def _get_cluster_response(self):
|
||||
response = {}
|
||||
try:
|
||||
response = (
|
||||
self.k8s_api_client.custom_object_api.get_namespaced_custom_object(
|
||||
VMRAY_GROUP,
|
||||
VMRAY_CRD_VER,
|
||||
self.namespace,
|
||||
VMRAYCLUSTER_PLURAL,
|
||||
self.cluster_name,
|
||||
)
|
||||
)
|
||||
return response
|
||||
except client.exceptions.ApiException as e:
|
||||
# If HTTP 404 received means the cluster is not yet created.
|
||||
logger.warning(
|
||||
f"Exception while getting {self.cluster_name}. Exception: {str(e)}"
|
||||
)
|
||||
if e.status == 404:
|
||||
logger.warning(f"{self.cluster_name} not available. Creating new one.")
|
||||
return response
|
||||
|
||||
def _get_node(self, node_id: str) -> Any:
|
||||
vmray_cluster_response = self._get_cluster_response()
|
||||
vmray_cluster_status = vmray_cluster_response.get("status", {})
|
||||
if not vmray_cluster_status:
|
||||
return {}
|
||||
head_node_status = vmray_cluster_status.get("head_node_status", {})
|
||||
current_workers = vmray_cluster_status.get("current_workers", {})
|
||||
# head node is found
|
||||
if head_node_status and node_id == self._get_head_name():
|
||||
return head_node_status
|
||||
# worker nodes found
|
||||
for worker in current_workers.keys():
|
||||
if worker == node_id:
|
||||
return current_workers.get(worker)
|
||||
# If worker not found in the current worker then it might be getting created
|
||||
# and not yet ready. So check if it is in the desired workers list.
|
||||
vmray_cluster_spec = vmray_cluster_response.get("spec", {})
|
||||
desired_workers = vmray_cluster_spec.get("autoscaler_desired_workers", {})
|
||||
for worker in desired_workers.keys():
|
||||
if worker == node_id:
|
||||
# set vm_status as VM in the desired workers' list will not
|
||||
# have vm_status field.
|
||||
node = {"vm_status": VMNodeStatus.INITIALIZED.value}
|
||||
return node
|
||||
logger.info(f"VM {node_id} not found")
|
||||
return {}
|
||||
|
||||
def safe_to_scale(self):
|
||||
"""
|
||||
It is safe to scale as long as total number of workers(desired + current)
|
||||
do not exceeds cluster level max_workers.
|
||||
This function should handle cases:
|
||||
1. If there are workers in the desired_workers list but not in the
|
||||
current_workers list that means few workers are not yet up and running.
|
||||
2. If there are workers in the current_workers list but not in a
|
||||
desired_workers list indicates workers are not yet deleted completely
|
||||
and we should wait.
|
||||
3. If workers are present in both the list shows stable state for the cluster.
|
||||
"""
|
||||
vmray_cluster_response = self._get_cluster_response()
|
||||
vmray_cluster_status = vmray_cluster_response.get("status", {})
|
||||
if not vmray_cluster_status:
|
||||
return False
|
||||
current_workers = vmray_cluster_status.get("current_workers", {})
|
||||
vmray_cluster_spec = vmray_cluster_response.get("spec", {})
|
||||
desired_workers = vmray_cluster_spec.get("autoscaler_desired_workers", {})
|
||||
logger.info(
|
||||
f"Checking is it safe to scale:\n"
|
||||
f"Current workers: {current_workers.keys()} \n"
|
||||
f"Autoscaler desired workers: {desired_workers}"
|
||||
)
|
||||
|
||||
# Do not scale until reaches desired state
|
||||
if len(desired_workers) != len(current_workers):
|
||||
return False
|
||||
|
||||
# Wait until all nodes are in a Running state
|
||||
for worker in current_workers.values():
|
||||
if worker.get("vm_status", None) != VMNodeStatus.RUNNING.value:
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
def _create_node_name(self, node_name_tag):
|
||||
"""Create name for a Ray node"""
|
||||
# The nodes are named as follows:
|
||||
# <cluster-name>-h-<random alphanumeric string> for the head node
|
||||
# <cluster-name>-w-<<random alphanumeric string>> for the worker nodes
|
||||
random_str = "".join(
|
||||
random.choice(string.ascii_lowercase + string.digits) for _ in range(8)
|
||||
)
|
||||
if "head" in node_name_tag:
|
||||
self.vmraycluster_nounce = random_str
|
||||
return f"{self.cluster_name}-h-" + self.vmraycluster_nounce
|
||||
return f"{self.cluster_name}-w-" + random_str
|
||||
|
||||
def _get_head_name(self):
|
||||
if not self.vmraycluster_nounce:
|
||||
vmray_cluster_response = self._get_cluster_response()
|
||||
self.vmraycluster_nounce = vmray_cluster_response["metadata"]["labels"][
|
||||
"vmray.kubernetes.io/head-nounce"
|
||||
]
|
||||
return f"{self.cluster_name}-h-" + self.vmraycluster_nounce
|
||||
|
||||
def _create_vmraycluster(self):
|
||||
|
||||
# Define vmraycluster config structure.
|
||||
ray_cluster_config = {}
|
||||
ray_cluster_config["apiVersion"] = VMRAY_GROUP + "/" + VMRAY_CRD_VER
|
||||
ray_cluster_config["kind"] = "VMRayCluster"
|
||||
ray_cluster_config["metadata"] = {}
|
||||
ray_cluster_config["spec"] = {}
|
||||
|
||||
# Start reading values from local bootstrap config.
|
||||
ray_cluster_config["metadata"]["name"] = self.cluster_name
|
||||
ray_cluster_config["metadata"]["labels"] = {
|
||||
"vmray.kubernetes.io/head-nounce": self.vmraycluster_nounce,
|
||||
"vmray.io/created-by": "ray-cli",
|
||||
}
|
||||
ray_cluster_config["metadata"]["namespace"] = self.namespace
|
||||
|
||||
ray_cluster_config["spec"]["api_server"] = {}
|
||||
ray_cluster_config["spec"]["api_server"]["location"] = self.vsphere_config.get(
|
||||
"api_server"
|
||||
)
|
||||
ray_cluster_config["spec"]["ray_docker_image"] = self.docker["image"]
|
||||
|
||||
# Set head node specific config.
|
||||
ray_cluster_config["spec"]["head_node"] = {}
|
||||
ray_cluster_config["spec"]["head_node"][
|
||||
"head_setup_commands"
|
||||
] = self.head_setup_commands
|
||||
ray_cluster_config["spec"]["head_node"][
|
||||
"port"
|
||||
] = 6379 # using default GCS port for now.
|
||||
|
||||
ray_cluster_config["spec"]["head_node"]["node_type"] = self.head_node_type
|
||||
|
||||
# Set common node config & available node types.
|
||||
ray_cluster_config["spec"]["common_node_config"] = {}
|
||||
ray_cluster_config["spec"]["common_node_config"][
|
||||
"vm_image"
|
||||
] = self.vsphere_config.get("vm_image")
|
||||
ray_cluster_config["spec"]["common_node_config"][
|
||||
"storage_class"
|
||||
] = self.vsphere_config.get("storage_class")
|
||||
ray_cluster_config["spec"]["common_node_config"][
|
||||
"vm_password_salt_hash"
|
||||
] = self.vsphere_config.get("vm_password_salt_hash", "")
|
||||
ray_cluster_config["spec"]["common_node_config"][
|
||||
"max_workers"
|
||||
] = self.max_worker_nodes
|
||||
available_node_types = {}
|
||||
for node_type, node_config in self.available_node_types.items():
|
||||
available_node_types[node_type] = {}
|
||||
available_node_types[node_type]["vm_class"] = node_config[
|
||||
"node_config"
|
||||
].get("vm_class", None)
|
||||
available_node_types[node_type]["resources"] = node_config.get(
|
||||
"resources", {}
|
||||
)
|
||||
available_node_types[node_type]["min_workers"] = node_config.get(
|
||||
"min_workers", 0
|
||||
)
|
||||
available_node_types[node_type]["max_workers"] = node_config.get(
|
||||
"max_workers", 2
|
||||
)
|
||||
|
||||
ray_cluster_config["spec"]["common_node_config"][
|
||||
"available_node_types"
|
||||
] = available_node_types
|
||||
ray_cluster_config["spec"]["common_node_config"][
|
||||
"vm_user"
|
||||
] = self.provider_auth["ssh_user"]
|
||||
if self.docker_config is not None:
|
||||
ray_cluster_config["spec"]["docker_config"] = self.docker_config
|
||||
|
||||
logger.info(f"Creating VmRayCluster \n{ray_cluster_config}")
|
||||
self.k8s_api_client.custom_object_api.create_namespaced_custom_object(
|
||||
VMRAY_GROUP,
|
||||
VMRAY_CRD_VER,
|
||||
self.namespace,
|
||||
VMRAYCLUSTER_PLURAL,
|
||||
ray_cluster_config,
|
||||
)
|
||||
|
||||
def _set_tags(self, node_id, node_kind, node_user_type, node_status, tags):
|
||||
new_tags = tags.copy()
|
||||
if node_status == VMNodeStatus.RUNNING.value:
|
||||
new_tags[TAG_RAY_NODE_STATUS] = STATUS_UP_TO_DATE
|
||||
elif node_status == VMNodeStatus.INITIALIZED.value:
|
||||
new_tags[TAG_RAY_NODE_STATUS] = STATUS_SETTING_UP
|
||||
else:
|
||||
new_tags[TAG_RAY_NODE_STATUS] = STATUS_UNINITIALIZED
|
||||
|
||||
new_tags[TAG_RAY_NODE_NAME] = node_id
|
||||
new_tags[TAG_RAY_NODE_KIND] = node_kind
|
||||
new_tags[TAG_RAY_USER_NODE_TYPE] = node_user_type
|
||||
return new_tags
|
||||
|
||||
def _set_max_worker_nodes(self):
|
||||
vmray_cluster_response = self._get_cluster_response()
|
||||
if not self.max_worker_nodes:
|
||||
vmray_cluster_spec = vmray_cluster_response.get("spec", {})
|
||||
common_node_config = vmray_cluster_spec.get("common_node_config", {})
|
||||
# If max_workers is not provided then default to 2
|
||||
# ref: https://docs.ray.io/en/latest/cluster/vms/references/
|
||||
# ray-cluster-configuration.html#max-workers
|
||||
self.max_worker_nodes = common_node_config.get("max_workers", 2)
|
||||
logger.info(f"Max worker is set to {self.max_worker_nodes}")
|
||||
|
||||
def _get_vm_service_ingress(self):
|
||||
response = self._get_vm_service()
|
||||
status = response.get("status", {})
|
||||
if not status:
|
||||
return []
|
||||
ingress = status["loadBalancer"].get("ingress", [])
|
||||
logger.info(f"VM service ingress is {ingress}")
|
||||
return ingress
|
||||
|
||||
def _get_vm_service(self):
|
||||
try:
|
||||
return self.k8s_api_client.custom_object_api.get_namespaced_custom_object(
|
||||
VMSERVICE_GROUP,
|
||||
VMSERVICE_CRD_VER,
|
||||
self.namespace,
|
||||
VMSERVICE_PLURAL,
|
||||
self._get_head_name(),
|
||||
)
|
||||
except client.exceptions.ApiException as e:
|
||||
logger.warning(
|
||||
f"Exception while getting vm service in namespace {self.namespace}."
|
||||
f"Exception: {str(e)}"
|
||||
)
|
||||
return {}
|
||||
|
||||
def _create_secret(self, namespace, name, kp):
|
||||
data = {}
|
||||
for k, v in kp.items():
|
||||
# Base64 encode the SSH key
|
||||
val = v
|
||||
if type(v) is str:
|
||||
val = v.encode("utf-8")
|
||||
data[k] = base64.b64encode(val).decode("utf-8")
|
||||
|
||||
# Create metadata
|
||||
metadata = client.V1ObjectMeta(name=name)
|
||||
|
||||
# Secret object
|
||||
secret = client.V1Secret(
|
||||
api_version="v1", kind="Secret", metadata=metadata, data=data, type="Opaque"
|
||||
)
|
||||
instance = client.CoreV1Api(self.k8s_api_client.client)
|
||||
|
||||
# Create the secret in the specified namespace.
|
||||
try:
|
||||
instance.create_namespaced_secret(namespace=namespace, body=secret)
|
||||
logger.info(f"Secret {name} created in namespace {namespace}")
|
||||
except client.exceptions.ApiException as e:
|
||||
print("Failure while creating Secret [`%s`] : %s\n" % (name, e))
|
||||
|
||||
def _create_ssh_secret(self):
|
||||
"""Create a K8s SSH secret using a local private SSH key
|
||||
specified in the config."""
|
||||
|
||||
# Define secret name
|
||||
secret_name = f"{self.cluster_name}-ssh-key"
|
||||
|
||||
pvt_key = self.provider_auth.get("ssh_private_key")
|
||||
private_key_path = os.path.expanduser(pvt_key)
|
||||
|
||||
# Read the SSH private key
|
||||
with open(private_key_path, "rb") as ssh_key:
|
||||
ssh_key_data = ssh_key.read()
|
||||
|
||||
kp = {"ssh-pvt-key": ssh_key_data}
|
||||
self._create_secret(self.namespace, secret_name, kp)
|
||||
|
||||
|
||||
def _is_ipv4(ip):
|
||||
try:
|
||||
ipaddress.IPv4Address(ip)
|
||||
return True
|
||||
except ipaddress.AddressValueError:
|
||||
return False
|
||||
@@ -0,0 +1,148 @@
|
||||
import copy
|
||||
import logging
|
||||
import os
|
||||
|
||||
from cryptography.hazmat.primitives import serialization as crypto_serialization
|
||||
from cryptography.hazmat.primitives.asymmetric import rsa
|
||||
|
||||
from ray.autoscaler._private.constants import DISABLE_NODE_UPDATERS_KEY
|
||||
from ray.autoscaler._private.event_system import CreateClusterEvent, global_event_system
|
||||
from ray.autoscaler._private.util import check_legacy_fields
|
||||
|
||||
PRIVATE_KEY_NAME = "ray-bootstrap-key.pem"
|
||||
PUBLIC_KEY_NAME = "ray_bootstrap_public_key.key"
|
||||
|
||||
PRIVATE_KEY_PATH = os.path.expanduser(f"~/{PRIVATE_KEY_NAME}")
|
||||
PUBLIC_KEY_PATH = os.path.expanduser(f"~/{PUBLIC_KEY_NAME}")
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def bootstrap_vsphere(config):
|
||||
|
||||
# create a copy of the input config to modify
|
||||
config = copy.deepcopy(config)
|
||||
|
||||
# Log warnings if user included deprecated `head_node` or `worker_nodes`
|
||||
# fields. Raise error if no `available_node_types`
|
||||
check_legacy_fields(config)
|
||||
|
||||
# Configure SSH access, using an existing key pair if possible.
|
||||
config = configure_key_pair(config)
|
||||
# Configure docker run command to be executed on head and wroker nodes
|
||||
config = configure_run_options(config)
|
||||
|
||||
global_event_system.execute_callback(
|
||||
CreateClusterEvent.ssh_keypair_downloaded,
|
||||
{"ssh_key_path": config["auth"]["ssh_private_key"]},
|
||||
)
|
||||
logger.info(f"{config}")
|
||||
return config
|
||||
|
||||
|
||||
def configure_key_pair(config):
|
||||
logger.info("Configuring keys for Ray Cluster Launcher to ssh into the head node.")
|
||||
|
||||
if not os.path.exists(PRIVATE_KEY_PATH):
|
||||
logger.warning(
|
||||
"Private key file at path {} was not found".format(PRIVATE_KEY_PATH)
|
||||
)
|
||||
_create_ssh_keys()
|
||||
logger.info(
|
||||
f"New SSH key pair {PRIVATE_KEY_PATH} and {PUBLIC_KEY_PATH} created."
|
||||
)
|
||||
|
||||
# updater.py file uses the following config to ssh onto the head node
|
||||
# Also, copies the file onto the head node
|
||||
config["auth"]["ssh_private_key"] = PRIVATE_KEY_PATH
|
||||
# The path where the public key should be copied onto the remote host
|
||||
public_key_remote_path = f"~/{PUBLIC_KEY_NAME}"
|
||||
|
||||
# Copy the public key to the remote host
|
||||
config["file_mounts"][public_key_remote_path] = PUBLIC_KEY_PATH
|
||||
|
||||
return config
|
||||
|
||||
|
||||
def configure_run_options(config):
|
||||
ssh_user = config["auth"]["ssh_user"]
|
||||
# By default enable TLS for Head-Worker grpc communication
|
||||
tls_enable = (
|
||||
1 if config["provider"]["vsphere_config"].get("tls_enable", True) else 0
|
||||
)
|
||||
# Configure common run options
|
||||
if "run_options" not in config["docker"]:
|
||||
config["docker"]["run_options"] = []
|
||||
config["docker"]["run_options"].append(f"--env RAY_USE_TLS={tls_enable}")
|
||||
|
||||
# Configure head_run_options
|
||||
if "head_run_options" not in config["docker"]:
|
||||
config["docker"]["head_run_options"] = []
|
||||
config["docker"]["head_run_options"].append(
|
||||
f"--env-file /home/{ssh_user}/svc-account-token.env"
|
||||
)
|
||||
|
||||
# Configure worker_run_options
|
||||
if "worker_run_options" not in config["docker"]:
|
||||
config["docker"]["worker_run_options"] = []
|
||||
|
||||
if tls_enable == 1:
|
||||
# Generate TLS cert and key for head and worker nodes.
|
||||
# This needs to be done before ray start command
|
||||
config["head_start_ray_commands"].insert(0, "sh /home/ray/gencert.sh")
|
||||
config["worker_start_ray_commands"].insert(0, "sh /home/ray/gencert.sh")
|
||||
|
||||
config["docker"]["run_options"].append(
|
||||
f"-v /home/{ssh_user}/ca.crt:/home/ray/ca.crt"
|
||||
)
|
||||
config["docker"]["run_options"].append(
|
||||
f"-v /home/{ssh_user}/ca.key:/home/ray/ca.key"
|
||||
)
|
||||
config["docker"]["run_options"].append(
|
||||
f"-v /home/{ssh_user}/gencert.sh:/home/ray/gencert.sh"
|
||||
)
|
||||
|
||||
config["docker"]["run_options"].append("--env RAY_TLS_CA_CERT=/home/ray/ca.crt")
|
||||
config["docker"]["run_options"].append(
|
||||
"--env RAY_TLS_SERVER_KEY=/home/ray/tls.key"
|
||||
)
|
||||
config["docker"]["run_options"].append(
|
||||
"--env RAY_TLS_SERVER_CERT=/home/ray/tls.crt"
|
||||
)
|
||||
|
||||
return config
|
||||
|
||||
|
||||
def disable_node_updater(config):
|
||||
logger.info(
|
||||
"Disabling NodeUpdater threads as Cluster Operator is "
|
||||
+ "responsible for Ray setup on nodes."
|
||||
)
|
||||
config["provider"][DISABLE_NODE_UPDATERS_KEY] = True
|
||||
return config
|
||||
|
||||
|
||||
def _create_ssh_keys():
|
||||
"""Create SSH keys as specified"""
|
||||
# Create a private key
|
||||
private_key = rsa.generate_private_key(public_exponent=65537, key_size=4096)
|
||||
# Encode it in PEM format
|
||||
unencrypted_private_key = private_key.private_bytes(
|
||||
encoding=crypto_serialization.Encoding.PEM,
|
||||
format=crypto_serialization.PrivateFormat.TraditionalOpenSSL,
|
||||
encryption_algorithm=crypto_serialization.NoEncryption(),
|
||||
)
|
||||
# Create a public key
|
||||
public_key = private_key.public_key().public_bytes(
|
||||
encoding=crypto_serialization.Encoding.PEM,
|
||||
format=crypto_serialization.PublicFormat.SubjectPublicKeyInfo,
|
||||
)
|
||||
# Write keys
|
||||
with open(PRIVATE_KEY_PATH, "wb") as pvt_key:
|
||||
# manage access mode for the pvt key
|
||||
os.chmod(PRIVATE_KEY_PATH, 0o600)
|
||||
pvt_key.write(unencrypted_private_key)
|
||||
|
||||
with open(PUBLIC_KEY_PATH, "wb") as pub_key:
|
||||
# manage access mode for the pvt key
|
||||
pub_key.write(public_key)
|
||||
+125
@@ -0,0 +1,125 @@
|
||||
import logging
|
||||
import threading
|
||||
from typing import Any, Dict
|
||||
|
||||
from ray.autoscaler._private.vsphere.cluster_operator_client import (
|
||||
ClusterOperatorClient,
|
||||
)
|
||||
from ray.autoscaler._private.vsphere.config import bootstrap_vsphere
|
||||
from ray.autoscaler.node_provider import NodeProvider
|
||||
from ray.autoscaler.tags import (
|
||||
STATUS_SETTING_UP,
|
||||
TAG_RAY_CLUSTER_NAME,
|
||||
TAG_RAY_NODE_NAME,
|
||||
TAG_RAY_NODE_STATUS,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class VsphereWcpNodeProvider(NodeProvider):
|
||||
max_terminate_nodes = 1000
|
||||
cluster_config = None
|
||||
|
||||
def __init__(self, provider_config, cluster_name):
|
||||
NodeProvider.__init__(self, provider_config, cluster_name)
|
||||
self.tag_cache = {}
|
||||
self.tag_cache_lock = threading.Lock()
|
||||
self.client = ClusterOperatorClient(
|
||||
cluster_name, provider_config, VsphereWcpNodeProvider.cluster_config
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def bootstrap_config(cluster_config):
|
||||
config = bootstrap_vsphere(cluster_config)
|
||||
VsphereWcpNodeProvider.cluster_config = config
|
||||
return config
|
||||
|
||||
def non_terminated_nodes(self, tag_filters):
|
||||
nodes, tag_cache = self.client.list_vms(tag_filters)
|
||||
with self.tag_cache_lock:
|
||||
for node_id in nodes:
|
||||
for k, v in tag_cache[node_id].items():
|
||||
if node_id in self.tag_cache.keys():
|
||||
self.tag_cache[node_id][k] = v
|
||||
else:
|
||||
self.tag_cache[node_id] = {}
|
||||
self.tag_cache[node_id][k] = v
|
||||
logger.info(f"Non terminated nodes' tags are {self.tag_cache}")
|
||||
return nodes
|
||||
|
||||
def is_running(self, node_id):
|
||||
return self.client.is_vm_power_on(node_id)
|
||||
|
||||
def is_terminated(self, node_id):
|
||||
if self.client.is_vm_power_on(node_id):
|
||||
return False
|
||||
else:
|
||||
# If the node is not powered on but has the creating tag, then it could
|
||||
# be under reconfiguration, such as plugging the GPU. In this case we
|
||||
# should consider the node is not terminated, it will be turned on later
|
||||
return not self.client.is_vm_creating(node_id)
|
||||
|
||||
def node_tags(self, node_id):
|
||||
with self.tag_cache_lock:
|
||||
return self.tag_cache[node_id]
|
||||
|
||||
def external_ip(self, node_id):
|
||||
return self.client.get_vm_external_ip(node_id)
|
||||
|
||||
def internal_ip(self, node_id):
|
||||
# Currently vSphere VMs do not show an internal IP. So we just return the
|
||||
# external IP
|
||||
return self.client.get_vm_external_ip(node_id)
|
||||
|
||||
def set_node_tags(self, node_id, tags):
|
||||
# This method gets called from the Ray and it passes
|
||||
# node_id. It updates old tags (if present) with new values.
|
||||
with self.tag_cache_lock:
|
||||
for k, v in tags.items():
|
||||
# update tags for node_id
|
||||
self.tag_cache[node_id][k] = v
|
||||
logger.info(f"Updated tags for {node_id} to: {self.tag_cache[node_id]}")
|
||||
|
||||
def create_node(self, node_config, tags, count) -> Dict[str, Any]:
|
||||
"""Creates instances.
|
||||
|
||||
Returns dict mapping instance id to VM object for the created
|
||||
instances.
|
||||
"""
|
||||
to_be_launched_node_count = count
|
||||
created_nodes_dict = {}
|
||||
if to_be_launched_node_count > 0:
|
||||
created_nodes_dict = self.client.create_nodes(
|
||||
tags, to_be_launched_node_count, node_config
|
||||
)
|
||||
# make sure to mark newly created nodes as ready
|
||||
# so autoscaler shouldn't provision new ones
|
||||
with self.tag_cache_lock:
|
||||
for node_id in created_nodes_dict.keys():
|
||||
self.tag_cache[node_id] = tags.copy()
|
||||
self.tag_cache[node_id][TAG_RAY_NODE_STATUS] = STATUS_SETTING_UP
|
||||
self.tag_cache[node_id][TAG_RAY_NODE_NAME] = node_id
|
||||
self.tag_cache[node_id][TAG_RAY_CLUSTER_NAME] = self.cluster_name
|
||||
logger.info(
|
||||
f"Node {node_id} created with tags: {self.tag_cache[node_id]}"
|
||||
)
|
||||
return created_nodes_dict
|
||||
|
||||
def terminate_node(self, node_id):
|
||||
if not node_id or self.client.is_vm_creating(node_id):
|
||||
return
|
||||
# Delete node iff it is either in a running or a failure state
|
||||
self.client.delete_node(node_id)
|
||||
with self.tag_cache_lock:
|
||||
if node_id in self.tag_cache:
|
||||
self.tag_cache.pop(node_id)
|
||||
|
||||
def terminate_nodes(self, node_ids):
|
||||
if not node_ids:
|
||||
return
|
||||
for node_id in node_ids:
|
||||
self.terminate_node(node_id)
|
||||
|
||||
def safe_to_scale(self) -> bool:
|
||||
return self.client.safe_to_scale()
|
||||
Reference in New Issue
Block a user