implement

Signed-off-by: Alex Wu <c.alexwu@gmail.com>

Author: popojk

actions/k8s/client.go

@@ -63,7 +63,8 @@ type ActionsClient struct {
 	bufferSize int
 	runClient  workflowconnect.InternalRunServiceClient
 	// recordedFilter deduplicates RecordAction calls across watch reconnects.
-	recordedFilter fastcheck.Filter
+	recordedFilter  fastcheck.Filter
+	recoveryMetrics *recoveryMetrics
 
 	// Watch management
 	mu sync.RWMutex
@@ -119,6 +120,7 @@ func NewActionsClient(k8sClient client.WithWatch, sharedCache ctrlcache.Cache, n
 		return nil, fmt.Errorf("actions: failed to create RecordAction dedup filter (size=%d): %w", recordFilterSize, err)
 	}
 	c.recordedFilter = filter
+	c.recoveryMetrics = newRecoveryMetrics(scope)
 
 	return c, nil
 }
@@ -138,17 +140,8 @@ func (c *ActionsClient) Enqueue(ctx context.Context, action *actions.Action, run
 			return fmt.Errorf("failed to ensure namespace %s: %w", c.namespace, err)
 		}
 		taskAction := c.newTaskActionCR(actionID, executorv1.ActionTypeTask, isRoot)
-		// Set OwnerReference to parent so K8s cascades deletion to children.
-		if !isRoot {
-			parentTaskAction, err := c.setParentOwnership(ctx, taskAction, actionID.Run, *action.ParentActionName)
-			if err != nil {
-				return err
-			}
-			// For child actions, inherit parent's run context
-			inheritRunContextFromParentTaskAction(taskAction, parentTaskAction)
-		} else {
-			// For root action, apply the RunSpec to TaskAction
-			applyRunSpecToTaskAction(taskAction, runSpec)
+		if err := c.applyRunContext(ctx, taskAction, action, runSpec, isRoot); err != nil {
+			return err
 		}
 
 		// Build and set the ActionSpec for the executor.
@@ -157,6 +150,7 @@ func (c *ActionsClient) Enqueue(ctx context.Context, action *actions.Action, run
 			return fmt.Errorf("failed to set action spec: %w", err)
 		}
 		taskAction.Spec.CacheKey = extractTaskCacheKey(action)
+		taskAction.Spec.RecoveredFrom = c.resolveRecoveredFrom(ctx, taskAction, action, isRoot)
 
 		// Embed the inline TaskTemplate if present.
 		if err := embedTaskTemplate(action, taskAction, runSpec); err != nil {
@@ -179,17 +173,16 @@ func (c *ActionsClient) Enqueue(ctx context.Context, action *actions.Action, run
 			return fmt.Errorf("failed to ensure namespace %s: %w", c.namespace, err)
 		}
 		taskAction := c.newTaskActionCR(actionID, executorv1.ActionTypeCondition, isRoot)
-		if !isRoot {
-			if _, err := c.setParentOwnership(ctx, taskAction, actionID.Run, *action.ParentActionName); err != nil {
-				return err
-			}
+		if err := c.applyRunContext(ctx, taskAction, action, runSpec, isRoot); err != nil {
+			return err
 		}
 
 		actionSpec := buildActionSpec(action, runSpec)
 		if err := taskAction.Spec.SetActionSpec(actionSpec); err != nil {
 			return fmt.Errorf("failed to set action spec: %w", err)
 		}
 		taskAction.Spec.ActionType = executorv1.ActionTypeCondition
+		taskAction.Spec.RecoveredFrom = c.resolveRecoveredFrom(ctx, taskAction, action, isRoot)
 		condBytes, err := proto.Marshal(cond)
 		if err != nil {
 			return fmt.Errorf("failed to marshal condition spec: %w", err)
@@ -866,6 +859,9 @@ func GetPhaseFromConditions(taskAction *executorv1.TaskAction) common.ActionPhas
 		switch cond.Type {
 		case string(executorv1.ConditionTypeSucceeded):
 			if cond.Status == "True" {
+				if cond.Reason == string(executorv1.ConditionReasonRecovered) {
+					return common.ActionPhase_ACTION_PHASE_RECOVERED
+				}
 				return common.ActionPhase_ACTION_PHASE_SUCCEEDED
 			}
 		case string(executorv1.ConditionTypeFailed):
@@ -995,6 +991,11 @@ func buildTaskActionName(actionID *common.ActionIdentifier) string {
 // It uses the same path structure as the executor's ComputeActionOutputPath so that
 // the SDK can find outputs written by the executor.
 func BuildOutputUri(ctx context.Context, ta *executorv1.TaskAction) string {
+	// A recovered action wrote nothing under this run's base; its result is the source run's.
+	// RecoveredFrom carries the outputs file, this returns the directory the SDK joins onto.
+	if ta.Spec.RecoveredFrom != nil {
+		return plugin.OutputPrefixOf(ta.Spec.RecoveredFrom.OutputUri)
+	}
 	if ta.Spec.RunOutputBase == "" {
 		return ""
 	}
@@ -1037,13 +1038,57 @@ func buildActionSpec(action *actions.Action, runSpec *task.RunSpec) *workflow.Ac
 	return actionSpec
 }
 
-func applyRunSpecToTaskAction(taskAction *executorv1.TaskAction, runSpec *task.RunSpec) {
+// applyRunContext sets the OwnerReference on the parent and applies the run context reaching
+// this action: from the RunSpec for the root, from the parent CR for everything else.
+func (c *ActionsClient) applyRunContext(
+	ctx context.Context,
+	taskAction *executorv1.TaskAction,
+	action *actions.Action,
+	runSpec *task.RunSpec,
+	isRoot bool,
+) error {
+	if isRoot {
+		return applyRunSpecToTaskAction(taskAction, runSpec)
+	}
+	parentTaskAction, err := c.setParentOwnership(ctx, taskAction, action.ActionId.Run, *action.ParentActionName)
+	if err != nil {
+		return err
+	}
+	// child actions inherit run context from parent
+	inheritRunContextFromParentTaskAction(taskAction, parentTaskAction)
+	return nil
+}
+
+// recoveryContextFromRunSpec returns nil for any run that is not a recovery.
+func recoveryContextFromRunSpec(runSpec *task.RunSpec) (*executorv1.RecoveryContext, error) {
+	relation := runSpec.GetRelation()
+	if relation.GetRelationType() != common.RelationType_RELATION_TYPE_RECOVER {
+		return nil, nil
+	}
+	relationBytes, err := proto.Marshal(relation)
+	if err != nil {
+		return nil, fmt.Errorf("failed to marshal recovery relation: %w", err)
+	}
+	return &executorv1.RecoveryContext{
+		Relation:          relationBytes,
+		ForceRerunActions: runSpec.GetRecover().GetForceRerunActions(),
+	}, nil
+}
+
+func applyRunSpecToTaskAction(taskAction *executorv1.TaskAction, runSpec *task.RunSpec) error {
 	if runSpec == nil {
 		taskAction.Spec.EnvVars = nil
 		taskAction.Spec.Interruptible = nil
-		return
+		taskAction.Spec.RecoveryContext = nil
+		return nil
 	}
 
+	recoveryContext, err := recoveryContextFromRunSpec(runSpec)
+	if err != nil {
+		return err
+	}
+	taskAction.Spec.RecoveryContext = recoveryContext
+
 	taskAction.Spec.EnvVars = keyValuePairsToMap(runSpec.GetEnvs().GetValues())
 	if runSpec.GetInterruptible() != nil {
 		value := runSpec.GetInterruptible().GetValue()
@@ -1066,12 +1111,15 @@ func applyRunSpecToTaskAction(taskAction *executorv1.TaskAction, runSpec *task.R
 			taskAction.Annotations[key] = value
 		}
 	}
+
+	return nil
 }
 
 func inheritRunContextFromParentTaskAction(taskAction *executorv1.TaskAction, parentTaskAction *executorv1.TaskAction) {
 	if taskAction == nil || parentTaskAction == nil {
 		return
 	}
+	taskAction.Spec.RecoveryContext = parentTaskAction.Spec.RecoveryContext.DeepCopy()
 	taskAction.Spec.EnvVars = cloneStringMap(parentTaskAction.Spec.EnvVars)
 	if len(parentTaskAction.Annotations) > 0 {
 		if taskAction.Annotations == nil {

actions/k8s/client_recovery_test.go

@@ -0,0 +1,149 @@
+package k8s
+
+import (
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"google.golang.org/protobuf/proto"
+	corev1 "k8s.io/api/core/v1"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/apimachinery/pkg/runtime"
+	"sigs.k8s.io/controller-runtime/pkg/client/fake"
+
+	executorv1 "github.com/flyteorg/flyte/v2/executor/api/v1"
+	"github.com/flyteorg/flyte/v2/gen/go/flyteidl2/common"
+	"github.com/flyteorg/flyte/v2/gen/go/flyteidl2/core"
+	"github.com/flyteorg/flyte/v2/gen/go/flyteidl2/task"
+)
+
+func recoverRunSpec(sourceRun string, forceRerun ...string) *task.RunSpec {
+	spec := &task.RunSpec{
+		Relation: &common.Relation{
+			RelatedTo: &common.RunIdentifier{
+				Org: "org1", Project: "proj", Domain: "dev", Name: sourceRun,
+			},
+			RelationType: common.RelationType_RELATION_TYPE_RECOVER,
+		},
+	}
+	if len(forceRerun) > 0 {
+		spec.Recover = &task.Recover{ForceRerunActions: forceRerun}
+	}
+	return spec
+}
+
+func TestApplyRunSpecToTaskAction_StampsRecoveryContext(t *testing.T) {
+	taskAction := &executorv1.TaskAction{Spec: executorv1.TaskActionSpec{}}
+
+	require.NoError(t, applyRunSpecToTaskAction(taskAction, recoverRunSpec("r1", "a3", "a7")))
+
+	recoveryContext := taskAction.Spec.RecoveryContext
+	require.NotNil(t, recoveryContext)
+	assert.Equal(t, []string{"a3", "a7"}, recoveryContext.ForceRerunActions)
+
+	relation := &common.Relation{}
+	require.NoError(t, proto.Unmarshal(recoveryContext.Relation, relation))
+	assert.Equal(t, common.RelationType_RELATION_TYPE_RECOVER, relation.GetRelationType())
+	assert.Equal(t, "r1", relation.GetRelatedTo().GetName())
+}
+
+// rerun and spawn share RunSpec.relation with recover; only the type makes it a recovery.
+func TestApplyRunSpecToTaskAction_NonRecoveryRelationStampsNothing(t *testing.T) {
+	for _, relationType := range []common.RelationType{
+		common.RelationType_RELATION_TYPE_RERUN,
+		common.RelationType_RELATION_TYPE_SPAWN,
+		common.RelationType_RELATION_TYPE_UNSPECIFIED,
+	} {
+		t.Run(relationType.String(), func(t *testing.T) {
+			spec := recoverRunSpec("r1")
+			spec.Relation.RelationType = relationType
+
+			taskAction := &executorv1.TaskAction{Spec: executorv1.TaskActionSpec{}}
+			require.NoError(t, applyRunSpecToTaskAction(taskAction, spec))
+			assert.Nil(t, taskAction.Spec.RecoveryContext)
+		})
+	}
+}
+
+func TestApplyRunSpecToTaskAction_NilRunSpecClearsRecoveryContext(t *testing.T) {
+	taskAction := &executorv1.TaskAction{
+		Spec: executorv1.TaskActionSpec{
+			RecoveryContext: &executorv1.RecoveryContext{Relation: []byte("stale")},
+		},
+	}
+
+	require.NoError(t, applyRunSpecToTaskAction(taskAction, nil))
+	assert.Nil(t, taskAction.Spec.RecoveryContext)
+}
+
+func TestInheritRunContextFromParentTaskAction_CopiesRecoveryContext(t *testing.T) {
+	parent := &executorv1.TaskAction{
+		Spec: executorv1.TaskActionSpec{
+			RecoveryContext: recoveryContextFor("source-run", "a3"),
+		},
+	}
+	child := &executorv1.TaskAction{Spec: executorv1.TaskActionSpec{}}
+
+	inheritRunContextFromParentTaskAction(child, parent)
+
+	require.NotNil(t, child.Spec.RecoveryContext)
+	assert.Equal(t, []string{"a3"}, child.Spec.RecoveryContext.ForceRerunActions)
+
+	child.Spec.RecoveryContext.ForceRerunActions[0] = "mutated"
+	child.Spec.RecoveryContext.Relation[0] = 'X'
+	assert.Equal(t, []string{"a3"}, parent.Spec.RecoveryContext.ForceRerunActions)
+	assert.Equal(t, recoveryContextFor("source-run").Relation, parent.Spec.RecoveryContext.Relation)
+}
+
+func TestInheritRunContextFromParentTaskAction_NoRecoveryContextOnParent(t *testing.T) {
+	child := &executorv1.TaskAction{Spec: executorv1.TaskActionSpec{}}
+	inheritRunContextFromParentTaskAction(child, &executorv1.TaskAction{})
+	assert.Nil(t, child.Spec.RecoveryContext)
+}
+
+// A condition action used to inherit nothing from its parent, so a subtree beneath one lost
+// the run context entirely.
+func TestEnqueueCondition_InheritsRunContextFromParent(t *testing.T) {
+	scheme := runtime.NewScheme()
+	require.NoError(t, corev1.AddToScheme(scheme))
+	require.NoError(t, executorv1.AddToScheme(scheme))
+
+	interruptible := true
+	parent := &executorv1.TaskAction{
+		ObjectMeta: metav1.ObjectMeta{
+			Name:        "run1-a0",
+			Namespace:   "flyte",
+			Annotations: map[string]string{"owner": "sdk"},
+			Labels:      map[string]string{"team": "platform"},
+		},
+		Spec: executorv1.TaskActionSpec{
+			EnvVars:         map[string]string{"TRACE_ID": "abc123"},
+			Interruptible:   &interruptible,
+			RecoveryContext: recoveryContextFor("source-run", "a3"),
+		},
+	}
+	c := &ActionsClient{
+		recordedFilter: testFilter(),
+		namespace:      "flyte",
+		k8sClient: fake.NewClientBuilder().
+			WithScheme(scheme).
+			WithObjects(parent).
+			WithStatusSubresource(&executorv1.TaskAction{}).
+			Build(),
+	}
+
+	action := newConditionAction(core.SimpleType_BOOLEAN)
+	require.NoError(t, c.Enqueue(context.Background(), action, nil))
+
+	created, err := c.GetTaskAction(context.Background(), action.ActionId)
+	require.NoError(t, err)
+	require.NotNil(t, created.Spec.RecoveryContext)
+	assert.Equal(t, recoveryContextFor("source-run", "a3").Relation, created.Spec.RecoveryContext.Relation)
+	assert.Equal(t, []string{"a3"}, created.Spec.RecoveryContext.ForceRerunActions)
+	assert.Equal(t, "abc123", created.Spec.EnvVars["TRACE_ID"])
+	require.NotNil(t, created.Spec.Interruptible)
+	assert.True(t, *created.Spec.Interruptible)
+	assert.Equal(t, "sdk", created.Annotations["owner"])
+	assert.Equal(t, "platform", created.Labels["team"])
+}