From 5769effd8db973baabdc9a69f286c13d3aaa4200 Mon Sep 17 00:00:00 2001
From: Dan McPherson <dmcphers@redhat.com>
Date: Thu, 27 Feb 2025 14:46:02 -0500
Subject: [PATCH] Add option to search multiple subnets with launch

Signed-off-by: Dan McPherson <dmcphers@redhat.com>
---
 scripts/infra/cloud-instance.sh               | 38 ++++++++++++-------
 .../infra/instructlab-cloud-config-example    |  4 ++
 2 files changed, 28 insertions(+), 14 deletions(-)

diff --git a/scripts/infra/cloud-instance.sh b/scripts/infra/cloud-instance.sh
index 0ed10d1e04..53b021a3be 100755
--- a/scripts/infra/cloud-instance.sh
+++ b/scripts/infra/cloud-instance.sh
@@ -15,20 +15,30 @@ else
 fi
 
 ec2__launch() {
-    local instance_id
-    instance_id="$(aws ec2 run-instances \
-        --image-id "$EC2_AMI_ID" \
-        --region "$EC2_REGION" \
-        --instance-type "${INSTANCE_TYPE:-$EC2_INSTANCE_TYPE}" \
-        --security-group-ids "$EC2_SECURITY_GROUP_ID" \
-        --subnet-id "$EC2_SUBNET_ID" \
-        --key-name "$EC2_KEY_NAME" \
-        --block-device-mappings '{"DeviceName": "/dev/sda1","Ebs": {"VolumeSize": 800}}' \
-        --associate-public-ip-address \
-        --tag-specifications "ResourceType=instance,Tags=[{Key=Name,Value=$INSTANCE_NAME}]" \
-        "ResourceType=volume,Tags=[{Key=Name,Value=$INSTANCE_NAME}]" \
-        --query 'Instances[0].InstanceId' \
-        --output text)" &> /dev/null
+    if [ -z ${EC2_SUBNET_VARS+x} ]; then
+        EC2_SUBNET_VARS=("EC2_SUBNET_ID")
+    fi
+
+    for subnet_var in "${EC2_SUBNET_VARS[@]}"; do
+        subnet="${!subnet_var}"
+        echo "Attempting to launch an instance of type $EC2_INSTANCE_TYPE in subnet $subnet named $INSTANCE_NAME"
+        local instance_id
+        if instance_id="$(aws ec2 run-instances \
+            --image-id "$EC2_AMI_ID" \
+            --region "$EC2_REGION" \
+            --instance-type "${INSTANCE_TYPE:-$EC2_INSTANCE_TYPE}" \
+            --security-group-ids "$EC2_SECURITY_GROUP_ID" \
+            --subnet-id "$subnet" \
+            --key-name "$EC2_KEY_NAME" \
+            --block-device-mappings '{"DeviceName": "/dev/sda1","Ebs": {"VolumeSize": 800}}' \
+            --associate-public-ip-address \
+            --tag-specifications "ResourceType=instance,Tags=[{Key=Name,Value=$INSTANCE_NAME}]" \
+            "ResourceType=volume,Tags=[{Key=Name,Value=$INSTANCE_NAME}]" \
+            --query 'Instances[0].InstanceId' \
+            --output text)" &> /dev/null; then
+            break
+        fi
+    done
 
     if [ -z "$instance_id" ]; then
         echo "Failed to launch instance"
diff --git a/scripts/infra/instructlab-cloud-config-example b/scripts/infra/instructlab-cloud-config-example
index c1dac2b612..fb3d9d6a31 100644
--- a/scripts/infra/instructlab-cloud-config-example
+++ b/scripts/infra/instructlab-cloud-config-example
@@ -6,7 +6,11 @@ EC2_INSTANCE_TYPE="m4.large"
 # Nvidia 1x A10 GPU, 24GB VRAM
 #EC2_INSTANCE_TYPE="g5.4xlarge"
 EC2_SECURITY_GROUP_ID="sg-12345678901234567"
+# EC2_SUBNET_ID is ignored if you specify EC2_SUBNET_VARS which searches across the specified subnets for an available instance on launch
 EC2_SUBNET_ID="subnet-12345678901234567"
+EC2_SUBNET_VARS=("EC2_SUBNET_ID_1" "EC2_SUBNET_ID_2")
+EC2_SUBNET_ID_1=subnet-12345678901234567
+EC2_SUBNET_ID_2=subnet-12345678901234568
 EC2_AMI_ID="ami-12345678901234567"
 EC2_REGION="us-east-2"
 EC2_INSTANCE_NAME=`whoami`