validate_model_deployment_quota.sh

#!/bin/bash

SUBSCRIPTION_ID=""
LOCATION=""
MODELS_PARAMETER=""

while [[ $# -gt 0 ]]; do
  case "$1" in
    --SubscriptionId)
      SUBSCRIPTION_ID="$2"
      shift 2
      ;;
    --Location)
      LOCATION="$2"
      shift 2
      ;;
    --ModelsParameter)
      MODELS_PARAMETER="$2"
      shift 2
      ;;
    *)
      echo "❌ ERROR: Unknown option: $1"
      exit 1
      ;;
  esac
done

AIFOUNDRY_NAME="${AZURE_AIFOUNDRY_NAME}"
RESOURCE_GROUP="${AZURE_RESOURCE_GROUP}"

# Validate required parameters
MISSING_PARAMS=()
[[ -z "$SUBSCRIPTION_ID" ]] && MISSING_PARAMS+=("SubscriptionId")
[[ -z "$LOCATION" ]] && MISSING_PARAMS+=("Location")
[[ -z "$MODELS_PARAMETER" ]] && MISSING_PARAMS+=("ModelsParameter")

if [[ ${#MISSING_PARAMS[@]} -ne 0 ]]; then
  echo "❌ ERROR: Missing required parameters: ${MISSING_PARAMS[*]}"
  echo "Usage: $0 --SubscriptionId <SUBSCRIPTION_ID> --Location <LOCATION> --ModelsParameter <MODELS_PARAMETER>"
  exit 1
fi

# Check if user is logged in to Azure
if ! az account show > /dev/null 2>&1; then
  echo "❌ ERROR: You are not logged in to Azure CLI."
  echo "👉 Please run 'az login' to continue."
  exit 1
fi

# Load model definitions
aiModelDeployments=$(jq -c ".parameters.$MODELS_PARAMETER.value[]" ./infra/main.parameters.json 2>/dev/null)
if [[ $? -ne 0 || -z "$aiModelDeployments" ]]; then
  echo "❌ ERROR: Failed to parse main.parameters.json or missing '$MODELS_PARAMETER'"
  exit 1
fi

# Try to discover AI Foundry name if not set
if [[ -z "$AIFOUNDRY_NAME" && -n "$RESOURCE_GROUP" ]]; then
  AIFOUNDRY_NAME=$(az cognitiveservices account list --resource-group "$RESOURCE_GROUP" \
    --query "sort_by([?kind=='AIServices'], &name)[0].name" -o tsv 2>/dev/null)
fi

# Check if AI Foundry exists
if [[ -n "$AIFOUNDRY_NAME" && -n "$RESOURCE_GROUP" ]]; then
  existing=$(az cognitiveservices account show --name "$AIFOUNDRY_NAME" \
    --resource-group "$RESOURCE_GROUP" --query "name" --output tsv 2>/dev/null)

  if [[ -n "$existing" ]]; then
    # adding into .env
    azd env set AZURE_AIFOUNDRY_NAME "$existing" > /dev/null

    # Check model deployments
    existing_deployments=$(az cognitiveservices account deployment list \
      --name "$AIFOUNDRY_NAME" \
      --resource-group "$RESOURCE_GROUP" \
      --query "[].name" --output tsv 2>/dev/null)

    required_models=$(jq -r ".parameters.$MODELS_PARAMETER.value[].name" ./infra/main.parameters.json)

    missing_models=()
    for model in $required_models; do
      if ! grep -q -w "$model" <<< "$existing_deployments"; then
        missing_models+=("$model")
      fi
    done

    if [[ ${#missing_models[@]} -eq 0 ]]; then
      echo "ℹ️ AI Foundry '$AIFOUNDRY_NAME' exists and all required model deployments are already provisioned."
      echo "⏭️ Skipping quota validation."
      exit 0
    else
      echo "🔍 AI Foundry exists, but the following model deployments are missing: ${missing_models[*]}"
      echo "➡️ Proceeding with quota validation for missing models..."
    fi
  fi
fi

# Run quota validation
az account set --subscription "$SUBSCRIPTION_ID"
echo "🎯 Active Subscription: $(az account show --query '[name, id]' --output tsv)"

quotaAvailable=true

while IFS= read -r deployment; do
  name=${AZURE_ENV_MODEL_NAME:-$(echo "$deployment" | jq -r '.name')}
  model=${AZURE_ENV_MODEL_NAME:-$(echo "$deployment" | jq -r '.model.name')}
  type=${AZURE_ENV_MODEL_DEPLOYMENT_TYPE:-$(echo "$deployment" | jq -r '.sku.name')}
  capacity=${AZURE_ENV_MODEL_CAPACITY:-$(echo "$deployment" | jq -r '.sku.capacity')}

  echo -e "🔍 Validation started for model deployment: \e[1m$name\e[0m in region \e[1m$LOCATION\e[0m with capacity \e[1m$capacity\e[0m tokens..."
  ./scripts/validate_model_quota.sh --location "$LOCATION" --model "$model" --capacity "$capacity" --deployment-type "$type"
  exit_code=$?

  if [[ $exit_code -ne 0 ]]; then
    if [[ $exit_code -eq 2 ]]; then
      exit 1
    fi
    echo -e "\n❌ ERROR: Quota validation failed for model deployment: \033[1m$name\033[0m"
    quotaAvailable=false
  fi
done <<< "$(echo "$aiModelDeployments")"

if [[ "$quotaAvailable" = false ]]; then
  echo -e "\n❌ ERROR: Quota validation failed — insufficient quota in all regions. Deployment cannot proceed."
  echo -e "\nℹ️  Please request a quota increase at https://portal.azure.com/#blade/Microsoft_Azure_Capacity/UsageAndQuota and try again."
  exit 1
else
  echo "✅ All model deployments passed quota validation successfully."
  exit 0
fi