From 6c14225310fdbf69c1090c2baee3571d4cbd3db7 Mon Sep 17 00:00:00 2001
From: Daniel Panzella <daniel.panzella@wandb.com>
Date: Tue, 3 Sep 2024 15:28:15 -0700
Subject: [PATCH] fix!: set default availability zones to null and find azs
 that support the instance type

---
 main.tf         | 10 +++++++++-
 variables.tf    |  2 +-
 vmtype_to_az.sh | 32 ++++++++++++++++++++++++++++++++
 3 files changed, 42 insertions(+), 2 deletions(-)
 create mode 100755 vmtype_to_az.sh

diff --git a/main.tf b/main.tf
index 959747d..5b816e4 100644
--- a/main.tf
+++ b/main.tf
@@ -113,6 +113,14 @@ module "app_lb" {
   tags = var.tags
 }
 
+data "external" "az_zones" {
+  program = ["bash", "${path.module}/vmtype_to_az.sh", local.kubernetes_instance_type, azurerm_resource_group.default.location]
+}
+
+locals {
+  node_pool_zones = (var.node_pool_zones == null) ? jsondecode(data.external.az_zones.result.zones) : var.node_pool_zones
+}
+
 module "app_aks" {
   source     = "./modules/app_aks"
   depends_on = [module.app_lb]
@@ -126,7 +134,7 @@ module "app_aks" {
   node_pool_min_vm_count = local.kubernetes_min_node_count
   node_pool_max_vm_count = local.kubernetes_max_node_count
   node_pool_vm_size      = local.kubernetes_instance_type
-  node_pool_zones        = var.node_pool_zones
+  node_pool_zones        = local.node_pool_zones
   public_subnet          = module.networking.public_subnet
   resource_group         = azurerm_resource_group.default
   sku_tier               = var.cluster_sku_tier
diff --git a/variables.tf b/variables.tf
index c7f0120..c339420 100644
--- a/variables.tf
+++ b/variables.tf
@@ -211,7 +211,7 @@ variable "cluster_sku_tier" {
 variable "node_pool_zones" {
   type        = list(string)
   description = "Availability zones for the node pool"
-  default     = ["1", "2"]
+  default     = null
 }
 
 variable "node_max_pods" {
diff --git a/vmtype_to_az.sh b/vmtype_to_az.sh
new file mode 100755
index 0000000..b69a196
--- /dev/null
+++ b/vmtype_to_az.sh
@@ -0,0 +1,32 @@
+#! /usr/bin/env bash
+
+# Given a Azure VM instance type and a region return the availability zones that support the instance type
+
+# Example:
+# ./vmtype_to_az.sh Standard_D2_v3 westeurope
+# 
+# Output:
+# ["1", "2", "3"]
+
+# Copy script arguments to named environment variables
+VM_TYPE="$1"
+REGION="$2"
+
+# Check if both arguments are provided
+if [ -z "$VM_TYPE" ] || [ -z "$REGION" ]; then
+    echo "Error: Both VM type and region must be provided." >&2
+    echo "Usage: $0 <vm_type> <region>" >&2
+    exit 1
+fi
+
+# Query Azure CLI for availability zones in the region for the specified VM type
+ZONES=$(az vm list-skus --location "$REGION" --size "$VM_TYPE" --query "[0].locationInfo[0].zones" -o json | jq -r -c 'sort | .[0:3]')
+
+# Check if the query returned any results
+if [ -z "$ZONES" ] || [ "$ZONES" == "null" ]; then
+    echo "Error: No availability zones found for VM type $VM_TYPE in region $REGION." >&2
+    exit 1
+fi
+
+# Output the result
+jq -n --arg zones "$ZONES" '{"zones":$zones}'