Merge pull request #128 from microsoft/feature/user-interactive-quota-check

Roopan-Microsoft · web-flow · commit 900ba0e73d92 · 2025-06-12T18:03:50.000+05:30
feat: implement interactive quota check script with user inputs
diff --git a/azure.yaml b/azure.yaml
@@ -27,13 +27,13 @@ hooks:
       shell: sh
       run: >
         chmod u+r+x ./scripts/validate_model_deployment_quota.sh; chmod u+r+x ./scripts/validate_model_quota.sh; ./scripts/validate_model_deployment_quota.sh --SubscriptionId "$AZURE_SUBSCRIPTION_ID" --Location "${AZURE_AISERVICE_LOCATION:-japaneast}" --ModelsParameter "aiModelDeployments"
-      interactive: false
+      interactive: true
       continueOnError: false
 
     windows:
       shell: pwsh
       run: >
         $location = if ($env:AZURE_AISERVICE_LOCATION) { $env:AZURE_AISERVICE_LOCATION } else { "japaneast" };
         ./scripts/validate_model_deployment_quota.ps1 -SubscriptionId $env:AZURE_SUBSCRIPTION_ID -Location $location -ModelsParameter "aiModelDeployments"
-      interactive: false
+      interactive: true
       continueOnError: false
diff --git a/scripts/validate_model_deployment_quota.ps1 b/scripts/validate_model_deployment_quota.ps1
@@ -15,58 +15,64 @@ if (-not $ModelsParameter) { $MissingParams += "ModelsParameter" }
 
 if ($MissingParams.Count -gt 0) {
     Write-Error "❌ ERROR: Missing required parameters: $($MissingParams -join ', ')"
+    Write-Host "Usage: validate_model_deployment_quota.ps1 -SubscriptionId <SUBSCRIPTION_ID> -Location <LOCATION> -ModelsParameter <MODELS_PARAMETER>"
     exit 1
 }
 
 # Load model deployments from parameter file
 $JsonContent = Get-Content -Path "./infra/main.parameters.json" -Raw | ConvertFrom-Json
-if (-not $JsonContent) {
-    Write-Error "❌ ERROR: Failed to parse main.parameters.json. Ensure the JSON file is valid."
-    exit 1
-}
-
 $aiModelDeployments = $JsonContent.parameters.$ModelsParameter.value
 if (-not $aiModelDeployments -or -not ($aiModelDeployments -is [System.Collections.IEnumerable])) {
-    Write-Error "❌ ERROR: The specified property '$ModelsParameter' does not exist or is not an array."
+    Write-Error "❌ ERROR: Failed to parse main.parameters.json or missing '$ModelsParameter'"
     exit 1
 }
 
-# Check if AI Foundry exists and has all required model deployments
-$existing = $null
+# Try to discover AI Foundry name if not set
+if (-not $AiFoundryName -and $ResourceGroup) {
+    $AiFoundryName = az cognitiveservices account list `
+        --resource-group $ResourceGroup `
+        --query "sort_by([?kind=='AIServices'], &name)[0].name" `
+        -o tsv 2>$null
+}
+
+# Check if AI Foundry exists
 if ($AiFoundryName -and $ResourceGroup) {
     $existing = az cognitiveservices account show `
         --name $AiFoundryName `
         --resource-group $ResourceGroup `
         --query "name" --output tsv 2>$null
-}
 
-if ($existing) {
-    $deployedModelsOutput = az cognitiveservices account deployment list `
-        --name $AiFoundryName `
-        --resource-group $ResourceGroup `
-        --query "[].name" --output tsv 2>$null
+    if ($existing) {
+        # adding into .env
+        azd env set AZURE_AIFOUNDRY_NAME $existing | Out-Null
 
-    $deployedModels = @()
-    if ($deployedModelsOutput -is [string]) {
-        $deployedModels += $deployedModelsOutput
-    } elseif ($deployedModelsOutput) {
-        $deployedModels = $deployedModelsOutput -split "`r?`n"
-    }
+        $deployedModelsOutput = az cognitiveservices account deployment list `
+            --name $AiFoundryName `
+            --resource-group $ResourceGroup `
+            --query "[].name" --output tsv 2>$null
 
-    $requiredDeployments = $aiModelDeployments | ForEach-Object { $_.name }
-    $missingDeployments = $requiredDeployments | Where-Object { $_ -notin $deployedModels }
+        $deployedModels = @()
+        if ($deployedModelsOutput -is [string]) {
+            $deployedModels += $deployedModelsOutput
+        } elseif ($deployedModelsOutput) {
+            $deployedModels = $deployedModelsOutput -split "`r?`n"
+        }
+
+        $requiredDeployments = $aiModelDeployments | ForEach-Object { $_.name }
+        $missingDeployments = $requiredDeployments | Where-Object { $_ -notin $deployedModels }
 
-    if ($missingDeployments.Count -eq 0) {
-        Write-Host "ℹ️ AI Foundry '$AiFoundryName' exists and all required model deployments are already provisioned."
-        Write-Host "⏭️ Skipping quota validation."
-        exit 0
-    } else {
-        Write-Host "🔍 AI Foundry exists, but the following model deployments are missing: $($missingDeployments -join ', ')"
-        Write-Host "➡️ Proceeding with quota validation for missing models..."
+        if ($missingDeployments.Count -eq 0) {
+            Write-Host "ℹ️ AI Foundry '$AiFoundryName' exists and all required model deployments are already provisioned."
+            Write-Host "⏭️ Skipping quota validation."
+            exit 0
+        } else {
+            Write-Host "🔍 AI Foundry exists, but the following model deployments are missing: $($missingDeployments -join ', ')"
+            Write-Host "➡️ Proceeding with quota validation for missing models..."
+        }
     }
 }
 
-# Run quota validation for all models
+# Run quota validation
 az account set --subscription $SubscriptionId
 Write-Host "🎯 Active Subscription: $(az account show --query '[name, id]' --output tsv)"
 
@@ -78,7 +84,8 @@ foreach ($deployment in $aiModelDeployments) {
     $type = if ($env:AZURE_ENV_MODEL_DEPLOYMENT_TYPE) { $env:AZURE_ENV_MODEL_DEPLOYMENT_TYPE } else { $deployment.sku.name }
     $capacity = if ($env:AZURE_ENV_MODEL_CAPACITY) { $env:AZURE_ENV_MODEL_CAPACITY } else { $deployment.sku.capacity }
 
-    Write-Host "`n🔍 Validating model deployment: $name ..."
+    Write-Host ""
+    Write-Host "🔍 Validating model deployment: $name ..."
     & .\scripts\validate_model_quota.ps1 -Location $Location -Model $model -Capacity $capacity -DeploymentType $type
     $exitCode = $LASTEXITCODE
 
diff --git a/scripts/validate_model_deployment_quota.sh b/scripts/validate_model_deployment_quota.sh
@@ -47,34 +47,44 @@ if [[ $? -ne 0 || -z "$aiModelDeployments" ]]; then
   exit 1
 fi
 
-# Check if AI Foundry exists and has all required model deployments
-existing=""
-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)
+# 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
 
-if [[ -n "$existing" ]]; then
-  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")
+# 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
-  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
 
diff --git a/scripts/validate_model_quota.ps1 b/scripts/validate_model_quota.ps1
@@ -10,7 +10,6 @@ $MissingParams = @()
 if (-not $Location) { $MissingParams += "location" }
 if (-not $Model) { $MissingParams += "model" }
 if (-not $Capacity) { $MissingParams += "capacity" }
-if (-not $DeploymentType) { $MissingParams += "deployment-type" }
 
 if ($MissingParams.Count -gt 0) {
     Write-Error "❌ ERROR: Missing required parameters: $($MissingParams -join ', ')"
@@ -35,7 +34,7 @@ function Check-Quota {
     try {
         $ModelInfoRaw = az cognitiveservices usage list --location $Region --query "[?name.value=='$ModelType']" --output json
         $ModelInfo = $ModelInfoRaw | ConvertFrom-Json
-        if (-not $ModelInfo) { return }
+        if (-not $ModelInfo) { return $null }
 
         $CurrentValue = ($ModelInfo | Where-Object { $_.name.value -eq $ModelType }).currentValue
         $Limit = ($ModelInfo | Where-Object { $_.name.value -eq $ModelType }).limit
@@ -52,11 +51,23 @@ function Check-Quota {
             Available = $Available
         }
     } catch {
-        return
+        return $null
     }
 }
 
-# First, check the user-specified region
+function Show-Table {
+    Write-Host "`n--------------------------------------------------------------------------------------------"
+    Write-Host "| No. | Region         | Model Name                          | Limit | Used  | Available |"
+    Write-Host "--------------------------------------------------------------------------------------------"
+    $count = 1
+    foreach ($entry in $AllResults) {
+        Write-Host ("| {0,-3} | {1,-14} | {2,-35} | {3,-5} | {4,-5} | {5,-9} |" -f $count, $entry.Region, $entry.Model, $entry.Limit, $entry.Used, $entry.Available)
+        $count++
+    }
+    Write-Host "--------------------------------------------------------------------------------------------"
+}
+
+# ----------- First check the user-specified region -----------
 Write-Host "`n🔍 Checking quota in the requested region '$Location'..."
 $PrimaryResult = Check-Quota -Region $Location
 
@@ -72,45 +83,53 @@ if ($PrimaryResult) {
     Write-Host "`n⚠️  Could not retrieve quota info for region '$Location'. Checking fallback regions..."
 }
 
-# Remove primary region from fallback list
+# ----------- Check all other fallback regions -----------
 $FallbackRegions = $PreferredRegions | Where-Object { $_ -ne $Location }
+$EligibleFallbacks = @()
 
 foreach ($region in $FallbackRegions) {
     $result = Check-Quota -Region $region
     if ($result) {
         $AllResults += $result
+        if ($result.Available -ge $Capacity) {
+            $EligibleFallbacks += $result
+        }
     }
 }
 
-# Display Results Table
-Write-Host "`n-------------------------------------------------------------------------------------------------------------"
-Write-Host "| No.  | Region            | Model Name                           | Limit   | Used    | Available |"
-Write-Host "-------------------------------------------------------------------------------------------------------------"
-
-$count = 1
-foreach ($entry in $AllResults) {
-    $modelShort = $entry.Model.Substring($entry.Model.LastIndexOf(".") + 1)
-    Write-Host ("| {0,-4} | {1,-16} | {2,-35} | {3,-7} | {4,-7} | {5,-9} |" -f $count, $entry.Region, $entry.Model, $entry.Limit, $entry.Used, $entry.Available)
-    $count++
-}
-Write-Host "-------------------------------------------------------------------------------------------------------------"
-
-# Suggest fallback regions
-$EligibleFallbacks = $AllResults | Where-Object { $_.Region -ne $Location -and $_.Available -ge $Capacity }
+# ----------- Show Table of All Regions Checked -----------
+$AllResults = $AllResults | Where-Object { $_.Available -gt 50 }
+Show-Table
 
+# ----------- If eligible fallback regions found, ask user -----------
 if ($EligibleFallbacks.Count -gt 0) {
     Write-Host "`n❌ Deployment cannot proceed in '$Location'."
-    Write-Host "➡️ You can retry using one of the following regions with sufficient quota:`n"
-    foreach ($region in $EligibleFallbacks) {
-        Write-Host "   • $($region.Region) (Available: $($region.Available))"
-    }
+    Write-Host "➡️  Found fallback regions with sufficient quota."
+    
+    while ($true) {
+        Write-Host "`nPlease enter a fallback region from the list above to proceed:"
+        $NewLocation = Read-Host "Enter region"
+
+        if (-not $NewLocation) {
+            Write-Host "❌ No location entered. Exiting."
+            exit 1
+        }
+
+        $UserResult = Check-Quota -Region $NewLocation
+        if (-not $UserResult) {
+            Write-Host "⚠️ Could not retrieve quota info for region '$NewLocation'. Try again."
+            continue
+        }
 
-    Write-Host "`n🔧 To proceed, run:"
-    Write-Host "    azd env set AZURE_AISERVICE_LOCATION '<region>'"
-    Write-Host "📌 To confirm it's set correctly, run:"
-    Write-Host "    azd env get-value AZURE_AISERVICE_LOCATION"
-    Write-Host "▶️  Once confirmed, re-run azd up to deploy the model in the new region."
-    exit 2
+        if ($UserResult.Available -ge $Capacity) {
+            Write-Host "✅ Sufficient quota found in '$NewLocation'. Proceeding with deployment."
+            azd env set AZURE_AISERVICE_LOCATION "$NewLocation" | Out-Null
+            Write-Host "➡️  Set AZURE_AISERVICE_LOCATION to '$NewLocation'."
+            exit 0
+        } else {
+            Write-Host "❌ Insufficient quota in '$NewLocation'. Try another."
+        }
+    }
 }
 
 Write-Error "`n❌ ERROR: No available quota found in any region."
diff --git a/scripts/validate_model_quota.sh b/scripts/validate_model_quota.sh
