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{}