feat(clustered): resume a preempted gang when Kueue readmits it

flyteidl2/plugins/clustered.proto

@@ -77,7 +77,8 @@ enum Interconnect {
   ROCE = 3;
 }
 
-// Controls how the JobSet responds to rank failures and host-maintenance evictions.
+// Controls how the JobSet responds to rank failures, host-maintenance evictions and
+// Kueue preemptions.
 message ClusterFailurePolicy {
   // Maximum number of whole-cluster restarts before the JobSet is marked failed.
   // Each restart increments JOBSET_RESTART_ATTEMPT and rotates the torchrun RDZV_ID.
@@ -86,4 +87,9 @@ message ClusterFailurePolicy {
   // If true, pods evicted due to host maintenance (DisruptionTarget condition)
   // trigger a free restart that does not count against max_restarts.
   bool restart_on_host_maintenance = 2;
+
+  // If true, a gang that Kueue preempts after it started keeps its attempt and JobSet
+  // and resumes when Kueue re-admits it, so a preemption costs no retry. The task's
+  // max_runtime keeps counting while it waits. If false, the preemption ends the attempt.
+  bool resume_after_preemption = 3;
 }

flyteplugins/go/tasks/pluginmachinery/gang/gang.go

@@ -57,7 +57,8 @@ func (e Eviction) ExecutionError() *core.ExecutionError {
 	}
 }
 
-func (e Eviction) describe() string {
+// Summary says who took the gang back and why, without what that costs the task.
+func (e Eviction) Summary() string {
 	var b strings.Builder
 	b.WriteString("gang evicted")
 	if e.Source != "" {
@@ -74,6 +75,12 @@ func (e Eviction) describe() string {
 		b.WriteString(": ")
 		b.WriteString(msg)
 	}
+	return b.String()
+}
+
+func (e Eviction) describe() string {
+	var b strings.Builder
+	b.WriteString(e.Summary())
 	switch {
 	case !e.UserCaused:
 		b.WriteString("; this does not count against the task's retries")

flyteplugins/go/tasks/pluginmachinery/gang/gang_test.go

@@ -24,6 +24,11 @@ func TestEvictionExecutionError_Kind(t *testing.T) {
 	assert.Equal(t, core.ExecutionError_SYSTEM, system.GetKind())
 }
 
+func TestEvictionSummary_LeavesOutTheCost(t *testing.T) {
+	eviction := Eviction{Source: kueueSource, Reason: preempted, Message: preemptedMessage, UserCaused: true, RetriesLeft: 2}
+	assert.Equal(t, "gang evicted by kueue (Preempted): "+preemptedMessage, eviction.Summary())
+}
+
 func TestEvictionExecutionError_Message(t *testing.T) {
 	tests := []struct {
 		name     string

flyteplugins/go/tasks/plugins/k8s/clustered/kueue.go

@@ -3,6 +3,7 @@ package clustered
 import (
 	"context"
 	"fmt"
+	"strings"
 
 	"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
 	"sigs.k8s.io/controller-runtime/pkg/client"
@@ -29,6 +30,10 @@ const (
 	evictedByPreemption           = "Preempted"
 	evictedByMaximumExecutionTime = "DeactivatedDueToMaximumExecutionTimeExceeded"
 
+	// evictedByDeactivationPrefix starts every Evicted reason Kueue gives when it deactivates
+	// a Workload, e.g. "DeactivatedDueToRequeuingLimitExceeded".
+	evictedByDeactivationPrefix = "Deactivated"
+
 	quotaReservedInadmissible = "Inadmissible"
 )
 
@@ -94,7 +99,15 @@ func readWorkloadCondition(wl *unstructured.Unstructured, conditionType string)
 func classifyEviction(wl *unstructured.Unstructured) gang.Eviction {
 	eviction := gang.Eviction{Source: evictionSource}
 	evicted, ok := readWorkloadCondition(wl, workloadConditionEvicted)
-	if !ok || evicted.status != "True" {
+	evictedNow := ok && evicted.status == "True"
+	// Deactivating a Workload that is already evicted, e.g. one waiting for re-admission
+	// after a preemption, can leave the earlier reason in place: the deactivation is the
+	// cause now.
+	if isInactive(wl) && !(evictedNow && strings.HasPrefix(evicted.reason, evictedByDeactivationPrefix)) {
+		eviction.Reason = evictedByDeactivationPrefix
+		return eviction
+	}
+	if !evictedNow {
 		return eviction
 	}
 	eviction.Reason = evicted.reason
@@ -111,6 +124,29 @@ func classifyEviction(wl *unstructured.Unstructured) gang.Eviction {
 	return eviction
 }
 
+// awaitsReadmission reports whether a gang Kueue took back was preempted, and so will be
+// readmitted once there is room again. Any other eviction, such as node failures or a
+// PodsReady timeout, may recur, so it is not worth waiting for; nor is a deactivated
+// Workload or a missing one. A JobSet can be suspended before Kueue has written the
+// Evicted condition, which is waited out until the condition says why.
+func awaitsReadmission(wl *unstructured.Unstructured) bool {
+	if wl == nil || isInactive(wl) {
+		return false
+	}
+	evicted, ok := readWorkloadCondition(wl, workloadConditionEvicted)
+	return !ok || evicted.status != "True" || evicted.reason == evictedByPreemption
+}
+
+// isInactive reports whether the Workload has been deactivated (spec.active: false), which
+// Kueue does on its own limits and an operator does by hand.
+func isInactive(wl *unstructured.Unstructured) bool {
+	if wl == nil {
+		return false
+	}
+	active, found, err := unstructured.NestedBool(wl.Object, "spec", "active")
+	return err == nil && found && !active
+}
+
 // missingLocalQueue reports the LocalQueue the JobSet names when Kueue cannot admit it because
 // that queue does not exist. Only that case is a user error: a queue that exists but is stopped
 // or inactive is the platform's, so the gang keeps waiting. The message is Kueue's own

flyteplugins/go/tasks/plugins/k8s/clustered/kueue_test.go

@@ -22,6 +22,7 @@ import (
 	plugink8s "github.com/flyteorg/flyte/v2/flyteplugins/go/tasks/pluginmachinery/k8s"
 	k8smocks "github.com/flyteorg/flyte/v2/flyteplugins/go/tasks/pluginmachinery/k8s/mocks"
 	"github.com/flyteorg/flyte/v2/gen/go/flyteidl2/core"
+	clusteredpb "github.com/flyteorg/flyte/v2/gen/go/flyteidl2/plugins"
 )
 
 const (
@@ -69,10 +70,25 @@ func readerWith(t *testing.T, objects ...client.Object) client.Reader {
 // kueuePluginCtx is a plugin context for an attempt of a task with maxAttempts attempts.
 func kueuePluginCtx(
 	reader client.Reader, state plugink8s.PluginState, retryAttempt, maxAttempts uint32,
+) *k8smocks.PluginContext {
+	return kueuePluginCtxFor(twoNodeSpec(), reader, state, retryAttempt, maxAttempts)
+}
+
+// resumingPluginCtx is kueuePluginCtx for a task that resumes after a preemption.
+func resumingPluginCtx(
+	reader client.Reader, state plugink8s.PluginState, retryAttempt, maxAttempts uint32,
+) *k8smocks.PluginContext {
+	spec := twoNodeSpecWithPolicy(&clusteredpb.ClusterFailurePolicy{ResumeAfterPreemption: true})
+	return kueuePluginCtxFor(spec, reader, state, retryAttempt, maxAttempts)
+}
+
+func kueuePluginCtxFor(
+	taskTemplate *core.TaskTemplate,
+	reader client.Reader, state plugink8s.PluginState, retryAttempt, maxAttempts uint32,
 ) *k8smocks.PluginContext {
 	pCtx := &k8smocks.PluginContext{}
 	taskReader := &coreMocks.TaskReader{}
-	taskReader.EXPECT().Read(mock.Anything).Return(twoNodeSpec(), nil)
+	taskReader.EXPECT().Read(mock.Anything).Return(taskTemplate, nil)
 	pCtx.EXPECT().TaskReader().Return(taskReader)
 	pCtx.EXPECT().K8sReader().Return(reader)
 
@@ -156,6 +172,15 @@ func TestGetTaskPhase_Evicted_ClassifiedByWorkloadReason(t *testing.T) {
 			wantKind:    core.ExecutionError_SYSTEM,
 			wantInMsg:   []string{"(Deactivated): The workload is deactivated", "does not count against the task's retries"},
 		},
+		{
+			name: "deactivating it while a stale preemption reason is in place is a system retry",
+			workload: deactivated(testJobSetUID,
+				workloadCond(workloadConditionEvicted, "True", evictedByPreemption, preemptedMessage),
+				workloadCond(workloadConditionPreempted, "True", "InClusterQueue", preemptedMessage)),
+			maxAttempts: 3,
+			wantKind:    core.ExecutionError_SYSTEM,
+			wantInMsg:   []string{"gang evicted by kueue (Deactivated); this does not count against the task's retries"},
+		},
 		{
 			name: "a stopped queue is a system retry",
 			workload: workload(testJobSetUID,
@@ -201,6 +226,176 @@ func TestGetTaskPhase_Evicted_ClassifiedByWorkloadReason(t *testing.T) {
 	}
 }
 
+// deactivated is a deactivated Workload (spec.active: false) with the given conditions.
+func deactivated(jobSetUID k8stypes.UID, conditions ...map[string]interface{}) *unstructured.Unstructured {
+	wl := workload(jobSetUID, conditions...)
+	wl.Object["spec"] = map[string]interface{}{"active": false}
+	return wl
+}
+
+func TestGetTaskPhase_Resume_WaitsForReadmission(t *testing.T) {
+	tests := []struct {
+		name       string
+		workload   *unstructured.Unstructured
+		wantReason string
+	}{
+		{
+			name: "preemption",
+			workload: workload(testJobSetUID,
+				workloadCond(workloadConditionEvicted, "True", evictedByPreemption, preemptedMessage),
+				workloadCond(workloadConditionPreempted, "True", "InClusterQueue", preemptedMessage)),
+			wantReason: "gang evicted by kueue (Preempted, InClusterQueue): " + preemptedMessage +
+				"; waiting for re-admission",
+		},
+		{
+			name:       "no Evicted condition yet",
+			workload:   workload(testJobSetUID),
+			wantReason: "gang evicted by kueue; waiting for re-admission",
+		},
+		{
+			name: "Evicted condition from an earlier eviction not True yet",
+			workload: workload(testJobSetUID,
+				workloadCond(workloadConditionEvicted, "False", "NodeFailures", "")),
+			wantReason: "gang evicted by kueue; waiting for re-admission",
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			// Given
+			pCtx := resumingPluginCtx(readerWith(t, tt.workload), startedState, 0, 1)
+
+			// When
+			phase, err := clusteredResourceHandler{}.GetTaskPhase(context.Background(), pCtx, suspendedJobSet())
+
+			// Then
+			require.NoError(t, err)
+			assert.Equal(t, pluginsCore.PhaseRunning, phase.Phase())
+			assert.Equal(t, tt.wantReason, phase.Reason())
+			assert.Nil(t, phase.Err())
+			assert.False(t, phase.CleanupOnFailure())
+		})
+	}
+}
+
+func TestGetTaskPhase_Resume_EndsAttemptUnlessPreempted(t *testing.T) {
+	tests := []struct {
+		name      string
+		workload  *unstructured.Unstructured
+		wantKind  core.ExecutionError_ErrorKind
+		wantInMsg string
+	}{
+		{
+			name: "pods ready timeout",
+			workload: workload(testJobSetUID,
+				workloadCond(workloadConditionEvicted, "True", "PodsReadyTimeout", "Exceeded the PodsReady timeout")),
+			wantKind:  core.ExecutionError_SYSTEM,
+			wantInMsg: "gang evicted by kueue (PodsReadyTimeout)",
+		},
+		{
+			name: "node failures",
+			workload: workload(testJobSetUID,
+				workloadCond(workloadConditionEvicted, "True", "NodeFailures", "node lost")),
+			wantKind:  core.ExecutionError_SYSTEM,
+			wantInMsg: "gang evicted by kueue (NodeFailures)",
+		},
+		{
+			name: "a stopped queue",
+			workload: workload(testJobSetUID,
+				workloadCond(workloadConditionEvicted, "True", "ClusterQueueStopped", "The ClusterQueue is stopped")),
+			wantKind:  core.ExecutionError_SYSTEM,
+			wantInMsg: "gang evicted by kueue (ClusterQueueStopped)",
+		},
+		{
+			name: "requeuing limit exceeded",
+			workload: deactivated(testJobSetUID,
+				workloadCond(workloadConditionEvicted, "True", "DeactivatedDueToRequeuingLimitExceeded", "")),
+			wantKind: core.ExecutionError_SYSTEM,
+		},
+		{
+			name: "the user's maximum execution time",
+			workload: deactivated(testJobSetUID,
+				workloadCond(workloadConditionEvicted, "True", evictedByMaximumExecutionTime, "")),
+			wantKind: core.ExecutionError_USER,
+		},
+		{
+			name: "deactivated while waiting after a preemption",
+			workload: deactivated(testJobSetUID,
+				workloadCond(workloadConditionEvicted, "True", evictedByPreemption, preemptedMessage),
+				workloadCond(workloadConditionPreempted, "True", "InClusterQueue", preemptedMessage)),
+			wantKind:  core.ExecutionError_SYSTEM,
+			wantInMsg: "gang evicted by kueue (Deactivated); this does not count against the task's retries",
+		},
+		{
+			name: "deactivated by an operator",
+			workload: workload(testJobSetUID,
+				workloadCond(workloadConditionEvicted, "True", "Deactivated", "The workload is deactivated")),
+			wantKind: core.ExecutionError_SYSTEM,
+		},
+		{
+			name:     "deactivated before Kueue wrote the condition",
+			workload: deactivated(testJobSetUID),
+			wantKind: core.ExecutionError_SYSTEM,
+		},
+		{
+			name:     "no Workload for this JobSet",
+			workload: workload("another-jobset"),
+			wantKind: core.ExecutionError_SYSTEM,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			// Given
+			pCtx := resumingPluginCtx(readerWith(t, tt.workload), startedState, 0, 3)
+
+			// When
+			phase, err := clusteredResourceHandler{}.GetTaskPhase(context.Background(), pCtx, suspendedJobSet())
+
+			// Then
+			require.NoError(t, err)
+			assert.Equal(t, pluginsCore.PhaseRetryableFailure, phase.Phase())
+			require.NotNil(t, phase.Err())
+			assert.Equal(t, gang.CodeGangEvicted, phase.Err().GetCode())
+			assert.Equal(t, tt.wantKind, phase.Err().GetKind())
+			assert.True(t, phase.CleanupOnFailure())
+			assert.Contains(t, phase.Err().GetMessage(), tt.wantInMsg)
+		})
+	}
+}
+
+func TestGetTaskPhase_Resume_EvictionWhileRunningBumpsVersion(t *testing.T) {
+	// Given a running gang that Kueue preempts
+	wl := workload(testJobSetUID, workloadCond(workloadConditionEvicted, "True", evictedByPreemption, ""))
+	running := plugink8s.PluginState{Phase: pluginsCore.PhaseRunning, PhaseVersion: 4, Reason: "running (restart attempt 0)"}
+	pCtx := resumingPluginCtx(readerWith(t, wl), running, 0, 1)
+
+	// When
+	phase, err := clusteredResourceHandler{}.GetTaskPhase(context.Background(), pCtx, suspendedJobSet())
+
+	// Then the new reason is not dropped as a duplicate of the last Running event
+	require.NoError(t, err)
+	assert.Equal(t, pluginsCore.PhaseRunning, phase.Phase())
+	assert.Equal(t, uint32(5), phase.Version())
+}
+
+func TestGetTaskPhase_Resume_ReadmittedGangRunsAgain(t *testing.T) {
+	// Given a gang waiting for re-admission that Kueue has resumed
+	waiting := plugink8s.PluginState{
+		Phase: pluginsCore.PhaseRunning, PhaseVersion: 5,
+		Reason: "gang evicted by kueue (Preempted); waiting for re-admission",
+	}
+	pCtx := resumingPluginCtx(readerWith(t), waiting, 0, 1)
+	js := makeJobSet("", "", false)
+
+	// When
+	phase, err := clusteredResourceHandler{}.GetTaskPhase(context.Background(), pCtx, js)
+
+	// Then
+	require.NoError(t, err)
+	assert.Equal(t, pluginsCore.PhaseRunning, phase.Phase())
+	assert.Equal(t, "running (restart attempt 0)", phase.Reason())
+	assert.Equal(t, uint32(6), phase.Version())
+}
+
 func TestGetTaskPhase_Held_MissingLocalQueue_FailsAsUserError(t *testing.T) {
 	wl := workload(testJobSetUID,
 		inadmissible("LocalQueue "+testQueue+" doesn't exist"))