Fixing ArrayNode integration with backoff controller (#4640)

* handle WaitingForResources phase from backoff controller Signed-off-by: Daniel Rammer <[email protected]> * added unit test Signed-off-by: Daniel Rammer <[email protected]> --------- Signed-off-by: Daniel Rammer <[email protected]>
flyteorg · Jan 2, 2024 · ba10600 · ba10600
1 parent cb37291
commit ba10600
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 1 deletion.
diff --git a/flytepropeller/pkg/controller/nodes/array/handler.go b/flytepropeller/pkg/controller/nodes/array/handler.go
@@ -633,7 +633,7 @@ func (a *arrayNodeHandler) buildArrayNodeContext(ctx context.Context, nCtx inter
 	// currently just mocking based on node phase -> which works for all k8s plugins
 	// we can not pre-allocated a bit array because max size is 256B and with 5k fanout node state = 1.28MB
 	pluginStateBytes := a.pluginStateBytesStarted
-	if taskPhase == int(core.PhaseUndefined) || taskPhase == int(core.PhaseRetryableFailure) {
+	if taskPhase == int(core.PhaseUndefined) || taskPhase == int(core.PhaseRetryableFailure) || taskPhase == int(core.PhaseWaitingForResources) {
 		pluginStateBytes = a.pluginStateBytesNotStarted
 	}
 

diff --git a/flytepropeller/pkg/controller/nodes/array/handler_test.go b/flytepropeller/pkg/controller/nodes/array/handler_test.go
@@ -507,6 +507,24 @@ func TestHandleArrayNodePhaseExecuting(t *testing.T) {
 			expectedTransitionPhase:        handler.EPhaseRunning,
 			expectedExternalResourcePhases: []idlcore.TaskExecution_Phase{idlcore.TaskExecution_RUNNING},
 		},
+		{
+			name: "StartSubNodesNewAttempts",
+			subNodePhases: []v1alpha1.NodePhase{
+				v1alpha1.NodePhaseQueued,
+				v1alpha1.NodePhaseQueued,
+			},
+			subNodeTaskPhases: []core.Phase{
+				core.PhaseRetryableFailure,
+				core.PhaseWaitingForResources,
+			},
+			subNodeTransitions: []handler.Transition{
+				handler.DoTransition(handler.TransitionTypeEphemeral, handler.PhaseInfoRunning(&handler.ExecutionInfo{})),
+				handler.DoTransition(handler.TransitionTypeEphemeral, handler.PhaseInfoRunning(&handler.ExecutionInfo{})),
+			},
+			expectedArrayNodePhase:         v1alpha1.ArrayNodePhaseExecuting,
+			expectedTransitionPhase:        handler.EPhaseRunning,
+			expectedExternalResourcePhases: []idlcore.TaskExecution_Phase{idlcore.TaskExecution_RUNNING, idlcore.TaskExecution_RUNNING},
+		},
 		{
 			name: "AllSubNodesSuccedeed",
 			subNodePhases: []v1alpha1.NodePhase{