From 07878e0f29802a3d33a4633f6d3037d85b09aea2 Mon Sep 17 00:00:00 2001 From: mao3267 Date: Sat, 4 Jan 2025 10:08:19 +0800 Subject: [PATCH] Add variable_map to GetData service response Signed-off-by: mao3267 --- flyteadmin/dataproxy/service.go | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/flyteadmin/dataproxy/service.go b/flyteadmin/dataproxy/service.go index 27d03e3eda..1aa3956575 100644 --- 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,6 +404,7 @@ func (s Service) GetDataFromNodeExecution(ctx context.Context, nodeExecID *core. Data: &service.GetDataResponse_LiteralMap{ LiteralMap: lm, }, + VariableMap: vm, }, nil } @@ -409,6 +412,7 @@ func (s Service) GetDataFromTaskExecution(ctx context.Context, taskExecID *core. *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 }