Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 16 additions & 0 deletions runs/service/run_service.go
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@ import (

"connectrpc.com/connect"
semver "github.com/Masterminds/semver/v3"
"github.com/flyteorg/flyte/v2/gen/go/flyteidl2/settings"
"golang.org/x/sync/errgroup"
"google.golang.org/protobuf/proto"
"google.golang.org/protobuf/types/known/timestamppb"
Expand All @@ -40,6 +41,7 @@ import (
// RunService implements the RunServiceHandler interface
type RunService struct {
repo interfaces.Repository
settingsRepo interfaces.SettingsRepo
actionsClient actionsconnect.ActionsServiceClient
projectClient projectconnect.ProjectServiceClient
storagePrefix string
Expand Down Expand Up @@ -143,6 +145,7 @@ func (s *RunService) WatchGroups(ctx context.Context, req *connect.Request[workf
// NewRunService creates a new RunService instance
func NewRunService(
repo interfaces.Repository,
settingsRepo interfaces.SettingsRepo,
actionsClient actionsconnect.ActionsServiceClient,
projectClient projectconnect.ProjectServiceClient,
storagePrefix string,
Expand All @@ -154,6 +157,7 @@ func NewRunService(
) *RunService {
return &RunService{
repo: repo,
settingsRepo: settingsRepo,
actionsClient: actionsClient,
projectClient: projectClient,
storagePrefix: storagePrefix,
Expand Down Expand Up @@ -276,6 +280,18 @@ func (s *RunService) CreateRun(
}
request.RunSpec = runSpec

// Settings sit between an explicit request value and the static config defaults
// applied below.
resolved, err := resolveSettings(ctx, s.settingsRepo, &settings.SettingsKey{
Org: runId.GetOrg(),

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The OSS decision was to not support orgs, ( #7157 removed it from every table and service, and the backend is meant to ignore it). The settings key shouldnt take org from the request. The SDK derives an org from the endpoint hostname, so a named run can look up v1:<host>:: and miss the row scheduled runs hit.
Maybe EncodeSettingsKey can ignore org, so settings are stored and looked up under the same key no matter what the client sends

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Makes sense, thanks for the review! Unnamed runs drop the org as well, so the fix covers that path too. I'll make EncodeSettingsKey ignore org in a separate PR and remove the org note from this one.

Domain: runId.GetDomain(),
Project: runId.GetProject(),
})
if err != nil {
return nil, connect.NewError(connect.CodeInternal, err)
}
applyRunSettings(runSpec, resolved)

// Stamp the run start time, but only for SDKs that understand it (>= 2.3.6) — older task
// templates have no {{.runStartTime}} placeholder, so leaving it unset keeps the executor from
// substituting anything. The scheduler sets CreateRunRequest.run_start_time to a trigger's
Expand Down
94 changes: 94 additions & 0 deletions runs/service/run_service_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -13,10 +13,12 @@ import (
"time"

"connectrpc.com/connect"
"github.com/flyteorg/flyte/v2/gen/go/flyteidl2/settings"
"github.com/golang/protobuf/proto"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/mock"
"github.com/stretchr/testify/require"
"google.golang.org/protobuf/encoding/protojson"
"google.golang.org/protobuf/types/known/timestamppb"

"github.com/flyteorg/flyte/v2/flytestdlib/storage"
Expand All @@ -36,6 +38,32 @@ import (
"github.com/flyteorg/flyte/v2/runs/repository/models"
)

// noSettings returns a settings repo reporting no stored rows, so run creation
// resolves to empty settings and the applier changes nothing.
func noSettings(t *testing.T) *repoMocks.SettingsRepo {
m := &repoMocks.SettingsRepo{}
m.On("GetSettingsByKeys", mock.Anything, mock.Anything).Return(nil, nil)
return m
}

// settingsWithQueue returns a settings repo holding one org-level row that sets the
// default queue. The row key must match what fetchLevels asks for, or the lookup
// aligns to nothing.
func settingsWithQueue(t *testing.T, org, queue string) *repoMocks.SettingsRepo {
t.Helper()
data, err := protojson.Marshal(&settings.Settings{
Run: &settings.RunSettings{
DefaultQueue: &settings.StringSetting{State: stateValue, StringValue: queue},
},
})
require.NoError(t, err)

m := &repoMocks.SettingsRepo{}
m.On("GetSettingsByKeys", mock.Anything, mock.Anything).
Return([]*models.Settings{{Key: models.EncodeSettingsKey(org, "", ""), Data: data, Version: 1}}, nil)
return m
}

// newMockProjectClientAlwaysOK returns a mock ProjectServiceClient whose GetProject always succeeds.
func newMockProjectClientAlwaysOK(t *testing.T) *projectMocks.ProjectServiceClient {
pc := projectMocks.NewProjectServiceClient(t)
Expand Down Expand Up @@ -546,6 +574,7 @@ func TestCreateRunResponseIncludesMetadataAndStatus(t *testing.T) {

svc := &RunService{
repo: repo,
settingsRepo: noSettings(t),
actionsClient: actionsClient,
projectClient: newMockProjectClientAlwaysOK(t),
storagePrefix: "s3://flyte-data",
Expand Down Expand Up @@ -1068,6 +1097,7 @@ func TestCreateRun_WritesEmptyInputsProto(t *testing.T) {

svc := &RunService{
repo: repo,
settingsRepo: noSettings(t),
actionsClient: actionsClient,
projectClient: newMockProjectClientAlwaysOK(t),
storagePrefix: "s3://flyte-data",
Expand Down Expand Up @@ -1126,6 +1156,7 @@ func TestCreateRun_ResponseUsesRunModel(t *testing.T) {

svc := &RunService{
repo: repo,
settingsRepo: noSettings(t),
actionsClient: actionsClient,
projectClient: newMockProjectClientAlwaysOK(t),
storagePrefix: "s3://flyte-data",
Expand Down Expand Up @@ -1190,6 +1221,7 @@ func TestCreateRun_TriggerFire_CarriesRunSpecEnvVars(t *testing.T) {

svc := &RunService{
repo: repo,
settingsRepo: noSettings(t),
actionsClient: actionsClient,
projectClient: newMockProjectClientAlwaysOK(t),
storagePrefix: "s3://flyte-data",
Expand Down Expand Up @@ -1287,6 +1319,7 @@ func TestCreateRun_ActionIDUsesRunName(t *testing.T) {

svc := &RunService{
repo: repo,
settingsRepo: noSettings(t),
actionsClient: actionsClient,
projectClient: newMockProjectClientAlwaysOK(t),
storagePrefix: "s3://flyte-data",
Expand Down Expand Up @@ -1377,6 +1410,7 @@ func TestCreateRun_PreservesInputContextAndRawDataPath(t *testing.T) {

svc := &RunService{
repo: repo,
settingsRepo: noSettings(t),
actionsClient: actionsClient,
projectClient: newMockProjectClientAlwaysOK(t),
storagePrefix: "s3://flyte-data",
Expand Down Expand Up @@ -1434,6 +1468,65 @@ func TestCreateRun_PreservesInputContextAndRawDataPath(t *testing.T) {
require.NoError(t, err)
}

// TestCreateRun_AppliesSettingsQueue proves the applier is actually wired into run
// creation: the request names no queue, and the value stored on the run comes from
// the org's settings row.
func TestCreateRun_AppliesSettingsQueue(t *testing.T) {
actionRepo := &repoMocks.ActionRepo{}
taskRepo := &repoMocks.TaskRepo{}
actionsClient := actionsconnectmocks.NewActionsServiceClient(t)
repo := &repoMocks.Repository{}
store := &storageMocks.ComposedProtobufStore{}
dataStore := &storage.DataStore{ComposedProtobufStore: store}

repo.On("ActionRepo").Return(actionRepo)
repo.On("TaskRepo").Return(taskRepo)

svc := &RunService{
repo: repo,
settingsRepo: settingsWithQueue(t, "org", "fast-queue"),
actionsClient: actionsClient,
projectClient: newMockProjectClientAlwaysOK(t),
storagePrefix: "s3://flyte-data",
dataStore: dataStore,
}

req := &workflow.CreateRunRequest{
Id: &workflow.CreateRunRequest_RunId{
RunId: &common.RunIdentifier{
Org: "org",
Project: "proj",
Domain: "dev",
Name: "rq-123",
},
},
InputWrapper: &workflow.CreateRunRequest_Inputs{Inputs: &task.Inputs{}},
Task: &workflow.CreateRunRequest_TaskSpec{
TaskSpec: &task.TaskSpec{},
},
}

store.On("WriteProtobuf", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nil).Once()
taskRepo.On("CreateTaskSpec", mock.Anything, mock.Anything).Return(nil).Once()

actionRepo.On("CreateAction", mock.Anything, mock.MatchedBy(func(m *models.Run) bool {
var rs task.RunSpec
_ = proto.Unmarshal(m.RunSpec, &rs)
return rs.GetQueue() == "fast-queue"
}), mock.Anything).Return(&models.Run{
Project: "proj",
Domain: "dev",
Name: "rq-123",
}, nil).Once()

actionsClient.On("Enqueue", mock.Anything, mock.MatchedBy(func(req *connect.Request[actions.EnqueueRequest]) bool {
return req.Msg.GetRunSpec().GetQueue() == "fast-queue"
})).Return(connect.NewResponse(&actions.EnqueueResponse{}), nil).Once()

_, err := svc.CreateRun(context.Background(), connect.NewRequest(req))
require.NoError(t, err)
}

func TestCreateRun_WithOffloadedInputData(t *testing.T) {
actionRepo := &repoMocks.ActionRepo{}
taskRepo := &repoMocks.TaskRepo{}
Expand All @@ -1447,6 +1540,7 @@ func TestCreateRun_WithOffloadedInputData(t *testing.T) {

svc := &RunService{
repo: repo,
settingsRepo: noSettings(t),
actionsClient: actionsClient,
projectClient: newMockProjectClientAlwaysOK(t),
storagePrefix: "s3://flyte-data",
Expand Down
28 changes: 28 additions & 0 deletions runs/service/settings_apply.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,28 @@
package service

import (
"github.com/flyteorg/flyte/v2/gen/go/flyteidl2/settings"
"github.com/flyteorg/flyte/v2/gen/go/flyteidl2/task"
)

// applyRunSettings fills fields the caller left unset with values resolved from
// settings. An explicit value in the request always wins, so this only ever fills
// gaps, and a setting that is INHERIT or UNSET contributes nothing.
func applyRunSettings(spec *task.RunSpec, resolved *settings.Settings) {
if spec == nil {
return
}

if spec.GetQueue() == "" && resolved.GetRun().GetDefaultQueue().GetState() ==
settings.SettingState_SETTING_STATE_VALUE {
spec.Queue = resolved.GetRun().GetDefaultQueue().GetStringValue()
}

// The proto defines 0 as unset for this field, so there is no explicit request for
// "unlimited" to override. The settings validator bounds the value to 0 or
// [2, MaxUint32], so narrowing to uint32 cannot overflow.
if concurrency := resolved.GetRun().GetMaxActionConcurrency(); spec.GetMaxActionConcurrency() == 0 &&
concurrency.GetState() == settings.SettingState_SETTING_STATE_VALUE {
spec.MaxActionConcurrency = uint32(concurrency.GetIntValue())
}
}
80 changes: 80 additions & 0 deletions runs/service/settings_apply_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,80 @@
package service

import (
"testing"

"github.com/flyteorg/flyte/v2/gen/go/flyteidl2/settings"
"github.com/flyteorg/flyte/v2/gen/go/flyteidl2/task"
"github.com/stretchr/testify/assert"
)

func runSettings(queue *settings.StringSetting, concurrency *settings.Int64Setting) *settings.Settings {
return &settings.Settings{
Run: &settings.RunSettings{DefaultQueue: queue, MaxActionConcurrency: concurrency},
}
}

func TestApplyRunSettings(t *testing.T) {
tests := []struct {
name string
spec *task.RunSpec
resolved *settings.Settings
wantQueue string
wantConcurrency uint32
}{
{
name: "empty queue takes the settings value",
spec: &task.RunSpec{},
resolved: runSettings(&settings.StringSetting{State: stateValue, StringValue: "fast-queue"}, nil),
wantQueue: "fast-queue",
},
{
name: "an explicit queue wins over settings",
spec: &task.RunSpec{Queue: "user-queue"},
resolved: runSettings(&settings.StringSetting{State: stateValue, StringValue: "fast-queue"}, nil),
wantQueue: "user-queue",
},
{
name: "a queue in INHERIT contributes nothing",
spec: &task.RunSpec{},
resolved: runSettings(&settings.StringSetting{State: stateInherit, StringValue: "fast-queue"}, nil),
wantQueue: "",
},
{
name: "zero concurrency takes the settings value",
spec: &task.RunSpec{},
resolved: runSettings(nil, &settings.Int64Setting{State: stateValue, IntValue: 5}),
wantConcurrency: 5,
},
{
name: "an explicit concurrency wins over settings",
spec: &task.RunSpec{MaxActionConcurrency: 3},
resolved: runSettings(nil, &settings.Int64Setting{State: stateValue, IntValue: 5}),
wantConcurrency: 3,
},
{
name: "concurrency in UNSET contributes nothing",
spec: &task.RunSpec{},
resolved: runSettings(nil, &settings.Int64Setting{State: stateUnset, IntValue: 5}),
wantConcurrency: 0,
},
{
name: "no settings at all",
spec: &task.RunSpec{},
resolved: &settings.Settings{},
},
{
name: "nil spec does not panic",
spec: nil,
resolved: runSettings(&settings.StringSetting{State: stateValue, StringValue: "fast-queue"}, nil),
},
}

for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
applyRunSettings(tt.spec, tt.resolved)
assert.Equal(t, tt.wantQueue, tt.spec.GetQueue())
assert.Equal(t, tt.wantConcurrency, tt.spec.GetMaxActionConcurrency())
})
}
}
6 changes: 4 additions & 2 deletions runs/setup.go
Original file line number Diff line number Diff line change
Expand Up @@ -97,6 +97,8 @@ func Setup(ctx context.Context, sc *app.SetupContext) error {
return fmt.Errorf("runs: failed to create repository: %w", err)
}

settingsRepo := impl.NewSettingsRepo(sc.DB)

// In unified mode, intra-service calls go through the same mux.
actionsServiceCfg := cfg.ActionsService
if sc.BaseURL != "" {
Expand Down Expand Up @@ -129,7 +131,7 @@ func Setup(ctx context.Context, sc *app.SetupContext) error {
return abortReconciler.Run(ctx)
})

runsSvc := service.NewRunService(repo, actionsClient, projectClient, cfg.StoragePrefix, sc.DataStore, abortReconciler, cfg.AuthMetadata.ExternalAuthServerBaseURL, cfg.TrustForwardedIdentityHeaders, cfg.IdentityHeaders)
runsSvc := service.NewRunService(repo, settingsRepo, actionsClient, projectClient, cfg.StoragePrefix, sc.DataStore, abortReconciler, cfg.AuthMetadata.ExternalAuthServerBaseURL, cfg.TrustForwardedIdentityHeaders, cfg.IdentityHeaders)
taskSvc := service.NewTaskService(repo, projectClient)

runsPath, runsHandler := workflowconnect.NewRunServiceHandler(runsSvc, connect.WithInterceptors(interceptors...))
Expand Down Expand Up @@ -178,7 +180,7 @@ func Setup(ctx context.Context, sc *app.SetupContext) error {
sc.Mux.Handle(projectPath, projectHandler)
logger.Infof(ctx, "Mounted ProjectService at %s", projectPath)

settingsSvc := service.NewSettingsService(impl.NewSettingsRepo(sc.DB))
settingsSvc := service.NewSettingsService(settingsRepo)
settingsPath, settingsHandler := settingsconnect.NewSettingsServiceHandler(settingsSvc, connect.WithInterceptors(otelInterceptor))
sc.Mux.Handle(settingsPath, settingsHandler)
logger.Infof(ctx, "Mounted SettingsService at %s", settingsPath)
Expand Down
7 changes: 5 additions & 2 deletions runs/test/api/setup_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -114,6 +114,9 @@ func TestMain(m *testing.M) {
exitCode = 1
return
}

settingsRepo := impl.NewSettingsRepo(testDB)

// Services validate project existence through the ProjectService client, so mount a
// real ProjectService on the same mux (mirrors production unified mode in setup.go).
endpointURL := fmt.Sprintf("http://localhost:%d", testPort)
Expand All @@ -123,7 +126,7 @@ func TestMain(m *testing.M) {

// Create RunService with a no-op actions client (points at test server; not used by watch tests)
actionsClient := actionsconnect.NewActionsServiceClient(http.DefaultClient, endpointURL)
runSvc := service.NewRunService(repo, actionsClient, projectClient, "", nil, nil, "", true, config.GetConfig().IdentityHeaders)
runSvc := service.NewRunService(repo, settingsRepo, actionsClient, projectClient, "", nil, nil, "", true, config.GetConfig().IdentityHeaders)

// Setup HTTP server
mux := http.NewServeMux()
Expand All @@ -139,7 +142,7 @@ func TestMain(m *testing.M) {
internalRunPath, internalRunHandler := workflowconnect.NewInternalRunServiceHandler(runSvc)
mux.Handle(internalRunPath, internalRunHandler)

settingsSvc := service.NewSettingsService(impl.NewSettingsRepo(testDB))
settingsSvc := service.NewSettingsService(settingsRepo)
settingsPath, settingsHandler := settingsconnect.NewSettingsServiceHandler(settingsSvc)
mux.Handle(settingsPath, settingsHandler)

Expand Down
Loading