Treat kubelet NodeAffinity status.reason as retryable system error (#6461)

Signed-off-by: Mike Hotan <mike@union.ai>

Author: mhotan

flyteplugins/go/tasks/pluginmachinery/flytek8s/pod_helper.go

@@ -38,9 +38,15 @@ const primaryContainerTemplateName = "primary"
 const primaryInitContainerTemplateName = "primary-init"
 const PrimaryContainerKey = "primary_container_name"
 
-// nodePreemptionStatusReasons are the status reasons that a pod's respective node
-// is preempted by the scheduler
-var nodePreemptionStatusReasons = sets.NewString("Shutdown", "Terminated", "NodeShutdown")
+var retryableStatusReasons = sets.NewString(
+	// Reasons that indicate the node was preempted aggressively.
+	// Kubelet can miss deleting the pod prior to the node being shutdown.
+	"Shutdown",
+	"Terminated",
+	"NodeShutdown",
+	// kubelet admission rejects the pod before the node gets assigned appropriate labels.
+	"NodeAffinity",
+)
 
 // AddRequiredNodeSelectorRequirements adds the provided v1.NodeSelectorRequirement
 // objects to an existing v1.Affinity object. If there are no existing required
@@ -1207,7 +1213,7 @@ func DemystifyFailure(ctx context.Context, status v1.PodStatus, info pluginsCore
 
 	var isSystemError bool
 	// In some versions of GKE the reason can also be "Terminated" or "NodeShutdown"
-	if nodePreemptionStatusReasons.Has(code) {
+	if retryableStatusReasons.Has(code) {
 		isSystemError = true
 	}
 

flyteplugins/go/tasks/pluginmachinery/flytek8s/pod_helper_test.go

@@ -1721,6 +1721,28 @@ func TestDemystifyFailure(t *testing.T) {
 			})
 		}
 	})
+
+	t.Run("Kubelet admission denies pod due to missing node label", func(t *testing.T) {
+		for _, reason := range []string{
+			"NodeAffinity",
+		} {
+			t.Run(reason, func(t *testing.T) {
+				message := "Pod was rejected: Predicate NodeAffinity failed: node(s) didn't match Pod's node affinity/selector"
+				phaseInfo, err := DemystifyFailure(ctx, v1.PodStatus{
+					Message: message,
+					Reason:  reason,
+					Phase:   v1.PodFailed,
+					// Can't always rely on GCP returining container statuses when node is preempted
+					ContainerStatuses: []v1.ContainerStatus{},
+				}, pluginsCore.TaskInfo{}, "")
+				assert.Nil(t, err)
+				assert.Equal(t, pluginsCore.PhaseRetryableFailure, phaseInfo.Phase())
+				assert.Equal(t, "Interrupted", phaseInfo.Err().GetCode())
+				assert.Equal(t, core.ExecutionError_SYSTEM, phaseInfo.Err().GetKind())
+				assert.Equal(t, message, phaseInfo.Err().GetMessage())
+			})
+		}
+	})
 }
 
 func TestDemystifyPending_testcases(t *testing.T) {