wehub-resource-sync
587 lines
22 KiB
Python
587 lines
22 KiB
Python
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
|
|
)
|