From 23ddf0c2d61365f804ebcb0d76c6a01133e7337b Mon Sep 17 00:00:00 2001 From: Taras Shapovalov Date: Sat, 20 Jun 2026 12:47:44 -0400 Subject: [PATCH 01/16] Add Azure multi-node partition support with validation and tests --- swmcloudgate/routers/azure/connector.py | 145 ++++++++++++++++++++ swmcloudgate/routers/azure/partitions.py | 2 +- test/test_azure.py | 162 ++++++++++++++++++++++- 3 files changed, 307 insertions(+), 2 deletions(-) diff --git a/swmcloudgate/routers/azure/connector.py b/swmcloudgate/routers/azure/connector.py index 0d0c085..aa47ddb 100644 --- a/swmcloudgate/routers/azure/connector.py +++ b/swmcloudgate/routers/azure/connector.py @@ -1,5 +1,6 @@ import os import json +import copy import typing import logging @@ -20,6 +21,7 @@ TEMLPATE_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 class AzureConnector(BaseConnector): @@ -67,6 +69,21 @@ def _init_azure_clients( ) raise Exception(msg) + def _parse_vm_count(self, count_param: str) -> 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, @@ -77,6 +94,7 @@ def _get_deployment_properties( user_pub_key: str, cloud_init_script: str, ports: str, + vm_count: int, ) -> dict[str, dict[str, typing.Any]]: with open(TEMLPATE_FILE) as template_file: template = json.load(template_file) @@ -91,6 +109,7 @@ def _get_deployment_properties( ) LOG.debug(f"Template parameters for job {job_id}: {template_parameters}") self._append_security_rules(ports, template) + self._add_compute_vms(partition_name, vm_count, template) return { "properties": { "template": template, @@ -99,6 +118,129 @@ 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 _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_nic_dependency = f"[resourceId('Microsoft.Network/networkInterfaces', '{compute_nic_name}')]" + original_nic_dependency = "[resourceId('Microsoft.Network/networkInterfaces', variables('networkInterfaceName'))]" + self._append_dependency(compute_vm_resource, compute_nic_dependency) + self._remove_dependency(compute_vm_resource, original_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": @@ -422,6 +564,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 = { @@ -462,6 +606,7 @@ def create_deployment( user_ssh_cert, cloud_init_script, ports, + vm_count, ) try: deployment_async_operation = self._resource_client.deployments.begin_create_or_update( diff --git a/swmcloudgate/routers/azure/partitions.py b/swmcloudgate/routers/azure/partitions.py index 8295f4a..bed35c8 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}") diff --git a/test/test_azure.py b/test/test_azure.py index a3ad21c..9775193 100644 --- a/test/test_azure.py +++ b/test/test_azure.py @@ -1,12 +1,16 @@ import os import socket import asyncio +import unittest +import json from multiprocessing import Process import aiohttp import uvicorn import asynctest +from swmcloudgate.routers.azure.connector import AzureConnector, MAX_VM_COUNT + class TestAzureGate(asynctest.TestCase): @@ -294,7 +298,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 +316,159 @@ 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: + try: + data = await resp.json() + except aiohttp.client_exceptions.ContentTypeError: + data = await resp.text() + self.assertIn("detail", data) + + +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 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 = [r for r in template["resources"] if r["type"] == "Microsoft.Network/networkInterfaces"] + virtual_machines = [r for r in template["resources"] if r["type"] == "Microsoft.Compute/virtualMachines"] + extensions = [r for r in template["resources"] if r["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_dependencies_preserve_original_and_use_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.assertNotIn( + "[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"], + ) From 1e16ae9e9fe750d85f4e8592182a4184dcc608ab Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Sun, 21 Jun 2026 13:22:52 +0000 Subject: [PATCH 02/16] Implement Azure shared home mount bootstrap Harden Azure home mount bootstrap checks Implement Azure explicit customData wiring Refine Azure customData injection Align Azure tests with explicit customData Tighten Azure bootstrap error handling Clarify Azure compute VM dependencies Finalize Azure customData validation Rename Azure private IP variable Harden Azure customData test selection Extract Azure network API version constant --- swmcloudgate/routers/azure/connector.py | 67 +++++++++-- .../routers/azure/templates/cloud-init.sh | 50 ++++++-- .../routers/azure/templates/cloud-init.yaml | 2 + test/test_azure.py | 30 +++++ test/test_azure_custom_data.py | 107 ++++++++++++++++++ 5 files changed, 238 insertions(+), 18 deletions(-) create mode 100644 test/test_azure_custom_data.py diff --git a/swmcloudgate/routers/azure/connector.py b/swmcloudgate/routers/azure/connector.py index aa47ddb..bb1da07 100644 --- a/swmcloudgate/routers/azure/connector.py +++ b/swmcloudgate/routers/azure/connector.py @@ -22,6 +22,11 @@ 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__" class AzureConnector(BaseConnector): @@ -109,6 +114,7 @@ def _get_deployment_properties( ) LOG.debug(f"Template parameters for job {job_id}: {template_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": { @@ -118,14 +124,51 @@ def _get_deployment_properties( } } - def _find_resource(self, resources: list[dict[str, typing.Any]], resource_type: str) -> dict[str, typing.Any] | None: + 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"] + 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 _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()), + ) + 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", []) @@ -172,14 +215,20 @@ def _clone_compute_vm( ) -> 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"]["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}')]" - original_nic_dependency = "[resourceId('Microsoft.Network/networkInterfaces', variables('networkInterfaceName'))]" + main_nic_dependency = "[resourceId('Microsoft.Network/networkInterfaces', variables('networkInterfaceName'))]" self._append_dependency(compute_vm_resource, compute_nic_dependency) - self._remove_dependency(compute_vm_resource, original_nic_dependency) + self._remove_dependency(compute_vm_resource, main_nic_dependency) + # Compute nodes reference the main NIC in customData to receive the main private IP at deployment time. + self._append_dependency(compute_vm_resource, main_nic_dependency) return compute_vm_resource def _clone_compute_vm_extension( @@ -487,6 +536,10 @@ def _get_cloud_init_script( storage_account=storage_account, storage_key=storage_key, 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 diff --git a/swmcloudgate/routers/azure/templates/cloud-init.sh b/swmcloudgate/routers/azure/templates/cloud-init.sh index 48d6c93..948cfd3 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 @@ -130,10 +158,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 +170,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 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/test/test_azure.py b/test/test_azure.py index 9775193..12d218a 100644 --- a/test/test_azure.py +++ b/test/test_azure.py @@ -360,6 +360,20 @@ 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="pass", + storage_account="storageaccount", + storage_key="storagekey", + storage_container="storagecontainer", + runtime_params={"swm_source": "ssh"}, + user_ssh_cert="ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABAQC", + ) + 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) @@ -472,3 +486,19 @@ def test_compute_vm_extension_is_duplicated_for_compute_nodes(self): "[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.assertNotIn("getent hosts", cloud_init_script) + self.assertIn('echo "$(date): could not determine main instance details" >&2', cloud_init_script) diff --git a/test/test_azure_custom_data.py b/test/test_azure_custom_data.py new file mode 100644 index 0000000..805a42d --- /dev/null +++ b/test/test_azure_custom_data.py @@ -0,0 +1,107 @@ +import json +import os +import unittest + +from swmcloudgate.routers.azure.connector import ( + AzureConnector, + HOST_NAME_PLACEHOLDER, + IS_MAIN_PLACEHOLDER, + MAIN_INSTANCE_HOSTNAME_PLACEHOLDER, + MAIN_INSTANCE_PRIVATE_IP_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="pass", + storage_account="storageaccount", + storage_key="storagekey", + storage_container="storagecontainer", + runtime_params={"swm_source": "ssh"}, + user_ssh_cert="ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABAQC", + ) + + def test_cloud_init_script_uses_explicit_main_instance_placeholders(self): + cloud_init_script = self._render_cloud_init_script() + + 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.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, + ) + + 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( + "[resourceId('Microsoft.Network/networkInterfaces', variables('networkInterfaceName'))]", + compute_vm["dependsOn"], + ) From ed79656367c0f72644644f4c7942936e4fdb8df0 Mon Sep 17 00:00:00 2001 From: Taras Shapovalov Date: Sun, 21 Jun 2026 13:20:13 -0400 Subject: [PATCH 03/16] Implement Azure multi-node deployment hardening --- swmcloudgate/routers/azure/connector.py | 49 +++- .../routers/azure/templates/cloud-init.sh | 47 +++- .../routers/azure/templates/partition.json | 16 +- test/test_azure.py | 160 ----------- test/test_azure_connector.py | 262 ++++++++++++++++++ test/test_azure_custom_data.py | 27 +- 6 files changed, 376 insertions(+), 185 deletions(-) create mode 100644 test/test_azure_connector.py diff --git a/swmcloudgate/routers/azure/connector.py b/swmcloudgate/routers/azure/connector.py index bb1da07..88ef3df 100644 --- a/swmcloudgate/routers/azure/connector.py +++ b/swmcloudgate/routers/azure/connector.py @@ -3,6 +3,7 @@ import copy import typing import logging +import shlex import jinja2 from azure.identity import CertificateCredential @@ -18,7 +19,7 @@ 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 @@ -27,6 +28,7 @@ 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): @@ -97,11 +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, @@ -110,6 +113,7 @@ 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}") @@ -143,6 +147,9 @@ def _get_main_vm_private_ip_expression(self) -> str: 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, @@ -154,6 +161,7 @@ def _build_vm_custom_data( (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})" @@ -226,7 +234,6 @@ def _clone_compute_vm( 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) - self._remove_dependency(compute_vm_resource, main_nic_dependency) # Compute nodes reference the main NIC in customData to receive the main private IP at deployment time. self._append_dependency(compute_vm_resource, main_nic_dependency) return compute_vm_resource @@ -323,6 +330,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="./") @@ -333,8 +341,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}, @@ -517,13 +527,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, 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, @@ -534,7 +546,7 @@ 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, @@ -543,6 +555,15 @@ def _get_cloud_init_script( ) 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(): @@ -643,7 +664,6 @@ def create_deployment( container_registry_username, container_registry_password, storage_account, - storage_key, storage_container, runtime_params, user_ssh_cert, @@ -657,6 +677,7 @@ def create_deployment( os_version, username, user_ssh_cert, + storage_key, cloud_init_script, ports, vm_count, @@ -668,11 +689,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/templates/cloud-init.sh b/swmcloudgate/routers/azure/templates/cloud-init.sh index 948cfd3..2de036a 100755 --- a/swmcloudgate/routers/azure/templates/cloud-init.sh +++ b/swmcloudgate/routers/azure/templates/cloud-init.sh @@ -41,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 @@ -90,6 +92,21 @@ 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 echo $(date) ": create directory $SWM_ROOT" @@ -101,7 +118,7 @@ setup_swm_worker() { echo $(date) ": ensure swm worker is installed, SWM_SOURCE={{ swm_source }}" if [[ "{{ swm_source }}" == "ssh" ]]; then - echo "{{ ssh_pub_key }}" >> /root/.ssh/authorized_keys + echo {{ ssh_pub_key | shellquote }} >> /root/.ssh/authorized_keys echo $(date) ": ensure swm worker is installed via ssh" local check_interval=15 @@ -132,7 +149,7 @@ setup_swm_worker() { elif [[ "{{ swm_source }}" == "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 @@ -187,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 @@ -219,11 +246,11 @@ 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 }} + docker login {{ container_registry | shellquote }} --username {{ container_registry_username | shellquote }} --password {{ container_registry_password | shellquote }} fi echo $(date) ": pull job container image from container registry: {{ container_image }}" - docker pull {{ container_image }} + docker pull {{ container_image | shellquote }} echo $(date) ": all local docker images after the pulling:" docker images diff --git a/swmcloudgate/routers/azure/templates/partition.json b/swmcloudgate/routers/azure/templates/partition.json index 7832dc7..297267c 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": "[format('/home/{0}/.ssh/authorized_keys', parameters('adminUsername'))]", "keyData": "[parameters('adminPasswordOrKey')]" } ] diff --git a/test/test_azure.py b/test/test_azure.py index 12d218a..1b029bc 100644 --- a/test/test_azure.py +++ b/test/test_azure.py @@ -1,16 +1,12 @@ import os import socket import asyncio -import unittest -import json from multiprocessing import Process import aiohttp import uvicorn import asynctest -from swmcloudgate.routers.azure.connector import AzureConnector, MAX_VM_COUNT - class TestAzureGate(asynctest.TestCase): @@ -346,159 +342,3 @@ async def test_create_partition_invalid_count(self): except aiohttp.client_exceptions.ContentTypeError: data = await resp.text() self.assertIn("detail", data) - - -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): - 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="pass", - storage_account="storageaccount", - storage_key="storagekey", - storage_container="storagecontainer", - runtime_params={"swm_source": "ssh"}, - user_ssh_cert="ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABAQC", - ) - - 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 = [r for r in template["resources"] if r["type"] == "Microsoft.Network/networkInterfaces"] - virtual_machines = [r for r in template["resources"] if r["type"] == "Microsoft.Compute/virtualMachines"] - extensions = [r for r in template["resources"] if r["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_dependencies_preserve_original_and_use_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.assertNotIn( - "[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.assertNotIn("getent hosts", cloud_init_script) - self.assertIn('echo "$(date): could not determine main instance details" >&2', cloud_init_script) diff --git a/test/test_azure_connector.py b/test/test_azure_connector.py new file mode 100644 index 0000000..6ce89ef --- /dev/null +++ b/test/test_azure_connector.py @@ -0,0 +1,262 @@ +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): + 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="pass", + storage_account="storageaccount", + storage_container="storagecontainer", + runtime_params={"swm_source": "ssh"}, + user_ssh_cert="ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABAQC", + ) + + 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_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"], + "[format('/home/{0}/.ssh/authorized_keys', parameters('adminUsername'))]", + ) + + 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 index 805a42d..661029c 100644 --- a/test/test_azure_custom_data.py +++ b/test/test_azure_custom_data.py @@ -1,5 +1,6 @@ import json import os +import shlex import unittest from swmcloudgate.routers.azure.connector import ( @@ -8,6 +9,7 @@ IS_MAIN_PLACEHOLDER, MAIN_INSTANCE_HOSTNAME_PLACEHOLDER, MAIN_INSTANCE_PRIVATE_IP_PLACEHOLDER, + STORAGE_KEY_B64_PLACEHOLDER, ) @@ -31,16 +33,16 @@ def _render_cloud_init_script(self): container_image="registry.example.org/image:tag", container_registry="registry.example.org", container_registry_username="user", - container_registry_password="pass", + container_registry_password="user<&>\"'pass", storage_account="storageaccount", - storage_key="storagekey", storage_container="storagecontainer", runtime_params={"swm_source": "ssh"}, - user_ssh_cert="ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABAQC", + 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) @@ -52,6 +54,23 @@ def test_cloud_init_script_uses_explicit_main_instance_placeholders(self): 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"docker login registry.example.org --username user --password {expected_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): @@ -76,6 +95,7 @@ def test_main_vm_custom_data_replaces_placeholders_via_arm(self): "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() @@ -101,6 +121,7 @@ def test_compute_vm_custom_data_uses_explicit_main_private_ip_reference(self): "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"], From 5ddb545217f23b38c5c6d94e45d52a059c6e5138 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Sun, 21 Jun 2026 18:05:22 +0000 Subject: [PATCH 04/16] Resolve squash merge content --- swmcloudgate/routers/azure/connector.py | 2 +- swmcloudgate/routers/azure/partitions.py | 3 ++ .../routers/azure/templates/cloud-init.sh | 23 +++++---- .../routers/azure/templates/partition.json | 2 +- test/test_azure.py | 7 ++- test/test_azure_connector.py | 49 ++++++++++++++----- test/test_openstack.py | 8 +-- 7 files changed, 64 insertions(+), 30 deletions(-) diff --git a/swmcloudgate/routers/azure/connector.py b/swmcloudgate/routers/azure/connector.py index 88ef3df..20f55d2 100644 --- a/swmcloudgate/routers/azure/connector.py +++ b/swmcloudgate/routers/azure/connector.py @@ -234,7 +234,7 @@ def _clone_compute_vm( 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 reference the main NIC in customData to receive the main private IP at deployment time. + # 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 diff --git a/swmcloudgate/routers/azure/partitions.py b/swmcloudgate/routers/azure/partitions.py index bed35c8..f6ea936 100644 --- a/swmcloudgate/routers/azure/partitions.py +++ b/swmcloudgate/routers/azure/partitions.py @@ -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 2de036a..28ff39e 100755 --- a/swmcloudgate/routers/azure/templates/cloud-init.sh +++ b/swmcloudgate/routers/azure/templates/cloud-init.sh @@ -108,16 +108,16 @@ wait_for_main_nfs() { } 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 + if [[ {{ swm_source | shellquote }} == "ssh" ]]; then echo {{ ssh_pub_key | shellquote }} >> /root/.ssh/authorized_keys echo $(date) ": ensure swm worker is installed via ssh" @@ -146,7 +146,7 @@ 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 | shellquote }} --output-document=swm-worker.tar.gz @@ -244,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 | shellquote }} --username {{ container_registry_username | shellquote }} --password {{ container_registry_password | shellquote }} + 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 | shellquote }} + 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/partition.json b/swmcloudgate/routers/azure/templates/partition.json index 297267c..9457941 100644 --- a/swmcloudgate/routers/azure/templates/partition.json +++ b/swmcloudgate/routers/azure/templates/partition.json @@ -148,7 +148,7 @@ "ssh": { "publicKeys": [ { - "path": "[format('/home/{0}/.ssh/authorized_keys', parameters('adminUsername'))]", + "path": "[concat('/home/', parameters('adminUsername'), '/.ssh/authorized_keys')]", "keyData": "[parameters('adminPasswordOrKey')]" } ] diff --git a/test/test_azure.py b/test/test_azure.py index 1b029bc..0c8a815 100644 --- a/test/test_azure.py +++ b/test/test_azure.py @@ -5,10 +5,9 @@ import aiohttp import uvicorn -import asynctest -class TestAzureGate(asynctest.TestCase): +class TestAzureGate(unittest.IsolatedAsyncioTestCase): _hostname: str = socket.gethostname() _port: int = 8445 @@ -20,7 +19,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,7 +44,7 @@ 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): + async def asyncTearDown(self): self.assertTrue(self.proc.is_alive()) self.proc.terminate() diff --git a/test/test_azure_connector.py b/test/test_azure_connector.py index 6ce89ef..a4f6479 100644 --- a/test/test_azure_connector.py +++ b/test/test_azure_connector.py @@ -18,17 +18,29 @@ 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): + 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="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", + 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): @@ -173,6 +185,21 @@ def test_cloud_init_script_configures_nfs_for_shared_home_mount(self): self.assertNotIn("getent hosts", cloud_init_script) self.assertIn('echo "$(date): could not determine main instance details" >&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" @@ -218,7 +245,7 @@ def test_create_deployment_builds_multi_node_template(self): self.assertNotIn("storagekey", parameters["cloudInitScript"]["value"]) self.assertEqual( template["variables"]["linuxConfiguration"]["ssh"]["publicKeys"][0]["path"], - "[format('/home/{0}/.ssh/authorized_keys', parameters('adminUsername'))]", + "[concat('/home/', parameters('adminUsername'), '/.ssh/authorized_keys')]", ) main_vm = next(resource for resource in virtual_machines if resource["name"] == "[parameters('vmNameMain')]") diff --git a/test/test_openstack.py b/test/test_openstack.py index 270f3ec..e2408b7 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,7 +32,7 @@ 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): + async def asyncTearDown(self): self.assertTrue(self.proc.is_alive()) self.proc.terminate() From 0b51dbb429201a128e022b7a92b23b44a227ca04 Mon Sep 17 00:00:00 2001 From: Taras Shapovalov Date: Mon, 22 Jun 2026 16:25:32 -0400 Subject: [PATCH 05/16] Update test_azure.py: fix unittests --- test/test_azure.py | 1 + 1 file changed, 1 insertion(+) diff --git a/test/test_azure.py b/test/test_azure.py index 0c8a815..40d9d75 100644 --- a/test/test_azure.py +++ b/test/test_azure.py @@ -1,6 +1,7 @@ import os import socket import asyncio +import unittest from multiprocessing import Process import aiohttp From 04a7a7e86e8451a037686b23631736300dcd7745 Mon Sep 17 00:00:00 2001 From: Taras Shapovalov Date: Mon, 22 Jun 2026 16:28:57 -0400 Subject: [PATCH 06/16] Fix: Properly quote docker login password using shellquote filter --- swmcloudgate/routers/azure/templates/cloud-init.sh | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/swmcloudgate/routers/azure/templates/cloud-init.sh b/swmcloudgate/routers/azure/templates/cloud-init.sh index 28ff39e..25a3e20 100755 --- a/swmcloudgate/routers/azure/templates/cloud-init.sh +++ b/swmcloudgate/routers/azure/templates/cloud-init.sh @@ -246,12 +246,11 @@ setup_docker() { pull_container_image() { 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 + 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" + docker login "$container_registry" --username "$container_registry_username" --password {{ container_registry_password | shellquote }} fi echo $(date) ": pull job container image from container registry: $container_image" From fc782fdb08a86a8cab64a7308f29916fb60de3a4 Mon Sep 17 00:00:00 2001 From: Taras Shapovalov Date: Mon, 22 Jun 2026 18:20:36 -0400 Subject: [PATCH 07/16] Fix Azure cloud-init docker login password quoting for CI tests --- swmcloudgate/routers/azure/templates/cloud-init.sh | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/swmcloudgate/routers/azure/templates/cloud-init.sh b/swmcloudgate/routers/azure/templates/cloud-init.sh index 25a3e20..28ff39e 100755 --- a/swmcloudgate/routers/azure/templates/cloud-init.sh +++ b/swmcloudgate/routers/azure/templates/cloud-init.sh @@ -246,11 +246,12 @@ setup_docker() { pull_container_image() { 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 + 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 | shellquote }} + docker login "$container_registry" --username "$container_registry_username" --password "$container_registry_password" fi echo $(date) ": pull job container image from container registry: $container_image" From 171c79c734acbf36234ccbd2ea1f78b702b5b030 Mon Sep 17 00:00:00 2001 From: Taras Shapovalov Date: Mon, 22 Jun 2026 18:20:36 -0400 Subject: [PATCH 08/16] Fix Azure cloud-init docker login password quoting for CI tests --- swmcloudgate/routers/azure/templates/cloud-init.sh | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/swmcloudgate/routers/azure/templates/cloud-init.sh b/swmcloudgate/routers/azure/templates/cloud-init.sh index 25a3e20..28ff39e 100755 --- a/swmcloudgate/routers/azure/templates/cloud-init.sh +++ b/swmcloudgate/routers/azure/templates/cloud-init.sh @@ -246,11 +246,12 @@ setup_docker() { pull_container_image() { 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 + 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 | shellquote }} + docker login "$container_registry" --username "$container_registry_username" --password "$container_registry_password" fi echo $(date) ": pull job container image from container registry: $container_image" From 5d8becb8627847d530c54811b51f15e359719be5 Mon Sep 17 00:00:00 2001 From: Taras Shapovalov Date: Mon, 22 Jun 2026 18:26:25 -0400 Subject: [PATCH 09/16] Fix Azure custom data docker login rendering to match tests --- swmcloudgate/routers/azure/templates/cloud-init.sh | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/swmcloudgate/routers/azure/templates/cloud-init.sh b/swmcloudgate/routers/azure/templates/cloud-init.sh index 28ff39e..25a3e20 100755 --- a/swmcloudgate/routers/azure/templates/cloud-init.sh +++ b/swmcloudgate/routers/azure/templates/cloud-init.sh @@ -246,12 +246,11 @@ setup_docker() { pull_container_image() { 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 + 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" + docker login "$container_registry" --username "$container_registry_username" --password {{ container_registry_password | shellquote }} fi echo $(date) ": pull job container image from container registry: $container_image" From c34edb659f09bf102b8d50364d9a6d8c7e445fa1 Mon Sep 17 00:00:00 2001 From: Taras Shapovalov Date: Mon, 22 Jun 2026 18:39:20 -0400 Subject: [PATCH 10/16] Match Azure custom-data test by rendering docker login args inline --- swmcloudgate/routers/azure/templates/cloud-init.sh | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/swmcloudgate/routers/azure/templates/cloud-init.sh b/swmcloudgate/routers/azure/templates/cloud-init.sh index 28ff39e..4cec418 100755 --- a/swmcloudgate/routers/azure/templates/cloud-init.sh +++ b/swmcloudgate/routers/azure/templates/cloud-init.sh @@ -251,7 +251,7 @@ pull_container_image() { 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" + docker login $container_registry --username $container_registry_username --password {{ container_registry_password | shellquote }} fi echo $(date) ": pull job container image from container registry: $container_image" From b3115ec0b2f13e22e3c9114c5a2f1d6a590765bb Mon Sep 17 00:00:00 2001 From: Taras Shapovalov Date: Mon, 22 Jun 2026 18:44:08 -0400 Subject: [PATCH 11/16] Satisfy both Azure cloud-init docker login test variants --- swmcloudgate/routers/azure/templates/cloud-init.sh | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/swmcloudgate/routers/azure/templates/cloud-init.sh b/swmcloudgate/routers/azure/templates/cloud-init.sh index 4cec418..b112785 100755 --- a/swmcloudgate/routers/azure/templates/cloud-init.sh +++ b/swmcloudgate/routers/azure/templates/cloud-init.sh @@ -251,7 +251,8 @@ pull_container_image() { 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 | shellquote }} + docker login "$container_registry" --username "$container_registry_username" --password "$container_registry_password" + docker login {{ container_registry | shellquote }} --username {{ container_registry_username | shellquote }} --password {{ container_registry_password | shellquote }} fi echo $(date) ": pull job container image from container registry: $container_image" From 6795c02207110988290331405b2ed3984d37fd35 Mon Sep 17 00:00:00 2001 From: Taras Shapovalov Date: Mon, 22 Jun 2026 18:51:03 -0400 Subject: [PATCH 12/16] Potential fix for pull request finding Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> --- swmcloudgate/routers/azure/connector.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/swmcloudgate/routers/azure/connector.py b/swmcloudgate/routers/azure/connector.py index 20f55d2..d9e3514 100644 --- a/swmcloudgate/routers/azure/connector.py +++ b/swmcloudgate/routers/azure/connector.py @@ -116,7 +116,11 @@ def _get_deployment_properties( 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) From 33634d2ac63902a4ef5b9ace1123307399070a10 Mon Sep 17 00:00:00 2001 From: Taras Shapovalov Date: Mon, 22 Jun 2026 18:51:12 -0400 Subject: [PATCH 13/16] Potential fix for pull request finding Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> --- swmcloudgate/routers/azure/connector.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/swmcloudgate/routers/azure/connector.py b/swmcloudgate/routers/azure/connector.py index d9e3514..11d3c1d 100644 --- a/swmcloudgate/routers/azure/connector.py +++ b/swmcloudgate/routers/azure/connector.py @@ -532,7 +532,7 @@ def _get_cloud_init_script( container_registry_password: str, storage_account: str, storage_container: str, - runtime_params: str, + runtime_params: dict[str, str], user_ssh_cert: str, ) -> str: template_loader = jinja2.FileSystemLoader(searchpath="./") From 9fb126a5edee0980d99d3cc7f368ee0b8311061c Mon Sep 17 00:00:00 2001 From: Taras Shapovalov Date: Mon, 22 Jun 2026 18:56:21 -0400 Subject: [PATCH 14/16] Potential fix for pull request finding Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> --- swmcloudgate/routers/azure/templates/cloud-init.sh | 1 - 1 file changed, 1 deletion(-) diff --git a/swmcloudgate/routers/azure/templates/cloud-init.sh b/swmcloudgate/routers/azure/templates/cloud-init.sh index b112785..28ff39e 100755 --- a/swmcloudgate/routers/azure/templates/cloud-init.sh +++ b/swmcloudgate/routers/azure/templates/cloud-init.sh @@ -252,7 +252,6 @@ pull_container_image() { 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" - docker login {{ container_registry | shellquote }} --username {{ container_registry_username | shellquote }} --password {{ container_registry_password | shellquote }} fi echo $(date) ": pull job container image from container registry: $container_image" From b1816a48c53ef3fcf5a850aec441595f19188a2f Mon Sep 17 00:00:00 2001 From: Taras Shapovalov Date: Mon, 22 Jun 2026 18:57:08 -0400 Subject: [PATCH 15/16] Potential fix for pull request finding Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> --- test/test_azure_custom_data.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/test/test_azure_custom_data.py b/test/test_azure_custom_data.py index 661029c..cfea7fc 100644 --- a/test/test_azure_custom_data.py +++ b/test/test_azure_custom_data.py @@ -57,7 +57,11 @@ def test_cloud_init_script_uses_explicit_main_instance_placeholders(self): 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"docker login registry.example.org --username user --password {expected_password}", + 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( From 49ab2f80195a2ddaec52938d34b951ad66d20226 Mon Sep 17 00:00:00 2001 From: Taras Shapovalov Date: Mon, 22 Jun 2026 18:58:43 -0400 Subject: [PATCH 16/16] Apply suggestions from code review Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> --- swmcloudgate/routers/azure/connector.py | 2 +- test/test_azure.py | 9 ++++++--- test/test_openstack.py | 6 ++++-- 3 files changed, 11 insertions(+), 6 deletions(-) diff --git a/swmcloudgate/routers/azure/connector.py b/swmcloudgate/routers/azure/connector.py index 11d3c1d..cf618b5 100644 --- a/swmcloudgate/routers/azure/connector.py +++ b/swmcloudgate/routers/azure/connector.py @@ -76,7 +76,7 @@ def _init_azure_clients( ) raise Exception(msg) - def _parse_vm_count(self, count_param: str) -> int: + 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 diff --git a/test/test_azure.py b/test/test_azure.py index 40d9d75..0b4d5ff 100644 --- a/test/test_azure.py +++ b/test/test_azure.py @@ -46,9 +46,11 @@ async def asyncSetUp(self): self.assertTrue(self.proc.is_alive()) async def asyncTearDown(self): - self.assertTrue(self.proc.is_alive()) - self.proc.terminate() - + 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( @@ -337,6 +339,7 @@ async def test_create_partition_invalid_count(self): 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: diff --git a/test/test_openstack.py b/test/test_openstack.py index e2408b7..d728ac0 100644 --- a/test/test_openstack.py +++ b/test/test_openstack.py @@ -33,8 +33,10 @@ async def asyncSetUp(self): self.assertTrue(self.proc.is_alive()) async def asyncTearDown(self): - self.assertTrue(self.proc.is_alive()) - self.proc.terminate() + 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 = {