Feat: Allow controlling in which task phases log links are shown

.github/workflows/single-binary.yml

@@ -174,7 +174,7 @@ jobs:
       - name: Install Python dependencies
         run: |
           python -m pip install --upgrade pip
-          pip install flytekit flytekitplugins-deck-standard
+          pip install flytekit flytekitplugins-deck-standard 'numpy<2'
           pip freeze
       - name: Checkout flytesnacks
         uses: actions/checkout@v4

flyteadmin/pkg/repositories/transformers/task_execution.go

@@ -372,6 +372,31 @@ func mergeMetadata(existing, latest *event.TaskExecutionMetadata) *event.TaskExe
 	return existing
 }
 
+func filterExternalResourceLogsByPhase(externalResources []*event.ExternalResourceInfo, phase core.TaskExecution_Phase) {
+	for _, externalResource := range externalResources {
+		externalResource.Logs = filterLogsByPhase(externalResource.Logs, phase)
+	}
+}
+
+func filterLogsByPhase(logs []*core.TaskLog, phase core.TaskExecution_Phase) []*core.TaskLog {
+	filteredLogs := make([]*core.TaskLog, 0, len(logs))
+
+	for _, l := range logs {
+		if common.IsTaskExecutionTerminal(phase) && l.HideOnceFinished {
+			continue
+		}
+		// Some plugins like e.g. Dask, Ray start with or very quickly transition to core.TaskExecution_INITIALIZING
+		// once the CR has been created even though the underlying pods are still pending. We thus treat queued and
+		// initializing the same here.
+		if (phase == core.TaskExecution_QUEUED || phase == core.TaskExecution_INITIALIZING) && !l.ShowWhilePending {
+			continue
+		}
+		filteredLogs = append(filteredLogs, l)
+
+	}
+	return filteredLogs
+}
+
 func UpdateTaskExecutionModel(ctx context.Context, request *admin.TaskExecutionEventRequest, taskExecutionModel *models.TaskExecution,
 	inlineEventDataPolicy interfaces.InlineEventDataPolicy, storageClient *storage.DataStore) error {
 	err := handleTaskExecutionInputs(ctx, taskExecutionModel, request, storageClient)
@@ -384,6 +409,7 @@ func UpdateTaskExecutionModel(ctx context.Context, request *admin.TaskExecutionE
 		return errors.NewFlyteAdminErrorf(codes.Internal,
 			"failed to unmarshal task execution closure with error: %+v", err)
 	}
+	isPhaseChange := taskExecutionModel.Phase != request.Event.Phase.String()
 	existingTaskPhase := taskExecutionModel.Phase
 	taskExecutionModel.Phase = request.Event.Phase.String()
 	taskExecutionModel.PhaseVersion = request.Event.PhaseVersion
@@ -393,7 +419,11 @@ func UpdateTaskExecutionModel(ctx context.Context, request *admin.TaskExecutionE
 		reportedAt = request.Event.OccurredAt
 	}
 	taskExecutionClosure.UpdatedAt = reportedAt
-	taskExecutionClosure.Logs = mergeLogs(taskExecutionClosure.Logs, request.Event.Logs)
+
+	mergedLogs := mergeLogs(taskExecutionClosure.Logs, request.Event.Logs)
+	filteredLogs := filterLogsByPhase(mergedLogs, request.Event.Phase)
+	taskExecutionClosure.Logs = filteredLogs
+
 	if len(request.Event.Reasons) > 0 {
 		for _, reason := range request.Event.Reasons {
 			taskExecutionClosure.Reasons = append(
@@ -437,6 +467,11 @@ func UpdateTaskExecutionModel(ctx context.Context, request *admin.TaskExecutionE
 		return errors.NewFlyteAdminErrorf(codes.Internal, "failed to merge task event custom_info with error: %v", err)
 	}
 	taskExecutionClosure.Metadata = mergeMetadata(taskExecutionClosure.Metadata, request.Event.Metadata)
+
+	if isPhaseChange && taskExecutionClosure.Metadata != nil && len(taskExecutionClosure.Metadata.ExternalResources) > 0 {
+		filterExternalResourceLogsByPhase(taskExecutionClosure.Metadata.ExternalResources, request.Event.Phase)
+	}
+
 	if request.Event.EventVersion > taskExecutionClosure.EventVersion {
 		taskExecutionClosure.EventVersion = request.Event.EventVersion
 	}

flyteadmin/pkg/repositories/transformers/task_execution_test.go

@@ -652,6 +652,183 @@ func TestUpdateTaskExecutionModelRunningToFailed(t *testing.T) {
 
 }
 
+func TestUpdateTaskExecutionModelFilterLogLinks(t *testing.T) {
+	existingClosure := &admin.TaskExecutionClosure{
+		Phase:     core.TaskExecution_QUEUED,
+		StartedAt: taskEventOccurredAtProto,
+		CreatedAt: taskEventOccurredAtProto,
+		UpdatedAt: taskEventOccurredAtProto,
+		Logs:      []*core.TaskLog{},
+		Reason:    "task submitted to k8s",
+		Reasons: []*admin.Reason{
+			{
+				OccurredAt: taskEventOccurredAtProto,
+				Message:    "task submitted to k8s",
+			},
+		},
+	}
+
+	closureBytes, err := proto.Marshal(existingClosure)
+	assert.Nil(t, err)
+
+	existingTaskExecution := models.TaskExecution{
+		TaskExecutionKey: models.TaskExecutionKey{
+			TaskKey: models.TaskKey{
+				Project: sampleTaskID.Project,
+				Domain:  sampleTaskID.Domain,
+				Name:    sampleTaskID.Name,
+				Version: sampleTaskID.Version,
+			},
+			NodeExecutionKey: models.NodeExecutionKey{
+				NodeID: sampleNodeExecID.NodeId,
+				ExecutionKey: models.ExecutionKey{
+					Project: sampleNodeExecID.ExecutionId.Project,
+					Domain:  sampleNodeExecID.ExecutionId.Domain,
+					Name:    sampleNodeExecID.ExecutionId.Name,
+				},
+			},
+			RetryAttempt: &retryAttemptValue,
+		},
+		Phase:                  "TaskExecutionPhase_TASK_PHASE_QUEUED",
+		InputURI:               "input uri",
+		Closure:                closureBytes,
+		StartedAt:              &taskEventOccurredAt,
+		TaskExecutionCreatedAt: &taskEventOccurredAt,
+		TaskExecutionUpdatedAt: &taskEventOccurredAt,
+	}
+
+	occuredAt := taskEventOccurredAt.Add(time.Minute)
+	occuredAtProto, err := ptypes.TimestampProto(occuredAt)
+	assert.Nil(t, err)
+
+	updatedEventRequest := &admin.TaskExecutionEventRequest{
+		Event: &event.TaskExecutionEvent{
+			TaskId:                sampleTaskID,
+			ParentNodeExecutionId: sampleNodeExecID,
+			Phase:                 core.TaskExecution_QUEUED,
+			OccurredAt:            occuredAtProto,
+			Logs: []*core.TaskLog{
+				{
+					Uri:              "uri-show-pending",
+					ShowWhilePending: true,
+				},
+				{
+					Uri: "uri-default",
+				},
+			},
+			Reason: "task update",
+		},
+	}
+
+	err = UpdateTaskExecutionModel(context.TODO(), updatedEventRequest, &existingTaskExecution,
+		interfaces.InlineEventDataPolicyStoreInline, commonMocks.GetMockStorageClient())
+	assert.Nil(t, err)
+
+	updatedClosure := &admin.TaskExecutionClosure{}
+	err = proto.Unmarshal(existingTaskExecution.Closure, updatedClosure)
+	assert.Nil(t, err)
+
+	assert.Equal(t, updatedClosure.Logs, []*core.TaskLog{
+		{
+			Uri:              "uri-show-pending",
+			ShowWhilePending: true,
+		},
+	},
+	)
+
+}
+
+func TestUpdateTaskExecutionModelFilterLogLinksArray(t *testing.T) {
+	existingClosure := &admin.TaskExecutionClosure{
+		Phase:     core.TaskExecution_RUNNING,
+		StartedAt: taskEventOccurredAtProto,
+		CreatedAt: taskEventOccurredAtProto,
+		UpdatedAt: taskEventOccurredAtProto,
+		Logs:      []*core.TaskLog{},
+		Reason:    "task started",
+		Reasons: []*admin.Reason{
+			{
+				OccurredAt: taskEventOccurredAtProto,
+				Message:    "task started",
+			},
+		},
+		Metadata: &event.TaskExecutionMetadata{
+			ExternalResources: []*event.ExternalResourceInfo{
+				{
+					Logs: []*core.TaskLog{
+						{
+							Uri: "uri-default",
+						},
+						{
+							Uri:              "uri-hide-finished",
+							HideOnceFinished: true,
+						},
+					},
+				},
+			},
+		},
+	}
+
+	closureBytes, err := proto.Marshal(existingClosure)
+	assert.Nil(t, err)
+
+	existingTaskExecution := models.TaskExecution{
+		TaskExecutionKey: models.TaskExecutionKey{
+			TaskKey: models.TaskKey{
+				Project: sampleTaskID.Project,
+				Domain:  sampleTaskID.Domain,
+				Name:    sampleTaskID.Name,
+				Version: sampleTaskID.Version,
+			},
+			NodeExecutionKey: models.NodeExecutionKey{
+				NodeID: sampleNodeExecID.NodeId,
+				ExecutionKey: models.ExecutionKey{
+					Project: sampleNodeExecID.ExecutionId.Project,
+					Domain:  sampleNodeExecID.ExecutionId.Domain,
+					Name:    sampleNodeExecID.ExecutionId.Name,
+				},
+			},
+			RetryAttempt: &retryAttemptValue,
+		},
+		Phase:                  "TaskExecutionPhase_TASK_PHASE_RUNNING",
+		InputURI:               "input uri",
+		Closure:                closureBytes,
+		StartedAt:              &taskEventOccurredAt,
+		TaskExecutionCreatedAt: &taskEventOccurredAt,
+		TaskExecutionUpdatedAt: &taskEventOccurredAt,
+	}
+
+	occuredAt := taskEventOccurredAt.Add(time.Minute)
+	occuredAtProto, err := ptypes.TimestampProto(occuredAt)
+	assert.Nil(t, err)
+
+	failedEventRequest := &admin.TaskExecutionEventRequest{
+		Event: &event.TaskExecutionEvent{
+			TaskId:                sampleTaskID,
+			ParentNodeExecutionId: sampleNodeExecID,
+			Phase:                 core.TaskExecution_FAILED,
+			OccurredAt:            occuredAtProto,
+			Reason:                "something went wrong",
+		},
+	}
+
+	err = UpdateTaskExecutionModel(context.TODO(), failedEventRequest, &existingTaskExecution,
+		interfaces.InlineEventDataPolicyStoreInline, commonMocks.GetMockStorageClient())
+	assert.Nil(t, err)
+
+	updatedClosure := &admin.TaskExecutionClosure{}
+	err = proto.Unmarshal(existingTaskExecution.Closure, updatedClosure)
+	assert.Nil(t, err)
+
+	assert.Equal(t, updatedClosure.Metadata.ExternalResources[0].Logs, []*core.TaskLog{
+		{
+			Uri: "uri-default",
+		},
+	},
+	)
+
+}
+
 func TestUpdateTaskExecutionModelSingleEvents(t *testing.T) {
 	existingClosure := &admin.TaskExecutionClosure{
 		Phase:     core.TaskExecution_RUNNING,
@@ -1208,6 +1385,125 @@ func TestMergeLogs(t *testing.T) {
 	}
 }
 
+func TestFilterLogsByPhase(t *testing.T) {
+	type testCase struct {
+		existing []*core.TaskLog
+		expected []*core.TaskLog
+		phase    core.TaskExecution_Phase
+		name     string
+	}
+
+	testCases := []testCase{
+		{
+			existing: []*core.TaskLog{
+				{
+					Uri:              "default-uri",
+					ShowWhilePending: false,
+					HideOnceFinished: false,
+				},
+				{
+					Uri:              "show-pending-uri",
+					ShowWhilePending: true,
+					HideOnceFinished: false,
+				},
+				{
+					Uri:              "hide-finished-uri",
+					ShowWhilePending: false,
+					HideOnceFinished: true,
+				},
+			},
+			expected: []*core.TaskLog{
+				{
+					Uri:              "show-pending-uri",
+					ShowWhilePending: true,
+					HideOnceFinished: false,
+				},
+			},
+			phase: core.TaskExecution_QUEUED,
+			name:  "Filtered logs in QUEUED phase",
+		},
+		{
+			existing: []*core.TaskLog{
+				{
+					Uri:              "default-uri",
+					ShowWhilePending: false,
+					HideOnceFinished: false,
+				},
+				{
+					Uri:              "show-pending-uri",
+					ShowWhilePending: true,
+					HideOnceFinished: false,
+				},
+				{
+					Uri:              "hide-finished-uri",
+					ShowWhilePending: false,
+					HideOnceFinished: true,
+				},
+			},
+			expected: []*core.TaskLog{
+				{
+					Uri:              "default-uri",
+					ShowWhilePending: false,
+					HideOnceFinished: false,
+				},
+				{
+					Uri:              "show-pending-uri",
+					ShowWhilePending: true,
+					HideOnceFinished: false,
+				},
+				{
+					Uri:              "hide-finished-uri",
+					ShowWhilePending: false,
+					HideOnceFinished: true,
+				},
+			},
+			phase: core.TaskExecution_RUNNING,
+			name:  "Filtered logs in RUNNING phase",
+		},
+		{
+			existing: []*core.TaskLog{
+				{
+					Uri:              "default-uri",
+					ShowWhilePending: false,
+					HideOnceFinished: false,
+				},
+				{
+					Uri:              "show-pending-uri",
+					ShowWhilePending: true,
+					HideOnceFinished: false,
+				},
+				{
+					Uri:              "hide-finished-uri",
+					ShowWhilePending: false,
+					HideOnceFinished: true,
+				},
+			},
+			expected: []*core.TaskLog{
+				{
+					Uri:              "default-uri",
+					ShowWhilePending: false,
+					HideOnceFinished: false,
+				},
+				{
+					Uri:              "show-pending-uri",
+					ShowWhilePending: true,
+					HideOnceFinished: false,
+				},
+			},
+			phase: core.TaskExecution_SUCCEEDED,
+			name:  "Filtered logs in terminated phase",
+		},
+	}
+	for _, filterTestCase := range testCases {
+		filteredLogs := filterLogsByPhase(filterTestCase.existing, filterTestCase.phase)
+
+		assert.Equal(t, len(filterTestCase.expected), len(filteredLogs), fmt.Sprintf("%s failed", filterTestCase.name))
+		for idx, expectedLog := range filterTestCase.expected {
+			assert.True(t, proto.Equal(expectedLog, filteredLogs[idx]), fmt.Sprintf("%s failed", filterTestCase.name))
+		}
+	}
+}
+
 func TestMergeCustoms(t *testing.T) {
 	t.Run("nothing to do", func(t *testing.T) {
 		custom, err := mergeCustom(nil, nil)