Kunapuli/manage job tcp

master/internal/command/command.go

@@ -277,6 +277,7 @@ func (c *Command) garbageCollect() {
 	go DefaultCmdService.unregisterCommand(c.taskID)
 }
 
+// TODO update me to include constraints
 func (c *Command) setNTSCPriority(priority int, forward bool) error {
 	if forward {
 		switch err := c.rm.SetGroupPriority(sproto.SetGroupPriority{

master/internal/command/command_job_service.go

@@ -1,13 +1,16 @@
 package command
 
 import (
+	"context"
 	"fmt"
 
 	"google.golang.org/protobuf/types/known/timestamppb"
 
 	"github.com/determined-ai/determined/master/internal/config"
+	"github.com/determined-ai/determined/master/internal/configpolicy"
 	"github.com/determined-ai/determined/master/internal/rm/rmerrors"
 	"github.com/determined-ai/determined/master/internal/sproto"
+	"github.com/determined-ai/determined/master/pkg/model"
 	"github.com/determined-ai/determined/proto/pkg/jobv1"
 )
 
@@ -47,7 +50,26 @@ func (c *Command) SetJobPriority(priority int) error {
 	if priority < 1 || priority > 99 {
 		return fmt.Errorf("priority must be between 1 and 99")
 	}
-	err := c.setNTSCPriority(priority, true)
+
+	// check if a priority limit has been set in task config policies
+	wkspID := int(c.GenericCommandSpec.Metadata.WorkspaceID)
+	priorityLimit, found, err := configpolicy.GetPriorityLimitPrecedence(context.TODO(), wkspID, model.NTSCType)
+
+	// returns err if RM does not have concept of priority
+	smallerIsHigher, rmErr := c.rm.SmallerValueIsHigherPriority()
+	if found && rmErr == nil {
+
+		ok := configpolicy.PriorityOK(priority, priorityLimit, smallerIsHigher)
+		if !ok {
+			return fmt.Errorf("priority exceeds task config policy's priority_limit: %d", priorityLimit)
+		}
+
+	} else if err != nil {
+		// TODO do we really want to block on this?
+		return err
+	}
+
+	err = c.setNTSCPriority(priority, true)
 	if err != nil {
 		c.syslog.WithError(err).Info("setting command job priority")
 	}

master/internal/configpolicy/postgres_task_config_policy.go

@@ -2,6 +2,7 @@ package configpolicy
 
 import (
 	"context"
+	"database/sql"
 	"fmt"
 	"strings"
 
@@ -88,6 +89,39 @@ func GetNTSCConfigPolicies(ctx context.Context,
 	return &ntscTCP, nil
 }
 
+// GetPriorityLimit reads the priority limit for the given scope and workload type. It returns found=false if no limit exists.
+func GetPriorityLimit(ctx context.Context, scope *int, workloadType string) (limit int, found bool, err error) {
+
+	if !ValidWorkloadType(workloadType) {
+		return 0, false, fmt.Errorf("invalid workload type: %s", workloadType)
+	}
+
+	wkspQuery := wkspIDQuery
+	if scope == nil {
+		wkspQuery = wkspIDGlobalQuery
+	}
+
+	var constraints model.Constraints
+	err = db.Bun().NewSelect().
+		Table("task_config_policies").
+		Column("constraints").
+		Where(wkspQuery, scope).
+		Where("workload_type = ?", workloadType).
+		Scan(ctx, &constraints)
+
+	if err == sql.ErrNoRows {
+		return 0, false, nil
+	} else if err != nil {
+		return 0, false, fmt.Errorf("error retrieving priority limit: %w", err)
+	}
+
+	if constraints.PriorityLimit != nil {
+		return *constraints.PriorityLimit, true, nil
+	}
+
+	return 0, false, nil
+}
+
 // DeleteConfigPolicies deletes the invariant experiment config and constraints for the
 // given scope (global or workspace-level) and workload type.
 func DeleteConfigPolicies(ctx context.Context,

master/internal/configpolicy/postgres_task_config_policy_intg_test.go

@@ -207,6 +207,115 @@ func TestSetNTSCConfigPolicies(t *testing.T) {
 	require.ErrorContains(t, err, "violates foreign key constraint")
 }
 
+func TestWorkspaceGetPriorityLimit(t *testing.T) {
+	ctx := context.Background()
+	require.NoError(t, etc.SetRootPath(db.RootFromDB))
+	pgDB, cleanup := db.MustResolveNewPostgresDatabase(t)
+	defer cleanup()
+	db.MustMigrateTestPostgres(t, pgDB, db.MigrationsFromDB)
+	user := db.RequireMockUser(t, pgDB)
+
+	// add a workspace to use
+	w := model.Workspace{Name: uuid.NewString(), UserID: user.ID}
+	_, err := db.Bun().NewInsert().Model(&w).Exec(ctx)
+	require.NoError(t, err)
+	defer func() {
+		err := db.CleanupMockWorkspace([]int32{int32(w.ID)})
+		if err != nil {
+			log.Errorf("error when cleaning up mock workspaces")
+		}
+	}()
+
+	// add priority limit for workspace NTSC
+	wkspLimit := 20
+	wkspInput := model.NTSCTaskConfigPolicies{
+		WorkloadType: model.NTSCType,
+		WorkspaceID:  &w.ID,
+		Constraints: model.Constraints{
+			PriorityLimit: &wkspLimit,
+		},
+		LastUpdatedBy: user.ID,
+	}
+	err = SetNTSCConfigPolicies(ctx, &wkspInput)
+	require.NoError(t, err)
+
+	// get priority limit; should match workspace
+	res, found, err := GetPriorityLimit(ctx, &w.ID, model.NTSCType)
+	require.NoError(t, err)
+	require.True(t, found)
+	require.Equal(t, wkspLimit, res)
+
+	// get limit for workspace that does not exist
+	wkspIDDoesNotExist := 404
+	res, found, err = GetPriorityLimit(ctx, &wkspIDDoesNotExist, model.NTSCType)
+	require.NoError(t, err)
+	require.False(t, found)
+
+	// read global
+	res, found, err = GetPriorityLimit(ctx, nil, model.NTSCType)
+	require.NoError(t, err)
+	require.False(t, found)
+
+	// read experiment
+	res, found, err = GetPriorityLimit(ctx, &w.ID, model.ExperimentType)
+	require.NoError(t, err)
+	require.False(t, found)
+
+	// read invalid workload type
+	res, found, err = GetPriorityLimit(ctx, &w.ID, "bogus")
+	require.Error(t, err)
+	require.False(t, found)
+}
+
+func TestGlobalGetPriorityLimit(t *testing.T) {
+	ctx := context.Background()
+	require.NoError(t, etc.SetRootPath(db.RootFromDB))
+	pgDB, cleanup := db.MustResolveNewPostgresDatabase(t)
+	defer cleanup()
+	db.MustMigrateTestPostgres(t, pgDB, db.MigrationsFromDB)
+	user := db.RequireMockUser(t, pgDB)
+
+	// add a workspace to use
+	w := model.Workspace{Name: uuid.NewString(), UserID: user.ID}
+	_, err := db.Bun().NewInsert().Model(&w).Exec(ctx)
+	require.NoError(t, err)
+	defer func() {
+		err := db.CleanupMockWorkspace([]int32{int32(w.ID)})
+		if err != nil {
+			log.Errorf("error when cleaning up mock workspaces")
+		}
+	}()
+
+	// add priority limit for global NTSC
+	globalLimit := 5
+	wkspInput := model.NTSCTaskConfigPolicies{
+		WorkloadType: model.NTSCType,
+		WorkspaceID:  nil,
+		Constraints: model.Constraints{
+			PriorityLimit: &globalLimit,
+		},
+		LastUpdatedBy: user.ID,
+	}
+	err = SetNTSCConfigPolicies(ctx, &wkspInput)
+	require.NoError(t, err)
+
+	// get priority limit
+	res, found, err := GetPriorityLimit(ctx, nil, model.NTSCType)
+	require.NoError(t, err)
+	require.True(t, found)
+	require.Equal(t, globalLimit, res)
+
+	// read experiment
+	res, found, err = GetPriorityLimit(ctx, nil, model.ExperimentType)
+	require.NoError(t, err)
+	require.False(t, found)
+
+	// read invalid workload type
+	res, found, err = GetPriorityLimit(ctx, nil, "bogus")
+	require.Error(t, err)
+	require.False(t, found)
+}
+
 // Test the enforcement of the primary key on the task_config_polciies table.
 func TestTaskConfigPoliciesUnique(t *testing.T) {
 	ctx := context.Background()

master/internal/configpolicy/utils.go

@@ -2,6 +2,7 @@ package configpolicy
 
 import (
 	"bytes"
+	"context"
 	"encoding/json"
 	"fmt"
 
@@ -22,6 +23,32 @@ func ValidWorkloadType(val string) bool {
 	}
 }
 
+// GetPriorityLimitPrecedence retrieves the priority limit using order of precedence
+func GetPriorityLimitPrecedence(ctx context.Context, workspace_id int, workload_type string) (limit int, found bool, err error) {
+	// highest precedence: get global limit
+	if limit, found, err = GetPriorityLimit(ctx, nil, workload_type); found {
+		return limit, found, err
+	}
+
+	// second precedence: get workspace limit
+	if limit, found, err = GetPriorityLimit(ctx, &workspace_id, workload_type); found {
+		return limit, found, err
+	}
+
+	// default
+	return 0, false, nil
+}
+
+// PriorityOK returns true if the current priority is acceptable given the priority limit and resource manager.
+func PriorityOK(currPriority int, priorityLimit int, smallerValueIsHigherPriority bool) bool {
+
+	if smallerValueIsHigherPriority {
+		return currPriority >= priorityLimit
+	}
+
+	return currPriority <= priorityLimit
+}
+
 // UnmarshalExperimentConfigPolicy unpacks a string into ExperimentConfigPolicy struct.
 func UnmarshalExperimentConfigPolicy(str string) (*ExperimentConfigPolicies, error) {
 	var expConfigPolicy ExperimentConfigPolicies