[WIP][Flyteadmin] Add variablemap in dataproxy for dataclass/pydantic

flyteadmin/dataproxy/service.go

@@ -367,10 +367,12 @@ func (s Service) GetDataFromNodeExecution(ctx context.Context, nodeExecID *core.
 	}
 
 	var lm *core.LiteralMap
+	var vm *core.VariableMap
 	if ioType == common.ArtifactTypeI {
 		lm = resp.GetFullInputs()
 	} else if ioType == common.ArtifactTypeO {
 		lm = resp.GetFullOutputs()
+		vm = resp.GetOutputVariableMap()
 	} else {
 		// Assume deck, and create a download link request
 		dlRequest := service.CreateDownloadLinkRequest{
@@ -402,13 +404,15 @@ func (s Service) GetDataFromNodeExecution(ctx context.Context, nodeExecID *core.
 		Data: &service.GetDataResponse_LiteralMap{
 			LiteralMap: lm,
 		},
+		VariableMap: vm,
 	}, nil
 }
 
 func (s Service) GetDataFromTaskExecution(ctx context.Context, taskExecID *core.TaskExecutionIdentifier, ioType common.ArtifactType, name string) (
 	*service.GetDataResponse, error) {
 
 	var lm *core.LiteralMap
+	var vm *core.VariableMap
 	reqT := &admin.TaskExecutionGetDataRequest{
 		Id: taskExecID,
 	}
@@ -421,6 +425,7 @@ func (s Service) GetDataFromTaskExecution(ctx context.Context, taskExecID *core.
 		lm = resp.GetFullInputs()
 	} else if ioType == common.ArtifactTypeO {
 		lm = resp.GetFullOutputs()
+		vm = resp.GetOutputVariableMap()
 	} else {
 		return nil, errors.NewFlyteAdminErrorf(codes.InvalidArgument, "deck type cannot be specified with a retry attempt, just use the node instead")
 	}
@@ -440,6 +445,7 @@ func (s Service) GetDataFromTaskExecution(ctx context.Context, taskExecID *core.
 		Data: &service.GetDataResponse_LiteralMap{
 			LiteralMap: lm,
 		},
+		VariableMap: vm,
 	}, nil
 
 }

flyteadmin/pkg/manager/impl/node_execution_manager.go

@@ -538,24 +538,39 @@ func (m *NodeExecutionManager) GetNodeExecutionData(
 
 	var outputs *core.LiteralMap
 	var outputURLBlob *admin.UrlBlob
+	var outputVariableMap *core.VariableMap
 	group.Go(func() error {
 		var err error
 		outputs, outputURLBlob, err = util.GetOutputs(groupCtx, m.urlData, m.config.ApplicationConfiguration().GetRemoteDataConfig(),
 			m.storageClient, nodeExecution.GetClosure())
 		return err
 	})
 
+	// TODO: Get the output variable map from the node execution model
+	// group.Go(func() error {
+	// 	var err error
+
+	// 	modelNode, err := m.db.NodeExecutionRepo().Get(groupCtx, repoInterfaces.NodeExecutionResource{
+	// 		NodeExecutionIdentifier: request.GetId(),
+	// 	})
+
+	// 	node, err := transformers.FromNodeExecutionModel(modelNode, transformers.DefaultExecutionTransformerOptions)
+
+	// 	return err
+	// })
+
 	err = group.Wait()
 	if err != nil {
 		return nil, err
 	}
 
 	response := &admin.NodeExecutionGetDataResponse{
-		Inputs:      inputURLBlob,
-		Outputs:     outputURLBlob,
-		FullInputs:  inputs,
-		FullOutputs: outputs,
-		FlyteUrls:   common.FlyteURLsFromNodeExecutionID(request.GetId(), nodeExecution.GetClosure() != nil && nodeExecution.GetClosure().GetDeckUri() != ""),
+		Inputs:            inputURLBlob,
+		Outputs:           outputURLBlob,
+		FullInputs:        inputs,
+		FullOutputs:       outputs,
+		FlyteUrls:         common.FlyteURLsFromNodeExecutionID(request.GetId(), nodeExecution.GetClosure() != nil && nodeExecution.GetClosure().GetDeckUri() != ""),
+		OutputVariableMap: outputVariableMap,
 	}
 
 	if len(nodeExecutionModel.DynamicWorkflowRemoteClosureReference) > 0 {

flyteadmin/pkg/manager/impl/task_execution_manager.go

@@ -328,24 +328,41 @@ func (m *TaskExecutionManager) GetTaskExecutionData(
 
 	var outputs *core.LiteralMap
 	var outputURLBlob *admin.UrlBlob
+	var outputVariableMap *core.VariableMap
 	group.Go(func() error {
 		var err error
 		outputs, outputURLBlob, err = util.GetOutputs(groupCtx, m.urlData, m.config.ApplicationConfiguration().GetRemoteDataConfig(),
 			m.storageClient, taskExecution.GetClosure())
 		return err
 	})
 
+	group.Go(func() error {
+		var err error
+		taskModel, err := m.db.TaskRepo().Get(groupCtx, repoInterfaces.Identifier{
+			Project: taskExecution.GetId().GetTaskId().GetProject(),
+			Domain:  taskExecution.GetId().GetTaskId().GetDomain(),
+			Name:    taskExecution.GetId().GetTaskId().GetName(),
+			Version: taskExecution.GetId().GetTaskId().GetVersion(),
+		})
+
+		task, err := transformers.FromTaskModel(taskModel)
+		outputVariableMap = task.GetClosure().GetCompiledTask().GetTemplate().GetInterface().Outputs
+
+		return err
+	})
+
 	err = group.Wait()
 	if err != nil {
 		return nil, err
 	}
 
 	response := &admin.TaskExecutionGetDataResponse{
-		Inputs:      inputURLBlob,
-		Outputs:     outputURLBlob,
-		FullInputs:  inputs,
-		FullOutputs: outputs,
-		FlyteUrls:   common.FlyteURLsFromTaskExecutionID(request.GetId(), false),
+		Inputs:            inputURLBlob,
+		Outputs:           outputURLBlob,
+		FullInputs:        inputs,
+		FullOutputs:       outputs,
+		FlyteUrls:         common.FlyteURLsFromTaskExecutionID(request.GetId(), false),
+		OutputVariableMap: outputVariableMap,
 	}
 
 	m.metrics.TaskExecutionInputBytes.Observe(float64(response.GetInputs().GetBytes()))