Add variable_map to GetData service response

Signed-off-by: mao3267 <chenvincent610@gmail.com>
flyteorg · Jan 4, 2025 · 07878e0 · 07878e0
1 parent 547c155
commit 07878e0
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/flyteadmin/dataproxy/service.go b/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
 
 }