Initial WIP on CustomInfo for agents

flyteadmin/pkg/repositories/transformers/task_execution.go

@@ -152,6 +152,7 @@ func CreateTaskExecutionModel(ctx context.Context, input CreateTaskExecutionMode
 		UpdatedAt:    reportedAt,
 		CreatedAt:    input.Request.Event.OccurredAt,
 		Logs:         input.Request.Event.Logs,
+		// TODO: does this help us pass CustomInfo at all?
 		CustomInfo:   input.Request.Event.CustomInfo,
 		TaskType:     input.Request.Event.TaskType,
 		Metadata:     metadata,
@@ -288,6 +289,8 @@ func mergeCustom(existing, latest *_struct.Struct) (*_struct.Struct, error) {
 	return &response, nil
 }
 
+// TODO: this merging seems to accumulate logs -- should it also  accumulate CustomInfo?
+
 // mergeExternalResource combines the latest ExternalResourceInfo proto with an existing instance
 // by updating fields and merging logs.
 func mergeExternalResource(existing, latest *event.ExternalResourceInfo) *event.ExternalResourceInfo {

flyteplugins/go/tasks/pluginmachinery/core/phase.go

@@ -89,6 +89,7 @@ type ExternalResource struct {
 	RetryAttempt uint32
 	// Phase (if exists) associated with the external resource
 	Phase Phase
+	// TODO: should there also be a CustomInfo here?
 }
 
 type ReasonInfo struct {

flyteplugins/go/tasks/plugins/array/core/state.go

@@ -178,8 +178,10 @@ func MapArrayStateToPluginPhase(_ context.Context, state *State, logLinks []*idl
 	t := time.Now()
 
 	nowTaskInfo := &core.TaskInfo{
-		OccurredAt:        &t,
-		Logs:              logLinks,
+		OccurredAt: &t,
+		Logs:       logLinks,
+		// TODO: is this missing CustomInfo ? Should it be copied from externalResources? Or passed as a param?
+		// CustomInfo:        &structpb.Struct{Fields: map[string]*structpb.Value{}},
 		ExternalResources: externalResources,
 	}
 

flyteplugins/go/tasks/plugins/array/k8s/management.go

@@ -273,6 +273,8 @@ func LaunchAndCheckSubTasksState(ctx context.Context, tCtx core.TaskExecutionCon
 			logLinks = phaseInfo.Info().Logs
 		}
 
+		// TODO: it looks like this is where .CustomInfo should be used
+
 		externalResources = append(externalResources, &core.ExternalResource{
 			ExternalID:   podName,
 			Index:        uint32(originalIdx),

flyteplugins/go/tasks/plugins/array/k8s/management_test.go

@@ -437,6 +437,7 @@ func TestCheckSubTasksState(t *testing.T) {
 		resourceManager.AssertNumberOfCalls(t, "AllocateResource", 0)
 		resourceManager.AssertNumberOfCalls(t, "ReleaseResource", 0)
 
+		// TODO: write tests for CustomInfo here?
 		assert.Equal(t, subtaskCount, len(externalResources))
 		for i := 0; i < subtaskCount; i++ {
 			externalResource := externalResources[i]

flyteplugins/go/tasks/plugins/webapi/agent/plugin.go

@@ -20,6 +20,7 @@ import (
 	"github.com/flyteorg/flyte/flyteplugins/go/tasks/pluginmachinery/webapi"
 	"github.com/flyteorg/flyte/flytestdlib/logger"
 	"github.com/flyteorg/flyte/flytestdlib/promutils"
+	structpb "github.com/golang/protobuf/ptypes/struct"
 )
 
 type Registry map[string]map[int32]*Agent // map[taskTypeName][taskTypeVersion] => Agent
@@ -38,6 +39,8 @@ type ResourceWrapper struct {
 	Outputs  *flyteIdl.LiteralMap
 	Message  string
 	LogLinks []*flyteIdl.TaskLog
+	// TODO: should there also be a CustomInfo here?
+	CustomInfo *structpb.Struct
 }
 
 // IsTerminal is used to avoid making network calls to the agent service if the resource is already in a terminal state.
@@ -184,10 +187,11 @@ func (p Plugin) ExecuteTaskSync(
 	}
 
 	return nil, ResourceWrapper{
-		Phase:    resource.Phase,
-		Outputs:  resource.Outputs,
-		Message:  resource.Message,
-		LogLinks: resource.LogLinks,
+		Phase:      resource.Phase,
+		Outputs:    resource.Outputs,
+		Message:    resource.Message,
+		LogLinks:   resource.LogLinks,
+		CustomInfo: resource.CustomInfo,
 	}, err
 }
 
@@ -213,11 +217,12 @@ func (p Plugin) Get(ctx context.Context, taskCtx webapi.GetContext) (latest weba
 	}
 
 	return ResourceWrapper{
-		Phase:    res.Resource.Phase,
-		State:    res.Resource.State,
-		Outputs:  res.Resource.Outputs,
-		Message:  res.Resource.Message,
-		LogLinks: res.Resource.LogLinks,
+		Phase:      res.Resource.Phase,
+		State:      res.Resource.State,
+		Outputs:    res.Resource.Outputs,
+		Message:    res.Resource.Message,
+		LogLinks:   res.Resource.LogLinks,
+		CustomInfo: res.Resource.CustomInfo,
 	}, nil
 }
 
@@ -246,7 +251,9 @@ func (p Plugin) Delete(ctx context.Context, taskCtx webapi.DeleteContext) error
 
 func (p Plugin) Status(ctx context.Context, taskCtx webapi.StatusContext) (phase core.PhaseInfo, err error) {
 	resource := taskCtx.Resource().(ResourceWrapper)
-	taskInfo := &core.TaskInfo{Logs: resource.LogLinks}
+	// taskInfo := &core.TaskInfo{Logs: resource.LogLinks}
+	// TODO: should this actually be something like this?
+	taskInfo := &core.TaskInfo{Logs: resource.LogLinks, CustomInfo: resource.CustomInfo}
 
 	switch resource.Phase {
 	case flyteIdl.TaskExecution_QUEUED:

flytepropeller/pkg/controller/nodes/task/transformer.go

@@ -178,6 +178,7 @@
 
 	if input.Info.Info() != nil {
 		tev.Logs = input.Info.Info().Logs
+		//  TODO: this is one place whre CustomInfo is actually used
 		tev.CustomInfo = input.Info.Info().CustomInfo
 	}