From 5b1380cf61c21b88748dff470afe2f3f23fe74fa Mon Sep 17 00:00:00 2001 From: Kristine Kunapuli <49738148+kkunapuli@users.noreply.github.com> Date: Mon, 7 Oct 2024 10:26:22 -0400 Subject: [PATCH] chore: check experiment constraints (#10018) --- master/internal/api_command.go | 4 +- master/internal/api_experiment.go | 13 ++ master/internal/api_experiment_intg_test.go | 12 +- master/internal/api_logretention_intg_test.go | 4 + .../configpolicy/task_config_policy.go | 89 +++++++-- .../task_config_policy_intg_test.go | 179 ++++++++++++++---- 6 files changed, 243 insertions(+), 58 deletions(-) diff --git a/master/internal/api_command.go b/master/internal/api_command.go index 3af94d5227d..4ef783db660 100644 --- a/master/internal/api_command.go +++ b/master/internal/api_command.go @@ -151,8 +151,8 @@ func (a *apiServer) getCommandLaunchParams(ctx context.Context, req *protoComman } // Check submitted config against task config policies. - valid, err := configpolicy.CheckNTSCConstraints(ctx, int(cmdSpec.Metadata.WorkspaceID), config, a.m.rm) - if !valid { + err = configpolicy.CheckNTSCConstraints(ctx, int(cmdSpec.Metadata.WorkspaceID), config, a.m.rm) + if err != nil { return nil, nil, status.Errorf(codes.InvalidArgument, "failed constraint check: %v", err) } diff --git a/master/internal/api_experiment.go b/master/internal/api_experiment.go index 0af79921847..6839ca878a5 100644 --- a/master/internal/api_experiment.go +++ b/master/internal/api_experiment.go @@ -30,6 +30,7 @@ import ( "github.com/determined-ai/determined/master/internal/api" "github.com/determined-ai/determined/master/internal/authz" + "github.com/determined-ai/determined/master/internal/configpolicy" "github.com/determined-ai/determined/master/internal/db" "github.com/determined-ai/determined/master/internal/db/bunutils" "github.com/determined-ai/determined/master/internal/experiment" @@ -1681,6 +1682,18 @@ func (a *apiServer) CreateExperiment( return nil, status.Errorf(codes.PermissionDenied, err.Error()) } + wkspIDs, err := workspace.WorkspaceIDsFromNames(ctx, []string{taskSpec.Workspace}) + if err != nil { + return nil, status.Errorf(codes.InvalidArgument, err.Error()) + } + if len(wkspIDs) != 1 { + return nil, status.Error(codes.InvalidArgument, "expected exactly one workspace") + } + err = configpolicy.CheckExperimentConstraints(ctx, int(wkspIDs[0]), activeConfig, a.m.rm) + if err != nil { + return nil, status.Errorf(codes.InvalidArgument, err.Error()) + } + if req.ValidateOnly { return &apiv1.CreateExperimentResponse{ Experiment: &experimentv1.Experiment{}, diff --git a/master/internal/api_experiment_intg_test.go b/master/internal/api_experiment_intg_test.go index 711f53ded8b..a9a9d620e91 100644 --- a/master/internal/api_experiment_intg_test.go +++ b/master/internal/api_experiment_intg_test.go @@ -583,7 +583,8 @@ func TestPutExperimentsRetainLogs(t *testing.T) { func TestParseAndMergeContinueConfig(t *testing.T) { // Blank config. - api, curUser, ctx := setupAPITest(t, nil) + mockRM := MockRM() + api, curUser, ctx := setupAPITest(t, nil, mockRM) exp := createTestExp(t, api, curUser) _, _, err := api.parseAndMergeContinueConfig(exp.ID, ``) @@ -652,6 +653,7 @@ resources: } // No checkpoint specified anywhere. + mockRM.On("SmallerValueIsHigherPriority", mock.Anything).Return(true, nil) resp, err := api.CreateExperiment(ctx, createReq) require.NoError(t, err) _, _, err = api.parseAndMergeContinueConfig(int(resp.Experiment.Id), `{}`) @@ -671,7 +673,8 @@ searcher: // nolint: exhaustruct func TestCreateExperimentCheckpointStorage(t *testing.T) { - api, _, ctx := setupAPITest(t, nil) + mockRM := MockRM() + api, _, ctx := setupAPITest(t, nil, mockRM) api.m.config.CheckpointStorage = expconf.CheckpointStorageConfig{} defer func() { api.m.config.CheckpointStorage = expconf.CheckpointStorageConfig{} @@ -694,6 +697,7 @@ resources: } // No checkpoint specified anywhere. + mockRM.On("SmallerValueIsHigherPriority", mock.Anything).Return(true, nil) _, err := api.CreateExperiment(ctx, createReq) require.ErrorContains(t, err, "checkpoint_storage: type is a required property") @@ -1589,7 +1593,8 @@ func TestAuthZGetExperimentLabels(t *testing.T) { } func TestAuthZCreateExperiment(t *testing.T) { - api, authZExp, _, curUser, ctx := setupExpAuthTest(t, nil) + mockRM := MockRM() + api, authZExp, _, curUser, ctx := setupExpAuthTest(t, nil, mockRM) forkFrom := createTestExp(t, api, curUser) workspaceID, projectID := createProjectAndWorkspace(ctx, t, api) @@ -1682,6 +1687,7 @@ func TestAuthZCreateExperiment(t *testing.T) { Return(nil).Once() authZExp.On("CanEditExperiment", mock.Anything, mockUserArg, mock.Anything, mock.Anything).Return( fmt.Errorf("canActivateExperimentError")).Once() + mockRM.On("SmallerValueIsHigherPriority", mock.Anything).Return(true, nil) _, err := api.CreateExperiment(ctx, &apiv1.CreateExperimentRequest{ Activate: true, Config: minExpConfToYaml(t), diff --git a/master/internal/api_logretention_intg_test.go b/master/internal/api_logretention_intg_test.go index 99afa0e35fb..0c868f8732c 100644 --- a/master/internal/api_logretention_intg_test.go +++ b/master/internal/api_logretention_intg_test.go @@ -12,6 +12,7 @@ import ( "github.com/go-co-op/gocron/v2" "github.com/jonboulle/clockwork" + "github.com/stretchr/testify/mock" "github.com/stretchr/testify/require" "github.com/uptrace/bun" @@ -108,6 +109,9 @@ searcher: } // No checkpoint specified anywhere. + mockRM := MockRM() + api.m.rm = mockRM + mockRM.On("SmallerValueIsHigherPriority", mock.Anything).Return(true, nil) resp, err := api.CreateExperiment(ctx, createReq) require.NoError(t, err) require.Empty(t, resp.Warnings) diff --git a/master/internal/configpolicy/task_config_policy.go b/master/internal/configpolicy/task_config_policy.go index 4eb62aed894..0600eb1f77a 100644 --- a/master/internal/configpolicy/task_config_policy.go +++ b/master/internal/configpolicy/task_config_policy.go @@ -36,43 +36,100 @@ var ( errResourceConstraintFailure = errors.New("submitted workload failed a resource constraint") ) -// CheckNTSCConstraints returns true if the NTSC config passes constraint checks. +// CheckNTSCConstraints returns an error if the NTSC config fails constraint checks. func CheckNTSCConstraints( ctx context.Context, workspaceID int, workloadConfig model.CommandConfig, resourceManager rm.ResourceManager, -) (bool, error) { +) error { constraints, err := GetMergedConstraints(ctx, workspaceID, model.NTSCType) if err != nil { - return false, err + return err + } + + if constraints.ResourceConstraints != nil && constraints.ResourceConstraints.MaxSlots != nil { + if err = checkSlotsConstraint(*constraints.ResourceConstraints.MaxSlots, workloadConfig.Resources.Slots, + workloadConfig.Resources.MaxSlots); err != nil { + return err + } } // For each submitted constraint, check if the workload config is within allowed values. // rm.SmallerValueIsHigherPriority only returns an error if task priority is not implemented for that resource manager. // In that case, there is no need to check if requested priority is within limits. smallerHigher, err := resourceManager.SmallerValueIsHigherPriority() - if err == nil && constraints.PriorityLimit != nil && workloadConfig.Resources.Priority != nil { - if !priorityWithinLimit(*workloadConfig.Resources.Priority, *constraints.PriorityLimit, smallerHigher) { - return false, fmt.Errorf("requested priority [%d] exceeds limit set by admin [%d]: %w", - *workloadConfig.Resources.Priority, *constraints.PriorityLimit, errPriorityConstraintFailure) + if err == nil { + if err = checkPriorityConstraint(smallerHigher, constraints.PriorityLimit, + workloadConfig.Resources.Priority); err != nil { + return err } } + return nil +} + +// CheckExperimentConstraints returns an error if the NTSC config fails constraint checks. +func CheckExperimentConstraints( + ctx context.Context, + workspaceID int, + workloadConfig expconf.ExperimentConfigV0, + resourceManager rm.ResourceManager, +) error { + constraints, err := GetMergedConstraints(ctx, workspaceID, model.ExperimentType) + if err != nil { + return err + } + if constraints.ResourceConstraints != nil && constraints.ResourceConstraints.MaxSlots != nil { - if workloadConfig.Resources.MaxSlots != nil { - if *constraints.ResourceConstraints.MaxSlots < *workloadConfig.Resources.MaxSlots { - return false, fmt.Errorf("requested resources.max_slots [%d] exceeds limit set by admin [%d]: %w", - *workloadConfig.Resources.MaxSlots, *constraints.ResourceConstraints.MaxSlots, errResourceConstraintFailure) - } + // users cannot specify number of slots for an experiment + slotsRequest := 0 + if err = checkSlotsConstraint(*constraints.ResourceConstraints.MaxSlots, slotsRequest, + workloadConfig.Resources().MaxSlots()); err != nil { + return err } - if *constraints.ResourceConstraints.MaxSlots < workloadConfig.Resources.Slots { - return false, fmt.Errorf("requested resources.slots [%d] exceeds limit set by admin [%d]: %w", - workloadConfig.Resources.Slots, *constraints.ResourceConstraints.MaxSlots, errResourceConstraintFailure) + } + + // For each submitted constraint, check if the workload config is within allowed values. + // rm.SmallerValueIsHigherPriority only returns an error if task priority is not implemented for that resource manager. + // In that case, there is no need to check if requested priority is within limits. + smallerHigher, err := resourceManager.SmallerValueIsHigherPriority() + if err == nil { + if err = checkPriorityConstraint(smallerHigher, constraints.PriorityLimit, + workloadConfig.Resources().Priority()); err != nil { + return err } } - return true, nil + return nil +} + +func checkPriorityConstraint(smallerHigher bool, priorityLimit *int, priorityRequest *int) error { + if priorityLimit == nil || priorityRequest == nil { + return nil + } + + if !priorityWithinLimit(*priorityRequest, *priorityLimit, smallerHigher) { + return fmt.Errorf("requested priority [%d] exceeds limit set by admin [%d]: %w", + *priorityRequest, *priorityLimit, errPriorityConstraintFailure) + } + return nil +} + +func checkSlotsConstraint(slotsLimit int, slotsRequest int, maxSlotsRequest *int) error { + if slotsLimit < slotsRequest { + return fmt.Errorf("requested resources.slots [%d] exceeds limit set by admin [%d]: %w", + slotsRequest, slotsLimit, errResourceConstraintFailure) + } + + if maxSlotsRequest != nil { + if slotsLimit < *maxSlotsRequest { + return fmt.Errorf("requested resources.max_slots [%d] exceeds limit set by admin [%d]: %w", + *maxSlotsRequest, slotsLimit, errResourceConstraintFailure) + } + } + + return nil } // GetMergedConstraints retrieves Workspace and Global constraints and returns a merged result. diff --git a/master/internal/configpolicy/task_config_policy_intg_test.go b/master/internal/configpolicy/task_config_policy_intg_test.go index cfe8bb8db4f..d059d9aea7d 100644 --- a/master/internal/configpolicy/task_config_policy_intg_test.go +++ b/master/internal/configpolicy/task_config_policy_intg_test.go @@ -13,6 +13,7 @@ import ( "github.com/determined-ai/determined/master/internal/mocks" "github.com/determined-ai/determined/master/pkg/etc" "github.com/determined-ai/determined/master/pkg/model" + "github.com/determined-ai/determined/master/pkg/schemas/expconf" ) func TestPriorityAllowed(t *testing.T) { @@ -28,7 +29,7 @@ func TestPriorityAllowed(t *testing.T) { wkspLimit := 50 user := db.RequireMockUser(t, pgDB) - w := addWorkspacePriorityLimit(t, user, wkspLimit) + w := addWorkspacePriorityLimit(t, user, wkspLimit, model.NTSCType) // Priority is outside workspace limit. smallerValueIsHigherPriority := true @@ -37,7 +38,7 @@ func TestPriorityAllowed(t *testing.T) { require.False(t, ok) globalLimit := 42 - addConstraints(t, user, nil, fmt.Sprintf(`{"priority_limit": %d}`, globalLimit)) + addConstraints(t, user, nil, fmt.Sprintf(`{"priority_limit": %d}`, globalLimit), model.NTSCType) // Priority is within global limit. ok, err = PriorityAllowed(w.ID, model.NTSCType, wkspLimit-1, true) @@ -50,7 +51,7 @@ func TestPriorityAllowed(t *testing.T) { require.False(t, ok) } -func TestValidateNTSCConstraints(t *testing.T) { +func TestCheckNTSCConstraints(t *testing.T) { require.NoError(t, etc.SetRootPath(db.RootFromDB)) pgDB, cleanup := db.MustResolveNewPostgresDatabase(t) defer cleanup() @@ -63,21 +64,20 @@ func TestValidateNTSCConstraints(t *testing.T) { resourceManager := mocks.ResourceManager{} resourceManager.On("SmallerValueIsHigherPriority", mock.Anything).Return(false, nil) - config := defaultConfig() - ok, err := CheckNTSCConstraints(context.Background(), 1, config, &resourceManager) + config := defaultNTSCConfig() + err := CheckNTSCConstraints(context.Background(), 1, config, &resourceManager) require.NoError(t, err) - require.True(t, ok) }) t.Run("running in wksp with constraints - not ok", func(t *testing.T) { - w := addWorkspacePriorityLimit(t, user, wkspPriorityLimit) + w := addWorkspacePriorityLimit(t, user, wkspPriorityLimit, model.NTSCType) resourceManager := mocks.ResourceManager{} resourceManager.On("SmallerValueIsHigherPriority", mock.Anything).Return(false, nil) ctx, cancel := context.WithCancel(context.Background()) defer cancel() - config := defaultConfig() - _, err := CheckNTSCConstraints(ctx, w.ID, config, &resourceManager) + config := defaultNTSCConfig() + err := CheckNTSCConstraints(ctx, w.ID, config, &resourceManager) require.Error(t, err) require.ErrorIs(t, err, errPriorityConstraintFailure) }) @@ -89,9 +89,8 @@ func TestValidateNTSCConstraints(t *testing.T) { _, err := db.Bun().NewInsert().Model(&w).Exec(context.Background()) require.NoError(t, err) - config := defaultConfig() - ok, err := CheckNTSCConstraints(context.Background(), w.ID, config, &resourceManager) - require.True(t, ok) + config := defaultNTSCConfig() + err = CheckNTSCConstraints(context.Background(), w.ID, config, &resourceManager) require.NoError(t, err) }) @@ -100,13 +99,13 @@ func TestValidateNTSCConstraints(t *testing.T) { w := model.Workspace{Name: uuid.NewString(), UserID: user.ID} _, err := db.Bun().NewInsert().Model(&w).Exec(context.Background()) require.NoError(t, err) - addConstraints(t, user, &w.ID, *constraints) + addConstraints(t, user, &w.ID, *constraints, model.NTSCType) resourceManager := mocks.ResourceManager{} resourceManager.On("SmallerValueIsHigherPriority", mock.Anything).Return(true, nil) - config := defaultConfig() - _, err = CheckNTSCConstraints(context.Background(), w.ID, config, &resourceManager) + config := defaultNTSCConfig() + err = CheckNTSCConstraints(context.Background(), w.ID, config, &resourceManager) require.Error(t, err) require.ErrorIs(t, err, errResourceConstraintFailure) }) @@ -116,52 +115,144 @@ func TestValidateNTSCConstraints(t *testing.T) { w := model.Workspace{Name: uuid.NewString(), UserID: user.ID} _, err := db.Bun().NewInsert().Model(&w).Exec(context.Background()) require.NoError(t, err) - addConstraints(t, user, &w.ID, *constraints) + addConstraints(t, user, &w.ID, *constraints, model.NTSCType) resourceManager := mocks.ResourceManager{} resourceManager.On("SmallerValueIsHigherPriority", mock.Anything).Return(true, nil) - config := defaultConfig() + config := defaultNTSCConfig() config.Resources.Slots = *config.Resources.MaxSlots config.Resources.MaxSlots = nil // ensure only slots is set - _, err = CheckNTSCConstraints(context.Background(), w.ID, config, &resourceManager) + err = CheckNTSCConstraints(context.Background(), w.ID, config, &resourceManager) require.Error(t, err) require.ErrorIs(t, err, errResourceConstraintFailure) }) t.Run("rm priority not supported - ok", func(t *testing.T) { - w := addWorkspacePriorityLimit(t, user, wkspPriorityLimit) + w := addWorkspacePriorityLimit(t, user, wkspPriorityLimit, model.NTSCType) rm1 := mocks.ResourceManager{} rm1.On("SmallerValueIsHigherPriority", mock.Anything).Return(false, nil).Once() - config := defaultConfig() - _, err := CheckNTSCConstraints(context.Background(), w.ID, config, &rm1) + config := defaultNTSCConfig() + err := CheckNTSCConstraints(context.Background(), w.ID, config, &rm1) require.Error(t, err) require.ErrorIs(t, err, errPriorityConstraintFailure) // Validate constraints again. This time, the RM does not support priority. rmNoPriority := mocks.ResourceManager{} rmNoPriority.On("SmallerValueIsHigherPriority", mock.Anything).Return(false, fmt.Errorf("not supported")).Once() - ok, err := CheckNTSCConstraints(context.Background(), w.ID, config, &rmNoPriority) - require.True(t, ok) + err = CheckNTSCConstraints(context.Background(), w.ID, config, &rmNoPriority) require.NoError(t, err) }) t.Run("no config set - ok", func(t *testing.T) { constraints := DefaultConstraints() - addConstraints(t, user, nil, *constraints) // add global constraints + addConstraints(t, user, nil, *constraints, model.NTSCType) // add global constraints resourceManager := mocks.ResourceManager{} resourceManager.On("SmallerValueIsHigherPriority", mock.Anything).Return(true, nil) - config := defaultConfig() - ok, err := CheckNTSCConstraints(context.Background(), 1, config, &resourceManager) - require.False(t, ok) + config := defaultNTSCConfig() + err := CheckNTSCConstraints(context.Background(), 1, config, &resourceManager) require.Error(t, err) emptyConfig := model.CommandConfig{} - ok, err = CheckNTSCConstraints(context.Background(), 1, emptyConfig, &resourceManager) - require.True(t, ok) + err = CheckNTSCConstraints(context.Background(), 1, emptyConfig, &resourceManager) + require.NoError(t, err) + }) +} + +func TestCheckExperimentConstraints(t *testing.T) { + require.NoError(t, etc.SetRootPath(db.RootFromDB)) + pgDB, cleanup := db.MustResolveNewPostgresDatabase(t) + defer cleanup() + db.MustMigrateTestPostgres(t, pgDB, db.MigrationsFromDB) + + wkspPriorityLimit := 7 + user := db.RequireMockUser(t, pgDB) + + t.Run("no constraints set - ok", func(t *testing.T) { + resourceManager := mocks.ResourceManager{} + resourceManager.On("SmallerValueIsHigherPriority", mock.Anything).Return(false, nil) + + config := defaultExperimentConfig() + err := CheckExperimentConstraints(context.Background(), 1, config, &resourceManager) + require.NoError(t, err) + }) + + t.Run("running in wksp with constraints - not ok", func(t *testing.T) { + w := addWorkspacePriorityLimit(t, user, wkspPriorityLimit, model.ExperimentType) + resourceManager := mocks.ResourceManager{} + resourceManager.On("SmallerValueIsHigherPriority", mock.Anything).Return(false, nil) + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + config := defaultExperimentConfig() + err := CheckExperimentConstraints(ctx, w.ID, config, &resourceManager) + require.Error(t, err) + require.ErrorIs(t, err, errPriorityConstraintFailure) + }) + + t.Run("running in wksp without constraints - ok", func(t *testing.T) { + resourceManager := mocks.ResourceManager{} + resourceManager.On("SmallerValueIsHigherPriority", mock.Anything).Return(false, nil) + w := model.Workspace{Name: uuid.NewString(), UserID: user.ID} + _, err := db.Bun().NewInsert().Model(&w).Exec(context.Background()) + require.NoError(t, err) + + config := defaultExperimentConfig() + err = CheckExperimentConstraints(context.Background(), w.ID, config, &resourceManager) + require.NoError(t, err) + }) + + t.Run("exceeds max slots - not ok", func(t *testing.T) { + constraints := DefaultConstraints() + w := model.Workspace{Name: uuid.NewString(), UserID: user.ID} + _, err := db.Bun().NewInsert().Model(&w).Exec(context.Background()) + require.NoError(t, err) + addConstraints(t, user, &w.ID, *constraints, model.ExperimentType) + + resourceManager := mocks.ResourceManager{} + resourceManager.On("SmallerValueIsHigherPriority", mock.Anything).Return(true, nil) + + config := defaultExperimentConfig() + err = CheckExperimentConstraints(context.Background(), w.ID, config, &resourceManager) + require.Error(t, err) + require.ErrorIs(t, err, errResourceConstraintFailure) + }) + + t.Run("rm priority not supported - ok", func(t *testing.T) { + w := addWorkspacePriorityLimit(t, user, wkspPriorityLimit, model.ExperimentType) + rm1 := mocks.ResourceManager{} + rm1.On("SmallerValueIsHigherPriority", mock.Anything).Return(false, nil).Once() + + config := defaultExperimentConfig() + err := CheckExperimentConstraints(context.Background(), w.ID, config, &rm1) + require.Error(t, err) + require.ErrorIs(t, err, errPriorityConstraintFailure) + + // Validate constraints again. This time, the RM does not support priority. + rmNoPriority := mocks.ResourceManager{} + rmNoPriority.On("SmallerValueIsHigherPriority", mock.Anything).Return(false, fmt.Errorf("not supported")).Once() + err = CheckExperimentConstraints(context.Background(), w.ID, config, &rmNoPriority) + require.NoError(t, err) + }) + + t.Run("no config set - ok", func(t *testing.T) { + constraints := DefaultConstraints() + addConstraints(t, user, nil, *constraints, model.ExperimentType) // add global constraints + + resourceManager := mocks.ResourceManager{} + resourceManager.On("SmallerValueIsHigherPriority", mock.Anything).Return(true, nil) + + config := defaultExperimentConfig() + err := CheckExperimentConstraints(context.Background(), 1, config, &resourceManager) + require.Error(t, err) + + emptyResources := expconf.ResourcesConfigV0{} + emptyConfig := expconf.ExperimentConfigV0{} + emptyConfig.SetResources(emptyResources) + err = CheckExperimentConstraints(context.Background(), 1, emptyConfig, &resourceManager) require.NoError(t, err) }) } @@ -181,7 +272,7 @@ func TestGetMergedConstraints(t *testing.T) { // Workspace priority limit set. wkspLimit := 42 user := db.RequireMockUser(t, pgDB) - w := addWorkspacePriorityLimit(t, user, wkspLimit) + w := addWorkspacePriorityLimit(t, user, wkspLimit, model.NTSCType) constraints, err = GetMergedConstraints(context.Background(), w.ID, model.NTSCType) require.NoError(t, err) require.Nil(t, constraints.ResourceConstraints) @@ -189,21 +280,21 @@ func TestGetMergedConstraints(t *testing.T) { // Global limit overrides workspace limit. globalLimit := 25 - addConstraints(t, user, nil, fmt.Sprintf(`{"priority_limit": %d}`, globalLimit)) + addConstraints(t, user, nil, fmt.Sprintf(`{"priority_limit": %d}`, globalLimit), model.NTSCType) constraints, err = GetMergedConstraints(context.Background(), w.ID, model.NTSCType) require.NoError(t, err) require.Nil(t, constraints.ResourceConstraints) require.Equal(t, globalLimit, *constraints.PriorityLimit) // Workspace max slots set. - addConstraints(t, user, &w.ID, *DefaultConstraints()) + addConstraints(t, user, &w.ID, *DefaultConstraints(), model.NTSCType) constraints, err = GetMergedConstraints(context.Background(), w.ID, model.NTSCType) require.NoError(t, err) require.Equal(t, 8, *constraints.ResourceConstraints.MaxSlots) // defined in DefaultConstraintsStr require.Equal(t, globalLimit, *constraints.PriorityLimit) // global constraint overrides workspace value } -func addWorkspacePriorityLimit(t *testing.T, user model.User, limit int) model.Workspace { +func addWorkspacePriorityLimit(t *testing.T, user model.User, limit int, workloadType string) model.Workspace { ctx := context.Background() // add a workspace to use @@ -213,7 +304,7 @@ func addWorkspacePriorityLimit(t *testing.T, user model.User, limit int) model.W constraints := fmt.Sprintf(`{"priority_limit": %d}`, limit) input := model.TaskConfigPolicies{ - WorkloadType: model.NTSCType, + WorkloadType: workloadType, WorkspaceID: &w.ID, Constraints: &constraints, LastUpdatedBy: user.ID, @@ -224,11 +315,11 @@ func addWorkspacePriorityLimit(t *testing.T, user model.User, limit int) model.W return w } -func addConstraints(t *testing.T, user model.User, wkspID *int, constraints string) { +func addConstraints(t *testing.T, user model.User, wkspID *int, constraints string, workloadType string) { ctx := context.Background() input := model.TaskConfigPolicies{ - WorkloadType: model.NTSCType, + WorkloadType: workloadType, WorkspaceID: wkspID, Constraints: &constraints, LastUpdatedBy: user.ID, @@ -237,7 +328,7 @@ func addConstraints(t *testing.T, user model.User, wkspID *int, constraints stri require.NoError(t, err) } -func defaultConfig() model.CommandConfig { +func defaultNTSCConfig() model.CommandConfig { config := model.DefaultConfig(nil) configPriority := 50 @@ -247,3 +338,17 @@ func defaultConfig() model.CommandConfig { return config } + +func defaultExperimentConfig() expconf.ExperimentConfigV0 { + config := expconf.ExperimentConfigV0{} + + configPriority := 50 + configMaxSlots := 12 + resources := expconf.ResourcesConfigV0{ + RawMaxSlots: &configMaxSlots, + RawPriority: &configPriority, + } + config.SetResources(resources) + + return config +}