Skip to content

Commit 0507fe4

Browse files
mao3267mao3267
committed
fix: get task variable_map (TODO: fix workflow)
Signed-off-by: mao3267 <chenvincent610@gmail.com>
1 parent 07878e0 commit 0507fe4

File tree

24 files changed

+321
-123
lines changed

24 files changed

+321
-123
lines changed

‎flyteadmin/dataproxy/service.go‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -370,6 +370,7 @@ func (s Service) GetDataFromNodeExecution(ctx context.Context, nodeExecID *core.
370370
var vm *core.VariableMap
371371
if ioType == common.ArtifactTypeI {
372372
lm = resp.GetFullInputs()
373+
vm = resp.GetInputVariableMap()
373374
} else if ioType == common.ArtifactTypeO {
374375
lm = resp.GetFullOutputs()
375376
vm = resp.GetOutputVariableMap()
@@ -423,6 +424,7 @@ func (s Service) GetDataFromTaskExecution(ctx context.Context, taskExecID *core.
423424

424425
if ioType == common.ArtifactTypeI {
425426
lm = resp.GetFullInputs()
427+
vm = resp.GetInputVariableMap()
426428
} else if ioType == common.ArtifactTypeO {
427429
lm = resp.GetFullOutputs()
428430
vm = resp.GetOutputVariableMap()

‎flyteadmin/go.mod‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -97,7 +97,7 @@ require (
9797
github.com/eapache/go-resiliency v1.2.0 // indirect
9898
github.com/eapache/go-xerial-snappy v0.0.0-20180814174437-776d5712da21 // indirect
9999
github.com/eapache/queue v1.1.0 // indirect
100-
github.com/emicklei/go-restful/v3 v3.11.0 // indirect
100+
github.com/emicklei/go-restful/v3 v3.12.0 // indirect
101101
github.com/evanphx/json-patch/v5 v5.6.0 // indirect
102102
github.com/fatih/color v1.13.0 // indirect
103103
github.com/felixge/httpsnoop v1.0.4 // indirect
@@ -147,7 +147,7 @@ require (
147147
github.com/lestrrat-go/option v1.0.1 // indirect
148148
github.com/mailru/easyjson v0.7.7 // indirect
149149
github.com/mattn/go-colorable v0.1.12 // indirect
150-
github.com/mattn/go-isatty v0.0.16 // indirect
150+
github.com/mattn/go-isatty v0.0.18 // indirect
151151
github.com/mattn/go-sqlite3 v2.0.3+incompatible // indirect
152152
github.com/mattn/goveralls v0.0.6 // indirect
153153
github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd // indirect
@@ -159,7 +159,7 @@ require (
159159
github.com/ory/go-convenience v0.1.0 // indirect
160160
github.com/ory/viper v1.7.5 // indirect
161161
github.com/pborman/uuid v1.2.0 // indirect
162-
github.com/pelletier/go-toml v1.9.4 // indirect
162+
github.com/pelletier/go-toml v1.9.5 // indirect
163163
github.com/pelletier/go-toml/v2 v2.0.0-beta.8 // indirect
164164
github.com/pierrec/lz4 v2.5.2+incompatible // indirect
165165
github.com/pkg/browser v0.0.0-20240102092130-5ac0b6a4141c // indirect
@@ -171,7 +171,7 @@ require (
171171
github.com/sendgrid/rest v2.6.9+incompatible // indirect
172172
github.com/shamaton/msgpack/v2 v2.2.2 // indirect
173173
github.com/sirupsen/logrus v1.9.3 // indirect
174-
github.com/spf13/afero v1.8.2 // indirect
174+
github.com/spf13/afero v1.9.2 // indirect
175175
github.com/spf13/cast v1.4.1 // indirect
176176
github.com/spf13/jwalterweatherman v1.1.0 // indirect
177177
github.com/spf13/viper v1.11.0 // indirect

‎flyteadmin/go.sum‎

Lines changed: 9 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -215,8 +215,8 @@ github.com/eapache/queue v1.1.0/go.mod h1:6eCeP0CKFpHLu8blIFXhExK/dRa7WDZfr6jVFP
215215
github.com/elastic/go-sysinfo v1.1.1/go.mod h1:i1ZYdU10oLNfRzq4vq62BEwD2fH8KaWh6eh0ikPT9F0=
216216
github.com/elastic/go-windows v1.0.0/go.mod h1:TsU0Nrp7/y3+VwE82FoZF8gC/XFg/Elz6CcloAxnPgU=
217217
github.com/elazarl/goproxy v0.0.0-20181003060214-f58a169a71a5/go.mod h1:/Zj4wYkgs4iZTTu3o/KG3Itv/qCCa8VVMlb3i9OVuzc=
218-
github.com/emicklei/go-restful/v3 v3.11.0 h1:rAQeMHw1c7zTmncogyy8VvRZwtkmkZ4FxERmMY4rD+g=
219-
github.com/emicklei/go-restful/v3 v3.11.0/go.mod h1:6n3XBCmQQb25CM2LCACGz8ukIrRry+4bhvbpWn3mrbc=
218+
github.com/emicklei/go-restful/v3 v3.12.0 h1:y2DdzBAURM29NFF94q6RaY4vjIH1rtwDapwQtU84iWk=
219+
github.com/emicklei/go-restful/v3 v3.12.0/go.mod h1:6n3XBCmQQb25CM2LCACGz8ukIrRry+4bhvbpWn3mrbc=
220220
github.com/envoyproxy/go-control-plane v0.6.9/go.mod h1:SBwIajubJHhxtWwsL9s8ss4safvEdbitLhGGK48rN6g=
221221
github.com/envoyproxy/go-control-plane v0.9.0/go.mod h1:YTl/9mNaCwkRvm6d1a2C3ymFceY/DCBVvsKhRF0iEA4=
222222
github.com/envoyproxy/go-control-plane v0.9.1-0.20191026205805-5f8ba28d4473/go.mod h1:YTl/9mNaCwkRvm6d1a2C3ymFceY/DCBVvsKhRF0iEA4=
@@ -970,8 +970,8 @@ github.com/mattn/go-isatty v0.0.9/go.mod h1:YNRxwqDuOph6SZLI9vUUz6OYw3QyUt7WiY2y
970970
github.com/mattn/go-isatty v0.0.11/go.mod h1:PhnuNfih5lzO57/f3n+odYbM4JtupLOxQOAqxQCu2WE=
971971
github.com/mattn/go-isatty v0.0.12/go.mod h1:cbi8OIDigv2wuxKPP5vlRcQ1OAZbq2CE4Kysco4FUpU=
972972
github.com/mattn/go-isatty v0.0.14/go.mod h1:7GGIvUiUoEMVVmxf/4nioHXj79iQHKdU27kJ6hsGG94=
973-
github.com/mattn/go-isatty v0.0.16 h1:bq3VjFmv/sOjHtdEhmkEV4x1AJtvUvOJ2PFAZ5+peKQ=
974-
github.com/mattn/go-isatty v0.0.16/go.mod h1:kYGgaQfpe5nmfYZH+SKPsOc2e4SrIfOl2e/yFXSvRLM=
973+
github.com/mattn/go-isatty v0.0.18 h1:DOKFKCQ7FNG2L1rbrmstDN4QVRdS89Nkh85u68Uwp98=
974+
github.com/mattn/go-isatty v0.0.18/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y=
975975
github.com/mattn/go-sqlite3 v1.9.0/go.mod h1:FPy6KqzDD04eiIsT53CuJW3U88zkxoIYsOqkbpncsNc=
976976
github.com/mattn/go-sqlite3 v1.10.0/go.mod h1:FPy6KqzDD04eiIsT53CuJW3U88zkxoIYsOqkbpncsNc=
977977
github.com/mattn/go-sqlite3 v1.11.0/go.mod h1:FPy6KqzDD04eiIsT53CuJW3U88zkxoIYsOqkbpncsNc=
@@ -1083,8 +1083,8 @@ github.com/pelletier/go-toml v1.4.0/go.mod h1:PN7xzY2wHTK0K9p34ErDQMlFxa51Fk0OUr
10831083
github.com/pelletier/go-toml v1.6.0/go.mod h1:5N711Q9dKgbdkxHL+MEfF31hpT7l0S0s/t2kKREewys=
10841084
github.com/pelletier/go-toml v1.7.0/go.mod h1:vwGMzjaWMwyfHwgIBhI2YUM4fB6nL6lVAvS1LBMMhTE=
10851085
github.com/pelletier/go-toml v1.8.0/go.mod h1:D6yutnOGMveHEPV7VQOuvI/gXY61bv+9bAOTRnLElKs=
1086-
github.com/pelletier/go-toml v1.9.4 h1:tjENF6MfZAg8e4ZmZTeWaWiT2vXtsoO6+iuOjFhECwM=
1087-
github.com/pelletier/go-toml v1.9.4/go.mod h1:u1nR/EPcESfeI/szUZKdtJ0xRNbUoANCkoOuaOx1Y+c=
1086+
github.com/pelletier/go-toml v1.9.5 h1:4yBQzkHv+7BHq2PQUZF3Mx0IYxG7LsP222s7Agd3ve8=
1087+
github.com/pelletier/go-toml v1.9.5/go.mod h1:u1nR/EPcESfeI/szUZKdtJ0xRNbUoANCkoOuaOx1Y+c=
10881088
github.com/pelletier/go-toml/v2 v2.0.0-beta.8 h1:dy81yyLYJDwMTifq24Oi/IslOslRrDSb3jwDggjz3Z0=
10891089
github.com/pelletier/go-toml/v2 v2.0.0-beta.8/go.mod h1:r9LEWfGN8R5k0VXJ+0BkIe7MYkRdwZOjgMj2KwnJFUo=
10901090
github.com/phayes/freeport v0.0.0-20180830031419-95f893ade6f2/go.mod h1:iIss55rKnNBTvrwdmkUpLnDpZoAHvWaiq5+iMmen4AE=
@@ -1213,8 +1213,8 @@ github.com/spf13/afero v1.1.2/go.mod h1:j4pytiNVoe2o6bmDsKpLACNPDBIoEAkihy7loJ1B
12131213
github.com/spf13/afero v1.2.0/go.mod h1:9ZxEEn6pIJ8Rxe320qSDBk6AsU0r9pR7Q4OcevTdifk=
12141214
github.com/spf13/afero v1.2.2/go.mod h1:9ZxEEn6pIJ8Rxe320qSDBk6AsU0r9pR7Q4OcevTdifk=
12151215
github.com/spf13/afero v1.3.2/go.mod h1:5KUK8ByomD5Ti5Artl0RtHeI5pTF7MIDuXL3yY520V4=
1216-
github.com/spf13/afero v1.8.2 h1:xehSyVa0YnHWsJ49JFljMpg1HX19V6NDZ1fkm1Xznbo=
1217-
github.com/spf13/afero v1.8.2/go.mod h1:CtAatgMJh6bJEIs48Ay/FOnkljP3WeGUG0MC1RfAqwo=
1216+
github.com/spf13/afero v1.9.2 h1:j49Hj62F0n+DaZ1dDCvhABaPNSGNkt32oRFxI33IEMw=
1217+
github.com/spf13/afero v1.9.2/go.mod h1:iUV7ddyEEZPO5gA3zD4fJt6iStLlL+Lg4m2cihcDf8Y=
12181218
github.com/spf13/cast v1.2.0/go.mod h1:r2rcYCSwa1IExKTDiTfzaxqT2FNHs8hODu4LnUfgKEg=
12191219
github.com/spf13/cast v1.3.0/go.mod h1:Qx5cxh0v+4UWYiBimWS+eyWzqEqokIECu5etghLkUJE=
12201220
github.com/spf13/cast v1.3.1/go.mod h1:Qx5cxh0v+4UWYiBimWS+eyWzqEqokIECu5etghLkUJE=
@@ -1650,10 +1650,10 @@ golang.org/x/sys v0.0.0-20211025201205-69cdffdb9359/go.mod h1:oPkhp1MJrh7nUepCBc
16501650
golang.org/x/sys v0.0.0-20220520151302-bc2c85ada10a/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
16511651
golang.org/x/sys v0.0.0-20220715151400-c0bba94af5f8/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
16521652
golang.org/x/sys v0.0.0-20220722155257-8c9f86f7a55f/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
1653-
golang.org/x/sys v0.0.0-20220811171246-fbc7d0a398ab/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
16541653
golang.org/x/sys v0.0.0-20220908164124-27713097b956/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
16551654
golang.org/x/sys v0.1.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
16561655
golang.org/x/sys v0.5.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
1656+
golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
16571657
golang.org/x/sys v0.8.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
16581658
golang.org/x/sys v0.17.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
16591659
golang.org/x/sys v0.18.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=

‎flyteadmin/pkg/manager/impl/node_execution_manager.go‎

Lines changed: 69 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -538,26 +538,84 @@ func (m *NodeExecutionManager) GetNodeExecutionData(
538538

539539
var outputs *core.LiteralMap
540540
var outputURLBlob *admin.UrlBlob
541-
var outputVariableMap *core.VariableMap
542541
group.Go(func() error {
543542
var err error
544543
outputs, outputURLBlob, err = util.GetOutputs(groupCtx, m.urlData, m.config.ApplicationConfiguration().GetRemoteDataConfig(),
545544
m.storageClient, nodeExecution.GetClosure())
546545
return err
547546
})
548547

549-
// TODO: Get the output variable map from the node execution model
550-
// group.Go(func() error {
551-
// var err error
548+
// Get the output variable map from workflow model TypedInterface
549+
var inputVariableMap, outputVariableMap *core.VariableMap
550+
group.Go(func() error {
551+
var err error
552+
553+
switch nodeExecution.GetClosure().GetTargetMetadata().(type) {
554+
case *admin.NodeExecutionClosure_WorkflowNodeMetadata:
555+
execID := nodeExecution.GetClosure().GetTargetMetadata().(*admin.NodeExecutionClosure_WorkflowNodeMetadata).WorkflowNodeMetadata.GetExecutionId()
556+
workflowModel, err := m.db.WorkflowRepo().Get(groupCtx, repoInterfaces.Identifier{
557+
Project: execID.GetProject(),
558+
Domain: execID.GetDomain(),
559+
Name: execID.GetName(),
560+
})
561+
562+
if err != nil {
563+
logger.Debugf(groupCtx, "Failed to get workflow model for node execution [%+v] with err %v", request.GetId(), err)
564+
return err
565+
}
566+
workflow, err := transformers.FromWorkflowModel(workflowModel)
567+
if err != nil {
568+
logger.Debugf(groupCtx, "Failed to transform workflow model for node execution [%+v] with err %v", request.GetId(), err)
569+
return err
570+
}
552571

553-
// modelNode, err := m.db.NodeExecutionRepo().Get(groupCtx, repoInterfaces.NodeExecutionResource{
554-
// NodeExecutionIdentifier: request.GetId(),
555-
// })
572+
inputVariableMap = workflow.GetClosure().GetCompiledWorkflow().GetPrimary().GetTemplate().GetInterface().GetInputs()
573+
outputVariableMap = workflow.GetClosure().GetCompiledWorkflow().GetPrimary().GetTemplate().GetInterface().GetOutputs()
556574

557-
// node, err := transformers.FromNodeExecutionModel(modelNode, transformers.DefaultExecutionTransformerOptions)
575+
case *admin.NodeExecutionClosure_TaskNodeMetadata:
576+
execID := nodeExecution.GetId().GetExecutionId()
577+
executionModel, err := m.db.ExecutionRepo().Get(groupCtx, repoInterfaces.Identifier{
578+
Project: execID.GetProject(),
579+
Domain: execID.GetDomain(),
580+
Name: execID.GetName(),
581+
})
582+
583+
if err != nil {
584+
logger.Debugf(groupCtx, "Failed to get execution model for node execution [%+v] with err %v", request.GetId(), err)
585+
return err
586+
}
587+
588+
execution, err := transformers.FromExecutionModel(groupCtx, executionModel, transformers.DefaultExecutionTransformerOptions)
589+
590+
if err != nil {
591+
logger.Debugf(groupCtx, "Failed to transform execution model for node execution [%+v] with err %v", request.GetId(), err)
592+
return err
593+
}
558594

559-
// return err
560-
// })
595+
taskModel, err := m.db.TaskRepo().Get(groupCtx, repoInterfaces.Identifier{
596+
Project: execID.GetProject(),
597+
Domain: execID.GetDomain(),
598+
Name: execution.GetSpec().GetLaunchPlan().GetName(),
599+
Version: execution.GetSpec().GetLaunchPlan().GetVersion(),
600+
})
601+
602+
if err != nil {
603+
logger.Debugf(groupCtx, "Failed to get task model for node execution [%+v] with err %v", request.GetId(), err)
604+
return err
605+
}
606+
607+
task, err := transformers.FromTaskModel(taskModel)
608+
609+
if err != nil {
610+
logger.Debugf(groupCtx, "Failed to transform task model for node execution [%+v] with err %v", request.GetId(), err)
611+
return err
612+
}
613+
614+
inputVariableMap = task.GetClosure().GetCompiledTask().GetTemplate().GetInterface().GetInputs()
615+
outputVariableMap = task.GetClosure().GetCompiledTask().GetTemplate().GetInterface().GetOutputs()
616+
}
617+
return err
618+
})
561619

562620
err = group.Wait()
563621
if err != nil {
@@ -570,6 +628,7 @@ func (m *NodeExecutionManager) GetNodeExecutionData(
570628
FullInputs: inputs,
571629
FullOutputs: outputs,
572630
FlyteUrls: common.FlyteURLsFromNodeExecutionID(request.GetId(), nodeExecution.GetClosure() != nil && nodeExecution.GetClosure().GetDeckUri() != ""),
631+
InputVariableMap: inputVariableMap,
573632
OutputVariableMap: outputVariableMap,
574633
}
575634

‎flyteadmin/pkg/manager/impl/task_execution_manager.go‎

Lines changed: 15 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -328,7 +328,7 @@ func (m *TaskExecutionManager) GetTaskExecutionData(
328328

329329
var outputs *core.LiteralMap
330330
var outputURLBlob *admin.UrlBlob
331-
var outputVariableMap *core.VariableMap
331+
var inputVariableMap, outputVariableMap *core.VariableMap
332332
group.Go(func() error {
333333
var err error
334334
outputs, outputURLBlob, err = util.GetOutputs(groupCtx, m.urlData, m.config.ApplicationConfiguration().GetRemoteDataConfig(),
@@ -345,8 +345,20 @@ func (m *TaskExecutionManager) GetTaskExecutionData(
345345
Version: taskExecution.GetId().GetTaskId().GetVersion(),
346346
})
347347

348+
if err != nil {
349+
logger.Debugf(groupCtx, "Failed to get task [%+v] with err %v", taskExecution.GetId().GetTaskId(), err)
350+
return err
351+
}
352+
348353
task, err := transformers.FromTaskModel(taskModel)
349-
outputVariableMap = task.GetClosure().GetCompiledTask().GetTemplate().GetInterface().Outputs
354+
355+
if err != nil {
356+
logger.Debugf(groupCtx, "Failed to transform task model [%+v] with err %v", taskModel, err)
357+
return err
358+
}
359+
360+
inputVariableMap = task.GetClosure().GetCompiledTask().GetTemplate().GetInterface().GetInputs()
361+
outputVariableMap = task.GetClosure().GetCompiledTask().GetTemplate().GetInterface().GetOutputs()
350362

351363
return err
352364
})
@@ -363,6 +375,7 @@ func (m *TaskExecutionManager) GetTaskExecutionData(
363375
FullOutputs: outputs,
364376
FlyteUrls: common.FlyteURLsFromTaskExecutionID(request.GetId(), false),
365377
OutputVariableMap: outputVariableMap,
378+
InputVariableMap: inputVariableMap,
366379
}
367380

368381
m.metrics.TaskExecutionInputBytes.Observe(float64(response.GetInputs().GetBytes()))

‎flyteidl/clients/go/admin/mocks/isGetDataResponse_VariableMap.go‎

Lines changed: 15 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎flyteidl/clients/go/assets/admin.swagger.json‎

Lines changed: 6 additions & 3 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎flyteidl/gen/pb-es/flyteidl/admin/node_execution_pb.ts‎

Lines changed: 8 additions & 2 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎flyteidl/gen/pb-es/flyteidl/admin/task_execution_pb.ts‎

Lines changed: 8 additions & 2 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

0 commit comments

Comments
 (0)