chore: check task config policies against slots and max_slots (#10015)

determined-ai · Oct 3, 2024 · 262b4a9 · 262b4a9
1 parent a0cc818
commit 262b4a9
Show file tree

Hide file tree

Showing 3 changed files with 30 additions and 7 deletions.
diff --git a/master/internal/api_command.go b/master/internal/api_command.go
@@ -153,7 +153,7 @@ 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 {
- return nil, nil, err
+ return nil, nil, status.Errorf(codes.InvalidArgument, "failed constraint check: %v", err)
  }
 
  token, err := getTaskSessionToken(ctx, userModel)

diff --git a/master/internal/configpolicy/task_config_policy.go b/master/internal/configpolicy/task_config_policy.go
@@ -55,15 +55,20 @@ func CheckNTSCConstraints(
  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",
- *constraints.PriorityLimit, *workloadConfig.Resources.Priority, errPriorityConstraintFailure)
+ *workloadConfig.Resources.Priority, *constraints.PriorityLimit, errPriorityConstraintFailure)
  }
  }
 
- if constraints.ResourceConstraints != nil && constraints.ResourceConstraints.MaxSlots != nil &&
- 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",
- *constraints.ResourceConstraints.MaxSlots, *workloadConfig.Resources.MaxSlots, errResourceConstraintFailure)
+ 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)
+ }
+ }
+ 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)
  }
  }
 

diff --git a/master/internal/configpolicy/task_config_policy_intg_test.go b/master/internal/configpolicy/task_config_policy_intg_test.go
@@ -111,6 +111,24 @@ func TestValidateNTSCConstraints(t *testing.T) {
  require.ErrorIs(t, err, errResourceConstraintFailure)
  })
 
+ t.Run("exceeds 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)
+
+ resourceManager := mocks.ResourceManager{}
+ resourceManager.On("SmallerValueIsHigherPriority", mock.Anything).Return(true, nil)
+
+ config := defaultConfig()
+ config.Resources.Slots = *config.Resources.MaxSlots
+ config.Resources.MaxSlots = nil // ensure only slots is set
+ _, 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)
  rm1 := mocks.ResourceManager{}