Skip to content
5 changes: 4 additions & 1 deletion cli/azd/grpc/proto/ai_model.proto
Original file line number Diff line number Diff line change
Expand Up @@ -194,7 +194,10 @@ message ListLocationsWithQuotaResponse {
message ModelLocationQuota {
// Location where model quota was evaluated.
Location location = 1;
double max_remaining_quota = 2; // max remaining quota across model SKUs
// Maximum remaining quota across model SKUs.
// A value of -1 indicates that usage data was unavailable (e.g. free-tier
// subscriptions) and the actual remaining quota is unknown.
double max_remaining_quota = 2;
}

message ListModelLocationsWithQuotaRequest {
Expand Down
2 changes: 2 additions & 0 deletions cli/azd/grpc/proto/prompt.proto
Original file line number Diff line number Diff line change
Expand Up @@ -301,5 +301,7 @@ message PromptAiModelLocationWithQuotaResponse {
// Selected location.
Location location = 1;
// Maximum remaining quota at the selected location across model SKUs.
// A value of -1 indicates that usage data was unavailable (e.g. free-tier
// subscriptions) and the actual remaining quota is unknown.
double max_remaining_quota = 2;
}
11 changes: 9 additions & 2 deletions cli/azd/internal/grpcserver/prompt_service.go
Original file line number Diff line number Diff line change
Expand Up @@ -1197,8 +1197,15 @@ func (s *promptService) PromptAiModelLocationWithQuota(
EnableFiltering: new(true),
}
for i, loc := range locations {
quotaLabel := output.WithGrayFormat("[up to %.0f quota available]", loc.MaxRemainingQuota)
label := fmt.Sprintf("%s %s", loc.Location, quotaLabel)
var label string
if loc.MaxRemainingQuota == ai.QuotaRemainingUnknown {
label = loc.Location
} else {
quotaLabel := output.WithGrayFormat(
"[up to %.0f quota available]",
loc.MaxRemainingQuota)
label = fmt.Sprintf("%s %s", loc.Location, quotaLabel)
}
Comment thread
vhvb1989 marked this conversation as resolved.
selectOpts.Choices[i] = &ux.SelectChoice{
Value: loc.Location,
Label: label,
Expand Down
4 changes: 2 additions & 2 deletions cli/azd/pkg/ai/model_helpers_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -305,10 +305,10 @@ func TestModelHasQuota(t *testing.T) {
expected: true,
},
{
name: "no usage entries for model",
name: "no usage entries for model assumes available",
usageMap: map[string]AiModelUsage{},
minRemaining: 1,
expected: false,
expected: true,
},
{
name: "remaining exactly equals min",
Expand Down
66 changes: 60 additions & 6 deletions cli/azd/pkg/ai/model_service.go
Original file line number Diff line number Diff line change
Expand Up @@ -250,6 +250,16 @@ func (s *AiModelService) ListLocationsWithQuota(

var results []string
sharedResults.Range(func(loc string, usages []*armcognitiveservices.Usage) bool {
// When the /usages API returns an empty list (e.g. free-tier subscriptions
// that have not yet provisioned Cognitive Services resources), treat the
// location as having full quota available. The AI Services account SKU
// (AIServices/S0) was already confirmed available in this region; empty
// usages means no consumption data exists, not that quota is zero.
if len(usages) == 0 {
results = append(results, loc)
return true
}

for _, req := range requirements {
minCap := req.MinCapacity
if minCap <= 0 {
Expand Down Expand Up @@ -332,8 +342,15 @@ func (s *AiModelService) ListModelLocationsWithQuota(
usageMap[usage.Name] = usage
}

maxRemainingAtLocation, found := maxModelRemainingQuota(*targetModel, usageMap)
if found && maxRemainingAtLocation >= minRemaining {
maxRemainingAtLocation, found := maxModelRemainingQuota(
*targetModel, usageMap)
// Include the location when the model has at least one
// deployable SKU and either: (a) usage data confirms
// sufficient remaining quota, or (b) usage data is
// unavailable (e.g. free-tier subscriptions).
if found &&
(maxRemainingAtLocation == QuotaRemainingUnknown ||
maxRemainingAtLocation >= minRemaining) {
Comment thread
vhvb1989 marked this conversation as resolved.
results = append(results, ModelLocationQuota{
Location: loc,
MaxRemainingQuota: maxRemainingAtLocation,
Expand Down Expand Up @@ -469,9 +486,10 @@ func (s *AiModelService) resolveDeployments(
continue
}

// Quota check
// Quota check — skip when usage data is empty (e.g. free-tier
// subscriptions where the /usages API returns no entries).
capacity := ResolveCapacity(sku, options.Capacity)
if quotaOpts != nil && usageMap != nil {
if quotaOpts != nil && usageMap != nil && len(usageMap) > 0 {
usage, ok := usageMap[sku.UsageName]
Comment thread
vhvb1989 marked this conversation as resolved.
if !ok {
continue
Expand Down Expand Up @@ -953,6 +971,17 @@ func ModelHasDefaultVersion(model AiModel) bool {
}

func modelHasQuota(model AiModel, usageMap map[string]AiModelUsage, minRemaining float64) bool {
// When usage data is empty (e.g. free-tier subscriptions), assume the
// model is eligible as long as it has at least one deployable SKU.
if len(usageMap) == 0 {
for _, version := range model.Versions {
if len(version.Skus) > 0 {
return true
}
}
return false
}

for _, version := range model.Versions {
for _, sku := range version.Skus {
usage, ok := usageMap[sku.UsageName]
Expand All @@ -971,6 +1000,19 @@ func modelHasQuota(model AiModel, usageMap map[string]AiModelUsage, minRemaining
}

func maxModelRemainingQuota(model AiModel, usageMap map[string]AiModelUsage) (float64, bool) {
// When usage data is empty (e.g. free-tier subscriptions), treat the
// model as available if it has at least one SKU. Return
// QuotaRemainingUnknown to signal that the actual remaining quota
// is unknown.
if len(usageMap) == 0 {
Comment thread
vhvb1989 marked this conversation as resolved.
for _, version := range model.Versions {
if len(version.Skus) > 0 {
return QuotaRemainingUnknown, true
}
}
return 0, false
Comment thread
vhvb1989 marked this conversation as resolved.
}

var maxRemaining float64
found := false
for _, version := range model.Versions {
Expand Down Expand Up @@ -1019,8 +1061,20 @@ func filterModelsByAnyLocationQuota(
) []AiModel {
eligible := map[string]struct{}{}

for _, usages := range usagesByLocation {
for _, model := range FilterModelsByQuota(models, usages, minRemaining) {
for loc, usages := range usagesByLocation {
// Only consider models that are actually available in this
// location. Without this filter, empty usages from an
// unrelated location could mark a model as eligible even
// when its quota is exhausted in its actual location.
locModels := make([]AiModel, 0, len(models))
for _, m := range models {
if slices.Contains(m.Locations, loc) {
locModels = append(locModels, m)
}
}

for _, model := range FilterModelsByQuota(
locModels, usages, minRemaining) {
eligible[model.Name] = struct{}{}
}
}
Expand Down
110 changes: 108 additions & 2 deletions cli/azd/pkg/ai/model_service_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -482,10 +482,10 @@ func TestFilterModelsByQuota(t *testing.T) {
expected: []string{"gpt-4o"},
},
{
name: "empty usages excludes all models",
name: "empty usages includes models with SKUs",
usages: []AiModelUsage{},
minRemaining: 1,
expected: []string{},
expected: []string{"gpt-4o", "gpt-4o-mini", "text-embedding-ada-002"},
},
}

Expand Down Expand Up @@ -775,6 +775,61 @@ func TestFilterModelsByAnyLocationQuota(t *testing.T) {
require.Equal(t, []string{"model-a", "model-b"}, filteredNames)
}

func TestFilterModelsByAnyLocationQuota_EmptyUsagesDoNotLeakAcrossLocations(t *testing.T) {
// Regression test: a model available only in location A with exhausted
// quota should NOT become eligible because location B returns empty
// usages (e.g. free-tier subscription).
models := []AiModel{
{
Name: "model-a",
Locations: []string{"eastus"},
Versions: []AiModelVersion{
{
Version: "1",
Skus: []AiModelSku{
{Name: "Standard", UsageName: "a_usage"},
},
},
},
},
{
Name: "model-b",
Locations: []string{"eastus", "westus"},
Versions: []AiModelVersion{
{
Version: "1",
Skus: []AiModelSku{
{Name: "Standard", UsageName: "b_usage"},
},
},
},
},
}

usagesByLocation := map[string][]AiModelUsage{
// eastus: model-a exhausted, model-b has quota
"eastus": {
{Name: "a_usage", CurrentValue: 10, Limit: 10},
{Name: "b_usage", CurrentValue: 0, Limit: 100},
},
// westus: empty usages (free-tier). model-a is NOT
// available here, so it must not become eligible.
"westus": {},
}

filtered := filterModelsByAnyLocationQuota(
models, usagesByLocation, 1)
names := make([]string, 0, len(filtered))
for _, m := range filtered {
names = append(names, m.Name)
}

// model-a should be excluded (exhausted in eastus,
// not available in westus).
// model-b should be included (has quota in eastus).
require.Equal(t, []string{"model-b"}, names)
}

func TestIsFinetuneUsageName(t *testing.T) {
tests := []struct {
usageName string
Expand All @@ -795,3 +850,54 @@ func TestIsFinetuneUsageName(t *testing.T) {
})
}
}

func TestModelHasQuota_EmptyUsages(t *testing.T) {
modelWithSkus := AiModel{
Name: "gpt-4o",
Versions: []AiModelVersion{
{
Version: "2024-05-13",
Skus: []AiModelSku{{Name: "Standard", UsageName: "OpenAI.Standard.gpt-4o"}},
},
},
}
modelNoSkus := AiModel{
Name: "empty-model",
Versions: []AiModelVersion{{Version: "1.0"}},
}

emptyUsages := map[string]AiModelUsage{}

// Model with SKUs should be considered eligible when usage data is empty
require.True(t, modelHasQuota(modelWithSkus, emptyUsages, 1))

// Model with no SKUs should not be considered eligible even with empty usage data
require.False(t, modelHasQuota(modelNoSkus, emptyUsages, 1))
}

func TestMaxModelRemainingQuota_EmptyUsages(t *testing.T) {
modelWithSkus := AiModel{
Name: "gpt-4o",
Versions: []AiModelVersion{
{
Version: "2024-05-13",
Skus: []AiModelSku{{Name: "Standard", UsageName: "OpenAI.Standard.gpt-4o"}},
},
},
}
modelNoSkus := AiModel{
Name: "empty-model",
Versions: []AiModelVersion{{Version: "1.0"}},
}

emptyUsages := map[string]AiModelUsage{}

// Model with SKUs: found=true, remaining=QuotaRemainingUnknown
remaining, found := maxModelRemainingQuota(modelWithSkus, emptyUsages)
require.True(t, found)
require.Equal(t, QuotaRemainingUnknown, remaining)

// Model with no SKUs: found=false
_, found = maxModelRemainingQuota(modelNoSkus, emptyUsages)
require.False(t, found)
}
7 changes: 7 additions & 0 deletions cli/azd/pkg/ai/types.go
Original file line number Diff line number Diff line change
Expand Up @@ -104,9 +104,16 @@ type ModelLocationQuota struct {
// Location is the Azure location name.
Location string
// MaxRemainingQuota is the maximum remaining quota across model SKUs with usage entries.
// A value of QuotaRemainingUnknown (-1) indicates that usage data was unavailable
// (e.g. free-tier subscriptions) and the actual remaining quota is unknown.
MaxRemainingQuota float64
}

// QuotaRemainingUnknown is a sentinel value for MaxRemainingQuota indicating that
// the /usages API returned no data (e.g. free-tier subscriptions that have not yet
// provisioned Cognitive Services resources) and the actual remaining quota is unknown.
const QuotaRemainingUnknown float64 = -1

// QuotaRequirement specifies a single quota check: the usage name to check
// and the minimum remaining capacity needed.
type QuotaRequirement struct {
Expand Down
10 changes: 10 additions & 0 deletions cli/azd/pkg/infra/provisioning/bicep/bicep_provider.go
Original file line number Diff line number Diff line change
Expand Up @@ -2996,6 +2996,16 @@ func (p *BicepProvider) checkAiModelQuota(
}

// Check aggregated capacity against remaining quota.
// Skip when the /usages API returned an empty list (e.g. free-tier
// subscriptions that have not yet provisioned Cognitive Services
// resources). Empty usages means no consumption data is
// available — not that quota is zero. The model-catalog
// validation above still runs so that ai_model_not_found
// warnings are surfaced regardless.
if len(usages) == 0 {
continue
}
Comment thread
vhvb1989 marked this conversation as resolved.

reportedUsage := map[string]bool{}
for _, r := range resolved {
if reportedUsage[r.usageName] {
Expand Down
Loading