Move scaling quota check to base BYOC provider (#1083)

* Move scaling quota check to base BYOC provider

* Update src/pkg/cli/client/byoc/baseclient.go

Co-authored-by: Jordan Stephens <[email protected]>

* rename quota to policy to be more encompassing of quotas and access

* update naming of storage of user CanIUse values

---------

Co-authored-by: Eric Liu <[email protected]>
Co-authored-by: Jordan Stephens <[email protected]>
Co-authored-by: Eric Liu <[email protected]>

Author: 3 people

src/cmd/cli/command/commands.go

@@ -60,14 +60,6 @@ func getCluster() string {
 	return org + "@" + cluster
 }
 
-const TIER_ERROR_MESSAGE = "current subscription tier does not allow this action: "
-
-type ErrNoPermission string
-
-func (e ErrNoPermission) Error() string {
-	return TIER_ERROR_MESSAGE + string(e)
-}
-
 func prettyError(err error) error {
 	// To avoid printing the internal gRPC error code
 	var cerr *connect.Error
@@ -1223,10 +1215,7 @@ func canIUseProvider(ctx context.Context, provider cliClient.Provider, projectNa
 	if err != nil {
 		return err
 	}
-	// Allow local override of the CD image
-	cdImage := pkg.Getenv("DEFANG_CD_IMAGE", resp.CdImage)
-	provider.SetCDImage(cdImage)
-
+	provider.SetCanIUseConfig(resp)
 	return nil
 }
 

src/cmd/cli/command/compose.go

@@ -95,34 +95,17 @@ func makeComposeUpCmd() *cobra.Command {
 				return err
 			}
 
+			// Check if the user has permission to use the provider
 			err = canIUseProvider(ctx, provider, project.Name)
 			if err != nil {
 				return err
 			}
 
 			managedServices, _ := cli.SplitManagedAndUnmanagedServices(project.Services)
-
 			if len(managedServices) > 0 {
 				term.Warnf("Defang cannot monitor status of the following managed service(s): %v.\n   To check if the managed service is up, check the status of the service which depends on it.", managedServices)
 			}
 
-			numGPUS := compose.GetNumOfGPUs(ctx, project)
-			if numGPUS > 0 {
-				req := &defangv1.CanIUseRequest{
-					Project:  project.Name,
-					Provider: providerID.EnumValue(),
-				}
-
-				resp, err := client.CanIUse(ctx, req)
-				if err != nil {
-					return err
-				}
-
-				if !resp.Gpu {
-					return ErrNoPermission("usage of GPUs. Please upgrade on https://s.defang.io/subscription")
-				}
-			}
-
 			deploy, project, err := cli.ComposeUp(ctx, project, client, provider, upload, mode.Value())
 
 			if err != nil {

src/pkg/cli/client/byoc/aws/byoc.go

@@ -210,28 +210,17 @@ func (b *ByocAws) deploy(ctx context.Context, req *defangv1.DeployRequest, cmd s
 		return nil, err
 	}
 
-	etag := pkg.RandomID()
-	quotaClient = NewServiceQuotasClient(ctx, cfg)
-	if err = ValidateGPUResources(ctx, project); err != nil {
+	quotaClient = NewServiceQuotasClient(cfg)
+	if err = validateGPUResources(ctx, project); err != nil {
 		return nil, err
 	}
 
+	etag := pkg.RandomID()
 	serviceInfos, err := b.GetServiceInfos(ctx, project.Name, req.DelegateDomain, etag, project.Services)
 	if err != nil {
 		return nil, err
 	}
 
-	// Ensure all service endpoints are unique
-	endpoints := make(map[string]bool)
-	for _, serviceInfo := range serviceInfos {
-		for _, endpoint := range serviceInfo.Endpoints {
-			if endpoints[endpoint] {
-				return nil, fmt.Errorf("duplicate endpoint: %s", endpoint) // CodeInvalidArgument
-			}
-			endpoints[endpoint] = true
-		}
-	}
-
 	data, err := proto.Marshal(&defangv1.ProjectUpdate{
 		CdVersion: b.CDImage,
 		Compose:   req.Compose,

src/pkg/cli/client/byoc/aws/validation.go

@@ -25,7 +25,7 @@ var ErrAWSNoConnection = errors.New("no connect to AWS service quotas")
 var ErrGPUQuotaZero = errors.New("no GPUs enabled. To resolve see https://docs.defang.io/docs/tutorials/deploy-with-gpu")
 var ErrNoQuotasReceived = errors.New("no service quotas received")
 
-func NewServiceQuotasClient(ctx context.Context, cfg aws.Config) *servicequotas.Client {
+func NewServiceQuotasClient(cfg aws.Config) *servicequotas.Client {
 	return servicequotas.NewFromConfig(cfg)
 }
 
@@ -69,12 +69,11 @@ func hasGPUQuota(ctx context.Context) (bool, error) {
 	return false, nil
 }
 
-func ValidateGPUResources(ctx context.Context, project *composeTypes.Project) error {
+func validateGPUResources(ctx context.Context, project *composeTypes.Project) error {
 	// return after checking if there are actually non-zero GPUs requested
 	gpusAvailable, quotaErr := hasGPUQuota(ctx)
 
-	gpuCount := compose.GetNumOfGPUs(ctx, project)
-
+	gpuCount := compose.GetNumOfGPUs(project.Services)
 	if gpuCount == 0 {
 		return nil
 	}

src/pkg/cli/client/byoc/aws/validation_test.go

@@ -57,7 +57,7 @@ func TestValidateGPUResources(t *testing.T) {
 		quotaClient = mockQuotaClient
 		mockQuotaClient.output = nil
 		mockQuotaClient.err = ErrNoQuotasReceived
-		err := ValidateGPUResources(ctx, &project)
+		err := validateGPUResources(ctx, &project)
 		if err != nil && errors.Is(err, ErrNoQuotasReceived) {
 			t.Fatalf("ValidateGPUResources() failed: Unexpected errors %v", err)
 		}
@@ -73,7 +73,7 @@ func TestValidateGPUResources(t *testing.T) {
 
 		quotaClient = nil
 		mockQuotaClient.err = nil
-		err := ValidateGPUResources(ctx, &project)
+		err := validateGPUResources(ctx, &project)
 		if err != nil {
 			t.Fatalf("ValidateGPUResources() failed: expected no errors but got %v", err)
 		}
@@ -97,7 +97,7 @@ func TestValidateGPUResources(t *testing.T) {
 				},
 			},
 		}
-		err := ValidateGPUResources(ctx, &project)
+		err := validateGPUResources(ctx, &project)
 		if err != nil && !errors.Is(err, ErrGPUQuotaZero) {
 			t.Fatalf("ValidateGPUResources() failed: Unexpected err %v", err)
 		}
@@ -113,7 +113,7 @@ func TestValidateGPUResources(t *testing.T) {
 		}
 
 		quotaClient = nil
-		err := ValidateGPUResources(ctx, &project)
+		err := validateGPUResources(ctx, &project)
 		if err != nil && !errors.Is(err, ErrAWSNoConnection) {
 			t.Fatalf("ValidateGPUResources() failed: Unexpected err %v", err)
 		}

src/pkg/cli/client/byoc/baseclient.go

@@ -36,14 +36,20 @@ type HasStackSupport interface {
 	GetStackName() string
 }
 
+type CanIUseConfig struct {
+	CDImage      string
+	AllowScaling bool
+	AllowGPU     bool
+}
+
 type ByocBaseClient struct {
 	client.RetryDelayer
 
 	PulumiStack             string
 	SetupDone               bool
 	ShouldDelegateSubdomain bool
 	TenantName              string
-	CDImage                 string
+	CanIUseConfig
 
 	projectBackend ProjectBackend
 }
@@ -78,8 +84,11 @@ func (b *ByocBaseClient) Debug(context.Context, *defangv1.DebugRequest) (*defang
 	return nil, client.ErrNotImplemented("AI debugging is not yet supported for BYOC")
 }
 
-func (b *ByocBaseClient) SetCDImage(image string) {
-	b.CDImage = image
+func (b *ByocBaseClient) SetCanIUseConfig(quotas *defangv1.CanIUseResponse) {
+	// Allow local override of the CD image
+	b.CDImage = pkg.Getenv("DEFANG_CD_IMAGE", quotas.CdImage)
+	b.AllowScaling = quotas.AllowScaling
+	b.AllowGPU = quotas.Gpu
 }
 
 func (b *ByocBaseClient) ServiceDNS(name string) string {
@@ -122,7 +131,18 @@ func (b *ByocBaseClient) GetProjectDomain(projectName, zone string) string {
 	return domain
 }
 
+type ErrNoPermission string
+
+func (e ErrNoPermission) Error() string {
+	return "current subscription tier does not allow this action: " + string(e)
+}
+
 func (b *ByocBaseClient) GetServiceInfos(ctx context.Context, projectName, delegateDomain, etag string, services map[string]composeTypes.ServiceConfig) ([]*defangv1.ServiceInfo, error) {
+	numGPUS := compose.GetNumOfGPUs(services)
+	if numGPUS > 0 && !b.AllowGPU {
+		return nil, ErrNoPermission("usage of GPUs. Please upgrade on https://s.defang.io/subscription")
+	}
+
 	serviceInfoMap := make(map[string]*Node)
 	for _, service := range services {
 		serviceInfo, err := b.update(ctx, projectName, delegateDomain, service)
@@ -138,7 +158,20 @@ func (b *ByocBaseClient) GetServiceInfos(ctx context.Context, projectName, deleg
 	}
 
 	// Reorder the serviceInfos to make sure the dependencies are created first
-	return topologicalSort(serviceInfoMap), nil
+	serviceInfos := topologicalSort(serviceInfoMap)
+
+	// Ensure all service endpoints are unique
+	endpoints := make(map[string]bool)
+	for _, serviceInfo := range serviceInfos {
+		for _, endpoint := range serviceInfo.Endpoints {
+			if endpoints[endpoint] {
+				return nil, fmt.Errorf("duplicate endpoint: %s", endpoint) // CodeInvalidArgument
+			}
+			endpoints[endpoint] = true
+		}
+	}
+
+	return serviceInfos, nil
 }
 
 // Simple DFS topological sort to make sure the dependencies are created first
@@ -180,10 +213,11 @@ func (b *ByocBaseClient) update(ctx context.Context, projectName, delegateDomain
 
 	pkg.Ensure(projectName != "", "ProjectName not set")
 	si := &defangv1.ServiceInfo{
-		Etag:       pkg.RandomID(), // TODO: could be hash for dedup/idempotency
-		Project:    projectName,    // was: tenant
-		Service:    &defangv1.Service{Name: service.Name},
-		Domainname: service.DomainName,
+		AllowScaling: b.AllowScaling,
+		Domainname:   service.DomainName,
+		Etag:         pkg.RandomID(), // TODO: could be hash for dedup/idempotency
+		Project:      projectName,    // was: tenant
+		Service:      &defangv1.Service{Name: service.Name},
 	}
 
 	hasHost := false

src/pkg/cli/client/byoc/do/byoc.go

@@ -146,23 +146,11 @@ func (b *ByocDo) deploy(ctx context.Context, req *defangv1.DeployRequest, cmd st
 	}
 
 	etag := pkg.RandomID()
-
 	serviceInfos, err := b.GetServiceInfos(ctx, project.Name, req.DelegateDomain, etag, project.Services)
 	if err != nil {
 		return nil, err
 	}
 
-	// Ensure all service endpoints are unique
-	endpoints := make(map[string]bool)
-	for _, serviceInfo := range serviceInfos {
-		for _, endpoint := range serviceInfo.Endpoints {
-			if endpoints[endpoint] {
-				return nil, fmt.Errorf("duplicate endpoint: %s", endpoint) // CodeInvalidArgument
-			}
-			endpoints[endpoint] = true
-		}
-	}
-
 	data, err := proto.Marshal(&defangv1.ProjectUpdate{
 		CdVersion: b.CDImage,
 		Compose:   req.Compose,

src/pkg/cli/client/byoc/gcp/byoc.go

@@ -410,7 +410,6 @@ func (b *ByocGcp) deploy(ctx context.Context, req *defangv1.DeployRequest, comma
 	}
 
 	etag := pkg.RandomID()
-
 	serviceInfos, err := b.GetServiceInfos(ctx, project.Name, req.DelegateDomain, etag, project.Services)
 	if err != nil {
 		return nil, err

src/pkg/cli/client/playground.go

@@ -128,7 +128,7 @@ func (g *PlaygroundProvider) QueryForDebug(ctx context.Context, req *defangv1.De
 func (g *PlaygroundProvider) PrepareDomainDelegation(ctx context.Context, req PrepareDomainDelegationRequest) (*PrepareDomainDelegationResponse, error) {
 	return nil, nil // Playground does not support delegate domains
 }
-func (g *PlaygroundProvider) SetCDImage(string) {}
+func (g *PlaygroundProvider) SetCanIUseConfig(*defangv1.CanIUseResponse) {}
 
 type PlaygroundAccountInfo struct{}
 

src/pkg/cli/client/provider.go

@@ -147,7 +147,7 @@ type Provider interface {
 	PutConfig(context.Context, *defangv1.PutConfigRequest) error
 	RemoteProjectName(context.Context) (string, error)
 	ServiceDNS(string) string
-	SetCDImage(string)
+	SetCanIUseConfig(*defangv1.CanIUseResponse)
 	Subscribe(context.Context, *defangv1.SubscribeRequest) (ServerStream[defangv1.SubscribeResponse], error)
 	TearDown(context.Context) error
 }