This is an automated email from the ASF dual-hosted git repository. wilfred-s pushed a commit to branch branch-1.9 in repository https://gitbox.apache.org/repos/asf/yunikorn-k8shim.git
commit 4e52440d57ffa709a179eee1eaef52898f72c5ff Author: Peter Bacsko <[email protected]> AuthorDate: Tue Jun 30 11:13:43 2026 +1000 [YUNIKORN-3089] Fix stale New apps when pods are deleted before core accept (#1044) Defer task releases until the application is accepted, then remove empty apps from the core and shim cache instead of sending a lost allocation release during the New/Submitted window. Closes: #1044 Signed-off-by: Wilfred Spiegelenburg <[email protected]> (cherry picked from commit f0c70d44cafd2477d0ac023f9adda3adb58720d7) --- pkg/cache/application.go | 65 ++++++++++++++++++++++++++++++++++++++++ pkg/cache/application_state.go | 4 +++ pkg/cache/application_test.go | 67 ++++++++++++++++++++++++++++++++++++++++++ pkg/cache/context.go | 5 ++++ pkg/cache/task.go | 21 ++++++++++--- pkg/cache/task_test.go | 2 ++ 6 files changed, 160 insertions(+), 4 deletions(-) diff --git a/pkg/cache/application.go b/pkg/cache/application.go index 5ea488f5..44d0dc55 100644 --- a/pkg/cache/application.go +++ b/pkg/cache/application.go @@ -59,6 +59,8 @@ type Application struct { placeholderTimeoutInSec int64 schedulingStyle string originatingTask *Task // Original Pod which creates the requests + releaseableTasks []*Task + context *Context } const transitionErr = "no transition" @@ -85,6 +87,7 @@ func NewApplication(appID, queueName, user string, groups []string, tags map[str schedulerAPI: scheduler, placeholderTimeoutInSec: 0, schedulingStyle: constants.SchedulingPolicyStyleParamDefault, + releaseableTasks: make([]*Task, 0), } return app } @@ -208,6 +211,10 @@ func (app *Application) setOriginatingTask(task *Task) { app.originatingTask = task } +func (app *Application) setContext(ctx *Context) { + app.context = ctx +} + func (app *Application) GetOriginatingTask() *Task { app.lock.RLock() defer app.lock.RUnlock() @@ -578,6 +585,7 @@ func (app *Application) onReservationStateChange() { func (app *Application) handleRejectApplicationEvent(reason string) { log.Log(log.ShimCacheApplication).Info("app is rejected by scheduler", zap.String("appID", app.applicationID)) + app.clearReleaseableTasks() // for rejected apps, we directly move them to failed state dispatcher.Dispatch(NewFailApplicationEvent(app.applicationID, fmt.Sprintf("%s: %s", constants.ApplicationRejectedFailure, reason))) @@ -609,6 +617,7 @@ func (app *Application) handleFailApplicationEvent(errMsg string) { go func() { getPlaceholderManager().cleanUp(app) }() + app.clearReleaseableTasks() log.Log(log.ShimCacheApplication).Info("failApplication reason", zap.String("applicationID", app.applicationID), zap.String("errMsg", errMsg)) // unallocated task states include New, Pending and Scheduling unalloc := app.getTasks(TaskStates().New) @@ -690,3 +699,59 @@ func (app *Application) removeCompletedTasks() { } } } + +func (app *Application) tryAddReleasableTask(task *Task) bool { + app.lock.Lock() + defer app.lock.Unlock() + + current := app.sm.Current() + if current == ApplicationStates().New || + current == ApplicationStates().Submitted { + for _, existing := range app.releaseableTasks { + if existing.taskID == task.taskID { + return true + } + } + app.releaseableTasks = append(app.releaseableTasks, task) + return true + } + + return false +} + +func (app *Application) clearReleaseableTasks() { + app.releaseableTasks = nil +} + +// flushReleaseableTasks replays deferred task releases after the application has been accepted +// by the scheduler core. Must be called while the application lock is held. +func (app *Application) flushReleaseableTasks() { + if len(app.releaseableTasks) == 0 { + return + } + tasks := app.releaseableTasks + app.releaseableTasks = nil + + if app.AreAllTasksTerminated() { + app.removeFromSchedulerCore() + if app.context != nil { + app.context.removeApplication(app.applicationID) + } + return + } + + for _, task := range tasks { + task.releaseAllocation(true) + } +} + +func (app *Application) removeFromSchedulerCore() { + log.Log(log.ShimCacheApplication).Info("removing application from scheduler core", + zap.String("appID", app.applicationID)) + request := common.CreateUpdateRequestForRemoveApplication(app.applicationID, app.partition) + if err := app.schedulerAPI.UpdateApplication(request); err != nil { + log.Log(log.ShimCacheApplication).Warn("failed to remove application from scheduler core", + zap.String("appID", app.applicationID), + zap.Error(err)) + } +} diff --git a/pkg/cache/application_state.go b/pkg/cache/application_state.go index d22fde61..14c71860 100644 --- a/pkg/cache/application_state.go +++ b/pkg/cache/application_state.go @@ -455,6 +455,10 @@ func newAppState() *fsm.FSM { //nolint:funlen zap.String("destination", event.Dst), zap.String("event", event.Event)) }, + states.Accepted: func(_ context.Context, event *fsm.Event) { + app := event.Args[0].(*Application) //nolint:errcheck + app.flushReleaseableTasks() + }, states.Reserving: func(_ context.Context, event *fsm.Event) { app := event.Args[0].(*Application) //nolint:errcheck app.onReserving() diff --git a/pkg/cache/application_test.go b/pkg/cache/application_test.go index f54bf79c..0fffcc8e 100644 --- a/pkg/cache/application_test.go +++ b/pkg/cache/application_test.go @@ -1328,7 +1328,74 @@ func TestTaskRemoval(t *testing.T) { assert.Equal(t, 0, len(app.getTasks(TaskStates().Completed))) } +func TestDeferredReleaseOnAccept(t *testing.T) { + context := initContextForTest() + mockedAPI, ok := context.apiProvider.(*client.MockedAPIProvider) + assert.Assert(t, ok, "expecting MockedAPIProvider") + + removeCalled := false + mockScheduler := newMockSchedulerAPI() + mockScheduler.UpdateApplicationFn = func(request *si.ApplicationRequest) error { + if len(request.New) > 0 { + return nil + } + if len(request.Remove) == 1 { + removeCalled = true + assert.Equal(t, request.Remove[0].ApplicationID, "deferred-app") + assert.Equal(t, request.Remove[0].PartitionName, constants.DefaultPartition) + } + return nil + } + + mockedAPI.MockSchedulerAPIUpdateAllocationFn(func(request *si.AllocationRequest) error { + t.Fatal("unexpected allocation update during deferred release flow") + return nil + }) + + pod := &v1.Pod{ + TypeMeta: apis.TypeMeta{ + Kind: "Pod", + APIVersion: "v1", + }, + ObjectMeta: apis.ObjectMeta{ + Name: "pod-deferred-release", + UID: "task-deferred-01", + }, + Spec: v1.PodSpec{}, + } + + app := NewApplication("deferred-app", "root.default", "testuser", testGroups, map[string]string{}, mockScheduler) + context.addApplicationToContext(app) + task := NewTask("task-deferred-01", app, context, pod) + app.addTask(task) + + err := app.handle(NewSubmitApplicationEvent(app.applicationID)) + assert.NilError(t, err) + assert.Equal(t, app.GetApplicationState(), ApplicationStates().Submitted) + + err = task.handle(NewSimpleTaskEvent(app.applicationID, task.taskID, CompleteTask)) + assert.NilError(t, err) + assert.Equal(t, task.GetTaskState(), TaskStates().Completed) + assert.Equal(t, mockedAPI.GetSchedulerAPIUpdateAllocationCount(), int32(0)) + + err = app.handle(NewSimpleApplicationEvent(app.applicationID, AcceptApplication)) + assert.NilError(t, err) + assert.Equal(t, app.GetApplicationState(), ApplicationStates().Accepted) + assert.Assert(t, removeCalled, "expected RemoveApplication request to scheduler core") + assert.Assert(t, context.GetApplication(app.applicationID) == nil, "app should be removed from shim cache") +} + +func TestTryAddReleasableTaskDedupe(t *testing.T) { + app := NewApplication("app-dedupe", "root.default", "testuser", testGroups, map[string]string{}, newMockSchedulerAPI()) + task := &Task{taskID: "task01"} + + assert.Assert(t, app.tryAddReleasableTask(task)) + assert.Assert(t, app.tryAddReleasableTask(task)) + assert.Equal(t, len(app.releaseableTasks), 1) +} + func (ctx *Context) addApplicationToContext(app *Application) { + app.setContext(ctx) ctx.lock.Lock() defer ctx.lock.Unlock() ctx.applications[app.applicationID] = app diff --git a/pkg/cache/context.go b/pkg/cache/context.go index fc742129..5257974f 100644 --- a/pkg/cache/context.go +++ b/pkg/cache/context.go @@ -969,6 +969,7 @@ func (ctx *Context) addApplication(request *AddApplicationRequest) *Application request.Metadata.Groups, request.Metadata.Tags, ctx.apiProvider.GetAPIs().SchedulerAPI) + app.setContext(ctx) app.setTaskGroups(request.Metadata.TaskGroups) app.setTaskGroupsDefinition(request.Metadata.Tags[constants.AnnotationTaskGroups]) app.setSchedulingParamsDefinition(request.Metadata.Tags[constants.AnnotationSchedulingPolicyParam]) @@ -1018,6 +1019,10 @@ func (ctx *Context) getApplication(appID string) *Application { func (ctx *Context) RemoveApplication(appID string) { ctx.lock.Lock() defer ctx.lock.Unlock() + ctx.removeApplication(appID) +} + +func (ctx *Context) removeApplication(appID string) { if _, exist := ctx.applications[appID]; !exist { log.Log(log.ShimContext).Debug("Attempted to remove non-existent application", zap.String("appID", appID)) return diff --git a/pkg/cache/task.go b/pkg/cache/task.go index 6c00077e..e68b7b54 100644 --- a/pkg/cache/task.go +++ b/pkg/cache/task.go @@ -404,7 +404,7 @@ func (task *Task) beforeTaskAllocated(eventSrc string, allocationKey string, nod zap.String("currentTaskState", eventSrc), zap.String("allocationKey", allocationKey), zap.String("allocatedNode", nodeID)) - task.releaseAllocation() + task.releaseAllocation(false) } } @@ -436,23 +436,30 @@ func (task *Task) beforeTaskFail() { events.GetRecorder().Eventf(task.pod.DeepCopy(), nil, v1.EventTypeNormal, "TaskFailed", "TaskFailed", "Task %s is failed", task.alias) - task.releaseAllocation() + task.releaseAllocation(false) } // beforeTaskCompleted releases the allocation or ask from scheduler core // this is done as a before hook because the releaseAllocation() call needs to // send different requests to scheduler-core, depending on current task state func (task *Task) beforeTaskCompleted() { + task.releaseAllocation(false) + events.GetRecorder().Eventf(task.pod.DeepCopy(), nil, v1.EventTypeNormal, "TaskCompleted", "TaskCompleted", "Task %s is completed", task.alias) - task.releaseAllocation() } // releaseAllocation sends the release request for the Allocation to the core. -func (task *Task) releaseAllocation() { +func (task *Task) releaseAllocation(force bool) { terminationType := common.GetTerminationTypeFromString(task.terminationType) + if !force && task.shouldAppRelease() { + log.Log(log.ShimCacheTask).Info("not releasing task right now, app has not been accepted", + zap.String("appState", task.application.sm.Current())) + return + } + // scheduler api might be nil in some tests if task.context.apiProvider.GetAPIs().SchedulerAPI != nil { log.Log(log.ShimCacheTask).Debug("prepare to send release request", @@ -498,6 +505,12 @@ func (task *Task) releaseAllocation() { } } +func (task *Task) shouldAppRelease() bool { + task.lock.Unlock() + defer task.lock.Lock() + return task.application.tryAddReleasableTask(task) +} + // some sanity checks before sending task for scheduling, // this reduces the scheduling overhead by blocking such // request away from the core scheduler. diff --git a/pkg/cache/task_test.go b/pkg/cache/task_test.go index 190055e1..c8a51609 100644 --- a/pkg/cache/task_test.go +++ b/pkg/cache/task_test.go @@ -224,6 +224,7 @@ func TestReleaseTaskAllocation(t *testing.T) { }) // complete + task.application.sm.SetState(ApplicationStates().Running) event4 := NewSimpleTaskEvent(app.applicationID, task.taskID, CompleteTask) err = task.handle(event4) assert.NilError(t, err, "failed to handle CompleteTask event") @@ -324,6 +325,7 @@ func TestReleaseTaskAsk(t *testing.T) { }) // complete + task.application.sm.SetState(ApplicationStates().Running) event4 := NewSimpleTaskEvent(app.applicationID, task.taskID, CompleteTask) err = task.handle(event4) assert.NilError(t, err, "failed to handle CompleteTask event") --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
