76 lines
2.9 KiB
Bash
76 lines
2.9 KiB
Bash
#!/bin/bash
|
|
# query_capacity.sh
|
|
# Queries available capacity for an Azure OpenAI model.
|
|
#
|
|
# Usage:
|
|
# ./query_capacity.sh <model-name> [model-version] [region] [sku]
|
|
# Examples:
|
|
# ./query_capacity.sh o3-mini # List versions
|
|
# ./query_capacity.sh o3-mini 2025-01-31 # All regions
|
|
# ./query_capacity.sh o3-mini 2025-01-31 eastus2 # Specific region
|
|
# ./query_capacity.sh o3-mini 2025-01-31 "" Standard # Different SKU
|
|
|
|
set -euo pipefail
|
|
|
|
MODEL_NAME="${1:?Usage: $0 <model-name> [model-version] [region] [sku]}"
|
|
MODEL_VERSION="${2:-}"
|
|
REGION="${3:-}"
|
|
SKU="${4:-GlobalStandard}"
|
|
|
|
SUB_ID=$(az account show --query id -o tsv)
|
|
|
|
# If no version, list available versions
|
|
if [ -z "$MODEL_VERSION" ]; then
|
|
LOC="${REGION:-eastus}"
|
|
echo "Available versions for $MODEL_NAME:"
|
|
az cognitiveservices model list --location "$LOC" \
|
|
--query "[?model.name=='$MODEL_NAME'].{Version:model.version, Format:model.format}" \
|
|
--output table 2>/dev/null
|
|
exit 0
|
|
fi
|
|
|
|
# Build URL
|
|
if [ -n "$REGION" ]; then
|
|
URL="https://management.azure.com/subscriptions/${SUB_ID}/providers/Microsoft.CognitiveServices/locations/${REGION}/modelCapacities"
|
|
else
|
|
URL="https://management.azure.com/subscriptions/${SUB_ID}/providers/Microsoft.CognitiveServices/modelCapacities"
|
|
fi
|
|
|
|
# Query capacity
|
|
CAPACITY_RESULT=$(az rest --method GET --url "$URL" \
|
|
--url-parameters api-version=2024-10-01 modelFormat=OpenAI modelName="$MODEL_NAME" modelVersion="$MODEL_VERSION" \
|
|
2>/dev/null)
|
|
|
|
# Get regions with capacity
|
|
REGIONS_WITH_CAP=$(echo "$CAPACITY_RESULT" | jq -r ".value[] | select(.properties.skuName==\"$SKU\" and .properties.availableCapacity > 0) | .location" 2>/dev/null | sort -u)
|
|
|
|
if [ -z "$REGIONS_WITH_CAP" ]; then
|
|
echo "No capacity found for $MODEL_NAME v$MODEL_VERSION ($SKU)"
|
|
echo "Try a different SKU or version."
|
|
exit 0
|
|
fi
|
|
|
|
echo "Capacity: $MODEL_NAME v$MODEL_VERSION ($SKU)"
|
|
echo ""
|
|
printf "%-22s %-12s %-15s %s\n" "Region" "Available" "Quota" "SKU"
|
|
printf -- '-%.0s' {1..60}; echo ""
|
|
|
|
for region in $REGIONS_WITH_CAP; do
|
|
avail=$(echo "$CAPACITY_RESULT" | jq -r ".value[] | select(.location==\"$region\" and .properties.skuName==\"$SKU\") | .properties.availableCapacity" 2>/dev/null | head -1)
|
|
|
|
# Check subscription quota
|
|
usage_json=$(az cognitiveservices usage list --location "$region" --subscription "$SUB_ID" -o json 2>/dev/null || echo "[]")
|
|
quota_avail=$(echo "$usage_json" | jq -r --arg name "OpenAI.$SKU.$MODEL_NAME" \
|
|
'[.[] | select(.name.value == $name)] | if length > 0 then .[0].limit - .[0].currentValue else 0 end' 2>/dev/null || echo "?")
|
|
|
|
if [ "$quota_avail" = "0" ]; then
|
|
quota_display="0 (none)"
|
|
elif [ "$quota_avail" = "?" ]; then
|
|
quota_display="?"
|
|
else
|
|
quota_display="${quota_avail}K"
|
|
fi
|
|
|
|
printf "%-22s %-12s %-15s %s\n" "$region" "${avail}K TPM" "$quota_display" "$SKU"
|
|
done
|