1#!/bin/bash
2# query_capacity.sh
3# Queries available capacity for an Azure OpenAI model.
4#
5# Usage:
6#   ./query_capacity.sh <model-name> [model-version] [region] [sku]
7# Examples:
8#   ./query_capacity.sh o3-mini                          # List versions
9#   ./query_capacity.sh o3-mini 2025-01-31               # All regions
10#   ./query_capacity.sh o3-mini 2025-01-31 eastus2       # Specific region
11#   ./query_capacity.sh o3-mini 2025-01-31 "" Standard   # Different SKU
12 
13set -euo pipefail
14 
15MODEL_NAME="${1:?Usage: $0 <model-name> [model-version] [region] [sku]}"
16MODEL_VERSION="${2:-}"
17REGION="${3:-}"
18SKU="${4:-GlobalStandard}"
19 
20SUB_ID=$(az account show --query id -o tsv)
21 
22# If no version, list available versions
23if [ -z "$MODEL_VERSION" ]; then
24    LOC="${REGION:-eastus}"
25    echo "Available versions for $MODEL_NAME:"
26    az cognitiveservices model list --location "$LOC" \
27        --query "[?model.name=='$MODEL_NAME'].{Version:model.version, Format:model.format}" \
28        --output table 2>/dev/null
29    exit 0
30fi
31 
32# Build URL
33if [ -n "$REGION" ]; then
34    URL="https://management.azure.com/subscriptions/${SUB_ID}/providers/Microsoft.CognitiveServices/locations/${REGION}/modelCapacities"
35else
36    URL="https://management.azure.com/subscriptions/${SUB_ID}/providers/Microsoft.CognitiveServices/modelCapacities"
37fi
38 
39# Query capacity
40CAPACITY_RESULT=$(az rest --method GET --url "$URL" \
41    --url-parameters api-version=2024-10-01 modelFormat=OpenAI modelName="$MODEL_NAME" modelVersion="$MODEL_VERSION" \
42    2>/dev/null)
43 
44# Get regions with capacity
45REGIONS_WITH_CAP=$(echo "$CAPACITY_RESULT" | jq -r ".value[] | select(.properties.skuName==\"$SKU\" and .properties.availableCapacity > 0) | .location" 2>/dev/null | sort -u)
46 
47if [ -z "$REGIONS_WITH_CAP" ]; then
48    echo "No capacity found for $MODEL_NAME v$MODEL_VERSION ($SKU)"
49    echo "Try a different SKU or version."
50    exit 0
51fi
52 
53echo "Capacity: $MODEL_NAME v$MODEL_VERSION ($SKU)"
54echo ""
55printf "%-22s %-12s %-15s %s\n" "Region" "Available" "Quota" "SKU"
56printf -- '-%.0s' {1..60}; echo ""
57 
58for region in $REGIONS_WITH_CAP; do
59    avail=$(echo "$CAPACITY_RESULT" | jq -r ".value[] | select(.location==\"$region\" and .properties.skuName==\"$SKU\") | .properties.availableCapacity" 2>/dev/null | head -1)
60 
61    # Check subscription quota
62    usage_json=$(az cognitiveservices usage list --location "$region" --subscription "$SUB_ID" -o json 2>/dev/null || echo "[]")
63    quota_avail=$(echo "$usage_json" | jq -r --arg name "OpenAI.$SKU.$MODEL_NAME" \
64        '[.[] | select(.name.value == $name)] | if length > 0 then .[0].limit - .[0].currentValue else 0 end' 2>/dev/null || echo "?")
65 
66    if [ "$quota_avail" = "0" ]; then
67        quota_display="0 (none)"
68    elif [ "$quota_avail" = "?" ]; then
69        quota_display="?"
70    else
71        quota_display="${quota_avail}K"
72    fi
73 
74    printf "%-22s %-12s %-15s %s\n" "$region" "${avail}K TPM" "$quota_display" "$SKU"
75done
76