Allow multi node jobs

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)}

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/tcp/${MAIN_INSTANCE_PRIVATE_IP}/2049" >/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,28 +187,38 @@ 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
         echo $(date) ": /etc/fstab:"
         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

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:

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')]"
           }
         ]

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)