diff --git a/swmcloudgate/routers/azure/connector.py b/swmcloudgate/routers/azure/connector.py index 0d0c085..cf618b5 100644 --- a/swmcloudgate/routers/azure/connector.py +++ b/swmcloudgate/routers/azure/connector.py @@ -1,7 +1,9 @@ import os import json +import copy import typing import logging +import shlex import jinja2 from azure.identity import CertificateCredential @@ -17,9 +19,16 @@ from ..baseconnector import BaseConnector LOG = logging.getLogger("swm") -TEMLPATE_FILE = "swmcloudgate/routers/azure/templates/partition.json" +TEMPLATE_FILE = "swmcloudgate/routers/azure/templates/partition.json" CLOUD_INIT_SCRIPT_FILE = "swmcloudgate/routers/azure/templates/cloud-init.sh" CLOUD_INIT_YAML = "swmcloudgate/routers/azure/templates/cloud-init.yaml" +MAX_VM_COUNT = 32 +AZURE_NETWORK_API_VERSION = "2021-05-01" +HOST_NAME_PLACEHOLDER = "__SWM_HOST_NAME__" +IS_MAIN_PLACEHOLDER = "__SWM_IS_MAIN__" +MAIN_INSTANCE_HOSTNAME_PLACEHOLDER = "__SWM_MAIN_INSTANCE_HOSTNAME__" +MAIN_INSTANCE_PRIVATE_IP_PLACEHOLDER = "__SWM_MAIN_INSTANCE_PRIVATE_IP__" +STORAGE_KEY_B64_PLACEHOLDER = "__SWM_STORAGE_KEY_B64__" class AzureConnector(BaseConnector): @@ -67,6 +76,21 @@ def _init_azure_clients( ) raise Exception(msg) + def _parse_vm_count(self, count_param: str | None) -> int: + """Parse and validate VM count parameter.""" + if count_param is None or count_param == "": + return 1 + try: + vm_count = int(count_param) + except (ValueError, TypeError) as exc: + raise ValueError(f"Invalid VM count '{count_param}': must be an integer") from exc + + if vm_count < 1: + raise ValueError(f"Invalid VM count {vm_count}: must be greater than or equal to 1") + if vm_count > MAX_VM_COUNT: + raise ValueError(f"Invalid VM count {vm_count}: maximum supported count is {MAX_VM_COUNT}") + return vm_count + def _get_deployment_properties( self, job_id: str, @@ -75,10 +99,12 @@ def _get_deployment_properties( os_version: str, username: str, user_pub_key: str, + storage_key: str, cloud_init_script: str, ports: str, + vm_count: int, ) -> dict[str, dict[str, typing.Any]]: - with open(TEMLPATE_FILE) as template_file: + with open(TEMPLATE_FILE) as template_file: template = json.load(template_file) template_parameters = self._get_template_parameters( job_id, @@ -87,10 +113,17 @@ def _get_deployment_properties( os_version, username, user_pub_key, + storage_key, cloud_init_script, ) - LOG.debug(f"Template parameters for job {job_id}: {template_parameters}") + sanitized_parameters = copy.deepcopy(template_parameters) + for secret_key in ("adminPasswordOrKey", "storageKey"): + if secret_key in sanitized_parameters: + sanitized_parameters[secret_key]["value"] = "***" + LOG.debug(f"Template parameters for job {job_id}: {sanitized_parameters}") self._append_security_rules(ports, template) + self._configure_main_vm_custom_data(template) + self._add_compute_vms(partition_name, vm_count, template) return { "properties": { "template": template, @@ -99,6 +132,175 @@ def _get_deployment_properties( } } + def _find_resource( + self, resources: list[dict[str, typing.Any]], resource_type: str + ) -> dict[str, typing.Any] | None: + for resource in resources: + if resource.get("type") == resource_type: + return resource + return None + + def _find_vm_extension_resources(self, resources: list[dict[str, typing.Any]]) -> list[dict[str, typing.Any]]: + return [ + resource for resource in resources if resource.get("type") == "Microsoft.Compute/virtualMachines/extensions" + ] + + def _get_main_vm_private_ip_expression(self) -> str: + return ( + "reference(resourceId('Microsoft.Network/networkInterfaces', variables('networkInterfaceName')), " + f"'{AZURE_NETWORK_API_VERSION}').ipConfigurations[0].properties.privateIPAddress" + ) + + def _get_main_vm_name(self, partition_name: str) -> str: + return f"{partition_name}-main" + + def _build_vm_custom_data( + self, + vm_name_expression: str, + is_main: bool, + ) -> str: + custom_data_expression = "parameters('cloudInitScript')" + replacements = ( + (HOST_NAME_PLACEHOLDER, vm_name_expression), + (IS_MAIN_PLACEHOLDER, "'true'" if is_main else "'false'"), + (MAIN_INSTANCE_HOSTNAME_PLACEHOLDER, "parameters('vmNameMain')"), + (MAIN_INSTANCE_PRIVATE_IP_PLACEHOLDER, self._get_main_vm_private_ip_expression()), + (STORAGE_KEY_B64_PLACEHOLDER, "base64(parameters('storageKey'))"), + ) + for placeholder, replacement in replacements: + custom_data_expression = f"replace({custom_data_expression}, '{placeholder}', {replacement})" + return f"[base64({custom_data_expression})]" + + def _configure_main_vm_custom_data(self, template: dict[str, typing.Any]) -> None: + resources = template.get("resources", []) + main_vm_resource = self._find_resource(resources, "Microsoft.Compute/virtualMachines") + if not main_vm_resource: + LOG.warning("Could not find main VM resource in template") + return + main_vm_resource["properties"]["osProfile"]["customData"] = self._build_vm_custom_data( + "parameters('vmNameMain')", + is_main=True, + ) + + def _append_dependency(self, resource: dict[str, typing.Any], dependency: str) -> None: + depends_on = resource.setdefault("dependsOn", []) + if dependency not in depends_on: + depends_on.append(dependency) + + def _remove_dependency(self, resource: dict[str, typing.Any], dependency: str) -> None: + depends_on = resource.get("dependsOn", []) + resource["dependsOn"] = [item for item in depends_on if item != dependency] + + def _replace_vm_reference(self, value: typing.Any, compute_vm_name: str) -> typing.Any: + if isinstance(value, str): + return value.replace("parameters('vmNameMain')", f"'{compute_vm_name}'") + if isinstance(value, list): + return [self._replace_vm_reference(item, compute_vm_name) for item in value] + if isinstance(value, dict): + return {key: self._replace_vm_reference(item, compute_vm_name) for key, item in value.items()} + return value + + def _strip_public_ip(self, nic_resource: dict[str, typing.Any]) -> None: + ip_configurations = nic_resource.get("properties", {}).get("ipConfigurations", []) + if not ip_configurations: + return + ip_configuration_properties = ip_configurations[0].get("properties", {}) + ip_configuration_properties.pop("publicIPAddress", None) + + def _clone_compute_nic( + self, + network_interface_resource: dict[str, typing.Any], + compute_nic_name: str, + ) -> dict[str, typing.Any]: + compute_nic_resource = copy.deepcopy(network_interface_resource) + compute_nic_resource["name"] = f"[format('{compute_nic_name}')]" + self._strip_public_ip(compute_nic_resource) + public_ip_dependency = "[resourceId('Microsoft.Network/publicIPAddresses', variables('publicIPAddressName'))]" + self._remove_dependency(compute_nic_resource, public_ip_dependency) + return compute_nic_resource + + def _clone_compute_vm( + self, + main_vm_resource: dict[str, typing.Any], + compute_vm_name: str, + compute_nic_name: str, + ) -> dict[str, typing.Any]: + compute_vm_resource = copy.deepcopy(main_vm_resource) + compute_vm_resource["name"] = f"[format('{compute_vm_name}')]" + compute_vm_resource["properties"]["networkProfile"]["networkInterfaces"][0][ + "id" + ] = f"[resourceId('Microsoft.Network/networkInterfaces', '{compute_nic_name}')]" + compute_vm_resource["properties"]["osProfile"]["computerName"] = f"[format('{compute_vm_name}')]" + compute_vm_resource["properties"]["osProfile"]["customData"] = self._build_vm_custom_data( + f"'{compute_vm_name}'", + is_main=False, + ) + compute_nic_dependency = f"[resourceId('Microsoft.Network/networkInterfaces', '{compute_nic_name}')]" + main_nic_dependency = "[resourceId('Microsoft.Network/networkInterfaces', variables('networkInterfaceName'))]" + self._append_dependency(compute_vm_resource, compute_nic_dependency) + # Compute nodes keep the inherited main NIC dependency because customData references its private IP. + self._append_dependency(compute_vm_resource, main_nic_dependency) + return compute_vm_resource + + def _clone_compute_vm_extension( + self, + extension_resource: dict[str, typing.Any], + compute_vm_name: str, + ) -> dict[str, typing.Any]: + compute_extension_resource = copy.deepcopy(extension_resource) + compute_extension_resource = self._replace_vm_reference(compute_extension_resource, compute_vm_name) + extension_name = extension_resource.get("name", "") + if "variables('extensionName')" in extension_name: + compute_extension_resource["name"] = f"[format('{compute_vm_name}/{{0}}', variables('extensionName'))]" + compute_vm_dependency = f"[resourceId('Microsoft.Compute/virtualMachines', '{compute_vm_name}')]" + original_vm_dependency = "[resourceId('Microsoft.Compute/virtualMachines', parameters('vmNameMain'))]" + self._append_dependency(compute_extension_resource, compute_vm_dependency) + self._remove_dependency(compute_extension_resource, original_vm_dependency) + return compute_extension_resource + + def _add_compute_vms( + self, + partition_name: str, + vm_count: int, + template: dict[str, typing.Any], + ) -> None: + """Add compute VMs and related resources to the template if vm_count > 1.""" + if vm_count <= 1: + return + + num_compute_vms = vm_count - 1 + LOG.debug(f"Adding {num_compute_vms} compute VMs to partition {partition_name}") + + resources = template.get("resources", []) + main_vm_resource = self._find_resource(resources, "Microsoft.Compute/virtualMachines") + if not main_vm_resource: + LOG.warning("Could not find main VM resource in template") + return + + network_interface_resource = self._find_resource(resources, "Microsoft.Network/networkInterfaces") + if not network_interface_resource: + LOG.warning("Could not find network interface resource in template") + return + + extension_resources = self._find_vm_extension_resources(resources) + + for i in range(num_compute_vms): + compute_index = i + 1 + compute_vm_name = f"{partition_name}-compute{compute_index}" + compute_nic_name = f"{partition_name}-compute{compute_index}-NetInt" + + compute_nic_resource = self._clone_compute_nic(network_interface_resource, compute_nic_name) + resources.append(compute_nic_resource) + + compute_vm_resource = self._clone_compute_vm(main_vm_resource, compute_vm_name, compute_nic_name) + resources.append(compute_vm_resource) + + for extension_resource in extension_resources: + compute_extension_resource = self._clone_compute_vm_extension(extension_resource, compute_vm_name) + resources.append(compute_extension_resource) + + LOG.debug(f"Added compute VM {compute_vm_name} to template") + def _append_security_rules(self, ports: str, template: dict[str, typing.Any]) -> None: for resource in template["resources"]: if resource["type"] == "Microsoft.Network/networkSecurityGroups": @@ -132,6 +334,7 @@ def _get_template_parameters( os_version: str, username: str, user_pub_key: str, + storage_key: str, cloud_init_script: str, ) -> dict[str, str]: template_loader = jinja2.FileSystemLoader(searchpath="./") @@ -142,8 +345,10 @@ def _get_template_parameters( ) return { "resourcePrefix": {"value": partition_name}, + "vmNameMain": {"value": self._get_main_vm_name(partition_name)}, "adminUsername": {"value": username}, "adminPasswordOrKey": {"value": user_pub_key}, + "storageKey": {"value": storage_key}, "osVersion": {"value": os_version}, "vmSize": {"value": flavor_name}, "cloudInitScript": {"value": cloud_init_yaml}, @@ -326,13 +531,15 @@ def _get_cloud_init_script( container_registry_username: str, container_registry_password: str, storage_account: str, - storage_key: str, storage_container: str, - runtime_params: str, + runtime_params: dict[str, str], user_ssh_cert: str, ) -> str: template_loader = jinja2.FileSystemLoader(searchpath="./") - template_env = jinja2.Environment(loader=template_loader, autoescape=True) + template_env = jinja2.Environment( # nosec B701 - shellquote secures interpolated shell values here + loader=template_loader, autoescape=False + ) + template_env.filters["shellquote"] = shlex.quote template = template_env.get_template(CLOUD_INIT_SCRIPT_FILE) script: str = template.render( job_id=job_id, @@ -343,11 +550,24 @@ def _get_cloud_init_script( container_registry_username=container_registry_username, container_registry_password=container_registry_password, storage_account=storage_account, - storage_key=storage_key, + storage_key_b64=STORAGE_KEY_B64_PLACEHOLDER, storage_container=storage_container, + host_name=HOST_NAME_PLACEHOLDER, + is_main=IS_MAIN_PLACEHOLDER, + main_instance_hostname=MAIN_INSTANCE_HOSTNAME_PLACEHOLDER, + main_instance_private_ip=MAIN_INSTANCE_PRIVATE_IP_PLACEHOLDER, ) return script + def _rollback_resource_group(self, resource_group_name: str) -> None: + LOG.warning(f"Rolling back failed deployment by deleting resource group {resource_group_name}") + try: + delete_operation = self._resource_client.resource_groups.begin_delete(resource_group_name) + if delete_operation: + delete_operation.result() + except Exception as rollback_error: + LOG.error(f"Failed to roll back resource group {resource_group_name}: {rollback_error}") + def get_resource_group(self, resource_group_name: str) -> typing.Dict[str, typing.Any]: if "resource_groups" in self._test_responses: for it in self.list_resource_groups(): @@ -422,6 +642,8 @@ def create_deployment( user_ssh_cert: str, ) -> tuple[DeploymentExtended, str]: resource_group_name = self._get_resource_group_name(partition_name) + vm_count = self._parse_vm_count(count) + LOG.info(f"Creating deployment with {vm_count} VM(s)") if self._test_responses: new_part = { @@ -446,7 +668,6 @@ def create_deployment( container_registry_username, container_registry_password, storage_account, - storage_key, storage_container, runtime_params, user_ssh_cert, @@ -460,8 +681,10 @@ def create_deployment( os_version, username, user_ssh_cert, + storage_key, cloud_init_script, ports, + vm_count, ) try: deployment_async_operation = self._resource_client.deployments.begin_create_or_update( @@ -470,11 +693,23 @@ def create_deployment( deployment_properties, ) except HttpResponseError as e: + self._rollback_resource_group(resource_group_name) LOG.error(e) raise e + except Exception: + self._rollback_resource_group(resource_group_name) + raise LOG.info(f"Deploying resource group {resource_group_name}, deployment: {deployment_name}") - return deployment_async_operation.result(), resource_group_name + try: + return deployment_async_operation.result(), resource_group_name + except HttpResponseError as e: + self._rollback_resource_group(resource_group_name) + LOG.error(e) + raise e + except Exception: + self._rollback_resource_group(resource_group_name) + raise def delete_resource_group(self, resource_group_name: str) -> str | None: if "resource_groups" in self._test_responses: diff --git a/swmcloudgate/routers/azure/partitions.py b/swmcloudgate/routers/azure/partitions.py index 8295f4a..f6ea936 100644 --- a/swmcloudgate/routers/azure/partitions.py +++ b/swmcloudgate/routers/azure/partitions.py @@ -42,7 +42,7 @@ async def create_partition( LOG.debug(f" * containerimage: {containerimage}") LOG.debug(f" * flavor: {flavorname}") LOG.debug(f" * username: {username}") - LOG.debug(f" * extra nodes: {count}") + LOG.debug(f" * vm count: {count}") LOG.debug(f" * runtime: {runtime}") LOG.debug(f" * location: {location}") LOG.debug(f" * ports: {ports}") @@ -127,6 +127,9 @@ async def create_partition( return {"error": f"Error from Azure: {'; '.join(error_messages)}", "partition": partition} + except ValueError as e: + raise HTTPException(status_code=http.HTTPStatus.BAD_REQUEST, detail=str(e)) from e + except Exception as e: LOG.error(traceback.format_exception(e)) return {"error": traceback.format_exception(e)} diff --git a/swmcloudgate/routers/azure/templates/cloud-init.sh b/swmcloudgate/routers/azure/templates/cloud-init.sh index 48d6c93..28ff39e 100755 --- a/swmcloudgate/routers/azure/templates/cloud-init.sh +++ b/swmcloudgate/routers/azure/templates/cloud-init.sh @@ -1,7 +1,35 @@ #!/bin/bash -ex -HOST_NAME=$(hostname) +HOST_NAME="{{ host_name }}" SWM_ROOT="/opt/swm" +PRIMARY_INTERFACE="" +PRIVATE_IP_CIDR="" +PRIVATE_SUBNET_CIDR="" +IS_MAIN={{ is_main }} +MAIN_INSTANCE_HOSTNAME="{{ main_instance_hostname }}" +MAIN_INSTANCE_PRIVATE_IP="{{ main_instance_private_ip }}" + +detect_vm_context() { + PRIMARY_INTERFACE=$(ip -4 route list 0/0 | awk 'NR==1 { print $5 }') + PRIVATE_IP_CIDR=$(ip -4 -o addr show "$PRIMARY_INTERFACE" | awk 'NR==1 { print $4 }') + PRIVATE_SUBNET_CIDR=$(python3 - "$PRIVATE_IP_CIDR" <<'PY' +import ipaddress +import sys + +network = ipaddress.ip_interface(sys.argv[1]).network +print(f"{network.network_address}/{network.netmask}") +PY +) +if [[ -z "$PRIVATE_SUBNET_CIDR" ]]; then + echo "$(date): could not determine private subnet CIDR" >&2 + return 1 +fi + +if [[ -z "$MAIN_INSTANCE_HOSTNAME" || -z "$MAIN_INSTANCE_PRIVATE_IP" ]]; then + echo "$(date): could not determine main instance details" >&2 + return 1 +fi +} mount_azure_storage() { echo Mount Azure storage @@ -13,9 +41,11 @@ mount_azure_storage() { apt-get install fuse3 blobfuse2 -y popd - local azure_storage_account={{ storage_account }} - local azure_storage_key={{ storage_key }} - local azure_storage_container={{ storage_container }} + local azure_storage_account={{ storage_account | shellquote }} + local azure_storage_key_b64="{{ storage_key_b64 }}" + local azure_storage_key + azure_storage_key=$(printf '%s' "$azure_storage_key_b64" | base64 -d) + local azure_storage_container={{ storage_container | shellquote }} local config_file=/etc/blobfuse2.yaml local cache_dir=/tmp/blobfuse2.cache @@ -62,18 +92,33 @@ EOF blobfuse2 mount $mount_dir --config-file=$config_file --read-only } +wait_for_main_nfs() { + local attempt=0 + local max_attempts=60 + + until timeout 2 bash -c "/dev/null 2>&1; do + (( attempt += 1 )) + if (( attempt >= max_attempts )); then + echo "$(date): NFS on ${MAIN_INSTANCE_PRIVATE_IP}:2049 did not become ready after ${max_attempts} attempts" >&2 + return 1 + fi + echo "$(date): waiting for NFS on ${MAIN_INSTANCE_PRIVATE_IP}:2049 (${attempt}/${max_attempts})" + sleep 5 + done +} + create_directories() { - if [[ "{{ swm_source }}" == "ssh" ]]; then + if [[ {{ swm_source | shellquote }} == "ssh" ]]; then echo $(date) ": create directory $SWM_ROOT" mkdir -p "$SWM_ROOT" fi } setup_swm_worker() { - echo $(date) ": ensure swm worker is installed, SWM_SOURCE={{ swm_source }}" + echo $(date) ": ensure swm worker is installed, SWM_SOURCE={{ swm_source | shellquote }}" - if [[ "{{ swm_source }}" == "ssh" ]]; then - echo "{{ ssh_pub_key }}" >> /root/.ssh/authorized_keys + if [[ {{ swm_source | shellquote }} == "ssh" ]]; then + echo {{ ssh_pub_key | shellquote }} >> /root/.ssh/authorized_keys echo $(date) ": ensure swm worker is installed via ssh" local check_interval=15 @@ -101,10 +146,10 @@ setup_swm_worker() { ${SWM_ROOT}/${SWM_VERSION}/scripts/setup-swm-core.py -v ${SWM_VERSION} -p ${SWM_ROOT} -c ${SWM_ROOT}/${SWM_VERSION}/priv/setup/setup.config - elif [[ "{{ swm_source }}" == "http://*.tar.gz" ]]; then + elif [[ {{ swm_source | shellquote }} == "http://*.tar.gz" ]]; then TMP_DIR=$(mktemp -d -t swm-worker-XXXXX) pushd $TMP_DIR - wget {{ swm_source }} --output-document=swm-worker.tar.gz + wget {{ swm_source | shellquote }} --output-document=swm-worker.tar.gz mkdir -p /opt/swm tar zfx ./swm-worker.tar.gz --directory /opt/swm/ popd @@ -130,10 +175,10 @@ setup_swm_worker() { } setup_network() { - GATEWAY_IP=$(ip -4 addr show $(ip -4 route list 0/0 | awk -F" " "{ print \$5 }") | grep -oP "(?<=inet\\s)\\d+(\\.\\d+){3}") - IS_MAIN=true - echo $(date) ": start VM initialization (HOST: $HOST_NAME, IP=$GATEWAY_IP, master: ${IS_MAIN})" - echo $GATEWAY_IP $HOST_NAME.openworkload.org $HOST_NAME >> /etc/hosts + detect_vm_context || exit 1 + VM_PRIVATE_IP="${PRIVATE_IP_CIDR%%/*}" + echo $(date) ": start VM initialization (HOST: $HOST_NAME, IP=$VM_PRIVATE_IP, master: ${IS_MAIN})" + echo $VM_PRIVATE_IP $HOST_NAME.openworkload.org $HOST_NAME >> /etc/hosts echo $(date) ": /etc/hosts:" cat /etc/hosts echo @@ -142,16 +187,16 @@ setup_network() { setup_mounts() { if [ $IS_MAIN == "true" ]; then - echo "/home $PRIVATE_SUBNET_CIDR(rw,async,no_root_squash)" | sed "s/\\/25/\\/255.255.255.0/g" >> /etc/exports + echo "/home $PRIVATE_SUBNET_CIDR(rw,async,no_root_squash)" >> /etc/exports echo $(date) ": /etc/exports:" cat /etc/exports echo - # Temporary disable for debug purposes - #systemctl enable nfs-kernel-server - #systemctl restart nfs-kernel-server - #echo $(date) ": systemctl | grep nfs:" - #systemctl | grep nfs + exportfs -ra + systemctl enable nfs-kernel-server + systemctl restart nfs-kernel-server + echo $(date) ": systemctl | grep nfs:" + systemctl | grep nfs else echo "$MAIN_INSTANCE_PRIVATE_IP:/home /home nfs rsize=32768,wsize=32768,hard,intr,async 0 0" >> /etc/fstab @@ -159,11 +204,21 @@ setup_mounts() { cat /etc/fstab echo + wait_for_main_nfs + + local count=0 + local max_mount_attempts=60 echo $(date) ": waiting for mount ..." - until mount -a || (( count++ >= 20 )); do sleep 5; done + until mount -a; do + (( count += 1 )) + if (( count >= max_mount_attempts )); then + echo "$(date): failed to mount shared /home after ${max_mount_attempts} attempts" >&2 + return 1 + fi + echo "$(date): mount attempt ${count}/${max_mount_attempts} failed, retrying in 5 seconds" + sleep 5 + done echo $(date) ": mounted." - - systemctl restart docker # fix rare "connection closed" issues fi echo @@ -189,13 +244,18 @@ setup_docker() { } pull_container_image() { - if [ "{{ container_registry_password }}" != "" ]; then - echo $(date) ": login to the registry: {{ container_registry }}" - docker login {{ container_registry }} --username {{ container_registry_username }} --password {{ container_registry_password }} + local container_registry={{ container_registry | shellquote }} + local container_registry_username={{ container_registry_username | shellquote }} + local container_registry_password={{ container_registry_password | shellquote }} + local container_image={{ container_image | shellquote }} + + if [ -n "$container_registry_password" ]; then + echo $(date) ": login to the registry: $container_registry" + docker login "$container_registry" --username "$container_registry_username" --password "$container_registry_password" fi - echo $(date) ": pull job container image from container registry: {{ container_image }}" - docker pull {{ container_image }} + echo $(date) ": pull job container image from container registry: $container_image" + docker pull "$container_image" echo $(date) ": all local docker images after the pulling:" docker images diff --git a/swmcloudgate/routers/azure/templates/cloud-init.yaml b/swmcloudgate/routers/azure/templates/cloud-init.yaml index e14cbfe..601b0b2 100644 --- a/swmcloudgate/routers/azure/templates/cloud-init.yaml +++ b/swmcloudgate/routers/azure/templates/cloud-init.yaml @@ -7,6 +7,8 @@ packages: - docker.io - cgroupfs-mount - net-tools + - nfs-common + - nfs-kernel-server # - blobfuse2 write_files: diff --git a/swmcloudgate/routers/azure/templates/partition.json b/swmcloudgate/routers/azure/templates/partition.json index 7832dc7..9457941 100644 --- a/swmcloudgate/routers/azure/templates/partition.json +++ b/swmcloudgate/routers/azure/templates/partition.json @@ -18,15 +18,17 @@ }, "vmNameMain": { "type": "string", - "defaultValue": "[format('{0}-main', parameters('resourcePrefix'))]", "metadata": { "description": "The main VM name." } }, "adminUsername": { "type": "string", + "minLength": 1, + "maxLength": 32, + "allowedPattern": "^[a-z][a-z0-9_-]*$", "metadata": { - "description": "Admin username for all VMs." + "description": "Admin username for all VMs. Restricted to Linux-compatible usernames to avoid invalid authorized_keys paths and deployment failures." } }, "authenticationType": { @@ -46,9 +48,15 @@ "description": "SSH Key or password." } }, + "storageKey": { + "type": "securestring", + "metadata": { + "description": "Azure storage access key." + } + }, "dnsLabelPrefix": { "type": "string", - "defaultValue": "[toLower(format('{0}-main', parameters('resourcePrefix')))]", + "defaultValue": "[toLower(parameters('vmNameMain'))]", "metadata": { "description": "Unique DNS Name for the Public IP used to access the VM." } @@ -140,7 +148,7 @@ "ssh": { "publicKeys": [ { - "path": "/home/taras/.ssh/authorized_keys", + "path": "[concat('/home/', parameters('adminUsername'), '/.ssh/authorized_keys')]", "keyData": "[parameters('adminPasswordOrKey')]" } ] diff --git a/test/test_azure.py b/test/test_azure.py index a3ad21c..0b4d5ff 100644 --- a/test/test_azure.py +++ b/test/test_azure.py @@ -1,14 +1,14 @@ import os import socket import asyncio +import unittest from multiprocessing import Process import aiohttp import uvicorn -import asynctest -class TestAzureGate(asynctest.TestCase): +class TestAzureGate(unittest.IsolatedAsyncioTestCase): _hostname: str = socket.gethostname() _port: int = 8445 @@ -20,7 +20,7 @@ class TestAzureGate(asynctest.TestCase): "extra": "location=test", } - async def setUp(self): + async def asyncSetUp(self): self.maxDiff = None os.environ["SWM_TEST_CONFIG"] = "test/data/responses.json" # Point routers at a test cloud-gate.yaml that provides non-empty @@ -45,10 +45,12 @@ async def setUp(self): await asyncio.sleep(0.5) # time for the server to start self.assertTrue(self.proc.is_alive()) - async def tearDown(self): - self.assertTrue(self.proc.is_alive()) - self.proc.terminate() - + async def asyncTearDown(self): + if self.proc.is_alive(): + self.proc.terminate() + self.proc.join(timeout=5) + os.environ.pop("SWM_TEST_CONFIG", None) + os.environ.pop("SWM_GATE_CONFIG", None) async def test_list_flavors(self): async with aiohttp.ClientSession(headers=self._default_headers) as session: async with session.get( @@ -294,7 +296,7 @@ async def test_create_partition(self): "containerimage": "swmregistry.azurecr.io/jupyter/datascience-notebook:hub-3.1.1", "flavorname": "Standard_B2s", "username": "user", - "count": "0", + "count": "1", "jobid": "3579a076-9924-11ee-ba53-a3132f7ae2fb", "partname": "part1", "runtime": "swm_source=ssh, ssh_pub_key=ssh-rsa AAAAB3NzaC1yc2EAAAABIwAAAQEA7GA", @@ -312,3 +314,34 @@ async def test_create_partition(self): data = await resp.text() self.assertEqual(list(data.keys()), ["partition"]) self.assertTrue(isinstance(data["partition"]["id"], str)) + + async def test_create_partition_invalid_count(self): + headers = { + "Accept": "application/json", + "subscriptionid": "test", + "tenantid": "test", + "appid": "test", + "containerregistryuser": "user", + "containerregistrypass": "pass", + "osversion": "ubuntu-22.04", + "containerimage": "swmregistry.azurecr.io/jupyter/datascience-notebook:hub-3.1.1", + "flavorname": "Standard_B2s", + "username": "user", + "count": "0", + "jobid": "3579a076-9924-11ee-ba53-a3132f7ae2fb", + "partname": "part-invalid", + "runtime": "swm_source=ssh, ssh_pub_key=ssh-rsa AAAAB3NzaC1yc2EAAAABIwAAAQEA7GA", + "location": "eastus", + "ports": "10001,10022", + } + async with aiohttp.ClientSession(headers=headers) as session: + async with session.post( + url=f"http://{self._hostname}:{self._port}/azure/partitions", + json={"pem_data": "test"}, + ) as resp: + self.assertEqual(resp.status, 400) + try: + data = await resp.json() + except aiohttp.client_exceptions.ContentTypeError: + data = await resp.text() + self.assertIn("detail", data) diff --git a/test/test_azure_connector.py b/test/test_azure_connector.py new file mode 100644 index 0000000..a4f6479 --- /dev/null +++ b/test/test_azure_connector.py @@ -0,0 +1,289 @@ +import json +import os +import unittest +from unittest.mock import Mock + +from swmcloudgate.routers.azure.connector import AzureConnector, MAX_VM_COUNT + + +class TestAzureConnectorMultiNode(unittest.TestCase): + def setUp(self): + os.environ["SWM_TEST_CONFIG"] = "test/data/responses.json" + self.connector = AzureConnector() + + def tearDown(self): + os.environ.pop("SWM_TEST_CONFIG", None) + + def _load_template(self): + with open("swmcloudgate/routers/azure/templates/partition.json") as template_file: + return json.load(template_file) + + def _render_cloud_init_script(self, **overrides): + params = { + "job_id": "job-1", + "container_image": "registry.example.org/image:tag", + "container_registry": "registry.example.org", + "container_registry_username": "user", + "container_registry_password": "pass", + "storage_account": "storageaccount", + "storage_container": "storagecontainer", + "runtime_params": {"swm_source": "ssh"}, + "user_ssh_cert": "ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABAQC", + } + params.update(overrides) + return self.connector._get_cloud_init_script( + job_id=params["job_id"], + container_image=params["container_image"], + container_registry=params["container_registry"], + container_registry_username=params["container_registry_username"], + container_registry_password=params["container_registry_password"], + storage_account=params["storage_account"], + storage_container=params["storage_container"], + runtime_params=params["runtime_params"], + user_ssh_cert=params["user_ssh_cert"], + ) + + def test_parse_vm_count_defaults_to_one_for_missing(self): + self.assertEqual(self.connector._parse_vm_count(None), 1) + self.assertEqual(self.connector._parse_vm_count(""), 1) + + def test_parse_vm_count_rejects_non_integer(self): + with self.assertRaises(ValueError): + self.connector._parse_vm_count("abc") + + def test_parse_vm_count_rejects_out_of_range_values(self): + with self.assertRaises(ValueError): + self.connector._parse_vm_count("0") + with self.assertRaises(ValueError): + self.connector._parse_vm_count(str(MAX_VM_COUNT + 1)) + + def test_add_compute_vms_keeps_single_vm_template_unchanged(self): + template = self._load_template() + original_resource_count = len(template["resources"]) + + self.connector._add_compute_vms("part1", 1, template) + + self.assertEqual(len(template["resources"]), original_resource_count) + + def test_add_compute_vms_adds_nics_vms_and_extensions(self): + template = self._load_template() + + self.connector._add_compute_vms("part1", 3, template) + + network_interfaces = [ + resource for resource in template["resources"] if resource["type"] == "Microsoft.Network/networkInterfaces" + ] + virtual_machines = [ + resource for resource in template["resources"] if resource["type"] == "Microsoft.Compute/virtualMachines" + ] + extensions = [ + resource + for resource in template["resources"] + if resource["type"] == "Microsoft.Compute/virtualMachines/extensions" + ] + + self.assertEqual(len(network_interfaces), 3) + self.assertEqual(len(virtual_machines), 3) + self.assertEqual(len(extensions), 3) + + def test_compute_nics_do_not_have_public_ip(self): + template = self._load_template() + + self.connector._add_compute_vms("part1", 2, template) + + compute_nic = next( + resource + for resource in template["resources"] + if resource["type"] == "Microsoft.Network/networkInterfaces" + and resource["name"] == "[format('part1-compute1-NetInt')]" + ) + nic_properties = compute_nic["properties"]["ipConfigurations"][0]["properties"] + self.assertNotIn("publicIPAddress", nic_properties) + + def test_compute_vm_preserves_main_nic_dependency_and_adds_compute_nic(self): + template = self._load_template() + + self.connector._add_compute_vms("part1", 2, template) + + compute_vm = next( + resource + for resource in template["resources"] + if resource["type"] == "Microsoft.Compute/virtualMachines" + and resource["name"] == "[format('part1-compute1')]" + ) + self.assertIn( + "[resourceId('Microsoft.Network/networkInterfaces', 'part1-compute1-NetInt')]", + compute_vm["dependsOn"], + ) + self.assertIn( + "[resourceId('Microsoft.Network/networkInterfaces', variables('networkInterfaceName'))]", + compute_vm["dependsOn"], + ) + + def test_compute_nic_dependencies_preserve_non_public_ip_dependencies(self): + template = self._load_template() + + self.connector._add_compute_vms("part1", 2, template) + + compute_nic = next( + resource + for resource in template["resources"] + if resource["type"] == "Microsoft.Network/networkInterfaces" + and resource["name"] == "[format('part1-compute1-NetInt')]" + ) + self.assertIn( + "[resourceId('Microsoft.Network/networkSecurityGroups', parameters('networkSecurityGroupName'))]", + compute_nic["dependsOn"], + ) + self.assertIn( + "[resourceId('Microsoft.Network/virtualNetworks/subnets', parameters('virtualNetworkName'), " + "parameters('subnetName'))]", + compute_nic["dependsOn"], + ) + self.assertNotIn( + "[resourceId('Microsoft.Network/publicIPAddresses', variables('publicIPAddressName'))]", + compute_nic["dependsOn"], + ) + + def test_compute_vm_extension_is_duplicated_for_compute_nodes(self): + template = self._load_template() + + self.connector._add_compute_vms("part1", 2, template) + + compute_extension = next( + resource + for resource in template["resources"] + if resource["type"] == "Microsoft.Compute/virtualMachines/extensions" + and resource["name"] == "[format('part1-compute1/{0}', variables('extensionName'))]" + ) + self.assertIn( + "[resourceId('Microsoft.Compute/virtualMachines', 'part1-compute1')]", + compute_extension["dependsOn"], + ) + self.assertNotIn( + "[resourceId('Microsoft.Compute/virtualMachines', parameters('vmNameMain'))]", + compute_extension["dependsOn"], + ) + + def test_cloud_init_script_detects_main_and_compute_vm_roles(self): + cloud_init_script = self._render_cloud_init_script() + + self.assertIn('HOST_NAME="__SWM_HOST_NAME__"', cloud_init_script) + self.assertIn("IS_MAIN=__SWM_IS_MAIN__", cloud_init_script) + self.assertIn('MAIN_INSTANCE_HOSTNAME="__SWM_MAIN_INSTANCE_HOSTNAME__"', cloud_init_script) + + def test_cloud_init_script_configures_nfs_for_shared_home_mount(self): + cloud_init_script = self._render_cloud_init_script() + + self.assertIn("exportfs -ra", cloud_init_script) + self.assertIn("systemctl enable nfs-kernel-server", cloud_init_script) + self.assertIn('echo "$MAIN_INSTANCE_PRIVATE_IP:/home /home nfs', cloud_init_script) + self.assertIn('timeout 2 bash -c "&2', cloud_init_script) + + def test_cloud_init_script_quotes_shell_special_characters(self): + cloud_init_script = self._render_cloud_init_script( + container_registry_password="pa&ss", + ) + + self.assertIn("pa&ss", cloud_init_script) + self.assertIn("local container_registry_password='pa&ss'", cloud_init_script) + self.assertIn('if [ -n "$container_registry_password" ]; then', cloud_init_script) + self.assertIn( + 'docker login "$container_registry" --username "$container_registry_username" --password "$container_registry_password"', + cloud_init_script, + ) + self.assertIn('docker pull "$container_image"', cloud_init_script) + self.assertNotIn("pa&ss<word>", cloud_init_script) + + def test_create_deployment_builds_multi_node_template(self): + self.connector._test_responses = {} + self.connector._subscription_id = "test-subscription" + self.connector._resource_client = Mock() + self.connector._resource_client.resource_groups.create_or_update.return_value = Mock( + id="/subscriptions/test/rg" + ) + + deployment_result = {"id": "/subscriptions/test/deployments/part1"} + deployment_operation = Mock() + deployment_operation.result.return_value = deployment_result + self.connector._resource_client.deployments.begin_create_or_update.return_value = deployment_operation + + result, resource_group_name = self.connector.create_deployment( + job_id="job-1", + partition_name="part1", + os_version="ubuntu-hpc/2204", + container_image="registry.example.org/image:tag", + container_registry_username="user", + container_registry_password="pass", + storage_account="storageaccount", + storage_key="storagekey", + storage_container="storagecontainer", + flavor_name="Standard_B2s", + username="azureuser", + count="3", + runtime="swm_source=ssh", + location="eastus", + ports="10001,10022", + user_ssh_cert="ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABAQC", + ) + + self.assertEqual(result, deployment_result) + self.assertEqual(resource_group_name, "part1-resource-group") + _, _, deployment_properties = self.connector._resource_client.deployments.begin_create_or_update.call_args[0] + template = deployment_properties["properties"]["template"] + parameters = deployment_properties["properties"]["parameters"] + virtual_machines = [r for r in template["resources"] if r["type"] == "Microsoft.Compute/virtualMachines"] + + self.assertEqual(len(virtual_machines), 3) + self.assertEqual(parameters["vmNameMain"]["value"], "part1-main") + self.assertEqual(parameters["storageKey"]["value"], "storagekey") + self.assertNotIn("storagekey", parameters["cloudInitScript"]["value"]) + self.assertEqual( + template["variables"]["linuxConfiguration"]["ssh"]["publicKeys"][0]["path"], + "[concat('/home/', parameters('adminUsername'), '/.ssh/authorized_keys')]", + ) + + main_vm = next(resource for resource in virtual_machines if resource["name"] == "[parameters('vmNameMain')]") + self.assertIn("base64(parameters('storageKey'))", main_vm["properties"]["osProfile"]["customData"]) + + def test_create_deployment_rolls_back_resource_group_on_failure(self): + self.connector._test_responses = {} + self.connector._subscription_id = "test-subscription" + self.connector._resource_client = Mock() + self.connector._resource_client.resource_groups.create_or_update.return_value = Mock( + id="/subscriptions/test/rg" + ) + + deployment_operation = Mock() + deployment_operation.result.side_effect = RuntimeError("deployment failed") + self.connector._resource_client.deployments.begin_create_or_update.return_value = deployment_operation + delete_operation = Mock() + self.connector._resource_client.resource_groups.begin_delete.return_value = delete_operation + + with self.assertRaises(RuntimeError): + self.connector.create_deployment( + job_id="job-1", + partition_name="part1", + os_version="ubuntu-hpc/2204", + container_image="registry.example.org/image:tag", + container_registry_username="user", + container_registry_password="pass", + storage_account="storageaccount", + storage_key="storagekey", + storage_container="storagecontainer", + flavor_name="Standard_B2s", + username="azureuser", + count="3", + runtime="swm_source=ssh", + location="eastus", + ports="10001,10022", + user_ssh_cert="ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABAQC", + ) + + self.connector._resource_client.resource_groups.begin_delete.assert_called_once_with("part1-resource-group") + delete_operation.result.assert_called_once() diff --git a/test/test_azure_custom_data.py b/test/test_azure_custom_data.py new file mode 100644 index 0000000..cfea7fc --- /dev/null +++ b/test/test_azure_custom_data.py @@ -0,0 +1,132 @@ +import json +import os +import shlex +import unittest + +from swmcloudgate.routers.azure.connector import ( + AzureConnector, + HOST_NAME_PLACEHOLDER, + IS_MAIN_PLACEHOLDER, + MAIN_INSTANCE_HOSTNAME_PLACEHOLDER, + MAIN_INSTANCE_PRIVATE_IP_PLACEHOLDER, + STORAGE_KEY_B64_PLACEHOLDER, +) + + +class TestAzureConnectorCustomDataInjection(unittest.TestCase): + _test_config = "test/data/responses.json" + + def setUp(self): + os.environ["SWM_TEST_CONFIG"] = self._test_config + self.connector = AzureConnector() + + def tearDown(self): + os.environ.pop("SWM_TEST_CONFIG", None) + + def _load_template(self): + with open("swmcloudgate/routers/azure/templates/partition.json") as template_file: + return json.load(template_file) + + def _render_cloud_init_script(self): + return self.connector._get_cloud_init_script( + job_id="job-1", + container_image="registry.example.org/image:tag", + container_registry="registry.example.org", + container_registry_username="user", + container_registry_password="user<&>\"'pass", + storage_account="storageaccount", + storage_container="storagecontainer", + runtime_params={"swm_source": "ssh"}, + user_ssh_cert="ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABAQC== user@example", + ) + + def test_cloud_init_script_uses_explicit_main_instance_placeholders(self): + cloud_init_script = self._render_cloud_init_script() + expected_password = shlex.quote("user<&>\"'pass") + + self.assertIn(f'HOST_NAME="{HOST_NAME_PLACEHOLDER}"', cloud_init_script) + self.assertIn(f"IS_MAIN={IS_MAIN_PLACEHOLDER}", cloud_init_script) + self.assertIn( + f'MAIN_INSTANCE_HOSTNAME="{MAIN_INSTANCE_HOSTNAME_PLACEHOLDER}"', + cloud_init_script, + ) + self.assertIn( + f'MAIN_INSTANCE_PRIVATE_IP="{MAIN_INSTANCE_PRIVATE_IP_PLACEHOLDER}"', + cloud_init_script, + ) + self.assertIn(f'local azure_storage_key_b64="{STORAGE_KEY_B64_PLACEHOLDER}"', cloud_init_script) + self.assertIn("azure_storage_key=$(printf '%s' \"$azure_storage_key_b64\" | base64 -d)", cloud_init_script) + self.assertIn( + f"local container_registry_password={expected_password}", + cloud_init_script, + ) + self.assertIn( + 'docker login "$container_registry" --username "$container_registry_username" --password "$container_registry_password"', + cloud_init_script, + ) + self.assertNotIn( + "docker login registry.example.org --username user --password user<&>\"'pass", + cloud_init_script, + ) + self.assertIn( + "echo 'ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABAQC== user@example' >> /root/.ssh/authorized_keys", + cloud_init_script, + ) + self.assertNotIn("<", cloud_init_script) + self.assertNotIn(">", cloud_init_script) + self.assertNotIn("&", cloud_init_script) + self.assertNotIn("getent hosts", cloud_init_script) + + def test_main_vm_custom_data_replaces_placeholders_via_arm(self): + template = self._load_template() + + self.connector._configure_main_vm_custom_data(template) + + main_vm = next( + resource + for resource in template["resources"] + if resource["type"] == "Microsoft.Compute/virtualMachines" + and resource["name"] == "[parameters('vmNameMain')]" + ) + custom_data = main_vm["properties"]["osProfile"]["customData"] + self.assertIn("[base64(", custom_data) + self.assertIn(f"'{HOST_NAME_PLACEHOLDER}'", custom_data) + self.assertIn("parameters('vmNameMain')", custom_data) + self.assertIn(f"'{IS_MAIN_PLACEHOLDER}'", custom_data) + self.assertIn("'true'", custom_data) + self.assertIn(f"'{MAIN_INSTANCE_PRIVATE_IP_PLACEHOLDER}'", custom_data) + self.assertIn( + "reference(resourceId('Microsoft.Network/networkInterfaces'", + custom_data, + ) + self.assertIn("base64(parameters('storageKey'))", custom_data) + + def test_compute_vm_custom_data_uses_explicit_main_private_ip_reference(self): + template = self._load_template() + + self.connector._configure_main_vm_custom_data(template) + self.connector._add_compute_vms("part1", 2, template) + + compute_vm = next( + resource + for resource in template["resources"] + if resource["type"] == "Microsoft.Compute/virtualMachines" + and resource["name"] == "[format('part1-compute1')]" + ) + custom_data = compute_vm["properties"]["osProfile"]["customData"] + self.assertIn(f"'{HOST_NAME_PLACEHOLDER}'", custom_data) + self.assertIn("'part1-compute1'", custom_data) + self.assertIn(f"'{IS_MAIN_PLACEHOLDER}'", custom_data) + self.assertIn("'false'", custom_data) + self.assertIn(f"'{MAIN_INSTANCE_HOSTNAME_PLACEHOLDER}'", custom_data) + self.assertIn("parameters('vmNameMain')", custom_data) + self.assertIn(f"'{MAIN_INSTANCE_PRIVATE_IP_PLACEHOLDER}'", custom_data) + self.assertIn( + "reference(resourceId('Microsoft.Network/networkInterfaces'", + custom_data, + ) + self.assertIn("base64(parameters('storageKey'))", custom_data) + self.assertIn( + "[resourceId('Microsoft.Network/networkInterfaces', variables('networkInterfaceName'))]", + compute_vm["dependsOn"], + ) diff --git a/test/test_openstack.py b/test/test_openstack.py index 270f3ec..d728ac0 100644 --- a/test/test_openstack.py +++ b/test/test_openstack.py @@ -1,19 +1,19 @@ import os import socket import asyncio +import unittest from multiprocessing import Process import aiohttp import uvicorn -import asynctest -class TestOpenstackGate(asynctest.TestCase): +class TestOpenstackGate(unittest.IsolatedAsyncioTestCase): _hostname: str = socket.gethostname() _port: int = 8445 - async def setUp(self): + async def asyncSetUp(self): self.maxDiff = None os.environ["SWM_TEST_CONFIG"] = "test/data/responses.json" self.proc = Process( @@ -32,9 +32,11 @@ async def setUp(self): await asyncio.sleep(0.5) # time for the server to start self.assertTrue(self.proc.is_alive()) - async def tearDown(self): - self.assertTrue(self.proc.is_alive()) - self.proc.terminate() + async def asyncTearDown(self): + if self.proc.is_alive(): + self.proc.terminate() + self.proc.join(timeout=5) + os.environ.pop("SWM_TEST_CONFIG", None) async def test_list_flavors(self): headers = {