chore: populate exp config priority with enforced priority before opting to rm default

master/internal/configpolicy/task_config_policy.go

@@ -231,10 +231,15 @@ func MergeWithInvariantExperimentConfigs(ctx context.Context, workspaceID int,
 	return &config, nil
 }
 
-func findAllowedPriority(scope *int, workloadType string) (limit int, exists bool, err error) {
+// FindAllowedPriority finds the optionally set priority limit in scope's invariant config
+// policies. Returns the invariant config priority if that's set, and otherwise returns
+// the priority_limit constraint. If neither of the two is set, returns nil limit.
+func FindAllowedPriority(scope *int, workloadType string) (limit *int, constraintExists bool,
+	err error,
+) {
 	configPolicies, err := GetTaskConfigPolicies(context.TODO(), scope, workloadType)
 	if err != nil {
-		return 0, false, fmt.Errorf("unable to fetch task config policies: %w", err)
+		return nil, false, fmt.Errorf("unable to fetch task config policies: %w", err)
 	}
 
 	// Cannot update priority if priority set in invariant config.
@@ -244,50 +249,51 @@ func findAllowedPriority(scope *int, workloadType string) (limit int, exists boo
 			var configs model.CommandConfig
 			err = json.Unmarshal([]byte(*configPolicies.InvariantConfig), &configs)
 			if err != nil {
-				return 0, false, fmt.Errorf("unable to unmarshal task config policies: %w", err)
+				return nil, false, fmt.Errorf("unable to unmarshal task config policies: %w", err)
 			}
 			if configs.Resources.Priority != nil {
 				adminPriority := *configs.Resources.Priority
-				return adminPriority, false,
+				return &adminPriority, false,
 					fmt.Errorf("priority set by invariant config: %w", errPriorityImmutable)
 			}
 		case model.ExperimentType:
 			var configs expconf.ExperimentConfigV0
 			err = json.Unmarshal([]byte(*configPolicies.InvariantConfig), &configs)
 			if err != nil {
-				return 0, false, fmt.Errorf("unable to unmarshal task config policies: %w", err)
+				return nil, false, fmt.Errorf("unable to unmarshal task config policies: %w", err)
 			}
 			if configs.RawResources != nil && configs.RawResources.RawPriority != nil {
 				adminPriority := *configs.RawResources.RawPriority
-				return adminPriority, false,
+				return &adminPriority, false,
 					fmt.Errorf("priority set by invariant config: %w", errPriorityImmutable)
 			}
 		default:
-			return 0, false, fmt.Errorf("workload type %s not supported", workloadType)
+			return nil, false, fmt.Errorf("workload type %s not supported", workloadType)
 		}
 	}
 
 	// Find priority constraint, if set.
 	var constraints model.Constraints
 	if configPolicies.Constraints != nil {
 		if err = json.Unmarshal([]byte(*configPolicies.Constraints), &constraints); err != nil {
-			return 0, false, fmt.Errorf("unable to unmarshal task config policies: %w", err)
+			return nil, false, fmt.Errorf("unable to unmarshal task config policies: %w", err)
 		}
 		if constraints.PriorityLimit != nil {
-			return *constraints.PriorityLimit, true, nil
+			return constraints.PriorityLimit, true, nil
 		}
 	}
 
-	return 0, false, nil
+	return nil, false, nil
 }
 
 // PriorityUpdateAllowed returns true if the desired priority is within the task config policy limit.
 func PriorityUpdateAllowed(wkspID int, workloadType string, priority int, smallerHigher bool) (bool, error) {
 	// Check if a priority limit has been set with a constraint policy.
 	// Global policies have highest precedence.
-	globalLimit, globalExists, err := findAllowedPriority(nil, workloadType)
+	globalEnforcedPriority, globalExists, err := FindAllowedPriority(nil, workloadType)
 
-	if errors.Is(err, errPriorityImmutable) && globalLimit == priority {
+	if errors.Is(err, errPriorityImmutable) && globalEnforcedPriority != nil &&
+		*globalEnforcedPriority == priority {
 		// If task config policies have updated since the workload was originally scheduled, allow users
 		// to update the priority to the new priority set by invariant config.
 		return true, nil
@@ -298,9 +304,9 @@ func PriorityUpdateAllowed(wkspID int, workloadType string, priority int, smalle
 
 	// TODO use COALESCE instead once postgres updates are complete.
 	// Workspace policies have second precedence.
-	wkspLimit, wkspExists, err := findAllowedPriority(&wkspID, workloadType)
-
-	if errors.Is(err, errPriorityImmutable) && wkspLimit == priority {
+	wkspEnforcedPriority, wkspExists, err := FindAllowedPriority(&wkspID, workloadType)
+	if errors.Is(err, errPriorityImmutable) && wkspEnforcedPriority != nil &&
+		*wkspEnforcedPriority == priority {
 		// If task config policies have updated since the workload was originally scheduled, allow users
 		// to update the priority to the new priority set by invariant config.
 		return true, nil
@@ -311,10 +317,10 @@ func PriorityUpdateAllowed(wkspID int, workloadType string, priority int, smalle
 
 	// No invariant configs. Check for constraints.
 	if globalExists {
-		return priorityWithinLimit(priority, globalLimit, smallerHigher), nil
+		return priorityWithinLimit(priority, *globalEnforcedPriority, smallerHigher), nil
 	}
 	if wkspExists {
-		return priorityWithinLimit(priority, wkspLimit, smallerHigher), nil
+		return priorityWithinLimit(priority, *wkspEnforcedPriority, smallerHigher), nil
 	}
 
 	// No priority limit has been set.

master/internal/configpolicy/task_config_policy_intg_test.go

@@ -25,34 +25,37 @@ func TestFindAllowedPriority(t *testing.T) {
 	db.MustMigrateTestPostgres(t, pgDB, db.MigrationsFromDB)
 
 	// No priority limit to find.
-	_, exists, err := findAllowedPriority(nil, model.ExperimentType)
+	_, exists, err := FindAllowedPriority(nil, model.ExperimentType)
 	require.NoError(t, err)
 	require.False(t, exists)
 
 	// Priority limit set.
 	globalLimit := 10
 	user := db.RequireMockUser(t, pgDB)
 	addConstraints(t, user, nil, fmt.Sprintf(`{"priority_limit": %d}`, globalLimit), model.ExperimentType)
-	limit, exists, err := findAllowedPriority(nil, model.ExperimentType)
-	require.Equal(t, globalLimit, limit)
+	limit, exists, err := FindAllowedPriority(nil, model.ExperimentType)
+	require.NotNil(t, limit)
+	require.Equal(t, globalLimit, *limit)
 	require.True(t, exists)
 	require.NoError(t, err)
 
 	// NTSC priority set.
 	configPriority := 15
 	invariantConfig := fmt.Sprintf(`{"resources": {"priority": %d}}`, configPriority)
 	addConfig(t, user, nil, invariantConfig, model.NTSCType)
-	limit, _, err = findAllowedPriority(nil, model.NTSCType)
+	limit, _, err = FindAllowedPriority(nil, model.NTSCType)
 	require.ErrorIs(t, err, errPriorityImmutable)
-	require.Equal(t, configPriority, limit)
+	require.NotNil(t, limit)
+	require.Equal(t, configPriority, *limit)
 
 	// Experiment priority set.
 	configPriority = 7
 	invariantConfig = fmt.Sprintf(`{"resources": {"priority": %d}}`, configPriority)
 	addConfig(t, user, nil, invariantConfig, model.ExperimentType)
-	limit, _, err = findAllowedPriority(nil, model.ExperimentType)
+	limit, _, err = FindAllowedPriority(nil, model.ExperimentType)
 	require.ErrorIs(t, err, errPriorityImmutable)
-	require.Equal(t, configPriority, limit)
+	require.NotNil(t, limit)
+	require.Equal(t, configPriority, *limit)
 }
 
 func TestPriorityUpdateAllowed(t *testing.T) {

master/internal/core_experiment.go

@@ -349,10 +349,22 @@ func (m *Master) parseCreateExperiment(ctx context.Context, req *apiv1.CreateExp
 		config.RawCheckpointStorage, &m.config.CheckpointStorage,
 	)
 
-	// Apply the scheduler's default priority.
+	// Apply the scheduler's default priority if priority is not set in an invariant config or
+	// constraint.
 	if config.Resources().Priority() == nil {
-		prio := masterConfig.DefaultPriorityForPool(poolName.String())
-		config.RawResources.RawPriority = &prio
+		enforcedPriority, constraint, err := configpolicy.FindAllowedPriority(
+			&workspaceID,
+			model.ExperimentType,
+		)
+		if err != nil {
+			return nil, nil, config, nil, nil, err
+		}
+		if enforcedPriority == nil || constraint {
+			prio := masterConfig.DefaultPriorityForPool(poolName.String())
+			config.RawResources.RawPriority = &prio
+		} else {
+			config.RawResources.RawPriority = enforcedPriority
+		}
 	}
 
 	// Lastly, apply any json-schema-defined defaults.